@axiom-lattice/gateway 2.1.89 → 2.1.91

package/jest.config.js

@@ -64,7 +64,7 @@ module.exports = {
     "^e2b$": "<rootDir>/src/__tests__/__mocks__/e2b.ts",
     // Strip .js extension from ESM-style imports (must come after all @-scoped mappings)
     "^(@.*)[.]js$": "$1",
-    "^(\..*)[.]js$": "$1",
+    "^(\\.{1,2}/.*)\\.js$": "$1",
   },
   collectCoverageFrom: ["src/**/*.ts", "!src/**/*.d.ts", "!src/__tests__/**/*"],
   coverageReporters: ["text", "lcov", "html"],

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/gateway",
-  "version": "2.1.89",
+  "version": "2.1.91",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
   "types": "dist/index.d.ts",
@@ -40,11 +40,11 @@
     "redis": "^5.0.1",
     "uuid": "^9.0.1",
     "zod": "3.25.76",
-    "@axiom-lattice/agent-eval": "2.1.71",
-    "@axiom-lattice/core": "2.1.77",
-    "@axiom-lattice/pg-stores": "1.0.68",
-    "@axiom-lattice/protocols": "2.1.39",
-    "@axiom-lattice/queue-redis": "1.0.38"
+    "@axiom-lattice/agent-eval": "2.1.73",
+    "@axiom-lattice/core": "2.1.79",
+    "@axiom-lattice/pg-stores": "1.0.70",
+    "@axiom-lattice/protocols": "2.1.41",
+    "@axiom-lattice/queue-redis": "1.0.40"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",

package/src/__tests__/a2a.test.ts

@@ -0,0 +1,346 @@
+import { FastifyInstance } from "fastify";
+import fastify from "fastify";
+import { registerA2ARoutes } from "../routes/a2a";
+
+const mockChunkStream = jest.fn();
+const mockAddMessage = jest.fn();
+const mockAbort = jest.fn();
+
+const mockAgent = {
+  addMessage: mockAddMessage,
+  chunkStream: mockChunkStream,
+  abort: mockAbort,
+};
+
+const mockGetAgent = jest.fn(() => mockAgent);
+
+jest.mock("@axiom-lattice/core", () => ({
+  agentInstanceManager: {
+    getAgent: (...args: unknown[]) => mockGetAgent(...args),
+  },
+}));
+
+describe("A2A Routes", () => {
+  let app: FastifyInstance;
+
+  beforeEach(() => {
+    mockGetAgent.mockClear();
+    mockAddMessage.mockClear();
+    mockChunkStream.mockClear();
+    mockAbort.mockClear();
+    mockGetAgent.mockReturnValue(mockAgent);
+
+    process.env.A2A_API_KEYS = "valid-key:tenant-a";
+    process.env.A2A_DEFAULT_AGENT_ID = "agent-1";
+    process.env.A2A_DEFAULT_TENANT_ID = "default-tenant";
+    delete process.env.A2A_DEFAULT_PROJECT_ID;
+    delete process.env.A2A_DEFAULT_WORKSPACE_ID;
+
+    app = fastify();
+    registerA2ARoutes(app);
+  });
+
+  afterEach(async () => {
+    await app.close();
+    delete process.env.A2A_API_KEYS;
+    delete process.env.A2A_DEFAULT_AGENT_ID;
+    delete process.env.A2A_DEFAULT_TENANT_ID;
+  });
+
+  // ─── Agent Card ──────────────────────────────────────────────────────
+
+  describe("GET /api/a2a/.well-known/agent.json", () => {
+    it("should return agent card without auth", async () => {
+      const response = await app.inject({
+        method: "GET",
+        url: "/api/a2a/.well-known/agent.json",
+      });
+
+      expect(response.statusCode).toBe(200);
+      const body = JSON.parse(response.body);
+      expect(body).toHaveProperty("name", "Axiom Lattice Agent");
+      expect(body).toHaveProperty("capabilities");
+      expect(body.capabilities).toHaveProperty("streaming", true);
+      expect(body).toHaveProperty("defaultInputModes");
+      expect(body).toHaveProperty("defaultOutputModes");
+      expect(body).toHaveProperty("url");
+    });
+  });
+
+  describe("GET /api/a2a/.well-known/agent-card.json", () => {
+    it("should support alternate path", async () => {
+      const response = await app.inject({
+        method: "GET",
+        url: "/api/a2a/.well-known/agent-card.json",
+      });
+
+      expect(response.statusCode).toBe(200);
+      const body = JSON.parse(response.body);
+      expect(body).toHaveProperty("name", "Axiom Lattice Agent");
+    });
+  });
+
+  // ─── Agent Card with env override ────────────────────────────────────
+
+  describe("Agent Card with custom env", () => {
+    it("should reflect custom agent name from env", async () => {
+      process.env.A2A_AGENT_NAME = "Custom Bot";
+      process.env.A2A_ORGANIZATION = "Custom Org";
+      process.env.A2A_VERSION = "2.0.0";
+
+      const response = await app.inject({
+        method: "GET",
+        url: "/api/a2a/.well-known/agent.json",
+      });
+
+      expect(response.statusCode).toBe(200);
+      const body = JSON.parse(response.body);
+      expect(body.name).toBe("Custom Bot");
+      expect(body.provider.organization).toBe("Custom Org");
+      expect(body.version).toBe("2.0.0");
+
+      delete process.env.A2A_AGENT_NAME;
+      delete process.env.A2A_ORGANIZATION;
+      delete process.env.A2A_VERSION;
+    });
+  });
+
+  // ─── Task Send ───────────────────────────────────────────────────────
+
+  describe("POST /api/a2a/tasks/send", () => {
+    it("should return 400 when message.parts is missing", async () => {
+      const response = await app.inject({
+        method: "POST",
+        url: "/api/a2a/tasks/send",
+        headers: {
+          authorization: "Bearer valid-key",
+          "content-type": "application/json",
+        },
+        payload: { message: {} },
+      });
+
+      expect(response.statusCode).toBe(400);
+      const body = JSON.parse(response.body);
+      expect(body.error).toBe("message.parts is required");
+    });
+
+    it("should return 400 when message text is empty", async () => {
+      const response = await app.inject({
+        method: "POST",
+        url: "/api/a2a/tasks/send",
+        headers: {
+          authorization: "Bearer valid-key",
+          "content-type": "application/json",
+        },
+        payload: {
+          message: {
+            role: "user",
+            parts: [{ type: "text", text: "" }],
+          },
+        },
+      });
+
+      expect(response.statusCode).toBe(400);
+      const body = JSON.parse(response.body);
+      expect(body.error).toBe("message must contain text content");
+    });
+
+    it("should return 500 when no agent is configured", async () => {
+      delete process.env.A2A_DEFAULT_AGENT_ID;
+
+      const response = await app.inject({
+        method: "POST",
+        url: "/api/a2a/tasks/send",
+        headers: {
+          authorization: "Bearer valid-key",
+          "content-type": "application/json",
+        },
+        payload: {
+          message: {
+            role: "user",
+            parts: [{ type: "text", text: "Hello" }],
+          },
+        },
+      });
+
+      expect(response.statusCode).toBe(500);
+      const body = JSON.parse(response.body);
+      expect(body.error).toBe("No agent configured");
+    });
+
+    it("should return 401 with invalid API key", async () => {
+      const response = await app.inject({
+        method: "POST",
+        url: "/api/a2a/tasks/send",
+        headers: {
+          authorization: "Bearer wrong-key",
+          "content-type": "application/json",
+        },
+        payload: {
+          message: {
+            role: "user",
+            parts: [{ type: "text", text: "Hello" }],
+          },
+        },
+      });
+
+      expect(response.statusCode).toBe(401);
+    });
+
+    it("should stream task result with valid key and message", async () => {
+      mockAddMessage.mockResolvedValueOnce({ messageId: "msg-1" });
+
+      const mockStream = (async function* () {
+        yield {
+          type: "ai",
+          data: { id: "c1", content: "Hello " },
+        };
+        yield {
+          type: "ai",
+          data: { id: "c2", content: "World" },
+        };
+      })();
+      mockChunkStream.mockReturnValueOnce(mockStream);
+
+      const response = await app.inject({
+        method: "POST",
+        url: "/api/a2a/tasks/send",
+        headers: {
+          authorization: "Bearer valid-key",
+          "content-type": "application/json",
+        },
+        payload: {
+          message: {
+            role: "user",
+            parts: [{ type: "text", text: "Hello World" }],
+          },
+        },
+      });
+
+      expect(response.statusCode).toBe(200);
+      expect(response.headers["content-type"]).toBe("text/event-stream");
+
+      // Verify SSE events were emitted
+      const body = response.body;
+      expect(body).toContain("event: status-update");
+      expect(body).toContain('"state":"working"');
+      expect(body).toContain("Hello ");
+      expect(body).toContain("World");
+      expect(body).toContain('"state":"completed"');
+      expect(body).toContain('"final":true');
+
+      // Verify agent was called with correct scope from API key
+      expect(mockGetAgent).toHaveBeenCalledWith(
+        expect.objectContaining({
+          assistant_id: "agent-1",
+          tenant_id: "tenant-a",
+        }),
+      );
+
+      // Verify agent was called with workspace/project from API key
+      // tenant-a has no project/workspace, so they should be undefined
+      const getAgentCall = mockGetAgent.mock.calls[0][0];
+      expect(getAgentCall.project_id).toBeUndefined();
+      expect(getAgentCall.workspace_id).toBeUndefined();
+    });
+
+    it("should propagate project scope from API key", async () => {
+      process.env.A2A_API_KEYS = "scoped-key:tenant-a:proj-x:ws-y";
+
+      mockAddMessage.mockResolvedValueOnce({ messageId: "msg-1" });
+      const mockStream = (async function* () {
+        yield { type: "ai", data: { id: "c1", content: "ok" } };
+      })();
+      mockChunkStream.mockReturnValueOnce(mockStream);
+
+      const response = await app.inject({
+        method: "POST",
+        url: "/api/a2a/tasks/send",
+        headers: {
+          authorization: "Bearer scoped-key",
+          "content-type": "application/json",
+        },
+        payload: {
+          message: {
+            role: "user",
+            parts: [{ type: "text", text: "Hello" }],
+          },
+        },
+      });
+
+      expect(response.statusCode).toBe(200);
+
+      const getAgentCall = mockGetAgent.mock.calls[0][0];
+      expect(getAgentCall.tenant_id).toBe("tenant-a");
+      expect(getAgentCall.project_id).toBe("proj-x");
+      expect(getAgentCall.workspace_id).toBe("ws-y");
+    });
+  });
+
+  // ─── Task Cancel ─────────────────────────────────────────────────────
+
+  describe("POST /api/a2a/tasks/:taskId/cancel", () => {
+    it("should return 404 for unknown task", async () => {
+      const response = await app.inject({
+        method: "POST",
+        url: "/api/a2a/tasks/nonexistent/cancel",
+        headers: {
+          authorization: "Bearer valid-key",
+          "content-type": "application/json",
+        },
+        payload: {},
+      });
+
+      expect(response.statusCode).toBe(404);
+    });
+  });
+
+  // ─── Task Get ────────────────────────────────────────────────────────
+
+  describe("GET /api/a2a/tasks/:taskId", () => {
+    it("should return 404 for unknown task", async () => {
+      const response = await app.inject({
+        method: "GET",
+        url: "/api/a2a/tasks/nonexistent",
+        headers: {
+          authorization: "Bearer valid-key",
+          "content-type": "application/json",
+        },
+      });
+
+      expect(response.statusCode).toBe(404);
+    });
+  });
+
+  // ─── No auth mode ────────────────────────────────────────────────────
+
+  describe("when no API keys configured", () => {
+    beforeEach(() => {
+      delete process.env.A2A_API_KEYS;
+    });
+
+    it("should accept requests without auth", async () => {
+      mockAddMessage.mockResolvedValueOnce({ messageId: "msg-1" });
+      const mockStream = (async function* () {
+        yield { type: "ai", data: { id: "c1", content: "Hi" } };
+      })();
+      mockChunkStream.mockReturnValueOnce(mockStream);
+
+      const response = await app.inject({
+        method: "POST",
+        url: "/api/a2a/tasks/send",
+        headers: {
+          "content-type": "application/json",
+        },
+        payload: {
+          message: {
+            role: "user",
+            parts: [{ type: "text", text: "Hello" }],
+          },
+        },
+      });
+
+      expect(response.statusCode).toBe(200);
+    });
+  });
+});

package/src/index.ts

@@ -131,6 +131,14 @@ app.addHook("preHandler", async (request, reply) => {
   });
 });
 
+// Add default header values for project/workspace isolation
+app.addHook("onRequest", (request, reply, done) => {
+  if (!request.headers["x-project-id"]) {
+    request.headers["x-project-id"] = "default";
+  }
+  done();
+});
+
 // Add custom logging hooks
 app.addHook("onRequest", (request, reply, done) => {
   const context = {
@@ -294,6 +302,17 @@ const start = async (config?: LatticeGatewayConfig) => {
       });
 
       channelDeps = { router, installationStore };
+
+      // Initialize A2A API key store
+      try {
+        const a2aKeyStore = getStoreLattice("default", "a2aApiKey").store as import("@axiom-lattice/protocols").A2AApiKeyStore;
+        const a2a = await import("./routes/a2a");
+        a2a.setA2AKeyStore(a2aKeyStore);
+        await a2a.refreshStoreKeyMap();
+        logger.info("A2A key store initialized");
+      } catch {
+        // A2A key store optional — env-based keys still work
+      }
     } catch {
       // Stores not registered — channel routes will be skipped gracefully
     }

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,
+  });
+}