feuerwerk/tools/mcp/relay.mjs

import { randomUUID } from 'node:crypto';
import { createServer, request as httpRequest } from 'node:http';
import { acceptWebSocket } from './websocket.mjs';

/**
 * The MCP relay (#429): an agent speaks MCP over Streamable HTTP, the user's editor tab is on a
 * WebSocket, and the relay passes the agent's tool calls to the tab of the same account. The tools
 * live in the editor; the relay knows their list (`tools`) and nothing of what they do.
 *
 * Who is who is the PHP API's call (`api`): an agent's token is asked about on every request, so a
 * revoked one stops at the next, and a tab hands in a ticket it got with its session. One tab per
 * account is on at a time; the next one takes over.
 *
 * With `backend` it also stands in front of a server of plain HTTP and passes everything but its own
 * path on, which is how CI puts it beside `php -S` on one origin (nginx does that on stein2).
 */

/** Newest first; an unknown version asked for gets the newest. */
export const PROTOCOL_VERSIONS = ['2025-11-25', '2025-06-18', '2025-03-26', '2024-11-05'];

export const VIEW_URI = 'feuerwerk://view';

const VIEW_RESOURCE = {
  uri: VIEW_URI,
  name: 'view',
  title: 'What the user sees',
  description:
    "The user's editor right now: open tab and document, selection, camera, play state, open dialog. " +
    'Subscribe to hear when the user switches tab or selection.',
  mimeType: 'application/json',
};

const INSTRUCTIONS = [
  "Feuerwerk is a 2D game editor in the browser. You work inside the user's open editor tab,",
  'live: the user sees every change you make and can take it back with Ctrl+Z.',
  'Start with get_view to learn what the user is looking at; screenshot shows it as a picture.',
  'When a tool answers that no editor is connected, ask the user to open the editor and switch on',
  '"Agent" in its account menu.',
].join(' ');

const NOT_CONNECTED =
  'No Feuerwerk editor is connected for this account. Ask the user to open the editor and switch on ' +
  '"Agent" in the account menu (top right), then try again.';

/** How a tab says goodbye on the WebSocket (4000–4999 are the application's). */
export const CLOSE = { replaced: 4000, badHello: 4400, refused: 4401, helloTimeout: 4408 };

/**
 * @param {object} options
 * @param {number} [options.port] 0 for any free one
 * @param {string} [options.host]
 * @param {{verifyToken(token: string, address: string): Promise<ApiAnswer>, redeemTicket(ticket: string): Promise<ApiAnswer>}} options.api
 * @param {() => Promise<object[]>} options.tools the tool list, asked for on every use
 * @param {string[]} options.origins the editor's origins: a tab has to come from one, an agent may
 *   send none, but not another
 * @param {boolean} [options.trustProxy] take the agent's address from `X-Real-IP` (nginx sets it)
 * @param {string} [options.backend] pass everything off the relay's path on to this server
 * @param {(line: string) => void} [options.log]
 */
export async function startRelay(options) {
  const {
    api,
    tools,
    origins,
    trustProxy = false,
    backend,
    log = () => {},
    callTimeoutMs = 30_000,
    helloTimeoutMs = 10_000,
    pingMs = 25_000,
    maxMessageBytes = 32 * 1024 * 1024,
    maxBodyBytes = 4 * 1024 * 1024,
    sessionIdleMs = 12 * 60 * 60 * 1000,
    agentRecentMs = 2 * 60 * 1000,
    maxSessionsPerAccount = 20,
  } = options;

  /** @type {Map<string, Tab>} by account, lower case */
  const tabs = new Map();
  /** @type {Map<string, Session>} by session id */
  const sessions = new Map();

  const accountKey = (user) => user.toLowerCase();
  const sessionsOf = (key) => [...sessions.values()].filter((session) => session.key === key);

  const server = createServer((req, res) => {
    handle(req, res).catch((error) => {
      log(`error: ${error instanceof Error ? error.message : String(error)}`);
      if (!res.headersSent) reply(res, 500, { error: 'Internal relay error' });
      else res.end();
    });
  });
  server.on('upgrade', (req, socket, head) => upgrade(req, socket, head));

  const keepAlive = setInterval(() => {
    for (const tab of tabs.values()) tab.conn.ping();
    const now = Date.now();
    for (const session of sessions.values()) {
      for (const stream of session.streams) stream.write(': ping\n\n');
      if (session.streams.size === 0 && now - session.seenAt > sessionIdleMs) endSession(session);
    }
    // An agent that went quiet is counted out with time, not with a request.
    for (const key of tabs.keys()) tellAgents(key);
  }, pingMs);

  async function handle(req, res) {
    if (!isRelayPath(req.url)) {
      if (backend) return passOn(req, res, backend);
      return reply(res, 404, { error: 'Not found' });
    }
    const origin = req.headers.origin;
    if (origin !== undefined && !origins.includes(origin)) {
      return reply(res, 403, { error: 'Requests from other sites are refused' });
    }
    if (req.method === 'POST') return post(req, res);
    if (req.method === 'GET') return stream(req, res);
    if (req.method === 'DELETE') return remove(req, res);
    return reply(res, 405, { error: 'Method not allowed' }, { Allow: 'GET, POST, DELETE' });
  }

  /** The account behind the agent's token, or null once the refusal has been sent. */
  async function agent(req, res) {
    const match = /^Bearer\s+(\S+)$/i.exec(String(req.headers.authorization ?? ''));
    if (!match) {
      reply(res, 401, { error: 'An agent token is needed: Authorization: Bearer fwa_…' }, BEARER);
      return null;
    }
    let answer;
    try {
      answer = await api.verifyToken(match[1], addressOf(req));
    } catch (error) {
      log(`api unreachable: ${error instanceof Error ? error.message : String(error)}`);
      reply(res, 502, { error: 'The editor server did not answer' });
      return null;
    }
    if (answer.status === 200 && typeof answer.user === 'string') return answer.user;
    if (answer.status === 429) {
      reply(
        res,
        429,
        { error: answer.error ?? 'Too many failed attempts' },
        { 'Retry-After': answer.retryAfter ?? '900' },
      );
      return null;
    }
    reply(res, 401, { error: 'Unknown, expired or revoked agent token' }, BEARER);
    return null;
  }

  function addressOf(req) {
    const forwarded = req.headers['x-real-ip'];
    return trustProxy && typeof forwarded === 'string' && forwarded !== ''
      ? forwarded
      : (req.socket.remoteAddress ?? '');
  }

  /** The session a request names, if it is this account's; 404 otherwise, which tells a client to start over. */
  function sessionOf(req, res, user) {
    const session = sessions.get(String(req.headers['mcp-session-id'] ?? ''));
    if (!session || session.key !== accountKey(user)) {
      reply(res, 404, { jsonrpc: '2.0', id: null, error: { code: -32001, message: 'Session not found' } });
      return null;
    }
    session.seenAt = Date.now();
    return session;
  }

  async function post(req, res) {
    const user = await agent(req, res);
    if (user === null) return;
    const body = await readBody(req, maxBodyBytes);
    if (body === null) return reply(res, 413, { error: 'Request body too large' });
    let parsed;
    try {
      parsed = JSON.parse(body);
    } catch {
      return reply(res, 400, rpcError(null, -32700, 'Parse error'));
    }
    const messages = Array.isArray(parsed) ? parsed : [parsed];
    if (messages.length === 0) return reply(res, 400, rpcError(null, -32600, 'Invalid request'));

    let session;
    if (messages.some((message) => message?.method === 'initialize')) {
      if (messages.length !== 1) return reply(res, 400, rpcError(null, -32600, 'initialize comes alone'));
      session = openSession(user);
    } else {
      session = sessionOf(req, res, user);
      if (!session) return;
    }

    const replies = [];
    for (const message of messages) {
      if (message === null || typeof message !== 'object' || message.jsonrpc !== '2.0') {
        replies.push(rpcError(null, -32600, 'Invalid request'));
      } else if (typeof message.method === 'string' && message.id !== undefined && message.id !== null) {
        replies.push(await dispatch(session, message));
      }
      // Notifications and the client's answers need nothing back.
    }
    const headers = { 'Mcp-Session-Id': session.id };
    if (replies.length === 0) return reply(res, 202, undefined, headers);
    reply(res, 200, Array.isArray(parsed) ? replies : replies[0], headers);
  }

  function openSession(user) {
    const key = accountKey(user);
    const own = sessionsOf(key).sort((a, b) => a.seenAt - b.seenAt);
    while (own.length >= maxSessionsPerAccount) endSession(own.shift());
    /** @type {Session} */
    const session = {
      id: randomUUID(),
      user,
      key,
      protocolVersion: PROTOCOL_VERSIONS[0],
      subscriptions: new Set(),
      streams: new Set(),
      seenAt: Date.now(),
      lastSummary: null,
    };
    sessions.set(session.id, session);
    log(`agent session for ${user}`);
    tellAgents(key);
    return session;
  }

  function endSession(session) {
    for (const stream of session.streams) stream.end();
    session.streams.clear();
    sessions.delete(session.id);
    tellAgents(session.key);
  }

  async function dispatch(session, message) {
    const { id, method } = message;
    const params = message.params !== null && typeof message.params === 'object' ? message.params : {};
    const ok = (result) => ({ jsonrpc: '2.0', id, result });
    switch (method) {
      case 'initialize': {
        const asked = params.protocolVersion;
        session.protocolVersion = PROTOCOL_VERSIONS.includes(asked) ? asked : PROTOCOL_VERSIONS[0];
        return ok({
          protocolVersion: session.protocolVersion,
          capabilities: { tools: {}, resources: { subscribe: true } },
          serverInfo: { name: 'feuerwerk', title: 'Feuerwerk', version: '1.0.0' },
          instructions: INSTRUCTIONS,
        });
      }
      case 'ping':
        return ok({});
      case 'tools/list':
        return ok({ tools: await tools() });
      case 'tools/call': {
        const name = params.name;
        const args = params.arguments ?? {};
        if (typeof name !== 'string' || !(await tools()).some((tool) => tool.name === name)) {
          return rpcError(id, -32602, `Unknown tool: ${String(name)}`);
        }
        if (args === null || typeof args !== 'object' || Array.isArray(args)) {
          return rpcError(id, -32602, 'arguments has to be an object');
        }
        return ok(await callTool(session, name, args));
      }
      case 'resources/list':
        return ok({ resources: [VIEW_RESOURCE] });
      case 'resources/templates/list':
        return ok({ resourceTemplates: [] });
      case 'resources/read': {
        if (params.uri !== VIEW_URI) return rpcError(id, -32002, 'Resource not found', { uri: params.uri });
        const tab = tabs.get(session.key);
        let view = tab?.view ?? null;
        if (tab && view === null) view = (await tab.call('get_view', {})).structuredContent ?? null;
        if (tab) session.lastSummary = tab.summary;
        const text = JSON.stringify(view ?? { connected: false, hint: NOT_CONNECTED });
        return ok({ contents: [{ uri: VIEW_URI, mimeType: 'application/json', text }] });
      }
      case 'resources/subscribe':
      case 'resources/unsubscribe':
        if (params.uri !== VIEW_URI) return rpcError(id, -32002, 'Resource not found', { uri: params.uri });
        if (method === 'resources/subscribe') session.subscriptions.add(VIEW_URI);
        else session.subscriptions.delete(VIEW_URI);
        return ok({});
      default:
        return rpcError(id, -32601, `Method not found: ${method}`);
    }
  }

  /**
   * Hands a call to the account's tab. An agent that has not looked since the user changed what they
   * see gets told so with the answer: not every client follows a subscription.
   */
  async function callTool(session, name, args) {
    const tab = tabs.get(session.key);
    if (!tab) return { content: [{ type: 'text', text: NOT_CONNECTED }], isError: true };
    const seen = session.lastSummary;
    const result = await tab.call(name, args);
    if (name === 'get_view') {
      session.lastSummary = tab.summary;
    } else if (tab.summary !== null && seen !== tab.summary) {
      session.lastSummary = tab.summary;
      result.content = [
        ...result.content,
        { type: 'text', text: `The user's view has changed since you last looked: ${tab.summary}` },
      ];
    }
    return result;
  }

  async function stream(req, res) {
    if (!String(req.headers.accept ?? '').includes('text/event-stream')) {
      return reply(res, 405, { error: 'Ask for text/event-stream' }, { Allow: 'GET, POST, DELETE' });
    }
    const user = await agent(req, res);
    if (user === null) return;
    const session = sessionOf(req, res, user);
    if (!session) return;
    res.writeHead(200, {
      'Content-Type': 'text/event-stream',
      'Cache-Control': 'no-cache, no-transform',
      'X-Accel-Buffering': 'no',
      'Mcp-Session-Id': session.id,
    });
    res.write(': open\n\n');
    session.streams.add(res);
    tellAgents(session.key);
    req.on('close', () => {
      session.streams.delete(res);
      tellAgents(session.key);
    });
  }

  async function remove(req, res) {
    const user = await agent(req, res);
    if (user === null) return;
    const session = sessionOf(req, res, user);
    if (!session) return;
    endSession(session);
    reply(res, 204);
  }

  /** Tells the subscribed agents of an account that its view changed. */
  function viewChanged(key) {
    const note = `event: message\ndata: ${JSON.stringify({
      jsonrpc: '2.0',
      method: 'notifications/resources/updated',
      params: { uri: VIEW_URI },
    })}\n\n`;
    for (const session of sessionsOf(key)) {
      if (!session.subscriptions.has(VIEW_URI)) continue;
      for (const out of session.streams) out.write(note);
    }
  }

  /**
   * How many agents an account has on, for its tab to show: a session counts while its event stream
   * is open or it asked something lately. A client that went away without a DELETE leaves its
   * session behind, and so does a revoked token until it is used again.
   */
  function tellAgents(key) {
    const tab = tabs.get(key);
    if (!tab) return;
    const now = Date.now();
    const count = sessionsOf(key).filter(
      (session) => session.streams.size > 0 || now - session.seenAt < agentRecentMs,
    ).length;
    if (count === tab.agents) return;
    tab.agents = count;
    tab.conn.send(JSON.stringify({ type: 'agents', count }));
  }

  function upgrade(req, socket, head) {
    socket.on('error', () => socket.destroy());
    if (!isRelayPath(req.url)) return socket.destroy();
    const origin = req.headers.origin;
    if (typeof origin !== 'string' || !origins.includes(origin)) {
      socket.end('HTTP/1.1 403 Forbidden\r\nConnection: close\r\n\r\n');
      return;
    }
    const conn = acceptWebSocket(req, socket, head, { maxMessageBytes });
    if (!conn) return;
    /** @type {Tab | null} */
    let tab = null;
    let greeted = false;
    const hello = setTimeout(() => conn.close(CLOSE.helloTimeout, 'hello first'), helloTimeoutMs);
    conn.onmessage = async (text) => {
      let message;
      try {
        message = JSON.parse(text);
      } catch {
        return conn.close(CLOSE.badHello, 'not JSON');
      }
      if (tab) return tab.receive(message);
      if (greeted) return;
      if (message?.type !== 'hello' || typeof message.ticket !== 'string') {
        return conn.close(CLOSE.badHello, 'hello first');
      }
      greeted = true;
      clearTimeout(hello);
      let answer;
      try {
        answer = await api.redeemTicket(message.ticket);
      } catch {
        answer = { status: 502 };
      }
      if (conn.closed) return;
      if (answer.status !== 200 || typeof answer.user !== 'string') {
        return conn.close(CLOSE.refused, 'ticket refused');
      }
      tab = attach(answer.user, conn);
    };
    conn.onclose = () => {
      clearTimeout(hello);
      if (tab) detach(tab, 'The editor tab went away');
    };
  }

  function attach(user, conn) {
    const key = accountKey(user);
    const old = tabs.get(key);
    if (old) {
      old.conn.send(JSON.stringify({ type: 'replaced' }));
      detach(old, 'Another editor tab took over');
      old.conn.close(CLOSE.replaced, 'replaced');
    }
    const tab = new Tab(user, conn, callTimeoutMs);
    tab.onView = () => viewChanged(key);
    tabs.set(key, tab);
    conn.send(JSON.stringify({ type: 'welcome', user }));
    log(`tab on for ${user}`);
    tellAgents(key);
    return tab;
  }

  function detach(tab, why) {
    const key = accountKey(tab.user);
    tab.abandon(why);
    if (tabs.get(key) !== tab) return;
    tabs.delete(key);
    log(`tab off for ${tab.user}`);
    viewChanged(key);
  }

  await new Promise((resolve, reject) => {
    server.once('error', reject);
    server.listen(options.port ?? 0, options.host ?? '127.0.0.1', () => resolve(undefined));
  });

  return {
    port: /** @type {import('node:net').AddressInfo} */ (server.address()).port,
    /** Accounts with a tab on, and sessions per account: for the tests. */
    stats: () => ({
      tabs: [...tabs.values()].map((tab) => tab.user),
      sessions: [...sessions.values()].map((session) => session.user),
    }),
    close: () =>
      new Promise((resolve) => {
        clearInterval(keepAlive);
        for (const tab of tabs.values()) tab.conn.close(1001, 'relay stops');
        for (const session of sessions.values()) for (const out of session.streams) out.end();
        server.close(() => resolve(undefined));
        server.closeAllConnections();
      }),
  };
}

/** One editor tab on the relay: the calls it has open and the last view it reported. */
class Tab {
  view = null;
  /** One line of what the user sees, which the agent is told about when it changes. */
  summary = null;
  agents = -1;
  onView = () => {};
  #pending = new Map();
  #next = 1;

  constructor(user, conn, timeoutMs) {
    this.user = user;
    this.conn = conn;
    this.timeoutMs = timeoutMs;
  }

  /** @returns {Promise<{content: object[], isError?: boolean, structuredContent?: object}>} */
  call(tool, args) {
    const id = this.#next++;
    return new Promise((resolve) => {
      const timer = setTimeout(() => {
        this.#pending.delete(id);
        resolve(failure(`The editor did not answer within ${Math.round(this.timeoutMs / 1000)} s`));
      }, this.timeoutMs);
      this.#pending.set(id, { resolve, timer });
      this.conn.send(JSON.stringify({ type: 'call', id, tool, args }));
    });
  }

  receive(message) {
    if (message?.type === 'result') {
      const waiting = this.#pending.get(message.id);
      if (!waiting) return;
      this.#pending.delete(message.id);
      clearTimeout(waiting.timer);
      waiting.resolve(toolResult(message));
    } else if (message?.type === 'view' && message.view !== null && typeof message.view === 'object') {
      this.view = message.view;
      this.summary = typeof message.summary === 'string' ? message.summary.slice(0, 500) : null;
      this.onView();
    }
  }

  abandon(why) {
    for (const { resolve, timer } of this.#pending.values()) {
      clearTimeout(timer);
      resolve(failure(why));
    }
    this.#pending.clear();
  }
}

/** What a tab answered, as a tool result: only content the protocol knows goes on to the agent. */
function toolResult(message) {
  const content = Array.isArray(message.content)
    ? message.content.filter(
        (item) =>
          (item?.type === 'text' && typeof item.text === 'string') ||
          (item?.type === 'image' && typeof item.data === 'string' && typeof item.mimeType === 'string'),
      )
    : [];
  if (content.length === 0) return failure('The editor answered with nothing');
  const result = { content, isError: message.isError === true };
  if (message.structuredContent !== null && typeof message.structuredContent === 'object') {
    result.structuredContent = message.structuredContent;
  }
  return result;
}

function failure(text) {
  return { content: [{ type: 'text', text }], isError: true };
}

const BEARER = { 'WWW-Authenticate': 'Bearer realm="feuerwerk"' };

function isRelayPath(url) {
  const path = new URL(url ?? '/', 'http://relay').pathname;
  return path === '/mcp' || path.endsWith('/mcp');
}

function rpcError(id, code, message, data) {
  return { jsonrpc: '2.0', id, error: data === undefined ? { code, message } : { code, message, data } };
}

function reply(res, status, body, headers = {}) {
  const text = body === undefined ? '' : JSON.stringify(body);
  res.writeHead(status, {
    ...(body === undefined ? {} : { 'Content-Type': 'application/json' }),
    'Cache-Control': 'no-store',
    ...headers,
  });
  res.end(text);
}

/** The request body as text, or null when it is larger than `max` bytes. */
function readBody(req, max) {
  return new Promise((resolve, reject) => {
    const chunks = [];
    let size = 0;
    req.on('data', (chunk) => {
      size += chunk.length;
      if (size > max) {
        resolve(null);
        req.removeAllListeners('data');
        req.resume();
        return;
      }
      chunks.push(chunk);
    });
    req.on('end', () => resolve(Buffer.concat(chunks).toString('utf8')));
    req.on('error', reject);
  });
}

/** Front mode: one request through to the backend, streamed both ways (Range requests included). */
function passOn(req, res, backend) {
  const target = new URL(req.url ?? '/', backend);
  const out = httpRequest(
    target,
    { method: req.method, headers: { ...req.headers, host: target.host } },
    (answer) => {
      res.writeHead(answer.statusCode ?? 502, answer.headers);
      answer.pipe(res);
    },
  );
  out.on('error', () => {
    if (!res.headersSent) res.writeHead(502);
    res.end();
  });
  req.pipe(out);
}

/**
 * @typedef {{status: number, user?: string, role?: string, error?: string, retryAfter?: string | null}} ApiAnswer
 * @typedef {{id: string, user: string, key: string, protocolVersion: string, subscriptions: Set<string>,
 *   streams: Set<import('node:http').ServerResponse>, seenAt: number, lastSummary: string | null}} Session
 */