292 lines
10 KiB
Python
292 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from langchain_core.messages import SystemMessage, ToolMessage
|
|
from langgraph.types import Command
|
|
|
|
import src.agents.common.middlewares.runtime_config_middleware as runtime_middleware
|
|
from src.agents.common.middlewares.runtime_config_middleware import RuntimeConfigMiddleware
|
|
|
|
|
|
@dataclass
|
|
class _FakeTool:
|
|
name: str
|
|
|
|
|
|
@dataclass
|
|
class _FakeRequest:
|
|
runtime: Any
|
|
tools: list[Any]
|
|
system_message: SystemMessage
|
|
state: dict[str, Any]
|
|
|
|
def override(self, **kwargs):
|
|
return _FakeRequest(
|
|
runtime=kwargs.get("runtime", self.runtime),
|
|
tools=kwargs.get("tools", self.tools),
|
|
system_message=kwargs.get("system_message", self.system_message),
|
|
state=kwargs.get("state", self.state),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _FakeToolCallRequest:
|
|
tool_call: dict[str, Any]
|
|
runtime: Any
|
|
state: dict[str, Any]
|
|
|
|
|
|
async def _echo_handler(request):
|
|
return request
|
|
|
|
|
|
def _build_request(
|
|
*,
|
|
skills: list[str],
|
|
tools: list[str],
|
|
system_prompt: str = "你是助手",
|
|
state: dict[str, Any] | None = None,
|
|
skill_snapshot: dict[str, Any] | None = None,
|
|
) -> _FakeRequest:
|
|
context = SimpleNamespace(
|
|
system_prompt=system_prompt,
|
|
skills=skills,
|
|
tools=[],
|
|
knowledges=[],
|
|
mcps=[],
|
|
)
|
|
if skill_snapshot is not None:
|
|
context.skill_session_snapshot = skill_snapshot
|
|
runtime = SimpleNamespace(context=context)
|
|
return _FakeRequest(
|
|
runtime=runtime,
|
|
tools=[_FakeTool(name=name) for name in tools],
|
|
system_message=SystemMessage(content=[{"type": "text", "text": "base"}]),
|
|
state=state or {},
|
|
)
|
|
|
|
|
|
def _build_tool_request(*, skills: list[str], visible_skills: list[str], file_path: str) -> _FakeToolCallRequest:
|
|
return _FakeToolCallRequest(
|
|
tool_call={"name": "read_file", "args": {"file_path": file_path}},
|
|
runtime=SimpleNamespace(
|
|
context=SimpleNamespace(
|
|
skills=skills,
|
|
skill_session_snapshot=_build_snapshot(selected=skills, visible=visible_skills),
|
|
)
|
|
),
|
|
state={},
|
|
)
|
|
|
|
|
|
def _extract_appended_prompt(request: _FakeRequest) -> str:
|
|
return request.system_message.content_blocks[-1]["text"]
|
|
|
|
|
|
def _build_middleware() -> RuntimeConfigMiddleware:
|
|
return RuntimeConfigMiddleware(
|
|
enable_model_override=False,
|
|
enable_tools_override=False,
|
|
enable_system_prompt_override=True,
|
|
enable_skills_prompt_override=True,
|
|
)
|
|
|
|
|
|
def _build_snapshot(
|
|
selected: list[str],
|
|
visible: list[str] | None = None,
|
|
metadata: dict[str, dict[str, str]] | None = None,
|
|
dependency_map: dict[str, dict[str, list[str]]] | None = None,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"selected_skills": selected,
|
|
"visible_skills": visible if visible is not None else selected,
|
|
"prompt_metadata": metadata or {},
|
|
"dependency_map": dependency_map or {},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_abefore_agent_resolves_visible_skills_and_preinjects_prompt(monkeypatch: pytest.MonkeyPatch):
|
|
async def fake_resolve(selected):
|
|
assert selected == ["deliver-prd"]
|
|
return _build_snapshot(
|
|
selected=["deliver-prd"],
|
|
visible=["deliver-prd", "brainstorming"],
|
|
metadata={
|
|
"deliver-prd": {
|
|
"name": "deliver-prd",
|
|
"description": "deliver prd",
|
|
"path": "/skills/deliver-prd/SKILL.md",
|
|
},
|
|
"brainstorming": {
|
|
"name": "brainstorming",
|
|
"description": "brainstorming desc",
|
|
"path": "/skills/brainstorming/SKILL.md",
|
|
},
|
|
},
|
|
dependency_map={
|
|
"deliver-prd": {"tools": [], "mcps": [], "skills": ["brainstorming"]},
|
|
"brainstorming": {"tools": [], "mcps": [], "skills": []},
|
|
},
|
|
)
|
|
|
|
monkeypatch.setattr(runtime_middleware, "resolve_session_snapshot", fake_resolve)
|
|
middleware = _build_middleware()
|
|
request = _build_request(skills=["deliver-prd"], tools=["read_file"], system_prompt="你是助手")
|
|
|
|
await middleware.abefore_agent(request.state, request.runtime)
|
|
|
|
snapshot = request.runtime.context.skill_session_snapshot
|
|
assert snapshot["selected_skills"] == ["deliver-prd"]
|
|
assert snapshot["visible_skills"] == ["deliver-prd", "brainstorming"]
|
|
assert snapshot["dependency_map"]["deliver-prd"]["skills"] == ["brainstorming"]
|
|
assert request.runtime.context._skills_prompt_injected is True
|
|
assert "## Skills System" in request.runtime.context.system_prompt
|
|
assert "- **deliver-prd**: deliver prd" in request.runtime.context.system_prompt
|
|
assert "- **brainstorming**: brainstorming desc" in request.runtime.context.system_prompt
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_abefore_agent_injection_is_idempotent(monkeypatch: pytest.MonkeyPatch):
|
|
async def fake_resolve(_selected):
|
|
return _build_snapshot(
|
|
selected=["alpha"],
|
|
metadata={"alpha": {"name": "alpha", "description": "alpha desc", "path": "/skills/alpha/SKILL.md"}},
|
|
)
|
|
|
|
monkeypatch.setattr(runtime_middleware, "resolve_session_snapshot", fake_resolve)
|
|
middleware = _build_middleware()
|
|
request = _build_request(skills=["alpha"], tools=["read_file"], system_prompt="base prompt")
|
|
|
|
await middleware.abefore_agent(request.state, request.runtime)
|
|
await middleware.abefore_agent(request.state, request.runtime)
|
|
|
|
assert request.runtime.context.system_prompt.count("## Skills System") == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_awrap_model_call_keeps_preinjected_skills_prompt(monkeypatch: pytest.MonkeyPatch):
|
|
middleware = _build_middleware()
|
|
request = _build_request(
|
|
skills=["alpha"],
|
|
tools=["read_file"],
|
|
system_prompt="base prompt\n\n## Skills System\n- **alpha**: alpha desc",
|
|
)
|
|
|
|
def raise_if_called(_skills_meta):
|
|
raise AssertionError("_build_skills_section should not be called in awrap_model_call")
|
|
|
|
monkeypatch.setattr(middleware, "_build_skills_section", raise_if_called)
|
|
result = await middleware.awrap_model_call(request, _echo_handler)
|
|
prompt = _extract_appended_prompt(result)
|
|
|
|
assert "当前时间：" in prompt
|
|
assert "## Skills System" in prompt
|
|
assert "- **alpha**: alpha desc" in prompt
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_abefore_agent_degrades_when_resolver_fails(monkeypatch: pytest.MonkeyPatch):
|
|
async def fake_resolve(_selected):
|
|
raise RuntimeError("boom")
|
|
|
|
monkeypatch.setattr(runtime_middleware, "resolve_session_snapshot", fake_resolve)
|
|
middleware = _build_middleware()
|
|
request = _build_request(skills=["alpha"], tools=["read_file"], system_prompt="base prompt")
|
|
|
|
await middleware.abefore_agent(request.state, request.runtime)
|
|
|
|
snapshot = request.runtime.context.skill_session_snapshot
|
|
assert snapshot["selected_skills"] == ["alpha"]
|
|
assert snapshot["visible_skills"] == []
|
|
assert "## Skills System" not in request.runtime.context.system_prompt
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_awrap_tool_call_activates_skill_when_visible_in_context():
|
|
middleware = _build_middleware()
|
|
request = _build_tool_request(
|
|
skills=["deliver-prd"],
|
|
visible_skills=["deliver-prd", "brainstorming"],
|
|
file_path="/skills/brainstorming/SKILL.md",
|
|
)
|
|
|
|
async def _handler(_request):
|
|
return ToolMessage(content="ok", tool_call_id="tc-1")
|
|
|
|
result = await middleware.awrap_tool_call(request, _handler)
|
|
assert isinstance(result, Command)
|
|
assert result.update["activated_skills"] == ["brainstorming"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_awrap_tool_call_denies_invisible_skill():
|
|
middleware = _build_middleware()
|
|
request = _build_tool_request(
|
|
skills=["deliver-prd"],
|
|
visible_skills=["deliver-prd"],
|
|
file_path="/skills/brainstorming/SKILL.md",
|
|
)
|
|
|
|
async def _handler(_request):
|
|
return ToolMessage(content="ok", tool_call_id="tc-1")
|
|
|
|
result = await middleware.awrap_tool_call(request, _handler)
|
|
assert isinstance(result, ToolMessage)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_call_injects_dependency_tools_and_mcps_after_activation(monkeypatch: pytest.MonkeyPatch):
|
|
monkeypatch.setattr(
|
|
runtime_middleware,
|
|
"get_buildin_tools",
|
|
lambda: [_FakeTool(name="calculator"), _FakeTool(name="dep-tool")],
|
|
)
|
|
monkeypatch.setattr(runtime_middleware, "get_kb_based_tools", lambda db_names=None: [])
|
|
|
|
snapshot = _build_snapshot(
|
|
selected=["alpha"],
|
|
visible=["alpha", "beta"],
|
|
dependency_map={"alpha": {"tools": ["dep-tool"], "mcps": ["mcp-a"], "skills": ["beta"]}},
|
|
)
|
|
|
|
def fake_build_dependency_bundle(received_snapshot, activated):
|
|
assert received_snapshot == snapshot
|
|
assert activated == ["alpha"]
|
|
return {"tools": ["dep-tool"], "mcps": ["mcp-a"], "skills": ["alpha", "beta"]}
|
|
|
|
monkeypatch.setattr(runtime_middleware, "build_dependency_bundle", fake_build_dependency_bundle)
|
|
|
|
async def fake_get_enabled_mcp_tools(server_name: str):
|
|
if server_name == "mcp-a":
|
|
return [_FakeTool(name="mcp_tool")]
|
|
return []
|
|
|
|
monkeypatch.setattr(runtime_middleware, "get_enabled_mcp_tools", fake_get_enabled_mcp_tools)
|
|
|
|
middleware = RuntimeConfigMiddleware(
|
|
extra_tools=[_FakeTool(name="mcp_tool")],
|
|
enable_model_override=False,
|
|
enable_tools_override=True,
|
|
enable_system_prompt_override=False,
|
|
enable_skills_prompt_override=False,
|
|
)
|
|
|
|
request = _build_request(
|
|
skills=["alpha"],
|
|
tools=["calculator", "dep-tool", "mcp_tool", "read_file"],
|
|
state={"activated_skills": ["alpha"]},
|
|
skill_snapshot=snapshot,
|
|
)
|
|
|
|
result = await middleware.awrap_model_call(request, _echo_handler)
|
|
tool_names = [t.name for t in result.tools]
|
|
assert "dep-tool" in tool_names
|
|
assert "mcp_tool" in tool_names
|
|
assert "calculator" not in tool_names
|