From 635666dec975548ee760131569384b629595165a Mon Sep 17 00:00:00 2001
From: cornmander <shikhman@google.com>
Date: Fri, 23 May 2025 17:19:30 -0400
Subject: [PATCH] MCP SSE support (#511)

Matches the config format used by other MCP clients.
---
 packages/server/src/config/config.ts       |  6 ++-
 packages/server/src/tools/tool-registry.ts | 47 +++++++++++++++-------
 2 files changed, 37 insertions(+), 16 deletions(-)

diff --git a/packages/server/src/config/config.ts b/packages/server/src/config/config.ts
index 0a9bc2bf..625b9b93 100644
--- a/packages/server/src/config/config.ts
+++ b/packages/server/src/config/config.ts
@@ -24,10 +24,14 @@ import { WebSearchTool } from '../tools/web-search.js';
 
 export class MCPServerConfig {
   constructor(
-    readonly command: string,
+    // For stdio transport
+    readonly command?: string,
     readonly args?: string[],
     readonly env?: Record<string, string>,
     readonly cwd?: string,
+    // For sse transport
+    readonly url?: string,
+    // Common
     readonly timeout?: number,
   ) {}
 }
diff --git a/packages/server/src/tools/tool-registry.ts b/packages/server/src/tools/tool-registry.ts
index ca59d666..7b75e0f2 100644
--- a/packages/server/src/tools/tool-registry.ts
+++ b/packages/server/src/tools/tool-registry.ts
@@ -12,6 +12,7 @@ import { spawn, execSync } from 'node:child_process';
 // TODO: remove this dependency once MCP support is built into genai SDK
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
+import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
 type ToolParams = Record<string, unknown>;
 
 const MCP_TOOL_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
@@ -206,14 +207,28 @@ export class ToolRegistry {
           name: 'mcp-client',
           version: '0.0.1',
         });
-        const transport = new StdioClientTransport({
-          ...mcpServerConfig,
-          env: {
-            ...process.env,
-            ...(mcpServerConfig.env || {}),
-          } as Record<string, string>,
-          stderr: 'pipe',
-        });
+        let transport;
+        if (mcpServerConfig.url) {
+          // SSE transport if URL is provided
+          transport = new SSEClientTransport(new URL(mcpServerConfig.url));
+        } else if (mcpServerConfig.command) {
+          // Stdio transport if command is provided
+          transport = new StdioClientTransport({
+            command: mcpServerConfig.command,
+            args: mcpServerConfig.args || [],
+            env: {
+              ...process.env,
+              ...(mcpServerConfig.env || {}),
+            } as Record<string, string>,
+            cwd: mcpServerConfig.cwd,
+            stderr: 'pipe',
+          });
+        } else {
+          console.error(
+            `MCP server '${mcpServerName}' has invalid configuration: missing both url (for SSE) and command (for stdio). Skipping.`,
+          );
+          return;
+        }
         try {
           await mcpClient.connect(transport);
         } catch (error) {
@@ -227,15 +242,17 @@ export class ToolRegistry {
         mcpClient.onerror = (error) => {
           console.error('MCP ERROR', error.toString());
         };
-        if (!transport.stderr) {
+        if (transport instanceof StdioClientTransport && !transport.stderr) {
           throw new Error('transport missing stderr stream');
         }
-        transport.stderr.on('data', (data) => {
-          // filter out INFO messages logged for each request received
-          if (!data.toString().includes('] INFO')) {
-            console.debug('MCP STDERR', data.toString());
-          }
-        });
+        if (transport instanceof StdioClientTransport) {
+          transport.stderr!.on('data', (data) => {
+            // filter out INFO messages logged for each request received
+            if (!data.toString().includes('] INFO')) {
+              console.debug('MCP STDERR', data.toString());
+            }
+          });
+        }
         const result = await mcpClient.listTools();
         for (const tool of result.tools) {
           // Recursively remove additionalProperties and $schema from the inputSchema