| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import base64 |
| 5 | import importlib |
| 6 | import sys |
| 7 | from dataclasses import dataclass |
| 8 | from pathlib import Path |
| 9 | from types import ModuleType, SimpleNamespace |
| 10 | |
| 11 | import pytest |
| 12 | |
| 13 | |
| 14 | PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| 15 | if str(PROJECT_ROOT) not in sys.path: |
| 16 | sys.path.insert(0, str(PROJECT_ROOT)) |
| 17 | |
| 18 | |
| 19 | @dataclass |
| 20 | class _StubResponse: |
| 21 | message: str |
| 22 | break_loop: bool |
| 23 | additional: dict | None = None |
| 24 | |
| 25 | |
| 26 | class _StubTool: |
| 27 | def __init__( |
| 28 | self, |
| 29 | agent=None, |
| 30 | name="", |
| 31 | method=None, |
| 32 | args=None, |
| 33 | message="", |
| 34 | loop_data=None, |
| 35 | **kwargs, |
| 36 | ): |
| 37 | self.agent = agent |
| 38 | self.name = name |
| 39 | self.method = method |
| 40 | self.args = args or {} |
| 41 | self.message = message |
| 42 | self.loop_data = loop_data |
| 43 | self.log = None |
| 44 | |
| 45 | def nice_key(self, key: str) -> str: |
| 46 | return key |
| 47 | |
| 48 | |
| 49 | class _FakeContent(SimpleNamespace): |
| 50 | pass |
| 51 | |
| 52 | |
| 53 | class _FakeCallToolResult(SimpleNamespace): |
| 54 | pass |
| 55 | |
| 56 | |
| 57 | class _TrackingLock: |
| 58 | def __init__(self): |
| 59 | self.held = False |
| 60 | |
| 61 | def __enter__(self): |
| 62 | assert self.held is False |
| 63 | self.held = True |
| 64 | return self |
| 65 | |
| 66 | def __exit__(self, exc_type, exc, tb): |
| 67 | self.held = False |
| 68 | return False |
| 69 | |
| 70 | |
| 71 | @pytest.fixture |
| 72 | def mcp_handler_module(monkeypatch, tmp_path): |
| 73 | monkeypatch.delitem(sys.modules, "helpers.mcp_handler", raising=False) |
| 74 | |
| 75 | agent_module = ModuleType("agent") |
| 76 | agent_module.AgentContext = type("AgentContext", (), {}) |
| 77 | agent_module.Agent = type("Agent", (), {}) |
| 78 | agent_module.LoopData = type("LoopData", (), {}) |
| 79 | monkeypatch.setitem(sys.modules, "agent", agent_module) |
| 80 | |
| 81 | tool_module = ModuleType("helpers.tool") |
| 82 | tool_module.Response = _StubResponse |
| 83 | tool_module.Tool = _StubTool |
| 84 | monkeypatch.setitem(sys.modules, "helpers.tool", tool_module) |
| 85 | |
| 86 | settings_module = ModuleType("helpers.settings") |
| 87 | monkeypatch.setitem(sys.modules, "helpers.settings", settings_module) |
| 88 | |
| 89 | history_module = ModuleType("helpers.history") |
| 90 | history_module.RawMessage = lambda **kwargs: dict(kwargs) |
| 91 | monkeypatch.setitem(sys.modules, "helpers.history", history_module) |
| 92 | |
| 93 | mcp_module = ModuleType("mcp") |
| 94 | mcp_module.ClientSession = type("ClientSession", (), {}) |
| 95 | mcp_module.StdioServerParameters = type("StdioServerParameters", (), {}) |
| 96 | monkeypatch.setitem(sys.modules, "mcp", mcp_module) |
| 97 | |
| 98 | mcp_client_stdio = ModuleType("mcp.client.stdio") |
| 99 | mcp_client_stdio.stdio_client = lambda *args, **kwargs: None |
| 100 | monkeypatch.setitem(sys.modules, "mcp.client.stdio", mcp_client_stdio) |
| 101 | |
| 102 | mcp_client_sse = ModuleType("mcp.client.sse") |
| 103 | mcp_client_sse.sse_client = lambda *args, **kwargs: None |
| 104 | monkeypatch.setitem(sys.modules, "mcp.client.sse", mcp_client_sse) |
| 105 | |
| 106 | mcp_client_streamable_http = ModuleType("mcp.client.streamable_http") |
| 107 | mcp_client_streamable_http.streamablehttp_client = lambda *args, **kwargs: None |
| 108 | monkeypatch.setitem( |
| 109 | sys.modules, |
| 110 | "mcp.client.streamable_http", |
| 111 | mcp_client_streamable_http, |
| 112 | ) |
| 113 | |
| 114 | mcp_shared_message = ModuleType("mcp.shared.message") |
| 115 | mcp_shared_message.SessionMessage = type("SessionMessage", (), {}) |
| 116 | monkeypatch.setitem(sys.modules, "mcp.shared.message", mcp_shared_message) |
| 117 | |
| 118 | mcp_types = ModuleType("mcp.types") |
| 119 | mcp_types.CallToolResult = _FakeCallToolResult |
| 120 | mcp_types.ListToolsResult = type("ListToolsResult", (), {}) |
| 121 | monkeypatch.setitem(sys.modules, "mcp.types", mcp_types) |
| 122 | |
| 123 | module = importlib.import_module("helpers.mcp_handler") |
| 124 | |
| 125 | class _SilentPrintStyle: |
| 126 | def __init__(self, *args, **kwargs): |
| 127 | pass |
| 128 | |
| 129 | def print(self, *args, **kwargs): |
| 130 | return self |
| 131 | |
| 132 | def stream(self, *args, **kwargs): |
| 133 | return self |
| 134 | |
| 135 | @staticmethod |
| 136 | def warning(*args, **kwargs): |
| 137 | return None |
| 138 | |
| 139 | def _fake_get_abs_path(*parts): |
| 140 | return str(tmp_path.joinpath(*parts)) |
| 141 | |
| 142 | def _fake_normalize_a0_path(path: str) -> str: |
| 143 | path_obj = Path(path) |
| 144 | try: |
| 145 | rel = path_obj.relative_to(tmp_path) |
| 146 | except ValueError: |
| 147 | return str(path_obj) |
| 148 | return "/a0/" + str(rel).replace("\\", "/") |
| 149 | |
| 150 | monkeypatch.setattr(module, "PrintStyle", _SilentPrintStyle) |
| 151 | monkeypatch.setattr(module.media_artifacts.files, "get_abs_path", _fake_get_abs_path) |
| 152 | monkeypatch.setattr(module.media_artifacts.files, "normalize_a0_path", _fake_normalize_a0_path) |
| 153 | return module, tmp_path |
| 154 | |
| 155 | |
| 156 | def _agent_recorder(context_id: str = "ctx-mcp"): |
| 157 | tool_results: list[tuple[tuple, dict]] = [] |
| 158 | messages: list[tuple[tuple, dict]] = [] |
| 159 | updates: list[dict] = [] |
| 160 | warnings: list[dict] = [] |
| 161 | agent = SimpleNamespace( |
| 162 | agent_name="Agent Zero", |
| 163 | context=SimpleNamespace( |
| 164 | id=context_id, |
| 165 | log=SimpleNamespace(log=lambda **kwargs: warnings.append(kwargs)), |
| 166 | ), |
| 167 | hist_add_tool_result=lambda *args, **kwargs: tool_results.append((args, kwargs)), |
| 168 | hist_add_message=lambda *args, **kwargs: messages.append((args, kwargs)), |
| 169 | ) |
| 170 | log = SimpleNamespace(id="mcp-log", update=lambda **kwargs: updates.append(kwargs)) |
| 171 | return agent, log, tool_results, messages, updates, warnings |
| 172 | |
| 173 | |
| 174 | def test_mcp_config_preserves_dotted_tool_names(mcp_handler_module): |
| 175 | module, _tmp_path = mcp_handler_module |
| 176 | called: list[tuple[str, dict]] = [] |
| 177 | |
| 178 | class _FakeServer: |
| 179 | name = "server" |
| 180 | description = "Fake MCP server" |
| 181 | type = "stdio" |
| 182 | scope = "global" |
| 183 | |
| 184 | def get_tools(self): |
| 185 | return [ |
| 186 | { |
| 187 | "name": "alpha.beta", |
| 188 | "description": "Dotted MCP tool", |
| 189 | "input_schema": {}, |
| 190 | } |
| 191 | ] |
| 192 | |
| 193 | def has_tool(self, tool_name): |
| 194 | return tool_name == "alpha.beta" |
| 195 | |
| 196 | async def call_tool(self, tool_name, input_data): |
| 197 | called.append((tool_name, input_data)) |
| 198 | return _FakeCallToolResult(content=[], isError=False) |
| 199 | |
| 200 | def get_error(self): |
| 201 | return "" |
| 202 | |
| 203 | def get_log(self): |
| 204 | return "" |
| 205 | |
| 206 | config = module.MCPConfig(servers_list=[]) |
| 207 | config.servers = [_FakeServer()] |
| 208 | |
| 209 | assert config.has_tool("server.alpha.beta") is True |
| 210 | asyncio.run(config.call_tool("server.alpha.beta", {"value": 7})) |
| 211 | |
| 212 | assert called == [("alpha.beta", {"value": 7})] |
| 213 | |
| 214 | |
| 215 | def test_mcp_config_resolves_advertised_responses_alias( |
| 216 | mcp_handler_module, monkeypatch |
| 217 | ): |
| 218 | module, _tmp_path = mcp_handler_module |
| 219 | canonical_name = "google_workspace.search_gmail_messages" |
| 220 | native_name = "google_workspace_search_gmail_messages_ecb900b9" |
| 221 | |
| 222 | class _FakeServer: |
| 223 | name = "google_workspace" |
| 224 | |
| 225 | def has_tool(self, tool_name): |
| 226 | return tool_name == "search_gmail_messages" |
| 227 | |
| 228 | config = module.MCPConfig(servers_list=[]) |
| 229 | config.servers = [_FakeServer()] |
| 230 | monkeypatch.setattr( |
| 231 | module.MCPConfig, |
| 232 | "get_for_agent", |
| 233 | classmethod(lambda cls, _agent: config), |
| 234 | ) |
| 235 | |
| 236 | agent = SimpleNamespace( |
| 237 | DATA_NAME_RESPONSES_TOOL_NAME_MAP="responses_tool_name_map", |
| 238 | get_data=lambda key: ( |
| 239 | {native_name: canonical_name} |
| 240 | if key == "responses_tool_name_map" |
| 241 | else None |
| 242 | ), |
| 243 | ) |
| 244 | |
| 245 | assert config.get_tool(agent, canonical_name).name == canonical_name |
| 246 | assert config.get_tool(agent, native_name).name == canonical_name |
| 247 | assert config.get_tool(agent, "local_tool") is None |
| 248 | |
| 249 | |
| 250 | def test_mcp_config_call_tool_releases_config_lock_before_await( |
| 251 | mcp_handler_module, monkeypatch |
| 252 | ): |
| 253 | module, _tmp_path = mcp_handler_module |
| 254 | lock = _TrackingLock() |
| 255 | observed_lock_state: list[bool] = [] |
| 256 | |
| 257 | monkeypatch.setattr(module.MCPConfig, "_MCPConfig__lock", lock, raising=False) |
| 258 | |
| 259 | class _FakeServer: |
| 260 | name = "server" |
| 261 | description = "Fake MCP server" |
| 262 | type = "stdio" |
| 263 | scope = "global" |
| 264 | |
| 265 | def has_tool(self, tool_name): |
| 266 | return tool_name == "run" |
| 267 | |
| 268 | async def call_tool(self, tool_name, input_data): |
| 269 | observed_lock_state.append(lock.held) |
| 270 | await asyncio.sleep(0) |
| 271 | return _FakeCallToolResult(content=[], isError=False) |
| 272 | |
| 273 | config = module.MCPConfig(servers_list=[]) |
| 274 | config.servers = [_FakeServer()] |
| 275 | |
| 276 | asyncio.run(config.call_tool("server.run", {})) |
| 277 | |
| 278 | assert observed_lock_state == [False] |
| 279 | |
| 280 | |
| 281 | def test_mcp_config_update_initializes_outside_config_lock( |
| 282 | mcp_handler_module, monkeypatch |
| 283 | ): |
| 284 | module, _tmp_path = mcp_handler_module |
| 285 | lock = _TrackingLock() |
| 286 | observed_lock_state: list[bool] = [] |
| 287 | original_init = module.MCPConfig.__init__ |
| 288 | |
| 289 | def tracking_init(self, *args, **kwargs): |
| 290 | observed_lock_state.append(lock.held) |
| 291 | original_init(self, *args, **kwargs) |
| 292 | |
| 293 | monkeypatch.setattr(module.MCPConfig, "_MCPConfig__lock", lock, raising=False) |
| 294 | monkeypatch.setattr(module.MCPConfig, "__init__", tracking_init) |
| 295 | |
| 296 | module.MCPConfig.update('{"mcpServers": {}}') |
| 297 | |
| 298 | assert observed_lock_state[-1] is False |
| 299 | |
| 300 | |
| 301 | def test_mcp_status_marks_servers_with_errors_disconnected(mcp_handler_module): |
| 302 | module, _tmp_path = mcp_handler_module |
| 303 | |
| 304 | class _FakeServer: |
| 305 | name = "broken" |
| 306 | description = "Broken MCP server" |
| 307 | type = "stdio" |
| 308 | scope = "global" |
| 309 | |
| 310 | def get_tools(self): |
| 311 | return [] |
| 312 | |
| 313 | def get_error(self): |
| 314 | return "Failed to initialize" |
| 315 | |
| 316 | def get_log(self): |
| 317 | return "stderr" |
| 318 | |
| 319 | config = module.MCPConfig(servers_list=[]) |
| 320 | config.servers = [_FakeServer()] |
| 321 | |
| 322 | status = config.get_servers_status() |
| 323 | |
| 324 | assert status[0]["connected"] is False |
| 325 | assert status[0]["error"] == "Failed to initialize" |
| 326 | assert status[0]["has_log"] is True |
| 327 | |
| 328 | |
| 329 | def test_mcp_disabled_tools_are_hidden_from_agent_paths_but_visible_in_detail(mcp_handler_module): |
| 330 | module, _tmp_path = mcp_handler_module |
| 331 | |
| 332 | server = module.MCPServerLocal( |
| 333 | { |
| 334 | "name": "files", |
| 335 | "command": "npx", |
| 336 | "disabled_tools": ["write_file"], |
| 337 | } |
| 338 | ) |
| 339 | client = getattr(server, "_MCPServerLocal__client") |
| 340 | client.tools = [ |
| 341 | { |
| 342 | "name": "read_file", |
| 343 | "description": "Read a file", |
| 344 | "input_schema": {}, |
| 345 | }, |
| 346 | { |
| 347 | "name": "write_file", |
| 348 | "description": "Write a file", |
| 349 | "input_schema": {}, |
| 350 | }, |
| 351 | ] |
| 352 | |
| 353 | config = module.MCPConfig(servers_list=[]) |
| 354 | config.servers = [server] |
| 355 | |
| 356 | assert [tool["name"] for tool in server.get_tools()] == ["read_file"] |
| 357 | assert server.has_tool("write_file") is False |
| 358 | assert config.has_tool("files.write_file") is False |
| 359 | assert config.get_servers_status()[0]["tool_count"] == 1 |
| 360 | assert "files.write_file" not in config.get_tools_prompt() |
| 361 | |
| 362 | detail_tools = config.get_server_detail("files")["tools"] |
| 363 | assert [(tool["name"], tool.get("disabled", False)) for tool in detail_tools] == [ |
| 364 | ("read_file", False), |
| 365 | ("write_file", True), |
| 366 | ] |
| 367 | |
| 368 | with pytest.raises(ValueError): |
| 369 | asyncio.run(server.call_tool("write_file", {})) |
| 370 | |
| 371 | malformed_config = module.MCPConfig( |
| 372 | servers_list=[ |
| 373 | { |
| 374 | "name": "malformed", |
| 375 | "command": "npx", |
| 376 | "disabled_tools": "write_file", |
| 377 | } |
| 378 | ] |
| 379 | ) |
| 380 | assert malformed_config.servers[0].disabled_tools == [] |
| 381 | |
| 382 | |
| 383 | def test_mcp_local_server_accepts_manager_style_command_lines(mcp_handler_module): |
| 384 | module, _tmp_path = mcp_handler_module |
| 385 | |
| 386 | server = module.MCPServerLocal( |
| 387 | { |
| 388 | "name": "google_workspace", |
| 389 | "command": "uvx workspace-mcp", |
| 390 | "args": [ |
| 391 | "--tool-tier core", |
| 392 | "/tmp/path with spaces", |
| 393 | "--label=Two Words", |
| 394 | ], |
| 395 | } |
| 396 | ) |
| 397 | |
| 398 | assert server.command == "uvx" |
| 399 | assert server.args == [ |
| 400 | "workspace-mcp", |
| 401 | "--tool-tier", |
| 402 | "core", |
| 403 | "/tmp/path with spaces", |
| 404 | "--label=Two Words", |
| 405 | ] |
| 406 | |
| 407 | |
| 408 | def test_mcp_client_call_tool_uses_server_tool_timeout(mcp_handler_module, monkeypatch): |
| 409 | module, _tmp_path = mcp_handler_module |
| 410 | session_timeouts = [] |
| 411 | call_timeouts = [] |
| 412 | |
| 413 | monkeypatch.setattr( |
| 414 | module.settings, |
| 415 | "get_settings", |
| 416 | lambda: {"mcp_client_init_timeout": 10, "mcp_client_tool_timeout": 120}, |
| 417 | raising=False, |
| 418 | ) |
| 419 | |
| 420 | class _FakeSession: |
| 421 | async def call_tool(self, tool_name, input_data, read_timeout_seconds=None): |
| 422 | call_timeouts.append(read_timeout_seconds) |
| 423 | return _FakeCallToolResult(content=[], isError=False) |
| 424 | |
| 425 | class _FakeClient(module.MCPClientBase): |
| 426 | async def _create_stdio_transport(self, current_exit_stack): |
| 427 | raise AssertionError("transport should be bypassed by fake session") |
| 428 | |
| 429 | async def _execute_with_session(self, coro_func, read_timeout_seconds=60): |
| 430 | session_timeouts.append(read_timeout_seconds) |
| 431 | return await coro_func(_FakeSession()) |
| 432 | |
| 433 | client = _FakeClient(SimpleNamespace(name="server", tool_timeout=7, init_timeout=0)) |
| 434 | client.tools = [{"name": "run"}] |
| 435 | |
| 436 | asyncio.run(client.call_tool("run", {"x": 1})) |
| 437 | |
| 438 | assert session_timeouts == [7] |
| 439 | assert call_timeouts[0].total_seconds() == 7 |
| 440 | |
| 441 | |
| 442 | def test_mcp_session_cleanup_timeout_does_not_mask_success( |
| 443 | mcp_handler_module, monkeypatch |
| 444 | ): |
| 445 | module, _tmp_path = mcp_handler_module |
| 446 | monkeypatch.setattr(module, "MCP_SESSION_CLEANUP_TIMEOUT_SECONDS", 0.01) |
| 447 | |
| 448 | class _HangingTransport: |
| 449 | async def __aenter__(self): |
| 450 | return "stdio", "write" |
| 451 | |
| 452 | async def __aexit__(self, exc_type, exc, tb): |
| 453 | await asyncio.sleep(60) |
| 454 | |
| 455 | class _FakeSession: |
| 456 | def __init__(self, *args, **kwargs): |
| 457 | pass |
| 458 | |
| 459 | async def __aenter__(self): |
| 460 | return self |
| 461 | |
| 462 | async def __aexit__(self, exc_type, exc, tb): |
| 463 | return False |
| 464 | |
| 465 | async def initialize(self): |
| 466 | pass |
| 467 | |
| 468 | class _FakeClient(module.MCPClientBase): |
| 469 | async def _create_stdio_transport(self, current_exit_stack): |
| 470 | return await current_exit_stack.enter_async_context(_HangingTransport()) |
| 471 | |
| 472 | async def operation(_session): |
| 473 | return "ok" |
| 474 | |
| 475 | monkeypatch.setattr(module, "ClientSession", _FakeSession) |
| 476 | client = _FakeClient(SimpleNamespace(name="server")) |
| 477 | |
| 478 | assert asyncio.run(client._execute_with_session(operation)) == "ok" |
| 479 | |
| 480 | |
| 481 | def test_mcp_isolated_operation_timeout_returns_control(mcp_handler_module): |
| 482 | module, _tmp_path = mcp_handler_module |
| 483 | |
| 484 | class _FakeClient(module.MCPClientBase): |
| 485 | async def _create_stdio_transport(self, current_exit_stack): |
| 486 | raise AssertionError("transport should not be used") |
| 487 | |
| 488 | async def never_finishes(): |
| 489 | await asyncio.sleep(60) |
| 490 | |
| 491 | client = _FakeClient(SimpleNamespace(name="server")) |
| 492 | |
| 493 | with pytest.raises(TimeoutError): |
| 494 | asyncio.run( |
| 495 | client._run_isolated_operation( |
| 496 | "wedged", |
| 497 | never_finishes, |
| 498 | timeout_seconds=0.01, |
| 499 | ) |
| 500 | ) |
| 501 | |
| 502 | assert "operation did not finish" in client.error |
| 503 | |
| 504 | |
| 505 | def test_mcp_image_content_becomes_history_image_attachment(mcp_handler_module, monkeypatch): |
| 506 | module, tmp_path = mcp_handler_module |
| 507 | agent, log, tool_results, messages, updates, warnings = _agent_recorder() |
| 508 | image_b64 = base64.b64encode(b"image-bytes").decode("ascii") |
| 509 | result = _FakeCallToolResult( |
| 510 | content=[_FakeContent(type="image", data=image_b64, mimeType="image/webp")], |
| 511 | isError=False, |
| 512 | ) |
| 513 | |
| 514 | class _FakeConfig: |
| 515 | async def call_tool(self, name, kwargs): |
| 516 | return result |
| 517 | |
| 518 | monkeypatch.setattr(module.MCPConfig, "get_for_agent", lambda agent: _FakeConfig()) |
| 519 | |
| 520 | tool = module.MCPTool( |
| 521 | agent=agent, |
| 522 | name="venice_image", |
| 523 | method=None, |
| 524 | args={}, |
| 525 | message="", |
| 526 | loop_data=None, |
| 527 | ) |
| 528 | tool.log = log |
| 529 | |
| 530 | response = asyncio.run(tool.execute()) |
| 531 | |
| 532 | assert "[Tool returned no textual content]" not in response.message |
| 533 | assert ( |
| 534 | "Saved MCP image attachment (image/webp, 11 bytes) to " |
| 535 | "/a0/tmp/mcp/ctx_mcp/venice_image/" |
| 536 | ) in response.message |
| 537 | assert response.additional is not None |
| 538 | image_path = response.additional["raw_content"][0]["image_url"]["url"] |
| 539 | assert image_path.startswith("/a0/tmp/mcp/ctx_mcp/venice_image/") |
| 540 | assert response.additional["attachments"] == [image_path] |
| 541 | assert response.additional["media_paths"] == [image_path] |
| 542 | assert (tmp_path / image_path.removeprefix("/a0/")).exists() |
| 543 | |
| 544 | asyncio.run(tool.after_execution(response)) |
| 545 | |
| 546 | assert tool_results[0][0] == ("venice_image", response.message) |
| 547 | assert tool_results[0][1]["attachments"] == [image_path] |
| 548 | assert tool_results[0][1]["media_paths"] == [image_path] |
| 549 | raw_message = messages[0][1]["content"] |
| 550 | assert raw_message["raw_content"][0]["image_url"]["url"] == image_path |
| 551 | assert messages[0][1]["tokens"] == module.MCP_MEDIA_TOKENS_ESTIMATE |
| 552 | assert updates[-1]["content"] == response.message |
| 553 | assert warnings == [] |
| 554 | |
| 555 | |
| 556 | def test_mcp_audio_content_is_saved_instead_of_discarded(mcp_handler_module, monkeypatch): |
| 557 | module, tmp_path = mcp_handler_module |
| 558 | agent, log, tool_results, messages, updates, warnings = _agent_recorder() |
| 559 | audio_b64 = base64.b64encode(b"audio-bytes").decode("ascii") |
| 560 | result = _FakeCallToolResult( |
| 561 | content=[_FakeContent(type="audio", data=audio_b64, mimeType="audio/mpeg")], |
| 562 | isError=False, |
| 563 | ) |
| 564 | |
| 565 | class _FakeConfig: |
| 566 | async def call_tool(self, name, kwargs): |
| 567 | return result |
| 568 | |
| 569 | monkeypatch.setattr(module.MCPConfig, "get_for_agent", lambda agent: _FakeConfig()) |
| 570 | |
| 571 | tool = module.MCPTool( |
| 572 | agent=agent, |
| 573 | name="venice_audio", |
| 574 | method=None, |
| 575 | args={}, |
| 576 | message="", |
| 577 | loop_data=None, |
| 578 | ) |
| 579 | tool.log = log |
| 580 | |
| 581 | response = asyncio.run(tool.execute()) |
| 582 | |
| 583 | assert response.additional is None |
| 584 | assert "[Tool returned no textual content]" not in response.message |
| 585 | assert "Saved MCP audio attachment (audio/mpeg, 11 bytes) to /a0/tmp/mcp/ctx_mcp/venice_audio/" in response.message |
| 586 | saved_path = response.message.split(" to ", 1)[1].rstrip(".") |
| 587 | assert (tmp_path / saved_path.removeprefix("/a0/")).exists() |
| 588 | |
| 589 | asyncio.run(tool.after_execution(response)) |
| 590 | |
| 591 | assert tool_results[0][0] == ("venice_audio", response.message) |
| 592 | assert messages == [] |
| 593 | assert updates[-1]["content"] == response.message |
| 594 | assert warnings == [] |
| 595 | |
| 596 | |
| 597 | def test_mcp_image_resource_blob_becomes_history_image_attachment(mcp_handler_module, monkeypatch): |
| 598 | module, tmp_path = mcp_handler_module |
| 599 | agent, log, tool_results, messages, updates, warnings = _agent_recorder() |
| 600 | image_b64 = base64.b64encode(b"resource-image").decode("ascii") |
| 601 | result = _FakeCallToolResult( |
| 602 | content=[ |
| 603 | _FakeContent( |
| 604 | type="resource", |
| 605 | resource=_FakeContent( |
| 606 | uri="memory://venice/image.webp", |
| 607 | mimeType="image/webp", |
| 608 | blob=image_b64, |
| 609 | ), |
| 610 | ) |
| 611 | ], |
| 612 | isError=False, |
| 613 | ) |
| 614 | |
| 615 | class _FakeConfig: |
| 616 | async def call_tool(self, name, kwargs): |
| 617 | return result |
| 618 | |
| 619 | monkeypatch.setattr(module.MCPConfig, "get_for_agent", lambda agent: _FakeConfig()) |
| 620 | |
| 621 | tool = module.MCPTool( |
| 622 | agent=agent, |
| 623 | name="venice_resource_image", |
| 624 | method=None, |
| 625 | args={}, |
| 626 | message="", |
| 627 | loop_data=None, |
| 628 | ) |
| 629 | tool.log = log |
| 630 | |
| 631 | response = asyncio.run(tool.execute()) |
| 632 | |
| 633 | assert ( |
| 634 | "Saved MCP resource image attachment (image/webp, 14 bytes) to " |
| 635 | "/a0/tmp/mcp/ctx_mcp/venice_resource_image/" |
| 636 | ) in response.message |
| 637 | assert response.additional is not None |
| 638 | image_path = response.additional["raw_content"][0]["image_url"]["url"] |
| 639 | assert image_path.startswith("/a0/tmp/mcp/ctx_mcp/venice_resource_image/") |
| 640 | assert response.additional["attachments"] == [image_path] |
| 641 | assert (tmp_path / image_path.removeprefix("/a0/")).exists() |
| 642 | |
| 643 | asyncio.run(tool.after_execution(response)) |
| 644 | |
| 645 | assert tool_results[0][0] == ("venice_resource_image", response.message) |
| 646 | raw_message = messages[0][1]["content"] |
| 647 | assert raw_message["raw_content"][0]["image_url"]["url"] == image_path |
| 648 | assert updates[-1]["content"] == response.message |
| 649 | assert warnings == [] |
| 650 | |
| 651 | |
| 652 | def test_mcp_resource_text_is_preserved(mcp_handler_module, monkeypatch): |
| 653 | module, _tmp_path = mcp_handler_module |
| 654 | agent, log, tool_results, messages, updates, warnings = _agent_recorder() |
| 655 | result = _FakeCallToolResult( |
| 656 | content=[ |
| 657 | _FakeContent( |
| 658 | type="resource", |
| 659 | resource=_FakeContent( |
| 660 | uri="memory://venice/caption.txt", |
| 661 | mimeType="text/plain", |
| 662 | text="Generated caption text", |
| 663 | ), |
| 664 | ) |
| 665 | ], |
| 666 | isError=False, |
| 667 | ) |
| 668 | |
| 669 | class _FakeConfig: |
| 670 | async def call_tool(self, name, kwargs): |
| 671 | return result |
| 672 | |
| 673 | monkeypatch.setattr(module.MCPConfig, "get_for_agent", lambda agent: _FakeConfig()) |
| 674 | |
| 675 | tool = module.MCPTool( |
| 676 | agent=agent, |
| 677 | name="venice_resource", |
| 678 | method=None, |
| 679 | args={}, |
| 680 | message="", |
| 681 | loop_data=None, |
| 682 | ) |
| 683 | tool.log = log |
| 684 | |
| 685 | response = asyncio.run(tool.execute()) |
| 686 | |
| 687 | assert response.additional is None |
| 688 | assert "Resource memory://venice/caption.txt:" in response.message |
| 689 | assert "Generated caption text" in response.message |
| 690 | |
| 691 | asyncio.run(tool.after_execution(response)) |
| 692 | |
| 693 | assert tool_results[0][0] == ("venice_resource", response.message) |
| 694 | assert messages == [] |
| 695 | assert updates[-1]["content"] == response.message |
| 696 | assert warnings == [] |