966 lines
34 KiB
Python
966 lines
34 KiB
Python
"""Agent 循环单元测试（流式版）：标准 ReAct/Agent loop 的核心场景 + 累加器单测。
|
||
|
||
mock 边界：
|
||
- ``app.application.agent.runtime.ProviderManagerFactory.create`` 的 ``.chat_stream()`` —— 不真实调 LLM，
|
||
按轮次返回预置的 chunk 序列（async generator）。
|
||
- ``tool_runtime.execute`` —— 不真实调沙箱（用 AsyncMock 返回 ToolResult 或抛异常）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
from typing import Any
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
import pytest
|
||
|
||
from app.ai_provider import TokenUsage
|
||
from app.ai_provider.llm import ToolCallAccumulator
|
||
from app.ai_provider.types.chat import ChatChunk
|
||
from app.ai_provider.types.messages import (
|
||
AIMessage,
|
||
BaseMessage,
|
||
HumanMessage,
|
||
SystemMessage,
|
||
ToolMessage,
|
||
)
|
||
from app.application.agent import (
|
||
AgentContext,
|
||
AgentEvent,
|
||
AgentRunContext,
|
||
AgentTurnResult,
|
||
ExtensionRegistry,
|
||
run_agent_turn,
|
||
)
|
||
from app.application.agent.confirmation import ConfirmationRegistry
|
||
from app.application.agent.tools import default_tools, to_ai_provider_tools
|
||
from app.application.tool_runtime.builtins import builtin_specs
|
||
from app.modules.safety.exceptions import RequiresConfirmation
|
||
from app.modules.tool_runtime.definition import ToolResult
|
||
|
||
|
||
@pytest.fixture
|
||
def base_messages() -> list[BaseMessage]:
|
||
return [
|
||
SystemMessage(content="你是测试助手。"),
|
||
HumanMessage(content="请用 echo 工具回显 'hello'。"),
|
||
]
|
||
|
||
|
||
@pytest.fixture
|
||
def tools() -> list[dict[str, Any]]:
|
||
return default_tools()
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 流式测试辅助：构造 chunk / chunk 序列 / 假 router / 假 tool_runtime
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
def _chunk(
|
||
*,
|
||
content: str = "",
|
||
reasoning: str = "",
|
||
tool_call_chunks: list[dict[str, Any]] | None = None,
|
||
usage: TokenUsage | None = None,
|
||
finish_reason: str | None = None,
|
||
) -> ChatChunk:
|
||
"""构造单个 ChatChunk。"""
|
||
return ChatChunk(
|
||
delta_content=content,
|
||
delta_reasoning=reasoning,
|
||
finish_reason=finish_reason,
|
||
usage=usage,
|
||
tool_call_chunks=tool_call_chunks or [],
|
||
)
|
||
|
||
|
||
async def _chunk_iter(chunks: list[ChatChunk]):
|
||
"""把一组 chunk 包装成 async generator（模拟 router.stream 的返回值）。"""
|
||
for c in chunks:
|
||
yield c
|
||
|
||
|
||
def _split_tool_call_stream(
|
||
*,
|
||
index: int = 0,
|
||
call_id: str = "call_echo",
|
||
name: str = "echo",
|
||
args: dict[str, Any],
|
||
) -> list[dict[str, Any]]:
|
||
"""把一次 tool_call 拆成分片 chunk，模拟真实流式 tool_call_chunks：
|
||
|
||
第 1 片带 id + name + 空 args；后续把 args JSON 拆两段追加。最后附 finish_reason 片。
|
||
"""
|
||
args_json = json.dumps(args)
|
||
mid = max(1, len(args_json) // 2)
|
||
return [
|
||
{"index": index, "id": call_id, "name": name, "args": ""},
|
||
{"index": index, "args": args_json[:mid]},
|
||
{"index": index, "args": args_json[mid:]},
|
||
]
|
||
|
||
|
||
def _make_turn(
|
||
*,
|
||
content: str = "",
|
||
reasoning: str = "",
|
||
tool_calls: list[dict[str, Any]] | None = None,
|
||
usage: TokenUsage | None = None,
|
||
) -> list[ChatChunk]:
|
||
"""构造一轮流式的 chunk 序列。
|
||
|
||
若提供 ``tool_calls``（``[{"id":..., "name":..., "args": {...}}]``），则把每个
|
||
tool_call 拆成分片并附在序列末尾；否则只产出 content/reasoning chunk。
|
||
"""
|
||
chunks: list[ChatChunk] = []
|
||
if reasoning:
|
||
chunks.append(_chunk(reasoning=reasoning))
|
||
if content:
|
||
chunks.append(_chunk(content=content))
|
||
if tool_calls:
|
||
for idx, tc in enumerate(tool_calls):
|
||
chunks.append(
|
||
_chunk(
|
||
tool_call_chunks=_split_tool_call_stream(
|
||
index=idx,
|
||
call_id=tc.get("id", f"call_{tc['name']}"),
|
||
name=tc["name"],
|
||
args=tc.get("args", {}),
|
||
)
|
||
)
|
||
)
|
||
# usage 放在最后一个 chunk（与 OpenAI streaming 的 final chunk 习惯一致）
|
||
if chunks:
|
||
last = chunks[-1]
|
||
chunks[-1] = ChatChunk(
|
||
delta_content=last.delta_content,
|
||
delta_reasoning=last.delta_reasoning,
|
||
finish_reason="tool_calls" if tool_calls else "stop",
|
||
usage=usage or TokenUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||
tool_call_chunks=last.tool_call_chunks,
|
||
)
|
||
else:
|
||
chunks.append(
|
||
_chunk(
|
||
usage=usage or TokenUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||
finish_reason="stop",
|
||
)
|
||
)
|
||
return chunks
|
||
|
||
|
||
def _make_tool_result(stdout: str = "hello\n", ok: bool = True) -> ToolResult:
|
||
return ToolResult(
|
||
ok=ok,
|
||
exit_code=0 if ok else 1,
|
||
stdout=stdout,
|
||
stderr="",
|
||
duration_ms=3,
|
||
timeout=False,
|
||
truncated=False,
|
||
trace_id="t1",
|
||
)
|
||
|
||
|
||
def _patch_router_stream(
|
||
monkeypatch: pytest.MonkeyPatch, turns: list[list[ChatChunk]]
|
||
) -> MagicMock:
|
||
"""让 ``ProviderManagerFactory.create().chat_stream()`` 按顺序返回 turns（每轮 chunk 列表）。
|
||
|
||
``chat_stream`` 是 ``async def`` 返回 ``AsyncIterator[ChatChunk]``，故这里用 ``AsyncMock``
|
||
返回值设为对应 async generator。
|
||
"""
|
||
call_count = {"n": 0}
|
||
|
||
async def fake_chat_stream(_messages, **_kwargs): # noqa: ANN001
|
||
idx = call_count["n"]
|
||
call_count["n"] += 1
|
||
if idx >= len(turns):
|
||
raise AssertionError(f"chat_stream 被调用了 {idx + 1} 次，但只预期 {len(turns)} 次")
|
||
async for chunk in _chunk_iter(turns[idx]):
|
||
yield chunk
|
||
|
||
manager = MagicMock()
|
||
manager.chat_stream = fake_chat_stream
|
||
monkeypatch.setattr(
|
||
"app.application.agent.runtime.ProviderManagerFactory.create",
|
||
lambda: manager,
|
||
)
|
||
return manager
|
||
|
||
|
||
def _make_tool_runtime(execute_result: ToolResult | Exception) -> MagicMock:
|
||
"""构造一个 ToolRuntime 替身；execute 要么返回 ToolResult，要么抛给定异常。"""
|
||
runtime = MagicMock()
|
||
if isinstance(execute_result, Exception):
|
||
runtime.execute = AsyncMock(side_effect=execute_result)
|
||
else:
|
||
runtime.execute = AsyncMock(return_value=execute_result)
|
||
return runtime
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 1. 正常闭环：第一轮流式 tool_call（分片），第二轮纯 content
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_message_lifecycle_covers_assistant_updates_and_tool_results(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
_patch_router_stream(
|
||
monkeypatch,
|
||
[
|
||
_make_turn(
|
||
reasoning="先调用工具",
|
||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "hi"}}],
|
||
),
|
||
_make_turn(content="完成"),
|
||
],
|
||
)
|
||
registry = ExtensionRegistry()
|
||
events: list[AgentEvent] = []
|
||
|
||
async def listen(event: AgentEvent, _run: AgentRunContext) -> None:
|
||
events.append(event)
|
||
|
||
registry.add_event_listener(listen)
|
||
await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=_make_tool_runtime(_make_tool_result()),
|
||
extensions=registry,
|
||
)
|
||
|
||
event_types = [event.type for event in events]
|
||
assert event_types.count("message.started") == 3
|
||
assert event_types.count("message.ended") == 3
|
||
assert "message.updated" in event_types
|
||
tool_ended = next(
|
||
event
|
||
for event in events
|
||
if event.type == "message.ended" and event.payload["role"] == "tool"
|
||
)
|
||
assert tool_ended.payload["message"]["tool_call_id"] == "call_echo"
|
||
|
||
|
||
async def test_normal_loop_completes_after_tool_call(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
first = _make_turn(
|
||
reasoning="我需要用 echo 工具。",
|
||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "hello"}}],
|
||
)
|
||
second = _make_turn(content="已为你回显：hello", reasoning="已收到工具结果，准备总结。")
|
||
_patch_router_stream(monkeypatch, [first, second])
|
||
|
||
tool_runtime = _make_tool_runtime(_make_tool_result(stdout="hello\n"))
|
||
|
||
# 收集 on_chunk / on_tool_event 调用
|
||
chunks_seen: list[ChatChunk] = []
|
||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||
|
||
async def on_chunk(chunk: ChatChunk) -> None:
|
||
chunks_seen.append(chunk)
|
||
|
||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||
tool_events.append((event_type, payload))
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
model_hint="test-model",
|
||
on_chunk=on_chunk,
|
||
on_tool_event=on_tool_event,
|
||
)
|
||
|
||
assert isinstance(result, AgentTurnResult)
|
||
assert result.iterations == 2
|
||
assert result.truncated is False
|
||
assert result.content == "已为你回显：hello"
|
||
# reasoning 用最后一轮（最终回复轮）的累积值
|
||
assert result.reasoning == "已收到工具结果，准备总结。"
|
||
assert len(result.tool_history) == 1
|
||
assert result.tool_history[0]["name"] == "echo"
|
||
assert result.tool_history[0]["status"] == "success"
|
||
assert result.tool_history[0]["args"] == {"text": "hello"}
|
||
assert "hello" in result.tool_history[0]["result"]
|
||
|
||
# 流式 chunk 被透传多次（至少两轮的 chunk 数）
|
||
assert len(chunks_seen) >= 2
|
||
|
||
# 工具事件：started + completed 各 1 次
|
||
assert len(tool_events) == 2
|
||
assert tool_events[0][0] == "tool.started"
|
||
assert tool_events[0][1]["tool_name"] == "echo"
|
||
assert tool_events[0][1]["tool_call_id"] == "call_echo"
|
||
assert tool_events[1][0] == "tool.completed"
|
||
assert tool_events[1][1]["status"] == "success"
|
||
assert "hello" in tool_events[1][1]["result"]
|
||
|
||
# messages 正确追加了 AIMessage(tool_calls) + ToolMessage + 最终 AIMessage
|
||
# base_messages 原长度为 2（SystemMessage + HumanMessage）
|
||
assert len(base_messages) == 5
|
||
ai_msg = base_messages[2]
|
||
assert isinstance(ai_msg, AIMessage)
|
||
assert ai_msg.tool_calls[0]["name"] == "echo"
|
||
tool_msg = base_messages[3]
|
||
assert isinstance(tool_msg, ToolMessage)
|
||
assert tool_msg.tool_call_id == "call_echo"
|
||
assert "hello" in tool_msg.content
|
||
final_msg = base_messages[4]
|
||
assert isinstance(final_msg, AIMessage)
|
||
assert final_msg.content == "已为你回显：hello"
|
||
|
||
# usage 累计两轮
|
||
assert result.usage.total_tokens == 30
|
||
|
||
# tool_runtime.execute 被调用了一次，参数正确
|
||
tool_runtime.execute.assert_awaited_once()
|
||
call_kwargs = tool_runtime.execute.call_args
|
||
assert call_kwargs.args[0] == "echo"
|
||
assert call_kwargs.args[1] == {"text": "hello"}
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 2. 工具失败不中断：execute 抛异常，stream 第二轮仍能正常返回 content
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_tool_failure_does_not_break_loop(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
first = _make_turn(
|
||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "hi"}}],
|
||
)
|
||
second = _make_turn(content="工具坏了，但我还是回复你。")
|
||
_patch_router_stream(monkeypatch, [first, second])
|
||
|
||
tool_runtime = _make_tool_runtime(RuntimeError("sandbox down"))
|
||
|
||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||
|
||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||
tool_events.append((event_type, payload))
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
on_tool_event=on_tool_event,
|
||
)
|
||
|
||
assert result.iterations == 2
|
||
assert result.truncated is False
|
||
assert result.content == "工具坏了，但我还是回复你。"
|
||
assert len(result.tool_history) == 1
|
||
assert result.tool_history[0]["status"] == "error"
|
||
assert "sandbox down" in result.tool_history[0]["result"]
|
||
|
||
# 错误信息作为 ToolMessage 回填
|
||
tool_msg = base_messages[3]
|
||
assert isinstance(tool_msg, ToolMessage)
|
||
assert "sandbox down" in tool_msg.content
|
||
|
||
# 工具事件 completed 带 status=error
|
||
assert [e[0] for e in tool_events] == ["tool.started", "tool.completed"]
|
||
assert tool_events[1][1]["status"] == "error"
|
||
assert "sandbox down" in tool_events[1][1]["result"]
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 3. 上限截断：stream 每次都返回 tool_call，max_iterations=3
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_loop_truncates_at_max_iterations(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
# 每次都返回 tool_call，让循环无法终止
|
||
looping_turn = _make_turn(
|
||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "loop"}}],
|
||
)
|
||
_patch_router_stream(monkeypatch, [looping_turn, looping_turn, looping_turn])
|
||
|
||
tool_runtime = _make_tool_runtime(_make_tool_result(stdout="loop\n"))
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
max_iterations=3,
|
||
)
|
||
|
||
assert result.iterations == 3
|
||
assert result.truncated is True
|
||
assert len(result.tool_history) == 3
|
||
assert all(h["status"] == "success" for h in result.tool_history)
|
||
# messages 累积：base 2 + 3 轮 × (AIMessage + ToolMessage) = 8
|
||
assert len(base_messages) == 8
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 4. 无工具直接返回：stream 第一次就是纯 content
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_no_tool_call_returns_immediately(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
direct = _make_turn(content="直接回答，不需要工具。")
|
||
_patch_router_stream(monkeypatch, [direct])
|
||
|
||
tool_runtime = _make_tool_runtime(_make_tool_result())
|
||
|
||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||
|
||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||
tool_events.append((event_type, payload))
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
on_tool_event=on_tool_event,
|
||
)
|
||
|
||
assert result.iterations == 1
|
||
assert result.truncated is False
|
||
assert result.content == "直接回答，不需要工具。"
|
||
assert result.tool_history == []
|
||
# messages 只追加了一条 AIMessage
|
||
assert len(base_messages) == 3
|
||
assert isinstance(base_messages[2], AIMessage)
|
||
# execute 从未被调用
|
||
tool_runtime.execute.assert_not_awaited()
|
||
# 无工具事件
|
||
assert tool_events == []
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 5. 工具转换辅助：to_ai_provider_tools / default_tools 基本正确
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
def test_to_ai_provider_tools_produces_openai_function_schema() -> None:
|
||
specs = list(builtin_specs().values())
|
||
tools = to_ai_provider_tools(specs)
|
||
assert len(tools) == len(specs)
|
||
for tool, spec in zip(tools, specs, strict=True):
|
||
assert tool["type"] == "function"
|
||
fn = tool["function"]
|
||
assert fn["name"] == spec.name
|
||
assert fn["description"] == spec.description
|
||
assert fn["parameters"]["type"] == "object"
|
||
|
||
|
||
def test_default_tools_returns_core_only_by_default() -> None:
|
||
"""default_tools 默认只返回 5 个核心沙箱工具（不含联网工具）。"""
|
||
tools = default_tools()
|
||
names = {t["function"]["name"] for t in tools}
|
||
assert names == {"datetime", "read_file", "list_dir", "write_file", "shell"}
|
||
assert len(tools) == 5
|
||
|
||
|
||
def test_default_tools_returns_core_only() -> None:
|
||
"""default_tools 返回 5 个核心沙箱工具。"""
|
||
tools = default_tools()
|
||
names = {t["function"]["name"] for t in tools}
|
||
assert names == {"datetime", "read_file", "list_dir", "write_file", "shell"}
|
||
assert len(tools) == 5
|
||
|
||
|
||
def test_run_agent_turn_rejects_non_positive_max_iterations(
|
||
base_messages: list[BaseMessage], tools: list[dict[str, Any]]
|
||
) -> None:
|
||
tool_runtime = _make_tool_runtime(_make_tool_result())
|
||
with pytest.raises(ValueError):
|
||
import asyncio
|
||
|
||
asyncio.get_event_loop().run_until_complete(
|
||
run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
max_iterations=0,
|
||
)
|
||
)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 6. ToolCallAccumulator 单元测试
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
def test_accumulator_empty_drain_returns_empty_list() -> None:
|
||
acc = ToolCallAccumulator()
|
||
drained = acc.drain()
|
||
assert drained == []
|
||
|
||
|
||
def test_accumulator_assembles_fragmented_args_across_chunks() -> None:
|
||
"""args 分片到达也能拼回完整 JSON。"""
|
||
acc = ToolCallAccumulator()
|
||
# 片段：先到 id+name，args 分 "hello"、""、" world" 三段
|
||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "id": "c1", "name": "echo", "args": ""}]))
|
||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": "hello"}]))
|
||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": ""}]))
|
||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": " world"}]))
|
||
|
||
drained = acc.drain()
|
||
assert len(drained) == 1
|
||
entry = drained[0]
|
||
assert entry["index"] == 0
|
||
assert entry["id"] == "c1"
|
||
assert entry["name"] == "echo"
|
||
assert entry["args"] == "hello world"
|
||
|
||
|
||
def test_accumulator_parallel_multiple_tool_calls() -> None:
|
||
"""index 0 和 1 的 tool_call 并行累积。"""
|
||
acc = ToolCallAccumulator()
|
||
acc.add_chunk(
|
||
_chunk(
|
||
tool_call_chunks=[
|
||
{"index": 0, "id": "c1", "name": "echo", "args": '{"text"'},
|
||
{"index": 1, "id": "c2", "name": "datetime", "args": "{}"},
|
||
]
|
||
)
|
||
)
|
||
acc.add_chunk(
|
||
_chunk(
|
||
tool_call_chunks=[
|
||
{"index": 0, "args": ': "a"}'},
|
||
]
|
||
)
|
||
)
|
||
|
||
drained = acc.drain()
|
||
assert len(drained) == 2
|
||
assert drained[0]["index"] == 0
|
||
assert drained[0]["name"] == "echo"
|
||
assert drained[0]["args"] == '{"text": "a"}'
|
||
assert drained[1]["index"] == 1
|
||
assert drained[1]["name"] == "datetime"
|
||
assert drained[1]["args"] == "{}"
|
||
|
||
|
||
def test_accumulator_args_arrive_before_name() -> None:
|
||
"""分片中先到 args 片段、后到 name，也能正确合并。"""
|
||
acc = ToolCallAccumulator()
|
||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": '{"x": '}]))
|
||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "id": "c1", "name": "calc", "args": "1"}]))
|
||
acc.add_chunk(_chunk(tool_call_chunks=[{"index": 0, "args": "}"}]))
|
||
|
||
drained = acc.drain()
|
||
assert len(drained) == 1
|
||
entry = drained[0]
|
||
assert entry["id"] == "c1"
|
||
assert entry["name"] == "calc"
|
||
assert entry["args"] == '{"x": 1}'
|
||
|
||
|
||
def test_accumulator_ignores_chunks_without_tool_call_chunks() -> None:
|
||
"""非 tool_call chunk（只有 content）不会影响累加器状态。"""
|
||
acc = ToolCallAccumulator()
|
||
acc.add_chunk(_chunk(content="hello"))
|
||
acc.add_chunk(_chunk(reasoning="thinking"))
|
||
acc.add_chunk(_chunk())
|
||
drained = acc.drain()
|
||
assert drained == []
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 7. 工具确认挂起/恢复（阶段3 步骤1）
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
def _make_confirmation_tool_runtime(
|
||
*, exc: Exception | None = None, result: ToolResult | None = None
|
||
) -> MagicMock:
|
||
"""构造 ToolRuntime：先抛 RequiresConfirmation（或给定 exc），再返回 result。
|
||
|
||
用于"批准后重试 confirmed=True 成功"的场景。result 缺省为成功 ToolResult。
|
||
"""
|
||
result = result or _make_tool_result(stdout="done\n")
|
||
if exc is not None:
|
||
runtime = MagicMock()
|
||
runtime.execute = AsyncMock(side_effect=[exc, result])
|
||
return runtime
|
||
raise ValueError("必须提供 exc")
|
||
|
||
|
||
def _make_deny_only_tool_runtime(exc: Exception) -> MagicMock:
|
||
"""构造 ToolRuntime：只抛一次 RequiresConfirmation（拒绝/超时不重试）。"""
|
||
runtime = MagicMock()
|
||
runtime.execute = AsyncMock(side_effect=exc)
|
||
return runtime
|
||
|
||
|
||
async def test_context_transformer_runs_for_every_model_request_without_persisting(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
first = _make_turn(
|
||
tool_calls=[{"id": "call_echo", "name": "echo", "args": {"text": "x"}}]
|
||
)
|
||
second = _make_turn(content="完成")
|
||
manager = _patch_router_stream(monkeypatch, [first, second])
|
||
runtime = MagicMock()
|
||
runtime.execute = AsyncMock(return_value=_make_tool_result())
|
||
registry = ExtensionRegistry()
|
||
transformed_calls: list[int] = []
|
||
|
||
async def transform(context: AgentContext, _run: AgentRunContext) -> AgentContext:
|
||
transformed_calls.append(len(context.messages))
|
||
return AgentContext(
|
||
(*context.messages, HumanMessage(content="临时注入")),
|
||
context.tools,
|
||
context.model_hint,
|
||
)
|
||
|
||
registry.add_context_transformer(transform)
|
||
await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=runtime,
|
||
extensions=registry,
|
||
)
|
||
|
||
assert transformed_calls == [2, 4]
|
||
assert all(message.content != "临时注入" for message in base_messages)
|
||
assert manager.chat_stream is not None
|
||
|
||
|
||
async def test_max_tokens_marks_agent_result_truncated(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
turn = _make_turn(content="未完成")
|
||
turn[-1].finish_reason = "length"
|
||
_patch_router_stream(monkeypatch, [turn])
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=MagicMock(),
|
||
)
|
||
|
||
assert result.content == "未完成"
|
||
assert result.truncated is True
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 7.1 确认 + 批准恢复：execute 第一次抛 RequiresConfirmation → 发事件 →
|
||
# resolve(cid, True) → confirmed=True 重试成功 → 第二轮纯 content
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_confirmation_approved_resumes_and_succeeds(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
first = _make_turn(
|
||
tool_calls=[{"id": "call_write", "name": "write_file", "args": {"path": "a.txt"}}],
|
||
)
|
||
second = _make_turn(content="已写入文件。")
|
||
_patch_router_stream(monkeypatch, [first, second])
|
||
|
||
tool_runtime = _make_confirmation_tool_runtime(
|
||
exc=RequiresConfirmation("策略 ask 需用户确认", risk_level="L2")
|
||
)
|
||
registry = ConfirmationRegistry()
|
||
|
||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||
|
||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||
tool_events.append((event_type, payload))
|
||
# 收到 requires_confirmation 事件后，手动触发批准恢复循环
|
||
if event_type == "tool.requires_confirmation":
|
||
cid = payload["confirmation_id"]
|
||
asyncio.create_task(registry.resolve(cid, True))
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
on_tool_event=on_tool_event,
|
||
confirmation_registry=registry,
|
||
run_id="test-run-1",
|
||
character_id="test-char-1",
|
||
)
|
||
|
||
assert result.iterations == 2
|
||
assert result.content == "已写入文件。"
|
||
assert len(result.tool_history) == 1
|
||
assert result.tool_history[0]["status"] == "success"
|
||
|
||
# 事件顺序：started（初次） → requires_confirmation → started（重试，带 confirmed=True）
|
||
# → completed(success)
|
||
types = [e[0] for e in tool_events]
|
||
assert types == [
|
||
"tool.started",
|
||
"tool.requires_confirmation",
|
||
"tool.started",
|
||
"tool.completed",
|
||
]
|
||
req_event = tool_events[1][1]
|
||
assert req_event["tool_name"] == "write_file"
|
||
assert req_event["risk_level"] == "L2"
|
||
assert req_event["confirmation_id"].startswith("confirm_")
|
||
# 重试时带 confirmed=True 标记
|
||
assert tool_events[2][1].get("confirmed") is True
|
||
assert tool_events[3][1]["status"] == "success"
|
||
|
||
# execute 被调用 2 次：第一次 confirmed=False（默认），第二次 confirmed=True
|
||
assert tool_runtime.execute.await_count == 2
|
||
first_call = tool_runtime.execute.call_args_list[0]
|
||
second_call = tool_runtime.execute.call_args_list[1]
|
||
assert first_call.kwargs["confirmed"] is False
|
||
assert second_call.kwargs["confirmed"] is True
|
||
|
||
# registry 已清理（resolve 后 Future 出列）
|
||
assert not registry.has(req_event["confirmation_id"])
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 7.2 确认 + 拒绝：resolve(cid, False) → 回填"用户拒绝" → 循环继续
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_confirmation_denied_backfills_and_continues(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
first = _make_turn(
|
||
tool_calls=[{"id": "call_shell", "name": "shell", "args": {"command": "rm -rf /"}}],
|
||
)
|
||
second = _make_turn(content="好的，我不执行了。")
|
||
_patch_router_stream(monkeypatch, [first, second])
|
||
|
||
tool_runtime = _make_deny_only_tool_runtime(
|
||
RequiresConfirmation("策略 always_ask 需用户确认", risk_level="L3")
|
||
)
|
||
registry = ConfirmationRegistry()
|
||
|
||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||
|
||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||
tool_events.append((event_type, payload))
|
||
if event_type == "tool.requires_confirmation":
|
||
cid = payload["confirmation_id"]
|
||
asyncio.create_task(registry.resolve(cid, False))
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
on_tool_event=on_tool_event,
|
||
confirmation_registry=registry,
|
||
run_id="test-run-2",
|
||
character_id="test-char-2",
|
||
confirmation_timeout=1.0,
|
||
)
|
||
|
||
assert result.iterations == 2
|
||
assert result.content == "好的，我不执行了。"
|
||
assert len(result.tool_history) == 1
|
||
assert result.tool_history[0]["status"] == "denied"
|
||
assert "用户拒绝" in result.tool_history[0]["result"]
|
||
|
||
# ToolMessage 回填的是"用户拒绝"
|
||
tool_msg = base_messages[3]
|
||
assert isinstance(tool_msg, ToolMessage)
|
||
assert "用户拒绝" in tool_msg.content
|
||
|
||
# 事件：started → requires_confirmation → completed(status=denied)
|
||
types = [e[0] for e in tool_events]
|
||
assert types == ["tool.started", "tool.requires_confirmation", "tool.completed"]
|
||
assert tool_events[-1][1]["status"] == "denied"
|
||
|
||
# execute 只被调用 1 次（拒绝不重试）
|
||
assert tool_runtime.execute.await_count == 1
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 7.3 确认超时按拒绝：confirmation_timeout 极短，不 resolve → 超时 → 回填 → 继续
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_confirmation_timeout_treated_as_denied(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
first = _make_turn(
|
||
tool_calls=[{"id": "call_write", "name": "write_file", "args": {"path": "x.txt"}}],
|
||
)
|
||
second = _make_turn(content="收到，跳过写入。")
|
||
_patch_router_stream(monkeypatch, [first, second])
|
||
|
||
tool_runtime = _make_deny_only_tool_runtime(RequiresConfirmation("需确认", risk_level="L2"))
|
||
registry = ConfirmationRegistry()
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
on_tool_event=None, # 不需要事件断言
|
||
confirmation_registry=registry,
|
||
confirmation_timeout=1.5, # > poll_interval（0.5s）确保至少完成一次轮询
|
||
)
|
||
|
||
assert result.iterations == 2
|
||
assert result.content == "收到，跳过写入。"
|
||
assert len(result.tool_history) == 1
|
||
assert result.tool_history[0]["status"] == "denied"
|
||
assert "超时" in result.tool_history[0]["result"]
|
||
# execute 只被调用 1 次（超时不重试）
|
||
assert tool_runtime.execute.await_count == 1
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 7.4 registry=None 降级：execute 抛 RequiresConfirmation → 回填"未启用" → 继续
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_confirmation_degraded_when_registry_none(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
base_messages: list[BaseMessage],
|
||
tools: list[dict[str, Any]],
|
||
) -> None:
|
||
first = _make_turn(
|
||
tool_calls=[{"id": "call_shell", "name": "shell", "args": {"command": "ls"}}],
|
||
)
|
||
second = _make_turn(content="好的，跳过了。")
|
||
_patch_router_stream(monkeypatch, [first, second])
|
||
|
||
tool_runtime = _make_deny_only_tool_runtime(RequiresConfirmation("需确认", risk_level="L3"))
|
||
|
||
tool_events: list[tuple[str, dict[str, Any]]] = []
|
||
|
||
async def on_tool_event(event_type: str, payload: dict[str, Any]) -> None:
|
||
tool_events.append((event_type, payload))
|
||
|
||
result = await run_agent_turn(
|
||
messages=base_messages,
|
||
tools=tools,
|
||
tool_runtime=tool_runtime,
|
||
on_tool_event=on_tool_event,
|
||
# 不传 confirmation_registry → 降级路径
|
||
)
|
||
|
||
assert result.iterations == 2
|
||
assert result.content == "好的，跳过了。"
|
||
assert len(result.tool_history) == 1
|
||
assert result.tool_history[0]["status"] == "skipped"
|
||
assert "需要用户确认" in result.tool_history[0]["result"]
|
||
assert "未启用" in result.tool_history[0]["result"]
|
||
|
||
# 事件：started → completed(status=skipped)（不发 requires_confirmation）
|
||
types = [e[0] for e in tool_events]
|
||
assert types == ["tool.started", "tool.completed"]
|
||
assert tool_events[-1][1]["status"] == "skipped"
|
||
|
||
# execute 只被调用 1 次
|
||
assert tool_runtime.execute.await_count == 1
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 7.5 ConfirmationRegistry 单元测试
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
async def test_confirmation_registry_resolve_returns_false_for_unknown() -> None:
|
||
registry = ConfirmationRegistry()
|
||
assert await registry.resolve("unknown_id", True) is False
|
||
assert registry.has("unknown_id") is False
|
||
|
||
|
||
async def test_confirmation_registry_resolves_only_requested_confirmation() -> None:
|
||
from app.db.session import AsyncSessionLocal
|
||
from app.models.safety import PendingConfirmationRecord
|
||
|
||
registry = ConfirmationRegistry()
|
||
async with AsyncSessionLocal() as session:
|
||
session.add_all(
|
||
[
|
||
PendingConfirmationRecord(
|
||
id="confirm_first",
|
||
run_id="run_shared",
|
||
tool="first",
|
||
input_json="{}",
|
||
risk_level="L2",
|
||
status="pending",
|
||
),
|
||
PendingConfirmationRecord(
|
||
id="confirm_second",
|
||
run_id="run_shared",
|
||
tool="second",
|
||
input_json="{}",
|
||
risk_level="L2",
|
||
status="pending",
|
||
),
|
||
]
|
||
)
|
||
await session.commit()
|
||
|
||
assert await registry.resolve(
|
||
"run_shared", True, confirmation_id="confirm_second"
|
||
) is True
|
||
|
||
async with AsyncSessionLocal() as session:
|
||
first = await session.get(PendingConfirmationRecord, "confirm_first")
|
||
second = await session.get(PendingConfirmationRecord, "confirm_second")
|
||
assert first is not None and first.status == "pending"
|
||
assert second is not None and second.status == "approved"
|
||
|
||
|
||
async def test_confirmation_registry_resolve_returns_false_after_done() -> None:
|
||
from app.db.session import AsyncSessionLocal
|
||
from app.models.safety import PendingConfirmationRecord
|
||
|
||
registry = ConfirmationRegistry()
|
||
assert registry.has("cid_1") is False
|
||
|
||
# 创建 PendingConfirmationRecord（wait/resolve 的 DB 前提）
|
||
async with AsyncSessionLocal() as session:
|
||
session.add(
|
||
PendingConfirmationRecord(
|
||
id="rec_cid_1",
|
||
run_id="cid_1",
|
||
tool="test",
|
||
input_json="{}",
|
||
risk_level="L2",
|
||
status="pending",
|
||
)
|
||
)
|
||
await session.commit()
|
||
|
||
running = asyncio.create_task(registry.wait("cid_1", poll_interval=0.1, timeout=0.5))
|
||
await asyncio.sleep(0.1)
|
||
assert registry.has("cid_1") is True
|
||
resolved = await registry.resolve("cid_1", True)
|
||
assert resolved is True
|
||
result = await running
|
||
assert result is True
|
||
# 已处理，再次 resolve 返回 False
|
||
assert await registry.resolve("cid_1", True) is False
|
||
assert registry.has("cid_1") is False
|