Add generic session broker core
This commit is contained in:
udo committed
1 parent
83a1049fe2
commit
cf1f9a998f
2 files changed
+805
No files matched your search
@@ -0,0 +1,465 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
|
||||
export type RuntimeState = "live" | "deleting" | "disposed";
|
||||
|
||||
export interface SessionPromptOptions {
|
||||
preflightResult?: (accepted: boolean) => void;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
export interface BrokerSession {
|
||||
readonly sessionId: string;
|
||||
readonly isIdle: boolean;
|
||||
subscribe(listener: (event: unknown) => void): () => void;
|
||||
prompt(text: string, options?: SessionPromptOptions): Promise<void>;
|
||||
steer(text: string): Promise<void>;
|
||||
followUp(text: string): Promise<void>;
|
||||
abort(): Promise<void>;
|
||||
dispose(): void;
|
||||
}
|
||||
|
||||
export interface BrokerFactory<C extends object = Record<string, unknown>, S extends BrokerSession = BrokerSession> {
|
||||
create(key: string, context: Readonly<C>): Promise<S>;
|
||||
}
|
||||
|
||||
export interface BrokerSerializer<C extends object = Record<string, unknown>, S extends BrokerSession = BrokerSession> {
|
||||
snapshot(runtime: BrokerRuntime<C, S>): unknown;
|
||||
}
|
||||
|
||||
export interface SocketSink {
|
||||
isOpen(): boolean;
|
||||
send(message: unknown): void;
|
||||
}
|
||||
|
||||
export interface BrokerClock {
|
||||
now(): number;
|
||||
}
|
||||
|
||||
export interface PromptRequest {
|
||||
requestId: string;
|
||||
text: string;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
export interface ChoiceRequest {
|
||||
choiceId: string;
|
||||
payload: unknown;
|
||||
}
|
||||
|
||||
export interface BrokerMessageHooks<C extends object = Record<string, unknown>, S extends BrokerSession = BrokerSession> {
|
||||
ready(runtime: BrokerRuntime<C, S>): unknown;
|
||||
stateSync(runtime: BrokerRuntime<C, S>, snapshot: unknown): unknown;
|
||||
event(event: unknown): unknown;
|
||||
choiceRequest(request: ChoiceRequest): unknown;
|
||||
promptAccepted(requestId: string): unknown;
|
||||
promptRejected(requestId: string): unknown;
|
||||
}
|
||||
|
||||
export interface BrokerOptions<C extends object = Record<string, unknown>, S extends BrokerSession = BrokerSession> {
|
||||
factory: BrokerFactory<C, S>;
|
||||
serializer: BrokerSerializer<C, S>;
|
||||
messages: BrokerMessageHooks<C, S>;
|
||||
clock?: BrokerClock;
|
||||
idleTtlMs?: number;
|
||||
activeMaxAgeMs?: number;
|
||||
maxRuntimes?: number;
|
||||
sweepAbortGraceMs?: number;
|
||||
promptStarter?: (session: S, request: PromptRequest, preflightResult: (accepted: boolean) => void) => void | Promise<void>;
|
||||
onDelete: (key: string, runtime: BrokerRuntime<C, S>) => Promise<void>;
|
||||
onError: (error: unknown) => void;
|
||||
}
|
||||
|
||||
interface PromptRecord {
|
||||
fingerprint: string;
|
||||
state: "starting" | "accepted" | "completed" | "rejected" | "failed";
|
||||
acceptance: Promise<void>;
|
||||
accept: () => void;
|
||||
reject: (error: Error) => void;
|
||||
}
|
||||
|
||||
interface Attachment {
|
||||
readonly token: symbol;
|
||||
readonly generation: number;
|
||||
readonly sink: SocketSink;
|
||||
syncing: boolean;
|
||||
attached: boolean;
|
||||
}
|
||||
|
||||
interface PendingChoice {
|
||||
request: ChoiceRequest;
|
||||
resolve: (value: unknown) => void;
|
||||
reject: (error: Error) => void;
|
||||
}
|
||||
|
||||
const noop = (): void => {};
|
||||
|
||||
export class BrokerRuntime<C extends object = Record<string, unknown>, S extends BrokerSession = BrokerSession> {
|
||||
readonly runtimeId = randomUUID();
|
||||
readonly createdAt: number;
|
||||
readonly context: Readonly<C>;
|
||||
state: RuntimeState = "live";
|
||||
private readonly attachments = new Map<symbol, Attachment>();
|
||||
private readonly prompts = new Map<string, PromptRecord>();
|
||||
private readonly choices = new Map<string, PendingChoice>();
|
||||
private unsubscribe: (() => void) | undefined;
|
||||
private invocationTail: Promise<void> = Promise.resolve();
|
||||
private generation = 0;
|
||||
private lastActivity: number;
|
||||
private ageAbortRequested = false;
|
||||
private ageAbortDeadline = 0;
|
||||
private abortRequested = false;
|
||||
private disposed = false;
|
||||
private queuedInvocations = 0;
|
||||
|
||||
constructor(
|
||||
readonly key: string,
|
||||
readonly session: S,
|
||||
context: Readonly<Record<string, unknown>>,
|
||||
private readonly serializer: BrokerSerializer<C, S>,
|
||||
private readonly clock: BrokerClock,
|
||||
private readonly onError: (error: unknown) => void,
|
||||
private readonly onDelete: (key: string, runtime: BrokerRuntime<C, S>) => Promise<void>,
|
||||
private readonly promptStarter: ((session: S, request: PromptRequest, preflightResult: (accepted: boolean) => void) => void | Promise<void>) | undefined,
|
||||
private readonly messages: BrokerMessageHooks<C, S>,
|
||||
) {
|
||||
this.context = Object.freeze({ ...(context as object) }) as Readonly<C>;
|
||||
this.createdAt = clock.now();
|
||||
this.lastActivity = this.createdAt;
|
||||
this.unsubscribe = session.subscribe((event) => this.broadcast(event));
|
||||
}
|
||||
|
||||
get attachmentCount(): number { return this.attachments.size; }
|
||||
get hasPendingChoices(): boolean { return this.choices.size > 0; }
|
||||
get hasQueuedInvocation(): boolean { return this.queuedInvocations > 0; }
|
||||
get lastUsedAt(): number { return this.lastActivity; }
|
||||
get isDetached(): boolean { return this.attachments.size === 0; }
|
||||
get isActive(): boolean { return !this.session.isIdle || this.hasPendingChoices || this.hasQueuedInvocation || this.promptsInFlight; }
|
||||
private get promptsInFlight(): boolean { return [...this.prompts.values()].some(({ state }) => state === "starting" || state === "accepted"); }
|
||||
|
||||
attach(sink: SocketSink): AttachmentHandle<C, S> {
|
||||
if (this.state !== "live") throw new Error(`Runtime ${this.key} is ${this.state}`);
|
||||
const attachment: Attachment = { token: Symbol("attachment"), generation: ++this.generation, sink, syncing: true, attached: true };
|
||||
this.attachments.set(attachment.token, attachment);
|
||||
try {
|
||||
if (!sink.isOpen()) throw new Error("Socket is closed");
|
||||
this.send(attachment, this.message(() => this.messages?.ready(this)));
|
||||
this.send(attachment, this.message(() => this.messages?.stateSync(this, this.serializer.snapshot(this))));
|
||||
for (const choice of this.choices.values()) this.send(attachment, this.message(() => this.messages?.choiceRequest(choice.request)));
|
||||
if (!attachment.attached) throw new Error("Socket closed during attach");
|
||||
attachment.syncing = false;
|
||||
this.lastActivity = this.clock.now();
|
||||
return new AttachmentHandle(this, attachment);
|
||||
} catch (error) {
|
||||
this.detachAttachment(attachment);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
detach(handle: AttachmentHandle<C, S>): void { this.detachAttachment(handle.attachment); }
|
||||
|
||||
isAttachmentValid(handle: AttachmentHandle<C, S>, generation = handle.generation): boolean {
|
||||
return this.state === "live" && handle.attachment.attached && handle.generation === generation && this.attachments.get(handle.token) === handle.attachment;
|
||||
}
|
||||
|
||||
private detachAttachment(attachment: Attachment): void {
|
||||
if (!attachment.attached || this.attachments.get(attachment.token) !== attachment) return;
|
||||
attachment.attached = false;
|
||||
this.attachments.delete(attachment.token);
|
||||
this.lastActivity = this.clock.now();
|
||||
}
|
||||
|
||||
private message(create: () => unknown): unknown {
|
||||
try { return create(); } catch (error) { this.onError(error); return undefined; }
|
||||
}
|
||||
|
||||
private send(attachment: Attachment, message: unknown): void {
|
||||
if (message === undefined || !attachment.attached) return;
|
||||
if (!attachment.sink.isOpen()) { this.detachAttachment(attachment); return; }
|
||||
try { attachment.sink.send(message); } catch (error) { this.onError(error); this.detachAttachment(attachment); }
|
||||
}
|
||||
|
||||
private broadcast(event: unknown): void {
|
||||
if ((event as { type?: string })?.type === "agent_settled") this.lastActivity = this.clock.now();
|
||||
const message = this.message(() => this.messages?.event(event));
|
||||
if (message === undefined) return;
|
||||
for (const attachment of this.attachments.values()) {
|
||||
if (!attachment.syncing) this.send(attachment, message);
|
||||
}
|
||||
}
|
||||
|
||||
private enqueue(handle: AttachmentHandle<C, S>, operation: () => Promise<void>): Promise<void> {
|
||||
if (!this.isAttachmentValid(handle)) return Promise.reject(new Error("Attachment is no longer valid"));
|
||||
this.queuedInvocations++;
|
||||
const run = this.invocationTail.then(() => {
|
||||
if (!this.isAttachmentValid(handle)) throw new Error("Attachment is no longer valid");
|
||||
this.lastActivity = this.clock.now();
|
||||
return operation();
|
||||
});
|
||||
this.invocationTail = run.catch(noop).finally(() => { this.queuedInvocations--; });
|
||||
return run;
|
||||
}
|
||||
|
||||
invoke(handle: AttachmentHandle<C, S>, operation: () => Promise<void>): Promise<void> { return this.enqueue(handle, operation); }
|
||||
steer(handle: AttachmentHandle<C, S>, text: string): Promise<void> { return this.enqueue(handle, () => this.session.steer(text)); }
|
||||
followUp(handle: AttachmentHandle<C, S>, text: string): Promise<void> { return this.enqueue(handle, () => this.session.followUp(text)); }
|
||||
abort(handle: AttachmentHandle<C, S>): Promise<void> { return this.enqueue(handle, () => this.abortOnce()); }
|
||||
prompt(handle: AttachmentHandle<C, S>, request: PromptRequest, requestFingerprint: string): Promise<void> {
|
||||
return this.enqueue(handle, () => this.startPrompt(request, requestFingerprint, handle));
|
||||
}
|
||||
|
||||
private startPrompt(request: PromptRequest, fingerprint: string, handle: AttachmentHandle<C, S>): Promise<void> {
|
||||
const existing = this.prompts.get(request.requestId);
|
||||
if (existing) {
|
||||
if (existing.fingerprint !== fingerprint) return Promise.reject(new Error(`Prompt request ${request.requestId} has a different payload`));
|
||||
this.send(handle.attachment, this.message(() => existing.state === "rejected" ? this.messages?.promptRejected(request.requestId) : this.messages?.promptAccepted(request.requestId)));
|
||||
return existing.acceptance;
|
||||
}
|
||||
let accept!: () => void;
|
||||
let reject!: (error: Error) => void;
|
||||
const acceptance = new Promise<void>((resolve, rejectPromise) => { accept = resolve; reject = rejectPromise; });
|
||||
const record: PromptRecord = { fingerprint, state: "starting", acceptance, accept, reject };
|
||||
this.prompts.set(request.requestId, record);
|
||||
let preflightCalled = false;
|
||||
const preflightResult = (accepted: boolean): void => {
|
||||
if (preflightCalled || record.state !== "starting") return;
|
||||
preflightCalled = true;
|
||||
if (this.state !== "live") {
|
||||
this.rejectPrompt(record, new Error(`Runtime ${this.key} is ${this.state}`));
|
||||
return;
|
||||
}
|
||||
if (accepted) {
|
||||
record.state = "accepted";
|
||||
record.accept();
|
||||
this.send(handle.attachment, this.message(() => this.messages?.promptAccepted(request.requestId)));
|
||||
} else {
|
||||
this.rejectPrompt(record, new Error(`Prompt request ${request.requestId} was rejected`));
|
||||
this.send(handle.attachment, this.message(() => this.messages?.promptRejected(request.requestId)));
|
||||
}
|
||||
};
|
||||
try {
|
||||
const started = this.promptStarter
|
||||
? this.promptStarter(this.session, request, preflightResult)
|
||||
: this.session.prompt(request.text, { ...request, preflightResult });
|
||||
void Promise.resolve(started).then(() => {
|
||||
if (!preflightCalled) preflightResult(true);
|
||||
if (record.state === "accepted") record.state = "completed";
|
||||
this.trimPrompts();
|
||||
}, (error: unknown) => {
|
||||
if (!preflightCalled) {
|
||||
this.rejectPrompt(record, error instanceof Error ? error : new Error(String(error)));
|
||||
this.send(handle.attachment, this.message(() => this.messages.promptRejected(request.requestId)));
|
||||
} else if (record.state === "accepted") record.state = "failed";
|
||||
this.onError(error);
|
||||
this.trimPrompts();
|
||||
});
|
||||
} catch (error) {
|
||||
this.rejectPrompt(record, error instanceof Error ? error : new Error(String(error)));
|
||||
this.send(handle.attachment, this.message(() => this.messages.promptRejected(request.requestId)));
|
||||
this.onError(error);
|
||||
}
|
||||
return acceptance;
|
||||
}
|
||||
|
||||
private rejectPrompt(record: PromptRecord, error: Error): void {
|
||||
if (record.state !== "starting") return;
|
||||
record.state = "rejected";
|
||||
record.reject(error);
|
||||
this.trimPrompts();
|
||||
}
|
||||
|
||||
private trimPrompts(): void {
|
||||
const settled = [...this.prompts].filter(([, record]) => record.state === "completed" || record.state === "rejected" || record.state === "failed");
|
||||
while (settled.length > 1000) this.prompts.delete(settled.shift()![0]);
|
||||
}
|
||||
|
||||
requestChoice(request: ChoiceRequest): Promise<unknown> {
|
||||
if (this.state !== "live") return Promise.reject(new Error("Runtime is not live"));
|
||||
if (this.choices.has(request.choiceId)) return Promise.reject(new Error(`Choice ${request.choiceId} already exists`));
|
||||
const promise = new Promise<unknown>((resolve, reject) => this.choices.set(request.choiceId, { request, resolve, reject }));
|
||||
const message = this.message(() => this.messages?.choiceRequest(request));
|
||||
if (message !== undefined) {
|
||||
for (const attachment of this.attachments.values()) {
|
||||
if (!attachment.syncing) this.send(attachment, message);
|
||||
}
|
||||
}
|
||||
return promise;
|
||||
}
|
||||
|
||||
resolveChoice(handle: AttachmentHandle<C, S>, choiceId: string, value: unknown): Promise<void> {
|
||||
return this.enqueue(handle, async () => {
|
||||
const choice = this.choices.get(choiceId);
|
||||
if (!choice) throw new Error(`Choice ${choiceId} is not pending`);
|
||||
this.choices.delete(choiceId);
|
||||
choice.resolve(value);
|
||||
});
|
||||
}
|
||||
|
||||
rejectChoice(choiceId: string, reason = "Choice cancelled"): void {
|
||||
const choice = this.choices.get(choiceId);
|
||||
if (!choice) return;
|
||||
this.choices.delete(choiceId);
|
||||
choice.reject(new Error(reason));
|
||||
}
|
||||
|
||||
async delete(): Promise<void> {
|
||||
if (this.state !== "live") throw new Error(`Runtime ${this.key} is ${this.state}`);
|
||||
if (this.isActive || this.attachmentCount > 1) throw new Error("Runtime is active or has multiple attachments");
|
||||
this.state = "deleting";
|
||||
await this.dispose();
|
||||
await this.onDelete(this.key, this);
|
||||
}
|
||||
|
||||
async requestAgeAbort(graceUntil: number): Promise<void> {
|
||||
if (!this.isDetached || this.state !== "live") return;
|
||||
if (!this.ageAbortRequested) {
|
||||
if (this.session.isIdle) return;
|
||||
this.ageAbortRequested = true;
|
||||
this.ageAbortDeadline = graceUntil;
|
||||
void this.abortOnce().catch((error) => this.onError(error));
|
||||
}
|
||||
if (this.clock.now() >= this.ageAbortDeadline && !this.session.isIdle) await this.dispose();
|
||||
}
|
||||
|
||||
async abortOnce(): Promise<void> {
|
||||
if (this.abortRequested) return;
|
||||
this.abortRequested = true;
|
||||
await this.session.abort();
|
||||
}
|
||||
|
||||
async shutdown(): Promise<void> {
|
||||
if (this.isActive) void this.abortOnce().catch((error) => this.onError(error));
|
||||
await this.dispose();
|
||||
}
|
||||
|
||||
async dispose(): Promise<void> {
|
||||
if (this.disposed) return;
|
||||
this.disposed = true;
|
||||
this.state = "disposed";
|
||||
this.unsubscribe?.();
|
||||
this.unsubscribe = undefined;
|
||||
for (const choice of this.choices.values()) choice.reject(new Error("Runtime disposed"));
|
||||
this.choices.clear();
|
||||
for (const record of this.prompts.values()) {
|
||||
if (record.state === "starting") this.rejectPrompt(record, new Error("Runtime disposed"));
|
||||
}
|
||||
for (const attachment of this.attachments.values()) attachment.attached = false;
|
||||
this.attachments.clear();
|
||||
try { this.session.dispose(); } catch (error) { this.onError(error); }
|
||||
}
|
||||
}
|
||||
|
||||
export class AttachmentHandle<C extends object = Record<string, unknown>, S extends BrokerSession = BrokerSession> {
|
||||
constructor(readonly runtime: BrokerRuntime<C, S>, readonly attachment: Attachment) {}
|
||||
detach(): void { this.runtime.detach(this); }
|
||||
get generation(): number { return this.attachment.generation; }
|
||||
get token(): symbol { return this.attachment.token; }
|
||||
}
|
||||
|
||||
export class SessionBroker<C extends object = Record<string, unknown>, S extends BrokerSession = BrokerSession> {
|
||||
private readonly runtimes = new Map<string, BrokerRuntime<C, S>>();
|
||||
private readonly creating = new Map<string, Promise<BrokerRuntime<C, S>>>();
|
||||
private readonly deleting = new Map<string, Promise<void>>();
|
||||
private reservations = 0;
|
||||
private readonly clock: BrokerClock;
|
||||
private shuttingDown = false;
|
||||
private shutdownPromise: Promise<void> | undefined;
|
||||
|
||||
constructor(private readonly options: BrokerOptions<C, S>) { this.clock = options.clock ?? { now: () => Date.now() }; }
|
||||
get(key: string): BrokerRuntime<C, S> | undefined {
|
||||
key = String(key);
|
||||
const runtime = this.runtimes.get(key);
|
||||
if (runtime?.state === "disposed") this.runtimes.delete(key);
|
||||
return runtime?.state === "disposed" ? undefined : runtime;
|
||||
}
|
||||
|
||||
async getOrCreate(key: string, context: C): Promise<BrokerRuntime<C, S>> {
|
||||
key = String(key);
|
||||
if (this.shuttingDown) throw new Error("Broker is shutting down");
|
||||
const deleting = this.deleting.get(key);
|
||||
if (deleting) { await deleting; return this.getOrCreate(key, context); }
|
||||
const live = this.get(key);
|
||||
if (live) {
|
||||
if (live.state === "live") return live;
|
||||
throw new Error(`Runtime ${key} is ${live.state}`);
|
||||
}
|
||||
const existing = this.creating.get(key);
|
||||
if (existing) return existing;
|
||||
const promise = this.createWithCapacity(key, context);
|
||||
this.creating.set(key, promise);
|
||||
try { return await promise; } finally { if (this.creating.get(key) === promise) this.creating.delete(key); }
|
||||
}
|
||||
|
||||
private async createWithCapacity(key: string, context: C): Promise<BrokerRuntime<C, S>> {
|
||||
this.reserveCapacity(key);
|
||||
try { return await this.create(key, context); } finally { this.reservations--; }
|
||||
}
|
||||
|
||||
private async create(key: string, context: C): Promise<BrokerRuntime<C, S>> {
|
||||
const immutableContext = Object.freeze({ ...context }) as Readonly<C>;
|
||||
const session = await this.options.factory.create(key, immutableContext);
|
||||
const runtime = new BrokerRuntime<C, S>(key, session, immutableContext, this.options.serializer, this.clock, this.options.onError, this.options.onDelete, this.options.promptStarter, this.options.messages);
|
||||
if (this.shuttingDown) { await runtime.dispose(); throw new Error("Broker shut down during runtime creation"); }
|
||||
this.runtimes.set(key, runtime);
|
||||
return runtime;
|
||||
}
|
||||
|
||||
async sweep(): Promise<void> {
|
||||
const now = this.clock.now();
|
||||
for (const [key, runtime] of this.runtimes) {
|
||||
if (!runtime.isDetached || runtime.state !== "live") continue;
|
||||
if (!runtime.isActive && now - runtime.lastUsedAt >= (this.options.idleTtlMs ?? 15 * 60_000)) {
|
||||
await runtime.dispose();
|
||||
if (this.runtimes.get(key) === runtime) this.runtimes.delete(key);
|
||||
} else if (runtime.isActive && now - runtime.createdAt >= (this.options.activeMaxAgeMs ?? 24 * 60 * 60_000)) {
|
||||
await runtime.requestAgeAbort(now + (this.options.sweepAbortGraceMs ?? 30_000));
|
||||
if (runtime.state !== "live" && this.runtimes.get(key) === runtime) this.runtimes.delete(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async delete(key: string): Promise<void> {
|
||||
key = String(key);
|
||||
const runtime = this.runtimes.get(key);
|
||||
if (!runtime) throw new Error(`Runtime ${key} does not exist`);
|
||||
const existing = this.deleting.get(key);
|
||||
if (existing) return existing;
|
||||
const deletion = runtime.delete();
|
||||
this.deleting.set(key, deletion);
|
||||
try { await deletion; } finally {
|
||||
if (this.deleting.get(key) === deletion) this.deleting.delete(key);
|
||||
if (this.runtimes.get(key) === runtime) this.runtimes.delete(key);
|
||||
}
|
||||
}
|
||||
|
||||
private reserveCapacity(key: string): void {
|
||||
const limit = this.options.maxRuntimes ?? 20;
|
||||
if (this.runtimes.size + this.reservations >= limit) {
|
||||
const candidate = [...this.runtimes.values()]
|
||||
.filter((runtime) => runtime.isDetached && !runtime.isActive)
|
||||
.sort((a, b) => a.lastUsedAt - b.lastUsedAt)[0];
|
||||
if (!candidate) throw new Error(`Runtime capacity reached for ${key}`);
|
||||
void candidate.dispose();
|
||||
if (this.runtimes.get(candidate.key) === candidate) this.runtimes.delete(candidate.key);
|
||||
}
|
||||
this.reservations++;
|
||||
}
|
||||
|
||||
async shutdown(timeoutMs = 30_000): Promise<void> {
|
||||
if (this.shutdownPromise) return this.shutdownPromise;
|
||||
this.shuttingDown = true;
|
||||
this.shutdownPromise = (async () => {
|
||||
const work = [
|
||||
...this.creating.values(),
|
||||
...this.runtimes.values(),
|
||||
].map((item) => item instanceof Promise
|
||||
? item.then(async (runtime) => { await runtime.shutdown(); }).catch(() => {})
|
||||
: item.shutdown());
|
||||
await Promise.race([
|
||||
Promise.all(work).then(() => undefined),
|
||||
new Promise<void>((resolve) => setTimeout(resolve, timeoutMs).unref()),
|
||||
]);
|
||||
this.runtimes.clear();
|
||||
})();
|
||||
return this.shutdownPromise;
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user