a2a-mcp/src/gateway.ts

/**
 * MCP gateway server. Fans out tools/list and tools/call to each enabled
 * upstream agent (deployed on a2acloud). Tool names are prefixed with the
 * agent slug so multiple agents can coexist in one MCP namespace.
 *
 * Naming: ``{agent}__{skill}``. Double underscore separator. Skill names
 * containing ``__`` are not supported.
 */
import { createRequire } from "node:module";
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
  CallToolRequestSchema,
  ElicitResultSchema,
  ListToolsRequestSchema,
  ResultSchema,
} from "@modelcontextprotocol/sdk/types.js";

import { loadConfig, type EnabledAgent } from "./config.js";
import { loadCredentials } from "./credentials.js";
import { UpstreamAgent, UpstreamError, type UpstreamTool } from "./upstream.js";

const require = createRequire(import.meta.url);
const pkg = require("../package.json") as { version: string };

export const SEP = "__";

export interface GatewayOptions {
  /** Override the enabled agents (otherwise read from ~/.a2a/mcp.json). */
  agents?: EnabledAgent[];
  /** Override the bearer token (otherwise read from ~/.a2a/credentials.json). */
  token?: string | null;
  /** Test/advanced override for one-shot upstream JSON-RPC calls. */
  upstreamRequestTimeoutMs?: number;
  /** Test/advanced override for idle SSE upstream tools/call streams. */
  upstreamStreamIdleTimeoutMs?: number;
  /** Optional override for the underlying MCP Server (tests). */
  server?: Server;
}

export interface GatewayHandle {
  server: Server;
  upstreams: Map<string, UpstreamAgent>;
}

export async function buildGateway(opts: GatewayOptions = {}): Promise<GatewayHandle> {
  const cfg = opts.agents ?? (await loadConfig()).agents;
  const creds = opts.token !== undefined ? null : await loadCredentials();
  const token = opts.token !== undefined ? opts.token : creds?.token ?? null;
  // Forward CP credentials + bucket as params._meta so upstream agents
  // can act on behalf of the caller (mirrors the platform-orchestrator
  // call shape). When the user isn't logged in, _meta is omitted and
  // upstream sees the same unauthenticated context as before.
  const cpJwt = creds?.token ?? null;
  const cpUrl = creds?.apiUrl ?? null;
  const bucket = creds?.bucket ?? null;

  const upstreams = new Map<string, UpstreamAgent>();
  for (const a of cfg) {
    upstreams.set(
      a.name,
      new UpstreamAgent({
        name: a.name, url: a.url, token,
        cpJwt, cpUrl, bucket,
        requestTimeoutMs: opts.upstreamRequestTimeoutMs,
        streamIdleTimeoutMs: opts.upstreamStreamIdleTimeoutMs,
      }),
    );
  }

  const server =
    opts.server ??
    new Server(
      { name: "a2amcp", version: pkg.version },
      { capabilities: { tools: { listChanged: false } } },
    );

  server.setRequestHandler(ListToolsRequestSchema, async () => {
    const tools: Array<UpstreamTool & { name: string }> = [];
    await Promise.all(
      [...upstreams.values()].map(async (u) => {
        try {
          const upstreamTools = await u.listTools();
          for (const t of upstreamTools) {
            tools.push({
              ...t,
              name: `${u.name}${SEP}${t.name}`,
              description: t.description
                ? `[${u.name}] ${t.description}`
                : `[${u.name}] ${t.name}`,
            });
          }
        } catch (err) {
          // Silently skip an upstream that's down — surfacing an error here
          // would break the client's whole tool list. The user can debug via
          // `a2amcp doctor`.
          process.stderr.write(
            `a2amcp: skipping ${u.name}: ${(err as Error).message}\n`,
          );
        }
      }),
    );
    return { tools };
  });

  server.setRequestHandler(CallToolRequestSchema, async (req, extra) => {
    const fullName = req.params.name;
    const idx = fullName.indexOf(SEP);
    if (idx < 0) {
      return _softError(`tool name missing agent prefix: ${fullName}`) as any;
    }
    const agentName = fullName.slice(0, idx);
    const toolName = fullName.slice(idx + SEP.length);
    const upstream = upstreams.get(agentName);
    if (!upstream) {
      return _softError(`unknown agent: ${agentName}`) as any;
    }
    // Honor the client's progressToken if it supplied one; else fall
    // back to the request id. Either way, when upstream emits
    // notifications/progress we forward them via the SDK's
    // sendNotification so clients like Claude Code reset their idle
    // tool-call timer and don't abort long builds.
    const clientToken = (req.params._meta as any)?.progressToken as
      | string
      | number
      | undefined;
    const progressToken = clientToken ?? (extra.requestId as string | number);
    try {
      const result = await upstream.callTool(
        toolName,
        (req.params.arguments as Record<string, unknown>) ?? {},
        {
          progressToken,
          onNotification: async (notif) => {
            try {
              await extra.sendNotification(notif as any);
            } catch {
              // best effort
            }
          },
          onRequest: async (upstreamReq) => {
            const schema = upstreamReq.method === "elicitation/create"
              ? ElicitResultSchema
              : ResultSchema;
            return extra.sendRequest(
              {
                method: upstreamReq.method,
                params: upstreamReq.params,
              } as any,
              schema as any,
            );
          },
        },
      );
      // Pass through the upstream's CallToolResult as-is. Upstreams already
      // populate content / structuredContent / isError per the MCP spec.
      // Cast through `unknown` because the SDK's ServerResult union has
      // additional task-shaped variants we don't produce.
      return result as unknown as any;
    } catch (err) {
      const msg =
        err instanceof UpstreamError
          ? err.message
          : `${(err as Error).name}: ${(err as Error).message}`;
      return _softError(msg) as any;
    }
  });

  return { server, upstreams };
}

export async function runStdio(opts: GatewayOptions = {}): Promise<void> {
  const { server } = await buildGateway(opts);
  const transport = new StdioServerTransport();
  await server.connect(transport);
}

function _softError(message: string) {
  return {
    content: [{ type: "text" as const, text: message }],
    isError: true,
  };
}