709 lines
24 KiB
Python
709 lines
24 KiB
Python
"""Concurrent tool dispatch tests.
|
|
|
|
Verifies that same-turn safe tools run concurrently via asyncio.gather
|
|
and mutex tools remain serial. All tests use mock LLM and mock tool
|
|
handlers; no real network or provider calls are made.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine import Agent, AgentConfig, ToolResult
|
|
from opensquilla.engine.runtime import _SAFE_TOOL_NAMES, _get_tool_concurrency_policy
|
|
from opensquilla.engine.types import ToolCall
|
|
from opensquilla.provider import (
|
|
ChatConfig,
|
|
Message,
|
|
ToolDefinition,
|
|
ToolInputSchema,
|
|
)
|
|
from opensquilla.provider import DoneEvent as ProviderDone
|
|
from opensquilla.provider import TextDeltaEvent as ProviderTextDelta
|
|
from opensquilla.provider import ToolUseEndEvent as ProviderToolUseEnd
|
|
from opensquilla.provider import ToolUseStartEvent as ProviderToolUseStart
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_TOOL_SLEEP_S = 0.2
|
|
_SCHEDULER_TOLERANCE_S = 0.05
|
|
|
|
|
|
def _tool_def(name: str) -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name=name,
|
|
description=f"Mock tool {name}",
|
|
input_schema=ToolInputSchema(properties={}, required=[]),
|
|
)
|
|
|
|
|
|
class _FixedToolCallProvider:
|
|
"""Fake LLM that returns a fixed list of tool calls on the first turn,
|
|
then a plain text done on the second turn."""
|
|
|
|
provider_name = "fake"
|
|
|
|
def __init__(self, tool_names: list[str]) -> None:
|
|
self._tool_names = tool_names
|
|
self.calls: list[list[Message]] = []
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls.append(messages)
|
|
call_number = len(self.calls)
|
|
return self._stream(call_number)
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number > 1:
|
|
# Second call: return a non-empty text response so the agent
|
|
# does not trigger the empty-response retry loop.
|
|
yield ProviderTextDelta(text="done")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
return
|
|
# First call: emit all tool uses then done
|
|
for i, name in enumerate(self._tool_names):
|
|
tid = f"tool-{i}"
|
|
yield ProviderToolUseStart(tool_use_id=tid, tool_name=name)
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id=tid,
|
|
tool_name=name,
|
|
arguments={},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
class _FixedToolCallArgsProvider:
|
|
"""Fake LLM that returns fixed tool calls with explicit arguments."""
|
|
|
|
provider_name = "fake"
|
|
|
|
def __init__(self, tool_calls: list[tuple[str, dict[str, Any]]]) -> None:
|
|
self._tool_calls = tool_calls
|
|
self.calls: list[list[Message]] = []
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls.append(messages)
|
|
call_number = len(self.calls)
|
|
return self._stream(call_number)
|
|
|
|
async def _stream(self, call_number: int) -> AsyncIterator[Any]:
|
|
if call_number > 1:
|
|
yield ProviderTextDelta(text="done")
|
|
yield ProviderDone(stop_reason="stop", input_tokens=1, output_tokens=1)
|
|
return
|
|
for i, (name, arguments) in enumerate(self._tool_calls):
|
|
tid = f"tool-{i}"
|
|
yield ProviderToolUseStart(tool_use_id=tid, tool_name=name)
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id=tid,
|
|
tool_name=name,
|
|
arguments=arguments,
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
async def _collect(agent: Agent, message: str = "go") -> list[Any]:
|
|
return [e async for e in agent.run_turn(message)]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pick two safe tool names that are guaranteed to be in _SAFE_TOOL_NAMES
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_SAFE_SAMPLE = sorted(_SAFE_TOOL_NAMES)[:6] # pick 6 safe names for the concurrent test
|
|
|
|
|
|
def test_web_search_is_safe_for_same_turn_concurrency() -> None:
|
|
"""web_search is read-only network I/O and should batch with safe tools."""
|
|
assert "web_search" in _SAFE_TOOL_NAMES
|
|
|
|
|
|
def test_sessions_spawn_policy_keys_by_parent_session() -> None:
|
|
"""Spawn policy is scoped to the parent session rather than global state."""
|
|
policy = _get_tool_concurrency_policy(
|
|
"sessions_spawn",
|
|
{},
|
|
parent_session_key="agent:main:parent-a",
|
|
)
|
|
|
|
assert policy.mode == "keyed"
|
|
assert policy.key == ("sessions_spawn", "agent:main:parent-a")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_feishu_read_only_tools_run_concurrent() -> None:
|
|
"""Read-only Feishu tools should batch without enabling Feishu mutators."""
|
|
tool_calls = [
|
|
("feishu_doc_read_raw", {"document_id": "doc-a"}),
|
|
("feishu_drive_meta", {"token": "doc-b"}),
|
|
("feishu_wiki_get_node", {"token": "node-c"}),
|
|
]
|
|
intervals: list[tuple[str, float, float]] = []
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
start = time.monotonic()
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
intervals.append((tc.tool_name, start, time.monotonic()))
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallArgsProvider(tool_calls)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n, _ in tool_calls],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
t0 = time.monotonic()
|
|
await _collect(agent)
|
|
elapsed = time.monotonic() - t0
|
|
|
|
assert elapsed < 0.45, (
|
|
f"Expected Feishu read-only tools to overlap, got {elapsed:.3f} s."
|
|
)
|
|
assert len(intervals) == len(tool_calls)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_feishu_read_only_tools_have_independent_inflight_cap() -> None:
|
|
"""Feishu reads should batch, but keep their own platform-facing cap."""
|
|
tool_calls = [
|
|
("feishu_doc_read_raw", {"document_id": f"doc-{i}"})
|
|
for i in range(6)
|
|
]
|
|
in_flight = 0
|
|
max_in_flight = 0
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
nonlocal in_flight, max_in_flight
|
|
in_flight += 1
|
|
max_in_flight = max(max_in_flight, in_flight)
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
in_flight -= 1
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallArgsProvider(tool_calls)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2, max_safe_tool_concurrency=6),
|
|
tool_definitions=[_tool_def(n) for n, _ in tool_calls],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
t0 = time.monotonic()
|
|
await _collect(agent)
|
|
elapsed = time.monotonic() - t0
|
|
|
|
assert max_in_flight == 4
|
|
assert 2 * _TOOL_SLEEP_S - _SCHEDULER_TOLERANCE_S <= elapsed < 3 * _TOOL_SLEEP_S
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_feishu_mutating_tools_remain_serial() -> None:
|
|
"""Feishu writes must remain mutex even when read-only Feishu tools batch."""
|
|
tool_calls = [
|
|
("feishu_chat_send", {"chat_id": "chat-a", "text": "one"}),
|
|
("feishu_chat_edit", {"message_id": "msg-a", "text": "two"}),
|
|
]
|
|
intervals: list[tuple[str, float, float]] = []
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
start = time.monotonic()
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
intervals.append((tc.tool_name, start, time.monotonic()))
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallArgsProvider(tool_calls)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n, _ in tool_calls],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
await _collect(agent)
|
|
|
|
assert len(intervals) == len(tool_calls)
|
|
(_, _, first_end), (_, second_start, _) = intervals
|
|
assert second_start >= first_end - 0.01
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sessions_send_different_targets_run_concurrent() -> None:
|
|
"""session sends to different target sessions should not globally serialize."""
|
|
tool_calls = [
|
|
("sessions_send", {"session_key": "agent:main:child-a", "message": "one"}),
|
|
("sessions_send", {"session_key": "agent:main:child-b", "message": "two"}),
|
|
]
|
|
intervals: list[tuple[str, str, float, float]] = []
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
target = str(tc.arguments["session_key"])
|
|
start = time.monotonic()
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
intervals.append((tc.tool_name, target, start, time.monotonic()))
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallArgsProvider(tool_calls)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n, _ in tool_calls],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
t0 = time.monotonic()
|
|
await _collect(agent)
|
|
elapsed = time.monotonic() - t0
|
|
|
|
assert elapsed < 0.35, (
|
|
f"Expected sends to different sessions to overlap, got {elapsed:.3f} s."
|
|
)
|
|
assert len(intervals) == len(tool_calls)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sessions_send_same_target_serializes() -> None:
|
|
"""session sends to one target session must preserve per-target order."""
|
|
tool_calls = [
|
|
("sessions_send", {"session_key": "agent:main:child-a", "message": "one"}),
|
|
("sessions_send", {"session_key": "agent:main:child-a", "message": "two"}),
|
|
]
|
|
intervals: list[tuple[str, float, float]] = []
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
start = time.monotonic()
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
intervals.append((str(tc.arguments["message"]), start, time.monotonic()))
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallArgsProvider(tool_calls)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n, _ in tool_calls],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
await _collect(agent)
|
|
|
|
assert len(intervals) == len(tool_calls)
|
|
(_, _, first_end), (_, second_start, _) = intervals
|
|
assert second_start >= first_end - 0.01
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 1: 6 safe tools run concurrently
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_six_safe_tools_run_concurrent() -> None:
|
|
"""Six safe tools should complete in roughly one sleep period, not six."""
|
|
assert len(_SAFE_SAMPLE) == 6, "Need at least 6 safe tools in _SAFE_TOOL_NAMES"
|
|
|
|
call_order: list[str] = []
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
call_order.append(tc.tool_name)
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallProvider(_SAFE_SAMPLE)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n in _SAFE_SAMPLE],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
t0 = time.monotonic()
|
|
await _collect(agent)
|
|
elapsed = time.monotonic() - t0
|
|
|
|
# Concurrent: should be ~0.2 s; serial would be ~1.2 s
|
|
assert elapsed < 0.60, (
|
|
f"Expected concurrent execution (<0.60 s), got {elapsed:.3f} s. "
|
|
"Safe tools may still be running serially."
|
|
)
|
|
# Speed-up vs serial lower bound
|
|
serial_estimate = len(_SAFE_SAMPLE) * _TOOL_SLEEP_S
|
|
assert elapsed * 3 < serial_estimate, (
|
|
f"Expected at least 3x speedup over serial ({serial_estimate:.1f} s), "
|
|
f"got {elapsed:.3f} s"
|
|
)
|
|
# All 6 tools were called
|
|
assert sorted(call_order) == sorted(_SAFE_SAMPLE)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_safe_tool_concurrency_limit_caps_in_flight_tasks() -> None:
|
|
"""Safe tools should run concurrently, but not without an upper bound."""
|
|
tool_names = _SAFE_SAMPLE
|
|
max_in_flight = 0
|
|
in_flight = 0
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
nonlocal in_flight, max_in_flight
|
|
in_flight += 1
|
|
max_in_flight = max(max_in_flight, in_flight)
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
in_flight -= 1
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallProvider(tool_names)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2, max_safe_tool_concurrency=2),
|
|
tool_definitions=[_tool_def(n) for n in tool_names],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
t0 = time.monotonic()
|
|
await _collect(agent)
|
|
elapsed = time.monotonic() - t0
|
|
|
|
assert max_in_flight == 2
|
|
assert 2 * _TOOL_SLEEP_S <= elapsed < 4 * _TOOL_SLEEP_S
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 2: 2 mutex tools stay serial
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Pick two tool names that are NOT in _SAFE_TOOL_NAMES
|
|
_MUTEX_NAMES = ["write_file", "exec_command"]
|
|
assert all(n not in _SAFE_TOOL_NAMES for n in _MUTEX_NAMES), (
|
|
"Test assumes write_file and exec_command are mutex (not in _SAFE_TOOL_NAMES)"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_two_mutex_tools_serialize() -> None:
|
|
"""Two mutex tools must not run concurrently."""
|
|
execution_intervals: list[tuple[float, float]] = []
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
start = time.monotonic()
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
execution_intervals.append((start, time.monotonic()))
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallProvider(_MUTEX_NAMES)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n in _MUTEX_NAMES],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
t0 = time.monotonic()
|
|
await _collect(agent)
|
|
elapsed = time.monotonic() - t0
|
|
|
|
# Serial: must take at least 2 * sleep
|
|
assert elapsed >= 2 * _TOOL_SLEEP_S * 0.9, (
|
|
f"Expected serial execution (>= {2 * _TOOL_SLEEP_S * 0.9:.2f} s), "
|
|
f"got {elapsed:.3f} s. Mutex tools may be running concurrently."
|
|
)
|
|
|
|
# Verify non-overlapping: second start >= first end
|
|
assert len(execution_intervals) == 2
|
|
(s1, e1), (s2, e2) = execution_intervals
|
|
assert s2 >= e1 - 0.01, (
|
|
f"Execution intervals overlap: [{s1:.3f}, {e1:.3f}] and [{s2:.3f}, {e2:.3f}]. "
|
|
"Mutex tools ran concurrently."
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 3: one failing safe tool does not block the others
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_SAFE_TRIO = sorted(_SAFE_TOOL_NAMES)[:3] # three safe tools (also reused for cancellation test)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_failing_tool_does_not_block_others() -> None:
|
|
"""Tool A raises; tools B and C must still complete and results stay ordered."""
|
|
assert len(_SAFE_TRIO) == 3, "Need at least 3 safe tools in _SAFE_TOOL_NAMES"
|
|
|
|
tool_a, tool_b, tool_c = _SAFE_TRIO
|
|
completed: list[str] = []
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
if tc.tool_name == tool_a:
|
|
raise RuntimeError("simulated failure")
|
|
await asyncio.sleep(0.05)
|
|
completed.append(tc.tool_name)
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content=f"{tc.tool_name} ok",
|
|
)
|
|
|
|
provider = _FixedToolCallProvider([tool_a, tool_b, tool_c])
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n in _SAFE_TRIO],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
events = await _collect(agent)
|
|
|
|
# B and C completed despite A failing
|
|
assert tool_b in completed, "tool_b did not complete"
|
|
assert tool_c in completed, "tool_c did not complete"
|
|
|
|
# All three tool result events were yielded (A as error, B+C as success)
|
|
from opensquilla.engine.types import ToolResultEvent
|
|
|
|
result_events = [e for e in events if isinstance(e, ToolResultEvent)]
|
|
assert len(result_events) == 3, (
|
|
f"Expected 3 ToolResultEvents, got {len(result_events)}. "
|
|
"A failing tool may have blocked sibling results."
|
|
)
|
|
|
|
# Ordering: results appear in original tool_calls order (A, B, C)
|
|
names_in_order = [e.tool_name for e in result_events]
|
|
assert names_in_order == [tool_a, tool_b, tool_c], (
|
|
f"Results out of original order: {names_in_order}"
|
|
)
|
|
|
|
# A's result is an error
|
|
assert result_events[0].is_error, "Expected tool_a result to be an error"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 4: cancellation propagates to in-flight safe tools
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancellation_propagates() -> None:
|
|
"""Cancelling the outer Task must reach all in-flight safe tool coroutines."""
|
|
cancel_received: list[str] = []
|
|
started: list[str] = []
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
started.append(tc.tool_name)
|
|
try:
|
|
await asyncio.sleep(2.0) # long sleep — will be cancelled
|
|
except asyncio.CancelledError:
|
|
cancel_received.append(tc.tool_name)
|
|
raise
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
provider = _FixedToolCallProvider(_SAFE_TRIO)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n in _SAFE_TRIO],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
async def _run() -> None:
|
|
await _collect(agent)
|
|
|
|
task = asyncio.create_task(_run())
|
|
|
|
# Let tools start
|
|
await asyncio.sleep(0.05)
|
|
task.cancel()
|
|
|
|
t0 = time.monotonic()
|
|
with pytest.raises((asyncio.CancelledError, Exception)):
|
|
await task
|
|
elapsed = time.monotonic() - t0
|
|
|
|
# All in-flight tools should have received CancelledError
|
|
assert len(cancel_received) > 0, (
|
|
"No tool received CancelledError — cancellation did not propagate."
|
|
)
|
|
# Should exit well under 2 s (the tool sleep duration)
|
|
assert elapsed < 1.0, (
|
|
f"Cancellation took {elapsed:.3f} s — tools may not be cancelling promptly."
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 5: http_request is not in the safe set
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_http_request_not_in_safe_set() -> None:
|
|
"""http_request writes to .fetch/ directory and must not be dispatched concurrently."""
|
|
assert "http_request" not in _SAFE_TOOL_NAMES, (
|
|
"http_request has write side-effects (_save_http_response_body) and must not "
|
|
"be in _SAFE_TOOL_NAMES."
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 6: safe tool after mutex must wait for mutex to finish
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_safe_after_mutex_waits_for_mutex() -> None:
|
|
"""A safe tool that comes AFTER a mutex must not start until the mutex finishes."""
|
|
# call sequence: [write_file (mutex), read_file (safe)]
|
|
mutex_first = "write_file"
|
|
safe_second = "read_file"
|
|
assert mutex_first not in _SAFE_TOOL_NAMES
|
|
assert safe_second in _SAFE_TOOL_NAMES
|
|
|
|
intervals: dict[str, tuple[float, float]] = {}
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
start = time.monotonic()
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
end = time.monotonic()
|
|
intervals[tc.tool_name] = (start, end)
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
tool_names = [mutex_first, safe_second]
|
|
provider = _FixedToolCallProvider(tool_names)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n in tool_names],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
await _collect(agent)
|
|
|
|
assert mutex_first in intervals, "mutex tool did not run"
|
|
assert safe_second in intervals, "safe tool did not run"
|
|
|
|
mutex_end = intervals[mutex_first][1]
|
|
safe_start = intervals[safe_second][0]
|
|
|
|
assert safe_start >= mutex_end - 0.01, (
|
|
f"Safe tool started at {safe_start:.3f} before mutex finished at {mutex_end:.3f}. "
|
|
"Ordering barrier was not respected."
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Test 7: mutex after safe batch waits for all safe tools, safe tools overlap
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mutex_after_safe_waits_for_safe_batch() -> None:
|
|
"""A mutex tool after safe tools must not start until all safe tools finish,
|
|
and the safe tools must run concurrently (overlapping intervals)."""
|
|
safe_a = "read_file"
|
|
safe_b = "glob_search"
|
|
mutex_c = "write_file"
|
|
assert safe_a in _SAFE_TOOL_NAMES
|
|
assert safe_b in _SAFE_TOOL_NAMES
|
|
assert mutex_c not in _SAFE_TOOL_NAMES
|
|
|
|
intervals: dict[str, tuple[float, float]] = {}
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
start = time.monotonic()
|
|
await asyncio.sleep(_TOOL_SLEEP_S)
|
|
end = time.monotonic()
|
|
intervals[tc.tool_name] = (start, end)
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
tool_names = [safe_a, safe_b, mutex_c]
|
|
provider = _FixedToolCallProvider(tool_names)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(max_iterations=2),
|
|
tool_definitions=[_tool_def(n) for n in tool_names],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
await _collect(agent)
|
|
|
|
assert safe_a in intervals
|
|
assert safe_b in intervals
|
|
assert mutex_c in intervals
|
|
|
|
# safe tools must have overlapping execution intervals
|
|
(sa_s, sa_e) = intervals[safe_a]
|
|
(sb_s, sb_e) = intervals[safe_b]
|
|
assert sa_s < sb_e and sb_s < sa_e, (
|
|
f"safe_a [{sa_s:.3f}, {sa_e:.3f}] and safe_b [{sb_s:.3f}, {sb_e:.3f}] "
|
|
"did not overlap — they may have run serially."
|
|
)
|
|
|
|
# mutex must start after both safe tools finish
|
|
mutex_start = intervals[mutex_c][0]
|
|
safe_batch_end = max(sa_e, sb_e)
|
|
assert mutex_start >= safe_batch_end - 0.01, (
|
|
f"Mutex started at {mutex_start:.3f} before safe batch ended at {safe_batch_end:.3f}."
|
|
)
|