"""Tests for omnigent.tools.mcp (MCP connections and tools).""" from __future__ import annotations import asyncio import json import logging from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass, field from datetime import timedelta from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from cachetools import TTLCache from mcp.shared.exceptions import McpError from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, ImageContent, TextContent from mcp.types import Tool as McpToolDef from omnigent.spec.types import MCPServerConfig, RetryPolicy from omnigent.tools.mcp import ( _CIRCUIT_BREAKER_COOLDOWN_SECONDS, _CIRCUIT_BREAKER_THRESHOLD, _MCP_RECONNECT_DEFAULTS, McpElicitationRequired, McpServerConnection, McpServerDisabledError, _backoff_delay, _cache_key, _CircuitBreaker, _collect_problematic_keywords, _discovery_cache, _format_call_result, _is_connection_error, _normalize_input_schema, clear_discovery_cache, ) @pytest.fixture(autouse=True) def _clean_cache() -> None: """ Clear the module-level discovery cache before each test. """ clear_discovery_cache() def _make_http_config( name: str = "test-server", url: str = "http://localhost:9000/mcp", ) -> MCPServerConfig: """ Create a minimal HTTP MCP server config. :param name: Server name identifier. :param url: Server endpoint URL. :returns: An ``MCPServerConfig`` for HTTP transport. """ return MCPServerConfig( name=name, url=url, ) def _make_mcp_tool_def( name: str = "test_tool", description: str = "A test tool.", ) -> MagicMock: """ Create a mock MCP tool definition matching ``mcp.types.Tool``. Uses a MagicMock because we only read ``.name``, ``.description``, and ``.inputSchema`` — these are plain attribute reads, not isinstance checks. :param name: Tool name. :param description: Tool description. :returns: A mock with name, description, and inputSchema. """ tool_def = MagicMock() tool_def.name = name tool_def.description = description tool_def.inputSchema = { "type": "object", "properties": { "query": {"type": "string"}, }, "required": ["query"], } return tool_def @contextmanager def _mock_mcp_transport( tools: list[MagicMock] | None = None, ) -> Iterator[AsyncMock]: """ Mock the MCP transport and session for ``connect()`` tests. Patches ``sse_client`` and ``ClientSession`` so that ``McpServerConnection.connect()`` can run without a real MCP server. The mock session's ``list_tools()`` returns the provided tool definitions. :param tools: Mock tool definitions to return from ``list_tools()``. Defaults to an empty list. :yields: The mock ``ClientSession`` instance. """ mock_session = AsyncMock() mock_tools_result = MagicMock() mock_tools_result.tools = tools or [] mock_session.list_tools.return_value = mock_tools_result mock_session.initialize = AsyncMock() mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aexit__ = AsyncMock(return_value=False) mock_ctx = AsyncMock() # streamablehttp_client yields (read, write, get_session_id) mock_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock(), MagicMock())) mock_ctx.__aexit__ = AsyncMock(return_value=False) with patch( "omnigent.tools.mcp.streamablehttp_client", return_value=mock_ctx, ): with patch( "omnigent.tools.mcp.ClientSession", return_value=mock_session, ): yield mock_session # ── _cache_key ─────────────────────────────────────────── def test_cache_key_includes_name_and_url() -> None: """ Cache key includes the server name and URL. """ config = _make_http_config() key = _cache_key(config) assert "http" in key assert "test-server" in key assert "localhost:9000" in key def test_cache_key_different_configs_differ() -> None: """ Different server configs produce different cache keys. """ key1 = _cache_key(_make_http_config("server-a")) key2 = _cache_key(_make_http_config("server-b")) assert key1 != key2 def test_cache_key_stdio_includes_command_and_args() -> None: """ Stdio MCP cache key includes the transport tag, name, command, and joined args — so two stdio configs that share a name but point at different subprocesses don't collide in the module-level ``_discovery_cache``. What breaks if this fails: two agents that both declare a ``glean`` MCP — one pointing at ``--profile dogfood`` and another at ``--profile prod`` — would share one cache entry and one would see the other's discovered tools. """ config = MCPServerConfig( name="stdio-server", transport="stdio", command="npx", args=["-y", "@modelcontextprotocol/server-github"], ) key = _cache_key(config) # Tag identifies the transport so stdio + http with the same # name never collide. assert key.startswith("stdio:") assert "stdio-server" in key assert "npx" in key # Arg content matters — different args means different cache # entries, not a silent share. assert "@modelcontextprotocol/server-github" in key def test_cache_key_stdio_args_changes_key() -> None: """ Different args on the same stdio command produce different cache keys — covers the realistic "same binary, different CLI flags" case (e.g. databricks MCPs differentiated by ``--profile``). What breaks if this fails: the cache treats ``my-mcp --profile prod`` and ``my-mcp --profile dev`` as the same server, silently cross-contaminating discovery results between agent runs. """ base = { "name": "my-mcp", "transport": "stdio", "command": "my-mcp", } key_prod = _cache_key(MCPServerConfig(args=["--profile", "prod"], **base)) key_dev = _cache_key(MCPServerConfig(args=["--profile", "dev"], **base)) assert key_prod != key_dev def test_cache_key_stdio_and_http_do_not_collide() -> None: """ A stdio server named ``my-mcp`` and an HTTP server named ``my-mcp`` produce different keys — the transport tag is the first segment of the key so same-name, different- transport never shares a cache entry. What breaks if this fails: an HTTP MCP named ``foo``'s tool list would get served to an agent declaring a stdio MCP also named ``foo`` (or vice versa). """ http = MCPServerConfig(name="foo", url="http://mcp.example.com") stdio = MCPServerConfig(name="foo", transport="stdio", command="foo") assert _cache_key(http) != _cache_key(stdio) # ── McpServerConnection caching ────────────────────────── @pytest.mark.asyncio() async def test_connect_skips_list_tools_when_cache_fresh() -> None: """ ``connect()`` skips the ``list_tools()`` round-trip when the cache is fresh, but still opens a live session so ``call_tool()`` works. """ config = _make_http_config() tool_def = _make_mcp_tool_def() # Pre-populate the cache — TTLCache uses dict assignment. _discovery_cache[_cache_key(config)] = [tool_def] with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) tools = await conn.connect() assert len(tools) == 1 assert tools[0].name == "test_tool" # Session was created (for invocation), but list_tools # was NOT called (served from cache). mock_session.initialize.assert_awaited_once() mock_session.list_tools.assert_not_awaited() await conn.close() @pytest.mark.asyncio() async def test_cached_connect_has_live_session() -> None: """ When discovery is served from cache, the connection still has a live session that can invoke tools. """ config = _make_http_config() _discovery_cache[_cache_key(config)] = [_make_mcp_tool_def()] with _mock_mcp_transport() as mock_session: # Set up call_tool to return a mock result. mock_result = MagicMock() mock_result.content = [TextContent(type="text", text="cached ok")] mock_result.isError = False mock_session.call_tool.return_value = mock_result conn = McpServerConnection(config=config) await conn.connect() result = await conn.call_tool("test_tool", {"query": "hi"}) assert result == "cached ok" mock_session.call_tool.assert_awaited_once() await conn.close() @pytest.mark.asyncio() async def test_connect_skips_expired_cache() -> None: """ ``connect()`` ignores cache entries older than the TTL and performs a live discovery via ``list_tools()``. """ config = _make_http_config() # Use a TTLCache with a controllable timer so we can # simulate expiry without sleeping. Start at t=0, insert # the entry, then advance past the TTL. current_time = [0.0] expired_cache: TTLCache[str, list[MagicMock]] = TTLCache( maxsize=64, ttl=300, timer=lambda: current_time[0], ) expired_cache[_cache_key(config)] = [_make_mcp_tool_def()] # Advance time past the 300s TTL. current_time[0] = 1000.0 fresh_tool = _make_mcp_tool_def("fresh_tool") with _mock_mcp_transport([fresh_tool]) as mock_session: with patch("omnigent.tools.mcp._discovery_cache", expired_cache): conn = McpServerConnection(config=config) tools = await conn.connect() assert len(tools) == 1 assert tools[0].name == "fresh_tool" mock_session.initialize.assert_awaited_once() mock_session.list_tools.assert_awaited_once() await conn.close() @pytest.mark.asyncio() async def test_connect_populates_cache() -> None: """ A successful live ``connect()`` stores results in the module-level cache. """ config = _make_http_config() tool_def = _make_mcp_tool_def("cached_tool") with _mock_mcp_transport([tool_def]): conn = McpServerConnection(config=config) await conn.connect() key = _cache_key(config) assert key in _discovery_cache cached = _discovery_cache.get(key) assert cached is not None assert len(cached) == 1 assert cached[0].name == "cached_tool" await conn.close() # ── McpServerConnection.call_tool ──────────────────────── @pytest.mark.asyncio() async def test_call_tool_raises_without_connect() -> None: """ ``call_tool()`` raises RuntimeError when ``connect()`` was never called. """ conn = McpServerConnection(config=_make_http_config()) with pytest.raises(RuntimeError, match="no live session"): await conn.call_tool("test_tool", {"query": "hi"}) # ── McpServerConnection.close ──────────────────────────── @pytest.mark.asyncio() async def test_close_is_safe_when_never_connected() -> None: """ ``close()`` does not raise if ``connect()`` was never called. """ conn = McpServerConnection(config=_make_http_config()) await conn.close() def test_mcp_server_config_repr_redacts_headers() -> None: """ MCPServerConfig.__repr__ replaces header values with ``[REDACTED]`` so credentials don't leak in logs or exception tracebacks. If the real Authorization token appeared in repr, it would be captured by ``_logger.exception()`` in manager.py when MCP connections fail. """ config = MCPServerConfig( name="secret-svc", url="http://example.com/sse", headers={ "Authorization": "Bearer sk-SUPER-SECRET-TOKEN", "X-Custom": "also-secret", }, ) r = repr(config) # Header keys are visible (useful for debugging which headers are set) assert "Authorization" in r assert "X-Custom" in r # Actual secret values must NOT appear assert "sk-SUPER-SECRET-TOKEN" not in r assert "also-secret" not in r # Redaction marker is present assert "[REDACTED]" in r # Non-sensitive fields are still visible assert "secret-svc" in r assert "http://example.com/sse" in r def test_mcp_server_config_repr_empty_headers() -> None: """ Repr works correctly when there are no headers — no crash, no ``[REDACTED]`` in the output. """ config = MCPServerConfig(name="plain", url="http://localhost/sse") r = repr(config) assert "plain" in r assert "http://localhost/sse" in r assert "[REDACTED]" not in r # ── _normalize_input_schema ─────────────────────────────── def test_normalize_none_schema_returns_empty_object() -> None: """ ``None`` inputSchema (tool has no parameters) is normalized to a valid empty object schema. """ result = _normalize_input_schema(None, "no_args_tool") assert result == {"type": "object", "properties": {}} def test_normalize_missing_properties_injects_empty() -> None: """ A schema with ``type: object`` but no ``properties`` key gets ``properties: {}`` injected. OpenAI rejects schemas without this key (openai/openai-agents-python#449). """ schema = {"type": "object"} result = _normalize_input_schema(schema, "bare_object_tool") assert result["properties"] == {} assert result["type"] == "object" def test_normalize_preserves_existing_properties() -> None: """ A schema that already has ``properties`` is returned as-is (no double-injection). """ schema = { "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], } result = _normalize_input_schema(schema, "normal_tool") assert result["properties"] == {"query": {"type": "string"}} assert result["required"] == ["query"] def test_normalize_does_not_mutate_original_schema() -> None: """ ``_normalize_input_schema`` returns a new dict when modifying — it does not mutate the original schema dict. """ original = {"type": "object"} result = _normalize_input_schema(original, "test") # Result has properties injected. assert "properties" in result # Original is untouched. assert "properties" not in original def test_normalize_non_object_schema_unchanged() -> None: """ A schema with a non-object type (e.g. ``array``) is not modified — ``properties`` injection only applies to objects. """ schema = {"type": "array", "items": {"type": "string"}} result = _normalize_input_schema(schema, "array_tool") assert result == schema assert "properties" not in result def test_normalize_warns_on_ref(caplog: pytest.LogCaptureFixture) -> None: """ A schema containing ``$ref`` triggers a warning log. """ schema = { "type": "object", "properties": { "item": {"$ref": "#/$defs/Item"}, }, "$defs": { "Item": { "type": "object", "properties": {"name": {"type": "string"}}, }, }, } _normalize_input_schema(schema, "ref_tool") assert any("$ref" in msg for msg in caplog.messages) def test_normalize_warns_on_oneof(caplog: pytest.LogCaptureFixture) -> None: """ A schema containing ``oneOf`` triggers a warning log. """ schema = { "type": "object", "properties": { "value": { "oneOf": [ {"type": "string"}, {"type": "integer"}, ], }, }, } _normalize_input_schema(schema, "oneof_tool") assert any("oneOf" in msg for msg in caplog.messages) def test_normalize_no_warning_for_clean_schema( caplog: pytest.LogCaptureFixture, ) -> None: """ A clean schema with no problematic keywords produces no warnings. """ schema = { "type": "object", "properties": {"q": {"type": "string"}}, } _normalize_input_schema(schema, "clean_tool") assert not any("reject" in msg or "inconsistent" in msg for msg in caplog.messages) # ── _collect_problematic_keywords ───────────────────────── def test_collect_finds_ref_in_properties() -> None: """ ``$ref`` nested inside a property is detected. """ schema = { "type": "object", "properties": { "item": {"$ref": "#/$defs/Item"}, }, } assert "$ref" in _collect_problematic_keywords(schema) def test_collect_finds_oneof_in_nested_property() -> None: """ ``oneOf`` inside a nested property is detected. """ schema = { "type": "object", "properties": { "value": { "oneOf": [ {"type": "string"}, {"type": "integer"}, ], }, }, } assert "oneOf" in _collect_problematic_keywords(schema) def test_collect_finds_keywords_in_array_items() -> None: """ Problematic keywords inside ``items`` of an array are detected. """ schema = { "type": "array", "items": { "type": "object", "properties": { "x": {"$ref": "#/$defs/X"}, }, }, } assert "$ref" in _collect_problematic_keywords(schema) def test_collect_finds_keywords_in_defs() -> None: """ Problematic keywords inside ``$defs`` are detected. """ schema = { "type": "object", "properties": {}, "$defs": { "Node": { "type": "object", "properties": { "child": {"$ref": "#/$defs/Node"}, }, }, }, } assert "$ref" in _collect_problematic_keywords(schema) def test_collect_returns_empty_for_clean_schema() -> None: """ A clean schema returns an empty set. """ schema = { "type": "object", "properties": { "name": {"type": "string"}, "count": {"type": "integer"}, }, } assert _collect_problematic_keywords(schema) == set() # ── _format_call_result ────────────────────────────────── def test_format_call_result_text_content() -> None: """ Text content blocks are extracted and joined. """ block = TextContent(type="text", text="Hello world") result = MagicMock() result.content = [block] result.isError = False assert _format_call_result(result) == "Hello world" def test_format_call_result_multiple_blocks() -> None: """ Multiple text blocks are joined with newlines. """ block1 = TextContent(type="text", text="Line 1") block2 = TextContent(type="text", text="Line 2") result = MagicMock() result.content = [block1, block2] result.isError = False assert _format_call_result(result) == "Line 1\nLine 2" def test_format_call_result_error_prefix() -> None: """ Error results are prefixed with "Error: ". """ block = TextContent(type="text", text="something went wrong") result = MagicMock() result.content = [block] result.isError = True formatted = _format_call_result(result) assert formatted.startswith("Error: ") assert "something went wrong" in formatted def test_format_call_result_non_text_content() -> None: """ Non-text content (e.g. images) is serialized as JSON via ``model_dump()``. """ block = ImageContent(type="image", data="base64data", mimeType="image/png") result = MagicMock() result.content = [block] result.isError = False formatted = _format_call_result(result) parsed = json.loads(formatted) assert parsed["type"] == "image" assert parsed["data"] == "base64data" def test_format_call_result_empty_content() -> None: """ An empty content list returns ``"(empty response)"`` instead of a blank string. """ result = MagicMock() result.content = [] result.isError = False assert _format_call_result(result) == "(empty response)" def test_format_call_result_empty_content_with_error() -> None: """ An empty content list with ``isError=True`` returns ``"Error: (empty response)"``. """ result = MagicMock() result.content = [] result.isError = True assert _format_call_result(result) == "Error: (empty response)" # ── clear_discovery_cache ──────────────────────────────── def test_clear_discovery_cache() -> None: """ ``clear_discovery_cache()`` empties the module-level cache. """ config = _make_http_config() _discovery_cache[_cache_key(config)] = [] assert len(_discovery_cache) > 0 clear_discovery_cache() assert len(_discovery_cache) == 0 def test_discovery_cache_evicts_lru_when_full() -> None: """ The discovery cache evicts the least-recently-used entry when it reaches ``maxsize``. Uses a small TTLCache (maxsize=2) to verify that inserting a third entry evicts the oldest one. """ small_cache: TTLCache[str, list[MagicMock]] = TTLCache( maxsize=2, ttl=300, ) small_cache["server-a"] = [_make_mcp_tool_def("tool_a")] small_cache["server-b"] = [_make_mcp_tool_def("tool_b")] # Inserting a third entry should evict the LRU (server-a). small_cache["server-c"] = [_make_mcp_tool_def("tool_c")] assert "server-a" not in small_cache assert "server-b" in small_cache assert "server-c" in small_cache assert len(small_cache) == 2 # ── _run_async ─────────────────────────────────────────── # ── _is_connection_error ───────────────────────────────── def test_is_connection_error_eof() -> None: """ EOFError is classified as a connection error. """ assert _is_connection_error(EOFError()) is True def test_is_connection_error_broken_pipe() -> None: """ BrokenPipeError is classified as a connection error. """ assert _is_connection_error(BrokenPipeError()) is True def test_is_connection_error_connection_reset() -> None: """ ConnectionResetError (subclass of ConnectionError) is classified as a connection error. """ assert _is_connection_error(ConnectionResetError()) is True def test_is_connection_error_mcp_connection_closed() -> None: """ McpError with CONNECTION_CLOSED code is classified as a connection error. """ exc = McpError( ErrorData( code=CONNECTION_CLOSED, message="Connection closed", ) ) assert _is_connection_error(exc) is True def test_is_connection_error_mcp_other_code() -> None: """ McpError with a non-connection code (e.g. INVALID_PARAMS) is NOT classified as a connection error. """ exc = McpError( ErrorData( code=-32602, # INVALID_PARAMS message="Invalid params", ) ) assert _is_connection_error(exc) is False def test_is_connection_error_value_error() -> None: """ ValueError is NOT classified as a connection error. """ assert _is_connection_error(ValueError("bad")) is False # ── _backoff_delay ──────────────────────────────────────── def test_backoff_delay_increases_with_attempt() -> None: """ ``_backoff_delay`` increases with each attempt (exponential backoff) and applies jitter (0.5–1.5x). """ retry = RetryPolicy(backoff_base_s=2.0, backoff_max_s=30.0) # Attempt 0 → retry_index=1: base * 2^0 = 2.0, jitter [1.0, 3.0]. delay_0 = _backoff_delay(0, retry) assert 1.0 <= delay_0 <= 3.0 # Attempt 1 → retry_index=2: base * 2^1 = 4.0, jitter [2.0, 6.0]. delay_1 = _backoff_delay(1, retry) assert 2.0 <= delay_1 <= 6.0 # Attempt 2 → retry_index=3: base * 2^2 = 8.0, jitter [4.0, 12.0]. delay_2 = _backoff_delay(2, retry) assert 4.0 <= delay_2 <= 12.0 def test_backoff_delay_capped_at_max() -> None: """ ``_backoff_delay`` never exceeds ``backoff_max_s`` (before jitter). """ retry = RetryPolicy(backoff_base_s=10.0, backoff_max_s=5.0) # 10 * 2^0 = 10, capped to 5.0; jitter[0.5, 1.5] → [2.5, 7.5]. delay = _backoff_delay(0, retry) assert delay <= 7.5 # ── Reconnection on server death ───────────────────────── @pytest.mark.asyncio() async def test_call_tool_reconnects_on_connection_error() -> None: """ When a tool call fails with a connection error, the connection reconnects with backoff and retries. """ config = _make_http_config() with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) await conn.connect() # First call_tool raises a connection error (server died). # Second call_tool (after reconnect) succeeds. ok_result = MagicMock() ok_result.content = [TextContent(type="text", text="recovered")] ok_result.isError = False mock_session.call_tool.side_effect = [ EOFError("server died"), ok_result, ] # Patch _reconnect to re-establish the mock session # (in production this opens a new transport). # Patch the _sleep indirection so retry backoff is instant. with patch.object(conn, "_reconnect", new_callable=AsyncMock) as mock_reconnect: with patch("omnigent.tools.mcp._sleep", new_callable=AsyncMock): result = await conn.call_tool("test_tool", {"query": "hi"}) assert result == "recovered" mock_reconnect.assert_awaited_once() await conn.close() @pytest.mark.asyncio() async def test_call_tool_does_not_reconnect_on_tool_error() -> None: """ When a tool call fails with a non-connection error (e.g. McpError for invalid params), no reconnect is attempted. """ config = _make_http_config() with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) await conn.connect() mock_session.call_tool.side_effect = McpError( ErrorData( code=-32602, message="Invalid params", ) ) with pytest.raises(McpError): await conn.call_tool("test_tool", {"query": "hi"}) await conn.close() @pytest.mark.asyncio() async def test_call_tool_exhausts_all_retries_then_raises() -> None: """ When all reconnect attempts fail, the last connection error is propagated to the caller. """ # max_retries=2 means 2 retries beyond the first attempt, # so 3 invoke calls total (1 initial + 2 retries). config = MCPServerConfig( name="test-retry-exhaust", url="http://localhost:9000/mcp", retry=RetryPolicy( max_retries=2, backoff_base_s=1.0, backoff_max_s=10.0, ), ) with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) await conn.connect() mock_session.call_tool.side_effect = [ EOFError("attempt 1"), EOFError("attempt 2"), EOFError("attempt 3"), ] with patch.object(conn, "_reconnect", new_callable=AsyncMock): with patch("omnigent.tools.mcp._sleep", new_callable=AsyncMock): with pytest.raises(EOFError, match="attempt 3"): await conn.call_tool("test_tool", {"query": "hi"}) # 3 invoke calls, 2 reconnects (no reconnect after last failure). assert mock_session.call_tool.await_count == 3 @pytest.mark.asyncio() async def test_call_tool_uses_config_retry_policy() -> None: """ ``call_tool()`` uses the per-server ``config.retry`` when set, rather than the module-level default. """ # max_retries=1 means 1 retry beyond first attempt — 2 invoke calls total. config = MCPServerConfig( name="test-custom-retry", url="http://localhost:9000/mcp", retry=RetryPolicy( max_retries=1, backoff_base_s=0.5, backoff_max_s=5.0, ), ) with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) await conn.connect() mock_session.call_tool.side_effect = [ EOFError("attempt 1"), EOFError("attempt 2"), ] with patch.object(conn, "_reconnect", new_callable=AsyncMock): with patch("omnigent.tools.mcp._sleep", new_callable=AsyncMock): with pytest.raises(EOFError, match="attempt 2"): await conn.call_tool("test_tool", {"query": "hi"}) assert mock_session.call_tool.await_count == 2 @pytest.mark.asyncio() async def test_call_tool_sleeps_between_retries() -> None: """ ``call_tool()`` sleeps with backoff between reconnect attempts. Verifies that the ``_sleep`` indirection is called with increasing delays. """ # max_retries=2 → 3 total attempts (2 errors + 1 success). config = MCPServerConfig( name="test-backoff", url="http://localhost:9000/mcp", retry=RetryPolicy( max_retries=2, backoff_base_s=2.0, backoff_max_s=30.0, ), ) with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) await conn.connect() ok_result = MagicMock() ok_result.content = [TextContent(type="text", text="ok")] ok_result.isError = False mock_session.call_tool.side_effect = [ EOFError("attempt 1"), EOFError("attempt 2"), ok_result, ] with patch.object(conn, "_reconnect", new_callable=AsyncMock): with patch( "omnigent.tools.mcp._sleep", new_callable=AsyncMock, ) as mock_sleep: result = await conn.call_tool("test_tool", {"query": "hi"}) assert result == "ok" # Two sleeps: before retry 1 and before retry 2. assert mock_sleep.await_count == 2 # Delays computed by RetryPolicy.compute_backoff_delay: # retry_index=1 → 2.0 * 2^0 = 2.0; retry_index=2 → 2.0 * 2^1 = 4.0. # Jitter is uniform[0.5, 1.5], so delay 1 ∈ [1.0, 3.0], # delay 2 ∈ [2.0, 6.0]. delay_1 = mock_sleep.await_args_list[0].args[0] delay_2 = mock_sleep.await_args_list[1].args[0] assert 1.0 <= delay_1 <= 3.0 assert 2.0 <= delay_2 <= 6.0 @pytest.mark.asyncio() async def test_call_tool_default_retry_has_three_attempts() -> None: """ When ``config.retry`` is ``None``, ``call_tool()`` falls back to ``_MCP_RECONNECT_DEFAULTS`` which allows 3 attempts. """ # No retry config — should use default (3 attempts). config = MCPServerConfig( name="test-default-retry", url="http://localhost:9000/mcp", # retry defaults to None ) with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) await conn.connect() mock_session.call_tool.side_effect = [ EOFError("attempt 1"), EOFError("attempt 2"), EOFError("attempt 3"), ] with patch.object(conn, "_reconnect", new_callable=AsyncMock): with patch("omnigent.tools.mcp._sleep", new_callable=AsyncMock): with pytest.raises(EOFError, match="attempt 3"): await conn.call_tool("test_tool", {"query": "hi"}) # total attempts = max_retries + 1 (initial + retries). assert mock_session.call_tool.await_count == _MCP_RECONNECT_DEFAULTS.max_retries + 1 # ── Timeout propagation ────────────────────────────────── @pytest.mark.asyncio() async def test_connect_passes_timeout_to_client_session() -> None: """ When ``MCPServerConfig.timeout`` is set, ``connect()`` must pass ``read_timeout_seconds=timedelta(seconds=timeout)`` to ``ClientSession``. """ config = MCPServerConfig( name="test-timeout", url="http://localhost:9000/mcp", timeout=60, ) captured_kwargs: dict[str, Any] = {} def _capturing_session( *args: Any, **kwargs: Any, ) -> AsyncMock: """ Fake ``ClientSession`` constructor that records kwargs. :param args: Positional args (read_stream, write_stream). :param kwargs: Keyword args including read_timeout_seconds. :returns: A mock session with working initialize/list_tools. """ captured_kwargs.update(kwargs) mock_session = AsyncMock() mock_session.initialize = AsyncMock() mock_tools = MagicMock() mock_tools.tools = [] mock_session.list_tools.return_value = mock_tools mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aexit__ = AsyncMock(return_value=False) return mock_session mock_ctx = AsyncMock() mock_ctx.__aenter__ = AsyncMock( return_value=(MagicMock(), MagicMock(), MagicMock()), ) mock_ctx.__aexit__ = AsyncMock(return_value=False) with patch( "omnigent.tools.mcp.ClientSession", side_effect=_capturing_session, ): with patch( "omnigent.tools.mcp.streamablehttp_client", return_value=mock_ctx, ): conn = McpServerConnection(config=config) await conn.connect() # timeout=60 must be converted to timedelta(seconds=60) for the # MCP SDK's ClientSession read_timeout_seconds parameter. assert captured_kwargs.get("read_timeout_seconds") == timedelta(seconds=60), ( "ClientSession must receive read_timeout_seconds as a " "timedelta matching the config timeout" ) await conn.close() @pytest.mark.asyncio() async def test_connect_passes_none_timeout_to_client_session() -> None: """ When ``MCPServerConfig.timeout`` is ``None`` (default), ``connect()`` must pass ``read_timeout_seconds=None`` so the MCP SDK uses its built-in default. """ config = MCPServerConfig( name="test-no-timeout", url="http://localhost:9000/mcp", # timeout defaults to None ) captured_kwargs: dict[str, Any] = {} def _capturing_session( *args: Any, **kwargs: Any, ) -> AsyncMock: """ Fake ``ClientSession`` constructor that records kwargs. :param args: Positional args (read_stream, write_stream). :param kwargs: Keyword args including read_timeout_seconds. :returns: A mock session with working initialize/list_tools. """ captured_kwargs.update(kwargs) mock_session = AsyncMock() mock_session.initialize = AsyncMock() mock_tools = MagicMock() mock_tools.tools = [] mock_session.list_tools.return_value = mock_tools mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aexit__ = AsyncMock(return_value=False) return mock_session mock_ctx = AsyncMock() mock_ctx.__aenter__ = AsyncMock( return_value=(MagicMock(), MagicMock(), MagicMock()), ) mock_ctx.__aexit__ = AsyncMock(return_value=False) with patch( "omnigent.tools.mcp.ClientSession", side_effect=_capturing_session, ): with patch( "omnigent.tools.mcp.streamablehttp_client", return_value=mock_ctx, ): conn = McpServerConnection(config=config) await conn.connect() # When timeout is None, read_timeout_seconds must be None so the # MCP SDK falls back to its own default (no timeout). assert captured_kwargs.get("read_timeout_seconds") is None, ( "ClientSession must receive read_timeout_seconds=None when config timeout is unset" ) await conn.close() @pytest.mark.asyncio() async def test_connect_http_passes_timeout_to_transport() -> None: """ When ``MCPServerConfig(timeout=60)``, ``connect()`` must pass ``timeout=60.0`` and ``sse_read_timeout=60.0`` to the Streamable HTTP transport client. If the timeout is not forwarded, the transport uses its default (30s), and long-running MCP tool calls would time out prematurely. """ config = MCPServerConfig( name="test-http-timeout", url="http://localhost:9000/mcp", timeout=60, ) with _mock_http_transport() as captured: conn = McpServerConnection(config=config) await conn.connect() assert captured.transport_kwargs["timeout"] == 60.0, ( "transport timeout must equal the config timeout as float" ) assert captured.transport_kwargs["sse_read_timeout"] == 60.0, ( "transport sse_read_timeout must equal the config timeout as float" ) await conn.close() @pytest.mark.asyncio() async def test_connect_http_uses_default_timeouts_when_none() -> None: """ When ``MCPServerConfig(timeout=None)``, ``connect()`` must pass SDK defaults: ``timeout=30`` and ``sse_read_timeout=300`` to the Streamable HTTP transport. If the defaults are wrong, connections may hang indefinitely or fail immediately. """ config = MCPServerConfig( name="test-http-default", url="http://localhost:9000/mcp", # timeout defaults to None ) with _mock_http_transport() as captured: conn = McpServerConnection(config=config) await conn.connect() # Streamable HTTP default: 30s for request timeout. assert captured.transport_kwargs["timeout"] == 30, ( "transport timeout must default to 30 when config timeout is None" ) # SDK default: 300s (5 min) for SSE event read. assert captured.transport_kwargs["sse_read_timeout"] == 300, ( "transport sse_read_timeout must default to 300 when config timeout is None" ) # ── HTTP transport: connection, headers, discovery ──────── @dataclass class CapturedHttpArgs: """ Container for kwargs captured from ``streamablehttp_client`` and ``ClientSession`` during HTTP transport tests. :param transport_kwargs: Keyword arguments passed to ``streamablehttp_client``, e.g. ``{"url": "...", "headers": {...}, "timeout": 30}``. :param session_kwargs: Keyword arguments passed to ``ClientSession``, e.g. ``{"read_timeout_seconds": ...}``. :param mock_session: The mock ``ClientSession`` instance for setting up ``call_tool()`` side effects. """ transport_kwargs: dict[str, Any] = field(default_factory=dict) session_kwargs: dict[str, Any] = field(default_factory=dict) mock_session: AsyncMock = field(default_factory=AsyncMock) @contextmanager def _mock_http_transport( tools: list[MagicMock] | None = None, ) -> Iterator[CapturedHttpArgs]: """ Mock the Streamable HTTP transport and session for HTTP tests. Patches ``streamablehttp_client`` and ``ClientSession`` so that ``McpServerConnection.connect()`` can run without a real HTTP server. Captures the kwargs passed to both so tests can verify URL, headers, timeout, and other arguments. :param tools: Mock tool definitions to return from ``list_tools()``. Defaults to an empty list. :yields: A :class:`CapturedHttpArgs` with captured kwargs and the mock session. """ captured = CapturedHttpArgs() mock_streamable_ctx = AsyncMock() mock_streamable_ctx.__aenter__ = AsyncMock( return_value=(MagicMock(), MagicMock(), MagicMock()), ) mock_streamable_ctx.__aexit__ = AsyncMock(return_value=False) def _capturing_streamable_client(**kwargs: Any) -> AsyncMock: """ Fake ``streamablehttp_client`` that records kwargs. :param kwargs: Keyword args including url, headers, timeout, and sse_read_timeout. :returns: An async context manager yielding mock streams and a session-id getter. """ captured.transport_kwargs.update(kwargs) return mock_streamable_ctx def _capturing_session( *args: Any, **kwargs: Any, ) -> AsyncMock: """ Fake ``ClientSession`` constructor that records kwargs. :param args: Positional args (read_stream, write_stream). :param kwargs: Keyword args including read_timeout_seconds. :returns: A mock session with working initialize/list_tools. """ captured.session_kwargs.update(kwargs) mock_session = AsyncMock() mock_session.initialize = AsyncMock() mock_tools_result = MagicMock() mock_tools_result.tools = tools or [] mock_session.list_tools.return_value = mock_tools_result mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aexit__ = AsyncMock(return_value=False) captured.mock_session = mock_session return mock_session with patch( "omnigent.tools.mcp.streamablehttp_client", side_effect=_capturing_streamable_client, ): with patch( "omnigent.tools.mcp.ClientSession", side_effect=_capturing_session, ): yield captured @pytest.mark.asyncio() async def test_http_connect_passes_url_to_transport() -> None: """ HTTP ``connect()`` passes the config URL to the transport client. """ config = MCPServerConfig( name="test-http", url="https://mcp.example.com/mcp", ) with _mock_http_transport() as captured: conn = McpServerConnection(config=config) await conn.connect() assert captured.transport_kwargs["url"] == "https://mcp.example.com/mcp" await conn.close() @pytest.mark.asyncio() async def test_http_connect_passes_headers_to_transport() -> None: """ HTTP ``connect()`` propagates auth headers from the config to the transport client. """ config = MCPServerConfig( name="test-http-headers", url="http://localhost:9000/mcp", headers={ "Authorization": "Bearer tok_xyz", "X-Custom": "value", }, ) with _mock_http_transport() as captured: conn = McpServerConnection(config=config) await conn.connect() assert captured.transport_kwargs["headers"] == { "Authorization": "Bearer tok_xyz", "X-Custom": "value", } await conn.close() @pytest.mark.asyncio() async def test_http_connect_passes_none_headers_when_empty() -> None: """ HTTP ``connect()`` passes ``headers=None`` when the config has no headers, so the transport client uses its default. """ config = MCPServerConfig( name="test-http-no-headers", url="http://localhost:9000/mcp", # headers defaults to empty dict ) with _mock_http_transport() as captured: conn = McpServerConnection(config=config) await conn.connect() # Empty dict is converted to None via `or None`. assert captured.transport_kwargs["headers"] is None await conn.close() @pytest.mark.asyncio() async def test_http_falls_back_to_sse_when_streamable_fails() -> None: """ When ``streamablehttp_client`` raises (e.g. the server only speaks legacy SSE), ``_open_http_transport`` falls back to ``sse_client`` and the connection succeeds. If the fallback were removed, any legacy SSE-only MCP server would fail to connect. If ``streamablehttp_client`` is not tried first, Streamable HTTP servers (e.g. Databricks MCP gateways) would get the wrong transport. A non-``/sse`` URL is used on purpose: an ``…/sse`` URL is routed straight to the SSE client by ``_is_sse_endpoint`` and would bypass Streamable HTTP entirely, so it would not exercise the fallback this test guards. """ config = MCPServerConfig( name="test-sse-fallback", url="http://legacy-mcp.example.com/mcp", headers={"Authorization": "Bearer tok"}, ) captured_sse_kwargs: dict[str, Any] = {} mock_sse_ctx = AsyncMock() mock_sse_ctx.__aenter__ = AsyncMock( return_value=(MagicMock(), MagicMock()), ) mock_sse_ctx.__aexit__ = AsyncMock(return_value=False) def _capturing_sse(**kwargs: Any) -> AsyncMock: """ Fake ``sse_client`` that records kwargs. :param kwargs: Keyword args including url, headers. :returns: An async context manager yielding mock streams. """ captured_sse_kwargs.update(kwargs) return mock_sse_ctx mock_session = AsyncMock() mock_session.initialize = AsyncMock() mock_tools_result = MagicMock() mock_tools_result.tools = [_make_mcp_tool_def("legacy_tool")] mock_session.list_tools.return_value = mock_tools_result mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aexit__ = AsyncMock(return_value=False) with patch( "omnigent.tools.mcp.streamablehttp_client", side_effect=RuntimeError("server returned text/html, not application/json"), ) as mock_streamable: with patch( "omnigent.tools.mcp.sse_client", side_effect=_capturing_sse, ): with patch( "omnigent.tools.mcp.ClientSession", return_value=mock_session, ): conn = McpServerConnection(config=config) tools = await conn.connect() # Streamable HTTP was tried first and failed, so the fallback ran. assert mock_streamable.called, ( "Streamable HTTP must be tried first for a non-/sse URL; if it was " "skipped, the SSE fallback this test guards was never exercised" ) # Fallback reached sse_client with the correct URL and headers. assert captured_sse_kwargs["url"] == "http://legacy-mcp.example.com/mcp", ( "SSE fallback must receive the same URL as the original config" ) assert captured_sse_kwargs["headers"] == {"Authorization": "Bearer tok"}, ( "SSE fallback must forward headers from the config" ) # Tool discovery succeeded through the fallback path. assert len(tools) == 1, ( "SSE fallback must discover tools; if 0, the session was not initialized" ) assert tools[0].name == "legacy_tool" await conn.close() @pytest.mark.asyncio() async def test_http_connect_discovers_tools() -> None: """ HTTP ``connect()`` discovers tools via ``list_tools()`` and returns them. """ config = MCPServerConfig( name="test-http-discovery", url="http://localhost:9000/mcp", ) tool_def = _make_mcp_tool_def("http_tool") with _mock_http_transport([tool_def]) as captured: conn = McpServerConnection(config=config) tools = await conn.connect() assert len(tools) == 1 assert tools[0].name == "http_tool" captured.mock_session.initialize.assert_awaited_once() captured.mock_session.list_tools.assert_awaited_once() await conn.close() @pytest.mark.asyncio() async def test_http_call_tool_invokes_session() -> None: """ ``call_tool()`` on an HTTP connection delegates to the session's ``call_tool`` method and returns formatted results. """ config = MCPServerConfig( name="test-http-invoke", url="http://localhost:9000/mcp", ) with _mock_http_transport([_make_mcp_tool_def("http_tool")]) as captured: conn = McpServerConnection(config=config) await conn.connect() # Set up call_tool to return a mock result. mock_result = MagicMock() mock_result.content = [ TextContent(type="text", text="HTTP result"), ] mock_result.isError = False captured.mock_session.call_tool.return_value = mock_result result = await conn.call_tool("http_tool", {"query": "test"}) assert result == "HTTP result" captured.mock_session.call_tool.assert_awaited_once_with( name="http_tool", arguments={"query": "test"}, ) await conn.close() @pytest.mark.asyncio() async def test_http_reconnect_on_connection_error() -> None: """ HTTP ``call_tool()`` reconnects and retries on a connection error. """ config = MCPServerConfig( name="test-http-reconnect", url="http://localhost:9000/mcp", ) with _mock_http_transport([_make_mcp_tool_def("http_tool")]) as captured: conn = McpServerConnection(config=config) await conn.connect() ok_result = MagicMock() ok_result.content = [ TextContent(type="text", text="recovered via HTTP"), ] ok_result.isError = False captured.mock_session.call_tool.side_effect = [ ConnectionError("HTTP connection lost"), ok_result, ] with patch.object(conn, "_reconnect", new_callable=AsyncMock) as mock_reconnect: with patch("omnigent.tools.mcp._sleep", new_callable=AsyncMock): result = await conn.call_tool("http_tool", {"query": "retry"}) assert result == "recovered via HTTP" mock_reconnect.assert_awaited_once() await conn.close() @pytest.mark.asyncio() async def test_http_connect_uses_cache() -> None: """ HTTP ``connect()`` uses the discovery cache when fresh, skipping ``list_tools()`` while still opening a live session. """ config = MCPServerConfig( name="test-http-cached", url="http://localhost:9000/mcp", ) tool_def = _make_mcp_tool_def("cached_http_tool") _discovery_cache[_cache_key(config)] = [tool_def] with _mock_http_transport() as captured: conn = McpServerConnection(config=config) tools = await conn.connect() assert len(tools) == 1 assert tools[0].name == "cached_http_tool" # Session opened (for invocation), but list_tools skipped. captured.mock_session.initialize.assert_awaited_once() captured.mock_session.list_tools.assert_not_awaited() await conn.close() # ── Circuit breaker ────────────────────────────────────────── def test_circuit_breaker_allows_calls_when_closed() -> None: """ A fresh breaker in CLOSED state allows calls without raising. """ breaker = _CircuitBreaker(failure_threshold=3, cooldown_seconds=10.0) # Should not raise. breaker.pre_call("test-server") def test_circuit_breaker_trips_after_threshold_failures() -> None: """ The breaker trips after ``failure_threshold`` consecutive failures and blocks subsequent calls. """ breaker = _CircuitBreaker(failure_threshold=3, cooldown_seconds=60.0) for _ in range(3): breaker.record_failure("test-server") assert breaker.is_tripped is True with pytest.raises(McpServerDisabledError) as exc_info: breaker.pre_call("my-server") assert exc_info.value.server_name == "my-server" assert exc_info.value.consecutive_failures == 3 def test_circuit_breaker_does_not_trip_below_threshold() -> None: """ Fewer failures than the threshold do not trip the breaker. """ breaker = _CircuitBreaker(failure_threshold=5, cooldown_seconds=10.0) for _ in range(4): breaker.record_failure("test-server") assert breaker.is_tripped is False # Should not raise. breaker.pre_call("test-server") def test_circuit_breaker_resets_on_success() -> None: """ A successful call resets the failure counter and un-trips the breaker. """ breaker = _CircuitBreaker(failure_threshold=3, cooldown_seconds=60.0) for _ in range(3): breaker.record_failure("test-server") assert breaker.is_tripped is True breaker.record_success() assert breaker.is_tripped is False assert breaker.consecutive_failures == 0 # Should not raise after reset. breaker.pre_call("test-server") def test_circuit_breaker_half_open_after_cooldown( monkeypatch: pytest.MonkeyPatch, ) -> None: """ After the cooldown period elapses, the breaker enters half-open state and allows one probe call. """ import time as time_module breaker = _CircuitBreaker(failure_threshold=2, cooldown_seconds=10.0) breaker.record_failure("test-server") breaker.record_failure("test-server") assert breaker.is_tripped is True # Advance time past the cooldown. original_monotonic = time_module.monotonic monkeypatch.setattr( time_module, "monotonic", lambda: original_monotonic() + 15.0, ) # Cooldown elapsed — half-open state allows one probe. assert breaker.is_tripped is False breaker.pre_call("test-server") # Should not raise. def test_circuit_breaker_re_trips_on_half_open_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: """ If the half-open probe fails, the breaker re-trips. """ import time as time_module breaker = _CircuitBreaker(failure_threshold=2, cooldown_seconds=10.0) breaker.record_failure("test-server") breaker.record_failure("test-server") # Advance time past the cooldown. original_monotonic = time_module.monotonic monkeypatch.setattr( time_module, "monotonic", lambda: original_monotonic() + 15.0, ) # Half-open probe allowed. breaker.pre_call("test-server") # Probe fails — re-trip. breaker.record_failure("test-server") assert breaker.is_tripped is True def test_circuit_breaker_cooldown_remaining_in_error() -> None: """ The ``McpServerDisabledError`` includes the approximate cooldown remaining. """ breaker = _CircuitBreaker(failure_threshold=1, cooldown_seconds=30.0) breaker.record_failure("test-server") with pytest.raises(McpServerDisabledError) as exc_info: breaker.pre_call("test-server") # Cooldown just started, so remaining should be close to 30s. assert exc_info.value.cooldown_remaining > 25.0 def test_circuit_breaker_failure_count_resets_on_success() -> None: """ Interspersed successes prevent the breaker from tripping. """ breaker = _CircuitBreaker(failure_threshold=3, cooldown_seconds=10.0) breaker.record_failure("test-server") breaker.record_failure("test-server") # Success resets the counter. breaker.record_success() breaker.record_failure("test-server") breaker.record_failure("test-server") # Only 2 consecutive failures — not at threshold. assert breaker.is_tripped is False def test_circuit_breaker_trip_log_includes_server_name( caplog: pytest.LogCaptureFixture, ) -> None: """ When the breaker trips, the warning log message includes the server name so operators can identify which MCP server failed. :param caplog: Pytest fixture that captures log records. """ breaker = _CircuitBreaker(failure_threshold=2, cooldown_seconds=10.0) with caplog.at_level(logging.WARNING, logger="omnigent.tools.mcp"): breaker.record_failure("my-flaky-server") breaker.record_failure("my-flaky-server") assert breaker.is_tripped is True # The trip log must name the specific server. trip_messages = [r.message for r in caplog.records if "tripped" in r.message] assert len(trip_messages) == 1, f"Expected 1 trip log, got {len(trip_messages)}" assert "my-flaky-server" in trip_messages[0] def test_circuit_breaker_default_constants() -> None: """ Module-level circuit breaker constants have expected values. """ assert _CIRCUIT_BREAKER_THRESHOLD == 5 assert _CIRCUIT_BREAKER_COOLDOWN_SECONDS == 30.0 @pytest.mark.asyncio async def test_call_tool_trips_breaker_after_repeated_failures() -> None: """ ``McpServerConnection.call_tool()`` records failures in the circuit breaker. After ``failure_threshold`` exhausted invocations, subsequent calls raise ``McpServerDisabledError``. """ config = _make_http_config() with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) await conn.connect() # Override breaker threshold to 2 for a quick test. conn._breaker = _CircuitBreaker( failure_threshold=2, cooldown_seconds=60.0, ) mock_session.call_tool = AsyncMock(side_effect=EOFError("dead")) with patch.object(conn, "_reconnect", new_callable=AsyncMock): with patch("omnigent.tools.mcp._sleep", new_callable=AsyncMock): # First call: exhausts retries, records failure. with pytest.raises(EOFError): await conn.call_tool("my_tool", {"x": 1}) # Second call: exhausts retries, trips breaker. with pytest.raises(EOFError): await conn.call_tool("my_tool", {"x": 1}) # Third call: breaker is tripped, fails immediately. with pytest.raises(McpServerDisabledError) as exc_info: await conn.call_tool("my_tool", {"x": 1}) assert exc_info.value.server_name == config.name assert exc_info.value.consecutive_failures == 2 await conn.close() @pytest.mark.asyncio async def test_call_tool_resets_breaker_on_success() -> None: """ A successful ``call_tool()`` resets the circuit breaker so that prior failures don't accumulate across successes. """ config = _make_http_config() with _mock_mcp_transport() as mock_session: conn = McpServerConnection(config=config) await conn.connect() conn._breaker = _CircuitBreaker( failure_threshold=2, cooldown_seconds=60.0, ) ok_result = MagicMock() ok_result.content = [TextContent(type="text", text="ok")] ok_result.isError = False # 3 fails then 1 success (reconnect retries 3 per call). mock_session.call_tool = AsyncMock( side_effect=[ EOFError("dead"), EOFError("dead"), EOFError("dead"), EOFError("dead"), EOFError("dead"), ok_result, ] ) with patch.object(conn, "_reconnect", new_callable=AsyncMock): with patch("omnigent.tools.mcp._sleep", new_callable=AsyncMock): # First invocation: all 3 retries fail → failure (count=1). with pytest.raises(EOFError): await conn.call_tool("my_tool", {}) # Second invocation: 3rd retry succeeds → reset. result = await conn.call_tool("my_tool", {}) assert result == "ok" # Breaker should be reset — no accumulated failures. assert conn._breaker.consecutive_failures == 0 await conn.close() # ── Circuit breaker half-open atomic gate ───────────────── def test_circuit_breaker_half_open_clears_tripped_at( monkeypatch: pytest.MonkeyPatch, ) -> None: """ Entering half-open state clears ``_tripped_at`` so that concurrent callers see CLOSED (not half-open) and don't also enter the probe path. :param monkeypatch: Pytest monkeypatch for time.monotonic. """ import time as time_module breaker = _CircuitBreaker(failure_threshold=2, cooldown_seconds=10.0) breaker.record_failure("test-server") breaker.record_failure("test-server") assert breaker.is_tripped is True # Advance time past the cooldown. original_monotonic = time_module.monotonic monkeypatch.setattr( time_module, "monotonic", lambda: original_monotonic() + 15.0, ) # First pre_call enters half-open and clears the tripped state. breaker.pre_call("test-server") # Breaker should no longer report as tripped — a concurrent # caller sees CLOSED, not half-open. assert breaker.is_tripped is False # Second pre_call should not raise (sees CLOSED state). breaker.pre_call("test-server") # ── EventLoopThread ────────────────────────────────────── def _make_stdio_config( *, name: str = "stdio-server", command: str = "fake-mcp", args: list[str] | None = None, env: dict[str, str] | None = None, ) -> MCPServerConfig: """ Create a stdio MCP server config for transport-branch tests. :param name: Server name identifier. :param command: Executable command, e.g. ``"npx"``. :param args: Arguments to *command*, defaulting to ``[]``. :param env: Environment overlay, defaulting to ``{}``. :returns: An ``MCPServerConfig`` with ``transport="stdio"``. """ return MCPServerConfig( name=name, transport="stdio", command=command, args=list(args) if args is not None else [], env=dict(env) if env is not None else {}, ) def test_open_stdio_transport_spawns_unwrapped() -> None: """ The stdio branch spawns the MCP subprocess directly — the ``StdioServerParameters`` command equals ``config.command`` and the args equal ``config.args``, with no ``srt`` prefix inserted. Step 7 of the harness contract migration removed the ``MCPServerConfig.sandbox`` field and the ``wrap_with_srt`` call that gated this spawn — srt's default policy blocked outbound network and broke every useful MCP. This test pins the post-step-7 behavior: the unwrap is not optional. What breaks if this fails: a regression that re-introduces ``srt -c `` would silently hang every stdio MCP that needs outbound HTTPS (essentially all of them — Glean, Slack, GitHub, UC, etc.). """ config = _make_stdio_config( command="fake-mcp", args=["--flag", "value"], ) captured: dict[str, Any] = {} def _capture_stdio_client(params: Any) -> Any: """Record params, then yield a mock transport.""" captured["params"] = params mock_ctx = AsyncMock() mock_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) mock_ctx.__aexit__ = AsyncMock(return_value=False) return mock_ctx conn = McpServerConnection(config=config) with patch("omnigent.tools.mcp.stdio_client", side_effect=_capture_stdio_client): with patch( "omnigent.tools.mcp.ClientSession", return_value=_mock_session(), ): asyncio.run(conn.connect()) params = captured["params"] # command + args pass through unchanged — any wrap insertion # here means the srt path came back. assert params.command == "fake-mcp" assert params.args == ["--flag", "value"] def test_open_stdio_transport_overlays_env_on_parent() -> None: """ ``config.env`` is overlaid on ``os.environ`` so the spawned subprocess inherits ``PATH``, ``HOME``, etc., plus the MCP-specific overrides. What breaks if this fails: an MCP subprocess that needs ``GITHUB_TOKEN=ghp_xyz`` but also needs ``PATH`` to resolve its own binaries would lose PATH (getting only the caller's env overlay) and fail at spawn with "fake-mcp: command not found". """ config = _make_stdio_config( command="fake-mcp", env={"GITHUB_TOKEN": "ghp_xyz"}, ) captured: dict[str, Any] = {} def _capture_stdio_client(params: Any) -> Any: captured["params"] = params mock_ctx = AsyncMock() mock_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) mock_ctx.__aexit__ = AsyncMock(return_value=False) return mock_ctx conn = McpServerConnection(config=config) with patch("omnigent.tools.mcp.stdio_client", side_effect=_capture_stdio_client): with patch( "omnigent.tools.mcp.ClientSession", return_value=_mock_session(), ): asyncio.run(conn.connect()) env = captured["params"].env # Config overlay reached the subprocess assert env["GITHUB_TOKEN"] == "ghp_xyz" # Parent env merged — PATH is always in os.environ on any # realistic test runner, so its presence confirms the dict # union didn't wipe the inherited environment. assert "PATH" in env def test_open_stdio_transport_empty_env_inherits_fully() -> None: """ When ``config.env`` is empty, the transport passes ``env=None`` to :class:`StdioServerParameters` — which makes the MCP SDK inherit the parent environment wholesale. An explicit empty ``{}`` would wipe PATH etc. and break MCPs that don't set their own. What breaks if this fails: a common MCP YAML that omits ``env:`` (because the subprocess has no secret to inject) would spawn with an empty environment and fail at PATH resolution. """ config = _make_stdio_config(env={}) captured: dict[str, Any] = {} def _capture_stdio_client(params: Any) -> Any: captured["params"] = params mock_ctx = AsyncMock() mock_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) mock_ctx.__aexit__ = AsyncMock(return_value=False) return mock_ctx conn = McpServerConnection(config=config) with patch("omnigent.tools.mcp.stdio_client", side_effect=_capture_stdio_client): with patch( "omnigent.tools.mcp.ClientSession", return_value=_mock_session(), ): asyncio.run(conn.connect()) # None tells stdio_client to inherit the parent env as-is — # matches its documented behavior. assert captured["params"].env is None def _mock_session() -> AsyncMock: """ Build a mock ClientSession — minimum needed for the stdio transport tests to reach connect() without exercising the real MCP handshake. """ mock_session = AsyncMock() mock_tools_result = MagicMock() mock_tools_result.tools = [] mock_session.list_tools.return_value = mock_tools_result mock_session.initialize = AsyncMock() mock_session.__aenter__ = AsyncMock(return_value=mock_session) mock_session.__aexit__ = AsyncMock(return_value=False) return mock_session def _make_pool_mcp_tool(name: str) -> McpToolDef: """ Create a real MCP tool definition for session-pool tests. :param name: Tool name, e.g. ``"search_tool"``. :returns: A real :class:`mcp.types.Tool`. """ return McpToolDef( name=name, description=f"{name} description", inputSchema={ "type": "object", "properties": {}, }, ) # ── McpElicitationRequired & MRTR detection in _invoke_tool ──── def _make_input_required_result( input_requests: dict[str, Any] | None = None, request_state: str = "opaque_state", ) -> CallToolResult: """ Build a real ``CallToolResult`` with ``resultType: "input_required"`` via ``model_validate``, the same way the MCP SDK deserializes it over the wire. ``CallToolResult``'s parent ``Result`` has ``extra="allow"``, so extra fields like ``resultType``, ``inputRequests``, and ``requestState`` survive in ``model_extra``. :param input_requests: The ``inputRequests`` map keyed by server-assigned elicitation id. :param request_state: The opaque ``requestState`` string. :returns: A real ``CallToolResult`` with MRTR extras. """ if input_requests is None: input_requests = { "eid_1": { "method": "elicitation/create", "params": { "mode": "form", "message": "approve?", "requestedSchema": { "type": "object", "properties": { "decision": { "type": "string", "oneOf": [ {"const": "allow"}, {"const": "deny"}, ], }, }, }, }, }, } return CallToolResult.model_validate( { "content": [], "isError": False, "resultType": "input_required", "inputRequests": input_requests, "requestState": request_state, } ) def test_call_tool_result_model_extra_preserves_mrtr_fields() -> None: """ ``CallToolResult.model_validate`` preserves MRTR fields (``resultType``, ``inputRequests``, ``requestState``) in ``model_extra`` because the parent ``Result`` class has ``extra="allow"``. If this fails: the MCP SDK changed its Pydantic model config to ``extra="forbid"`` or ``extra="ignore"``, which would silently drop MRTR data and break the entire elicitation flow. """ result = _make_input_required_result() extras = result.model_extra or {} # resultType signals that the server needs user input. assert extras.get("resultType") == "input_required", ( "resultType must survive in model_extra; if missing, the MCP SDK " "is dropping extra fields and MRTR detection will silently fail" ) # inputRequests carries the elicitation payloads keyed by id. assert "eid_1" in extras.get("inputRequests", {}), ( "inputRequests must include the server's elicitation entries; " "if empty, the elicitation UI will show nothing to the user" ) # requestState must round-trip verbatim on retry. assert extras.get("requestState") == "opaque_state", ( "requestState must be preserved exactly; the server uses it " "to correlate the retry with the original request" ) @pytest.mark.asyncio async def test_invoke_tool_raises_elicitation_on_input_required() -> None: """ ``_invoke_tool`` raises ``McpElicitationRequired`` when the MCP session returns a ``CallToolResult`` with ``resultType == "input_required"`` in ``model_extra``. If this fails: the MRTR detection in ``_invoke_tool`` is broken and the runner will treat an ``InputRequiredResult`` as a normal (empty) tool result, silently skipping the elicitation flow. """ config = _make_http_config() mrtr_result = _make_input_required_result( input_requests={ "eid_abc": { "method": "elicitation/create", "params": {"mode": "form", "message": "confirm deploy?"}, }, }, request_state="state_xyz", ) with _mock_mcp_transport() as mock_session: # Stub session.call_tool to return the MRTR result. mock_session.call_tool.return_value = mrtr_result conn = McpServerConnection(config=config) await conn.connect() with pytest.raises(McpElicitationRequired) as exc_info: await conn._invoke_tool("deploy_tool", {"env": "prod"}) exc = exc_info.value # input_requests carries the full elicitation payloads. assert "eid_abc" in exc.input_requests, ( "input_requests must include the elicitation id from the server; " "if missing, the Omnigent server can't surface the elicitation to the user" ) # request_state must be echoed back verbatim on retry. assert exc.request_state == "state_xyz", ( "request_state must match the server's opaque value; " "if wrong, the retry will be rejected by the server" ) # tool_name and arguments are preserved for the retry call. assert exc.tool_name == "deploy_tool", ( "tool_name must be preserved so the retry knows which tool to call" ) assert exc.arguments == {"env": "prod"}, ( "arguments must be preserved so the retry uses the same parameters" ) await conn.close() @pytest.mark.asyncio async def test_invoke_tool_returns_normally_without_mrtr() -> None: """ ``_invoke_tool`` returns the formatted result string when ``model_extra`` does NOT contain ``resultType == "input_required"``. If this fails: normal (non-MRTR) tool calls are broken — the function is raising ``McpElicitationRequired`` when it shouldn't. """ config = _make_http_config() # Normal result — no extra fields triggering MRTR. normal_result = CallToolResult.model_validate( { "content": [{"type": "text", "text": "tool output here"}], "isError": False, } ) with _mock_mcp_transport() as mock_session: mock_session.call_tool.return_value = normal_result conn = McpServerConnection(config=config) await conn.connect() result = await conn._invoke_tool("normal_tool", {"x": 1}) # Normal path: formatted text returned, no exception. assert result == "tool output here", ( "Normal tool results must be returned as formatted text; " "if McpElicitationRequired was raised instead, the MRTR " "detection has a false positive" ) await conn.close() @pytest.mark.asyncio async def test_call_tool_with_elicitation_returns_result() -> None: """ ``call_tool_with_elicitation`` sends a retry with ``inputResponses`` and ``requestState`` and returns the formatted result when the server responds normally. If this fails: the elicitation retry path is broken — the runner can't complete the tool call after the user approves. """ config = _make_http_config() # The retry returns a normal result. retry_result = CallToolResult.model_validate( { "content": [{"type": "text", "text": "deploy succeeded"}], "isError": False, } ) with _mock_mcp_transport() as mock_session: mock_session.send_request = AsyncMock(return_value=retry_result) conn = McpServerConnection(config=config) await conn.connect() result = await conn.call_tool_with_elicitation( name="deploy_tool", arguments={"env": "prod"}, input_responses={"eid_1": {"action": "accept", "content": {"decision": "allow"}}}, request_state="state_xyz", ) # The retry result is formatted and returned. assert result == "deploy succeeded", ( "Elicitation retry must return the formatted tool result; " "if empty or wrong, the retry path lost the server's response" ) await conn.close() @pytest.mark.asyncio async def test_call_tool_with_elicitation_raises_on_second_mrtr() -> None: """ ``call_tool_with_elicitation`` raises ``McpElicitationRequired`` when the retry itself returns another ``InputRequiredResult`` (multi-round MRTR). If this fails: multi-round elicitation is broken — the second elicitation round silently returns an empty result instead of surfacing the next approval request. """ config = _make_http_config() # The retry returns yet another InputRequiredResult. second_mrtr = _make_input_required_result( input_requests={ "eid_2": {"method": "elicitation/create", "params": {"message": "round 2"}} }, request_state="state_round2", ) with _mock_mcp_transport() as mock_session: mock_session.send_request = AsyncMock(return_value=second_mrtr) conn = McpServerConnection(config=config) await conn.connect() with pytest.raises(McpElicitationRequired) as exc_info: await conn.call_tool_with_elicitation( name="deploy_tool", arguments={"env": "prod"}, input_responses={"eid_1": {"action": "accept", "content": {}}}, request_state="state_xyz", ) exc = exc_info.value # Second round's elicitation data must surface. assert "eid_2" in exc.input_requests, ( "Multi-round MRTR must surface the second elicitation's requests; " "if missing, the Omnigent server can't show the next approval form" ) assert exc.request_state == "state_round2", ( "Multi-round MRTR must carry the new requestState for the next retry" ) await conn.close() # ── _is_sse_endpoint / legacy-SSE routing ──────────────── def test_is_sse_endpoint_detects_sse_paths() -> None: """ URLs whose path ends in a ``/sse`` segment are legacy-SSE endpoints (e.g. crawl4ai's ``/mcp/sse``). The Streamable HTTP client hangs in teardown against such a server, so the transport router must detect these and use the SSE client directly. """ from omnigent.tools.mcp import _is_sse_endpoint assert _is_sse_endpoint("http://h:1/mcp/sse") assert _is_sse_endpoint("http://h:1/mcp/sse/") # trailing slash assert _is_sse_endpoint("http://h:1/sse") assert not _is_sse_endpoint("http://h:1/mcp") assert not _is_sse_endpoint("http://h:1/") assert not _is_sse_endpoint("http://h:1") assert not _is_sse_endpoint("http://h:1/mcp/sse-events") # not a /sse segment @pytest.mark.asyncio() async def test_open_http_transport_routes_sse_url_straight_to_sse() -> None: """ An ``…/sse`` URL goes directly to the SSE transport. The Streamable HTTP client hangs in teardown against an SSE-only server (crawl4ai), so it must be SKIPPED — not merely tried-then-fallen-back-from, because the hang prevents the fallback from ever running. """ from contextlib import AsyncExitStack conn = McpServerConnection(config=MCPServerConfig(name="c", url="http://h:1/mcp/sse")) calls: list[str] = [] async def fake_sse(stack, timeout, headers): calls.append("sse") return ("r", "w") async def fake_streamable(stack, timeout, headers): calls.append("streamable") return ("r", "w") conn._open_sse_transport = fake_sse # type: ignore[method-assign] conn._open_streamable_http_transport = fake_streamable # type: ignore[method-assign] async with AsyncExitStack() as stack: await conn._open_http_transport(stack) assert calls == ["sse"], "…/sse URL must skip Streamable HTTP entirely" @pytest.mark.asyncio() async def test_open_http_transport_uses_streamable_for_non_sse_url() -> None: """ A plain HTTP MCP URL still tries Streamable HTTP first (with the existing SSE fallback on failure) — the routing change must not regress modern Streamable-HTTP servers. """ from contextlib import AsyncExitStack conn = McpServerConnection(config=MCPServerConfig(name="c", url="http://h:1/mcp")) calls: list[str] = [] async def fake_sse(stack, timeout, headers): calls.append("sse") return ("r", "w") async def fake_streamable(stack, timeout, headers): calls.append("streamable") return ("r", "w") conn._open_sse_transport = fake_sse # type: ignore[method-assign] conn._open_streamable_http_transport = fake_streamable # type: ignore[method-assign] async with AsyncExitStack() as stack: await conn._open_http_transport(stack) assert calls == ["streamable"], "plain URL must try Streamable HTTP first"