From 9d9cc149c5e7740494e72b5395e02a2b8723d227 Mon Sep 17 00:00:00 2001 From: Mathew Benjamin <90913413+mathewtbenjamin@users.noreply.github.com> Date: Wed, 12 Aug 2026 01:51:46 +1000 Subject: [PATCH] fix: honor agent model routing --- HACKING.md | 5 +- extensions/piolium/agent-runner.ts | 80 ++++++++++++++++++++++-- extensions/piolium/index.ts | 8 ++- extensions/piolium/modes/balanced.ts | 10 ++- extensions/piolium/modes/phase-runner.ts | 9 ++- extensions/piolium/retry.ts | 20 ++++++ test/agent-runner.test.ts | 51 ++++++++++++++- test/retry.test.ts | 24 +++++++ 8 files changed, 196 insertions(+), 11 deletions(-) diff --git a/HACKING.md b/HACKING.md index d856465..438a6b2 100644 --- a/HACKING.md +++ b/HACKING.md @@ -306,6 +306,8 @@ Piolium records retry metadata in `piolium/audit-state.json`. Retry counts are r | Lite Q0/Q1 overrides | `2` retries, `5000` ms base backoff, `120000` ms max backoff | `--plm-lite-retries`, `--plm-lite-backoff`, `--plm-lite-backoff-max` | | Command reruns | `3` retries, `5000` ms base backoff, `120000` ms max backoff | `--plm-command-retries`, `--plm-command-backoff`, `--plm-command-backoff-max` | +Provider policy refusals that explicitly flag cybersecurity content are treated as non-retryable: repeating an unchanged phase prompt cannot succeed and only wastes requests. The phase still records the refusal in `audit-state.json` so the operator can resume with another model. + Example: ```bash @@ -397,7 +399,8 @@ Sub-agents under `agents/` are package-private because Pi has no first-class `ag - Canonical phase order is in `extensions/piolium/modes/modes.ts`. - Per-phase retry, state transitions, and heartbeat tracking are owned by `extensions/piolium/modes/phase-runner.ts`. - Sub-agents run through `extensions/piolium/agent-runner.ts`, which creates child Pi sessions in-process with `createAgentSession`. -- Agent runs write transcripts to `/piolium/tmp/piolium/runs//`. +- Agent frontmatter model families (`haiku`, `sonnet`, `opus`) resolve against authenticated Pi models. Resolution prefers the parent's provider when it offers that family, then direct Anthropic and Anthropic Vertex, and selects the newest matching model. Unresolvable declarations fall back to the parent model. +- Agent runs write transcripts to `/piolium/tmp/piolium/runs//` and record both requested and resolved models in `prompt.md`. - Durable state lives at `/piolium/audit-state.json`. - File writes to audit state should go through helpers in `extensions/piolium/audit-state.ts`. diff --git a/extensions/piolium/agent-runner.ts b/extensions/piolium/agent-runner.ts index 4f347c2..03ae412 100644 --- a/extensions/piolium/agent-runner.ts +++ b/extensions/piolium/agent-runner.ts @@ -24,7 +24,7 @@ import { type WriteStream, createWriteStream } from "node:fs"; import { writeFileSync } from "node:fs"; import { join } from "node:path"; import type { ThinkingLevel } from "@earendil-works/pi-agent-core"; -import type { ImageContent, Model, TextContent } from "@earendil-works/pi-ai"; +import type { Api, ImageContent, Model, TextContent } from "@earendil-works/pi-ai"; import { type AgentSessionEvent, DefaultResourceLoader, @@ -66,8 +66,7 @@ export interface RunAgentOptions { * Optional model override. When omitted the child boots with the * settings-derived default, matching the parent. */ - // biome-ignore lint/suspicious/noExplicitAny: pi-ai Model is generic over provider api - model?: Model; + model?: Model; /** * Optional parent registry. This preserves extension-registered providers * and their request headers for child sessions without loading extensions @@ -119,6 +118,76 @@ export class AgentRunError extends Error { } const TRANSCRIPT_STRING_LIMIT = 8_000; +const AGENT_MODEL_FAMILIES = new Set(["haiku", "sonnet", "opus"]); + +function modelVersionParts(model: Model): number[] { + return model.id.match(/\d+/g)?.map(Number) ?? []; +} + +function compareVersionParts(a: number[], b: number[]): number { + const length = Math.max(a.length, b.length); + for (let index = 0; index < length; index++) { + const difference = (a[index] ?? 0) - (b[index] ?? 0); + if (difference !== 0) return difference; + } + return 0; +} + +function providerPreference(model: Model, parentProvider: string | undefined): number { + if (parentProvider && model.provider === parentProvider) return 3; + if (model.provider === "anthropic") return 2; + if (model.provider === "anthropic-vertex") return 1; + return 0; +} + +function choosePreferredModel( + models: Model[], + parentProvider: string | undefined, +): Model | undefined { + let preferred: Model | undefined; + for (const candidate of models) { + if (!preferred) { + preferred = candidate; + continue; + } + const providerDifference = + providerPreference(candidate, parentProvider) - providerPreference(preferred, parentProvider); + if ( + providerDifference > 0 || + (providerDifference === 0 && + compareVersionParts(modelVersionParts(candidate), modelVersionParts(preferred)) > 0) + ) { + preferred = candidate; + } + } + return preferred; +} + +/** Resolve a Claude Code-style agent model declaration against authenticated Pi models. */ +export function resolveAgentModel( + requested: string | undefined, + parentModel: Model | undefined, + modelRegistry: ModelRegistry | undefined, +): Model | undefined { + const normalized = requested?.trim().toLowerCase(); + if (!normalized || !modelRegistry) return parentModel; + + const available = modelRegistry.getAvailable(); + let matches: Model[]; + if (AGENT_MODEL_FAMILIES.has(normalized)) { + matches = available.filter((model) => { + const searchable = `${model.provider} ${model.id} ${model.name ?? ""}`.toLowerCase(); + return searchable.includes(normalized); + }); + } else { + matches = available.filter((model) => { + const qualified = `${model.provider}/${model.id}`.toLowerCase(); + return qualified === normalized || model.id.toLowerCase() === normalized; + }); + } + + return choosePreferredModel(matches, parentModel?.provider) ?? parentModel; +} export function buildRuntimeHeader(runtime: RuntimeContext): string { const lines: string[] = ["# piolium Runtime", ""]; @@ -196,6 +265,7 @@ export async function runAgent(options: RunAgentOptions): Promise( maxRetries: readNonNegativeIntEnv("PIOLIUM_COMMAND_MAX_RETRIES", 3), backoffBaseMs: readPositiveIntEnv("PIOLIUM_COMMAND_BACKOFF_BASE_MS", 5000), backoffMaxMs: readPositiveIntEnv("PIOLIUM_COMMAND_BACKOFF_MAX_MS", 120_000), + shouldRetry: (err) => !isNonRetryableAgentError(err), onRetry: (info) => { ui.notify( `${label} attempt ${info.attempt}/${info.maxAttempts} failed; retrying in ${Math.ceil(info.backoffMs / 1000)}s.`, diff --git a/extensions/piolium/modes/balanced.ts b/extensions/piolium/modes/balanced.ts index eeab92f..519dd1b 100644 --- a/extensions/piolium/modes/balanced.ts +++ b/extensions/piolium/modes/balanced.ts @@ -39,6 +39,7 @@ import { runCandidateScanAsync } from "../candidate-scan.ts"; import { consolidateDrafts, findingsDraftDir, listFindingDirs } from "../findings.ts"; import { ingestKnowledgeBaseForRun } from "../knowledge-base-input.ts"; import { runReconAsync } from "../recon.ts"; +import { findNonRetryableRejection, isNonRetryableAgentError } from "../retry.ts"; import { Scheduler } from "../scheduler.ts"; import { cleanupConfirmArtifacts } from "./confirm.ts"; import { type PhaseUiHooks, runAgentPhase } from "./phase-runner.ts"; @@ -350,6 +351,8 @@ async function runL3PlusL4Parallel( }), ]); scheduler.dispose(); + const nonRetryable = findNonRetryableRejection(settled); + if (nonRetryable) throw nonRetryable.reason; return { failed: settled.some((s) => s.status === "rejected") }; } @@ -409,10 +412,12 @@ async function runPerFindingPhase( ); scheduler.dispose(); const failed = results.some((r) => r.status === "rejected"); + const nonRetryable = findNonRetryableRejection(results); await applyPhaseStatus(cwd, audit, phaseName, { status: failed ? "failed" : "complete", ...(failed ? { error: `Some per-finding ${phaseName} runs failed.` } : {}), }); + if (nonRetryable) throw nonRetryable.reason; return { failed }; } @@ -558,6 +563,7 @@ export async function runBalancedAudit(opts: RunBalancedOptions): Promise { return; } - if (opts.signal?.aborted || attempt >= maxAttempts) { + const nonRetryable = isNonRetryableAgentError(err); + if (opts.signal?.aborted || attempt >= maxAttempts || nonRetryable) { await applyPhaseStatus(cwd, audit, phaseName, { status: "failed", - error: - attempt >= maxAttempts && maxRetries > 0 + error: nonRetryable + ? `Non-retryable provider policy error: ${message}` + : attempt >= maxAttempts && maxRetries > 0 ? `Failed after ${maxRetries} retries: ${message}` : message, attempt, diff --git a/extensions/piolium/retry.ts b/extensions/piolium/retry.ts index 0cf9245..71dfbad 100644 --- a/extensions/piolium/retry.ts +++ b/extensions/piolium/retry.ts @@ -42,6 +42,26 @@ export function errorMessage(err: unknown): string { return err instanceof Error ? err.message : typeof err === "string" ? err : "Unknown error"; } +const NON_RETRYABLE_AGENT_ERROR_PATTERNS = [ + /This content was flagged for possible cybersecurity risk/i, + /join the Trusted Access for Cyber program/i, +]; + +/** Provider policy refusals are deterministic for an unchanged phase prompt. */ +export function isNonRetryableAgentError(err: unknown): boolean { + const message = errorMessage(err); + return NON_RETRYABLE_AGENT_ERROR_PATTERNS.some((pattern) => pattern.test(message)); +} + +export function findNonRetryableRejection( + results: readonly PromiseSettledResult[], +): PromiseRejectedResult | undefined { + for (const result of results) { + if (result.status === "rejected" && isNonRetryableAgentError(result.reason)) return result; + } + return undefined; +} + export function retryBackoffMs(attempt: number, baseMs: number, maxMs: number): number { const exponent = Math.max(0, attempt - 1); const raw = baseMs * 2 ** exponent; diff --git a/test/agent-runner.test.ts b/test/agent-runner.test.ts index 4f733b3..ddb6a3b 100644 --- a/test/agent-runner.test.ts +++ b/test/agent-runner.test.ts @@ -1,5 +1,54 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import type { ModelRegistry } from "@earendil-works/pi-coding-agent"; import { describe, expect, it } from "vitest"; -import { buildRuntimeHeader, compactTranscriptEvent } from "../extensions/piolium/agent-runner.ts"; +import { + buildRuntimeHeader, + compactTranscriptEvent, + resolveAgentModel, +} from "../extensions/piolium/agent-runner.ts"; + +function fakeModel(provider: string, id: string, name = id): Model { + return { provider, id, name } as Model; +} + +function fakeRegistry(models: Model[]): ModelRegistry { + return { getAvailable: () => models } as unknown as ModelRegistry; +} + +describe("resolveAgentModel", () => { + it("routes a family alias away from the parent model when an authenticated match exists", () => { + const parent = fakeModel("openai-codex", "gpt-5.6-sol"); + const sonnet45 = fakeModel("anthropic", "claude-sonnet-4-5"); + const sonnet5 = fakeModel("anthropic", "claude-sonnet-5"); + const registry = fakeRegistry([parent, sonnet45, sonnet5]); + + expect(resolveAgentModel("sonnet", parent, registry)).toBe(sonnet5); + }); + + it("prefers a matching model from the parent's provider", () => { + const parent = fakeModel("anthropic-vertex", "claude-opus-4-6@default"); + const direct = fakeModel("anthropic", "claude-sonnet-5"); + const vertex = fakeModel("anthropic-vertex", "claude-sonnet-4-5@20250929"); + const registry = fakeRegistry([direct, vertex]); + + expect(resolveAgentModel("sonnet", parent, registry)).toBe(vertex); + }); + + it("resolves a fully-qualified model id", () => { + const parent = fakeModel("openai-codex", "gpt-5.6-sol"); + const requested = fakeModel("anthropic", "claude-sonnet-4-6"); + const registry = fakeRegistry([parent, requested]); + + expect(resolveAgentModel("anthropic/claude-sonnet-4-6", parent, registry)).toBe(requested); + }); + + it("falls back to the parent model when the declaration cannot be resolved", () => { + const parent = fakeModel("openai-codex", "gpt-5.6-sol"); + const registry = fakeRegistry([parent]); + + expect(resolveAgentModel("sonnet", parent, registry)).toBe(parent); + }); +}); describe("buildRuntimeHeader", () => { it("includes the audit cwd, mode, and phase when provided", () => { diff --git a/test/retry.test.ts b/test/retry.test.ts index 881bcd0..006e8fb 100644 --- a/test/retry.test.ts +++ b/test/retry.test.ts @@ -1,5 +1,7 @@ import { afterEach, describe, expect, it } from "vitest"; import { + findNonRetryableRejection, + isNonRetryableAgentError, readNonNegativeIntEnv, readPositiveIntEnv, runWithRetry, @@ -59,6 +61,28 @@ describe("runWithRetry", () => { }); }); +describe("isNonRetryableAgentError", () => { + it("recognizes the Codex cybersecurity policy refusal", () => { + expect( + isNonRetryableAgentError(new Error("This content was flagged for possible cybersecurity risk.")), + ).toBe(true); + }); + + it("does not suppress retries for ordinary security-audit failures", () => { + expect(isNonRetryableAgentError(new Error("cybersecurity scanner timed out"))).toBe(false); + }); + + it("finds policy refusals captured by Promise.allSettled", () => { + const policyError = new Error("Join the Trusted Access for Cyber program to continue."); + const results: PromiseSettledResult[] = [ + { status: "fulfilled", value: undefined }, + { status: "rejected", reason: policyError }, + ]; + + expect(findNonRetryableRejection(results)?.reason).toBe(policyError); + }); +}); + describe("env readers", () => { it("reads PIOLIUM_* values", () => { process.env.PIOLIUM_TEST_LIMIT = "7";