feat(server): stream replayable harness events

This commit is contained in:
Wenxin Duan
2026-09-07 16:43:14 +08:00
parent 5e28e74ca1
commit acc5e26b1c
9 changed files with 314 additions and 6 deletions
+19 -6
View File
@@ -1,6 +1,7 @@
import type { import type {
ApprovalBroker, ApprovalBroker,
AttachmentStore, AttachmentStore,
EventStore,
ModelConnectionStore, ModelConnectionStore,
ResolvedModelConnection, ResolvedModelConnection,
SessionRecord, SessionRecord,
@@ -13,9 +14,11 @@ import { installErrorHandler } from "./errors/error-handler.js";
import { registerApprovalRoutes } from "./routes/approvals.js"; import { registerApprovalRoutes } from "./routes/approvals.js";
import { registerAttachmentRoutes } from "./routes/attachments.js"; import { registerAttachmentRoutes } from "./routes/attachments.js";
import { registerHealthRoute } from "./routes/health.js"; import { registerHealthRoute } from "./routes/health.js";
import { registerEventRoute } from "./routes/events.js";
import { registerModelRoutes } from "./routes/models.js"; import { registerModelRoutes } from "./routes/models.js";
import { registerSessionRoutes } from "./routes/sessions.js"; import { registerSessionRoutes } from "./routes/sessions.js";
import { registerSkillsRoute } from "./routes/skills.js"; import { registerSkillsRoute } from "./routes/skills.js";
import { SseRegistry, type SseRegistryPort } from "./sse/sse-registry.js";
export type ApiSession = SessionRecord & { export type ApiSession = SessionRecord & {
startRun(input: { text: string; attachmentIds: string[] }): Promise<{ runId: string }>; startRun(input: { text: string; attachmentIds: string[] }): Promise<{ runId: string }>;
@@ -34,6 +37,8 @@ export type AppDependencies = {
modelStore: ModelConnectionStore; modelStore: ModelConnectionStore;
attachmentStore: Pick<AttachmentStore, "add" | "remove">; attachmentStore: Pick<AttachmentStore, "add" | "remove">;
approvalBroker: Pick<ApprovalBroker, "approve" | "deny">; approvalBroker: Pick<ApprovalBroker, "approve" | "deny">;
eventStore?: EventStore;
sseRegistry?: SseRegistryPort;
testModelConnection?: ( testModelConnection?: (
resolved: ResolvedModelConnection, resolved: ResolvedModelConnection,
signal: AbortSignal, signal: AbortSignal,
@@ -41,6 +46,8 @@ export type AppDependencies = {
}; };
export function buildApp(dependencies: AppDependencies) { export function buildApp(dependencies: AppDependencies) {
const sseRegistry = dependencies.sseRegistry ?? new SseRegistry();
const effectiveDependencies: AppDependencies = { ...dependencies, sseRegistry };
const app = Fastify({ const app = Fastify({
logger: false, logger: false,
bodyLimit: 6 * 1024 * 1024, 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 } }); void app.register(multipart, { limits: { files: 1, fields: 0, fileSize: 5 * 1024 * 1024 } });
installErrorHandler(app); installErrorHandler(app);
registerHealthRoute(app); registerHealthRoute(app);
registerSkillsRoute(app, dependencies); registerSkillsRoute(app, effectiveDependencies);
registerSessionRoutes(app, dependencies); registerSessionRoutes(app, effectiveDependencies);
registerModelRoutes(app, dependencies); registerModelRoutes(app, effectiveDependencies);
registerAttachmentRoutes(app, dependencies); registerAttachmentRoutes(app, effectiveDependencies);
registerApprovalRoutes(app, dependencies); registerApprovalRoutes(app, effectiveDependencies);
app.addHook("onClose", async () => dependencies.sessionRegistry.disposeAll()); if (dependencies.eventStore !== undefined) {
registerEventRoute(app, effectiveDependencies, dependencies.eventStore, sseRegistry);
}
app.addHook("onClose", async () => {
sseRegistry.closeAll();
await dependencies.sessionRegistry.disposeAll();
});
return app; return app;
} }
+1
View File
@@ -94,6 +94,7 @@ export async function buildDefaultApp(environment: NodeJS.ProcessEnv = process.e
modelStore, modelStore,
attachmentStore, attachmentStore,
approvalBroker, approvalBroker,
eventStore,
}); });
} }
+45
View File
@@ -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();
});
}
+1
View File
@@ -14,6 +14,7 @@ export function registerSessionRoutes(app: FastifyInstance, dependencies: AppDep
app.delete("/api/sessions/:sessionId", async (request) => { app.delete("/api/sessions/:sessionId", async (request) => {
const { sessionId } = SessionParams.parse(request.params); const { sessionId } = SessionParams.parse(request.params);
dependencies.sseRegistry?.closeForSession(sessionId);
const deleted = await dependencies.sessionRegistry.delete(sessionId); const deleted = await dependencies.sessionRegistry.delete(sessionId);
if (!deleted) throw new HarnessError("SESSION_NOT_FOUND", "Session not found"); if (!deleted) throw new HarnessError("SESSION_NOT_FOUND", "Session not found");
return { deleted: true, requestId: request.id }; return { deleted: true, requestId: request.id };
+99
View File
@@ -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<typeof setInterval> | 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();
}
}
+35
View File
@@ -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<string, Set<SseConnection>>();
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);
}
}
+26
View File
@@ -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"]);
});
});
+84
View File
@@ -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<string, unknown>) => { 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<string, unknown>) => { 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<string, unknown>) => { 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<string, unknown>) => { 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);
});
});
@@ -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 { #state(sessionId: string): SessionEvents {
const existing = this.#sessions.get(sessionId); const existing = this.#sessions.get(sessionId);
if (existing !== undefined) return existing; if (existing !== undefined) return existing;