fix: store the background jobs with the transcript that depends on them

2026-09-06 00:02:13 +00:00 · 2026-09-03 16:30:09 +02:00
parent f9252aad78
commit af344c1a9a
2 changed files with 89 additions and 32 deletions
@@ -813,19 +813,23 @@ export class AIChatManager {
 	// turn-end save.
 	#maskPersistQueue: Promise<void> = Promise.resolve()
 	#persistModifiedItems(): Promise<void> {
-		this.#maskPersistQueue = this.#maskPersistQueue.then(() =>
-			this.historyManager
-				.saveChat(
-					this.#interruptedSnapshot(),
-					this.messages,
-					this.contextUsage,
-					this.modifiedItems ? [...this.modifiedItems] : undefined
-				)
-				// Swallow (and log) a failed write so it can't wedge the queue as a
-				// rejected link — the next persist snapshots the full current set, so
-				// a lost write self-heals on the next mutation or turn-end save.
-				.catch((e) => console.error('Failed to persist modified-items mask', e))
-		)
+		this.#maskPersistQueue = this.#maskPersistQueue.then(() => {
+			const { display, jobs } = this.#interruptedSnapshot()
+			return (
+				this.historyManager
+					.saveChat(
+						display,
+						this.messages,
+						this.contextUsage,
+						this.modifiedItems ? [...this.modifiedItems] : undefined,
+						jobs
+					)
+					// Swallow (and log) a failed write so it can't wedge the queue as a
+					// rejected link — the next persist snapshots the full current set, so
+					// a lost write self-heals on the next mutation or turn-end save.
+					.catch((e) => console.error('Failed to persist modified-items mask', e))
+			)
+		})
 		return this.#maskPersistQueue
 	}

@@ -1137,17 +1141,12 @@ export class AIChatManager {
 	// (saveChat keeps the prior mask when it is undefined).
 	#jobPersistQueue: Promise<void> = Promise.resolve()
 	#persistBackgroundJobs(): Promise<void> {
-		this.#jobPersistQueue = this.#jobPersistQueue.then(() =>
-			this.historyManager
-				.saveChat(
-					this.#interruptedSnapshot(),
-					this.messages,
-					this.contextUsage,
-					undefined,
-					$state.snapshot(this.backgroundJobs)
-				)
+		this.#jobPersistQueue = this.#jobPersistQueue.then(() => {
+			const { display, jobs } = this.#interruptedSnapshot()
+			return this.historyManager
+				.saveChat(display, this.messages, this.contextUsage, undefined, jobs)
 				.catch((e) => console.error('Failed to persist background jobs', e))
-		)
+		})
 		return this.#jobPersistQueue
 	}

@@ -3385,7 +3384,7 @@ export class AIChatManager {
 			)
 			if (messages.length === this.messages.length) return
 			checkpointedShape = shape
-			const display = this.#interruptedSnapshot()
+			const { display, jobs } = this.#interruptedSnapshot()
 			// onMessageEnd is what gives streamed text its bubble, and it clears
 			// currentReply doing so — text still there has none, and without one the
 			// reply returns as context the reader cannot see.
@@ -3405,7 +3404,8 @@ export class AIChatManager {
 					// partial turn — enough to skip the compaction its next send needs.
 					// Omitting drops the field, which is the "readers estimate" fallback.
 					undefined,
-					this.modifiedItems ? [...this.modifiedItems] : undefined
+					this.modifiedItems ? [...this.modifiedItems] : undefined,
+					jobs
 				)
 			} catch (e) {
 				console.error('Failed to checkpoint chat mid-turn', e)
@@ -4668,18 +4668,25 @@ export class AIChatManager {
 	 * "Interrupted" error, and the patch a completed job merges in carries no error to
 	 * clear it with. Which cards those are is loadPastChat's question, so ask it the same
 	 * way — a job still running inline is detached by the restore and polled like any
-	 * other. */
-	#interruptedSnapshot = (): DisplayMessage[] => {
+	 * other.
+	 *
+	 * That exemption is a bet on the poller, and the poller only knows the jobs stored
+	 * in the same record — registering one does not write it. So the jobs come back
+	 * with the transcript that depends on them, and both go into the same saveChat. */
+	#interruptedSnapshot = (): { display: DisplayMessage[]; jobs: ChatJob[] } => {
 		const polled = new Set(
 			this.backgroundJobs
 				.filter((j) => j.detached || this.isJobNonTerminal(j.status))
 				.map((j) => j.toolCallId)
 		)
-		return this.settledToolDisplay(
-			this.displayMessages,
-			'Interrupted',
-			(message) => !polled.has(message.tool_call_id)
-		)
+		return {
+			display: this.settledToolDisplay(
+				this.displayMessages,
+				'Interrupted',
+				(message) => !polled.has(message.tool_call_id)
+			),
+			jobs: $state.snapshot(this.backgroundJobs) as ChatJob[]
+		}
 	}
 }

@@ -2289,6 +2289,56 @@ describe('AIChatManager queued messages', () => {
 		])
 	})

+	// A checkpoint that leaves a card loading is betting the poller resolves it after
+	// the reload, and the poller only knows the jobs stored in the same record —
+	// registering one does not write it.
+	it('stores the job behind a card the checkpoint leaves loading', async () => {
+		const leavePage = stubHidingPage()
+		const manager = createManager()
+		const saveChat = vi.spyOn(manager.historyManager, 'saveChat').mockResolvedValue(undefined)
+
+		mocks.runChatLoop.mockImplementationOnce(async (config: any) => {
+			config.addedMessages.push({
+				role: 'assistant' as const,
+				content: '',
+				tool_calls: [
+					{ id: 't1', type: 'function' as const, function: { name: 'run_script', arguments: '{}' } }
+				]
+			})
+			// Inside the inline wait: the job is registered and still running, so no
+			// persist path has run for it yet.
+			manager.registerJob({
+				jobId: 'job-1',
+				toolCallId: 't1',
+				kind: 'script',
+				label: 'f/a/b',
+				workspace: 'ws'
+			})
+			config.callbacks.setToolStatus('t1', { content: 'Running...', isLoading: true })
+			leavePage.forEach((fn) => fn())
+			// The wait ends normally, so the only save that stored this card loading is
+			// the checkpoint that landed inside it.
+			config.callbacks.setToolStatus('t1', { content: 'Ran', isLoading: false })
+			manager.updateJob('job-1', { status: 'success' })
+			config.addedMessages.push({ role: 'tool' as const, tool_call_id: 't1', content: 'ran' })
+			return {
+				addedMessages: config.addedMessages,
+				tokenUsage: { prompt: 0, completion: 0, total: 0 },
+				hitMaxIterations: false
+			}
+		})
+
+		await manager.sendRequest({ instructions: 'run it' })
+
+		const checkpoint = saveChat.mock.calls.find(([display]) =>
+			(display as DisplayMessage[]).some(
+				(m) => m.role === 'tool' && m.tool_call_id === 't1' && m.isLoading
+			)
+		)
+		expect(checkpoint).toBeDefined()
+		expect(checkpoint?.[4]).toEqual([expect.objectContaining({ jobId: 'job-1' })])
+	})
+
 	it('stops checkpointing once the turn commits, so the transcript is never doubled', async () => {
 		const leavePage = stubHidingPage()
 		const manager = createManager()