agent-relay-server 0.9.0 → 0.10.1

package/src/bus-outbox.ts

@@ -0,0 +1,75 @@
+import { getDb } from "./db";
+
+export interface BusEvent {
+  seq: number;
+  eventType: string;
+  source: string;
+  subject?: string;
+  data: Record<string, unknown>;
+  timestamp: number;
+}
+
+interface OutboxRow {
+  seq: number;
+  event_type: string;
+  source: string;
+  subject: string | null;
+  data: string;
+  timestamp: number;
+}
+
+export function appendEvent(type: string, source: string, data: unknown, subject?: string): number {
+  const timestamp = Date.now();
+  const result = getDb().prepare(`
+    INSERT INTO bus_outbox (event_type, source, subject, data, timestamp)
+    VALUES (?, ?, ?, ?, ?)
+  `).run(type, source, subject ?? null, JSON.stringify(data ?? {}), timestamp);
+  return Number(result.lastInsertRowid);
+}
+
+export function replayEvents(since: number, limit = 500): BusEvent[] {
+  const rows = getDb().prepare(`
+    SELECT seq, event_type, source, subject, data, timestamp
+    FROM bus_outbox
+    WHERE seq > ?
+    ORDER BY seq ASC
+    LIMIT ?
+  `).all(since, limit) as OutboxRow[];
+  return rows.map(rowToEvent);
+}
+
+export function pruneOutbox(retentionMs = 60 * 60 * 1000): number {
+  const threshold = Date.now() - retentionMs;
+  const result = getDb().prepare("DELETE FROM bus_outbox WHERE timestamp < ?").run(threshold);
+  return result.changes;
+}
+
+export function getOutboxCursor(): number {
+  const row = getDb().prepare("SELECT coalesce(max(seq), 0) AS cursor FROM bus_outbox").get() as { cursor: number };
+  return row.cursor;
+}
+
+export function getOldestOutboxCursor(): number {
+  const row = getDb().prepare("SELECT min(seq) AS cursor FROM bus_outbox").get() as { cursor: number | null };
+  return row.cursor ?? 0;
+}
+
+function rowToEvent(row: OutboxRow): BusEvent {
+  return {
+    seq: row.seq,
+    eventType: row.event_type,
+    source: row.source,
+    subject: row.subject ?? undefined,
+    data: parseData(row.data),
+    timestamp: row.timestamp,
+  };
+}
+
+function parseData(raw: string): Record<string, unknown> {
+  try {
+    const parsed = JSON.parse(raw);
+    return parsed && typeof parsed === "object" && !Array.isArray(parsed) ? parsed : { value: parsed };
+  } catch {
+    return {};
+  }
+}

package/src/bus.ts

@@ -0,0 +1,439 @@
+import type { Server, ServerWebSocket } from "bun";
+import { getAgent, getDb, heartbeat, markReady, orphanTasksForAgent, setStatus, upsertAgent, validateAgentSession } from "./db";
+import { getOldestOutboxCursor, getOutboxCursor, replayEvents, type BusEvent } from "./bus-outbox";
+import { emitRelayEvent, subscribeRelayEvents, type RelayEvent } from "./events";
+import { createCommand, getCommand, updateCommand } from "./commands-db";
+import { applyCommandToRecipe } from "./recipe-runner";
+import {
+  BusProtocolError,
+  parseBusFrame,
+  validateClientFrame,
+  type BusFrame,
+  type RegisterFrame,
+} from "agent-relay-sdk/protocol";
+import { getComponentAuth, hasComponentScope, isAuthorized, isOriginAllowed, unauthorized } from "./security";
+import type { AgentCard, Command, Message, Task } from "./types";
+
+interface BusSocketData {
+  id: string;
+  componentScopes?: string[];
+}
+
+type BusWebSocket = ServerWebSocket<BusSocketData>;
+
+interface BusConnection {
+  id: string;
+  componentId: string;
+  agentId?: string;
+  role: string;
+  componentScopes?: string[];
+  instanceId: string;
+  epoch: number;
+  ws: BusWebSocket;
+  subscriptions: Set<string>;
+  lastAck: number;
+  registeredAt: number;
+  transportReconnect: boolean;
+}
+
+const busConnections = new Map<string, BusConnection>();
+
+subscribeRelayEvents((event) => broadcastRelayEvent(event));
+
+export function handleBusUpgrade(req: Request, server: Server<BusSocketData>): Response | undefined {
+  if (!isOriginAllowed(req)) {
+    return Response.json({ error: "origin not allowed" }, { status: 403 });
+  }
+  if (!isAuthorized(req)) return unauthorized(req);
+  const component = getComponentAuth(req);
+  const upgraded = server.upgrade(req, {
+    data: {
+      id: crypto.randomUUID(),
+      componentScopes: component?.scope,
+    },
+  });
+  if (!upgraded) return new Response("WebSocket upgrade failed", { status: 400 });
+  return undefined;
+}
+
+export function busHandleOpen(_ws: BusWebSocket): void {
+  // Registration is explicit. The socket is not visible to projections until
+  // a valid register frame establishes identity and subscriptions.
+}
+
+export function busHandleMessage(ws: BusWebSocket, data: string | Buffer): void {
+  let frame: ReturnType<typeof validateClientFrame>;
+  try {
+    frame = validateClientFrame(parseBusFrame(data));
+  } catch (error) {
+    sendError(ws, undefined, protocolCode(error), error instanceof Error ? error.message : "invalid frame");
+    return;
+  }
+
+  try {
+    handleFrame(ws, frame);
+  } catch (error) {
+    sendError(ws, frame.id, "FRAME_FAILED", error instanceof Error ? error.message : String(error));
+  }
+}
+
+export function busHandleClose(ws: BusWebSocket): void {
+  const conn = busConnections.get(ws.data.id);
+  busConnections.delete(ws.data.id);
+  if (conn?.agentId) {
+    if (!conn.transportReconnect) {
+      setStatus(conn.agentId, "stale", { instanceId: conn.instanceId, epoch: conn.epoch });
+      emitAgentStatusEvent(conn.agentId);
+    }
+  }
+}
+
+export function getBusConnectionCount(): number {
+  return busConnections.size;
+}
+
+export function expireStaleBusAgents(graceMs = Number(process.env.AGENT_RELAY_STALE_GRACE_MS) || 120_000): { agentIds: string[]; orphanedTasks: Task[] } {
+  const cutoff = Date.now() - graceMs;
+  const rows = getDb()
+    .prepare("SELECT id FROM agents WHERE status = 'stale' AND last_seen < ? AND id NOT IN ('user', 'system')")
+    .all(cutoff) as Array<{ id: string }>;
+  const orphanedTasks: Task[] = [];
+  for (const row of rows) {
+    setStatus(row.id, "offline");
+    emitAgentStatusEvent(row.id);
+    const tasks = orphanTasksForAgent(row.id);
+    for (const task of tasks) {
+      orphanedTasks.push(task);
+      emitRelayEvent({
+        type: "task.orphaned",
+        source: "agent-relay",
+        subject: String(task.id),
+        data: task as unknown as Record<string, unknown>,
+      });
+    }
+  }
+  return { agentIds: rows.map((row) => row.id), orphanedTasks };
+}
+
+function handleFrame(ws: BusWebSocket, frame: ReturnType<typeof validateClientFrame>): void {
+  switch (frame.type) {
+    case "register":
+      handleRegister(ws, frame);
+      return;
+    case "heartbeat":
+      withRegistered(ws, frame, (conn) => {
+        if (conn.agentId) {
+          heartbeat(conn.agentId, { instanceId: conn.instanceId, epoch: conn.epoch });
+          if (frame.payload.status) setStatus(conn.agentId, frame.payload.status, { instanceId: conn.instanceId, epoch: conn.epoch });
+          emitAgentStatusEvent(conn.agentId);
+        }
+      });
+      return;
+    case "subscribe":
+      withRegistered(ws, frame, (conn) => {
+        for (const event of frame.payload.events) conn.subscriptions.add(event);
+        send(ws, { type: "ack", payload: { frameId: frame.id } });
+      });
+      return;
+    case "status":
+      withRegistered(ws, frame, (conn) => {
+        if (!conn.agentId) return sendError(ws, frame.id, "AGENT_REQUIRED", "status frames require an agent registration");
+        const guard = { instanceId: conn.instanceId, epoch: conn.epoch };
+        if (!setStatus(conn.agentId, frame.payload.agentStatus, guard)) {
+          return sendError(ws, frame.id, "AGENT_NOT_FOUND", "agent not found");
+        }
+        if (frame.payload.ready !== undefined && !markReady(conn.agentId, frame.payload.ready, guard)) {
+          return sendError(ws, frame.id, "AGENT_NOT_FOUND", "agent not found");
+        }
+        emitAgentStatusEvent(conn.agentId);
+      });
+      return;
+    case "ack":
+      withRegistered(ws, frame, (conn) => {
+        const parsed = Number(frame.payload.frameId);
+        if (Number.isSafeInteger(parsed)) conn.lastAck = Math.max(conn.lastAck, parsed);
+      });
+      return;
+    case "resume":
+      withRegistered(ws, frame, () => handleResume(ws, frame.id, frame.payload.since));
+      return;
+    case "command":
+      withRegistered(ws, frame, (conn) => handleCommandFrame(ws, conn, frame.id, frame.payload.commandType, frame.payload.target, frame.payload.params));
+      return;
+  }
+}
+
+function handleCommandFrame(
+  ws: BusWebSocket,
+  conn: BusConnection,
+  frameId: string,
+  commandType: string,
+  target: string,
+  params: Record<string, unknown>,
+): void {
+  if (conn.componentScopes && !hasComponentScope({ sub: conn.componentId, role: conn.role, scope: conn.componentScopes, iat: 0 }, "command:*")) {
+    sendCommandResult(ws, frameId, "rejected", undefined, "component token lacks command scope");
+    return;
+  }
+
+  if (commandType === "command.update") {
+    const command = getCommand(target);
+    if (!command) {
+      sendCommandResult(ws, frameId, "failed", undefined, "command not found");
+      return;
+    }
+    const status = typeof params.status === "string" ? params.status : undefined;
+    const updated = updateCommand(command.id, {
+      status: status as any,
+      result: isRecord(params.result) ? params.result : undefined,
+      error: typeof params.error === "string" ? params.error : undefined,
+    });
+    if (updated) {
+      applyCommandToRecipe(updated);
+      emitCommandEvent(updated, `command.${updated.status}`);
+    }
+    sendCommandResult(ws, frameId, "succeeded", updated ? { command: updated } : undefined);
+    return;
+  }
+
+  const command = createCommand({
+    type: commandType,
+    source: conn.agentId ?? conn.componentId,
+    target,
+    params,
+  });
+  emitCommandEvent(command, "command.requested");
+  sendCommandResult(ws, frameId, "succeeded", { command });
+}
+
+function handleRegister(ws: BusWebSocket, frame: RegisterFrame): void {
+  const payload = frame.payload;
+  const runnerManaged = payload.meta?.runnerManaged === true || typeof payload.meta?.runnerId === "string";
+  if (runnerManaged && !hasRunnerLifecycleCapabilities(payload.capabilities, payload.meta)) {
+    sendError(ws, frame.id, "INVALID_RUNNER_REGISTRATION", "runner agents must advertise hard restart/shutdown and semantic status capabilities");
+    return;
+  }
+  let epoch = 0;
+  if (payload.agentId) {
+    const label = stringMeta(payload.meta, "label");
+    const agent = upsertAgent({
+      id: payload.agentId,
+      name: stringMeta(payload.meta, "name") ?? payload.componentId,
+      kind: payload.role === "integration" ? "provider" : payload.role,
+      ...(label ? { label } : {}),
+      tags: payload.tags,
+      machine: payload.machine,
+      capabilities: payload.capabilities,
+      status: runnerManaged ? "idle" : "online",
+      ready: true,
+      instanceId: payload.instanceId,
+      meta: payload.meta,
+    });
+    epoch = agent.epoch;
+    emitAgentStatusEvent(agent.id);
+  }
+
+  busConnections.set(ws.data.id, {
+    id: ws.data.id,
+    componentId: payload.componentId,
+    agentId: payload.agentId,
+    role: payload.role,
+    componentScopes: ws.data.componentScopes,
+    instanceId: payload.instanceId,
+    epoch,
+    ws,
+    subscriptions: new Set(["message.*", "agent.status", "task.*", "command.*"]),
+    lastAck: 0,
+    registeredAt: Date.now(),
+    transportReconnect: runnerManaged && isRecord(payload.meta.lifecycleCapabilities) && payload.meta.lifecycleCapabilities.transportReconnect === true,
+  });
+
+  send(ws, {
+    type: "registered",
+    payload: {
+      epoch,
+      cursor: getOutboxCursor(),
+      sessionId: ws.data.id,
+    },
+  });
+}
+
+function hasRunnerLifecycleCapabilities(capabilities: string[], meta: Record<string, unknown>): boolean {
+  const lifecycle = isRecord(meta.lifecycleCapabilities) ? meta.lifecycleCapabilities : {};
+  return capabilities.includes("lifecycle.shutdown.hard") &&
+    capabilities.includes("lifecycle.restart.hard") &&
+    capabilities.includes("lifecycle.status.semantic") &&
+    lifecycle.shutdownHard === true &&
+    lifecycle.restartHard === true &&
+    lifecycle.semanticStatus === true;
+}
+
+function handleResume(ws: BusWebSocket, frameId: string, since: number): void {
+  const oldest = getOldestOutboxCursor();
+  if (since > 0 && oldest > 0 && oldest > since + 1) {
+    sendError(ws, frameId, "REPLAY_GAP", "requested cursor is older than retained bus outbox events");
+    return;
+  }
+  const events = replayEvents(since).map(outboxToPayload);
+  send(ws, {
+    type: "resumed",
+    payload: {
+      events,
+      fromSeq: events[0]?.seq ?? since,
+      toSeq: events.at(-1)?.seq ?? since,
+    },
+  });
+}
+
+function withRegistered(
+  ws: BusWebSocket,
+  frame: BusFrame,
+  fn: (conn: BusConnection) => void,
+): void {
+  const conn = busConnections.get(ws.data.id);
+  if (!conn) {
+    sendError(ws, frame.id, "NOT_REGISTERED", "register before sending bus frames");
+    return;
+  }
+  if (conn.agentId) {
+    const session = validateAgentSession(conn.agentId, { instanceId: conn.instanceId, epoch: conn.epoch });
+    if (!session.ok) {
+      sendError(ws, frame.id, "STALE_SESSION", session.error ?? "stale agent session");
+      return;
+    }
+  }
+  fn(conn);
+}
+
+function broadcastRelayEvent(event: RelayEvent): void {
+  for (const conn of busConnections.values()) {
+    if (!connectionWantsEvent(conn, event)) continue;
+    send(conn.ws, {
+      type: "event",
+      payload: {
+        seq: event.seq,
+        eventType: event.type,
+        source: event.source,
+        subject: event.subject,
+        data: event.data,
+        timestamp: event.timestamp,
+      },
+    });
+  }
+}
+
+function connectionWantsEvent(conn: BusConnection, event: RelayEvent): boolean {
+  if (!matchesSubscription(conn.subscriptions, event.type)) return false;
+  if (event.type === "message.new" && conn.agentId) {
+    return messageMatchesAgent(event.data as unknown as Message, conn.agentId);
+  }
+  if (event.type.startsWith("task.") && conn.agentId) {
+    const target = typeof event.data.target === "string" ? event.data.target : "";
+    return targetMatchesAgent(target, conn.agentId);
+  }
+  if (event.type.startsWith("command.")) {
+    const command = isRecord(event.data.command) ? event.data.command : undefined;
+    if (!command) return true;
+    const target = typeof command.target === "string" ? command.target : "";
+    const source = typeof command.source === "string" ? command.source : "";
+    return target === conn.componentId || target === conn.agentId || source === conn.componentId || source === conn.agentId;
+  }
+  return true;
+}
+
+function matchesSubscription(subscriptions: Set<string>, eventType: string): boolean {
+  if (subscriptions.size === 0) return false;
+  for (const sub of subscriptions) {
+    if (sub === "*" || sub === eventType) return true;
+    if (sub.endsWith("*") && eventType.startsWith(sub.slice(0, -1))) return true;
+  }
+  return false;
+}
+
+function messageMatchesAgent(msg: Message, agentId: string): boolean {
+  const agent = getAgent(agentId);
+  if (!agent) return false;
+  if (msg.claimable && msg.claimedBy && msg.claimedBy !== agentId) return false;
+  if (msg.to === agentId || msg.from === agentId) return true;
+  return targetMatchesAgent(msg.to, agentId, agent);
+}
+
+function targetMatchesAgent(target: string, agentId: string, knownAgent?: AgentCard): boolean {
+  const agent = knownAgent ?? getAgent(agentId);
+  if (!agent) return false;
+  if (target === agentId || target === "broadcast") return true;
+  if (target.startsWith("tag:") && agent.tags.includes(target.slice(4))) return true;
+  if (target.startsWith("cap:") && agent.capabilities.includes(target.slice(4))) return true;
+  if (target.startsWith("label:") && agent.label === target.slice(6)) return true;
+  return false;
+}
+
+function outboxToPayload(event: BusEvent) {
+  return {
+    seq: event.seq,
+    eventType: event.eventType,
+    source: event.source,
+    subject: event.subject,
+    data: event.data,
+    timestamp: event.timestamp,
+  };
+}
+
+function emitAgentStatusEvent(agentId: string): void {
+  const agent = getAgent(agentId);
+  const data = agent ?? { id: agentId, status: "offline" };
+  emitRelayEvent({
+    type: "agent.status",
+    source: "server",
+    subject: agentId,
+    data: data as unknown as Record<string, unknown>,
+  });
+}
+
+function emitCommandEvent(command: Command, type: string): void {
+  emitRelayEvent({
+    type,
+    source: command.source,
+    subject: command.id,
+    data: { command },
+  });
+}
+
+function sendCommandResult(
+  ws: BusWebSocket,
+  commandId: string,
+  status: "succeeded" | "failed" | "rejected" | "timed_out",
+  result?: Record<string, unknown>,
+  error?: string,
+): void {
+  send(ws, {
+    type: "command.result",
+    payload: {
+      commandId,
+      status,
+      result,
+      error,
+    },
+  });
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
+}
+
+function send(ws: BusWebSocket, frame: Record<string, unknown>): void {
+  ws.send(JSON.stringify(frame));
+}
+
+function sendError(ws: BusWebSocket, frameId: string | undefined, code: string, message: string): void {
+  send(ws, { type: "error", payload: { frameId, code, message } });
+}
+
+function protocolCode(error: unknown): string {
+  return error instanceof BusProtocolError ? error.code : "INVALID_FRAME";
+}
+
+function stringMeta(meta: Record<string, unknown>, key: string): string | undefined {
+  const value = meta[key];
+  return typeof value === "string" && value.length > 0 ? value : undefined;
+}