Per-User MCP Access

import { Agent, Cogitator, CogitatorError, ErrorCode, tool } from '@cogitator-ai/core';
import { MCPClient, MCPServer, type MCPAuthFunction } from '@cogitator-ai/mcp';
import { z } from 'zod';

const apiKey = process.env.GOOGLE_API_KEY;
if (!apiKey) throw new Error('Set GOOGLE_API_KEY');

interface Order {
  id: string;
  item: string;
  status: 'processing' | 'shipped' | 'cancelled';
}

const accessTokens: Record<string, string> = { 'tok-alice-7f3a': 'alice', 'tok-bob-91c2': 'bob' };
const ordersByUser: Record<string, Order[]> = {
  alice: [{ id: 'A-1001', item: 'standing desk', status: 'shipped' }],
  bob: [{ id: 'B-2001', item: 'mechanical keyboard', status: 'processing' }],
};

const bearerAuth: MCPAuthFunction = (request) => {
  const token = request.headers.authorization?.replace(/^Bearer /, '');
  const userId = token ? accessTokens[token] : undefined;
  return userId ? { userId } : undefined;
};

const listMyOrders = tool({
  name: 'list_my_orders',
  description: "List the calling customer's orders with their status.",
  parameters: z.object({}),
  execute: async (_args, context) => ordersByUser[context.userId ?? ''] ?? [],
});

const cancelOrder = tool({
  name: 'cancel_order',
  description: "Cancel one of the calling customer's orders that has not shipped yet.",
  parameters: z.object({ orderId: z.string() }),
  execute: async ({ orderId }, context) => {
    const order = ordersByUser[context.userId ?? '']?.find((o) => o.id === orderId);
    if (!order) throw new Error(`No order ${orderId} for this customer`);
    if (order.status !== 'processing') throw new Error(`Order ${orderId} has already shipped`);
    order.status = 'cancelled';
    return order;
  },
});

const shop = new MCPServer({
  name: 'shop',
  version: '1.0.0',
  transport: 'http',
  host: '127.0.0.1',
  port: 0,
  auth: bearerAuth,
});
shop.registerTools([listMyOrders, cancelOrder]);
await shop.start();
const shopUrl = `http://127.0.0.1:${shop.getPort()}/mcp`;

const cog = new Cogitator({
  llm: { providers: { google: { apiKey } } },
  memory: { adapter: 'memory' },
});
const support = new Agent({
  name: 'support',
  model: 'google/gemini-3.5-flash-lite',
  instructions: "You are a shop's support agent. Use the tools to look up and change orders. Be brief.",
  temperature: 0.2,
});

const clients: MCPClient[] = [];

/** The support agent with the shop's tools, acting with this user's own token. */
async function supportFor(token: string): Promise<Agent> {
  const client = await MCPClient.connect({
    transport: 'http',
    url: shopUrl,
    headers: { Authorization: `Bearer ${token}` },
  });
  clients.push(client);
  return support.clone({ id: support.id, tools: await client.getTools() });
}

try {
  const forAlice = await supportFor('tok-alice-7f3a');
  const forBob = await supportFor('tok-bob-91c2');

  const alice = await cog.run(forAlice, { input: 'What did I order?', userId: 'alice', threadId: 'chat-alice' });
  console.log('Alice:', alice.output);

  const cancel = await cog.run(forBob, {
    input: 'Please cancel the order that has not shipped yet.',
    userId: 'bob',
    threadId: 'chat-bob',
  });
  console.log('Bob:', cancel.output, `(B-2001 is now ${ordersByUser.bob[0].status})`);

  try {
    await cog.run(forBob, { input: 'What did we talk about?', userId: 'bob', threadId: 'chat-alice' });
  } catch (error) {
    if (!(error instanceof CogitatorError && error.code === ErrorCode.THREAD_ACCESS_DENIED)) throw error;
    console.log(`Refused: ${error.message} (HTTP ${error.statusCode})`);
  }
} finally {
  await Promise.all(clients.map((client) => client.close()));
  await shop.stop();
  await cog.close();
}

pnpm add @cogitator-ai/core @cogitator-ai/mcp zod

GOOGLE_API_KEY=your-key npx tsx per-user-mcp.ts

npx tsx examples/mcp/03-per-user-mcp.ts