chore(mcp-server): add support for session id, forward client info

Author: stainless-app[bot]

packages/mcp-server/src/docs-search-tool.ts

@@ -76,17 +76,18 @@ async function searchLocal(args: Record<string, unknown>): Promise<unknown> {
   }).results;
 }
 
-async function searchRemote(
-  args: Record<string, unknown>,
-  stainlessApiKey: string | undefined,
-): Promise<unknown> {
+async function searchRemote(args: Record<string, unknown>, reqContext: McpRequestContext): Promise<unknown> {
   const body = args as any;
   const query = new URLSearchParams(body).toString();
 
   const startTime = Date.now();
   const result = await fetch(`${docsSearchURL}?${query}`, {
     headers: {
-      ...(stainlessApiKey && { Authorization: stainlessApiKey }),
+      ...(reqContext.stainlessApiKey && { Authorization: reqContext.stainlessApiKey }),
+      ...(reqContext.mcpSessionId && { 'x-stainless-mcp-session-id': reqContext.mcpSessionId }),
+      ...(reqContext.mcpClientInfo && {
+        'x-stainless-mcp-client-info': JSON.stringify(reqContext.mcpClientInfo),
+      }),
     },
   });
 
@@ -105,7 +106,7 @@ async function searchRemote(
       'Got error response from docs search tool',
     );
 
-    if (result.status === 404 && !stainlessApiKey) {
+    if (result.status === 404 && !reqContext.stainlessApiKey) {
       throw new Error(
         'Could not find docs for this project. You may need to provide a Stainless API key via the STAINLESS_API_KEY environment variable, the --stainless-api-key flag, or the x-stainless-api-key HTTP header.',
       );
@@ -140,7 +141,7 @@ export const handler = async ({
     return asTextContentResult(await searchLocal(body));
   }
 
-  return asTextContentResult(await searchRemote(body, reqContext.stainlessApiKey));
+  return asTextContentResult(await searchRemote(body, reqContext));
 };
 
 export default { metadata, tool, handler };

packages/mcp-server/src/http.ts

@@ -78,6 +78,11 @@ const newServer = async ({
     },
     stainlessApiKey: stainlessApiKey,
     upstreamClientEnvs,
+    mcpSessionId: (req as any).mcpSessionId,
+    mcpClientInfo:
+      typeof req.body?.params?.clientInfo?.name === 'string' ?
+        { name: req.body.params.clientInfo.name, version: String(req.body.params.clientInfo.version ?? '') }
+      : undefined,
   });
 
   return server;
@@ -135,6 +140,17 @@ export const streamableHTTPApp = ({
   const app = express();
   app.set('query parser', 'extended');
   app.use(express.json());
+  app.use((req: express.Request, res: express.Response, next: express.NextFunction) => {
+    const existing = req.headers['mcp-session-id'];
+    const sessionId = (Array.isArray(existing) ? existing[0] : existing) || crypto.randomUUID();
+    (req as any).mcpSessionId = sessionId;
+    const origWriteHead = res.writeHead.bind(res);
+    res.writeHead = function (statusCode: number, ...rest: any[]) {
+      res.setHeader('mcp-session-id', sessionId);
+      return origWriteHead(statusCode, ...rest);
+    } as typeof res.writeHead;
+    next();
+  });
   app.use(
     pinoHttp({
       logger: getLogger(),

packages/mcp-server/src/server.ts

@@ -47,6 +47,8 @@ export async function initMcpServer(params: {
   mcpOptions?: McpOptions;
   stainlessApiKey?: string | undefined;
   upstreamClientEnvs?: Record<string, string> | undefined;
+  mcpSessionId?: string | undefined;
+  mcpClientInfo?: { name: string; version: string } | undefined;
 }) {
   const server = params.server instanceof McpServer ? params.server.server : params.server;
 
@@ -136,6 +138,8 @@ export async function initMcpServer(params: {
         client,
         stainlessApiKey: params.stainlessApiKey ?? params.mcpOptions?.stainlessApiKey,
         upstreamClientEnvs: params.upstreamClientEnvs,
+        mcpSessionId: params.mcpSessionId,
+        mcpClientInfo: params.mcpClientInfo,
       },
       args,
     });

packages/mcp-server/src/types.ts

@@ -46,6 +46,8 @@ export type McpRequestContext = {
   client: Finch;
   stainlessApiKey?: string | undefined;
   upstreamClientEnvs?: Record<string, string> | undefined;
+  mcpSessionId?: string | undefined;
+  mcpClientInfo?: { name: string; version: string } | undefined;
 };
 
 export type HandlerFunction = ({