fix(coding-agent): complete usage-aware fallback integration

This commit is contained in:
Brent
2026-08-03 16:33:46 +00:00
parent 01c1f91ff5
commit db97103c32
20 changed files with 1369 additions and 241 deletions
+1
View File
@@ -5,6 +5,7 @@
## [17.2.6] - 2026-08-03
### 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.
- Fixed an issue where peer-IRC interrupts (such as subagent messages) incorrectly skipped non-interruptible tool calls queued in the same batch.
- Improved interruption messaging to clearly distinguish between parent-agent steering and system-advisory interruptions.
+5 -5
View File
@@ -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];
+79 -20
View File
@@ -426,6 +426,8 @@ export class Agent {
#asideMessageProvider?: () => AsideMessage[] | Promise<AsideMessage[]>;
#telemetry?: AgentLoopConfig["telemetry"];
#appendOnlyContext?: AppendOnlyContextManager;
#beforeQueuedMessageDequeueHooks = new Set<(signal?: AbortSignal) => Promise<void> | void>();
#beforeModelCallHooks = new Set<(signal?: AbortSignal) => Promise<void> | 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): () => 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): () => void {
const registration = (signal?: AbortSignal) => hook(signal);
this.#beforeModelCallHooks.add(registration);
return () => this.#beforeModelCallHooks.delete(registration);
}
async #runBeforeModelCallHooks(signal?: AbortSignal): Promise<void> {
for (const hook of this.#beforeModelCallHooks) await hook(signal);
}
async #runBeforeQueuedMessageDequeueHooks(signal?: AbortSignal): Promise<void> {
for (const hook of this.#beforeQueuedMessageDequeueHooks) await hook(signal);
}
async #dequeueSteeringMessagesAfterHooks(signal?: AbortSignal): Promise<AgentMessage[]> {
if (signal?.aborted || this.#steeringQueue.length === 0) return [];
await this.#runBeforeQueuedMessageDequeueHooks(signal);
return signal?.aborted ? [] : this.#dequeueSteeringMessages();
}
async #dequeueFollowUpMessagesAfterHooks(signal?: AbortSignal): Promise<AgentMessage[]> {
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,11 +1173,26 @@ export class Agent {
/**
* Continue from current context (used for retries and resuming queued messages).
*/
async continue() {
#continuationDequeueSignal(signal?: AbortSignal): AbortSignal | undefined {
if (this.#deadline === undefined) return signal;
const delay = this.#deadline - Date.now();
let deadlineSignal: AbortSignal;
if (delay <= 0) {
const controller = new AbortController();
controller.abort(new DOMException("Deadline exceeded", "TimeoutError"));
deadlineSignal = controller.signal;
} else {
deadlineSignal = AbortSignal.timeout(delay);
}
return signal ? AbortSignal.any([signal, deadlineSignal]) : deadlineSignal;
}
async continue(signal?: AbortSignal) {
if (this.#state.isStreaming) {
throw new AgentBusyError();
}
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
@@ -1150,35 +1201,35 @@ export class Agent {
// 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();
const queuedSteering = await this.#dequeueSteeringMessagesAfterHooks(dequeueSignal);
if (queuedSteering.length > 0) {
await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true });
await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true }, signal);
return;
}
const queuedFollowUp = this.#dequeueFollowUpMessages();
const queuedFollowUp = await this.#dequeueFollowUpMessagesAfterHooks(dequeueSignal);
if (queuedFollowUp.length > 0) {
await this.#runLoop(queuedFollowUp);
await this.#runLoop(queuedFollowUp, undefined, signal);
return;
}
throw new Error("No messages to continue from");
}
if (messages[messages.length - 1].role === "assistant") {
const queuedSteering = this.#dequeueSteeringMessages();
const queuedSteering = await this.#dequeueSteeringMessagesAfterHooks(dequeueSignal);
if (queuedSteering.length > 0) {
await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true });
await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true }, signal);
return;
}
const queuedFollowUp = this.#dequeueFollowUpMessages();
const queuedFollowUp = await this.#dequeueFollowUpMessagesAfterHooks(dequeueSignal);
if (queuedFollowUp.length > 0) {
await this.#runLoop(queuedFollowUp);
await this.#runLoop(queuedFollowUp, undefined, signal);
return;
}
throw new Error("Cannot continue from message role: assistant");
}
await this.#runLoop(undefined);
await this.#runLoop(undefined, undefined, signal);
}
/**
@@ -1186,7 +1237,11 @@ 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,
) {
const model = this.#state.model;
if (!model) throw new Error("No model configured");
@@ -1197,6 +1252,9 @@ export class Agent {
this.#resolveRunningPrompt = resolve;
this.#abortController = new AbortController();
const loopSignal = continuationSignal
? AbortSignal.any([this.#abortController.signal, continuationSignal])
: this.#abortController.signal;
this.#state.isStreaming = true;
this.#state.streamMessage = null;
this.#state.error = undefined;
@@ -1315,7 +1373,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 +1421,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 +1451,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 +1463,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 +1531,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);
+3 -3
View File
@@ -240,7 +240,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
* mid-batch interrupt poll uses {@link hasSteeringMessages} instead and
* never consumes the queue.
*/
getSteeringMessages?: () => Promise<AgentMessage[]>;
getSteeringMessages?: (signal?: AbortSignal) => Promise<AgentMessage[]>;
/**
* 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<AgentMessage[]>;
getFollowUpMessages?: (signal?: AbortSignal) => Promise<AgentMessage[]>;
/**
* 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<void>;
syncContextBeforeModelCall?: (context: AgentContext, signal?: AbortSignal) => void | Promise<void>;
/**
* Asked after the complete provider context has been built, including
+117 -1
View File
@@ -240,6 +240,106 @@ 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<AbortSignal | undefined> = [];
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<void>();
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<void>();
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("classifies an in-flight continuation cancellation as aborted", async () => {
const providerStarted = Promise.withResolvers<AbortSignal>();
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 +376,12 @@ describe("Agent", () => {
responses: [{ content: ["Processed 1"] }, { content: ["Processed 2"] }],
});
const agent = new Agent({ streamFn: mock.stream });
let dequeueHooks = 0;
const dequeueSignals: Array<AbortSignal | undefined> = [];
agent.addBeforeQueuedMessageDequeueHook(signal => {
dequeueHooks++;
dequeueSignals.push(signal);
});
agent.replaceMessages([
{
@@ -297,11 +403,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 +967,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 +990,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 () => {
+1
View File
@@ -9,6 +9,7 @@
- Added profile-aware Bedrock Mantle region selection, authenticated model discovery, bearer-token or SigV4 authentication, and credential refresh handling for OpenAI Responses models.
### Fixed
- 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.
- Fixed an issue where Ollama requests without a user-role message would fail to generate output or silently fail with a misleading error.
+12 -1
View File
@@ -3856,13 +3856,15 @@ export class AuthStorage {
: undefined;
const rankingContext: CredentialRankingContext = { modelId: options.modelId };
const planRequirement = resolveOpenAICodexPlanRequirement(provider, options.modelId);
const planEligibilityByCredential = new Map<number, boolean | undefined>();
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<ModelUsageAccountHealth> => {
const credentialType = entry.credential.type;
const providerKey = this.#getProviderTypeKey(provider, credentialType);
@@ -3890,6 +3892,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);
@@ -3942,6 +3947,12 @@ export class AuthStorage {
};
}),
);
if (
planRequirement !== "none" &&
accounts.some(account => planEligibilityByCredential.get(account.credentialId) === true)
) {
accounts = accounts.filter(account => planEligibilityByCredential.get(account.credentialId) === true);
}
if (selectedCredentialId !== undefined) {
const selectedAccount = accounts.find(account => account.credentialId === selectedCredentialId);
if (selectedAccount) selectedAccount.selected = true;
@@ -2027,6 +2027,50 @@ 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("routes codex spark to a single Plus account when no Pro is connected", async () => {
if (!authStorage) throw new Error("test setup failed");
+1
View File
@@ -9,6 +9,7 @@
- Added the `/reset` slash command to reset the conversation context in place: it drops the live messages, queued turns, and pending tool calls (and cancels the turn's async jobs, post-prompt continuations, and checkpoint/plan runtime state) while keeping the session id, title, cwd, model, and on-disk transcript. It records a durable reset boundary so the live transcript stays cleared across rebuilds (theme change, focus attach, `/shake`, resume) instead of resurrecting the pre-reset messages, while the full pre-reset history stays on disk ([#3580](https://github.com/can1357/oh-my-pi/issues/3580)).
### Fixed
- 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.
- Fixed extension slash commands appearing as user prompts after being handled locally.
- Preserved explicit session titles when branching from an earlier conversation turn.
@@ -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 },
);
});
}
@@ -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 },
);
});
+11 -5
View File
@@ -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;
}
@@ -92,6 +92,8 @@ export interface UsageFallbackConfirmation {
remainingPercent: number | undefined;
}
export type UsageFallbackConfirmer = (confirmation: UsageFallbackConfirmation, signal: AbortSignal) => Promise<boolean>;
/** Identifies a retry fallback chain already entered during startup model resolution. */
export interface InitialRetryFallbackState {
/** Role whose configured primary was unavailable. */
+129 -181
View File
@@ -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";
@@ -558,9 +556,13 @@ export class AgentSession {
// Model registry for API key resolution
#modelRegistry: ModelRegistry;
#usageFallbackConfirmer: ((confirmation: UsageFallbackConfirmation) => Promise<boolean>) | undefined;
#usageReserveApprovedSelector: string | undefined;
#usageFallbackConfirmer: UsageFallbackConfirmer | undefined;
#usagePreflightAbortControllers = new Set<AbortController>();
#queuedMessageDrainBlocked = false;
#usagePreflightReadyForNextModelCall = false;
#usagePreflightReadyModel: Model | undefined;
#detachUsageBeforeQueueDequeue: (() => void) | undefined;
#detachUsageBeforeModelCall: (() => void) | undefined;
#transformContext: (messages: AgentMessage[], signal?: AbortSignal) => AgentMessage[] | Promise<AgentMessage[]>;
#onPayload: SimpleStreamOptions["onPayload"] | undefined;
@@ -768,8 +770,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<void>) | undefined;
try {
@@ -805,6 +809,7 @@ export class AgentSession {
[...this.agent.peekSteeringQueue()],
[...parkedFollowUps, ...this.agent.peekFollowUpQueue()],
);
this.#queuedMessageDrainBlocked ||= parkedQueueDrainBlocked;
}
this.#endInFlight(async () => {
try {
@@ -830,6 +835,7 @@ export class AgentSession {
steering.filter(m => !isAdvisorCard(m)),
followUp.filter(m => !isAdvisorCard(m)),
);
this.#reconcileQueuedMessageDrain();
return cards;
}
@@ -1050,6 +1056,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,
@@ -2967,19 +2998,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),
@@ -2987,6 +3016,7 @@ export class AgentSession {
});
options?.onError?.(error);
} finally {
this.#usagePreflightReadyForNextModelCall = false;
this.#endInFlight();
}
},
@@ -3605,6 +3635,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();
@@ -3927,176 +3963,63 @@ export class AgentSession {
}
/** Install the interactive decision surface for reserve-triggered model changes. */
setUsageFallbackConfirmer(
confirmer: ((confirmation: UsageFallbackConfirmation) => Promise<boolean>) | undefined,
): void {
setUsageFallbackConfirmer(confirmer: UsageFallbackConfirmer | undefined): void {
this.#usageFallbackConfirmer = confirmer;
}
async #runUsageAwarePreflight(): Promise<boolean> {
#allowQueuedMessageDrainRetry(): void {
this.#queuedMessageDrainBlocked = false;
}
#reconcileQueuedMessageDrain(): void {
if (!this.agent.hasQueuedMessages()) {
this.#queuedMessageDrainBlocked = false;
}
}
async #runQueuedUsageAwarePreflight(signal?: AbortSignal): Promise<boolean> {
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<boolean> {
const allowed = await this.#runUsageAwarePreflight(signal);
this.#usagePreflightReadyForNextModelCall = allowed;
this.#usagePreflightReadyModel = allowed ? this.model : undefined;
return allowed;
}
async #runUsageAwarePreflight(signal?: AbortSignal): Promise<boolean> {
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;
const fallbackCommitted = await this.#recovery.maybeApplyUsageAwareFallback(
controller.signal,
this.#usageFallbackConfirmer,
);
return fallbackCommitted || (!controller.signal.aborted && this.#promptGeneration === generation);
} catch (error) {
if (controller.signal.aborted || this.#promptGeneration !== generation) return false;
throw error;
} finally {
signal?.removeEventListener("abort", onAbort);
this.#usagePreflightAbortControllers.delete(controller);
}
}
async #confirmUsageFallback(confirmation: UsageFallbackConfirmation, signal: AbortSignal): Promise<boolean> {
const confirmer = this.#usageFallbackConfirmer;
if (!confirmer || signal.aborted) return false;
const aborted = Promise.withResolvers<boolean>();
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<void> {
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<number | undefined>(
(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;
@@ -5027,7 +4950,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) {
@@ -5122,7 +5045,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);
}
@@ -5132,7 +5055,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);
}
@@ -5169,7 +5092,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();
@@ -5380,6 +5303,7 @@ export class AgentSession {
await this.#waitForPostPromptRecovery(generation);
}
} finally {
this.#usagePreflightReadyForNextModelCall = false;
this.#endInFlight();
}
}
@@ -5545,7 +5469,6 @@ export class AgentSession {
}
const expandedText = expandPromptTemplate(text, [...this.#promptTemplates]);
if (!(await this.#runUsageAwarePreflight())) return;
await this.#queueUserMessage(expandedText, images, "steer");
}
@@ -5563,7 +5486,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;
@@ -5580,6 +5502,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",
@@ -5609,6 +5532,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({
@@ -5635,7 +5559,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;
@@ -5649,6 +5578,7 @@ export class AgentSession {
},
onError: () => {
this.#queuedMessageDrainScheduled = false;
this.#queuedMessageDrainBlocked = this.agent.hasQueuedMessages();
},
});
}
@@ -5785,7 +5715,7 @@ export class AgentSession {
): Promise<void> {
this.#beginInFlight();
try {
if (!(await this.#runUsageAwarePreflight())) return;
if (!(await this.#runUsageAwarePreflightForNextModelCall())) return;
const acceptTerminalEmptyStop = options?.acceptTerminalEmptyStop === true;
if (acceptTerminalEmptyStop) {
this.#resetPromptMaintenanceState();
@@ -5794,6 +5724,7 @@ export class AgentSession {
await this.agent.prompt(message);
await this.#waitForPostPromptRecovery();
} finally {
this.#usagePreflightReadyForNextModelCall = false;
this.#recovery.setAcceptTerminalEmptyStop(false);
this.#endInFlight();
}
@@ -5825,6 +5756,7 @@ export class AgentSession {
timestamp: Date.now(),
};
const normalizedAppMessage = await this.#normalizeAgentMessageImages(appMessage);
this.#allowQueuedMessageDrainRetry();
if (deliverAs === "followUp") {
this.agent.followUp(normalizedAppMessage);
} else {
@@ -5881,7 +5813,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);
@@ -5964,8 +5896,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;
@@ -6005,6 +5935,7 @@ export class AgentSession {
? isAdvisorCard
: m => !isUserQueuedMessage(m) && !isHiddenUserCompanion(m);
this.agent.replaceQueues(steeringAll.filter(keep), followUpAll.filter(keep));
this.#reconcileQueuedMessageDrain();
return { steering, followUp };
}
@@ -6054,12 +5985,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;
@@ -6349,6 +6282,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());
@@ -7373,6 +7308,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;
@@ -7397,6 +7335,9 @@ export class AgentSession {
this.agent.clearAllQueues();
this.#pendingNextTurnMessages = [];
this.#scheduledHiddenNextTurnGeneration = undefined;
this.#queuedMessageDrainBlocked = false;
this.#usagePreflightReadyForNextModelCall = false;
this.#usagePreflightReadyModel = undefined;
try {
if (switchingToDifferentSession) {
@@ -7565,6 +7506,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;
@@ -7653,6 +7597,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
@@ -7780,6 +7726,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();
@@ -13,6 +13,7 @@ import type {
CodexCompactionContext,
Effort,
Model,
ModelUsageHealth,
TextContent,
ToolChoice,
} from "@oh-my-pi/pi-ai";
@@ -23,6 +24,7 @@ import { isFireworksFastModelId, toFireworksBaseModelId } from "@oh-my-pi/pi-cat
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" };
@@ -35,7 +37,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,
@@ -64,6 +70,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);
@@ -172,6 +179,7 @@ export class TurnRecovery {
#retryPromise: Promise<void> | undefined;
#retryResolve: (() => void) | undefined;
#activeRetryFallback: ActiveRetryFallbackState | undefined;
#usageReserveApprovedSelector: string | undefined;
#pendingRecoveredRetryErrors: PendingRecoveredRetryError[] = [];
#usageLimitOutcomes = new WeakMap<AssistantMessage, Promise<UsageLimitOutcome>>();
#emptyStopRetryCount = 0;
@@ -298,6 +306,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<boolean> {
return this.#maybeApplyUsageAwareFallback(signal, confirmer);
}
/** Applies automatic retry, credential rotation, and model fallback policy. */
handleRetryableError(
message: AssistantMessage,
@@ -840,6 +853,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
@@ -883,6 +899,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
@@ -1008,7 +1025,6 @@ export class TurnRecovery {
modelLookup: this.#host.modelRegistry,
};
}
#getRetryFallbackChains(): RetryFallbackChains {
return getRetryFallbackChains(this.#host.settings);
}
@@ -1071,23 +1087,185 @@ export class TurnRecovery {
);
}
async #maybeApplyUsageAwareFallback(signal: AbortSignal, confirmer?: UsageFallbackConfirmer): Promise<boolean> {
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) return false;
logger.debug("Usage-aware runtime preflight failed open", {
provider: currentModel.provider,
model: currentModel.id,
error: String(error),
});
return false;
}
if (signal.aborted) 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;
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;
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) 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) return false;
// Unknown usage fails open for an otherwise valid fallback.
}
if (signal.aborted) return false;
let apiKey: string | undefined;
try {
apiKey = await this.#host.modelRegistry.getApiKey(candidateModel, this.#host.sessionId(), { signal });
} catch {
if (signal.aborted) return false;
continue;
}
if (signal.aborted) 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<number | undefined>((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) 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<boolean> {
if (signal.aborted) return false;
const aborted = Promise.withResolvers<boolean>();
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<void> {
): Promise<boolean> {
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.
@@ -1101,7 +1279,14 @@ 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;
}
this.#host.sessionManager.appendModelChange(candidateSelector, EPHEMERAL_MODEL_CHANGE_ROLE);
this.#host.settings.getStorage()?.recordModelUsage(candidateSelector);
this.#host.setThinkingLevel(nextThinkingLevel);
@@ -1123,6 +1308,7 @@ export class TurnRecovery {
to: selector.raw,
role,
});
return true;
}
async #tryRetryModelFallback(currentSelector: string, options?: { pinFallback?: boolean }): Promise<boolean> {
@@ -1140,8 +1326,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;
@@ -1167,6 +1352,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.
@@ -1190,6 +1376,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");
+13 -4
View File
@@ -166,9 +166,11 @@ function resolveSubagentRetryFallbackCandidates(
): SubagentRetryFallbackCandidate[] {
const candidates: SubagentRetryFallbackCandidate[] = [];
const seen = new Set<string>();
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);
@@ -179,7 +181,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) ||
@@ -188,7 +193,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: {
@@ -2766,7 +2775,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
const configuredModelPatterns = resolveConfiguredModelPatterns(modelPatterns, settings);
const defaultRetryFallbackChain =
configuredModelPatterns.length === 1
? resolveSubagentDefaultRetryFallbackChain(subagentSettings)
? resolveSubagentDefaultRetryFallbackChain(subagentSettings, modelRegistry)
: undefined;
const {
model,
@@ -2800,7 +2809,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
const retryFallbackRole = installSubagentRetryFallbackChain({
settings: subagentSettings,
id,
candidates: resolveSubagentRetryFallbackCandidates(modelPatterns, modelRegistry, settings),
candidates: resolveSubagentRetryFallbackCandidates(modelPatterns, modelRegistry, subagentSettings),
defaultFallbackChain: defaultRetryFallbackChain,
model,
authFallbackUsed,
@@ -1,13 +1,14 @@
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test";
import * as path from "node:path";
import { scheduler } from "node:timers/promises";
import { Agent } from "@oh-my-pi/pi-agent-core";
import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core";
import {
type AssistantMessage,
Effort,
type Model,
type ModelUsageHealth,
type ProviderSessionState,
z,
} from "@oh-my-pi/pi-ai";
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
import { buildModel } from "@oh-my-pi/pi-catalog/build";
@@ -96,6 +97,7 @@ describe("AgentSession retry fallback", () => {
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,565 @@ 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<void>();
const pendingConfirmation = Promise.withResolvers<boolean>();
const confirmationAborted = Promise.withResolvers<void>();
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<void>();
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<void>();
const pendingConfirmation = Promise.withResolvers<boolean>();
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<void>();
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<void>();
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<typeof toolSchema, { value: string }> = {
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("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");
@@ -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();
}
@@ -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,
@@ -95,6 +95,98 @@ 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<void>();
const releaseReconciliation = Promise.withResolvers<void>();
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("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<void>();
const releaseEvent = Promise.withResolvers<void>();
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);