From 8fc43e187ae9d5fe8a0b5d96c4b646c565bb6b65 Mon Sep 17 00:00:00 2001 From: Shakker <165377636+shakkernerd@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:45:45 +0100 Subject: [PATCH] fix(agents): preserve persisted turns across restart (#116924) Accept recorder-confirmed durable transcript identity when the active-session copy is absent, while preserving orphan repair for mismatched or unconfirmed turns.\n\nCloses #116894 --- .../run/attempt-session-boundary.test.ts | 168 ++++++++++++++++++ .../run/attempt-session-boundary.ts | 23 +-- .../run/pre-persisted-user-turn.ts | 17 +- 3 files changed, 192 insertions(+), 16 deletions(-) diff --git a/src/agents/embedded-agent-runner/run/attempt-session-boundary.test.ts b/src/agents/embedded-agent-runner/run/attempt-session-boundary.test.ts index 1bba44ab8da8..50d1d89eba5c 100644 --- a/src/agents/embedded-agent-runner/run/attempt-session-boundary.test.ts +++ b/src/agents/embedded-agent-runner/run/attempt-session-boundary.test.ts @@ -318,6 +318,174 @@ describe("prepareEmbeddedAttemptSessionBoundary", () => { }, ); + it("preserves the admitted current user when its active-session copy is absent", () => { + const currentUser = { + role: "user" as const, + content: "current prompt", + idempotencyKey: "current-run:user", + timestamp: 1, + }; + const { activeSession } = createActiveSession([]); + const branch = vi.fn(); + const resetLeaf = vi.fn(); + const clearNextUserMessagePersistenceSuppression = vi.fn(); + const onUserMessagePersistenceInvalidated = vi.fn(); + const sessionManager = createSessionManager({ + branch, + resetLeaf, + clearNextUserMessagePersistenceSuppression, + getLeafEntry: () => ({ + id: "current-user", + parentId: "previous-assistant", + timestamp: "2026-07-13T00:00:00.000Z", + type: "message", + message: currentUser, + }), + }); + const recorder = { + hasPersisted: () => true, + } as NonNullable< + Parameters< + typeof prepareEmbeddedAttemptSessionBoundary + >[0]["attempt"]["userTurnTranscriptRecorder"] + >; + + const boundary = prepareEmbeddedAttemptSessionBoundary({ + activeSession, + attempt: { + onUserMessagePersistenceInvalidated, + prompt: "current prompt", + trigger: "user", + userTurnTranscriptRecorder: recorder, + }, + getUserTranscriptContexts: () => undefined, + isRawModelRun: false, + preparedUserTurnMessage: currentUser, + sessionManager, + setActiveSessionSystemPrompt: vi.fn(), + }); + + expect(boundary.orphanRepair).toBeUndefined(); + expect(activeSession.agent.state.messages).toEqual([]); + expect(branch).not.toHaveBeenCalled(); + expect(resetLeaf).not.toHaveBeenCalled(); + expect(clearNextUserMessagePersistenceSuppression).not.toHaveBeenCalled(); + expect(onUserMessagePersistenceInvalidated).not.toHaveBeenCalled(); + }); + + it("repairs an exact durable user leaf when persistence is not recorder-confirmed", () => { + const currentUser = { + role: "user" as const, + content: "current prompt", + idempotencyKey: "current-run:user", + timestamp: 1, + }; + const repairedMessages: AgentMessage[] = [currentUser]; + const { activeSession } = createActiveSession([]); + const branch = vi.fn(); + const clearNextUserMessagePersistenceSuppression = vi.fn(); + const onUserMessagePersistenceInvalidated = vi.fn(); + const sessionManager = createSessionManager({ + getLeafEntry: () => ({ + id: "unconfirmed-user", + parentId: "previous-assistant", + timestamp: "2026-07-13T00:00:00.000Z", + type: "message", + message: currentUser, + }), + branch, + clearNextUserMessagePersistenceSuppression, + buildSessionContext: () => ({ messages: repairedMessages }), + }); + const recorder = { + hasPersisted: () => false, + } as NonNullable< + Parameters< + typeof prepareEmbeddedAttemptSessionBoundary + >[0]["attempt"]["userTurnTranscriptRecorder"] + >; + + const boundary = prepareEmbeddedAttemptSessionBoundary({ + activeSession, + attempt: { + onUserMessagePersistenceInvalidated, + prompt: "current prompt", + trigger: "user", + userTurnTranscriptRecorder: recorder, + }, + getUserTranscriptContexts: () => undefined, + isRawModelRun: false, + preparedUserTurnMessage: currentUser, + sessionManager, + setActiveSessionSystemPrompt: vi.fn(), + }); + + expect(boundary.orphanRepair?.removeLeaf).toBe(true); + expect(branch).toHaveBeenCalledWith("previous-assistant"); + expect(clearNextUserMessagePersistenceSuppression).toHaveBeenCalledOnce(); + expect(onUserMessagePersistenceInvalidated).toHaveBeenCalledOnce(); + expect(activeSession.agent.state.messages).toBe(repairedMessages); + }); + + it("repairs a durable user leaf that does not match the admitted current user", () => { + const currentUser = { + role: "user" as const, + content: "current prompt", + idempotencyKey: "current-run:user", + timestamp: 2, + }; + const repairedMessages: AgentMessage[] = [currentUser]; + const { activeSession } = createActiveSession([]); + const branch = vi.fn(); + const clearNextUserMessagePersistenceSuppression = vi.fn(); + const onUserMessagePersistenceInvalidated = vi.fn(); + const sessionManager = createSessionManager({ + getLeafEntry: () => ({ + id: "orphan-user", + parentId: "previous-assistant", + timestamp: "2026-07-13T00:00:00.000Z", + type: "message", + message: { + role: "user", + content: "old prompt", + idempotencyKey: "previous-run:user", + timestamp: 1, + }, + }), + branch, + clearNextUserMessagePersistenceSuppression, + buildSessionContext: () => ({ messages: repairedMessages }), + }); + const recorder = { + hasPersisted: () => true, + } as NonNullable< + Parameters< + typeof prepareEmbeddedAttemptSessionBoundary + >[0]["attempt"]["userTurnTranscriptRecorder"] + >; + + const boundary = prepareEmbeddedAttemptSessionBoundary({ + activeSession, + attempt: { + onUserMessagePersistenceInvalidated, + prompt: "current prompt", + trigger: "user", + userTurnTranscriptRecorder: recorder, + }, + getUserTranscriptContexts: () => undefined, + isRawModelRun: false, + preparedUserTurnMessage: currentUser, + sessionManager, + setActiveSessionSystemPrompt: vi.fn(), + }); + + expect(boundary.orphanRepair?.removeLeaf).toBe(true); + expect(branch).toHaveBeenCalledWith("previous-assistant"); + expect(clearNextUserMessagePersistenceSuppression).toHaveBeenCalledOnce(); + expect(onUserMessagePersistenceInvalidated).toHaveBeenCalledOnce(); + expect(activeSession.agent.state.messages).toBe(repairedMessages); + }); + it("repairs an orphaned user leaf before rebuilding active session messages", () => { const repairedMessages: AgentMessage[] = [ { role: "user", content: [{ type: "text", text: "repaired" }], timestamp: 2 }, diff --git a/src/agents/embedded-agent-runner/run/attempt-session-boundary.ts b/src/agents/embedded-agent-runner/run/attempt-session-boundary.ts index 25096d9b8744..2c2593e8dde9 100644 --- a/src/agents/embedded-agent-runner/run/attempt-session-boundary.ts +++ b/src/agents/embedded-agent-runner/run/attempt-session-boundary.ts @@ -9,7 +9,7 @@ import { resolveOrphanRepairPlan, } from "./attempt-orphan-repair.js"; import { normalizeMessagesForLlmBoundary } from "./attempt.llm-boundary.js"; -import { detachPrePersistedCurrentUserTurn } from "./pre-persisted-user-turn.js"; +import { reconcilePrePersistedCurrentUserTurn } from "./pre-persisted-user-turn.js"; import type { EmbeddedRunAttemptParams } from "./types.js"; type SessionBoundaryAttempt = Pick< @@ -50,21 +50,22 @@ export function prepareEmbeddedAttemptSessionBoundary(input: { input.setActiveSessionSystemPrompt(""); } - const detachedCurrentUser = + const orphanRepairCandidate = preserveExactPrompt + ? undefined + : resolveOrphanRepairPlan({ + sessionManager, + prompt: attempt.prompt, + trigger: attempt.trigger, + }); + const reconciledCurrentUser = !preserveExactPrompt && - detachPrePersistedCurrentUserTurn({ + reconcilePrePersistedCurrentUserTurn({ activeSession, + durableUserTurnMessage: orphanRepairCandidate?.messageEntry.message, preparedUserTurnMessage: input.preparedUserTurnMessage, userTurnAlreadyPersisted: attempt.userTurnTranscriptRecorder?.hasPersisted() === true, }); - const orphanRepair = - preserveExactPrompt || detachedCurrentUser - ? undefined - : resolveOrphanRepairPlan({ - sessionManager, - prompt: attempt.prompt, - trigger: attempt.trigger, - }); + const orphanRepair = reconciledCurrentUser ? undefined : orphanRepairCandidate; if (orphanRepair?.removeLeaf) { if (orphanRepair.messageEntry.parentId) { sessionManager.branch(orphanRepair.messageEntry.parentId); diff --git a/src/agents/embedded-agent-runner/run/pre-persisted-user-turn.ts b/src/agents/embedded-agent-runner/run/pre-persisted-user-turn.ts index dd5bd8b34911..f24334a9ad24 100644 --- a/src/agents/embedded-agent-runner/run/pre-persisted-user-turn.ts +++ b/src/agents/embedded-agent-runner/run/pre-persisted-user-turn.ts @@ -11,8 +11,9 @@ export function sessionMessagesContainIdempotencyKey( ); } -export function detachPrePersistedCurrentUserTurn(params: { +export function reconcilePrePersistedCurrentUserTurn(params: { activeSession: { agent: { state: { messages: AgentMessage[] } } }; + durableUserTurnMessage: AgentMessage | undefined; preparedUserTurnMessage: AgentMessage | undefined; userTurnAlreadyPersisted: boolean; }): boolean { @@ -25,13 +26,19 @@ export function detachPrePersistedCurrentUserTurn(params: { if (typeof idempotencyKey !== "string" || idempotencyKey.length === 0) { return false; } + const durableIdempotencyKey = ( + params.durableUserTurnMessage as { idempotencyKey?: unknown } | undefined + )?.idempotencyKey; const messages = params.activeSession.agent.state.messages; const tail = messages.at(-1) as (AgentMessage & { idempotencyKey?: unknown }) | undefined; - if (tail?.role !== "user" || tail.idempotencyKey !== idempotencyKey) { + const activeTailMatches = tail?.role === "user" && tail.idempotencyKey === idempotencyKey; + if (!activeTailMatches && durableIdempotencyKey !== idempotencyKey) { return false; } - // The durable transcript remains authoritative. Remove only its exact active - // tail copy so Agent.prompt() submits the current user turn once to the model. - params.activeSession.agent.state.messages = messages.slice(0, -1); + if (activeTailMatches) { + // Persistence is recorder-owned; either synchronized representation can + // prove identity. Remove the active copy when present so the model sees it once. + params.activeSession.agent.state.messages = messages.slice(0, -1); + } return true; }