@axiom-lattice/gateway 2.1.89 → 2.1.90

package/src/routes/a2a.ts

@@ -0,0 +1,779 @@
+import type { FastifyInstance, FastifyRequest, FastifyReply } from "fastify";
+import { v4 } from "uuid";
+import {
+  agentInstanceManager,
+} from "@axiom-lattice/core";
+import { MessageChunkTypes } from "@axiom-lattice/protocols";
+import type {
+  AgentCard,
+  A2ATask,
+  A2ATaskState,
+  A2ATaskSendRequest,
+  A2AMessage,
+  A2APart,
+  A2AAuthContext,
+  A2AApiKeyStore,
+  A2AApiKeyRecord,
+  CreateA2AApiKeyInput,
+  A2AApiKeyEntry,
+} from "@axiom-lattice/protocols";
+
+// ─── Task Tracking ────────────────────────────────────────────────────────
+
+interface TaskRecord {
+  id: string;
+  threadId: string;
+  messageId?: string;
+  assistantId: string;
+  tenantId: string;
+  projectId?: string;
+  workspaceId?: string;
+  status: A2ATaskState;
+  createdAt: string;
+  updatedAt: string;
+  history: A2AMessage[];
+}
+
+const taskStore = new Map<string, TaskRecord>();
+
+// ─── Key Store ────────────────────────────────────────────────────────────
+
+let _a2aKeyStore: A2AApiKeyStore | null = null;
+let _storeKeyMap: Map<string, A2AApiKeyEntry> = new Map();
+
+export function setA2AKeyStore(store: A2AApiKeyStore): void {
+  _a2aKeyStore = store;
+}
+
+/**
+ * Reload the cached key map from the store.
+ * Call after any CRUD mutation or at startup.
+ */
+export async function refreshStoreKeyMap(): Promise<void> {
+  if (!_a2aKeyStore) return;
+  _storeKeyMap = await _a2aKeyStore.loadIntoMap();
+}
+
+// ─── Config ───────────────────────────────────────────────────────────────
+
+function parseEnvKeyMap(): Map<string, A2AApiKeyEntry> {
+  const keysEnv = process.env.A2A_API_KEYS || "";
+  const defaultTenantId = process.env.A2A_DEFAULT_TENANT_ID || "a2a-default-tenant";
+  const defaultProjectId = process.env.A2A_DEFAULT_PROJECT_ID || "";
+  const defaultWorkspaceId = process.env.A2A_DEFAULT_WORKSPACE_ID || "";
+  const map = new Map<string, A2AApiKeyEntry>();
+
+  if (keysEnv) {
+    keysEnv.split(",").forEach((entry) => {
+      const trimmed = entry.trim();
+      if (!trimmed) return;
+      const parts = trimmed.split(":");
+      const key = parts[0];
+      if (!key) return;
+      map.set(key, {
+        key,
+        tenantId: parts[1]?.trim() || defaultTenantId,
+        projectId: parts[2]?.trim() || defaultProjectId || undefined,
+        workspaceId: parts[3]?.trim() || defaultWorkspaceId || undefined,
+      });
+    });
+  }
+
+  return map;
+}
+
+function getA2AConfig() {
+  const envMap = parseEnvKeyMap();
+  // Store keys take priority; env keys serve as fallback
+  const apiKeyMap = new Map([..._storeKeyMap, ...envMap]);
+
+  return {
+    agentName: process.env.A2A_AGENT_NAME || "Axiom Lattice Agent",
+    agentDescription:
+      process.env.A2A_AGENT_DESCRIPTION ||
+      "AI agent powered by Axiom Lattice framework",
+    organization: process.env.A2A_ORGANIZATION || "Axiom Lattice",
+    organizationUrl: process.env.A2A_ORGANIZATION_URL || "https://axiom-lattice.ai",
+    version: process.env.A2A_VERSION || "1.0.0",
+    defaultAssistantId: process.env.A2A_DEFAULT_AGENT_ID || "",
+    defaultTenantId: process.env.A2A_DEFAULT_TENANT_ID || "a2a-default-tenant",
+    defaultProjectId: process.env.A2A_DEFAULT_PROJECT_ID || "",
+    defaultWorkspaceId: process.env.A2A_DEFAULT_WORKSPACE_ID || "",
+    apiKeyMap,
+  };
+}
+
+// ─── Auth ─────────────────────────────────────────────────────────────────
+
+function authenticateA2A(
+  request: FastifyRequest,
+): A2AAuthContext {
+  const config = getA2AConfig();
+  const noAuthRequired = config.apiKeyMap.size === 0;
+
+  const token =
+    request.headers.authorization?.startsWith("Bearer ")
+      ? request.headers.authorization.substring(7)
+      : (request.headers["x-api-key"] as string | undefined);
+
+  if (!token) {
+    return {
+      authenticated: noAuthRequired,
+      tenantId: config.defaultTenantId,
+      projectId: config.defaultProjectId || undefined,
+      workspaceId: config.defaultWorkspaceId || undefined,
+    };
+  }
+
+  const entry = config.apiKeyMap.get(token);
+  if (entry) {
+    return {
+      authenticated: true,
+      apiKey: token,
+      tenantId: entry.tenantId,
+      projectId: entry.projectId,
+      workspaceId: entry.workspaceId,
+      source: request.headers.authorization?.startsWith("Bearer ") ? "bearer" : "x-api-key",
+    };
+  }
+
+  return { authenticated: false };
+}
+
+function requireAuth(
+  request: FastifyRequest,
+  reply: FastifyReply,
+): A2AAuthContext {
+  const auth = authenticateA2A(request);
+  if (!auth.authenticated) {
+    reply.status(401).send({
+      error: "Unauthorized",
+      message: "Valid API key required via Bearer token or X-API-Key header",
+    });
+  }
+  return auth;
+}
+
+// ─── Agent Card ───────────────────────────────────────────────────────────
+
+function buildAgentCard(baseUrl: string): AgentCard {
+  const config = getA2AConfig();
+  return {
+    name: config.agentName,
+    description: config.agentDescription,
+    url: `${baseUrl}/api/a2a`,
+    provider: {
+      organization: config.organization,
+      url: config.organizationUrl,
+    },
+    version: config.version,
+    capabilities: {
+      streaming: true,
+      pushNotifications: false,
+      stateTransitionHistory: false,
+    },
+    defaultInputModes: ["text", "text/plain"],
+    defaultOutputModes: ["text", "text/plain", "text/markdown"],
+    skills: [],
+  };
+}
+
+// ─── Helpers ──────────────────────────────────────────────────────────────
+
+function extractTextFromParts(parts: A2APart[]): string {
+  return parts
+    .filter((p): p is { type: "text"; text: string } => p.type === "text")
+    .map((p) => p.text)
+    .join("\n");
+}
+
+function buildTextPart(text: string): A2APart {
+  return { type: "text", text };
+}
+
+function isoNow(): string {
+  return new Date().toISOString();
+}
+
+type SendEventFn = (event: string, data: Record<string, unknown>) => void;
+
+async function streamChunksAsA2A(
+  stream: AsyncIterable<{ type: unknown; data?: { content?: string } }>,
+  taskId: string,
+  sendEvent: SendEventFn,
+  onInterrupt?: () => void,
+): Promise<string> {
+  let accumulatedText = "";
+
+  for await (const chunk of stream) {
+    const chunkText = (chunk as { data?: { content?: string } }).data?.content || "";
+    const chunkType = (chunk as { type: string }).type;
+
+    if (chunkType === MessageChunkTypes.INTERRUPT) {
+      sendEvent("status-update", {
+        id: taskId,
+        status: {
+          state: "input-required",
+          message: {
+            role: "agent",
+            parts: [buildTextPart(accumulatedText)],
+          },
+          timestamp: isoNow(),
+        },
+        final: false,
+      });
+      onInterrupt?.();
+      continue;
+    }
+
+    if (
+      chunkType === MessageChunkTypes.AI ||
+      chunkType === MessageChunkTypes.TOOL
+    ) {
+      if (chunkText) {
+        accumulatedText += chunkText;
+        sendEvent("status-update", {
+          id: taskId,
+          status: {
+            state: "working",
+            message: {
+              role: "agent",
+              parts: [buildTextPart(chunkText)],
+            },
+            timestamp: isoNow(),
+          },
+          final: false,
+        });
+      }
+    }
+  }
+
+  return accumulatedText;
+}
+
+// ─── Task Status ──────────────────────────────────────────────────────────
+
+async function getTaskStatus(taskId: string): Promise<A2ATask | null> {
+  const record = taskStore.get(taskId);
+  if (!record) return null;
+
+  return {
+    id: record.id,
+    status: {
+      state: record.status,
+      timestamp: record.updatedAt,
+    },
+    artifacts: [],
+    history: record.history,
+  };
+}
+
+// ─── Handlers ─────────────────────────────────────────────────────────────
+
+async function handleAgentCard(
+  request: FastifyRequest,
+  reply: FastifyReply,
+): Promise<void> {
+  const protocol =
+    (request.headers["x-forwarded-proto"] as string) ||
+    request.protocol;
+  const host = request.hostname;
+  const baseUrl = `${protocol}://${host}`;
+  const card = buildAgentCard(baseUrl);
+  reply.status(200).send(card);
+}
+
+async function handleTasksGet(
+  request: FastifyRequest<{ Params: { taskId: string } }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+  const task = await getTaskStatus(request.params.taskId);
+  if (!task) {
+    reply.status(404).send({ error: "Task not found" });
+    return;
+  }
+  reply.status(200).send(task);
+}
+
+async function handleTasksSend(
+  request: FastifyRequest<{ Body: A2ATaskSendRequest }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+
+  const config = getA2AConfig();
+  const body = request.body as A2ATaskSendRequest;
+
+  if (!body.message || !body.message.parts || body.message.parts.length === 0) {
+    reply.status(400).send({ error: "message.parts is required" });
+    return;
+  }
+
+  const text = extractTextFromParts(body.message.parts);
+  if (!text.trim()) {
+    reply.status(400).send({ error: "message must contain text content" });
+    return;
+  }
+
+  const assistantId =
+    (request.headers["x-a2a-agent-id"] as string) ||
+    config.defaultAssistantId;
+
+  if (!assistantId) {
+    reply.status(500).send({
+      error: "No agent configured",
+      message:
+        "Set A2A_DEFAULT_AGENT_ID env var or provide x-a2a-agent-id header",
+    });
+    return;
+  }
+
+  const taskId = body.id || v4();
+  const threadId = v4();
+
+  const tenantId = auth.tenantId || config.defaultTenantId;
+
+  const now = isoNow();
+
+  const record: TaskRecord = {
+    id: taskId,
+    threadId,
+    assistantId,
+    tenantId,
+    projectId: auth.projectId || undefined,
+    workspaceId: auth.workspaceId || undefined,
+    status: "working",
+    createdAt: now,
+    updatedAt: now,
+    history: [
+      {
+        role: "user",
+        parts: body.message.parts,
+      },
+    ],
+  };
+
+  taskStore.set(taskId, record);
+
+  let agent;
+  try {
+    agent = agentInstanceManager.getAgent({
+      assistant_id: assistantId,
+      thread_id: threadId,
+      tenant_id: tenantId,
+      workspace_id: auth.workspaceId || undefined,
+      project_id: auth.projectId || undefined,
+    });
+  } catch (err) {
+    record.status = "failed";
+    record.updatedAt = isoNow();
+    reply.status(500).send({
+      id: taskId,
+      error: "Failed to initialize agent",
+      message: err instanceof Error ? err.message : String(err),
+    });
+    return;
+  }
+
+  reply.hijack();
+  reply.raw.writeHead(200, {
+    "Content-Type": "text/event-stream",
+    "Cache-Control": "no-cache",
+    Connection: "keep-alive",
+    "Access-Control-Allow-Origin": "*",
+  });
+
+  const sendEvent = (event: string, data: Record<string, unknown>): void => {
+    reply.raw.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`);
+  };
+
+  try {
+    const result = await agent.addMessage({
+      input: { message: text },
+    });
+
+    record.messageId = result.messageId;
+
+    const stream = agent.chunkStream(result.messageId, [
+      MessageChunkTypes.MESSAGE_COMPLETED,
+    ]);
+
+    const accumulatedText = await streamChunksAsA2A(stream, taskId, sendEvent, () => {
+      record.status = "input-required";
+      record.updatedAt = isoNow();
+    });
+
+    record.status = "completed";
+    record.updatedAt = isoNow();
+    record.history.push({
+      role: "agent",
+      parts: [buildTextPart(accumulatedText)],
+    });
+
+    sendEvent("status-update", {
+      id: taskId,
+      status: {
+        state: "completed",
+        timestamp: isoNow(),
+      },
+      final: true,
+    });
+  } catch (err) {
+    record.status = "failed";
+    record.updatedAt = isoNow();
+
+    sendEvent("error", {
+      code: "EXECUTION_ERROR",
+      message: err instanceof Error ? err.message : String(err),
+    });
+
+    sendEvent("status-update", {
+      id: taskId,
+      status: {
+        state: "failed",
+        timestamp: isoNow(),
+      },
+      final: true,
+    });
+  } finally {
+    reply.raw.end();
+  }
+}
+
+async function handleTasksCancel(
+  request: FastifyRequest<{ Params: { taskId: string } }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+  const { taskId } = request.params;
+  const record = taskStore.get(taskId);
+
+  if (!record) {
+    reply.status(404).send({ error: "Task not found" });
+    return;
+  }
+
+  if (record.status === "completed" || record.status === "failed" || record.status === "canceled") {
+    reply.status(409).send({
+      error: "Task is already in a terminal state",
+      task: { id: record.id, status: record.status },
+    });
+    return;
+  }
+
+  try {
+    const agent = agentInstanceManager.getAgent({
+      assistant_id: record.assistantId,
+      thread_id: record.threadId,
+      tenant_id: record.tenantId,
+      workspace_id: record.workspaceId,
+      project_id: record.projectId,
+    });
+    await agent.abort();
+  } catch (err) {
+    // Agent session may have already been cleaned up
+    console.warn({
+      event: "a2a:cancel:no_agent",
+      taskId,
+      threadId: record.threadId,
+      error: err instanceof Error ? err.message : String(err),
+    });
+  }
+
+  record.status = "canceled";
+  record.updatedAt = isoNow();
+
+  reply.status(200).send({
+    id: taskId,
+    status: {
+      state: "canceled",
+      timestamp: record.updatedAt,
+    },
+  });
+}
+
+async function handleTasksStream(
+  request: FastifyRequest<{ Params: { taskId: string } }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+
+  const { taskId } = request.params;
+  const record = taskStore.get(taskId);
+
+  if (!record) {
+    reply.status(404).send({ error: "Task not found" });
+    return;
+  }
+
+  let agent;
+  try {
+    agent = agentInstanceManager.getAgent({
+      assistant_id: record.assistantId,
+      thread_id: record.threadId,
+      tenant_id: record.tenantId,
+      workspace_id: record.workspaceId,
+      project_id: record.projectId,
+    });
+  } catch {
+    reply.status(404).send({
+      error: "Agent session not found. The task may have already completed.",
+    });
+    return;
+  }
+
+  reply.hijack();
+  reply.raw.writeHead(200, {
+    "Content-Type": "text/event-stream",
+    "Cache-Control": "no-cache",
+    Connection: "keep-alive",
+    "Access-Control-Allow-Origin": "*",
+  });
+
+  const sendEvent = (event: string, data: Record<string, unknown>): void => {
+    reply.raw.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`);
+  };
+
+  if (!record.messageId) {
+    sendEvent("status-update", {
+      id: taskId,
+      status: {
+        state: record.status,
+        timestamp: record.updatedAt,
+      },
+      final: true,
+    });
+    reply.raw.end();
+    return;
+  }
+
+  try {
+    const stream = agent.chunkStream(record.messageId, [
+      MessageChunkTypes.MESSAGE_COMPLETED,
+    ]);
+
+    await streamChunksAsA2A(stream, taskId, sendEvent);
+
+    sendEvent("status-update", {
+      id: taskId,
+      status: {
+        state: record.status,
+        timestamp: record.updatedAt,
+      },
+      final: true,
+    });
+  } catch (err) {
+    sendEvent("error", {
+      code: "STREAM_ERROR",
+      message: err instanceof Error ? err.message : String(err),
+    });
+  } finally {
+    reply.raw.end();
+  }
+}
+
+// ─── API Key Management Handlers ─────────────────────────────────────────
+
+async function handleApiKeyList(
+  request: FastifyRequest<{ Querystring: { tenantId?: string; limit?: string; offset?: string } }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+
+  if (!_a2aKeyStore) {
+    reply.status(501).send({ error: "Key store not configured" });
+    return;
+  }
+
+  const records = await _a2aKeyStore.list({
+    tenantId: request.query.tenantId,
+    limit: request.query.limit ? parseInt(request.query.limit, 10) : undefined,
+    offset: request.query.offset ? parseInt(request.query.offset, 10) : undefined,
+  });
+
+  reply.status(200).send({
+    success: true,
+    data: {
+      records: records.map((r) => ({ ...r, key: r.key.slice(0, 8) + "..." })),
+      total: records.length,
+    },
+  });
+}
+
+async function handleApiKeyCreate(
+  request: FastifyRequest<{ Body: Partial<CreateA2AApiKeyInput> }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+
+  if (!_a2aKeyStore) {
+    reply.status(501).send({ error: "Key store not configured" });
+    return;
+  }
+
+  const body = request.body as Partial<CreateA2AApiKeyInput>;
+  // Scope from request body / headers; fallback to auth context
+  const tenantId =
+    body.tenantId ||
+    (request.headers["x-tenant-id"] as string) ||
+    auth.tenantId;
+  const projectId =
+    body.projectId || (request.headers["x-project-id"] as string) || auth.projectId || "default";
+  const workspaceId =
+    body.workspaceId || (request.headers["x-workspace-id"] as string) || auth.workspaceId;
+
+  if (!tenantId) {
+    reply.status(400).send({
+      error: "tenantId is required",
+      message: "Provide via body.tenantId, x-tenant-id header, or authenticate with a scoped API key",
+    });
+    return;
+  }
+
+  const record = await _a2aKeyStore.create({
+    tenantId,
+    projectId,
+    workspaceId,
+    label: body.label,
+  });
+  await refreshStoreKeyMap();
+
+  reply.status(201).send({ success: true, data: record });
+}
+
+async function handleApiKeyDelete(
+  request: FastifyRequest<{ Params: { id: string } }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+
+  if (!_a2aKeyStore) {
+    reply.status(501).send({ error: "Key store not configured" });
+    return;
+  }
+
+  await _a2aKeyStore.delete(request.params.id);
+  await refreshStoreKeyMap();
+
+  reply.status(200).send({ success: true });
+}
+
+async function handleApiKeyDisable(
+  request: FastifyRequest<{ Params: { id: string } }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+
+  if (!_a2aKeyStore) {
+    reply.status(501).send({ error: "Key store not configured" });
+    return;
+  }
+
+  await _a2aKeyStore.disable(request.params.id);
+  await refreshStoreKeyMap();
+
+  reply.status(200).send({ success: true });
+}
+
+async function handleApiKeyEnable(
+  request: FastifyRequest<{ Params: { id: string } }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+
+  if (!_a2aKeyStore) {
+    reply.status(501).send({ error: "Key store not configured" });
+    return;
+  }
+
+  await _a2aKeyStore.enable(request.params.id);
+  await refreshStoreKeyMap();
+
+  reply.status(200).send({ success: true });
+}
+
+async function handleApiKeyRotate(
+  request: FastifyRequest<{ Params: { id: string } }>,
+  reply: FastifyReply,
+): Promise<void> {
+  const auth = requireAuth(request, reply);
+  if (!auth.authenticated) return;
+
+  if (!_a2aKeyStore) {
+    reply.status(501).send({ error: "Key store not configured" });
+    return;
+  }
+
+  const record = await _a2aKeyStore.rotate(request.params.id);
+  await refreshStoreKeyMap();
+
+  reply.status(200).send({ success: true, data: record });
+}
+
+// ─── Route Registration ───────────────────────────────────────────────────
+
+export function registerA2ARoutes(app: FastifyInstance): void {
+  app.get("/api/a2a/.well-known/agent.json", handleAgentCard);
+  app.get("/api/a2a/.well-known/agent-card.json", handleAgentCard);
+
+  app.post<{ Body: A2ATaskSendRequest }>(
+    "/api/a2a/tasks/send",
+    handleTasksSend,
+  );
+
+  app.get<{ Params: { taskId: string } }>(
+    "/api/a2a/tasks/:taskId",
+    handleTasksGet,
+  );
+
+  app.get<{ Params: { taskId: string } }>(
+    "/api/a2a/tasks/:taskId/stream",
+    handleTasksStream,
+  );
+
+  app.post<{ Params: { taskId: string } }>(
+    "/api/a2a/tasks/:taskId/cancel",
+    handleTasksCancel,
+  );
+
+  // A2A API Key management
+  app.get<{ Querystring: { tenantId?: string; limit?: string; offset?: string } }>(
+    "/api/a2a/keys",
+    handleApiKeyList,
+  );
+
+  app.post<{ Body: Partial<CreateA2AApiKeyInput> }>(
+    "/api/a2a/keys",
+    handleApiKeyCreate,
+  );
+
+  app.delete<{ Params: { id: string } }>(
+    "/api/a2a/keys/:id",
+    handleApiKeyDelete,
+  );
+
+  app.post<{ Params: { id: string } }>(
+    "/api/a2a/keys/:id/disable",
+    handleApiKeyDisable,
+  );
+
+  app.post<{ Params: { id: string } }>(
+    "/api/a2a/keys/:id/enable",
+    handleApiKeyEnable,
+  );
+
+  app.post<{ Params: { id: string } }>(
+    "/api/a2a/keys/:id/rotate",
+    handleApiKeyRotate,
+  );
+}