diff --git a/packages/opencode/src/session/processor.ts b/packages/opencode/src/session/processor.ts index 38dac41b0589..9927969d4058 100644 --- a/packages/opencode/src/session/processor.ts +++ b/packages/opencode/src/session/processor.ts @@ -364,7 +364,7 @@ export namespace SessionProcessor { error, }) } else { - const retry = SessionRetry.retryable(error) + const retry = SessionRetry.retryable(error, input.abort) if (retry !== undefined) { attempt++ const delay = SessionRetry.delay(attempt, error.name === "APIError" ? error : undefined) diff --git a/packages/opencode/src/session/retry.ts b/packages/opencode/src/session/retry.ts index 6d057f539f81..1ef9933a447b 100644 --- a/packages/opencode/src/session/retry.ts +++ b/packages/opencode/src/session/retry.ts @@ -58,9 +58,13 @@ export namespace SessionRetry { return Math.min(RETRY_INITIAL_DELAY * Math.pow(RETRY_BACKOFF_FACTOR, attempt - 1), RETRY_MAX_DELAY_NO_HEADERS) } - export function retryable(error: ReturnType) { + export function retryable(error: ReturnType, abort?: AbortSignal) { // context overflow errors should not be retried if (MessageV2.ContextOverflowError.isInstance(error)) return undefined + if (MessageV2.AbortedError.isInstance(error)) { + if (abort?.aborted) return undefined + return error.data.message + } if (MessageV2.APIError.isInstance(error)) { if (!error.data.isRetryable) return undefined if (error.data.responseBody?.includes("FreeUsageLimitError")) diff --git a/packages/opencode/test/session/llm.test.ts b/packages/opencode/test/session/llm.test.ts index b44191b6ba2d..5202c06dd934 100644 --- a/packages/opencode/test/session/llm.test.ts +++ b/packages/opencode/test/session/llm.test.ts @@ -368,8 +368,8 @@ describe("session.llm.stream", () => { await Instance.provide({ directory: tmp.path, fn: async () => { - const resolved = await Provider.getModel(providerID, model.id) - const sessionID = "session-test-tools" + const resolved = await Provider.getModel(ProviderID.make(providerID), ModelID.make(model.id)) + const sessionID = SessionID.make("session-test-tools") const agent = { name: "test", mode: "primary", @@ -378,12 +378,12 @@ describe("session.llm.stream", () => { } satisfies Agent.Info const user = { - id: "user-tools", + id: MessageID.make("user-tools"), sessionID, role: "user", time: { created: Date.now() }, agent: agent.name, - model: { providerID, modelID: resolved.id }, + model: { providerID: ProviderID.make(providerID), modelID: resolved.id }, tools: { question: true }, } satisfies MessageV2.User diff --git a/packages/opencode/test/session/processor.test.ts b/packages/opencode/test/session/processor.test.ts new file mode 100644 index 000000000000..33d01592f8a2 --- /dev/null +++ b/packages/opencode/test/session/processor.test.ts @@ -0,0 +1,184 @@ +import { describe, expect, test, spyOn } from "bun:test" +import type { Agent } from "../../src/agent/agent" +import { Agent as AgentApi } from "../../src/agent/agent" +import { Instance } from "../../src/project/instance" +import type { Provider } from "../../src/provider/provider" +import { ModelID, ProviderID } from "../../src/provider/schema" +import { Session } from "../../src/session" +import { MessageV2 } from "../../src/session/message-v2" +import * as LLMModule from "../../src/session/llm" +import { SessionProcessor } from "../../src/session/processor" +import * as RetryModule from "../../src/session/retry" +import { MessageID } from "../../src/session/schema" +import { Log } from "../../src/util/log" +import { tmpdir } from "../fixture/fixture" + +Log.init({ print: false }) + +type Stream = Awaited> + +function model() { + return { + id: ModelID.make("test-model"), + providerID: ProviderID.make("test"), + } as unknown as Provider.Model +} + +function user(sessionID: Session.Info["id"]) { + return { + id: MessageID.ascending(), + sessionID, + role: "user" as const, + time: { created: Date.now() }, + agent: "build", + model: { + providerID: ProviderID.make("test"), + modelID: ModelID.make("test-model"), + }, + tools: {}, + } as MessageV2.User +} + +function assistant(sessionID: Session.Info["id"], parentID: MessageV2.User["id"], dir: string) { + return { + id: MessageID.ascending(), + sessionID, + parentID, + role: "assistant" as const, + mode: "build", + agent: "build", + path: { + cwd: dir, + root: dir, + }, + cost: 0, + tokens: { + input: 0, + output: 0, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + modelID: ModelID.make("test-model"), + providerID: ProviderID.make("test"), + time: { + created: Date.now(), + }, + } as MessageV2.Assistant +} + +function output(values: unknown[]) { + return { + fullStream: (async function* () { + for (const value of values) yield value + })(), + } as unknown as Stream +} + +async function run( + fn: (input: { + agent: Agent.Info + msg: MessageV2.Assistant + session: Session.Info + usr: MessageV2.User + }) => Promise, +) { + await using tmp = await tmpdir({ git: true }) + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const session = await Session.create({}) + const agent = await AgentApi.get("build") + const usr = (await Session.updateMessage(user(session.id))) as MessageV2.User + const msg = (await Session.updateMessage(assistant(session.id, usr.id, tmp.path))) as MessageV2.Assistant + await fn({ agent, msg, session, usr }) + await Session.remove(session.id) + }, + }) +} + +describe("session.processor", () => { + test("retries unexpected aborts when the session is still active", async () => { + await run(async ({ agent, msg, session, usr }) => { + const ctl = new AbortController() + const stream = spyOn(LLMModule.LLM, "stream") + const sleep = spyOn(RetryModule.SessionRetry, "sleep").mockResolvedValue() + let calls = 0 + + try { + stream.mockImplementation(async (): Promise => { + calls++ + if (calls === 1) throw new DOMException("The operation was aborted.", "AbortError") + return output([{ type: "start" }]) + }) + + const result = await SessionProcessor.create({ + assistantMessage: msg, + sessionID: session.id, + model: model(), + abort: ctl.signal, + }).process({ + user: usr, + agent, + abort: ctl.signal, + sessionID: session.id, + system: [], + messages: [], + tools: {}, + model: model(), + }) + + const stored = await MessageV2.get({ sessionID: session.id, messageID: msg.id }) + if (stored.info.role !== "assistant") throw new Error("expected assistant") + expect(result).toBe("continue") + expect(stream).toHaveBeenCalledTimes(2) + expect(sleep).toHaveBeenCalledTimes(1) + expect(stored.info.error).toBeUndefined() + expect(stored.info.time.completed).toBeDefined() + } finally { + stream.mockRestore() + sleep.mockRestore() + } + }) + }) + + test("stops on aborts after session cancellation", async () => { + await run(async ({ agent, msg, session, usr }) => { + const ctl = new AbortController() + ctl.abort() + const stream = spyOn(LLMModule.LLM, "stream") + const sleep = spyOn(RetryModule.SessionRetry, "sleep").mockResolvedValue() + + try { + stream.mockImplementation(async (): Promise => { + throw new DOMException("The operation was aborted.", "AbortError") + }) + + const result = await SessionProcessor.create({ + assistantMessage: msg, + sessionID: session.id, + model: model(), + abort: ctl.signal, + }).process({ + user: usr, + agent, + abort: ctl.signal, + sessionID: session.id, + system: [], + messages: [], + tools: {}, + model: model(), + }) + + const stored = await MessageV2.get({ sessionID: session.id, messageID: msg.id }) + if (stored.info.role !== "assistant") throw new Error("expected assistant") + expect(result).toBe("stop") + expect(stream).toHaveBeenCalledTimes(1) + expect(sleep).not.toHaveBeenCalled() + expect(stored.info.error?.name).toBe("MessageAbortedError") + } finally { + stream.mockRestore() + sleep.mockRestore() + } + }) + }) +}) diff --git a/packages/opencode/test/session/retry.test.ts b/packages/opencode/test/session/retry.test.ts index 621ad99e9b47..14419f0e4562 100644 --- a/packages/opencode/test/session/retry.test.ts +++ b/packages/opencode/test/session/retry.test.ts @@ -125,6 +125,21 @@ describe("session.retry.retryable", () => { expect(SessionRetry.retryable(error)).toBeUndefined() }) + + test("retries aborted errors when the session is still active", () => { + const ctl = new AbortController() + const error = new MessageV2.AbortedError({ message: "The operation was aborted." }).toObject() + + expect(SessionRetry.retryable(error, ctl.signal)).toBe("The operation was aborted.") + }) + + test("does not retry aborted errors after session cancellation", () => { + const ctl = new AbortController() + ctl.abort() + const error = new MessageV2.AbortedError({ message: "The operation was aborted." }).toObject() + + expect(SessionRetry.retryable(error, ctl.signal)).toBeUndefined() + }) }) describe("session.message-v2.fromError", () => {