mirror of https://github.com/aliasrobotics/cai.git
114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
"""Tests for syncing MCP tools to the REPL session agent (factory clone vs module singleton)."""
|
|
|
|
from unittest.mock import AsyncMock, Mock
|
|
|
|
import pytest
|
|
|
|
from cai.repl.commands.mcp import (
|
|
_AGENT_MCP_ASSOCIATIONS,
|
|
_GLOBAL_MCP_SERVERS,
|
|
_MCP_TOOL_NAME_TO_SERVER,
|
|
merge_mcp_tools_into_session_agent,
|
|
register_mcp_tool_name,
|
|
strip_mcp_server_from_session_agents,
|
|
unregister_mcp_tools_for_server,
|
|
)
|
|
from cai.sdk.agents import Agent
|
|
from cai.sdk.agents.simple_agent_manager import AGENT_MANAGER
|
|
from cai.sdk.agents.tool import FunctionTool
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clear_mcp_globals():
|
|
_GLOBAL_MCP_SERVERS.clear()
|
|
_AGENT_MCP_ASSOCIATIONS.clear()
|
|
_MCP_TOOL_NAME_TO_SERVER.clear()
|
|
yield
|
|
_GLOBAL_MCP_SERVERS.clear()
|
|
_AGENT_MCP_ASSOCIATIONS.clear()
|
|
_MCP_TOOL_NAME_TO_SERVER.clear()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_agent_manager():
|
|
old_ref = AGENT_MANAGER._active_agent
|
|
old_name = AGENT_MANAGER._active_agent_name
|
|
AGENT_MANAGER._active_agent = None
|
|
AGENT_MANAGER._active_agent_name = None
|
|
yield
|
|
AGENT_MANAGER._active_agent = old_ref
|
|
AGENT_MANAGER._active_agent_name = old_name
|
|
|
|
|
|
def _noop_tool(name: str) -> FunctionTool:
|
|
async def _invoke(_cfg, _ctx, _inp):
|
|
return "ok"
|
|
|
|
return FunctionTool(
|
|
name=name,
|
|
description="d",
|
|
params_json_schema={},
|
|
on_invoke_tool=_invoke,
|
|
strict_json_schema=False,
|
|
)
|
|
|
|
|
|
def test_merge_mcp_tools_into_session_agent_when_types_match():
|
|
session = Agent(name="Red Team", model=Mock(), tools=[_noop_tool("bash_tool")])
|
|
session.model.agent_type = "redteam_agent"
|
|
|
|
mcp_ft = _noop_tool("mcp_navigate")
|
|
setattr(mcp_ft, "_mcp_server", "devtools")
|
|
setattr(mcp_ft, "_is_mcp_tool", True)
|
|
|
|
AGENT_MANAGER.set_active_agent(session, "Red Team")
|
|
|
|
merge_mcp_tools_into_session_agent("redteam_agent", [mcp_ft])
|
|
|
|
names = {t.name for t in session.tools}
|
|
assert "bash_tool" in names
|
|
assert "mcp_navigate" in names
|
|
|
|
|
|
def test_merge_mcp_tools_skips_when_session_agent_type_differs():
|
|
session = Agent(name="Blue", model=Mock(), tools=[_noop_tool("keep")])
|
|
session.model.agent_type = "blueteam_agent"
|
|
|
|
AGENT_MANAGER.set_active_agent(session, "Blue")
|
|
|
|
mcp_ft = _noop_tool("mcp_only")
|
|
merge_mcp_tools_into_session_agent("redteam_agent", [mcp_ft])
|
|
|
|
assert [t.name for t in session.tools] == ["keep"]
|
|
|
|
|
|
def test_strip_mcp_server_removes_tools_from_singleton_and_active(monkeypatch):
|
|
singleton = Agent(name="Red Team", model=Mock(), tools=[])
|
|
singleton.model.agent_type = "redteam_agent"
|
|
keep = _noop_tool("keep")
|
|
mcp_t = _noop_tool("navigate_page")
|
|
setattr(mcp_t, "_mcp_server", "devtools")
|
|
singleton.tools = [keep, mcp_t]
|
|
register_mcp_tool_name("navigate_page", "devtools")
|
|
|
|
session = Agent(name="Red Team", model=Mock(), tools=[])
|
|
session.model.agent_type = "redteam_agent"
|
|
keep2 = _noop_tool("keep2")
|
|
click_t = _noop_tool("click")
|
|
setattr(click_t, "_mcp_server", "devtools")
|
|
session.tools = [keep2, click_t]
|
|
|
|
AGENT_MANAGER.set_active_agent(session, "Red Team")
|
|
|
|
monkeypatch.setattr(
|
|
"cai.repl.commands.mcp.get_available_agents",
|
|
lambda: {"redteam_agent": singleton},
|
|
)
|
|
|
|
strip_mcp_server_from_session_agents("devtools")
|
|
unregister_mcp_tools_for_server("devtools")
|
|
|
|
assert {t.name for t in singleton.tools} == {"keep"}
|
|
assert {t.name for t in session.tools} == {"keep2"}
|
|
assert "navigate_page" not in _MCP_TOOL_NAME_TO_SERVER
|