Selaa lähdekoodia

fix: align tool mode execution safety

Problem: Direct provider calls could reuse IDs across active history, execute batches sequentially, and dual text events could still run after the shared event budget was exhausted.

Risk: Late duplicate detection can repeat side effects or create invalid transcripts, while naive concurrency can reorder replies and asymmetric budget handling can execute extra handlers. Validate before transcript construction, use ordered asyncio.gather results, and apply budget acceptance to both modes.
zhenyu.hu 2 viikkoa sitten
vanhempi
commit
721ab4cb5e
2 muutettua tiedostoa jossa 291 lisäystä ja 15 poistoa
  1. 41 15
      src/agent_lab/application/runtime.py
  2. 250 0
      tests/test_debug_runtime.py

+ 41 - 15
src/agent_lab/application/runtime.py

@@ -293,6 +293,9 @@ class DebugRuntime:
                 raw_chunks=raw_chunks,
             )
 
+            if events and request.tool_invocation_mode == "chat_agent_tools":
+                self._validate_provider_tool_call_ids(messages, events)
+
             if assistant_content or events:
                 assistant_message = ChatMessage(
                     role="assistant",
@@ -618,6 +621,9 @@ class DebugRuntime:
                 raw_chunks=raw_chunks,
             )
 
+            if events and request.tool_invocation_mode == "chat_agent_tools":
+                self._validate_provider_tool_call_ids(messages, events)
+
             if assistant_content or events:
                 assistant_message = ChatMessage(
                     role="assistant",
@@ -868,10 +874,11 @@ class DebugRuntime:
         *,
         budget_available: bool,
     ) -> ToolCallEvent | None:
+        if not budget_available:
+            return None
         if request.tool_invocation_mode == "chat_agent_tools":
             if (
-                budget_available
-                and item.kind == "provider_tool_call"
+                item.kind == "provider_tool_call"
                 and item.event is not None
             ):
                 return item.event
@@ -880,27 +887,46 @@ class DebugRuntime:
             return self._event_name_only(item.event)
         return None
 
+    def _validate_provider_tool_call_ids(
+        self,
+        messages: list[ChatMessage],
+        events: list[ToolCallEvent],
+    ) -> None:
+        declared_ids = {
+            tool_call.id
+            for message in messages
+            for tool_call in message.tool_calls
+        }
+        batch_ids: set[str] = set()
+        for event in events:
+            if event.id in declared_ids or event.id in batch_ids:
+                raise ValueError(f"duplicate provider tool-call ID: {event.id}")
+            batch_ids.add(event.id)
+
     async def _execute_provider_tools(
         self,
         events: list[ToolCallEvent],
         *,
         enabled_names: list[str],
     ) -> list[ChatMessage]:
-        replies: list[ChatMessage] = []
-        for event in events:
-            payload = await self.registry.execute_async(
-                event,
-                enabled_names=enabled_names,
-            )
-            replies.append(
-                ChatMessage(
-                    role="tool",
-                    content=json.dumps(payload, ensure_ascii=False),
-                    name=event.name,
-                    tool_call_id=event.id,
+        payloads = await asyncio.gather(
+            *[
+                self.registry.execute_async(
+                    event,
+                    enabled_names=enabled_names,
                 )
+                for event in events
+            ]
+        )
+        return [
+            ChatMessage(
+                role="tool",
+                content=json.dumps(payload, ensure_ascii=False),
+                name=event.name,
+                tool_call_id=event.id,
             )
-        return replies
+            for event, payload in zip(events, payloads, strict=True)
+        ]
 
     def _chat_messages_for_round(
         self,

+ 250 - 0
tests/test_debug_runtime.py

@@ -1972,3 +1972,253 @@ async def test_direct_mode_session_turns_reset_budget_and_snapshot_mode(tmp_path
         audit["details"]["tool_invocation_mode"] == "chat_agent_tools"
         for audit in request_audits
     )
+
+
+def _single_tool_registry(
+    handler: Any,
+    *,
+    name: str = "safe_tool",
+) -> ToolRegistry:
+    return ToolRegistry(
+        [
+            ToolDefinition(
+                name=name,
+                description="Tool used by execution-safety tests.",
+                parameters={"type": "object"},
+                handler=handler,
+            )
+        ]
+    )
+
+
+def _tool_call(call_id: str, *, name: str = "safe_tool") -> ToolCallEvent:
+    return ToolCallEvent(
+        id=call_id,
+        name=name,
+        arguments={},
+        raw_arguments="{}",
+    )
+
+
+@pytest.mark.asyncio
+async def test_direct_mode_rejects_provider_call_id_reused_across_one_shot_rounds():
+    executed: list[str] = []
+    registry = _single_tool_registry(
+        lambda event: executed.append(event.id) or {"tool": event.name}
+    )
+    client = ScriptedChatClient(
+        [
+            [StreamItem.provider_tool_call(_tool_call("duplicate-call"))],
+            [StreamItem.provider_tool_call(_tool_call("duplicate-call"))],
+        ]
+    )
+
+    outputs = await _collect_outputs(
+        DebugRuntime(client, registry=registry).run(
+            _direct_request(enabled_tools=["safe_tool"], max_event_loops=2)
+        )
+    )
+
+    assert executed == ["duplicate-call"]
+    assert [
+        message["message"]["tool_call_id"]
+        for message in outputs
+        if message["type"] == "tool_result"
+    ] == ["duplicate-call"]
+    assert outputs[-1]["type"] == "error"
+    assert "duplicate provider tool-call ID: duplicate-call" in outputs[-1]["message"]
+
+
+@pytest.mark.asyncio
+async def test_direct_mode_rejects_provider_call_id_reused_in_session_rounds():
+    executed: list[str] = []
+    registry = _single_tool_registry(
+        lambda event: executed.append(event.id) or {"tool": event.name}
+    )
+    client = ScriptedChatClient(
+        [
+            [StreamItem.provider_tool_call(_tool_call("session-duplicate"))],
+            [StreamItem.provider_tool_call(_tool_call("session-duplicate"))],
+        ]
+    )
+    runtime = DebugRuntime(client, registry=registry)
+
+    queues = runtime.start_session(
+        _direct_request(enabled_tools=["safe_tool"], max_event_loops=2)
+    )
+    outputs: list[dict[str, Any]] = []
+    while not any(message["type"] == "error" for message in outputs):
+        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
+    await runtime.aclose()
+
+    assert executed == ["session-duplicate"]
+    assert [
+        message["message"]["tool_call_id"]
+        for message in outputs
+        if message["type"] == "tool_result"
+    ] == ["session-duplicate"]
+    assert "duplicate provider tool-call ID: session-duplicate" in outputs[-1][
+        "message"
+    ]
+
+
+@pytest.mark.asyncio
+async def test_direct_mode_rejects_duplicate_provider_call_ids_within_batch():
+    executed: list[str] = []
+    registry = _single_tool_registry(
+        lambda event: executed.append(event.id) or {"tool": event.name}
+    )
+    client = ScriptedChatClient(
+        [
+            [
+                StreamItem.provider_tool_call(_tool_call("same-batch")),
+                StreamItem.provider_tool_call(_tool_call("same-batch")),
+            ]
+        ]
+    )
+
+    outputs = await _collect_outputs(
+        DebugRuntime(client, registry=registry).run(
+            _direct_request(enabled_tools=["safe_tool"])
+        )
+    )
+
+    assert executed == []
+    assert outputs[-1]["type"] == "error"
+    assert "duplicate provider tool-call ID: same-batch" in outputs[-1]["message"]
+
+
+@pytest.mark.asyncio
+async def test_direct_mode_executes_batch_concurrently_but_replies_in_provider_order():
+    both_started = asyncio.Event()
+    first_release = asyncio.Event()
+    second_release = asyncio.Event()
+    second_finished = asyncio.Event()
+    started: list[str] = []
+
+    async def handler(event: ToolCallEvent) -> dict[str, Any]:
+        started.append(event.id)
+        if len(started) == 2:
+            both_started.set()
+        if event.id == "first-call":
+            await first_release.wait()
+        else:
+            await second_release.wait()
+            second_finished.set()
+        return {"tool": event.name, "call_id": event.id}
+
+    registry = _single_tool_registry(handler)
+    client = ScriptedChatClient(
+        [
+            [
+                StreamItem.provider_tool_call(_tool_call("first-call")),
+                StreamItem.provider_tool_call(_tool_call("second-call")),
+            ],
+            [StreamItem.message_delta("finished")],
+        ]
+    )
+    runtime = DebugRuntime(client, registry=registry)
+    run_task = asyncio.create_task(
+        _collect_outputs(
+            runtime.run(_direct_request(enabled_tools=["safe_tool"]))
+        )
+    )
+    overlapped = False
+    try:
+        try:
+            await asyncio.wait_for(both_started.wait(), timeout=1)
+        except TimeoutError:
+            pass
+        else:
+            overlapped = True
+            second_release.set()
+            await asyncio.wait_for(second_finished.wait(), timeout=1)
+            assert not run_task.done()
+    finally:
+        first_release.set()
+        second_release.set()
+    outputs = await asyncio.wait_for(run_task, timeout=1)
+
+    assert overlapped is True
+    assert started == ["first-call", "second-call"]
+    assert [
+        message["message"]["tool_call_id"]
+        for message in outputs
+        if message["type"] == "tool_result"
+    ] == ["first-call", "second-call"]
+
+
+def _dual_budget_request() -> DebugRunRequest:
+    return DebugRunRequest(
+        user_message="dual budget",
+        system_prompts=[],
+        pre_messages=[],
+        chat_agent=AgentParams(model="chat-model"),
+        event_agent=EventAgentParams(
+            enabled_tools=["safe_tool"],
+            max_event_loops=1,
+        ),
+    )
+
+
+@pytest.mark.asyncio
+async def test_dual_mode_ignores_text_event_after_one_shot_budget_exhaustion():
+    executed: list[str] = []
+    registry = _single_tool_registry(
+        lambda event: executed.append(event.id) or {"tool": event.name}
+    )
+    client = ScriptedChatClient(
+        [
+            [StreamItem.text_event(_tool_call("accepted-text-event"))],
+            [
+                StreamItem.text_event(_tool_call("ignored-text-event")),
+                StreamItem.message_delta("final dual answer"),
+            ],
+        ]
+    )
+
+    outputs = await _collect_outputs(
+        DebugRuntime(client, registry=registry).run(_dual_budget_request())
+    )
+
+    assert executed == ["accepted-text-event"]
+    assert [
+        message["event"]["id"]
+        for message in outputs
+        if message["type"] == "event"
+    ] == ["accepted-text-event"]
+    assert outputs[-1] == {"type": "done"}
+
+
+@pytest.mark.asyncio
+async def test_dual_mode_session_main_path_ignores_text_event_after_budget_exhaustion():
+    executed: list[str] = []
+    registry = _single_tool_registry(
+        lambda event: executed.append(event.id) or {"tool": event.name}
+    )
+    client = ScriptedChatClient(
+        [
+            [StreamItem.text_event(_tool_call("session-accepted"))],
+            [
+                StreamItem.text_event(_tool_call("session-ignored")),
+                StreamItem.message_delta("session final"),
+            ],
+        ]
+    )
+    runtime = DebugRuntime(client, registry=registry)
+
+    queues = runtime.start_session(_dual_budget_request())
+    outputs: list[dict[str, Any]] = []
+    while not any(
+        message["type"] in {"turn_completed", "error"} for message in outputs
+    ):
+        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
+    await runtime.aclose()
+
+    assert executed == ["session-accepted"]
+    assert [
+        message["event"]["id"]
+        for message in outputs
+        if message["type"] == "event"
+    ] == ["session-accepted"]
+    assert outputs[-1]["type"] == "turn_completed"