diff --git a/packages/opencode/src/v2/session-entry-stepper.ts b/packages/opencode/src/v2/session-entry-stepper.ts index 3d642579d0..3fe4266c04 100644 --- a/packages/opencode/src/v2/session-entry-stepper.ts +++ b/packages/opencode/src/v2/session-entry-stepper.ts @@ -1,4 +1,4 @@ -import { castDraft, produce, type WritableDraft } from "immer" +import { produce, type WritableDraft } from "immer" import { SessionEvent } from "./session-event" import { SessionEntry } from "./session-entry" @@ -235,7 +235,15 @@ export function stepWith(adapter: Adapter, event: SessionEvent.E ) } }, - retried: () => {}, + retried: (event) => { + if (currentAssistant) { + adapter.updateAssistant( + produce(currentAssistant, (draft) => { + draft.retries = [...(draft.retries ?? []), SessionEntry.AssistantRetry.fromEvent(event)] + }), + ) + } + }, compacted: (event) => { adapter.appendEntry(SessionEntry.Compaction.fromEvent(event)) }, diff --git a/packages/opencode/src/v2/session-entry.ts b/packages/opencode/src/v2/session-entry.ts index 97c5fc7ce9..b261d8b5b2 100644 --- a/packages/opencode/src/v2/session-entry.ts +++ b/packages/opencode/src/v2/session-entry.ts @@ -104,6 +104,24 @@ export class AssistantReasoning extends Schema.Class("Sessio text: Schema.String, }) {} +export class AssistantRetry extends Schema.Class("Session.Entry.Assistant.Retry")({ + attempt: Schema.Number, + error: SessionEvent.RetryError, + time: Schema.Struct({ + created: Schema.DateTimeUtc, + }), +}) { + static fromEvent(event: SessionEvent.Retried) { + return new AssistantRetry({ + attempt: event.attempt, + error: event.error, + time: { + created: event.timestamp, + }, + }) + } +} + export const AssistantContent = Schema.Union([AssistantText, AssistantReasoning, AssistantTool]).pipe( Schema.toTaggedUnion("type"), ) @@ -113,6 +131,7 @@ export class Assistant extends Schema.Class("Session.Entry.Assistant" ...Base, type: Schema.Literal("assistant"), content: AssistantContent.pipe(Schema.Array), + retries: AssistantRetry.pipe(Schema.Array, Schema.optional), cost: Schema.Number.pipe(Schema.optional), tokens: Schema.Struct({ input: Schema.Number, @@ -137,6 +156,7 @@ export class Assistant extends Schema.Class("Session.Entry.Assistant" created: event.timestamp, }, content: [], + retries: [], }) } } diff --git a/packages/opencode/src/v2/session-event.ts b/packages/opencode/src/v2/session-event.ts index 11d4a5db2d..f922becf3a 100644 --- a/packages/opencode/src/v2/session-event.ts +++ b/packages/opencode/src/v2/session-event.ts @@ -53,6 +53,15 @@ export namespace SessionEvent { source: Source.pipe(Schema.optional), }) {} + export class RetryError extends Schema.Class("Session.Event.Retry.Error")({ + message: Schema.String, + statusCode: Schema.Number.pipe(Schema.optional), + isRetryable: Schema.Boolean, + responseHeaders: Schema.Record(Schema.String, Schema.String).pipe(Schema.optional), + responseBody: Schema.String.pipe(Schema.optional), + metadata: Schema.Record(Schema.String, Schema.String).pipe(Schema.optional), + }) {} + export class Prompt extends Schema.Class("Session.Event.Prompt")({ ...Base, type: Schema.Literal("prompt"), @@ -386,14 +395,16 @@ export namespace SessionEvent { export class Retried extends Schema.Class("Session.Event.Retried")({ ...Base, type: Schema.Literal("retried"), - error: Schema.String, + attempt: Schema.Number, + error: RetryError, }) { - static create(input: BaseInput & { error: string }) { + static create(input: BaseInput & { attempt: number; error: RetryError }) { return new Retried({ id: input.id ?? ID.create(), type: "retried", timestamp: input.timestamp ?? DateTime.makeUnsafe(Date.now()), metadata: input.metadata, + attempt: input.attempt, error: input.error, }) } diff --git a/packages/opencode/test/session/session-entry-stepper.test.ts b/packages/opencode/test/session/session-entry-stepper.test.ts index 5c7df2dbad..32036cb1e8 100644 --- a/packages/opencode/test/session/session-entry-stepper.test.ts +++ b/packages/opencode/test/session/session-entry-stepper.test.ts @@ -27,6 +27,24 @@ function assistant() { type: "assistant", time: { created: time(0) }, content: [], + retries: [], + }) +} + +function retryError(message: string) { + return new SessionEvent.RetryError({ + message, + isRetryable: true, + }) +} + +function retry(attempt: number, message: string, created: number) { + return new SessionEntry.AssistantRetry({ + attempt, + error: retryError(message), + time: { + created: time(created), + }, }) } @@ -78,6 +96,12 @@ function tool(state: SessionEntryStepper.MemoryState, callID: string) { return tools(state).find((x) => x.callID === callID) } +function retriesOf(state: SessionEntryStepper.MemoryState) { + const entry = last(state) + if (!entry) return [] + return entry.retries ?? [] +} + function adapterStore() { return { committed: [] as SessionEntry.Entry[], @@ -168,6 +192,33 @@ describe("session-entry-stepper", () => { ]) expect(store.committed[0].time.completed).toEqual(time(7)) }) + + test("aggregates retry events onto the current assistant", () => { + const store = adapterStore() + store.committed.push(assistant()) + + SessionEntryStepper.stepWith( + adapterFor(store), + SessionEvent.Retried.create({ + attempt: 1, + error: retryError("rate limited"), + timestamp: time(1), + }), + ) + SessionEntryStepper.stepWith( + adapterFor(store), + SessionEvent.Retried.create({ + attempt: 2, + error: retryError("provider overloaded"), + timestamp: time(2), + }), + ) + + expect(store.committed[0]?.type).toBe("assistant") + if (store.committed[0]?.type !== "assistant") return + + expect(store.committed[0].retries).toEqual([retry(1, "rate limited", 1), retry(2, "provider overloaded", 2)]) + }) }) describe("memory", () => { @@ -231,6 +282,21 @@ describe("session-entry-stepper", () => { expect(reasons(state)).toEqual([{ type: "reasoning", text: "final" }]) }) + + test("stepWith through memory records retries", () => { + const state = active() + + SessionEntryStepper.stepWith( + SessionEntryStepper.memory(state), + SessionEvent.Retried.create({ + attempt: 1, + error: retryError("rate limited"), + timestamp: time(1), + }), + ) + + expect(retriesOf(state)).toEqual([retry(1, "rate limited", 1)]) + }) }) describe("step", () => { @@ -481,6 +547,27 @@ describe("session-entry-stepper", () => { }) }) + test("records retries on the pending assistant", () => { + const next = run( + [ + SessionEvent.Retried.create({ + attempt: 1, + error: retryError("rate limited"), + timestamp: time(1), + }), + SessionEvent.Retried.create({ + attempt: 2, + error: retryError("provider overloaded"), + timestamp: time(2), + }), + ], + active(), + ) + + expect(retriesOf(next)).toEqual([retry(1, "rate limited", 1), retry(2, "provider overloaded", 2)]) + }) + }) + describe("known reducer gaps", () => { test("prompt appends immutably when no assistant is pending", () => { FastCheck.assert(