From 73f75ae13311ef68c3f6e037293913f99cce9d20 Mon Sep 17 00:00:00 2001 From: Grimmer Kang Date: Tue, 29 Jul 2025 15:35:25 +0800 Subject: [PATCH 1/4] feat: add custom context support to MCP SDK - Add customContext field to MessageExtraInfo interface - Pass customContext through RequestHandlerExtra to handlers - Update InMemoryTransport to support custom context - Add unit test to verify custom context propagation - Enable transport implementations to inject arbitrary context data This allows MCP server implementations to access custom context (e.g., tenant ID, user info, feature flags) passed from the transport layer to tool/resource/prompt handlers. --- src/inMemory.ts | 51 +++++++++++++++++++++++++------ src/server/mcp.test.ts | 58 ++++++++++++++++++++++++++++++++++++ src/server/sse.ts | 15 +++++++++- src/server/stdio.ts | 18 +++++++++-- src/server/streamableHttp.ts | 22 ++++++++++++-- src/shared/protocol.ts | 10 ++++++- src/shared/transport.ts | 6 ++++ src/types.ts | 7 +++++ 8 files changed, 171 insertions(+), 16 deletions(-) diff --git a/src/inMemory.ts b/src/inMemory.ts index 5dd6e81e0..b997965dd 100644 --- a/src/inMemory.ts +++ b/src/inMemory.ts @@ -1,10 +1,10 @@ import { Transport } from "./shared/transport.js"; -import { JSONRPCMessage, RequestId } from "./types.js"; +import { JSONRPCMessage, RequestId, MessageExtraInfo } from "./types.js"; import { AuthInfo } from "./server/auth/types.js"; interface QueuedMessage { message: JSONRPCMessage; - extra?: { authInfo?: AuthInfo }; + extra?: MessageExtraInfo; } /** @@ -13,10 +13,11 @@ interface QueuedMessage { export class InMemoryTransport implements Transport { private _otherTransport?: InMemoryTransport; private _messageQueue: QueuedMessage[] = []; + private _customContext?: Record; onclose?: () => void; onerror?: (error: Error) => void; - onmessage?: (message: JSONRPCMessage, extra?: { authInfo?: AuthInfo }) => void; + onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; sessionId?: string; /** @@ -34,7 +35,12 @@ export class InMemoryTransport implements Transport { // Process any messages that were queued before start was called while (this._messageQueue.length > 0) { const queuedMessage = this._messageQueue.shift()!; - this.onmessage?.(queuedMessage.message, queuedMessage.extra); + // Merge custom context with queued extra info + const enhancedExtra: MessageExtraInfo = { + ...queuedMessage.extra, + customContext: this._customContext + }; + this.onmessage?.(queuedMessage.message, enhancedExtra); } } @@ -46,18 +52,45 @@ export class InMemoryTransport implements Transport { } /** - * Sends a message with optional auth info. - * This is useful for testing authentication scenarios. + * Sends a message with optional extra info. + * This is useful for testing authentication scenarios and custom context. + * + * @deprecated The authInfo parameter is deprecated. Use MessageExtraInfo instead. */ - async send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId, authInfo?: AuthInfo }): Promise { + async send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId, authInfo?: AuthInfo } | MessageExtraInfo): Promise { if (!this._otherTransport) { throw new Error("Not connected"); } + // Handle both old and new API formats + let extra: MessageExtraInfo | undefined; + if (options && 'authInfo' in options && !('requestInfo' in options)) { + // Old API format - convert to new format + extra = { authInfo: options.authInfo }; + } else if (options && ('requestInfo' in options || 'customContext' in options || 'authInfo' in options)) { + // New API format + extra = options as MessageExtraInfo; + } else if (options && 'authInfo' in options) { + // Old API with authInfo + extra = { authInfo: options.authInfo }; + } + if (this._otherTransport.onmessage) { - this._otherTransport.onmessage(message, { authInfo: options?.authInfo }); + // Merge the other transport's custom context with the extra info + const enhancedExtra: MessageExtraInfo = { + ...extra, + customContext: this._otherTransport._customContext + }; + this._otherTransport.onmessage(message, enhancedExtra); } else { - this._otherTransport._messageQueue.push({ message, extra: { authInfo: options?.authInfo } }); + this._otherTransport._messageQueue.push({ message, extra }); } } + + /** + * Sets custom context data that will be passed to all message handlers. + */ + setCustomContext(context: Record): void { + this._customContext = context; + } } diff --git a/src/server/mcp.test.ts b/src/server/mcp.test.ts index 10e550df4..4a7629af2 100644 --- a/src/server/mcp.test.ts +++ b/src/server/mcp.test.ts @@ -1433,6 +1433,64 @@ describe("tool()", () => { expect(result.content && result.content[0].text).toContain("Received request ID:"); }); + /*** + * Test: Pass Custom Context to Tool Callback + */ + test("should pass customContext to tool callback via RequestHandlerExtra", async () => { + const mcpServer = new McpServer({ + name: "test server", + version: "1.0", + }); + + const client = new Client({ + name: "test client", + version: "1.0", + }); + + let receivedCustomContext: Record | undefined; + mcpServer.tool("custom-context-test", async (extra) => { + receivedCustomContext = extra.customContext; + return { + content: [ + { + type: "text", + text: `Custom context: ${JSON.stringify(extra.customContext)}`, + }, + ], + }; + }); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + + // Use the new setCustomContext method to inject custom context + serverTransport.setCustomContext({ + tenantId: "test-tenant-123", + featureFlags: { newFeature: true }, + customData: "test-value" + }); + + await Promise.all([ + client.connect(clientTransport), + mcpServer.server.connect(serverTransport), + ]); + + const result = await client.request( + { + method: "tools/call", + params: { + name: "custom-context-test", + }, + }, + CallToolResultSchema, + ); + + expect(receivedCustomContext).toBeDefined(); + expect(receivedCustomContext?.tenantId).toBe("test-tenant-123"); + expect(receivedCustomContext?.featureFlags).toEqual({ newFeature: true }); + expect(receivedCustomContext?.customData).toBe("test-value"); + expect(result.content && result.content[0].text).toContain("test-tenant-123"); + }); + /*** * Test: Send Notification within Tool Call */ diff --git a/src/server/sse.ts b/src/server/sse.ts index e07256867..af623cc66 100644 --- a/src/server/sse.ts +++ b/src/server/sse.ts @@ -41,6 +41,7 @@ export class SSEServerTransport implements Transport { private _sseResponse?: ServerResponse; private _sessionId: string; private _options: SSEServerTransportOptions; + private _customContext?: Record; onclose?: () => void; onerror?: (error: Error) => void; onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; @@ -191,7 +192,12 @@ export class SSEServerTransport implements Transport { throw error; } - this.onmessage?.(parsedMessage, extra); + // Merge custom context with the extra info + const enhancedExtra: MessageExtraInfo = { + ...extra, + customContext: this._customContext + }; + this.onmessage?.(parsedMessage, enhancedExtra); } async close(): Promise { @@ -218,4 +224,11 @@ export class SSEServerTransport implements Transport { get sessionId(): string { return this._sessionId; } + + /** + * Sets custom context data that will be passed to all message handlers. + */ + setCustomContext(context: Record): void { + this._customContext = context; + } } diff --git a/src/server/stdio.ts b/src/server/stdio.ts index 30c80012e..42411df49 100644 --- a/src/server/stdio.ts +++ b/src/server/stdio.ts @@ -1,7 +1,7 @@ import process from "node:process"; import { Readable, Writable } from "node:stream"; import { ReadBuffer, serializeMessage } from "../shared/stdio.js"; -import { JSONRPCMessage } from "../types.js"; +import { JSONRPCMessage, MessageExtraInfo } from "../types.js"; import { Transport } from "../shared/transport.js"; /** @@ -12,6 +12,7 @@ import { Transport } from "../shared/transport.js"; export class StdioServerTransport implements Transport { private _readBuffer: ReadBuffer = new ReadBuffer(); private _started = false; + private _customContext?: Record; constructor( private _stdin: Readable = process.stdin, @@ -20,7 +21,7 @@ export class StdioServerTransport implements Transport { onclose?: () => void; onerror?: (error: Error) => void; - onmessage?: (message: JSONRPCMessage) => void; + onmessage?: (message: JSONRPCMessage, extra?: MessageExtraInfo) => void; // Arrow functions to bind `this` properly, while maintaining function identity. _ondata = (chunk: Buffer) => { @@ -54,7 +55,11 @@ export class StdioServerTransport implements Transport { break; } - this.onmessage?.(message); + // Pass custom context to message handlers + const extra: MessageExtraInfo = { + customContext: this._customContext + }; + this.onmessage?.(message, extra); } catch (error) { this.onerror?.(error as Error); } @@ -89,4 +94,11 @@ export class StdioServerTransport implements Transport { } }); } + + /** + * Sets custom context data that will be passed to all message handlers. + */ + setCustomContext(context: Record): void { + this._customContext = context; + } } diff --git a/src/server/streamableHttp.ts b/src/server/streamableHttp.ts index 3bf84e430..02283ba44 100644 --- a/src/server/streamableHttp.ts +++ b/src/server/streamableHttp.ts @@ -143,6 +143,7 @@ export class StreamableHTTPServerTransport implements Transport { private _allowedHosts?: string[]; private _allowedOrigins?: string[]; private _enableDnsRebindingProtection: boolean; + private _customContext?: Record; sessionId?: string; onclose?: () => void; @@ -487,7 +488,12 @@ export class StreamableHTTPServerTransport implements Transport { // handle each message for (const message of messages) { - this.onmessage?.(message, { authInfo, requestInfo }); + const enhancedExtra: MessageExtraInfo = { + authInfo, + requestInfo, + customContext: this._customContext + }; + this.onmessage?.(message, enhancedExtra); } } else if (hasRequests) { // The default behavior is to use SSE streaming @@ -522,7 +528,12 @@ export class StreamableHTTPServerTransport implements Transport { // handle each message for (const message of messages) { - this.onmessage?.(message, { authInfo, requestInfo }); + const enhancedExtra: MessageExtraInfo = { + authInfo, + requestInfo, + customContext: this._customContext + }; + this.onmessage?.(message, enhancedExtra); } // The server SHOULD NOT close the SSE stream before sending all JSON-RPC responses // This will be handled by the send() method when responses are ready @@ -748,5 +759,12 @@ export class StreamableHTTPServerTransport implements Transport { } } } + + /** + * Sets custom context data that will be passed to all message handlers. + */ + setCustomContext(context: Record): void { + this._customContext = context; + } } diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 6142140dd..8df0e879c 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -141,6 +141,13 @@ export type RequestHandlerExtra; + /** * Sends a notification that relates to the current request being handled. * @@ -399,7 +406,8 @@ export abstract class Protocol< this.request(r, resultSchema, { ...options, relatedRequestId: request.id }), authInfo: extra?.authInfo, requestId: request.id, - requestInfo: extra?.requestInfo + requestInfo: extra?.requestInfo, + customContext: extra?.customContext }; // Starting with Promise.resolve() puts any synchronous errors into the monad as well. diff --git a/src/shared/transport.ts b/src/shared/transport.ts index 386b6bae5..8dc671a39 100644 --- a/src/shared/transport.ts +++ b/src/shared/transport.ts @@ -82,4 +82,10 @@ export interface Transport { * Sets the protocol version used for the connection (called when the initialize response is received). */ setProtocolVersion?: (version: string) => void; + + /** + * Sets custom context data that will be passed to all message handlers. + * This context will be included in the MessageExtraInfo passed to handlers. + */ + setCustomContext?: (context: Record) => void; } diff --git a/src/types.ts b/src/types.ts index 323e37389..ef87f3665 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1512,6 +1512,13 @@ export interface MessageExtraInfo { * The authentication information. */ authInfo?: AuthInfo; + + /** + * Custom context data that can be passed through the message handling pipeline. + * This allows transport implementations to attach arbitrary data that will be + * available to request handlers. + */ + customContext?: Record; } /* JSON-RPC types */ From 5bee24e92cefae506b238654f17309d6c2f0c6ca Mon Sep 17 00:00:00 2001 From: Grimmer Kang Date: Fri, 15 Aug 2025 00:05:52 +0800 Subject: [PATCH 2/4] feat: Add custom context example implementation - Add server demonstrating API key authentication and context injection - Add interactive client for testing custom context features - Add comprehensive documentation with walkthrough - Update examples README to reference new examples The example demonstrates all three MCP handler types with context: - Tool (get_user) accesses context via extra.customContext - Prompt (user-dashboard) personalizes content via extra.customContext - Resource (user-profile) returns user data via extra.customContext Shows how to build secure, multi-tenant MCP applications with user authentication and permission-based access control. --- src/examples/README.md | 37 ++ src/examples/client/customContextClient.ts | 342 ++++++++++++++++ src/examples/custom-context-example.md | 228 +++++++++++ src/examples/server/customContextServer.ts | 440 +++++++++++++++++++++ 4 files changed, 1047 insertions(+) create mode 100644 src/examples/client/customContextClient.ts create mode 100644 src/examples/custom-context-example.md create mode 100644 src/examples/server/customContextServer.ts diff --git a/src/examples/README.md b/src/examples/README.md index ac92e8ded..fc80e021a 100644 --- a/src/examples/README.md +++ b/src/examples/README.md @@ -6,6 +6,7 @@ This directory contains example implementations of MCP clients and servers using - [Client Implementations](#client-implementations) - [Streamable HTTP Client](#streamable-http-client) + - [Custom Context Client](#custom-context-client) - [Backwards Compatible Client](#backwards-compatible-client) - [Server Implementations](#server-implementations) - [Single Node Deployment](#single-node-deployment) @@ -39,6 +40,26 @@ Example client with OAuth: npx tsx src/examples/client/simpleOAuthClient.js ``` +### Custom Context Client + +An interactive client that demonstrates the custom context feature, showing how to: + +- Authenticate using API keys that map to user contexts +- Pass user context (identity, permissions, organization) through the transport layer +- Access context in MCP tool handlers for authorization and personalization +- Implement multi-tenant data isolation +- Track requests with unique IDs for auditing + +```bash +# Start the server first: +npx tsx src/examples/server/customContextServer.ts + +# Then run the client: +npx tsx src/examples/client/customContextClient.ts +``` + +See [custom-context-example.md](custom-context-example.md) for a detailed walkthrough. + ### Backwards Compatible Client A client that implements backwards compatibility according to the [MCP specification](https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#backwards-compatibility), allowing it to work with both new and legacy servers. This client demonstrates: @@ -106,6 +127,22 @@ A server that demonstrates server notifications using Streamable HTTP. npx tsx src/examples/server/standaloneSseWithGetStreamableHttp.ts ``` +##### Custom Context Server + +A server that demonstrates how to inject custom context (user authentication, permissions, tenant data) into MCP tool handlers. + +- API key authentication with user context extraction +- Context injection via `transport.setCustomContext()` +- Permission-based access control in tools +- Multi-tenant data isolation +- Request tracking with unique IDs + +```bash +npx tsx src/examples/server/customContextServer.ts +``` + +This example is essential for building secure, multi-tenant MCP applications. See [custom-context-example.md](custom-context-example.md) for implementation details. + #### Deprecated SSE Transport A server that implements the deprecated HTTP+SSE transport (protocol version 2024-11-05). This example only used for testing backwards compatibility for clients. diff --git a/src/examples/client/customContextClient.ts b/src/examples/client/customContextClient.ts new file mode 100644 index 000000000..7c7666b88 --- /dev/null +++ b/src/examples/client/customContextClient.ts @@ -0,0 +1,342 @@ +import { Client } from '../../client/index.js'; +import { StreamableHTTPClientTransport } from '../../client/streamableHttp.js'; +import { SSEClientTransport } from '../../client/sse.js'; +import { createInterface } from 'node:readline'; +import { + CallToolResultSchema, + ListToolsResultSchema, + GetPromptResultSchema, + ListResourcesResultSchema, + ReadResourceResultSchema, +} from '../../types.js'; + +/** + * Interactive client demonstrating custom context feature. + * + * This client shows how API keys are used to authenticate and + * how the server uses the context to provide user-specific responses. + */ + +// Create readline interface for user input +const readline = createInterface({ + input: process.stdin, + output: process.stdout +}); + +// Global state +let client: Client | null = null; +let transport: StreamableHTTPClientTransport | SSEClientTransport | null = null; +const serverUrl = 'http://localhost:3000/mcp'; +let currentUser: { + name: string; + organization: { name: string }; + role: string; + permissions: string[]; +} | null = null; + +// Available API keys for testing +const API_KEYS: Record = { + 'alice': 'sk-alice-admin-key', + 'bob': 'sk-bob-dev-key', + 'charlie': 'sk-charlie-user-key', + 'dana': 'sk-dana-admin-key', +}; + +async function main(): Promise { + console.log('=============================================='); + console.log('MCP Custom Context Demo Client'); + console.log('=============================================='); + console.log('\nThis client demonstrates how custom context works:'); + console.log('1. Authenticate with an API key'); + console.log('2. The server fetches user context from the API key'); + console.log('3. Tools receive the context and respond based on user permissions\n'); + + printHelp(); + commandLoop(); +} + +function printHelp(): void { + console.log('\nšŸ“‹ Available commands:'); + console.log(' auth - Authenticate as user (alice/bob/charlie/dana)'); + console.log(' auth-key - Authenticate with custom API key'); + console.log(' whoami - Get current user info from context'); + console.log(' dashboard [format] - Get personalized dashboard (brief/detailed)'); + console.log(' profile - Read user profile resource'); + console.log(' list-tools - List available tools'); + console.log(' disconnect - Disconnect from server'); + console.log(' help - Show this help'); + console.log(' quit - Exit the program'); + console.log('\nšŸ”‘ Quick start: Try "auth alice" then "whoami"'); + console.log('\nāš ļø Note: Only the get_user tool is available in this simplified demo.'); +} + +function commandLoop(): void { + const prompt = currentUser ? `[${currentUser!.name}]> ` : '> '; + + readline.question(prompt, async (input) => { + const args = input.trim().split(/\s+/); + const command = args[0]?.toLowerCase(); + + try { + switch (command) { + case 'auth': { + const userName = args[1] as keyof typeof API_KEYS; + if (args.length < 2 || !API_KEYS[userName]) { + console.log('āŒ Usage: auth '); + console.log(' Available users:'); + console.log(' - alice: TechCorp Admin (all permissions)'); + console.log(' - bob: TechCorp Developer (code/docs permissions)'); + console.log(' - charlie: StartupIO User (limited permissions)'); + console.log(' - dana: StartupIO Admin (org admin)'); + } else { + await authenticateAs(userName); + } + break; + } + + case 'auth-key': + if (args.length < 2) { + console.log('āŒ Usage: auth-key '); + } else { + await authenticateWithKey(args[1]); + } + break; + + case 'whoami': + await getCurrentUser(); + break; + + + case 'dashboard': + await getDashboard(args[1] || 'brief'); + break; + + case 'profile': + await readProfile(); + break; + + case 'list-tools': + await listTools(); + break; + + case 'disconnect': + await disconnect(); + break; + + case 'help': + printHelp(); + break; + + case 'quit': + case 'exit': + await cleanup(); + return; + + default: + if (command) { + console.log(`ā“ Unknown command: ${command}`); + } + break; + } + } catch (error) { + console.error(`āŒ Error: ${error}`); + } + + // Continue the command loop + commandLoop(); + }); +} + +async function authenticateAs(userName: string): Promise { + const apiKey = API_KEYS[userName as keyof typeof API_KEYS]; + await authenticateWithKey(apiKey); +} + +async function authenticateWithKey(apiKey: string): Promise { + // Disconnect existing connection + if (client) { + await disconnect(); + } + + // Store the API key for this session (used in fetch) + console.log(`\nšŸ” Authenticating with API key: ${apiKey.substring(0, 15)}...`); + + // Create transport with API key in headers + transport = new StreamableHTTPClientTransport( + new URL(serverUrl), + { + fetch: async (url: string | URL, options?: RequestInit) => { + // Add API key to all requests + // Handle Headers object or plain object + let headers: HeadersInit; + if (options?.headers instanceof Headers) { + headers = new Headers(options.headers); + (headers as Headers).set('X-API-Key', apiKey); + } else { + headers = { + ...(options?.headers || {}), + 'X-API-Key': apiKey, + }; + } + return fetch(url, { ...options, headers }); + } + } + ); + + // Create and connect client + client = new Client({ + name: 'custom-context-demo-client', + version: '1.0.0' + }); + + try { + await client.connect(transport); + console.log('āœ… Connected to server'); + + // Get user info immediately after connecting + const result = await client.request({ + method: 'tools/call', + params: { + name: 'get_user', + arguments: {} + } + }, CallToolResultSchema); + + if (result.content && result.content[0]?.type === 'text') { + const text = result.content[0].text; + try { + // Parse user info from response + const userMatch = text.match(/User Profile:\n([\s\S]*)/); + if (userMatch) { + currentUser = JSON.parse(userMatch[1]); + console.log(`\nšŸ‘¤ Authenticated as: ${currentUser!.name}`); + console.log(` Organization: ${currentUser!.organization.name}`); + console.log(` Role: ${currentUser!.role}`); + console.log(` Permissions: ${currentUser!.permissions.length} permission(s)`); + } + } catch { + console.log('āœ… Authenticated (could not parse user details)'); + } + } + } catch (error) { + console.error(`āŒ Failed to connect: ${error}`); + client = null; + transport = null; + } +} + +async function getCurrentUser(): Promise { + if (!client) { + console.log('āŒ Not connected. Use "auth " first.'); + return; + } + + console.log('\nšŸ” Fetching user information from context...'); + + const result = await client.request({ + method: 'tools/call', + params: { + name: 'get_user', + arguments: {} + } + }, CallToolResultSchema); + + if (result.content && result.content[0]?.type === 'text') { + console.log('\n' + result.content[0].text); + } +} + + +async function getDashboard(format: string): Promise { + if (!client) { + console.log('āŒ Not connected. Use "auth " first.'); + return; + } + + console.log(`\nšŸ“Š Getting ${format} dashboard...`); + + const result = await client.request({ + method: 'prompts/get', + params: { + name: 'user-dashboard', + arguments: { format } + } + }, GetPromptResultSchema); + + if (result.messages && result.messages[0]?.content?.type === 'text') { + console.log('\n' + result.messages[0].content.text); + } +} + +async function readProfile(): Promise { + if (!client) { + console.log('āŒ Not connected. Use "auth " first.'); + return; + } + + console.log('\nšŸ“„ Reading user profile resource...'); + + try { + // Read the resource directly using the known URI + const result = await client.request({ + method: 'resources/read', + params: { uri: 'user://profile' } + }, ReadResourceResultSchema); + + if (result.contents && result.contents[0]) { + const content = result.contents[0]; + console.log(`\nšŸ“„ Resource: ${content.uri}`); + if (content.mimeType) { + console.log(`Type: ${content.mimeType}`); + } + console.log('Content:'); + console.log(content.text || content.blob); + } + } catch (error) { + console.log(`āŒ Error reading profile: ${error}`); + } +} + +async function listTools(): Promise { + if (!client) { + console.log('āŒ Not connected. Use "auth " first.'); + return; + } + + const result = await client.request({ + method: 'tools/list', + params: {} + }, ListToolsResultSchema); + + console.log('\nšŸ”§ Available tools:'); + for (const tool of result.tools) { + console.log(` - ${tool.name}: ${tool.description}`); + } +} + +async function disconnect(): Promise { + if (client) { + await client.close(); + client = null; + transport = null; + currentUser = null; + console.log('āœ… Disconnected from server'); + } else { + console.log('āŒ Not connected'); + } +} + +async function cleanup(): Promise { + await disconnect(); + console.log('\nšŸ‘‹ Goodbye!'); + readline.close(); + process.exit(0); +} + +// Handle ctrl+c +process.on('SIGINT', async () => { + await cleanup(); +}); + +// Start the client +main().catch(console.error); \ No newline at end of file diff --git a/src/examples/custom-context-example.md b/src/examples/custom-context-example.md new file mode 100644 index 000000000..871c5edb0 --- /dev/null +++ b/src/examples/custom-context-example.md @@ -0,0 +1,228 @@ +# Custom Context Feature Demo + +This example demonstrates the **Custom Context** feature that allows MCP servers to inject contextual information (like user authentication, permissions, tenant data) into tool handlers. + +## What is Custom Context? + +Custom Context allows transport implementations to attach arbitrary data that will be available to all request handlers (tools, prompts, resources). This is essential for: + +- **Authentication**: Pass user identity from API keys or tokens +- **Multi-tenancy**: Isolate data between different organizations +- **Permissions**: Enforce access control based on user roles +- **Request tracking**: Add request IDs for debugging and auditing + +## Running the Demo + +### 1. Start the Server + +```bash +# From the typescript-sdk directory +npm run build +node dist/examples/server/customContextServer.js +``` + +The server will start on port 3000 and display available API keys for testing. + +### 2. Start the Client + +In a new terminal: + +```bash +# From the typescript-sdk directory +node dist/examples/client/customContextClient.js +``` + +## Demo Walkthrough + +### Step 1: Authenticate with an API Key + +The server simulates a database of API keys that map to user contexts. Try authenticating as different users: + +```bash +> auth alice +šŸ” Authenticating with API key: sk-alice-admin... +āœ… Connected to server +šŸ‘¤ Authenticated as: Alice Anderson + Organization: TechCorp Industries + Role: admin + Permissions: 4 permission(s) +``` + +### Step 2: Get User Information + +The `get_user` tool retrieves the authenticated user's information from the context: + +```bash +[Alice Anderson]> whoami +šŸ” Fetching user information from context... + +User Profile: +{ + "userId": "user-001", + "name": "Alice Anderson", + "email": "alice@techcorp.com", + "role": "admin", + "organization": { + "id": "org-techcorp", + "name": "TechCorp Industries" + }, + "permissions": ["read:all", "write:all", "delete:all", "admin:users"], + "accountCreated": "2024-01-15T08:00:00Z", + "lastActive": "2024-07-29T12:34:56Z" +} +``` + +### Step 3: Get Personalized Dashboard + +The `dashboard` command uses the prompt feature to generate personalized content: + +```bash +> auth alice +[Alice Anderson]> dashboard +šŸ“Š Getting brief dashboard... + +Welcome back, Alice Anderson! You have access to 3 projects in TechCorp Industries. + +[Alice Anderson]> dashboard detailed +šŸ“Š Getting detailed dashboard... + +Dashboard for Alice Anderson + +Organization: TechCorp Industries +Role: admin +Plan: enterprise +Projects: 3 +Permissions: read:all, write:all, delete:all, admin:users +Member since: 2024-01-15T08:00:00Z + +Your organization has 2 members and is on the enterprise plan. +``` + +### Step 4: Access User Profile Resource + +The `profile` command demonstrates resource access with context: + +```bash +> auth bob +[Bob Builder]> profile +šŸ“„ Reading user profile resource... + +šŸ“„ Resource: user://profile/user-002 +Type: application/json +Content: +{ + "user": { + "id": "user-002", + "name": "Bob Builder", + "email": "bob@techcorp.com", + "role": "developer", + "createdAt": "2024-02-20T10:30:00Z" + }, + "organization": { + "id": "org-techcorp", + "name": "TechCorp Industries", + "plan": "enterprise", + "memberCount": 2, + "projectCount": 3 + }, + "permissions": [ + "read:code", + "write:code", + "read:docs", + "write:docs" + ], + "apiKey": { + "id": "sk-bob-dev-key", + "lastUsed": "2025-08-14T16:12:57.561Z" + } +} +``` + +## How It Works + +### Server Side + +1. **API Key Extraction**: The server extracts the API key from request headers +2. **Context Fetching**: Uses the API key to fetch user context from a "database" +3. **Context Injection**: Calls `transport.setCustomContext(userContext)` +4. **Tool Access**: Tools receive context via `extra.customContext` + +```typescript +// In the server +const context = fetchUserContext(apiKey); +transport.setCustomContext(context); + +// In tool handlers +async (params, extra) => { + const context = extra.customContext as UserContext; + if (!context.permissions.includes('required:permission')) { + return { content: [{ type: 'text', text: 'Access denied' }] }; + } + // ... perform action +} +``` + +### Client Side + +The client adds the API key to all requests: + +```typescript +const transport = new StreamableHTTPClientTransport({ + url: serverUrl, + fetch: async (url, options) => { + const headers = { + ...options?.headers, + 'X-API-Key': apiKey, + }; + return fetch(url, { ...options, headers }); + } +}); +``` + +## Available Test Users + +| User | API Key | Organization | Role | Key Permissions | +|------|---------|--------------|------|-----------------| +| Alice | sk-alice-admin-key | TechCorp | admin | All permissions | +| Bob | sk-bob-dev-key | TechCorp | developer | Code & docs access | +| Charlie | sk-charlie-user-key | StartupIO | user | Limited read/write | +| Dana | sk-dana-admin-key | StartupIO | admin | Org administration | + +## Key Features Demonstrated + +1. **Authentication**: API key-based user authentication +2. **Tool Context**: `get_user` tool accesses user context +3. **Prompt Context**: `user-dashboard` prompt personalizes content based on context +4. **Resource Context**: `user-profile` resource returns context-aware data +5. **Multi-tenancy**: Data isolation between organizations +6. **Request Tracking**: Unique request IDs for auditing + +## Real-World Applications + +This pattern is essential for: + +- **SaaS Applications**: Isolate customer data in multi-tenant systems +- **Enterprise Tools**: Enforce role-based permissions +- **API Services**: Track usage and enforce rate limits per user +- **Audit Logging**: Track who did what and when +- **Personalization**: Customize responses based on user preferences + +## Code Structure + +- `customContextServer.ts`: Server with API key authentication and context injection +- `customContextClient.ts`: Interactive REPL client that sends API keys +- Tool (`get_user`) demonstrates context access +- Prompt (`user-dashboard`) shows personalized content +- Resource (`user-profile`) returns user-specific data + +## Next Steps + +To integrate custom context in your own MCP server: + +1. Define your context interface +2. Extract authentication info from requests (API keys, JWT tokens, etc.) +3. Call `transport.setCustomContext(context)` with user data +4. Access context in handlers via `extra.customContext` +5. Implement permission checking and data filtering based on context + +The custom context feature enables building secure, multi-tenant MCP applications with proper authentication and authorization. \ No newline at end of file diff --git a/src/examples/server/customContextServer.ts b/src/examples/server/customContextServer.ts new file mode 100644 index 000000000..00d0cf7d4 --- /dev/null +++ b/src/examples/server/customContextServer.ts @@ -0,0 +1,440 @@ +import express, { Request, Response } from 'express'; +import { randomUUID } from 'node:crypto'; +import { z } from 'zod'; +import { McpServer } from '../../server/mcp.js'; +import { StreamableHTTPServerTransport } from '../../server/streamableHttp.js'; +import { CallToolResult, GetPromptResult } from '../../types.js'; +import cors from 'cors'; + +/** + * Example server demonstrating custom context feature. + * + * This server simulates API key authentication where: + * - Each request includes an API key + * - The API key is used to fetch user context from a "database" + * - Tools can access the authenticated user's information + * - Different users have different permissions and data access + * + * The custom context includes: + * - userId: The authenticated user + * - email: User's email + * - organizationId: User's organization + * - role: User's role (admin, developer, user) + * - permissions: What the user is allowed to do + * - apiKeyId: The API key used + * - requestId: For tracking and logging + */ + +// Custom context interface for authenticated users +// Using type instead of interface to satisfy Record constraint +type UserContext = { + userId: string; + email: string; + name: string; + organizationId: string; + organizationName: string; + role: 'admin' | 'developer' | 'user'; + permissions: string[]; + apiKeyId: string; + requestId: string; + createdAt: string; + lastActive: string; + [key: string]: unknown; // Index signature for Record compatibility +} + +// Simulated API key database - maps API keys to user contexts +const apiKeyDatabase: Record = { + 'sk-alice-admin-key': { + userId: 'user-001', + email: 'alice@techcorp.com', + name: 'Alice Anderson', + organizationId: 'org-techcorp', + organizationName: 'TechCorp Industries', + role: 'admin', + permissions: ['read:all', 'write:all', 'delete:all', 'admin:users'], + apiKeyId: 'sk-alice-admin-key', + requestId: '', // Will be set per request + createdAt: '2024-01-15T08:00:00Z', + lastActive: new Date().toISOString(), + }, + 'sk-bob-dev-key': { + userId: 'user-002', + email: 'bob@techcorp.com', + name: 'Bob Builder', + organizationId: 'org-techcorp', + organizationName: 'TechCorp Industries', + role: 'developer', + permissions: ['read:code', 'write:code', 'read:docs', 'write:docs'], + apiKeyId: 'sk-bob-dev-key', + requestId: '', + createdAt: '2024-02-20T10:30:00Z', + lastActive: new Date().toISOString(), + }, + 'sk-charlie-user-key': { + userId: 'user-003', + email: 'charlie@startup.io', + name: 'Charlie Chen', + organizationId: 'org-startup', + organizationName: 'StartupIO', + role: 'user', + permissions: ['read:public', 'write:own'], + apiKeyId: 'sk-charlie-user-key', + requestId: '', + createdAt: '2024-03-10T14:15:00Z', + lastActive: new Date().toISOString(), + }, + 'sk-dana-admin-key': { + userId: 'user-004', + email: 'dana@startup.io', + name: 'Dana Davis', + organizationId: 'org-startup', + organizationName: 'StartupIO', + role: 'admin', + permissions: ['read:all', 'write:all', 'admin:organization'], + apiKeyId: 'sk-dana-admin-key', + requestId: '', + createdAt: '2024-01-01T00:00:00Z', + lastActive: new Date().toISOString(), + }, +}; + +// Simulated organization data +const organizationData: Record; + usage: { apiCalls: number; storage: number; }; +}> = { + 'org-techcorp': { + name: 'TechCorp Industries', + plan: 'enterprise', + members: ['user-001', 'user-002'], + projects: [ + { id: 'proj-001', name: 'Main Platform', visibility: 'private' }, + { id: 'proj-002', name: 'Public API', visibility: 'public' }, + { id: 'proj-003', name: 'Internal Tools', visibility: 'private' }, + ], + usage: { apiCalls: 150000, storage: 2048 }, + }, + 'org-startup': { + name: 'StartupIO', + plan: 'pro', + members: ['user-003', 'user-004'], + projects: [ + { id: 'proj-004', name: 'MVP Product', visibility: 'private' }, + { id: 'proj-005', name: 'Documentation', visibility: 'public' }, + ], + usage: { apiCalls: 25000, storage: 512 }, + }, +}; + +// Create an MCP server with custom context support +const getServer = () => { + const server = new McpServer({ + name: 'custom-context-demo-server', + version: '1.0.0' + }, { + capabilities: { + logging: {}, + prompts: {}, + resources: {}, + tools: {} + } + }); + + // Tool: Get current user information from context + server.registerTool( + 'get_user', + { + title: 'Get Current User', + description: 'Returns information about the currently authenticated user from the context', + inputSchema: {}, + }, + async (_, extra): Promise => { + const context = extra.customContext as UserContext | undefined; + + if (!context) { + return { + content: [{ + type: 'text', + text: 'Error: No authentication context found. Please provide a valid API key.', + }], + }; + } + + console.log(`[${context.requestId}] User ${context.name} (${context.userId}) accessed their profile`); + + // Return the user's context information + const userInfo = { + userId: context.userId, + name: context.name, + email: context.email, + role: context.role, + organization: { + id: context.organizationId, + name: context.organizationName, + }, + permissions: context.permissions, + accountCreated: context.createdAt, + lastActive: new Date().toISOString(), + }; + + return { + content: [{ + type: 'text', + text: `User Profile:\n${JSON.stringify(userInfo, null, 2)}`, + }], + }; + } + ); + + // Prompts can also access context for personalization + server.registerPrompt( + 'user-dashboard', + { + title: 'User Dashboard Summary', + description: 'Generates a personalized dashboard summary based on user context', + argsSchema: { + format: z.enum(['detailed', 'brief']).optional().describe('Summary format'), + }, + }, + async ({ format = 'brief' }, extra): Promise => { + const context = extra.customContext as UserContext | undefined; + + if (!context) { + return { + messages: [{ + role: 'user', + content: { + type: 'text', + text: 'Please authenticate to view your dashboard.', + }, + }], + }; + } + + const org = organizationData[context.organizationId]; + const projectCount = org?.projects.length || 0; + const plan = org?.plan || 'unknown'; + + let message: string; + if (format === 'detailed') { + message = `Dashboard for ${context.name}\n\n` + + `Organization: ${context.organizationName}\n` + + `Role: ${context.role}\n` + + `Plan: ${plan}\n` + + `Projects: ${projectCount}\n` + + `Permissions: ${context.permissions.join(', ')}\n` + + `Member since: ${context.createdAt}\n\n` + + `Your organization has ${org?.members.length || 0} members and is on the ${plan} plan.`; + } else { + message = `Welcome back, ${context.name}! You have access to ${projectCount} projects in ${context.organizationName}.`; + } + + return { + messages: [{ + role: 'user', + content: { + type: 'text', + text: message, + }, + }], + }; + } + ); + + // Resource example - testing context support + server.registerResource( + 'user-profile', + 'user://profile', + { + title: 'User Profile Resource', + description: 'View authenticated user profile from context', + mimeType: 'application/json' + }, + async (uri, extra) => { + const context = extra.customContext as UserContext | undefined; + + if (!context) { + return { + contents: [{ + uri: 'user://profile/error', + text: 'Authentication required', + mimeType: 'text/plain', + }], + }; + } + + const org = organizationData[context.organizationId]; + const profile = { + user: { + id: context.userId, + name: context.name, + email: context.email, + role: context.role, + createdAt: context.createdAt, + }, + organization: { + id: context.organizationId, + name: context.organizationName, + plan: org?.plan, + memberCount: org?.members.length, + projectCount: org?.projects.length, + }, + permissions: context.permissions, + apiKey: { + id: context.apiKeyId, + lastUsed: new Date().toISOString(), + }, + }; + + return { + contents: [{ + uri: `user://profile/${context.userId}`, + text: JSON.stringify(profile, null, 2), + mimeType: 'application/json', + }], + }; + } + ); + + return server; +}; + +// Express app setup +const MCP_PORT = process.env.MCP_PORT ? parseInt(process.env.MCP_PORT, 10) : 3000; +const app = express(); +app.use(express.json()); +app.use(cors({ + origin: '*', + exposedHeaders: ["Mcp-Session-Id"] +})); + +// Map to store transports by session ID +const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; + +// Function to fetch user context from API key +const fetchUserContext = (apiKey: string): UserContext | undefined => { + // In a real application, this would query a database + const baseContext = apiKeyDatabase[apiKey]; + + if (!baseContext) { + console.log(`Invalid API key: ${apiKey?.substring(0, 10)}...`); + return undefined; + } + + // Create a fresh context with new requestId and updated lastActive + return { + ...baseContext, + requestId: randomUUID(), + lastActive: new Date().toISOString(), + }; +}; + +// Middleware to extract user context from API key +const extractUserContext = (req: Request): UserContext | undefined => { + // Check for API key in various places + const apiKey = + req.headers['x-api-key'] as string || + req.headers['authorization']?.replace('Bearer ', '') as string || + (req.query?.api_key as string); + + if (!apiKey) { + console.log('No API key provided'); + return undefined; + } + + return fetchUserContext(apiKey); +}; + +// MCP endpoints with custom context injection +app.post('/mcp', async (req: Request, res: Response) => { + const sessionId = req.headers['mcp-session-id'] as string | undefined; + const context = extractUserContext(req); + + console.log(`Received request: Session=${sessionId}, User=${context?.name}, Org=${context?.organizationName}`); + + try { + let transport: StreamableHTTPServerTransport; + + if (sessionId && transports[sessionId]) { + transport = transports[sessionId]; + // Update context for existing session + if (context) { + transport.setCustomContext(context as Record); + } + } else { + // New session - create transport with initial context + transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: (sid) => { + console.log(`Session initialized: ${sid}`); + transports[sid] = transport; + } + }); + + // Set custom context if available + if (context) { + transport.setCustomContext(context as Record); + } + + transport.onclose = () => { + const sid = transport.sessionId; + if (sid && transports[sid]) { + console.log(`Session closed: ${sid}`); + delete transports[sid]; + } + }; + + const server = getServer(); + await server.connect(transport); + } + + await transport.handleRequest(req, res, req.body); + } catch (error) { + console.error('Error handling request:', error); + if (!res.headersSent) { + res.status(500).json({ + jsonrpc: '2.0', + error: { + code: -32603, + message: 'Internal server error', + }, + id: null, + }); + } + } +}); + +app.listen(MCP_PORT, () => { + console.log(`Custom Context MCP Server running on port ${MCP_PORT}`); + console.log('\nThis server demonstrates custom context features:'); + console.log('- API key authentication'); + console.log('- User context injection from API key'); + console.log('- Permission-based access control'); + console.log('- Organization data isolation'); + console.log('- Request tracking with unique IDs'); + console.log('\nAvailable API keys for testing:'); + console.log(' sk-alice-admin-key (Alice - TechCorp Admin)'); + console.log(' sk-bob-dev-key (Bob - TechCorp Developer)'); + console.log(' sk-charlie-user-key (Charlie - StartupIO User)'); + console.log(' sk-dana-admin-key (Dana - StartupIO Admin)'); + console.log('\nSend API key via:'); + console.log(' Header: X-API-Key: '); + console.log(' Header: Authorization: Bearer '); + console.log(' Query: ?api_key='); +}); + +process.on('SIGINT', async () => { + console.log('\nShutting down server...'); + for (const sessionId in transports) { + try { + await transports[sessionId].close(); + delete transports[sessionId]; + } catch (error) { + console.error(`Error closing transport: ${error}`); + } + } + process.exit(0); +}); \ No newline at end of file From 40739fd1f6276907842e10429b1e02374a8a41f0 Mon Sep 17 00:00:00 2001 From: Grimmer Kang Date: Fri, 15 Aug 2025 02:14:05 +0800 Subject: [PATCH 3/4] fix: Remove unused import to fix lint error Remove unused ListResourcesResultSchema import from customContextClient.ts --- src/examples/client/customContextClient.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/src/examples/client/customContextClient.ts b/src/examples/client/customContextClient.ts index 7c7666b88..c817599be 100644 --- a/src/examples/client/customContextClient.ts +++ b/src/examples/client/customContextClient.ts @@ -6,7 +6,6 @@ import { CallToolResultSchema, ListToolsResultSchema, GetPromptResultSchema, - ListResourcesResultSchema, ReadResourceResultSchema, } from '../../types.js'; From fa6265ec185cf311d9d494c638bdd41eda7762f8 Mon Sep 17 00:00:00 2001 From: Grimmer Kang Date: Sat, 16 Aug 2025 09:34:24 +0900 Subject: [PATCH 4/4] docs: clarify custom context supports MCP access tokens alongside API keys - Updated documentation to explicitly mention MCP access tokens (OAuth flow) - Simplified examples to focus on API keys and MCP tokens only - Removed mentions of other auth methods to avoid confusion - Emphasized that the pattern works for both API keys and MCP OAuth tokens --- python-sdk-custom-context-analysis.md | 253 +++++++++++++++++++++ src/examples/client/customContextClient.ts | 15 +- src/examples/custom-context-example.md | 35 ++- src/examples/server/customContextServer.ts | 23 +- 4 files changed, 305 insertions(+), 21 deletions(-) create mode 100644 python-sdk-custom-context-analysis.md diff --git a/python-sdk-custom-context-analysis.md b/python-sdk-custom-context-analysis.md new file mode 100644 index 000000000..dc88083cd --- /dev/null +++ b/python-sdk-custom-context-analysis.md @@ -0,0 +1,253 @@ +# Python MCP SDK Custom Context Analysis + +## Executive Summary + +The Python MCP SDK **lacks built-in support for custom context injection** similar to what was added to the TypeScript SDK. While it does provide access to the raw HTTP request object in handlers, there's no clean mechanism to inject processed custom context (e.g., user authentication data, permissions, tenant information) that can be accessed by tool/prompt/resource handlers. + +## Current State of Python SDK + +### Context Architecture + +1. **RequestContext Class** (`src/mcp/shared/context.py`): +```python +@dataclass +class RequestContext(Generic[SessionT, LifespanContextT, RequestT]): + request_id: RequestId + meta: RequestParams.Meta | None + session: SessionT + lifespan_context: LifespanContextT + request: RequestT | None = None # This is where custom data could go +``` + +2. **Context Access in Handlers**: +```python +@app.call_tool() +async def my_tool(name: str, arguments: dict) -> list[types.ContentBlock]: + ctx = app.request_context # Access context + # ctx.request contains the Starlette Request object + # ctx.session, ctx.request_id, ctx.lifespan_context are available +``` + +3. **Transport Layer** (`src/mcp/server/streamable_http.py`): + - Line 385-386: Creates `ServerMessageMetadata` with `request_context=request` + - The raw Starlette Request object is passed as the context + - No mechanism to inject processed custom data + +### Key Differences from TypeScript SDK + +| Aspect | TypeScript SDK | Python SDK | +|--------|---------------|------------| +| Custom Context Method | `transport.setCustomContext()` | None | +| Context Access | `extra.customContext` | `app.request_context.request` | +| Context Type | Arbitrary object | Starlette Request object | +| Processing | Transport can inject processed data | Only raw HTTP request available | +| Type Safety | Can define custom types | Limited to Request type | + +## Problems with Current Python SDK + +1. **No Clean Context Injection**: Handlers receive the raw HTTP request but there's no way to inject processed context +2. **Authentication Complexity**: Every handler would need to extract and validate authentication from headers +3. **No Abstraction**: Tight coupling to HTTP transport (Starlette Request) +4. **Repeated Logic**: Authentication/authorization logic must be duplicated in each handler +5. **Limited Flexibility**: Can't easily inject tenant data, user permissions, or other contextual information + +## Proposed Fix Plan + +### Option 1: Minimal Change - Add Custom Context Field (Recommended) + +Add a `custom_context` field to `RequestContext` and provide a way for transports to set it: + +#### 1. Update RequestContext (`src/mcp/shared/context.py`): +```python +@dataclass +class RequestContext(Generic[SessionT, LifespanContextT, RequestT]): + request_id: RequestId + meta: RequestParams.Meta | None + session: SessionT + lifespan_context: LifespanContextT + request: RequestT | None = None + custom_context: Any | None = None # NEW: Custom context field +``` + +#### 2. Update ServerMessageMetadata (`src/mcp/shared/message.py`): +```python +@dataclass +class ServerMessageMetadata: + related_request_id: RequestId | None = None + request_context: Any | None = None + custom_context: Any | None = None # NEW: Custom context field +``` + +#### 3. Update StreamableHTTPServerTransport (`src/mcp/server/streamable_http.py`): +Add a method to set custom context and use it in request handling: + +```python +class StreamableHTTPServerTransport: + def __init__(self, ...): + # ... existing code ... + self._custom_context: Any | None = None + + def set_custom_context(self, context: Any) -> None: + """Set custom context to be passed to handlers.""" + self._custom_context = context + + async def _handle_post_request(self, ...): + # ... existing code ... + # Line ~385, update metadata creation: + metadata = ServerMessageMetadata( + request_context=request, + custom_context=self._custom_context # NEW: Include custom context + ) +``` + +#### 4. Update Server (`src/mcp/server/lowlevel/server.py`): +Pass custom context to RequestContext: + +```python +async def _handle_request(self, ...): + # ... existing code ... + # Extract custom context from metadata + custom_context = None + if message.message_metadata and isinstance(message.message_metadata, ServerMessageMetadata): + request_data = message.message_metadata.request_context + custom_context = message.message_metadata.custom_context # NEW + + # Set context with custom data + token = request_ctx.set( + RequestContext( + message.request_id, + message.request_meta, + session, + lifespan_context, + request=request_data, + custom_context=custom_context # NEW: Pass custom context + ) + ) +``` + +#### 5. Add Middleware Support in StreamableHTTPSessionManager: +```python +class StreamableHTTPSessionManager: + def __init__(self, ..., context_middleware: Callable[[Request], Awaitable[Any]] | None = None): + self.context_middleware = context_middleware + + async def _handle_stateful_request(self, ...): + # ... existing code ... + # Before creating transport, process context + custom_context = None + if self.context_middleware: + custom_context = await self.context_middleware(request) + + # Pass custom context to transport + http_transport = StreamableHTTPServerTransport(...) + if custom_context: + http_transport.set_custom_context(custom_context) +``` + +### Option 2: Full Middleware Architecture + +Create a more comprehensive middleware system similar to Express.js or FastAPI: + +1. Define middleware interface +2. Allow chaining of middleware functions +3. Support both sync and async middleware +4. Provide built-in authentication middleware + +This is more complex but provides greater flexibility. + +### Option 3: Subclass-Based Approach + +Allow users to subclass `StreamableHTTPServerTransport` and override context extraction: + +```python +class CustomHTTPTransport(StreamableHTTPServerTransport): + async def extract_context(self, request: Request) -> Any: + # Custom logic to extract and process context + api_key = request.headers.get("X-API-Key") + return await fetch_user_context(api_key) +``` + +## Implementation Priority + +1. **Phase 1**: Implement Option 1 (Minimal Change) - Adds basic custom context support +2. **Phase 2**: Add helper utilities for common patterns (auth extraction, validation) +3. **Phase 3**: Consider full middleware architecture if needed + +## Example Usage After Fix + +```python +# Server setup with custom context +async def context_middleware(request: Request) -> dict: + """Extract and validate user context from request.""" + api_key = request.headers.get("X-API-Key") + if not api_key: + return None + + # Fetch user data from database + user_data = await fetch_user_by_api_key(api_key) + return { + "user_id": user_data["id"], + "email": user_data["email"], + "permissions": user_data["permissions"], + "organization_id": user_data["org_id"] + } + +# Initialize session manager with middleware +session_manager = StreamableHTTPSessionManager( + app=app, + context_middleware=context_middleware +) + +# In tool handlers +@app.call_tool() +async def my_tool(name: str, arguments: dict) -> list[types.ContentBlock]: + ctx = app.request_context + user_context = ctx.custom_context # Access custom context + + if not user_context: + return [types.TextContent(type="text", text="Not authenticated")] + + if "admin" not in user_context.get("permissions", []): + return [types.TextContent(type="text", text="Permission denied")] + + # Proceed with tool logic + return [types.TextContent( + type="text", + text=f"Hello {user_context['email']}!" + )] +``` + +## Benefits of Proposed Solution + +1. **Clean Separation**: Authentication logic separated from business logic +2. **Type Safety**: Can use TypedDict or dataclasses for context types +3. **Reusability**: Context extraction logic in one place +4. **Transport Agnostic**: Works with any transport that supports context +5. **Backward Compatible**: Existing code continues to work +6. **Minimal Changes**: Small, focused changes to core SDK + +## Migration Path + +1. Changes are backward compatible - existing handlers continue working +2. New `custom_context` field is optional +3. Gradual adoption - handlers can migrate to use custom context as needed +4. Documentation and examples to guide migration + +## Testing Strategy + +1. Unit tests for context injection and retrieval +2. Integration tests with authentication middleware +3. Example server demonstrating custom context usage +4. Performance tests to ensure no regression + +## Conclusion + +The Python MCP SDK currently lacks the custom context injection capability that was recently added to the TypeScript SDK. The proposed fix (Option 1) provides a minimal, backward-compatible solution that brings feature parity between the two SDKs while maintaining the Python SDK's design principles. + +The implementation is straightforward and can be completed in a few hours, with most of the work involving: +1. Adding fields to existing dataclasses +2. Passing context through the call chain +3. Adding a context middleware hook +4. Creating examples and documentation + +This would enable Python MCP servers to properly handle authentication, multi-tenancy, and permission-based access control in a clean, maintainable way. \ No newline at end of file diff --git a/src/examples/client/customContextClient.ts b/src/examples/client/customContextClient.ts index c817599be..14ab9cad5 100644 --- a/src/examples/client/customContextClient.ts +++ b/src/examples/client/customContextClient.ts @@ -12,7 +12,10 @@ import { /** * Interactive client demonstrating custom context feature. * - * This client shows how API keys are used to authenticate and + * This example uses API keys for authentication, but the same pattern works + * with MCP access tokens from the OAuth flow. + * + * The client shows how authentication credentials are sent with requests and * how the server uses the context to provide user-specific responses. */ @@ -46,8 +49,8 @@ async function main(): Promise { console.log('MCP Custom Context Demo Client'); console.log('=============================================='); console.log('\nThis client demonstrates how custom context works:'); - console.log('1. Authenticate with an API key'); - console.log('2. The server fetches user context from the API key'); + console.log('1. Authenticate with credentials (API key or MCP access token)'); + console.log('2. The server validates credentials and fetches user context'); console.log('3. Tools receive the context and respond based on user permissions\n'); printHelp(); @@ -160,12 +163,14 @@ async function authenticateWithKey(apiKey: string): Promise { // Store the API key for this session (used in fetch) console.log(`\nšŸ” Authenticating with API key: ${apiKey.substring(0, 15)}...`); - // Create transport with API key in headers + // Create transport with authentication credentials in headers + // This example uses API key, but you could also use MCP access tokens transport = new StreamableHTTPClientTransport( new URL(serverUrl), { fetch: async (url: string | URL, options?: RequestInit) => { - // Add API key to all requests + // Add authentication credentials to all requests + // For MCP access token: use Authorization header instead of X-API-Key // Handle Headers object or plain object let headers: HeadersInit; if (options?.headers instanceof Headers) { diff --git a/src/examples/custom-context-example.md b/src/examples/custom-context-example.md index 871c5edb0..d732e6ce1 100644 --- a/src/examples/custom-context-example.md +++ b/src/examples/custom-context-example.md @@ -6,7 +6,7 @@ This example demonstrates the **Custom Context** feature that allows MCP servers Custom Context allows transport implementations to attach arbitrary data that will be available to all request handlers (tools, prompts, resources). This is essential for: -- **Authentication**: Pass user identity from API keys or tokens +- **Authentication**: Pass user identity from API keys or MCP access tokens - **Multi-tenancy**: Isolate data between different organizations - **Permissions**: Enforce access control based on user roles - **Request tracking**: Add request IDs for debugging and auditing @@ -142,17 +142,23 @@ Content: ### Server Side -1. **API Key Extraction**: The server extracts the API key from request headers -2. **Context Fetching**: Uses the API key to fetch user context from a "database" +1. **Authentication Extraction**: The server extracts authentication credentials from request headers +2. **Context Fetching**: Uses the credentials to fetch/validate user context 3. **Context Injection**: Calls `transport.setCustomContext(userContext)` 4. **Tool Access**: Tools receive context via `extra.customContext` ```typescript -// In the server -const context = fetchUserContext(apiKey); +// Example with API key (shown in demo) +const apiKey = request.headers.get('x-api-key'); +const context = await fetchUserContextByApiKey(apiKey); transport.setCustomContext(context); -// In tool handlers +// Example with MCP access token (OAuth flow) +const accessToken = request.headers.get('authorization')?.replace('Bearer ', ''); +const context = await validateMcpAccessToken(accessToken); +transport.setCustomContext(context); + +// In tool handlers (same regardless of auth method) async (params, extra) => { const context = extra.customContext as UserContext; if (!context.permissions.includes('required:permission')) { @@ -164,9 +170,10 @@ async (params, extra) => { ### Client Side -The client adds the API key to all requests: +The client adds authentication credentials to all requests: ```typescript +// Example with API key (shown in demo) const transport = new StreamableHTTPClientTransport({ url: serverUrl, fetch: async (url, options) => { @@ -177,6 +184,18 @@ const transport = new StreamableHTTPClientTransport({ return fetch(url, { ...options, headers }); } }); + +// Example with MCP access token (OAuth flow) +const transport = new StreamableHTTPClientTransport({ + url: serverUrl, + fetch: async (url, options) => { + const headers = { + ...options?.headers, + 'Authorization': `Bearer ${mcpAccessToken}`, + }; + return fetch(url, { ...options, headers }); + } +}); ``` ## Available Test Users @@ -190,7 +209,7 @@ const transport = new StreamableHTTPClientTransport({ ## Key Features Demonstrated -1. **Authentication**: API key-based user authentication +1. **Authentication**: Supports both API keys and MCP access tokens 2. **Tool Context**: `get_user` tool accesses user context 3. **Prompt Context**: `user-dashboard` prompt personalizes content based on context 4. **Resource Context**: `user-profile` resource returns context-aware data diff --git a/src/examples/server/customContextServer.ts b/src/examples/server/customContextServer.ts index 00d0cf7d4..bbd50e142 100644 --- a/src/examples/server/customContextServer.ts +++ b/src/examples/server/customContextServer.ts @@ -9,9 +9,12 @@ import cors from 'cors'; /** * Example server demonstrating custom context feature. * - * This server simulates API key authentication where: - * - Each request includes an API key - * - The API key is used to fetch user context from a "database" + * This example uses API key authentication for simplicity, but the same pattern + * works with MCP access tokens from the OAuth flow. + * + * The authentication flow: + * - Each request includes authentication credentials (API key or MCP access token) + * - The credentials are used to fetch/validate user context * - Tools can access the authenticated user's information * - Different users have different permissions and data access * @@ -21,7 +24,7 @@ import cors from 'cors'; * - organizationId: User's organization * - role: User's role (admin, developer, user) * - permissions: What the user is allowed to do - * - apiKeyId: The API key used + * - apiKeyId: The credential identifier * - requestId: For tracking and logging */ @@ -314,9 +317,12 @@ app.use(cors({ // Map to store transports by session ID const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; -// Function to fetch user context from API key +// Function to fetch user context from authentication credentials +// This example uses API keys, but you could also validate MCP access tokens const fetchUserContext = (apiKey: string): UserContext | undefined => { - // In a real application, this would query a database + // In a real application, this would: + // - For API keys: query your database + // - For MCP access tokens: validate with MCP OAuth server const baseContext = apiKeyDatabase[apiKey]; if (!baseContext) { @@ -334,10 +340,11 @@ const fetchUserContext = (apiKey: string): UserContext | undefined => { // Middleware to extract user context from API key const extractUserContext = (req: Request): UserContext | undefined => { - // Check for API key in various places + // Extract authentication credentials from the request + // This example uses API keys, but you could also extract MCP access tokens const apiKey = req.headers['x-api-key'] as string || - req.headers['authorization']?.replace('Bearer ', '') as string || + req.headers['authorization']?.replace('Bearer ', '') as string || // MCP access token would be here (req.query?.api_key as string); if (!apiKey) {