hunjing/backend/tests/application/tool_runtime/test_runtime.py

"""Tool Runtime 编排闭环测试：司命评估 → 沙箱执行 → 一炁事件回写。

沙箱用 respx 模拟，事件存储用 InMemoryEventStore。
"""

from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator
from unittest.mock import AsyncMock, MagicMock

import httpx
import pytest
import respx

from app.application.tool_runtime.runtime import ToolRuntime
from app.core.config import SandboxSettings
from app.integrations.sandbox.client import SandboxClient
from app.modules.runtime.events import InMemoryEventStore
from app.modules.safety.exceptions import PermissionDenied, RequiresConfirmation
from app.modules.safety.governor import PermissionPolicy, SafetyGovernor, ToolPolicy


def _settings() -> SandboxSettings:
    return SandboxSettings(
        rust_sandbox_url="http://sandbox.test",
        sandbox_auth_enabled=True,
        sandbox_auth_secret="shared-secret",
        sandbox_default_timeout_ms=30000,
        sandbox_auth_ttl_seconds=30,
    )


def _ok_response() -> httpx.Response:
    return httpx.Response(
        200,
        json={
            "ok": True,
            "exit_code": 0,
            "stdout": "42\n",
            "stderr": "",
            "duration_ms": 3,
            "timeout": False,
            "truncated": False,
            "trace_id": "t1",
        },
    )


@pytest.fixture
async def sandbox_client() -> AsyncIterator[SandboxClient]:
    async with httpx.AsyncClient() as http:
        yield SandboxClient(http, _settings())


@respx.mock
async def test_allow_executes_and_records_started_completed(
    sandbox_client: SandboxClient,
) -> None:
    respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok_response())
    store = InMemoryEventStore()
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), store)

    result = await runtime.execute("datetime", {}, run_id="run_1", trace_id="t1")

    assert result.ok
    assert result.stdout == "42\n"
    assert [e.type for e in store.all()] == ["tool.started", "tool.completed"]
    assert all(e.run_id == "run_1" for e in store.all())


async def test_cancelled_sandbox_execute_requests_remote_cancel() -> None:
    sandbox = MagicMock()
    started = asyncio.Event()

    async def exec_forever(*_args, **_kwargs):
        started.set()
        await asyncio.Event().wait()

    sandbox.exec = AsyncMock(side_effect=exec_forever)
    sandbox.cancel = AsyncMock(return_value=True)
    runtime = ToolRuntime(sandbox, SafetyGovernor(), InMemoryEventStore())

    task = asyncio.create_task(runtime.execute("datetime", {}, trace_id="trace_cancel"))
    await started.wait()
    task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await task
    sandbox.cancel.assert_awaited_once()
    assert sandbox.cancel.await_args.args[0] == "trace_cancel"


@respx.mock
async def test_high_risk_requires_confirmation(sandbox_client: SandboxClient) -> None:
    respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok_response())
    store = InMemoryEventStore()
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), store)

    with pytest.raises(RequiresConfirmation):
        await runtime.execute("shell", {"command": "echo hi"}, run_id="run_1")

    # 未确认不应进入执行，故无 started 事件
    assert all(e.type != "tool.started" for e in store.all())


@respx.mock
async def test_confirmed_high_risk_executes(sandbox_client: SandboxClient) -> None:
    respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok_response())
    store = InMemoryEventStore()
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), store)

    result = await runtime.execute("shell", {"command": "echo hi"}, run_id="run_1", confirmed=True)

    assert result.ok
    assert [e.type for e in store.all()] == ["tool.started", "tool.completed"]


@respx.mock
async def test_forbidden_records_failed_event(sandbox_client: SandboxClient) -> None:
    respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok_response())
    governor = SafetyGovernor(
        {"shell": ToolPolicy(risk_level="L3", permission_policy=PermissionPolicy.forbidden)}
    )
    store = InMemoryEventStore()
    runtime = ToolRuntime(sandbox_client, governor, store)

    with pytest.raises(PermissionDenied):
        await runtime.execute("shell", {"command": "x"}, run_id="run_1", trace_id="t1")

    failed = [e for e in store.all() if e.type == "tool.failed"]
    assert len(failed) == 1
    assert failed[0].payload["error_code"] == "permission_denied"
    assert failed[0].trace_id == "t1"


async def test_unknown_tool_denied(sandbox_client: SandboxClient) -> None:
    store = InMemoryEventStore()
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), store)

    with pytest.raises(PermissionDenied):
        await runtime.execute("nope", {}, run_id="run_1")

    assert any(e.type == "tool.failed" for e in store.all())


@respx.mock
async def test_file_tool_isolates_workdir_by_character(
    sandbox_client: SandboxClient,
) -> None:
    """文件类工具的 workdir 应按角色隔离，同一角色跨 run 共享 workspace。"""
    route = respx.post("http://sandbox.test/v1/exec").mock(return_value=_ok_response())
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), InMemoryEventStore())

    await runtime.execute(
        "write_file",
        {"path": "a.txt", "content": "x"},
        character_id="char_01HABC",
        run_id="run_42",
        confirmed=True,
    )

    body = json.loads(route.calls.last.request.content)
    # workdir 按 character_id 隔离（而非 run_id），保证同一角色跨 run 文件互通
    assert body["workdir"] == "workspace/char_01HABC"


async def test_isolate_workdir_sanitizes_and_falls_back(
    sandbox_client: SandboxClient,
) -> None:
    """隔离键经 sanitize 杜绝路径逃逸；缺主体/run/trace 时逐级回退。"""
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), InMemoryEventStore())

    # 路径分隔符与 .. 被 sanitize 清除，无法逃逸沙箱根
    traversed = runtime._isolate_workdir(
        "workspace", subject_id="../../etc", run_id=None, trace_id=None
    )
    key = traversed.removeprefix("workspace/")
    assert key == "etc"
    assert "/" not in key and ".." not in key
    # 正常 subject_id（character_id）原样保留
    assert (
        runtime._isolate_workdir("workspace", subject_id="char_01HABC", run_id=None, trace_id=None)
        == "workspace/char_01HABC"
    )
    # 缺 subject_id 时回退 run_id
    assert (
        runtime._isolate_workdir("workspace", subject_id=None, run_id="run_01HABC", trace_id=None)
        == "workspace/run_01HABC"
    )
    # 缺 subject_id/run_id/trace 时回退 adhoc 前缀
    assert runtime._isolate_workdir(
        "workspace", subject_id=None, run_id=None, trace_id=None
    ).startswith("workspace/adhoc_")


@respx.mock
async def test_execute_stream_forwards_stdout_and_records_events(
    sandbox_client: SandboxClient,
) -> None:
    """流式执行：转发 stdout/done，started/completed 落一炁，stdout 增量不落库。"""
    sse = (
        b"event: stdout\n"
        b"data: aGVsbG8K\n"
        b"\n"
        b"event: done\n"
        b'data: {"exit_code":0,"duration_ms":3,"timeout":false,"cancelled":false,"truncated":false}\n'
        b"\n"
    )
    respx.post("http://sandbox.test/v1/exec/stream").mock(
        return_value=httpx.Response(200, content=sse)
    )
    store = InMemoryEventStore()
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), store)

    events = [ev async for ev in runtime.execute_stream("datetime", {}, run_id="run_1")]

    assert [e.kind for e in events] == ["stdout", "done"]
    assert events[-1].payload["exit_code"] == 0
    assert [e.type for e in store.all()] == ["tool.started", "tool.completed"]
    assert all(e.run_id == "run_1" for e in store.all())


@respx.mock
async def test_execute_stream_requires_confirmation_for_high_risk(
    sandbox_client: SandboxClient,
) -> None:
    """流式调用未确认的高风险工具：抛 RequiresConfirmation（端点据此发 requires_confirmation 事件）。"""
    respx.post("http://sandbox.test/v1/exec/stream").mock(return_value=_ok_response())
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), InMemoryEventStore())

    with pytest.raises(RequiresConfirmation):
        async for _ in runtime.execute_stream("shell", {"command": "echo hi"}):
            pass


# --------------------------------------------------------------------------- #
# executor 工具集成测试（fetch_url 走直接执行路径）
# --------------------------------------------------------------------------- #


@respx.mock
async def test_executor_tool_runs_directly_and_records_events(
    sandbox_client: SandboxClient,
) -> None:
    """executor 工具（fetch_url）走直接执行路径，跳过沙箱，started/completed 落库。"""
    respx.get("https://example.test").mock(
        return_value=httpx.Response(200, text="<html><body>Hello</body></html>")
    )
    async with httpx.AsyncClient() as tools_http:
        store = InMemoryEventStore()
        runtime = ToolRuntime(
            sandbox_client,
            SafetyGovernor(),
            store,
            http_client=tools_http,
        )
        # fetch_url permission=ask，confirmed=True 绕过确认。
        result = await runtime.execute(
            "fetch_url",
            {"url": "https://example.test"},
            run_id="run_fu",
            trace_id="t_fu",
            confirmed=True,
        )

    assert result.ok
    assert "Hello" in result.stdout
    assert result.trace_id == "t_fu"
    # 沙箱未被调用（直接执行路径）
    events = store.all()
    assert [e.type for e in events] == ["tool.started", "tool.completed"]
    assert all(e.run_id == "run_fu" for e in events)


@respx.mock
async def test_executor_tool_requires_confirmation_when_ask(
    sandbox_client: SandboxClient,
) -> None:
    """executor 工具的 permission_policy=ask 同样触发 RequiresConfirmation。"""
    async with httpx.AsyncClient() as tools_http:
        runtime = ToolRuntime(
            sandbox_client,
            SafetyGovernor(),
            InMemoryEventStore(),
            http_client=tools_http,
        )
        with pytest.raises(RequiresConfirmation):
            await runtime.execute(
                "fetch_url",
                {"url": "https://example.test"},
            )


async def test_executor_tool_without_http_client_returns_error(
    sandbox_client: SandboxClient,
) -> None:
    """未注入 http_client 时 executor 返回 ok=False（由 executor 自行兜底）。"""
    runtime = ToolRuntime(
        sandbox_client,
        SafetyGovernor(),
        InMemoryEventStore(),
        # 故意不传 http_client
    )
    result = await runtime.execute(
        "fetch_url",
        {"url": "https://example.test"},
        confirmed=True,
    )

    assert not result.ok
    assert "http_client" in result.stderr


# --------------------------------------------------------------------------- #
# 热刷新：builtin_specs() 的变更对已存在的 ToolRuntime 实例立即生效，无需重启
# --------------------------------------------------------------------------- #


def test_builtin_changes_reflected_in_list_specs_without_reinit(
    sandbox_client: SandboxClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """删除/修改 builtin 工具后，同一个 ToolRuntime 实例的 list_specs 立即反映变化。"""
    from app.application.tool_runtime import builtins as builtins_module

    original = builtins_module.builtin_specs()
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), InMemoryEventStore())

    # 基线：包含全部 builtin 工具
    baseline_names = {s.name for s in runtime.list_specs()}
    assert "datetime" in baseline_names
    assert "shell" in baseline_names

    # 模拟 builtin_specs 删除了 shell（例如配置关闭或被下线）
    pruned = {k: v for k, v in original.items() if k != "shell"}

    def _pruned_specs() -> dict:
        return dict(pruned)

    monkeypatch.setattr(builtins_module, "builtin_specs", _pruned_specs)

    after_names = {s.name for s in runtime.list_specs()}
    assert "datetime" in after_names
    assert "shell" not in after_names  # 删除立即生效，未重新构造 ToolRuntime


def test_register_overrides_builtin_and_survives_builtin_change(
    sandbox_client: SandboxClient,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """register 的工具覆盖 builtin 同名；builtin 变更不影响已注册的覆盖层。"""
    from app.application.tool_runtime import builtins as builtins_module
    from app.modules.tool_runtime.definition import ToolSpec

    custom_shell = ToolSpec(
        name="shell",
        description="自定义 shell（覆盖 builtin）",
        risk_level="L1",
        permission_policy=PermissionPolicy.allow,
        build_request=lambda _input: _ok_exec_request(),
        side_effect=False,
    )
    runtime = ToolRuntime(sandbox_client, SafetyGovernor(), InMemoryEventStore())
    runtime.register(custom_shell)

    registered = {s.name: s for s in runtime.list_specs()}
    assert registered["shell"].description == "自定义 shell（覆盖 builtin）"

    # builtin 变更（如整体被清空）不应影响 register 叠加的工具
    monkeypatch.setattr(builtins_module, "builtin_specs", lambda: {})

    after = {s.name: s for s in runtime.list_specs()}
    assert after["shell"].description == "自定义 shell（覆盖 builtin）"


def _ok_exec_request():
    """供 register 测试构造 ToolSpec.build_request 用的占位 ExecRequest。"""
    from app.integrations.sandbox.models import ExecRequest

    return ExecRequest(argv=["sh", "-c", "echo ok"])