fix(gateway): stream chat events for agent runs

src/infra/agent-events.test.ts

@@ -1,7 +1,22 @@
 import { describe, expect, test } from "vitest";
-import { emitAgentEvent, onAgentEvent } from "./agent-events.js";
+import {
+  emitAgentEvent,
+  onAgentEvent,
+  registerAgentRunContext,
+  getAgentRunContext,
+  clearAgentRunContext,
+  resetAgentRunContextForTest,
+} from "./agent-events.js";
 
 describe("agent-events sequencing", () => {
+  test("stores and clears run context", async () => {
+    resetAgentRunContextForTest();
+    registerAgentRunContext("run-1", { sessionKey: "main" });
+    expect(getAgentRunContext("run-1")?.sessionKey).toBe("main");
+    clearAgentRunContext("run-1");
+    expect(getAgentRunContext("run-1")).toBeUndefined();
+  });
+
   test("maintains monotonic seq per runId", async () => {
     const seen: Record<string, number[]> = {};
     const stop = onAgentEvent((evt) => {

src/infra/agent-events.ts

@@ -13,9 +13,41 @@ export type AgentEventPayload = {
   data: Record<string, unknown>;
 };
 
+export type AgentRunContext = {
+  sessionKey?: string;
+};
+
 // Keep per-run counters so streams stay strictly monotonic per runId.
 const seqByRun = new Map<string, number>();
 const listeners = new Set<(evt: AgentEventPayload) => void>();
+const runContextById = new Map<string, AgentRunContext>();
+
+export function registerAgentRunContext(
+  runId: string,
+  context: AgentRunContext,
+) {
+  if (!runId) return;
+  const existing = runContextById.get(runId);
+  if (!existing) {
+    runContextById.set(runId, { ...context });
+    return;
+  }
+  if (context.sessionKey && existing.sessionKey !== context.sessionKey) {
+    existing.sessionKey = context.sessionKey;
+  }
+}
+
+export function getAgentRunContext(runId: string) {
+  return runContextById.get(runId);
+}
+
+export function clearAgentRunContext(runId: string) {
+  runContextById.delete(runId);
+}
+
+export function resetAgentRunContextForTest() {
+  runContextById.clear();
+}
 
 export function emitAgentEvent(event: Omit<AgentEventPayload, "seq" | "ts">) {
   const nextSeq = (seqByRun.get(event.runId) ?? 0) + 1;