import type { UserPromptOrigin } from '#/agent/contextMemory/types'; /* oxlint-disable typescript-eslint/no-unsafe-declaration-merging, eslint-plugin-import/namespace -- Event2 class+payload-interface declaration merging is the sanctioned event-declaration idiom. */ import { type IDisposable } from '#/_base/di/lifecycle'; import { Service } from '#/_base/di/service'; import { BugIndicatingError } from '#/_base/errors/errors'; import { LifecycleScope } from '#/app/scopes'; import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; import { ILogService } from '#/_base/log/log'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; import { IAgentConversationUndoParticipantRegistry } from '#/agent/contextMemory/conversationUndoParticipants'; import { computeUndoCut, formatUndoUnavailableMessage, precheckUndo, } from '#/agent/contextMemory/contextOps'; import { isUndoAnchor, isValidUndoCount, } from '#/agent/contextMemory/conversationTime'; import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompaction'; import { IAgentLoopService } from '#/agent/loop/loop'; import { turnKey } from '#/agent/loop/turnOps'; import { promptMetadataTextFromContentParts } from '#/agent/prompt/promptMetadataText'; import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext'; import { IAgentStateService } from '#/agent/state/agentState'; import { IEventService } from '#/app/event/event'; import { AgentEvent2 } from '#/app/event/event2'; import { ITelemetryService } from '#/app/telemetry/telemetry'; import { ErrorCodes, Error2 } from '#/errors'; import { MAIN_AGENT_ID } from '#/session/agentLifecycle/agentLifecycle'; import { ISessionContext } from '#/session/sessionContext/sessionContext'; import { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata'; import { SessionMetaUpdated } from '#/session/sessionMetadata/sessionMetaEvents'; import { ISessionTokenCountingService } from '#/session/tokenCounting/sessionTokenCounting'; import { IEventDispatcher } from '#/state/eventDispatcher'; import { ForkLineError, IWireService } from '#/wire/wire'; import { IAgentConversationUndoService, type UndoAvailability } from './undo'; export class ContextUndone extends AgentEvent2<{ readonly agentId: string; readonly turns: number; readonly fromTurnId?: number; }> { static override readonly type = 'context.undone'; static override readonly observable = true; } export interface ContextUndone { readonly agentId: string; readonly turns: number; readonly fromTurnId?: number; } export class AgentConversationUndoService extends Service implements IAgentConversationUndoService { declare readonly _serviceBrand: undefined; private undoQueue: Promise = Promise.resolve(); constructor( @IAgentLoopService private readonly loop: IAgentLoopService, @IAgentFullCompactionService private readonly fullCompaction: IAgentFullCompactionService, @IAgentContextMemoryService private readonly context: IAgentContextMemoryService, @IAgentConversationUndoParticipantRegistry private readonly participants: IAgentConversationUndoParticipantRegistry, @IAgentScopeContext private readonly agentCtx: IAgentScopeContext, @ISessionContext private readonly session: ISessionContext, @ISessionMetadata private readonly metadata: ISessionMetadata, @IEventService private readonly eventService: IEventService, @ITelemetryService private readonly telemetry: ITelemetryService, @IEventDispatcher private readonly dispatcher: IEventDispatcher, @IAgentStateService private readonly agentState: IAgentStateService, @ISessionTokenCountingService private readonly tokenCounting: ISessionTokenCountingService, @IWireService private readonly wire: IWireService, @ILogService private readonly log: ILogService, ) { super(); } availability(): UndoAvailability { const cut = computeUndoCut(this.context.get(), Number.MAX_SAFE_INTEGER); return { maxTurns: cut.removedCount, stoppedAtCompaction: cut.stoppedAtCompaction, }; } async undo(turns: number): Promise { if (!isValidUndoCount(turns)) { throw new Error2( ErrorCodes.REQUEST_INVALID, 'Undo count must be a positive safe integer', { details: { field: 'count' } }, ); } const run = this.undoQueue.then(() => this.undoNow(turns)); this.undoQueue = run.then( () => undefined, () => undefined, ); return run; } private async undoNow(turns: number): Promise { let quiescence: IDisposable | undefined; try { quiescence = this.loop.tryAcquireQuiescence(); if (quiescence === undefined) { throw this.busyError('loop'); } if (this.fullCompaction.compacting !== null) { throw this.busyError('compaction'); } if (this.dispatcher.restorePhase !== 'ready') { throw new BugIndicatingError( `Conversation undo requires a restored dispatcher (phase '${this.dispatcher.restorePhase}')`, ); } this.assertUndoAvailable(turns); const fromTurnId = this.removedFromTurnId(turns); try { await this.wire.switchBranch({ turns, fromTurnId }); } catch (error) { throw this.forkError(error, turns); } try { await this.dispatcher.restore(); } catch (error) { this.log.warn('undo state restore failed; retrying once', { error }); await this.dispatcher.restore(); } await this.loop.resetMachineEngine(); this.tokenCounting.recordTruncation( this.agentCtx.agentContext, this.context.get().length, ); await this.reconcileParticipants(); await this.flushAfterReconcile(); await this.reconcileParticipants('after-flush'); await this.reconcileLastPromptSafely(); this.telemetry.track2('conversation_undo', { count: turns }); await this.dispatcher.dispatch( new ContextUndone({ agentId: this.agentCtx.agentId, turns, fromTurnId }), ); return turns; } finally { quiescence?.dispose(); } } private forkError(error: unknown, turns: number): unknown { if (!(error instanceof ForkLineError)) return error; return new Error2( ErrorCodes.SESSION_UNDO_UNAVAILABLE, formatUndoUnavailableMessage({ ok: false, reason: error.reason, requested: turns, undoable: error.available, }), { details: { reason: error.reason, requestedCount: turns, undoableCount: error.available, }, }, ); } private removedFromTurnId(turns: number): number | undefined { if (!this.agentState.has(turnKey)) return undefined; const anchorTurnIds = this.agentState.get(turnKey).anchorTurnIds; if (anchorTurnIds.length < turns) return undefined; const totalAnchors = computeUndoCut(this.context.get(), Number.MAX_SAFE_INTEGER).removedCount; if (totalAnchors !== anchorTurnIds.length) return undefined; return anchorTurnIds[anchorTurnIds.length - turns]; } private busyError(reason: 'loop' | 'compaction'): Error2 { const message = reason === 'loop' ? 'Cannot undo while a turn is active or queued. Wait for it to finish, then retry.' : 'Cannot undo while conversation compaction is running. Wait for it to finish, then retry.'; return new Error2(ErrorCodes.SESSION_BUSY, message, { details: { reason } }); } private assertUndoAvailable(turns: number): void { const check = precheckUndo(this.context.get(), turns); if (check.ok) return; throw new Error2( ErrorCodes.SESSION_UNDO_UNAVAILABLE, formatUndoUnavailableMessage(check), { details: { reason: check.reason, requestedCount: check.requested, undoableCount: check.undoable, }, }, ); } private async reconcileParticipants(phase?: 'after-flush'): Promise { const participants = this.participants.list().filter((participant) => participant.phase === phase); const results = await Promise.allSettled( participants.map((participant) => participant.reconcileAfterUndo()), ); results.forEach((result, index) => { if (result.status === 'fulfilled') return; this.log.error('undo participant reconciliation failed', { participantId: participants[index]?.id, error: result.reason, }); }); } private async reconcileLastPromptSafely(): Promise { try { await this.reconcileLastPrompt(); } catch (error) { this.log.error('undo lastPrompt reconciliation failed', { error }); } } private async flushAfterReconcile(): Promise { try { await this.dispatcher.flush(); } catch (error) { this.log.error('undo wire flush failed after in-memory commit', { stage: 'state reconciliation', error, }); throw error; } } private async reconcileLastPrompt(): Promise { if (this.agentCtx.agentId !== MAIN_AGENT_ID) return; const pending = this.loop.snapshot().queue.filter((item) => item.meta?.tracked === true).at(-1); let lastPrompt = pending === undefined ? undefined : promptMetadataTextFromContentParts(pending.message.content, (pending.meta?.origin as UserPromptOrigin | undefined)?.clientMetadata); if (lastPrompt === undefined) { const history = this.context.get(); for (let i = history.length - 1; i >= 0; i--) { const message = history[i]!; if (!isUndoAnchor(message)) continue; lastPrompt = promptMetadataTextFromContentParts(message.content, message.origin?.kind === 'user' || message.origin?.kind === 'skill_activation' ? message.origin.clientMetadata : undefined); if (lastPrompt !== undefined) break; } } await this.metadata.update({ lastPrompt }); this.eventService.publish( new SessionMetaUpdated({ payload: { agentId: MAIN_AGENT_ID, sessionId: this.session.sessionId, patch: { lastPrompt }, }, }), ); } } registerScopedService( LifecycleScope.Agent, IAgentConversationUndoService, AgentConversationUndoService, ScopeActivation.OnScopeCreated, 'undo', );