@axiom-lattice/gateway 2.1.88 → 2.1.90

package/src/routes/a2a-bridge.ts

@@ -0,0 +1,266 @@
+/**
+ * A2A Bridge — WebSocket-based reverse proxy for agents behind NAT.
+ *
+ * Agents open an outbound WebSocket to the gateway. The gateway stores
+ * the connection and exposes an HTTP proxy endpoint per agent. When the
+ * orchestrator calls the agent, the gateway proxies the HTTP request
+ * over the persistent WebSocket instead of trying to reach a private IP.
+ *
+ * Agent flow:
+ *   1. Connect ws://gateway/api/a2a/bridge
+ *   2. Send { type: "agent:register", agentId: "...", agentCard: {...} }
+ *   3. Wait for { type: "http:request", requestId, method, path, headers, body }
+ *   4. Process locally, respond with { type: "http:response", requestId, status, headers, body }
+ *
+ * Gateway flow:
+ *   1. Accept WebSocket, store connection by agentId
+ *   2. Expose GET/POST /api/a2a/bridge/proxy/{agentId}/*
+ *   3. When HTTP request arrives, push over WebSocket, await response, relay
+ */
+
+import type { FastifyInstance, FastifyRequest, FastifyReply } from 'fastify';
+import type { WebSocket } from 'ws';
+import { v4 } from 'uuid';
+
+// ─── Logging ────────────────────────────────────────────────────────────────
+
+const log = {
+  info(msg: string, ctx?: Record<string, unknown>) {
+    console.log(JSON.stringify({ level: 'info', msg, ...ctx }));
+  },
+  warn(msg: string, ctx?: Record<string, unknown>) {
+    console.warn(JSON.stringify({ level: 'warn', msg, ...ctx }));
+  },
+  error(msg: string, ctx?: Record<string, unknown>) {
+    console.error(JSON.stringify({ level: 'error', msg, ...ctx }));
+  },
+};
+
+// ─── Types ──────────────────────────────────────────────────────────────────
+
+interface PendingRequest {
+  resolve: (response: BridgeHttpResponse) => void;
+  reject: (error: Error) => void;
+  timer: ReturnType<typeof setTimeout>;
+  /** The agent owning this pending request — used for cleanup on disconnect */
+  agentId: string;
+}
+
+interface BridgeHttpResponse {
+  status: number;
+  headers: Record<string, string>;
+  body: string;
+}
+
+interface AgentConnection {
+  ws: WebSocket;
+  agentId: string;
+  agentCard: Record<string, unknown>;
+  connectedAt: number;
+}
+
+// ─── Connection Registry ────────────────────────────────────────────────────
+
+const connections = new Map<string, AgentConnection>();
+const pending = new Map<string, PendingRequest>();
+
+function getAgentIdFromPath(path: string): string | null {
+  // Path: /api/a2a/bridge/proxy/{agentId}/...
+  const match = path.match(/^\/api\/a2a\/bridge\/proxy\/([^/]+)/);
+  return match ? match[1] : null;
+}
+
+// ─── WebSocket Handler ──────────────────────────────────────────────────────
+
+function handleBridgeConnection(socket: WebSocket, _request: FastifyRequest): void {
+  let agentId: string | null = null;
+
+  log.info('Bridge WebSocket connected');
+
+  socket.on('message', (raw) => {
+    let msg: Record<string, unknown>;
+    try {
+      msg = JSON.parse(raw.toString());
+    } catch {
+      log.warn('Bridge invalid message');
+      return;
+    }
+
+    switch (msg.type) {
+      case 'agent:register': {
+        agentId = msg.agentId as string;
+        if (!agentId) {
+          socket.send(JSON.stringify({ type: 'error', message: 'agentId required' }));
+          socket.close();
+          return;
+        }
+        connections.set(agentId, {
+          ws: socket,
+          agentId,
+          agentCard: msg.agentCard as Record<string, unknown> ?? {},
+          connectedAt: Date.now(),
+        });
+        log.info('Agent registered via bridge', { agentId });
+        socket.send(JSON.stringify({ type: 'agent:registered', agentId }));
+        break;
+      }
+
+      case 'http:response': {
+        const requestId = msg.requestId as string;
+        const p = pending.get(requestId);
+        if (p) {
+          clearTimeout(p.timer);
+          pending.delete(requestId);
+          p.resolve({
+            status: (msg.status as number) ?? 200,
+            headers: (msg.headers as Record<string, string>) ?? {},
+            body: (msg.body as string) ?? '',
+          });
+        }
+        break;
+      }
+
+      case 'http:error': {
+        const requestId = msg.requestId as string;
+        const p = pending.get(requestId);
+        if (p) {
+          clearTimeout(p.timer);
+          pending.delete(requestId);
+          p.reject(new Error((msg.message as string) ?? 'Agent error'));
+        }
+        break;
+      }
+    }
+  });
+
+  socket.on('close', () => {
+    if (agentId) {
+      connections.delete(agentId);
+      log.info('Bridge WebSocket disconnected', { agentId });
+      // Reject only this agent's pending requests
+      for (const [reqId, p] of pending) {
+        if (p.agentId === agentId) {
+          clearTimeout(p.timer);
+          p.reject(new Error('Agent disconnected'));
+          pending.delete(reqId);
+        }
+      }
+    }
+  });
+
+  socket.on('error', (err) => {
+    log.error('Bridge WebSocket error', { agentId, error: err.message });
+  });
+}
+
+// ─── HTTP Proxy Handler ─────────────────────────────────────────────────────
+
+async function handleBridgeProxy(
+  request: FastifyRequest,
+  reply: FastifyReply,
+): Promise<void> {
+  const agentId = getAgentIdFromPath(request.url);
+  if (!agentId) {
+    reply.status(400).send({ error: 'Agent ID not found in path' });
+    return;
+  }
+
+  const conn = connections.get(agentId);
+  if (!conn) {
+    reply.status(503).send({
+      error: 'Agent not connected',
+      message: `Agent "${agentId}" is not connected via WebSocket bridge`,
+    });
+    return;
+  }
+
+  const requestId = v4();
+
+  // Build sub-path (strip the proxy prefix)
+  const subPath = request.url.replace(`/api/a2a/bridge/proxy/${agentId}`, '') || '/';
+
+  // Read body
+  let body: string | undefined;
+  if (request.method === 'POST' || request.method === 'PUT') {
+    body = typeof request.body === 'string'
+      ? request.body
+      : JSON.stringify(request.body ?? {});
+  }
+
+  const bridgeRequest = {
+    type: 'http:request',
+    requestId,
+    method: request.method,
+    path: subPath,
+    headers: {
+      'content-type': request.headers['content-type'] ?? 'application/json',
+      'accept': request.headers['accept'] ?? '*/*',
+    },
+    body: body ?? '',
+  };
+
+  const timeout = 300_000; // 5 min
+
+  try {
+    const response = await new Promise<BridgeHttpResponse>((resolve, reject) => {
+      const timer = setTimeout(() => {
+        pending.delete(requestId);
+        reject(new Error(`Bridge request timeout after ${timeout}ms`));
+      }, timeout);
+
+      pending.set(requestId, { resolve, reject, timer, agentId });
+
+      try {
+        if (conn.ws.readyState !== conn.ws.OPEN) {
+          pending.delete(requestId);
+          reject(new Error('WebSocket not open'));
+          return;
+        }
+        conn.ws.send(JSON.stringify(bridgeRequest));
+      } catch (err) {
+        clearTimeout(timer);
+        pending.delete(requestId);
+        reject(err as Error);
+      }
+    });
+
+    // Relay response headers
+    for (const [key, value] of Object.entries(response.headers)) {
+      if (!['content-length', 'transfer-encoding', 'connection'].includes(key.toLowerCase())) {
+        reply.header(key, value);
+      }
+    }
+
+    reply.status(response.status);
+
+    // Try to send as JSON if possible
+    try {
+      const parsed = JSON.parse(response.body);
+      reply.send(parsed);
+    } catch {
+      reply.send(response.body);
+    }
+  } catch (err) {
+    log.error('Bridge proxy error', { agentId, error: (err as Error).message });
+    reply.status(502).send({
+      error: 'Bridge proxy error',
+      message: (err as Error).message,
+    });
+  }
+}
+
+// ─── Route Registration ─────────────────────────────────────────────────────
+
+export function registerA2ABridgeRoutes(app: FastifyInstance): void {
+  // WebSocket endpoint for agents to connect
+  app.get('/api/a2a/bridge', { websocket: true }, (socket, req) => {
+    handleBridgeConnection(socket, req);
+  });
+
+  // HTTP proxy endpoint — all paths under /api/a2a/bridge/proxy/{agentId}/
+  app.route({
+    method: ['GET', 'POST', 'PUT', 'DELETE', 'PATCH'],
+    url: '/api/a2a/bridge/proxy/*',
+    handler: handleBridgeProxy,
+  });
+}