import { produce } from 'immer'; import { BugIndicatingError } from '#/_base/errors/errors'; import { onUnexpectedError } from '#/_base/errors/unexpectedError'; import { ILogService } from '#/_base/log/log'; import { Service } from '#/_base/di/service'; import { toDisposable, type IDisposable } from '#/_base/di/lifecycle'; import { type CollectionView } from '#/_base/di/collection'; import { LifecycleScope } from '#/app/scopes'; import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; import { AgentSpaceImpl, type AgentSpaceHost } from '#/agent/agentContext/agentSpace'; import { IAgentBlobService } from '#/agent/blob/agentBlobService'; import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext'; import { IAgentStateService } from '#/agent/state/agentState'; import { event2FromRecord, type AgentDomainTrait, type Event2, type Event2Class, } from '#/app/event/event2'; import { IEventBus } from '#/app/event/eventBus'; import type { ContentPart } from '#human/llm/message'; import { OrderedHookSlot } from '#/hooks'; import { IWireService } from '#/wire/wire'; import { WireError, WireErrors } from '#/wire/errors'; import { isHumanRecordType } from '#/wire/human'; import { AGENT_SWITCHED_TYPE } from '#/wire/tree/index'; import type { PartsTransformer } from '#/wire/record'; import { AgentModelContribution, agentModelDefinitions, type AgentModel, type AgentModelDefinition, } from './agentModel'; import { IEventDispatcher, type DurableAgentRuntimeParticipant, type RestorePhase } from './eventDispatcher'; import { StateError, StateErrors } from './errors'; import { expandedModelAppliers, expandedRuntimeFolds, type EventApplier, type StateFold, type FoldContext, type ReplayableStateKey, } from './state'; import { EventStateContribution, foldEventStateContributions, type EventStateContributionRecord, type FoldedEventStateRegistry, } from './stateContribution'; const MAX_DRAIN = 100; const UNREPORTED_WIRE_RECORD_TYPES: ReadonlySet = new Set([ 'staleGuard.recorded', 'staleGuard.cleared', AGENT_SWITCHED_TYPE, 'context.undone', ]); export class CycleError extends StateError { constructor(readonly depth: number, readonly eventTypes: readonly string[]) { super( StateErrors.codes.STATE_CYCLE, `Event dispatch cascade exceeded MAX_DRAIN (${depth}); possible event cycle`, { details: { depth, eventTypes: eventTypes.slice(0, 20) } }, ); this.name = 'CycleError'; } } interface StateMeta { checkpoints: unknown[]; } interface QueuedEvent { readonly event: Event2; readonly resolve: () => void; readonly reject: (error: unknown) => void; } interface PreparedFold { readonly key: ReplayableStateKey; readonly meta: StateMeta; readonly ctx: FoldContextImpl; readonly next: any; } type ParticipantApplier = ( state: any, event: Event2, ctx: FoldContextImpl, ) => unknown; interface ParticipantAttachment { readonly id: string; readonly appliers: ReadonlyMap, ParticipantApplier>; readonly meta: StateMeta; readonly undoable: boolean; readonly initial: unknown; readonly getState: () => any; readonly commit: (state: any) => void; } interface PreparedParticipant { readonly attachment: ParticipantAttachment; readonly ctx: FoldContextImpl; readonly next: any; } class FoldContextImpl implements FoldContext { pendingCheckpoint = false; pendingClear = false; pendingUndo: number | undefined; constructor( private readonly owner: EventDispatcherService, readonly silent: boolean, ) {} checkpoint(): void { if (!this.silent) return; this.pendingCheckpoint = true; } clearCheckpoints(): void { if (!this.silent) return; this.pendingClear = true; } undoToCheckpoint(count: number): void { if (!this.silent) return; this.pendingUndo = count; } emit(event: Event2): void { if (this.silent) return; this.owner.enqueue(event); } } function sanitizePendingUndo(ctx: FoldContextImpl, meta: StateMeta): void { if ( ctx.pendingUndo !== undefined && (!Number.isSafeInteger(ctx.pendingUndo) || ctx.pendingUndo <= 0 || meta.checkpoints.length < ctx.pendingUndo) ) { ctx.pendingUndo = undefined; } } export class EventDispatcherService extends Service implements IEventDispatcher { declare readonly _serviceBrand: undefined; readonly hooks: IEventDispatcher['hooks'] = { onDidRestore: new OrderedHookSlot(), }; private readonly metas = new Map, StateMeta>(); private folded: FoldedEventStateRegistry; private activeModelDefs = new Map>(); private readonly withdrawnModelIds = new Set(); private modelTargets = new Map[]>(); private readonly modelAttachments = new Map< AgentModelDefinition, ParticipantAttachment >(); private readonly participantTargets = new Map(); private readonly participantAttachments = new Map(); private readonly spaceHost: AgentSpaceHost = { isActiveModelDefinition: (definition) => this.activeModelDefs.get(definition.id) === definition, registerModel: (definition, model) => this.registerModel(definition, model), dispatchModelEvent: (event) => this.dispatch(event), readLegacyState: (key) => this.agentState.get(key), }; restorePhase: RestorePhase = 'new'; private dispatching = false; private disposed = false; private queue: QueuedEvent[] = []; private drainDepth = 0; private didRunRestoreHooks = false; private lateAttachments: Array<{ readonly participant: DurableAgentRuntimeParticipant; readonly resolve: (disposable: IDisposable) => void; readonly reject: (error: unknown) => void; }> = []; constructor( @IWireService private readonly wire: IWireService, @IEventBus private readonly eventBus: IEventBus, @IAgentScopeContext private readonly agentScope: IAgentScopeContext | undefined, @IAgentBlobService private readonly blobService: IAgentBlobService, @IAgentStateService private readonly agentState: IAgentStateService, @ILogService private readonly logger: ILogService, @EventStateContribution view: CollectionView, @AgentModelContribution modelView: CollectionView>, ) { super(); this.folded = this.foldContributions(view); this._register( view.onDidChange(() => { this.folded = this.foldContributions(view); }), ); this._register( this.agentState.onDidContributeReplayable((key) => { if (this.restorePhase !== 'new') { throw new BugIndicatingError( `Replayable state '${key.name}' contributed while the event dispatcher is in phase '${this.restorePhase}'; replayable state owners must contribute before restore`, ); } this.folded = this.foldContributions(view); }), ); this._register( this.agentState.onDidWithdrawReplayable((key) => { this.metas.delete(key); this.folded = this.foldContributions(view); }), ); this.refoldModels(modelView.items); this._register( modelView.onDidChange(({ added, removed }) => { for (const definition of removed) { this.withdrawnModelIds.add(definition.id); const attachment = this.modelAttachments.get(definition); if (attachment !== undefined) { this.modelAttachments.delete(definition); this.detachParticipant(attachment); this.space()?.retireModel(definition); } } for (const definition of added) { this.withdrawnModelIds.delete(definition.id); } this.refoldModels(modelView.items); this.materializeUndoableModels(); }), ); this.space()?._attachHost(this.spaceHost); this.materializeUndoableModels(); } private space(): AgentSpaceImpl | undefined { const space = this.agentScope?.agentContext.space; return space instanceof AgentSpaceImpl ? space : undefined; } private foldContributions( view: CollectionView, ): FoldedEventStateRegistry { return foldEventStateContributions(view.items, this.agentState.replayableKeys()); } attach(participant: DurableAgentRuntimeParticipant): IDisposable { if (this.restorePhase !== 'new') { throw new BugIndicatingError( `Agent runtime participant '${participant.id}' attached while the event dispatcher is in phase '${this.restorePhase}'; durable runtime owners must attach before restore`, ); } const attachment = this.buildParticipantAttachment(participant); this.attachParticipant(attachment); return toDisposable(() => { this.detachParticipant(attachment); }); } async attachLate(participant: DurableAgentRuntimeParticipant): Promise { if (this.restorePhase === 'restoring') { return new Promise((resolve, reject) => { this.lateAttachments.push({ participant, resolve, reject }); }); } if (this.restorePhase !== 'ready') { throw new BugIndicatingError( `Agent runtime participant '${participant.id}' late-attached while the event dispatcher is in phase '${this.restorePhase}'; late attach requires a restored dispatcher`, ); } return this.attachLateNow(participant); } private async attachLateNow(participant: DurableAgentRuntimeParticipant): Promise { if (this.disposed) { throw new Error(`Agent runtime participant '${participant.id}' late-attached to a disposed event dispatcher`); } const attachment = this.buildParticipantAttachment(participant); this.dispatching = true; try { await this.wire.flush(); const stream = participant.undoable ? this.wire.readRestorable() : this.wire.readJournal(); for await (const record of stream) { if (record.type === 'metadata') continue; const cls = this.folded.events.get(record.type); if (cls === undefined) continue; let eventRecord = record; if (cls.agentDomain) { if (this.agentScope === undefined) continue; const recordAgentId = record['agentId']; if (recordAgentId === undefined) eventRecord = { ...record, agentId: this.agentScope.agentId }; else if (recordAgentId !== this.agentScope.agentId) continue; } const event = event2FromRecord(cls, eventRecord); if (event === undefined) continue; const applier = attachment.appliers.get(event.constructor as Event2Class); if (applier === undefined) continue; const ctx = new FoldContextImpl(this, true); const next = produce( attachment.getState(), (draft: any) => applier(draft, event, ctx), ); if (ctx.pendingUndo !== undefined && next !== attachment.getState()) { throw new BugIndicatingError( `Fold of event '${event.type}' on durable participant '${attachment.id}' both mutates and undoes to a checkpoint`, ); } sanitizePendingUndo(ctx, attachment.meta); this.commitParticipant(attachment, ctx, next); } this.attachParticipant(attachment); this.drainQueue(); } catch (error) { for (const entry of this.queue.splice(0)) entry.reject(error); throw error; } finally { this.queue.length = 0; this.dispatching = false; this.drainDepth = 0; } return toDisposable(() => { this.detachParticipant(attachment); }); } private buildParticipantAttachment( participant: DurableAgentRuntimeParticipant, ): ParticipantAttachment { const base = new Map, StateFold>(); for (const cls of participant.events) base.set(cls, participant.transition); const folds = expandedRuntimeFolds(participant.id, participant.undoable, base); const appliers = new Map, ParticipantApplier>(); for (const [cls, fold] of folds) { appliers.set(cls, (state, event, ctx) => fold(state, event, ctx)); } return { id: participant.id, appliers, meta: { checkpoints: [] }, undoable: participant.undoable, initial: participant.getState(), getState: () => participant.getState(), commit: (state) => { participant.commit(state); }, }; } private attachParticipant(attachment: ParticipantAttachment): void { if (this.participantAttachments.has(attachment.id)) { throw new BugIndicatingError(`Durable participant '${attachment.id}' is already attached`); } this.participantAttachments.set(attachment.id, attachment); for (const cls of attachment.appliers.keys()) { const list = this.participantTargets.get(cls.type) ?? []; list.push(attachment); this.participantTargets.set(cls.type, list); } } private detachParticipant(attachment: ParticipantAttachment): void { if (this.participantAttachments.get(attachment.id) !== attachment) return; this.participantAttachments.delete(attachment.id); for (const cls of attachment.appliers.keys()) { const list = this.participantTargets.get(cls.type); if (list === undefined) continue; const next = list.filter((candidate) => candidate !== attachment); if (next.length === 0) this.participantTargets.delete(cls.type); else this.participantTargets.set(cls.type, next); } } private refoldModels(records: readonly AgentModelDefinition[]): void { const defs = new Map>(); for (const definition of agentModelDefinitions()) { if (!this.withdrawnModelIds.has(definition.id)) defs.set(definition.id, definition); } for (const definition of records) defs.set(definition.id, definition); this.activeModelDefs = defs; this.rebuildModelTargets(); } private rebuildModelTargets(): void { const targets = new Map[]>(); const add = (type: string, definition: AgentModelDefinition): void => { const list = targets.get(type); if (list === undefined) { targets.set(type, [definition]); return; } if (!list.includes(definition)) list.push(definition); }; const domainOwners = new Map>(); for (const definition of this.activeModelDefs.values()) { for (const cls of definition.events) { const owner = domainOwners.get(cls.type); if (owner !== undefined && owner !== definition) { throw new BugIndicatingError( `Event '${cls.type}' is applied by both agent models '${owner.id}' and '${definition.id}'`, ); } domainOwners.set(cls.type, definition); add(cls.type, definition); } } for (const [definition, attachment] of this.modelAttachments) { if (this.activeModelDefs.get(definition.id) !== definition) continue; for (const cls of attachment.appliers.keys()) add(cls.type, definition); } this.modelTargets = targets; } private materializeUndoableModels(): void { const space = this.space(); if (space === undefined) return; for (const definition of this.activeModelDefs.values()) { if (!definition.undoable || this.modelAttachments.has(definition)) continue; space.ensureModel(definition); } } private registerModel( definition: AgentModelDefinition, model: AgentModel, ): void { if (this.modelAttachments.has(definition)) return; const domainAppliers = new Map, EventApplier>(); for (const [cls, applier] of model._appliersTable()) { domainAppliers.set(cls, (event) => applier.call(model, event)); } const customUndo = model.onUndo === undefined ? undefined : (count: number): void => model.onUndo!(count); const expanded = expandedModelAppliers( definition.id, definition.undoable, domainAppliers, customUndo, ); const appliers = new Map, ParticipantApplier>(); for (const [cls, applier] of expanded) { appliers.set(cls, (state, event, ctx) => { model._enterWindow(state, ctx); let windowResult: ReturnType['_exitWindow']>; try { applier(event, ctx); } finally { windowResult = model._exitWindow(); } return windowResult.replaced ? windowResult.replacement : undefined; }); } const attachment: ParticipantAttachment = { id: definition.id, appliers, meta: { checkpoints: [] }, undoable: definition.undoable, initial: model._state(), getState: () => model._state(), commit: (state) => { model._commitState(state); }, }; this.attachParticipant(attachment); this.modelAttachments.set(definition, attachment); this.rebuildModelTargets(); } private materializeModel(definition: AgentModelDefinition): ParticipantAttachment { const space = this.space(); if (space === undefined) { throw new BugIndicatingError( `Agent model '${definition.id}' cannot materialize without an agent space`, ); } space.ensureModel(definition); const attachment = this.modelAttachments.get(definition); if (attachment === undefined) { throw new BugIndicatingError(`Agent model '${definition.id}' failed to attach`); } return attachment; } dispatch(event: Event2): Promise { const cls = event.constructor as Event2Class; if ( cls.agentDomain && (this.agentScope === undefined || (event as Event2 & AgentDomainTrait).agentId !== this.agentScope.agentId) ) { return Promise.reject( new Error(`Agent event '${event.type}' does not match dispatcher lifecycle context`), ); } if (this.dispatching) { return new Promise((resolve, reject) => { this.queue.push({ event, resolve, reject }); }); } this.dispatching = true; try { this.runDispatch(event); this.drainQueue(); return Promise.resolve(); } catch (error) { for (const entry of this.queue.splice(0)) { entry.reject(error); } return Promise.reject(error); } finally { this.queue.length = 0; this.dispatching = false; this.drainDepth = 0; } } private drainQueue(): void { while (this.queue.length > 0) { if (++this.drainDepth > MAX_DRAIN) { throw new CycleError( this.drainDepth, this.queue.map((entry) => entry.event.type), ); } const entry = this.queue.shift()!; try { this.runDispatch(entry.event); entry.resolve(); } catch (error) { entry.reject(error); throw error; } } } enqueue(event: Event2): void { this.queue.push({ event, resolve: () => {}, reject: (error: unknown) => onUnexpectedError(error), }); } private runDispatch(event: Event2): void { this.executeEvent(event, false); } private executeEvent(event: Event2, silent: boolean, replayUndoable?: boolean): void { const folds = this.folded.folds.get(event.type); const prepared: PreparedFold[] = []; if (folds !== undefined) { for (const { key, fold } of folds) { if ( replayUndoable !== undefined && (key.replayable.undoable !== undefined) !== replayUndoable ) { continue; } const meta = this.ensureMeta(key); const ctx = new FoldContextImpl(this, silent); const next = produce( this.agentState.get(key), (draft: any) => fold(draft, event, ctx), ); if (ctx.pendingUndo !== undefined && next !== this.agentState.get(key)) { throw new BugIndicatingError( `Fold of event '${event.type}' on state '${key.name}' both mutates and undoes to a checkpoint`, ); } sanitizePendingUndo(ctx, meta); prepared.push({ key, meta, ctx, next }); } } const modelTargets = this.modelTargets.get(event.type); if (modelTargets !== undefined) { for (const definition of modelTargets) { if (replayUndoable !== undefined && definition.undoable !== replayUndoable) continue; if (!this.modelAttachments.has(definition)) this.materializeModel(definition); } } const participantTargets = this.participantTargets.get(event.type); const preparedParticipants: PreparedParticipant[] = []; if (participantTargets !== undefined) { for (const attachment of participantTargets) { if (replayUndoable !== undefined && attachment.undoable !== replayUndoable) continue; const applier = attachment.appliers.get(event.constructor as Event2Class); if (applier === undefined) continue; const ctx = new FoldContextImpl(this, silent); const next = produce( attachment.getState(), (draft: any) => applier(draft, event, ctx), ); if (ctx.pendingUndo !== undefined && next !== attachment.getState()) { throw new BugIndicatingError( `Fold of event '${event.type}' on durable participant '${attachment.id}' both mutates and undoes to a checkpoint`, ); } sanitizePendingUndo(ctx, attachment.meta); preparedParticipants.push({ attachment, ctx, next }); } } for (const p of prepared) { this.commit(p.key, p.meta, p.ctx, p.next); } for (const p of preparedParticipants) { this.commitParticipant(p.attachment, p.ctx, p.next); } if (silent) return; const cls = event.constructor as Event2Class; if (cls.durable) { const dehydrator = folds?.find(({ key }) => key.replayable.blobs !== undefined)?.key .replayable.blobs?.dehydrate; this.wire.appendRecord(event.serialize(), dehydrator); } if (cls.observable && !this.disposed) { this.eventBus.publish(event, this.agentScope?.agentContext); } } override dispose(): void { this.disposed = true; const pending = this.lateAttachments.splice(0); if (pending.length > 0) { const error = new Error('Event dispatcher disposed while a late attach was pending'); for (const entry of pending) entry.reject(error); } this.space()?._detachHost(this.spaceHost); super.dispose(); } private commit( key: ReplayableStateKey, meta: StateMeta, ctx: FoldContextImpl, next: any, ): void { if (ctx.pendingUndo !== undefined) { const targetIndex = meta.checkpoints.length - ctx.pendingUndo; const snapshot = meta.checkpoints[targetIndex]!; this.agentState.set(key, snapshot); meta.checkpoints.length = targetIndex; return; } this.agentState.set(key, next); if (ctx.pendingClear) { meta.checkpoints.length = 0; } if (ctx.pendingCheckpoint) { meta.checkpoints.push(next); } } private commitParticipant( attachment: ParticipantAttachment, ctx: FoldContextImpl, next: any, ): void { const meta = attachment.meta; if (ctx.pendingUndo !== undefined) { const targetIndex = meta.checkpoints.length - ctx.pendingUndo; const snapshot = meta.checkpoints[targetIndex]!; attachment.commit(snapshot); meta.checkpoints.length = targetIndex; return; } attachment.commit(next); if (ctx.pendingClear) { meta.checkpoints.length = 0; } if (ctx.pendingCheckpoint) { meta.checkpoints.push(next); } } private ensureMeta(key: ReplayableStateKey): StateMeta { let meta = this.metas.get(key); if (meta === undefined) { meta = { checkpoints: [] }; this.metas.set(key, meta); } return meta; } async restore(): Promise { if (this.restorePhase === 'restoring') { throw new BugIndicatingError( `Agent state restore called while phase is ${this.restorePhase}`, ); } const rerun = this.restorePhase !== 'new'; this.restorePhase = 'restoring'; if (rerun) this.dispatching = true; try { if (rerun) { await this.wire.flush(); this.resetReplayState(); } await this.replayRecords(true); await this.replayRecords(false); await this.rehydrateStates(); this.restorePhase = 'ready'; if (!this.didRunRestoreHooks) { await this.hooks.onDidRestore.run({}); this.didRunRestoreHooks = true; } if (rerun) { this.drainQueue(); } await this.drainLateAttachments(); } catch (error) { this.restorePhase = 'failed'; for (const pending of this.lateAttachments.splice(0)) pending.reject(error); if (rerun) { for (const entry of this.queue.splice(0)) entry.reject(error); } throw error; } finally { if (rerun) { this.queue.length = 0; this.dispatching = false; this.drainDepth = 0; } } } private async drainLateAttachments(): Promise { for (const pending of this.lateAttachments.splice(0)) { try { pending.resolve(await this.attachLateNow(pending.participant)); } catch (error) { pending.reject(error); } } } private resetReplayState(): void { for (const key of this.agentState.replayableKeys()) { this.agentState.set(key, key.initial()); } for (const attachment of this.participantAttachments.values()) { attachment.commit(attachment.initial); attachment.meta.checkpoints.length = 0; } this.metas.clear(); } private async replayRecords(undoable: boolean): Promise { const stream = undoable ? this.wire.readRestorable() : this.wire.readJournal(); let recordIndex = 0; for await (const record of stream) { if (record.type === 'metadata') continue; const cls = this.folded.events.get(record.type); if (cls === undefined) { if ( !undoable && !UNREPORTED_WIRE_RECORD_TYPES.has(record.type) && !isHumanRecordType(record.type) ) { this.reportSkippedRecord(record.type, recordIndex, false); } recordIndex++; continue; } let eventRecord = record; if (cls.agentDomain) { if (this.agentScope === undefined) { if (!undoable) this.reportSkippedRecord(record.type, recordIndex, true); recordIndex++; continue; } const recordAgentId = record['agentId']; if (recordAgentId === undefined) { eventRecord = { ...record, agentId: this.agentScope.agentId }; } else if (recordAgentId !== this.agentScope.agentId) { if (!undoable) this.reportSkippedRecord(record.type, recordIndex, true); recordIndex++; continue; } } const event = event2FromRecord(cls, eventRecord); if (event === undefined) { if (!undoable) this.reportSkippedRecord(record.type, recordIndex, true); recordIndex++; continue; } this.executeEvent(event, true, undoable); recordIndex++; } } private reportSkippedRecord(type: string, index: number, malformed: boolean): void { const message = malformed ? `Malformed wire record type '${type}' skipped during restore` : `Unknown wire record type '${type}' skipped during restore`; if (malformed) { onUnexpectedError( new WireError(WireErrors.codes.WIRE_UNKNOWN_RECORD, message, { details: { type, index } }), ); return; } this.logger.warn(message, { code: WireErrors.codes.WIRE_UNKNOWN_RECORD, type, index }); } private async rehydrateStates(): Promise { const transform: PartsTransformer = (parts) => this.blobService.loadParts(parts as readonly ContentPart[]) as Promise; for (const key of this.folded.states) { const codec = key.replayable.blobs; if (codec?.rehydrate === undefined) continue; this.agentState.set(key, Object.freeze(await codec.rehydrate(this.agentState.get(key), transform))); } } async flush(): Promise { await this.wire.flush(); } } registerScopedService( LifecycleScope.Agent, IEventDispatcher, EventDispatcherService, ScopeActivation.OnScopeCreated, 'state', );