diff --git a/apps/server/src/app.ts b/apps/server/src/app.ts index 6a20523..db01930 100644 --- a/apps/server/src/app.ts +++ b/apps/server/src/app.ts @@ -1,6 +1,7 @@ import type { ApprovalBroker, AttachmentStore, + EventStore, ModelConnectionStore, ResolvedModelConnection, SessionRecord, @@ -13,9 +14,11 @@ import { installErrorHandler } from "./errors/error-handler.js"; import { registerApprovalRoutes } from "./routes/approvals.js"; import { registerAttachmentRoutes } from "./routes/attachments.js"; import { registerHealthRoute } from "./routes/health.js"; +import { registerEventRoute } from "./routes/events.js"; import { registerModelRoutes } from "./routes/models.js"; import { registerSessionRoutes } from "./routes/sessions.js"; import { registerSkillsRoute } from "./routes/skills.js"; +import { SseRegistry, type SseRegistryPort } from "./sse/sse-registry.js"; export type ApiSession = SessionRecord & { startRun(input: { text: string; attachmentIds: string[] }): Promise<{ runId: string }>; @@ -34,6 +37,8 @@ export type AppDependencies = { modelStore: ModelConnectionStore; attachmentStore: Pick; approvalBroker: Pick; + eventStore?: EventStore; + sseRegistry?: SseRegistryPort; testModelConnection?: ( resolved: ResolvedModelConnection, signal: AbortSignal, @@ -41,6 +46,8 @@ export type AppDependencies = { }; export function buildApp(dependencies: AppDependencies) { + const sseRegistry = dependencies.sseRegistry ?? new SseRegistry(); + const effectiveDependencies: AppDependencies = { ...dependencies, sseRegistry }; const app = Fastify({ logger: false, bodyLimit: 6 * 1024 * 1024, @@ -49,11 +56,17 @@ export function buildApp(dependencies: AppDependencies) { void app.register(multipart, { limits: { files: 1, fields: 0, fileSize: 5 * 1024 * 1024 } }); installErrorHandler(app); registerHealthRoute(app); - registerSkillsRoute(app, dependencies); - registerSessionRoutes(app, dependencies); - registerModelRoutes(app, dependencies); - registerAttachmentRoutes(app, dependencies); - registerApprovalRoutes(app, dependencies); - app.addHook("onClose", async () => dependencies.sessionRegistry.disposeAll()); + registerSkillsRoute(app, effectiveDependencies); + registerSessionRoutes(app, effectiveDependencies); + registerModelRoutes(app, effectiveDependencies); + registerAttachmentRoutes(app, effectiveDependencies); + registerApprovalRoutes(app, effectiveDependencies); + if (dependencies.eventStore !== undefined) { + registerEventRoute(app, effectiveDependencies, dependencies.eventStore, sseRegistry); + } + app.addHook("onClose", async () => { + sseRegistry.closeAll(); + await dependencies.sessionRegistry.disposeAll(); + }); return app; } diff --git a/apps/server/src/index.ts b/apps/server/src/index.ts index d1b15fb..1744ee4 100644 --- a/apps/server/src/index.ts +++ b/apps/server/src/index.ts @@ -94,6 +94,7 @@ export async function buildDefaultApp(environment: NodeJS.ProcessEnv = process.e modelStore, attachmentStore, approvalBroker, + eventStore, }); } diff --git a/apps/server/src/routes/events.ts b/apps/server/src/routes/events.ts new file mode 100644 index 0000000..f381c3c --- /dev/null +++ b/apps/server/src/routes/events.ts @@ -0,0 +1,45 @@ +import { HarnessError, type EventStore } from "@agent-studio/harness"; +import type { FastifyInstance } from "fastify"; +import { z } from "zod"; +import type { AppDependencies } from "../app.js"; +import { SseConnection } from "../sse/sse-connection.js"; +import type { SseRegistryPort } from "../sse/sse-registry.js"; +import { requireSession } from "./sessions.js"; + +const Params = z.object({ sessionId: z.string().regex(/^ses_[a-zA-Z0-9_-]+$/) }).strict(); +const Query = z.object({ afterSequence: z.coerce.number().int().nonnegative().optional() }).strict(); + +export function registerEventRoute( + app: FastifyInstance, + dependencies: AppDependencies, + eventStore: EventStore, + registry: SseRegistryPort, +): void { + app.get("/api/sessions/:sessionId/events", (request, reply) => { + const { sessionId } = Params.parse(request.params); + requireSession(dependencies, sessionId); + const query = Query.parse(request.query); + const header = request.headers["last-event-id"]; + const eventId = Array.isArray(header) ? header[0] : header; + if (eventId !== undefined && eventId.trim() === "") { + throw new HarnessError("INVALID_EVENT_CURSOR", "Last-Event-ID is invalid"); + } + reply.hijack(); + reply.raw.writeHead(200, { + "Content-Type": "text/event-stream; charset=utf-8", + "Cache-Control": "no-cache, no-transform", + Connection: "keep-alive", + "X-Accel-Buffering": "no", + }); + const connection = new SseConnection({ + sessionId, + eventStore, + response: reply.raw, + cursor: eventId === undefined ? { sequence: query.afterSequence ?? 0 } : { eventId }, + onClose: () => registry.remove?.(sessionId, connection), + }); + registry.add(sessionId, connection); + request.raw.once("close", () => connection.close()); + connection.start(); + }); +} diff --git a/apps/server/src/routes/sessions.ts b/apps/server/src/routes/sessions.ts index 3cf56cb..7d88752 100644 --- a/apps/server/src/routes/sessions.ts +++ b/apps/server/src/routes/sessions.ts @@ -14,6 +14,7 @@ export function registerSessionRoutes(app: FastifyInstance, dependencies: AppDep app.delete("/api/sessions/:sessionId", async (request) => { const { sessionId } = SessionParams.parse(request.params); + dependencies.sseRegistry?.closeForSession(sessionId); const deleted = await dependencies.sessionRegistry.delete(sessionId); if (!deleted) throw new HarnessError("SESSION_NOT_FOUND", "Session not found"); return { deleted: true, requestId: request.id }; diff --git a/apps/server/src/sse/sse-connection.ts b/apps/server/src/sse/sse-connection.ts new file mode 100644 index 0000000..ba6fbe3 --- /dev/null +++ b/apps/server/src/sse/sse-connection.ts @@ -0,0 +1,99 @@ +import type { EventCursor, EventStore } from "@agent-studio/harness"; +import type { HarnessEvent } from "@agent-studio/shared"; + +export type SseWritable = { + write(chunk: string): boolean; + end(): void; + once(event: "close" | "drain", listener: () => void): unknown; + removeListener(event: "close" | "drain", listener: () => void): unknown; +}; + +export type SseConnectionOptions = { + sessionId: string; + eventStore: EventStore; + response: SseWritable; + cursor?: EventCursor; + heartbeatMs?: number; + maxQueuedBytes?: number; + onClose?: () => void; +}; + +export class SseConnection { + readonly #sessionId: string; + readonly #eventStore: EventStore; + readonly #response: SseWritable; + readonly #cursor: EventCursor | undefined; + readonly #heartbeatMs: number; + readonly #maxQueuedBytes: number; + readonly #onClose: (() => void) | undefined; + #unsubscribe: (() => void) | undefined; + #heartbeat: ReturnType | undefined; + #closed = false; + #started = false; + #backpressuredBytes = 0; + + public constructor(options: SseConnectionOptions) { + this.#sessionId = options.sessionId; + this.#eventStore = options.eventStore; + this.#response = options.response; + this.#cursor = options.cursor; + this.#heartbeatMs = options.heartbeatMs ?? 15_000; + this.#maxQueuedBytes = options.maxQueuedBytes ?? 256 * 1024; + this.#onClose = options.onClose; + } + + public start(): void { + if (this.#started || this.#closed) return; + this.#started = true; + const buffered: HarnessEvent[] = []; + let replaying = true; + this.#unsubscribe = this.#eventStore.subscribe(this.#sessionId, (event) => { + if (replaying) buffered.push(event); + else this.#writeEvent(event); + }); + const replay = this.#eventStore.listAfter(this.#sessionId, this.#cursor); + if (replay.gap) { + this.#writeFrame("event: stream.reset\ndata: {\"reason\":\"history_gap\"}\n\n"); + } + for (const event of replay.events) this.#writeEvent(event); + replaying = false; + const latestReplayed = replay.events.at(-1)?.sequence ?? this.#cursor?.sequence ?? 0; + for (const event of buffered) { + if (event.sequence > latestReplayed) this.#writeEvent(event); + } + this.#heartbeat = setInterval(() => this.#writeFrame(": heartbeat\n\n"), this.#heartbeatMs); + this.#response.once("close", this.#handleClose); + } + + public close(): void { + if (this.#closed) return; + this.#closed = true; + this.#unsubscribe?.(); + this.#unsubscribe = undefined; + if (this.#heartbeat !== undefined) clearInterval(this.#heartbeat); + this.#heartbeat = undefined; + this.#response.removeListener("close", this.#handleClose); + this.#response.removeListener("drain", this.#handleDrain); + this.#response.end(); + this.#onClose?.(); + } + + readonly #handleClose = (): void => this.close(); + readonly #handleDrain = (): void => { + this.#backpressuredBytes = 0; + }; + + #writeEvent(event: HarnessEvent): void { + this.#writeFrame(`id: ${event.eventId}\nevent: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`); + } + + #writeFrame(frame: string): void { + if (this.#closed) return; + const writable = this.#response.write(frame); + if (writable) return; + this.#backpressuredBytes += Buffer.byteLength(frame, "utf8"); + this.#response.removeListener("drain", this.#handleDrain); + this.#response.once("drain", this.#handleDrain); + if (this.#backpressuredBytes > this.#maxQueuedBytes) this.close(); + } +} diff --git a/apps/server/src/sse/sse-registry.ts b/apps/server/src/sse/sse-registry.ts new file mode 100644 index 0000000..33c47ea --- /dev/null +++ b/apps/server/src/sse/sse-registry.ts @@ -0,0 +1,35 @@ +import type { SseConnection } from "./sse-connection.js"; + +export interface SseRegistryPort { + add(sessionId: string, connection: SseConnection): void; + remove?(sessionId: string, connection: SseConnection): void; + closeForSession(sessionId: string): void; + closeAll(): void; +} + +export class SseRegistry implements SseRegistryPort { + readonly #connections = new Map>(); + + public add(sessionId: string, connection: SseConnection): void { + const existing = this.#connections.get(sessionId); + if (existing !== undefined) existing.add(connection); + else this.#connections.set(sessionId, new Set([connection])); + } + + public remove(sessionId: string, connection: SseConnection): void { + const existing = this.#connections.get(sessionId); + existing?.delete(connection); + if (existing?.size === 0) this.#connections.delete(sessionId); + } + + public closeForSession(sessionId: string): void { + const connections = this.#connections.get(sessionId); + if (connections === undefined) return; + this.#connections.delete(sessionId); + for (const connection of [...connections]) connection.close(); + } + + public closeAll(): void { + for (const sessionId of [...this.#connections.keys()]) this.closeForSession(sessionId); + } +} diff --git a/apps/server/tests/shutdown.test.ts b/apps/server/tests/shutdown.test.ts new file mode 100644 index 0000000..a28aed0 --- /dev/null +++ b/apps/server/tests/shutdown.test.ts @@ -0,0 +1,26 @@ +import { describe, expect, it, vi } from "vitest"; +import { buildApp } from "../src/app.js"; + +describe("shutdown lifecycle", () => { + it("closes SSE before disposing all sessions", async () => { + const order: string[] = []; + const app = buildApp({ + sessionRegistry: { + create: vi.fn(), get: vi.fn(), delete: vi.fn(), + disposeAll: () => { order.push("sessions"); return Promise.resolve(); }, + }, + skillCatalog: { list: () => [] }, + modelStore: {} as never, + attachmentStore: {} as never, + approvalBroker: {} as never, + sseRegistry: { + add: vi.fn(), closeForSession: vi.fn(), + closeAll: () => { order.push("sse"); }, + }, + }); + + await app.close(); + + expect(order).toEqual(["sse", "sessions"]); + }); +}); diff --git a/apps/server/tests/sse.test.ts b/apps/server/tests/sse.test.ts new file mode 100644 index 0000000..f2b0670 --- /dev/null +++ b/apps/server/tests/sse.test.ts @@ -0,0 +1,84 @@ +import { EventEmitter } from "node:events"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { InMemoryEventStore } from "@agent-studio/harness"; +import * as sseModule from "../src/sse/sse-connection.js"; + +class FakeResponse extends EventEmitter { + chunks: string[] = []; + ended = false; + shouldBackpressure = false; + + write(chunk: string): boolean { + this.chunks.push(chunk); + return !this.shouldBackpressure; + } + + end(): void { + this.ended = true; + this.emit("close"); + } +} + +afterEach(() => vi.useRealTimers()); + +describe("SseConnection", () => { + it("replays after Last-Event-ID and gives it precedence over sequence", () => { + const store = new InMemoryEventStore(); + const first = store.append("ses_1", { type: "session.started", payload: { cwd: "." } }); + const second = store.append("ses_1", { type: "run.started", runId: "run_1", payload: { modelId: "demo" } }); + const third = store.append("ses_1", { type: "assistant.delta", runId: "run_1", payload: { messageId: "msg_1", delta: "hi" } }); + const Response = Reflect.get(sseModule, "SseConnection") as unknown as new (options: Record) => { start(): void; close(): void }; + expect(Response, "SseConnection must be exported").toBeDefined(); + const response = new FakeResponse(); + const connection = new Response({ sessionId: "ses_1", eventStore: store, response, cursor: { eventId: second.eventId, sequence: first.sequence } }); + + connection.start(); + + const output = response.chunks.join(""); + expect(output).toContain(`id: ${third.eventId}`); + expect(output).not.toContain(`id: ${first.eventId}`); + connection.close(); + }); + + it("emits reset for a retained-history gap and streams live events", () => { + const store = new InMemoryEventStore(1); + store.append("ses_1", { type: "session.started", payload: { cwd: "." } }); + store.append("ses_1", { type: "run.started", runId: "run_1", payload: { modelId: "demo" } }); + const response = new FakeResponse(); + const Response = Reflect.get(sseModule, "SseConnection") as unknown as new (options: Record) => { start(): void; close(): void }; + const connection = new Response({ sessionId: "ses_1", eventStore: store, response, cursor: { sequence: 0 } }); + connection.start(); + store.append("ses_1", { type: "assistant.delta", runId: "run_1", payload: { messageId: "msg_1", delta: "live" } }); + + const output = response.chunks.join(""); + expect(output).toContain("event: stream.reset"); + expect(output).toContain('"delta":"live"'); + connection.close(); + }); + + it("clears its subscriber and heartbeat when closed", () => { + vi.useFakeTimers(); + const store = new InMemoryEventStore() as InMemoryEventStore & { listenerCount(sessionId: string): number }; + const response = new FakeResponse(); + const Response = Reflect.get(sseModule, "SseConnection") as unknown as new (options: Record) => { start(): void; close(): void }; + const connection = new Response({ sessionId: "ses_1", eventStore: store, response, heartbeatMs: 15_000 }); + connection.start(); + expect(store.listenerCount("ses_1")).toBe(1); + + connection.close(); + + expect(store.listenerCount("ses_1")).toBe(0); + expect(vi.getTimerCount()).toBe(0); + }); + + it("closes a persistently backpressured client at the configured bound", () => { + const store = new InMemoryEventStore(); + const response = new FakeResponse(); + response.shouldBackpressure = true; + const Response = Reflect.get(sseModule, "SseConnection") as unknown as new (options: Record) => { start(): void }; + const connection = new Response({ sessionId: "ses_1", eventStore: store, response, maxQueuedBytes: 20 }); + connection.start(); + store.append("ses_1", { type: "assistant.delta", runId: "run_1", payload: { messageId: "msg_1", delta: "a long live chunk" } }); + expect(response.ended).toBe(true); + }); +}); diff --git a/packages/harness/src/events/in-memory-event-store.ts b/packages/harness/src/events/in-memory-event-store.ts index f65753f..875f996 100644 --- a/packages/harness/src/events/in-memory-event-store.ts +++ b/packages/harness/src/events/in-memory-event-store.ts @@ -87,6 +87,10 @@ export class InMemoryEventStore implements EventStore { }; } + public listenerCount(sessionId: string): number { + return this.#sessions.get(sessionId)?.listeners.size ?? 0; + } + #state(sessionId: string): SessionEvents { const existing = this.#sessions.get(sessionId); if (existing !== undefined) return existing;