feat(mcp): forward client metadata

Co-Authored-By: Miguel Brandão <miguel@exa.ai>

Author: devin-ai-integration[bot]

api/mcp.ts

@@ -6,6 +6,13 @@ import { initializeMcpServer } from '../src/mcp-handler.js';
 import { Ratelimit } from '@upstash/ratelimit';
 import { Redis } from '@upstash/redis';
 import { isJwtToken, verifyOAuthToken } from '../src/utils/auth.js';
+import {
+  buildMcpClientMetadata,
+  extractInitializeClientInfo,
+  MCP_CLIENT_SESSION_TTL_SECONDS,
+  sanitizeMcpClientMetadata,
+  type McpClientMetadata,
+} from '../src/utils/mcpClientMetadata.js';
 
 // Origin: '*' is safe — auth is per-request via headers/query, never cookies.
 const CORS_HEADERS: Record<string, string> = {
@@ -50,6 +57,59 @@ let dailyLimiter: Ratelimit | null = null;
 let rateLimitersInitialized = false;
 let redisClient: Redis | null = null;
 
+function getMcpClientSessionKey(sessionId: string): string {
+  return `exa-mcp:client:${sessionId}`;
+}
+
+async function saveMcpClientMetadata(sessionId: string | undefined, metadata: McpClientMetadata | undefined, debug: boolean): Promise<void> {
+  if (!sessionId || !metadata?.clientInfo) {
+    return;
+  }
+
+  initializeRateLimiters();
+
+  if (!redisClient) {
+    return;
+  }
+
+  try {
+    await redisClient.set(getMcpClientSessionKey(sessionId), JSON.stringify(metadata), {
+      ex: MCP_CLIENT_SESSION_TTL_SECONDS,
+    });
+  } catch (error) {
+    if (debug) {
+      console.error('[EXA-MCP] Failed to save MCP client metadata:', error);
+    }
+  }
+}
+
+async function loadMcpClientMetadata(sessionId: string | undefined, debug: boolean): Promise<McpClientMetadata | undefined> {
+  if (!sessionId) {
+    return undefined;
+  }
+
+  initializeRateLimiters();
+
+  if (!redisClient) {
+    return undefined;
+  }
+
+  try {
+    const value = await redisClient.get<string>(getMcpClientSessionKey(sessionId));
+    if (typeof value !== 'string') {
+      return undefined;
+    }
+
+    const parsed: unknown = JSON.parse(value);
+    return sanitizeMcpClientMetadata(parsed);
+  } catch (error) {
+    if (debug) {
+      console.error('[EXA-MCP] Failed to load MCP client metadata:', error);
+    }
+    return undefined;
+  }
+}
+
 function initializeRateLimiters(): boolean {
   if (rateLimitersInitialized) {
     return qpsLimiter !== null;
@@ -299,6 +359,7 @@ interface RequestConfig {
   authMethod: 'oauth' | 'api_key' | 'free_tier';
   exaSource?: string;
   mcpSessionId?: string;
+  mcpClient?: McpClientMetadata;
   defaultSearchType?: 'auto' | 'fast';
   /** True when a Bearer token was a JWT but failed OAuth verification (expired, bad sig, wrong issuer/audience). */
   invalidOAuthJwt: boolean;
@@ -417,7 +478,7 @@ async function getConfigFromRequest(request: Request): Promise<RequestConfig> {
  * configuration (tools and API key). This prevents API key leakage between
  * different users who might pass different keys via URL.
  */
-function createHandler(config: { exaApiKey?: string; enabledTools?: string[]; debug: boolean; userProvidedApiKey: boolean; exaSource?: string; mcpSessionId?: string; defaultSearchType?: 'auto' | 'fast' }) {
+function createHandler(config: { exaApiKey?: string; enabledTools?: string[]; debug: boolean; userProvidedApiKey: boolean; exaSource?: string; mcpSessionId?: string; mcpClient?: McpClientMetadata; defaultSearchType?: 'auto' | 'fast' }) {
   return createMcpHandler(
     (server: any) => {
       initializeMcpServer(server, config);
@@ -484,8 +545,9 @@ async function handleRequest(request: Request, options?: { forceOAuth?: boolean
  */
 async function processRequest(request: Request, options?: { forceOAuth?: boolean }): Promise<Response> {
   const debug = process.env.DEBUG === 'true';
-  const isInitializeRequest =
-    request.method === 'POST' ? isInitializeMethod(await request.clone().text()) : false;
+  const body = request.method === 'POST' ? await request.clone().text() : undefined;
+  const isInitializeRequest = isInitializeMethod(body ?? '');
+  const initializeClientInfo = extractInitializeClientInfo(body);
 
   // Check user-agent bypass BEFORE the 401 gate so bypass clients never see auth prompts
   const userAgent = request.headers.get('user-agent') || '';
@@ -518,6 +580,15 @@ async function processRequest(request: Request, options?: { forceOAuth?: boolean
     return create401Response();
   }
 
+  const storedMcpClient = isInitializeRequest ? undefined : await loadMcpClientMetadata(config.mcpSessionId, config.debug);
+  config.mcpClient = buildMcpClientMetadata({
+    source: config.exaSource,
+    sessionId: config.mcpSessionId,
+    clientInfo: initializeClientInfo,
+    stored: storedMcpClient,
+    userAgent,
+  });
+  
   if (config.debug) {
     console.log(`[EXA-MCP] Request URL: ${request.url}`);
     console.log(`[EXA-MCP] Enabled tools: ${config.enabledTools?.join(', ') || 'default'}`);
@@ -536,12 +607,8 @@ async function processRequest(request: Request, options?: { forceOAuth?: boolean
   // Rate limit users who didn't provide their own API key (including bypass users)
   // Only rate limit actual tool calls (tools/call), not protocol methods like tools/list
   if (!config.userProvidedApiKey && request.method === 'POST') {
-    // Clone the request to read the body without consuming it
-    const clonedRequest = request.clone();
-    const body = await clonedRequest.text();
-    
     // Only rate limit actual tool calls, not protocol methods
-    if (isRateLimitedMethod(body)) {
+    if (isRateLimitedMethod(body ?? '')) {
       // Initialize rate limiters on first request (lazy init)
       initializeRateLimiters();
 
@@ -591,9 +658,22 @@ async function processRequest(request: Request, options?: { forceOAuth?: boolean
 
   const response = withCors(await handler(request));
 
-  if (isInitializeRequest && response.ok && !response.headers.has('Mcp-Session-Id')) {
+  if (isInitializeRequest && response.ok) {
+    const responseSessionId = response.headers.get('Mcp-Session-Id') ?? config.mcpSessionId ?? randomUUID();
+    const metadata = buildMcpClientMetadata({
+      source: config.exaSource,
+      sessionId: responseSessionId,
+      clientInfo: initializeClientInfo,
+      userAgent,
+    });
+    await saveMcpClientMetadata(responseSessionId, metadata, config.debug);
+
+    if (response.headers.has('Mcp-Session-Id')) {
+      return response;
+    }
+
     const headers = new Headers(response.headers);
-    headers.set('Mcp-Session-Id', randomUUID());
+    headers.set('Mcp-Session-Id', responseSessionId);
     return new Response(response.body, {
       status: response.status,
       statusText: response.statusText,

src/mcp-handler.ts

@@ -36,6 +36,7 @@ export interface McpConfig {
   userProvidedApiKey?: boolean;
   exaSource?: string;
   mcpSessionId?: string;
+  mcpClient?: unknown;
   defaultSearchType?: 'auto' | 'fast';
 }
 

src/tools/config.ts

@@ -1,7 +1,10 @@
+import { serializeMcpClientMetadata } from '../utils/mcpClientMetadata.js';
+
 // Build Exa reporting headers, appending x-exa-source if present
 export function integrationHeaders(tool: string, config?: Record<string, unknown>) {
   const source = config?.exaSource;
   const mcpSessionId = config?.mcpSessionId;
+  const mcpClient = serializeMcpClientMetadata(config?.mcpClient);
   const headers: Record<string, string> = {
     'x-exa-integration': typeof source === 'string' ? `${tool}:${source}` : tool,
   };
@@ -10,6 +13,10 @@ export function integrationHeaders(tool: string, config?: Record<string, unknown
     headers['x-exa-mcp-session-id'] = mcpSessionId;
   }
 
+  if (mcpClient) {
+    headers['x-exa-mcp-client'] = mcpClient;
+  }
+
   return headers;
 }
 

src/utils/mcpClientMetadata.ts

@@ -0,0 +1,152 @@
+export const MCP_CLIENT_SESSION_TTL_SECONDS = 24 * 60 * 60;
+
+const MCP_CLIENT_FIELD_MAX_LENGTH = 256;
+const MCP_CLIENT_USER_AGENT_MAX_LENGTH = 512;
+const MCP_CLIENT_HEADER_MAX_LENGTH = 2048;
+const UNKNOWN_MCP_CLIENT_NAME = 'unknown';
+
+export interface McpClientInfo {
+  name?: string;
+  title?: string;
+  version?: string;
+}
+
+export interface McpClientMetadata {
+  source?: string;
+  sessionId?: string;
+  clientInfo?: McpClientInfo;
+  userAgent?: string;
+}
+
+function unknownMcpClientInfo(): McpClientInfo {
+  return { name: UNKNOWN_MCP_CLIENT_NAME };
+}
+
+function sanitizeMcpClientField(value: unknown, maxLength = MCP_CLIENT_FIELD_MAX_LENGTH): string | undefined {
+  if (typeof value !== 'string') {
+    return undefined;
+  }
+
+  let withoutControlCharacters = '';
+  for (const character of value) {
+    const codePoint = character.charCodeAt(0);
+    withoutControlCharacters += codePoint <= 31 || codePoint === 127 ? ' ' : character;
+  }
+
+  const sanitized = withoutControlCharacters.trim();
+  if (!sanitized) {
+    return undefined;
+  }
+
+  return sanitized.slice(0, maxLength);
+}
+
+function compactMcpClientMetadata(metadata: McpClientMetadata): McpClientMetadata | undefined {
+  const compact: McpClientMetadata = {};
+
+  if (metadata.source) compact.source = metadata.source;
+  if (metadata.sessionId) compact.sessionId = metadata.sessionId;
+  if (metadata.clientInfo && Object.keys(metadata.clientInfo).length > 0) {
+    compact.clientInfo = metadata.clientInfo;
+  }
+  if (metadata.userAgent) compact.userAgent = metadata.userAgent;
+
+  return Object.keys(compact).length > 0 ? compact : undefined;
+}
+
+function isMcpClientRecord(value: unknown): value is Record<string, unknown> {
+  return !!value && typeof value === 'object' && !Array.isArray(value);
+}
+
+function sanitizeMcpClientInfo(clientInfo: unknown): McpClientInfo | undefined {
+  if (!isMcpClientRecord(clientInfo)) {
+    return undefined;
+  }
+
+  const compact: McpClientInfo = {};
+  const name = sanitizeMcpClientField(clientInfo.name);
+  const title = sanitizeMcpClientField(clientInfo.title);
+  const version = sanitizeMcpClientField(clientInfo.version);
+
+  if (name) compact.name = name;
+  if (title) compact.title = title;
+  if (version) compact.version = version;
+
+  return Object.keys(compact).length > 0 ? compact : undefined;
+}
+
+export function extractInitializeClientInfo(body: string | undefined): McpClientInfo | undefined {
+  if (!body) {
+    return undefined;
+  }
+
+  try {
+    const parsed: unknown = JSON.parse(body);
+    if (!isMcpClientRecord(parsed) || parsed.method !== 'initialize') {
+      return undefined;
+    }
+
+    if (!isMcpClientRecord(parsed.params)) {
+      return unknownMcpClientInfo();
+    }
+
+    return sanitizeMcpClientInfo(parsed.params.clientInfo) ?? unknownMcpClientInfo();
+  } catch {
+    return unknownMcpClientInfo();
+  }
+}
+
+export function sanitizeMcpClientMetadata(metadata: unknown): McpClientMetadata | undefined {
+  if (!isMcpClientRecord(metadata)) {
+    return undefined;
+  }
+
+  return compactMcpClientMetadata({
+    source: sanitizeMcpClientField(metadata.source),
+    sessionId: sanitizeMcpClientField(metadata.sessionId),
+    clientInfo: sanitizeMcpClientInfo(metadata.clientInfo),
+    userAgent: sanitizeMcpClientField(metadata.userAgent, MCP_CLIENT_USER_AGENT_MAX_LENGTH),
+  });
+}
+
+export function buildMcpClientMetadata(input: {
+  source?: string;
+  sessionId?: string;
+  clientInfo?: McpClientInfo;
+  stored?: McpClientMetadata;
+  userAgent?: string;
+}): McpClientMetadata {
+  try {
+    const metadata = compactMcpClientMetadata({
+      source: sanitizeMcpClientField(input.source ?? input.stored?.source),
+      sessionId: sanitizeMcpClientField(input.sessionId ?? input.stored?.sessionId),
+      clientInfo: sanitizeMcpClientInfo(input.clientInfo ?? input.stored?.clientInfo),
+      userAgent: sanitizeMcpClientField(input.userAgent ?? input.stored?.userAgent, MCP_CLIENT_USER_AGENT_MAX_LENGTH),
+    }) ?? { clientInfo: unknownMcpClientInfo() };
+
+    if (!metadata.clientInfo && !metadata.userAgent) {
+      metadata.clientInfo = unknownMcpClientInfo();
+    }
+
+    return metadata;
+  } catch {
+    return { clientInfo: unknownMcpClientInfo() };
+  }
+}
+
+export function serializeMcpClientMetadata(value: unknown): string | undefined {
+  if (!value || typeof value !== 'object' || Array.isArray(value)) {
+    return undefined;
+  }
+
+  try {
+    const serialized = JSON.stringify(value);
+    if (serialized === '{}' || serialized.length > MCP_CLIENT_HEADER_MAX_LENGTH) {
+      return undefined;
+    }
+
+    return serialized;
+  } catch {
+    return JSON.stringify({ clientInfo: unknownMcpClientInfo() });
+  }
+}