fix(FN-745): improve fallback model recovery

This commit is contained in:
gsxdsm
2026-04-02 21:50:41 -07:00
parent 160a370117
commit f9dd3983c0
25 changed files with 949 additions and 129 deletions

View File

@@ -5,6 +5,13 @@ import { AgentSemaphore } from "./concurrency.js";
vi.mock("./pi.js", () => ({
createKbAgent: vi.fn(),
describeModel: vi.fn().mockReturnValue("mock-provider/mock-model"),
promptWithFallback: vi.fn(async (session, prompt, options) => {
if (options === undefined) {
await session.prompt(prompt);
} else {
await session.prompt(prompt, options);
}
}),
}));
vi.mock("./reviewer.js", () => ({
reviewStep: vi.fn(),

View File

@@ -5,7 +5,7 @@ import type { TaskStore, Task, TaskDetail, StepStatus, Settings, WorkflowStep, M
import { findWorktreeUser } from "./merger.js";
import { generateWorktreeName, slugify } from "./worktree-names.js";
import { Type, type Static } from "@mariozechner/pi-ai";
import { createKbAgent, describeModel } from "./pi.js";
import { createKbAgent, describeModel, promptWithFallback } from "./pi.js";
import { reviewStep, type ReviewVerdict } from "./reviewer.js";
import type { ToolDefinition, AgentSession, SessionManager } from "@mariozechner/pi-coding-agent";
import { PRIORITY_EXECUTE, type AgentSemaphore } from "./concurrency.js";
@@ -578,6 +578,8 @@ export class TaskExecutor {
const executorModelId = detail.modelProvider && detail.modelId
? detail.modelId
: settings.defaultModelId;
const executorFallbackProvider = settings.fallbackProvider;
const executorFallbackModelId = settings.fallbackModelId;
const { session } = await createKbAgent({
cwd: worktreePath,
@@ -590,6 +592,8 @@ export class TaskExecutor {
onToolEnd: agentLogger.onToolEnd,
defaultProvider: executorProvider,
defaultModelId: executorModelId,
fallbackProvider: executorFallbackProvider,
fallbackModelId: executorFallbackModelId,
defaultThinkingLevel: settings.defaultThinkingLevel,
});
@@ -616,7 +620,7 @@ export class TaskExecutor {
const agentPrompt = buildExecutionPrompt(detail, this.rootDir, settings);
// Record activity on prompt start (heartbeat for stuck detection)
stuckDetector?.recordActivity(task.id);
await session.prompt(agentPrompt);
await promptWithFallback(session, agentPrompt);
// Re-raise errors that pi-coding-agent swallowed after exhausting retries.
// session.prompt() resolves normally even when retries are exhausted —
@@ -1014,10 +1018,14 @@ export class TaskExecutor {
onText: (delta) => options.onAgentText?.(taskId, delta),
defaultProvider: settings.defaultProvider,
defaultModelId: settings.defaultModelId,
fallbackProvider: settings.fallbackProvider,
fallbackModelId: settings.fallbackModelId,
defaultThinkingLevel: settings.defaultThinkingLevel,
// Per-task validator overrides take precedence over global validator settings
validatorModelProvider: detail.validatorModelProvider ?? settings.validatorProvider,
validatorModelId: detail.validatorModelId ?? settings.validatorModelId,
validatorFallbackModelProvider: settings.validatorFallbackProvider,
validatorFallbackModelId: settings.validatorFallbackModelId,
store,
taskId,
},
@@ -1369,6 +1377,8 @@ If issues are found that need attention, describe them clearly.`;
tools: "readonly",
defaultProvider: settings.defaultProvider,
defaultModelId: settings.defaultModelId,
fallbackProvider: settings.fallbackProvider,
fallbackModelId: settings.fallbackModelId,
defaultThinkingLevel: settings.defaultThinkingLevel,
});
@@ -1394,7 +1404,8 @@ If issues are found that need attention, describe them clearly.`;
}
});
await session.prompt(
await promptWithFallback(
session,
`Execute the workflow step "${workflowStep.name}" for task ${task.id}.\n\n` +
`Review the work done in this worktree and evaluate it against the criteria in your instructions.`,
);

View File

@@ -3,6 +3,13 @@ import { describe, it, expect, vi, beforeEach } from "vitest";
// Mock external dependencies
vi.mock("./pi.js", () => ({
createKbAgent: vi.fn(),
promptWithFallback: vi.fn(async (session, prompt, options) => {
if (options === undefined) {
await session.prompt(prompt);
} else {
await session.prompt(prompt, options);
}
}),
}));
vi.mock("node:child_process", () => ({

View File

@@ -1,7 +1,7 @@
import { execSync } from "node:child_process";
import { existsSync, readFileSync } from "node:fs";
import type { TaskStore, Task, MergeResult } from "@fusion/core";
import { createKbAgent } from "./pi.js";
import { createKbAgent, promptWithFallback } from "./pi.js";
import type { WorktreePool } from "./worktree-pool.js";
import { AgentLogger } from "./agent-logger.js";
import { mergerLog } from "./logger.js";
@@ -1155,7 +1155,7 @@ async function runAiAgentForCommit(params: AiAgentParams): Promise<{ success: bo
buildCommand,
});
await withRateLimitRetry(async () => {
await session.prompt(prompt);
await promptWithFallback(session, prompt);
checkSessionError(session);
}, {
onRetry: (attempt, delayMs, error) => {

View File

@@ -22,6 +22,24 @@ export interface AgentResult {
session: AgentSession;
}
export interface PromptableSession extends AgentSession {
promptWithFallback: (prompt: string, options?: unknown) => Promise<void>;
}
export async function promptWithFallback(session: AgentSession, prompt: string, options?: unknown): Promise<void> {
const maybePromptable = session as Partial<PromptableSession>;
if (typeof maybePromptable.promptWithFallback === "function") {
await maybePromptable.promptWithFallback(prompt, options);
return;
}
if (options === undefined) {
await session.prompt(prompt);
} else {
await (session.prompt as any)(prompt, options);
}
}
/**
* Extract a human-readable model description from an AgentSession.
* Returns `"<provider>/<modelId>"` (e.g. `"anthropic/claude-sonnet-4-5"`)
@@ -46,10 +64,26 @@ export interface AgentOptions {
defaultProvider?: string;
/** Default model ID within the provider (e.g. "claude-sonnet-4-5"). Used with `defaultProvider`. */
defaultModelId?: string;
/** Optional fallback model provider used when the primary selected model hits
* a retryable provider-side failure such as rate limiting or overload. */
fallbackProvider?: string;
/** Optional fallback model ID used with `fallbackProvider`. */
fallbackModelId?: string;
/** Default thinking effort level (e.g. "medium", "high"). When provided, sets the session's thinking level after creation. */
defaultThinkingLevel?: string;
}
function isRetryableModelSelectionError(message: string): boolean {
const normalized = message.toLowerCase();
return normalized.includes("rate limit")
|| normalized.includes("too many requests")
|| normalized.includes("429")
|| normalized.includes("overloaded")
|| normalized.includes("quota")
|| normalized.includes("capacity")
|| normalized.includes("temporarily unavailable");
}
/**
* Create a pi agent session configured for kb.
* Reuses the user's existing pi auth and model configuration.
@@ -72,6 +106,9 @@ export async function createKbAgent(options: AgentOptions): Promise<AgentResult>
const selectedModel = options.defaultProvider && options.defaultModelId
? modelRegistry.find(options.defaultProvider, options.defaultModelId)
: undefined;
const fallbackModel = options.fallbackProvider && options.fallbackModelId
? modelRegistry.find(options.fallbackProvider, options.fallbackModelId)
: undefined;
const resourceLoader = new DefaultResourceLoader({
cwd: options.cwd,
@@ -81,25 +118,98 @@ export async function createKbAgent(options: AgentOptions): Promise<AgentResult>
});
await resourceLoader.reload();
const { session } = await createAgentSession({
cwd: options.cwd,
authStorage,
modelRegistry,
resourceLoader,
tools,
customTools: options.customTools,
sessionManager: SessionManager.inMemory(),
settingsManager,
...(selectedModel ? { model: selectedModel } : {}),
});
const createSessionWithModel = async (modelOverride?: typeof selectedModel) => {
return createAgentSession({
cwd: options.cwd,
authStorage,
modelRegistry,
resourceLoader,
tools,
customTools: options.customTools,
sessionManager: SessionManager.inMemory(),
settingsManager,
...(modelOverride ? { model: modelOverride } : {}),
});
};
let sessionResult;
let usingFallback = false;
try {
sessionResult = await createSessionWithModel(selectedModel);
} catch (err: any) {
if (!fallbackModel || !selectedModel || !isRetryableModelSelectionError(err?.message || "")) {
throw err;
}
usingFallback = true;
sessionResult = await createSessionWithModel(fallbackModel);
}
const { session } = sessionResult;
const promptableSession = session as PromptableSession;
promptableSession.promptWithFallback = async (prompt: string, promptOptions?: unknown) => {
try {
if (promptOptions === undefined) {
await session.prompt(prompt);
} else {
await (session.prompt as any)(prompt, promptOptions);
}
return;
} catch (err: any) {
if (!fallbackModel || usingFallback || !isRetryableModelSelectionError(err?.message || "")) {
throw err;
}
usingFallback = true;
try {
session.dispose();
} catch {
// ignore dispose errors while swapping sessions
}
const fallbackSessionResult = await createSessionWithModel(fallbackModel);
const fallbackSession = fallbackSessionResult.session as PromptableSession;
if (options.defaultThinkingLevel) {
fallbackSession.setThinkingLevel(options.defaultThinkingLevel as any);
}
fallbackSession.subscribe((event) => {
if (event.type === "message_update") {
const msgEvent = event.assistantMessageEvent;
if (msgEvent.type === "text_delta") {
options.onText?.(msgEvent.delta);
} else if (msgEvent.type === "thinking_delta") {
options.onThinking?.(msgEvent.delta);
}
}
if (event.type === "tool_execution_start") {
options.onToolStart?.(event.toolName, event.args as Record<string, unknown> | undefined);
}
if (event.type === "tool_execution_end") {
options.onToolEnd?.(event.toolName, event.isError, event.result);
}
});
Object.setPrototypeOf(promptableSession, Object.getPrototypeOf(fallbackSession));
Object.assign(promptableSession, fallbackSession);
promptableSession.promptWithFallback = fallbackSession.promptWithFallback ?? promptableSession.promptWithFallback;
if (promptOptions === undefined) {
await fallbackSession.prompt(prompt);
} else {
await (fallbackSession.prompt as any)(prompt, promptOptions);
}
}
};
// Apply thinking level if specified
if (options.defaultThinkingLevel) {
session.setThinkingLevel(options.defaultThinkingLevel as any);
promptableSession.setThinkingLevel(options.defaultThinkingLevel as any);
}
// Wire up event listeners
session.subscribe((event) => {
promptableSession.subscribe((event) => {
if (event.type === "message_update") {
const msgEvent = event.assistantMessageEvent;
if (msgEvent.type === "text_delta") {
@@ -116,5 +226,5 @@ export async function createKbAgent(options: AgentOptions): Promise<AgentResult>
}
});
return { session };
return { session: promptableSession };
}

View File

@@ -17,6 +17,13 @@ import { AgentSemaphore } from "./concurrency.js";
vi.mock("./pi.js", () => ({
createKbAgent: vi.fn(),
describeModel: vi.fn().mockReturnValue("mock-provider/mock-model"),
promptWithFallback: vi.fn(async (session, prompt, options) => {
if (options === undefined) {
await session.prompt(prompt);
} else {
await session.prompt(prompt, options);
}
}),
}));
vi.mock("./reviewer.js", () => ({
reviewStep: vi.fn(),

View File

@@ -3,6 +3,13 @@ import { describe, it, expect, vi, beforeEach } from "vitest";
vi.mock("./pi.js", () => ({
createKbAgent: vi.fn(),
describeModel: vi.fn().mockReturnValue("mock-provider/mock-model"),
promptWithFallback: vi.fn(async (session, prompt, options) => {
if (options === undefined) {
await session.prompt(prompt);
} else {
await session.prompt(prompt, options);
}
}),
}));
import { reviewStep } from "./reviewer.js";

View File

@@ -9,7 +9,7 @@
*/
import type { TaskStore } from "@fusion/core";
import { createKbAgent, describeModel } from "./pi.js";
import { createKbAgent, describeModel, promptWithFallback } from "./pi.js";
import { AgentLogger } from "./agent-logger.js";
import { reviewerLog } from "./logger.js";
import { checkSessionError } from "./usage-limit-detector.js";
@@ -149,6 +149,14 @@ export interface ReviewOptions {
validatorModelProvider?: string;
/** Validator model ID override. When both `validatorModelProvider` and `validatorModelId` are set, they take precedence over `defaultProvider`/`defaultModelId`. */
validatorModelId?: string;
/** Fallback model provider used when the primary reviewer model hits a retryable provider-side error. */
fallbackProvider?: string;
/** Fallback model ID used with `fallbackProvider`. */
fallbackModelId?: string;
/** Validator fallback model provider override. When both validator fallback fields are set, they take precedence over fallbackProvider/fallbackModelId. */
validatorFallbackModelProvider?: string;
/** Validator fallback model ID override. When both validator fallback fields are set, they take precedence over fallbackProvider/fallbackModelId. */
validatorFallbackModelId?: string;
/** Default thinking effort level for the reviewer agent session. */
defaultThinkingLevel?: string;
/** Task store for persisting agent log entries. When provided with `taskId`, enables full conversation logging. */
@@ -195,6 +203,12 @@ export async function reviewStep(
const validatorModelId = options.validatorModelProvider && options.validatorModelId
? options.validatorModelId
: options.defaultModelId;
const validatorFallbackProvider = options.validatorFallbackModelProvider && options.validatorFallbackModelId
? options.validatorFallbackModelProvider
: options.fallbackProvider;
const validatorFallbackModelId = options.validatorFallbackModelProvider && options.validatorFallbackModelId
? options.validatorFallbackModelId
: options.fallbackModelId;
// Spawn a reviewer agent with read-only tools
const { session } = await createKbAgent({
@@ -207,6 +221,8 @@ export async function reviewStep(
onToolEnd: agentLogger?.onToolEnd,
defaultProvider: validatorProvider,
defaultModelId: validatorModelId,
fallbackProvider: validatorFallbackProvider,
fallbackModelId: validatorFallbackModelId,
defaultThinkingLevel: options.defaultThinkingLevel,
});
@@ -225,7 +241,7 @@ export async function reviewStep(
});
try {
await session.prompt(request);
await promptWithFallback(session, request);
// Re-raise errors that pi-coding-agent swallowed after exhausting retries.
// The caller (executor's createReviewStepTool) catches errors and returns

View File

@@ -10,6 +10,14 @@ import { fileURLToPath } from "node:url";
import { dirname, join } from "node:path";
import { mkdir, writeFile, rm } from "node:fs/promises";
const { mockReviewStep } = vi.hoisted(() => ({
mockReviewStep: vi.fn(),
}));
vi.mock("./reviewer.js", () => ({
reviewStep: mockReviewStep,
}));
const __dirname = dirname(fileURLToPath(import.meta.url));
function createMockStore(overrides: Partial<TaskStore> = {}): TaskStore {
@@ -344,6 +352,7 @@ describe("TriageProcessor", () => {
beforeEach(() => {
store = createMockStore();
processor = new TriageProcessor(store, rootDir);
mockReviewStep.mockReset();
});
it("creates processor with default options", () => {
@@ -372,6 +381,73 @@ describe("TriageProcessor", () => {
expect(store.on).toHaveBeenCalledWith("settings:updated", expect.any(Function));
});
it("re-reads settings when review_spec runs so reviewer uses the latest validator model", async () => {
const taskId = "FN-001";
const testRootDir = join(__dirname, "__test_triage_review_spec__");
const promptPath = `.fusion/tasks/${taskId}/PROMPT.md`;
const taskDir = join(testRootDir, ".fusion", "tasks", taskId);
await mkdir(taskDir, { recursive: true });
await writeFile(join(taskDir, "PROMPT.md"), "# Spec\n\nCurrent prompt");
const freshSettings: Settings = {
maxConcurrent: 2,
maxWorktrees: 4,
pollIntervalMs: 10000,
groupOverlappingFiles: false,
autoMerge: true,
defaultProvider: "openai-codex",
defaultModelId: "gpt-5.4",
validatorProvider: "zai",
validatorModelId: "glm-5.1",
};
store = createMockStore({
getSettings: vi.fn().mockResolvedValue(freshSettings),
});
processor = new TriageProcessor(store, testRootDir);
mockReviewStep.mockResolvedValue({
verdict: "APPROVE",
review: "Looks good.",
summary: "approved",
});
const tool = (processor as any).createReviewSpecTool(
taskId,
promptPath,
{ current: null },
{ current: null },
{ current: null },
{
defaultProvider: "anthropic",
defaultModelId: "claude-opus-4-6",
validatorProvider: "anthropic",
validatorModelId: "claude-opus-4-6",
},
);
await tool.execute({});
expect(store.getSettings).toHaveBeenCalled();
expect(mockReviewStep).toHaveBeenCalledWith(
testRootDir,
taskId,
0,
"Specification",
"spec",
"# Spec\n\nCurrent prompt",
undefined,
expect.objectContaining({
defaultProvider: "openai-codex",
defaultModelId: "gpt-5.4",
validatorModelProvider: "zai",
validatorModelId: "glm-5.1",
}),
);
await rm(testRootDir, { recursive: true, force: true });
});
});
describe("Re-specification flow", () => {

View File

@@ -11,7 +11,7 @@ import type {
ToolDefinition,
AgentSession,
} from "@mariozechner/pi-coding-agent";
import { createKbAgent, describeModel } from "./pi.js";
import { createKbAgent, describeModel, promptWithFallback } from "./pi.js";
import { reviewStep, type ReviewVerdict } from "./reviewer.js";
import { PRIORITY_SPECIFY, type AgentSemaphore } from "./concurrency.js";
import { AgentLogger } from "./agent-logger.js";
@@ -474,6 +474,12 @@ export class TriageProcessor {
defaultModelId: settings.planningProvider && settings.planningModelId
? settings.planningModelId
: settings.defaultModelId,
fallbackProvider: settings.planningFallbackProvider && settings.planningFallbackModelId
? settings.planningFallbackProvider
: settings.fallbackProvider,
fallbackModelId: settings.planningFallbackProvider && settings.planningFallbackModelId
? settings.planningFallbackModelId
: settings.fallbackModelId,
defaultThinkingLevel: settings.defaultThinkingLevel,
});
@@ -523,7 +529,8 @@ export class TriageProcessor {
existingPrompt,
feedback,
);
await session.prompt(
await promptWithFallback(
session,
agentPrompt,
imageContents.length > 0 ? { images: imageContents } : undefined,
);
@@ -934,6 +941,10 @@ export class TriageProcessor {
};
}
// Re-read settings at review time so long-lived triage sessions pick up
// model changes made after the session started.
const currentSettings = await store.getSettings();
const result = await reviewStep(
rootDir,
taskId,
@@ -944,11 +955,11 @@ export class TriageProcessor {
undefined,
{
onText: (delta) => options.onAgentText?.(taskId, delta),
defaultProvider: settings.defaultProvider,
defaultModelId: settings.defaultModelId,
validatorModelProvider: settings.validatorProvider,
validatorModelId: settings.validatorModelId,
defaultThinkingLevel: settings.defaultThinkingLevel,
defaultProvider: currentSettings.defaultProvider,
defaultModelId: currentSettings.defaultModelId,
validatorModelProvider: currentSettings.validatorProvider,
validatorModelId: currentSettings.validatorModelId,
defaultThinkingLevel: currentSettings.defaultThinkingLevel,
store,
taskId,
},