feat(server): stream replayable harness events
This commit is contained in:
+19
-6
@@ -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<AttachmentStore, "add" | "remove">;
|
||||
approvalBroker: Pick<ApprovalBroker, "approve" | "deny">;
|
||||
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;
|
||||
}
|
||||
|
||||
@@ -94,6 +94,7 @@ export async function buildDefaultApp(environment: NodeJS.ProcessEnv = process.e
|
||||
modelStore,
|
||||
attachmentStore,
|
||||
approvalBroker,
|
||||
eventStore,
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
}
|
||||
@@ -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 };
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user