From 5aa276c117a94a368868931eed967001c3aaa2f9 Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Fri, 24 Jul 2026 12:11:31 -0400 Subject: [PATCH] refactor(core): mint assistant message identity before the step runs (#38717) --- packages/core/src/session/runner/llm.ts | 33 ++++++++++++------- .../src/session/runner/publish-llm-event.ts | 11 +++---- packages/core/src/session/runner/retry.ts | 5 ++- .../test/session-runner-tool-events.test.ts | 1 + 4 files changed, 28 insertions(+), 22 deletions(-) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 94d5c7580e..7c96a1853d 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -30,7 +30,7 @@ import { SessionUsage } from "../usage" /** How one model call ended: settled, awaiting a scheduled retry, or restarted by compaction. */ type CallOutcome = Data.TaggedEnum<{ Completed: { readonly needsContinuation: boolean; readonly step: number } - Retry: { readonly step: number; readonly assistantMessageID: SessionMessage.ID } + Retry: { readonly step: number } Restart: { readonly step: number; readonly recoveredOverflow: boolean } }> const CallOutcome = Data.taggedEnum() @@ -120,20 +120,26 @@ const layer = Layer.effect( promotable: SessionPending.Promotable, step: number, ) { - const retry = yield* Schedule.toStepWithSleep(SessionRunnerRetry.schedule(events, sessionID)) + // Minting message identity before any attempt lets retries resume the same durable + // message. A compaction restart re-mints: the old message is stranded behind the new + // compaction boundary, so the rebuilt step needs identity inside the new epoch. + let assistantMessageID = SessionMessage.ID.create() + const retry = yield* Schedule.toStepWithSleep( + SessionRunnerRetry.schedule(events, sessionID, () => assistantMessageID), + ) /** - * Consumes one retry allowance: sleeps the scheduled backoff and reports what the next - * attempt should reuse, or publishes Step.Failed and fails once attempts are exhausted. - * The step loop performs the retry itself on the next iteration. + * Consumes one retry allowance: sleeps the scheduled backoff, or publishes + * Step.Failed and fails once attempts are exhausted. The step loop performs + * the retry itself on the next iteration. */ const waitForRetry = (failure: SessionRunnerRetry.RetryableFailure) => retry(failure).pipe( - Effect.as(CallOutcome.Retry({ step: failure.step, assistantMessageID: failure.assistantMessageID })), + Effect.as(CallOutcome.Retry({ step: failure.step })), Pull.catchDone(() => events .publish(SessionEvent.Step.Failed, { sessionID, - assistantMessageID: failure.assistantMessageID, + assistantMessageID, error: failure.error, }) .pipe(Effect.andThen(Effect.fail(failure.cause))), @@ -141,7 +147,6 @@ const layer = Layer.effect( ) let currentPromotable: SessionPending.Promotable | undefined = promotable let currentStep = step - let assistantMessageID: SessionMessage.ID | undefined // Overflow recovery is one-shot: a call after recovery must not recover another overflow. let recoverOverflow = true while (true) { @@ -153,8 +158,10 @@ const layer = Layer.effect( assistantMessageID, ).pipe(Effect.catchTag("SessionRunner.RetryableFailure", waitForRetry)) if (outcome._tag === "Completed") return { needsContinuation: outcome.needsContinuation, step: outcome.step } - if (outcome._tag === "Retry") assistantMessageID = outcome.assistantMessageID - if (outcome._tag === "Restart" && outcome.recoveredOverflow) recoverOverflow = false + if (outcome._tag === "Restart") { + if (outcome.recoveredOverflow) recoverOverflow = false + assistantMessageID = SessionMessage.ID.create() + } // Neither a retry nor a compaction restart re-promotes input. currentPromotable = undefined currentStep = outcome.step @@ -170,7 +177,7 @@ const layer = Layer.effect( promotable: SessionPending.Promotable | undefined, step: number, recoverOverflow: boolean, - assistantMessageID?: SessionMessage.ID, + assistantMessageID: SessionMessage.ID, ) { const selected = yield* context.select(sessionID) // Establish what the model knows before admitting what the user said, so @@ -318,9 +325,11 @@ const layer = Layer.effect( if (llmFailure && !publisher.hasProviderError()) { const error = toSessionError(llmFailure) if (SessionRunnerRetry.isRetryable(llmFailure) && !publisher.hasRetryEvidence()) { + // RetryScheduled and Step.Failed fold onto an existing assistant message, so + // Step.Started must be durable before the failure escapes. + yield* serialized(publisher.startAssistant()) return yield* new SessionRunnerRetry.RetryableFailure({ cause: llmFailure, - assistantMessageID: yield* publisher.startAssistant(), error, step: currentStep, }) diff --git a/packages/core/src/session/runner/publish-llm-event.ts b/packages/core/src/session/runner/publish-llm-event.ts index 06fedf6fcf..bb00c540c6 100644 --- a/packages/core/src/session/runner/publish-llm-event.ts +++ b/packages/core/src/session/runner/publish-llm-event.ts @@ -20,7 +20,7 @@ type Input = { readonly model: ModelV2.Ref readonly providerMetadataKey: string readonly snapshot?: Snapshot.ID - readonly assistantMessageID?: SessionMessage.ID + readonly assistantMessageID: SessionMessage.ID } const record = (value: unknown): Record => @@ -50,7 +50,7 @@ export const createLLMEventPublisher = (events: Pick() const failureSnapshot = (tool: { readonly progress?: ToolRegistry.Progress }) => tool.progress === undefined ? {} : { metadata: tool.progress } - let assistantMessageID = input.assistantMessageID + const assistantMessageID = input.assistantMessageID let stepStarted = false let stepFailed = false let providerFailed = false @@ -64,8 +64,7 @@ export const createLLMEventPublisher = (events: Pick - assistantMessageID === undefined - ? Effect.die(new Error("Tool event before assistant step start")) - : Effect.succeed(assistantMessageID) + stepStarted ? Effect.succeed(assistantMessageID) : Effect.die(new Error("Tool event before assistant step start")) const providerState = (metadata: ProviderMetadata | undefined) => metadata?.[input.providerMetadataKey] const fragments = ( name: string, diff --git a/packages/core/src/session/runner/retry.ts b/packages/core/src/session/runner/retry.ts index 94401a35d6..de94d36ec6 100644 --- a/packages/core/src/session/runner/retry.ts +++ b/packages/core/src/session/runner/retry.ts @@ -10,7 +10,6 @@ import { SessionSchema } from "../schema" export class RetryableFailure extends Data.TaggedError("SessionRunner.RetryableFailure")<{ readonly cause: LLMError - readonly assistantMessageID: SessionMessage.ID readonly error: SessionError.Error readonly step: number }> {} @@ -42,7 +41,7 @@ const retryAfter = (failure: RetryableFailure) => { return undefined } -export const schedule = (events: EventV2.Interface, sessionID: SessionSchema.ID) => +export const schedule = (events: EventV2.Interface, sessionID: SessionSchema.ID, assistantMessageID: () => SessionMessage.ID) => Schedule.max([Schedule.exponential("2 seconds"), Schedule.recurs(4)]).pipe( Schedule.setInputType(), Schedule.modifyDelay(({ input: failure, duration: delay }) => { @@ -52,7 +51,7 @@ export const schedule = (events: EventV2.Interface, sessionID: SessionSchema.ID) Schedule.tap((metadata) => events.publish(SessionEvent.RetryScheduled, { sessionID, - assistantMessageID: metadata.input.assistantMessageID, + assistantMessageID: assistantMessageID(), attempt: metadata.attempt + 1, at: metadata.now + Duration.toMillis(metadata.duration), error: metadata.input.error, diff --git a/packages/core/test/session-runner-tool-events.test.ts b/packages/core/test/session-runner-tool-events.test.ts index 4d6dc3c440..ac90c1dd4f 100644 --- a/packages/core/test/session-runner-tool-events.test.ts +++ b/packages/core/test/session-runner-tool-events.test.ts @@ -45,6 +45,7 @@ const capture = (providerMetadataKey = "anthropic", options?: { readonly interru providerID: ProviderV2.ID.opencode, }, providerMetadataKey, + assistantMessageID: SessionMessage.ID.create(), }), } }