feuerwerk/tests/mcp/relay.test.ts

import { createServer, type Server } from 'node:http';
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
import { type ApiAnswer, CLOSE, type Relay, startRelay, VIEW_URI } from '../../tools/mcp/relay.mjs';

/**
 * The MCP relay (#429) between a fake API, agents speaking MCP over HTTP, and fake editor tabs on
 * WebSockets: no browser, no PHP.
 */

const ORIGIN = 'https://editor.test';
const TOOLS = [
  { name: 'get_view', description: 'view', inputSchema: { type: 'object' } },
  { name: 'screenshot', description: 'picture', inputSchema: { type: 'object' } },
];
const TOKENS: Record<string, string> = { fwa_alice: 'alice', fwa_bob: 'bob', fwa_alice2: 'Alice' };

let relay!: Relay;
let running = false;
let tickets: Map<string, string>;
let addresses: string[];
let throttled: boolean;
/** How long the fake API takes to know a token. */
let verifyMs: number;
const open: WebSocket[] = [];

async function start(extra: Partial<Parameters<typeof startRelay>[0]> = {}): Promise<void> {
  if (running) await relay.close();
  relay = await startRelay({
    api: {
      verifyToken: async (token: string, address: string): Promise<ApiAnswer> => {
        addresses.push(address);
        if (verifyMs > 0) await new Promise((resolve) => setTimeout(resolve, verifyMs));
        if (throttled) return { status: 429, retryAfter: '600', error: 'Too many' };
        const user = TOKENS[token];
        return user ? { status: 200, user } : { status: 401 };
      },
      redeemTicket: (ticket: string): Promise<ApiAnswer> => {
        const user = tickets.get(ticket);
        tickets.delete(ticket);
        return Promise.resolve(user ? { status: 200, user } : { status: 401 });
      },
    },
    tools: () => Promise.resolve(TOOLS),
    origins: [ORIGIN],
    ...extra,
  });
  running = true;
}

beforeEach(async () => {
  tickets = new Map();
  addresses = [];
  throttled = false;
  verifyMs = 0;
  await start();
});

afterEach(async () => {
  for (const socket of open.splice(0)) socket.close();
  await relay.close();
  running = false;
});

const url = () => `http://127.0.0.1:${relay.port}/mcp`;

interface Content {
  type: string;
  text?: string;
  data?: string;
  mimeType?: string;
}
interface ToolResult {
  content: Content[];
  isError?: boolean;
  structuredContent?: object;
}
interface Reply {
  result?: Record<string, unknown>;
  error?: { code: number; message: string };
}
/** What a fake tab hears from the relay. */
interface TabMessage {
  type: string;
  user?: string;
  count?: number;
  id?: number;
  tool?: string;
}

const toolResult = (body: Reply | null) => body?.result as unknown as ToolResult;
const textOf = (body: Reply | null, at = 0) => toolResult(body).content[at]?.text;
const viewText = (body: Reply | null) =>
  (body?.result as unknown as { contents: { text: string }[] }).contents[0]?.text ?? '';

async function rpc(
  token: string | null,
  method: string,
  params: object = {},
  session?: string,
  headers: Record<string, string> = {},
) {
  const res = await fetch(url(), {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      Accept: 'application/json, text/event-stream',
      ...(token ? { Authorization: `Bearer ${token}` } : {}),
      ...(session ? { 'Mcp-Session-Id': session } : {}),
      ...headers,
    },
    body: JSON.stringify({ jsonrpc: '2.0', id: 1, method, params }),
  });
  const text = await res.text();
  return {
    status: res.status,
    headers: res.headers,
    body: (text ? JSON.parse(text) : null) as Reply | null,
  };
}

async function initialize(token: string, protocolVersion = '2025-06-18'): Promise<string> {
  const answer = await rpc(token, 'initialize', {
    protocolVersion,
    capabilities: {},
    clientInfo: { name: 'test' },
  });
  expect(answer.status).toBe(200);
  return answer.headers.get('mcp-session-id') ?? '';
}

interface FakeTab {
  socket: WebSocket;
  received: TabMessage[];
  closed: Promise<number>;
  /** Answers every call with what `answer` makes of it; null leaves it unanswered. */
  answer: (call: TabMessage) => object | null;
  next(type: string): Promise<TabMessage>;
}

function connectTab(
  user: string,
  origin: string | null = ORIGIN,
  ticket = `ticket-${Math.random()}`,
): FakeTab {
  tickets.set(ticket, user);
  const socket = new WebSocket(url().replace('http', 'ws'), {
    headers: origin ? { Origin: origin } : {},
  } as unknown as string[]);
  open.push(socket);
  const tab: FakeTab = {
    socket,
    received: [],
    // Node's client reports a refused handshake as an error alone, without the close a browser adds.
    closed: new Promise((resolve) => {
      socket.addEventListener('close', (event) => resolve(event.code));
      socket.addEventListener('error', () => resolve(1006));
    }),
    answer: (call) => ({ content: [{ type: 'text', text: `${user} did ${call.tool ?? ''}` }] }),
    next: (type) =>
      new Promise((resolve) => {
        const seen = tab.received.find((message) => message.type === type);
        if (seen) {
          tab.received.splice(tab.received.indexOf(seen), 1);
          resolve(seen);
          return;
        }
        const listener = (event: MessageEvent) => {
          const message = JSON.parse(String(event.data)) as TabMessage;
          if (message.type !== type) return;
          tab.received.splice(tab.received.indexOf(message), 1);
          socket.removeEventListener('message', listener);
          resolve(message);
        };
        socket.addEventListener('message', listener);
      }),
  };
  socket.addEventListener('open', () => socket.send(JSON.stringify({ type: 'hello', ticket })));
  socket.addEventListener('message', (event) => {
    const message = JSON.parse(String(event.data)) as TabMessage;
    if (message.type === 'call') {
      const result = tab.answer(message);
      if (result) socket.send(JSON.stringify({ type: 'result', id: message.id, ...result }));
      return;
    }
    tab.received.push(message);
  });
  return tab;
}

async function onTab(user: string): Promise<FakeTab> {
  const tab = connectTab(user);
  expect((await tab.next('welcome')).user).toBe(user);
  return tab;
}

const call = (token: string, session: string, name: string, args: object = {}) =>
  rpc(token, 'tools/call', { name, arguments: args }, session);

describe('agents', () => {
  it('need a token the API knows, and hear 429 while the API brakes them', async () => {
    const none = await rpc(null, 'initialize');
    const wrong = await rpc('fwa_nobody', 'initialize');
    throttled = true;
    const braked = await rpc('fwa_alice', 'initialize');
    expect([none.status, wrong.status, braked.status]).toEqual([401, 401, 429]);
    expect(none.headers.get('www-authenticate')).toMatch(/^Bearer/);
    expect(braked.headers.get('retry-after')).toBe('600');
  });

  it('pass their own address to the brake, or nginx’s word for it only when told to trust it', async () => {
    await rpc('fwa_alice', 'initialize', {}, undefined, { 'X-Real-IP': '203.0.113.9' });
    await start({ trustProxy: true });
    await rpc('fwa_alice', 'initialize', {}, undefined, { 'X-Real-IP': '203.0.113.9' });
    expect(addresses).toHaveLength(2);
    expect(addresses[0]).not.toBe('203.0.113.9');
    expect(addresses[1]).toBe('203.0.113.9');
  });

  it('start a session in the protocol version they ask for, or the newest', async () => {
    const asked = await rpc('fwa_alice', 'initialize', { protocolVersion: '2025-03-26' });
    const unknown = await rpc('fwa_alice', 'initialize', { protocolVersion: '1999-01-01' });
    expect(asked.body?.result).toMatchObject({
      protocolVersion: '2025-03-26',
      capabilities: { tools: {}, resources: { subscribe: true } },
      serverInfo: { name: 'feuerwerk' },
    });
    expect(asked.headers.get('mcp-session-id')).toMatch(/^[0-9a-f-]{36}$/);
    expect(unknown.body?.result?.protocolVersion).toBe('2025-11-25');
  });

  it('get 202 for a notification and the tool list on asking', async () => {
    const session = await initialize('fwa_alice');
    const res = await fetch(url(), {
      method: 'POST',
      headers: {
        Authorization: 'Bearer fwa_alice',
        'Mcp-Session-Id': session,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({ jsonrpc: '2.0', method: 'notifications/initialized' }),
    });
    expect(res.status).toBe(202);
    expect((await rpc('fwa_alice', 'tools/list', {}, session)).body?.result).toEqual({ tools: TOOLS });
  });

  it('cannot use a session of another account, nor one that is not there', async () => {
    const session = await initialize('fwa_alice');
    expect((await rpc('fwa_bob', 'tools/list', {}, session)).status).toBe(404);
    expect((await rpc('fwa_alice', 'tools/list', {}, 'no-such-session')).status).toBe(404);
    expect((await rpc('fwa_alice2', 'tools/list', {}, session)).status).toBe(200);
  });

  it('are refused from a page of another site, not from the editor’s own', async () => {
    const foreign = await rpc('fwa_alice', 'initialize', {}, undefined, { Origin: 'https://evil.test' });
    const own = await rpc('fwa_alice', 'initialize', {}, undefined, { Origin: ORIGIN });
    expect([foreign.status, own.status]).toEqual([403, 200]);
  });

  it('hear about unknown tools and methods in JSON-RPC terms', async () => {
    const session = await initialize('fwa_alice');
    expect((await call('fwa_alice', session, 'rm_rf')).body?.error?.code).toBe(-32602);
    expect((await rpc('fwa_alice', 'prompts/list', {}, session)).body?.error?.code).toBe(-32601);
  });

  it('end a session with DELETE', async () => {
    const session = await initialize('fwa_alice');
    const res = await fetch(url(), {
      method: 'DELETE',
      headers: { Authorization: 'Bearer fwa_alice', 'Mcp-Session-Id': session },
    });
    expect(res.status).toBe(204);
    expect((await rpc('fwa_alice', 'tools/list', {}, session)).status).toBe(404);
  });

  it('send no body larger than the limit', async () => {
    await start({ maxBodyBytes: 1000 });
    const res = await fetch(url(), {
      method: 'POST',
      headers: { Authorization: 'Bearer fwa_alice', 'Content-Type': 'application/json' },
      body: JSON.stringify({
        jsonrpc: '2.0',
        id: 1,
        method: 'initialize',
        params: { pad: 'x'.repeat(2000) },
      }),
    });
    expect(res.status).toBe(413);
  });
});

describe('tabs', () => {
  it('get a call through to the tab of the same account and its answer back', async () => {
    await onTab('alice');
    const session = await initialize('fwa_alice');
    const answer = await call('fwa_alice', session, 'get_view');
    expect(answer.body?.result).toEqual({
      content: [{ type: 'text', text: 'alice did get_view' }],
      isError: false,
    });
  });

  it('without a tab, a call says so and how to fix it', async () => {
    const session = await initialize('fwa_alice');
    const answer = await call('fwa_alice', session, 'get_view');
    expect(toolResult(answer.body).isError).toBe(true);
    expect(textOf(answer.body)).toMatch(/No Feuerwerk editor is connected.*"Agent"/);
  });

  it('never reach the tab of another account', async () => {
    const bobs = await onTab('bob');
    const calls: string[] = [];
    bobs.answer = (message) => {
      calls.push(message.tool ?? '');
      return { content: [{ type: 'text', text: 'bob' }] };
    };
    const session = await initialize('fwa_alice');
    const answer = await call('fwa_alice', session, 'get_view');
    expect(toolResult(answer.body).isError).toBe(true);
    expect(calls).toEqual([]);
  });

  it('are refused from another origin or none, before a ticket is asked for', async () => {
    const foreign = connectTab('alice', 'https://evil.test');
    const none = connectTab('alice', null);
    expect(await foreign.closed).toBe(1006);
    expect(await none.closed).toBe(1006);
    expect(tickets.size).toBe(2);
  });

  it('need a ticket the API takes, once', async () => {
    const tab = connectTab('alice', ORIGIN, 'shared');
    await tab.next('welcome');
    tickets.delete('shared');
    const again = connectTab('alice', ORIGIN, 'other');
    tickets.delete('other');
    expect(await again.closed).toBe(CLOSE.refused);
    expect(relay.stats().tabs).toEqual(['alice']);
  });

  it('that say no hello in time are let go', async () => {
    await start({ helloTimeoutMs: 100 });
    const socket = new WebSocket(url().replace('http', 'ws'), {
      headers: { Origin: ORIGIN },
    } as unknown as string[]);
    open.push(socket);
    const code = await new Promise((resolve) =>
      socket.addEventListener('close', (event) => resolve(event.code)),
    );
    expect(code).toBe(CLOSE.helloTimeout);
  });

  it('one per account: the next takes over, and a call open on the first ends', async () => {
    const first = await onTab('alice');
    first.answer = () => null;
    const session = await initialize('fwa_alice');
    const pending = call('fwa_alice', session, 'get_view');
    await new Promise((resolve) => setTimeout(resolve, 50));
    const second = await onTab('Alice');
    expect((await first.next('replaced')).type).toBe('replaced');
    expect(await first.closed).toBe(CLOSE.replaced);
    expect(textOf((await pending).body)).toBe('Another editor tab took over');
    expect(textOf((await call('fwa_alice', session, 'get_view')).body)).toBe('Alice did get_view');
    second.socket.close();
  });

  it('that do not answer in time end the call with an error', async () => {
    await start({ callTimeoutMs: 100 });
    const tab = await onTab('alice');
    tab.answer = () => null;
    const session = await initialize('fwa_alice');
    const answer = await call('fwa_alice', session, 'screenshot');
    expect(answer.body?.result).toEqual({
      content: [{ type: 'text', text: 'The editor did not answer within 0 s' }],
      isError: true,
    });
  });

  it('that ask the user something get the longer time, once', async () => {
    await start({ callTimeoutMs: 100, waitTimeoutMs: 400 });
    const tab = await onTab('alice');
    tab.answer = (message) => {
      tab.socket.send(JSON.stringify({ type: 'wait', id: message.id }));
      if (message.tool === 'get_view') {
        setTimeout(() => {
          tab.socket.send(
            JSON.stringify({ type: 'result', id: message.id, content: [{ type: 'text', text: 'yes' }] }),
          );
        }, 250);
      } else {
        tab.socket.send(JSON.stringify({ type: 'wait', id: message.id }));
      }
      return null;
    };
    const session = await initialize('fwa_alice');
    expect(textOf((await call('fwa_alice', session, 'get_view')).body)).toBe('yes');
    const started = Date.now();
    const unanswered = await call('fwa_alice', session, 'screenshot');
    const took = Date.now() - started;
    expect(toolResult(unanswered.body).isError).toBe(true);
    expect([took >= 350, took < 1500]).toEqual([true, true]);
  });

  it('pass a picture of several hundred kilobytes through, and only content the protocol knows', async () => {
    const tab = await onTab('alice');
    const data = Buffer.alloc(300_000, 7).toString('base64');
    tab.answer = () => ({
      content: [
        { type: 'image', data, mimeType: 'image/png' },
        { type: 'script', text: 'x' },
      ],
      structuredContent: { width: 10 },
    });
    const session = await initialize('fwa_alice');
    const answer = await call('fwa_alice', session, 'screenshot');
    expect(toolResult(answer.body).content).toEqual([{ type: 'image', data, mimeType: 'image/png' }]);
    expect(toolResult(answer.body).structuredContent).toEqual({ width: 10 });
  });

  it('learn how many agents their account has on, and an agent gone quiet counts out', async () => {
    await start({ agentRecentMs: 150, pingMs: 50 });
    await call('fwa_alice', await initialize('fwa_alice'), 'get_view');
    await new Promise((resolve) => setTimeout(resolve, 200));
    const tab = await onTab('alice');
    expect((await tab.next('agents')).count).toBe(0);
    await call('fwa_alice', await initialize('fwa_alice'), 'get_view');
    expect((await tab.next('agents')).count).toBe(1);
    expect((await tab.next('agents')).count).toBe(0);
  });

  it('do not count the spare sessions a client opens and never calls a tool in', async () => {
    const tab = await onTab('alice');
    expect((await tab.next('agents')).count).toBe(0);
    await initialize('fwa_alice');
    await initialize('fwa_alice');
    await call('fwa_alice', await initialize('fwa_alice'), 'get_view');
    expect((await tab.next('agents')).count).toBe(1);
    expect(tab.received.filter((message) => message.type === 'agents')).toEqual([]);
  });

  it('count out when the event stream of an agent drops, without a DELETE', async () => {
    await start({ agentRecentMs: 150, pingMs: 50 });
    const tab = await onTab('alice');
    expect((await tab.next('agents')).count).toBe(0);
    const session = await initialize('fwa_alice');
    const gone = new AbortController();
    const stream = await fetch(url(), {
      headers: { Authorization: 'Bearer fwa_alice', 'Mcp-Session-Id': session, Accept: 'text/event-stream' },
      signal: gone.signal,
    });
    expect(stream.status).toBe(200);
    expect((await tab.next('agents')).count).toBe(1);
    await new Promise((resolve) => setTimeout(resolve, 400));
    expect(tab.received.filter((message) => message.type === 'agents')).toEqual([]);
    gone.abort();
    expect((await tab.next('agents')).count).toBe(0);
  });

  it('do not count an agent that drops its event stream while its token is still being checked', async () => {
    const tab = await onTab('alice');
    expect((await tab.next('agents')).count).toBe(0);
    const session = await initialize('fwa_alice');
    verifyMs = 100;
    const gone = new AbortController();
    const stream = fetch(url(), {
      headers: { Authorization: 'Bearer fwa_alice', 'Mcp-Session-Id': session, Accept: 'text/event-stream' },
      signal: gone.signal,
    }).catch(() => null);
    await new Promise((resolve) => setTimeout(resolve, 30));
    gone.abort();
    await stream;
    await new Promise((resolve) => setTimeout(resolve, 200));
    expect(tab.received.filter((message) => message.type === 'agents')).toEqual([]);
  });
});

describe('the view', () => {
  it('is read as a resource, and a subscribed agent hears when it changes', async () => {
    const tab = await onTab('alice');
    const session = await initialize('fwa_alice');
    expect((await rpc('fwa_alice', 'resources/list', {}, session)).body?.result).toMatchObject({
      resources: [{ uri: VIEW_URI }],
    });
    expect((await rpc('fwa_alice', 'resources/subscribe', { uri: VIEW_URI }, session)).body?.result).toEqual(
      {},
    );

    const stream = await fetch(url(), {
      headers: { Authorization: 'Bearer fwa_alice', 'Mcp-Session-Id': session, Accept: 'text/event-stream' },
    });
    expect(stream.headers.get('content-type')).toBe('text/event-stream');
    const reader = stream.body!.getReader();
    const heard = (async () => {
      let text = '';
      while (!text.includes('notifications/resources/updated')) {
        const { value, done } = await reader.read();
        if (done) break;
        text += new TextDecoder().decode(value);
      }
      return text;
    })();
    tab.socket.send(JSON.stringify({ type: 'view', view: { tab: 'graph' }, summary: 'tab graph' }));
    expect(await heard).toContain(`"params":{"uri":"${VIEW_URI}"}`);
    await reader.cancel();

    const read = await rpc('fwa_alice', 'resources/read', { uri: VIEW_URI }, session);
    expect(JSON.parse(viewText(read.body))).toEqual({ tab: 'graph' });
  });

  it('read without a tab says it is not connected', async () => {
    const session = await initialize('fwa_alice');
    const read = await rpc('fwa_alice', 'resources/read', { uri: VIEW_URI }, session);
    expect(JSON.parse(viewText(read.body))).toMatchObject({ connected: false });
  });

  it('changed since the agent last looked is told with the next answer, once', async () => {
    const tab = await onTab('alice');
    const session = await initialize('fwa_alice');
    tab.socket.send(JSON.stringify({ type: 'view', view: { tab: 'scene' }, summary: 'tab scene' }));
    await new Promise((resolve) => setTimeout(resolve, 50));
    await call('fwa_alice', session, 'get_view');
    const same = await call('fwa_alice', session, 'screenshot');
    tab.socket.send(JSON.stringify({ type: 'view', view: { tab: 'graph' }, summary: 'tab graph' }));
    await new Promise((resolve) => setTimeout(resolve, 50));
    const changed = await call('fwa_alice', session, 'screenshot');
    const after = await call('fwa_alice', session, 'screenshot');
    expect(toolResult(same.body).content).toHaveLength(1);
    expect(textOf(changed.body, 1)).toBe("The user's view has changed since you last looked: tab graph");
    expect(toolResult(after.body).content).toHaveLength(1);
  });
});

describe('in front of a plain server', () => {
  let backend: Server;
  afterEach(() => new Promise<void>((resolve) => backend.close(() => resolve())));

  it('passes everything off its own path on', async () => {
    backend = createServer((req, res) =>
      res.end(`backend ${req.method} ${req.url} ${req.headers.range ?? ''}`),
    );
    await new Promise<void>((resolve) => backend.listen(0, '127.0.0.1', () => resolve()));
    const port = (backend.address() as { port: number }).port;
    await start({ backend: `http://127.0.0.1:${port}` });
    const page = await fetch(`http://127.0.0.1:${relay.port}/feuerwerk/data/x.mp4?y=1`, {
      headers: { Range: 'bytes=0-9' },
    });
    expect(await page.text()).toBe('backend GET /feuerwerk/data/x.mp4?y=1 bytes=0-9');
    expect((await fetch(`http://127.0.0.1:${relay.port}/feuerwerk/mcp`, { method: 'POST' })).status).toBe(
      401,
    );
  });
});