diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 844d6f388..de1366ec0 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Preserved queued steering and follow-up messages when a continuation is cancelled before or during pre-dequeue hooks, and propagated the caller's cancellation signal through every continuation model-call loop. + ## [17.2.6] - 2026-08-03 ### Fixed diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index fbce2f756..9cd8b1671 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -1016,7 +1016,7 @@ async function runLoopBody( // Skip when the run is already externally aborted — dequeuing would strand // the messages in a run that is about to die. try { - pendingMessages = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || []; + pendingMessages = signal?.aborted ? [] : (await config.getSteeringMessages?.(signal)) || []; } catch (error) { stream.push({ type: "turn_start" }); emitInputMessages(stream, messagesToEmit); @@ -1075,7 +1075,7 @@ async function runLoopBody( let gateResult: AgentPreModelCallResult; try { if (config.syncContextBeforeModelCall) { - await config.syncContextBeforeModelCall(currentContext); + await config.syncContextBeforeModelCall(currentContext, signal); } if (!directiveResolvedForTurn) { @@ -1421,7 +1421,7 @@ async function runLoopBody( // instantly aborts — message lands in history, agent never responds. The // mid-batch interrupt poll only peeks (hasSteeringMessages), so the queue // still owns every message until this dequeue. - const steering = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || []; + const steering = signal?.aborted ? [] : (await config.getSteeringMessages?.(signal)) || []; if (hasMoreToolCalls) { // Mid-work: fold any non-interrupting asides into the next turn alongside steering. const asides = signal?.aborted ? [] : resolveAsides(await config.getAsideMessages?.()); @@ -1450,9 +1450,9 @@ async function runLoopBody( // Re-poll steering too: a steer can land between the stop-boundary dequeue // above and this yield point (e.g. queued while onBeforeYield ran). Without // this poll it would strand in the queue until the next manual prompt. - const lateSteering = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || []; + const lateSteering = signal?.aborted ? [] : (await config.getSteeringMessages?.(signal)) || []; const asideMessages = signal?.aborted ? [] : resolveAsides(await config.getAsideMessages?.()); - const followUpMessages = signal?.aborted ? [] : (await config.getFollowUpMessages?.()) || []; + const followUpMessages = signal?.aborted ? [] : (await config.getFollowUpMessages?.(signal)) || []; if (lateSteering.length > 0 || asideMessages.length > 0 || followUpMessages.length > 0) { // Set as pending so the inner loop processes them before stopping. pendingMessages = [...lateSteering, ...asideMessages, ...followUpMessages]; diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 8a17735c0..c2e779a7e 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -426,6 +426,8 @@ export class Agent { #asideMessageProvider?: () => AsideMessage[] | Promise; #telemetry?: AgentLoopConfig["telemetry"]; #appendOnlyContext?: AppendOnlyContextManager; + #beforeQueuedMessageDequeueHooks = new Set<(signal?: AbortSignal) => Promise | void>(); + #beforeModelCallHooks = new Set<(signal?: AbortSignal) => Promise | void>(); /** Buffered Cursor tool results with text length at time of call (for correct ordering) */ #cursorToolResultBuffer: CursorToolResultEntry[] = []; @@ -784,6 +786,40 @@ export class Agent { return () => this.#listeners.delete(fn); } + /** Register an independently removable hook that runs before queued messages are consumed. */ + addBeforeQueuedMessageDequeueHook(hook: (signal?: AbortSignal) => Promise | void): () => void { + const registration = (signal?: AbortSignal) => hook(signal); + this.#beforeQueuedMessageDequeueHooks.add(registration); + return () => this.#beforeQueuedMessageDequeueHooks.delete(registration); + } + + /** Register an independently removable hook that runs immediately before each model call. */ + addBeforeModelCallHook(hook: (signal?: AbortSignal) => Promise | void): () => void { + const registration = (signal?: AbortSignal) => hook(signal); + this.#beforeModelCallHooks.add(registration); + return () => this.#beforeModelCallHooks.delete(registration); + } + + async #runBeforeModelCallHooks(signal?: AbortSignal): Promise { + for (const hook of this.#beforeModelCallHooks) await hook(signal); + } + + async #runBeforeQueuedMessageDequeueHooks(signal?: AbortSignal): Promise { + for (const hook of this.#beforeQueuedMessageDequeueHooks) await hook(signal); + } + + async #dequeueSteeringMessagesAfterHooks(signal?: AbortSignal): Promise { + if (signal?.aborted || this.#steeringQueue.length === 0) return []; + await this.#runBeforeQueuedMessageDequeueHooks(signal); + return signal?.aborted ? [] : this.#dequeueSteeringMessages(); + } + + async #dequeueFollowUpMessagesAfterHooks(signal?: AbortSignal): Promise { + if (signal?.aborted || this.#followUpQueue.length === 0) return []; + await this.#runBeforeQueuedMessageDequeueHooks(signal); + return signal?.aborted ? [] : this.#dequeueFollowUpMessages(); + } + setProviderResponseInterceptor(fn: SimpleStreamOptions["onResponse"] | undefined): void { this.#onResponse = fn; } @@ -1137,48 +1173,90 @@ export class Agent { /** * Continue from current context (used for retries and resuming queued messages). */ - async continue() { + #continuationDequeueSignal(signal?: AbortSignal): AbortSignal | undefined { + const signals: AbortSignal[] = []; + if (this.#abortController) signals.push(this.#abortController.signal); + if (signal) signals.push(signal); + if (this.#deadline !== undefined) { + const delay = this.#deadline - Date.now(); + if (delay <= 0) { + const controller = new AbortController(); + controller.abort(new DOMException("Deadline exceeded", "TimeoutError")); + signals.push(controller.signal); + } else { + signals.push(AbortSignal.timeout(delay)); + } + } + if (signals.length === 0) return undefined; + return signals.length === 1 ? signals[0] : AbortSignal.any(signals); + } + + async continue(signal?: AbortSignal) { if (this.#state.isStreaming) { throw new AgentBusyError(); } - const messages = this.#state.messages; - if (messages.length === 0) { - // An empty transcript has nothing to resume, but a queued steer/follow-up - // must still be delivered as the opening turn — mirroring the assistant-tail - // branch below. Throwing here leaves the message undeliverable, and idle-drain - // callers (AgentSession#scheduleQueuedMessageDrain) re-arm continue() on every - // microtask because hasQueuedMessages() never clears, spinning an unbounded - // allocation loop until OOM (issue #6344). - const queuedSteering = this.#dequeueSteeringMessages(); - if (queuedSteering.length > 0) { - await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true }); - return; + const { promise, resolve } = Promise.withResolvers(); + this.#runningPrompt = promise; + this.#resolveRunningPrompt = resolve; + const continuationAbortController = new AbortController(); + this.#abortController = continuationAbortController; + this.#state.isStreaming = true; + this.#state.streamMessage = null; + this.#state.error = undefined; + + try { + const dequeueSignal = this.#continuationDequeueSignal(signal); + const messages = this.#state.messages; + if (messages.length === 0) { + // An empty transcript has nothing to resume, but a queued steer/follow-up + // must still be delivered as the opening turn — mirroring the assistant-tail + // branch below. Throwing here leaves the message undeliverable, and idle-drain + // callers (AgentSession#scheduleQueuedMessageDrain) re-arm continue() on every + // microtask because hasQueuedMessages() never clears, spinning an unbounded + // allocation loop until OOM (issue #6344). + const queuedSteering = await this.#dequeueSteeringMessagesAfterHooks(dequeueSignal); + if (queuedSteering.length > 0) { + await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true }, signal, true); + return; + } + const queuedFollowUp = await this.#dequeueFollowUpMessagesAfterHooks(dequeueSignal); + if (queuedFollowUp.length > 0) { + await this.#runLoop(queuedFollowUp, undefined, signal, true); + return; + } + throw new Error("No messages to continue from"); } - const queuedFollowUp = this.#dequeueFollowUpMessages(); - if (queuedFollowUp.length > 0) { - await this.#runLoop(queuedFollowUp); - return; + if (messages[messages.length - 1].role === "assistant") { + const queuedSteering = await this.#dequeueSteeringMessagesAfterHooks(dequeueSignal); + if (queuedSteering.length > 0) { + await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true }, signal, true); + return; + } + + const queuedFollowUp = await this.#dequeueFollowUpMessagesAfterHooks(dequeueSignal); + if (queuedFollowUp.length > 0) { + await this.#runLoop(queuedFollowUp, undefined, signal, true); + return; + } + + throw new Error("Cannot continue from message role: assistant"); + } + + await this.#runLoop(undefined, undefined, signal, true); + } finally { + resolve(); + if (this.#abortController === continuationAbortController) { + this.#state.isStreaming = false; + this.#state.streamMessage = null; + this.#state.pendingToolCalls.clear(); + this.#abortController = undefined; + if (this.#runningPrompt === promise) { + this.#runningPrompt = undefined; + this.#resolveRunningPrompt = undefined; + } } - throw new Error("No messages to continue from"); } - if (messages[messages.length - 1].role === "assistant") { - const queuedSteering = this.#dequeueSteeringMessages(); - if (queuedSteering.length > 0) { - await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true }); - return; - } - - const queuedFollowUp = this.#dequeueFollowUpMessages(); - if (queuedFollowUp.length > 0) { - await this.#runLoop(queuedFollowUp); - return; - } - - throw new Error("Cannot continue from message role: assistant"); - } - - await this.#runLoop(undefined); } /** @@ -1186,17 +1264,29 @@ export class Agent { * If messages are provided, starts a new conversation turn with those messages. * Otherwise, continues from existing context. */ - async #runLoop(messages?: AgentMessage[], options?: AgentPromptOptions & { skipInitialSteeringPoll?: boolean }) { + async #runLoop( + messages?: AgentMessage[], + options?: AgentPromptOptions & { skipInitialSteeringPoll?: boolean }, + continuationSignal?: AbortSignal, + runStateClaimed = false, + ) { const model = this.#state.model; if (!model) throw new Error("No model configured"); let skipInitialSteeringPoll = options?.skipInitialSteeringPoll === true; using _ = new EventLoopKeepalive(); - const { promise, resolve } = Promise.withResolvers(); - this.#runningPrompt = promise; - this.#resolveRunningPrompt = resolve; - - this.#abortController = new AbortController(); + if (!runStateClaimed) { + const { promise, resolve } = Promise.withResolvers(); + this.#runningPrompt = promise; + this.#resolveRunningPrompt = resolve; + this.#abortController = new AbortController(); + } + const resolveRun = this.#resolveRunningPrompt; + const loopAbortController = this.#abortController; + if (!loopAbortController) throw new Error("Agent run state was not initialized"); + const loopSignal = continuationSignal + ? AbortSignal.any([loopAbortController.signal, continuationSignal]) + : loopAbortController.signal; this.#state.isStreaming = true; this.#state.streamMessage = null; this.#state.error = undefined; @@ -1315,7 +1405,8 @@ export class Agent { onSseEvent: this.#onSseEvent, getApiKey: this.getApiKey, getToolContext: this.#getToolContext, - syncContextBeforeModelCall: async context => { + syncContextBeforeModelCall: async (context, signal) => { + await this.#runBeforeModelCallHooks(signal); if (this.#listeners.size > 0) { await Bun.sleep(0); } @@ -1362,12 +1453,12 @@ export class Agent { getReasoning: () => this.#state.thinkingLevel, getDisableReasoning: () => this.#state.disableReasoning, getServiceTier: this.#serviceTierResolver, - getSteeringMessages: async () => { + getSteeringMessages: async signal => { if (skipInitialSteeringPoll) { skipInitialSteeringPoll = false; return []; } - return this.#dequeueSteeringMessages(); + return this.#dequeueSteeringMessagesAfterHooks(signal); }, hasSteeringMessages: () => { if (this.#steeringQueue.length === 0) { @@ -1392,7 +1483,7 @@ export class Agent { }, waitForSteeringMessages: signal => this.#waitForSteeringMessages(signal), hasIrcInterrupts: this.hasIrcInterrupts, - getFollowUpMessages: async () => this.#dequeueFollowUpMessages(), + getFollowUpMessages: signal => this.#dequeueFollowUpMessagesAfterHooks(signal), getAsideMessages: async () => (await this.#asideMessageProvider?.()) ?? [], onBeforeYield: () => this.#onBeforeYield?.(), telemetry: this.#telemetry, @@ -1404,8 +1495,8 @@ export class Agent { try { const stream = messages - ? agentLoop(messages, context, config, this.#abortController.signal, this.streamFn) - : agentLoopContinue(context, config, this.#abortController.signal, this.streamFn); + ? agentLoop(messages, context, config, loopSignal, this.streamFn) + : agentLoopContinue(context, config, loopSignal, this.streamFn); for await (const event of stream) { if (event.type === "turn_start") turnOpen = true; @@ -1472,15 +1563,15 @@ export class Agent { if (!onlyEmpty) { this.appendMessage(partial); } else { - if (this.#abortController?.signal.aborted) { + if (loopSignal.aborted) { throw new Error("Request was aborted"); } } } } catch (err) { - const stoppedForAbort = this.#abortController?.signal.aborted === true; + const stoppedForAbort = loopSignal.aborted; const errorMessage = stoppedForAbort - ? abortReasonText(this.#abortController?.signal) + ? abortReasonText(loopSignal) : err instanceof Error ? err.message : String(err); @@ -1582,13 +1673,15 @@ export class Agent { this.#emit({ type: "agent_end", messages: [errorMsg] }); } } finally { - this.#state.isStreaming = false; - this.#state.streamMessage = null; - this.#state.pendingToolCalls.clear(); - this.#abortController = undefined; - this.#resolveRunningPrompt?.(); - this.#runningPrompt = undefined; - this.#resolveRunningPrompt = undefined; + resolveRun?.(); + if (this.#abortController === loopAbortController) { + this.#state.isStreaming = false; + this.#state.streamMessage = null; + this.#state.pendingToolCalls.clear(); + this.#abortController = undefined; + this.#runningPrompt = undefined; + this.#resolveRunningPrompt = undefined; + } } } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index de90a681e..198c9a4a3 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -240,7 +240,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions { * mid-batch interrupt poll uses {@link hasSteeringMessages} instead and * never consumes the queue. */ - getSteeringMessages?: () => Promise; + getSteeringMessages?: (signal?: AbortSignal) => Promise; /** * Peeks whether steering messages are queued, without consuming them. @@ -285,7 +285,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions { * If messages are returned, they're added to the context and the agent * continues with another turn. */ - getFollowUpMessages?: () => Promise; + getFollowUpMessages?: (signal?: AbortSignal) => Promise; /** * Returns non-interrupting "aside" messages to inject at a step boundary. * @@ -319,7 +319,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions { * Mutate the agent context here; use `beforeModelCall` to inspect the * provider-bound context. */ - syncContextBeforeModelCall?: (context: AgentContext) => void | Promise; + syncContextBeforeModelCall?: (context: AgentContext, signal?: AbortSignal) => void | Promise; /** * Asked after the complete provider context has been built, including diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index a887758c8..1f4604efe 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "bun:test"; -import { Agent, type AgentEvent, type AgentTool, ThinkingLevel } from "@oh-my-pi/pi-agent-core"; +import { Agent, AgentBusyError, type AgentEvent, type AgentTool, ThinkingLevel } from "@oh-my-pi/pi-agent-core"; import { type SimpleStreamOptions, type ToolResultMessage, z } from "@oh-my-pi/pi-ai"; import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; import { kCursorExecResolved } from "@oh-my-pi/pi-ai/utils/block-symbols"; @@ -240,6 +240,213 @@ describe("Agent", () => { } }); + it("removes duplicate queued-message hooks independently", async () => { + const mock = createMockModel({ responses: [{ content: ["first"] }, { content: ["second"] }] }); + const agent = new Agent({ streamFn: mock.stream }); + agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]); + let calls = 0; + const signals: Array = []; + const hook = (signal?: AbortSignal) => { + calls++; + signals.push(signal); + }; + const removeFirst = agent.addBeforeQueuedMessageDequeueHook(hook); + const removeSecond = agent.addBeforeQueuedMessageDequeueHook(hook); + + const controller = new AbortController(); + removeFirst(); + agent.followUp({ role: "user", content: "first turn", timestamp: Date.now() }); + await agent.continue(controller.signal); + expect(calls).toBe(1); + expect(signals).toEqual([controller.signal]); + + removeSecond(); + agent.followUp({ role: "user", content: "second turn", timestamp: Date.now() }); + await agent.continue(); + expect(calls).toBe(1); + }); + + it("continue() leaves queued messages owned when its signal is already aborted", async () => { + const agent = new Agent(); + agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]); + agent.followUp({ role: "user", content: "stay queued", timestamp: Date.now() }); + const controller = new AbortController(); + controller.abort(); + + await expect(agent.continue(controller.signal)).rejects.toThrow("Cannot continue from message role: assistant"); + expect(agent.peekFollowUpQueue()).toHaveLength(1); + }); + it("keeps follow-up ownership when the deadline expires during a dequeue hook", async () => { + const mock = createMockModel({ responses: [{ content: ["done"] }] }); + const agent = new Agent({ streamFn: mock.stream, deadline: Date.now() + 25 }); + let hookSignal: AbortSignal | undefined; + agent.addBeforeQueuedMessageDequeueHook(async signal => { + if (!signal) throw new Error("Expected the active loop signal"); + hookSignal = signal; + if (signal.aborted) return; + const { promise, resolve } = Promise.withResolvers(); + signal.addEventListener("abort", () => resolve(), { once: true }); + await promise; + }); + agent.followUp({ role: "user", content: "stay queued after deadline", timestamp: Date.now() }); + + await agent.prompt("start"); + + expect(hookSignal?.aborted).toBe(true); + expect(agent.peekFollowUpQueue()).toHaveLength(1); + }); + it("keeps queued work when continue() reaches its deadline inside a dequeue hook", async () => { + const agent = new Agent({ deadline: Date.now() + 25 }); + agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]); + agent.addBeforeQueuedMessageDequeueHook(async signal => { + if (!signal) throw new Error("Expected the deadline-aware dequeue signal"); + if (signal.aborted) return; + const { promise, resolve } = Promise.withResolvers(); + signal.addEventListener("abort", () => resolve(), { once: true }); + await promise; + }); + agent.followUp({ role: "user", content: "stay queued before run loop", timestamp: Date.now() }); + + await expect(agent.continue()).rejects.toThrow("Cannot continue from message role: assistant"); + + expect(agent.peekFollowUpQueue()).toHaveLength(1); + }); + + it("claims an abortable busy state while continue() awaits dequeue hooks", async () => { + const agent = new Agent(); + agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]); + agent.followUp({ role: "user", content: "stay queued", timestamp: Date.now() }); + const hookStarted = Promise.withResolvers(); + agent.addBeforeQueuedMessageDequeueHook(async signal => { + if (!signal) throw new Error("Expected continuation dequeue signal"); + hookStarted.resolve(); + if (signal.aborted) return; + const { promise, resolve } = Promise.withResolvers(); + signal.addEventListener("abort", () => resolve(), { once: true }); + await promise; + }); + + const continuing = agent.continue(); + await hookStarted.promise; + let idleResolved = false; + const idle = agent.waitForIdle().then(() => { + idleResolved = true; + }); + await Promise.resolve(); + + expect(agent.state.isStreaming).toBe(true); + expect(idleResolved).toBe(false); + await expect(agent.prompt("must not overlap")).rejects.toBeInstanceOf(AgentBusyError); + + agent.abort("cancel dequeue"); + await expect(continuing).rejects.toThrow("Cannot continue from message role: assistant"); + await idle; + expect(idleResolved).toBe(true); + expect(agent.state.isStreaming).toBe(false); + expect(agent.peekFollowUpQueue()).toHaveLength(1); + }); + + it("does not clear a successor prompt after continue() releases idle waiters", async () => { + const firstStarted = Promise.withResolvers(); + const releaseFirst = Promise.withResolvers(); + const secondStarted = Promise.withResolvers(); + const releaseSecond = Promise.withResolvers(); + const mock = createMockModel({ + responses: [ + async () => { + firstStarted.resolve(); + await releaseFirst.promise; + return { content: ["continued"] }; + }, + async () => { + secondStarted.resolve(); + await releaseSecond.promise; + return { content: ["successor"] }; + }, + ], + }); + const agent = new Agent({ streamFn: mock.stream }); + agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]); + agent.followUp({ role: "user", content: "continue", timestamp: Date.now() }); + + const continuing = agent.continue(); + await firstStarted.promise; + const successor = agent.waitForIdle().then(() => agent.prompt("next prompt")); + releaseFirst.resolve(); + await secondStarted.promise; + await continuing; + + expect(agent.state.isStreaming).toBe(true); + releaseSecond.resolve(); + await successor; + expect(agent.state.isStreaming).toBe(false); + }); + + it("resolves a predecessor idle waiter when agent_end starts a successor", async () => { + const secondStarted = Promise.withResolvers(); + const releaseSecond = Promise.withResolvers(); + const mock = createMockModel({ + responses: [ + { content: ["first"] }, + async () => { + secondStarted.resolve(); + await releaseSecond.promise; + return { content: ["second"] }; + }, + ], + }); + const agent = new Agent({ streamFn: mock.stream }); + let successor: Promise | undefined; + agent.subscribe(event => { + if (event.type === "agent_end" && !successor) { + successor = agent.prompt("successor"); + } + }); + + const predecessor = agent.prompt("predecessor"); + let predecessorIdleResolved = false; + void agent.waitForIdle().then(() => { + predecessorIdleResolved = true; + }); + await secondStarted.promise; + await predecessor; + expect(agent.state.isStreaming).toBe(true); + + releaseSecond.resolve(); + await successor; + await Promise.resolve(); + expect(predecessorIdleResolved).toBe(true); + expect(agent.state.isStreaming).toBe(false); + }); + + it("classifies an in-flight continuation cancellation as aborted", async () => { + const providerStarted = Promise.withResolvers(); + const agent = new Agent({ + streamFn: (_model, _context, options) => { + const signal = options?.signal; + if (!signal) throw new Error("Expected provider abort signal"); + providerStarted.resolve(signal); + const stream = new AssistantMessageEventStream(); + signal.addEventListener("abort", () => stream.fail(new Error("provider aborted")), { once: true }); + return stream; + }, + }); + agent.replaceMessages([createAssistantMessage([{ type: "text", text: "ready" }])]); + agent.followUp({ role: "user", content: "cancel this continuation", timestamp: Date.now() }); + const controller = new AbortController(); + + const running = agent.continue(controller.signal); + await providerStarted.promise; + controller.abort("caller cancelled"); + await running; + + const finalMessage = agent.state.messages.at(-1); + expect(finalMessage?.role).toBe("assistant"); + if (finalMessage?.role !== "assistant") throw new Error("Expected aborted assistant message"); + expect(finalMessage.stopReason).toBe("aborted"); + expect(finalMessage.errorMessage).toBe("caller cancelled"); + }); + it("continue() should process queued follow-up messages after an assistant turn", async () => { const mock = createMockModel({ responses: [{ content: ["Processed"] }] }); const agent = new Agent({ streamFn: mock.stream }); @@ -276,6 +483,12 @@ describe("Agent", () => { responses: [{ content: ["Processed 1"] }, { content: ["Processed 2"] }], }); const agent = new Agent({ streamFn: mock.stream }); + let dequeueHooks = 0; + const dequeueSignals: Array = []; + agent.addBeforeQueuedMessageDequeueHook(signal => { + dequeueHooks++; + dequeueSignals.push(signal); + }); agent.replaceMessages([ { @@ -297,11 +510,16 @@ describe("Agent", () => { timestamp: Date.now() + 1, }); - await expect(agent.continue()).resolves.toBeUndefined(); + const controller = new AbortController(); + await expect(agent.continue(controller.signal)).resolves.toBeUndefined(); const recentMessages = agent.state.messages.slice(-4); expect(recentMessages.map(m => m.role)).toEqual(["user", "assistant", "user", "assistant"]); expect(mock.calls.length).toBe(2); + expect(dequeueHooks).toBe(2); + expect(dequeueSignals).toHaveLength(2); + controller.abort(); + expect(dequeueSignals.every(signal => signal?.aborted === true)).toBe(true); }); it("delivers a steer that lands at the yield boundary instead of stranding it", async () => { @@ -856,6 +1074,10 @@ describe("Agent", () => { }, streamFn: mock.stream, }); + let beforeModelCalls = 0; + agent.addBeforeModelCallHook(() => { + beforeModelCalls++; + }); const unsubscribe = agent.subscribe(event => { if (event.type === "message_end" && event.message.role === "toolResult") { @@ -875,6 +1097,7 @@ describe("Agent", () => { { systemPrompt: "prompt-one", toolNames: ["alpha"] }, { systemPrompt: "prompt-two", toolNames: ["alpha", "beta"] }, ]); + expect(beforeModelCalls).toBe(2); }); it("prompt() drops stale forced toolChoice after same-turn tool refresh", async () => { diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index bd71a5086..3e22f7630 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -17,6 +17,7 @@ - Fixed OpenAI-Codex (ChatGPT OAuth) requests failing with an `Unsupported service_tier: auto` error on default or legacy sessions by omitting the implicit `auto` service tier on the wire. - Fixed an issue where Cursor `kimi-k3` sessions would break permanently when a same-model assistant turn was persisted without thinking blocks, replacing hard errors with graceful warnings. +- Made model-scoped usage health ignore Codex accounts that cannot use the requested plan-gated model while retaining conservative unknown-state handling and independent usage-window resets. ## [17.2.6] - 2026-08-03 diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 6f8605fcb..762e737af 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -3846,13 +3846,15 @@ export class AuthStorage { : undefined; const rankingContext: CredentialRankingContext = { modelId: options.modelId }; + const planRequirement = resolveOpenAICodexPlanRequirement(provider, options.modelId); + const planEligibilityByCredential = new Map(); const blockScope = strategy.blockScope?.(rankingContext); const blockScopes = strategy.blockScopes?.(rankingContext) ?? (blockScope ? [blockScope] : []); const reserveFraction = Number.isFinite(options.reserveFraction) ? Math.max(0, Math.min(1, options.reserveFraction)) : 0; const nowMs = Date.now(); - const accounts = await Promise.all( + let accounts = await Promise.all( pool.map(async ({ entry, index }): Promise => { const credentialType = entry.credential.type; const providerKey = this.#getProviderTypeKey(provider, credentialType); @@ -3880,6 +3882,9 @@ export class AuthStorage { if (options.signal?.aborted) throw error; report = null; } + if (planRequirement !== "none") { + planEligibilityByCredential.set(entry.id, getOpenAICodexPlanEligibility(report, planRequirement)); + } if (provider === "openai-codex") { blockedUntil = this.#getCredentialBlockedUntil(provider, providerKey, index, blockScopes); @@ -3932,6 +3937,9 @@ export class AuthStorage { }; }), ); + if (planRequirement !== "none") { + accounts = accounts.filter(account => planEligibilityByCredential.get(account.credentialId) !== false); + } if (selectedCredentialId !== undefined) { const selectedAccount = accounts.find(account => account.credentialId === selectedCredentialId); if (selectedAccount) selectedAccount.selected = true; diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 365e7da0b..fe769f1c4 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -2027,6 +2027,90 @@ describe("AuthStorage codex oauth ranking", () => { expect(apiKey).toBe("api-acct-pro"); }); + test("ignores plan-ineligible headroom when reporting Spark model health", async () => { + if (!authStorage) throw new Error("test setup failed"); + + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-free", "free@example.com") }, + { type: "oauth", ...createCredential("acct-pro", "pro@example.com") }, + ]); + usageByAccount.set( + "acct-free", + addSparkUsage( + createCodexUsageReport({ + accountId: "acct-free", + primary: { usedFraction: 0.05, resetInMs: 30 * 60 * 1000 }, + secondary: { usedFraction: 0.05, resetInMs: 6 * 24 * 60 * 60 * 1000 }, + metadata: { planType: "free", email: "free@example.com" }, + }), + 0.05, + 0.05, + ), + ); + usageByAccount.set( + "acct-pro", + addSparkUsage( + createCodexUsageReport({ + accountId: "acct-pro", + primary: { usedFraction: 1, resetInMs: 2 * HOUR_MS }, + secondary: { usedFraction: 1, resetInMs: 6 * 24 * 60 * 60 * 1000 }, + metadata: { planType: "pro", email: "pro@example.com", limitReached: true }, + }), + 1, + 1, + ), + ); + + const health = await authStorage.getModelUsageHealth("openai-codex", { + modelId: "gpt-5.3-codex-spark", + reserveFraction: 0.1, + }); + + expect(health.state).toBe("depleted"); + expect(health.accounts).toHaveLength(1); + expect(health.accounts[0]?.state).toBe("depleted"); + }); + + test("reports an all-plan-ineligible Codex pool as depleted", async () => { + if (!authStorage) throw new Error("test setup failed"); + + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-free", "free@example.com") }, + { type: "oauth", ...createCredential("acct-plus", "plus@example.com") }, + ]); + usageByAccount.set( + "acct-free", + createCodexUsageReport({ + accountId: "acct-free", + primary: { usedFraction: 0.05, resetInMs: 30 * 60 * 1000 }, + secondary: { usedFraction: 0.05, resetInMs: 6 * 24 * 60 * 60 * 1000 }, + metadata: { planType: "free", email: "free@example.com" }, + }), + ); + usageByAccount.set( + "acct-plus", + createCodexUsageReport({ + accountId: "acct-plus", + primary: { usedFraction: 0.05, resetInMs: 30 * 60 * 1000 }, + secondary: { usedFraction: 0.05, resetInMs: 6 * 24 * 60 * 60 * 1000 }, + metadata: { planType: "plus", email: "plus@example.com" }, + }), + ); + + const paidHealth = await authStorage.getModelUsageHealth("openai-codex", { + modelId: "gpt-5.6-sol", + reserveFraction: 0.1, + }); + const proHealth = await authStorage.getModelUsageHealth("openai-codex", { + modelId: "gpt-5.3-codex-spark", + reserveFraction: 0.1, + }); + + expect(paidHealth.state).toBe("healthy"); + expect(paidHealth.accounts).toHaveLength(1); + expect(proHealth).toEqual({ state: "depleted", accounts: [] }); + }); + test("routes codex spark to a single Plus account when no Pro is connected", async () => { if (!authStorage) throw new Error("test setup failed"); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 61e7000ba..752f310c0 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -73,6 +73,7 @@ - Fixed the built-in `tail` printing `tail: Broken pipe` and failing when a downstream pipeline reader exited early (e.g. `tail -c N file.jsonl | jq …` with jq aborting on a parse error); it now exits silently with 141 (128+SIGPIPE) like a real tail, in every output path including `--follow`. - Fixed the in-process ps shell builtin rejecting common procps/BSD format specifiers (`ps -o tpgid,...` failed with `unknown output format specifier`); added `tpgid`, `pri`, `flags`, real/effective user and group columns, `wchan`, fault counters, `sz`, and the STAT `+` foreground flag. - Fixed Herdr rejecting the macOS development launcher because its foreground process was reported as `bun` instead of `omp`. +- Completed usage-aware model fallback across startup, queued turns, same-turn tool continuations, ACP/TUI confirmation cancellation, eligible account reselection, cooldown restoration, and isolated subagent settings so low-usage handoffs remain lossless and cannot consume cancelled queued work. ## [17.2.6] - 2026-08-03 diff --git a/packages/coding-agent/src/modes/acp/acp-agent.ts b/packages/coding-agent/src/modes/acp/acp-agent.ts index 8ee0db5ee..d3d480564 100644 --- a/packages/coding-agent/src/modes/acp/acp-agent.ts +++ b/packages/coding-agent/src/modes/acp/acp-agent.ts @@ -2312,7 +2312,7 @@ export class AcpAgent implements Agent { this.#clientCapabilities, ); if (this.#clientCapabilities?.elicitation?.form != null) { - record.session.setUsageFallbackConfirmer(confirmation => { + record.session.setUsageFallbackConfirmer((confirmation, signal) => { const reserve = confirmation.remainingPercent === undefined ? "inside the configured reserve margin" @@ -2320,6 +2320,7 @@ export class AcpAgent implements Agent { return uiContext.confirm( "Coding-plan reserve reached", `${confirmation.from} has ${reserve}. Switch to ${confirmation.to}? Choose No to keep using the current plan.`, + { signal }, ); }); } diff --git a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts index 750bdd193..942b03da8 100644 --- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts +++ b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts @@ -128,7 +128,7 @@ export class ExtensionUiController { }; this.ctx.setToolUIContext(uiContext, true); this.#toolUIContext = uiContext; - this.ctx.session.setUsageFallbackConfirmer?.(confirmation => { + this.ctx.session.setUsageFallbackConfirmer?.((confirmation, signal) => { const reserve = confirmation.remainingPercent === undefined ? "inside the configured reserve margin" @@ -136,6 +136,7 @@ export class ExtensionUiController { return this.showHookConfirm( "Coding-plan reserve reached", `${confirmation.from} has ${reserve}. Switch to ${confirmation.to}? Choose No to keep using the current plan.`, + { signal }, ); }); diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 4ee9a24c5..3f3ae2bcc 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -2283,9 +2283,10 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro } } const usageReservePolicy = settings.get("retry.usageReservePolicy"); + const modelFallbackEnabled = settings.get("retry.modelFallback"); if ( - (hasUsageFallbackCandidate || usageReservePolicy === "fail-closed") && - settings.get("retry.modelFallback") && + ((modelFallbackEnabled && (hasUsageFallbackCandidate || usageFallbackTriggered)) || + usageReservePolicy === "fail-closed") && settings.get("retry.usageAwareFallback") ) { let usageHealth: ModelUsageHealth | undefined; @@ -2308,8 +2309,10 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro `Usage depleted for ${primary.model.provider}/${primary.model.id}; reserve policy is fail-closed.`, ); } - usageFallbackTriggered = true; - continue; + if (modelFallbackEnabled) { + usageFallbackTriggered = true; + continue; + } } if (usageHealth?.state === "reserve") { if (usageReservePolicy === "fail-closed") { @@ -2317,7 +2320,10 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro `Usage reserve reached for ${primary.model.provider}/${primary.model.id}; reserve policy is fail-closed.`, ); } - if (usageReservePolicy === "auto" || (!options.hasUI && !options.deferUsageReserveConfirmation)) { + if ( + modelFallbackEnabled && + (usageReservePolicy === "auto" || (!options.hasUI && !options.deferUsageReserveConfirmation)) + ) { usageFallbackTriggered = true; continue; } diff --git a/packages/coding-agent/src/session/agent-session-types.ts b/packages/coding-agent/src/session/agent-session-types.ts index 51a2e6fc8..6d7241b3a 100644 --- a/packages/coding-agent/src/session/agent-session-types.ts +++ b/packages/coding-agent/src/session/agent-session-types.ts @@ -92,6 +92,14 @@ export interface UsageFallbackConfirmation { remainingPercent: number | undefined; } +/** + * Confirms whether a reserve-triggered model fallback may proceed. + * + * Interactive callers use the confirmation details to present the pending + * route change; aborting `signal` cancels that pending confirmation. + */ +export type UsageFallbackConfirmer = (confirmation: UsageFallbackConfirmation, signal: AbortSignal) => Promise; + /** Identifies a retry fallback chain already entered during startup model resolution. */ export interface InitialRetryFallbackState { /** Role whose configured primary was unavailable. */ diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index e247468e2..f9d24773e 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -58,7 +58,6 @@ import type { ImageContent, Message, Model, - ModelUsageHealth, OAuthAccountIdentity, ProviderSessionState, ResetCreditAccountStatus, @@ -101,7 +100,7 @@ import { type AdvisorConfig, type AdvisorRuntimeStatus, loadAdvisorTranscriptCos import { type AsyncJob, AsyncJobManager } from "../async"; import { shouldEnableAppendOnlyContext } from "../config/append-only-context-mode"; import type { ModelRegistry } from "../config/model-registry"; -import { type ResolvedModelRoleValue, resolveModelOverride } from "../config/model-resolver"; +import type { ResolvedModelRoleValue } from "../config/model-resolver"; import { expandPromptTemplate, type PromptTemplate } from "../config/prompt-templates"; import { buildServiceTierByFamily } from "../config/service-tier"; import type { Settings, SkillsSettings } from "../config/settings"; @@ -230,7 +229,7 @@ import type { SessionHandoffOptions, SessionOAuthAccountList, SessionStats, - UsageFallbackConfirmation, + UsageFallbackConfirmer, } from "./agent-session-types"; import { ASYNC_INLINE_RESULT_MAX_CHARS, @@ -312,7 +311,6 @@ import { queueChipText, toRestoredQueuedMessage, } from "./queued-messages"; -import { formatRetryFallbackSelector, type RetryFallbackSelector } from "./retry-fallback-chains"; import { type AdvisorStats, SessionAdvisors, type SessionAdvisorsHost } from "./session-advisors"; import type { BuildSessionContextOptions, SessionContext } from "./session-context"; import { getRestorableSessionModels } from "./session-context"; @@ -559,9 +557,13 @@ export class AgentSession { // Model registry for API key resolution #modelRegistry: ModelRegistry; - #usageFallbackConfirmer: ((confirmation: UsageFallbackConfirmation) => Promise) | undefined; - #usageReserveApprovedSelector: string | undefined; + #usageFallbackConfirmer: UsageFallbackConfirmer | undefined; #usagePreflightAbortControllers = new Set(); + #queuedMessageDrainBlocked = false; + #usagePreflightReadyForNextModelCall = false; + #usagePreflightReadyModel: Model | undefined; + #detachUsageBeforeQueueDequeue: (() => void) | undefined; + #detachUsageBeforeModelCall: (() => void) | undefined; #transformContext: (messages: AgentMessage[], signal?: AbortSignal) => AgentMessage[] | Promise; #onPayload: SimpleStreamOptions["onPayload"] | undefined; @@ -769,8 +771,10 @@ export class AgentSession { !this.#canAutoContinueForFollowUp() ? [...this.agent.peekFollowUpQueue()] : []; + const parkedQueueDrainBlocked = parkedFollowUps.length > 0 && this.#queuedMessageDrainBlocked; if (parkedFollowUps.length > 0) { this.agent.replaceQueues([...this.agent.peekSteeringQueue()], []); + if (parkedQueueDrainBlocked) this.#queuedMessageDrainBlocked = false; } let finishObservation: ((error?: unknown) => void | Promise) | undefined; try { @@ -806,6 +810,7 @@ export class AgentSession { [...this.agent.peekSteeringQueue()], [...parkedFollowUps, ...this.agent.peekFollowUpQueue()], ); + this.#queuedMessageDrainBlocked ||= parkedQueueDrainBlocked; } this.#endInFlight(async () => { try { @@ -831,6 +836,7 @@ export class AgentSession { steering.filter(m => !isAdvisorCard(m)), followUp.filter(m => !isAdvisorCard(m)), ); + this.#reconcileQueuedMessageDrain(); return cards; } @@ -1052,6 +1058,31 @@ export class AgentSession { withBashBranchTransition: operation => this.#bash.withBranchTransition(operation), }; this.#recovery = new TurnRecovery(recoveryHost, { initialRetryFallback: config.initialRetryFallback }); + this.#detachUsageBeforeQueueDequeue = this.agent.addBeforeQueuedMessageDequeueHook(async signal => { + if ( + !this.settings.get("retry.usageAwareFallback") || + (this.#usagePreflightReadyForNextModelCall && this.#usagePreflightReadyModel === this.model) + ) { + return; + } + if (!(await this.#runQueuedUsageAwarePreflight(signal))) { + signal?.throwIfAborted(); + throw new DOMException("Usage preflight cancelled", "AbortError"); + } + }); + this.#detachUsageBeforeModelCall = this.agent.addBeforeModelCallHook(async signal => { + if (!this.settings.get("retry.usageAwareFallback")) return; + if (this.#usagePreflightReadyForNextModelCall) { + const checkedModel = this.#usagePreflightReadyModel; + this.#usagePreflightReadyForNextModelCall = false; + this.#usagePreflightReadyModel = undefined; + if (checkedModel === this.model) return; + } + if (!(await this.#runUsageAwarePreflight(signal))) { + signal?.throwIfAborted(); + throw new DOMException("Usage preflight cancelled", "AbortError"); + } + }); const statsHost: SessionStatsTrackerHost = { session: this, agent: this.agent, @@ -2969,19 +3000,17 @@ export class AgentSession { this.#beginInFlight(); try { await this.#recovery.maybeRestoreRetryFallbackPrimary(); - if ( - this.settings.get("retry.modelFallback") && - this.settings.get("retry.usageAwareFallback") && - !(await this.#runUsageAwarePreflight()) - ) { - this.#skipAgentContinue("session-unavailable", options); - return; - } if (signal.aborted || this.#isDisposed) { this.#skipAgentContinue("post-restore-unavailable", options); return; } - await this.agent.continue(); + if (this.settings.get("retry.usageAwareFallback")) { + if (!(await this.#runQueuedUsageAwarePreflight(signal))) { + this.#skipAgentContinue("session-unavailable", options); + return; + } + } + await this.agent.continue(signal); } catch (error) { logger.warn("agent.continue failed after scheduling", { error: error instanceof Error ? error.message : String(error), @@ -2989,6 +3018,7 @@ export class AgentSession { }); options?.onError?.(error); } finally { + this.#usagePreflightReadyForNextModelCall = false; this.#endInFlight(); } }, @@ -3607,6 +3637,12 @@ export class AgentSession { */ beginDispose(): void { this.#isDisposed = true; + this.#queuedMessageDrainBlocked = false; + this.#usagePreflightReadyForNextModelCall = false; + this.#detachUsageBeforeQueueDequeue?.(); + this.#detachUsageBeforeQueueDequeue = undefined; + this.#detachUsageBeforeModelCall?.(); + this.#detachUsageBeforeModelCall = undefined; this.#memory.cancelLocalMemoryStartup(); this.#titleGenerationAbortController.abort(); this.#abortAutolearnCapture(); @@ -3928,176 +3964,71 @@ export class AgentSession { } /** Install the interactive decision surface for reserve-triggered model changes. */ - setUsageFallbackConfirmer( - confirmer: ((confirmation: UsageFallbackConfirmation) => Promise) | undefined, - ): void { + setUsageFallbackConfirmer(confirmer: UsageFallbackConfirmer | undefined): void { this.#usageFallbackConfirmer = confirmer; } - async #runUsageAwarePreflight(): Promise { + #allowQueuedMessageDrainRetry(): void { + this.#queuedMessageDrainBlocked = false; + } + + #reconcileQueuedMessageDrain(): void { + if (!this.agent.hasQueuedMessages()) { + this.#queuedMessageDrainBlocked = false; + } + } + + async #runQueuedUsageAwarePreflight(signal?: AbortSignal): Promise { + try { + const allowed = await this.#runUsageAwarePreflight(signal); + this.#usagePreflightReadyForNextModelCall = allowed; + this.#usagePreflightReadyModel = allowed ? this.model : undefined; + this.#queuedMessageDrainBlocked = !allowed && this.agent.hasQueuedMessages(); + return allowed; + } catch (error) { + this.#queuedMessageDrainBlocked = this.agent.hasQueuedMessages(); + throw error; + } + } + + async #runUsageAwarePreflightForNextModelCall(signal?: AbortSignal): Promise { + const allowed = await this.#runUsageAwarePreflight(signal); + this.#usagePreflightReadyForNextModelCall = allowed; + this.#usagePreflightReadyModel = allowed ? this.model : undefined; + return allowed; + } + + async #runUsageAwarePreflight(signal?: AbortSignal): Promise { + if (signal?.aborted) return false; const generation = this.#promptGeneration; + const controller = new AbortController(); + const onAbort = () => controller.abort(signal?.reason); + signal?.addEventListener("abort", onAbort, { once: true }); this.#usagePreflightAbortControllers.add(controller); try { - await this.#maybeApplyUsageAwareFallback(controller.signal); - return !controller.signal.aborted && this.#promptGeneration === generation; - } catch (error) { - if (controller.signal.aborted || this.#promptGeneration !== generation) return false; - throw error; + while (true) { + const model = this.model; + try { + const fallbackCommitted = await this.#recovery.maybeApplyUsageAwareFallback( + controller.signal, + this.#usageFallbackConfirmer, + ); + if (fallbackCommitted) return true; + if (controller.signal.aborted || this.#promptGeneration !== generation) return false; + if (this.model === model || modelsAreEqual(this.model, model)) return true; + } catch (error) { + if (controller.signal.aborted || this.#promptGeneration !== generation) return false; + if (this.model !== model && !modelsAreEqual(this.model, model)) continue; + throw error; + } + } } finally { + signal?.removeEventListener("abort", onAbort); this.#usagePreflightAbortControllers.delete(controller); } } - async #confirmUsageFallback(confirmation: UsageFallbackConfirmation, signal: AbortSignal): Promise { - const confirmer = this.#usageFallbackConfirmer; - if (!confirmer || signal.aborted) return false; - const aborted = Promise.withResolvers(); - const onAbort = () => aborted.resolve(false); - signal.addEventListener("abort", onAbort, { once: true }); - try { - return await Promise.race([confirmer(confirmation), aborted.promise]); - } finally { - signal.removeEventListener("abort", onAbort); - } - } - - async #maybeApplyUsageAwareFallback(signal: AbortSignal): Promise { - if (!this.settings.get("retry.modelFallback") || !this.settings.get("retry.usageAwareFallback")) return; - const currentModel = this.model; - if (!currentModel) return; - const currentSelector = formatRetryFallbackSelector(currentModel, this.thinkingLevel); - let health: ModelUsageHealth; - try { - health = await this.#modelRegistry.authStorage.getModelUsageHealth(currentModel.provider, { - modelId: currentModel.id, - sessionId: this.sessionId, - baseUrl: currentModel.baseUrl, - reserveFraction: this.settings.get("retry.usageReservePct") / 100, - signal, - }); - } catch (error) { - logger.debug("Usage-aware runtime preflight failed open", { - provider: currentModel.provider, - model: currentModel.id, - error: String(error), - }); - return; - } - if (signal.aborted) return; - - if (health.state === "healthy") { - this.#usageReserveApprovedSelector = undefined; - const selected = health.accounts.find(account => account.selected); - if (selected && selected.state !== "healthy" && health.accounts.some(account => account.state === "healthy")) { - this.#modelRegistry.authStorage.releaseSessionCredentialForReselection( - currentModel.provider, - this.sessionId, - ); - } - return; - } - if (health.state === "unknown") { - this.#usageReserveApprovedSelector = undefined; - return; - } - const reservePolicy = this.settings.get("retry.usageReservePolicy"); - if (reservePolicy === "fail-closed") { - const condition = health.state === "reserve" ? "reserve reached" : "usage depleted"; - throw new Error(`${condition} for ${currentSelector}; reserve policy is fail-closed.`); - } - - const role = this.#recovery.resolveRetryFallbackRole(currentSelector, currentModel); - if (!role) return; - let fallback: { selector: RetryFallbackSelector; apiKey: string } | undefined; - for (const candidate of this.#recovery.findRetryFallbackCandidates(role, currentSelector, currentModel)) { - if (this.#recovery.isRetryFallbackSelectorSuppressed(candidate)) continue; - const resolved = resolveModelOverride([candidate.raw], this.#modelRegistry, this.settings); - const candidateModel = resolved.model ?? this.#modelRegistry.find(candidate.provider, candidate.id); - if (!candidateModel) continue; - if (!this.#modelRegistry.hasConfiguredAuth(candidateModel)) continue; - try { - const candidateHealth = await this.#modelRegistry.authStorage.getModelUsageHealth(candidateModel.provider, { - modelId: candidateModel.id, - sessionId: this.sessionId, - baseUrl: candidateModel.baseUrl, - reserveFraction: this.settings.get("retry.usageReservePct") / 100, - signal, - }); - if (signal.aborted) return; - if (candidateHealth.state === "depleted" || candidateHealth.state === "reserve") continue; - if (candidateHealth.state === "healthy") { - const selected = candidateHealth.accounts.find(account => account.selected); - if ( - selected && - selected.state !== "healthy" && - candidateHealth.accounts.some(account => account.state === "healthy") - ) { - this.#modelRegistry.authStorage.releaseSessionCredentialForReselection( - candidateModel.provider, - this.sessionId, - ); - } - } - } catch { - if (signal.aborted) return; - // Unknown usage fails open for an otherwise valid fallback. - } - if (signal.aborted) return; - let apiKey: string | undefined; - try { - apiKey = await this.#modelRegistry.getApiKey(candidateModel, this.sessionId, { signal }); - } catch { - if (signal.aborted) return; - continue; - } - if (signal.aborted) return; - if (!apiKey) continue; - fallback = { selector: candidate, apiKey }; - break; - } - if (!fallback) return; - - if (health.state === "reserve") { - if (reservePolicy === "confirm" && this.#usageFallbackConfirmer) { - if (this.#usageReserveApprovedSelector === currentSelector) return; - const selected = health.accounts.find(account => account.selected); - const remainingFraction = - selected?.remainingFraction ?? - health.accounts.reduce( - (minimum, account) => - account.remainingFraction === undefined - ? minimum - : minimum === undefined - ? account.remainingFraction - : Math.min(minimum, account.remainingFraction), - undefined, - ); - const shouldFallback = await this.#confirmUsageFallback( - { - from: currentSelector, - to: fallback.selector.raw, - remainingPercent: remainingFraction === undefined ? undefined : Math.max(0, remainingFraction * 100), - }, - signal, - ); - if (signal.aborted) return; - if (!shouldFallback) { - this.#usageReserveApprovedSelector = currentSelector; - return; - } - } - } - - if (signal.aborted) return; - this.#usageReserveApprovedSelector = undefined; - await this.#recovery.applyRetryFallbackCandidate(role, fallback.selector, currentSelector, { - pinFallback: true, - apiKey: fallback.apiKey, - signal, - }); - } - /** Effective thinking level applied to the agent (the resolved level when `auto`). */ get thinkingLevel(): ThinkingLevel | undefined { return this.#models.thinkingLevel; @@ -5033,7 +4964,7 @@ export class AgentSession { if (this.isStreaming) { const streamingBehavior = options?.streamingBehavior; if (!streamingBehavior) throw new AgentBusyError(); - if (!(await this.#runUsageAwarePreflight())) return false; + // Steer/follow-up the keyword notices BEFORE the queued user message so the // model reads the steering notice ahead of the prompt it modifies. for (const notice of keywordNotices) { @@ -5128,7 +5059,7 @@ export class AgentSession { if (options?.queueOnly) { const streamingBehavior = options?.streamingBehavior; if (!streamingBehavior) throw new AgentBusyError(); - if (!(await this.#runUsageAwarePreflight())) return; + for (const notice of keywordNotices) { await this.#queueCustomMessage(notice, streamingBehavior); } @@ -5138,7 +5069,7 @@ export class AgentSession { if (this.isStreaming) { const streamingBehavior = options?.streamingBehavior; if (!streamingBehavior) throw new AgentBusyError(); - if (!(await this.#runUsageAwarePreflight())) return; + for (const notice of keywordNotices) { await this.#queueCustomMessage(notice, streamingBehavior); } @@ -5175,7 +5106,7 @@ export class AgentSession { const generation = this.#promptGeneration; try { await this.#recovery.maybeRestoreRetryFallbackPrimary(); - if (!(await this.#runUsageAwarePreflight())) return; + if (!(await this.#runUsageAwarePreflightForNextModelCall())) return; // Flush any pending bash messages before the new prompt await this.#bash.flushPending(); this.#eval.flushPending(); @@ -5386,6 +5317,7 @@ export class AgentSession { await this.#waitForPostPromptRecovery(generation); } } finally { + this.#usagePreflightReadyForNextModelCall = false; this.#endInFlight(); } } @@ -5551,7 +5483,6 @@ export class AgentSession { } const expandedText = expandPromptTemplate(text, [...this.#promptTemplates]); - if (!(await this.#runUsageAwarePreflight())) return; await this.#queueUserMessage(expandedText, images, "steer"); } @@ -5569,7 +5500,6 @@ export class AgentSession { const expandedText = options?.expandPromptTemplates === false ? text : expandPromptTemplate(text, [...this.#promptTemplates]); - if (!(await this.#runUsageAwarePreflight())) return; if (!options?.synthetic) { await this.#queueUserMessage(expandedText, images, "followUp"); return; @@ -5586,6 +5516,7 @@ export class AgentSession { const imageDescriptionNotice = normalizedImages?.length ? await this.#buildImageDescriptionNotice(normalizedImages) : undefined; + this.#allowQueuedMessageDrainRetry(); if (imageDescriptionNotice) this.agent.followUp(imageDescriptionNotice); this.agent.followUp({ role: "developer", @@ -5615,6 +5546,7 @@ export class AgentSession { const imageDescriptionNotice = normalizedImages?.length ? await this.#buildImageDescriptionNotice(normalizedImages) : undefined; + this.#allowQueuedMessageDrainRetry(); if (mode === "followUp") { if (imageDescriptionNotice) this.agent.followUp(imageDescriptionNotice); this.agent.followUp({ @@ -5641,7 +5573,12 @@ export class AgentSession { } #scheduleQueuedMessageDrain(): void { - if (this.#queuedMessageDrainScheduled || !this.#canAutoContinueForFollowUp() || !this.agent.hasQueuedMessages()) { + if ( + this.#queuedMessageDrainScheduled || + this.#queuedMessageDrainBlocked || + !this.#canAutoContinueForFollowUp() || + !this.agent.hasQueuedMessages() + ) { return; } this.#queuedMessageDrainScheduled = true; @@ -5655,6 +5592,7 @@ export class AgentSession { }, onError: () => { this.#queuedMessageDrainScheduled = false; + this.#queuedMessageDrainBlocked = this.agent.hasQueuedMessages(); }, }); } @@ -5791,7 +5729,7 @@ export class AgentSession { ): Promise { this.#beginInFlight(); try { - if (!(await this.#runUsageAwarePreflight())) return; + if (!(await this.#runUsageAwarePreflightForNextModelCall())) return; const acceptTerminalEmptyStop = options?.acceptTerminalEmptyStop === true; if (acceptTerminalEmptyStop) { this.#resetPromptMaintenanceState(); @@ -5800,6 +5738,7 @@ export class AgentSession { await this.agent.prompt(message); await this.#waitForPostPromptRecovery(); } finally { + this.#usagePreflightReadyForNextModelCall = false; this.#recovery.setAcceptTerminalEmptyStop(false); this.#endInFlight(); } @@ -5831,6 +5770,7 @@ export class AgentSession { timestamp: Date.now(), }; const normalizedAppMessage = await this.#normalizeAgentMessageImages(appMessage); + this.#allowQueuedMessageDrainRetry(); if (deliverAs === "followUp") { this.agent.followUp(normalizedAppMessage); } else { @@ -5887,7 +5827,7 @@ export class AgentSession { this.#queueHiddenNextTurnMessage(normalizedAppMessage, options?.triggerTurn ?? false); return false; } - if (!(await this.#runUsageAwarePreflight())) return false; + this.#allowQueuedMessageDrainRetry(); if (options?.deliverAs === "followUp") { this.agent.followUp(normalizedAppMessage); @@ -5970,8 +5910,6 @@ export class AgentSession { if (images.length === 0) images = undefined; } - if (options?.deliverAs && !(await this.#runUsageAwarePreflight())) return; - if (options?.deliverAs === "followUp") { await this.#queueUserMessage(text, images, "followUp"); return; @@ -6011,6 +5949,7 @@ export class AgentSession { ? isAdvisorCard : m => !isUserQueuedMessage(m) && !isHiddenUserCompanion(m); this.agent.replaceQueues(steeringAll.filter(keep), followUpAll.filter(keep)); + this.#reconcileQueuedMessageDrain(); return { steering, followUp }; } @@ -6060,12 +5999,14 @@ export class AgentSession { if (fromSteer >= 0) { const removed = steering[fromSteer]; this.agent.replaceQueues(removeWithCompanions(steering, fromSteer), followUp.slice()); + this.#reconcileQueuedMessageDrain(); return toRestoredQueuedMessage(removed); } const fromFollowUp = lastUserIndex(followUp); if (fromFollowUp >= 0) { const removed = followUp[fromFollowUp]; this.agent.replaceQueues(steering.slice(), removeWithCompanions(followUp, fromFollowUp)); + this.#reconcileQueuedMessageDrain(); return toRestoredQueuedMessage(removed); } return undefined; @@ -6355,6 +6296,8 @@ export class AgentSession { await this.#memory.resetContextForNewTranscript(); this.#pendingNextTurnMessages = []; this.#scheduledHiddenNextTurnGeneration = undefined; + this.#queuedMessageDrainBlocked = false; + this.#usagePreflightReadyForNextModelCall = false; this.sessionManager.appendThinkingLevelChange(this.thinkingLevel, this.configuredThinkingLevel()); this.sessionManager.appendServiceTierChange(this.#models.serviceTierEntry()); @@ -7379,6 +7322,9 @@ export class AgentSession { const previousFollowUpMessages = [...this.agent.peekFollowUpQueue()]; const previousPendingNextTurnMessages = [...this.#pendingNextTurnMessages]; const previousScheduledHiddenNextTurnGeneration = this.#scheduledHiddenNextTurnGeneration; + const previousQueuedMessageDrainBlocked = this.#queuedMessageDrainBlocked; + const previousUsagePreflightReadyForNextModelCall = this.#usagePreflightReadyForNextModelCall; + const previousUsagePreflightReadyModel = this.#usagePreflightReadyModel; const previousModel = this.model; const previousThinkingLevel = this.thinkingLevel; const previousAutoThinking = this.isAutoThinking; @@ -7403,6 +7349,9 @@ export class AgentSession { this.agent.clearAllQueues(); this.#pendingNextTurnMessages = []; this.#scheduledHiddenNextTurnGeneration = undefined; + this.#queuedMessageDrainBlocked = false; + this.#usagePreflightReadyForNextModelCall = false; + this.#usagePreflightReadyModel = undefined; try { if (switchingToDifferentSession) { @@ -7571,6 +7520,9 @@ export class AgentSession { this.agent.replaceQueues(previousSteeringMessages, previousFollowUpMessages); this.#pendingNextTurnMessages = previousPendingNextTurnMessages; this.#scheduledHiddenNextTurnGeneration = previousScheduledHiddenNextTurnGeneration; + this.#queuedMessageDrainBlocked = previousQueuedMessageDrainBlocked; + this.#usagePreflightReadyForNextModelCall = previousUsagePreflightReadyForNextModelCall; + this.#usagePreflightReadyModel = previousUsagePreflightReadyModel; this.#inheritedProviderPromptCacheKey = previousInheritedProviderPromptCacheKey; this.#checkpointState = previousCheckpointState; this.#pendingRewindReport = previousPendingRewindReport; @@ -7659,6 +7611,8 @@ export class AgentSession { // Clear pending messages (bound to old session state) this.#pendingNextTurnMessages = []; this.#scheduledHiddenNextTurnGeneration = undefined; + this.#queuedMessageDrainBlocked = false; + this.#usagePreflightReadyForNextModelCall = false; await this.#bash.flushPending(); // Flush pending writes before branching @@ -7786,6 +7740,8 @@ export class AgentSession { this.#pendingNextTurnMessages = []; this.#scheduledHiddenNextTurnGeneration = undefined; this.agent.replaceQueues([], []); + this.#queuedMessageDrainBlocked = false; + this.#usagePreflightReadyForNextModelCall = false; await this.#bash.flushPending(); await this.sessionManager.flush(); const bashTransition = this.#bash.beginSessionTransition(); diff --git a/packages/coding-agent/src/session/turn-recovery.ts b/packages/coding-agent/src/session/turn-recovery.ts index 5c02e86b8..37393ab81 100644 --- a/packages/coding-agent/src/session/turn-recovery.ts +++ b/packages/coding-agent/src/session/turn-recovery.ts @@ -13,6 +13,7 @@ import type { CodexCompactionContext, Effort, Model, + ModelUsageHealth, TextContent, ThinkingContent, ToolChoice, @@ -21,9 +22,11 @@ import { calculateRateLimitBackoffMs, parseRateLimitReason } from "@oh-my-pi/pi- import * as AIError from "@oh-my-pi/pi-ai/error"; import { kCursorExecResolved } from "@oh-my-pi/pi-ai/utils/block-symbols"; import { isFireworksFastModelId, toFireworksBaseModelId } from "@oh-my-pi/pi-catalog/fireworks-model-id"; +import { modelsAreEqual } from "@oh-my-pi/pi-catalog/models"; import { extractRetryHint, logger, prompt } from "@oh-my-pi/pi-utils"; import type { ModelRegistry } from "../config/model-registry"; import { formatModelStringWithRouting, resolveModelOverride } from "../config/model-resolver"; + import type { Settings } from "../config/settings"; import type { RecoveredRetryError } from "../extensibility/shared-events"; import emptyStopRetryTemplate from "../prompts/system/empty-stop-retry.md" with { type: "text" }; @@ -36,7 +39,11 @@ import { modelSupportsEffortCeiling, } from "../thinking"; import type { AgentSessionEvent } from "./agent-session-events"; -import type { InitialRetryFallbackState } from "./agent-session-types"; +import type { + InitialRetryFallbackState, + UsageFallbackConfirmation, + UsageFallbackConfirmer, +} from "./agent-session-types"; import { isEmptyErrorTurn } from "./messages"; import { type ActiveRetryFallbackState, @@ -65,6 +72,7 @@ const UNEXPECTED_STOP_TIMEOUT_MS = 4000; const EMPTY_STOP_MAX_RETRIES = 3; const SIBLING_UNBLOCK_BUFFER_MS = 1_000; const NON_WHITESPACE_RE = /\S/; +const USAGE_PREFLIGHT_BLOCKED_PREFIX = "Usage preflight blocked:"; function hasNonWhitespace(value: string): boolean { return NON_WHITESPACE_RE.test(value); @@ -175,6 +183,7 @@ export class TurnRecovery { #retryPromise: Promise | undefined; #retryResolve: (() => void) | undefined; #activeRetryFallback: ActiveRetryFallbackState | undefined; + #usageReserveApprovedSelector: string | undefined; #pendingRecoveredRetryErrors: PendingRecoveredRetryError[] = []; #usageLimitOutcomes = new WeakMap>(); #emptyStopRetryCount = 0; @@ -301,6 +310,11 @@ export class TurnRecovery { return this.#maybeRestoreRetryFallbackPrimary(); } + /** Applies model fallback policy from live usage health before a turn starts. */ + maybeApplyUsageAwareFallback(signal: AbortSignal, confirmer?: UsageFallbackConfirmer): Promise { + return this.#maybeApplyUsageAwareFallback(signal, confirmer); + } + /** Applies automatic retry, credential rotation, and model fallback policy. */ handleRetryableError( message: AssistantMessage, @@ -851,6 +865,9 @@ export class TurnRecovery { return id; } + #isUsagePreflightBlocked(message: AssistantMessage): boolean { + return message.errorMessage?.startsWith(USAGE_PREFLIGHT_BLOCKED_PREFIX) === true; + } /** * Retry an empty, reason-less provider abort: a turn with no content that * carries the generic sentinel (bare `abort()`), whether the provider @@ -894,6 +911,7 @@ export class TurnRecovery { */ isRetryableError(message: AssistantMessage): boolean { if (message.stopReason !== "error") return false; + if (this.#isUsagePreflightBlocked(message)) return false; const id = this.#classifyRetryMessage(message); // Context overflow is handled by compaction, not retry @@ -1019,7 +1037,6 @@ export class TurnRecovery { modelLookup: this.#host.modelRegistry, }; } - #getRetryFallbackChains(): RetryFallbackChains { return getRetryFallbackChains(this.#host.settings); } @@ -1082,23 +1099,187 @@ export class TurnRecovery { ); } + async #maybeApplyUsageAwareFallback(signal: AbortSignal, confirmer?: UsageFallbackConfirmer): Promise { + if (!this.#host.settings.get("retry.usageAwareFallback")) return false; + const currentModel = this.#host.model(); + if (!currentModel) return false; + const currentSelector = formatRetryFallbackSelector(currentModel, this.#host.thinkingLevel()); + let health: ModelUsageHealth; + try { + health = await this.#host.modelRegistry.authStorage.getModelUsageHealth(currentModel.provider, { + modelId: currentModel.id, + sessionId: this.#host.sessionId(), + baseUrl: currentModel.baseUrl, + reserveFraction: this.#host.settings.get("retry.usageReservePct") / 100, + signal, + }); + } catch (error) { + if (signal.aborted || !modelsAreEqual(this.#host.model(), currentModel)) return false; + logger.debug("Usage-aware runtime preflight failed open", { + provider: currentModel.provider, + model: currentModel.id, + error: String(error), + }); + return false; + } + if (signal.aborted || !modelsAreEqual(this.#host.model(), currentModel)) return false; + const selectedAccount = health.accounts.find(account => account.selected); + if (health.state === "healthy") { + this.#usageReserveApprovedSelector = undefined; + if ( + selectedAccount && + selectedAccount.state !== "healthy" && + health.accounts.some(account => account.state === "healthy") + ) { + this.#host.modelRegistry.authStorage.releaseSessionCredentialForReselection( + currentModel.provider, + this.#host.sessionId(), + ); + } + return false; + } + if (health.state === "unknown") { + this.#usageReserveApprovedSelector = undefined; + return false; + } + if (health.state !== "reserve") this.#usageReserveApprovedSelector = undefined; + + const reservePolicy = this.#host.settings.get("retry.usageReservePolicy"); + if (reservePolicy === "fail-closed") { + const condition = health.state === "reserve" ? "reserve reached" : "usage depleted"; + throw new Error( + `${USAGE_PREFLIGHT_BLOCKED_PREFIX} ${condition} for ${currentSelector}; reserve policy is fail-closed.`, + ); + } + if ( + reservePolicy === "confirm" && + health.state === "reserve" && + this.#usageReserveApprovedSelector === currentSelector + ) { + return false; + } + if (!this.#host.settings.get("retry.modelFallback")) return false; + + const role = this.#activeRetryFallback?.role ?? this.resolveRetryFallbackRole(currentSelector, currentModel); + if (!role) return false; + let fallback: { selector: RetryFallbackSelector; apiKey: string } | undefined; + const ceiling = this.#host.thinkingLevelCeiling(); + for (const candidate of this.findRetryFallbackCandidates(role, currentSelector, currentModel)) { + if (this.isRetryFallbackSelectorSuppressed(candidate)) continue; + const resolved = resolveModelOverride([candidate.raw], this.#host.modelRegistry, this.#host.settings); + const candidateModel = resolved.model ?? this.#host.modelRegistry.find(candidate.provider, candidate.id); + if (!candidateModel || !this.#host.modelRegistry.hasConfiguredAuth(candidateModel)) continue; + if (ceiling !== undefined && !modelSupportsEffortCeiling(candidateModel, ceiling)) continue; + try { + const candidateHealth = await this.#host.modelRegistry.authStorage.getModelUsageHealth( + candidateModel.provider, + { + modelId: candidateModel.id, + sessionId: this.#host.sessionId(), + baseUrl: candidateModel.baseUrl, + reserveFraction: this.#host.settings.get("retry.usageReservePct") / 100, + signal, + }, + ); + if (signal.aborted || !modelsAreEqual(this.#host.model(), currentModel)) return false; + if (candidateHealth.state === "depleted" || candidateHealth.state === "reserve") continue; + if (candidateHealth.state === "healthy") { + const selected = candidateHealth.accounts.find(account => account.selected); + if ( + selected && + selected.state !== "healthy" && + candidateHealth.accounts.some(account => account.state === "healthy") + ) { + this.#host.modelRegistry.authStorage.releaseSessionCredentialForReselection( + candidateModel.provider, + this.#host.sessionId(), + ); + } + } + } catch { + if (signal.aborted || !modelsAreEqual(this.#host.model(), currentModel)) return false; + // Unknown usage fails open for an otherwise valid fallback. + } + if (signal.aborted || !modelsAreEqual(this.#host.model(), currentModel)) return false; + let apiKey: string | undefined; + try { + apiKey = await this.#host.modelRegistry.getApiKey(candidateModel, this.#host.sessionId(), { signal }); + } catch { + if (signal.aborted || !modelsAreEqual(this.#host.model(), currentModel)) return false; + continue; + } + if (signal.aborted || !modelsAreEqual(this.#host.model(), currentModel)) return false; + if (!apiKey) continue; + fallback = { selector: candidate, apiKey }; + break; + } + if (!fallback) return false; + + let shouldFallback = health.state === "depleted" || reservePolicy === "auto" || !confirmer; + if (!shouldFallback && health.state === "reserve" && confirmer) { + const remainingFraction = + selectedAccount?.remainingFraction ?? + health.accounts.reduce((minimum, account) => { + if (account.remainingFraction === undefined) return minimum; + return minimum === undefined ? account.remainingFraction : Math.min(minimum, account.remainingFraction); + }, undefined); + shouldFallback = await this.#confirmUsageFallback( + confirmer, + { + from: currentSelector, + to: fallback.selector.raw, + remainingPercent: remainingFraction === undefined ? undefined : Math.max(0, remainingFraction * 100), + }, + signal, + ); + if (signal.aborted || !modelsAreEqual(this.#host.model(), currentModel)) return false; + } + if (!shouldFallback) { + this.#usageReserveApprovedSelector = currentSelector; + return false; + } + this.#usageReserveApprovedSelector = undefined; + return this.applyRetryFallbackCandidate(role, fallback.selector, currentSelector, { + pinFallback: true, + apiKey: fallback.apiKey, + signal, + }); + } + + async #confirmUsageFallback( + confirmer: UsageFallbackConfirmer, + confirmation: UsageFallbackConfirmation, + signal: AbortSignal, + ): Promise { + if (signal.aborted) return false; + const aborted = Promise.withResolvers(); + const onAbort = () => aborted.resolve(false); + signal.addEventListener("abort", onAbort, { once: true }); + try { + return await Promise.race([confirmer(confirmation, signal), aborted.promise]); + } finally { + signal.removeEventListener("abort", onAbort); + } + } + async applyRetryFallbackCandidate( role: string, selector: RetryFallbackSelector, currentSelector: string, options?: { pinFallback?: boolean; apiKey?: string; signal?: AbortSignal }, - ): Promise { + ): Promise { const resolved = resolveModelOverride([selector.raw], this.#host.modelRegistry, this.#host.settings); const candidate = resolved.model ?? this.#host.modelRegistry.find(selector.provider, selector.id); if (!candidate) { throw new Error(`Retry fallback model not found: ${selector.raw}`); } const apiKey = - options?.apiKey ?? (await this.#host.modelRegistry.getApiKey(candidate, this.#host.sessionId(), options)); + options?.apiKey ?? + (await this.#host.modelRegistry.getApiKey(candidate, this.#host.sessionId(), { signal: options?.signal })); if (!apiKey) { throw new Error(`No API key for retry fallback ${selector.raw}`); } - if (options?.signal?.aborted) return; + if (options?.signal?.aborted) return false; // Capture the configured selector (auto-aware) so a fallback chain preserves // `auto` instead of collapsing it to the level it resolved to this turn. @@ -1112,7 +1293,15 @@ export class TurnRecovery { ? requestedThinkingLevel : clampThinkingLevelToCeiling(candidate, requestedThinkingLevel, this.#host.thinkingLevelCeiling()); const candidateSelector = formatModelStringWithRouting(candidate); + const previousModel = this.#host.model(); await this.#host.setModelWithProviderSessionReset(candidate); + if (options?.signal?.aborted) { + if (previousModel && this.#host.model() === candidate) { + await this.#host.setModelWithProviderSessionReset(previousModel); + } + return false; + } + if (this.#host.model() !== candidate) return false; this.#host.sessionManager.appendModelChange(candidateSelector, EPHEMERAL_MODEL_CHANGE_ROLE); this.#host.settings.getStorage()?.recordModelUsage(candidateSelector); this.#host.setThinkingLevel(nextThinkingLevel); @@ -1134,6 +1323,7 @@ export class TurnRecovery { to: selector.raw, role, }); + return true; } async #tryRetryModelFallback(currentSelector: string, options?: { pinFallback?: boolean }): Promise { @@ -1151,8 +1341,7 @@ export class TurnRecovery { if (ceiling !== undefined && !modelSupportsEffortCeiling(candidate, ceiling)) continue; const apiKey = await this.#host.modelRegistry.getApiKey(candidate, this.#host.sessionId()); if (!apiKey) continue; - await this.applyRetryFallbackCandidate(role, selector, currentSelector, options); - return true; + return this.applyRetryFallbackCandidate(role, selector, currentSelector, options); } return false; @@ -1178,6 +1367,7 @@ export class TurnRecovery { const model = this.#activeFireworksFastModel(); if (!model) return false; if (message.stopReason !== "error") return false; + if (this.#isUsagePreflightBlocked(message)) return false; if (this.#hasReplayUnsafeOutput(message)) return false; // A content refusal/sensitivity stop is the model's decision, not a route // failure — switching to the base model would just re-trigger it. @@ -1201,6 +1391,7 @@ export class TurnRecovery { */ isHardErrorFallbackEligible(message: AssistantMessage): boolean { if (message.stopReason !== "error") return false; + if (this.#isUsagePreflightBlocked(message)) return false; const model = this.#host.model(); if (!model) return false; const retrySettings = this.#host.settings.getGroup("retry"); diff --git a/packages/coding-agent/src/task/executor.ts b/packages/coding-agent/src/task/executor.ts index 610696357..904e21d23 100644 --- a/packages/coding-agent/src/task/executor.ts +++ b/packages/coding-agent/src/task/executor.ts @@ -168,9 +168,11 @@ function resolveSubagentRetryFallbackCandidates( ): SubagentRetryFallbackCandidate[] { const candidates: SubagentRetryFallbackCandidate[] = []; const seen = new Set(); + const disabledProviders = new Set(settings.get("disabledProviders")); for (const pattern of modelPatterns) { const resolved = resolveModelOverride([pattern], modelRegistry, settings); if (!resolved.model) continue; + if (disabledProviders.has(resolved.model.provider)) continue; const selector = resolved.explicitThinkingLevel ? formatModelSelectorValue(formatModelStringWithRouting(resolved.model), resolved.thinkingLevel) : formatModelStringWithRouting(resolved.model); @@ -181,7 +183,10 @@ function resolveSubagentRetryFallbackCandidates( return candidates; } -function resolveSubagentDefaultRetryFallbackChain(settings: Settings): string[] | undefined { +function resolveSubagentDefaultRetryFallbackChain( + settings: Settings, + modelRegistry: ModelRegistry, +): string[] | undefined { const fallbackChain = settings.get("retry.fallbackChains")?.default; if ( !Array.isArray(fallbackChain) || @@ -190,7 +195,11 @@ function resolveSubagentDefaultRetryFallbackChain(settings: Settings): string[] ) { return undefined; } - return fallbackChain; + const disabledProviders = new Set(settings.get("disabledProviders")); + return fallbackChain.filter(entry => { + const resolved = resolveModelOverride([entry], modelRegistry, settings); + return !resolved.model || !disabledProviders.has(resolved.model.provider); + }); } function installSubagentRetryFallbackChain(args: { @@ -2783,7 +2792,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise { authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); authStorage.setRuntimeApiKey("anthropic", "anthropic-test-key"); authStorage.setRuntimeApiKey("openai", "openai-test-key"); + authStorage.setRuntimeApiKey("fireworks", "fireworks-test-key"); authStorage.setRuntimeApiKey("google", "google-test-key"); authStorage.setRuntimeApiKey("google-vertex", "google-vertex-test-key"); authStorage.setRuntimeApiKey("openrouter", "openrouter-test-key"); @@ -290,9 +292,15 @@ describe("AgentSession retry fallback", () => { { credentialId: 1, credentialType: "oauth", + state: "reserve", + remainingFraction: 0.08, + }, + { + credentialId: 2, + credentialType: "oauth", selected: true, state: "reserve", - remainingFraction: 0.05, + remainingFraction: 0.02, }, ], } @@ -308,15 +316,72 @@ describe("AgentSession retry fallback", () => { session.setUsageFallbackConfirmer(confirmFallback); await session.prompt("Keep working on the same task"); await session.waitForIdle(); - expect(confirmFallback).toHaveBeenCalledWith({ - from: `${primaryModel.provider}/${primaryModel.id}`, - to: `${fallbackModel.provider}/${fallbackModel.id}`, - remainingPercent: 5, - }); + expect(confirmFallback).toHaveBeenCalledWith( + { + from: `${primaryModel.provider}/${primaryModel.id}`, + to: `${fallbackModel.provider}/${fallbackModel.id}`, + remainingPercent: 2, + }, + expect.any(AbortSignal), + ); expect(requestedModels).toEqual([`${fallbackModel.provider}/${fallbackModel.id}`]); expect(session.messages.some(message => message.role === "user")).toBe(true); }); + it("honors a live fail-closed policy after reserve spending was approved", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + const fallbackModel = getBundledModel("openai", "gpt-4o-mini"); + if (!primaryModel || !fallbackModel) throw new Error("Expected bundled reserve policy models"); + const mock = createMockModel({ responses: [{ content: ["stayed on primary"] }] }); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: mock.stream, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "confirm", + "retry.fallbackChains": { + default: [`${fallbackModel.provider}/${fallbackModel.id}`], + }, + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + const usageHealth = vi + .spyOn(modelRegistry.authStorage, "getModelUsageHealth") + .mockImplementation(async provider => + provider === primaryModel.provider + ? { + state: "reserve", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "reserve", + remainingFraction: 0.05, + }, + ], + } + : { state: "healthy", accounts: [] }, + ); + const confirmFallback = vi.fn(async () => false); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + session.setUsageFallbackConfirmer(confirmFallback); + + await session.prompt("Stay on the primary"); + await session.waitForIdle(); + settings.override("retry.usageReservePolicy", "fail-closed"); + expect(settings.get("retry.usageReservePolicy")).toBe("fail-closed"); + + await expect(session.prompt("Do not spend reserve")).rejects.toThrow("reserve policy is fail-closed"); + expect(confirmFallback).toHaveBeenCalledTimes(1); + expect(usageHealth).toHaveBeenCalledTimes(3); + }); it("reselects a healthy same-provider account before considering a model fallback", async () => { const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); const fallbackModel = getBundledModel("openai", "gpt-4o-mini"); @@ -485,6 +550,647 @@ describe("AgentSession retry fallback", () => { expect(requestedModels).toEqual([]); }); + it("cancels a pending reserve confirmation without dispatching the prompt", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + const fallbackModel = getBundledModel("openai", "gpt-4o-mini"); + if (!primaryModel || !fallbackModel) throw new Error("Expected bundled confirmation cancellation models"); + const requestedModels: string[] = []; + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return createMockModel().stream(model, context, options); + }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.fallbackChains": { + default: [`${fallbackModel.provider}/${fallbackModel.id}`], + }, + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + vi.spyOn(modelRegistry.authStorage, "getModelUsageHealth").mockImplementation(async provider => + provider === primaryModel.provider + ? { + state: "reserve", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "reserve", + remainingFraction: 0.05, + }, + ], + } + : { state: "healthy", accounts: [] }, + ); + const confirmationStarted = Promise.withResolvers(); + const pendingConfirmation = Promise.withResolvers(); + const confirmationAborted = Promise.withResolvers(); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + session.setUsageFallbackConfirmer(async (_confirmation, signal) => { + confirmationStarted.resolve(); + signal.addEventListener("abort", () => confirmationAborted.resolve(), { once: true }); + return pendingConfirmation.promise; + }); + + const prompt = session.prompt("Do not send after confirmation cancellation"); + await confirmationStarted.promise; + await session.abort(); + await confirmationAborted.promise; + await prompt; + + expect(requestedModels).toEqual([]); + }); + + it("defers usage fallback for a queued steer until the active stream finishes", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + const fallbackModel = getBundledModel("openai", "gpt-4o-mini"); + if (!primaryModel || !fallbackModel) throw new Error("Expected bundled queued fallback models"); + const requestedModels: string[] = []; + const streamStarted = Promise.withResolvers(); + const firstResponse = Promise.withResolvers<{ content: string[] }>(); + const mock = createMockModel({ + responses: [ + async () => { + streamStarted.resolve(); + return firstResponse.promise; + }, + { content: ["queued steer completed"] }, + ], + }); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return mock.stream(model, context, options); + }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "auto", + "retry.fallbackChains": { + default: [`${fallbackModel.provider}/${fallbackModel.id}`], + }, + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + let useReserve = false; + const usageHealth = vi + .spyOn(modelRegistry.authStorage, "getModelUsageHealth") + .mockImplementation(async provider => + provider === primaryModel.provider + ? useReserve + ? { + state: "reserve", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "reserve", + remainingFraction: 0.05, + }, + ], + } + : { + state: "healthy", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "healthy", + remainingFraction: 0.8, + }, + ], + } + : { state: "healthy", accounts: [] }, + ); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + + const firstPrompt = session.prompt("Keep the primary stream active"); + await streamStarted.promise; + useReserve = true; + await session.sendUserMessage("Queue this steer", { deliverAs: "steer" }); + + expect(usageHealth).toHaveBeenCalledTimes(1); + expect(session.model?.id).toBe(primaryModel.id); + + firstResponse.resolve({ content: ["primary stream completed"] }); + await firstPrompt; + await session.waitForIdle(); + + expect(requestedModels).toEqual([ + `${primaryModel.provider}/${primaryModel.id}`, + `${fallbackModel.provider}/${fallbackModel.id}`, + ]); + }); + + it("cancels queued-turn usage confirmation when post-prompt work is disposed", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + const fallbackModel = getBundledModel("openai", "gpt-4o-mini"); + if (!primaryModel || !fallbackModel) throw new Error("Expected bundled queued cancellation models"); + const requestedModels: string[] = []; + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return createMockModel().stream(model, context, options); + }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.fallbackChains": { + default: [`${fallbackModel.provider}/${fallbackModel.id}`], + }, + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + vi.spyOn(modelRegistry.authStorage, "getModelUsageHealth").mockImplementation(async provider => + provider === primaryModel.provider + ? { + state: "reserve", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "reserve", + remainingFraction: 0.05, + }, + ], + } + : { state: "healthy", accounts: [] }, + ); + const confirmationStarted = Promise.withResolvers(); + const pendingConfirmation = Promise.withResolvers(); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + session.setUsageFallbackConfirmer(async () => { + confirmationStarted.resolve(); + return pendingConfirmation.promise; + }); + + await session.sendUserMessage("Queue this turn", { deliverAs: "steer" }); + await confirmationStarted.promise; + await session.dispose(); + session = undefined; + + expect(requestedModels).toEqual([]); + }); + + it("does not reschedule a queued drain after a dequeue hook rejects", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!primaryModel) throw new Error("Expected bundled queued-drain model"); + const requestedModels: string[] = []; + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return createMockModel().stream(model, context, options); + }, + }); + const settings = Settings.isolated({ "compaction.enabled": false }); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + const hookRan = Promise.withResolvers(); + let attempts = 0; + const failingHook = vi.fn(() => { + hookRan.resolve(); + if (++attempts === 1) throw new Error("blocked before dequeue"); + }); + agent.addBeforeQueuedMessageDequeueHook(failingHook); + + await session.sendUserMessage("Keep this queued", { deliverAs: "steer" }); + await hookRan.promise; + await session.waitForIdle(); + + expect(failingHook).toHaveBeenCalledTimes(1); + expect(agent.hasQueuedMessages()).toBe(true); + expect(requestedModels).toEqual([]); + }); + + it("enforces fail-closed usage health when model fallback is disabled", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!primaryModel) throw new Error("Expected bundled fail-closed model"); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: createMockModel().stream, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.modelFallback": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "fail-closed", + }); + vi.spyOn(modelRegistry.authStorage, "getModelUsageHealth").mockResolvedValue({ + state: "reserve", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "reserve", + remainingFraction: 0.05, + }, + ], + }); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + + await expect(session.prompt("Do not spend reserve")).rejects.toThrow("reserve policy is fail-closed"); + }); + + it("does not degrade Fireworks Fast or retry a chain after queued fail-closed preflight", async () => { + const primaryModel = getBundledModel("fireworks", "kimi-k2.6-fast"); + const fallbackModel = getBundledModel("openai", "gpt-4o-mini"); + if (!primaryModel || !fallbackModel) throw new Error("Expected bundled queued fail-closed models"); + const requestedModels: string[] = []; + const streamStarted = Promise.withResolvers(); + const firstResponse = Promise.withResolvers<{ content: string[] }>(); + const mock = createMockModel({ + responses: [ + async () => { + streamStarted.resolve(); + return firstResponse.promise; + }, + { content: ["must not run"] }, + ], + }); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return mock.stream(model, context, options); + }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "fail-closed", + "retry.fallbackChains": { + default: [`${fallbackModel.provider}/${fallbackModel.id}`], + }, + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + let useReserve = false; + const usageHealth = vi.spyOn(modelRegistry.authStorage, "getModelUsageHealth").mockImplementation(async () => + useReserve + ? { + state: "reserve", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "reserve", + remainingFraction: 0.05, + }, + ], + } + : { + state: "healthy", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "healthy", + remainingFraction: 0.8, + }, + ], + }, + ); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + + const firstPrompt = session.prompt("Keep the primary stream active"); + await streamStarted.promise; + useReserve = true; + await session.sendUserMessage("Queue blocked work", { deliverAs: "steer" }); + firstResponse.resolve({ content: ["primary stream completed"] }); + await firstPrompt; + await session.waitForIdle(); + + expect(usageHealth).toHaveBeenCalledTimes(2); + expect(requestedModels).toEqual([`${primaryModel.provider}/${primaryModel.id}`]); + expect(session.model?.id).toBe(primaryModel.id); + expect(agent.hasQueuedMessages()).toBe(true); + }); + + it("rechecks fail-closed usage health before an internally scheduled continuation", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!primaryModel) throw new Error("Expected bundled scheduled continuation model"); + const requestedModels: string[] = []; + let useReserve = false; + const mock = createMockModel({ + responses: [ + async () => { + useReserve = true; + return { content: [], stopReason: "stop" }; + }, + { content: ["must not run"] }, + ], + }); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return mock.stream(model, context, options); + }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "fail-closed", + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + const usageHealth = vi.spyOn(modelRegistry.authStorage, "getModelUsageHealth").mockImplementation(async () => + useReserve + ? { + state: "reserve", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "reserve", + remainingFraction: 0.05, + }, + ], + } + : { + state: "healthy", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "healthy", + remainingFraction: 0.8, + }, + ], + }, + ); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + + await session.prompt("Retry this empty response"); + await session.waitForIdle(); + + expect(usageHealth).toHaveBeenCalledTimes(2); + expect(requestedModels).toEqual([`${primaryModel.provider}/${primaryModel.id}`]); + }); + + it("rechecks fail-closed usage health before a same-turn tool continuation", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + if (!primaryModel) throw new Error("Expected bundled tool-continuation model"); + const requestedModels: string[] = []; + let useReserve = false; + const toolSchema = z.object({ value: z.string() }); + const tool: AgentTool = { + name: "consume", + label: "Consume", + description: "Consume plan quota", + parameters: toolSchema, + async execute(_toolCallId, params) { + useReserve = true; + return { content: [{ type: "text", text: params.value }], details: params }; + }, + }; + const mock = createMockModel({ + responses: [ + { content: [{ type: "toolCall", id: "tool-1", name: "consume", arguments: { value: "done" } }] }, + { content: ["must not run"] }, + ], + }); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [tool], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return mock.stream(model, context, options); + }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "fail-closed", + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + const usageHealth = vi.spyOn(modelRegistry.authStorage, "getModelUsageHealth").mockImplementation(async () => + useReserve + ? { + state: "reserve", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "reserve", + remainingFraction: 0.05, + }, + ], + } + : { + state: "healthy", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: "healthy", + remainingFraction: 0.8, + }, + ], + }, + ); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + + await session.prompt("Use the tool"); + await session.waitForIdle(); + + expect(usageHealth).toHaveBeenCalledTimes(2); + expect(requestedModels).toEqual([`${primaryModel.provider}/${primaryModel.id}`]); + }); + it("rechecks fail-closed usage health when prompt setup changes the model", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + const setupTarget = getBundledModel("anthropic", "claude-sonnet-4-6"); + if (!primaryModel || !setupTarget) throw new Error("Expected bundled setup-handoff models"); + const requestedModels: string[] = []; + const usageChecks: string[] = []; + const mock = createMockModel({ responses: [{ content: ["must not run"] }] }); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return mock.stream(model, context, options); + }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "fail-closed", + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + const usageHealth = vi + .spyOn(modelRegistry.authStorage, "getModelUsageHealth") + .mockImplementation(async (_provider, options) => { + usageChecks.push(options.modelId ?? ""); + const reserve = options.modelId === setupTarget.id; + return { + state: reserve ? "reserve" : "healthy", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: reserve ? "reserve" : "healthy", + remainingFraction: reserve ? 0.05 : 0.8, + }, + ], + }; + }); + const extensionRunner = { + emit: vi.fn().mockResolvedValue(undefined), + hasHandlers: vi.fn().mockReturnValue(false), + emitBeforeAgentStart: vi.fn(async () => { + if (!session) throw new Error("Expected active session"); + await session.setModelTemporary(setupTarget, undefined, { ephemeral: true }); + return undefined; + }), + } as unknown as ExtensionRunner; + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + extensionRunner, + }); + + await session.prompt("Change models during setup"); + + expect(usageHealth).toHaveBeenCalledTimes(2); + expect(usageChecks).toEqual([primaryModel.id, setupTarget.id]); + expect(requestedModels).toEqual([]); + }); + + it("restarts usage preflight when the model changes during a health request", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + const selectedModel = getBundledModel("anthropic", "claude-sonnet-4-6"); + if (!primaryModel || !selectedModel) throw new Error("Expected bundled preflight race models"); + const requestedModels: string[] = []; + const usageChecks: string[] = []; + const healthStarted = Promise.withResolvers(); + const releaseHealth = Promise.withResolvers(); + const mock = createMockModel({ responses: [{ content: ["must not run"] }] }); + const agent = new Agent({ + getApiKey: model => `${model.provider}-test-key`, + initialState: { model: primaryModel, systemPrompt: ["Test"], tools: [], messages: [] }, + streamFn: (model, context, options) => { + requestedModels.push(`${model.provider}/${model.id}`); + return mock.stream(model, context, options); + }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "fail-closed", + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + const usageHealth = vi + .spyOn(modelRegistry.authStorage, "getModelUsageHealth") + .mockImplementation(async (_provider, options) => { + usageChecks.push(options.modelId ?? ""); + if (options.modelId === primaryModel.id) { + healthStarted.resolve(); + await releaseHealth.promise; + } + const reserve = options.modelId === selectedModel.id; + return { + state: reserve ? "reserve" : "healthy", + accounts: [ + { + credentialId: 1, + credentialType: "oauth", + state: reserve ? "reserve" : "healthy", + remainingFraction: reserve ? 0.05 : 0.8, + }, + ], + }; + }); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + + const prompting = session.prompt("Change models during preflight"); + await healthStarted.promise; + await session.setModelTemporary(selectedModel, undefined, { ephemeral: true }); + releaseHealth.resolve(); + await expect(prompting).rejects.toThrow(`reserve reached for ${selectedModel.provider}/${selectedModel.id}`); + + expect(usageHealth).toHaveBeenCalledTimes(2); + expect(usageChecks).toEqual([primaryModel.id, selectedModel.id]); + expect(session.model?.id).toBe(selectedModel.id); + expect(requestedModels).toEqual([]); + }); + + it("finishes usage preflight when no model is selected", async () => { + const agent = new Agent({ + initialState: { model: undefined, systemPrompt: ["Test"], tools: [], messages: [] }, + }); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + }); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + }); + + await expect(session.prompt("No model configured")).rejects.toThrow("No model selected"); + expect(agent.state.isStreaming).toBe(false); + }); + it("continues a startup-owned role fallback chain from the active fallback", async () => { const firstFallback = getBundledModel("openai", "gpt-4o-mini"); const secondFallback = getBundledModel("openai", "gpt-4o"); @@ -2870,6 +3576,54 @@ describe("AgentSession retry fallback", () => { expect(session.thinkingLevel).toBe(Effort.Low); }); + it("skips usage fallbacks whose effort floor exceeds the session ceiling", async () => { + const primaryModel = getBundledModel("anthropic", "claude-sonnet-4-5"); + const incompatibleFallback = getBundledModel("fireworks", "deepseek-v4-pro"); + const compatibleFallback = getBundledModel("openai", "gpt-4o-mini"); + if (!primaryModel || !incompatibleFallback || !compatibleFallback) { + throw new Error("Expected bundled usage fallback effort models"); + } + const requestedModels: string[] = []; + const usageChecks: string[] = []; + const agent = createFallbackAgent(primaryModel, requestedModels); + const settings = Settings.isolated({ + "compaction.enabled": false, + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "auto", + "retry.fallbackChains": { + default: [ + `${incompatibleFallback.provider}/${incompatibleFallback.id}`, + `${compatibleFallback.provider}/${compatibleFallback.id}`, + ], + }, + }); + settings.setModelRole("default", `${primaryModel.provider}/${primaryModel.id}`); + vi.spyOn(modelRegistry.authStorage, "getModelUsageHealth").mockImplementation(async (_provider, options) => { + usageChecks.push(options.modelId ?? ""); + return options.modelId === primaryModel.id + ? { + state: "depleted", + accounts: [{ credentialId: 1, credentialType: "oauth", state: "depleted" }], + } + : { state: "healthy", accounts: [] }; + }); + session = new AgentSession({ + agent, + sessionManager: SessionManager.inMemory(), + settings, + modelRegistry, + thinkingLevel: Effort.Low, + thinkingLevelCeiling: Effort.Low, + }); + + await session.prompt("Use an effort-compatible fallback"); + await session.waitForIdle(); + + expect(usageChecks).toEqual([primaryModel.id, compatibleFallback.id]); + expect(requestedModels).toEqual([`${compatibleFallback.provider}/${compatibleFallback.id}`]); + expect(session.model?.id).toBe(compatibleFallback.id); + }); + it("accepts cached Ollama Cloud fallback selectors during startup validation", () => { const primaryModel = getBundledModel("openai", "gpt-4o-mini"); if (!primaryModel) { diff --git a/packages/coding-agent/test/issue-905-repro.test.ts b/packages/coding-agent/test/issue-905-repro.test.ts index 8d7d1aa94..e9901670c 100644 --- a/packages/coding-agent/test/issue-905-repro.test.ts +++ b/packages/coding-agent/test/issue-905-repro.test.ts @@ -208,10 +208,8 @@ test("omp models prints invalid models.yml schema errors before listing output", `providers: myprovider: baseUrl: http://localhost:8000/v1 - api: openai-completions + api: invalid-api auth: none - compat: - thinkingFormat: deepseek models: - id: my-model name: My Model @@ -248,8 +246,8 @@ test("omp models prints invalid models.yml schema errors before listing output", const output = captured.join(""); expect(output).toContain("Warning: models.yml validation failed — custom providers disabled"); - expect(output).toContain("providers.myprovider.compat.thinkingFormat"); - expect(output).toContain("deepseek"); + expect(output).toContain("providers.myprovider.api"); + expect(output).toContain("invalid-api"); } finally { authStorage.close(); } diff --git a/packages/coding-agent/test/sdk-model-selection.test.ts b/packages/coding-agent/test/sdk-model-selection.test.ts index 87ea06ffd..14f8fa2ab 100644 --- a/packages/coding-agent/test/sdk-model-selection.test.ts +++ b/packages/coding-agent/test/sdk-model-selection.test.ts @@ -512,6 +512,33 @@ describe("createAgentSession deferred model pattern resolution", () => { } }); + test("rejects a depleted terminal fallback after startup skips the primary", async () => { + const settings = Settings.isolated({ + "retry.usageAwareFallback": true, + "retry.usageReservePolicy": "confirm", + }); + settings.setModelRole("task", "runtime-provider/runtime-model,runtime-provider/runtime-reasoning-model"); + const options = await buildSessionOptions("task"); + const usageHealth = vi.spyOn(options.authStorage, "getModelUsageHealth").mockResolvedValue({ + state: "depleted", + accounts: [{ credentialId: 1, credentialType: "oauth", state: "depleted" }], + }); + + const { session, modelFallbackMessage } = await createAgentSession({ + ...options, + modelPatternFallbackRole: "subagent:usage-aware-terminal", + settings, + hasUI: false, + }); + try { + expect(usageHealth).toHaveBeenCalledTimes(2); + expect(session.model).toBeUndefined(); + expect(modelFallbackMessage).toContain("not found"); + } finally { + await session.dispose(); + } + }); + test("defers ACP reserve fallback until prompt-time capabilities are configured", async () => { const settings = Settings.isolated({ "retry.usageAwareFallback": true, diff --git a/packages/coding-agent/test/turn-recovery-replay-unsafe.test.ts b/packages/coding-agent/test/turn-recovery-replay-unsafe.test.ts index 98f5c1891..2d5f91eee 100644 --- a/packages/coding-agent/test/turn-recovery-replay-unsafe.test.ts +++ b/packages/coding-agent/test/turn-recovery-replay-unsafe.test.ts @@ -99,6 +99,149 @@ describe("TurnRecovery replay-unsafe output classification", () => { tempDir.removeSync(); }); + it("rolls back a usage fallback cancelled during model reconciliation", async () => { + const fallback = getBundledModel("openai", "gpt-4o-mini"); + if (!fallback) throw new Error("Expected bundled fallback model"); + let activeModel = model; + const fallbackApplied = Promise.withResolvers(); + const releaseReconciliation = Promise.withResolvers(); + const modelChanges: string[] = []; + const emittedEvents: string[] = []; + const host = createHost(model, modelRegistry); + host.model = () => activeModel; + host.sessionManager = { + appendModelChange: (selector: string) => modelChanges.push(selector), + } as never; + host.setModelWithProviderSessionReset = async nextModel => { + activeModel = nextModel; + if (nextModel.provider === fallback.provider && nextModel.id === fallback.id) { + fallbackApplied.resolve(); + await releaseReconciliation.promise; + } + }; + host.emitSessionEvent = async event => { + emittedEvents.push(event.type); + }; + const recovery = new TurnRecovery(host); + const controller = new AbortController(); + const applying = recovery.applyRetryFallbackCandidate( + "default", + { + raw: `${fallback.provider}/${fallback.id}`, + provider: fallback.provider, + id: fallback.id, + thinkingLevel: undefined, + }, + `${model.provider}/${model.id}`, + { pinFallback: true, apiKey: "test-key", signal: controller.signal }, + ); + + await fallbackApplied.promise; + controller.abort(); + releaseReconciliation.resolve(); + const committed = await applying; + + expect(committed).toBe(false); + expect(activeModel).toBe(model); + expect(modelChanges).toEqual([]); + expect(emittedEvents).toEqual([]); + }); + + it("does not commit a fallback superseded during model reconciliation", async () => { + const fallback = getBundledModel("openai", "gpt-4o-mini"); + if (!fallback) throw new Error("Expected bundled fallback race model"); + const selectedModel = { ...fallback, baseUrl: "https://user-selected-route.example" }; + let activeModel = model; + const fallbackApplied = Promise.withResolvers(); + const releaseReconciliation = Promise.withResolvers(); + const modelChanges: string[] = []; + const emittedEvents: string[] = []; + const thinkingChanges: unknown[] = []; + const host = createHost(model, modelRegistry); + host.model = () => activeModel; + host.sessionManager = { + appendModelChange: (selector: string) => modelChanges.push(selector), + } as never; + host.setThinkingLevel = level => thinkingChanges.push(level); + host.setModelWithProviderSessionReset = async nextModel => { + activeModel = nextModel; + if (nextModel.provider === fallback.provider && nextModel.id === fallback.id) { + fallbackApplied.resolve(); + await releaseReconciliation.promise; + } + }; + host.emitSessionEvent = async event => { + emittedEvents.push(event.type); + }; + const recovery = new TurnRecovery(host); + const applying = recovery.applyRetryFallbackCandidate( + "default", + { + raw: `${fallback.provider}/${fallback.id}`, + provider: fallback.provider, + id: fallback.id, + thinkingLevel: undefined, + }, + `${model.provider}/${model.id}`, + { pinFallback: true, apiKey: "test-key" }, + ); + + await fallbackApplied.promise; + activeModel = selectedModel; + releaseReconciliation.resolve(); + const committed = await applying; + + expect(committed).toBe(false); + expect(activeModel).toBe(selectedModel); + expect(modelChanges).toEqual([]); + expect(thinkingChanges).toEqual([]); + expect(emittedEvents).toEqual([]); + }); + it("keeps a committed fallback when cancellation arrives during applied-event delivery", async () => { + const fallback = getBundledModel("openai", "gpt-4o-mini"); + if (!fallback) throw new Error("Expected bundled fallback model"); + let activeModel = model; + const eventStarted = Promise.withResolvers(); + const releaseEvent = Promise.withResolvers(); + const modelChanges: string[] = []; + const host = createHost(model, modelRegistry); + host.model = () => activeModel; + host.sessionManager = { + appendModelChange: (selector: string) => modelChanges.push(selector), + } as never; + host.setModelWithProviderSessionReset = async nextModel => { + activeModel = nextModel; + }; + host.emitSessionEvent = async event => { + if (event.type !== "retry_fallback_applied") return; + eventStarted.resolve(); + await releaseEvent.promise; + }; + const recovery = new TurnRecovery(host); + const controller = new AbortController(); + const applying = recovery.applyRetryFallbackCandidate( + "default", + { + raw: `${fallback.provider}/${fallback.id}`, + provider: fallback.provider, + id: fallback.id, + thinkingLevel: undefined, + }, + `${model.provider}/${model.id}`, + { pinFallback: true, apiKey: "test-key", signal: controller.signal }, + ); + + await eventStarted.promise; + controller.abort(); + releaseEvent.resolve(); + const committed = await applying; + + expect(committed).toBe(true); + expect(activeModel.provider).toBe(fallback.provider); + expect(activeModel.id).toBe(fallback.id); + expect(modelChanges).toEqual([`${fallback.provider}/${fallback.id}`]); + }); + it("treats a failed turn with partial non-whitespace text as NOT retriable", () => { const recovery = new TurnRecovery(createHost(model, modelRegistry)); const message = makeMessage([{ type: "text", text: "Here is the first part of my answer" }], model);