From 29d16087e89d71d22846c9ffd548ff4149e73c50 Mon Sep 17 00:00:00 2001 From: Ali Mohammad Date: Tue, 29 Sep 2026 02:03:24 +0300 Subject: [PATCH] LOU-K5: restore checkpointing after resumeAfterApproval() resumeAfterApproval() correctly deleted the stale pre-pause checkpoint under snapshot.sessionId, but then forced sessionId/checkpointStore to undefined on the follow-up AgentExecutor.execute() call as a belt-and-suspenders guard against re-rehydrating that stale state. The guard was overbroad: it also silently disabled checkpointing for the entire remainder of the resumed run, so a crash a few tool-calls after a human approval would lose all progress made since the approval. Now the real sessionId/checkpointStore are threaded through instead. This stays safe because AgentExecutor's rehydration branch only fires when checkpointStore.load(sessionId) finds something, and the delete above guarantees it won't - so the resumed run always falls into the normal "build from scratch" path using the reconstructed messages/initialSteps, and AgentExecutor's per-tool-result checkpoint.save() calls resume writing durable progress from there on. Adds tests proving: a new checkpoint is written for a tool-call step taken after a successful resume, the stale pre-pause checkpoint is still never rehydrated, and resuming without a checkpointStore remains fully backward compatible. Co-Authored-By: Claude Sonnet 5 --- src/execution/resume.test.ts | 178 +++++++++++++++++++++++++++++++++++ src/execution/resume.ts | 115 +++++++++++++++------- 2 files changed, 258 insertions(+), 35 deletions(-) diff --git a/src/execution/resume.test.ts b/src/execution/resume.test.ts index 219405db3..6d709b145 100644 --- a/src/execution/resume.test.ts +++ b/src/execution/resume.test.ts @@ -552,6 +552,184 @@ describe('Execution - resumeAfterApproval', () => { expect(await checkpointStore.load(sessionId)).toBeNull(); }); + it('LOU-K5: writes a NEW checkpoint for a tool-call step taken after a successful resume, when a checkpointStore is supplied', async () => { + // Reproduces/verifies the LOU-K5 fix: resumeAfterApproval() used to + // force sessionId/checkpointStore to `undefined` on the follow-up + // execute() call, which silently disabled checkpointing for the whole + // remainder of the resumed run. This test drives the resumed run + // through one MORE (non-approval-gated) tool call after the approved + // 'chargeCard' call, and asserts AgentExecutor's normal per-tool-result + // checkpoint.save() fires for that later step - proving durability is + // restored for the post-resume portion of the run, not just re-armed + // and immediately turned back off. + const chargeExecute = vi.fn().mockResolvedValue({ charged: true }); + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { description: 'Charge a card', parameters: {}, execute: chargeExecute } as any, + needsApproval: true, + }); + + const lookupExecute = vi.fn().mockResolvedValue({ found: true }); + toolRegistry.register('lookup', { + displayName: 'Lookup', + tool: { description: 'Look something up', parameters: {}, execute: lookupExecute } as any, + needsApproval: false, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .addTool('lookup', { tool: 'lookup', options: {} }) + .build(); + + // Generation 1: calls the approval-gated chargeCard tool (pauses). + // Generation 2 (post-resume): calls the unguarded lookup tool - this is + // the "one more tool-call step after resume" that must get checkpointed. + // Generation 3: stops with no further tool calls. + let call = 0; + const scriptedProvider = { + name: 'scripted', + supportsTools: () => true, + supportsStreaming: () => false, + getModels: async () => ['scripted'], + stream: async () => { + throw new Error('not implemented'); + }, + generate: async () => { + call++; + if (call === 1) { + return { + text: 'charging', + finishReason: 'tool_calls' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + toolCalls: [ + { id: 'call-charge', type: 'function' as const, function: { name: 'chargeCard', arguments: '{}' } }, + ], + }; + } + if (call === 2) { + return { + text: 'looking up', + finishReason: 'tool_calls' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + toolCalls: [ + { id: 'call-lookup', type: 'function' as const, function: { name: 'lookup', arguments: '{}' } }, + ], + }; + } + return { + text: 'all done', + finishReason: 'stop' as const, + usage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }; + }, + }; + + const approvalStore = createInMemoryApprovalStore(); + const checkpointStore = createInMemoryCheckpointStore(); + const saveSpy = vi.spyOn(checkpointStore, 'save'); + const sessionId = 'post-resume-checkpoint-session'; + + const paused = await AgentExecutor.execute({ + agent, + input: 'go', + provider: scriptedProvider as any, + toolRegistry, + approvalStore, + sessionId, + checkpointStore, + }); + + expect(paused.finishReason).toBe('awaiting-approval'); + expect(chargeExecute).not.toHaveBeenCalled(); + saveSpy.mockClear(); + + const resumed = await resumeAfterApproval( + { id: paused.approvalId!, approved: true }, + approvalStore, + toolRegistry, + scriptedProvider as any, + {}, + checkpointStore + ); + + expect(resumed.finishReason).toBe('stop'); + expect(chargeExecute).toHaveBeenCalledTimes(1); + expect(lookupExecute).toHaveBeenCalledTimes(1); + + // A NEW checkpoint must have been saved (under the same sessionId) for + // the post-resume 'lookup' step - this is the forward-checkpointing + // this fix restores. Before the fix, sessionId/checkpointStore were + // forced to undefined on the resumed execute() call, so this save() + // would never have been reached at all. + expect(saveSpy).toHaveBeenCalled(); + const savedCheckpoint = saveSpy.mock.calls[saveSpy.mock.calls.length - 1][1]; + expect(savedCheckpoint.sessionId).toBe(sessionId); + // The saved checkpoint's messages must include BOTH the deferred + // chargeCard result (reconstructed by resume.ts) and the post-resume + // lookup result (checkpointed by AgentExecutor's normal loop) - proof + // this checkpoint reflects genuinely new, post-resume progress rather + // than a rehydrated/duplicated stale snapshot. + expect(savedCheckpoint.messages.some((m: any) => m.toolName === 'chargeCard')).toBe(true); + expect(savedCheckpoint.messages.some((m: any) => m.toolName === 'lookup')).toBe(true); + + // The run reached a terminal state, so AgentExecutor's own + // "clear the checkpoint when finished" cleanup (unrelated to this fix) + // removes it afterwards - this is expected and not a regression. + expect(await checkpointStore.load(sessionId)).toBeNull(); + }); + + it('LOU-K5: resuming without a checkpointStore still works exactly as before (no crash, no checkpoint attempted)', async () => { + // Guards the "fully optional/backward compatible" requirement: a + // caller that never passes a checkpointStore to resumeAfterApproval() + // must see identical behavior to before this fix - sessionId and + // checkpointStore simply stay undefined on the follow-up execute() + // call, so AgentExecutor's rehydration/save/delete branches (all + // gated on `sessionId && checkpointStore`) are never entered. + const execute = vi.fn().mockResolvedValue({ charged: true }); + toolRegistry.register('chargeCard', { + displayName: 'Charge Card', + tool: { description: 'Charge a card', parameters: {}, execute } as any, + needsApproval: true, + }); + + const agent = AgentBuilder.create() + .setType(AgentType.SmartAssistant) + .setName('Test Agent') + .addTool('chargeCard', { tool: 'chargeCard', options: {} }) + .build(); + + const provider = createMockProvider({ + name: 'mock', + responses: ['Charging now', 'All done'], + }); + + const approvalStore = createInMemoryApprovalStore(); + + const paused = await AgentExecutor.execute({ + agent, + input: 'Please call chargeCard now', + provider, + toolRegistry, + approvalStore, + // Deliberately no sessionId/checkpointStore anywhere in this test. + }); + + expect(paused.finishReason).toBe('awaiting-approval'); + + // No checkpointStore argument at all (not even undefined explicitly). + const resumed = await resumeAfterApproval( + { id: paused.approvalId!, approved: true }, + approvalStore, + toolRegistry, + provider + ); + + expect(execute).toHaveBeenCalledTimes(1); + expect(resumed.finishReason).toBe('stop'); + }); + it('should throw a clear error for an unknown or already-resolved approval id', async () => { const approvalStore = createInMemoryApprovalStore(); const provider = createMockProvider({ name: 'mock' }); diff --git a/src/execution/resume.ts b/src/execution/resume.ts index 2d90515ee..47f490f1e 100644 --- a/src/execution/resume.ts +++ b/src/execution/resume.ts @@ -16,17 +16,24 @@ import { executeToolWithSandboxGuard } from './sandboxGuard'; * once the deferred tool result (or rejection) has been appended to the * conversation. * - * `sessionId`/`checkpointStore` are also omitted here (in addition to the - * fields already omitted below): the resume path always fully reconstructs - * `messages`/`steps` from the ApprovalGate's ExecutionSnapshot, so it must - * never delegate to AgentExecutor's checkpoint-rehydration mechanism. Any - * checkpoint that exists under a given sessionId at this point was written - * BEFORE the pause (checkpoints are only saved after a tool *result*, and a - * pause happens before the deferred tool runs), so it is stale by - * construction and would silently clobber the just-reconstructed messages - * (including the deferred tool's own result) and step count if allowed - * through. See resumeAfterApproval() below for the belt-and-suspenders - * runtime guard that backs this type-level omission. + * `sessionId`/`checkpointStore` are omitted here (in addition to the other + * fields already omitted below) for the same reason `agent`/`input` are: + * resumeAfterApproval() derives them itself rather than taking them from the + * caller's `executeOptions`. `sessionId` comes from `snapshot.sessionId` + * (the ExecutionSnapshot the paused run was paused with) and `checkpointStore` + * comes from this function's own dedicated `checkpointStore` parameter - see + * resumeAfterApproval() below for exactly what gets threaded through and why + * that's safe. + * + * LOU-K5 update: earlier, this file forced BOTH to `undefined` on the + * follow-up execute() call, as a "belt-and-suspenders" guard against + * re-rehydrating stale pre-pause state. That guard was overbroad: it also + * disabled checkpointing for the entire remainder of the resumed run, so a + * crash a few tool-calls after a human approval would lose all progress + * since the approval. resumeAfterApproval() now threads the real + * `sessionId`/`checkpointStore` through instead, once it has already + * deleted the stale pre-pause checkpoint - see the "Defense in depth" + * comment below for why that ordering still makes rehydration impossible. */ export type ResumeExecuteOptions = Omit< ExecuteOptions, @@ -49,14 +56,24 @@ export type ResumeExecuteOptions = Omit< * so this takes the LLMProvider needed to keep generating instead of an * `executor: AgentExecutor` instance. * - * @param checkpointStore Optional durable-execution CheckpointStore. This - * is NOT threaded through to the follow-up execute() call (see - * ResumeExecuteOptions) - it is used only, when the paused run's - * ExecutionSnapshot carries a sessionId, to proactively delete the stale - * checkpoint left behind under that sessionId from before the pause. This - * closes the gap for a LATER, unrelated execute() call against the same - * sessionId+checkpointStore (one that doesn't go through resume at all) - * so it can't silently rehydrate that stale pre-pause state either. + * @param checkpointStore Optional durable-execution CheckpointStore. When + * the paused run's ExecutionSnapshot carries a `sessionId`, this is used + * two ways: first, to proactively delete the stale checkpoint left behind + * under that sessionId from before the pause (this closes the gap for a + * LATER, unrelated execute() call against the same sessionId+checkpointStore + * that doesn't go through resume at all, so it can't silently rehydrate + * that stale pre-pause state either); second - LOU-K5 - it is then threaded + * through, together with `snapshot.sessionId`, to the follow-up + * AgentExecutor.execute() call that continues the run, so checkpointing + * resumes for the remainder of the run instead of staying silently + * disabled. This is safe because the delete happens first: by the time + * execute() runs its rehydration check (`checkpointStore.load(sessionId)`), + * there is nothing under that key, so it falls into the normal + * "build from scratch" path using the reconstructed `messages`/ + * `initialSteps` below, exactly as if no checkpointStore had been passed at + * all - it never rehydrates. If `checkpointStore` is omitted entirely, + * both the delete and the pass-through are skipped and resume behaves + * exactly as it did before durable execution existed. */ export async function resumeAfterApproval( decision: ApprovalDecision, @@ -78,6 +95,24 @@ export async function resumeAfterApproval( // written before this pause and is therefore stale by construction - // clear it now so it can't be loaded by this resume or by some later, // unrelated execute() call against the same sessionId+checkpointStore. + // + // Doing this BEFORE the AgentExecutor.execute() call below (rather than + // just relying on `skipSystemPromptInjection`/`initialSteps` to make a + // rehydration "harmless") is what makes it safe to pass sessionId and + // checkpointStore through to that call: its rehydration branch only + // triggers when `checkpointStore.load(sessionId)` finds something, and + // by the time it runs (synchronously after this delete resolves, with no + // intervening await back to caller code) there is nothing left to find. + // + // Accepted race: if some OTHER process/caller writes a new checkpoint + // under this exact sessionId in the narrow window between this delete + // and the execute() call's load(), that checkpoint would be picked up + // instead of the freshly-reconstructed messages. This mirrors how + // ApprovalStore.resolve() above is documented as delete-on-read without + // an additional distributed lock (see the catch block below) - this + // codebase accepts same-sessionId-concurrent-caller races as an existing + // caller-responsibility invariant (a sessionId identifies a single + // logical run) rather than adding cross-process locking to CheckpointStore. if (snapshot.sessionId && checkpointStore) { await checkpointStore.delete(snapshot.sessionId); } @@ -159,25 +194,35 @@ export async function resumeAfterApproval( input: messages, provider, toolRegistry, - // Belt-and-suspenders runtime guard backing the ResumeExecuteOptions - // type-level Omit above: `Omit<...>` only stops *type-checked* callers - // from passing sessionId/checkpointStore through executeOptions - it - // does nothing to stop a caller who bypasses the type (e.g. an `as - // any` cast, or a plain-JS caller) from putting them on the object at - // runtime, where `...executeOptions` would otherwise spread them - // straight into this call. Explicitly forcing both to undefined here, - // after the spread, guarantees AgentExecutor.execute() can never fall - // into its checkpoint-rehydration branch for a resume, regardless of - // what executeOptions actually contains. - sessionId: undefined, - checkpointStore: undefined, + // LOU-K5: thread sessionId/checkpointStore through so AgentExecutor + // resumes writing per-tool-result checkpoints for the rest of this run + // (previously both were forced to `undefined` here, which also killed + // forward checkpointing for the whole remainder of the resumed run - + // not just the resume step itself). `snapshot.sessionId` and this + // function's own `checkpointStore` parameter are used explicitly here + // rather than whatever (if anything) `executeOptions` carries, since + // `ResumeExecuteOptions` omits both fields - see that type's doc + // comment. This still can't rehydrate stale state: the stale + // pre-pause checkpoint under `snapshot.sessionId` was just deleted + // above, so AgentExecutor.execute()'s `checkpointStore.load(sessionId)` + // rehydration check finds nothing and falls into its normal + // "build from scratch" path, using exactly the `input`/ + // `skipSystemPromptInjection`/`initialSteps` reconstructed below - it + // never re-triggers the rehydration branch this guard used to worry + // about. When no `checkpointStore` was passed to resumeAfterApproval() + // at all, this is `undefined` and AgentExecutor.execute() behaves + // exactly as it always has for callers that don't use durable + // execution. + sessionId: snapshot.sessionId, + checkpointStore, // `messages` was reconstructed from the ExecutionSnapshot's // currentMessages, which already include the original system message // (if any) that AgentExecutor.buildMessages() built the first time - // this agent ran. Since no checkpoint/sessionId is threaded through - // this call, execute() would otherwise fall into its "build from - // scratch" path and prepend a second, duplicate system message built - // fresh from agent.prompt. + // this agent ran. Because the stale checkpoint was just deleted above, + // execute() always falls into its "build from scratch" path here + // (never the rehydration path) - without skipSystemPromptInjection, + // that path would prepend a second, duplicate system message built + // fresh from agent.prompt on top of the one already in `messages`. skipSystemPromptInjection: true, // Continue step-budget accounting from where the paused run left off, // rather than silently resetting to a full fresh maxSteps allowance.