Skip to content

Commit 81982cd

Browse files
grimmerkgrimmerk
committed
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: - Tool handler access to context via extra.customContext - Prompt handler access to context via extra.customContext - API key to user context mapping - Permission-based access control - Multi-tenant data isolation Note: Resources also support customContext through RequestHandlerExtra, but not included in this simplified example.
1 parent 80bf790 commit 81982cd

File tree

4 files changed

+923
-0
lines changed

4 files changed

+923
-0
lines changed

‎src/examples/README.md

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ This directory contains example implementations of MCP clients and servers using
66

77
- [Client Implementations](#client-implementations)
88
- [Streamable HTTP Client](#streamable-http-client)
9+
- [Custom Context Client](#custom-context-client)
910
- [Backwards Compatible Client](#backwards-compatible-client)
1011
- [Server Implementations](#server-implementations)
1112
- [Single Node Deployment](#single-node-deployment)
@@ -39,6 +40,26 @@ Example client with OAuth:
3940
npx tsx src/examples/client/simpleOAuthClient.js
4041
```
4142

43+
### Custom Context Client
44+
45+
An interactive client that demonstrates the custom context feature, showing how to:
46+
47+
- Authenticate using API keys that map to user contexts
48+
- Pass user context (identity, permissions, organization) through the transport layer
49+
- Access context in MCP tool handlers for authorization and personalization
50+
- Implement multi-tenant data isolation
51+
- Track requests with unique IDs for auditing
52+
53+
```bash
54+
# Start the server first:
55+
npx tsx src/examples/server/customContextServer.ts
56+
57+
# Then run the client:
58+
npx tsx src/examples/client/customContextClient.ts
59+
```
60+
61+
See [custom-context-example.md](custom-context-example.md) for a detailed walkthrough.
62+
4263
### Backwards Compatible Client
4364

4465
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.
106127
npx tsx src/examples/server/standaloneSseWithGetStreamableHttp.ts
107128
```
108129

130+
##### Custom Context Server
131+
132+
A server that demonstrates how to inject custom context (user authentication, permissions, tenant data) into MCP tool handlers.
133+
134+
- API key authentication with user context extraction
135+
- Context injection via `transport.setCustomContext()`
136+
- Permission-based access control in tools
137+
- Multi-tenant data isolation
138+
- Request tracking with unique IDs
139+
140+
```bash
141+
npx tsx src/examples/server/customContextServer.ts
142+
```
143+
144+
This example is essential for building secure, multi-tenant MCP applications. See [custom-context-example.md](custom-context-example.md) for implementation details.
145+
109146
#### Deprecated SSE Transport
110147

111148
A server that implements the deprecated HTTP+SSE transport (protocol version 2024-11-05). This example only used for testing backwards compatibility for clients.

‎src/examples/client/customContextClient.ts

Lines changed: 306 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,306 @@
1+
import { Client } from '../../client/index.js';
2+
import { StreamableHTTPClientTransport } from '../../client/streamableHttp.js';
3+
import { SSEClientTransport } from '../../client/sse.js';
4+
import { createInterface } from 'node:readline';
5+
import {
6+
CallToolResultSchema,
7+
ListToolsResultSchema,
8+
GetPromptResultSchema,
9+
} from '../../types.js';
10+
11+
/**
12+
* Interactive client demonstrating custom context feature.
13+
*
14+
* This client shows how API keys are used to authenticate and
15+
* how the server uses the context to provide user-specific responses.
16+
*/
17+
18+
// Create readline interface for user input
19+
const readline = createInterface({
20+
input: process.stdin,
21+
output: process.stdout
22+
});
23+
24+
// Global state
25+
let client: Client | null = null;
26+
let transport: StreamableHTTPClientTransport | SSEClientTransport | null = null;
27+
const serverUrl = 'http://localhost:3000/mcp';
28+
let currentUser: {
29+
name: string;
30+
organization: { name: string };
31+
role: string;
32+
permissions: string[];
33+
} | null = null;
34+
35+
// Available API keys for testing
36+
const API_KEYS: Record<string, string> = {
37+
'alice': 'sk-alice-admin-key',
38+
'bob': 'sk-bob-dev-key',
39+
'charlie': 'sk-charlie-user-key',
40+
'dana': 'sk-dana-admin-key',
41+
};
42+
43+
async function main(): Promise<void> {
44+
console.log('==============================================');
45+
console.log('MCP Custom Context Demo Client');
46+
console.log('==============================================');
47+
console.log('\nThis client demonstrates how custom context works:');
48+
console.log('1. Authenticate with an API key');
49+
console.log('2. The server fetches user context from the API key');
50+
console.log('3. Tools receive the context and respond based on user permissions\n');
51+
52+
printHelp();
53+
commandLoop();
54+
}
55+
56+
function printHelp(): void {
57+
console.log('\n📋 Available commands:');
58+
console.log(' auth <user> - Authenticate as user (alice/bob/charlie/dana)');
59+
console.log(' auth-key <key> - Authenticate with custom API key');
60+
console.log(' whoami - Get current user info from context');
61+
console.log(' dashboard [format] - Get personalized dashboard (brief/detailed)');
62+
console.log(' list-tools - List available tools');
63+
console.log(' disconnect - Disconnect from server');
64+
console.log(' help - Show this help');
65+
console.log(' quit - Exit the program');
66+
console.log('\n🔑 Quick start: Try "auth alice" then "whoami"');
67+
console.log('\n⚠️ Note: Only the get_user tool is available in this simplified demo.');
68+
}
69+
70+
function commandLoop(): void {
71+
const prompt = currentUser ? `[${currentUser!.name}]> ` : '> ';
72+
73+
readline.question(prompt, async (input) => {
74+
const args = input.trim().split(/\s+/);
75+
const command = args[0]?.toLowerCase();
76+
77+
try {
78+
switch (command) {
79+
case 'auth': {
80+
const userName = args[1] as keyof typeof API_KEYS;
81+
if (args.length < 2 || !API_KEYS[userName]) {
82+
console.log('❌ Usage: auth <alice|bob|charlie|dana>');
83+
console.log(' Available users:');
84+
console.log(' - alice: TechCorp Admin (all permissions)');
85+
console.log(' - bob: TechCorp Developer (code/docs permissions)');
86+
console.log(' - charlie: StartupIO User (limited permissions)');
87+
console.log(' - dana: StartupIO Admin (org admin)');
88+
} else {
89+
await authenticateAs(userName);
90+
}
91+
break;
92+
}
93+
94+
case 'auth-key':
95+
if (args.length < 2) {
96+
console.log('❌ Usage: auth-key <api-key>');
97+
} else {
98+
await authenticateWithKey(args[1]);
99+
}
100+
break;
101+
102+
case 'whoami':
103+
await getCurrentUser();
104+
break;
105+
106+
107+
case 'dashboard':
108+
await getDashboard(args[1] || 'brief');
109+
break;
110+
111+
case 'list-tools':
112+
await listTools();
113+
break;
114+
115+
case 'disconnect':
116+
await disconnect();
117+
break;
118+
119+
case 'help':
120+
printHelp();
121+
break;
122+
123+
case 'quit':
124+
case 'exit':
125+
await cleanup();
126+
return;
127+
128+
default:
129+
if (command) {
130+
console.log(`❓ Unknown command: ${command}`);
131+
}
132+
break;
133+
}
134+
} catch (error) {
135+
console.error(`❌ Error: ${error}`);
136+
}
137+
138+
// Continue the command loop
139+
commandLoop();
140+
});
141+
}
142+
143+
async function authenticateAs(userName: string): Promise<void> {
144+
const apiKey = API_KEYS[userName as keyof typeof API_KEYS];
145+
await authenticateWithKey(apiKey);
146+
}
147+
148+
async function authenticateWithKey(apiKey: string): Promise<void> {
149+
// Disconnect existing connection
150+
if (client) {
151+
await disconnect();
152+
}
153+
154+
// Store the API key for this session (used in fetch)
155+
console.log(`\n🔐 Authenticating with API key: ${apiKey.substring(0, 15)}...`);
156+
157+
// Create transport with API key in headers
158+
transport = new StreamableHTTPClientTransport(
159+
new URL(serverUrl),
160+
{
161+
fetch: async (url: string | URL, options?: RequestInit) => {
162+
// Add API key to all requests
163+
// Handle Headers object or plain object
164+
let headers: HeadersInit;
165+
if (options?.headers instanceof Headers) {
166+
headers = new Headers(options.headers);
167+
(headers as Headers).set('X-API-Key', apiKey);
168+
} else {
169+
headers = {
170+
...(options?.headers || {}),
171+
'X-API-Key': apiKey,
172+
};
173+
}
174+
return fetch(url, { ...options, headers });
175+
}
176+
}
177+
);
178+
179+
// Create and connect client
180+
client = new Client({
181+
name: 'custom-context-demo-client',
182+
version: '1.0.0'
183+
});
184+
185+
try {
186+
await client.connect(transport);
187+
console.log('✅ Connected to server');
188+
189+
// Get user info immediately after connecting
190+
const result = await client.request({
191+
method: 'tools/call',
192+
params: {
193+
name: 'get_user',
194+
arguments: {}
195+
}
196+
}, CallToolResultSchema);
197+
198+
if (result.content && result.content[0]?.type === 'text') {
199+
const text = result.content[0].text;
200+
try {
201+
// Parse user info from response
202+
const userMatch = text.match(/User Profile:\n([\s\S]*)/);
203+
if (userMatch) {
204+
currentUser = JSON.parse(userMatch[1]);
205+
console.log(`\n👤 Authenticated as: ${currentUser!.name}`);
206+
console.log(` Organization: ${currentUser!.organization.name}`);
207+
console.log(` Role: ${currentUser!.role}`);
208+
console.log(` Permissions: ${currentUser!.permissions.length} permission(s)`);
209+
}
210+
} catch {
211+
console.log('✅ Authenticated (could not parse user details)');
212+
}
213+
}
214+
} catch (error) {
215+
console.error(`❌ Failed to connect: ${error}`);
216+
client = null;
217+
transport = null;
218+
}
219+
}
220+
221+
async function getCurrentUser(): Promise<void> {
222+
if (!client) {
223+
console.log('❌ Not connected. Use "auth <user>" first.');
224+
return;
225+
}
226+
227+
console.log('\n🔍 Fetching user information from context...');
228+
229+
const result = await client.request({
230+
method: 'tools/call',
231+
params: {
232+
name: 'get_user',
233+
arguments: {}
234+
}
235+
}, CallToolResultSchema);
236+
237+
if (result.content && result.content[0]?.type === 'text') {
238+
console.log('\n' + result.content[0].text);
239+
}
240+
}
241+
242+
243+
async function getDashboard(format: string): Promise<void> {
244+
if (!client) {
245+
console.log('❌ Not connected. Use "auth <user>" first.');
246+
return;
247+
}
248+
249+
console.log(`\n📊 Getting ${format} dashboard...`);
250+
251+
const result = await client.request({
252+
method: 'prompts/get',
253+
params: {
254+
name: 'user-dashboard',
255+
arguments: { format }
256+
}
257+
}, GetPromptResultSchema);
258+
259+
if (result.messages && result.messages[0]?.content?.type === 'text') {
260+
console.log('\n' + result.messages[0].content.text);
261+
}
262+
}
263+
264+
async function listTools(): Promise<void> {
265+
if (!client) {
266+
console.log('❌ Not connected. Use "auth <user>" first.');
267+
return;
268+
}
269+
270+
const result = await client.request({
271+
method: 'tools/list',
272+
params: {}
273+
}, ListToolsResultSchema);
274+
275+
console.log('\n🔧 Available tools:');
276+
for (const tool of result.tools) {
277+
console.log(` - ${tool.name}: ${tool.description}`);
278+
}
279+
}
280+
281+
async function disconnect(): Promise<void> {
282+
if (client) {
283+
await client.close();
284+
client = null;
285+
transport = null;
286+
currentUser = null;
287+
console.log('✅ Disconnected from server');
288+
} else {
289+
console.log('❌ Not connected');
290+
}
291+
}
292+
293+
async function cleanup(): Promise<void> {
294+
await disconnect();
295+
console.log('\n👋 Goodbye!');
296+
readline.close();
297+
process.exit(0);
298+
}
299+
300+
// Handle ctrl+c
301+
process.on('SIGINT', async () => {
302+
await cleanup();
303+
});
304+
305+
// Start the client
306+
main().catch(console.error);

0 commit comments

Comments
 (0)