Mcp optimization

packages/api/src/chat/mcpChatSession.ts

@@ -0,0 +1,165 @@
+import type * as mcpSDKClientLibraryTypes from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import type * as mcpSDKClientTypes from '@modelcontextprotocol/sdk/client/index.js';
+import type * as aiTypes from 'ai';
+import type * as mcpTypes from '@ai-sdk/mcp';
+import { createLogger } from '@unchainedshop/logger';
+
+const logger = createLogger('unchained:api:chat');
+
+let createMCPClient: typeof mcpTypes.createMCPClient;
+let StreamableHTTPClientTransport: typeof mcpSDKClientLibraryTypes.StreamableHTTPClientTransport;
+let Client: typeof mcpSDKClientTypes.Client;
+
+try {
+  const mcpTools = await import('@ai-sdk/mcp');
+  const mcpSDKClientLibrary = await import('@modelcontextprotocol/sdk/client/streamableHttp.js');
+  const mcpSDKClient = await import('@modelcontextprotocol/sdk/client/index.js');
+
+  StreamableHTTPClientTransport = mcpSDKClientLibrary.StreamableHTTPClientTransport;
+  Client = mcpSDKClient.Client;
+  createMCPClient = mcpTools.createMCPClient;
+} catch {
+  // Handled at handler level
+}
+
+const { UNCHAINED_COOKIE_NAME = 'unchained_token' } = process.env;
+
+interface CachedSession {
+  tools: aiTypes.ToolSet;
+  resourceContext: string;
+  timestamp: number;
+  toolClient: Awaited<ReturnType<typeof mcpTypes.createMCPClient>>;
+}
+
+const sessionCache = new Map<string, CachedSession>();
+const inflightRequests = new Map<string, Promise<{ tools: aiTypes.ToolSet; resourceContext: string }>>();
+
+const CACHE_TTL_MS = 5 * 60 * 1000;
+const MAX_SESSIONS = 50;
+
+function extractAuthToken(cookie: string): string | null {
+  const match = cookie.match(new RegExp(`(?:^|;\\s*)${UNCHAINED_COOKIE_NAME}=([^;]+)`));
+  return match?.[1] || null;
+}
+
+function evictStale() {
+  const now = Date.now();
+  for (const [key, session] of sessionCache) {
+    if (now - session.timestamp > CACHE_TTL_MS) {
+      session.toolClient.close().catch(() => undefined);
+      sessionCache.delete(key);
+    }
+  }
+}
+
+function evictOldest() {
+  if (sessionCache.size < MAX_SESSIONS) return;
+  let oldestKey: string | null = null;
+  let oldestTime = Infinity;
+  for (const [key, session] of sessionCache) {
+    if (session.timestamp < oldestTime) {
+      oldestTime = session.timestamp;
+      oldestKey = key;
+    }
+  }
+  if (oldestKey) {
+    sessionCache
+      .get(oldestKey)!
+      .toolClient.close()
+      .catch(() => undefined);
+    sessionCache.delete(oldestKey);
+  }
+}
+
+async function fetchResources(mcpUrl: string, cookie: string): Promise<string> {
+  const transport = new StreamableHTTPClientTransport(new URL(mcpUrl), {
+    requestInit: { headers: { Cookie: cookie } },
+  });
+  const sdkClient = new Client({ name: 'unchained-chat-client', version: '1.0.0' });
+
+  try {
+    await sdkClient.connect(transport as any);
+    const resources = await sdkClient.listResources();
+    if (!resources?.resources?.length) return '';
+
+    const resourceTexts = await Promise.all(
+      resources.resources.map(async (resource) => {
+        try {
+          const content = await sdkClient.readResource({ uri: resource.uri });
+          if ((content?.contents?.[0] as any)?.text) {
+            return `${resource.name}:\n${(content.contents[0] as any).text}`;
+          }
+        } catch (e: any) {
+          logger.error(`Failed to read resource ${resource.uri}: ${e.message}`);
+        }
+        return null;
+      }),
+    );
+
+    const joined = resourceTexts.filter(Boolean).join('\n\n');
+    return joined ? '\n\nAVAILABLE SHOP CONFIGURATION:\n' + joined : '';
+  } finally {
+    await sdkClient.close().catch(() => undefined);
+  }
+}
+
+async function createSession(
+  mcpUrl: string,
+  cookie: string,
+  cacheKey: string,
+): Promise<{ tools: aiTypes.ToolSet; resourceContext: string }> {
+  const transport = new StreamableHTTPClientTransport(new URL(mcpUrl), {
+    requestInit: { headers: { Cookie: cookie } },
+  });
+
+  const toolClient = await createMCPClient({ transport });
+  const tools = await toolClient.tools();
+  const resourceContext = await fetchResources(mcpUrl, cookie);
+
+  evictOldest();
+  sessionCache.set(cacheKey, {
+    tools,
+    resourceContext,
+    timestamp: Date.now(),
+    toolClient,
+  });
+
+  return { tools, resourceContext };
+}
+
+export async function getMCPSession(
+  mcpUrl: string,
+  cookie: string,
+): Promise<{ tools: aiTypes.ToolSet; resourceContext: string }> {
+  evictStale();
+
+  const authToken = extractAuthToken(cookie);
+  if (!authToken) {
+    throw new Error('No authentication token found');
+  }
+
+  const cacheKey = authToken;
+  const cached = sessionCache.get(cacheKey);
+
+  if (cached && Date.now() - cached.timestamp < CACHE_TTL_MS) {
+    return { tools: cached.tools, resourceContext: cached.resourceContext };
+  }
+
+  if (cached) {
+    cached.toolClient.close().catch(() => undefined);
+    sessionCache.delete(cacheKey);
+  }
+
+  // Deduplicate concurrent requests for the same session
+  const inflight = inflightRequests.get(cacheKey);
+  if (inflight) return inflight;
+
+  const promise = createSession(mcpUrl, cookie, cacheKey).finally(() => {
+    inflightRequests.delete(cacheKey);
+  });
+  inflightRequests.set(cacheKey, promise);
+
+  return promise;
+}
+
+export { createMCPClient };

packages/api/src/chat/utils.ts

@@ -22,7 +22,7 @@ export const errorHandler = (error: any): string => {
   if (NoSuchToolError.isInstance(error)) return 'NoSuchToolError';
   if (InvalidArgumentError.isInstance(error)) return 'InvalidToolArgumentsError';
   /* if (ToolExecutionError.isInstance(error)) return 'ToolExecutionError'; */
-  if (error?.message?.toLowerCase()?.includes('forbidden')) return 'NetworkError';
+  if (error?.message?.toLowerCase()?.includes('forbidden')) return 'ForbiddenError';
   if (error?.message?.toLowerCase()?.includes('limit')) return 'LimitExceeded';
   return `Failed to stream response: ${error?.message || 'Unknown error'}`;
 };

packages/api/src/express/chatHandler.ts

@@ -1,36 +1,24 @@
 import express from 'express';
 import type { Express, Request, RequestHandler, Response } from 'express';
 import type * as aiTypes from 'ai';
-import type * as mcpTypes from '@ai-sdk/mcp';
-import type * as mcpSDKClientLibraryTypes from '@modelcontextprotocol/sdk/client/streamableHttp.js';
-import type * as mcpSDKClientTypes from '@modelcontextprotocol/sdk/client/index.js';
 import { type ChatConfiguration, errorHandler } from '../chat/utils.ts';
 import generateImageHandler from '../chat/generateImageHandler.ts';
 import defaultSystemPrompt from '../chat/defaultSystemPrompt.ts';
 import normalizeToolsIndex from '../chat/normalizeToolsIndex.ts';
+import { getMCPSession, createMCPClient } from '../chat/mcpChatSession.ts';
 import { createLogger } from '@unchainedshop/logger';
 
 const logger = createLogger('unchained:api:chat');
 
 let convertToModelMessages: typeof aiTypes.convertToModelMessages;
 let stepCountIs: typeof aiTypes.stepCountIs;
 let streamText: typeof aiTypes.streamText;
-let createMCPClient: typeof mcpTypes.createMCPClient;
-let StreamableHTTPClientTransport: typeof mcpSDKClientLibraryTypes.StreamableHTTPClientTransport;
-let Client: typeof mcpSDKClientTypes.Client;
 
 try {
   const aiTools = await import('ai');
-  const mcpTools = await import('@ai-sdk/mcp');
-  const mcpSDKClientLibrary = await import('@modelcontextprotocol/sdk/client/streamableHttp.js');
-  const mcpSDKClient = await import('@modelcontextprotocol/sdk/client/index.js');
-
-  StreamableHTTPClientTransport = mcpSDKClientLibrary.StreamableHTTPClientTransport;
-  Client = mcpSDKClient.Client;
   convertToModelMessages = aiTools.convertToModelMessages;
   stepCountIs = aiTools.stepCountIs;
   streamText = aiTools.streamText;
-  createMCPClient = mcpTools.createMCPClient;
 } catch {
   logger.warn(
     `optional peer npm packages 'ai', '@ai-sdk/mcp' and '@modelcontextprotocol/sdk' not installed, chat will not work`,
@@ -63,56 +51,13 @@ const setupMCPChatHandler = (chatConfiguration: ChatConfiguration & any): Reques
       res.status(405).json({ error: 'Method Not Allowed. Use POST.' });
       return;
     }
-    const resourceTransport = new StreamableHTTPClientTransport(new URL(unchainedMCPUrl), {
-      requestInit: {
-        headers: {
-          Cookie: req.headers.cookie || '',
-        },
-      },
-    });
-
-    const sdkClient = new Client({ name: 'unchained-chat-client', version: '1.0.0' });
-    await sdkClient.connect(resourceTransport as any);
-    const transport = new StreamableHTTPClientTransport(new URL(unchainedMCPUrl), {
-      requestInit: {
-        headers: {
-          Cookie: req.headers.cookie || '',
-        },
-      },
-    });
-
-    const client = await createMCPClient({
-      transport,
-    });
 
     try {
-      const defaultUnchainedTools = await client.tools();
-      let resourceContext = '';
-      try {
-        const resources = await sdkClient.listResources();
-        if (resources?.resources) {
-          const resourceTexts = await Promise.all(
-            resources.resources.map(async (resource) => {
-              try {
-                const content = await sdkClient.readResource({ uri: resource.uri });
-                if ((content?.contents?.[0] as any)?.text) {
-                  return `${resource.name}:\n${(content.contents[0] as any).text}`;
-                }
-              } catch (e) {
-                logger.error(`Failed to read resource ${resource.uri}: ${e.message}`);
-              }
-              return null;
-            }),
-          );
-          resourceContext =
-            '\n\nAVAILABLE SHOP CONFIGURATION:\n' + resourceTexts.filter(Boolean).join('\n\n');
-        }
-      } catch (e) {
-        logger.error(`Failed to fetch MCP resources: ${e.message}`);
-      }
+      const cookie = req.headers.cookie || '';
+      const { tools: mcpTools, resourceContext } = await getMCPSession(unchainedMCPUrl, cookie);
 
       const tools: aiTypes.ToolSet = {
-        ...defaultUnchainedTools,
+        ...mcpTools,
         ...additionalTools,
       };
       if (imageGenerationTool) {
@@ -122,7 +67,7 @@ const setupMCPChatHandler = (chatConfiguration: ChatConfiguration & any): Reques
       if (req.method === 'GET') {
         res.status(200).json({
           tools: normalizeToolsIndex(tools),
-          cached: false,
+          cached: true,
         });
         return;
       }
@@ -162,16 +107,13 @@ const setupMCPChatHandler = (chatConfiguration: ChatConfiguration & any): Reques
       const messagesToInclude = normalizedMessages.slice(startIndex);
 
       const result = streamText({
-        stopWhen: stepCountIs(10),
+        stopWhen: stepCountIs(500),
         temperature: 0.2,
         maxRetries: 3,
         ...restChatConfig,
         system: system + resourceContext,
         model,
         tools: cacheControlledTools,
-        onFinish: async () => {
-          await client?.close();
-        },
         messages: messagesToInclude,
         providerOptions: {
           anthropic: {
@@ -186,8 +128,7 @@ const setupMCPChatHandler = (chatConfiguration: ChatConfiguration & any): Reques
         onError: errorHandler,
       });
     } catch (err) {
-      await client?.close();
-      await sdkClient?.close();
+      logger.error(err);
       res.status(500).json({ error: errorHandler(err) });
     }
   };