@plosson/agentio 0.7.1 → 0.7.3

package/src/server/mcp-http.test.ts

@@ -0,0 +1,364 @@
+import { afterEach, describe, expect, test } from 'bun:test';
+
+import {
+  _getSessionCount,
+  _resetSessionsForTests,
+  handleMcpRequest,
+  parseServicesQuery,
+} from './mcp-http';
+
+/**
+ * Pure-ish unit tests for the session manager. `parseServicesQuery` is
+ * a pure function — tested exhaustively. `handleMcpRequest` touches the
+ * module-level `sessions` map, so tests wipe it in afterEach.
+ */
+
+afterEach(() => {
+  _resetSessionsForTests();
+});
+
+/* ------------------------------------------------------------------ */
+/* parseServicesQuery                                                  */
+/* ------------------------------------------------------------------ */
+
+describe('parseServicesQuery — happy path', () => {
+  test('null → empty services (valid)', () => {
+    const r = parseServicesQuery(null);
+    expect(r.ok).toBe(true);
+    if (r.ok) expect(r.services).toEqual([]);
+  });
+
+  test('empty string → empty services', () => {
+    const r = parseServicesQuery('');
+    expect(r.ok).toBe(true);
+    if (r.ok) expect(r.services).toEqual([]);
+  });
+
+  test('whitespace-only → empty services', () => {
+    const r = parseServicesQuery('   ,  , ');
+    expect(r.ok).toBe(true);
+    if (r.ok) expect(r.services).toEqual([]);
+  });
+
+  test('single service without profile', () => {
+    const r = parseServicesQuery('rss');
+    expect(r.ok).toBe(true);
+    if (r.ok)
+      expect(r.services).toEqual([{ service: 'rss', profile: undefined }]);
+  });
+
+  test('single service:profile pair', () => {
+    const r = parseServicesQuery('gchat:default');
+    expect(r.ok).toBe(true);
+    if (r.ok)
+      expect(r.services).toEqual([{ service: 'gchat', profile: 'default' }]);
+  });
+
+  test('multiple services with mixed profile/no-profile', () => {
+    const r = parseServicesQuery('rss,gchat:default,gmail:work');
+    expect(r.ok).toBe(true);
+    if (r.ok)
+      expect(r.services).toEqual([
+        { service: 'rss', profile: undefined },
+        { service: 'gchat', profile: 'default' },
+        { service: 'gmail', profile: 'work' },
+      ]);
+  });
+
+  test('trims whitespace around each entry', () => {
+    const r = parseServicesQuery('  rss , gchat:default  ');
+    expect(r.ok).toBe(true);
+    if (r.ok)
+      expect(r.services).toEqual([
+        { service: 'rss', profile: undefined },
+        { service: 'gchat', profile: 'default' },
+      ]);
+  });
+
+  test('accepts all valid registered services', () => {
+    const all =
+      'discourse,gcal,gchat,gdocs,gdrive,github,gmail,gsheets,gtasks,jira,rss,slack,sql,telegram,whatsapp';
+    const r = parseServicesQuery(all);
+    expect(r.ok).toBe(true);
+    if (r.ok) expect(r.services.length).toBe(15);
+  });
+});
+
+describe('parseServicesQuery — adversarial', () => {
+  test('unknown service name → 400 with known-services hint', () => {
+    const r = parseServicesQuery('nope');
+    expect(r.ok).toBe(false);
+    if (!r.ok) {
+      expect(r.status).toBe(400);
+      expect(r.message).toContain('unknown service');
+      expect(r.message).toContain('"nope"');
+      expect(r.message).toContain('rss'); // sample of known services
+    }
+  });
+
+  test('one unknown among valid → still rejected', () => {
+    const r = parseServicesQuery('rss,nope,gmail:x');
+    expect(r.ok).toBe(false);
+    if (!r.ok) expect(r.message).toContain('nope');
+  });
+
+  test('service name with empty profile (trailing colon) → 400', () => {
+    const r = parseServicesQuery('gmail:');
+    expect(r.ok).toBe(false);
+    if (!r.ok) {
+      expect(r.status).toBe(400);
+      expect(r.message).toContain('profile name is empty');
+    }
+  });
+
+  test('leading colon (empty service name) → 400', () => {
+    const r = parseServicesQuery(':foo');
+    expect(r.ok).toBe(false);
+    if (!r.ok) expect(r.message).toContain('service name is empty');
+  });
+
+  test('bare colon → 400', () => {
+    const r = parseServicesQuery(':');
+    expect(r.ok).toBe(false);
+  });
+
+  test('case-sensitive: RSS (uppercase) → unknown', () => {
+    const r = parseServicesQuery('RSS');
+    expect(r.ok).toBe(false);
+  });
+
+  test('service name with unexpected characters → unknown', () => {
+    const r = parseServicesQuery('rss-feed');
+    expect(r.ok).toBe(false);
+  });
+});
+
+/* ------------------------------------------------------------------ */
+/* handleMcpRequest — routing (no real MCP protocol yet)               */
+/* ------------------------------------------------------------------ */
+
+describe('handleMcpRequest — routing', () => {
+  test('unknown mcp-session-id → 404 JSON-RPC error', async () => {
+    const req = new Request('http://localhost:9999/mcp', {
+      method: 'POST',
+      headers: {
+        'mcp-session-id': 'does-not-exist',
+        'content-type': 'application/json',
+        accept: 'application/json, text/event-stream',
+      },
+      body: JSON.stringify({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'tools/list',
+        params: {},
+      }),
+    });
+    const res = await handleMcpRequest(req);
+    expect(res.status).toBe(404);
+    const body = (await res.json()) as Record<string, unknown>;
+    expect(body.jsonrpc).toBe('2.0');
+    const error = body.error as Record<string, unknown>;
+    expect(error.message).toContain('does-not-exist');
+  });
+
+  test('no session id + invalid services → 400 BEFORE transport runs', async () => {
+    const req = new Request(
+      'http://localhost:9999/mcp?services=nope',
+      {
+        method: 'POST',
+        headers: {
+          'content-type': 'application/json',
+          accept: 'application/json, text/event-stream',
+        },
+        body: JSON.stringify({
+          jsonrpc: '2.0',
+          id: 1,
+          method: 'initialize',
+          params: {
+            protocolVersion: '2024-11-05',
+            capabilities: {},
+            clientInfo: { name: 'test', version: '0.0.0' },
+          },
+        }),
+      }
+    );
+    const res = await handleMcpRequest(req);
+    expect(res.status).toBe(400);
+    const body = (await res.json()) as Record<string, unknown>;
+    expect(body.error).toBe('invalid_request');
+    expect(body.error_description).toContain('unknown service');
+    // And no session was leaked into the map.
+    expect(_getSessionCount()).toBe(0);
+  });
+
+  test('no session id + empty services → transport runs, session created on initialize', async () => {
+    const req = new Request('http://localhost:9999/mcp', {
+      method: 'POST',
+      headers: {
+        'content-type': 'application/json',
+        accept: 'application/json, text/event-stream',
+      },
+      body: JSON.stringify({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'initialize',
+        params: {
+          protocolVersion: '2024-11-05',
+          capabilities: {},
+          clientInfo: { name: 'test', version: '0.0.0' },
+        },
+      }),
+    });
+    const res = await handleMcpRequest(req);
+    // The SDK accepts the initialize and assigns a session id.
+    expect(res.status).toBe(200);
+    expect(res.headers.get('mcp-session-id')).toBeDefined();
+    expect(_getSessionCount()).toBe(1);
+  });
+
+  test('session survives across multiple requests routed by mcp-session-id', async () => {
+    // First: initialize (creates session).
+    const initReq = new Request('http://localhost:9999/mcp?services=rss', {
+      method: 'POST',
+      headers: {
+        'content-type': 'application/json',
+        accept: 'application/json, text/event-stream',
+      },
+      body: JSON.stringify({
+        jsonrpc: '2.0',
+        id: 1,
+        method: 'initialize',
+        params: {
+          protocolVersion: '2024-11-05',
+          capabilities: {},
+          clientInfo: { name: 'test', version: '0.0.0' },
+        },
+      }),
+    });
+    const initRes = await handleMcpRequest(initReq);
+    const sid = initRes.headers.get('mcp-session-id');
+    expect(sid).toBeTruthy();
+    // Drain the body so the stream doesn't leak.
+    await initRes.text();
+
+    // Per MCP spec: client sends notifications/initialized after the
+    // initialize response before using other methods.
+    const notifReq = new Request('http://localhost:9999/mcp?services=rss', {
+      method: 'POST',
+      headers: {
+        'mcp-session-id': sid!,
+        'content-type': 'application/json',
+        accept: 'application/json, text/event-stream',
+        'mcp-protocol-version': '2024-11-05',
+      },
+      body: JSON.stringify({
+        jsonrpc: '2.0',
+        method: 'notifications/initialized',
+      }),
+    });
+    const notifRes = await handleMcpRequest(notifReq);
+    expect(notifRes.status).toBeLessThan(500);
+    await notifRes.text();
+
+    // Second request: tools/list against the same session.
+    const listReq = new Request('http://localhost:9999/mcp?services=rss', {
+      method: 'POST',
+      headers: {
+        'mcp-session-id': sid!,
+        'content-type': 'application/json',
+        accept: 'application/json, text/event-stream',
+        'mcp-protocol-version': '2024-11-05',
+      },
+      body: JSON.stringify({
+        jsonrpc: '2.0',
+        id: 2,
+        method: 'tools/list',
+        params: {},
+      }),
+    });
+    const listRes = await handleMcpRequest(listReq);
+    expect(listRes.status).toBe(200);
+    const listBody = (await listRes.json()) as Record<string, unknown>;
+    expect(listBody.jsonrpc).toBe('2.0');
+    const result = listBody.result as Record<string, unknown>;
+    expect(Array.isArray(result.tools)).toBe(true);
+    // rss has at least one tool (rss_articles).
+    const tools = result.tools as Array<{ name: string }>;
+    expect(tools.some((t) => t.name.startsWith('rss_'))).toBe(true);
+  });
+
+  test('DELETE on a live session removes it from the map', async () => {
+    // Create a session first.
+    const init = await handleMcpRequest(
+      new Request('http://localhost:9999/mcp?services=rss', {
+        method: 'POST',
+        headers: {
+          'content-type': 'application/json',
+          accept: 'application/json, text/event-stream',
+        },
+        body: JSON.stringify({
+          jsonrpc: '2.0',
+          id: 1,
+          method: 'initialize',
+          params: {
+            protocolVersion: '2024-11-05',
+            capabilities: {},
+            clientInfo: { name: 'test', version: '0.0.0' },
+          },
+        }),
+      })
+    );
+    const sid = init.headers.get('mcp-session-id')!;
+    await init.text();
+    expect(_getSessionCount()).toBe(1);
+
+    const del = await handleMcpRequest(
+      new Request('http://localhost:9999/mcp', {
+        method: 'DELETE',
+        headers: {
+          'mcp-session-id': sid,
+          'mcp-protocol-version': '2024-11-05',
+        },
+      })
+    );
+    // DELETE is handled by the transport; either 200 or 204 is fine.
+    expect([200, 204]).toContain(del.status);
+    expect(_getSessionCount()).toBe(0);
+  });
+
+  test('two concurrent initialize requests create two independent sessions', async () => {
+    const mkInitReq = () =>
+      new Request('http://localhost:9999/mcp?services=rss', {
+        method: 'POST',
+        headers: {
+          'content-type': 'application/json',
+          accept: 'application/json, text/event-stream',
+        },
+        body: JSON.stringify({
+          jsonrpc: '2.0',
+          id: 1,
+          method: 'initialize',
+          params: {
+            protocolVersion: '2024-11-05',
+            capabilities: {},
+            clientInfo: { name: 'test', version: '0.0.0' },
+          },
+        }),
+      });
+
+    const [a, b] = await Promise.all([
+      handleMcpRequest(mkInitReq()),
+      handleMcpRequest(mkInitReq()),
+    ]);
+    expect(a.status).toBe(200);
+    expect(b.status).toBe(200);
+    const sidA = a.headers.get('mcp-session-id');
+    const sidB = b.headers.get('mcp-session-id');
+    expect(sidA).toBeTruthy();
+    expect(sidB).toBeTruthy();
+    expect(sidA).not.toBe(sidB);
+    expect(_getSessionCount()).toBe(2);
+    await a.text();
+    await b.text();
+  });
+});

package/src/server/mcp-http.ts

@@ -0,0 +1,339 @@
+import { randomUUID } from 'crypto';
+
+import { Server } from '@modelcontextprotocol/sdk/server/index.js';
+import { WebStandardStreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js';
+import {
+  CallToolRequestSchema,
+  ListToolsRequestSchema,
+} from '@modelcontextprotocol/sdk/types.js';
+
+import {
+  buildProgram,
+  executeCommand,
+  parseServiceProfiles,
+  SERVICE_REGISTRATIONS,
+  type ServiceProfilePair,
+} from '../mcp/server';
+import { collectMcpTools, type McpToolDefinition } from '../mcp/tools';
+
+/**
+ * Per-session bookkeeping for the Streamable HTTP MCP transport.
+ *
+ * The SDK transport is session-oriented: each `WebStandardStreamableHTTPServerTransport`
+ * instance has exactly one `sessionId` set during the `initialize` request.
+ * That means we need one Server + one Transport pair per active MCP
+ * session, not one shared instance for the whole daemon.
+ *
+ * The session is keyed by the SDK-generated `mcp-session-id` header value.
+ * On a new connection (no session id header), we mint a fresh
+ * Server+Transport pair, parse the URL's `?services=` to determine which
+ * tools to expose, and let the SDK assign a session id during initialize
+ * via our `onsessioninitialized` callback.
+ *
+ * The service set is FROZEN for the session's lifetime — once initialize
+ * has run, we ignore any new `?services=` on subsequent requests because
+ * Claude Code (and any sane MCP client) won't send them anyway, and
+ * letting them mutate the tool surface mid-session would break the
+ * client's cached `tools/list`.
+ */
+
+interface McpSession {
+  server: Server;
+  transport: WebStandardStreamableHTTPServerTransport;
+  services: ServiceProfilePair[];
+  toolNames: Set<string>;
+  /**
+   * Per-session "previously checked" tracking for gchat_list. The plan
+   * keys this by `sessionId:service:space`, but since we already have one
+   * Map *per session*, the key inside the Map only needs `service:space`.
+   */
+  lastChecked: Map<string, Date>;
+}
+
+const sessions = new Map<string, McpSession>();
+
+/**
+ * For tests only — drop all in-memory session state. Tests that exercise
+ * the session manager should call this in afterEach so leaked state from
+ * one test never bleeds into another.
+ */
+export function _resetSessionsForTests(): void {
+  for (const session of sessions.values()) {
+    try {
+      session.transport.close();
+    } catch {
+      /* ignore */
+    }
+  }
+  sessions.clear();
+}
+
+/* ------------------------------------------------------------------ */
+/* services parsing + validation                                       */
+/* ------------------------------------------------------------------ */
+
+export interface ParseServicesResult {
+  ok: true;
+  services: ServiceProfilePair[];
+}
+export interface ParseServicesError {
+  ok: false;
+  status: number;
+  message: string;
+}
+
+/**
+ * Parse the `?services=gmail:work,slack:team` query string into
+ * ServiceProfilePair[]. Returns a structured result so the caller can
+ * decide how to surface failures (HTTP 400, JSON-RPC error, etc.).
+ *
+ * Empty input → empty array (valid; the session just exposes no tools).
+ * Unknown service name → error with the offending service.
+ * Empty profile after `:` → error.
+ */
+export function parseServicesQuery(
+  servicesParam: string | null
+): ParseServicesResult | ParseServicesError {
+  if (!servicesParam) {
+    return { ok: true, services: [] };
+  }
+
+  const parts = servicesParam
+    .split(',')
+    .map((p) => p.trim())
+    .filter((p) => p.length > 0);
+
+  if (parts.length === 0) {
+    return { ok: true, services: [] };
+  }
+
+  // Reject anything with an empty service name (e.g. ":foo" or "foo::bar").
+  for (const part of parts) {
+    if (part.startsWith(':') || part === ':') {
+      return {
+        ok: false,
+        status: 400,
+        message: `invalid services entry "${part}": service name is empty`,
+      };
+    }
+    const colonIdx = part.indexOf(':');
+    if (colonIdx !== -1 && colonIdx === part.length - 1) {
+      return {
+        ok: false,
+        status: 400,
+        message: `invalid services entry "${part}": profile name is empty after ":"`,
+      };
+    }
+  }
+
+  const pairs = parseServiceProfiles(parts);
+
+  // Validate every service exists in the registry.
+  for (const pair of pairs) {
+    if (!(pair.service in SERVICE_REGISTRATIONS)) {
+      const known = Object.keys(SERVICE_REGISTRATIONS).sort().join(', ');
+      return {
+        ok: false,
+        status: 400,
+        message: `unknown service "${pair.service}". known services: ${known}`,
+      };
+    }
+  }
+
+  return { ok: true, services: pairs };
+}
+
+/* ------------------------------------------------------------------ */
+/* session creation                                                    */
+/* ------------------------------------------------------------------ */
+
+/**
+ * Build the Server + Transport pair for a new MCP session, register
+ * tool handlers, and connect them. The session is added to the global
+ * `sessions` map by the `onsessioninitialized` callback (fired during
+ * the first `transport.handleRequest()` call from the initialize
+ * request).
+ */
+async function createSession(
+  pairs: ServiceProfilePair[]
+): Promise<McpSession> {
+  const serviceNames = [...new Set(pairs.map((p) => p.service))];
+  const profileMap = new Map<string, string | undefined>();
+  for (const pair of pairs) {
+    profileMap.set(pair.service, pair.profile);
+  }
+
+  // Build the program once to collect the tool list. The same program is
+  // NOT reused for executeCommand — we build a fresh one per call to
+  // avoid Commander state leaks across invocations.
+  const toolProgram = buildProgram(serviceNames);
+  const allTools: McpToolDefinition[] = [];
+  for (const service of serviceNames) {
+    allTools.push(...collectMcpTools(toolProgram, service));
+  }
+  const toolNames = new Set(allTools.map((t) => t.name));
+
+  const server = new Server(
+    { name: 'agentio', version: '1.0.0' },
+    { capabilities: { tools: {} } }
+  );
+
+  // Allocated up front so the closures below + onsessioninitialized can
+  // close over the same object.
+  const session: McpSession = {
+    server,
+    // transport is filled in below; the field is non-null after the
+    // assignment but TypeScript doesn't know that mid-construction.
+    transport: undefined as unknown as WebStandardStreamableHTTPServerTransport,
+    services: pairs,
+    toolNames,
+    lastChecked: new Map(),
+  };
+
+  server.setRequestHandler(ListToolsRequestSchema, async () => {
+    return {
+      tools: allTools.map((tool) => ({
+        name: tool.name,
+        description: tool.description,
+        inputSchema: tool.inputSchema,
+      })),
+    };
+  });
+
+  server.setRequestHandler(CallToolRequestSchema, async (request) => {
+    const { name, arguments: args } = request.params;
+    const tool = allTools.find((t) => t.name === name);
+
+    if (!tool) {
+      return {
+        content: [{ type: 'text' as const, text: `Unknown tool: ${name}` }],
+        isError: true,
+      };
+    }
+
+    const service = tool.commandPath[0];
+    const profile = profileMap.get(service);
+    const input = (args as Record<string, unknown>) || {};
+
+    // Fresh program per call — Commander mutates parser state, and two
+    // overlapping calls in the same session would otherwise stomp on
+    // each other.
+    const execProgram = buildProgram(serviceNames);
+
+    try {
+      const output = await executeCommand(execProgram, tool, input, profile);
+
+      let result = output || '(no output)';
+      if (name === 'gchat_list' && typeof input.space === 'string') {
+        const key = `gchat:${input.space}`;
+        const last = session.lastChecked.get(key);
+        if (last) {
+          result += `\n\nPreviously checked: ${last.toISOString()}`;
+        }
+        session.lastChecked.set(key, new Date());
+      }
+
+      return {
+        content: [{ type: 'text' as const, text: result }],
+      };
+    } catch (error: unknown) {
+      const message =
+        error instanceof Error ? error.message : String(error);
+      return {
+        content: [{ type: 'text' as const, text: `Error: ${message}` }],
+        isError: true,
+      };
+    }
+  });
+
+  const transport = new WebStandardStreamableHTTPServerTransport({
+    sessionIdGenerator: () => randomUUID(),
+    enableJsonResponse: true, // simpler for hand-rolled HTTP testing
+    onsessioninitialized: (sid) => {
+      sessions.set(sid, session);
+    },
+    onsessionclosed: (sid) => {
+      sessions.delete(sid);
+    },
+  });
+
+  session.transport = transport;
+  await server.connect(transport);
+
+  return session;
+}
+
+/* ------------------------------------------------------------------ */
+/* request entry point                                                 */
+/* ------------------------------------------------------------------ */
+
+/**
+ * Top-level handler for /mcp requests, called from `http.ts` after the
+ * bearer middleware has run.
+ *
+ * - Existing session id in `mcp-session-id` header → look up the session
+ *   and forward to its transport. Unknown id → 404.
+ * - No session id → create a new session, validate `?services=`, then
+ *   forward to the freshly-built transport. The SDK will assign a
+ *   session id during initialize and our `onsessioninitialized` callback
+ *   will register it in the map.
+ */
+export async function handleMcpRequest(req: Request): Promise<Response> {
+  const sessionId = req.headers.get('mcp-session-id');
+
+  if (sessionId) {
+    const session = sessions.get(sessionId);
+    if (!session) {
+      return new Response(
+        JSON.stringify({
+          jsonrpc: '2.0',
+          error: {
+            code: -32001,
+            message: `unknown session id "${sessionId}"`,
+          },
+          id: null,
+        }),
+        {
+          status: 404,
+          headers: { 'content-type': 'application/json' },
+        }
+      );
+    }
+    return session.transport.handleRequest(req);
+  }
+
+  // No session id — this should be an initialize request. Validate
+  // services BEFORE building the session so a typo in ?services= gives
+  // a clean 400 instead of an opaque MCP error.
+  const url = new URL(req.url);
+  const parsed = parseServicesQuery(url.searchParams.get('services'));
+  if (!parsed.ok) {
+    return new Response(
+      JSON.stringify({
+        error: 'invalid_request',
+        error_description: parsed.message,
+      }),
+      {
+        status: parsed.status,
+        headers: { 'content-type': 'application/json' },
+      }
+    );
+  }
+
+  const session = await createSession(parsed.services);
+  return session.transport.handleRequest(req);
+}
+
+/**
+ * Test/diagnostic helper: how many sessions are currently active?
+ */
+export function _getSessionCount(): number {
+  return sessions.size;
+}
+
+/**
+ * Test helper: peek at the live sessions map. Do not mutate.
+ */
+export function _peekSessionIds(): string[] {
+  return [...sessions.keys()];
+}