diff --git a/docs/design/2026-08-08-selective-session-restore.md b/docs/design/2026-08-08-selective-session-restore.md index d835cb07e0..d82913dfcc 100644 --- a/docs/design/2026-08-08-selective-session-restore.md +++ b/docs/design/2026-08-08-selective-session-restore.md @@ -475,11 +475,13 @@ segments once: malformed-context, turn-reentry, and truncation decisions without retaining evidence content. Add only the selected evidence UUIDs to the union, then feed their materialized records to the shared accumulator and retain the resulting - window in the projection. This two-stage selection must preserve both the - existing production helper's result and its fail-closed errors; it must not - select every active record, perform a second scan, or copy Goal precedence. - Deferred Goal activation consumes that window instead of reading the - transcript again. + window in the projection. This two-stage selection must preserve the existing + production helper's valid result. When its evidence source is unavailable or + invalid, omit the projected window so deferred Goal activation falls back to + the existing runtime path and its established degradation behavior instead of + rejecting the whole session restore. It must not select every active record, + perform a second scan, or copy Goal precedence. Deferred Goal activation + consumes a valid projected window instead of reading the transcript again. 5. **File history.** Read every active `file_history_snapshot` record in chronological order and feed each batch through the existing whole-batch deserializer. This preserves today's behavior where one malformed item skips diff --git a/packages/acp-bridge/src/bridge.test.ts b/packages/acp-bridge/src/bridge.test.ts index 3fa6c717f0..5053f187ba 100644 --- a/packages/acp-bridge/src/bridge.test.ts +++ b/packages/acp-bridge/src/bridge.test.ts @@ -4273,6 +4273,7 @@ describe('createAcpSessionBridge', () => { keep: 'state', 'qwen.session.loadReplay': { v: 1, + anchorRecordId: 'record-anchor', hasMore: true, partial: true, replayError: 'replay boom', @@ -4317,6 +4318,7 @@ describe('createAcpSessionBridge', () => { expect(loaded.partial).toBe(true); expect(loaded.replayError).toBe('replay boom'); expect(loaded.historyHasMore).toBe(true); + expect(loaded.historyAnchorRecordId).toBe('record-anchor'); expect(loaded.lastEventId).toBe(2); expect(loaded.compactedReplay).toHaveLength(2); expect(loaded.liveJournal).toEqual([]); diff --git a/packages/acp-bridge/src/bridge.ts b/packages/acp-bridge/src/bridge.ts index 9003e1f7e6..669a35db71 100644 --- a/packages/acp-bridge/src/bridge.ts +++ b/packages/acp-bridge/src/bridge.ts @@ -137,6 +137,7 @@ import { DAEMON_PROMPT_DISPLAY_TEXT_META_KEY, LOAD_REPLAY_BULK_MODE, LOAD_REPLAY_HIDE_INHERITED_META_KEY, + LOAD_REPLAY_MAX_UPDATES, LOAD_REPLAY_META_KEY, LOAD_REPLAY_MODE_META_KEY, LOAD_REPLAY_PAGE_SIZE_META_KEY, @@ -252,7 +253,6 @@ const KNOWN_SESSION_UPDATE_TYPES = new Set([ 'session_info_update', 'usage_update', ]); -const MAX_BULK_REPLAY_UPDATES = 10_000; function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); @@ -661,6 +661,7 @@ function describeLoadReplayValue(value: unknown): string { function extractLoadReplayResponse(state: BridgeSessionState): { state: BridgeSessionState; updates: SessionUpdate[]; + anchorRecordId?: string; partial?: true; replayError?: string; hasMore?: boolean; @@ -682,10 +683,10 @@ function extractLoadReplayResponse(state: BridgeSessionState): { `(version=${LOAD_REPLAY_VERSION}, count=not-array)`, ); } - if (rawUpdates.length > MAX_BULK_REPLAY_UPDATES) { + if (rawUpdates.length > LOAD_REPLAY_MAX_UPDATES) { throw new Error( `qwen.session.loadReplay updates exceed limit ` + - `(${rawUpdates.length} > ${MAX_BULK_REPLAY_UPDATES})`, + `(${rawUpdates.length} > ${LOAD_REPLAY_MAX_UPDATES})`, ); } const partial = replay['partial']; @@ -709,6 +710,13 @@ function extractLoadReplayResponse(state: BridgeSessionState): { `(version=${LOAD_REPLAY_VERSION}, hasMore=${describeLoadReplayValue(hasMore)})`, ); } + const anchorRecordId = replay['anchorRecordId']; + if (anchorRecordId !== undefined && typeof anchorRecordId !== 'string') { + throw new Error( + `Invalid qwen.session.loadReplay anchorRecordId ` + + `(version=${LOAD_REPLAY_VERSION}, anchorRecordId=${describeLoadReplayValue(anchorRecordId)})`, + ); + } const invalidUpdateIndex = rawUpdates.findIndex( (update) => !isBulkReplayUpdate(update), ); @@ -735,6 +743,7 @@ function extractLoadReplayResponse(state: BridgeSessionState): { return { state: cleanState, updates: rawUpdates, + ...(typeof anchorRecordId === 'string' ? { anchorRecordId } : {}), ...(partial === true ? { partial: true as const } : {}), ...(typeof replayError === 'string' ? { replayError } : {}), ...(hasMore === true ? { hasMore: true } : {}), @@ -1130,6 +1139,7 @@ interface SessionEntry { restoreReplayPartial?: true; restoreReplayError?: string; restoreHistoryHasMore?: true; + restoreHistoryAnchorRecordId?: string; /** * Most recent heartbeat across any client on this session (Date.now() * epoch ms). Set on every `recordHeartbeat` call regardless of whether @@ -5379,6 +5389,7 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { | 'restoreReplayPartial' | 'restoreReplayError' | 'restoreHistoryHasMore' + | 'restoreHistoryAnchorRecordId' | 'activePromptId' >, action: 'load' | 'resume', @@ -5392,6 +5403,7 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { | 'partial' | 'replayError' | 'historyHasMore' + | 'historyAnchorRecordId' > => { const replayStatus = action === 'load' && entry.restoreReplayPartial === true @@ -5412,6 +5424,10 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { lastEventId: entry.events.lastEventId, eventEpoch, ...replayStatus, + ...(action === 'load' && + entry.restoreHistoryAnchorRecordId !== undefined + ? { historyAnchorRecordId: entry.restoreHistoryAnchorRecordId } + : {}), }; } if (action === 'load') { @@ -5437,6 +5453,9 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { ...(entry.restoreHistoryHasMore === true ? { historyHasMore: true } : {}), + ...(entry.restoreHistoryAnchorRecordId !== undefined + ? { historyAnchorRecordId: entry.restoreHistoryAnchorRecordId } + : {}), }; } return { lastEventId: snapshot.lastEventId, eventEpoch, ...replayStatus }; @@ -6108,6 +6127,7 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { let replayPartial: true | undefined; let replayError: string | undefined; let replayHasMore: true | undefined; + let replayAnchorRecordId: string | undefined; try { const rawRestore = telemetry.withSpan( 'session.restore', @@ -6232,6 +6252,7 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { replayPartial = extracted.partial; replayError = extracted.replayError; replayHasMore = extracted.hasMore === true ? true : undefined; + replayAnchorRecordId = extracted.anchorRecordId; } } catch (err) { if (err instanceof SessionRestoreTimeoutError) throw err; @@ -6353,6 +6374,9 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { if (replayHasMore === true) { entry.restoreHistoryHasMore = true; } + if (replayAnchorRecordId !== undefined) { + entry.restoreHistoryAnchorRecordId = replayAnchorRecordId; + } seedSnapshotCaches(entry, publicState); const artifactRestoreWarnings = await entry.artifacts.restore( restoredArtifactSnapshot, diff --git a/packages/acp-bridge/src/bridgeTypes.ts b/packages/acp-bridge/src/bridgeTypes.ts index 0eec489117..0966b85d65 100644 --- a/packages/acp-bridge/src/bridgeTypes.ts +++ b/packages/acp-bridge/src/bridgeTypes.ts @@ -188,6 +188,8 @@ export const LOAD_REPLAY_HIDE_INHERITED_META_KEY = 'qwen.session.loadReplayHideInherited'; export const LOAD_REPLAY_BULK_MODE = 'bulk'; export const LOAD_REPLAY_VERSION = 1 as const; +export const LOAD_REPLAY_MAX_BYTES = 32 * 1024 * 1024; +export const LOAD_REPLAY_MAX_UPDATES = 10_000; export const REQUESTED_SESSION_ID_META_KEY = 'qwen-code/sessionId'; @@ -338,6 +340,7 @@ export interface ChannelStartupProfileV1 { export interface BridgeLoadReplayEnvelope { v: typeof LOAD_REPLAY_VERSION; updates: SessionUpdate[]; + anchorRecordId?: string; hasMore?: boolean; partial?: true; replayError?: string; diff --git a/packages/channels/base/src/AcpBridge.test.ts b/packages/channels/base/src/AcpBridge.test.ts index 4387a551a8..d2994f2ac0 100644 --- a/packages/channels/base/src/AcpBridge.test.ts +++ b/packages/channels/base/src/AcpBridge.test.ts @@ -109,6 +109,7 @@ type TestableAcpBridge = AcpBridge & { extMethod: ReturnType; newSession?: ReturnType; loadSession?: ReturnType; + unstable_resumeSession?: ReturnType; prompt?: ReturnType; }; knownSessionIds: Set; @@ -418,6 +419,34 @@ describe('AcpBridge', () => { expect(extMethod).toHaveBeenCalledOnce(); }); + it('restores channel sessions through resume without replaying history', async () => { + const bridge = new AcpBridge({ + cliEntryPath: '/tmp/qwen', + cwd: '/tmp', + }) as unknown as TestableAcpBridge; + const resumeSession = vi.fn().mockResolvedValue({}); + bridge.child = { killed: false, exitCode: null }; + bridge.connection = { + extMethod: vi.fn(), + unstable_resumeSession: resumeSession, + } as TestableAcpBridge['connection']; + const bindingToken = {}; + + await expect( + bridge.loadSession('restored-session', '/tmp', undefined, bindingToken), + ).resolves.toBe('restored-session'); + + expect(resumeSession).toHaveBeenCalledWith({ + sessionId: 'restored-session', + cwd: '/tmp', + mcpServers: [], + }); + expect(bridge.knownSessionIds.has('restored-session')).toBe(true); + expect(bridge.sessionBindingTokens.get('restored-session')).toBe( + bindingToken, + ); + }); + it('returns only the final turn text after tool calls', async () => { const bridge = new AcpBridge({ cliEntryPath: '/tmp/qwen', diff --git a/packages/channels/base/src/AcpBridge.ts b/packages/channels/base/src/AcpBridge.ts index 3eccc018f2..8fd41a2441 100644 --- a/packages/channels/base/src/AcpBridge.ts +++ b/packages/channels/base/src/AcpBridge.ts @@ -247,7 +247,7 @@ export class AcpBridge extends EventEmitter implements ChannelAgentBridge { ): Promise { const conn = this.ensureConnection(); await this.registerChannelLoopMcpServer(); - await conn.loadSession({ + await conn.unstable_resumeSession({ sessionId, cwd, mcpServers: [], diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index 1f3d585cce..e9ae574f6c 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -90,6 +90,10 @@ const { mockPreloadContentGenerator } = vi.hoisted(() => ({ mockPreloadContentGenerator: vi.fn().mockResolvedValue(undefined), })); +const { mockAddDaemonRequestAttribute } = vi.hoisted(() => ({ + mockAddDaemonRequestAttribute: vi.fn(), +})); + const { mockExtractDaemonTraceContext, mockSessionStartSpan, @@ -227,6 +231,7 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => ({ : undefined, ), initializeTelemetry: vi.fn().mockResolvedValue(undefined), + addDaemonRequestAttribute: mockAddDaemonRequestAttribute, preloadContentGenerator: mockPreloadContentGenerator, createDebugLogger: () => mockDebugLogger, extractDaemonTraceContext: mockExtractDaemonTraceContext, @@ -309,6 +314,12 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => ({ isReplayTurnStartType: ( await importOriginal() ).isReplayTurnStartType, + parseGoalSnapshotV2: ( + await importOriginal() + ).parseGoalSnapshotV2, + parseGoalStateCause: ( + await importOriginal() + ).parseGoalStateCause, findBoundaryAtOrBefore: ( await importOriginal() ).findBoundaryAtOrBefore, @@ -676,6 +687,12 @@ const { mockHistoryReplay } = vi.hoisted(() => ({ const { mockHistoryReplayPage } = vi.hoisted(() => ({ mockHistoryReplayPage: vi.fn(), })); +const { mockHistoryV2GoalBootstrap } = vi.hoisted(() => ({ + mockHistoryV2GoalBootstrap: vi.fn(), +})); +const { mockRenderPreparedGoalUpdate } = vi.hoisted(() => ({ + mockRenderPreparedGoalUpdate: vi.fn(), +})); type MockPendingToolCall = { callId: string; toolName: string; @@ -685,33 +702,46 @@ type MockPendingToolCall = { const { mockHistoryPendingToolCalls } = vi.hoisted(() => ({ mockHistoryPendingToolCalls: vi.fn((): MockPendingToolCall[] => []), })); -vi.mock('./session/history-replayer.js', () => ({ - HistoryReplayer: vi.fn().mockImplementation( - (context: { - cumulativeUsage: { - promptTokens: number; - cachedTokens: number; - candidateTokens: number; - apiTimeMs: number; - }; - }) => ({ - replay: (messages: unknown, gaps: unknown) => - mockHistoryReplay(context, messages, gaps), - replayPage: (messages: unknown, options: unknown) => - mockHistoryReplayPage(context, messages, options), - getPendingToolCalls: () => mockHistoryPendingToolCalls(), - getReplayState: () => ({ - v: 1, - pendingToolCalls: mockHistoryPendingToolCalls().map((call) => ({ - callId: call.callId, - toolName: call.toolName, - sourceRecordId: call.recordId, - ...(call.timestamp ? { sourceTimestamp: call.timestamp } : {}), - })), - cumulativeUsage: { ...context.cumulativeUsage }, +vi.mock('./session/history-replayer.js', () => { + const HistoryReplayer = Object.assign( + vi.fn().mockImplementation( + (context: { + cumulativeUsage: { + promptTokens: number; + cachedTokens: number; + candidateTokens: number; + apiTimeMs: number; + }; + }) => ({ + replay: ( + messages: unknown, + gaps: unknown, + options: Record, + ) => + Object.keys(options).length === 0 + ? mockHistoryReplay(context, messages, gaps) + : mockHistoryReplay(context, messages, gaps, options), + replayPage: (messages: unknown, options: unknown) => + mockHistoryReplayPage(context, messages, options), + getPendingToolCalls: () => mockHistoryPendingToolCalls(), + getReplayState: () => ({ + v: 1, + pendingToolCalls: mockHistoryPendingToolCalls().map((call) => ({ + callId: call.callId, + toolName: call.toolName, + sourceRecordId: call.recordId, + ...(call.timestamp ? { sourceTimestamp: call.timestamp } : {}), + })), + cumulativeUsage: { ...context.cumulativeUsage }, + }), }), - }), - ), + ), + { v2GoalBootstrap: mockHistoryV2GoalBootstrap }, + ); + return { HistoryReplayer }; +}); +vi.mock('./session/recovered-goal-update.js', () => ({ + renderPreparedGoalUpdate: mockRenderPreparedGoalUpdate, })); vi.mock('./runtimeOutputDirContext.js', () => ({ @@ -884,6 +914,8 @@ import { unregisterGoalHook, getActiveGoal, registerGoalHook, + findBoundaryAtOrBefore, + isReplayTurnStartType, startEventLoopLagMonitor, registerAcpEventLoopLagGauge, SESSION_ARTIFACT_PERSISTENCE_VERSION, @@ -3371,9 +3403,9 @@ describe('QwenAgent MCP SSE/HTTP support', () => { initialize: vi.fn().mockResolvedValue(undefined), shutdown: vi.fn().mockResolvedValue(undefined), closeSessionWriter: vi.fn().mockResolvedValue(undefined), - setSessionSource: vi.fn(), setSessionWriterReclaimPolicy: vi.fn(), setSessionWriterTakeoverPolicy: vi.fn(), + setSessionSource: vi.fn(), waitForMcpReady: vi.fn().mockResolvedValue(undefined), getModelsConfig: vi.fn().mockReturnValue({ getCurrentAuthType: vi.fn().mockReturnValue('api-key'), @@ -14673,10 +14705,15 @@ describe('QwenAgent unstable_listSessions cursor parsing', () => { describe('QwenAgent loadSession / unstable_resumeSession', () => { let capturedAgentFactory: | ((conn: { closed: Promise }) => { + initialize: (args: Record) => Promise; loadSession: (args: Record) => Promise; unstable_resumeSession: ( args: Record, ) => Promise; + beginManagedShutdown: () => { + configs: Config[]; + writerShutdown: Promise; + }; cancel: (args: Record) => Promise; }) | undefined; @@ -14690,7 +14727,8 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { replayHistory: ReturnType; primeTurnFromHistory: ReturnType; publishRecoveredGoalState: ReturnType; - renderRecoveredGoalUpdates: ReturnType; + primeRecoveredGoalPublication: ReturnType; + primeTurnState: ReturnType; cumulativeUsage: { promptTokens: number; cachedTokens: number; @@ -14719,6 +14757,19 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { beforeEach(() => { vi.clearAllMocks(); + mockHistoryV2GoalBootstrap.mockImplementation( + (goalState: unknown, goalCause: unknown) => + goalState && goalCause + ? { + goalStatus: { + kind: 'set', + condition: 'restored goal', + }, + goalState, + } + : undefined, + ); + mockRenderPreparedGoalUpdate.mockResolvedValue({ updates: [] }); mockExtractDaemonTraceContext.mockReturnValue(undefined); vi.mocked(Storage.getRuntimeBaseDir).mockReturnValue( '/tmp/qwen-runtime-test', @@ -14738,6 +14789,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { mockConfig = { initialize: vi.fn().mockResolvedValue(undefined), + shutdown: vi.fn().mockResolvedValue(undefined), waitForMcpReady: vi.fn().mockResolvedValue(undefined), getHookSystem: vi.fn().mockReturnValue(undefined), getDisableAllHooks: vi.fn().mockReturnValue(false), @@ -14747,6 +14799,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { getCurrentAuthType: vi.fn().mockReturnValue('api-key'), }), refreshAuth: vi.fn().mockResolvedValue(undefined), + closeSessionWriter: vi.fn().mockResolvedValue(undefined), getWorkspaceContext: vi.fn().mockReturnValue({}), getDebugMode: vi.fn().mockReturnValue(false), } as unknown as Config; @@ -14786,6 +14839,9 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { return { initialize: vi.fn().mockResolvedValue(undefined), shutdown: vi.fn().mockResolvedValue(undefined), + closeSessionWriter: vi.fn().mockResolvedValue(undefined), + setSessionWriterReclaimPolicy: vi.fn(), + setSessionWriterTakeoverPolicy: vi.fn(), setSessionSource: vi.fn(), waitForMcpReady: vi.fn().mockResolvedValue(undefined), getModelsConfig: vi.fn().mockReturnValue({ @@ -14820,10 +14876,13 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { getSessionRuntimeBaseDir: vi .fn() .mockReturnValue('/tmp/qwen-runtime-test'), + hydrateSessionRestoreFileHistory: vi.fn(), + finalizeSessionRestore: vi.fn(), loadPausedBackgroundAgents: vi.fn().mockResolvedValue([]), consumePendingRecoveredAgentsNotice: vi.fn().mockReturnValue(null), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), getSessionService: vi.fn(), + consumeSessionRestoreProjection: vi.fn(), // load path reads back the persisted conversation here and feeds // it to `session.replayHistory`. resume path doesn't read this. getResumedSessionData: vi @@ -14844,29 +14903,177 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { } as unknown as LoadedSettings; } + function selectRestoreReplay( + messages: Array>, + restoreOptions: { + replay: + | { kind: 'none' } + | { kind: 'all'; hideInheritedHistory: boolean } + | { kind: 'recent'; limit: number; hideInheritedHistory: boolean }; + }, + ) { + if (restoreOptions.replay.kind === 'none') return undefined; + const visible = selectVisibleHistoryRecords( + messages as never[], + restoreOptions.replay.hideInheritedHistory, + ) as unknown as Array>; + if ( + restoreOptions.replay.kind === 'all' || + visible.length <= restoreOptions.replay.limit + ) { + return { records: visible, gaps: [], hasMore: false }; + } + const limit = restoreOptions.replay.limit; + const isTurnStart = (record: Record) => + isReplayTurnStartType( + record['type'] as never, + record['subtype'] as never, + ); + const isPageStart = (record: Record) => + record['subtype'] !== 'realtime_message' && + (record['type'] === 'assistant' || isTurnStart(record)); + let start = visible.length - limit; + for (let i = start; i < visible.length; i++) { + if (isTurnStart(visible[i]!)) { + start = i; + break; + } + } + const aligned = findBoundaryAtOrBefore( + visible, + start, + Math.max(0, visible.length - 2 * limit), + isTurnStart, + ); + if (isTurnStart(visible[aligned]!)) start = aligned; + const startsOnOrphan = (() => { + for (let i = start; i < visible.length; i++) { + if (visible[i]?.['type'] !== 'tool_result') continue; + for (let owner = i - 1; owner >= start; owner--) { + if (isPageStart(visible[owner]!)) return false; + } + return true; + } + return false; + })(); + if (start > 0 && startsOnOrphan) { + const owner = findBoundaryAtOrBefore( + visible, + start, + Math.max(0, start - limit), + isPageStart, + ); + if (isPageStart(visible[owner]!)) start = owner; + } + return { + records: visible.slice(start), + gaps: [], + hasMore: start > 0, + }; + } + function bindRestoreMocks(opts: { sessionExists: boolean; resumedConversation?: { messages: unknown[] }; replayHistoryImpl?: (...args: unknown[]) => Promise; primeTurnFromHistoryImpl?: (...args: unknown[]) => unknown; recoveredGoalUpdates?: unknown[]; + recoveredGoalError?: Error; + recoveredGoalSendError?: Error; + primeTurnStateImpl?: (...args: unknown[]) => unknown; }) { const innerConfig = makeRestoreInnerConfig({ resumedConversation: opts.resumedConversation, }); + mockRenderPreparedGoalUpdate.mockImplementation(async () => { + if (opts.recoveredGoalError) throw opts.recoveredGoalError; + return { + updates: opts.recoveredGoalUpdates ?? [], + suppressedGoalId: 'hidden-goal', + }; + }); const loadSession = vi .fn() .mockImplementation(() => innerConfig.getResumedSessionData()); - innerConfig.getSessionService.mockReturnValue({ loadSession }); - vi.mocked(loadSettings).mockReturnValue(makeRestoreSettings()); - vi.mocked(loadCliConfig).mockResolvedValue( - innerConfig as unknown as Config, + const readRestoreProjection = vi.fn( + async ( + sessionId: string, + restoreOptions: Parameters[1], + ) => { + const data = innerConfig.getResumedSessionData(); + if (!data?.conversation) return undefined; + const messages = data.conversation.messages as Array< + Record + >; + return { + sessionId, + filePath: '/tmp/session.jsonl', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:00.000Z', + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { + lastCompletedUuid: + (messages.at(-1)?.['uuid'] as string | undefined) ?? '', + turnParentUuids: [], + }, + artifactSnapshot: data.artifactSnapshot, + goalRecords: messages, + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + replay: selectRestoreReplay(messages, restoreOptions), + }; + }, ); + const readLiveRestoreProjection = vi.fn( + async ( + sessionId: string, + restoreOptions: Parameters[1], + ) => { + const data = innerConfig.getResumedSessionData(); + if (!data?.conversation) return undefined; + const messages = data.conversation.messages as Array< + Record + >; + return { + sessionId, + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:00.000Z', + replay: selectRestoreReplay(messages, restoreOptions), + artifactSnapshot: data.artifactSnapshot, + }; + }, + ); + innerConfig.getSessionService.mockReturnValue({ + loadSession, + readRestoreProjection, + readLiveRestoreProjection, + }); + vi.mocked(loadSettings).mockReturnValue(makeRestoreSettings()); + vi.mocked(loadCliConfig).mockImplementation(async (...args: unknown[]) => { + const hostPolicy = args[9] as + | { + sessionRestore?: { + projectionSource: (sessionId: string) => Promise; + }; + } + | undefined; + const argv = args[1] as CliArgs; + const projection = argv.resume + ? await hostPolicy?.sessionRestore?.projectionSource(argv.resume) + : undefined; + innerConfig.consumeSessionRestoreProjection.mockReturnValue(projection); + return innerConfig as unknown as Config; + }); vi.mocked(SessionService).mockImplementation( () => ({ sessionExists: vi.fn().mockResolvedValue(opts.sessionExists), loadSession, + readRestoreProjection, + readLiveRestoreProjection, }) as unknown as InstanceType, ); vi.mocked(Session).mockImplementation(() => { @@ -14882,9 +15089,8 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { ), primeTurnFromHistory: vi.fn(opts.primeTurnFromHistoryImpl), publishRecoveredGoalState: vi.fn().mockResolvedValue(undefined), - renderRecoveredGoalUpdates: vi - .fn() - .mockResolvedValue(opts.recoveredGoalUpdates ?? []), + primeRecoveredGoalPublication: vi.fn(), + primeTurnState: vi.fn(opts.primeTurnStateImpl), cumulativeUsage: { promptTokens: 7, cachedTokens: 3, @@ -14901,7 +15107,9 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { waitForActiveTurnsToSettle: vi.fn().mockResolvedValue(undefined), cancelPendingPrompt: vi.fn().mockResolvedValue(undefined), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), - sendUpdate: vi.fn().mockResolvedValue(undefined), + sendUpdate: opts.recoveredGoalSendError + ? vi.fn().mockRejectedValue(opts.recoveredGoalSendError) + : vi.fn().mockResolvedValue(undefined), clearActiveTodoPlanRevision: vi.fn(), dispose: vi.fn(), }; @@ -14911,11 +15119,12 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { return innerConfig; } - async function spawnAgent() { + async function spawnAgent(privateParentCapability?: string) { const agentPromise = runAcpAgent( mockConfig, makeRestoreSettings(), mockArgv, + privateParentCapability ? { privateParentCapability } : undefined, ); await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); const agent = capturedAgentFactory!({ @@ -14923,6 +15132,14 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { return mockConnectionState.promise; }, }); + if (privateParentCapability) { + await agent.initialize({ + clientCapabilities: {}, + _meta: { + 'qwen-code/private-parent-capability': privateParentCapability, + }, + }); + } return { agent, agentPromise }; } @@ -14952,7 +15169,12 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { it.each(['load', 'resume'] as const)( 'profiles %s restore stages under the daemon trace context', async (action) => { - bindRestoreMocks({ sessionExists: true }); + bindRestoreMocks({ + sessionExists: true, + resumedConversation: { + messages: [{ role: 'user', parts: [{ text: 'restored' }] }], + }, + }); const parentContext = { trace: 'restore-parent' }; mockExtractDaemonTraceContext.mockReturnValue(parentContext); const { agent, agentPromise } = await spawnAgent(); @@ -14990,12 +15212,23 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { 'auth', 'file_system_setup', 'session_register', + 'runtime_initialize', 'response_build', + 'post_replay_services', ]) { expect( attributes[`qwen-code.daemon.session_restore.${stage}_ms`], ).toEqual(expect.any(Number)); } + if (action === 'load') { + expect( + attributes['qwen-code.daemon.session_restore.history_replay_ms'], + ).toEqual(expect.any(Number)); + } else { + expect( + attributes['qwen-code.daemon.session_restore.history_replay_ms'], + ).toBeUndefined(); + } // Mirror of the live-path test's `existence_check_ms` absence check. // Hoisting `live_restore` out of its `if (liveSession)` guard would make // cold loads report a stage that implies a live session existed, @@ -15116,11 +15349,16 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { const loadSession = vi .fn() .mockImplementation(() => innerConfig.getResumedSessionData()); + const projectionService = innerConfig.getSessionService(); vi.mocked(SessionService).mockImplementation( () => - ({ sessionExists, loadSession }) as unknown as InstanceType< - typeof SessionService - >, + ({ + sessionExists, + loadSession, + readRestoreProjection: projectionService.readRestoreProjection, + readLiveRestoreProjection: + projectionService.readLiveRestoreProjection, + }) as unknown as InstanceType, ); const { agent, agentPromise } = await spawnAgent(); vi.mocked(loadSettings).mockClear(); @@ -15302,13 +15540,9 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { // load semantic: history MUST be replayed so SSE subscribers see // the persisted turns. Second arg is the detected history gaps // (undefined here — this fixture session has an intact chain). - expect(lastSessionMock?.replayHistory).toHaveBeenCalledWith( - messages, - undefined, - ); + expect(lastSessionMock?.replayHistory).toHaveBeenCalledWith(messages, []); - const recording = lastSessionMock?.getConfig().getChatRecordingService(); - expect(recording?.rebuildTurnBoundaries).toHaveBeenCalledWith(messages); + expect(innerConfig.getSessionService().loadSession).not.toHaveBeenCalled(); expect(innerConfig.loadPausedBackgroundAgents).toHaveBeenCalledWith( 'persisted-1', ); @@ -15406,15 +15640,23 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { apiTimeMs: 0, }); expect(lastSessionMock?.replayHistory).not.toHaveBeenCalled(); - expect(lastSessionMock?.primeTurnFromHistory).toHaveBeenCalledWith( - messages, - ); + expect(lastSessionMock?.primeTurnState).toHaveBeenCalledWith(0, []); expect(mockHistoryReplay).toHaveBeenCalledTimes(1); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.partial_replay', + false, + ); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.partial_replay', + false, + ); - expect( - lastSessionMock!.primeTurnFromHistory.mock.invocationCallOrder[0], - ).toBeLessThan(mockHistoryReplay.mock.invocationCallOrder[0]!); expect(mockHistoryReplay.mock.invocationCallOrder[0]!).toBeLessThan( + lastSessionMock!.primeTurnState.mock.invocationCallOrder[0], + ); + expect( + lastSessionMock!.primeTurnState.mock.invocationCallOrder[0], + ).toBeLessThan( innerConfig.loadPausedBackgroundAgents.mock.invocationCallOrder[0]!, ); expect( @@ -15473,13 +15715,12 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { })) as { _meta?: Record }; const updates = response._meta?.['qwen.session.loadReplay']?.updates; - // Order is the fix: the authoritative card must be the LAST goal card. expect(updates).toEqual([{ ...replayUpdate, timestamp: 4242 }, goalUpdate]); - expect(lastSessionMock?.renderRecoveredGoalUpdates).toHaveBeenCalledWith( - messages, + expect(mockRenderPreparedGoalUpdate).toHaveBeenCalledOnce(); + expect(mockRenderPreparedGoalUpdate).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ replayedRecords: messages }), ); - // Rendered, never streamed — a streamed card would sort ahead of the - // envelope the bridge seeds after this call returns. expect(lastSessionMock?.sendUpdate).not.toHaveBeenCalled(); expect(lastSessionMock?.publishRecoveredGoalState).not.toHaveBeenCalled(); @@ -15487,28 +15728,110 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { await agentPromise; }); - // The other half of the same gap: resume also runs with - // `replayHistory: false`, so without an explicit publication the client is - // never told what the recovered goal is. Streaming is correct here — - // resume replays nothing for the card to sort against. - it('publishes the recovered Goal state on unstable_resumeSession', async () => { + it('streams the unrestorable Goal correction after cold history replay', async () => { const messages = [{ role: 'user', parts: [{ text: 'hi' }] }]; + const goalUpdate = { + sessionUpdate: 'agent_message_chunk', + content: { type: 'text', text: '' }, + _meta: { + goalStatus: { + kind: 'cleared', + condition: 'ship the thing', + }, + }, + }; bindRestoreMocks({ sessionExists: true, resumedConversation: { messages }, + recoveredGoalUpdates: [goalUpdate], }); const { agent, agentPromise } = await spawnAgent(); - await agent.unstable_resumeSession({ + await agent.loadSession({ cwd: '/tmp', sessionId: 'persisted-1', mcpServers: [], }); - expect(lastSessionMock?.publishRecoveredGoalState).toHaveBeenCalledWith( - messages, + expect(lastSessionMock?.replayHistory).toHaveBeenCalledWith(messages, []); + expect(mockRenderPreparedGoalUpdate).toHaveBeenCalledWith( + expect.any(Function), + expect.objectContaining({ replayedRecords: messages }), ); - expect(lastSessionMock?.replayHistory).not.toHaveBeenCalled(); + expect(lastSessionMock?.sendUpdate).toHaveBeenCalledWith(goalUpdate); + expect( + lastSessionMock!.replayHistory.mock.invocationCallOrder[0], + ).toBeLessThan(lastSessionMock!.sendUpdate.mock.invocationCallOrder[0]); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('keeps a cold load open when the Goal correction cannot be sent', async () => { + const messages = [{ role: 'user', parts: [{ text: 'hi' }] }]; + bindRestoreMocks({ + sessionExists: true, + resumedConversation: { messages }, + recoveredGoalUpdates: [{ sessionUpdate: 'agent_message_chunk' }], + recoveredGoalSendError: new Error('connection closed'), + }); + const { agent, agentPromise } = await spawnAgent(); + + await expect( + agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + }), + ).resolves.toMatchObject({ + modes: expect.anything(), + models: expect.anything(), + configOptions: expect.anything(), + }); + + expect(lastSessionMock?.sendUpdate).toHaveBeenCalledOnce(); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('does not publish a Session when managed shutdown starts during bulk replay preparation', async () => { + bindRestoreMocks({ + sessionExists: true, + resumedConversation: { + messages: [{ role: 'user', parts: [{ text: 'hi' }] }], + }, + }); + let markReplayStarted!: () => void; + const replayStarted = new Promise((resolve) => { + markReplayStarted = resolve; + }); + let releaseReplay!: () => void; + const replayGate = new Promise((resolve) => { + releaseReplay = resolve; + }); + mockHistoryReplay.mockReset(); + mockHistoryReplay.mockImplementation(async () => { + markReplayStarted(); + await replayGate; + }); + const { agent, agentPromise } = await spawnAgent('expected-capability'); + + const load = agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { 'qwen.session.loadReplayMode': 'bulk' }, + }); + await replayStarted; + await agent.beginManagedShutdown().writerShutdown; + releaseReplay(); + + await expect(load).rejects.toMatchObject({ + code: -32023, + data: { errorKind: 'session_writer_unavailable' }, + }); + expect(Session).not.toHaveBeenCalled(); mockConnectionState.resolve(); await agentPromise; @@ -15521,9 +15844,13 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { }); innerConfig.getApprovalMode.mockReturnValue('plan'); innerConfig.getSessionService.mockReturnValue({ - loadSession: vi - .fn() - .mockImplementation(() => innerConfig.getResumedSessionData()), + loadSession: vi.fn(), + readLiveRestoreProjection: vi.fn().mockResolvedValue({ + sessionId: 'persisted-1', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:00.000Z', + replay: { records: messages, gaps: [], hasMore: false }, + }), }); vi.mocked(loadSettings).mockReturnValue(makeRestoreSettings()); const replayUpdate = { @@ -15619,8 +15946,397 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { }; expect(response._meta?.['qwen.session.loadReplay']?.hasMore).toBe(true); - expect(lastSessionMock?.primeTurnFromHistory).toHaveBeenCalledWith( - messages, + expect(lastSessionMock?.primeTurnState).toHaveBeenCalledWith(0, []); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('bootstraps a recent replay from the latest goal state before the page', async () => { + const activeGoalState = { + v: 2, + activity: 'idle', + goal: { + goalId: 'goal-recent', + revision: 1, + objective: 'restore the visible goal card', + status: 'active', + evidenceCursor: { recordId: 'goal-state' }, + turnCount: 0, + activeTimeMs: 0, + createdAt: 1, + updatedAt: 1, + }, + }; + const goalCause = 'create'; + const goalRecord = { + uuid: 'goal-state', + parentUuid: null, + sessionId: 'persisted-1', + timestamp: '2026-07-16T00:00:00.000Z', + type: 'system', + subtype: 'goal_state', + cwd: '/tmp', + version: 'test', + systemPayload: { + v: 2, + cause: goalCause, + snapshot: activeGoalState, + }, + }; + const recentRecords = [ + { + uuid: 'u1', + parentUuid: 'goal-state', + sessionId: 'persisted-1', + timestamp: '2026-07-16T00:00:01.000Z', + type: 'user', + cwd: '/tmp', + version: 'test', + message: { role: 'user', parts: [] }, + }, + { + uuid: 'a1', + parentUuid: 'u1', + sessionId: 'persisted-1', + timestamp: '2026-07-16T00:00:02.000Z', + type: 'assistant', + cwd: '/tmp', + version: 'test', + message: { role: 'model', parts: [] }, + }, + ]; + const innerConfig = bindRestoreMocks({ + sessionExists: true, + resumedConversation: { messages: [goalRecord, ...recentRecords] }, + }); + const bootstrap = { + goalStatus: { + kind: 'set' as const, + condition: 'restore the visible goal card', + }, + goalState: activeGoalState, + }; + mockHistoryV2GoalBootstrap.mockReturnValueOnce(bootstrap); + const projection = { + sessionId: 'persisted-1', + filePath: '/tmp/session.jsonl', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:02.000Z', + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { + lastCompletedUuid: 'a1', + turnParentUuids: [], + }, + goalRecords: [goalRecord], + goalRecoverySourceUuid: 'goal-state', + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + replay: { + records: recentRecords, + gaps: [], + hasMore: true, + anchorRecordId: 'u1', + goalRecoverySourceUuid: 'goal-state', + goalBootstrapRecords: [goalRecord], + }, + }; + innerConfig + .getSessionService() + .readRestoreProjection.mockResolvedValue(projection); + const { agent, agentPromise } = await spawnAgent(); + + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { + 'qwen.session.loadReplayMode': 'bulk', + 'qwen.session.loadReplayPageSize': 2, + }, + }); + + expect(mockHistoryV2GoalBootstrap).toHaveBeenCalledWith( + activeGoalState, + goalCause, + ); + expect(mockHistoryReplay).toHaveBeenCalledWith( + expect.anything(), + recentRecords, + [], + { goalBootstrap: bootstrap }, + ); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('keeps a cold legacy Goal bootstrap inside visible history', async () => { + const recentRecords = [ + { uuid: 'u2', role: 'user', parts: [{ text: 'latest' }] }, + { uuid: 'a2', role: 'model', parts: [{ text: 'answer' }] }, + ]; + const visibleGoalRecord = { + uuid: 'visible-goal', + type: 'system', + subtype: 'slash_command', + systemPayload: { + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'visible goal', + iterations: 0, + setAt: 123, + }, + ], + }, + }; + const hiddenGoalRecord = { + uuid: 'hidden-goal', + type: 'system', + subtype: 'slash_command', + systemPayload: { + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'hidden inherited goal', + iterations: 0, + }, + ], + }, + }; + const innerConfig = bindRestoreMocks({ + sessionExists: true, + resumedConversation: { messages: recentRecords }, + }); + innerConfig.getSessionService().readRestoreProjection.mockResolvedValue({ + sessionId: 'persisted-1', + filePath: '/tmp/session.jsonl', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:02.000Z', + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { + lastCompletedUuid: 'a2', + turnParentUuids: [], + }, + goalRecords: [visibleGoalRecord, hiddenGoalRecord], + goalRecoverySourceUuid: 'hidden-goal', + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + replay: { + records: recentRecords, + gaps: [], + hasMore: true, + anchorRecordId: 'u2', + goalRecoverySourceUuid: 'visible-goal', + goalBootstrapRecords: [visibleGoalRecord], + }, + }); + let replayOptions: unknown; + mockHistoryReplay.mockImplementation( + async (_context, _history, _gaps, options) => { + replayOptions = options; + }, + ); + const { agent, agentPromise } = await spawnAgent(); + + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { + 'qwen.session.loadReplayMode': 'bulk', + 'qwen.session.loadReplayPageSize': 2, + 'qwen.session.loadReplayHideInherited': true, + }, + }); + + expect(replayOptions).toEqual({ + goalBootstrap: { + goalStatus: { + kind: 'set', + condition: 'visible goal', + iterations: 0, + setAt: 123, + }, + }, + }); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('keeps a streamed hide-inherited load on the visible Goal', async () => { + const visibleGoalRecord = { + uuid: 'visible-goal', + type: 'system', + subtype: 'slash_command', + systemPayload: { + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'visible goal', + iterations: 0, + }, + ], + }, + }; + const visibleRecords = [ + visibleGoalRecord, + { uuid: 'u2', role: 'user', parts: [{ text: 'latest' }] }, + { uuid: 'a2', role: 'model', parts: [{ text: 'answer' }] }, + ]; + const hiddenGoalRecord = { + uuid: 'hidden-goal', + type: 'system', + subtype: 'slash_command', + systemPayload: { + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'hidden inherited goal', + iterations: 0, + }, + ], + }, + }; + const innerConfig = bindRestoreMocks({ + sessionExists: true, + resumedConversation: { messages: visibleRecords }, + }); + innerConfig.getSessionService().readRestoreProjection.mockResolvedValue({ + sessionId: 'persisted-1', + filePath: '/tmp/session.jsonl', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:02.000Z', + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { + lastCompletedUuid: 'a2', + turnParentUuids: [], + }, + goalRecords: [visibleGoalRecord, hiddenGoalRecord], + goalRecoverySourceUuid: 'hidden-goal', + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + replay: { + records: visibleRecords, + gaps: [], + hasMore: false, + goalRecoverySourceUuid: 'visible-goal', + }, + }); + const { agent, agentPromise } = await spawnAgent(); + + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { 'qwen.session.loadReplayHideInherited': true }, + }); + + expect(lastSessionMock?.replayHistory).toHaveBeenCalledWith( + visibleRecords, + [], + ); + expect(mockRenderPreparedGoalUpdate).toHaveBeenCalledOnce(); + expect(lastSessionMock?.primeRecoveredGoalPublication).toHaveBeenCalledWith( + undefined, + 'hidden-goal', + ); + expect(lastSessionMock?.sendUpdate).not.toHaveBeenCalled(); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('fails a hide-inherited load closed when Goal filtering fails', async () => { + const innerConfig = bindRestoreMocks({ + sessionExists: true, + resumedConversation: { + messages: [{ uuid: 'u1', role: 'user', parts: [{ text: 'visible' }] }], + }, + recoveredGoalError: new Error('goal projection failed'), + }); + innerConfig.getSessionService().readRestoreProjection.mockResolvedValue({ + sessionId: 'persisted-1', + filePath: '/tmp/session.jsonl', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:02.000Z', + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { lastCompletedUuid: 'u1', turnParentUuids: [] }, + goalRecords: [], + goalRecoverySourceUuid: 'hidden-goal', + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + replay: { records: [], gaps: [], hasMore: false }, + }); + const { agent, agentPromise } = await spawnAgent(); + + await expect( + agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { 'qwen.session.loadReplayHideInherited': true }, + }), + ).rejects.toThrow('goal projection failed'); + + expect(lastSessionMock).toBeUndefined(); + expect(innerConfig.hydrateSessionRestoreFileHistory).not.toHaveBeenCalled(); + expect(innerConfig.finalizeSessionRestore).not.toHaveBeenCalled(); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('builds bulk Goal corrections before constructing or hydrating a Session', async () => { + const innerConfig = bindRestoreMocks({ + sessionExists: true, + resumedConversation: { + messages: [{ uuid: 'u1', role: 'user', parts: [{ text: 'visible' }] }], + }, + recoveredGoalUpdates: [ + { + sessionUpdate: 'agent_message_chunk', + content: { type: 'text', text: '' }, + }, + ], + }); + const { agent, agentPromise } = await spawnAgent(); + + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { 'qwen.session.loadReplayMode': 'bulk' }, + }); + + expect(mockRenderPreparedGoalUpdate).toHaveBeenCalledOnce(); + expect( + mockRenderPreparedGoalUpdate.mock.invocationCallOrder[0], + ).toBeLessThan(vi.mocked(Session).mock.invocationCallOrder[0]!); + expect(vi.mocked(Session).mock.invocationCallOrder[0]).toBeLessThan( + innerConfig.hydrateSessionRestoreFileHistory.mock.invocationCallOrder[0]!, ); mockConnectionState.resolve(); @@ -15675,9 +16391,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { expect( response._meta?.['qwen.session.loadReplay']?.hasMore, ).toBeUndefined(); - expect(lastSessionMock?.primeTurnFromHistory).toHaveBeenCalledWith( - messages, - ); + expect(lastSessionMock?.primeTurnState).toHaveBeenCalledWith(0, []); mockConnectionState.resolve(); await agentPromise; @@ -15925,7 +16639,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { resumedConversation: { messages, }, - primeTurnFromHistoryImpl: () => { + primeTurnStateImpl: () => { primeCalls++; if (primeCalls === 1) { throw new Error('prime boom'); @@ -15977,11 +16691,12 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { resumedConversation: { messages: [{ role: 'user', parts: [{ text: 'hi' }] }], }, - primeTurnFromHistoryImpl: () => { + primeTurnStateImpl: () => { throw setupError; }, }); const recording = innerConfig.getChatRecordingService(); + innerConfig.shutdown.mockImplementation(async () => recording.close()); recording.close.mockRejectedValue(cleanupError); recording.hasWriteOwnership.mockReturnValue(true); mockHistoryReplay.mockReset(); @@ -16353,7 +17068,10 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { }); expect(Session).toHaveBeenCalledTimes(1); - expect(innerConfig.getSessionService().loadSession).toHaveBeenCalledOnce(); + expect(innerConfig.getSessionService().loadSession).not.toHaveBeenCalled(); + expect( + innerConfig.getSessionService().readLiveRestoreProjection, + ).toHaveBeenCalledOnce(); mockConnectionState.resolve(); await agentPromise; @@ -16414,6 +17132,235 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { await agentPromise; }); + it('bootstraps a live recent page from an older legacy Goal record', async () => { + const recentRecords = [ + { uuid: 'u2', role: 'user', parts: [{ text: 'latest' }] }, + { uuid: 'a2', role: 'model', parts: [{ text: 'answer' }] }, + ]; + bindRestoreMocks({ + sessionExists: true, + resumedConversation: { messages: recentRecords }, + }); + const { agent, agentPromise } = await spawnAgent(); + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + }); + const firstSession = lastSessionMock!; + const legacyGoalRecord = { + uuid: 'goal', + type: 'system', + subtype: 'slash_command', + systemPayload: { + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'keep the live goal visible', + iterations: 0, + }, + ], + }, + }; + firstSession + .getConfig() + .getSessionService() + .readLiveRestoreProjection.mockResolvedValue({ + sessionId: 'persisted-1', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:02.000Z', + replay: { + records: recentRecords, + gaps: [], + hasMore: true, + anchorRecordId: 'u2', + }, + goalRecords: [legacyGoalRecord], + goalRecoverySourceUuid: 'goal', + }); + let replayOptions: unknown; + mockHistoryReplay.mockImplementation( + async (_context, _history, _gaps, options) => { + replayOptions = options; + }, + ); + + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { + 'qwen.session.loadReplayMode': 'bulk', + 'qwen.session.loadReplayPageSize': 2, + }, + }); + + expect(replayOptions).toEqual({ + goalBootstrap: { + goalStatus: { + kind: 'set', + condition: 'keep the live goal visible', + iterations: 0, + }, + }, + }); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('does not fall back to a legacy Goal behind the determining malformed v2 record', async () => { + const recentRecords = [ + { uuid: 'u2', role: 'user', parts: [{ text: 'latest' }] }, + { uuid: 'a2', role: 'model', parts: [{ text: 'answer' }] }, + ]; + bindRestoreMocks({ + sessionExists: true, + resumedConversation: { messages: recentRecords }, + }); + const { agent, agentPromise } = await spawnAgent(); + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + }); + const firstSession = lastSessionMock!; + firstSession + .getConfig() + .getSessionService() + .readLiveRestoreProjection.mockResolvedValue({ + sessionId: 'persisted-1', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:02.000Z', + replay: { + records: recentRecords, + gaps: [], + hasMore: true, + anchorRecordId: 'u2', + }, + goalRecords: [ + { + uuid: 'legacy-goal', + type: 'system', + subtype: 'slash_command', + systemPayload: { + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'stale legacy goal', + iterations: 0, + }, + ], + }, + }, + { + uuid: 'malformed-goal', + type: 'system', + subtype: 'goal_state', + systemPayload: null, + }, + ], + goalRecoverySourceUuid: 'malformed-goal', + }); + let replayOptions: unknown; + mockHistoryReplay.mockImplementation( + async (_context, _history, _gaps, options) => { + replayOptions = options; + }, + ); + + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { + 'qwen.session.loadReplayMode': 'bulk', + 'qwen.session.loadReplayPageSize': 2, + }, + }); + + expect(replayOptions).toBeUndefined(); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('does not bootstrap a filtered live Goal candidate without a visible source', async () => { + const recentRecords = [ + { uuid: 'u2', role: 'user', parts: [{ text: 'latest' }] }, + { uuid: 'a2', role: 'model', parts: [{ text: 'answer' }] }, + ]; + bindRestoreMocks({ + sessionExists: true, + resumedConversation: { messages: recentRecords }, + }); + const { agent, agentPromise } = await spawnAgent(); + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + }); + const firstSession = lastSessionMock!; + firstSession + .getConfig() + .getSessionService() + .readLiveRestoreProjection.mockResolvedValue({ + sessionId: 'persisted-1', + startTime: '2026-07-16T00:00:00.000Z', + lastUpdated: '2026-07-16T00:00:02.000Z', + replay: { + records: recentRecords, + gaps: [], + hasMore: true, + anchorRecordId: 'u2', + }, + goalRecords: [ + { + uuid: 'filtered-goal', + type: 'system', + subtype: 'slash_command', + systemPayload: { + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'must remain hidden', + iterations: 0, + }, + ], + }, + }, + ], + }); + let replayOptions: unknown; + mockHistoryReplay.mockImplementation( + async (_context, _history, _gaps, options) => { + replayOptions = options; + }, + ); + + await agent.loadSession({ + cwd: '/tmp', + sessionId: 'persisted-1', + mcpServers: [], + _meta: { + 'qwen.session.loadReplayMode': 'bulk', + 'qwen.session.loadReplayPageSize': 2, + 'qwen.session.loadReplayHideInherited': true, + }, + }); + + expect(replayOptions).toBeUndefined(); + + mockConnectionState.resolve(); + await agentPromise; + }); + it('live resume refreshes artifacts without replaying UI history', async () => { const messages = [{ role: 'user', parts: [{ text: 'first' }] }]; bindRestoreMocks({ @@ -16497,7 +17444,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { firstSession .getConfig() .getSessionService() - .loadSession.mockResolvedValue(undefined); + .readLiveRestoreProjection.mockResolvedValue(undefined); await expect( agent.loadSession({ @@ -16564,8 +17511,7 @@ describe('QwenAgent loadSession / unstable_resumeSession', () => { // the SSE stream stays clean for clients that already have the // history rendered. expect(lastSessionMock?.replayHistory).not.toHaveBeenCalled(); - const recording = lastSessionMock?.getConfig().getChatRecordingService(); - expect(recording?.rebuildTurnBoundaries).toHaveBeenCalledWith(messages); + expect(innerConfig.getSessionService().loadSession).not.toHaveBeenCalled(); expect( innerConfig.loadPausedBackgroundAgents.mock.invocationCallOrder[0]!, ).toBeLessThan( @@ -16640,6 +17586,7 @@ describe('QwenAgent extMethod runtime MCP add/remove (T2.8)', () => { mockConfig = { initialize: vi.fn().mockResolvedValue(undefined), + shutdown: vi.fn().mockResolvedValue(undefined), waitForMcpReady: vi.fn().mockResolvedValue(undefined), getHookSystem: vi.fn().mockReturnValue(undefined), getDisableAllHooks: vi.fn().mockReturnValue(false), diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index 2b1e7dfa46..60b514ee27 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -72,11 +72,11 @@ import { SessionTranscriptSnapshotUnavailableError, SessionTranscriptTooLargeError, encodeSessionTranscriptCursor, - findBoundaryAtOrBefore, - isReplayTurnStartType, subagentGenerator, redactUrlCredentials, computeUniqueBranchTitle, + parseGoalSnapshotV2, + parseGoalStateCause, ToolNames, FORK_SUBAGENT_TYPE, runManagedAutoMemoryDream, @@ -92,6 +92,7 @@ import { refreshMemoryInstruction, applyReasoningEffort, REASONING_EFFORT_TIERS, + addDaemonRequestAttribute, extractDaemonTraceContext, withDaemonSpan, emptyGoalSnapshot, @@ -114,7 +115,10 @@ import { type ProviderSetupInputs, type ReasoningEffort, type ResumedSessionData, + type SelectiveSessionRestoreOptions, type SendSdkMcpMessage, + type SessionLiveRestoreProjection, + type SessionRestoreProjection, type SessionArtifactEventRecordPayload, type SessionArtifactSnapshotRecordPayload, type WorkspaceRememberContextMode, @@ -154,8 +158,8 @@ import type { ResumeSessionResponse, SessionConfigOption, SessionInfo, - SessionModeState, SessionUpdate, + SessionModeState, SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, SetSessionModelRequest, @@ -229,6 +233,8 @@ import { isInactiveExtensionSkill, } from './extension-skills.js'; import { Session, buildAvailableCommandsSnapshot } from './session/Session.js'; +import { HistoryReplayer } from './session/history-replayer.js'; +import { renderPreparedGoalUpdate } from './session/recovered-goal-update.js'; import { ActiveWorkReporter } from './active-work-reporter.js'; import { getModelConfiguration, @@ -239,6 +245,7 @@ import { collectHistoryReplayUpdates, copyCumulativeUsage, createReplayCumulativeUsage, + HistoryReplayLimitError, replayTranscriptRecordPage, } from './session/history-replay-page.js'; import { @@ -338,6 +345,8 @@ import { DAEMON_PROMPT_DISPLAY_TEXT_META_KEY, LOAD_REPLAY_BULK_MODE, LOAD_REPLAY_HIDE_INHERITED_META_KEY, + LOAD_REPLAY_MAX_BYTES, + LOAD_REPLAY_MAX_UPDATES, LOAD_REPLAY_META_KEY, LOAD_REPLAY_MODE_META_KEY, LOAD_REPLAY_PAGE_SIZE_META_KEY, @@ -365,6 +374,10 @@ import { } from '../ui/commands/contextCommand.js'; import type { HistoryItemContextUsage } from '../ui/types.js'; import { fireSessionDeleteHook } from '../hooks/session-delete-hook.js'; +import { + collectGoalStatusItemsFromRecords, + findGoalToRestore, +} from '../ui/utils/restoreGoal.js'; import { writeStderrLineSafe } from '../utils/stdioHelpers.js'; import { executeGeneration, @@ -402,7 +415,10 @@ type AcpSessionProfileStage = | 'auth' | 'file_system_setup' | 'session_register' - | 'response_build'; + | 'runtime_initialize' + | 'response_build' + | 'history_replay' + | 'post_replay_services'; interface AcpSessionProfileSpan { setAttribute(name: string, value: string | number | boolean): unknown; @@ -706,6 +722,196 @@ function shouldHideInheritedHistory(params: LoadSessionRequest): boolean { return meta?.[LOAD_REPLAY_HIDE_INHERITED_META_KEY] === true; } +function loadRestoreOptions( + params: LoadSessionRequest, +): SelectiveSessionRestoreOptions { + const hideInheritedHistory = shouldHideInheritedHistory(params); + if (!isBulkLoadReplayRequest(params)) { + return { replay: { kind: 'all', hideInheritedHistory } }; + } + const limit = getLoadReplayPageSize(params); + return limit === undefined + ? { replay: { kind: 'all', hideInheritedHistory } } + : { replay: { kind: 'recent', limit, hideInheritedHistory } }; +} + +const RESUME_RESTORE_OPTIONS: SelectiveSessionRestoreOptions = { + replay: { kind: 'none' }, +}; + +function mapSessionRestoreRequestError( + error: unknown, + sessionId: string, +): unknown { + const mappedWriterError = mapSessionWriterRequestError(error); + if (mappedWriterError !== error) return mappedWriterError; + if (error instanceof SessionTranscriptSnapshotUnavailableError) { + return new RequestError(-32010, error.message, { + errorKind: 'transcript_snapshot_unavailable', + sessionId, + }); + } + if (error instanceof SessionTranscriptTooLargeError) { + return new RequestError(-32011, error.message, { + errorKind: 'transcript_too_large', + sessionId, + snapshotSize: error.snapshotSize, + maxBytes: error.maxBytes, + }); + } + if (error instanceof SessionTranscriptPageTooLargeError) { + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.envelope_limit_reason', + 'bytes', + ); + return new RequestError(-32012, error.message, { + errorKind: 'transcript_page_too_large', + sessionId, + pageBytes: error.pageBytes, + maxBytes: error.maxBytes, + }); + } + if (error instanceof HistoryReplayLimitError) { + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.envelope_limit_reason', + error.reason, + ); + return new RequestError(-32012, error.message, { + errorKind: 'transcript_page_too_large', + sessionId, + reason: error.reason, + observed: error.observed, + limit: error.limit, + }); + } + return error; +} + +function validateLoadReplayEnvelope( + sessionId: string, + envelope: BridgeLoadReplayEnvelope, + enforceLimits: boolean, +): void { + if (!enforceLimits) return; + if (envelope.updates.length > LOAD_REPLAY_MAX_UPDATES) { + throw new HistoryReplayLimitError( + sessionId, + 'updates', + envelope.updates.length, + LOAD_REPLAY_MAX_UPDATES, + ); + } + const pageBytes = Buffer.byteLength(JSON.stringify(envelope), 'utf8'); + if (pageBytes > LOAD_REPLAY_MAX_BYTES) { + throw new HistoryReplayLimitError( + sessionId, + 'bytes', + pageBytes, + LOAD_REPLAY_MAX_BYTES, + ); + } +} + +function replayGoalBootstrap( + projection: + | SessionRestoreProjection + | SessionLiveRestoreProjection + | undefined, +): ReturnType { + if (projection && !('runtime' in projection)) { + const sourceUuid = projection.goalRecoverySourceUuid; + if (!sourceUuid) return undefined; + if ( + projection.replay?.records.some((record) => record.uuid === sourceUuid) + ) { + return undefined; + } + const replay = projection.replay?.replay; + if (isObjectRecord(replay)) { + const bootstrap = HistoryReplayer.v2GoalBootstrap( + replay['goalState'], + replay['goalCause'], + ); + if (bootstrap) return bootstrap; + } + const sourceRecord = projection.goalRecords?.find( + (record) => record.uuid === sourceUuid, + ); + if (sourceRecord?.subtype === 'goal_state') { + const payload = isObjectRecord(sourceRecord.systemPayload) + ? sourceRecord.systemPayload + : undefined; + return HistoryReplayer.v2GoalBootstrap( + payload?.['snapshot'], + payload?.['cause'], + ); + } + const active = findGoalToRestore( + collectGoalStatusItemsFromRecords(projection.goalRecords ?? []), + ); + return active + ? { + goalStatus: { + kind: active.iterations > 0 ? 'checking' : 'set', + condition: active.condition, + iterations: active.iterations, + ...(active.setAt !== undefined ? { setAt: active.setAt } : {}), + }, + } + : undefined; + } + const sourceUuid = projection?.replay?.goalRecoverySourceUuid; + if (!projection?.replay || !sourceUuid) return undefined; + if (projection.replay.records.some((record) => record.uuid === sourceUuid)) { + return undefined; + } + const goalBootstrapRecords = projection.replay.goalBootstrapRecords ?? []; + const sourceRecord = goalBootstrapRecords.find( + (record) => record.uuid === sourceUuid, + ); + if (sourceRecord?.subtype === 'goal_state') { + const payload = isObjectRecord(sourceRecord.systemPayload) + ? sourceRecord.systemPayload + : undefined; + return HistoryReplayer.v2GoalBootstrap( + payload?.['snapshot'], + payload?.['cause'], + ); + } + const active = findGoalToRestore( + collectGoalStatusItemsFromRecords(goalBootstrapRecords), + ); + if (!active) return undefined; + return { + goalStatus: { + kind: active.iterations > 0 ? 'checking' : 'set', + condition: active.condition, + iterations: active.iterations, + ...(active.setAt !== undefined ? { setAt: active.setAt } : {}), + }, + }; +} + +function replayInitialGoalState( + projection: SessionRestoreProjection | undefined, +): { + initialGoalState?: NonNullable< + Parameters[2] + >['initialGoalState']; + initialGoalCause?: NonNullable< + Parameters[2] + >['initialGoalCause']; +} { + const replay = projection?.replay?.replay; + if (!isObjectRecord(replay)) return {}; + const initialGoalState = parseGoalSnapshotV2(replay['goalState']); + const initialGoalCause = parseGoalStateCause(replay['goalCause']); + return { + ...(initialGoalState ? { initialGoalState } : {}), + ...(initialGoalCause ? { initialGoalCause } : {}), + }; +} + export function selectVisibleHistoryRecords( records: ChatRecord[], hideInheritedHistory: boolean, @@ -772,91 +978,6 @@ function getLoadReplayPageSize(params: LoadSessionRequest): number | undefined { return value as number; } -function isHistoryTurnStart(record: ChatRecord): boolean { - return isReplayTurnStartType(record.type, record.subtype); -} - -// A bulk page can safely start at a turn start or at the assistant record -// owning any following tool results, mirroring the core reader's page-start -// rule. Realtime records interleave at wall-clock time and own no tool -// results, so the pair walk passes through them instead of splitting them. -function isHistoryPageStart(record: ChatRecord): boolean { - return ( - record.subtype !== 'realtime_message' && - (record.type === 'assistant' || isHistoryTurnStart(record)) - ); -} - -// True when the first tool_result in records[start, end) lost its owning -// call below `start`, i.e. the selection begins mid-pair. Only the first -// result needs checking: later results belong to calls at or after it, all -// inside the page once the first pair is whole. -function historyOrphansToolResult( - records: ChatRecord[], - start: number, - end: number, -): boolean { - for (let i = start; i < end; i++) { - if (records[i]!.type !== 'tool_result') continue; - for (let owner = i - 1; owner >= start; owner--) { - if (isHistoryPageStart(records[owner]!)) return false; - } - return true; - } - return false; -} - -function selectRecentHistoryRecords( - records: ChatRecord[], - pageSize: number | undefined, -): { records: ChatRecord[]; hasMore: boolean } { - if (pageSize === undefined || records.length <= pageSize) { - return { records, hasMore: false }; - } - let start = records.length - pageSize; - for (let i = start; i < records.length; i++) { - if (isHistoryTurnStart(records[i]!)) { - start = i; - break; - } - } - // Turn-boundary alignment may expand the page past the requested window, - // but never without bound: a history dominated by a single long in-flight - // turn would otherwise replay the WHOLE history in one payload and report - // hasMore=false, leaving the client unable to page backward. Allow at - // most one extra window of expansion, and only when it reaches a real - // turn start; otherwise keep the requested window so pages inside a long - // turn stay bounded and chainable. - const expansionFloor = Math.max(0, records.length - 2 * pageSize); - const aligned = findBoundaryAtOrBefore( - records, - start, - expansionFloor, - isHistoryTurnStart, - ); - if (isHistoryTurnStart(records[aligned]!)) { - start = aligned; - } - // Backward replay renders each page independently, so a page that starts - // mid-pair (on a tool_result whose owning call lies below) would show the - // completed call as failed and its result as an orphan block. When the - // bounded selection starts on an orphaned tool_result, extend down to the - // owning record within one further window, mirroring the core reader. - if (start > 0 && historyOrphansToolResult(records, start, records.length)) { - const pairFloor = Math.max(0, start - pageSize); - const owner = findBoundaryAtOrBefore( - records, - start, - pairFloor, - isHistoryPageStart, - ); - if (isHistoryPageStart(records[owner]!)) { - start = owner; - } - } - return { records: records.slice(start), hasMore: start > 0 }; -} - function createHiddenWorkspaceMemoryConfig(config: Config): Config { return new Proxy(config, { get(target, prop) { @@ -4221,7 +4342,11 @@ class QwenAgent implements Agent { private async withLiveSessionRestore( sessionId: string, session: Session, - operation: (config: Config, data: ResumedSessionData) => Promise, + options: SelectiveSessionRestoreOptions, + operation: ( + config: Config, + projection: SessionLiveRestoreProjection | undefined, + ) => Promise, ): Promise { await session.assertCanStartTurn(); const config = session.getConfig(); @@ -4233,15 +4358,17 @@ class QwenAgent implements Agent { 'restore', ); const recorder = config.getChatRecordingService(); - const loadAuthoritative = () => - config.getSessionService().loadSession(sessionId); - const data = recorder - ? await recorder.runWithWriteBarrier(loadAuthoritative) - : await loadAuthoritative(); - if (!data) throw new SessionWriterUnavailableError(); - return await operation(config, data); + const readProjection = () => + config + .getSessionService() + .readLiveRestoreProjection(sessionId, options); + const projection = recorder + ? await recorder.runWithWriteBarrier(readProjection) + : await readProjection(); + if (!projection) throw new SessionWriterUnavailableError(); + return await operation(config, projection); } catch (error) { - throw mapSessionWriterRequestError(error); + throw mapSessionRestoreRequestError(error, sessionId); } finally { releaseGate(); } @@ -4273,6 +4400,7 @@ class QwenAgent implements Agent { private async cleanupAfterRequestFailure( error: unknown, cleanup: () => Promise, + sessionId?: string, ): Promise { try { await cleanup(); @@ -4281,7 +4409,9 @@ class QwenAgent implements Agent { `Session cleanup failed while preserving the original request error: ${cleanupError instanceof Error ? cleanupError.message : String(cleanupError)}`, ); } - throw mapSessionWriterRequestError(error); + throw sessionId + ? mapSessionRestoreRequestError(error, sessionId) + : mapSessionWriterRequestError(error); } private pendingConfigCleanupKey( @@ -4868,6 +4998,7 @@ class QwenAgent implements Agent { ): Promise { let sessionId = initialSessionId; const sessionSource = getSessionSource(params); + const restoreOptions = loadRestoreOptions(params); const liveSession = this.sessions.get(sessionId); if (liveSession) { const settings = profiler.timeSync('settings_load', () => @@ -4879,7 +5010,8 @@ class QwenAgent implements Agent { return this.withLiveSessionRestore( sessionId, liveSession, - async (config, sessionData) => { + restoreOptions, + async (config, projection) => { const response = profiler.timeSync( 'response_build', () => @@ -4887,33 +5019,39 @@ class QwenAgent implements Agent { modes: this.buildModesData(config), models: this.buildAvailableModels(config), configOptions: this.buildConfigOptions(config), - ...(sessionData.artifactSnapshot - ? { artifactSnapshot: sessionData.artifactSnapshot } + ...(projection?.artifactSnapshot + ? { artifactSnapshot: projection.artifactSnapshot } : {}), }) as LoadSessionResponse, ); - const records = sessionData.conversation.messages; - const visibleRecords = selectVisibleHistoryRecords( - records, - shouldHideInheritedHistory(params), - ); - if (visibleRecords.length === 0) return response; + const replayPage = projection?.replay; + if (!replayPage || replayPage.records.length === 0) return response; const bulkReplay = isBulkLoadReplayRequest(params); - const replayPage = bulkReplay - ? selectRecentHistoryRecords( - visibleRecords, - getLoadReplayPageSize(params), - ) - : { records: visibleRecords, hasMore: false }; - const replay = await collectHistoryReplayUpdates({ - sessionId, - config, - records: replayPage.records, - gaps: sessionData.historyGaps, - cumulativeUsage: createReplayCumulativeUsage(), - logger: debugLogger, - }); + const replay = await profiler.time('history_replay', () => + collectHistoryReplayUpdates({ + sessionId, + config, + records: replayPage.records, + gaps: replayPage.gaps, + cumulativeUsage: createReplayCumulativeUsage(), + replayState: replayPage.replay, + goalBootstrap: replayGoalBootstrap(projection), + ...(restoreOptions.replay.kind === 'recent' + ? { + limits: { + maxBytes: LOAD_REPLAY_MAX_BYTES, + maxUpdates: LOAD_REPLAY_MAX_UPDATES, + }, + } + : {}), + logger: debugLogger, + }), + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.partial_replay', + replay.replayError !== undefined, + ); if (!bulkReplay) { try { for (const update of replay.updates) { @@ -4934,20 +5072,29 @@ class QwenAgent implements Agent { return response; } + const envelope: BridgeLoadReplayEnvelope = { + v: LOAD_REPLAY_VERSION, + updates: replay.updates, + ...(replayPage.anchorRecordId + ? { anchorRecordId: replayPage.anchorRecordId } + : {}), + ...(replay.replayError !== undefined + ? { + partial: true, + replayError: replay.replayError, + } + : {}), + ...(replayPage.hasMore ? { hasMore: true } : {}), + }; + validateLoadReplayEnvelope( + sessionId, + envelope, + restoreOptions.replay.kind === 'recent', + ); return { ...response, _meta: { - [LOAD_REPLAY_META_KEY]: { - v: LOAD_REPLAY_VERSION, - updates: replay.updates, - ...(replay.replayError !== undefined - ? { - partial: true as const, - replayError: replay.replayError, - } - : {}), - ...(replayPage.hasMore ? { hasMore: true as const } : {}), - }, + [LOAD_REPLAY_META_KEY]: envelope, }, }; }, @@ -4993,129 +5140,222 @@ class QwenAgent implements Agent { sessionSource, sessionId, true, + {}, + undefined, + restoreOptions, ), ); - const sessionData = config.getResumedSessionData(); + const projection = config.consumeSessionRestoreProjection?.(); + const suppressRecoveredGoalPresentation = + projection?.runtime.goalRecoverySourceUuid !== undefined && + projection.runtime.goalRecoverySourceUuid !== + projection.replay?.goalRecoverySourceUuid; const bulkReplay = isBulkLoadReplayRequest(params); - const replayPageSize = bulkReplay - ? getLoadReplayPageSize(params) - : undefined; let replayEnvelope: BridgeLoadReplayEnvelope | undefined; + const replayUsage = createReplayCumulativeUsage(); + let recoveredGoalPublicationKey: string | undefined; + let suppressedRecoveredGoalId: string | undefined; + let streamGoalUpdates: SessionUpdate[] = []; + let response: LoadSessionResponse | undefined; + const buildResponse = () => + profiler.timeSync('response_build', () => ({ + modes: this.buildModesData(config), + models: this.buildAvailableModels(config), + configOptions: this.buildConfigOptions(config), + ...(projection?.runtime.artifactSnapshot + ? { artifactSnapshot: projection.runtime.artifactSnapshot } + : {}), + ...(replayEnvelope + ? { + _meta: { + [LOAD_REPLAY_META_KEY]: replayEnvelope, + }, + } + : {}), + })) as LoadSessionResponse; try { await profiler.time('auth', () => this.ensureAuthenticated(config)); profiler.timeSync('file_system_setup', () => this.setupFileSystem(config), ); await profiler.time('session_register', () => - this.createAndStoreSession(config, settings, sessionData, { + this.createAndStoreSession(config, settings, undefined, { enableLiveScreenContext: isCompatibleLiveSessionSource( sessionSource ?? {}, ), - ...(bulkReplay ? { replayHistory: false } : {}), - beforeStartPostReplayServices: async (createdSession) => { - if (bulkReplay) { - const records = sessionData?.conversation.messages; - let replayUpdates: SessionUpdate[] = []; - if (records) { - createdSession.primeTurnFromHistory(records); - const visibleRecords = selectVisibleHistoryRecords( - records, - shouldHideInheritedHistory(params), - ); - const replayPage = selectRecentHistoryRecords( - visibleRecords, - replayPageSize, - ); - const replayUsage = createReplayCumulativeUsage(); - const replay = await collectHistoryReplayUpdates({ + replayHistory: false, + prepareBeforeSessionCreate: async () => { + if (bulkReplay && projection?.replay) { + const replay = await profiler.time('history_replay', () => + collectHistoryReplayUpdates({ sessionId, config, - records: replayPage.records, - gaps: sessionData?.historyGaps, + records: projection.replay!.records, + gaps: projection.replay!.gaps, cumulativeUsage: replayUsage, + replayState: projection.replay!.replay, + goalBootstrap: replayGoalBootstrap(projection), + ...(restoreOptions.replay.kind === 'recent' + ? { + limits: { + maxBytes: LOAD_REPLAY_MAX_BYTES, + maxUpdates: LOAD_REPLAY_MAX_UPDATES, + }, + } + : {}), logger: debugLogger, - }); - replayUpdates = replay.updates; - copyCumulativeUsage( - createdSession.cumulativeUsage, - replayUsage, - ); - // Strictly after the replay page, mirroring the streamed - // path in `createAndStoreSession`: the page re-emits the - // pre-migration `set` card, so the authoritative state has - // to be the newest goal card or the client keeps showing a - // phantom running goal. It rides *inside* the envelope - // because this path hands its updates to the client rather - // than streaming them — see - // `Session.renderRecoveredGoalUpdates`. Never fatal: a - // session that cannot publish its goal state must still - // open. - try { - replayUpdates = replayUpdates.concat( - await createdSession.renderRecoveredGoalUpdates(records), - ); - } catch (error) { - debugLogger.debug( - `Failed to render recovered Goal state: ${ - error instanceof Error ? error.message : String(error) - }`, - ); - } - if (replay.replayError !== undefined) { - replayEnvelope = { - v: LOAD_REPLAY_VERSION, - updates: replayUpdates, - partial: true, - replayError: replay.replayError, - ...(replayPage.hasMore ? { hasMore: true } : {}), - }; - } - replayEnvelope ??= { - v: LOAD_REPLAY_VERSION, - updates: replayUpdates, - ...(replayPage.hasMore ? { hasMore: true } : {}), - }; - } + }), + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.partial_replay', + replay.replayError !== undefined, + ); + replayEnvelope = { + v: LOAD_REPLAY_VERSION, + updates: replay.updates, + ...(projection.replay.anchorRecordId + ? { anchorRecordId: projection.replay.anchorRecordId } + : {}), + ...(replay.replayError !== undefined + ? { + partial: true, + replayError: replay.replayError, + } + : {}), + ...(projection.replay.hasMore ? { hasMore: true } : {}), + }; + validateLoadReplayEnvelope( + sessionId, + replayEnvelope, + restoreOptions.replay.kind === 'recent', + ); + } + const goalBootstrap = replayGoalBootstrap(projection); + const rendered = await renderPreparedGoalUpdate( + () => config.getGoalRuntimePrepared(), + { + ...(projection?.replay?.records + ? { replayedRecords: projection.replay.records } + : {}), + ...(suppressRecoveredGoalPresentation + ? { hideRuntimeGoal: true } + : {}), + ...(goalBootstrap ? { bootstrap: goalBootstrap } : {}), + }, + ).catch((error) => { + if (suppressRecoveredGoalPresentation) throw error; + debugLogger.debug( + `Failed to render recovered Goal state: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return { + publicationKey: undefined, + suppressedGoalId: undefined, + updates: [], + }; + }); + recoveredGoalPublicationKey = rendered.publicationKey; + suppressedRecoveredGoalId = rendered.suppressedGoalId; + if (!bulkReplay) { + streamGoalUpdates = rendered.updates; + return; + } + const goalUpdates = rendered.updates; + if (goalUpdates.length > 0) { replayEnvelope ??= { v: LOAD_REPLAY_VERSION, - updates: replayUpdates, + updates: [], }; + replayEnvelope.updates.push(...goalUpdates); + validateLoadReplayEnvelope( + sessionId, + replayEnvelope, + restoreOptions.replay.kind === 'recent', + ); } - - await this.#restoreWorktreeOnResume(config, createdSession); - await this.#restoreBackgroundAgentsOnResume( - config, - createdSession, - ); + }, + beforeSessionCreate: () => { + response = buildResponse(); + }, + primeSession: (createdSession) => { + profiler.timeSync('runtime_initialize', () => { + if (!projection) return; + createdSession.primeTurnState( + projection.runtime.initialTurn, + projection.runtime.backgroundNotificationTaskIds, + ); + copyCumulativeUsage( + createdSession.cumulativeUsage, + replayUsage, + ); + createdSession.primeRecoveredGoalPublication( + recoveredGoalPublicationKey, + suppressedRecoveredGoalId, + ); + }); + }, + beforeStartPostReplayServices: async (createdSession) => { + if (!bulkReplay && projection?.replay) { + await profiler.time('history_replay', async () => { + const goalBootstrap = replayGoalBootstrap(projection); + const initialGoalState = replayInitialGoalState(projection); + const hasGoalReplayState = + goalBootstrap !== undefined || + initialGoalState.initialGoalState !== undefined || + initialGoalState.initialGoalCause !== undefined; + if (hasGoalReplayState) { + await createdSession.replayHistory( + projection.replay!.records, + projection.replay!.gaps, + { + ...(goalBootstrap ? { goalBootstrap } : {}), + ...initialGoalState, + }, + ); + } else { + await createdSession.replayHistory( + projection.replay!.records, + projection.replay!.gaps, + ); + } + }); + try { + for (const update of streamGoalUpdates) { + await createdSession.sendUpdate(update); + } + } catch (error) { + debugLogger.debug( + `Failed to publish recovered Goal state: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } + } + await profiler.time('post_replay_services', async () => { + await this.#restoreWorktreeOnResume(config, createdSession); + await this.#restoreBackgroundAgentsOnResume( + config, + createdSession, + ); + }); }, }), ); } catch (error) { - return this.cleanupAfterRequestFailure(error, async () => { - if ( - this.sessions.get(config.getSessionId())?.getConfig() !== config - ) { - await this.cleanupUnstoredConfig(config); - } - }); - } - return profiler.timeSync('response_build', () => { - const response: LoadSessionResponse = { - modes: this.buildModesData(config), - models: this.buildAvailableModels(config), - configOptions: this.buildConfigOptions(config), - ...(sessionData?.artifactSnapshot - ? { artifactSnapshot: sessionData.artifactSnapshot } - : {}), - } as LoadSessionResponse; - if (!replayEnvelope) return response; - return { - ...response, - _meta: { - [LOAD_REPLAY_META_KEY]: replayEnvelope, + return this.cleanupAfterRequestFailure( + error, + async () => { + if ( + this.sessions.get(config.getSessionId())?.getConfig() !== config + ) { + await this.cleanupUnstoredConfig(config); + } }, - }; - }); + sessionId, + ); + } + return response!; } finally { releaseStartingSessionId(); } @@ -5161,7 +5401,8 @@ class QwenAgent implements Agent { return this.withLiveSessionRestore( sessionId, liveSession, - async (config, sessionData) => + RESUME_RESTORE_OPTIONS, + async (config, projection) => profiler.timeSync( 'response_build', () => @@ -5169,8 +5410,8 @@ class QwenAgent implements Agent { modes: this.buildModesData(config), models: this.buildAvailableModels(config), configOptions: this.buildConfigOptions(config), - ...(sessionData.artifactSnapshot - ? { artifactSnapshot: sessionData.artifactSnapshot } + ...(projection?.artifactSnapshot + ? { artifactSnapshot: projection.artifactSnapshot } : {}), }) as ResumeSessionResponse, ), @@ -5208,73 +5449,69 @@ class QwenAgent implements Agent { sessionSource, sessionId, true, + {}, + undefined, + RESUME_RESTORE_OPTIONS, ), ); + const projection = config.consumeSessionRestoreProjection?.(); + let response: ResumeSessionResponse | undefined; try { await profiler.time('auth', () => this.ensureAuthenticated(config)); profiler.timeSync('file_system_setup', () => this.setupFileSystem(config), ); await profiler.time('session_register', () => - this.createAndStoreSession( - config, - settings, - config.getResumedSessionData(), - { - enableLiveScreenContext: isCompatibleLiveSessionSource( - sessionSource ?? {}, - ), - replayHistory: false, - beforeStartPostReplayServices: async (createdSession) => { - // `replayHistory: false` skips the publication in - // `createAndStoreSession`, so without this a resumed session - // never tells the client what the recovered goal actually is. - // Safe to stream here, unlike the bulk load path: resume - // replays nothing, so there is no envelope this card could - // sort ahead of. Never fatal. - try { - await createdSession.publishRecoveredGoalState( - config.getResumedSessionData()?.conversation.messages, - ); - } catch (error) { - debugLogger.debug( - `Failed to publish recovered Goal state: ${ - error instanceof Error ? error.message : String(error) - }`, - ); - } + this.createAndStoreSession(config, settings, undefined, { + enableLiveScreenContext: isCompatibleLiveSessionSource( + sessionSource ?? {}, + ), + replayHistory: false, + beforeSessionCreate: () => { + response = profiler.timeSync('response_build', () => ({ + modes: this.buildModesData(config), + models: this.buildAvailableModels(config), + configOptions: this.buildConfigOptions(config), + ...(projection?.runtime.artifactSnapshot + ? { artifactSnapshot: projection.runtime.artifactSnapshot } + : {}), + })) as ResumeSessionResponse; + }, + primeSession: (createdSession) => { + profiler.timeSync('runtime_initialize', () => { + if (!projection) return; + createdSession.primeTurnState( + projection.runtime.initialTurn, + projection.runtime.backgroundNotificationTaskIds, + ); + }); + }, + beforeStartPostReplayServices: async (createdSession) => { + await profiler.time('post_replay_services', async () => { await this.#restoreWorktreeOnResume(config, createdSession); await this.#restoreBackgroundAgentsOnResume( config, createdSession, ); - }, + }); }, - ), + }), ); } catch (error) { - return this.cleanupAfterRequestFailure(error, async () => { - if ( - this.sessions.get(config.getSessionId())?.getConfig() !== config - ) { - await this.cleanupUnstoredConfig(config); - } - }); + return this.cleanupAfterRequestFailure( + error, + async () => { + if ( + this.sessions.get(config.getSessionId())?.getConfig() !== config + ) { + await this.cleanupUnstoredConfig(config); + } + }, + sessionId, + ); } - const sessionData = config.getResumedSessionData(); - return profiler.timeSync( - 'response_build', - () => - ({ - modes: this.buildModesData(config), - models: this.buildAvailableModels(config), - configOptions: this.buildConfigOptions(config), - ...(sessionData?.artifactSnapshot - ? { artifactSnapshot: sessionData.artifactSnapshot } - : {}), - }) as ResumeSessionResponse, - ); + return response!; } finally { releaseStartingSessionId(); } @@ -11806,6 +12043,7 @@ class QwenAgent implements Agent { resume?: boolean, initializeOptions: ConfigInitializeOptions = {}, chatRecording?: boolean, + restoreOptions?: SelectiveSessionRestoreOptions, ): Promise { try { this.assertManagedSessionAdmission(); @@ -11823,6 +12061,7 @@ class QwenAgent implements Agent { resume, initializeOptions, chatRecording, + restoreOptions, ); }); } catch (error) { @@ -11838,7 +12077,9 @@ class QwenAgent implements Agent { errorKind: writerError.errorKind, }); } - throw error; + throw sessionId && restoreOptions + ? mapSessionRestoreRequestError(error, sessionId) + : error; } } @@ -11851,6 +12092,7 @@ class QwenAgent implements Agent { resume?: boolean, initializeOptions: ConfigInitializeOptions = {}, chatRecording?: boolean, + restoreOptions?: SelectiveSessionRestoreOptions, ): Promise { // ACP/IDE-injected servers are session-level: they must outrank a project // `.mcp.json` and stay un-gated. Collect them separately and pass them as @@ -11963,8 +12205,25 @@ class QwenAgent implements Agent { // not process.exit(1) the shared ACP child and every session on its // channel. newSessionConfig maps the throw to a RequestError. true, - this.managedToolInvocationGuard - ? { toolInvocationGuard: this.managedToolInvocationGuard } + this.managedToolInvocationGuard || restoreOptions + ? { + ...(this.managedToolInvocationGuard + ? { toolInvocationGuard: this.managedToolInvocationGuard } + : {}), + ...(restoreOptions && sessionId + ? { + sessionRestore: { + projectionSource: (restoreSessionId) => + new SessionService(cwd, { + runtimeBaseDir: Storage.getRuntimeBaseDir(), + }).readRestoreProjection( + restoreSessionId, + restoreOptions, + ), + }, + } + : {}), + } : undefined, ); if (sessionSource) { @@ -12149,6 +12408,9 @@ class QwenAgent implements Agent { options: { replayHistory?: boolean; enableLiveScreenContext?: boolean; + prepareBeforeSessionCreate?: () => Promise; + beforeSessionCreate?: () => void; + primeSession?: (session: Session) => void; beforeStartPostReplayServices?: (session: Session) => Promise; } = {}, ): Promise { @@ -12170,6 +12432,17 @@ class QwenAgent implements Agent { ); } + await options.prepareBeforeSessionCreate?.(); + this.assertManagedSessionAdmission(); + if (this.sessions.has(sessionId)) { + throw new RequestError( + ACP_ERROR_CODES.INVALID_PARAMS, + `Session ${sessionId} is already active.`, + { errorKind: 'session_id_conflict', sessionId }, + ); + } + options.beforeSessionCreate?.(); + const session = new Session( sessionId, config, @@ -12177,12 +12450,16 @@ class QwenAgent implements Agent { settings, () => this.activeWorkReporter?.notifyChanged(), ); - this.sessions.set(sessionId, session); - // The Session set itself is part of the snapshot: publish so the daemon - // learns about this Session from a report rather than inferring it. - this.activeWorkReporter?.notifyChanged(); - this.initializingConfigs.delete(config); + let published = false; try { + options.primeSession?.(session); + config.hydrateSessionRestoreFileHistory?.(); + this.sessions.set(sessionId, session); + published = true; + // The Session set itself is part of the snapshot: publish so the daemon + // learns about this Session from a report rather than inferring it. + this.activeWorkReporter?.notifyChanged(); + this.initializingConfigs.delete(config); if (options.enableLiveScreenContext) { await session.enableLiveScreenContext(); } @@ -12231,6 +12508,8 @@ class QwenAgent implements Agent { // Install rewriter AFTER history replay to avoid rewriting historical messages session.installRewriter(); + config.finalizeSessionRestore?.(); + // After replay and resume-state restoration so a durable cron fire can't // interleave with either. session.startCronScheduler(); @@ -12240,6 +12519,16 @@ class QwenAgent implements Agent { }, 0); return session; } catch (error) { + if (!published) { + try { + session.dispose(); + } catch (disposeError) { + debugLogger.warn( + `Failed to dispose unpublished session ${sessionId}: ${disposeError instanceof Error ? disposeError.message : String(disposeError)}`, + ); + } + throw error; + } try { await this.discardStoredSessionIfCurrent(sessionId, session, { shutdownConfig: false, diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 76e57e632c..9def7d5189 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -769,6 +769,7 @@ describe('Session', () => { getGeminiClient: vi.fn().mockReturnValue(mockGeminiClient), getGoalRuntime: vi.fn().mockReturnValue(mockGoalRuntime), getGoalRuntimeReady: vi.fn().mockResolvedValue(mockGoalRuntime), + getGoalRuntimePrepared: vi.fn().mockResolvedValue(mockGoalRuntime), bindGoalTurnHost: vi.fn().mockImplementation((host) => { boundGoalHost = host; return () => { @@ -14583,6 +14584,64 @@ describe('Session', () => { expect(mockClient.sessionUpdate).not.toHaveBeenCalled(); }); + it('suppresses a hidden recovered Goal until a different Goal replaces it', async () => { + const listener = mockGoalRuntime.subscribe.mock.calls[0]?.[0] as ( + snapshot: core.GoalSnapshotV2, + cause?: core.GoalStateCause, + ) => void; + session.primeRecoveredGoalPublication(undefined, 'goal-hidden'); + const hidden: core.GoalSnapshotV2 = { + v: 2, + activity: 'idle', + goal: { + ...migratedSnapshot.goal!, + goalId: 'goal-hidden', + revision: 1, + objective: 'hidden inherited goal', + status: 'active', + }, + }; + + listener(hidden, 'create'); + listener({ ...hidden, activity: 'running' }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(mockClient.sessionUpdate).not.toHaveBeenCalled(); + + const progressed = { + ...hidden, + activity: 'idle' as const, + goal: { + ...hidden.goal!, + revision: 2, + objective: 'still hidden', + }, + }; + listener(progressed, 'edit'); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(mockClient.sessionUpdate).not.toHaveBeenCalled(); + + const replacement = { + ...progressed, + goal: { + ...progressed.goal!, + goalId: 'goal-visible', + revision: 1, + objective: 'visible replacement', + }, + }; + listener(replacement, 'replace'); + await vi.waitFor(() => + expect(mockClient.sessionUpdate).toHaveBeenCalledOnce(), + ); + expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: expect.objectContaining({ + _meta: expect.objectContaining({ goalState: replacement }), + }), + }); + }); + it('returns nothing when no Goal was recovered', async () => { mockGoalRuntime.getRecoveryCause.mockReturnValue(undefined); expect(await session.renderRecoveredGoalUpdates([])).toEqual([]); diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index e79cceaca4..68ffeef476 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -182,6 +182,8 @@ import { runWithInvocationContext, truncateNotificationLabel, buildBackgroundEntryLabel, + collectSessionTurnState, + computeInitialTurnFromHistory as computeInitialTurnFromHistoryCore, } from '@qwen-code/qwen-code-core'; import { NOT_CURRENTLY_GENERATING_CANCEL_MESSAGE } from '@qwen-code/acp-bridge/bridgeErrors'; import { CHANNEL_PROMPT_META_KEY } from '@qwen-code/channel-base'; @@ -297,12 +299,12 @@ import { projectAcpToolResultUpdate } from './acp-tool-result-text-projection.js import { ToolCallEmitter } from './emitters/tool-call-emitter.js'; import { ToolCallPreparationTracker } from './tool-call-preparation-tracker.js'; import { PlanEmitter } from './emitters/PlanEmitter.js'; -import { - MessageEmitter, - buildGoalStateUpdate, - buildGoalStatusUpdate, -} from './emitters/MessageEmitter.js'; +import { MessageEmitter } from './emitters/MessageEmitter.js'; import type { HistoryItemGoalStatus } from '../../ui/types.js'; +import { + goalPublicationKey, + renderPreparedGoalUpdate, +} from './recovered-goal-update.js'; import { SubAgentTracker } from './SubAgentTracker.js'; import { buildPermissionRequestContent, @@ -1271,30 +1273,7 @@ export function computeInitialTurnFromHistory( records: ChatRecord[], sessionId: string, ): number { - let maxPromptTurn = 0; - let userMessageCount = 0; - const promptIdPrefix = `${sessionId}########`; - - for (const record of records) { - if (record.sessionId === sessionId && isUserPromptRecord(record)) { - userMessageCount += 1; - } - - for (const promptId of getRecordPromptIds(record)) { - if (!promptId.startsWith(promptIdPrefix)) { - continue; - } - - const suffix = promptId.slice(promptIdPrefix.length); - if (!/^\d+$/.test(suffix)) { - continue; - } - - maxPromptTurn = Math.max(maxPromptTurn, Number(suffix)); - } - } - - return maxPromptTurn > 0 ? maxPromptTurn : userMessageCount; + return computeInitialTurnFromHistoryCore(records, sessionId); } export async function fireSessionPermissionDeniedForAutoMode( @@ -1329,42 +1308,6 @@ export async function fireSessionPermissionDeniedForAutoMode( } } -function getRecordPromptIds(record: ChatRecord): string[] { - const promptIds: string[] = []; - const recordPromptId = (record as { promptId?: unknown }).promptId; - if (typeof recordPromptId === 'string') { - promptIds.push(recordPromptId); - } - const telemetryPromptId = readTelemetryPromptId(record.systemPayload); - if (telemetryPromptId) { - promptIds.push(telemetryPromptId); - } - return promptIds; -} - -function readTelemetryPromptId(payload: unknown): string | undefined { - if (!payload || typeof payload !== 'object' || !('uiEvent' in payload)) { - return undefined; - } - const uiEvent = (payload as { uiEvent?: unknown }).uiEvent; - if (!uiEvent || typeof uiEvent !== 'object' || !('prompt_id' in uiEvent)) { - return undefined; - } - const promptId = (uiEvent as { prompt_id?: unknown }).prompt_id; - return typeof promptId === 'string' ? promptId : undefined; -} - -function isUserPromptRecord(record: ChatRecord): boolean { - if (record.type !== 'user' || record.subtype === 'realtime_message') { - return false; - } - return ( - record.message?.parts?.some( - (part) => typeof part.text === 'string' && part.text.trim().length > 0, - ) ?? false - ); -} - const AT_TOKEN_RE = /@([^\s,;!?()[\]{}]+)/g; function collectExtensionMentionRefs( @@ -1633,6 +1576,9 @@ export class Session implements SessionContext { private goalRuntimeUnsubscribe?: () => void; private lastGoalSnapshot?: GoalSnapshotV2; private lastGoalPublicationKey?: string; + // Set only when runtime recovery selected a Goal that initial replay hid. + // Keep that Goal private through activation and later progress updates. + private suppressedRecoveredGoalId?: string; private goalPublicationTail: Promise = Promise.resolve(); // Set true in dispose(). Guards #drainCronQueue and #drainNotificationQueue @@ -1811,6 +1757,7 @@ export class Session implements SessionContext { this.goalHostUnbind = undefined; this.lastGoalSnapshot = undefined; this.lastGoalPublicationKey = undefined; + this.suppressedRecoveredGoalId = undefined; this.#bindGoalRuntime(); } @@ -1853,53 +1800,46 @@ export class Session implements SessionContext { await this.#queueGoalState(runtime.getSnapshot(), cause); } - /** - * Render the recovered-Goal cards instead of streaming them. - * - * The bulk load-replay path (`historyReplay: 'response'`) does not stream - * its replay: `loadSession` collects the page into the `LOAD_REPLAY` - * envelope and the bridge seeds those updates onto the session's event bus - * *after* the ACP `session/load` call returns. A card streamed from inside - * that call therefore lands on the bus **before** the replayed - * pre-migration `set` card — the reverse of the ordering - * {@link publishRecoveredGoalState} exists to produce, leaving the phantom - * running goal exactly as it was. Returning the cards lets the caller - * append them to the envelope, after the replay page. - * - * Appending after a truncated page (`hasMore`) is still correct: paging - * drops the oldest records, so the authoritative state belongs last either - * way. - * - * Marks the publication as delivered, so the runtime subscription cannot - * emit a duplicate card for the same `(cause, snapshot)` once the session - * goes live. - */ async renderRecoveredGoalUpdates( replayedRecords?: readonly ChatRecord[], ): Promise { if (this.disposed || this.closing) return []; - let runtime; - try { - runtime = await this.config.getGoalRuntimeReady(); - } catch (error) { - if (!(error instanceof GoalPersistenceUnavailableError)) throw error; - const status = this.#unrestorableGoalStatus(replayedRecords); - return status ? [buildGoalStatusUpdate(status)] : []; + const rendered = await renderPreparedGoalUpdate( + () => this.config.getGoalRuntimeReady(), + { + ...(replayedRecords ? { replayedRecords } : {}), + previousGoal: this.lastGoalSnapshot?.goal ?? null, + }, + ); + if ( + rendered.publicationKey && + rendered.publicationKey === this.lastGoalPublicationKey + ) { + return []; } - const cause = runtime.getRecoveryCause?.(); - // Nothing was recovered, so the replay already told the whole story. - if (!cause) return []; - const snapshot = runtime.getSnapshot(); - const publicationKey = this.#goalPublicationKey(snapshot, cause); - if (publicationKey === this.lastGoalPublicationKey) return []; - this.lastGoalPublicationKey = publicationKey; - return [ - buildGoalStateUpdate( - snapshot, - cause, - this.lastGoalSnapshot?.goal ?? null, - ), - ]; + this.primeRecoveredGoalPublication(rendered.publicationKey); + return rendered.updates; + } + + primeRecoveredGoalPublication( + publicationKey: string | undefined, + suppressedGoalId?: string, + ): void { + if (publicationKey) this.lastGoalPublicationKey = publicationKey; + this.suppressedRecoveredGoalId = suppressedGoalId; + } + + #suppressRecoveredGoalUpdate(snapshot: GoalSnapshotV2): boolean { + const suppressedGoalId = this.suppressedRecoveredGoalId; + if (!suppressedGoalId) return false; + const goal = snapshot.goal; + if (goal?.goalId === suppressedGoalId) return true; + if (goal === null) { + this.suppressedRecoveredGoalId = undefined; + return true; + } + this.suppressedRecoveredGoalId = undefined; + return false; } /** @@ -1940,19 +1880,13 @@ export class Session implements SessionContext { }; } - #goalPublicationKey( - snapshot: GoalSnapshotV2, - cause?: GoalStateCause, - ): string | undefined { - return cause ? `${cause}:${JSON.stringify(snapshot)}` : undefined; - } - async #publishGoalState( snapshot: GoalSnapshotV2, cause?: GoalStateCause, previousGoal: GoalRecord | null = this.lastGoalSnapshot?.goal ?? null, ): Promise { - const publicationKey = this.#goalPublicationKey(snapshot, cause); + if (this.#suppressRecoveredGoalUpdate(snapshot)) return; + const publicationKey = goalPublicationKey(snapshot, cause); if (publicationKey && publicationKey === this.lastGoalPublicationKey) { return; } @@ -3148,30 +3082,34 @@ export class Session implements SessionContext { * Delegates to HistoryReplayer for consistent event emission. */ primeTurnFromHistory(records: ChatRecord[]): void { - for (const record of records) { - if (record.subtype !== 'notification') continue; - const backgroundTask = ( - record.systemPayload as - | { backgroundTask?: { taskId?: unknown } } - | undefined - )?.backgroundTask; - if (typeof backgroundTask?.taskId === 'string') { - this.persistedBackgroundNotificationTaskIds.add(backgroundTask.taskId); - } - } - this.turn = Math.max( - this.turn, - computeInitialTurnFromHistory(records, this.config.getSessionId()), + const turnState = collectSessionTurnState( + records, + this.config.getSessionId(), ); + this.primeTurnState( + turnState.initialTurn, + turnState.backgroundNotificationTaskIds, + ); + } + + primeTurnState( + initialTurn: number, + backgroundNotificationTaskIds: readonly string[], + ): void { + for (const taskId of backgroundNotificationTaskIds) { + this.persistedBackgroundNotificationTaskIds.add(taskId); + } + this.turn = Math.max(this.turn, initialTurn); } async replayHistory( records: ChatRecord[], gaps?: HistoryGap[], + options?: Parameters[2], ): Promise { this.primeTurnFromHistory(records); try { - await this.historyReplayer.replay(records, gaps); + await this.historyReplayer.replay(records, gaps, options); } finally { // Replayed plan updates re-stamp the revision via sendUpdate, but they // belong to finished cycles; only live updates may bind the next diff --git a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts index 06c467ed53..ee2d303410 100644 --- a/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts +++ b/packages/cli/src/acp-integration/session/emitters/MessageEmitter.ts @@ -121,8 +121,13 @@ export class MessageEmitter extends BaseEmitter { async emitGoalStatus( status: Omit, + goalState?: unknown, ): Promise { - await this.sendUpdate(buildGoalStatusUpdate(status)); + const update = buildGoalStatusUpdate(status); + if (goalState) { + update._meta = { ...update._meta, goalState }; + } + await this.sendUpdate(update); } async emitGoalState( diff --git a/packages/cli/src/acp-integration/session/history-replay-page.test.ts b/packages/cli/src/acp-integration/session/history-replay-page.test.ts index 277e353f08..bf0c7d9b70 100644 --- a/packages/cli/src/acp-integration/session/history-replay-page.test.ts +++ b/packages/cli/src/acp-integration/session/history-replay-page.test.ts @@ -241,6 +241,37 @@ describe('history replay page', () => { ]); }); + it('fails incrementally before collecting an update above the count limit', async () => { + await expect( + collectHistoryReplayUpdates({ + sessionId: SESSION_ID, + records: [userRecord()], + cumulativeUsage: createReplayCumulativeUsage(), + limits: { maxBytes: Number.MAX_SAFE_INTEGER, maxUpdates: 0 }, + }), + ).rejects.toMatchObject({ + name: 'HistoryReplayLimitError', + reason: 'updates', + observed: 1, + limit: 0, + }); + }); + + it('fails incrementally before retaining serialized updates above the byte limit', async () => { + await expect( + collectHistoryReplayUpdates({ + sessionId: SESSION_ID, + records: [userRecord()], + cumulativeUsage: createReplayCumulativeUsage(), + limits: { maxBytes: 2, maxUpdates: 1 }, + }), + ).rejects.toMatchObject({ + name: 'HistoryReplayLimitError', + reason: 'bytes', + limit: 2, + }); + }); + it('filters malformed replay state before encoding the next cursor', async () => { const logger = { warn: vi.fn() }; const encodeCursor = vi.fn(() => 'next-cursor'); diff --git a/packages/cli/src/acp-integration/session/history-replay-page.ts b/packages/cli/src/acp-integration/session/history-replay-page.ts index 20295cf93c..715a6c5ba1 100644 --- a/packages/cli/src/acp-integration/session/history-replay-page.ts +++ b/packages/cli/src/acp-integration/session/history-replay-page.ts @@ -17,6 +17,7 @@ import { } from '@qwen-code/qwen-code-core'; import type { SessionUpdate } from '@agentclientprotocol/sdk'; import type { TranscriptReplayStateV1 } from '@qwen-code/acp-bridge/transcriptReplay'; +import { Buffer } from 'node:buffer'; import { projectAcpToolResultUpdate } from './acp-tool-result-text-projection.js'; import { HistoryReplayer } from './history-replayer.js'; import type { PendingReplayToolCall } from './history-replayer.js'; @@ -26,6 +27,25 @@ interface ReplayLogger { warn(message: string, ...args: unknown[]): void; } +export class HistoryReplayLimitError extends Error { + constructor( + readonly sessionId: string, + readonly reason: 'bytes' | 'updates', + readonly observed: number, + readonly limit: number, + ) { + super( + `Transcript replay for session ${sessionId} exceeds the ${reason} limit (${observed}, max ${limit})`, + ); + this.name = 'HistoryReplayLimitError'; + } +} + +export interface HistoryReplayLimits { + maxBytes: number; + maxUpdates: number; +} + export function createReplayCumulativeUsage(): CumulativeUsage { return { promptTokens: 0, @@ -164,22 +184,46 @@ function replayContext( updates: SessionUpdate[], cumulativeUsage: CumulativeUsage, config?: Config, + limits?: HistoryReplayLimits, ): SessionEmitterContext { let activeRecordId: string | null = null; + let serializedUpdateBytes = 2; return { sessionId, sendUpdate: async (update) => { const projectedUpdate = projectAcpToolResultUpdate(update); - if (activeRecordId === null) { - updates.push(projectedUpdate); - return; + const updateWithRecordId = (() => { + if (activeRecordId === null) return projectedUpdate; + const record = projectedUpdate as unknown as Record; + const meta = isObjectRecord(record['_meta']) ? record['_meta'] : {}; + return { + ...record, + _meta: { ...meta, 'qwen.session.recordId': activeRecordId }, + } as unknown as SessionUpdate; + })(); + if (limits) { + const updateCount = updates.length + 1; + if (updateCount > limits.maxUpdates) { + throw new HistoryReplayLimitError( + sessionId, + 'updates', + updateCount, + limits.maxUpdates, + ); + } + serializedUpdateBytes += + (updates.length === 0 ? 0 : 1) + + Buffer.byteLength(JSON.stringify(updateWithRecordId), 'utf8'); + if (serializedUpdateBytes > limits.maxBytes) { + throw new HistoryReplayLimitError( + sessionId, + 'bytes', + serializedUpdateBytes, + limits.maxBytes, + ); + } } - const record = projectedUpdate as unknown as Record; - const meta = isObjectRecord(record['_meta']) ? record['_meta'] : {}; - updates.push({ - ...record, - _meta: { ...meta, 'qwen.session.recordId': activeRecordId }, - } as unknown as SessionUpdate); + updates.push(updateWithRecordId); }, setActiveRecordId: (recordId: string | null) => { activeRecordId = recordId; @@ -196,6 +240,9 @@ export async function collectHistoryReplayUpdates({ gaps, cumulativeUsage, logger, + replayState, + goalBootstrap, + limits, }: { sessionId: string; config?: Config; @@ -203,13 +250,22 @@ export async function collectHistoryReplayUpdates({ gaps?: HistoryGap[]; cumulativeUsage: CumulativeUsage; logger?: ReplayLogger; + replayState?: unknown; + goalBootstrap?: import('./history-replayer.js').HistoryReplayGoalBootstrap; + limits?: HistoryReplayLimits; }): Promise<{ updates: SessionUpdate[]; replayError?: string }> { const updates: SessionUpdate[] = []; try { + const initial = parseTranscriptReplayState(replayState, logger); await new HistoryReplayer( - replayContext(sessionId, updates, cumulativeUsage, config), - ).replay(records, gaps); + replayContext(sessionId, updates, cumulativeUsage, config, limits), + ).replay(records, gaps, { + ...(initial.goalState ? { initialGoalState: initial.goalState } : {}), + ...(initial.goalCause ? { initialGoalCause: initial.goalCause } : {}), + ...(goalBootstrap ? { goalBootstrap } : {}), + }); } catch (error) { + if (error instanceof HistoryReplayLimitError) throw error; const replayError = error instanceof Error ? error.message : String(error); logger?.warn( '[historyReplay] History replay failed for session %s (partial updates: %d):', diff --git a/packages/cli/src/acp-integration/session/history-replayer.ts b/packages/cli/src/acp-integration/session/history-replayer.ts index 75541d556e..4554578577 100644 --- a/packages/cli/src/acp-integration/session/history-replayer.ts +++ b/packages/cli/src/acp-integration/session/history-replayer.ts @@ -10,6 +10,11 @@ import type { GoalStateCause, HistoryGap, } from '@qwen-code/qwen-code-core'; +import { + parseGoalSnapshotV2, + parseGoalStateCause, + projectGoalStateToLegacy, +} from '@qwen-code/qwen-code-core'; import { createTranscriptReplayMachine, MISSING_TRANSCRIPT_TOOL_RESULT_MESSAGE, @@ -49,6 +54,18 @@ export interface HistoryReplayPageState { replay: TranscriptReplayStateV1; } +export interface HistoryReplayGoalBootstrap { + goalStatus: { + kind: 'set' | 'checking'; + condition: string; + iterations?: number; + setAt?: number; + durationMs?: number; + lastReason?: string; + }; + goalState?: GoalSnapshotV2; +} + /** * Handles replaying session history on session load. * @@ -65,17 +82,65 @@ export class HistoryReplayer { this.machine = this.createMachine(); } - async replay(records: ChatRecord[], gaps?: HistoryGap[]): Promise { + async replay( + records: ChatRecord[], + gaps?: HistoryGap[], + options: { + initialGoalState?: GoalSnapshotV2; + initialGoalCause?: GoalStateCause; + goalBootstrap?: HistoryReplayGoalBootstrap; + } = {}, + ): Promise { try { + if (options.goalBootstrap) { + const update = { + sessionUpdate: 'agent_message_chunk' as const, + content: { type: 'text' as const, text: '' }, + _meta: { + ...(options.goalBootstrap.goalState + ? { goalState: options.goalBootstrap.goalState } + : {}), + goalStatus: options.goalBootstrap.goalStatus, + }, + }; + await this.sendUpdate(update); + } await this.replayPage(records, { finalizeDangling: true, gaps, + ...(options.initialGoalState + ? { goalState: options.initialGoalState } + : {}), + ...(options.initialGoalCause + ? { goalCause: options.initialGoalCause } + : {}), }); } finally { this.setActiveRecordId(null); } } + static v2GoalBootstrap( + rawGoalState: unknown, + rawGoalCause: unknown, + ): HistoryReplayGoalBootstrap | undefined { + const goalState = parseGoalSnapshotV2(rawGoalState); + const goalCause = parseGoalStateCause(rawGoalCause); + if (!goalState?.goal || goalState.goal.status !== 'active' || !goalCause) { + return undefined; + } + const projection = projectGoalStateToLegacy({ + v: 2, + cause: goalCause, + snapshot: goalState, + }); + const { type: _type, kind, ...goalStatus } = projection.goalStatus; + if (kind !== 'set' && kind !== 'checking') { + return undefined; + } + return { goalStatus: { ...goalStatus, kind }, goalState }; + } + async replayPage( records: ChatRecord[], options: HistoryReplayPageOptions = {}, diff --git a/packages/cli/src/acp-integration/session/recovered-goal-update.test.ts b/packages/cli/src/acp-integration/session/recovered-goal-update.test.ts new file mode 100644 index 0000000000..c587c83a4d --- /dev/null +++ b/packages/cli/src/acp-integration/session/recovered-goal-update.test.ts @@ -0,0 +1,217 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it, vi } from 'vitest'; +import { + GoalPersistenceUnavailableError, + type GoalRuntime, + type GoalSnapshotV2, +} from '@qwen-code/qwen-code-core'; +import { renderPreparedGoalUpdate } from './recovered-goal-update.js'; + +const hiddenSnapshot: GoalSnapshotV2 = { + v: 2, + activity: 'idle', + goal: { + goalId: 'hidden-goal', + revision: 1, + objective: 'hidden objective', + status: 'active', + evidenceCursor: { recordId: 'hidden-record' }, + turnCount: 1, + activeTimeMs: 10, + createdAt: 1, + updatedAt: 2, + }, +}; + +function runtime(): GoalRuntime { + return { + getSnapshot: vi.fn(() => hiddenSnapshot), + getRecoveryCause: vi.fn(() => 'create'), + } as unknown as GoalRuntime; +} + +describe('renderPreparedGoalUpdate', () => { + it('renders the prepared runtime state for an ordinary load', async () => { + const result = await renderPreparedGoalUpdate(async () => runtime()); + + expect(result.publicationKey).toContain('hidden-goal'); + expect(result.updates).toEqual([ + expect.objectContaining({ + _meta: expect.objectContaining({ goalState: hiddenSnapshot }), + }), + ]); + }); + + it('does not duplicate the visible bootstrap for hidden-inherited history', async () => { + const bootstrap = { + goalStatus: { kind: 'set' as const, condition: 'visible objective' }, + }; + + const result = await renderPreparedGoalUpdate(async () => runtime(), { + hideRuntimeGoal: true, + bootstrap, + }); + + expect(result.publicationKey).toContain('hidden-goal'); + expect(result.suppressedGoalId).toBe('hidden-goal'); + expect(result.updates).toEqual([]); + }); + + it('does not duplicate a v2 bootstrap that matches the runtime', async () => { + const result = await renderPreparedGoalUpdate(async () => runtime(), { + bootstrap: { + goalStatus: { kind: 'set', condition: 'hidden objective' }, + goalState: hiddenSnapshot, + }, + }); + + expect(result.updates).toEqual([]); + }); + + it('appends the runtime correction after a legacy bootstrap', async () => { + const result = await renderPreparedGoalUpdate(async () => runtime(), { + bootstrap: { + goalStatus: { kind: 'set', condition: 'hidden objective' }, + }, + }); + + expect(result.updates).toEqual([ + expect.objectContaining({ + _meta: expect.objectContaining({ goalState: hiddenSnapshot }), + }), + ]); + }); + + it('clears a visible legacy bootstrap when recovery is unavailable', async () => { + const result = await renderPreparedGoalUpdate( + async () => { + throw new GoalPersistenceUnavailableError('unsupported record'); + }, + { + bootstrap: { + goalStatus: { + kind: 'checking', + condition: 'visible objective', + iterations: 2, + setAt: 123, + }, + }, + }, + ); + + expect(result.updates).toEqual([ + expect.objectContaining({ + _meta: { + goalStatus: expect.objectContaining({ + kind: 'cleared', + condition: 'visible objective', + iterations: 2, + setAt: 123, + }), + }, + }), + ]); + }); + + it('clears a replayed legacy Goal when recovery is unavailable', async () => { + const result = await renderPreparedGoalUpdate( + async () => { + throw new GoalPersistenceUnavailableError('unsupported record'); + }, + { + replayedRecords: [ + { + uuid: 'goal-result', + parentUuid: null, + sessionId: 'session-1', + timestamp: new Date(0).toISOString(), + type: 'system', + subtype: 'slash_command', + cwd: '/tmp', + version: 'test', + systemPayload: { + phase: 'result', + rawCommand: '/goal', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'replayed objective', + iterations: 3, + setAt: 456, + }, + ], + }, + }, + ], + }, + ); + + expect(result.updates).toEqual([ + expect.objectContaining({ + _meta: { + goalStatus: expect.objectContaining({ + kind: 'cleared', + condition: 'replayed objective', + iterations: 3, + setAt: 456, + }), + }, + }), + ]); + }); + + it('falls back to a page-out bootstrap when replay has no Goal card', async () => { + const result = await renderPreparedGoalUpdate( + async () => { + throw new GoalPersistenceUnavailableError('unsupported record'); + }, + { + replayedRecords: [ + { + uuid: 'user-1', + parentUuid: null, + sessionId: 'session-1', + timestamp: new Date(0).toISOString(), + type: 'user', + cwd: '/tmp', + version: 'test', + message: { role: 'user', parts: [{ text: 'continue' }] }, + }, + ], + bootstrap: { + goalStatus: { + kind: 'set', + condition: 'page-out objective', + iterations: 1, + }, + }, + }, + ); + + expect(result.updates).toEqual([ + expect.objectContaining({ + _meta: { + goalStatus: expect.objectContaining({ + kind: 'cleared', + condition: 'page-out objective', + iterations: 1, + }), + }, + }), + ]); + }); + + it('propagates unexpected runtime failures', async () => { + await expect( + renderPreparedGoalUpdate(async () => { + throw new Error('snapshot failed'); + }), + ).rejects.toThrow('snapshot failed'); + }); +}); diff --git a/packages/cli/src/acp-integration/session/recovered-goal-update.ts b/packages/cli/src/acp-integration/session/recovered-goal-update.ts new file mode 100644 index 0000000000..76452895da --- /dev/null +++ b/packages/cli/src/acp-integration/session/recovered-goal-update.ts @@ -0,0 +1,106 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { SessionUpdate } from '@agentclientprotocol/sdk'; +import { + GoalPersistenceUnavailableError, + type ChatRecord, + type GoalRecord, + type GoalRuntime, + type GoalSnapshotV2, + type GoalStateCause, +} from '@qwen-code/qwen-code-core'; +import type { HistoryItemGoalStatus } from '../../ui/types.js'; +import { + collectGoalStatusItemsFromRecords, + findGoalToRestore, +} from '../../ui/utils/restoreGoal.js'; +import type { HistoryReplayGoalBootstrap } from './history-replayer.js'; +import { + buildGoalStateUpdate, + buildGoalStatusUpdate, +} from './emitters/MessageEmitter.js'; + +export interface RecoveredGoalUpdate { + publicationKey?: string; + suppressedGoalId?: string; + updates: SessionUpdate[]; +} + +export async function renderPreparedGoalUpdate( + getRuntime: () => Promise, + options: { + replayedRecords?: readonly ChatRecord[]; + hideRuntimeGoal?: boolean; + bootstrap?: HistoryReplayGoalBootstrap; + previousGoal?: GoalRecord | null; + } = {}, +): Promise { + let runtime; + try { + runtime = await getRuntime(); + } catch (error) { + if (!(error instanceof GoalPersistenceUnavailableError)) throw error; + const status = unrestorableGoalStatus( + options.replayedRecords, + options.bootstrap, + ); + return { updates: status ? [buildGoalStatusUpdate(status)] : [] }; + } + const cause = runtime.getRecoveryCause?.(); + if (!cause) return { updates: [] }; + const snapshot = runtime.getSnapshot(); + const publicationKey = goalPublicationKey(snapshot, cause); + if (options.hideRuntimeGoal) { + return { + publicationKey, + ...(snapshot.goal + ? { + suppressedGoalId: snapshot.goal.goalId, + } + : {}), + updates: [], + }; + } + const bootstrapGoal = options.bootstrap?.goalState?.goal; + const bootstrapMatchesRuntime = + bootstrapGoal != null && + snapshot.goal?.goalId === bootstrapGoal.goalId && + snapshot.goal?.revision === bootstrapGoal.revision; + return { + publicationKey, + updates: + options.bootstrap && bootstrapMatchesRuntime + ? [] + : [buildGoalStateUpdate(snapshot, cause, options.previousGoal ?? null)], + }; +} + +function unrestorableGoalStatus( + replayedRecords?: readonly ChatRecord[], + bootstrap?: HistoryReplayGoalBootstrap, +): Omit | undefined { + const active = + (replayedRecords?.length + ? findGoalToRestore(collectGoalStatusItemsFromRecords(replayedRecords)) + : undefined) ?? bootstrap?.goalStatus; + if (!active) return undefined; + return { + kind: 'cleared', + condition: active.condition, + iterations: active.iterations, + ...(active.setAt !== undefined ? { setAt: active.setAt } : {}), + lastReason: + 'Goal not restored: its saved state could not be read, so this session is not driving it.', + }; +} + +export function goalPublicationKey( + snapshot: GoalSnapshotV2, + cause?: GoalStateCause, +): string | undefined { + return cause ? `${cause}:${JSON.stringify(snapshot)}` : undefined; +} diff --git a/packages/cli/src/commands/channel/daemon-worker.test.ts b/packages/cli/src/commands/channel/daemon-worker.test.ts index 6709044463..33b5bed08c 100644 --- a/packages/cli/src/commands/channel/daemon-worker.test.ts +++ b/packages/cli/src/commands/channel/daemon-worker.test.ts @@ -121,7 +121,7 @@ const mockDefaultDaemonClient = vi.hoisted(() => ); const mockDefaultDaemonSessionClient = vi.hoisted(() => ({ createOrAttach: vi.fn(), - load: vi.fn(), + resume: vi.fn(), })); const mockBridgeStart = vi.hoisted(() => vi.fn()); @@ -337,7 +337,7 @@ function createSdk() { setModel: vi.fn(), respondToPermission: vi.fn(), }), - load: vi.fn().mockResolvedValue({ + resume: vi.fn().mockResolvedValue({ sessionId: 'loaded-session', workspaceCwd: '/workspace', prompt: vi.fn(), @@ -437,7 +437,7 @@ describe('createDaemonSessionFactory', () => { }, 'qwen-channel-worker', ); - expect(sdk.DaemonSessionClient.load).toHaveBeenCalledWith( + expect(sdk.DaemonSessionClient.resume).toHaveBeenCalledWith( sdk.client, 'existing-session', { @@ -477,7 +477,7 @@ describe('createDaemonSessionFactory', () => { }, 'qwen-channel-worker', ); - expect(sdk.DaemonSessionClient.load).toHaveBeenCalledWith( + expect(sdk.DaemonSessionClient.resume).toHaveBeenCalledWith( sdk.client, 'existing-session', { @@ -516,7 +516,7 @@ describe('createDaemonSessionFactory', () => { ); // The load branch never re-stamps creation attribution: no sourceId in the // load request even when the factory request carried one. - expect(sdk.DaemonSessionClient.load).toHaveBeenCalledWith( + expect(sdk.DaemonSessionClient.resume).toHaveBeenCalledWith( sdk.client, 'existing-session', { diff --git a/packages/cli/src/commands/channel/daemon-worker.ts b/packages/cli/src/commands/channel/daemon-worker.ts index d23a71cc4f..d3a802a711 100644 --- a/packages/cli/src/commands/channel/daemon-worker.ts +++ b/packages/cli/src/commands/channel/daemon-worker.ts @@ -137,7 +137,7 @@ interface DaemonSessionClientStaticLike { }, clientId?: string, ): Promise; - load( + resume( client: DaemonClientLike, sessionId: string, req: { @@ -210,7 +210,7 @@ export function createDaemonSessionFactory({ sessionScope: 'thread' as const, }; if (req.sessionId) { - return await DaemonSessionClient.load( + return await DaemonSessionClient.resume( client, req.sessionId, daemonReq, diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index c995ac20bc..2022b3156e 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -1742,6 +1742,107 @@ describe('loadCliConfig', () => { ); }); + it('rebinds a selective restore projection to the forked session', async () => { + const sourceSessionId = '123e4567-e89b-42d3-a456-426614174000'; + const projectionSource = vi.fn(async (sessionId: string) => ({ + sessionId, + filePath: `/mock/${sessionId}.jsonl`, + startTime: '2026-08-13T00:00:00.000Z', + lastUpdated: '2026-08-13T00:00:00.000Z', + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { lastCompletedUuid: 'leaf', turnParentUuids: [] }, + goalRecords: [], + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + })); + + const config = await loadCliConfig( + {}, + { resume: sourceSessionId, forkSession: true } as CliArgs, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + false, + { sessionRestore: { projectionSource } }, + ); + + const forkedSessionId = config.getSessionId(); + expect(mockSessionServiceInstance.forkSession).toHaveBeenCalledWith( + sourceSessionId, + forkedSessionId, + ); + expect(projectionSource).toHaveBeenCalledOnce(); + expect(projectionSource).toHaveBeenCalledWith(forkedSessionId); + expect(mockSessionServiceInstance.loadSession).not.toHaveBeenCalled(); + const configParams = mockConfigConstructorParams.mock.calls.at(-1)?.[0]; + expect(configParams).toEqual( + expect.objectContaining({ + sessionId: forkedSessionId, + sessionData: undefined, + sessionRestoreProjection: expect.objectContaining({ + sessionId: forkedSessionId, + }), + sessionRestoreProjectionSource: expect.any(Function), + }), + ); + + const deferredProjection = + await configParams.sessionRestoreProjectionSource(); + expect(deferredProjection).toEqual( + expect.objectContaining({ sessionId: forkedSessionId }), + ); + expect(projectionSource).toHaveBeenNthCalledWith(2, forkedSessionId); + }); + + it('preloads a selective projection when a non-ACP host cannot acquire a writer lease', async () => { + const sourceSessionId = '123e4567-e89b-42d3-a456-426614174000'; + const projectionSource = vi.fn(async (sessionId: string) => ({ + sessionId, + filePath: `/mock/${sessionId}.jsonl`, + startTime: '2026-08-13T00:00:00.000Z', + lastUpdated: '2026-08-13T00:00:00.000Z', + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { lastCompletedUuid: 'leaf', turnParentUuids: [] }, + goalRecords: [], + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + })); + + await loadCliConfig( + { experimental: { sessionWriterLease: true } }, + { resume: sourceSessionId } as CliArgs, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + false, + { sessionRestore: { projectionSource } }, + ); + + expect(projectionSource).toHaveBeenCalledOnce(); + expect(projectionSource).toHaveBeenCalledWith(sourceSessionId); + expect(mockConfigConstructorParams).toHaveBeenLastCalledWith( + expect.objectContaining({ + experimentalZedIntegration: false, + sessionWriterLeaseEnabled: true, + sessionRestoreProjection: expect.objectContaining({ + sessionId: sourceSessionId, + }), + }), + ); + }); + it('should explain when --fork-session fails to copy the source session', async () => { const sourceSessionId = '123e4567-e89b-42d3-a456-426614174000'; const sourceData = { diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index d27226420e..bf538d2595 100755 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -23,6 +23,7 @@ import { SessionService, ideContextStore, type ResumedSessionData, + type SessionRestoreProjection, type LspClient, type ToolName, type ToolInvocationGuard, @@ -42,6 +43,7 @@ import { type SkillLevel, type WebSearchSettings, MAX_SUBAGENT_DEPTH_LIMIT, + addDaemonRequestAttribute, } from '@qwen-code/qwen-code-core'; import { extensionsCommand } from '../commands/extensions.js'; import { hooksCommand } from '../commands/hooks.js'; @@ -1557,6 +1559,11 @@ export async function loadCliConfig( */ hostPolicy?: { toolInvocationGuard?: ToolInvocationGuard; + sessionRestore?: { + projectionSource: ( + sessionId: string, + ) => Promise; + }; }, ): Promise { const debugMode = isDebugMode(argv); @@ -1975,6 +1982,10 @@ export async function loadCliConfig( let sessionId: string | undefined; let sessionData: ResumedSessionData | undefined; + let sessionRestoreProjection: SessionRestoreProjection | undefined; + const sessionRestoreProjectionSource = + hostPolicy?.sessionRestore?.projectionSource; + let deferProjectionUntilWriterLease = false; if (argv.continue || argv.resume) { const sessionService = new SessionService(cwd); @@ -1995,8 +2006,24 @@ export async function loadCliConfig( // session UUID by gemini.tsx (which handles custom title lookup and // the interactive picker for ambiguous matches). sessionId = argv.resume; - sessionData = await sessionService.loadSession(argv.resume); - if (!sessionData) { + deferProjectionUntilWriterLease = + sessionRestoreProjectionSource !== undefined && + (argv.chatRecording ?? settings.general?.chatRecording ?? true) && + isAcpMode === true && + settings.experimental?.sessionWriterLease === true; + if (sessionRestoreProjectionSource) { + if (!deferProjectionUntilWriterLease && !argv.forkSession) { + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.projection_acquisition', + 'preloaded', + ); + sessionRestoreProjection = + await sessionRestoreProjectionSource(sessionId); + } + } else { + sessionData = await sessionService.loadSession(argv.resume); + } + if (!sessionRestoreProjectionSource && !sessionData) { const message = `No saved session found with ID ${argv.resume}. Run \`qwen --resume\` without an ID to choose from existing sessions.`; writeStderrLine(message); process.exit(1); @@ -2015,10 +2042,22 @@ export async function loadCliConfig( process.exit(1); } sessionId = forkedSessionId; - sessionData = await sessionService.loadSession(forkedSessionId); - if (!sessionData) { - writeStderrLine(`Failed to load forked session ${forkedSessionId}.`); - process.exit(1); + if (sessionRestoreProjectionSource) { + sessionData = undefined; + if (!deferProjectionUntilWriterLease) { + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.projection_acquisition', + 'preloaded', + ); + sessionRestoreProjection = + await sessionRestoreProjectionSource(forkedSessionId); + } + } else { + sessionData = await sessionService.loadSession(forkedSessionId); + if (!sessionData) { + writeStderrLine(`Failed to load forked session ${forkedSessionId}.`); + process.exit(1); + } } } } else if (argv.sandboxSessionId) { @@ -2047,6 +2086,11 @@ export async function loadCliConfig( const modelProvidersConfig = settings.modelProviders; const providerProtocolConfig = settings.providerProtocol; + const restoreSessionId = sessionId; + const boundSessionRestoreProjectionSource = + sessionRestoreProjectionSource && restoreSessionId + ? () => sessionRestoreProjectionSource(restoreSessionId) + : undefined; // Assemble MCP servers across all sources in precedence order (user/default // settings < project `.mcp.json` < workspace/system settings < `--mcp-config`) @@ -2083,6 +2127,8 @@ export async function loadCliConfig( const configParams: ConfigParameters = { sessionId, sessionData, + sessionRestoreProjection, + sessionRestoreProjectionSource: boundSessionRestoreProjectionSource, embeddingModel: DEFAULT_QWEN_EMBEDDING_MODEL, sandbox: sandboxConfig, targetDir: cwd, diff --git a/packages/cli/src/serve/scheduled-task-keepalive.test.ts b/packages/cli/src/serve/scheduled-task-keepalive.test.ts index 9ee626446e..f8f6ff3727 100644 --- a/packages/cli/src/serve/scheduled-task-keepalive.test.ts +++ b/packages/cli/src/serve/scheduled-task-keepalive.test.ts @@ -45,6 +45,9 @@ describe('scheduled-task keepalive', () => { loadSession: async (req: { sessionId: string }) => { loads.push(req.sessionId); }, + resumeSession: async (req: { sessionId: string }) => { + loads.push(req.sessionId); + }, spawnOrAttach: async () => { throw new Error('spawnOrAttach not mocked'); }, @@ -131,7 +134,7 @@ describe('scheduled-task keepalive', () => { } beats.push(id); }, - loadSession: async (req: { sessionId: string }) => { + resumeSession: async (req: { sessionId: string }) => { loads.push(req.sessionId); }, spawnOrAttach: async () => { @@ -173,7 +176,7 @@ describe('scheduled-task keepalive', () => { } beats.push(id); }, - loadSession: async (req: { sessionId: string }) => { + resumeSession: async (req: { sessionId: string }) => { loads.push(req.sessionId); }, spawnOrAttach: async () => { @@ -282,7 +285,7 @@ describe('scheduled-task keepalive', () => { if (id === 'sess-1') throw new Error('not resident'); beats.push(id); }, - loadSession: async (req: { sessionId: string }) => { + resumeSession: async (req: { sessionId: string }) => { loads.push(req.sessionId); loadRequests.push(req); }, @@ -309,7 +312,6 @@ describe('scheduled-task keepalive', () => { { sessionId: 'sess-1', workspaceCwd: workspace, - historyReplay: 'response', sourceType: 'scheduled_task', sourceId: 'a', }, @@ -327,7 +329,7 @@ describe('scheduled-task keepalive', () => { if (id === 'sess-1') throw new Error('not resident'); beats.push(id); }, - loadSession: async (req: { sessionId: string }) => { + resumeSession: async (req: { sessionId: string }) => { loads.push(req.sessionId); if (req.sessionId === 'sess-1') throw new Error('transcript gone'); }, @@ -358,7 +360,7 @@ describe('scheduled-task keepalive', () => { recordHeartbeat: () => { throw new Error('not resident'); }, - loadSession: async (req: { sessionId: string }) => { + resumeSession: async (req: { sessionId: string }) => { loads.push(req.sessionId); throw new Error('transcript gone'); }, @@ -390,7 +392,7 @@ describe('scheduled-task keepalive', () => { recordHeartbeat: () => { throw new Error('not resident'); }, - loadSession: async (req: { sessionId: string }) => { + resumeSession: async (req: { sessionId: string }) => { loads.push(req.sessionId); // Hang: loadSession isn't abortable, so it keeps running past the timeout. await new Promise((resolve) => { @@ -435,7 +437,7 @@ describe('scheduled-task keepalive', () => { recordHeartbeat: () => { throw new Error('not resident'); }, - loadSession: async () => { + resumeSession: async () => { markStarted?.(); await new Promise((resolve) => { releaseLoad = resolve; @@ -492,7 +494,7 @@ describe('scheduled-task keepalive', () => { }> = []; const res = await rehydrateScheduledTaskSessions({ bridge: { - loadSession: async (req) => { + resumeSession: async (req) => { loaded.push(req); }, }, @@ -525,7 +527,7 @@ describe('scheduled-task keepalive', () => { const errors: string[] = []; const res = await rehydrateScheduledTaskSessions({ bridge: { - loadSession: async (req) => { + resumeSession: async (req) => { if (req.sessionId === 'gone') throw new Error('missing transcript'); }, }, @@ -540,7 +542,7 @@ describe('scheduled-task keepalive', () => { it('rehydrate is a no-op when there are no tasks', async () => { const res = await rehydrateScheduledTaskSessions({ bridge: { - loadSession: async () => { + resumeSession: async () => { throw new Error('should not be called'); }, }, @@ -561,7 +563,7 @@ describe('scheduled-task keepalive', () => { let maxInFlight = 0; const res = await rehydrateScheduledTaskSessions({ bridge: { - loadSession: async () => { + resumeSession: async () => { inFlight++; maxInFlight = Math.max(maxInFlight, inFlight); await new Promise((r) => setTimeout(r, 5)); @@ -590,7 +592,7 @@ describe('scheduled-task keepalive', () => { const res = await rehydrateScheduledTaskSessions({ bridge: { // Never resolves — a genuinely hung, non-abortable load. - loadSession: () => { + resumeSession: () => { started++; return new Promise(() => {}); }, @@ -616,7 +618,7 @@ describe('scheduled-task keepalive', () => { }); const rehydrate = rehydrateScheduledTaskSessions({ bridge: { - loadSession: async () => { + resumeSession: async () => { markStarted?.(); await new Promise((resolve) => { releaseLoad = resolve; @@ -648,7 +650,7 @@ describe('scheduled-task keepalive', () => { ]); const res = await rehydrateScheduledTaskSessions({ bridge: { - loadSession: async () => { + resumeSession: async () => { await new Promise((resolve) => setTimeout(resolve, 20)); }, }, diff --git a/packages/cli/src/serve/scheduled-task-keepalive.ts b/packages/cli/src/serve/scheduled-task-keepalive.ts index ff8906506b..392b07acc1 100644 --- a/packages/cli/src/serve/scheduled-task-keepalive.ts +++ b/packages/cli/src/serve/scheduled-task-keepalive.ts @@ -71,17 +71,16 @@ function collectBoundSessionIds(tasks: readonly DurableCronTask[]): string[] { } /** The slice of the bridge the keepalive needs — narrowed for testability. - * `recordHeartbeat` keeps a live session resident; `loadSession` revives one + * `recordHeartbeat` keeps a live session resident; `resumeSession` revives one * the reaper already let go (a re-enabled task's session). `spawnOrAttach` * and `updateSessionMetadata` bind unbound durable tasks to dedicated * sessions — the same flow the POST /scheduled-tasks route uses for * UI-created tasks, applied retroactively to cron_create tool tasks. */ export interface KeepaliveBridge { recordHeartbeat(sessionId: string): unknown; - loadSession(req: { + resumeSession(req: { sessionId: string; workspaceCwd: string; - historyReplay?: 'stream' | 'response'; sourceType?: string; sourceId?: string; }): Promise; @@ -286,9 +285,9 @@ export function startScheduledTaskKeepalive( string, { failures: number; nextAttemptAt: number } >(); - // Sessions with a revive in flight. loadSession isn't abortable, so a + // Sessions with a revive in flight. resumeSession isn't abortable, so a // timed-out revive keeps running in the background; without this guard a later - // tick would spawn a SECOND loadSession (a duplicate child) for it. Cleared on + // tick would spawn a SECOND resumeSession (a duplicate child) for it. Cleared on // the load's TRUE settlement, not the timeout. const reviving = new Set(); @@ -339,21 +338,24 @@ export function startScheduledTaskKeepalive( const metadata = await new SessionService( boundWorkspace, ).readCreationMetadata(sessionId); - const load = bridge.loadSession({ + const resume = bridge.resumeSession({ sessionId, workspaceCwd: boundWorkspace, - historyReplay: 'response', ...metadata, }); - // Clear the in-flight guard on the load's TRUE settlement (not the + // Clear the in-flight guard on the resume's TRUE settlement (not the // timeout below) so a still-running load keeps blocking a duplicate. - void load + void resume .catch(() => {}) .finally(() => { reviving.delete(sessionId); }); try { - await withTimeout(load, reviveTimeoutMs, `loadSession(${sessionId})`); + await withTimeout( + resume, + reviveTimeoutMs, + `resumeSession(${sessionId})`, + ); log.debug('keepalive: revived non-resident session', sessionId); reviveState.delete(sessionId); } catch (loadErr) { @@ -408,7 +410,7 @@ export function startScheduledTaskKeepalive( // In-flight guard: a pass can outlast the interval (each revive awaits up to // the revive timeout), so skip a tick while the previous is still running — - // overlapping passes would issue duplicate concurrent loadSession spawns for + // overlapping passes would issue duplicate concurrent resumeSession spawns for // the same dead sessions. let running = false; const timer: ReturnType = setInterval(() => { @@ -476,10 +478,9 @@ export function startScheduledTaskKeepalive( /** The slice of the bridge rehydration needs — narrowed for testability. */ export interface RehydrateBridge { - loadSession(req: { + resumeSession(req: { sessionId: string; workspaceCwd: string; - historyReplay?: 'stream' | 'response'; sourceType?: string; sourceId?: string; }): Promise; @@ -497,12 +498,12 @@ export interface RehydrateResult { * lock owner deliberately never fires a bound task) until something loaded it. * * Best-effort: a session whose transcript is gone (deleted out-of-band) fails - * its `loadSession` and is skipped rather than aborting the sweep. Distinct + * its `resumeSession` and is skipped rather than aborting the sweep. Distinct * session ids only; unbound tasks are ignored (they fire via the lock owner). */ /** Default caller headroom above the bridge's 60-second restore deadline. */ -const REHYDRATE_LOAD_TIMEOUT_MS = 70_000; -/** Max sessions rehydrated at once. Each `loadSession` forks a real agent +const REHYDRATE_RESUME_TIMEOUT_MS = 70_000; +/** Max sessions rehydrated at once. Each `resumeSession` forks a real agent * child, so loading all of them (up to MAX_JOBS = 50) in one shot would spike * CPU/memory on boot and, on constrained hosts, hit spawn failures * (EAGAIN/ENOMEM) that strand healthy tasks. Load in small batches instead. */ @@ -517,7 +518,7 @@ export async function rehydrateScheduledTaskSessions(deps: { onTasksRead?: (tasks: readonly DurableCronTask[]) => void; }): Promise { const { bridge, boundWorkspace } = deps; - const timeoutMs = deps.loadTimeoutMs ?? REHYDRATE_LOAD_TIMEOUT_MS; + const timeoutMs = deps.loadTimeoutMs ?? REHYDRATE_RESUME_TIMEOUT_MS; let tasks; try { tasks = await readCronTasks(boundWorkspace); @@ -540,18 +541,17 @@ export async function rehydrateScheduledTaskSessions(deps: { const metadata = await new SessionService( boundWorkspace, ).readCreationMetadata(sessionId); - const load = bridge.loadSession({ + const resume = bridge.resumeSession({ sessionId, workspaceCwd: boundWorkspace, - historyReplay: 'response', ...metadata, }); - // loadSession isn't abortable, so a timed-out load keeps forking/replaying + // resumeSession isn't abortable, so a timed-out resume keeps running // in the background. Swallow its eventual settlement up front so it can't // raise an unhandled rejection once we've stopped awaiting it below. - void load.catch(() => {}); + void resume.catch(() => {}); try { - await withTimeout(load, timeoutMs, `loadSession(${sessionId})`); + await withTimeout(resume, timeoutMs, `resumeSession(${sessionId})`); loaded.push(sessionId); } catch (err) { // Timed out (or the load rejected). Do NOT await the raw `load` here: a @@ -559,7 +559,7 @@ export async function rehydrateScheduledTaskSessions(deps: { // enough loads hang, the whole boot sweep never completes (`Promise.all` // never settles) — later task sessions would then never rehydrate. Record // it as failed and free the worker to pull the next queued session; the - // background load, if it ever settles, just warms that session late. + // background resume, if it ever settles, just warms that session late. failed.push(sessionId); // The onError callback must never abort the sweep: if it throws (e.g. a // stderr EPIPE during log rotation) the rejection would escape loadOne, diff --git a/packages/cli/src/ui/utils/restoreGoal.ts b/packages/cli/src/ui/utils/restoreGoal.ts index cbe7f66285..0b98151e01 100644 --- a/packages/cli/src/ui/utils/restoreGoal.ts +++ b/packages/cli/src/ui/utils/restoreGoal.ts @@ -9,8 +9,8 @@ import { setGoalTerminalObserver, setLastGoalTerminal, unregisterGoalHook, - type ChatRecord, type Config, + type GoalRecoveryRecord, type GoalTerminalEvent, type GoalTerminalKind, type SlashCommandRecordPayload, @@ -178,7 +178,7 @@ export function parseGoalStatusItem(item: unknown): GoalStatusItem | null { * exists, so `findGoalToRestore` / `findLastTerminalGoal` are fed from here. */ export function collectGoalStatusItemsFromRecords( - records: readonly ChatRecord[], + records: readonly GoalRecoveryRecord[], ): GoalStatusItem[] { const items: GoalStatusItem[] = []; for (const record of records) { diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index df73200254..b44a452c9f 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -2522,6 +2522,83 @@ describe('Server Config (config.ts)', () => { expect(replacement.getSnapshot().goal?.status).toBe('active'); }); + it('holds selective Goal readiness and autonomous work until finalization', async () => { + const record = resumedGoalSession('active').conversation.messages[0]!; + const config = new Config({ + ...baseParams, + chatRecording: true, + sessionRestoreProjection: { + sessionId: 'resumed-session', + filePath: '/tmp/resumed-session.jsonl', + startTime: new Date(0).toISOString(), + lastUpdated: new Date(0).toISOString(), + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { + lastCompletedUuid: record.uuid, + turnParentUuids: [], + }, + goalRecords: [record], + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + }, + }); + const started: string[] = []; + config.bindGoalTurnHost({ + startGoalTurn: vi.fn(async ({ permit }) => { + started.push(permit.goalId); + }), + preemptGoalTurn: vi.fn(), + }); + let ready = false; + void config.getGoalRuntimeReady().then(() => { + ready = true; + }); + + await Promise.resolve(); + expect(ready).toBe(false); + expect(started).toEqual([]); + + config.finalizeSessionRestore(); + + await expect(config.getGoalRuntimeReady()).resolves.toBe( + config.getGoalRuntime(), + ); + await vi.waitFor(() => expect(started).toEqual(['g-resumed'])); + }); + + it('rejects selective Goal readiness when restore is abandoned', async () => { + const record = resumedGoalSession('active').conversation.messages[0]!; + const config = new Config({ + ...baseParams, + chatRecording: true, + sessionRestoreProjection: { + sessionId: 'resumed-session', + filePath: '/tmp/resumed-session.jsonl', + startTime: new Date(0).toISOString(), + lastUpdated: new Date(0).toISOString(), + runtime: { + apiHistory: [], + uiTelemetryEvents: [], + recording: { + lastCompletedUuid: record.uuid, + turnParentUuids: [], + }, + goalRecords: [record], + initialTurn: 0, + backgroundNotificationTaskIds: [], + }, + }, + }); + const readiness = config.getGoalRuntimeReady(); + + config.startNewSession('replacement-session'); + + await expect(readiness).rejects.toThrow('Session restore was abandoned'); + }); + it('owns one durable Goal runtime per canonical session', async () => { const config = new Config({ ...baseParams, chatRecording: true }); const first = config.getGoalRuntime(); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 7fd5381bb5..ada64e4dce 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -130,6 +130,7 @@ import { SENSITIVE_SPAN_ATTRIBUTE_MAX_LENGTH_LIMIT, isValidSensitiveSpanAttributeMaxLength, isTelemetrySdkInitialized, + addDaemonRequestAttribute, initializeTelemetry, shutdownTelemetry, refreshSessionContext, @@ -172,6 +173,7 @@ import { type GoalRuntime, type GoalTurnHost, } from '../goals/goal-runtime.js'; +import type { GoalRecoveryRecord } from '../goals/goal-persistence.js'; import { createGoalCheckpointVerifier } from '../goals/goal-checkpoint-verifier.js'; import { createGoalVerifier } from '../goals/goal-verifier.js'; import type { ToolInvocationGuard } from '../core/tool-invocation-guard.js'; @@ -210,7 +212,6 @@ import { ChatRecordingService, type ChatRecordingFailureEvent, type ChatRecordingFailureListener, - type ChatRecord, } from '../services/chatRecordingService.js'; import { CHARS_PER_TOKEN } from '../services/tokenEstimation.js'; import { @@ -221,6 +222,10 @@ import { SessionService, type ResumedSessionData, } from '../services/sessionService.js'; +import type { + SessionRestoreProjection, + SessionRuntimeResumeState, +} from '../services/session-transcript-reader.js'; import { SessionTranscriptChangedError, SessionWriterError, @@ -944,6 +949,10 @@ export interface AgentsCollabSettings { export interface ConfigParameters { sessionId?: string; sessionData?: ResumedSessionData; + sessionRestoreProjection?: SessionRestoreProjection; + sessionRestoreProjectionSource?: () => Promise< + SessionRestoreProjection | undefined + >; embeddingModel?: string; sandbox?: SandboxConfig; targetDir: string; @@ -1735,6 +1744,14 @@ export class Config { private sessionSourceType?: string; private sessionSourceId?: string; private sessionData?: ResumedSessionData; + private pendingSessionRestoreProjection?: SessionRestoreProjection; + private sessionRestoreRuntime?: SessionRuntimeResumeState; + private readonly sessionRestoreProjectionSource?: () => Promise< + SessionRestoreProjection | undefined + >; + private restoredFileHistory = false; + private goalRestoreActivation?: () => Promise; + private rejectGoalRestoreActivation?: (reason?: unknown) => void; private readonly sessionRuntimeBaseDir: string; private sessionProjectDirRegistered = false; private pendingSessionWriterLease?: SessionWriterLease; @@ -2134,6 +2151,8 @@ export class Config { sessionEnvClaimed = true; } this.sessionData = params.sessionData; + this.sessionRestoreProjectionSource = params.sessionRestoreProjectionSource; + this.setSessionRestoreProjection(params.sessionRestoreProjection); setDebugLogSession(this); this.debugLogger = createDebugLogger(); this.embeddingModel = params.embeddingModel ?? DEFAULT_QWEN_EMBEDDING_MODEL; @@ -2523,7 +2542,17 @@ export class Config { this.chatRecordingService = this.chatRecordingEnabled ? this.createChatRecordingService() : undefined; - this.initializeGoalRuntime(this.sessionData?.conversation.messages); + if ( + !this.sessionRestoreProjectionSource || + this.sessionRestoreRuntime || + !this.sessionWriterLeaseEnabled + ) { + this.initializeGoalRuntime( + this.sessionRestoreRuntime?.goalRecords ?? + this.sessionData?.conversation.messages, + this.sessionRestoreRuntime, + ); + } this.extensionManager = new ExtensionManager({ workspaceDir: this.targetDir, enabledExtensionOverrides: this.overrideExtensions, @@ -2621,6 +2650,7 @@ export class Config { this.sessionProjectDirRegistered = true; await this.initializeInternal(options); } catch (error) { + this.clearSessionRestoreProjection(); if (this.sessionProjectDirRegistered) { unregisterSessionProjectDir(this.sessionId); this.sessionProjectDirRegistered = false; @@ -3183,7 +3213,15 @@ export class Config { throw new SessionTranscriptChangedError(); } let authoritative: ResumedSessionData | undefined; - if (this.sessionData || lease.transcriptExistedAtAcquire) { + let projection: SessionRestoreProjection | undefined; + if (this.sessionRestoreProjectionSource) { + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.projection_acquisition', + 'after_writer_lease', + ); + projection = await this.sessionRestoreProjectionSource(); + this.setSessionRestoreProjection(projection); + } else if (this.sessionData || lease.transcriptExistedAtAcquire) { authoritative = await this.getSessionService().loadSession( this.sessionId, ); @@ -3199,7 +3237,18 @@ export class Config { throw new SessionWriterShutdownError(); } this.sessionData = authoritative; - recorder.activate(lease, authoritative, persistedTitleInfo); + recorder.activate( + lease, + authoritative, + persistedTitleInfo, + projection?.runtime.recording, + ); + if (this.sessionRestoreProjectionSource) { + this.initializeGoalRuntime( + projection?.runtime.goalRecords, + projection?.runtime, + ); + } this.pendingSessionWriterLease = undefined; lease = undefined; // The recorder can take writes now, so the restore the constructor @@ -3778,6 +3827,77 @@ export class Config { return this.sessionId; } + getSessionRestoreRuntime(): SessionRuntimeResumeState | undefined { + return this.sessionRestoreRuntime; + } + + consumeSessionRestoreProjection(): SessionRestoreProjection | undefined { + const projection = this.pendingSessionRestoreProjection; + this.pendingSessionRestoreProjection = undefined; + return projection; + } + + hydrateSessionRestoreFileHistory(): void { + if (this.restoredFileHistory) return; + const snapshots = this.sessionRestoreRuntime?.fileHistorySnapshots; + if (!snapshots?.length) return; + const service = this.getFileHistoryService(); + if (!service.isEnabled()) return; + service.restoreFromSnapshots(snapshots); + this.restoredFileHistory = true; + } + + finalizeSessionRestore(): void { + const runtime = this.sessionRestoreRuntime; + if (!runtime) return; + this.sessionRestoreRuntime = undefined; + + if (runtime.attributionSnapshot) { + try { + CommitAttributionService.getInstance().restoreFromSnapshot( + runtime.attributionSnapshot, + ); + } catch (error) { + this.debugLogger.error( + `Session restore attribution activation failed: ${error}`, + ); + } + } + + const activateGoal = this.goalRestoreActivation; + this.goalRestoreActivation = undefined; + this.rejectGoalRestoreActivation = undefined; + if (activateGoal) { + try { + void activateGoal().catch((error) => { + this.debugLogger.error( + `Session restore goal activation failed: ${error}`, + ); + }); + } catch (error) { + this.debugLogger.error( + `Session restore goal activation failed: ${error}`, + ); + } + } + + if (this.restoredFileHistory && this.fileHistoryService) { + try { + void this.fileHistoryService + .validateRestoredSnapshots() + .catch((error) => { + this.debugLogger.error( + `FileHistory: validateRestoredSnapshots failed: ${error}`, + ); + }); + } catch (error) { + this.debugLogger.error( + `FileHistory: validateRestoredSnapshots failed: ${error}`, + ); + } + } + } + setSessionSource(sourceType: string, sourceId?: string): void { this.sessionSourceType = sourceType; this.sessionSourceId = sourceId; @@ -3863,6 +3983,7 @@ export class Config { unregisterSessionModel(previousSessionId); this.publishModelEnv(); this.sessionData = sessionData; + this.clearSessionRestoreProjection(); this.pendingRecoveredAgentsNotice = null; this.getOwnActiveTodoReminders().clear(); this.getOwnActiveTodoWorkChainOwners().clear(); @@ -5097,6 +5218,7 @@ export class Config { private async shutdownResourcesOnce(): Promise { try { + this.clearSessionRestoreProjection(); // Drop this session's project-dir registry entry. It is registered during // initialization, so it is released here whenever that step completed — // in daemon mode, where one process serves many sessions, an unreleased @@ -5111,6 +5233,11 @@ export class Config { unregisterSessionModel(this.sessionId); if (Object.hasOwn(this, 'goalRuntime')) { + this.rejectGoalRestoreActivation?.( + new GoalPersistenceUnavailableError('Goal runtime disposed'), + ); + this.goalRestoreActivation = undefined; + this.rejectGoalRestoreActivation = undefined; this.goalTurnHostUnbind?.(); this.goalTurnHostUnbind = undefined; // Shutting down before the writer arrived: nothing will ever run @@ -7482,6 +7609,12 @@ export class Config { return this.goalRuntimeReady.then(() => runtime); } + getGoalRuntimePrepared(): Promise { + const runtime = this.getGoalRuntime(); + if (!this.sessionRestoreRuntime) return this.getGoalRuntimeReady(); + return runtime.getPreparedRestore().then(() => runtime); + } + async rebaseGoalRuntimeFromActiveTranscript(): Promise { const runtime = this.getGoalRuntime(); const recordingService = this.chatRecordingService; @@ -7533,10 +7666,19 @@ export class Config { this.notifyChatRecordingFailure(event); }, this.sessionWriterLeaseEnabled, + this.sessionRestoreRuntime?.recording, ); } - private initializeGoalRuntime(records?: readonly ChatRecord[]): void { + private initializeGoalRuntime( + records?: readonly GoalRecoveryRecord[], + restoreRuntime?: SessionRuntimeResumeState, + ): void { + this.rejectGoalRestoreActivation?.( + new GoalPersistenceUnavailableError('Goal runtime replaced'), + ); + this.goalRestoreActivation = undefined; + this.rejectGoalRestoreActivation = undefined; this.goalTurnHostUnbind?.(); this.goalTurnHostUnbind = undefined; // A runtime built here supersedes any restore still waiting on the @@ -7569,7 +7711,30 @@ export class Config { // failure as `recoveryError` for the life of the runtime — the // migrated goal is dropped and goal persistence is bricked for the // whole resumed session. Wait for the writer instead. - if (this.sessionWriterLeaseEnabled && !recorder.hasWriteOwnership()) { + if (restoreRuntime) { + const preparation = runtime.prepareRestore( + records ?? [], + restoreRuntime.goalCheckpointWindow, + ); + let resolveActivation!: () => void; + let rejectActivation!: (reason?: unknown) => void; + const activation = new Promise((resolve, reject) => { + resolveActivation = resolve; + rejectActivation = reject; + }); + this.rejectGoalRestoreActivation = rejectActivation; + this.goalRestoreActivation = () => { + const started = runtime.activateRestoredWork(); + void started.then(resolveActivation, rejectActivation); + return started; + }; + this.goalRuntimeReady = Promise.all([preparation, activation]).then( + () => runtime, + ); + } else if ( + this.sessionWriterLeaseEnabled && + !recorder.hasWriteOwnership() + ) { const ready = new Promise((resolve, reject) => { this.pendingGoalRestore = { runtime, resolve, reject }; }); @@ -7624,6 +7789,25 @@ export class Config { pending.reject(error); } + private setSessionRestoreProjection( + projection: SessionRestoreProjection | undefined, + ): void { + this.pendingSessionRestoreProjection = projection; + this.sessionRestoreRuntime = projection?.runtime; + this.restoredFileHistory = false; + } + + private clearSessionRestoreProjection(): void { + this.pendingSessionRestoreProjection = undefined; + this.sessionRestoreRuntime = undefined; + this.restoredFileHistory = false; + this.rejectGoalRestoreActivation?.( + new GoalPersistenceUnavailableError('Session restore was abandoned'), + ); + this.goalRestoreActivation = undefined; + this.rejectGoalRestoreActivation = undefined; + } + private notifyChatRecordingFailure(event: ChatRecordingFailureEvent): void { for (const listener of [...this.chatRecordingFailureListeners]) { try { diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index f46ad2d0ae..24ac8b58e0 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -633,6 +633,7 @@ describe('Gemini Client (client.ts)', () => { getChatRecordingService: vi.fn().mockReturnValue(undefined), getFileHistoryService: vi.fn().mockReturnValue(mockFileHistoryService), getResumedSessionData: vi.fn().mockReturnValue(undefined), + getSessionRestoreRuntime: vi.fn().mockReturnValue(undefined), getArenaAgentClient: vi.fn().mockReturnValue(null), getManagedAutoMemoryEnabled: vi.fn().mockReturnValue(true), isManagedMemoryAvailable: vi.fn().mockReturnValue(true), @@ -705,6 +706,46 @@ describe('Gemini Client (client.ts)', () => { }); describe('initialize', () => { + it('initializes from the selective runtime projection without the full transcript', async () => { + const seedResumeTokenCountsSpy = vi.spyOn( + GeminiChat.prototype, + 'seedResumeTokenCounts', + ); + const apiHistory = [ + { role: 'user' as const, parts: [{ text: 'projected history' }] }, + ]; + const uiEvent = { type: 'projected-event' }; + vi.mocked(mockConfig.getSessionRestoreRuntime).mockReturnValue({ + apiHistory, + resumeTokenCounts: { + promptTokenCount: 321, + outputTokenCount: 45, + isEstimated: false, + }, + uiTelemetryEvents: [uiEvent], + recording: { + lastCompletedUuid: 'record-1', + turnParentUuids: [], + }, + goalRecords: [], + initialTurn: 0, + backgroundNotificationTaskIds: [], + } as unknown as ReturnType); + + const resumedClient = new GeminiClient(mockConfig); + await resumedClient.initialize(); + + expect(resumedClient.getHistory().at(-1)).toEqual(apiHistory[0]); + expect(uiTelemetryService.resetSession).toHaveBeenCalledWith( + 'test-session-id', + ); + expect(uiTelemetryService.addEvent).toHaveBeenCalledWith( + uiEvent, + 'test-session-id', + ); + expect(seedResumeTokenCountsSpy).toHaveBeenCalledWith(321, 45, false); + }); + it('seeds resumed chat with replayed prompt token count', async () => { vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({ conversation: { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 18cb53216d..52a7c3add7 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -428,7 +428,32 @@ export class GeminiClient { // Check if we're resuming from a previous session const resumedSessionData = this.config.getResumedSessionData(); - if (resumedSessionData) { + const restoreRuntime = this.config.getSessionRestoreRuntime?.(); + if (restoreRuntime) { + uiTelemetryService.resetSession(sessionId); + for (const event of restoreRuntime.uiTelemetryEvents) { + uiTelemetryService.addEvent(event, sessionId); + } + this.seedRecentCompletedToolNamesFromHistory(restoreRuntime.apiHistory); + await this.startChat( + restoreRuntime.apiHistory, + sessionStartSource ?? SessionStartSource.Resume, + ); + const chat = this.getChat(); + if (restoreRuntime.resumeTokenCounts) { + const counts = restoreRuntime.resumeTokenCounts; + uiTelemetryService.setLastPromptTokenCount(counts.promptTokenCount); + chat.seedResumeTokenCounts( + counts.promptTokenCount, + counts.outputTokenCount, + counts.isEstimated, + ); + } else { + chat.setLastPromptTokenCount( + uiTelemetryService.getLastPromptTokenCount(), + ); + } + } else if (resumedSessionData) { const resumeTokenCounts = replayUiTelemetryFromConversation( resumedSessionData.conversation, this.config.getSessionId(), diff --git a/packages/core/src/goals/goal-evidence.ts b/packages/core/src/goals/goal-evidence.ts index 5c08c55a42..3fae3e8e88 100644 --- a/packages/core/src/goals/goal-evidence.ts +++ b/packages/core/src/goals/goal-evidence.ts @@ -14,7 +14,10 @@ import { type GoalTerminalProposal, type GoalTurnPermit, } from './goal-protocol.js'; -import { projectUserTranscriptForDisplay } from '../utils/transcript-records.js'; +import { + isUserPromptSubmitContextPartText, + projectUserTranscriptForDisplay, +} from '../utils/transcript-records.js'; const CATALOG_PREVIEW_LIMIT = 240; const CATALOG_ENTRY_LIMIT = 100; @@ -158,6 +161,378 @@ interface ParsedGoalContext { turnId: string; } +export interface GoalEvidenceRecordIndexHint { + uuid: string; + parsedGoalContext?: { + goalId: string; + revision: number; + turnId: string; + }; + claimedGoalId?: string; + claimedRevision?: number; + provenance?: GoalEvidenceProvenance; + hasCatalogEligibleContent: boolean; + hasRawEligibleContent: boolean; + catalogEntryBytes?: number; +} + +export class GoalEvidenceRecordIndexAccumulator { + private readonly uuid: string; + private readonly parsedGoalContext?: ParsedGoalContext; + private readonly claimedGoalId?: string; + private readonly claimedRevision?: number; + private readonly provenance?: GoalEvidenceProvenance; + private readonly hasObjectSystemPayload: boolean; + private readonly displayText?: string; + private readonly hasHookContext: boolean; + private prefixPreview = ''; + private lastPartPreviewValues: string[] = []; + private lastPartIsHookContext = false; + private partCount = 0; + private hasRawEligibleContent = false; + + constructor(record: GoalEvidenceRecord) { + this.uuid = record.uuid; + this.parsedGoalContext = parseGoalContext(record.goalContext); + const claimed = isRecord(record.goalContext) + ? record.goalContext + : undefined; + this.claimedGoalId = + typeof claimed?.['goalId'] === 'string' ? claimed['goalId'] : undefined; + this.claimedRevision = + typeof claimed?.['revision'] === 'number' + ? claimed['revision'] + : undefined; + this.provenance = this.parsedGoalContext + ? coherentEvidenceProvenance(record) + : undefined; + const systemPayload = isRecord(record.systemPayload) + ? record.systemPayload + : undefined; + this.hasObjectSystemPayload = systemPayload !== undefined; + this.displayText = + typeof systemPayload?.['displayText'] === 'string' + ? systemPayload['displayText'].slice(0, CATALOG_PREVIEW_LIMIT) + : undefined; + this.hasHookContext = typeof systemPayload?.['hookContext'] === 'string'; + this.addFragment(record); + } + + addFragment(record: GoalEvidenceRecord): void { + if (!this.provenance) return; + for (const part of record.message?.parts ?? []) { + this.finishPreviousPart(); + const previewValues: string[] = []; + if (part.thought !== true && typeof part.text === 'string') { + previewValues.push(part.text.slice(0, CATALOG_PREVIEW_LIMIT)); + if (part.text.trim()) this.hasRawEligibleContent = true; + } + if (this.provenance === 'tool_result' && part.functionResponse) { + previewValues.push(renderToolResponsePreview(part.functionResponse)); + if (part.functionResponse.response !== undefined) { + this.hasRawEligibleContent = true; + } + } + this.lastPartPreviewValues = previewValues; + this.lastPartIsHookContext = + typeof part.text === 'string' && + isUserPromptSubmitContextPartText(part.text); + this.partCount++; + } + } + + finish(): GoalEvidenceRecordIndexHint { + let preview: string; + const hasFinalHookContextPart = + this.partCount > 1 && this.lastPartIsHookContext; + if ( + this.provenance === 'real_user' && + (this.hasHookContext || hasFinalHookContextPart) && + this.displayText !== undefined + ) { + preview = this.displayText.slice(0, CATALOG_PREVIEW_LIMIT).trim(); + } else if ( + this.provenance === 'real_user' && + !this.hasObjectSystemPayload && + hasFinalHookContextPart + ) { + preview = this.prefixPreview.trim(); + } else { + preview = appendPreviewValues( + this.prefixPreview, + this.lastPartPreviewValues, + ).trim(); + } + const catalogEntry = + this.provenance && this.parsedGoalContext && preview + ? { + uuid: this.uuid, + provenance: this.provenance, + turnId: this.parsedGoalContext.turnId, + preview, + proofKind: proofKindOf(this.provenance), + } + : undefined; + return { + uuid: this.uuid, + ...(this.parsedGoalContext + ? { parsedGoalContext: this.parsedGoalContext } + : {}), + ...(this.claimedGoalId !== undefined + ? { claimedGoalId: this.claimedGoalId } + : {}), + ...(this.claimedRevision !== undefined + ? { claimedRevision: this.claimedRevision } + : {}), + ...(this.provenance ? { provenance: this.provenance } : {}), + hasCatalogEligibleContent: catalogEntry !== undefined, + hasRawEligibleContent: this.hasRawEligibleContent, + ...(catalogEntry + ? { + catalogEntryBytes: Buffer.byteLength( + JSON.stringify(catalogEntry), + 'utf8', + ), + } + : {}), + }; + } + + private finishPreviousPart(): void { + if (this.partCount === 0) return; + this.prefixPreview = appendPreviewValues( + this.prefixPreview, + this.lastPartPreviewValues, + ); + } +} + +function appendPreviewValues( + current: string, + values: readonly string[], +): string { + let preview = current; + for (const value of values) { + if (!value || preview.length >= CATALOG_PREVIEW_LIMIT) continue; + const separator = preview ? '\n' : ''; + const remaining = CATALOG_PREVIEW_LIMIT - preview.length; + preview += `${separator}${value}`.slice(0, remaining); + } + return preview; +} + +export class GoalEvidenceCheckpointAccumulator { + private readonly candidateUuids: string[] = []; + private readonly candidateUuidSet = new Set(); + private readonly captured = new Map(); + private readonly checkpointEntries: GoalEvidenceCatalogEntry[]; + private readonly truncated: boolean; + private readonly shouldCheckpoint: boolean; + + constructor( + hints: readonly GoalEvidenceRecordIndexHint[], + private readonly goal: GoalRecord, + permit: GoalTurnPermit, + ) { + if ( + permit.goalId !== goal.goalId || + permit.revision !== goal.revision || + !isNonEmptyString(permit.turnId) + ) { + throw new EvidenceSourceUnavailableError( + 'permit_goal_mismatch', + 'The current Goal permit does not match the Goal evidence revision.', + ); + } + const indexByUuid = new Map(); + for (let index = 0; index < hints.length; index++) { + const uuid = hints[index]!.uuid; + if (indexByUuid.has(uuid)) { + throw new EvidenceSourceUnavailableError( + 'duplicate_record_uuid', + `The active transcript chain contains duplicate record UUID ${uuid}.`, + ); + } + indexByUuid.set(uuid, index); + } + const cursorId = goal.evidenceCursor.recordId; + if (cursorId === null) { + throw new EvidenceSourceUnavailableError( + 'cursor_unset', + 'The Goal evidence cursor is not available.', + ); + } + const cursorIndex = indexByUuid.get(cursorId); + if (cursorIndex === undefined) { + throw new EvidenceSourceUnavailableError( + 'cursor_not_found', + `The Goal evidence cursor ${cursorId} is not in the active transcript chain.`, + ); + } + + const lineageTurnIds: string[] = []; + const seenTurnIds = new Set(); + let currentTurnId: string | undefined; + for (let index = cursorIndex + 1; index < hints.length; index++) { + const hint = hints[index]!; + const context = hint.parsedGoalContext; + if (!context) { + if ( + hint.claimedGoalId === goal.goalId && + hint.claimedRevision === goal.revision + ) { + throw new EvidenceSourceUnavailableError( + 'malformed_turn_context', + `Goal-owned transcript record ${hint.uuid} has malformed turn context.`, + ); + } + continue; + } + if ( + context.goalId !== goal.goalId || + context.revision !== goal.revision + ) { + continue; + } + if (context.turnId === currentTurnId) continue; + if (seenTurnIds.has(context.turnId)) { + throw new EvidenceSourceUnavailableError( + 'turn_reentry', + `Goal turn ${context.turnId} re-enters the active transcript lineage.`, + ); + } + seenTurnIds.add(context.turnId); + lineageTurnIds.push(context.turnId); + currentTurnId = context.turnId; + } + if (lineageTurnIds.at(-1) !== permit.turnId) { + throw new EvidenceSourceUnavailableError( + 'current_turn_not_tail', + 'The current Goal permit is not the tail of the active transcript lineage.', + ); + } + + this.checkpointEntries = checkpointCatalogEntries(goal); + const checkpointBytes = this.checkpointEntries.reduce( + (total, entry) => + total + Buffer.byteLength(JSON.stringify(entry), 'utf8'), + 0, + ); + let truncated = + this.checkpointEntries.length >= CATALOG_ENTRY_LIMIT || + checkpointBytes > CATALOG_BYTE_LIMIT; + const rawEntryLimit = Math.max( + 0, + CATALOG_ENTRY_LIMIT - this.checkpointEntries.length, + ); + let catalogBytes = checkpointBytes; + for ( + let index = hints.length - 1; + !truncated && index > cursorIndex; + index-- + ) { + const hint = hints[index]!; + const context = hint.parsedGoalContext; + if ( + !hint.provenance || + !context || + context.goalId !== goal.goalId || + context.revision !== goal.revision + ) { + continue; + } + if (this.candidateUuids.length >= rawEntryLimit) { + if (hint.hasRawEligibleContent) { + truncated = true; + break; + } + continue; + } + if (!hint.hasCatalogEligibleContent) continue; + const entryBytes = hint.catalogEntryBytes; + if ( + entryBytes === undefined || + catalogBytes + entryBytes > CATALOG_BYTE_LIMIT + ) { + truncated = true; + break; + } + this.candidateUuids.push(hint.uuid); + this.candidateUuidSet.add(hint.uuid); + catalogBytes += entryBytes; + } + this.truncated = truncated; + this.shouldCheckpoint = + !truncated && + this.candidateUuids.length > 0 && + (this.checkpointEntries.length + this.candidateUuids.length >= + CHECKPOINT_ENTRY_THRESHOLD || + catalogBytes >= CHECKPOINT_BYTE_THRESHOLD); + } + + getCandidateUuids(): readonly string[] { + return this.shouldCheckpoint ? this.candidateUuids : []; + } + + capture(record: GoalEvidenceRecord): void { + if (!this.shouldCheckpoint || !this.candidateUuidSet.has(record.uuid)) { + return; + } + const provenance = coherentEvidenceProvenance(record); + if (!provenance) return; + const context = parseGoalContext(record.goalContext); + if ( + !context || + context.goalId !== this.goal.goalId || + context.revision !== this.goal.revision + ) { + return; + } + const preview = evidencePreview(record, provenance); + const content = evidenceContent(record, provenance); + if (!preview || !content) return; + this.captured.set(record.uuid, { + uuid: record.uuid, + provenance, + turnId: context.turnId, + preview, + proofKind: proofKindOf(provenance), + content: capCheckpointContent(content), + }); + } + + finish(): GoalEvidenceCheckpointWindow { + const selected = this.shouldCheckpoint + ? this.candidateUuids.map((uuid) => { + const entry = this.captured.get(uuid); + if (!entry) { + throw new InvalidGoalEvidenceReferenceError( + 'ineligible_reference', + `Transcript record ${uuid} has no eligible evidence content.`, + uuid, + ); + } + return entry; + }) + : []; + selected.reverse(); + return { + previousClaims: structuredClone( + this.goal.evidenceCheckpoint?.claims ?? [], + ), + evidence: selected, + truncated: this.truncated, + shouldCheckpoint: this.shouldCheckpoint, + }; + } +} + +export function getGoalEvidenceRecordIndexHint( + record: GoalEvidenceRecord, +): GoalEvidenceRecordIndexHint { + return new GoalEvidenceRecordIndexAccumulator(record).finish(); +} + export function buildGoalEvidenceCatalog( input: GoalEvidenceContext, ): GoalEvidenceCatalog { @@ -172,37 +547,19 @@ export function buildGoalEvidenceCatalog( export function buildGoalEvidenceCheckpointWindow( input: GoalEvidenceContext, ): GoalEvidenceCheckpointWindow { - const analysis = analyzeEvidence(input); - const rawEntries = analysis.catalog.filter( - (entry) => entry.provenance !== 'goal_checkpoint', + const accumulator = new GoalEvidenceCheckpointAccumulator( + input.records.map(getGoalEvidenceRecordIndexHint), + input.goal, + input.permit, ); - const shouldCheckpoint = - !analysis.catalogTruncated && - rawEntries.length > 0 && - (analysis.catalog.length >= CHECKPOINT_ENTRY_THRESHOLD || - analysis.catalogBytes >= CHECKPOINT_BYTE_THRESHOLD); - const evidence = (shouldCheckpoint ? rawEntries : []).map((entry) => { - const recordIndex = analysis.indexByUuid.get(entry.uuid); - const record = - recordIndex === undefined ? undefined : input.records[recordIndex]; - const content = record ? evidenceContent(record, entry.provenance) : ''; - if (!content) { - throw new InvalidGoalEvidenceReferenceError( - 'ineligible_reference', - `Transcript record ${entry.uuid} has no eligible evidence content.`, - entry.uuid, - ); - } - return { ...entry, content: capCheckpointContent(content) }; - }); - return { - previousClaims: structuredClone( - input.goal.evidenceCheckpoint?.claims ?? [], - ), - evidence, - truncated: analysis.catalogTruncated, - shouldCheckpoint, - }; + const recordsByUuid = new Map( + input.records.map((record) => [record.uuid, record]), + ); + for (const uuid of accumulator.getCandidateUuids()) { + const record = recordsByUuid.get(uuid); + if (record) accumulator.capture(record); + } + return accumulator.finish(); } export function validateGoalEvidenceReferences( diff --git a/packages/core/src/goals/goal-persistence.ts b/packages/core/src/goals/goal-persistence.ts index 09874deeef..8ba9b05101 100644 --- a/packages/core/src/goals/goal-persistence.ts +++ b/packages/core/src/goals/goal-persistence.ts @@ -25,6 +25,11 @@ export type GoalRecoveryRecord = Pick & { systemPayload?: unknown; }; +export interface GoalRecoverySelection { + recovery: GoalRecovery; + sourceUuid?: string; +} + const LEGACY_ACTIVE_KINDS = new Set(['set', 'checking']); const LEGACY_STOPPED_KINDS = new Set([ 'achieved', @@ -37,7 +42,14 @@ const LEGACY_STOPPED_KINDS = new Set([ export function recoverGoalFromRecords( records: readonly GoalRecoveryRecord[], ): GoalRecovery { + return selectGoalRecoveryFromRecords(records).recovery; +} + +export function selectGoalRecoveryFromRecords( + records: readonly GoalRecoveryRecord[], +): GoalRecoverySelection { let unsupported: GoalRecovery | undefined; + let unsupportedSourceUuid: string | undefined; for (let index = records.length - 1; index >= 0; index -= 1) { const record = records[index]; if (record?.subtype !== 'goal_state') continue; @@ -45,19 +57,26 @@ export function recoverGoalFromRecords( record.type === 'system' ? parseGoalStateRecordPayloadV2(record.systemPayload) : undefined; - if (payload) return { kind: 'v2', payload }; - unsupported ??= { - kind: 'unsupported', - reason: `Goal lifecycle record ${record.uuid} is malformed or uses an unsupported version`, - }; + if (payload) { + return { recovery: { kind: 'v2', payload }, sourceUuid: record.uuid }; + } + if (!unsupported) { + unsupported = { + kind: 'unsupported', + reason: `Goal lifecycle record ${record.uuid} is malformed or uses an unsupported version`, + }; + unsupportedSourceUuid = record.uuid; + } } - return unsupported ?? recoverLegacyGoal(records); + return unsupported + ? { recovery: unsupported, sourceUuid: unsupportedSourceUuid } + : recoverLegacyGoal(records); } function recoverLegacyGoal( records: readonly GoalRecoveryRecord[], -): GoalRecovery { +): GoalRecoverySelection { for ( let recordIndex = records.length - 1; recordIndex >= 0; @@ -86,16 +105,81 @@ function recoverLegacyGoal( const kind = value['kind']; const condition = value['condition']; if (typeof kind !== 'string' || typeof condition !== 'string') { - return unsupportedLegacy(record.uuid); + return { + recovery: unsupportedLegacy(record.uuid), + sourceUuid: record.uuid, + }; + } + if (LEGACY_STOPPED_KINDS.has(kind)) { + return { recovery: { kind: 'none' }, sourceUuid: record.uuid }; } - if (LEGACY_STOPPED_KINDS.has(kind)) return { kind: 'none' }; if (!LEGACY_ACTIVE_KINDS.has(kind) || condition.trim().length === 0) { - return unsupportedLegacy(record.uuid); + return { + recovery: unsupportedLegacy(record.uuid), + sourceUuid: record.uuid, + }; } - return { kind: 'legacy', objective: condition.trim() }; + return { + recovery: { kind: 'legacy', objective: condition.trim() }, + sourceUuid: record.uuid, + }; } } - return { kind: 'none' }; + return { recovery: { kind: 'none' } }; +} + +export function normalizeGoalRecoveryRecord( + record: GoalRecoveryRecord, +): GoalRecoveryRecord | undefined { + if (record.subtype === 'goal_state') { + return { + uuid: record.uuid, + type: record.type, + subtype: record.subtype, + systemPayload: + record.type === 'system' + ? (parseGoalStateRecordPayloadV2(record.systemPayload) ?? null) + : null, + }; + } + if (record.type !== 'system' || record.subtype !== 'slash_command') { + return undefined; + } + const payload = record.systemPayload as SlashCommandRecordPayload | undefined; + if ( + payload?.phase !== 'result' || + !Array.isArray(payload.outputHistoryItems) + ) { + return undefined; + } + const goalStatusItems = payload.outputHistoryItems.filter( + (value) => isObjectRecord(value) && value['type'] === 'goal_status', + ); + if (goalStatusItems.length === 0) return undefined; + return { + uuid: record.uuid, + type: record.type, + subtype: record.subtype, + systemPayload: { + phase: 'result', + outputHistoryItems: goalStatusItems, + }, + }; +} + +export function isGoalRecoveryCandidate(record: GoalRecoveryRecord): boolean { + if (record.subtype === 'goal_state') return true; + if (record.type !== 'system' || record.subtype !== 'slash_command') { + return false; + } + const payload = record.systemPayload as SlashCommandRecordPayload | undefined; + return ( + payload?.phase === 'result' && + Array.isArray(payload.outputHistoryItems) && + payload.outputHistoryItems.some( + (value) => isObjectRecord(value) && value['type'] === 'goal_status', + ) + ); } function unsupportedLegacy(recordUuid: string): GoalRecovery { diff --git a/packages/core/src/goals/goal-runtime.test.ts b/packages/core/src/goals/goal-runtime.test.ts index ee7f1505fd..f041e2e407 100644 --- a/packages/core/src/goals/goal-runtime.test.ts +++ b/packages/core/src/goals/goal-runtime.test.ts @@ -3670,6 +3670,82 @@ describe('goal runtime', () => { }); }); + it('prepares an active restore without broadcasting or starting work', async () => { + const host = fakeGoalTurnHost(); + const runtime = createGoalRuntime({ journal: fakeGoalJournal() }); + const listener = vi.fn(); + runtime.bindHost(host); + runtime.subscribe(listener); + const record = goalStateRecord({ + v: 2, + activity: 'idle', + goal: { + goalId: 'g-selective', + revision: 1, + objective: 'resume selectively', + status: 'active', + evidenceCursor: { recordId: 'restore-record' }, + turnCount: 1, + activeTimeMs: 10, + createdAt: 1, + updatedAt: 2, + }, + }); + + await runtime.prepareRestore([record]); + + expect(runtime.getSnapshot().goal?.status).toBe('active'); + expect(listener).not.toHaveBeenCalled(); + expect(host.started).toEqual([]); + + await runtime.activateRestoredWork(); + + expect(listener).toHaveBeenCalledTimes(2); + expect(host.started).toHaveLength(1); + }); + + it('coalesces preparation and activation and rejects activation before preparation', async () => { + const runtime = createGoalRuntime({ journal: fakeGoalJournal() }); + await expect(runtime.activateRestoredWork()).rejects.toThrow( + 'preparation has not started', + ); + const record = goalStateRecord({ + v: 2, + activity: 'idle', + goal: null, + }); + + const firstPreparation = runtime.prepareRestore([record]); + const secondPreparation = runtime.prepareRestore([record]); + await Promise.all([firstPreparation, secondPreparation]); + const firstActivation = runtime.activateRestoredWork(); + const secondActivation = runtime.activateRestoredWork(); + + await expect( + Promise.all([firstActivation, secondActivation]), + ).resolves.toEqual([undefined, undefined]); + }); + + it('prevents unfinished restore preparation from committing after disposal', async () => { + let releaseAppend!: () => void; + const appendGate = new Promise((resolve) => { + releaseAppend = resolve; + }); + const runtime = createGoalRuntime({ + journal: fakeGoalJournal({ beforeAppend: () => appendGate }), + }); + const preparing = runtime.prepareRestore([legacyGoalRecord()]); + + await Promise.resolve(); + runtime.dispose(); + releaseAppend(); + + await expect(preparing).rejects.toThrow('Goal runtime has been disposed'); + await expect(runtime.activateRestoredWork()).rejects.toThrow( + 'Goal runtime has been disposed', + ); + }); + it('commits paused legacy recovery before a reentrant resume', async () => { const journal = fakeGoalJournal({ appendErrors: [new Error('migration write failed'), undefined], diff --git a/packages/core/src/goals/goal-runtime.ts b/packages/core/src/goals/goal-runtime.ts index a9f649e2f5..8cae6c229c 100644 --- a/packages/core/src/goals/goal-runtime.ts +++ b/packages/core/src/goals/goal-runtime.ts @@ -12,6 +12,7 @@ import { InvalidGoalEvidenceReferenceError, validateGoalEvidenceReferences, type GoalEvidenceCatalog, + type GoalEvidenceCheckpointWindow, type GoalEvidenceRecord, } from './goal-evidence.js'; import { @@ -126,6 +127,12 @@ export interface GoalRuntime { listener: (snapshot: GoalSnapshotV2, cause?: GoalStateCause) => void, ): () => void; restore(records: readonly GoalRecoveryRecord[]): Promise; + prepareRestore( + records: readonly GoalRecoveryRecord[], + checkpointWindow?: GoalEvidenceCheckpointWindow, + ): Promise; + getPreparedRestore(): Promise; + activateRestoredWork(): Promise; dispatch(request: GoalControlRequest): Promise; bindHost(host: GoalTurnHost): () => void; beginTurn(turnKey: string): GoalTurnPermit | undefined; @@ -214,6 +221,12 @@ export function createGoalRuntime( let nextVerifierFeedback: string | undefined; let currentTurnFeedback: string | undefined; let restored = false; + let restoreActivationPending = false; + let restorePreparation: Promise | undefined; + let restoreActivation: Promise | undefined; + let preparedRestoreCause: GoalStateCause | undefined; + let preparedRestoreHasSnapshot = false; + let preparedCheckpointWindow: GoalEvidenceCheckpointWindow | undefined; let disposed = false; let recoveryError: Error | undefined; /** @@ -349,6 +362,7 @@ export function createGoalRuntime( const queueContinuation = (cause?: GoalStateCause) => { if ( + restoreActivationPending || snapshot.goal?.status !== 'active' || currentPermit || pendingProposal || @@ -800,10 +814,13 @@ export function createGoalRuntime( }); }; - const runCheckpoint = async (attempt: CheckpointAttempt): Promise => { + const runCheckpoint = async ( + attempt: CheckpointAttempt, + preparedWindow?: GoalEvidenceCheckpointWindow, + ): Promise => { const evidenceSource = options.evidenceSource; const checkpointVerifier = options.checkpointVerifier; - if (!evidenceSource || !checkpointVerifier) { + if ((!preparedWindow && !evidenceSource) || !checkpointVerifier) { await recordCheckpointFailure( attempt, 'Goal checkpoint recovery dependencies are unavailable', @@ -812,15 +829,18 @@ export function createGoalRuntime( } try { - await evidenceSource.flush(); - if (attempt.controller.signal.aborted) return; - const records = await evidenceSource.readActiveTranscriptChain(); - if (attempt.controller.signal.aborted) return; - const window = buildGoalEvidenceCheckpointWindow({ - records, - goal: attempt.goal, - permit: attempt.permit, - }); + let window = preparedWindow; + if (!window) { + await evidenceSource!.flush(); + if (attempt.controller.signal.aborted) return; + const records = await evidenceSource!.readActiveTranscriptChain(); + if (attempt.controller.signal.aborted) return; + window = buildGoalEvidenceCheckpointWindow({ + records, + goal: attempt.goal, + permit: attempt.permit, + }); + } if (window.truncated) { await recordCheckpointFailure( attempt, @@ -900,8 +920,14 @@ export function createGoalRuntime( listeners.add(listener); return () => listeners.delete(listener); }, - restore(records: readonly GoalRecoveryRecord[]): Promise { - const restoring = enqueue( + prepareRestore( + records: readonly GoalRecoveryRecord[], + checkpointWindow?: GoalEvidenceCheckpointWindow, + ): Promise { + if (restorePreparation) return restorePreparation.then(() => undefined); + restoreActivationPending = true; + preparedCheckpointWindow = checkpointWindow; + const preparation = enqueue( async (): Promise => { assertAvailable(); if (restored) return; @@ -961,14 +987,15 @@ export function createGoalRuntime( recoveredSnapshot = structuredClone(payload.snapshot); recoveredCause = payload.cause; } + assertAvailable(); if (recoveredSnapshot) snapshot = recoveredSnapshot; recoveryError = undefined; restored = true; if (recoveredSnapshot) { recoveryCause = recoveredCause; - broadcast(recoveredCause); } - if (!checkpointAttempt) queueContinuation(); + preparedRestoreHasSnapshot = recoveredSnapshot !== undefined; + preparedRestoreCause = recoveredCause; return checkpointAttempt; } catch (error) { if (!disposed) { @@ -979,10 +1006,57 @@ export function createGoalRuntime( } }, ); - return restoring.then(async (attempt) => { - if (!attempt) return; + restorePreparation = preparation; + return preparation.then( + () => undefined, + (error) => { + if (!restored && restorePreparation === preparation) { + restorePreparation = undefined; + restoreActivation = undefined; + restoreActivationPending = false; + preparedCheckpointWindow = undefined; + } + throw error; + }, + ); + }, + getPreparedRestore(): Promise { + if (!restorePreparation) { + return Promise.reject( + new GoalPersistenceUnavailableError( + 'Goal restore preparation has not started', + ), + ); + } + return restorePreparation.then(() => undefined); + }, + activateRestoredWork(): Promise { + try { + assertAvailable(); + } catch (error) { + return Promise.reject(error); + } + if (!restorePreparation) { + return Promise.reject( + new GoalPersistenceUnavailableError( + 'Goal restore preparation has not started', + ), + ); + } + if (restoreActivation) return restoreActivation; + restoreActivation = restorePreparation.then(async (attempt) => { + assertAvailable(); + restoreActivationPending = false; + if (preparedRestoreHasSnapshot) broadcast(preparedRestoreCause); + if (!attempt) { + await enqueue(async () => { + assertAvailable(); + queueContinuation(); + }); + return; + } try { - await runCheckpoint(attempt); + await runCheckpoint(attempt, preparedCheckpointWindow); } catch { // Recovery committed before the replay began, so a failed replay // degrades instead of bricking the runtime: drop the pending @@ -990,6 +1064,11 @@ export function createGoalRuntime( await settleDanglingAttempt(attempt.permit); } }); + return restoreActivation; + }, + async restore(records: readonly GoalRecoveryRecord[]): Promise { + await this.prepareRestore(records); + await this.activateRestoredWork(); }, bindHost(nextHost: GoalTurnHost): () => void { assertOperational(); diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index fdca6770cf..9ae7c40aef 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -295,6 +295,10 @@ export * from './services/sessionRecap.js'; export * from './services/session-artifact-persistence.js'; export * from './services/session-reference-service.js'; export * from './services/sessionService.js'; +export { + collectSessionTurnState, + computeInitialTurnFromHistory, +} from './services/session-turn-state.js'; export * from './services/session-writer-lease.js'; export { decodeSessionTranscriptCursor, @@ -315,6 +319,12 @@ export { SessionTranscriptTooLargeError, } from './services/session-transcript-reader.js'; export type { + SelectiveSessionRestoreOptions, + SessionLiveRestoreProjection, + SessionRestoreProjection, + SessionRestoreReplayPage, + SessionRestoreReplaySelection, + SessionRuntimeResumeState, SessionTranscriptCursorState, SessionTranscriptReadPageOptions, SessionTranscriptRecordPage, diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 26f8d7d63e..83558f351b 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -1756,6 +1756,48 @@ describe('ChatRecordingService', () => { }); describe('legacy recorder', () => { + it('restores reduced recorder state without the full conversation', async () => { + const service = new ChatRecordingService(mockConfig, undefined, false, { + lastCompletedUuid: 'projected-leaf', + turnParentUuids: [null, 'projected-parent'], + customTitle: 'Projected title', + titleSource: 'manual', + parentSessionId: 'parent-session', + sourceType: 'channel', + sourceId: 'channel-main', + }); + + service.recordUserMessage([{ text: 'next' }]); + await service.flush(); + + const record = vi.mocked(jsonl.writeLine).mock.calls[0][1] as ChatRecord; + expect(record.parentUuid).toBe('projected-leaf'); + expect(service.getCurrentCustomTitle()).toBe('Projected title'); + expect(service.getCurrentTitleSource()).toBe('manual'); + vi.mocked(jsonl.writeLine).mockClear(); + await expect(service.recordParentSession('parent-session')).resolves.toBe( + true, + ); + await expect( + service.recordSessionSource('channel', 'channel-main'), + ).resolves.toBe(true); + expect(jsonl.writeLine).not.toHaveBeenCalled(); + }); + + it('activates a leased recorder from reduced state', async () => { + const service = new ChatRecordingService(mockConfig); + service.activate(mockLease, undefined, undefined, { + lastCompletedUuid: 'leased-projected-leaf', + turnParentUuids: [null], + }); + + service.recordUserMessage([{ text: 'next' }]); + await service.flush(); + + const record = vi.mocked(jsonl.writeLine).mock.calls[0][1] as ChatRecord; + expect(record.parentUuid).toBe('leased-projected-leaf'); + }); + it('uses the effective session writer lease gate by default', async () => { mockConfig.getExperimentalZedIntegration = vi.fn().mockReturnValue(true); mockConfig.isSessionWriterLeaseEnabled = vi.fn().mockReturnValue(false); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 613f7fd9e8..9fc990f81f 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -574,6 +574,16 @@ export type ChatRecordingFailureListener = ( event: ChatRecordingFailureEvent, ) => void | Promise; +export interface ChatRecordingRestoreState { + lastCompletedUuid: string; + turnParentUuids: Array; + customTitle?: string; + titleSource?: TitleSource; + parentSessionId?: string; + sourceType?: string; + sourceId?: string; +} + /** * Service for recording the current chat session to disk. * @@ -716,25 +726,31 @@ export class ChatRecordingService { writerLeaseRequired = config.isSessionWriterLeaseEnabled?.() ?? config.getExperimentalZedIntegration?.() ?? true, + restoreState?: ChatRecordingRestoreState, ) { this.config = config; this.writerLeaseRequired = writerLeaseRequired; const resumed = config.getResumedSessionData(); if (writerLeaseRequired) { - this.lastRecordUuid = resumed?.lastCompletedUuid ?? null; + this.lastRecordUuid = + restoreState?.lastCompletedUuid ?? resumed?.lastCompletedUuid ?? null; this.lastPersistedRecordUuid = this.lastRecordUuid; } else { this.state = 'active'; this.acceptingWrites = true; - this.restoreSessionState( - resumed - ? { - conversation: resumed.conversation ?? { messages: [] }, - lastCompletedUuid: resumed.lastCompletedUuid, - } - : undefined, - resumed ? this.readPersistedTitleInfo() : undefined, - ); + if (restoreState) { + this.restoreProjectedState(restoreState); + } else { + this.restoreSessionState( + resumed + ? { + conversation: resumed.conversation ?? { messages: [] }, + lastCompletedUuid: resumed.lastCompletedUuid, + } + : undefined, + resumed ? this.readPersistedTitleInfo() : undefined, + ); + } } } @@ -866,6 +882,20 @@ export class ChatRecordingService { } } + private restoreProjectedState(state: ChatRecordingRestoreState): void { + this.lastRecordUuid = state.lastCompletedUuid; + this.lastPersistedRecordUuid = state.lastCompletedUuid; + this.turnParentUuids = [...state.turnParentUuids]; + this.currentCustomTitle = state.customTitle; + this.currentTitleSource = state.titleSource; + this.currentParentSessionId = state.parentSessionId; + this.currentSourceType = state.sourceType; + this.currentSourceId = state.sourceId; + if (this.currentCustomTitle) { + this.bytesSinceTitleAnchor = TITLE_REANCHOR_BYTES; + } + } + activate( lease: SessionWriterLease, sessionData?: { @@ -873,6 +903,7 @@ export class ChatRecordingService { lastCompletedUuid: string | null; }, persistedTitleInfo?: { title?: string; source?: TitleSource }, + restoreState?: ChatRecordingRestoreState, ): void { if ( !this.writerLeaseRequired || @@ -882,7 +913,11 @@ export class ChatRecordingService { throw new SessionWriterUnavailableError(); } this.binding = { sessionId: lease.sessionId, lease }; - this.restoreSessionState(sessionData, persistedTitleInfo); + if (restoreState) { + this.restoreProjectedState(restoreState); + } else { + this.restoreSessionState(sessionData, persistedTitleInfo); + } this.state = 'active'; this.acceptingWrites = true; } diff --git a/packages/core/src/services/session-api-history.ts b/packages/core/src/services/session-api-history.ts new file mode 100644 index 0000000000..f8d9a74898 --- /dev/null +++ b/packages/core/src/services/session-api-history.ts @@ -0,0 +1,130 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Content, Part } from '@google/genai'; +import type { + ChatCompressionRecordPayload, + ChatRecord, +} from './chatRecordingService.js'; + +export interface BuildApiHistoryOptions { + /** + * Whether to strip thought parts from the history. + * Thought parts are content parts that have `thought: true`. + * Keeping thoughts ensures `reasoning_content` from reasoning models + * (e.g. DeepSeek) is properly passed back in subsequent API calls. + * @default false + */ + stripThoughtsFromHistory?: boolean; +} + +function stripThoughtsFromContent(content: Content): Content | null { + if (!content.parts) return content; + + const filteredParts = content.parts.filter((part) => !(part as Part).thought); + if (filteredParts.length === 0) return null; + return { ...content, parts: filteredParts }; +} + +function copyContentForApiHistory(content: Content): Content { + return { + ...content, + parts: content.parts?.map((part) => { + if ('functionCall' in part && part.functionCall) { + return { + ...part, + functionCall: { + ...part.functionCall, + args: part.functionCall.args + ? { ...part.functionCall.args } + : part.functionCall.args, + }, + }; + } + if ('functionResponse' in part && part.functionResponse) { + return { + ...part, + functionResponse: { ...part.functionResponse }, + }; + } + return { ...part }; + }), + }; +} + +function appendApiHistoryRecord(history: Content[], record: ChatRecord): void { + if (!record.message || record.subtype === 'realtime_message') return; + + const message = copyContentForApiHistory(record.message); + if (record.subtype === 'mid_turn_user_message') { + const previous = history.at(-1); + if (previous?.role === 'user') { + previous.parts = [...(previous.parts ?? []), ...(message.parts ?? [])]; + return; + } + } + + history.push(message); +} + +export class SessionApiHistoryAccumulator { + private history: Content[] = []; + private compressionCandidate: unknown; + + add(record: ChatRecord): void { + if (record.type === 'system') { + if (!isApiHistoryCompressionCandidate(record)) return; + const payload = record.systemPayload as ChatCompressionRecordPayload; + this.compressionCandidate = payload.compressedHistory; + this.history = Array.isArray(payload.compressedHistory) + ? payload.compressedHistory.map(copyContentForApiHistory) + : []; + return; + } + + if ( + this.compressionCandidate !== undefined && + !Array.isArray(this.compressionCandidate) + ) { + return; + } + appendApiHistoryRecord(this.history, record); + } + + finish(options: BuildApiHistoryOptions = {}): Content[] { + if ( + this.compressionCandidate !== undefined && + !Array.isArray(this.compressionCandidate) + ) { + return (this.compressionCandidate as Content[]).map( + copyContentForApiHistory, + ); + } + if (!options.stripThoughtsFromHistory) return this.history; + return this.history + .map(stripThoughtsFromContent) + .filter((content): content is Content => content !== null); + } +} + +export function isApiHistoryCompressionCandidate(record: ChatRecord): boolean { + if (record.type !== 'system' || record.subtype !== 'chat_compression') { + return false; + } + const payload = record.systemPayload as + | ChatCompressionRecordPayload + | undefined; + return Boolean(payload?.compressedHistory); +} + +export function buildApiHistoryFromConversation( + conversation: { messages: readonly ChatRecord[] }, + options: BuildApiHistoryOptions = {}, +): Content[] { + const accumulator = new SessionApiHistoryAccumulator(); + for (const record of conversation.messages) accumulator.add(record); + return accumulator.finish(options); +} diff --git a/packages/core/src/services/session-artifact-persistence.ts b/packages/core/src/services/session-artifact-persistence.ts index 0cb3fe792a..04c224551b 100644 --- a/packages/core/src/services/session-artifact-persistence.ts +++ b/packages/core/src/services/session-artifact-persistence.ts @@ -159,6 +159,84 @@ export function isSessionArtifactRecord( return isTranscriptArtifactRecord(record); } +export function selectActiveSideArtifactRecordUuids( + records: ReadonlyArray< + SessionArtifactChatRecordLike & { + uuid: string; + parentUuid: string | null; + } + >, + activeRecordUuids: readonly string[], +): string[] { + const activeUuids = new Set(activeRecordUuids); + const firstActiveUuid = activeRecordUuids[0]; + const firstActiveIndex = + firstActiveUuid === undefined + ? -1 + : records.findIndex((record) => record.uuid === firstActiveUuid); + const nextActiveUuidByIndex = new Map(); + const nextBlockingUuidByIndex = new Map(); + let nextActiveUuid: string | undefined; + let nextBlockingUuid: string | undefined; + for (let index = records.length - 1; index >= 0; index--) { + if (nextActiveUuid !== undefined) { + nextActiveUuidByIndex.set(index, nextActiveUuid); + } + if (nextBlockingUuid !== undefined) { + nextBlockingUuidByIndex.set(index, nextBlockingUuid); + } + const record = records[index]!; + if (activeUuids.has(record.uuid)) { + nextActiveUuid = record.uuid; + nextBlockingUuid = undefined; + } else if ( + !isSessionArtifactRecord(record) && + !(record.type === 'system' && record.subtype === 'custom_title') + ) { + nextBlockingUuid = record.uuid; + } + } + + const selected: string[] = []; + const includedSideArtifactUuids = new Set(); + let previousActiveUuid: string | undefined; + for (let index = 0; index < records.length; index++) { + const record = records[index]!; + if (activeUuids.has(record.uuid)) { + previousActiveUuid = record.uuid; + continue; + } + if (!isSessionArtifactRecord(record)) continue; + + const nextUuid = nextActiveUuidByIndex.get(index); + const isInActiveSegment = + !nextBlockingUuidByIndex.has(index) && + (nextUuid !== undefined + ? activeUuids.has(nextUuid) + : previousActiveUuid !== undefined && + activeUuids.has(previousActiveUuid)); + if ( + record.parentUuid !== null && + (activeUuids.has(record.parentUuid) || + includedSideArtifactUuids.has(record.parentUuid)) && + isInActiveSegment && + (record.parentUuid === previousActiveUuid || + includedSideArtifactUuids.has(record.parentUuid)) + ) { + selected.push(record.uuid); + includedSideArtifactUuids.add(record.uuid); + } else if ( + record.parentUuid === null && + index < firstActiveIndex && + isInActiveSegment + ) { + selected.push(record.uuid); + includedSideArtifactUuids.add(record.uuid); + } + } + return selected; +} + export function stableSessionArtifactId( sessionId: string, identityKey: string, @@ -181,80 +259,90 @@ export function sessionArtifactIdentityKey( return undefined; } -export function rebuildSessionArtifactSnapshot( - records: readonly SessionArtifactChatRecordLike[], - fallbackSessionId?: string, -): RebuiltSessionArtifactSnapshot | undefined { - const artifacts = new Map(); - const tombstonedIds = new Set(); - const stickyEphemeralIds = new Set(); - const markerArtifacts = new Map(); - const warnings: string[] = []; - let sequence = 0; - let lastSnapshotSequence = 0; - let sessionId = fallbackSessionId; - let sawRecord = false; +export class SessionArtifactSnapshotAccumulator { + private readonly artifacts = new Map(); + private readonly tombstonedIds = new Set(); + private readonly stickyEphemeralIds = new Set(); + private readonly markerArtifacts = new Map< + string, + PersistedSessionArtifact + >(); + private readonly warnings: string[] = []; + private sequence = 0; + private lastSnapshotSequence = 0; + private sessionId: string | undefined; + private sawRecord = false; - for (const record of records) { - if (!isSessionArtifactRecord(record)) continue; + constructor(fallbackSessionId?: string) { + this.sessionId = fallbackSessionId; + } + + add(record: SessionArtifactChatRecordLike): void { + if (!isSessionArtifactRecord(record)) return; if (record.subtype === 'session_artifact_snapshot') { - const payload = normalizeSnapshotPayload(record.systemPayload, warnings); - if (!payload) continue; - sawRecord = true; - sessionId = payload.sessionId; - sequence = Math.max(sequence, payload.sequence); - lastSnapshotSequence = payload.sequence; - artifacts.clear(); - tombstonedIds.clear(); - stickyEphemeralIds.clear(); - markerArtifacts.clear(); - for (const id of payload.tombstonedIds ?? []) tombstonedIds.add(id); + const payload = normalizeSnapshotPayload( + record.systemPayload, + this.warnings, + ); + if (!payload) return; + this.sawRecord = true; + this.sessionId = payload.sessionId; + this.sequence = Math.max(this.sequence, payload.sequence); + this.lastSnapshotSequence = payload.sequence; + this.artifacts.clear(); + this.tombstonedIds.clear(); + this.stickyEphemeralIds.clear(); + this.markerArtifacts.clear(); + for (const id of payload.tombstonedIds ?? []) this.tombstonedIds.add(id); for (const id of payload.stickyEphemeralIds ?? []) { - stickyEphemeralIds.add(id); + this.stickyEphemeralIds.add(id); } - const markerIds = new Set([...tombstonedIds, ...stickyEphemeralIds]); + const markerIds = new Set([ + ...this.tombstonedIds, + ...this.stickyEphemeralIds, + ]); for (const artifact of payload.markerArtifacts ?? []) { if (markerIds.has(artifact.id)) { - markerArtifacts.set(artifact.id, artifact); + this.markerArtifacts.set(artifact.id, artifact); } } for (const artifact of payload.artifacts) { if (artifact.retention === 'ephemeral') continue; - artifacts.set(artifact.id, artifact); - markerArtifacts.delete(artifact.id); + this.artifacts.set(artifact.id, artifact); + this.markerArtifacts.delete(artifact.id); } - continue; + return; } - const payload = normalizeEventPayload(record.systemPayload, warnings); - if (!payload) continue; - sawRecord = true; - sessionId = payload.sessionId; - if (payload.sequence <= lastSnapshotSequence) { - warnings.push( - `skipped stale event sequence ${payload.sequence} at or before snapshot sequence ${lastSnapshotSequence}`, + const payload = normalizeEventPayload(record.systemPayload, this.warnings); + if (!payload) return; + this.sawRecord = true; + this.sessionId = payload.sessionId; + if (payload.sequence <= this.lastSnapshotSequence) { + this.warnings.push( + `skipped stale event sequence ${payload.sequence} at or before snapshot sequence ${this.lastSnapshotSequence}`, ); - continue; + return; } - sequence = Math.max(sequence, payload.sequence); + this.sequence = Math.max(this.sequence, payload.sequence); for (const change of payload.changes) { if (change.action === 'removed') { - artifacts.delete(change.artifactId); + this.artifacts.delete(change.artifactId); if (change.reason === 'explicit') { - tombstonedIds.add(change.artifactId); - stickyEphemeralIds.delete(change.artifactId); + this.tombstonedIds.add(change.artifactId); + this.stickyEphemeralIds.delete(change.artifactId); if (change.artifact) { - markerArtifacts.set(change.artifactId, change.artifact); + this.markerArtifacts.set(change.artifactId, change.artifact); } } if (change.reason === 'eviction') { - stickyEphemeralIds.delete(change.artifactId); - markerArtifacts.delete(change.artifactId); + this.stickyEphemeralIds.delete(change.artifactId); + this.markerArtifacts.delete(change.artifactId); } if (change.reason === 'unpin_to_ephemeral') { - stickyEphemeralIds.add(change.artifactId); + this.stickyEphemeralIds.add(change.artifactId); if (change.artifact) { - markerArtifacts.set(change.artifactId, change.artifact); + this.markerArtifacts.set(change.artifactId, change.artifact); } } continue; @@ -262,29 +350,38 @@ export function rebuildSessionArtifactSnapshot( if (!change.artifact || change.artifact.retention === 'ephemeral') { continue; } - artifacts.set(change.artifact.id, change.artifact); - tombstonedIds.delete(change.artifact.id); - stickyEphemeralIds.delete(change.artifact.id); - markerArtifacts.delete(change.artifact.id); + this.artifacts.set(change.artifact.id, change.artifact); + this.tombstonedIds.delete(change.artifact.id); + this.stickyEphemeralIds.delete(change.artifact.id); + this.markerArtifacts.delete(change.artifact.id); } } - if (!sawRecord || !sessionId) { - return undefined; - } + finish(): RebuiltSessionArtifactSnapshot | undefined { + if (!this.sawRecord || !this.sessionId) return undefined; - return { - v: SESSION_ARTIFACT_PERSISTENCE_VERSION, - sessionId, - sequence, - artifacts: Array.from(artifacts.values()), - tombstonedIds: Array.from(tombstonedIds), - stickyEphemeralIds: Array.from(stickyEphemeralIds), - ...(markerArtifacts.size > 0 - ? { markerArtifacts: Array.from(markerArtifacts.values()) } - : {}), - warnings, - }; + return { + v: SESSION_ARTIFACT_PERSISTENCE_VERSION, + sessionId: this.sessionId, + sequence: this.sequence, + artifacts: Array.from(this.artifacts.values()), + tombstonedIds: Array.from(this.tombstonedIds), + stickyEphemeralIds: Array.from(this.stickyEphemeralIds), + ...(this.markerArtifacts.size > 0 + ? { markerArtifacts: Array.from(this.markerArtifacts.values()) } + : {}), + warnings: this.warnings, + }; + } +} + +export function rebuildSessionArtifactSnapshot( + records: readonly SessionArtifactChatRecordLike[], + fallbackSessionId?: string, +): RebuiltSessionArtifactSnapshot | undefined { + const accumulator = new SessionArtifactSnapshotAccumulator(fallbackSessionId); + for (const record of records) accumulator.add(record); + return accumulator.finish(); } export function remapSessionArtifactPayloadForFork( diff --git a/packages/core/src/services/session-file-history-state.test.ts b/packages/core/src/services/session-file-history-state.test.ts new file mode 100644 index 0000000000..81eb39508e --- /dev/null +++ b/packages/core/src/services/session-file-history-state.test.ts @@ -0,0 +1,90 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import type { + ChatRecord, + FileHistorySnapshotRecordPayload, +} from './chatRecordingService.js'; +import { MAX_SNAPSHOTS } from './fileHistoryService.js'; +import { SessionFileHistoryAccumulator } from './session-file-history-state.js'; + +function snapshotRecord( + snapshots: FileHistorySnapshotRecordPayload['snapshots'], +): Pick { + return { + type: 'system', + subtype: 'file_history_snapshot', + systemPayload: { snapshots }, + }; +} + +function snapshot(promptId: string, timestamp: string) { + return { + promptId, + timestamp, + trackedFileBackups: {}, + }; +} + +describe('SessionFileHistoryAccumulator', () => { + it('keeps the final 100 first-insertion slots with retained replacements', () => { + const accumulator = new SessionFileHistoryAccumulator(); + accumulator.add( + snapshotRecord( + Array.from({ length: MAX_SNAPSHOTS + 1 }, (_, index) => + snapshot( + `prompt-${index}`, + `2026-01-01T00:00:${String(index % 60).padStart(2, '0')}.000Z`, + ), + ), + ), + ); + accumulator.add( + snapshotRecord([ + snapshot('prompt-0', '2026-02-01T00:00:00.000Z'), + snapshot('prompt-50', '2026-03-01T00:00:00.000Z'), + ]), + ); + + const restored = accumulator.finish(); + + expect(restored).toHaveLength(MAX_SNAPSHOTS); + expect(restored?.map((item) => item.promptId)).toEqual( + Array.from( + { length: MAX_SNAPSHOTS }, + (_, index) => `prompt-${index + 1}`, + ), + ); + expect( + restored?.find((item) => item.promptId === 'prompt-50')?.timestamp, + ).toEqual(new Date('2026-03-01T00:00:00.000Z')); + }); + + it('does not partially apply a malformed snapshot batch', () => { + const accumulator = new SessionFileHistoryAccumulator(); + accumulator.add( + snapshotRecord([snapshot('before', '2026-01-01T00:00:00.000Z')]), + ); + + expect(() => + accumulator.add( + snapshotRecord([ + snapshot('partial', '2026-01-02T00:00:00.000Z'), + { + promptId: 'malformed', + timestamp: '2026-01-03T00:00:00.000Z', + trackedFileBackups: null, + } as never, + ]), + ), + ).toThrow(); + + expect(accumulator.finish()?.map((item) => item.promptId)).toEqual([ + 'before', + ]); + }); +}); diff --git a/packages/core/src/services/session-file-history-state.ts b/packages/core/src/services/session-file-history-state.ts new file mode 100644 index 0000000000..8d0d7ff5c8 --- /dev/null +++ b/packages/core/src/services/session-file-history-state.ts @@ -0,0 +1,56 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + ChatRecord, + FileHistorySnapshotRecordPayload, +} from './chatRecordingService.js'; +import { + deserializeSnapshots, + MAX_SNAPSHOTS, + type FileHistorySnapshot, +} from './fileHistoryService.js'; + +export class SessionFileHistoryAccumulator { + private readonly seenPromptIds = new Set(); + private readonly retainedPromptIds: string[] = []; + private readonly snapshotsByPromptId = new Map(); + + add(record: Pick): void { + if ( + record.type !== 'system' || + record.subtype !== 'file_history_snapshot' || + !record.systemPayload + ) { + return; + } + const payload = record.systemPayload as FileHistorySnapshotRecordPayload; + if (!Array.isArray(payload.snapshots)) return; + const deserialized = deserializeSnapshots(payload.snapshots); + for (const snapshot of deserialized) { + if (this.seenPromptIds.has(snapshot.promptId)) { + if (this.snapshotsByPromptId.has(snapshot.promptId)) { + this.snapshotsByPromptId.set(snapshot.promptId, snapshot); + } + continue; + } + this.seenPromptIds.add(snapshot.promptId); + this.retainedPromptIds.push(snapshot.promptId); + this.snapshotsByPromptId.set(snapshot.promptId, snapshot); + if (this.retainedPromptIds.length > MAX_SNAPSHOTS) { + const evictedPromptId = this.retainedPromptIds.shift()!; + this.snapshotsByPromptId.delete(evictedPromptId); + } + } + } + + finish(): FileHistorySnapshot[] | undefined { + const snapshots = this.retainedPromptIds.map( + (promptId) => this.snapshotsByPromptId.get(promptId)!, + ); + return snapshots.length > 0 ? snapshots : undefined; + } +} diff --git a/packages/core/src/services/session-resume-token-counts.ts b/packages/core/src/services/session-resume-token-counts.ts new file mode 100644 index 0000000000..50bc956028 --- /dev/null +++ b/packages/core/src/services/session-resume-token-counts.ts @@ -0,0 +1,81 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + ChatCompressionRecordPayload, + ChatRecord, +} from './chatRecordingService.js'; +import { getUsageOutputTokenCountForPromptEstimate } from './tokenEstimation.js'; + +export interface ResumeTokenCounts { + promptTokenCount: number; + outputTokenCount: number; + isEstimated: boolean; +} + +export class ResumeTokenCountsAccumulator { + private value: ResumeTokenCounts | undefined; + + add(record: ChatRecord): void { + if (record.type === 'assistant') { + const usage = record.usageMetadata; + const candidate = usage?.promptTokenCount ?? usage?.totalTokenCount; + if (candidate) { + this.value = { + promptTokenCount: candidate, + outputTokenCount: getUsageOutputTokenCountForPromptEstimate(usage), + isEstimated: false, + }; + } + return; + } + + if (record.type === 'system' && record.subtype === 'chat_compression') { + const payload = record.systemPayload as + | ChatCompressionRecordPayload + | undefined; + if (payload?.info) { + this.value = { + promptTokenCount: payload.info.newTokenCount, + outputTokenCount: 0, + isEstimated: payload.info.newTokenCountIsEstimated ?? true, + }; + } + } + } + + finish(): ResumeTokenCounts | undefined { + return this.value; + } +} + +export function isResumeTokenCountsCandidate(record: ChatRecord): boolean { + if (record.type === 'assistant') { + const usage = record.usageMetadata; + return Boolean(usage?.promptTokenCount ?? usage?.totalTokenCount); + } + if (record.type !== 'system' || record.subtype !== 'chat_compression') { + return false; + } + const payload = record.systemPayload as + | ChatCompressionRecordPayload + | undefined; + return payload?.info !== undefined; +} + +export function getResumeTokenCounts(conversation: { + messages: readonly ChatRecord[]; +}): ResumeTokenCounts | undefined { + const accumulator = new ResumeTokenCountsAccumulator(); + for (const record of conversation.messages) accumulator.add(record); + return accumulator.finish(); +} + +export function getResumePromptTokenCount(conversation: { + messages: readonly ChatRecord[]; +}): number | undefined { + return getResumeTokenCounts(conversation)?.promptTokenCount; +} diff --git a/packages/core/src/services/session-transcript-reader.test.ts b/packages/core/src/services/session-transcript-reader.test.ts index e7cf487f8c..bebf4a0389 100644 --- a/packages/core/src/services/session-transcript-reader.test.ts +++ b/packages/core/src/services/session-transcript-reader.test.ts @@ -9,11 +9,12 @@ import * as os from 'node:os'; import * as path from 'node:path'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; -const { mockDebugLogger } = vi.hoisted(() => ({ +const { mockDebugLogger, mockAddDaemonRequestAttribute } = vi.hoisted(() => ({ mockDebugLogger: { debug: vi.fn(), warn: vi.fn(), }, + mockAddDaemonRequestAttribute: vi.fn(), })); const { statFault } = vi.hoisted(() => ({ @@ -24,6 +25,10 @@ vi.mock('../utils/debugLogger.js', () => ({ createDebugLogger: () => mockDebugLogger, })); +vi.mock('../telemetry/daemon-tracing.js', () => ({ + addDaemonRequestAttribute: mockAddDaemonRequestAttribute, +})); + vi.mock('node:fs/promises', async (importOriginal) => { const actual = await importOriginal(); return { @@ -41,6 +46,20 @@ vi.mock('node:fs/promises', async (importOriginal) => { import { Storage } from '../config/storage.js'; import type { ChatRecord } from './chatRecordingService.js'; import { + buildApiHistoryFromConversation, + getResumeTokenCounts, + SessionService, +} from './sessionService.js'; +import { collectSessionTurnState } from './session-turn-state.js'; +import { recoverGoalFromRecords } from '../goals/goal-persistence.js'; +import type { GoalStateRecordPayloadV2 } from '../goals/goal-protocol.js'; +import { buildGoalEvidenceCheckpointWindow } from '../goals/goal-evidence.js'; +import { + SESSION_ARTIFACT_PERSISTENCE_VERSION, + stableSessionArtifactId, +} from './session-artifact-persistence.js'; +import { + clearSessionTranscriptIndexCacheEntriesForTest, encodeSessionTranscriptCursor, getSessionTranscriptIndexCacheStatsForTest, InvalidSessionTranscriptCursorError, @@ -48,8 +67,11 @@ import { SESSION_TRANSCRIPT_MAX_INDEX_BYTES, SESSION_TRANSCRIPT_MAX_LIMIT, resetSessionTranscriptIndexCacheForTest, + setSessionTranscriptCooperativeReadBudgetForTest, setSessionTranscriptExpandedPageBytesForTest, + setSessionTranscriptIndexBuildCompleteHookForTest, setSessionTranscriptIndexCacheMaxBytesForTest, + setSessionTranscriptSelectedLineReadHookForTest, SessionTranscriptCursorCodec, SessionTranscriptSnapshotUnavailableError, SessionTranscriptReader, @@ -600,6 +622,1392 @@ describe('SessionTranscriptReader', () => { expect(page.hasMore).toBe(false); }); + it('derives the side-task boundary from the active chain', async () => { + const sessionSource = { + ...record('source', null, 'session source'), + type: 'system' as const, + subtype: 'session_source' as const, + systemPayload: { + sourceType: 'side_task', + sourceId: 'parent-session', + }, + }; + const inheritedUser = { + ...record('parent-u1', 'source', 'parent prompt'), + forkedFrom: { + sessionId: 'parent-session', + messageUuid: 'parent-u1', + }, + }; + const deadBranchSource = { + ...record('dead-source', 'parent-u1', 'dead source'), + type: 'system' as const, + subtype: 'session_source' as const, + systemPayload: { + sourceType: 'side_task', + sourceId: 'abandoned-parent', + }, + }; + await writeRecords([ + sessionSource, + inheritedUser, + record('side-u1', 'parent-u1', 'side prompt'), + deadBranchSource, + record('side-a1', 'side-u1', 'side answer'), + ]); + + const page = await new SessionTranscriptReader(workspaceDir).readPage( + sessionId, + { direction: 'backward', limit: 100 }, + ); + + expect(page.records.map((item) => item.uuid)).toEqual([ + 'source', + 'side-u1', + 'side-a1', + ]); + }); + + it('derives a fragmented session source from the first fragment', async () => { + const source: ChatRecord = { + ...record('source', null, ''), + type: 'system', + subtype: 'session_source', + message: undefined, + systemPayload: { sourceType: 'daemon' }, + }; + const conflictingFragment: ChatRecord = { + ...source, + systemPayload: { sourceType: 'side_task', sourceId: 'parent-session' }, + }; + const inherited = { + ...record('parent-u1', 'source', 'parent prompt'), + forkedFrom: { + sessionId: 'parent-session', + messageUuid: 'parent-u1', + }, + }; + await writeRecords([ + source, + conflictingFragment, + inherited, + record('u1', 'parent-u1', 'current prompt'), + ]); + + const page = await new SessionTranscriptReader(workspaceDir).readPage( + sessionId, + ); + + expect(page.records.map(({ uuid }) => uuid)).toEqual([ + 'source', + 'parent-u1', + 'u1', + ]); + }); + + it('builds a cold runtime projection with full-loader parity', async () => { + const source: ChatRecord = { + ...record('source', null, 'source'), + type: 'system', + subtype: 'session_source', + message: undefined, + systemPayload: { sourceType: 'daemon', sourceId: 'restore-test' }, + }; + const firstUser = record('u1', 'source', 'first prompt') as ChatRecord & { + promptId: string; + }; + firstUser.promptId = `${sessionId}########3`; + const firstAssistant: ChatRecord = { + ...record('a1', 'u1', 'first answer'), + usageMetadata: { + promptTokenCount: 30, + candidatesTokenCount: 4, + totalTokenCount: 34, + }, + }; + const compression: ChatRecord = { + ...record('compression', 'a1', ''), + type: 'system', + subtype: 'chat_compression', + message: undefined, + systemPayload: { + compressedHistory: [ + { role: 'user', parts: [{ text: 'compressed prompt' }] }, + { role: 'model', parts: [{ text: 'compressed answer' }] }, + ], + info: { newTokenCount: 20, newTokenCountIsEstimated: false }, + } as ChatRecord['systemPayload'], + }; + const uiEvent = { prompt_id: `${sessionId}########7`, duration_ms: 12 }; + const telemetry: ChatRecord = { + ...record('telemetry', 'compression', ''), + type: 'system', + subtype: 'ui_telemetry', + message: undefined, + systemPayload: { uiEvent } as ChatRecord['systemPayload'], + }; + const attributionSnapshot = { v: 1, commits: [] }; + const attribution: ChatRecord = { + ...record('attribution', 'telemetry', ''), + type: 'system', + subtype: 'attribution_snapshot', + message: undefined, + systemPayload: { + snapshot: attributionSnapshot, + } as unknown as ChatRecord['systemPayload'], + }; + const fileHistory: ChatRecord = { + ...record('files', 'attribution', ''), + type: 'system', + subtype: 'file_history_snapshot', + message: undefined, + systemPayload: { + snapshots: [ + { + promptId: `${sessionId}########3`, + timestamp: '2026-01-01T00:00:00.000Z', + trackedFileBackups: {}, + }, + ], + }, + }; + const goalPayload: GoalStateRecordPayloadV2 = { + v: 2, + cause: 'create', + snapshot: { + v: 2, + activity: 'idle', + goal: { + goalId: 'goal-1', + revision: 1, + objective: 'finish the restore', + status: 'active', + evidenceCursor: { recordId: 'goal' }, + turnCount: 1, + activeTimeMs: 0, + createdAt: 1, + updatedAt: 1, + }, + }, + }; + const goal: ChatRecord = { + ...record('goal', 'files', ''), + type: 'system', + subtype: 'goal_state', + message: undefined, + systemPayload: goalPayload, + }; + const notification: ChatRecord = { + ...record('notification', 'goal', 'background result'), + subtype: 'notification', + systemPayload: { + backgroundTask: { taskId: 'task-1' }, + } as ChatRecord['systemPayload'], + }; + const artifactId = stableSessionArtifactId( + sessionId, + 'url:https://example.com/report', + ); + const artifact: ChatRecord = { + ...record('artifact', 'notification', ''), + type: 'system', + subtype: 'session_artifact_event', + message: undefined, + systemPayload: { + v: SESSION_ARTIFACT_PERSISTENCE_VERSION, + sessionId, + sequence: 1, + recordedAt: '2026-01-01T00:00:01.000Z', + changes: [ + { + action: 'created', + artifactId, + artifact: { + id: artifactId, + kind: 'link', + storage: 'external_url', + source: 'client', + status: 'available', + title: 'Report', + url: 'https://example.com/report', + retention: 'restorable', + clientRetained: true, + createdAt: '2026-01-01T00:00:01.000Z', + updatedAt: '2026-01-01T00:00:01.000Z', + persistedAt: '2026-01-01T00:00:01.000Z', + }, + }, + ], + }, + }; + const secondAssistant = record('a2', 'notification', 'after compression'); + const title: ChatRecord = { + ...record('title', 'a2', ''), + type: 'system', + subtype: 'custom_title', + message: undefined, + systemPayload: { customTitle: 'Restored', titleSource: 'manual' }, + }; + await writeRecords([ + source, + firstUser, + firstAssistant, + compression, + telemetry, + attribution, + fileHistory, + goal, + notification, + artifact, + secondAssistant, + title, + ]); + + const service = new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }); + const [loaded, projection] = await Promise.all([ + service.loadSession(sessionId), + service.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ]); + + expect(loaded).toBeDefined(); + expect(projection).toBeDefined(); + expect(projection?.replay).toBeUndefined(); + expect(projection?.runtime.apiHistory).toEqual( + buildApiHistoryFromConversation(loaded!.conversation), + ); + expect(projection?.runtime.resumeTokenCounts).toEqual( + getResumeTokenCounts(loaded!.conversation), + ); + expect(projection?.runtime.fileHistorySnapshots).toEqual( + loaded?.fileHistorySnapshots, + ); + expect(projection?.runtime.artifactSnapshot).toEqual( + loaded?.artifactSnapshot, + ); + expect(projection?.runtime.recording.lastCompletedUuid).toBe( + loaded?.lastCompletedUuid, + ); + const expectedTurnState = collectSessionTurnState( + loaded!.conversation.messages, + sessionId, + ); + expect(projection?.runtime.recording.turnParentUuids).toEqual( + expectedTurnState.turnParentUuids, + ); + expect(projection?.runtime.initialTurn).toBe(expectedTurnState.initialTurn); + expect(projection?.runtime.backgroundNotificationTaskIds).toEqual( + expectedTurnState.backgroundNotificationTaskIds, + ); + expect(projection?.runtime.uiTelemetryEvents).toEqual([uiEvent]); + expect(projection?.runtime.attributionSnapshot).toEqual( + attributionSnapshot, + ); + expect(projection?.runtime.recording).toMatchObject({ + customTitle: 'Restored', + titleSource: 'manual', + sourceType: 'daemon', + sourceId: 'restore-test', + }); + expect(recoverGoalFromRecords(projection!.runtime.goalRecords)).toEqual( + recoverGoalFromRecords(loaded!.conversation.messages), + ); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.transcript_index_ms', + expect.any(Number), + ); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.resume_state_select_ms', + expect.any(Number), + ); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.selected_record_read_ms', + expect.any(Number), + ); + for (const attribute of [ + 'transcript_bytes', + 'records_indexed', + 'active_records', + 'selected_records', + 'selected_bytes', + 'replay_records', + 'replay_bytes', + ]) { + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + `qwen-code.daemon.session_restore.${attribute}`, + expect.any(Number), + ); + } + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.index_cache_state', + 'fresh', + ); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.replay_mode', + 'none', + ); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.compression_selected', + expect.any(Boolean), + ); + }); + + it('uses the bounded persisted-title picker instead of the full active chain', async () => { + const largeText = 'x'.repeat(70 * 1024); + const title: ChatRecord = { + ...record('title', 'u1', ''), + type: 'system', + subtype: 'custom_title', + message: undefined, + systemPayload: { customTitle: 'Middle title', titleSource: 'manual' }, + }; + await writeRecords([ + record('u1', null, largeText), + title, + record('a1', 'title', largeText), + ]); + + const service = new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }); + const projection = await service.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }); + + expect(service.getSessionTitleInfo(sessionId)).toEqual({}); + expect(projection?.runtime.recording.customTitle).toBeUndefined(); + expect(projection?.runtime.recording.titleSource).toBeUndefined(); + }); + + it('matches full-loader artifact selection across an abandoned branch', async () => { + const activeArtifactId = stableSessionArtifactId( + sessionId, + 'url:https://example.com/active', + ); + const abandonedArtifactId = stableSessionArtifactId( + sessionId, + 'url:https://example.com/abandoned', + ); + const artifactRecord = ( + uuid: string, + parentUuid: string, + sequence: number, + artifactId: string, + url: string, + ): ChatRecord => ({ + ...record(uuid, parentUuid, ''), + type: 'system', + subtype: 'session_artifact_event', + message: undefined, + systemPayload: { + v: SESSION_ARTIFACT_PERSISTENCE_VERSION, + sessionId, + sequence, + recordedAt: `2026-01-01T00:00:0${sequence}.000Z`, + changes: [ + { + action: 'created', + artifactId, + artifact: { + id: artifactId, + kind: 'link', + storage: 'external_url', + source: 'client', + status: 'available', + title: url, + url, + retention: 'restorable', + clientRetained: true, + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + persistedAt: '2026-01-01T00:00:00.000Z', + }, + }, + ], + }, + }); + await writeRecords([ + record('u1', null, 'prompt'), + record('a1', 'u1', 'answer'), + record('abandoned', 'a1', 'dead branch'), + artifactRecord( + 'abandoned-artifact', + 'abandoned', + 2, + abandonedArtifactId, + 'https://example.com/abandoned', + ), + artifactRecord( + 'active-artifact', + 'a1', + 1, + activeArtifactId, + 'https://example.com/active', + ), + record('u2', 'a1', 'next prompt'), + record('a2', 'u2', 'next answer'), + ]); + + const service = new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }); + const [loaded, projection] = await Promise.all([ + service.loadSession(sessionId), + service.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ]); + + expect(projection?.runtime.artifactSnapshot).toEqual( + loaded?.artifactSnapshot, + ); + expect( + projection?.runtime.artifactSnapshot?.artifacts.map(({ id }) => id), + ).toEqual([activeArtifactId]); + }); + + it('uses a leading artifact record as the restore start time', async () => { + const leadingTimestamp = '2025-12-31T23:59:59.000Z'; + const artifactId = stableSessionArtifactId( + sessionId, + 'url:https://example.com/leading', + ); + const leadingArtifact: ChatRecord = { + ...record('leading-artifact', null, ''), + timestamp: leadingTimestamp, + type: 'system', + subtype: 'session_artifact_event', + message: undefined, + systemPayload: { + v: SESSION_ARTIFACT_PERSISTENCE_VERSION, + sessionId, + sequence: 1, + recordedAt: leadingTimestamp, + changes: [ + { + action: 'created', + artifactId, + artifact: { + id: artifactId, + kind: 'link', + storage: 'external_url', + source: 'client', + status: 'available', + title: 'Leading artifact', + url: 'https://example.com/leading', + retention: 'restorable', + clientRetained: true, + createdAt: leadingTimestamp, + updatedAt: leadingTimestamp, + persistedAt: leadingTimestamp, + }, + }, + ], + }, + }; + await writeRecords([leadingArtifact, record('u1', null, 'prompt')]); + + const service = new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }); + const [loaded, projection] = await Promise.all([ + service.loadSession(sessionId), + service.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ]); + + expect(projection?.startTime).toBe(loaded?.conversation.startTime); + expect(projection?.startTime).toBe(leadingTimestamp); + }); + + it('preserves physical active-fragment markers for artifact selection', async () => { + const artifactId = stableSessionArtifactId( + sessionId, + 'url:https://example.com/fragmented', + ); + const artifact: ChatRecord = { + ...record('artifact', 'u1', ''), + type: 'system', + subtype: 'session_artifact_event', + message: undefined, + systemPayload: { + v: SESSION_ARTIFACT_PERSISTENCE_VERSION, + sessionId, + sequence: 1, + recordedAt: '2026-01-01T00:00:01.000Z', + changes: [ + { + action: 'created', + artifactId, + artifact: { + id: artifactId, + kind: 'link', + storage: 'external_url', + source: 'client', + status: 'available', + title: 'Fragmented', + url: 'https://example.com/fragmented', + retention: 'restorable', + clientRetained: true, + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-01T00:00:00.000Z', + persistedAt: '2026-01-01T00:00:00.000Z', + }, + }, + ], + }, + }; + await writeRecords([ + record('u1', null, 'first fragment'), + artifact, + record('u1', null, 'second fragment'), + record('dead', null, 'abandoned blocker'), + record('u1', null, 'final fragment'), + ]); + + const service = new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }); + const [loaded, projection] = await Promise.all([ + service.loadSession(sessionId), + service.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ]); + + expect(projection?.runtime.artifactSnapshot).toEqual( + loaded?.artifactSnapshot, + ); + expect( + projection?.runtime.artifactSnapshot?.artifacts.map(({ id }) => id), + ).toEqual([artifactId]); + }); + + it('selects recent restore replay with an older Goal bootstrap', async () => { + const goalPayload: GoalStateRecordPayloadV2 = { + v: 2, + cause: 'create', + snapshot: { + v: 2, + activity: 'idle', + goal: { + goalId: 'goal-1', + revision: 1, + objective: 'continue', + status: 'active', + evidenceCursor: { recordId: 'goal' }, + turnCount: 1, + activeTimeMs: 0, + createdAt: 1, + updatedAt: 1, + }, + }, + }; + await writeRecords([ + record('u1', null, 'one'), + record('a1', 'u1', 'one answer'), + { + ...record('goal', 'a1', ''), + type: 'system', + subtype: 'goal_state', + message: undefined, + systemPayload: goalPayload, + }, + record('u2', 'goal', 'two'), + record('a2', 'u2', 'two answer'), + record('u3', 'a2', 'three'), + record('a3', 'u3', 'three answer'), + ]); + + const projection = await new SessionTranscriptReader( + workspaceDir, + ).readRestoreProjection(sessionId, { + replay: { + kind: 'recent', + limit: 2, + hideInheritedHistory: false, + }, + }); + + expect(projection?.replay).toMatchObject({ + records: [ + expect.objectContaining({ uuid: 'u3' }), + expect.objectContaining({ uuid: 'a3' }), + ], + hasMore: true, + anchorRecordId: 'u3', + replay: { + goalState: goalPayload.snapshot, + goalCause: goalPayload.cause, + }, + }); + }); + + it('applies inherited-history filtering only to replay', async () => { + const inheritedUser = { + ...record('u1', null, 'inherited prompt'), + forkedFrom: { sessionId: 'parent', messageUuid: 'u1' }, + }; + const inheritedAssistant = { + ...record('a1', 'u1', 'inherited answer'), + forkedFrom: { sessionId: 'parent', messageUuid: 'a1' }, + }; + const inheritedGoal: ChatRecord = { + ...record('goal', 'a1', ''), + type: 'system', + subtype: 'slash_command', + message: undefined, + forkedFrom: { sessionId: 'parent', messageUuid: 'goal' }, + systemPayload: { + rawCommand: '/goal inherited goal', + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'inherited goal', + }, + ], + }, + }; + await writeRecords([ + inheritedUser, + inheritedAssistant, + inheritedGoal, + record('u2', 'goal', 'branch prompt'), + record('a2', 'u2', 'branch answer'), + ]); + const reader = new SessionTranscriptReader(workspaceDir); + + const visible = await reader.readRestoreProjection(sessionId, { + replay: { kind: 'all', hideInheritedHistory: false }, + }); + const hidden = await reader.readRestoreProjection(sessionId, { + replay: { kind: 'all', hideInheritedHistory: true }, + }); + + expect(visible?.replay?.records.map((item) => item.uuid)).toEqual([ + 'u1', + 'a1', + 'goal', + 'u2', + 'a2', + ]); + expect(hidden?.replay?.records.map((item) => item.uuid)).toEqual([ + 'u2', + 'a2', + ]); + expect(hidden?.runtime.apiHistory).toEqual(visible?.runtime.apiHistory); + expect(visible?.runtime.goalRecoverySourceUuid).toBe('goal'); + expect(hidden?.runtime.goalRecoverySourceUuid).toBe('goal'); + expect(hidden?.replay?.goalRecoverySourceUuid).toBeUndefined(); + expect(hidden?.runtime.goalRecords).toEqual([ + expect.objectContaining({ uuid: 'goal' }), + ]); + + const hiddenLive = await reader.readLiveRestoreProjection(sessionId, { + replay: { kind: 'all', hideInheritedHistory: true }, + }); + expect(hiddenLive?.replay?.records.map((item) => item.uuid)).toEqual([ + 'u2', + 'a2', + ]); + expect(hiddenLive?.goalRecoverySourceUuid).toBeUndefined(); + expect(hiddenLive?.goalRecords).toBeUndefined(); + }); + + it('selects Goal bootstrap precedence from filtered visible history', async () => { + const visibleGoal: ChatRecord = { + ...record('visible-goal', null, ''), + type: 'system', + subtype: 'slash_command', + message: undefined, + systemPayload: { + rawCommand: '/goal visible goal', + phase: 'result', + outputHistoryItems: [ + { type: 'goal_status', kind: 'set', condition: 'visible goal' }, + ], + }, + }; + const hiddenGoal: ChatRecord = { + ...record('hidden-goal', 'visible-goal', ''), + type: 'system', + subtype: 'slash_command', + message: undefined, + forkedFrom: { sessionId: 'parent', messageUuid: 'hidden-goal' }, + systemPayload: { + rawCommand: '/goal hidden goal', + phase: 'result', + outputHistoryItems: [ + { type: 'goal_status', kind: 'set', condition: 'hidden goal' }, + ], + }, + }; + await writeRecords([ + record('u0', null, 'older prompt'), + record('a0', 'u0', 'older answer'), + { ...visibleGoal, parentUuid: 'a0' }, + hiddenGoal, + record('u1', 'hidden-goal', 'branch prompt'), + record('a1', 'u1', 'branch answer'), + record('u2', 'a1', 'latest prompt'), + record('a2', 'u2', 'latest answer'), + ]); + const reader = new SessionTranscriptReader(workspaceDir); + const options = { + replay: { + kind: 'recent' as const, + limit: 2, + hideInheritedHistory: true, + }, + }; + + const cold = await reader.readRestoreProjection(sessionId, options); + const live = await reader.readLiveRestoreProjection(sessionId, options); + + expect(cold?.replay?.records.map((record) => record.uuid)).toEqual([ + 'u2', + 'a2', + ]); + expect(cold?.runtime.goalRecoverySourceUuid).toBe('hidden-goal'); + expect(cold?.replay?.goalRecoverySourceUuid).toBe('visible-goal'); + expect(cold?.replay?.goalBootstrapRecords).toEqual([ + expect.objectContaining({ uuid: 'visible-goal' }), + ]); + expect(live?.goalRecoverySourceUuid).toBe('visible-goal'); + expect(live?.goalRecords?.map((record) => record.uuid)).toEqual([ + 'visible-goal', + ]); + }); + + it('uses the aggregated final usage metadata when selecting resume tokens', async () => { + const earlier = record('a1', 'u1', 'earlier answer'); + earlier.usageMetadata = { + promptTokenCount: 20, + candidatesTokenCount: 3, + }; + const latest = record('a2', 'u2', 'latest answer'); + latest.usageMetadata = { + promptTokenCount: 10, + candidatesTokenCount: 2, + }; + const latestFragment = record('a2', 'u2', 'latest tail'); + latestFragment.usageMetadata = { + promptTokenCount: 0, + totalTokenCount: 0, + candidatesTokenCount: 0, + }; + await writeRecords([ + record('u1', null, 'first'), + earlier, + record('u2', 'a1', 'second'), + latest, + latestFragment, + ]); + + const service = new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }); + const [loaded, projection] = await Promise.all([ + service.loadSession(sessionId), + service.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ]); + + expect(projection?.runtime.resumeTokenCounts).toEqual( + getResumeTokenCounts(loaded!.conversation), + ); + expect(projection?.runtime.resumeTokenCounts).toMatchObject({ + promptTokenCount: 20, + }); + }); + + it('ignores a later non-snapshot attribution record like the full loader', async () => { + const snapshot = { + type: 'attribution-snapshot' as const, + version: 1, + surface: 'cli', + fileStates: {}, + promptCount: 2, + promptCountAtLastCommit: 1, + }; + const valid: ChatRecord = { + ...record('attribution', 'u1', ''), + type: 'system', + subtype: 'attribution_snapshot', + message: undefined, + systemPayload: { snapshot }, + }; + const malformed: ChatRecord = { + ...record('malformed-attribution', 'attribution', ''), + type: 'system', + subtype: 'attribution_snapshot', + message: undefined, + systemPayload: { + ignored: true, + } as unknown as ChatRecord['systemPayload'], + }; + await writeRecords([ + record('u1', null, 'prompt'), + valid, + malformed, + record('a1', 'malformed-attribution', 'answer'), + ]); + + const projection = await new SessionTranscriptReader( + workspaceDir, + ).readRestoreProjection(sessionId, { replay: { kind: 'none' } }); + + expect(projection?.runtime.attributionSnapshot).toEqual(snapshot); + }); + + it('rejects a cold projection when the frozen transcript changes', async () => { + const filePath = await writeRecords([record('u1', null, 'one')]); + let appended = false; + setSessionTranscriptIndexBuildCompleteHookForTest(async (builtPath) => { + if (builtPath !== filePath || appended) return; + appended = true; + await fs.appendFile( + filePath, + `${JSON.stringify(record('a1', 'u1', 'late'))}\n`, + 'utf8', + ); + }); + + await expect( + new SessionTranscriptReader(workspaceDir).readRestoreProjection( + sessionId, + { replay: { kind: 'none' } }, + ), + ).rejects.toBeInstanceOf(SessionTranscriptSnapshotUnavailableError); + expect(getSessionTranscriptIndexCacheStatsForTest()).toEqual({ + entries: 0, + byteSize: 0, + }); + }); + + it('reports a removed frozen transcript as snapshot-unavailable', async () => { + const filePath = await writeRecords([record('u1', null, 'one')]); + setSessionTranscriptIndexBuildCompleteHookForTest(async (builtPath) => { + if (builtPath === filePath) await fs.unlink(filePath); + }); + + await expect( + new SessionTranscriptReader(workspaceDir).readRestoreProjection( + sessionId, + { replay: { kind: 'none' } }, + ), + ).rejects.toBeInstanceOf(SessionTranscriptSnapshotUnavailableError); + expect(getSessionTranscriptIndexCacheStatsForTest()).toEqual({ + entries: 0, + byteSize: 0, + }); + }); + + it('rejects an invalid recent projection before scanning', async () => { + const filePath = await writeRecords([record('u1', null, 'one')]); + let buildCount = 0; + setSessionTranscriptIndexBuildCompleteHookForTest((builtPath) => { + if (builtPath === filePath) buildCount++; + }); + + await expect( + new SessionTranscriptReader(workspaceDir).readRestoreProjection( + sessionId, + { + replay: { + kind: 'recent', + limit: 0, + hideInheritedHistory: false, + }, + }, + ), + ).rejects.toBeInstanceOf(RangeError); + expect(buildCount).toBe(0); + }); + + it('returns no cold projection for empty sessions and rejects foreign ones', async () => { + await writeRawTranscript(''); + const reader = new SessionTranscriptReader(workspaceDir); + await expect( + reader.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ).resolves.toBeUndefined(); + + await writeRawTranscript('{not-json}\n'); + await expect( + reader.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ).resolves.toBeUndefined(); + + await writeRecords([record('u1', null, 'foreign')]); + await expect( + reader.readRestoreProjection( + sessionId, + { replay: { kind: 'none' } }, + { validateFirstRecord: () => false }, + ), + ).rejects.toBeInstanceOf(SessionTranscriptSnapshotUnavailableError); + + await expect( + reader.readLiveRestoreProjection( + sessionId, + { replay: { kind: 'none' } }, + { validateFirstRecord: () => false }, + ), + ).rejects.toBeInstanceOf(SessionTranscriptSnapshotUnavailableError); + }); + + it('fails closed when a cold restore transcript is missing', async () => { + await expect( + new SessionTranscriptReader(workspaceDir).readRestoreProjection( + sessionId, + { replay: { kind: 'none' } }, + ), + ).rejects.toBeInstanceOf(SessionTranscriptSnapshotUnavailableError); + }); + + it('preserves the real leaf for a metadata-only session', async () => { + const source: ChatRecord = { + ...record('source', null, ''), + type: 'system', + subtype: 'session_source', + message: undefined, + systemPayload: { sourceType: 'daemon', sourceId: 'metadata-only' }, + }; + await writeRecords([source]); + + const projection = await new SessionTranscriptReader( + workspaceDir, + ).readRestoreProjection(sessionId, { replay: { kind: 'none' } }); + + expect(projection?.runtime.apiHistory).toEqual([]); + expect(projection?.runtime.recording).toMatchObject({ + lastCompletedUuid: 'source', + sourceType: 'daemon', + sourceId: 'metadata-only', + }); + }); + + it('preserves malformed compression failure behavior', async () => { + await writeRecords([ + record('u1', null, 'prompt'), + { + ...record('compression', 'u1', ''), + type: 'system', + subtype: 'chat_compression', + message: undefined, + systemPayload: { + compressedHistory: { malformed: true }, + } as unknown as ChatRecord['systemPayload'], + }, + record('a1', 'compression', 'answer'), + ]); + const loaded = await new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }).loadSession(sessionId); + expect(() => buildApiHistoryFromConversation(loaded!.conversation)).toThrow( + TypeError, + ); + + await expect( + new SessionTranscriptReader(workspaceDir).readRestoreProjection( + sessionId, + { replay: { kind: 'none' } }, + ), + ).rejects.toBeInstanceOf(TypeError); + }); + + it('normalizes Goal candidates without changing recovery precedence', async () => { + const validGoal: GoalStateRecordPayloadV2 = { + v: 2, + cause: 'create', + snapshot: { + v: 2, + activity: 'idle', + goal: { + goalId: 'goal-1', + revision: 1, + objective: 'keep the valid v2 state', + status: 'active', + evidenceCursor: { recordId: 'u1' }, + turnCount: 0, + activeTimeMs: 0, + createdAt: 1, + updatedAt: 1, + }, + }, + }; + const legacy: ChatRecord = { + ...record('legacy', 'u1', ''), + type: 'system', + subtype: 'slash_command', + message: undefined, + systemPayload: { + phase: 'result', + rawCommand: '/goal legacy', + outputHistoryItems: [ + { type: 'message', text: 'discard me' }, + { type: 'goal_status', kind: 'set', condition: 'legacy goal' }, + ], + }, + }; + const valid: ChatRecord = { + ...record('valid-goal', 'legacy', ''), + type: 'system', + subtype: 'goal_state', + message: undefined, + systemPayload: validGoal, + }; + const malformed: ChatRecord = { + ...record('malformed-goal', 'valid-goal', ''), + type: 'system', + subtype: 'goal_state', + message: undefined, + systemPayload: { + v: 2, + cause: 'create', + snapshot: null, + } as unknown as ChatRecord['systemPayload'], + }; + await writeRecords([ + record('u1', null, 'prompt'), + legacy, + valid, + malformed, + record('a1', 'malformed-goal', 'answer'), + ]); + + const projection = await new SessionTranscriptReader( + workspaceDir, + ).readRestoreProjection(sessionId, { replay: { kind: 'none' } }); + + expect(recoverGoalFromRecords(projection!.runtime.goalRecords)).toEqual({ + kind: 'v2', + payload: validGoal, + }); + expect(projection?.runtime.goalRecoverySourceUuid).toBe('valid-goal'); + expect(projection?.runtime.goalRecords).toEqual([ + expect.objectContaining({ + uuid: 'legacy', + systemPayload: { + phase: 'result', + outputHistoryItems: [ + { type: 'goal_status', kind: 'set', condition: 'legacy goal' }, + ], + }, + }), + expect.objectContaining({ uuid: 'valid-goal', systemPayload: validGoal }), + expect.objectContaining({ uuid: 'malformed-goal', systemPayload: null }), + ]); + }); + + it('dispatches a pre-read malformed Goal record to other consumers once', async () => { + const malformedGoal: ChatRecord = { + ...record('a-goal', 'u1', 'model payload'), + subtype: 'goal_state', + systemPayload: { + malformed: true, + } as unknown as ChatRecord['systemPayload'], + }; + await writeRecords([record('u1', null, 'prompt'), malformedGoal]); + + const service = new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }); + const [loaded, projection] = await Promise.all([ + service.loadSession(sessionId), + service.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ]); + + expect(projection?.runtime.apiHistory).toEqual( + buildApiHistoryFromConversation(loaded!.conversation), + ); + expect(projection?.runtime.apiHistory).toHaveLength(2); + expect(recoverGoalFromRecords(projection!.runtime.goalRecords)).toEqual({ + kind: 'unsupported', + reason: expect.stringContaining('a-goal'), + }); + }); + + it('reads a narrow live projection without cold runtime state', async () => { + const artifactId = stableSessionArtifactId( + sessionId, + 'url:https://example.com/live', + ); + await writeRecords([ + record('u1', null, 'one'), + { + ...record('artifact', 'u1', ''), + type: 'system', + subtype: 'session_artifact_event', + message: undefined, + systemPayload: { + v: SESSION_ARTIFACT_PERSISTENCE_VERSION, + sessionId, + sequence: 1, + recordedAt: '2026-01-01T00:00:01.000Z', + changes: [ + { + action: 'created', + artifactId, + artifact: { + id: artifactId, + kind: 'link', + storage: 'external_url', + source: 'client', + status: 'available', + title: 'Live', + url: 'https://example.com/live', + retention: 'restorable', + clientRetained: true, + createdAt: '2026-01-01T00:00:01.000Z', + updatedAt: '2026-01-01T00:00:01.000Z', + persistedAt: '2026-01-01T00:00:01.000Z', + }, + }, + ], + }, + }, + record('a1', 'u1', 'one answer'), + record('u2', 'a1', 'two'), + record('a2', 'u2', 'two answer'), + ]); + const reader = new SessionTranscriptReader(workspaceDir); + + const liveLoad = await reader.readLiveRestoreProjection(sessionId, { + replay: { + kind: 'recent', + limit: 2, + hideInheritedHistory: false, + }, + }); + const liveResume = await reader.readLiveRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }); + + expect(liveLoad?.replay?.records.map((item) => item.uuid)).toEqual([ + 'u2', + 'a2', + ]); + expect(liveLoad?.artifactSnapshot?.artifacts).toHaveLength(1); + expect(liveResume?.replay).toBeUndefined(); + expect(liveResume?.artifactSnapshot).toEqual(liveLoad?.artifactSnapshot); + expect(liveResume).not.toHaveProperty('runtime'); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.index_cache_state', + 'miss', + ); + expect(mockAddDaemonRequestAttribute).toHaveBeenCalledWith( + 'qwen-code.daemon.session_restore.index_cache_state', + 'hit', + ); + }); + + it('retains the determining legacy Goal candidate for a recent live page', async () => { + const legacyGoal: ChatRecord = { + ...record('goal', null, ''), + type: 'system', + subtype: 'slash_command', + message: undefined, + systemPayload: { + rawCommand: '/goal keep the live goal visible', + phase: 'result', + outputHistoryItems: [ + { + type: 'goal_status', + kind: 'set', + condition: 'keep the live goal visible', + iterations: 0, + }, + ], + }, + }; + await writeRecords([ + legacyGoal, + record('u1', 'goal', 'one'), + record('a1', 'u1', 'one answer'), + record('u2', 'a1', 'two'), + record('a2', 'u2', 'two answer'), + ]); + + const projection = await new SessionTranscriptReader( + workspaceDir, + ).readLiveRestoreProjection(sessionId, { + replay: { + kind: 'recent', + limit: 2, + hideInheritedHistory: false, + }, + }); + + expect(projection?.replay?.records.map((item) => item.uuid)).toEqual([ + 'u2', + 'a2', + ]); + expect(projection?.goalRecoverySourceUuid).toBe('goal'); + expect(projection?.goalRecords).toEqual([ + expect.objectContaining({ uuid: 'goal', subtype: 'slash_command' }), + ]); + }); + + it('projects a pending Goal checkpoint window without a full-loader fallback', async () => { + const permit = { goalId: 'goal-1', revision: 1, turnId: 'turn-1' }; + const cursor: ChatRecord = { + ...record('cursor', null, ''), + type: 'system', + subtype: 'goal_runtime', + message: undefined, + }; + const evidence = Array.from({ length: 80 }, (_, index) => ({ + ...record( + `a-evidence-${index}`, + index === 0 ? 'cursor' : `a-evidence-${index - 1}`, + `evidence ${index}`, + ), + provenance: 'assistant_output' as const, + goalContext: permit, + })); + const fragmentedEvidence: ChatRecord = { + ...evidence[0]!, + message: { role: 'model', parts: [{ text: 'fragment tail' }] }, + }; + const compression: ChatRecord = { + ...record('compression', evidence.at(-1)!.uuid, ''), + type: 'system', + subtype: 'chat_compression', + message: undefined, + systemPayload: { + compressedHistory: [ + { role: 'user', parts: [{ text: 'summary' }] }, + { role: 'model', parts: [{ text: 'summary result' }] }, + ], + } as ChatRecord['systemPayload'], + }; + const goalPayload: GoalStateRecordPayloadV2 = { + v: 2, + cause: 'turn_finished', + snapshot: { + v: 2, + activity: 'idle', + goal: { + goalId: permit.goalId, + revision: permit.revision, + objective: 'verify the result', + status: 'active', + evidenceCursor: { recordId: 'cursor' }, + turnCount: 1, + activeTimeMs: 0, + createdAt: 1, + updatedAt: 2, + }, + }, + checkpointPending: { + permit, + recordUuid: evidence.at(-1)!.uuid, + }, + }; + const goalState: ChatRecord = { + ...record('goal-state', 'compression', ''), + type: 'system', + subtype: 'goal_state', + message: undefined, + systemPayload: goalPayload, + }; + const filePath = await writeRecords([ + cursor, + evidence[0]!, + fragmentedEvidence, + ...evidence.slice(1), + compression, + goalState, + ]); + let buildCount = 0; + setSessionTranscriptIndexBuildCompleteHookForTest((builtPath) => { + if (builtPath === filePath) buildCount++; + }); + + const service = new SessionService(workspaceDir, { + runtimeBaseDir: runtimeDir, + }); + const [loaded, projection] = await Promise.all([ + service.loadSession(sessionId), + service.readRestoreProjection(sessionId, { + replay: { kind: 'none' }, + }), + ]); + const expected = buildGoalEvidenceCheckpointWindow({ + records: loaded!.conversation.messages, + goal: goalPayload.snapshot.goal!, + permit, + }); + + expect(projection?.runtime.goalCheckpointWindow).toEqual(expected); + expect(projection?.runtime.goalCheckpointWindow).toMatchObject({ + shouldCheckpoint: true, + truncated: false, + }); + expect(projection?.runtime.goalCheckpointWindow?.evidence).toHaveLength(80); + expect(projection?.runtime.goalCheckpointWindow?.evidence[0]?.content).toBe( + 'evidence 0\nfragment tail', + ); + expect(buildCount).toBe(1); + }); + + it('defers unavailable pending Goal evidence to runtime recovery', async () => { + const permit = { goalId: 'goal-1', revision: 1, turnId: 'turn-1' }; + const evidence = { + ...record('evidence', null, 'result'), + provenance: 'assistant_output' as const, + goalContext: permit, + }; + const goalState: ChatRecord = { + ...record('goal-state', 'evidence', ''), + type: 'system', + subtype: 'goal_state', + message: undefined, + systemPayload: { + v: 2, + cause: 'turn_finished', + snapshot: { + v: 2, + activity: 'idle', + goal: { + goalId: permit.goalId, + revision: permit.revision, + objective: 'verify the result', + status: 'active', + evidenceCursor: { recordId: 'missing-cursor' }, + turnCount: 1, + activeTimeMs: 0, + createdAt: 1, + updatedAt: 2, + }, + }, + checkpointPending: { permit, recordUuid: evidence.uuid }, + } satisfies GoalStateRecordPayloadV2, + }; + await writeRecords([evidence, goalState]); + + const projection = await new SessionTranscriptReader( + workspaceDir, + ).readRestoreProjection(sessionId, { replay: { kind: 'none' } }); + + expect(projection?.runtime.goalCheckpointWindow).toBeUndefined(); + expect(projection?.runtime.goalRecords).toHaveLength(1); + expect(projection?.runtime.goalRecords[0]?.uuid).toBe('goal-state'); + }); + it('keeps backward pages within a normal user turn boundary', async () => { const toolCall = record('a-tool', 'u1', 'call tool'); const toolResult = { @@ -1204,6 +2612,63 @@ describe('SessionTranscriptReader', () => { } }); + it('yields cooperatively after complete scan lines and selected records', async () => { + const yielded: number[] = []; + let readSettled = false; + let siblingRanBeforeSettlement = false; + let siblingScheduled = false; + setSessionTranscriptCooperativeReadBudgetForTest( + 1, + Number.POSITIVE_INFINITY, + () => { + yielded.push(yielded.length + 1); + if (siblingScheduled) return; + siblingScheduled = true; + setImmediate(() => { + siblingRanBeforeSettlement = !readSettled; + }); + }, + ); + await writeRecords([ + record('u1', null, 'first'), + record('a1', 'u1', 'second'), + record('u2', 'a1', 'third'), + ]); + + const page = await new SessionTranscriptReader(workspaceDir) + .readPage(sessionId) + .finally(() => { + readSettled = true; + }); + + expect(page.records.map((item) => item.uuid)).toEqual(['u1', 'a1', 'u2']); + expect(yielded.length).toBeGreaterThanOrEqual(6); + expect(siblingRanBeforeSettlement).toBe(true); + }); + + it('reuses one glued physical line across projection pre-read and dispatch', async () => { + const first = record('u1', null, 'first'); + const second = record('a1', 'u1', 'second'); + const gluedLine = `${JSON.stringify(first)}${JSON.stringify(second)}`; + await writeRawTranscript(`${gluedLine}\n`); + const selectedReads: Array<{ offset: number; length: number }> = []; + setSessionTranscriptSelectedLineReadHookForTest((offset, length) => { + selectedReads.push({ offset, length }); + }); + + const projection = await new SessionTranscriptReader( + workspaceDir, + ).readRestoreProjection(sessionId, { replay: { kind: 'none' } }); + + expect(projection?.runtime.apiHistory).toEqual([ + first.message, + second.message, + ]); + expect(selectedReads).toEqual([ + { offset: 0, length: Buffer.byteLength(gluedLine) }, + ]); + }); + it('rejects oversized snapshots before indexing', async () => { const filePath = await writeRecords([record('u1', null, 'hello')]); await fs.truncate(filePath, SESSION_TRANSCRIPT_MAX_INDEX_BYTES + 1); @@ -1217,6 +2682,31 @@ describe('SessionTranscriptReader', () => { }); }); + it('shares an oversized in-flight index without retaining its completion', async () => { + const filePath = await writeRecords([ + record('u1', null, 'hello'), + record('a1', 'u1', 'reply'), + ]); + setSessionTranscriptIndexCacheMaxBytesForTest(1); + let buildCount = 0; + setSessionTranscriptIndexBuildCompleteHookForTest((builtPath) => { + if (builtPath === filePath) buildCount++; + }); + const reader = new SessionTranscriptReader(workspaceDir); + + const [first, second] = await Promise.all([ + reader.readPage(sessionId), + reader.readPage(sessionId), + ]); + + expect(buildCount).toBe(1); + expect(first.records).toEqual(second.records); + expect(getSessionTranscriptIndexCacheStatsForTest()).toEqual({ + entries: 0, + byteSize: 0, + }); + }); + it('does not evict cached indexes when a new index exceeds the byte budget alone', async () => { await writeRecords([ record('u1', null, 'hello'), @@ -1252,6 +2742,141 @@ describe('SessionTranscriptReader', () => { expect(second.records.map((r) => r.uuid)).toEqual(['a1']); }); + it('accounts for retained projection hints in the cache byte estimate', async () => { + const reader = new SessionTranscriptReader(workspaceDir); + const makeHintRecords = ( + targetSessionId: string, + suffix: string, + ): ChatRecord[] => { + const goalContext = { + goalId: `goal-${suffix}`, + revision: 1, + turnId: `turn-${suffix}`, + }; + return [ + { + ...record('notification', null, 'notice', targetSessionId), + subtype: 'notification', + goalContext, + systemPayload: { + displayText: 'notice', + backgroundTask: { + taskId: `task-${suffix}`, + status: 'completed', + kind: 'agent', + }, + }, + }, + { + ...record('assistant', 'notification', 'result', targetSessionId), + provenance: 'assistant_output', + goalContext, + }, + ]; + }; + const shortSessionId = '710e8400-e29b-41d4-a716-446655440000'; + await writeRecords( + makeHintRecords(shortSessionId, 'short'), + shortSessionId, + ); + await reader.readPage(shortSessionId); + const shortEstimate = getSessionTranscriptIndexCacheStatsForTest().byteSize; + + clearSessionTranscriptIndexCacheEntriesForTest(); + const longSessionId = '720e8400-e29b-41d4-a716-446655440000'; + const longSuffix = 'x'.repeat(8 * 1024); + await writeRecords( + makeHintRecords(longSessionId, longSuffix), + longSessionId, + ); + await reader.readPage(longSessionId); + const longEstimate = getSessionTranscriptIndexCacheStatsForTest().byteSize; + + expect(longEstimate - shortEstimate).toBeGreaterThan(80 * 1024); + }); + + it('does not let an evicted pending build overwrite a newer cache entry', async () => { + const initial = `${JSON.stringify(record('u1', null, 'hello'))}\n`; + const filePath = await writeRawTranscript(initial); + const fixed = new Date('2026-02-02T02:02:02.000Z'); + await fs.utimes(filePath, fixed, fixed); + + let buildCount = 0; + let releaseFirstBuild: (() => void) | undefined; + const firstBuildBlocked = new Promise((resolve) => { + releaseFirstBuild = resolve; + }); + setSessionTranscriptIndexBuildCompleteHookForTest(async (builtPath) => { + if (builtPath !== filePath || buildCount++ !== 0) return; + await firstBuildBlocked; + }); + + const reader = new SessionTranscriptReader(workspaceDir); + const staleRead = reader.readPage(sessionId); + await vi.waitFor(() => expect(buildCount).toBe(1)); + + clearSessionTranscriptIndexCacheEntriesForTest(); + await fs.writeFile( + filePath, + initial.replace('"uuid":"u1"', '"uuid":"x1"'), + 'utf8', + ); + await fs.utimes(filePath, fixed, fixed); + + await expect(reader.readPage(sessionId)).resolves.toMatchObject({ + records: [expect.objectContaining({ uuid: 'x1' })], + }); + + releaseFirstBuild?.(); + await expect(staleRead).rejects.toBeInstanceOf( + SessionTranscriptSnapshotUnavailableError, + ); + await expect(reader.readPage(sessionId)).resolves.toMatchObject({ + records: [expect.objectContaining({ uuid: 'x1' })], + }); + }); + + it('does not let a fresh cold projection replace a cached pending build', async () => { + const filePath = await writeRecords([ + record('u1', null, 'hello'), + record('a1', 'u1', 'reply'), + ]); + let buildCount = 0; + let releaseCachedBuild: (() => void) | undefined; + const cachedBuildBlocked = new Promise((resolve) => { + releaseCachedBuild = resolve; + }); + setSessionTranscriptIndexBuildCompleteHookForTest(async (builtPath) => { + if (builtPath !== filePath) return; + buildCount++; + if (buildCount === 1) await cachedBuildBlocked; + }); + + const reader = new SessionTranscriptReader(workspaceDir); + const cachedRead = reader.readPage(sessionId); + await vi.waitFor(() => expect(buildCount).toBe(1)); + + await expect( + reader.readRestoreProjection(sessionId, { replay: { kind: 'none' } }), + ).resolves.toBeDefined(); + expect(buildCount).toBe(2); + expect(getSessionTranscriptIndexCacheStatsForTest()).toEqual({ + entries: 1, + byteSize: 0, + }); + + releaseCachedBuild?.(); + await expect(cachedRead).resolves.toMatchObject({ + records: [ + expect.objectContaining({ uuid: 'u1' }), + expect.objectContaining({ uuid: 'a1' }), + ], + }); + expect( + getSessionTranscriptIndexCacheStatsForTest().byteSize, + ).toBeGreaterThan(0); + }); + it('evicts the least-recently-used index after 32 cached sessions', async () => { const reader = new SessionTranscriptReader(workspaceDir); for (let index = 0; index < 33; index++) { @@ -1284,6 +2909,43 @@ describe('SessionTranscriptReader', () => { }); }); + it('rejects selected records from a different session', async () => { + const foreignSessionId = '660e8400-e29b-41d4-a716-446655440000'; + await writeRecords([ + record('u1', null, 'local'), + record('a1', 'u1', 'foreign', foreignSessionId), + ]); + + await expect( + new SessionTranscriptReader(workspaceDir).readPage(sessionId), + ).rejects.toBeInstanceOf(SessionTranscriptSnapshotUnavailableError); + }); + + it('rejects foreign dead-branch metadata even when no consumer selects it', async () => { + const foreignSessionId = '660e8400-e29b-41d4-a716-446655440000'; + const foreignMetadata: ChatRecord = { + ...record('metadata', 'u1', '', foreignSessionId), + type: 'system', + subtype: 'rewind', + message: undefined, + systemPayload: { truncatedCount: 1 }, + }; + await writeRecords([ + record('u1', null, 'local'), + record('a1', 'u1', 'local reply'), + foreignMetadata, + record('u2', 'a1', 'next local prompt'), + record('a2', 'u2', 'next local reply'), + ]); + + await expect( + new SessionTranscriptReader(workspaceDir).readRestoreProjection( + sessionId, + { replay: { kind: 'none' } }, + ), + ).rejects.toBeInstanceOf(SessionTranscriptSnapshotUnavailableError); + }); + it('rejects tampered cursor snapshots before cache lookup', async () => { await writeRecords([ record('u1', null, 'hello'), @@ -1599,7 +3261,10 @@ describe('SessionTranscriptReader', () => { records.push(record('af1', 'ar1', 'filler')); // A distinct session per variant: rewriting one file in place can // keep the inode and byte length, which the index cache keys on. - await writeRecords(records, targetSessionId); + await writeRecords( + records.map((item) => ({ ...item, sessionId: targetSessionId })), + targetSessionId, + ); const page = await new SessionTranscriptReader(workspaceDir).readPage( targetSessionId, diff --git a/packages/core/src/services/session-transcript-reader.ts b/packages/core/src/services/session-transcript-reader.ts index fbb061d728..29706dcfca 100644 --- a/packages/core/src/services/session-transcript-reader.ts +++ b/packages/core/src/services/session-transcript-reader.ts @@ -8,13 +8,62 @@ import * as crypto from 'node:crypto'; import * as fs from 'node:fs'; import * as fsp from 'node:fs/promises'; import * as path from 'node:path'; +import { performance } from 'node:perf_hooks'; +import type { Content } from '@google/genai'; import { Storage } from '../config/storage.js'; import * as jsonl from '../utils/jsonl-utils.js'; import { createDebugLogger } from '../utils/debugLogger.js'; +import { addDaemonRequestAttribute } from '../telemetry/daemon-tracing.js'; +import { readSessionTitleInfoFromFileSync } from '../utils/sessionStorageUtils.js'; import type { HistoryGap } from '../utils/conversation-chain.js'; import { parseGoalStateRecordPayloadV2 } from '../goals/goal-reducer.js'; import type { GoalStateRecordPayloadV2 } from '../goals/goal-protocol.js'; -import type { ChatRecord } from './chatRecordingService.js'; +import type { + AttributionSnapshotPayload, + ChatRecord, + ParentSessionRecordPayload, + SessionSourceRecordPayload, + TitleSource, + UiTelemetryRecordPayload, +} from './chatRecordingService.js'; +import { + isApiHistoryCompressionCandidate, + SessionApiHistoryAccumulator, +} from './session-api-history.js'; +import { + isResumeTokenCountsCandidate, + ResumeTokenCountsAccumulator, + type ResumeTokenCounts, +} from './session-resume-token-counts.js'; +import { + getSessionTurnRecordHint, + type SessionTurnRecordHint, + SessionTurnStateAccumulator, +} from './session-turn-state.js'; +import { + isGoalRecoveryCandidate, + normalizeGoalRecoveryRecord, + selectGoalRecoveryFromRecords, + type GoalRecoveryRecord, + type GoalRecoverySelection, +} from '../goals/goal-persistence.js'; +import { + EvidenceSourceUnavailableError, + GoalEvidenceCheckpointAccumulator, + GoalEvidenceRecordIndexAccumulator, + InvalidGoalEvidenceReferenceError, + type GoalEvidenceCheckpointWindow, + type GoalEvidenceRecordIndexHint, +} from '../goals/goal-evidence.js'; +import type { UiEvent } from '../telemetry/uiTelemetry.js'; +import type { AttributionSnapshot } from './commitAttribution.js'; +import type { FileHistorySnapshot } from './fileHistoryService.js'; +import { SessionFileHistoryAccumulator } from './session-file-history-state.js'; +import { + selectActiveSideArtifactRecordUuids, + SessionArtifactSnapshotAccumulator, + type RebuiltSessionArtifactSnapshot, +} from './session-artifact-persistence.js'; import { aggregateTranscriptRecordFragments, isTranscriptConversationRecord, @@ -53,6 +102,8 @@ export class SessionTranscriptSnapshotUnavailableError extends Error { } } +class EmptySessionTranscriptError extends SessionTranscriptSnapshotUnavailableError {} + export class SessionTranscriptTooLargeError extends Error { constructor( readonly sessionId: string, @@ -116,6 +167,76 @@ export interface SessionTranscriptRecordPage { lastUpdated: string; } +export type SessionRestoreReplaySelection = + | { kind: 'none' } + | { kind: 'all'; hideInheritedHistory: boolean } + | { + kind: 'recent'; + limit: number; + hideInheritedHistory: boolean; + }; + +export interface SelectiveSessionRestoreOptions { + replay: SessionRestoreReplaySelection; +} + +export interface SessionRestoreReplayPage { + records: ChatRecord[]; + gaps: HistoryGap[]; + hasMore: boolean; + anchorRecordId?: string; + replay?: unknown; + goalRecoverySourceUuid?: string; + goalBootstrapRecords?: GoalRecoveryRecord[]; +} + +export interface SessionRuntimeResumeState { + apiHistory: Content[]; + resumeTokenCounts?: ResumeTokenCounts; + uiTelemetryEvents: UiEvent[]; + attributionSnapshot?: AttributionSnapshot; + historyGaps?: HistoryGap[]; + recording: { + lastCompletedUuid: string; + turnParentUuids: Array; + customTitle?: string; + titleSource?: TitleSource; + parentSessionId?: string; + sourceType?: string; + sourceId?: string; + }; + fileHistorySnapshots?: FileHistorySnapshot[]; + artifactSnapshot?: RebuiltSessionArtifactSnapshot; + goalRecords: GoalRecoveryRecord[]; + goalRecoverySourceUuid?: string; + goalCheckpointWindow?: GoalEvidenceCheckpointWindow; + initialTurn: number; + backgroundNotificationTaskIds: string[]; +} + +export interface SessionRestoreProjection { + sessionId: string; + filePath: string; + startTime: string; + lastUpdated: string; + runtime: SessionRuntimeResumeState; + replay?: SessionRestoreReplayPage; +} + +export interface SessionLiveRestoreProjection { + sessionId: string; + startTime: string; + lastUpdated: string; + replay?: SessionRestoreReplayPage; + artifactSnapshot?: RebuiltSessionArtifactSnapshot; + goalRecords?: GoalRecoveryRecord[]; + goalRecoverySourceUuid?: string; +} + +interface RestoreProjectionReadOptions { + validateFirstRecord?: (record: ChatRecord) => boolean | Promise; +} + interface SessionTranscriptFileIdentity { dev: number; ino: number; @@ -128,11 +249,39 @@ interface RecordSegment { fragmentIndex: number; } +interface CachedPhysicalLine { + offset: number; + length: number; + records: ChatRecord[]; +} + +interface PhysicalRecordHint { + uuid: string; + parentUuid: string | null; + type: TranscriptRecordInput['type']; + subtype?: TranscriptRecordInput['subtype']; +} + +interface AggregatedRecordReadContext { + handle: fsp.FileHandle; + scheduler: CooperativeReadScheduler; + lineCache: { value?: CachedPhysicalLine }; + preloadedRecords?: Map; +} + interface UuidIndexEntry { parentUuid: string | null; + sessionIdMatchesFile: boolean; type: ChatRecord['type']; subtype?: TranscriptRecordInput['subtype']; inherited: boolean; + sideTaskSource: boolean; + apiHistoryCompressionCandidate: boolean; + resumeTokenCountsCandidate: boolean; + attributionSnapshotCandidate: boolean; + goalRecoveryCandidate: boolean; + goalEvidenceHint: GoalEvidenceRecordIndexHint; + turnHint: SessionTurnRecordHint; segments: RecordSegment[]; } @@ -141,9 +290,13 @@ interface TranscriptIndex { fileIdentity: SessionTranscriptFileIdentity; snapshotSize: number; leafUuid: string; - activeUuids: string[]; + firstRecordUuid: string; + physicalRecords: PhysicalRecordHint[]; + runtimeUuids: string[]; + replayUuids: string[]; goalStatePositions: number[]; gaps: HistoryGap[]; + restoreStartTime: string; startTime: string; lastUpdated: string; byUuid: Map; @@ -162,13 +315,85 @@ const INDEX_CACHE_TTL_MS = 5 * 60 * 1000; const INDEX_ENTRY_BASE_BYTES = 256; const INDEX_SEGMENT_BYTES = 64; const INDEX_STRING_BYTES = 2; +const INDEX_HINT_BASE_BYTES = 64; +const INDEX_CONTAINER_BASE_BYTES = 64; +const INDEX_CONTAINER_SLOT_BYTES = 8; +const INDEX_MAP_ENTRY_BYTES = 48; const READ_CHUNK_SIZE = 64 * 1024; +const COOPERATIVE_READ_BYTE_BUDGET = 2 * 1024 * 1024; +const COOPERATIVE_READ_TIME_BUDGET_MS = 8; const CURSOR_HMAC_KEY_BYTES = 32; const CURSOR_HMAC_KEY_FILENAME = 'session-transcript-cursor-key'; const SESSION_TRANSCRIPT_SESSION_ID_PATTERN = /^[0-9a-fA-F-]{32,36}$/; const debugLogger = createDebugLogger('SESSION_TRANSCRIPT'); +function recordRestoreStage(stage: string, startedAt: number): void { + const durationMs = Math.round((performance.now() - startedAt) * 100) / 100; + if (Number.isFinite(durationMs) && durationMs >= 0) { + addDaemonRequestAttribute( + `qwen-code.daemon.session_restore.${stage}_ms`, + durationMs, + ); + } +} + +function recordRestoreIndexAttributes(index: TranscriptIndex): void { + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.transcript_bytes', + index.snapshotSize, + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.records_indexed', + index.byUuid.size, + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.active_records', + index.runtimeUuids.length, + ); +} + +function selectedRecordBytes( + index: TranscriptIndex, + uuids: Iterable, +): number { + const offsets = new Set(); + let bytes = 0; + for (const uuid of uuids) { + const entry = index.byUuid.get(uuid); + if (!entry) continue; + for (const segment of entry.segments) { + if (offsets.has(segment.offset)) continue; + offsets.add(segment.offset); + bytes += segment.length; + } + } + return bytes; +} + +function recordRestoreSelectionAttributes( + index: TranscriptIndex, + selectedUuids: ReadonlySet, + replayUuids: ReadonlySet, +): void { + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.selected_records', + selectedUuids.size, + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.selected_bytes', + selectedRecordBytes(index, selectedUuids), + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.replay_records', + replayUuids.size, + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.replay_bytes', + selectedRecordBytes(index, replayUuids), + ); +} + const indexCache = new Map(); // Per-workspace HMAC signing keys are cached for the daemon's lifetime (keyed by // key-file path). Rotating a key file externally therefore requires a daemon @@ -179,6 +404,39 @@ const indexCache = new Map(); const cursorHmacKeys = new Map(); let indexCacheMaxBytesForTest: number | undefined; let expandedPageBytesForTest: number | undefined; +let cooperativeReadByteBudgetForTest: number | undefined; +let cooperativeReadTimeBudgetMsForTest: number | undefined; +let cooperativeYieldHookForTest: (() => void) | undefined; +let selectedLineReadHookForTest: + | ((offset: number, length: number) => void) + | undefined; +let indexBuildCompleteHookForTest: + | ((filePath: string) => void | Promise) + | undefined; + +class CooperativeReadScheduler { + private processedBytes = 0; + private startedAt = performance.now(); + + async afterUnit(sourceBytes: number): Promise { + this.processedBytes += sourceBytes; + const byteBudget = + cooperativeReadByteBudgetForTest ?? COOPERATIVE_READ_BYTE_BUDGET; + const timeBudgetMs = + cooperativeReadTimeBudgetMsForTest ?? COOPERATIVE_READ_TIME_BUDGET_MS; + if ( + this.processedBytes < byteBudget && + performance.now() - this.startedAt < timeBudgetMs + ) { + return; + } + + cooperativeYieldHookForTest?.(); + await new Promise((resolve) => setImmediate(resolve)); + this.processedBytes = 0; + this.startedAt = performance.now(); + } +} function getExpandedPageBytes(): number { return expandedPageBytesForTest ?? SESSION_TRANSCRIPT_MAX_EXPANDED_PAGE_BYTES; @@ -200,6 +458,23 @@ function isObjectRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null && !Array.isArray(value); } +function isFileMissingError(error: unknown): boolean { + return ( + typeof error === 'object' && + error !== null && + 'code' in error && + error.code === 'ENOENT' + ); +} + +function isAttributionSnapshotCandidate(record: ChatRecord): boolean { + return ( + record.subtype === 'attribution_snapshot' && + isObjectRecord(record.systemPayload) && + 'snapshot' in record.systemPayload + ); +} + function isFiniteNonNegativeInteger(value: unknown): value is number { return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0; } @@ -479,7 +754,7 @@ function selectPageUuids( limit: number, maxBytes: number | undefined, ): string[] { - const candidates = index.activeUuids.slice(position, position + limit); + const candidates = index.replayUuids.slice(position, position + limit); if (maxBytes === undefined) return candidates; const selected: string[] = []; @@ -565,7 +840,7 @@ function findReplayBoundaryAtOrBefore( floor: number, isBoundary: (index: TranscriptIndex, uuid: string) => boolean, ): number { - return findBoundaryAtOrBefore(index.activeUuids, from, floor, (uuid) => + return findBoundaryAtOrBefore(index.replayUuids, from, floor, (uuid) => isBoundary(index, uuid), ); } @@ -578,7 +853,7 @@ function backwardPageBytesFit( ): boolean { let total = 0; for (let i = start; i < end; i++) { - total += recordSegmentBytes(index, index.activeUuids[i]!); + total += recordSegmentBytes(index, index.replayUuids[i]!); if (total > budget) return false; } return true; @@ -594,11 +869,11 @@ function selectionOrphansToolResult( end: number, ): boolean { for (let i = start; i < end; i++) { - if (index.byUuid.get(index.activeUuids[i]!)?.type !== 'tool_result') { + if (index.byUuid.get(index.replayUuids[i]!)?.type !== 'tool_result') { continue; } for (let owner = i - 1; owner >= start; owner--) { - if (isReplayPageStart(index, index.activeUuids[owner]!)) return false; + if (isReplayPageStart(index, index.replayUuids[owner]!)) return false; } return true; } @@ -630,7 +905,7 @@ function selectBackwardPageUuids( let start = Math.max(0, position - limit); for (let i = start; i < position; i++) { - if (isReplayTurnStart(index, index.activeUuids[i]!)) { + if (isReplayTurnStart(index, index.replayUuids[i]!)) { start = i; break; } @@ -650,14 +925,14 @@ function selectBackwardPageUuids( expansionFloor, isReplayTurnStart, ); - if (isReplayTurnStart(index, index.activeUuids[expandedStart]!)) { + if (isReplayTurnStart(index, index.replayUuids[expandedStart]!)) { start = expandedStart; } let selectedStart = position; let selectedBytes = 0; for (let i = position - 1; i >= start; i--) { - const uuid = index.activeUuids[i]!; + const uuid = index.replayUuids[i]!; const bytes = recordSegmentBytes(index, uuid); // Always take at least one record so backward pagination cannot // dead-end. @@ -679,12 +954,12 @@ function selectBackwardPageUuids( const logTurnExpansionSkipped = (reason: string): void => { debugLogger.debug( `backward turn expansion skipped session=${sessionId} ` + - `start=${index.activeUuids[selectedStart]!} reason=${reason}`, + `start=${index.replayUuids[selectedStart]!} reason=${reason}`, ); }; let alignedToReplayBoundary = false; for (let i = selectedStart; i < position; i++) { - if (isReplayTurnStart(index, index.activeUuids[i]!)) { + if (isReplayTurnStart(index, index.replayUuids[i]!)) { selectedStart = i; alignedToReplayBoundary = true; break; @@ -723,7 +998,7 @@ function selectBackwardPageUuids( expansionFloor, isReplayTurnStart, ); - if (isReplayTurnStart(index, index.activeUuids[candidate]!)) { + if (isReplayTurnStart(index, index.replayUuids[candidate]!)) { if ( backwardPageBytesFit(index, candidate, position, expansionByteBudget) ) { @@ -768,10 +1043,10 @@ function selectBackwardPageUuids( pairFloor, isReplayPageStart, ); - if (!isReplayPageStart(index, index.activeUuids[pairStart]!)) { + if (!isReplayPageStart(index, index.replayUuids[pairStart]!)) { debugLogger.debug( `backward pair extension skipped session=${sessionId} ` + - `start=${index.activeUuids[selectedStart]!} reason=record-budget`, + `start=${index.replayUuids[selectedStart]!} reason=record-budget`, ); } else if ( !backwardPageBytesFit( @@ -783,7 +1058,7 @@ function selectBackwardPageUuids( ) { debugLogger.debug( `backward pair extension skipped session=${sessionId} ` + - `start=${index.activeUuids[selectedStart]!} reason=byte-budget`, + `start=${index.replayUuids[selectedStart]!} reason=byte-budget`, ); } else { selectedStart = pairStart; @@ -791,7 +1066,7 @@ function selectBackwardPageUuids( } return { - uuids: index.activeUuids.slice(selectedStart, position), + uuids: index.replayUuids.slice(selectedStart, position), nextPosition: selectedStart, }; } @@ -848,12 +1123,33 @@ function estimateStringBytes(value: string | null | undefined): number { function estimateIndexCacheBytes(index: TranscriptIndex): number { let total = INDEX_ENTRY_BASE_BYTES + + INDEX_CONTAINER_BASE_BYTES * 6 + + (index.physicalRecords.length + + index.runtimeUuids.length + + index.replayUuids.length + + index.goalStatePositions.length + + index.gaps.length) * + INDEX_CONTAINER_SLOT_BYTES + estimateStringBytes(index.filePath) + estimateStringBytes(index.leafUuid) + + estimateStringBytes(index.firstRecordUuid) + + estimateStringBytes(index.restoreStartTime) + estimateStringBytes(index.startTime) + estimateStringBytes(index.lastUpdated); - for (const uuid of index.activeUuids) { + for (const record of index.physicalRecords) { + total += + INDEX_HINT_BASE_BYTES + + estimateStringBytes(record.uuid) + + estimateStringBytes(record.parentUuid) + + estimateStringBytes(record.type) + + estimateStringBytes(record.subtype); + } + + for (const uuid of index.runtimeUuids) { + total += estimateStringBytes(uuid); + } + for (const uuid of index.replayUuids) { total += estimateStringBytes(uuid); } total += index.goalStatePositions.length * 8; @@ -866,8 +1162,21 @@ function estimateIndexCacheBytes(index: TranscriptIndex): number { for (const [uuid, entry] of index.byUuid) { total += INDEX_ENTRY_BASE_BYTES + + INDEX_HINT_BASE_BYTES * 2 + + (entry.goalEvidenceHint.parsedGoalContext ? INDEX_HINT_BASE_BYTES : 0) + + INDEX_MAP_ENTRY_BYTES + + INDEX_CONTAINER_BASE_BYTES + + entry.segments.length * INDEX_CONTAINER_SLOT_BYTES + estimateStringBytes(uuid) + estimateStringBytes(entry.parentUuid) + + estimateStringBytes(entry.type) + + estimateStringBytes(entry.subtype) + + estimateStringBytes(entry.turnHint.turnParentUuid) + + estimateStringBytes(entry.turnHint.backgroundNotificationTaskId) + + estimateStringBytes(entry.goalEvidenceHint.parsedGoalContext?.goalId) + + estimateStringBytes(entry.goalEvidenceHint.parsedGoalContext?.turnId) + + estimateStringBytes(entry.goalEvidenceHint.claimedGoalId) + + estimateStringBytes(entry.goalEvidenceHint.provenance) + entry.segments.length * INDEX_SEGMENT_BYTES; } @@ -911,9 +1220,14 @@ function pruneCache(now = Date.now()): void { async function forEachLineInSnapshot( filePath: string, snapshotSize: number, - onLine: (line: Buffer, offset: number, length: number) => void, + onLine: ( + line: Buffer, + offset: number, + length: number, + ) => void | Promise, ): Promise { if (snapshotSize === 0) return; + const scheduler = new CooperativeReadScheduler(); let pending: Buffer[] = []; let pendingLength = 0; let pendingOffset = 0; @@ -947,7 +1261,8 @@ async function forEachLineInSnapshot( rawLine.length > 0 && rawLine[rawLine.length - 1] === 0x0d ? rawLine.subarray(0, rawLine.length - 1) : rawLine; - onLine(line, lineOffset, line.length); + await onLine(line, lineOffset, line.length); + await scheduler.afterUnit(line.length); pending = []; pendingLength = 0; lineStart = lineEnd + 1; @@ -970,7 +1285,8 @@ async function forEachLineInSnapshot( rawLine[rawLine.length - 1] === 0x0d ? rawLine.subarray(0, rawLine.length - 1) : rawLine; - onLine(line, pendingOffset, line.length); + await onLine(line, pendingOffset, line.length); + await scheduler.afterUnit(line.length); } } @@ -979,20 +1295,33 @@ async function readSegmentRecords( filePath: string, segment: RecordSegment, uuid: string, + lineCache: { value?: CachedPhysicalLine }, ): Promise { if (segment.length === 0) return []; - const buffer = Buffer.alloc(segment.length); - await handle.read(buffer, 0, segment.length, segment.offset); - const line = buffer.toString('utf8').trim(); - if (line.length === 0) return []; - const records = jsonl - .parseLineTolerant(line, filePath) - .flatMap((value): ChatRecord[] => { - const record = validateTranscriptRecord(value).record; - return record && isTranscriptConversationRecord(record) - ? [record as unknown as ChatRecord] - : []; - }); + let records: ChatRecord[]; + if ( + lineCache.value?.offset === segment.offset && + lineCache.value.length === segment.length + ) { + records = lineCache.value.records; + } else { + selectedLineReadHookForTest?.(segment.offset, segment.length); + const buffer = Buffer.alloc(segment.length); + await handle.read(buffer, 0, segment.length, segment.offset); + const line = buffer.toString('utf8').trim(); + if (line.length === 0) return []; + records = jsonl + .parseLineTolerant(line, filePath) + .flatMap((value): ChatRecord[] => { + const record = validateTranscriptRecord(value).record; + return record ? [record as unknown as ChatRecord] : []; + }); + lineCache.value = { + offset: segment.offset, + length: segment.length, + records, + }; + } const anomalySessionId = path.basename(filePath, '.jsonl'); const record = records[segment.fragmentIndex]; if (!record) { @@ -1013,33 +1342,94 @@ async function readSegmentRecords( ); throw new SessionTranscriptSnapshotUnavailableError(anomalySessionId); } + if (record.sessionId !== anomalySessionId) { + debugLogger.warn( + `segment read anomaly: session mismatch session=${anomalySessionId} ` + + `recordSession=${record.sessionId} uuid=${uuid} offset=${segment.offset}`, + ); + throw new SessionTranscriptSnapshotUnavailableError(anomalySessionId); + } return [record]; } +async function forEachAggregatedRecord( + index: TranscriptIndex, + uuids: string[], + onRecord: (record: ChatRecord) => void | Promise, + context?: AggregatedRecordReadContext, +): Promise { + if (!context) { + await withAggregatedRecordReadContext(index, (readContext) => + forEachAggregatedRecord(index, uuids, onRecord, readContext), + ); + return; + } + for (const uuid of uuids) { + const entry = index.byUuid.get(uuid); + if (!entry) continue; + const preloadedRecord = context.preloadedRecords?.get(uuid); + if (preloadedRecord) { + context.preloadedRecords?.delete(uuid); + await onRecord(preloadedRecord); + continue; + } + const physicalRecords: ChatRecord[] = []; + for (const segment of entry.segments) { + physicalRecords.push( + ...(await readSegmentRecords( + context.handle, + index.filePath, + segment, + uuid, + context.lineCache, + )), + ); + } + if (physicalRecords.length > 0) { + await onRecord(aggregateTranscriptRecordFragments(physicalRecords)); + } + await context.scheduler.afterUnit( + entry.segments.reduce((total, segment) => total + segment.length, 0), + ); + } +} + +async function withAggregatedRecordReadContext( + index: TranscriptIndex, + callback: (context: AggregatedRecordReadContext) => Promise, +): Promise { + let handle: fsp.FileHandle; + try { + handle = await fsp.open(index.filePath, 'r'); + } catch (error) { + if (isFileMissingError(error)) { + throw new SessionTranscriptSnapshotUnavailableError( + path.basename(index.filePath, '.jsonl'), + ); + } + throw error; + } + const context: AggregatedRecordReadContext = { + handle, + scheduler: new CooperativeReadScheduler(), + lineCache: {}, + }; + try { + return await callback(context); + } finally { + await context.handle.close(); + } +} + async function readAggregatedRecords( index: TranscriptIndex, uuids: string[], ): Promise { - const handle = await fsp.open(index.filePath, 'r'); - try { - const records: ChatRecord[] = []; - for (const uuid of uuids) { - const entry = index.byUuid.get(uuid); - if (!entry) continue; - const physicalRecords: ChatRecord[] = []; - for (const segment of entry.segments) { - physicalRecords.push( - ...(await readSegmentRecords(handle, index.filePath, segment, uuid)), - ); - } - if (physicalRecords.length > 0) { - records.push(aggregateTranscriptRecordFragments(physicalRecords)); - } - } - return records; - } finally { - await handle.close(); - } + const records: ChatRecord[] = []; + await forEachAggregatedRecord(index, uuids, (record) => { + records.push(record); + }); + return records; } async function readGoalStatePayloadBeforePosition( @@ -1058,7 +1448,7 @@ async function readGoalStatePayloadBeforePosition( } const goalStatePosition = index.goalStatePositions[low - 1]; if (goalStatePosition === undefined) return undefined; - const uuid = index.activeUuids[goalStatePosition]!; + const uuid = index.replayUuids[goalStatePosition]!; const [record] = await readAggregatedRecords(index, [uuid]); return parseGoalStateRecordPayloadV2(record?.systemPayload); } @@ -1086,69 +1476,144 @@ async function buildIndex(params: { `index build start session=${sessionId} snapshotSize=${snapshotSize}`, ); const byUuid = new Map(); + const goalEvidenceAccumulators = new Map< + string, + GoalEvidenceRecordIndexAccumulator + >(); let sequence = 0; + const physicalRecords: PhysicalRecordHint[] = []; let leafUuid: string | undefined; + let firstRecordUuid: string | undefined; + let firstRecordTimestamp: string | undefined; let startTime: string | undefined; - let sideTaskSourceUuid: string | undefined; - await forEachLineInSnapshot( - filePath, - snapshotSize, - (line, offset, length) => { - const text = line.toString('utf8').trim(); - if (text.length === 0) return; - let fragmentIndex = 0; - for (const value of jsonl.parseLineTolerant(text, filePath)) { - const record = validateTranscriptRecord(value).record; - if (!record || !isTranscriptConversationRecord(record)) { - continue; - } - if ( - record.type === 'system' && - record.subtype === 'session_source' && - isObjectRecord(record.systemPayload) && - record.systemPayload['sourceType'] === 'side_task' - ) { - sideTaskSourceUuid = record.uuid; - } - if (record.timestamp) startTime ??= record.timestamp; - leafUuid = record.uuid; - const existing = byUuid.get(record.uuid); - const segment = { - offset, - length, - sequence: sequence++, - fragmentIndex, - }; - fragmentIndex++; - if (existing) { - existing.segments.push(segment); - } else { - byUuid.set(record.uuid, { + try { + await forEachLineInSnapshot( + filePath, + snapshotSize, + (line, offset, length) => { + const text = line.toString('utf8').trim(); + if (text.length === 0) return; + let fragmentIndex = 0; + for (const value of jsonl.parseLineTolerant(text, filePath)) { + const record = validateTranscriptRecord(value).record; + if (!record) { + continue; + } + if (firstRecordUuid === undefined) { + firstRecordUuid = record.uuid; + firstRecordTimestamp = record.timestamp; + } + const sideTaskSource = + record.type === 'system' && + record.subtype === 'session_source' && + isObjectRecord(record.systemPayload) && + record.systemPayload['sourceType'] === 'side_task'; + if (isTranscriptConversationRecord(record)) { + if (record.timestamp) startTime ??= record.timestamp; + leafUuid = record.uuid; + } + const existing = byUuid.get(record.uuid); + physicalRecords.push({ + uuid: record.uuid, parentUuid: record.parentUuid, type: record.type, ...(record.subtype !== undefined ? { subtype: record.subtype } : {}), - inherited: record.forkedFrom !== undefined, - segments: [segment], }); + const segment = { + offset, + length, + sequence: sequence++, + fragmentIndex, + }; + fragmentIndex++; + if (existing) { + existing.segments.push(segment); + existing.sessionIdMatchesFile &&= record.sessionId === sessionId; + if (existing.type === 'assistant' && record.usageMetadata) { + existing.resumeTokenCountsCandidate = + isResumeTokenCountsCandidate({ + ...(record as unknown as ChatRecord), + type: existing.type, + subtype: existing.subtype, + } as unknown as ChatRecord); + } + existing.turnHint.countsAsUserPrompt ||= getSessionTurnRecordHint( + record as unknown as ChatRecord, + sessionId, + ).countsAsUserPrompt; + goalEvidenceAccumulators + .get(record.uuid) + ?.addFragment(record as unknown as ChatRecord); + } else { + const chatRecord = record as unknown as ChatRecord; + const goalEvidenceAccumulator = + new GoalEvidenceRecordIndexAccumulator(chatRecord); + const goalEvidenceHint = goalEvidenceAccumulator.finish(); + if (goalEvidenceHint.provenance) { + goalEvidenceAccumulators.set( + record.uuid, + goalEvidenceAccumulator, + ); + } + byUuid.set(record.uuid, { + parentUuid: record.parentUuid, + sessionIdMatchesFile: record.sessionId === sessionId, + type: record.type, + ...(record.subtype !== undefined + ? { subtype: record.subtype } + : {}), + inherited: record.forkedFrom !== undefined, + sideTaskSource, + apiHistoryCompressionCandidate: + isApiHistoryCompressionCandidate(chatRecord), + resumeTokenCountsCandidate: + isResumeTokenCountsCandidate(chatRecord), + attributionSnapshotCandidate: + isAttributionSnapshotCandidate(chatRecord), + goalRecoveryCandidate: isGoalRecoveryCandidate(chatRecord), + goalEvidenceHint, + turnHint: getSessionTurnRecordHint(chatRecord, sessionId), + segments: [segment], + }); + } } - } - }, - ); - - if (!leafUuid) { - debugLogger.warn( - `index build failed: no transcript records session=${sessionId}`, + }, ); - throw new SessionTranscriptSnapshotUnavailableError(sessionId); + } catch (error) { + if (isFileMissingError(error)) { + throw new SessionTranscriptSnapshotUnavailableError(sessionId); + } + throw error; + } + + for (const [uuid, accumulator] of goalEvidenceAccumulators) { + byUuid.get(uuid)!.goalEvidenceHint = accumulator.finish(); + } + + for (const [uuid, entry] of byUuid) { + if (!entry.sessionIdMatchesFile) { + debugLogger.warn( + `transcript session mismatch session=${sessionId} uuid=${uuid}`, + ); + throw new SessionTranscriptSnapshotUnavailableError(sessionId); + } + } + + if (!leafUuid || !firstRecordUuid) { + debugLogger.warn( + `index build failed: no active transcript records session=${sessionId}`, + ); + throw new EmptySessionTranscriptError(sessionId); } startTime ??= lastUpdated; + const restoreStartTime = firstRecordTimestamp ?? startTime; const chain = walkTranscriptUuidChain(leafUuid, (uuid) => { const entry = byUuid.get(uuid); - return entry + return entry && isTranscriptConversationRecord(entry) ? { uuid, parentUuid: entry.parentUuid, @@ -1158,18 +1623,19 @@ async function buildIndex(params: { } : undefined; }); - const sourceBoundary = sideTaskSourceUuid - ? chain.uuids.indexOf(sideTaskSourceUuid) - : -1; - const activeUuids = + const runtimeUuids = [...chain.uuids]; + const sourceBoundary = runtimeUuids.findIndex( + (uuid) => byUuid.get(uuid)?.sideTaskSource === true, + ); + const replayUuids = sourceBoundary >= 0 - ? chain.uuids + ? runtimeUuids .slice(sourceBoundary) .filter((uuid) => byUuid.get(uuid)?.inherited !== true) - : [...chain.uuids]; + : [...runtimeUuids]; const goalStatePositions: number[] = []; - for (let position = 0; position < activeUuids.length; position++) { - const uuid = activeUuids[position]!; + for (let position = 0; position < replayUuids.length; position++) { + const uuid = replayUuids[position]!; const entry = byUuid.get(uuid); if (entry?.type === 'system' && entry.subtype === 'goal_state') { goalStatePositions.push(position); @@ -1184,17 +1650,24 @@ async function buildIndex(params: { debugLogger.debug( `index build complete session=${sessionId} records=${byUuid.size} ` + - `active=${activeUuids.length} gaps=${gaps.length}`, + `runtime=${runtimeUuids.length} replay=${replayUuids.length} ` + + `gaps=${gaps.length}`, ); + await indexBuildCompleteHookForTest?.(filePath); + return { filePath, fileIdentity, snapshotSize, leafUuid, - activeUuids, + firstRecordUuid, + physicalRecords, + runtimeUuids, + replayUuids, goalStatePositions, gaps, + restoreStartTime, startTime, lastUpdated, byUuid, @@ -1206,6 +1679,7 @@ async function getCachedIndex(params: { fileIdentity: SessionTranscriptFileIdentity; snapshotSize: number; lastUpdated: string; + onCacheState?: (state: 'hit' | 'pending' | 'miss') => void; }): Promise { const now = Date.now(); pruneCache(now); @@ -1220,14 +1694,17 @@ async function getCachedIndex(params: { indexCache.delete(key); indexCache.set(key, cached); debugLogger.debug(`index cache hit ${key}`); + params.onCacheState?.('hit'); return cached.value; } if (cached?.pending && cached.expiresAt > now) { debugLogger.debug(`index cache pending hit ${key}`); + params.onCacheState?.('pending'); return cached.pending; } debugLogger.debug(`index cache miss ${key}`); + params.onCacheState?.('miss'); const pending = buildIndex(params); indexCache.set(key, { pending, @@ -1245,6 +1722,17 @@ async function getCachedIndex(params: { ); return value; } + if (indexCache.get(key)?.pending !== pending) { + debugLogger.debug(`index cache skipped stale completion ${key}`); + return value; + } + if (getIndexCacheBytes() + byteSize > getIndexCacheMaxBytes()) { + indexCache.delete(key); + debugLogger.debug( + `index cache skipped byte-budget admission ${key} byteSize=${byteSize}`, + ); + return value; + } indexCache.set(key, { value, byteSize, @@ -1253,7 +1741,9 @@ async function getCachedIndex(params: { pruneCache(); return value; } catch (error) { - indexCache.delete(key); + if (indexCache.get(key)?.pending === pending) { + indexCache.delete(key); + } debugLogger.debug( `index cache build failed ${key}: ${ error instanceof Error ? error.message : String(error) @@ -1263,14 +1753,165 @@ async function getCachedIndex(params: { } } +function makeReplayIndex( + index: TranscriptIndex, + hideInheritedHistory: boolean, +): TranscriptIndex { + const replayUuids = hideInheritedHistory + ? index.replayUuids.filter( + (uuid) => index.byUuid.get(uuid)?.inherited !== true, + ) + : index.replayUuids; + if (replayUuids === index.replayUuids) return index; + const goalStatePositions: number[] = []; + for (let position = 0; position < replayUuids.length; position++) { + const entry = index.byUuid.get(replayUuids[position]!); + if (entry?.type === 'system' && entry.subtype === 'goal_state') { + goalStatePositions.push(position); + } + } + return { ...index, replayUuids, goalStatePositions }; +} + +function selectRestoreReplayUuids( + index: TranscriptIndex, + sessionId: string, + replay: SessionRestoreReplaySelection, +): + | { + index: TranscriptIndex; + uuids: string[]; + hasMore: boolean; + nextPosition: number; + } + | undefined { + if (replay.kind === 'none') return undefined; + const replayIndex = makeReplayIndex(index, replay.hideInheritedHistory); + if (replay.kind === 'all') { + return { + index: replayIndex, + uuids: replayIndex.replayUuids, + hasMore: false, + nextPosition: 0, + }; + } + const limit = normalizeLimit(replay.limit); + const selected = selectBackwardPageUuids( + replayIndex, + sessionId, + replayIndex.replayUuids.length, + limit, + SESSION_TRANSCRIPT_MAX_PAGE_BYTES, + ); + return { + index: replayIndex, + uuids: selected.uuids, + hasMore: selected.nextPosition > 0, + nextPosition: selected.nextPosition, + }; +} + +function validateRestoreReplaySelection( + replay: SessionRestoreReplaySelection, +): void { + if (replay.kind === 'recent') normalizeLimit(replay.limit); +} + +async function assertIndexSnapshotUnchanged( + index: TranscriptIndex, + sessionId: string, +): Promise { + if ( + !(await hasSnapshotSignature( + index.filePath, + index.fileIdentity, + index.snapshotSize, + index.lastUpdated, + )) + ) { + throw new SessionTranscriptSnapshotUnavailableError(sessionId); + } +} + +async function hasSnapshotSignature( + filePath: string, + fileIdentity: SessionTranscriptFileIdentity, + snapshotSize: number, + lastUpdated: string, +): Promise { + let stats: fs.Stats; + try { + stats = await fsp.stat(filePath); + } catch (error) { + if (isFileMissingError(error)) return false; + throw error; + } + return ( + stats.size === snapshotSize && + sameFileIdentity(fileIdentityFromStats(stats), fileIdentity) && + new Date(stats.mtimeMs).toISOString() === lastUpdated + ); +} + +function offerFreshIndexToCache(index: TranscriptIndex): void { + pruneCache(); + const key = makeCacheKey( + index.filePath, + index.fileIdentity, + index.snapshotSize, + index.lastUpdated, + ); + if (indexCache.has(key)) return; + const byteSize = estimateIndexCacheBytes(index); + if ( + byteSize > getIndexCacheMaxBytes() || + getIndexCacheBytes() + byteSize > getIndexCacheMaxBytes() || + indexCache.size >= INDEX_CACHE_MAX_ENTRIES + ) { + debugLogger.debug( + `fresh index cache offer skipped ${key} byteSize=${byteSize}`, + ); + return; + } + indexCache.set(key, { + value: index, + byteSize, + expiresAt: Date.now() + INDEX_CACHE_TTL_MS, + }); +} + +function lastUuidMatching( + index: TranscriptIndex, + predicate: (entry: UuidIndexEntry) => boolean, +): string | undefined { + for ( + let position = index.runtimeUuids.length - 1; + position >= 0; + position-- + ) { + const uuid = index.runtimeUuids[position]!; + const entry = index.byUuid.get(uuid); + if (entry && predicate(entry)) return uuid; + } + return undefined; +} + +function selectArtifactUuids(index: TranscriptIndex): string[] { + return selectActiveSideArtifactRecordUuids( + index.physicalRecords, + index.runtimeUuids, + ); +} + export class SessionTranscriptReader { private readonly storage: Storage; constructor( private readonly workspaceCwd: string, private readonly cursorCodec?: SessionTranscriptCursorCodec, + runtimeBaseDir?: string, ) { - this.storage = new Storage(workspaceCwd); + this.storage = new Storage(workspaceCwd, runtimeBaseDir); } getSessionFilePath(sessionId: string): string { @@ -1285,6 +1926,655 @@ export class SessionTranscriptReader { ); } + async readRestoreProjection( + sessionId: string, + options: SelectiveSessionRestoreOptions, + readOptions: RestoreProjectionReadOptions = {}, + ): Promise { + const filePath = this.getSessionFilePath(sessionId); + validateRestoreReplaySelection(options.replay); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.replay_mode', + options.replay.kind, + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.index_cache_state', + 'fresh', + ); + let stats: fs.Stats; + try { + stats = await fsp.stat(filePath); + } catch (error) { + if (isFileMissingError(error)) { + throw new SessionTranscriptSnapshotUnavailableError(sessionId); + } + throw error; + } + const fileIdentity = fileIdentityFromStats(stats); + const lastUpdated = new Date(stats.mtimeMs).toISOString(); + let index: TranscriptIndex; + const indexStartedAt = performance.now(); + try { + index = await buildIndex({ + filePath, + fileIdentity, + snapshotSize: stats.size, + lastUpdated, + }); + } catch (error) { + if (error instanceof EmptySessionTranscriptError) { + if ( + await hasSnapshotSignature( + filePath, + fileIdentity, + stats.size, + lastUpdated, + ) + ) { + return undefined; + } + } + throw error; + } finally { + recordRestoreStage('transcript_index', indexStartedAt); + } + recordRestoreIndexAttributes(index); + const selectionStartedAt = performance.now(); + let replaySelection: ReturnType; + try { + replaySelection = selectRestoreReplayUuids( + index, + sessionId, + options.replay, + ); + } finally { + recordRestoreStage('resume_state_select', selectionStartedAt); + } + const replaySet = new Set(replaySelection?.uuids ?? []); + const modelSet = new Set(); + let compressionPosition = -1; + for (let position = 0; position < index.runtimeUuids.length; position++) { + const entry = index.byUuid.get(index.runtimeUuids[position]!); + if (entry?.apiHistoryCompressionCandidate) { + compressionPosition = position; + } + } + for (let position = 0; position < index.runtimeUuids.length; position++) { + const uuid = index.runtimeUuids[position]!; + const entry = index.byUuid.get(uuid); + if ( + position === compressionPosition || + (entry?.type !== 'system' && + (compressionPosition < 0 || position > compressionPosition)) + ) { + modelSet.add(uuid); + } + } + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.compression_selected', + compressionPosition >= 0, + ); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.legacy_full_model_history', + compressionPosition < 0, + ); + + const tokenUuid = lastUuidMatching( + index, + (entry) => entry.resumeTokenCountsCandidate, + ); + const attributionUuid = lastUuidMatching( + index, + (entry) => entry.attributionSnapshotCandidate, + ); + const parentSessionUuid = lastUuidMatching( + index, + (entry) => entry.type === 'system' && entry.subtype === 'parent_session', + ); + const sessionSourceUuid = lastUuidMatching( + index, + (entry) => entry.type === 'system' && entry.subtype === 'session_source', + ); + const uiTelemetrySet = new Set( + index.runtimeUuids.filter((uuid) => { + const entry = index.byUuid.get(uuid); + return entry?.type === 'system' && entry.subtype === 'ui_telemetry'; + }), + ); + const fileHistorySet = new Set( + index.runtimeUuids.filter((uuid) => { + const entry = index.byUuid.get(uuid); + return ( + entry?.type === 'system' && entry.subtype === 'file_history_snapshot' + ); + }), + ); + const goalSet = new Set( + index.runtimeUuids.filter( + (uuid) => index.byUuid.get(uuid)?.goalRecoveryCandidate === true, + ), + ); + const replayGoalSet = new Set( + replaySelection + ? replaySelection.index.replayUuids.filter( + (uuid) => index.byUuid.get(uuid)?.goalRecoveryCandidate === true, + ) + : [], + ); + const artifactUuids = selectArtifactUuids(index); + const artifactSet = new Set(artifactUuids); + const metadataSet = new Set( + [parentSessionUuid, sessionSourceUuid].filter( + (uuid): uuid is string => uuid !== undefined, + ), + ); + const apiHistory = new SessionApiHistoryAccumulator(); + const resumeTokenCounts = new ResumeTokenCountsAccumulator(); + const turnState = new SessionTurnStateAccumulator(sessionId); + for (const uuid of index.runtimeUuids) { + const hint = index.byUuid.get(uuid)?.turnHint; + if (hint) turnState.addHint(hint); + } + const uiTelemetryEvents: UiEvent[] = []; + const fileHistory = new SessionFileHistoryAccumulator(); + const artifacts = new SessionArtifactSnapshotAccumulator(sessionId); + const goalRecords: GoalRecoveryRecord[] = []; + const goalStatePayloads = new Map< + string, + GoalStateRecordPayloadV2 | undefined + >(); + const replayRecordsByUuid = new Map(); + let attributionSnapshot: AttributionSnapshot | undefined; + let parentSessionId: string | undefined; + let sourceType: string | undefined; + let sourceId: string | undefined; + let firstRecord: ChatRecord | undefined; + let firstRecordSeen = false; + let goalCheckpointAccumulator: + | GoalEvidenceCheckpointAccumulator + | undefined; + let goalEvidenceSet = new Set(); + const deferredPreReadRecords = new Map(); + const dispatchRecord = (record: ChatRecord): void => { + if (modelSet.has(record.uuid)) apiHistory.add(record); + if (record.uuid === tokenUuid) resumeTokenCounts.add(record); + if (uiTelemetrySet.has(record.uuid)) { + const uiEvent = ( + record.systemPayload as UiTelemetryRecordPayload | undefined + )?.uiEvent; + if (uiEvent) uiTelemetryEvents.push(uiEvent); + } + if (record.uuid === attributionUuid) { + const snapshot = ( + record.systemPayload as AttributionSnapshotPayload | undefined + )?.snapshot; + if (snapshot && typeof snapshot === 'object') { + attributionSnapshot = snapshot; + } + } + if (record.uuid === parentSessionUuid) { + parentSessionId = ( + record.systemPayload as ParentSessionRecordPayload | undefined + )?.parentSessionId; + } else if (record.uuid === sessionSourceUuid) { + const payload = record.systemPayload as + | SessionSourceRecordPayload + | undefined; + sourceType = payload?.sourceType; + sourceId = payload?.sourceId; + } + if (fileHistorySet.has(record.uuid)) { + try { + fileHistory.add(record); + } catch (error) { + debugLogger.warn( + `restore projection: skipping malformed file_history_snapshot: ${error}`, + ); + } + } + if (artifactSet.has(record.uuid)) artifacts.add(record); + if (goalEvidenceSet.has(record.uuid)) { + goalCheckpointAccumulator?.capture(record); + } + if (replaySet.has(record.uuid)) { + replayRecordsByUuid.set(record.uuid, record); + } + }; + const preReadUuids = Array.from( + new Set([ + index.firstRecordUuid, + ...index.runtimeUuids.filter((uuid) => goalSet.has(uuid)), + ]), + ); + const selectedReadsStartedAt = performance.now(); + const selectedReadSet = new Set(preReadUuids); + let goalRecovery: { + selectedGoalRecovery: GoalRecoverySelection; + goalCheckpointWindow: GoalEvidenceCheckpointWindow | undefined; + }; + try { + goalRecovery = await withAggregatedRecordReadContext( + index, + async (readContext) => { + await forEachAggregatedRecord( + index, + preReadUuids, + async (record) => { + if (record.uuid === index.firstRecordUuid) { + if ( + readOptions.validateFirstRecord && + !(await readOptions.validateFirstRecord(record)) + ) { + throw new SessionTranscriptSnapshotUnavailableError( + sessionId, + ); + } + firstRecordSeen = true; + if (!goalSet.has(record.uuid)) firstRecord = record; + } + if (goalSet.has(record.uuid)) { + const normalized = normalizeGoalRecoveryRecord(record); + if (normalized) { + goalRecords.push(normalized); + if (record.subtype === 'goal_state') { + goalStatePayloads.set( + record.uuid, + parseGoalStateRecordPayloadV2(normalized.systemPayload), + ); + } + } + } + if (replaySet.has(record.uuid)) { + replayRecordsByUuid.set(record.uuid, record); + } + const needsDeferredDispatch = + modelSet.has(record.uuid) || + record.uuid === tokenUuid || + record.uuid === attributionUuid || + metadataSet.has(record.uuid) || + uiTelemetrySet.has(record.uuid) || + fileHistorySet.has(record.uuid) || + artifactSet.has(record.uuid); + if (needsDeferredDispatch) { + if ( + record.uuid === index.firstRecordUuid && + artifactSet.has(record.uuid) + ) { + artifacts.add(record); + } else { + deferredPreReadRecords.set(record.uuid, record); + } + } + }, + readContext, + ); + if (!firstRecordSeen) { + throw new SessionTranscriptSnapshotUnavailableError(sessionId); + } + + const selectedGoalRecovery = + selectGoalRecoveryFromRecords(goalRecords); + const pendingGoal = + selectedGoalRecovery.recovery.kind === 'v2' + ? selectedGoalRecovery.recovery.payload + : undefined; + const pendingCheckpoint = pendingGoal?.checkpointPending; + if (pendingCheckpoint && pendingGoal.snapshot.goal) { + try { + goalCheckpointAccumulator = new GoalEvidenceCheckpointAccumulator( + index.runtimeUuids.map( + (uuid) => index.byUuid.get(uuid)!.goalEvidenceHint, + ), + pendingGoal.snapshot.goal, + pendingCheckpoint.permit, + ); + } catch (error) { + if (!(error instanceof EvidenceSourceUnavailableError)) { + throw error; + } + debugLogger.warn( + `restore projection: deferring unavailable Goal checkpoint evidence: ${error.message}`, + ); + } + } + goalEvidenceSet = new Set( + goalCheckpointAccumulator?.getCandidateUuids() ?? [], + ); + if (firstRecord && goalEvidenceSet.has(firstRecord.uuid)) { + goalCheckpointAccumulator?.capture(firstRecord); + } + firstRecord = undefined; + const selectedRuntimeUuids = index.runtimeUuids.filter( + (uuid) => + modelSet.has(uuid) || + uuid === tokenUuid || + uuid === attributionUuid || + metadataSet.has(uuid) || + uiTelemetrySet.has(uuid) || + fileHistorySet.has(uuid) || + replaySet.has(uuid) || + goalEvidenceSet.has(uuid), + ); + const preReadSet = new Set(preReadUuids); + readContext.preloadedRecords = deferredPreReadRecords; + const remainingUuids = Array.from( + new Set([...selectedRuntimeUuids, ...artifactUuids]), + ).filter( + (uuid) => !preReadSet.has(uuid) || deferredPreReadRecords.has(uuid), + ); + for (const uuid of remainingUuids) selectedReadSet.add(uuid); + await forEachAggregatedRecord( + index, + remainingUuids, + dispatchRecord, + readContext, + ); + + let goalCheckpointWindow: GoalEvidenceCheckpointWindow | undefined; + try { + goalCheckpointWindow = goalCheckpointAccumulator?.finish(); + } catch (error) { + if (!(error instanceof InvalidGoalEvidenceReferenceError)) { + throw error; + } + debugLogger.warn( + `restore projection: deferring invalid Goal checkpoint evidence: ${error.message}`, + ); + } + return { selectedGoalRecovery, goalCheckpointWindow }; + }, + ); + } finally { + recordRestoreStage('selected_record_read', selectedReadsStartedAt); + } + recordRestoreSelectionAttributes(index, selectedReadSet, replaySet); + + let persistedTitle: ReturnType; + try { + persistedTitle = readSessionTitleInfoFromFileSync(index.filePath); + } catch (error) { + if (isFileMissingError(error)) { + throw new SessionTranscriptSnapshotUnavailableError(sessionId); + } + throw error; + } + const customTitle = persistedTitle.title; + const titleSource = persistedTitle.source; + + const turnStateValue = turnState.finish(); + const replayRecords = replaySelection + ? replaySelection.uuids + .map((uuid) => replayRecordsByUuid.get(uuid)) + .filter((record): record is ChatRecord => record !== undefined) + : []; + const replayGoalRecords = replaySelection + ? goalRecords.filter((record) => replayGoalSet.has(record.uuid)) + : []; + const replayGoalRecoverySourceUuid = replaySelection + ? selectGoalRecoveryFromRecords(replayGoalRecords).sourceUuid + : goalRecovery.selectedGoalRecovery.sourceUuid; + let replay: SessionRestoreReplayPage | undefined; + if (replaySelection) { + const goalStatePosition = + replaySelection.index.goalStatePositions.findLast( + (position) => position < replaySelection.nextPosition, + ); + const goalStateUuid = + goalStatePosition === undefined + ? undefined + : replaySelection.index.replayUuids[goalStatePosition]; + const goalState = goalStateUuid + ? goalStatePayloads.get(goalStateUuid) + : undefined; + replay = { + records: replayRecords, + gaps: index.gaps, + hasMore: replaySelection.hasMore, + ...(replaySelection.hasMore && replayRecords[0] + ? { anchorRecordId: replayRecords[0].uuid } + : {}), + ...(replayGoalRecoverySourceUuid && + !replaySet.has(replayGoalRecoverySourceUuid) + ? { goalBootstrapRecords: replayGoalRecords } + : {}), + ...(replayGoalRecoverySourceUuid + ? { goalRecoverySourceUuid: replayGoalRecoverySourceUuid } + : {}), + ...(goalState + ? { + replay: { + goalState: goalState.snapshot, + goalCause: goalState.cause, + }, + } + : {}), + }; + } + + const restoredTokenCounts = resumeTokenCounts.finish(); + const restoredFileHistory = fileHistory.finish(); + const artifactSnapshot = artifacts.finish(); + const runtime: SessionRuntimeResumeState = { + apiHistory: apiHistory.finish(), + ...(restoredTokenCounts + ? { resumeTokenCounts: restoredTokenCounts } + : {}), + uiTelemetryEvents, + ...(attributionSnapshot ? { attributionSnapshot } : {}), + ...(index.gaps.length > 0 ? { historyGaps: index.gaps } : {}), + recording: { + lastCompletedUuid: index.leafUuid, + turnParentUuids: turnStateValue.turnParentUuids, + ...(customTitle !== undefined ? { customTitle } : {}), + ...(titleSource !== undefined ? { titleSource } : {}), + ...(parentSessionId !== undefined ? { parentSessionId } : {}), + ...(sourceType !== undefined ? { sourceType } : {}), + ...(sourceId !== undefined ? { sourceId } : {}), + }, + ...(restoredFileHistory + ? { fileHistorySnapshots: restoredFileHistory } + : {}), + ...(artifactSnapshot ? { artifactSnapshot } : {}), + goalRecords, + ...(goalRecovery.selectedGoalRecovery.sourceUuid + ? { + goalRecoverySourceUuid: + goalRecovery.selectedGoalRecovery.sourceUuid, + } + : {}), + ...(goalRecovery.goalCheckpointWindow + ? { goalCheckpointWindow: goalRecovery.goalCheckpointWindow } + : {}), + initialTurn: turnStateValue.initialTurn, + backgroundNotificationTaskIds: + turnStateValue.backgroundNotificationTaskIds, + }; + + await assertIndexSnapshotUnchanged(index, sessionId); + offerFreshIndexToCache(index); + return { + sessionId, + filePath, + startTime: index.restoreStartTime, + lastUpdated: index.lastUpdated, + runtime, + ...(replay ? { replay } : {}), + }; + } + + async readLiveRestoreProjection( + sessionId: string, + options: SelectiveSessionRestoreOptions, + readOptions: RestoreProjectionReadOptions = {}, + ): Promise { + const filePath = this.getSessionFilePath(sessionId); + validateRestoreReplaySelection(options.replay); + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.replay_mode', + options.replay.kind, + ); + let stats: fs.Stats; + try { + stats = await fsp.stat(filePath); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return undefined; + throw error; + } + let index: TranscriptIndex; + const fileIdentity = fileIdentityFromStats(stats); + const lastUpdated = new Date(stats.mtimeMs).toISOString(); + const indexStartedAt = performance.now(); + try { + index = await getCachedIndex({ + filePath, + fileIdentity, + snapshotSize: stats.size, + lastUpdated, + onCacheState: (state) => + addDaemonRequestAttribute( + 'qwen-code.daemon.session_restore.index_cache_state', + state, + ), + }); + } catch (error) { + if (error instanceof EmptySessionTranscriptError) { + if ( + await hasSnapshotSignature( + filePath, + fileIdentity, + stats.size, + lastUpdated, + ) + ) { + return undefined; + } + } + throw error; + } finally { + recordRestoreStage('transcript_index', indexStartedAt); + } + recordRestoreIndexAttributes(index); + const selectionStartedAt = performance.now(); + const replaySelection = selectRestoreReplayUuids( + index, + sessionId, + options.replay, + ); + const replaySet = new Set(replaySelection?.uuids ?? []); + const goalSet = new Set( + replaySelection + ? replaySelection.index.replayUuids.filter( + (uuid) => index.byUuid.get(uuid)?.goalRecoveryCandidate === true, + ) + : [], + ); + const goalStatePosition = replaySelection + ? replaySelection.index.goalStatePositions.findLast( + (position) => position < replaySelection.nextPosition, + ) + : undefined; + const goalStateUuid = + goalStatePosition === undefined || !replaySelection + ? undefined + : replaySelection.index.replayUuids[goalStatePosition]; + const goalStateSet = new Set(goalStateUuid ? [goalStateUuid] : []); + const artifactUuids = selectArtifactUuids(index); + const artifactSet = new Set(artifactUuids); + const selectedRuntimeUuids = index.runtimeUuids.filter( + (uuid) => + replaySet.has(uuid) || goalStateSet.has(uuid) || goalSet.has(uuid), + ); + if (!selectedRuntimeUuids.includes(index.firstRecordUuid)) { + selectedRuntimeUuids.unshift(index.firstRecordUuid); + } + const replayRecords: ChatRecord[] = []; + const goalRecords: GoalRecoveryRecord[] = []; + const goalStatePayloads = new Map< + string, + GoalStateRecordPayloadV2 | undefined + >(); + const artifacts = new SessionArtifactSnapshotAccumulator(sessionId); + + recordRestoreStage('resume_state_select', selectionStartedAt); + const selectedReadSet = new Set([ + ...selectedRuntimeUuids, + ...artifactUuids, + ]); + const selectedReadsStartedAt = performance.now(); + try { + await forEachAggregatedRecord( + index, + [...selectedReadSet], + async (record) => { + if ( + record.uuid === index.firstRecordUuid && + readOptions.validateFirstRecord && + !(await readOptions.validateFirstRecord(record)) + ) { + throw new SessionTranscriptSnapshotUnavailableError(sessionId); + } + if (replaySet.has(record.uuid)) replayRecords.push(record); + if (goalStateSet.has(record.uuid)) { + goalStatePayloads.set( + record.uuid, + parseGoalStateRecordPayloadV2(record.systemPayload), + ); + } + if (goalSet.has(record.uuid)) { + const normalized = normalizeGoalRecoveryRecord(record); + if (normalized) goalRecords.push(normalized); + } + if (artifactSet.has(record.uuid)) artifacts.add(record); + }, + ); + } finally { + recordRestoreStage('selected_record_read', selectedReadsStartedAt); + } + recordRestoreSelectionAttributes(index, selectedReadSet, replaySet); + + let replay: SessionRestoreReplayPage | undefined; + if (replaySelection) { + const goalState = goalStateUuid + ? goalStatePayloads.get(goalStateUuid) + : undefined; + replay = { + records: replayRecords, + gaps: index.gaps, + hasMore: replaySelection.hasMore, + ...(replaySelection.hasMore && replayRecords[0] + ? { anchorRecordId: replayRecords[0].uuid } + : {}), + ...(goalState + ? { + replay: { + goalState: goalState.snapshot, + goalCause: goalState.cause, + }, + } + : {}), + }; + } + const artifactSnapshot = artifacts.finish(); + const goalRecovery = selectGoalRecoveryFromRecords(goalRecords); + const replayGoalRecoverySourceUuid = + goalRecovery.sourceUuid && + replaySelection?.index.replayUuids.includes(goalRecovery.sourceUuid) + ? goalRecovery.sourceUuid + : undefined; + await assertIndexSnapshotUnchanged(index, sessionId); + return { + sessionId, + startTime: index.restoreStartTime, + lastUpdated: index.lastUpdated, + ...(replay ? { replay } : {}), + ...(artifactSnapshot ? { artifactSnapshot } : {}), + ...(goalRecords.length > 0 ? { goalRecords } : {}), + ...(replayGoalRecoverySourceUuid + ? { goalRecoverySourceUuid: replayGoalRecoverySourceUuid } + : {}), + }; + } + async readPage( sessionId: string, options: SessionTranscriptReadPageOptions = {}, @@ -1347,20 +2637,20 @@ export class SessionTranscriptReader { (options.beforeRecordId !== undefined ? 'backward' : 'forward'); let position = cursor?.position ?? - (direction === 'backward' ? index.activeUuids.length : 0); + (direction === 'backward' ? index.replayUuids.length : 0); if (!cursor && options.beforeRecordId !== undefined) { if (options.beforeRecordId.length === 0) { throw new InvalidSessionTranscriptCursorError(); } - position = index.activeUuids.indexOf(options.beforeRecordId); + position = index.replayUuids.indexOf(options.beforeRecordId); if (position < 0) { throw new InvalidSessionTranscriptCursorError(); } } - if (position > index.activeUuids.length) { + if (position > index.replayUuids.length) { debugLogger.debug( `cursor position out of range session=${sessionId} ` + - `position=${position} active=${index.activeUuids.length}`, + `position=${position} replay=${index.replayUuids.length}`, ); throw new InvalidSessionTranscriptCursorError(); } @@ -1380,7 +2670,7 @@ export class SessionTranscriptReader { const hasMore = direction === 'backward' ? nextPosition > 0 - : nextPosition < index.activeUuids.length; + : nextPosition < index.replayUuids.length; const nextCursorState: SessionTranscriptCursorState | undefined = hasMore ? { v: SESSION_TRANSCRIPT_CURSOR_VERSION, @@ -1432,6 +2722,37 @@ export function resetSessionTranscriptIndexCacheForTest(): void { cursorHmacKeys.clear(); indexCacheMaxBytesForTest = undefined; expandedPageBytesForTest = undefined; + cooperativeReadByteBudgetForTest = undefined; + cooperativeReadTimeBudgetMsForTest = undefined; + cooperativeYieldHookForTest = undefined; + selectedLineReadHookForTest = undefined; + indexBuildCompleteHookForTest = undefined; +} + +export function clearSessionTranscriptIndexCacheEntriesForTest(): void { + indexCache.clear(); +} + +export function setSessionTranscriptCooperativeReadBudgetForTest( + byteBudget: number, + timeBudgetMs: number, + onYield?: () => void, +): void { + cooperativeReadByteBudgetForTest = byteBudget; + cooperativeReadTimeBudgetMsForTest = timeBudgetMs; + cooperativeYieldHookForTest = onYield; +} + +export function setSessionTranscriptIndexBuildCompleteHookForTest( + hook: (filePath: string) => void | Promise, +): void { + indexBuildCompleteHookForTest = hook; +} + +export function setSessionTranscriptSelectedLineReadHookForTest( + hook: (offset: number, length: number) => void, +): void { + selectedLineReadHookForTest = hook; } export function setSessionTranscriptIndexCacheMaxBytesForTest( diff --git a/packages/core/src/services/session-turn-state.ts b/packages/core/src/services/session-turn-state.ts new file mode 100644 index 0000000000..4d812afcc3 --- /dev/null +++ b/packages/core/src/services/session-turn-state.ts @@ -0,0 +1,154 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ChatRecord } from './chatRecordingService.js'; + +export interface SessionTurnState { + initialTurn: number; + turnParentUuids: Array; + backgroundNotificationTaskIds: string[]; +} + +export interface SessionTurnRecordHint { + promptTurn?: number; + countsAsUserPrompt: boolean; + turnParentUuid?: string | null; + backgroundNotificationTaskId?: string; +} + +export class SessionTurnStateAccumulator { + private maxPromptTurn = 0; + private userMessageCount = 0; + private readonly turnParentUuids: Array = []; + private readonly backgroundNotificationTaskIds = new Set(); + + constructor(private readonly sessionId: string) {} + + add(record: ChatRecord): void { + this.addHint(getSessionTurnRecordHint(record, this.sessionId)); + } + + addHint(hint: SessionTurnRecordHint): void { + if (hint.countsAsUserPrompt) { + this.userMessageCount += 1; + } + if (hint.promptTurn !== undefined) { + this.maxPromptTurn = Math.max(this.maxPromptTurn, hint.promptTurn); + } + if (hint.turnParentUuid !== undefined) { + this.turnParentUuids.push(hint.turnParentUuid); + } + if (hint.backgroundNotificationTaskId !== undefined) { + this.backgroundNotificationTaskIds.add(hint.backgroundNotificationTaskId); + } + } + + finish(): SessionTurnState { + return { + initialTurn: + this.maxPromptTurn > 0 ? this.maxPromptTurn : this.userMessageCount, + turnParentUuids: [...this.turnParentUuids], + backgroundNotificationTaskIds: [...this.backgroundNotificationTaskIds], + }; + } +} + +export function getSessionTurnRecordHint( + record: ChatRecord, + sessionId: string, +): SessionTurnRecordHint { + let promptTurn: number | undefined; + for (const promptId of getRecordPromptIds(record)) { + const candidate = parseSessionPromptTurn(promptId, sessionId); + if (candidate !== undefined) { + promptTurn = Math.max(promptTurn ?? 0, candidate); + } + } + const turnParentUuid = + record.type === 'user' && + record.subtype !== 'goal_runtime' && + record.subtype !== 'notification' && + record.subtype !== 'cron' && + record.subtype !== 'mid_turn_user_message' && + record.subtype !== 'realtime_message' + ? (record.parentUuid ?? null) + : undefined; + const backgroundTask = + record.subtype === 'notification' + ? ( + record.systemPayload as + | { backgroundTask?: { taskId?: unknown } } + | undefined + )?.backgroundTask + : undefined; + return { + ...(promptTurn !== undefined ? { promptTurn } : {}), + countsAsUserPrompt: + record.sessionId === sessionId && isUserPromptRecord(record), + ...(turnParentUuid !== undefined ? { turnParentUuid } : {}), + ...(typeof backgroundTask?.taskId === 'string' + ? { backgroundNotificationTaskId: backgroundTask.taskId } + : {}), + }; +} + +export function collectSessionTurnState( + records: readonly ChatRecord[], + sessionId: string, +): SessionTurnState { + const accumulator = new SessionTurnStateAccumulator(sessionId); + for (const record of records) accumulator.add(record); + return accumulator.finish(); +} + +export function computeInitialTurnFromHistory( + records: readonly ChatRecord[], + sessionId: string, +): number { + return collectSessionTurnState(records, sessionId).initialTurn; +} + +function getRecordPromptIds(record: ChatRecord): string[] { + const promptIds: string[] = []; + const recordPromptId = (record as { promptId?: unknown }).promptId; + if (typeof recordPromptId === 'string') promptIds.push(recordPromptId); + const telemetryPromptId = readTelemetryPromptId(record.systemPayload); + if (telemetryPromptId) promptIds.push(telemetryPromptId); + return promptIds; +} + +function readTelemetryPromptId(payload: unknown): string | undefined { + if (!payload || typeof payload !== 'object' || !('uiEvent' in payload)) { + return undefined; + } + const uiEvent = (payload as { uiEvent?: unknown }).uiEvent; + if (!uiEvent || typeof uiEvent !== 'object' || !('prompt_id' in uiEvent)) { + return undefined; + } + const promptId = (uiEvent as { prompt_id?: unknown }).prompt_id; + return typeof promptId === 'string' ? promptId : undefined; +} + +function parseSessionPromptTurn( + promptId: string, + sessionId: string, +): number | undefined { + const promptIdPrefix = `${sessionId}########`; + if (!promptId.startsWith(promptIdPrefix)) return undefined; + const suffix = promptId.slice(promptIdPrefix.length); + return /^\d+$/.test(suffix) ? Number(suffix) : undefined; +} + +function isUserPromptRecord(record: ChatRecord): boolean { + if (record.type !== 'user' || record.subtype === 'realtime_message') { + return false; + } + return ( + record.message?.parts?.some( + (part) => typeof part.text === 'string' && part.text.trim().length > 0, + ) ?? false + ); +} diff --git a/packages/core/src/services/sessionService.ts b/packages/core/src/services/sessionService.ts index 0e0cd0e8cb..dd22d2eeab 100644 --- a/packages/core/src/services/sessionService.ts +++ b/packages/core/src/services/sessionService.ts @@ -16,7 +16,6 @@ import * as jsonl from '../utils/jsonl-utils.js'; import type { HistoryGap } from '../utils/conversation-chain.js'; import { prepareTranscriptRecords } from '../utils/transcript-records.js'; import type { - ChatCompressionRecordPayload, ChatRecord, FileHistorySnapshotRecordPayload, TitleSource, @@ -27,31 +26,50 @@ import type { FileHistorySnapshot } from './fileHistoryService.js'; import { deserializeSnapshots, FILE_HISTORY_DIR, - MAX_SNAPSHOTS, serializeSnapshot, } from './fileHistoryService.js'; +import { SessionFileHistoryAccumulator } from './session-file-history-state.js'; import { uiTelemetryService } from '../telemetry/uiTelemetry.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import { readRuntimeStatus } from '../utils/runtimeStatus.js'; import { LITE_READ_BUF_SIZE, readLastJsonStringFieldSync, - readLastJsonStringFieldsSync, + readSessionTitleInfoFromFileSync, } from '../utils/sessionStorageUtils.js'; -import { getUsageOutputTokenCountForPromptEstimate } from './tokenEstimation.js'; import { isSessionArtifactRecord, rebuildSessionArtifactSnapshot, remapSessionArtifactPayloadForFork, + selectActiveSideArtifactRecordUuids, type RebuiltSessionArtifactSnapshot, } from './session-artifact-persistence.js'; import { SessionOrganizationService } from './session-organization-service.js'; -import { SessionTranscriptTooLargeError } from './session-transcript-reader.js'; +import { + SessionTranscriptReader, + SessionTranscriptTooLargeError, + type SelectiveSessionRestoreOptions, + type SessionLiveRestoreProjection, + type SessionRestoreProjection, +} from './session-transcript-reader.js'; import { SessionWriterLease, SessionWriterUnavailableError, type SessionWriterProcessKind, } from './session-writer-lease.js'; +export { + buildApiHistoryFromConversation, + type BuildApiHistoryOptions, +} from './session-api-history.js'; +import { + getResumeTokenCounts, + type ResumeTokenCounts, +} from './session-resume-token-counts.js'; +export { + getResumePromptTokenCount, + getResumeTokenCounts, + type ResumeTokenCounts, +} from './session-resume-token-counts.js'; const debugLogger = createDebugLogger('SESSION'); @@ -332,12 +350,18 @@ export class SessionService { private readonly projectHash: string; private readonly projectRoot: string; private readonly onWarning: ((message: string) => void) | undefined; + private readonly transcriptReader: SessionTranscriptReader; constructor(cwd: string, options: SessionServiceOptions = {}) { this.storage = new Storage(cwd, options.runtimeBaseDir); this.projectRoot = cwd; this.projectHash = getProjectHash(cwd); this.onWarning = options.onWarning; + this.transcriptReader = new SessionTranscriptReader( + cwd, + undefined, + options.runtimeBaseDir, + ); } /** The workspace root this service is bound to (the cwd it was constructed @@ -696,19 +720,7 @@ export class SessionService { title?: string; source?: TitleSource; } { - const hit = readLastJsonStringFieldsSync( - filePath, - 'customTitle', - ['titleSource'], - '"subtype":"custom_title"', - tailBuffer, - ); - const title = hit['customTitle']; - if (!title) return {}; - const rawSource = hit['titleSource']; - const source = - rawSource === 'auto' || rawSource === 'manual' ? rawSource : undefined; - return { title, source }; + return readSessionTitleInfoFromFileSync(filePath, tailBuffer); } /** @@ -1273,6 +1285,26 @@ export class SessionService { return this.loadSessionFromState(sessionId, 'active'); } + async readRestoreProjection( + sessionId: string, + options: SelectiveSessionRestoreOptions, + ): Promise { + return this.transcriptReader.readRestoreProjection(sessionId, options, { + validateFirstRecord: (record) => + this.sessionBelongsToCurrentProject(record.sessionId, record.cwd), + }); + } + + async readLiveRestoreProjection( + sessionId: string, + options: SelectiveSessionRestoreOptions, + ): Promise { + return this.transcriptReader.readLiveRestoreProjection(sessionId, options, { + validateFirstRecord: (record) => + this.sessionBelongsToCurrentProject(record.sessionId, record.cwd), + }); + } + /** * Reads an archived session without changing its archive state. * Daemon load/resume paths must continue to use {@link loadSession}. @@ -1362,40 +1394,17 @@ export class SessionService { }; // Extract file history snapshots for /rewind across resume - const fileHistorySnapshots: FileHistorySnapshot[] = []; - const seenPromptIds = new Map(); + const fileHistoryAccumulator = new SessionFileHistoryAccumulator(); for (const msg of messages) { - if ( - msg.type === 'system' && - msg.subtype === 'file_history_snapshot' && - msg.systemPayload - ) { - const payload = msg.systemPayload as FileHistorySnapshotRecordPayload; - if (!Array.isArray(payload?.snapshots)) continue; - let deserialized: FileHistorySnapshot[]; - try { - deserialized = deserializeSnapshots(payload.snapshots); - } catch (e) { - debugLogger.warn( - `loadSession: skipping malformed file_history_snapshot: ${e}`, - ); - continue; - } - for (const s of deserialized) { - const existingIdx = seenPromptIds.get(s.promptId); - if (existingIdx !== undefined) { - fileHistorySnapshots[existingIdx] = s; - } else { - seenPromptIds.set(s.promptId, fileHistorySnapshots.length); - fileHistorySnapshots.push(s); - } - } + try { + fileHistoryAccumulator.add(msg); + } catch (e) { + debugLogger.warn( + `loadSession: skipping malformed file_history_snapshot: ${e}`, + ); } } - const cappedSnapshots = - fileHistorySnapshots.length > MAX_SNAPSHOTS - ? fileHistorySnapshots.slice(-MAX_SNAPSHOTS) - : fileHistorySnapshots; + const fileHistorySnapshots = fileHistoryAccumulator.finish(); const activeBranchRecords = includeActiveSideArtifactRecords( records, messages, @@ -1409,8 +1418,7 @@ export class SessionService { conversation, filePath, lastCompletedUuid: lastMessage.uuid, - fileHistorySnapshots: - cappedSnapshots.length > 0 ? cappedSnapshots : undefined, + fileHistorySnapshots, ...(artifactSnapshot ? { artifactSnapshot } : {}), historyGaps: gaps.length > 0 ? gaps : undefined, }; @@ -2192,151 +2200,6 @@ export class SessionService { } } -/** - * Options for building API history from conversation. - */ -export interface BuildApiHistoryOptions { - /** - * Whether to strip thought parts from the history. - * Thought parts are content parts that have `thought: true`. - * Keeping thoughts ensures `reasoning_content` from reasoning models - * (e.g. DeepSeek) is properly passed back in subsequent API calls. - * @default false - */ - stripThoughtsFromHistory?: boolean; -} - -/** - * Strips thought parts from a Content object. - * Thought parts are identified by having `thought: true`. - * Returns null if the content only contained thought parts. - */ -function stripThoughtsFromContent(content: Content): Content | null { - if (!content.parts) return content; - - const filteredParts = content.parts.filter((part) => !(part as Part).thought); - - // If all parts were thoughts, remove the entire content - if (filteredParts.length === 0) { - return null; - } - - return { - ...content, - parts: filteredParts, - }; -} - -function copyContentForApiHistory(content: Content): Content { - return { - ...content, - parts: content.parts?.map((part) => { - if ('functionCall' in part && part.functionCall) { - return { - ...part, - functionCall: { - ...part.functionCall, - args: part.functionCall.args - ? { ...part.functionCall.args } - : part.functionCall.args, - }, - }; - } - if ('functionResponse' in part && part.functionResponse) { - return { - ...part, - functionResponse: { - ...part.functionResponse, - }, - }; - } - return { ...part }; - }), - }; -} - -function appendApiHistoryRecord(history: Content[], record: ChatRecord): void { - if (!record.message || record.subtype === 'realtime_message') return; - - const message = copyContentForApiHistory(record.message as Content); - if (record.subtype === 'mid_turn_user_message') { - const previous = history.at(-1); - if (previous?.role === 'user') { - previous.parts = [...(previous.parts ?? []), ...(message.parts ?? [])]; - return; - } - } - - history.push(message); -} - -/** - * Builds the model-facing chat history (Content[]) from a reconstructed - * conversation. This keeps UI history intact while applying chat compression - * checkpoints for the API history used on resume. - * - * Strategy: - * - Find the latest system/chat_compression record (if any). - * - Use its compressedHistory snapshot as the base history. - * - Append all messages after that checkpoint (skipping system records). - * - If no checkpoint exists, return the linear message list (message field only). - */ -export function buildApiHistoryFromConversation( - conversation: ConversationRecord, - options: BuildApiHistoryOptions = {}, -): Content[] { - const { stripThoughtsFromHistory = false } = options; - const { messages } = conversation; - - let lastCompressionIndex = -1; - let compressedHistory: Content[] | undefined; - - messages.forEach((record, index) => { - if (record.type === 'system' && record.subtype === 'chat_compression') { - const payload = record.systemPayload as - | ChatCompressionRecordPayload - | undefined; - if (payload?.compressedHistory) { - lastCompressionIndex = index; - compressedHistory = payload.compressedHistory; - } - } - }); - - if (compressedHistory && lastCompressionIndex >= 0) { - const baseHistory: Content[] = compressedHistory.map( - copyContentForApiHistory, - ); - - // Append everything after the compression record (newer turns) - for (let i = lastCompressionIndex + 1; i < messages.length; i++) { - const record = messages[i]; - if (record.type === 'system') continue; - appendApiHistoryRecord(baseHistory, record); - } - - if (stripThoughtsFromHistory) { - return baseHistory - .map(stripThoughtsFromContent) - .filter((content): content is Content => content !== null); - } - return baseHistory; - } - - // Fallback: return linear messages as Content[] - const result: Content[] = []; - for (const record of messages) { - appendApiHistoryRecord(result, record); - } - - if (stripThoughtsFromHistory) { - return result - .map(stripThoughtsFromContent) - .filter((content): content is Content => content !== null); - } - return result; -} - function remapSnapshotPromptId( snapshot: FileHistorySnapshot, sourceSessionId: string, @@ -2408,83 +2271,25 @@ function includeActiveSideArtifactRecords( const activeByUuid = new Map( activeRecords.map((record) => [record.uuid, record]), ); - const activeUuids = new Set(activeByUuid.keys()); - const firstActiveUuid = activeRecords[0]?.uuid; - const firstActiveIndex = - firstActiveUuid === undefined - ? -1 - : records.findIndex((record) => record.uuid === firstActiveUuid); - const nextActiveUuidByIndex = new Map(); - const nextBlockingUuidByIndex = new Map(); - let nextActiveUuid: string | undefined; - let nextBlockingUuid: string | undefined; - for (let index = records.length - 1; index >= 0; index--) { - if (nextActiveUuid !== undefined) { - nextActiveUuidByIndex.set(index, nextActiveUuid); - } - if (nextBlockingUuid !== undefined) { - nextBlockingUuidByIndex.set(index, nextBlockingUuid); - } - if (activeUuids.has(records[index]!.uuid)) { - nextActiveUuid = records[index]!.uuid; - nextBlockingUuid = undefined; - } else if ( - !isSessionArtifactRecord(records[index]!) && - !isTailNeutralSideRecord(records[index]!) - ) { - nextBlockingUuid = records[index]!.uuid; - } - } + const artifactUuids = new Set( + selectActiveSideArtifactRecordUuids( + records, + activeRecords.map((record) => record.uuid), + ), + ); const selected: ChatRecord[] = []; - const includedSideArtifactUuids = new Set(); - let previousActiveUuid: string | undefined; - for (let index = 0; index < records.length; index++) { - const record = records[index]!; + for (const record of records) { const activeRecord = activeByUuid.get(record.uuid); if (activeRecord) { selected.push(activeRecord); activeByUuid.delete(record.uuid); - previousActiveUuid = record.uuid; continue; } - if (!isSessionArtifactRecord(record)) { - continue; - } - const nextUuid = nextActiveUuidByIndex.get(index); - const hasBlockingRecordBeforeNextActive = - nextBlockingUuidByIndex.has(index); - const isInActiveSegment = - !hasBlockingRecordBeforeNextActive && - (nextUuid !== undefined - ? activeUuids.has(nextUuid) - : previousActiveUuid !== undefined && - activeUuids.has(previousActiveUuid)); - if ( - record.parentUuid !== null && - (activeUuids.has(record.parentUuid) || - includedSideArtifactUuids.has(record.parentUuid)) && - isInActiveSegment && - (record.parentUuid === previousActiveUuid || - includedSideArtifactUuids.has(record.parentUuid)) - ) { - selected.push(record); - includedSideArtifactUuids.add(record.uuid); - } else if ( - record.parentUuid === null && - index < firstActiveIndex && - isInActiveSegment - ) { - selected.push(record); - includedSideArtifactUuids.add(record.uuid); - } + if (artifactUuids.has(record.uuid)) selected.push(record); } return selected; } -function isTailNeutralSideRecord(record: ChatRecord): boolean { - return record.type === 'system' && record.subtype === 'custom_title'; -} - function collectFileHistorySnapshotPromptIds( records: ChatRecord[], ): Set { @@ -2546,68 +2351,6 @@ export function replayUiTelemetryFromConversation( return resumeTokenCounts; } -export interface ResumeTokenCounts { - promptTokenCount: number; - outputTokenCount: number; - isEstimated: boolean; -} - -/** - * Returns the best available prompt token count for resuming telemetry. - * Walks backward through messages and returns the first valid value: - * - The latest assistant's non-zero usage (promptTokenCount ?? totalTokenCount). - * - The most recent chat compression checkpoint's newTokenCount. - */ -export function getResumePromptTokenCount( - conversation: ConversationRecord, -): number | undefined { - return getResumeTokenCounts(conversation)?.promptTokenCount; -} - -/** - * Returns the prompt and previous-response output token counts used to seed a - * resumed chat. The prompt count restores the context anchor; the output - * count preserves the output tokens appended after that prompt count was - * reported, matching steady-state prompt estimation on the next send. - */ -export function getResumeTokenCounts( - conversation: ConversationRecord, -): ResumeTokenCounts | undefined { - for (let i = conversation.messages.length - 1; i >= 0; i--) { - const record = conversation.messages[i]; - - if (record.type === 'assistant') { - const usage = record.usageMetadata; - const candidate = usage?.promptTokenCount ?? usage?.totalTokenCount; - if (candidate) { - return { - promptTokenCount: candidate, - outputTokenCount: getUsageOutputTokenCountForPromptEstimate(usage), - isEstimated: false, - }; - } - } - - if (record.type === 'system' && record.subtype === 'chat_compression') { - const payload = record.systemPayload as - | ChatCompressionRecordPayload - | undefined; - if (payload?.info) { - return { - promptTokenCount: payload.info.newTokenCount, - outputTokenCount: 0, - // Checkpoints created before provenance was persisted are safest to - // treat as estimates: this keeps the output clamp's overhead pad on - // and prevents an optimistic resume from overflowing the window. - isEstimated: payload.info.newTokenCountIsEstimated ?? true, - }; - } - } - } - - return undefined; -} - const MAX_BRANCH_COLLISION_SCAN = 99; export async function computeUniqueBranchTitle( diff --git a/packages/core/src/utils/sessionStorageUtils.ts b/packages/core/src/utils/sessionStorageUtils.ts index 862a62b606..58dacc99e3 100644 --- a/packages/core/src/utils/sessionStorageUtils.ts +++ b/packages/core/src/utils/sessionStorageUtils.ts @@ -486,3 +486,22 @@ export function readLastJsonStringFieldsSync( } } } + +export function readSessionTitleInfoFromFileSync( + filePath: string, + scratchBuffer?: Buffer, +): { title?: string; source?: 'auto' | 'manual' } { + const hit = readLastJsonStringFieldsSync( + filePath, + 'customTitle', + ['titleSource'], + '"subtype":"custom_title"', + scratchBuffer, + ); + const title = hit['customTitle']; + if (!title) return {}; + const rawSource = hit['titleSource']; + const source = + rawSource === 'auto' || rawSource === 'manual' ? rawSource : undefined; + return { title, source }; +}