from __future__ import annotations import asyncio import json import pytest from opensquilla.engine.tool_result_store import ToolResultStore from opensquilla.engine.types import ToolCall from opensquilla.result_budget import ( DEFAULT_TOOL_RUN_BUDGET_POLICY, ToolResultBudgetPolicy, ToolRunBudgetExceededError, ToolRunBudgetPolicy, build_web_retrieval_tool_run_budget_policy, ) from opensquilla.tools import dispatch as dispatch_module from opensquilla.tools.dispatch import build_tool_handler from opensquilla.tools.registry import ToolRegistry from opensquilla.tools.types import ( CallerKind, InteractionMode, ToolContext, ToolSpec, current_tool_context, ) def _build_registry() -> ToolRegistry: registry = ToolRegistry() async def boom() -> str: raise ValueError("bad argument") async def echo(value: str = "") -> str: return value async def required_echo(value: str) -> str: return value async def pending() -> str: return json.dumps( { "status": "approval_required", "approval_id": "abc123", "command": "rm secret", "warning": "destructive", "message": "Resolve this approval via exec.approval.resolve", } ) registry.register(ToolSpec(name="boom", description="boom", parameters={}), boom) registry.register( ToolSpec( name="echo", description="echo", parameters={"value": {"type": "string"}}, ), echo, ) registry.register( ToolSpec( name="required_echo", description="required echo", parameters={"value": {"type": "string"}}, required=["value"], ), required_echo, ) registry.register(ToolSpec(name="pending", description="pending", parameters={}), pending) return registry def _strict_preview_chars(content: str) -> int: payload = json.loads(content) assert payload["result_truncated"] is True assert "tool_result_budget_applied" not in payload assert "result_returned_chars" not in payload assert "budget_class" not in payload preview = payload.get("preview", "") assert isinstance(preview, str) tail = payload.get("tail", "") assert isinstance(tail, str) return len(preview) + len(tail) def _assert_invalid_attempt_result( result, *, tool: str, reason_code: str, received_keys: list[str], ) -> dict[str, object]: assert result.is_error is True assert result.execution_status is not None assert result.execution_status["status"] == "error" assert result.execution_status["reason"] == "invalid_tool_arguments" assert result.execution_status["preflight_rejected"] is True assert result.execution_status["reason_code"] == reason_code payload = json.loads(result.content) assert payload["status"] == "rejected" assert payload["reason_code"] == reason_code assert payload["tool"] == tool assert payload["received_keys"] == received_keys assert payload["retry_allowed"] is True assert payload["error_class"] == "InvalidToolArgumentsError" return payload def test_web_retrieval_budget_profile_builder_preserves_unlimited_call_defaults() -> None: policy = build_web_retrieval_tool_run_budget_policy() assert policy.max_web_search_calls_per_turn is None assert policy.max_web_fetch_calls_per_turn is None assert policy.max_external_text_chars_per_turn is None assert policy.max_single_fetch_chars == 50_000 assert policy.max_web_search_results == 10 def test_default_tool_run_budget_leaves_room_for_complex_turns() -> None: search_calls = DEFAULT_TOOL_RUN_BUDGET_POLICY.max_web_search_calls_per_turn fetch_calls = DEFAULT_TOOL_RUN_BUDGET_POLICY.max_web_fetch_calls_per_turn external_chars = DEFAULT_TOOL_RUN_BUDGET_POLICY.max_external_text_chars_per_turn single_fetch_chars = DEFAULT_TOOL_RUN_BUDGET_POLICY.max_single_fetch_chars search_results = DEFAULT_TOOL_RUN_BUDGET_POLICY.max_web_search_results assert search_calls is None assert fetch_calls is None assert external_chars is None assert single_fetch_chars is not None and single_fetch_chars >= 50_000 assert search_results is not None and search_results >= 10 @pytest.mark.asyncio async def test_dispatch_missing_tool_returns_five_field_error_envelope() -> None: # Use a trusted CLI ctx so the descriptive ``ToolNotFound`` branch is # exercised. Anonymous/CHANNEL callers receive an opaque ``PolicyDenied`` # envelope to prevent registry enumeration; that branch is covered in # ``test_dispatch_surface_hardening``. handler = build_tool_handler( _build_registry(), ToolContext( is_owner=True, caller_kind=CallerKind.CLI, agent_id="main", session_key="cli:main:envelope", ), ) result = await handler( ToolCall( tool_use_id="tc-1", tool_name="nope", arguments={}, ) ) assert result.is_error is True payload = json.loads(result.content) assert set(payload.keys()) == { "status", "tool", "error_class", "user_message", "retry_allowed", } assert payload["status"] == "error" assert payload["tool"] == "nope" assert payload["error_class"] == "ToolNotFound" assert "Do not retry unavailable tools" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_unknown_bash_tool_points_to_exec_command() -> None: handler = build_tool_handler( _build_registry(), ToolContext( is_owner=True, caller_kind=CallerKind.CLI, agent_id="main", session_key="cli:main:envelope", ), ) result = await handler( ToolCall( tool_use_id="tc-bash", tool_name="bash", arguments={"cmd": "echo hi"}, ) ) assert result.is_error is True payload = json.loads(result.content) assert payload["error_class"] == "ToolNotFound" assert "Use exec_command with a command string instead" in payload["user_message"] assert "do not retry bash as a tool" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_tool_exception_envelope_is_canonical_five_key_shape() -> None: handler = build_tool_handler(_build_registry()) result = await handler( ToolCall( tool_use_id="tc-2", tool_name="boom", arguments={}, ) ) assert result.is_error is True payload = json.loads(result.content) assert payload["tool"] == "boom" assert payload["error_class"] == "ValueError" assert set(payload.keys()) == { "status", "tool", "error_class", "user_message", "retry_allowed", } @pytest.mark.asyncio async def test_dispatch_redacts_secret_like_tool_result_content() -> None: handler = build_tool_handler(_build_registry()) secret = "sk-or-v1-abcdefghijklmnopqrstuvwxyz" result = await handler( ToolCall( tool_use_id="tc-secret", tool_name="echo", arguments={"value": f"env.OPENROUTER_API_KEY={secret}"}, ) ) assert result.is_error is False assert secret not in result.content assert result.content == "env.OPENROUTER_API_KEY=[REDACTED]" @pytest.mark.asyncio async def test_dispatch_rejects_unparsed_raw_tool_arguments_before_handler() -> None: handler = build_tool_handler(_build_registry()) result = await handler( ToolCall( tool_use_id="tc-raw", tool_name="echo", arguments={"_raw": '{"value": "unescaped " quote"}'}, ) ) payload = _assert_invalid_attempt_result( result, tool="echo", reason_code="unparsed_raw_arguments", received_keys=["_raw"], ) assert "valid JSON" in payload["user_message"] assert "apply_patch" not in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_unwraps_schema_valid_raw_json_arguments_before_handler() -> None: handler = build_tool_handler(_build_registry()) result = await handler( ToolCall( tool_use_id="tc-valid-raw", tool_name="echo", arguments={"_raw": '{"value": "ok"}'}, ) ) assert result.is_error is False assert result.content == "ok" @pytest.mark.asyncio async def test_dispatch_rejects_missing_required_arguments_before_handler() -> None: handler = build_tool_handler(_build_registry()) result = await handler( ToolCall( tool_use_id="tc-missing-required", tool_name="required_echo", arguments={}, ) ) payload = _assert_invalid_attempt_result( result, tool="required_echo", reason_code="missing_required_arguments", received_keys=[], ) assert "missing required argument" in payload["user_message"] assert "`value`" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_missing_edit_file_old_text_explains_valid_shapes() -> None: registry = ToolRegistry() async def edit_file(path: str, old_text: str, new_text: str) -> str: return f"{path}:{old_text}:{new_text}" registry.register( ToolSpec( name="edit_file", description="edit file", parameters={ "path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}, "edits": {"type": "array"}, }, required=["path", "old_text", "new_text"], ), edit_file, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-edit-missing-old", tool_name="edit_file", arguments={"path": "demo.py", "new_text": "new"}, ) ) payload = _assert_invalid_attempt_result( result, tool="edit_file", reason_code="missing_required_arguments", received_keys=["new_text", "path"], ) assert "new_text alone cannot identify where to edit" in payload["user_message"] assert '{"path":"...","old_text":"...","new_text":"..."}' in payload["user_message"] assert '{"path":"...","edits":[{"old_text":"...","new_text":"..."}]}' in payload[ "user_message" ] @pytest.mark.asyncio async def test_dispatch_missing_write_file_content_shows_minimal_json() -> None: registry = ToolRegistry() async def write_file(path: str, content: str) -> str: return f"{path}:{content}" registry.register( ToolSpec( name="write_file", description="write file", parameters={ "path": {"type": "string"}, "content": {"type": "string"}, }, required=["path", "content"], ), write_file, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-write-missing-content", tool_name="write_file", arguments={"path": "demo.py"}, ) ) payload = json.loads(result.content) assert result.is_error is True assert '{"path":"...","content":"..."}' in payload["user_message"] assert "existing source files" in payload["user_message"] assert "prefer edit_file or apply_patch" in payload["user_message"] assert "You supplied argument(s)" not in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_missing_edit_file_guidance_omits_hidden_apply_patch() -> None: registry = ToolRegistry() async def edit_file( path: str, old_text: str | None = None, new_text: str | None = None, edits: list[dict[str, object]] | None = None, ) -> str: return f"{path}:{old_text}:{new_text}:{edits}" registry.register( ToolSpec( name="edit_file", description="edit file", parameters={ "path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}, "edits": {"type": "array"}, }, required=["path", "old_text", "new_text"], ), edit_file, ) handler = build_tool_handler( registry, ToolContext(allowed_tools={"edit_file"}), ) result = await handler( ToolCall( tool_use_id="tc-edit-hidden-patch", tool_name="edit_file", arguments={"path": "demo.py", "new_text": "new"}, ) ) payload = _assert_invalid_attempt_result( result, tool="edit_file", reason_code="missing_required_arguments", received_keys=["new_text", "path"], ) assert "new_text alone cannot identify where to edit" in payload["user_message"] assert "apply_patch" not in payload["user_message"] assert "split the edit into smaller edit_file calls" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_missing_write_file_guidance_omits_hidden_edit_tools() -> None: registry = ToolRegistry() async def write_file(path: str, content: str) -> str: return f"{path}:{content}" registry.register( ToolSpec( name="write_file", description="write file", parameters={ "path": {"type": "string"}, "content": {"type": "string"}, }, required=["path", "content"], ), write_file, ) handler = build_tool_handler( registry, ToolContext(allowed_tools={"write_file"}), ) result = await handler( ToolCall( tool_use_id="tc-write-hidden-edit-tools", tool_name="write_file", arguments={"path": "demo.py"}, ) ) payload = _assert_invalid_attempt_result( result, tool="write_file", reason_code="missing_required_arguments", received_keys=["path"], ) assert '{"path":"...","content":"..."}' in payload["user_message"] assert "edit_file" not in payload["user_message"] assert "apply_patch" not in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_missing_write_file_guidance_uses_visible_edit_file_only() -> None: registry = ToolRegistry() async def write_file(path: str, content: str) -> str: return f"{path}:{content}" registry.register( ToolSpec( name="write_file", description="write file", parameters={ "path": {"type": "string"}, "content": {"type": "string"}, }, required=["path", "content"], ), write_file, ) handler = build_tool_handler( registry, ToolContext(allowed_tools={"write_file", "edit_file"}), ) result = await handler( ToolCall( tool_use_id="tc-write-visible-edit-only", tool_name="write_file", arguments={"path": "demo.py"}, ) ) payload = _assert_invalid_attempt_result( result, tool="write_file", reason_code="missing_required_arguments", received_keys=["path"], ) assert "prefer edit_file" in payload["user_message"] assert "apply_patch" not in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_missing_write_file_path_can_include_supplied_argument_keys() -> None: registry = ToolRegistry() async def write_file(path: str, content: str) -> str: return f"{path}:{content}" registry.register( ToolSpec( name="write_file", description="write file", parameters={ "path": {"type": "string"}, "content": {"type": "string"}, }, required=["path", "content"], ), write_file, ) handler = build_tool_handler( registry, ToolContext(missing_required_argument_shape_guidance=True), ) result = await handler( ToolCall( tool_use_id="tc-write-missing-path", tool_name="write_file", arguments={"content": "print('scratch')\n"}, ) ) payload = json.loads(result.content) assert result.is_error is True assert "You supplied argument(s): `content`." in payload["user_message"] assert "Missing argument(s): `path`." in payload["user_message"] assert '{"path":"...","content":"..."}' in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_missing_required_shape_guidance_can_be_enabled_by_env( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("OPENSQUILLA_MISSING_REQUIRED_ARGUMENT_SHAPE_GUIDANCE", "1") registry = ToolRegistry() async def write_file(path: str, content: str) -> str: return f"{path}:{content}" registry.register( ToolSpec( name="write_file", description="write file", parameters={ "path": {"type": "string"}, "content": {"type": "string"}, }, required=["path", "content"], ), write_file, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-write-env-missing-path", tool_name="write_file", arguments={"content": "print('scratch')\n"}, ) ) payload = json.loads(result.content) assert result.is_error is True assert "You supplied argument(s): `content`." in payload["user_message"] assert "Missing argument(s): `path`." in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_missing_apply_patch_patch_shows_patch_shape() -> None: registry = ToolRegistry() async def apply_patch(patch: str) -> str: return patch registry.register( ToolSpec( name="apply_patch", description="apply patch", parameters={"patch": {"type": "string"}}, required=["patch"], ), apply_patch, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-patch-missing-patch", tool_name="apply_patch", arguments={}, ) ) payload = json.loads(result.content) assert result.is_error is True assert '{"patch":"*** Begin Patch' in payload["user_message"] assert "*** Update File: ..." in payload["user_message"] assert "*** End Patch" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_missing_exec_command_command_shows_minimal_json() -> None: registry = ToolRegistry() async def exec_command(command: str) -> str: return command registry.register( ToolSpec( name="exec_command", description="execute command", parameters={"command": {"type": "string"}}, required=["command"], ), exec_command, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-exec-missing-command", tool_name="exec_command", arguments={"new_text": "python test.py"}, ) ) payload = json.loads(result.content) assert result.is_error is True assert '{"command":"..."}' in payload["user_message"] assert "Do not put shell text in `new_text`, `path`, or `_raw`" in payload[ "user_message" ] @pytest.mark.asyncio async def test_dispatch_missing_execute_code_code_shows_minimal_json() -> None: registry = ToolRegistry() async def execute_code(code: str) -> str: return code registry.register( ToolSpec( name="execute_code", description="execute Python code", parameters={"code": {"type": "string"}}, required=["code"], ), execute_code, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-code-missing-code", tool_name="execute_code", arguments={"command": "print('hi')"}, ) ) payload = json.loads(result.content) assert result.is_error is True assert '{"code":"..."}' in payload["user_message"] assert "Use exec_command for shell commands" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_strips_provider_replay_arguments_before_handler() -> None: handler = build_tool_handler(_build_registry()) result = await handler( ToolCall( tool_use_id="tc-replay", tool_name="echo", arguments={ "_opensquilla_replay_index": 2, "_opensquilla_replay_verbosity": "1", "value": "ok", }, ) ) assert result.is_error is False assert result.content == "ok" @pytest.mark.asyncio async def test_dispatch_unwraps_nested_json_arguments_before_handler() -> None: handler = build_tool_handler(_build_registry()) result = await handler( ToolCall( tool_use_id="tc-nested-arguments", tool_name="echo", arguments={ "_opensquilla_replay_nonce": "history-replay-1", "arguments": '{"value": "ok"}', }, ) ) assert result.is_error is False assert result.content == "ok" @pytest.mark.asyncio async def test_dispatch_normalizes_common_edit_file_aliases_before_required_check() -> None: registry = ToolRegistry() async def edit_file(path: str, old_text: str, new_text: str) -> str: return json.dumps( { "path": path, "old_text": old_text, "new_text": new_text, } ) registry.register( ToolSpec( name="edit_file", description="edit file", parameters={ "path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}, }, required=["path", "old_text", "new_text"], ), edit_file, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-edit-aliases", tool_name="edit_file", arguments={ "file_path": "/tmp/example.py", "old_string": "old", "new_string": "new", }, ) ) assert result.is_error is False assert json.loads(result.content) == { "path": "/tmp/example.py", "old_text": "old", "new_text": "new", } @pytest.mark.asyncio async def test_dispatch_normalizes_camel_case_edit_file_aliases_before_required_check() -> None: registry = ToolRegistry() async def edit_file(path: str, old_text: str, new_text: str) -> str: return json.dumps( { "path": path, "old_text": old_text, "new_text": new_text, } ) registry.register( ToolSpec( name="edit_file", description="edit file", parameters={ "path": {"type": "string"}, "old_text": {"type": "string"}, "new_text": {"type": "string"}, }, required=["path", "old_text", "new_text"], ), edit_file, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-edit-camel-aliases", tool_name="edit_file", arguments={ "filePath": "/tmp/example.py", "oldText": "old", "newText": "new", }, ) ) assert result.is_error is False assert json.loads(result.content) == { "path": "/tmp/example.py", "old_text": "old", "new_text": "new", } @pytest.mark.asyncio @pytest.mark.parametrize( ("marker_key", "marker_value"), [ ("_opensquilla_compacted_tool_arguments", True), ("_opensquilla_compacted_tool_input", True), ("_opensquilla_compacted_tool_arguments", "true"), ("_opensquilla_compacted_tool_input", "true"), ("_invalid_provider_context_arguments", "true"), ], ) async def test_dispatch_rejects_provider_compacted_tool_arguments_before_handler( marker_key: str, marker_value: object, ) -> None: handler = build_tool_handler(_build_registry()) result = await handler( ToolCall( tool_use_id="tc-compacted", tool_name="echo", arguments={ marker_key: marker_value, "head": '{"value": "large', "tail": 'payload"}', }, ) ) assert result.is_error is True assert result.execution_status is not None assert result.execution_status["reason"] == "provider_context_projection_reused" payload = json.loads(result.content) assert payload["tool"] == "echo" assert payload["error_class"] == "ProjectedToolArgumentsError" assert payload["retry_allowed"] is False assert "compacted" in payload["user_message"] @pytest.mark.asyncio @pytest.mark.parametrize( "arguments", [ { "value": ( "[provider_request_tool_input_compacted: original_chars=987; " "sha256=aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa]" ) }, { "value": { "nested": [ "safe prefix", ( "[provider_request_tool_input_compacted: original_chars=987; " "sha256=aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa]" ), ] } }, ], ) async def test_dispatch_rejects_provider_request_compacted_argument_strings_before_handler( arguments: dict[str, object], ) -> None: calls: list[object] = [] registry = ToolRegistry() async def echo(value: object = "") -> str: calls.append(value) return json.dumps(value) registry.register( ToolSpec( name="echo", description="echo", parameters={"value": {}}, ), echo, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-provider-request-compacted", tool_name="echo", arguments=arguments, ) ) assert calls == [] assert result.is_error is True assert result.execution_status is not None assert result.execution_status["reason"] == "provider_context_projection_reused" payload = json.loads(result.content) assert payload["tool"] == "echo" assert payload["error_class"] == "ProjectedToolArgumentsError" assert payload["retry_allowed"] is False assert "compacted" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_rejects_mid_string_compacted_marker_before_handler() -> None: calls: list[object] = [] registry = ToolRegistry() async def edit_file( path: str = "", old_text: str = "", new_text: str = "", ) -> str: calls.append((path, old_text, new_text)) return json.dumps({"path": path}) registry.register( ToolSpec( name="edit_file", description="edit", parameters={"path": {}, "old_text": {}, "new_text": {}}, ), edit_file, ) handler = build_tool_handler(registry) result = await handler( ToolCall( tool_use_id="tc-mid-string-compacted", tool_name="edit_file", arguments={ "path": "src/example.py", "old_text": ( "def load(self):\n" " [provider_request_tool_result_compacted: omitted 4210 chars; " "original_chars=4655; sha256=" + "b" * 64 + "]\n return data\n" ), "new_text": "def load(self):\n return data\n", }, ) ) assert calls == [] assert result.is_error is True assert result.execution_status is not None assert result.execution_status["reason"] == "provider_context_projection_reused" payload = json.loads(result.content) assert payload["tool"] == "edit_file" assert payload["error_class"] == "ProjectedToolArgumentsError" assert payload["retry_allowed"] is False assert "compacted" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_unsupported_surface_approval_payload_is_pending_status() -> None: handler = build_tool_handler(_build_registry()) token = current_tool_context.set( ToolContext( is_owner=True, caller_kind=CallerKind.CHANNEL, interaction_mode=InteractionMode.UNATTENDED, session_key="agent:main:demo", agent_id="main", ) ) try: result = await handler( ToolCall( tool_use_id="tc-3", tool_name="pending", arguments={}, ) ) finally: current_tool_context.reset(token) assert result.is_error is False assert result.execution_status is not None assert result.execution_status["status"] == "unknown" assert result.execution_status["reason"] == "approval_pending" assert result.execution_status["preservation_class"] == "ephemeral" payload = json.loads(result.content) assert payload["status"] == "error" assert payload["tool"] == "pending" assert payload["error_class"] == "UnsupportedSurface" assert payload["retry_allowed"] is False @pytest.mark.asyncio async def test_dispatch_unattended_cli_approval_payload_is_pending_status() -> None: handler = build_tool_handler(_build_registry()) token = current_tool_context.set( ToolContext( is_owner=True, caller_kind=CallerKind.CLI, interaction_mode=InteractionMode.UNATTENDED, session_key="agent:main:demo", agent_id="main", ) ) try: result = await handler( ToolCall( tool_use_id="tc-4", tool_name="pending", arguments={}, ) ) finally: current_tool_context.reset(token) assert result.is_error is False assert result.execution_status is not None assert result.execution_status["status"] == "unknown" assert result.execution_status["reason"] == "approval_pending" payload = json.loads(result.content) assert set(payload.keys()) == { "status", "tool", "error_class", "user_message", "retry_allowed", } assert payload["status"] == "error" assert payload["tool"] == "pending" assert payload["error_class"] == "UnsupportedSurface" assert payload["retry_allowed"] is False @pytest.mark.asyncio async def test_dispatch_unknown_tool_in_skill_name_context_raises_unsupported_surface() -> None: known_skill_names = {"shell"} # Use a trusted CLI ctx so the descriptive ``UnsupportedSurface`` skill # branch is exercised. CHANNEL/anonymous callers receive an opaque # ``PolicyDenied`` envelope to prevent skill-name enumeration; that # branch is covered in ``test_dispatch_surface_hardening``. handler = build_tool_handler( _build_registry(), ToolContext( is_owner=True, caller_kind=CallerKind.CLI, agent_id="main", session_key="cli:main:envelope", ), known_skill_names=known_skill_names, ) result = await handler( ToolCall( tool_use_id="tc-4", tool_name="shell", arguments={}, ) ) payload = json.loads(result.content) assert result.is_error is True assert payload["error_class"] == "UnsupportedSurface" assert payload["tool"] == "shell" assert "skill" in payload["user_message"].lower() @pytest.mark.asyncio async def test_dispatch_attaches_published_artifacts_to_tool_result() -> None: registry = ToolRegistry() artifact = { "id": "art-dispatch", "kind": "artifact_ref", "name": "report.txt", "mime": "text/plain", "size": 4, "sha256": "1" * 64, "session_id": "session-1", "session_key": "agent:main:demo", "source": "publish_artifact", "created_at": "2026-05-06T12:00:00Z", "download_url": "/api/v1/artifacts/art-dispatch", } async def publish() -> str: ctx = current_tool_context.get() assert ctx is not None ctx.published_artifacts.append(artifact) return "published" registry.register(ToolSpec(name="publish", description="publish", parameters={}), publish) handler = build_tool_handler(registry) ctx = ToolContext(session_key="agent:main:demo") token = current_tool_context.set(ctx) try: result = await handler( ToolCall( tool_use_id="tc-5", tool_name="publish", arguments={}, ) ) finally: current_tool_context.reset(token) assert result.content == "published" assert result.artifacts == [artifact] @pytest.mark.asyncio async def test_dispatch_does_not_rewrite_artifact_result_text_when_over_budget() -> None: registry = ToolRegistry() artifact = { "id": "art-large", "kind": "artifact_ref", "name": "report.txt", "mime": "text/plain", "size": 4, "sha256": "2" * 64, "session_id": "session-1", "session_key": "agent:main:demo", "source": "publish_artifact", "created_at": "2026-05-06T12:00:00Z", "download_url": "/api/v1/artifacts/art-large", } async def publish() -> str: ctx = current_tool_context.get() assert ctx is not None ctx.published_artifacts.append(artifact) return "x" * 1000 registry.register(ToolSpec(name="publish", description="publish", parameters={}), publish) handler = build_tool_handler( registry, ToolContext( session_key="agent:main:demo", tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=120, ), ), ) result = await handler( ToolCall(tool_use_id="tc-art-large", tool_name="publish", arguments={}) ) assert result.content == "x" * 1000 assert result.artifacts == [artifact] @pytest.mark.asyncio async def test_dispatch_leaves_under_budget_result_unchanged() -> None: registry = ToolRegistry() async def echo() -> str: return '{"status":"ok","value":"unchanged"}' registry.register(ToolSpec(name="echo", description="echo", parameters={}), echo) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=10_000 ) ), ) result = await handler(ToolCall(tool_use_id="tc-under", tool_name="echo", arguments={})) assert result.content == '{"status":"ok","value":"unchanged"}' assert result.is_error is False @pytest.mark.asyncio async def test_dispatch_leaves_default_huge_tool_result_unchanged() -> None: registry = ToolRegistry() async def huge() -> str: return "x" * 1000 registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) handler = build_tool_handler(registry) result = await handler(ToolCall(tool_use_id="tc-huge-default", tool_name="huge", arguments={})) assert result.content == "x" * 1000 assert result.artifacts == [] @pytest.mark.asyncio async def test_dispatch_strict_policy_bounds_unknown_huge_tool_result() -> None: registry = ToolRegistry() async def huge() -> str: return "x" * 1000 registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=120, max_tool_result_chars_per_turn=200, ) ), ) result = await handler(ToolCall(tool_use_id="tc-huge", tool_name="huge", arguments={})) payload = json.loads(result.content) assert payload["result_truncated"] is True assert payload["result_original_chars"] == 1000 assert payload["result_omitted_chars"] == 880 assert len(payload["preview"]) + len(payload["tail"]) <= 120 assert payload["preview"] == "x" * 78 assert payload["tail"] == "x" * 42 assert "tool_result_budget_applied" not in payload assert "result_returned_chars" not in payload assert "budget_class" not in payload assert result.artifacts == [] assert len(result.content) < 400 @pytest.mark.asyncio async def test_dispatch_strict_policy_preserves_tail_of_huge_tool_result() -> None: registry = ToolRegistry() async def huge() -> str: return "HEAD-" + ("x" * 200) + "-TRACEBACK" registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=40, ) ), ) result = await handler( ToolCall(tool_use_id="tc-huge-tail", tool_name="huge", arguments={}) ) payload = json.loads(result.content) assert payload["preview"].startswith("HEAD-") assert payload["tail"].endswith("-TRACEBACK") assert len(payload["preview"]) + len(payload["tail"]) <= 40 @pytest.mark.asyncio async def test_dispatch_execution_policy_stores_raw_snapshot_for_truncated_exec_result( tmp_path, ) -> None: registry = ToolRegistry() raw_output = "HEAD\n" + ("x" * 9_000_000) + "\nTAIL" async def exec_command(command: str) -> str: assert command == "pytest -q" return raw_output registry.register( ToolSpec( name="exec_command", description="exec", parameters={"command": {"type": "string"}}, required=["command"], ), exec_command, ) ctx = ToolContext( session_key="agent:main:session-1", tool_result_store_dir=str(tmp_path / "tool-results"), tool_result_store_session_id="session-1", tool_result_budget_policy=ToolResultBudgetPolicy( max_single_execution_result_chars=10_000, ), ) handler = build_tool_handler(registry, ctx) result = await handler( ToolCall( tool_use_id="tc-large-exec", tool_name="exec_command", arguments={"command": "pytest -q"}, ) ) payload = json.loads(result.content) assert payload["result_truncated"] is True assert payload["result_original_chars"] == len(raw_output) assert payload["preview"].startswith("HEAD") assert payload["tail"].endswith("TAIL") assert len(payload["preview"]) + len(payload["tail"]) <= 10_000 assert payload["tool_result_handle"].startswith("tr-") assert "retrieve_tool_result" in payload["retrieve_hint"] assert len(result.content) < 12_000 stored = ToolResultStore(tmp_path / "tool-results").read( payload["tool_result_handle"], session_id="session-1", ) assert stored.content == raw_output assert stored.storage_encoding == "gzip+utf-8" assert stored.stored_size_bytes is not None assert stored.stored_size_bytes < 8 * 1024 * 1024 @pytest.mark.asyncio async def test_dispatch_execution_policy_bounds_broad_exec_result_only() -> None: registry = ToolRegistry() async def exec_command(command: str) -> str: return f"command={command}\n" + ("x" * 200) + "\nTRACEBACK" async def read_file() -> str: return "source\n" + ("x" * 200) registry.register( ToolSpec( name="exec_command", description="exec", parameters={"command": {"type": "string"}}, required=["command"], ), exec_command, ) registry.register(ToolSpec(name="read_file", description="read", parameters={}), read_file) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_execution_result_chars=40, ) ), ) broad = await handler( ToolCall( tool_use_id="tc-broad", tool_name="exec_command", arguments={"command": "pytest -q"}, ) ) preserved = await handler( ToolCall(tool_use_id="tc-read", tool_name="read_file", arguments={}) ) broad_payload = json.loads(broad.content) assert broad_payload["result_truncated"] is True assert broad_payload["tail"].endswith("TRACEBACK") assert preserved.content == "source\n" + ("x" * 200) @pytest.mark.asyncio async def test_dispatch_execution_policy_preserves_semantic_exec_results() -> None: registry = ToolRegistry() async def exec_command(command: str) -> str: return f"command={command}\n" + ("important diff\n" * 40) registry.register( ToolSpec( name="exec_command", description="exec", parameters={"command": {"type": "string"}}, required=["command"], ), exec_command, ) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_execution_result_chars=40, ) ), ) git_diff = await handler( ToolCall( tool_use_id="tc-diff", tool_name="exec_command", arguments={"command": "git diff"}, ) ) source_read = await handler( ToolCall( tool_use_id="tc-source", tool_name="exec_command", arguments={"command": "sed -n '1,80p' src/lib.rs"}, ) ) markdown_read = await handler( ToolCall( tool_use_id="tc-markdown", tool_name="exec_command", arguments={"command": "cat README.md"}, ) ) assert "result_truncated" not in git_diff.content assert "important diff" in git_diff.content assert "result_truncated" not in source_read.content assert "important diff" in source_read.content assert "result_truncated" not in markdown_read.content assert "important diff" in markdown_read.content @pytest.mark.asyncio async def test_dispatch_preserves_error_preview_after_turn_budget_is_exhausted() -> None: registry = ToolRegistry() async def huge() -> str: return "x" * 1000 async def missing_capability() -> str: return json.dumps( { "status": "blocked", "message": "Skill not found: nano-banana", } ) registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) registry.register( ToolSpec(name="missing_capability", description="missing", parameters={}), missing_capability, ) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=1000, max_tool_result_chars_per_turn=5, ) ), ) await handler(ToolCall(tool_use_id="tc-huge", tool_name="huge", arguments={})) result = await handler( ToolCall( tool_use_id="tc-missing", tool_name="missing_capability", arguments={}, ) ) payload = json.loads(result.content) assert result.is_error is True assert payload["status"] == "blocked" assert payload["message"] == "Skill not found: nano-banana" @pytest.mark.asyncio async def test_dispatch_preserves_control_status_after_turn_budget_is_exhausted() -> None: registry = ToolRegistry() async def huge() -> str: return "x" * 1000 async def control_error() -> str: return json.dumps( { "status": "error", "user_message": "Missing required path.", "retry_allowed": False, } ) registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) registry.register( ToolSpec( name="control_error", description="control", parameters={}, result_budget_class="control", ), control_error, ) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=1000, max_tool_result_chars_per_turn=5, ) ), ) await handler(ToolCall(tool_use_id="tc-huge", tool_name="huge", arguments={})) result = await handler( ToolCall(tool_use_id="tc-control", tool_name="control_error", arguments={}) ) payload = json.loads(result.content) assert payload["status"] == "error" assert payload["user_message"] == "Missing required path." assert payload["retry_allowed"] is False @pytest.mark.asyncio async def test_dispatch_clamps_web_fetch_max_chars_before_handler() -> None: registry = ToolRegistry() seen: dict[str, object] = {} async def web_fetch(url: str, max_chars: int | None = None) -> str: seen["url"] = url seen["max_chars"] = max_chars return "ok" registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_policy=ToolRunBudgetPolicy(max_single_fetch_chars=12_000) ), ) result = await handler( ToolCall( tool_use_id="tc-fetch", tool_name="web_fetch", arguments={"url": "https://example.com", "max_chars": 1_000_000}, ) ) assert result.content == "ok" assert seen == {"url": "https://example.com", "max_chars": 12_000} @pytest.mark.asyncio async def test_dispatch_clamps_web_search_results_before_handler() -> None: registry = ToolRegistry() seen: dict[str, object] = {} async def web_search( query: str, max_results: int | None = None, fetch_top_k: int | None = None, max_chars_per_source: int | None = None, ) -> str: seen["query"] = query seen["max_results"] = max_results seen["fetch_top_k"] = fetch_top_k seen["max_chars_per_source"] = max_chars_per_source return json.dumps({"results": []}) registry.register( ToolSpec( name="web_search", description="search", parameters={"query": {"type": "string"}, "max_results": {"type": "integer"}}, result_budget_class="external", ), web_search, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_policy=ToolRunBudgetPolicy(max_web_search_results=10) ), ) result = await handler( ToolCall( tool_use_id="tc-search", tool_name="web_search", arguments={"query": "test", "max_results": 1000}, ) ) assert json.loads(result.content) == {"results": []} assert seen == { "query": "test", "max_results": 10, "fetch_top_k": 3, "max_chars_per_source": 1500, } @pytest.mark.asyncio async def test_dispatch_clamps_canonical_web_search_fetch_arguments_before_handler() -> None: registry = ToolRegistry() seen: dict[str, object] = {} async def web_search( query: str, max_results: int | None = None, fetch_top_k: int | None = None, max_chars_per_source: int | None = None, ) -> str: seen["query"] = query seen["max_results"] = max_results seen["fetch_top_k"] = fetch_top_k seen["max_chars_per_source"] = max_chars_per_source return json.dumps({"results": []}) registry.register( ToolSpec( name="web_search", description="source-backed search", parameters={ "query": {"type": "string"}, "max_results": {"type": "integer"}, "fetch_top_k": {"type": "integer"}, "max_chars_per_source": {"type": "integer"}, }, result_budget_class="external", ), web_search, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_policy=ToolRunBudgetPolicy( max_web_search_results=8, max_web_search_fetch_top_k=2, max_web_search_chars_per_source=900, ) ), ) result = await handler( ToolCall( tool_use_id="tc-web-search-canonical-clamp", tool_name="web_search", arguments={ "query": "test", "max_results": 1000, "fetch_top_k": 1000, "max_chars_per_source": 1000, }, ) ) assert json.loads(result.content) == {"results": []} assert seen == { "query": "test", "max_results": 8, "fetch_top_k": 2, "max_chars_per_source": 900, } @pytest.mark.asyncio async def test_dispatch_run_budget_blocks_exhausted_external_call_before_handler() -> None: registry = ToolRegistry() calls = 0 async def web_fetch(url: str, max_chars: int | None = None) -> str: nonlocal calls calls += 1 return f"{url}:{max_chars}" registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_key="dispatch-test-fetch-limit", tool_run_budget_policy=ToolRunBudgetPolicy( max_web_fetch_calls_per_turn=1, max_single_fetch_chars=500, max_external_text_chars_per_turn=1_000, ), ), ) first = await handler( ToolCall( tool_use_id="tc-fetch-1", tool_name="web_fetch", arguments={"url": "https://example.com", "max_chars": 10_000}, ) ) second = await handler( ToolCall( tool_use_id="tc-fetch-2", tool_name="web_fetch", arguments={"url": "https://example.com/again", "max_chars": 10_000}, ) ) assert first.content == "https://example.com:500" assert calls == 1 assert second.is_error is False assert second.execution_status is not None assert second.execution_status["status"] == "unknown" assert second.execution_status["reason"] == "tool_run_budget_exhausted" payload = json.loads(second.content) assert payload["status"] == "control" assert payload["reason"] == "tool_run_budget_exhausted" assert payload["retry_allowed"] is False assert "larger budget" not in payload["user_message"] assert "runtime resource guard" in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_run_budget_blocks_repeated_web_search_before_handler() -> None: registry = ToolRegistry() calls = 0 async def web_search( query: str, max_results: int | None = None, fetch_top_k: int | None = None, max_chars_per_source: int | None = None, ) -> str: nonlocal calls del max_results, fetch_top_k, max_chars_per_source calls += 1 return json.dumps({"query": query, "results": []}) registry.register( ToolSpec( name="web_search", description="source-backed search", parameters={"query": {"type": "string"}}, result_budget_class="external", ), web_search, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_key="dispatch-test-web-search-repeat", tool_run_budget_policy=ToolRunBudgetPolicy( max_repeated_retrievals_per_turn=1, ), ), ) first = await handler( ToolCall( tool_use_id="tc-web-search-repeat-1", tool_name="web_search", arguments={"query": "Python Release"}, ) ) second = await handler( ToolCall( tool_use_id="tc-web-search-repeat-2", tool_name="web_search", arguments={"query": " python release "}, ) ) assert json.loads(first.content) == {"query": "Python Release", "results": []} assert calls == 1 assert second.is_error is False assert second.execution_status is not None assert second.execution_status["status"] == "unknown" assert second.execution_status["reason"] == "tool_run_budget_exhausted" payload = json.loads(second.content) assert payload["status"] == "control" assert payload["reason"] == "tool_run_budget_exhausted" assert payload["retry_allowed"] is False @pytest.mark.asyncio async def test_dispatch_run_budget_abort_releases_failed_external_reservation() -> None: registry = ToolRegistry() fail_next = True seen_max_chars: list[int | None] = [] async def web_fetch(url: str, max_chars: int | None = None) -> str: nonlocal fail_next seen_max_chars.append(max_chars) if fail_next: fail_next = False raise RuntimeError("temporary failure") return "ok" registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_key="dispatch-test-fetch-abort", tool_run_budget_policy=ToolRunBudgetPolicy( max_web_fetch_calls_per_turn=1, max_single_fetch_chars=400, max_external_text_chars_per_turn=400, ), ), ) failed = await handler( ToolCall( tool_use_id="tc-fetch-fail", tool_name="web_fetch", arguments={"url": "https://example.com", "max_chars": 10_000}, ) ) retried = await handler( ToolCall( tool_use_id="tc-fetch-retry", tool_name="web_fetch", arguments={"url": "https://example.com", "max_chars": 10_000}, ) ) assert failed.is_error is True assert retried.content == "ok" assert seen_max_chars == [400, 400] @pytest.mark.asyncio async def test_dispatch_run_budget_exception_after_reservation_is_control() -> None: registry = ToolRegistry() calls = 0 async def web_fetch(url: str, max_chars: int | None = None) -> str: nonlocal calls del url, max_chars calls += 1 raise ToolRunBudgetExceededError("web_fetch", "internal run budget") registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_key="dispatch-test-fetch-internal-budget", tool_run_budget_policy=ToolRunBudgetPolicy( max_web_fetch_calls_per_turn=1, max_single_fetch_chars=400, max_external_text_chars_per_turn=400, ), ), ) first = await handler( ToolCall( tool_use_id="tc-fetch-internal-budget-1", tool_name="web_fetch", arguments={"url": "https://example.com"}, ) ) second = await handler( ToolCall( tool_use_id="tc-fetch-internal-budget-2", tool_name="web_fetch", arguments={"url": "https://example.com"}, ) ) assert calls == 2 for result in (first, second): assert result.is_error is False assert result.execution_status is not None assert result.execution_status["status"] == "unknown" assert result.execution_status["reason"] == "tool_run_budget_exhausted" assert json.loads(result.content)["status"] == "control" @pytest.mark.asyncio async def test_dispatch_run_budget_limits_concurrent_external_calls_atomically() -> None: registry = ToolRegistry() started = 0 async def web_fetch(url: str, max_chars: int | None = None) -> str: nonlocal started started += 1 await asyncio.sleep(0) return url registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_key="dispatch-test-fetch-concurrent", tool_run_budget_policy=ToolRunBudgetPolicy( max_web_fetch_calls_per_turn=1, max_single_fetch_chars=400, max_external_text_chars_per_turn=1_000, ), ), ) results = await asyncio.gather( handler( ToolCall( tool_use_id="tc-fetch-a", tool_name="web_fetch", arguments={"url": "https://example.com/a"}, ) ), handler( ToolCall( tool_use_id="tc-fetch-b", tool_name="web_fetch", arguments={"url": "https://example.com/b"}, ) ), ) assert started == 1 assert sum(result.is_error for result in results) == 0 control_payloads = [] for result in results: try: payload = json.loads(result.content) except ValueError: continue if payload.get("status") == "control": control_payloads.append(payload) assert len(control_payloads) == 1 assert any( result.execution_status and result.execution_status["reason"] == "tool_run_budget_exhausted" for result in results ) @pytest.mark.asyncio async def test_dispatch_run_budget_allows_oversized_result_then_controls_retry() -> None: registry = ToolRegistry() async def web_fetch(url: str, max_chars: int | None = None) -> str: return "x" * 250 registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_key="dispatch-test-fetch-result-budget", tool_run_budget_policy=ToolRunBudgetPolicy( max_web_fetch_calls_per_turn=2, max_single_fetch_chars=200, max_external_text_chars_per_turn=200, ), ), ) result = await handler( ToolCall( tool_use_id="tc-fetch-large", tool_name="web_fetch", arguments={"url": "https://example.com", "max_chars": 200}, ) ) retry = await handler( ToolCall( tool_use_id="tc-fetch-after-large", tool_name="web_fetch", arguments={"url": "https://example.com", "max_chars": 200}, ) ) assert result.is_error is False assert result.execution_status is None assert result.content == "x" * 250 assert retry.is_error is False assert retry.execution_status is not None assert retry.execution_status["status"] == "unknown" assert retry.execution_status["reason"] == "tool_run_budget_exhausted" payload = json.loads(retry.content) assert payload["status"] == "control" assert payload["reason"] == "tool_run_budget_exhausted" assert "larger budget" not in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_logs_web_retrieval_run_diagnostics_without_default_call_caps( monkeypatch, ) -> None: registry = ToolRegistry() async def web_search( query: str, max_results: int | None = None, fetch_top_k: int | None = None, max_chars_per_source: int | None = None, ) -> str: del max_results, fetch_top_k, max_chars_per_source return json.dumps({"query": query, "results": ["one", "two"]}) registry.register( ToolSpec( name="web_search", description="search", parameters={"query": {"type": "string"}, "max_results": {"type": "integer"}}, result_budget_class="external", ), web_search, ) handler = build_tool_handler( registry, ToolContext(tool_run_budget_key="dispatch-test-search-diagnostics"), ) log_events: list[tuple[str, dict[str, object]]] = [] monkeypatch.setattr( dispatch_module.log, "debug", lambda event, **payload: log_events.append((event, payload)), ) result = await handler( ToolCall( tool_use_id="tc-search-diagnostics", tool_name="web_search", arguments={"query": "test", "max_results": 3}, ) ) assert result.is_error is False assert json.loads(result.content)["results"] == ["one", "two"] event = next( payload for name, payload in log_events if name == "dispatch.web_retrieval_tool_run_diagnostics" ) assert event["tool"] == "web_search" assert event["tool_use_id"] == "tc-search-diagnostics" assert event["web_search_calls_used"] == 1 assert event["web_fetch_calls_used"] == 0 external_text_chars_used = event["external_text_chars_used"] tool_wall_time_ms = event["tool_wall_time_ms"] assert isinstance(external_text_chars_used, int) assert isinstance(tool_wall_time_ms, int | float) assert external_text_chars_used >= len(result.content) assert tool_wall_time_ms >= 0 @pytest.mark.asyncio async def test_dispatch_run_budget_charges_web_search_external_text() -> None: registry = ToolRegistry() async def web_search( query: str, max_results: int | None = None, fetch_top_k: int | None = None, max_chars_per_source: int | None = None, ) -> str: del max_results, fetch_top_k, max_chars_per_source return json.dumps({"query": query, "results": ["x" * 250]}) registry.register( ToolSpec( name="web_search", description="search", parameters={"query": {"type": "string"}, "max_results": {"type": "integer"}}, result_budget_class="external", ), web_search, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_key="dispatch-test-search-result-budget", tool_run_budget_policy=ToolRunBudgetPolicy( max_web_search_calls_per_turn=2, max_external_text_chars_per_turn=200, ), ), ) result = await handler( ToolCall( tool_use_id="tc-search-large", tool_name="web_search", arguments={"query": "test", "max_results": 10}, ) ) retry = await handler( ToolCall( tool_use_id="tc-search-after-large", tool_name="web_search", arguments={"query": "test", "max_results": 10}, ) ) assert result.is_error is False assert result.execution_status is None assert json.loads(result.content)["results"] == ["x" * 250] assert retry.is_error is False assert retry.execution_status is not None assert retry.execution_status["status"] == "unknown" assert retry.execution_status["reason"] == "tool_run_budget_exhausted" payload = json.loads(retry.content) assert payload["status"] == "control" assert payload["reason"] == "tool_run_budget_exhausted" assert "larger budget" not in payload["user_message"] @pytest.mark.asyncio async def test_dispatch_run_budget_is_fresh_for_separate_current_contexts() -> None: registry = ToolRegistry() async def web_fetch(url: str, max_chars: int | None = None) -> str: return f"{url}:{max_chars}" registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler(registry) async def call_with_new_turn(turn_id: str) -> str: token = current_tool_context.set( ToolContext( session_key=f"agent:main:{turn_id}", tool_run_budget_key=f"agent:main:{turn_id}:turn-budget", tool_run_budget_policy=ToolRunBudgetPolicy( max_web_fetch_calls_per_turn=1, max_single_fetch_chars=300, max_external_text_chars_per_turn=500, ), ) ) try: first = await handler( ToolCall( tool_use_id=f"tc-{turn_id}-1", tool_name="web_fetch", arguments={"url": f"https://example.com/{turn_id}/first"}, ) ) second = await handler( ToolCall( tool_use_id=f"tc-{turn_id}-2", tool_name="web_fetch", arguments={"url": f"https://example.com/{turn_id}/second"}, ) ) finally: current_tool_context.reset(token) assert first.is_error is False assert second.is_error is False assert second.execution_status is not None assert second.execution_status["status"] == "unknown" assert second.execution_status["reason"] == "tool_run_budget_exhausted" return first.content first_turn = await call_with_new_turn("turn-a") second_turn = await call_with_new_turn("turn-b") assert first_turn == "https://example.com/turn-a/first:300" assert second_turn == "https://example.com/turn-b/first:300" @pytest.mark.asyncio async def test_dispatch_run_budget_without_key_does_not_leak_across_handler_reuse() -> None: registry = ToolRegistry() async def web_fetch(url: str, max_chars: int | None = None) -> str: return f"{url}:{max_chars}" registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler( registry, ToolContext( tool_run_budget_policy=ToolRunBudgetPolicy( max_web_fetch_calls_per_turn=1, max_single_fetch_chars=300, max_external_text_chars_per_turn=500, ) ), ) first = await handler( ToolCall( tool_use_id="tc-reuse-a", tool_name="web_fetch", arguments={"url": "https://example.com/a"}, ) ) second = await handler( ToolCall( tool_use_id="tc-reuse-b", tool_name="web_fetch", arguments={"url": "https://example.com/b"}, ) ) assert first.content == "https://example.com/a:300" assert second.content == "https://example.com/b:300" @pytest.mark.asyncio async def test_dispatch_run_budget_applies_to_subagent_current_context() -> None: registry = ToolRegistry() calls = 0 async def web_fetch(url: str, max_chars: int | None = None) -> str: nonlocal calls calls += 1 return f"{url}:{max_chars}" registry.register( ToolSpec( name="web_fetch", description="fetch", parameters={"url": {"type": "string"}, "max_chars": {"type": "integer"}}, result_budget_class="external", ), web_fetch, ) handler = build_tool_handler(registry) token = current_tool_context.set( ToolContext( session_key="subagent:agent:main:webchat:demo", caller_kind=CallerKind.SUBAGENT, tool_run_budget_key="subagent:agent:main:webchat:demo:worker:1", tool_run_budget_policy=ToolRunBudgetPolicy( max_web_fetch_calls_per_turn=1, max_single_fetch_chars=300, max_external_text_chars_per_turn=500, ), ) ) try: first = await handler( ToolCall( tool_use_id="tc-subagent-a", tool_name="web_fetch", arguments={"url": "https://example.com/a"}, ) ) second = await handler( ToolCall( tool_use_id="tc-subagent-b", tool_name="web_fetch", arguments={"url": "https://example.com/b"}, ) ) finally: current_tool_context.reset(token) assert first.content == "https://example.com/a:300" assert calls == 1 assert second.is_error is False assert second.execution_status is not None assert second.execution_status["status"] == "unknown" assert second.execution_status["reason"] == "tool_run_budget_exhausted" @pytest.mark.asyncio async def test_dispatch_preserves_sessions_yield_control_json_when_bounding() -> None: registry = ToolRegistry() async def sessions_yield() -> str: return json.dumps( { "status": "yielded", "waited": False, "message": "Current turn yielded; wait for pushed session events.", "yield_message": "y" * 1000, } ) registry.register( ToolSpec( name="sessions_yield", description="yield", parameters={}, result_budget_class="control", ), sessions_yield, ) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy(max_single_tool_result_chars=160) ), ) result = await handler( ToolCall(tool_use_id="tc-yield", tool_name="sessions_yield", arguments={}) ) payload = json.loads(result.content) assert payload["status"] == "yielded" assert payload["waited"] is False assert payload["result_truncated"] is True assert "tool_result_budget_applied" not in payload assert "result_returned_chars" not in payload assert "budget_class" not in payload assert len(result.content) < 500 @pytest.mark.asyncio async def test_dispatch_tracker_budget_is_fresh_for_reused_tool_context() -> None: registry = ToolRegistry() async def huge() -> str: return "x" * 1000 registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) ctx = ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=120, max_tool_result_chars_per_turn=140, ) ) first_handler = build_tool_handler(registry, ctx) first = await first_handler( ToolCall(tool_use_id="tc-first", tool_name="huge", arguments={}) ) second_handler = build_tool_handler(registry, ctx) second = await second_handler( ToolCall(tool_use_id="tc-second", tool_name="huge", arguments={}) ) assert _strict_preview_chars(first.content) == _strict_preview_chars(second.content) assert _strict_preview_chars(first.content) > 0 @pytest.mark.asyncio async def test_dispatch_uses_current_tool_context_budget_for_handler_without_static_ctx() -> None: registry = ToolRegistry() async def huge() -> str: return "x" * 1000 registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) handler = build_tool_handler(registry) ctx = ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=120, max_tool_result_chars_per_turn=140, ) ) token = current_tool_context.set(ctx) try: result = await handler( ToolCall(tool_use_id="tc-current-ctx", tool_name="huge", arguments={}) ) finally: current_tool_context.reset(token) assert _strict_preview_chars(result.content) <= 120 @pytest.mark.asyncio async def test_dispatch_reused_handler_gets_fresh_budget_for_separate_current_contexts() -> None: registry = ToolRegistry() async def huge() -> str: return "x" * 1000 registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) handler = build_tool_handler(registry) async def call_with_context() -> int: ctx = ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=120, max_tool_result_chars_per_turn=140, ) ) token = current_tool_context.set(ctx) try: result = await handler( ToolCall(tool_use_id="tc-current-ctx", tool_name="huge", arguments={}) ) finally: current_tool_context.reset(token) return _strict_preview_chars(result.content) first = await call_with_context() second = await call_with_context() assert first == second assert first > 0 @pytest.mark.asyncio async def test_dispatch_tracker_limits_concurrent_tool_results_per_turn() -> None: registry = ToolRegistry() async def huge() -> str: await asyncio.sleep(0) return "x" * 1000 registry.register(ToolSpec(name="huge", description="huge", parameters={}), huge) handler = build_tool_handler( registry, ToolContext( tool_result_budget_policy=ToolResultBudgetPolicy( max_single_tool_result_chars=120, max_tool_result_chars_per_turn=180, ) ), ) results = await asyncio.gather( *[ handler(ToolCall(tool_use_id=f"tc-{idx}", tool_name="huge", arguments={})) for idx in range(3) ] ) returned_total = sum(_strict_preview_chars(result.content) for result in results) assert returned_total <= 180