keep track of sessions for chat

marvinjude · marvinjude · commit bf8eda2482d8 · 2025-07-09T07:58:47.000+01:00
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 # Integration App MCP Server
-<img width="1148" alt="Screenshot 2025-07-07 at 23 03 05" src="https://github.com/user-attachments/assets/39f6cc74-a689-4657-91f3-ee8358c05e31" />
 
+<img width="1148" alt="Screenshot 2025-07-07 at 23 03 05" src="https://github.com/user-attachments/assets/39f6cc74-a689-4657-91f3-ee8358c05e31" />
 
 The Integration App MCP Server is a [Model Context Protocol (MCP)](https://modelcontextprotocol.io/introduction) server, it provides actions for connected integrations on Integration.app membrane as tools.
 
@@ -107,13 +107,15 @@ await client.connect(
 );
 ```
 
-### ⚡ Static & Dynamic Mode
+### ⚡ Static vs Dynamic Mode
+
+By default, the MCP server runs in **static mode**, which means it returns **all available tools** (actions) for all connected integrations.
 
-By default, the MCP server is in `static` mode and will return all tools. In `dynamic` mode (`?mode=dynamic`) the MCP server will only return only a single tool: `enable-tools`, you can use this tool to enable tools for the session.
+If you switch to **dynamic mode** (by adding `?mode=dynamic` to your request), the server will only return **one tool**: `enable-tools`. You can use this tool to selectively enable the tools you actually need for that session.
 
-Your implementation needs to provide a way to find the most relevant tools to the user query, after which you can use the `enable-tools` tool to enable the tools for the session. Ideally you want to prompt LLM to call this tool
+In dynamic mode, your implementation should figure out which tools are most relevant to the user's query. Once you've identified them, prompt the LLM to call the `enable-tools` tool with the appropriate list.
 
-See an example implementation in our [AI Agent Example](https://github.com/integration-app/ai-agent-example)
+Want to see how this works in practice? Check out our [AI Agent Example](https://github.com/integration-app/ai-agent-example).
 
 ```ts
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
@@ -151,6 +153,36 @@ In static mode, the MCP server fetches tools from all active connections associa
 
 You can choose to only fetch tools for a specific integration by passing the `apps` query parameter: `/mcp?apps=google-calendar,google-docs`
 
+### 💬 Chat Session Management
+
+The MCP server supports persistent chat sessions. Include an `x-chat-id` header in your requests to automatically track sessions for that specific chat.
+
+**Starting a new chat session:**
+
+```http
+POST /mcp
+Authorization: Bearer YOUR_ACCESS_TOKEN
+x-chat-id: my-awesome-chat-123
+```
+
+**Retrieving your chat sessions:**
+
+```http
+GET /mcp/sessions
+Authorization: Bearer YOUR_ACCESS_TOKEN
+```
+
+**Response:**
+
+```json
+{
+  "my-awesome-chat-123": "session-uuid-1",
+  "another-chat-456": "session-uuid-2"
+}
+```
+
+This feature lets you maintain conversation context across multiple requests without creating new sessions each time. Perfect for multi-chat applications! Check out our [AI Agent Example](https://github.com/integration-app/ai-agent-example) to see how this works in practice.
+
 ### Configuring other MCP clients
 
 #### 📝 Cursor
diff --git a/src/server/index.ts b/src/server/index.ts
@@ -22,7 +22,7 @@ app.use(
   cors({
     origin: '*',
     methods: ['GET', 'POST', 'PUT', 'DELETE', 'OPTIONS'],
-    allowedHeaders: ['Content-Type', 'Authorization', 'mcp-session-id', 'last-event-id'],
+    allowedHeaders: ['Content-Type', 'Authorization', 'mcp-session-id', 'x-chat-id', 'last-event-id'],
   })
 );
 app.use(morgan(customMorganFormat));
diff --git a/src/server/middlewares/logging.ts b/src/server/middlewares/logging.ts
@@ -13,8 +13,9 @@ export const customMorganFormat = (tokens: any, req: any, res: any) => {
   }
 
   const sessionId = req.headers['mcp-session-id'];
+  const chatId = req.headers['x-chat-id'];
   const userId = req.userId;
 
   return `
-  ${colorize(method, getMethodColor(method))} ${colorize(truncatedUrl, colors.white)} ${colorize(status, getStatusColor(status))} ${colorize(`${responseTime}ms`, colors.gray)} 👤 User: ${colorize(userId, colors.green)} Session: ${colorize(sessionId, colors.cyan)} 🔧 Method: ${colorize(req.body.method, colors.magenta)} Mode: ${colorize(req.query.mode, colors.yellow)}`;
+  ${colorize(method, getMethodColor(method))} ${colorize(truncatedUrl, colors.white)} ${colorize(status, getStatusColor(status))} ${colorize(`${responseTime}ms`, colors.gray)} 👤 User: ${colorize(userId, colors.green)} 💬 Chat: ${colorize(chatId, colors.blue)} 🔑 Session: ${colorize(sessionId, colors.cyan)} 🔧 Method: ${colorize(req.body.method, colors.magenta)}`;
 };
diff --git a/src/server/routes/streamable-http.ts b/src/server/routes/streamable-http.ts
@@ -3,6 +3,7 @@ import { randomUUID } from 'crypto';
 import express, { Request, Response } from 'express';
 import { createMcpServer, CreateMcpServerParams } from '../utils/create-mcp-server';
 import { InMemoryEventStore } from '@modelcontextprotocol/sdk/examples/shared/inMemoryEventStore.js';
+import { addToUserSessions, getUserSessions, removeFromUserSessions } from '../utils/user-sessions';
 
 export const streamableHttpRouter = express.Router();
 
@@ -34,6 +35,11 @@ streamableHttpRouter.post('/', async (req, res) => {
           // Store the transport by session ID when session is initialized
           // This avoids race conditions where requests might come in before the session is stored
           transports.set(sessionId, transport);
+
+          const userId = req.userId;
+          const chatId = req.headers['x-chat-id'] as string | undefined;
+
+          addToUserSessions({ userId, chatId, sessionId });
         },
       });
 
@@ -145,6 +151,11 @@ streamableHttpRouter.delete('/', async (req: Request, res: Response) => {
   try {
     const transport = transports.get(sessionId);
     await transport!.handleRequest(req, res);
+
+    const userId = req.userId;
+    const chatId = req.headers['x-chat-id'] as string | undefined;
+
+    removeFromUserSessions({ userId, chatId, sessionId });
   } catch (error) {
     console.error('Error handling session termination:', error);
     if (!res.headersSent) {
@@ -160,3 +171,11 @@ streamableHttpRouter.delete('/', async (req: Request, res: Response) => {
     }
   }
 });
+
+// Get all the sessions user has opened
+streamableHttpRouter.get('/sessions', async (req: Request, res: Response) => {
+  const userId = req.userId;
+  const sessions = getUserSessions(userId);
+
+  res.json(sessions);
+});
diff --git a/src/server/utils/create-mcp-server.ts b/src/server/utils/create-mcp-server.ts
@@ -6,6 +6,7 @@ import { isTestEnvironment } from './constants';
 import { zodFromJsonSchema } from './json-schema-to-zod';
 import { ZodRawShape } from 'zod';
 import pkg from '../../../package.json';
+import { getActionsByKeys } from './membrane/get-actions-by-keys';
 
 /**
  * Register a tool to the MCP server while adding some standardization
@@ -81,11 +82,20 @@ export const createMcpServer = async ({
   apps,
   mode = 'static',
 }: CreateMcpServerParams) => {
-  const mcpServer = new McpServer({
-    name: 'Integration App MCP Server',
-    version: pkg.version,
-    description: pkg.description,
-  });
+  const mcpServer = new McpServer(
+    {
+      name: 'Integration App MCP Server',
+      version: pkg.version,
+      description: pkg.description,
+    },
+    {
+      capabilities: {
+        tools: {
+          listChanged: true,
+        },
+      },
+    }
+  );
 
   let membrane: IntegrationAppClient;
 
@@ -140,8 +150,39 @@ export const createMcpServer = async ({
 
         // Enable new tools
         for (const tool of args.tools) {
-          registeredTools[tool]?.enable();
-          enabledTools.push(tool);
+          /**
+           * There are cases where the tool client is trying to enable is not a registered tool.
+           * In that case, we need to get the action from membrane, register it and enable it.
+           */
+          const isToolRegistered = registeredTools[tool];
+
+          if (isToolRegistered) {
+            registeredTools[tool].enable();
+            enabledTools.push(tool);
+          } else {
+            const actions = await getActionsByKeys(
+              membrane,
+              args.tools.map((toolKey: string) => {
+                const [integrationKey, actionKey] = toolKey.split('_');
+
+                return {
+                  key: actionKey,
+                  integrationKey,
+                };
+              })
+            );
+
+            for (const action of actions) {
+              if (action) {
+                const addedTool = addServerTool({ mcpServer, action, membrane });
+                if (addedTool) {
+                  registeredTools[addedTool.toolName] = addedTool.tool;
+                  addedTool.tool.enable();
+                  enabledTools.push(addedTool.toolName);
+                }
+              }
+            }
+          }
         }
 
         return {
diff --git a/src/server/utils/membrane/get-actions-by-keys.ts b/src/server/utils/membrane/get-actions-by-keys.ts
@@ -0,0 +1,27 @@
+import { IntegrationAppClient } from '@integration-app/sdk';
+
+export const getActionsByKeys = async (
+  membrane: IntegrationAppClient,
+  _actions: {
+    key: string;
+    integrationKey: string;
+  }[]
+) => {
+  const results = await Promise.allSettled(
+    _actions.map(async action => {
+      return membrane
+        .action({
+          integrationKey: action.integrationKey,
+          key: action.key,
+        })
+        .get();
+    })
+  );
+
+  // Filter out failed promises and return only successful results
+  const successfulActions = results
+    .filter((result): result is PromiseFulfilledResult<any> => result.status === 'fulfilled')
+    .map(result => result.value);
+
+  return successfulActions;
+};
diff --git a/src/server/utils/user-sessions.ts b/src/server/utils/user-sessions.ts
@@ -0,0 +1,51 @@
+/**
+ * Stores user sessions per chat
+ *
+ * @example
+ * {
+ *  "userId": {
+ *    "chatId": "sessionId",
+ *    "chatId2": "sessionId2",
+ *  }
+ * }
+ */
+
+const userSessions: Record<string, Record<string, string>> = {};
+
+export const addToUserSessions = ({
+  userId,
+  chatId,
+  sessionId,
+}: {
+  userId: string;
+  chatId?: string;
+  sessionId: string;
+}) => {
+  if (chatId) {
+    const currentSessionsForUser = userSessions[userId] ?? {};
+    userSessions[userId] = {
+      ...currentSessionsForUser,
+      [chatId]: sessionId,
+    };
+  }
+};
+
+export const removeFromUserSessions = ({
+  userId,
+  chatId,
+  sessionId,
+}: {
+  userId: string;
+  chatId?: string;
+  sessionId: string;
+}) => {
+  const currentSessionsForUser = userSessions[userId];
+  if (currentSessionsForUser) {
+    delete currentSessionsForUser[chatId ?? sessionId];
+    userSessions[userId] = currentSessionsForUser;
+  }
+};
+
+export const getUserSessions = (userId: string) => {
+  return userSessions[userId];
+};
diff --git a/tests/mcp-server.test.ts b/tests/mcp-server.test.ts
@@ -53,6 +53,7 @@ describe('MCP Server Integration Tests', () => {
       requestInit: {
         headers: {
           Authorization: `Bearer ${TEST_USER_ACCESS_TOKEN}`,
+          'x-chat-id': TEST_CHAT_ID,
         },
       },
     });

Original file line number	Diff line number	Diff line change
`@@ -13,8 +13,9 @@ export const customMorganFormat = (tokens: any, req: any, res: any) => {`
`13`	`13`	`}`
`14`	`14`
`15`	`15`	`const sessionId = req.headers['mcp-session-id'];`
	`16`	`+ const chatId = req.headers['x-chat-id'];`
`16`	`17`	`const userId = req.userId;`
`17`	`18`
`18`	`19`	return `
`19`		- ${colorize(method, getMethodColor(method))} ${colorize(truncatedUrl, colors.white)} ${colorize(status, getStatusColor(status))} ${colorize(`${responseTime}ms`, colors.gray)} 👤 User: ${colorize(userId, colors.green)} Session: ${colorize(sessionId, colors.cyan)} 🔧 Method: ${colorize(req.body.method, colors.magenta)} Mode: ${colorize(req.query.mode, colors.yellow)}`;
	`20`	+ ${colorize(method, getMethodColor(method))} ${colorize(truncatedUrl, colors.white)} ${colorize(status, getStatusColor(status))} ${colorize(`${responseTime}ms`, colors.gray)} 👤 User: ${colorize(userId, colors.green)} 💬 Chat: ${colorize(chatId, colors.blue)} 🔑 Session: ${colorize(sessionId, colors.cyan)} 🔧 Method: ${colorize(req.body.method, colors.magenta)}`;
`20`	`21`	`};`