From de1658cd88a05b94e246b826e6428e38b52253c7 Mon Sep 17 00:00:00 2001 From: Manoj Aggarwal Date: Tue, 13 Jan 2026 14:06:40 -0800 Subject: [PATCH] fix: resolve alembic migration conflicts by using unique revision ID for MCP migration --- ..._allow_multiple_connectors_with_unique_.py | 4 +- ...=> a1b2c3d4e5f6_add_mcp_connector_type.py} | 8 ++-- .../app/agents/new_chat/tools/mcp_client.py | 44 ++++++++++--------- 3 files changed, 29 insertions(+), 27 deletions(-) rename surfsense_backend/alembic/versions/{60_add_mcp_connector_type.py => a1b2c3d4e5f6_add_mcp_connector_type.py} (89%) diff --git a/surfsense_backend/alembic/versions/5263aa4e7f94_allow_multiple_connectors_with_unique_.py b/surfsense_backend/alembic/versions/5263aa4e7f94_allow_multiple_connectors_with_unique_.py index 5a61af0cb..de9505e3a 100644 --- a/surfsense_backend/alembic/versions/5263aa4e7f94_allow_multiple_connectors_with_unique_.py +++ b/surfsense_backend/alembic/versions/5263aa4e7f94_allow_multiple_connectors_with_unique_.py @@ -1,7 +1,7 @@ """allow_multiple_connectors_with_unique_names Revision ID: 5263aa4e7f94 -Revises: ffd7445eb90a +Revises: a1b2c3d4e5f6 Create Date: 2026-01-13 12:23:31.481643 """ @@ -11,7 +11,7 @@ from alembic import op # revision identifiers, used by Alembic. revision: str = '5263aa4e7f94' -down_revision: str | None = 'ffd7445eb90a' +down_revision: str | None = 'a1b2c3d4e5f6' branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None diff --git a/surfsense_backend/alembic/versions/60_add_mcp_connector_type.py b/surfsense_backend/alembic/versions/a1b2c3d4e5f6_add_mcp_connector_type.py similarity index 89% rename from surfsense_backend/alembic/versions/60_add_mcp_connector_type.py rename to surfsense_backend/alembic/versions/a1b2c3d4e5f6_add_mcp_connector_type.py index f3ec2611c..2ec158981 100644 --- a/surfsense_backend/alembic/versions/60_add_mcp_connector_type.py +++ b/surfsense_backend/alembic/versions/a1b2c3d4e5f6_add_mcp_connector_type.py @@ -1,7 +1,7 @@ """Add MCP connector type -Revision ID: 60 -Revises: 59 +Revision ID: a1b2c3d4e5f6 +Revises: 60 Create Date: 2026-01-09 15:19:51.827647 """ @@ -10,8 +10,8 @@ from collections.abc import Sequence from alembic import op # revision identifiers, used by Alembic. -revision: str = '60' -down_revision: str | None = '59' +revision: str = 'a1b2c3d4e5f6' +down_revision: str | None = '60' branch_labels: str | Sequence[str] | None = None depends_on: str | Sequence[str] | None = None diff --git a/surfsense_backend/app/agents/new_chat/tools/mcp_client.py b/surfsense_backend/app/agents/new_chat/tools/mcp_client.py index 181cbb642..d91065661 100644 --- a/surfsense_backend/app/agents/new_chat/tools/mcp_client.py +++ b/surfsense_backend/app/agents/new_chat/tools/mcp_client.py @@ -53,24 +53,26 @@ class MCPClient: ) # Spawn server process and create session - async with ( - stdio_client(server=server_params) as (read, write), - ClientSession(read, write) as session, - ): - # Initialize the connection - await session.initialize() - self.session = session - logger.info( - f"Connected to MCP server: {self.command} {' '.join(self.args)}" - ) - yield session + # Note: Cannot combine these context managers because ClientSession + # needs the read/write streams from stdio_client + async with stdio_client(server=server_params) as (read, write): + async with ClientSession(read, write) as session: + # Initialize the connection + await session.initialize() + self.session = session + logger.info( + "Connected to MCP server: %s %s", + self.command, + " ".join(self.args), + ) + yield session except Exception as e: - logger.error(f"Failed to connect to MCP server: {e!s}", exc_info=True) + logger.error("Failed to connect to MCP server: %s", e, exc_info=True) raise finally: self.session = None - logger.info(f"Disconnected from MCP server: {self.command}") + logger.info("Disconnected from MCP server: %s", self.command) async def list_tools(self) -> list[dict[str, Any]]: """List all tools available from the MCP server. @@ -97,11 +99,11 @@ class MCPClient: "input_schema": tool.inputSchema if hasattr(tool, "inputSchema") else {}, }) - logger.info(f"Listed {len(tools)} tools from MCP server") + logger.info("Listed %d tools from MCP server", len(tools)) return tools except Exception as e: - logger.error(f"Failed to list tools from MCP server: {e!s}", exc_info=True) + logger.error("Failed to list tools from MCP server: %s", e, exc_info=True) raise async def call_tool(self, tool_name: str, arguments: dict[str, Any]) -> Any: @@ -122,7 +124,7 @@ class MCPClient: raise RuntimeError("Not connected to MCP server. Use 'async with client.connect():'") try: - logger.info(f"Calling MCP tool '{tool_name}' with arguments: {arguments}") + logger.info("Calling MCP tool '%s' with arguments: %s", tool_name, arguments) # Call tools/call RPC method response = await self.session.call_tool(tool_name, arguments=arguments) @@ -138,19 +140,19 @@ class MCPClient: result.append(str(content)) result_str = "\n".join(result) if result else "" - logger.info(f"MCP tool '{tool_name}' succeeded: {result_str[:200]}") + logger.info("MCP tool '%s' succeeded: %s", tool_name, result_str[:200]) return result_str except RuntimeError as e: # Handle validation errors from MCP server responses # Some MCP servers (like server-memory) return extra fields not in their schema if "Invalid structured content" in str(e): - logger.warning(f"MCP server returned data not matching its schema, but continuing: {e}") + logger.warning("MCP server returned data not matching its schema, but continuing: %s", e) # Try to extract result from error message or return a success message return "Operation completed (server returned unexpected format)" raise - except Exception as e: - logger.error(f"Failed to call MCP tool '{tool_name}': {e!s}", exc_info=True) + except (ValueError, TypeError, AttributeError, KeyError) as e: + logger.error("Failed to call MCP tool '%s': %s", tool_name, e, exc_info=True) return f"Error calling tool: {e!s}" @@ -178,7 +180,7 @@ async def test_mcp_connection( "message": f"Connected successfully. Found {len(tools)} tools.", "tools": tools, } - except Exception as e: + except (RuntimeError, ConnectionError, TimeoutError, OSError) as e: return { "status": "error", "message": f"Failed to connect: {e!s}",