diff --git a/api/core/mcp/session/client_session.py b/api/core/mcp/session/client_session.py index 3d5dac4b483c38..85fdc8933daa7b 100644 --- a/api/core/mcp/session/client_session.py +++ b/api/core/mcp/session/client_session.py @@ -188,23 +188,25 @@ def set_logging_level(self, level: types.LoggingLevel) -> types.EmptyResult: types.EmptyResult, ) - def list_resources(self) -> types.ListResourcesResult: + def list_resources(self, cursor: str | None = None) -> types.ListResourcesResult: """Send a resources/list request.""" return self.send_request( types.ClientRequest( types.ListResourcesRequest( method="resources/list", + params=types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None, ) ), types.ListResourcesResult, ) - def list_resource_templates(self) -> types.ListResourceTemplatesResult: + def list_resource_templates(self, cursor: str | None = None) -> types.ListResourceTemplatesResult: """Send a resources/templates/list request.""" return self.send_request( types.ClientRequest( types.ListResourceTemplatesRequest( method="resources/templates/list", + params=types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None, ) ), types.ListResourceTemplatesResult, @@ -265,12 +267,13 @@ def call_tool( request_read_timeout_seconds=read_timeout_seconds, ) - def list_prompts(self) -> types.ListPromptsResult: + def list_prompts(self, cursor: str | None = None) -> types.ListPromptsResult: """Send a prompts/list request.""" return self.send_request( types.ClientRequest( types.ListPromptsRequest( method="prompts/list", + params=types.PaginatedRequestParams(cursor=cursor) if cursor is not None else None, ) ), types.ListPromptsResult, diff --git a/api/tests/unit_tests/core/mcp/session/test_client_session.py b/api/tests/unit_tests/core/mcp/session/test_client_session.py index ab990a98e42686..a0cd78e4c1c3ba 100644 --- a/api/tests/unit_tests/core/mcp/session/test_client_session.py +++ b/api/tests/unit_tests/core/mcp/session/test_client_session.py @@ -914,6 +914,68 @@ def mock_server(): assert result.tools == [] t.join(timeout=1) + @staticmethod + @pytest.mark.parametrize( + ("method_name", "wire_method", "result_key"), + [ + ("list_tools", "tools/list", "tools"), + ("list_resources", "resources/list", "resources"), + ("list_resource_templates", "resources/templates/list", "resourceTemplates"), + ("list_prompts", "prompts/list", "prompts"), + ], + ) + @pytest.mark.parametrize("cursor", ["opaque-cursor-1", ""]) + def test_list_methods_send_cursor(streams, method_name, wire_method, result_key, cursor: str): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == wire_method + assert msg.message.root.params == {"cursor": cursor} + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={result_key: [], "nextCursor": "cursor-2"}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + method = { + "list_tools": session.list_tools, + "list_resources": session.list_resources, + "list_resource_templates": session.list_resource_templates, + "list_prompts": session.list_prompts, + }[method_name] + result = method(cursor=cursor) + assert result.nextCursor == "cursor-2" + t.join(timeout=1) + + @staticmethod + def test_list_tools_omits_params_without_cursor(streams): + read_stream, write_stream = streams + session = ClientSession(read_stream, write_stream) + + def mock_server(): + msg = write_stream.get(timeout=2) + assert msg.message.root.method == "tools/list" + assert msg.message.root.params is None + req_id = msg.message.root.id + resp = types.JSONRPCResponse(jsonrpc="2.0", id=req_id, result={"tools": []}) + read_stream.put(_client_session_SessionMessage(message=types.JSONRPCMessage(resp))) + + import threading + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + with session: + result = session.list_tools() + assert result.tools == [] + t.join(timeout=1) + @staticmethod def test_send_roots_list_changed(streams): read_stream, write_stream = streams diff --git a/api/tests/unit_tests/core/mcp/test_mcp_client.py b/api/tests/unit_tests/core/mcp/test_mcp_client.py index 6176feb13936fc..da9f392207ea06 100644 --- a/api/tests/unit_tests/core/mcp/test_mcp_client.py +++ b/api/tests/unit_tests/core/mcp/test_mcp_client.py @@ -417,6 +417,115 @@ def test_list_tools_success(self): assert result == expected_tools mock_session.list_tools.assert_called_once() + def _client_with_paginated_server(self, pages: list[dict], fail_on_extra_page: bool = False): + """Client whose session talks to an in-memory server answering tools/list from pages. + + Each entry in pages is the raw result dict for one request. The server + records the cursor of every request it sees. With fail_on_extra_page, + any request past the last page gets an error response so a runaway + pagination loop fails fast instead of hanging the test. + """ + import queue + import threading + + from core.mcp.session.base_session import SessionMessage + from core.mcp.session.client_session import ClientSession + from core.mcp.types import JSONRPCError, JSONRPCMessage, JSONRPCResponse + + read_stream, write_stream = queue.Queue(), queue.Queue() + session = ClientSession(read_stream, write_stream) + seen_cursors: list[str | None] = [] + + def mock_server(): + for page in pages: + msg = write_stream.get(timeout=5) + if msg is None: + return + params = msg.message.root.params or {} + seen_cursors.append(params.get("cursor")) + resp = JSONRPCResponse(jsonrpc="2.0", id=msg.message.root.id, result=page) + read_stream.put(SessionMessage(message=JSONRPCMessage(resp))) + if not fail_on_extra_page: + return + try: + msg = write_stream.get(timeout=2) + except queue.Empty: + return + if msg is None: + return + params = msg.message.root.params or {} + seen_cursors.append(params.get("cursor")) + err = JSONRPCError( + jsonrpc="2.0", + id=msg.message.root.id, + error={"code": -32000, "message": "client asked for a page the server never offered"}, + ) + read_stream.put(SessionMessage(message=JSONRPCMessage(err))) + + t = threading.Thread(target=mock_server, daemon=True) + t.start() + + client = MCPClient(server_url="http://test.example.com") + client._session = session + return client, session, t, seen_cursors + + @pytest.mark.parametrize("cursor", ["cursor-1", ""]) + def test_list_tools_paginates_until_cursor_exhausted(self, cursor: str): + """Follow opaque nextCursor values, including the empty string.""" + pages = [ + { + "tools": [{"name": "alpha", "inputSchema": {"type": "object", "properties": {}}}], + "nextCursor": cursor, + }, + {"tools": [{"name": "beta", "inputSchema": {"type": "object", "properties": {}}}]}, + ] + client, session, t, seen_cursors = self._client_with_paginated_server(pages) + + with session: + result = client.list_tools() + t.join(timeout=5) + + assert [tool.name for tool in result] == ["alpha", "beta"] + assert seen_cursors == [None, cursor] + + @pytest.mark.parametrize( + "cursors", + [["cursor-1", "cursor-1"], ["cursor-1", "cursor-2", "cursor-1"], ["", ""]], + ids=["repeated", "cycle", "empty-repeated"], + ) + def test_list_tools_raises_on_repeated_cursor(self, cursors: list[str]): + """Reject incomplete tool directories without requesting a page twice.""" + pages = [ + { + "tools": [{"name": f"tool-{index}", "inputSchema": {"type": "object", "properties": {}}}], + "nextCursor": cursor, + } + for index, cursor in enumerate(cursors) + ] + client, session, t, seen_cursors = self._client_with_paginated_server(pages, fail_on_extra_page=True) + + with session: + with pytest.raises(MCPConnectionError, match="repeated pagination cursor"): + client.list_tools() + t.join(timeout=5) + + assert seen_cursors == [None, *cursors[:-1]] + + def test_list_tools_empty_first_page_with_cursor(self): + """Test list_tools keeps paging when an early page returns no tools.""" + pages = [ + {"tools": [], "nextCursor": "cursor-1"}, + {"tools": [{"name": "late-tool", "inputSchema": {"type": "object", "properties": {}}}]}, + ] + client, session, t, seen_cursors = self._client_with_paginated_server(pages) + + with session: + result = client.list_tools() + t.join(timeout=5) + + assert [tool.name for tool in result] == ["late-tool"] + assert seen_cursors == [None, "cursor-1"] + def test_invoke_tool_not_initialized(self): """Test invoke_tool when session not initialized.""" client = MCPClient(server_url="http://test.example.com") diff --git a/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py b/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py index 95ef5320ed023f..dcc1367105c621 100644 --- a/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py +++ b/api/tests/unit_tests/services/tools/test_mcp_tools_manage_service.py @@ -1,16 +1,25 @@ -"""Exercise tenant-scoped MCP provider lookups against persisted rows. +"""Exercise tenant-scoped MCP provider lookups and paginated tool refreshes. Primary keys and server identifiers are distinct, even when an identifier looks like a UUID or collides with another provider's primary key. """ +import json +from datetime import datetime from typing import Literal +from uuid import uuid4 import pytest +from sqlalchemy import event from sqlalchemy.orm import Session +from core.entities.mcp_provider import MCPProviderEntity +from core.mcp.auth_client import MCPClientWithAuthRetry +from core.mcp.types import ErrorData, ListToolsResult, Tool +from models.account import Account from models.tools import MCPToolProvider from services.tools.mcp_tools_manage_service import MCPToolManageService +from tests.unit_tests.core.mcp.test_mcp_client import client_for_pages PROVIDER_UUID = "0b2bd1e4-3a7d-4a0b-9f2a-9b2d6f3c1e55" OTHER_PROVIDER_UUID = "00000000-0000-4000-8000-000000000002" @@ -158,3 +167,126 @@ def test_server_identifier_can_be_reused_across_tenants(service: MCPToolManageSe service.get_provider_by_server_identifier(server_identifier=SERVER_IDENTIFIER, tenant_id=TENANT_ID).id == provider.id ) + + +@pytest.fixture +def provider(sqlite_session: Session) -> MCPToolProvider: + owner = Account(name="MCP owner", email="mcp-owner@example.com") + sqlite_session.add(owner) + provider = MCPToolProvider( + name="Test server", + server_identifier="test-server", + server_url="https://example.com/mcp", + server_url_hash="hash", + icon="", + tenant_id=str(uuid4()), + user_id=owner.id, + authed=True, + tools=json.dumps([{"name": "saved-tool", "inputSchema": {}}]), + ) + provider.created_at = datetime(2024, 1, 1) + provider.updated_at = datetime(2024, 1, 1) + sqlite_session.add(provider) + sqlite_session.flush() + return provider + + +@pytest.fixture +def flushes(sqlite_session: Session, provider: MCPToolProvider) -> list[None]: + assert provider in sqlite_session + flushes: list[None] = [] + event.listen(sqlite_session, "after_flush", lambda *_args: flushes.append(None)) + return flushes + + +@pytest.fixture +def client(monkeypatch: pytest.MonkeyPatch) -> MCPClientWithAuthRetry: + client = MCPClientWithAuthRetry(server_url="https://example.com/mcp") + # Keep the real auth wrapper and session; in-memory JSON-RPC replaces transport setup. + monkeypatch.setattr(client, "_initialize", lambda: None) + monkeypatch.setattr("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry", lambda **_kwargs: client) + monkeypatch.setattr(MCPProviderEntity, "decrypt_server_url", lambda _self: "https://example.com/mcp") + return client + + +def _page(name: str, cursor: str | None) -> ListToolsResult: + return ListToolsResult(tools=[Tool(name=name, inputSchema={})], nextCursor=cursor) + + +@pytest.mark.parametrize("cursor", ["next", ""]) +def test_list_provider_tools_saves_all_pages( + sqlite_session: Session, + provider: MCPToolProvider, + flushes: list[None], + client: MCPClientWithAuthRetry, + cursor: str, +) -> None: + service = MCPToolManageService(sqlite_session) + with client_for_pages({None: _page("alpha", cursor), cursor: _page("beta", None)}, client) as (_, requests): + service.list_provider_tools(tenant_id=provider.tenant_id, provider_id=provider.id) + + assert [tool["name"] for tool in json.loads(provider.tools)] == ["alpha", "beta"] + assert all(tool["description"] == "" for tool in json.loads(provider.tools)) + assert provider.authed is True + assert len(flushes) == 1 + assert [request.params.cursor if request.params is not None else None for request in requests] == [None, cursor] + sqlite_session.expire(provider) + assert [tool["name"] for tool in json.loads(provider.tools)] == ["alpha", "beta"] + + +@pytest.mark.parametrize( + "cursors", + [["next", "next"], ["a", "b", "a"], ["", ""]], + ids=["repeated", "cycle", "empty-repeated"], +) +def test_list_provider_tools_preserves_saved_tools_on_cursor_cycle( + sqlite_session: Session, + provider: MCPToolProvider, + flushes: list[None], + client: MCPClientWithAuthRetry, + cursors: list[str], +) -> None: + original_tools, original_updated_at = provider.tools, provider.updated_at + pages = { + None if index == 0 else cursors[index - 1]: _page(f"tool-{index}", cursor) + for index, cursor in enumerate(cursors) + } + service = MCPToolManageService(sqlite_session) + with client_for_pages(pages, client) as (_, requests): + with pytest.raises(ValueError, match="Failed to connect to MCP server: .*repeated pagination cursor"): + service.list_provider_tools(tenant_id=provider.tenant_id, provider_id=provider.id) + + assert provider.tools == original_tools + assert provider.updated_at == original_updated_at + assert provider.authed is True + assert not flushes + assert [request.params.cursor if request.params is not None else None for request in requests] == [ + None, + *cursors[:-1], + ] + sqlite_session.expire(provider) + assert provider.tools == original_tools + assert provider.updated_at == original_updated_at + + +def test_list_provider_tools_preserves_saved_tools_on_later_page_error( + sqlite_session: Session, + provider: MCPToolProvider, + flushes: list[None], + client: MCPClientWithAuthRetry, +) -> None: + original_tools, original_updated_at = provider.tools, provider.updated_at + pages = {None: _page("alpha", "next"), "next": ErrorData(code=500, message="later page failed")} + service = MCPToolManageService(sqlite_session) + with client_for_pages(pages, client) as (_, requests): + with pytest.raises(ValueError, match="Failed to connect to MCP server: .*later page failed"): + service.list_provider_tools(tenant_id=provider.tenant_id, provider_id=provider.id) + + assert provider.tools == original_tools + assert provider.updated_at == original_updated_at + assert provider.authed is True + assert not flushes + assert [request.params.cursor if request.params is not None else None for request in requests] == [None, "next"] + sqlite_session.expire(provider) + assert provider.tools == original_tools + assert provider.updated_at == original_updated_at