opensquilla--opensquilla/tests/test_tools/test_dispatch_envelope.py

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