From 2ce704550b5f70d07335bc6f7bc564a9dbce856d Mon Sep 17 00:00:00 2001 From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:31:52 -0700 Subject: [PATCH] fix(agent): execute local tool batches through per-tool pipeline --- agent/tool_call_batches.py | 80 +++++++++++++ agent/turn_tool_round.py | 6 + tests/agent/test_local_tool_batches.py | 158 +++++++++++++++++++++++++ tools/tool_search.py | 8 +- 4 files changed, 248 insertions(+), 4 deletions(-) create mode 100644 agent/tool_call_batches.py create mode 100644 tests/agent/test_local_tool_batches.py diff --git a/agent/tool_call_batches.py b/agent/tool_call_batches.py new file mode 100644 index 0000000000..b071c356a4 --- /dev/null +++ b/agent/tool_call_batches.py @@ -0,0 +1,80 @@ +"""Normalize local bridge batches before persistence and per-tool execution.""" + +from copy import deepcopy +import json + +from agent.message_sanitization import coalesce_tool_call_id, deterministic_call_id, uniquify_tool_call_ids +from agent.transports.types import ToolCall +from tools.connectors import is_connector_name +from tools.connectors.gateway.config import MAX_CALLS_PER_DISPATCH +from tools.tool_search_catalog import TOOL_CALL_NAME +from tools.tool_search_validation import normalize_tool_call_entries + + +def expand_local_tool_batches(tool_calls: list, *, provider_data: dict | None = None) -> list: + """Give every local entry the ordinary agent execution and persistence path. + + Registry-only batch dispatch loses the live agent's inline tools, guardrails, + and per-server concurrency policy. Split the *new* assistant turn instead, + before it is persisted; never rewrite an earlier conversation message. + Connector-only batches keep their existing dispatcher and wire shape. + """ + expanded, replacements = [], {} + for call in tool_calls: + if call.function.name != TOOL_CALL_NAME: + expanded.append(call) + continue + try: + args = call.function.arguments + entries, error = normalize_tool_call_entries(json.loads(args) if isinstance(args, str) else args) + except (TypeError, ValueError, AttributeError): + entries, error = [], True + if (error or not 1 < len(entries) <= MAX_CALLS_PER_DISPATCH + or all(is_connector_name(entry["name"]) for entry in entries)): + # Leave malformed/oversized requests intact for the dispatcher's error. + expanded.append(call) + continue + parent_id = coalesce_tool_call_id(call) + children = [] + for index, entry in enumerate(entries): + arguments = json.dumps({"calls": [entry]}, ensure_ascii=False) + if index == 0: + child = deepcopy(call) + child.function.arguments = arguments + else: + # A new entry must not inherit the parent's Responses item id or + # provider signature. IDs are deterministic for stable replay. + child = ToolCall( + id=deterministic_call_id(TOOL_CALL_NAME, parent_id + arguments, index), + name=TOOL_CALL_NAME, arguments=arguments, + ) + expanded.append(child) + children.append(child) + replacements[parent_id] = children + uniquify_tool_call_ids(expanded) + # Native providers replay ordered blocks instead of rebuilding them from + # tool_calls. Replace the matching tool block in place, preserving signed + # thinking/text blocks and their order exactly. + if provider_data and replacements: + for key, replace in (("anthropic_content_blocks", _anthropic_children), + ("bedrock_content_blocks", _bedrock_children)): + if blocks := provider_data.get(key): + provider_data[key] = [child for block in blocks for child in replace(block, replacements)] + return expanded + + +def _anthropic_children(block, replacements): + children = replacements.get(block.get("id")) if block.get("type") == "tool_use" else None + if children is None: + return [block] + return [{**block, "id": coalesce_tool_call_id(child), + "input": json.loads(child.function.arguments)} for child in children] + + +def _bedrock_children(block, replacements): + tool = block.get("toolUse") + children = replacements.get(tool.get("toolUseId")) if isinstance(tool, dict) else None + if children is None: + return [block] + return [{**block, "toolUse": {**tool, "toolUseId": coalesce_tool_call_id(child), + "input": json.loads(child.function.arguments)}} for child in children] diff --git a/agent/turn_tool_round.py b/agent/turn_tool_round.py index 0a98b93a89..e089d47a76 100644 --- a/agent/turn_tool_round.py +++ b/agent/turn_tool_round.py @@ -84,6 +84,12 @@ def run_tool_round( if _tvv.action == "continue": return _verdict("continue") + # Normalize only this unpersisted turn. Each local bridge entry must pass + # through the same scope, approvals, state and scheduling as a singleton. + from agent.tool_call_batches import expand_local_tool_batches + assistant_message.tool_calls = expand_local_tool_batches( + assistant_message.tool_calls, provider_data=getattr(assistant_message, "provider_data", None)) + # Post-call guardrails. assistant_message.tool_calls = agent._deduplicate_tool_calls( agent._cap_delegate_task_calls(assistant_message.tool_calls) diff --git a/tests/agent/test_local_tool_batches.py b/tests/agent/test_local_tool_batches.py new file mode 100644 index 0000000000..ce9d8d40a0 --- /dev/null +++ b/tests/agent/test_local_tool_batches.py @@ -0,0 +1,158 @@ +"""Local batches use the live agent pipeline, not registry fan-out.""" + +import copy +import json +import sys +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + + +def _call(entries, call_id="batch"): + return SimpleNamespace(id=call_id, type="function", function=SimpleNamespace( + name="tool_call", arguments=json.dumps({"calls": entries}))) + + +def _response(calls=None): + return SimpleNamespace(choices=[SimpleNamespace( + message=SimpleNamespace(content="" if calls else "done", tool_calls=calls), + finish_reason="tool_calls" if calls else "stop")], model="test/model", usage=None) + + +@pytest.fixture +def mcp_server(tmp_path): + from tools.mcp_tool_discovery import register_mcp_servers + from tools.mcp_tool_lifecycle import shutdown_mcp_servers + + server = tmp_path / "server.py" + server.write_text('''import json, pathlib, sys +root = pathlib.Path(__file__).parent +for line in sys.stdin: + req = json.loads(line) + if "id" not in req: + continue + result = {} + if req["method"] == "initialize": + result = {"protocolVersion": req["params"]["protocolVersion"], + "capabilities": {"tools": {}}, "serverInfo": {"name": "batch", "version": "1"}} + elif req["method"] == "tools/list": + result = {"tools": [{"name": "read_" + n, "description": "Read " + n, + "inputSchema": {"type": "object", "properties": {"key": {"type": "string"}}, "required": ["key"]}, + "annotations": {"readOnlyHint": True}} for n in ("alpha", "beta")]} + elif req["method"] == "tools/call": + name = req["params"]["name"] + with (root / "calls.jsonl").open("a", encoding="utf-8") as log: + log.write(json.dumps(name) + "\\n") + result = {"content": [{"type": "text", "text": name + "-result"}]} + print(json.dumps({"jsonrpc": "2.0", "id": req["id"], "result": result}), flush=True) +''', encoding="utf-8") + names = sorted(register_mcp_servers({"batchfixture": { + "command": sys.executable, "args": [str(server)], "timeout": 10}})) + assert len(names) == 2 + try: + yield names, tmp_path / "calls.jsonl" + finally: + shutdown_mcp_servers() + + +@pytest.mark.parametrize("mode", ["allow", "block", "scope", "schema", "interrupt"]) +def test_local_batch_runs_once_per_entry_through_agent_and_persists_pairs(mcp_server, mode): + from run_agent import AIAgent + + names, call_log = mcp_server + with patch("agent.process_bootstrap.OpenAI"), patch("agent.model_metadata.fetch_model_metadata", return_value={}): + agent = AIAgent( + model="test/model", api_key="test-key", base_url="http://127.0.0.1:1/v1", + enabled_toolsets=["mcp-batchfixture"], quiet_mode=True, + skip_context_files=True, skip_memory=True, max_iterations=3) + agent.client = MagicMock() + agent._cached_system_prompt = "Use the available tools." + agent._use_prompt_caching = False + agent.compression_enabled = False + agent.save_trajectories = False + entries = [{"name": name, "arguments": {"key": "value"}} for name in names] + if mode == "schema": + entries[1]["arguments"] = {} + if mode == "scope": + agent.disabled_toolsets = ["mcp-batchfixture"] + agent.client.chat.completions.create.side_effect = [_response([_call(entries)]), _response()] + if mode == "interrupt": + def stop_after_first(event, name, *args, **kwargs): + if event == "tool.completed" and name == names[0]: + agent.interrupt("test stop after first result") + agent.tool_progress_callback = stop_after_first + + hooks, snapshots = [], [] + + def pre_hook(name, args, **kwargs): + hooks.append(name) + return ("blocked by test policy" if mode == "block" and name == names[1] else None), None + + flush = agent._flush_messages_to_session_db + + def capture_flush(messages, *args, **kwargs): + snapshots.append(copy.deepcopy(messages)) + return flush(messages, *args, **kwargs) + + with patch("hermes_cli.plugins._dispatch_pre_tool_call_hooks", side_effect=pre_hook), patch.object( + agent, "_flush_messages_to_session_db", side_effect=capture_flush + ): + result = agent.run_conversation("Read alpha and beta without changing anything.") + + if mode != "interrupt": + assert result["final_response"] == "done" + expected = {"allow": ["read_alpha", "read_beta"], "scope": []}.get(mode, ["read_alpha"]) + executed = [json.loads(line) for line in call_log.read_text(encoding="utf-8").splitlines()] if call_log.exists() else [] + assert executed == expected + assert hooks == ([] if mode == "scope" else names[:1] if mode in {"schema", "interrupt"} else names) + messages = result["messages"] + assistant = next(m for m in messages if m.get("tool_calls")) + results = [m for m in messages if m["role"] == "tool"] + ids = [c["id"] for c in assistant["tool_calls"]] + assert len(set(ids)) == 2 + assert [m["tool_call_id"] for m in results] == ids + assert ("not available" if mode == "scope" else "read_alpha-result") in results[0]["content"] + expected_second = {"allow": "read_beta-result", "block": "blocked by test policy", + "scope": "not available", "schema": "key", "interrupt": "skipped"}[mode] + assert expected_second in results[1]["content"] + # Both calls are durable before the first result, and every result is flushed. + before_results = next(s for s in snapshots if any(m.get("tool_calls") for m in s)) + assert not any(m["role"] == "tool" for m in before_results) + assert next(m for m in before_results if m.get("tool_calls"))["tool_calls"] == assistant["tool_calls"] + assert any(sum(m["role"] == "tool" for m in s) == 2 for s in snapshots) + + +def test_expansion_preserves_connector_batches_and_rejects_bad_envelopes(): + from agent.tool_call_batches import expand_local_tool_batches + from tools.connectors.gateway.config import MAX_CALLS_PER_DISPATCH + + local = {"name": "mcp__fixture__read", "arguments": {"path": "alpha"}} + remote = {"name": "connectors__drive__list", "arguments": {}} + for entries in ([remote, remote], [local], [local] * (MAX_CALLS_PER_DISPATCH + 1), [local, {}]): + call = _call(entries) + original = call.function.arguments + assert expand_local_tool_batches([call]) == [call] + assert call.function.arguments == original + call = _call([local, remote]) + original = call.function.arguments + expanded = expand_local_tool_batches([call]) + assert [json.loads(c.function.arguments)["calls"] for c in expanded] == [[local], [remote]] + assert call.function.arguments == original + assert [c.id for c in expand_local_tool_batches([call])] == [c.id for c in expanded] + thinking = {"type": "thinking", "thinking": "opaque", "signature": "signed"} + bedrock_thinking = {"reasoningContent": {"reasoningText": {"text": "opaque", "signature": "signed"}}} + provider_data = { + "anthropic_content_blocks": [thinking, {"type": "tool_use", "id": "batch", + "name": "tool_call", "input": json.loads(original)}], + "bedrock_content_blocks": [bedrock_thinking, {"toolUse": {"toolUseId": "batch", + "name": "tool_call", "input": json.loads(original)}}], + } + expanded = expand_local_tool_batches([call], provider_data=provider_data) + assert provider_data["anthropic_content_blocks"][0] is thinking + assert provider_data["bedrock_content_blocks"][0] is bedrock_thinking + for index, child in enumerate(expanded, 1): + anthropic = provider_data["anthropic_content_blocks"][index] + bedrock = provider_data["bedrock_content_blocks"][index]["toolUse"] + assert anthropic["id"] == bedrock["toolUseId"] == child.id + assert anthropic["input"] == bedrock["input"] == json.loads(child.function.arguments) diff --git a/tools/tool_search.py b/tools/tool_search.py index 9f7998dcfa..a460a4466d 100644 --- a/tools/tool_search.py +++ b/tools/tool_search.py @@ -311,9 +311,9 @@ def bridge_tool_schemas(deferred_count: int, listing: Optional[str] = None, TOOL_CALL_NAME, "Invoke deferred tools. Takes `calls`, an array of {name, arguments} " "— one entry per invocation; a single call is an array of one. " - "Local tools require one entry per tool_call. Only connectors__ names " - "may be batched together; mixed and multi-local batches are rejected. " - "Connector entries execute individually with results in input order. " + "Independent calls may be batched together. Local entries are split " + "into individual tool calls by the agent and follow each tool's " + "concurrency policy. Results are returned in input order. " f"Argument shapes match each tool's schema (see `{TOOL_DESCRIBE_NAME}`). " "Policy, hooks, and approvals run as for directly-listed tools.", { @@ -327,7 +327,7 @@ def bridge_tool_schemas(deferred_count: int, listing: Optional[str] = None, }, "required": ["name", "arguments"], }, - "description": "One local invocation, or one or more connector invocations. Never mix local and connector tools.", + "description": "One or more independent deferred tool invocations.", }, }, ["calls"],