From 3c3bf416205cb235e4b6b2b7671810d0364cc857 Mon Sep 17 00:00:00 2001 From: Brandon Leung Date: Tue, 8 Sep 2026 14:02:45 -0400 Subject: [PATCH 1/4] feat(gateway): re-mint once and resume on a 401 from an aged bearer Co-Authored-By: Claude Fable 5.1 --- src/lib/__tests__/agent-interface.test.ts | 211 +++++ src/lib/__tests__/gateway-session.test.ts | 27 + .../agent/__tests__/output-signals.test.ts | 12 + .../agent/__tests__/triage-provider.test.ts | 18 + src/lib/agent/agent-interface.ts | 778 ++++++++++-------- src/lib/agent/output-signals.ts | 12 + .../harness/pi/__tests__/gateway.test.ts | 138 +++- src/lib/agent/runner/harness/pi/gateway.ts | 69 +- src/lib/agent/runner/harness/pi/index.ts | 50 +- src/lib/agent/runner/harness/pi/task.ts | 50 +- src/lib/agent/runner/shared/bootstrap.ts | 31 +- src/lib/agent/signals.ts | 7 + src/lib/agent/triage-provider.ts | 36 +- src/lib/gateway-session.ts | 16 +- 14 files changed, 1048 insertions(+), 407 deletions(-) diff --git a/src/lib/__tests__/agent-interface.test.ts b/src/lib/__tests__/agent-interface.test.ts index e3f2926f8..f87979395 100644 --- a/src/lib/__tests__/agent-interface.test.ts +++ b/src/lib/__tests__/agent-interface.test.ts @@ -10,7 +10,9 @@ import { reportMcpSetup, } from '@lib/agent/agent-interface'; import { AgentOutputSignals } from '@lib/agent/output-signals'; +import { RESUME_INSTRUCTION } from '@lib/agent/signals'; import { analytics } from '@utils/analytics'; +import { wizardAbort } from '@utils/wizard-abort'; import { Sequence } from '@lib/constants'; import type { WizardRunOptions } from '@utils/types'; import type { SpinnerHandle } from '@ui'; @@ -22,6 +24,11 @@ import { // Mock dependencies vi.mock('../../utils/analytics'); vi.mock('../../utils/debug'); +// wizardAbort exits the process; the 401 tests below need it to just reject. +vi.mock('@utils/wizard-abort', async (importOriginal) => ({ + ...(await importOriginal()), + wizardAbort: vi.fn(), +})); // Mock the SDK module const mockQuery = vi.fn(); @@ -54,6 +61,7 @@ const mockUIInstance = { showBlockingOutage: vi.fn(), setReadinessWarnings: vi.fn(), showSettingsOverride: vi.fn(), + showAuthError: vi.fn(), startRun: vi.fn(), syncTodos: vi.fn(), groupMultiselect: vi.fn(), @@ -94,6 +102,7 @@ describe('runAgent', () => { // would make either source pass. gatewayUrl: 'https://gateway.test', token: 'phe_run_scoped_token', + refreshAtMs: Date.now() + 3600_000, }, }; @@ -632,6 +641,7 @@ describe('subprocess gateway credentials', () => { gatewayUrl: 'https://ai-gateway.us.posthog.com', token: 'phe_run_scoped_token', teamId: 42, + refreshAtMs: Date.now() + 3600_000, }, }; const options: WizardRunOptions = { @@ -683,6 +693,207 @@ describe('subprocess gateway credentials', () => { }); }); +describe('gateway re-mint on 401', () => { + const spinner = { start: vi.fn(), stop: vi.fn(), message: vi.fn() }; + const options: WizardRunOptions = { + debug: false, + installDir: '/test/dir', + signup: false, + ci: false, + benchmark: false, + yaraReport: false, + }; + const HOUR = 3600_000; + const auth = (token: string, refreshAtMs: number) => ({ + gatewayUrl: 'https://ai-gateway.us.posthog.com', + token, + teamId: 42, + refreshAtMs, + }); + const config = ( + gatewayAuth: ReturnType, + refreshGatewayAuth: () => Promise>, + ) => ({ + workingDirectory: '/test/dir', + mcpServers: {}, + model: 'claude-sonnet-4-6', + posthogApiKey: 'phx_user_oauth_token', + sequence: Sequence.linear, + triageProvider: () => Promise.resolve('false_positive'), + gatewayAuth, + refreshGatewayAuth, + }); + const run = (cfg: ReturnType) => + runAgent(cfg, 'test prompt', options, spinner as unknown as SpinnerHandle, { + successMessage: 'ok', + errorMessage: 'err', + }); + + function* rejectedSession(id: string) { + yield { + type: 'system', + subtype: 'init', + session_id: id, + model: 'm', + tools: [], + mcp_servers: [], + }; + yield { + type: 'assistant', + session_id: id, + message: { + role: 'assistant', + content: [ + { type: 'text', text: 'API Error: 401 {"detail":"token expired"}' }, + ], + }, + }; + // Not reached: the 401 handler leaves the loop before the SDK's result. + yield { + type: 'result', + subtype: 'success', + session_id: id, + is_error: true, + result: 'API Error: 401', + }; + } + function* completedSession(id: string) { + yield { + type: 'system', + subtype: 'init', + session_id: id, + model: 'm', + tools: [], + mcp_servers: [], + }; + yield { + type: 'result', + subtype: 'success', + session_id: id, + is_error: false, + result: 'done', + }; + } + + beforeEach(() => { + vi.clearAllMocks(); + mockUIInstance.spinner.mockReturnValue(spinner); + vi.mocked(wizardAbort).mockRejectedValue(new Error('wizardAbort: exit')); + }); + + it('mints once and resumes the session when an aged bearer is rejected', async () => { + mockQuery + .mockReturnValueOnce(rejectedSession('sess-1')) + .mockReturnValueOnce(completedSession('sess-2')); + const refresh = vi + .fn() + .mockResolvedValue(auth('phe_fresh', Date.now() + HOUR)); + const cfg = config(auth('phe_stale', Date.now() - 1), refresh); + + const result = await run(cfg); + + expect(result).toEqual({}); + expect(refresh).toHaveBeenCalledTimes(1); + expect(wizardAbort).not.toHaveBeenCalled(); + expect(mockQuery).toHaveBeenCalledTimes(2); + const [first, second] = mockQuery.mock.calls.map((c) => c[0]); + expect(first.options.resume).toBeUndefined(); + expect(second.options.resume).toBe('sess-1'); + // The new subprocess carries the new bearer and finds the transcript in + // the same config dir; the env is frozen at spawn, so a new one is the + // only way to hand it over. + expect(second.options.env.ANTHROPIC_AUTH_TOKEN).toBe('phe_fresh'); + expect(second.options.env.CLAUDE_CODE_OAUTH_TOKEN).toBe('phe_fresh'); + expect(second.options.env.CLAUDE_CONFIG_DIR).toBe( + first.options.env.CLAUDE_CONFIG_DIR, + ); + // The resumed session is told to pick up, not restarted from the prompt. + const resumed = await second.prompt.next(); + expect(resumed.value.message.content).toBe(RESUME_INSTRUCTION); + expect(cfg.gatewayAuth.token).toBe('phe_fresh'); + expect(analytics.wizardCapture).toHaveBeenCalledWith( + 'gateway token reminted', + { resumed: true }, + ); + }); + + it('fails the run on a second 401 after the re-mint', async () => { + mockQuery + .mockReturnValueOnce(rejectedSession('sess-1')) + .mockReturnValueOnce(rejectedSession('sess-2')); + // The new bearer is also past refresh (a slow run under a short TTL), so + // only the once-per-run rule stands between this and a second mint. + const refresh = vi + .fn() + .mockResolvedValue(auth('phe_fresh', Date.now() - 1)); + + const result = await run( + config(auth('phe_stale', Date.now() - 1), refresh), + ); + + expect(refresh).toHaveBeenCalledTimes(1); + expect(mockQuery).toHaveBeenCalledTimes(2); + expect(mockUIInstance.showAuthError).toHaveBeenCalledTimes(1); + expect(wizardAbort).toHaveBeenCalledTimes(1); + // In production wizardAbort exits; the mocked rejection surfaces as the + // run's API error. + expect(result.error).toBe('WIZARD_API_ERROR'); + }); + + it('judges a failed resumed session on its own error, not the old 401', async () => { + function* resumedThenFailed(id: string) { + yield { + type: 'system', + subtype: 'init', + session_id: id, + model: 'm', + tools: [], + mcp_servers: [], + }; + yield { + type: 'result', + subtype: 'success', + session_id: id, + is_error: true, + result: 'API Error: 500 upstream exploded', + }; + } + mockQuery + .mockReturnValueOnce(rejectedSession('sess-1')) + .mockReturnValueOnce(resumedThenFailed('sess-2')); + const refresh = vi + .fn() + .mockResolvedValue(auth('phe_fresh', Date.now() + HOUR)); + + const result = await run( + config(auth('phe_stale', Date.now() - 1), refresh), + ); + + // The 401 that triggered the re-mint is history; reporting it here would + // send the user to the auth screen for a 500. + expect(result.error).toBe('WIZARD_API_ERROR'); + expect(result.message).toContain('500'); + expect(result.message).not.toContain('401'); + expect(mockUIInstance.showAuthError).not.toHaveBeenCalled(); + }); + + it('does not re-mint when a fresh bearer is rejected', async () => { + mockQuery.mockReturnValueOnce(rejectedSession('sess-1')); + const refresh = vi.fn(); + + const result = await run( + config(auth('phe_fresh', Date.now() + HOUR), refresh), + ); + + // A fresh token the gateway rejects is a bad credential, not age. + expect(refresh).not.toHaveBeenCalled(); + expect(mockQuery).toHaveBeenCalledTimes(1); + expect(mockUIInstance.showAuthError).toHaveBeenCalledTimes(1); + expect(wizardAbort).toHaveBeenCalledTimes(1); + expect(result.error).toBe('WIZARD_API_ERROR'); + }); +}); + describe('auth error context', () => { // The 401 screen's region comes from whichever url it is handed, which is why // runAgent passes the run's resolved auth rather than the process global a diff --git a/src/lib/__tests__/gateway-session.test.ts b/src/lib/__tests__/gateway-session.test.ts index a2cfb977f..5e5213b7f 100644 --- a/src/lib/__tests__/gateway-session.test.ts +++ b/src/lib/__tests__/gateway-session.test.ts @@ -4,6 +4,7 @@ import { GatewayMintRefused, buildWizardPropertiesBlob, gatewayAuth, + isPastRefresh, isTrustedGatewayUrl, resetGatewaySession, } from '@lib/gateway-session'; @@ -70,6 +71,7 @@ describe('gatewayAuth', () => { gatewayUrl: 'https://gateway.us.posthog.com', token: 'phe_minted', teamId: 42, + refreshAtMs: expect.any(Number), }); expect(fetchMock).toHaveBeenCalledWith( 'https://us.posthog.com/api/wizard/gateway_token/', @@ -566,6 +568,31 @@ describe('gatewayAuth', () => { } }); + it('sets the refresh instant at the refresh fraction of the token life', async () => { + vi.useFakeTimers(); + try { + const ttlMs = 60 * 60 * 1000; + fetchMock.mockResolvedValue({ + ok: true, + json: () => + Promise.resolve({ + token: 'phe_minted', + expires_at: new Date(Date.now() + ttlMs).toISOString(), + gateway_url: 'https://gateway.us.posthog.com', + }), + }); + const auth = await gatewayAuth(host, 'pha_oauth', 'integration'); + expect(auth.refreshAtMs).toBe(Date.now() + ttlMs * 0.8); + // A 401 before this instant is a bad credential; after it, an aged + // bearer that one re-mint recovers. + expect(isPastRefresh(auth)).toBe(false); + vi.setSystemTime(Date.now() + ttlMs * 0.8); + expect(isPastRefresh(auth)).toBe(true); + } finally { + vi.useRealTimers(); + } + }); + it('retries cleanly after a failed mint rather than wedging the session', async () => { // A rejected resolve must leave neither a cached posture nor a claimed // in-flight slot behind, or one transient 503 wedges the run for the diff --git a/src/lib/agent/__tests__/output-signals.test.ts b/src/lib/agent/__tests__/output-signals.test.ts index 95e92270c..76eeaef6c 100644 --- a/src/lib/agent/__tests__/output-signals.test.ts +++ b/src/lib/agent/__tests__/output-signals.test.ts @@ -60,6 +60,18 @@ describe('AgentOutputSignals', () => { expect(signals.remark()).toBeUndefined(); }); + it('forgets API error lines after a re-mint but keeps every other signal', () => { + const signals = new AgentOutputSignals(); + signals.push('API Error: 401 token expired'); + signals.push('[ERROR-MCP-MISSING] could not reach MCP'); + + signals.forgetApiErrors(); + + expect(signals.hasApiError()).toBe(false); + expect(signals.hasApiErrorStatus(401)).toBe(false); + expect(signals.has('MCP_MISSING')).toBe(true); + }); + it('treats the API error status as a parameter, not a fixed marker', () => { const signals = new AgentOutputSignals(); signals.push('API Error: 503 service unavailable'); diff --git a/src/lib/agent/__tests__/triage-provider.test.ts b/src/lib/agent/__tests__/triage-provider.test.ts index 5eb63683d..d51d82b9e 100644 --- a/src/lib/agent/__tests__/triage-provider.test.ts +++ b/src/lib/agent/__tests__/triage-provider.test.ts @@ -117,6 +117,24 @@ describe('createTriageLLMProvider', () => { }); }); + it('re-reads auth on every call instead of closing over the first token', async () => { + // A run that re-mints mid-way must scan with the current bearer. + complete.mockResolvedValue(reply('false_positive')); + const tokens = ['tok-1', 'tok-2']; + const provider = createTriageLLMProvider( + () => Promise.resolve({ ...AUTH, authToken: tokens.shift() ?? 'tok-3' }), + Harness.anthropic, + ); + + await provider('first?'); + await provider('second?'); + + expect(complete.mock.calls.map((c) => c[2]?.apiKey)).toEqual([ + 'tok-1', + 'tok-2', + ]); + }); + it('keeps only text blocks in the verdict', async () => { complete.mockResolvedValue({ content: [ diff --git a/src/lib/agent/agent-interface.ts b/src/lib/agent/agent-interface.ts index 270f1a9ba..8d1098767 100644 --- a/src/lib/agent/agent-interface.ts +++ b/src/lib/agent/agent-interface.ts @@ -32,6 +32,7 @@ import type { HostResolution } from '@lib/host-resolution'; import { buildWizardPropertiesBlob, gatewayAuth, + isPastRefresh, type GatewayAuth, } from '@lib/gateway-session'; import { evaluateBashCommand } from './bash-fence'; @@ -46,7 +47,12 @@ import type { LLMProvider } from '@posthog/warlock'; import { assembleCommandments } from './runner/switchboard/commandments'; import { classifyToolToStage } from './agent-phase'; import type { PackageManagerDetector } from '@lib/detection/package-manager'; -import { AgentSignals, AgentErrorType, REMARK_INSTRUCTION } from './signals'; +import { + AgentSignals, + AgentErrorType, + REMARK_INSTRUCTION, + RESUME_INSTRUCTION, +} from './signals'; import { classifyAuthFailure } from '@lib/errors'; import { isGrantRevoked } from '@lib/auth-session-state'; import { AgentOutputSignals } from './output-signals'; @@ -337,6 +343,12 @@ type AgentRunConfig = { triageProvider: LLMProvider; /** The run's minted gateway auth: base url, bearer and team for the subprocess. */ gatewayAuth: GatewayAuth; + /** + * Resolve the run's gateway auth again: the cached token while it is fresh, + * a new mint once it is past its refresh instant. Recovers a 401 on an aged + * bearer. + */ + refreshGatewayAuth?: () => Promise; /** Program id, for the program-axis commandments. */ program?: string; /** Resolved sequence, for the sequence-axis commandments. */ @@ -525,11 +537,9 @@ export async function initializeAgent( // gatewayAuth mints for this run. // Disable experimental betas (like input_examples) the gateway doesn't support. process.env.CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS = 'true'; - const auth = await gatewayAuth( - config.host, - config.posthogApiKey, - config.programId, - ); + const currentGatewayAuth = () => + gatewayAuth(config.host, config.posthogApiKey, config.programId); + const auth = await currentGatewayAuth(); const gatewayUrl = auth.gatewayUrl; process.env.ANTHROPIC_BASE_URL = gatewayUrl; process.env.ANTHROPIC_AUTH_TOKEN = auth.token; @@ -537,23 +547,24 @@ export async function initializeAgent( // Use CLAUDE_CODE_OAUTH_TOKEN to override any stored /login credentials process.env.CLAUDE_CODE_OAUTH_TOKEN = auth.token; - // Same values the env vars above carry, handed over explicitly so triage - // never has to read them back out of the environment. The run tags ride - // along so scan spend bills to this program, with `call_type` keeping it - // separable from the agent's own calls. - const triageProvider = createTriageLLMProvider( - { - baseURL: gatewayUrl, - authToken: auth.token, - teamId: auth.teamId, - wizardMetadata: { - ...(config.wizardMetadata ?? {}), - call_type: CallType.yaraTriage, - }, + // Handed over explicitly so triage never reads the environment, and + // re-read per call so a long run's scans follow a re-mint. The run tags + // ride along so scan spend bills to this program, with `call_type` + // keeping it separable from the agent's own calls. + const triageMetadata = { + ...(config.wizardMetadata ?? {}), + call_type: CallType.yaraTriage, + }; + const triageProvider = createTriageLLMProvider(async () => { + const current = await currentGatewayAuth(); + return { + baseURL: current.gatewayUrl, + authToken: current.token, + teamId: current.teamId, + wizardMetadata: triageMetadata, wizardFlags: config.wizardFlags ?? {}, - }, - Harness.anthropic, - ); + }; + }, Harness.anthropic); logToFile('Configured LLM gateway:', gatewayUrl); logToFile( @@ -645,6 +656,7 @@ export async function initializeAgent( capture: config.capture, triageProvider, gatewayAuth: auth, + refreshGatewayAuth: currentGatewayAuth, program: config.integrationLabel, // A queue context is present only on a task run; that is the sequence. sequence: config.orchestrator ? Sequence.orchestrator : Sequence.linear, @@ -765,19 +777,22 @@ export async function runAgent( // the result is received, keeping the stdin stream alive for permission responses. // See: https://github.com/anthropics/claude-code/issues/4775 // See: https://github.com/anthropics/claude-agent-sdk-typescript/issues/41 - let signalDone: () => void; - const resultReceived = new Promise((resolve) => { - signalDone = resolve; - }); - - const createPromptStream = async function* () { - yield { - type: 'user', - session_id: '', - message: { role: 'user', content: prompt }, - parent_tool_use_id: null, - }; - await resultReceived; + // One stream and one done-promise per query(): a session resumed after a + // re-mint needs its own. A no-op until the first stream installs its own. + let signalDone: () => void = () => undefined; + const createPromptStream = (text: string) => { + const resultReceived = new Promise((resolve) => { + signalDone = resolve; + }); + return (async function* () { + yield { + type: 'user', + session_id: '', + message: { role: 'user', content: text }, + parent_tool_use_id: null, + }; + await resultReceived; + })(); }; // Helper to handle successful completion (used in normal path and race condition recovery) @@ -852,12 +867,18 @@ export async function runAgent( // Abort controller — lets us force-kill the SDK query when we detect an // [ABORT] signal in the agent's output. Also stashes the reason so the // runner can surface it via outroData after we unwind. - const abortController = new AbortController(); + let abortController = new AbortController(); let abortReason: string | null = null; // Set when a YARA hook detects a terminal violation. Returning `stopReason` // from a PostToolUse hook does NOT stop the SDK, so we abort the query and // surface a YARA_VIOLATION below — mirroring the [ABORT] mechanism. let yaraViolationReason: string | null = null; + // Re-mint state: the SDK session to resume, one re-mint per run, and the + // config dir a resumed subprocess must share to find the transcript. + let sessionId: string | undefined; + let reminted = false; + let remintRequested = false; + const agentConfigDir = createIsolatedAgentConfigDir(); try { // Per-program allow/disallow lists tweak BASE_ALLOWED_TOOLS. Skills are @@ -887,7 +908,7 @@ export async function runAgent( yaraViolationReason = reason; logToFile(`[YARA] terminating run: ${reason}`); abortController.abort(); - signalDone!(); + signalDone(); }; // Local/CI escape hatch for Warlock/YARA scanning (off by default — see @@ -905,337 +926,394 @@ export async function runAgent( // capture is disabled. agentConfig.capture?.setInitialPrompt(prompt); - const response = query({ - prompt: createPromptStream(), - options: { - abortController, - model: agentConfig.model, - cwd: agentConfig.workingDirectory, - permissionMode: 'acceptEdits', - betas: ['context-1m-2025-08-07'], - mcpServers: agentConfig.mcpServers, - agents: { - 'general-purpose': { - description: - "General-purpose subagent. Inherits the parent run's tools plus the PostHog and wizard-tools MCP servers, so it can call mcp__posthog-wizard__* directly instead of curling the REST API.", - prompt: - 'You are a general-purpose subagent for the PostHog wizard. Prefer the authenticated mcp__posthog-wizard__* MCP tools over raw HTTP — they are already authenticated for this project. Only fall back to other transports if no MCP tool covers the operation.', - mcpServers: inheritedMcpServerNames, - // SDK does not propagate the parent's disallowedTools to subagents - // (sdk.d.ts: AgentDefinition has its own disallowedTools, and - // `tools: undefined` means "inherit all"). Without this, a program - // that disallows wizard_ask still leaks it to dispatched subagents. - disallowedTools: agentConfig.disallowedTools - ? [...agentConfig.disallowedTools] - : undefined, - }, - }, - // Load skills from project's .claude/skills/ directory - settingSources: ['project'], - // Enable all discovered skills. Omitting this is NOT "skills off" — - // it just means no SDK auto-config — so we set 'all' explicitly to - // preserve the prior behavior where 'Skill' in allowedTools exposed - // everything under .claude/skills/. (SDK ≥0.2.133 deprecates passing - // 'Skill' in allowedTools in favor of this option.) - skills: 'all', - allowedTools, - sandbox: { - enabled: true, - // SDK 0.2.91 made failIfUnavailable default to true when enabled is - // set, which would abort wizard runs on hosts that lack sandbox - // dependencies (e.g. Linux without bubblewrap). Wizard targets a - // broad set of user machines, so prefer graceful degradation — - // commands still respect allowUnsandboxedCommands below. - failIfUnavailable: false, - allowUnsandboxedCommands: false, - filesystem: { - allowWrite: [ - '/' + agentConfig.workingDirectory, - '/' + agentConfig.workingDirectory + '/**', - '//tmp', - '//tmp/**', - '//private/tmp', - '//private/tmp/**', - // Package manager stores and toolchain installs — allow writes - // so pnpm/npm/yarn/bun and version managers (corepack, volta) - // can install packages and self-update without breaking the - // user's existing setup. - '~/Library/pnpm/**', // pnpm root (macOS) — store + .tools/ for packageManager pinning - '~/.local/share/pnpm/**', // pnpm root (Linux) - '~/.pnpm-store/**', // pnpm alternate store - '~/.npm/**', // npm cache (covers _npx too) - '~/.yarn/**', // yarn classic + berry cache - '~/.bun/install/**', // bun cache + global installs - '~/.cache/node/corepack/**', // corepack version downloads (Linux/macOS) - '~/Library/Caches/node/corepack/**', // corepack on older macOS layouts - '~/.volta/**', // Volta toolchain (referenced by workbench package.json) - // Python — used by django/flask/fastapi wizards - '~/.cache/pip/**', - '~/Library/Caches/pip/**', - '~/.cache/uv/**', - '~/Library/Caches/uv/**', - '~/.cache/pypoetry/**', - '~/Library/Caches/pypoetry/**', - // Ruby — used by rails wizard - '~/.bundle/**', - '~/.gem/**', - ], - }, - network: { - allowedDomains: [ - 'github.com', - 'api.github.com', - 'raw.githubusercontent.com', - 'release-assets.githubusercontent.com', - 'objects.githubusercontent.com', - ], + const runQuery = async (resume?: string): Promise<'done' | 'remint'> => { + const response = query({ + prompt: createPromptStream(resume ? RESUME_INSTRUCTION : prompt), + options: { + abortController, + resume, + model: agentConfig.model, + cwd: agentConfig.workingDirectory, + permissionMode: 'acceptEdits', + betas: ['context-1m-2025-08-07'], + mcpServers: agentConfig.mcpServers, + agents: { + 'general-purpose': { + description: + "General-purpose subagent. Inherits the parent run's tools plus the PostHog and wizard-tools MCP servers, so it can call mcp__posthog-wizard__* directly instead of curling the REST API.", + prompt: + 'You are a general-purpose subagent for the PostHog wizard. Prefer the authenticated mcp__posthog-wizard__* MCP tools over raw HTTP — they are already authenticated for this project. Only fall back to other transports if no MCP tool covers the operation.', + mcpServers: inheritedMcpServerNames, + // SDK does not propagate the parent's disallowedTools to subagents + // (sdk.d.ts: AgentDefinition has its own disallowedTools, and + // `tools: undefined` means "inherit all"). Without this, a program + // that disallows wizard_ask still leaks it to dispatched subagents. + disallowedTools: agentConfig.disallowedTools + ? [...agentConfig.disallowedTools] + : undefined, + }, }, - }, - env: { - // Drop the ENTIRE ANTHROPIC_*/CLAUDE_CODE_* namespace from the - // inherited env so no shell/settings value can leak into or outrank - // the agent's routing; the wizard's own gateway routing is injected - // fresh below. See agent-env-isolation.ts. - ...sanitizeAgentSubprocessEnv(process.env), - // Gateway routing — injected explicitly (initializeAgent set these on - // process.env for in-process readers; the strip above removed them - // from the inherited copy, so re-add the wizard's own values here). - // From this run's resolved auth, not process.env: concurrent task - // runs each write those globals, so re-reading them here would hand - // a subprocess whichever run initialized last. - ANTHROPIC_BASE_URL: agentConfig.gatewayAuth.gatewayUrl, - ANTHROPIC_AUTH_TOKEN: agentConfig.gatewayAuth.token, - CLAUDE_CODE_OAUTH_TOKEN: agentConfig.gatewayAuth.token, - // Point the binary at an empty config dir so it cannot resolve a - // stored Claude login (a `~/.claude/.credentials.json`) and send that - // to the gateway, which 401s it. The env token above is then the only - // credential it can find. See stored-login.ts. - CLAUDE_CONFIG_DIR: createIsolatedAgentConfigDir(), - CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS: 'true', - // The MCP config resolves this in the child; sending the value would - // put it on the CLI's argv. - POSTHOG_MCP_TOKEN: agentConfig.posthogApiKey, - // SDK 0.3.142 made MCP servers connect in the background by default; - // the agent may start its first turn before posthog-wizard is ready - // (audit programs call audit_seed_checks on turn 1, integration - // programs call load_skill_menu / install_skill). Restore the prior - // blocking behavior so the SDK waits up to 5s for MCP connect before - // turn 1. - MCP_CONNECTION_NONBLOCKING: '0', - // PostHog gateway headers: this run's properties blob. - ANTHROPIC_CUSTOM_HEADERS: buildAgentEnv( - agentConfig.wizardMetadata ?? {}, - agentConfig.wizardFlags ?? {}, - agentConfig.gatewayAuth.teamId, - ), - }, - canUseTool: (toolName: string, input: unknown) => { - logToFile('canUseTool called:', { toolName, input }); - const result = wizardCanUseTool( - toolName, - input as Record, - { - wizardAskPending: agentConfig.getPendingQuestion?.() != null, - disallowedTools: agentConfig.disallowedTools, + // Load skills from project's .claude/skills/ directory + settingSources: ['project'], + // Enable all discovered skills. Omitting this is NOT "skills off" — + // it just means no SDK auto-config — so we set 'all' explicitly to + // preserve the prior behavior where 'Skill' in allowedTools exposed + // everything under .claude/skills/. (SDK ≥0.2.133 deprecates passing + // 'Skill' in allowedTools in favor of this option.) + skills: 'all', + allowedTools, + sandbox: { + enabled: true, + // SDK 0.2.91 made failIfUnavailable default to true when enabled is + // set, which would abort wizard runs on hosts that lack sandbox + // dependencies (e.g. Linux without bubblewrap). Wizard targets a + // broad set of user machines, so prefer graceful degradation — + // commands still respect allowUnsandboxedCommands below. + failIfUnavailable: false, + allowUnsandboxedCommands: false, + filesystem: { + allowWrite: [ + '/' + agentConfig.workingDirectory, + '/' + agentConfig.workingDirectory + '/**', + '//tmp', + '//tmp/**', + '//private/tmp', + '//private/tmp/**', + // Package manager stores and toolchain installs — allow writes + // so pnpm/npm/yarn/bun and version managers (corepack, volta) + // can install packages and self-update without breaking the + // user's existing setup. + '~/Library/pnpm/**', // pnpm root (macOS) — store + .tools/ for packageManager pinning + '~/.local/share/pnpm/**', // pnpm root (Linux) + '~/.pnpm-store/**', // pnpm alternate store + '~/.npm/**', // npm cache (covers _npx too) + '~/.yarn/**', // yarn classic + berry cache + '~/.bun/install/**', // bun cache + global installs + '~/.cache/node/corepack/**', // corepack version downloads (Linux/macOS) + '~/Library/Caches/node/corepack/**', // corepack on older macOS layouts + '~/.volta/**', // Volta toolchain (referenced by workbench package.json) + // Python — used by django/flask/fastapi wizards + '~/.cache/pip/**', + '~/Library/Caches/pip/**', + '~/.cache/uv/**', + '~/Library/Caches/uv/**', + '~/.cache/pypoetry/**', + '~/Library/Caches/pypoetry/**', + // Ruby — used by rails wizard + '~/.bundle/**', + '~/.gem/**', + ], }, - ); - logToFile('canUseTool result:', result); - return Promise.resolve(result); - }, - systemPrompt: { - type: 'preset', - preset: 'claude_code', - // Append the run's commandments rather than replacing the preset so - // we keep default Claude Code behaviors. An orchestrator context is - // present only on a task run — that is what picks the sequence. - append: assembleCommandments({ - program: agentConfig.program, - sequence: agentConfig.sequence, - harness: Harness.anthropic, - }), - }, - tools: { type: 'preset', preset: 'claude_code' }, - // Capture stderr from CLI subprocess for debugging - stderr: (data: string) => { - logToFile('CLI stderr:', data); - if (options.debug) { - debug('CLI stderr:', data); - } - }, - // Stop hook: drain additional feature queue, then collect remark, then allow stop - hooks: { - PreToolUse: warlockDisabled - ? [] - : createPreToolUseYaraHooks(triageProvider, onYaraTerminate), - PostToolUse: warlockDisabled - ? [] - : createPostToolUseYaraHooks(triageProvider, onYaraTerminate), - Stop: [ - { - hooks: [ - createStopHook( - config?.additionalFeatureQueue ?? [], - signals, - config?.requestRemark ?? true, - ), + network: { + allowedDomains: [ + 'github.com', + 'api.github.com', + 'raw.githubusercontent.com', + 'release-assets.githubusercontent.com', + 'objects.githubusercontent.com', ], - timeout: 30, }, - ], + }, + env: { + // Drop the ENTIRE ANTHROPIC_*/CLAUDE_CODE_* namespace from the + // inherited env so no shell/settings value can leak into or outrank + // the agent's routing; the wizard's own gateway routing is injected + // fresh below. See agent-env-isolation.ts. + ...sanitizeAgentSubprocessEnv(process.env), + // Gateway routing — injected explicitly (initializeAgent set these on + // process.env for in-process readers; the strip above removed them + // from the inherited copy, so re-add the wizard's own values here). + // From this run's resolved auth, not process.env: concurrent task + // runs each write those globals, so re-reading them here would hand + // a subprocess whichever run initialized last. + ANTHROPIC_BASE_URL: agentConfig.gatewayAuth.gatewayUrl, + ANTHROPIC_AUTH_TOKEN: agentConfig.gatewayAuth.token, + CLAUDE_CODE_OAUTH_TOKEN: agentConfig.gatewayAuth.token, + // Point the binary at an empty config dir so it cannot resolve a + // stored Claude login (a `~/.claude/.credentials.json`) and send that + // to the gateway, which 401s it. The env token above is then the only + // credential it can find. See stored-login.ts. Shared by a resumed + // query so it finds the transcript. + CLAUDE_CONFIG_DIR: agentConfigDir, + CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS: 'true', + // The MCP config resolves this in the child; sending the value would + // put it on the CLI's argv. + POSTHOG_MCP_TOKEN: agentConfig.posthogApiKey, + // SDK 0.3.142 made MCP servers connect in the background by default; + // the agent may start its first turn before posthog-wizard is ready + // (audit programs call audit_seed_checks on turn 1, integration + // programs call load_skill_menu / install_skill). Restore the prior + // blocking behavior so the SDK waits up to 5s for MCP connect before + // turn 1. + MCP_CONNECTION_NONBLOCKING: '0', + // PostHog gateway headers: this run's properties blob. + ANTHROPIC_CUSTOM_HEADERS: buildAgentEnv( + agentConfig.wizardMetadata ?? {}, + agentConfig.wizardFlags ?? {}, + agentConfig.gatewayAuth.teamId, + ), + }, + canUseTool: (toolName: string, input: unknown) => { + logToFile('canUseTool called:', { toolName, input }); + const result = wizardCanUseTool( + toolName, + input as Record, + { + wizardAskPending: agentConfig.getPendingQuestion?.() != null, + disallowedTools: agentConfig.disallowedTools, + }, + ); + logToFile('canUseTool result:', result); + return Promise.resolve(result); + }, + systemPrompt: { + type: 'preset', + preset: 'claude_code', + // Append the run's commandments rather than replacing the preset so + // we keep default Claude Code behaviors. An orchestrator context is + // present only on a task run — that is what picks the sequence. + append: assembleCommandments({ + program: agentConfig.program, + sequence: agentConfig.sequence, + harness: Harness.anthropic, + }), + }, + tools: { type: 'preset', preset: 'claude_code' }, + // Capture stderr from CLI subprocess for debugging + stderr: (data: string) => { + logToFile('CLI stderr:', data); + if (options.debug) { + debug('CLI stderr:', data); + } + }, + // Stop hook: drain additional feature queue, then collect remark, then allow stop + hooks: { + PreToolUse: warlockDisabled + ? [] + : createPreToolUseYaraHooks(triageProvider, onYaraTerminate), + PostToolUse: warlockDisabled + ? [] + : createPostToolUseYaraHooks(triageProvider, onYaraTerminate), + Stop: [ + { + hooks: [ + createStopHook( + config?.additionalFeatureQueue ?? [], + signals, + config?.requestRemark ?? true, + ), + ], + timeout: 30, + }, + ], + }, }, - }, - }); + }); - // Process the async generator - for await (const message of response) { - // Log initial context size on the first assistant response so we can - // detect sudden shifts in starting context (e.g. MCP schema bloat). - if (!loggedInitialContext && message.type === 'assistant') { - const usage = message.message?.usage as - | { - input_tokens?: number; - cache_creation_input_tokens?: number; - cache_read_input_tokens?: number; + // Process the async generator + try { + for await (const message of response) { + if (typeof message.session_id === 'string' && message.session_id) { + sessionId = message.session_id; + } + // Log initial context size on the first assistant response so we can + // detect sudden shifts in starting context (e.g. MCP schema bloat). + if (!loggedInitialContext && message.type === 'assistant') { + const usage = message.message?.usage as + | { + input_tokens?: number; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; + } + | undefined; + if (usage) { + const input = usage.input_tokens ?? 0; + const cacheCreation = usage.cache_creation_input_tokens ?? 0; + const cacheRead = usage.cache_read_input_tokens ?? 0; + const initialTokens = input + cacheCreation + cacheRead; + logToFile( + `Initial context: ${initialTokens} tokens (input=${input}, cache_creation=${cacheCreation}, cache_read=${cacheRead})`, + ); + analytics.wizardCapture('agent initial context', { + initial_tokens: initialTokens, + input_tokens: input, + cache_creation_input_tokens: cacheCreation, + cache_read_input_tokens: cacheRead, + }); } - | undefined; - if (usage) { - const input = usage.input_tokens ?? 0; - const cacheCreation = usage.cache_creation_input_tokens ?? 0; - const cacheRead = usage.cache_read_input_tokens ?? 0; - const initialTokens = input + cacheCreation + cacheRead; - logToFile( - `Initial context: ${initialTokens} tokens (input=${input}, cache_creation=${cacheCreation}, cache_read=${cacheRead})`, - ); - analytics.wizardCapture('agent initial context', { - initial_tokens: initialTokens, - input_tokens: input, - cache_creation_input_tokens: cacheCreation, - cache_read_input_tokens: cacheRead, - }); - } - loggedInitialContext = true; - } + loggedInitialContext = true; + } - // Mirror the assistant turn into the authenticated project's AIO tab. - // No-op when `--capture-aio` is off (dev/test builds only). Fire-and- - // forget: failures are debug-logged inside the module and never touch - // the stream loop. - agentConfig.capture?.captureFromAnthropicSDKMessage(message); - - // Pass receivedSuccessResult so handleSDKMessage can suppress user-facing error - // output for post-success cleanup errors while still logging them to file - handleSDKMessage( - message, - options, - spinner, - signals, - receivedSuccessResult, - tasks, - agentConfig.suppressTaskRender ?? false, - emitStepEvents, - resolveStepKey, - ); + // Mirror the assistant turn into the authenticated project's AIO tab. + // No-op when `--capture-aio` is off (dev/test builds only). Fire-and- + // forget: failures are debug-logged inside the module and never touch + // the stream loop. + agentConfig.capture?.captureFromAnthropicSDKMessage(message); + + // Pass receivedSuccessResult so handleSDKMessage can suppress user-facing error + // output for post-success cleanup errors while still logging them to file + handleSDKMessage( + message, + options, + spinner, + signals, + receivedSuccessResult, + tasks, + agentConfig.suppressTaskRender ?? false, + emitStepEvents, + resolveStepKey, + ); - // [ABORT] detection: the skill emits "[ABORT] " when it - // cannot complete the program. Kill the SDK query immediately — - // the prompt doesn't need to cooperate with "and exit" because the - // abort is enforced here. The reason is surfaced via the returned - // AgentErrorType.ABORT so the runner can render a custom screen. - if ( - abortCases.length > 0 && - !abortReason && - message.type === 'assistant' - ) { - const content = message.message?.content; - if (Array.isArray(content)) { - for (const block of content) { - if (block.type === 'text' && typeof block.text === 'string') { - const match = block.text.match(/\[ABORT\]\s*(.+?)(?:\n|$)/); - if (match) { - abortReason = match[1].trim(); - logToFile(`Agent emitted [ABORT]: ${abortReason}`); - abortController.abort(); - signalDone!(); - break; + // [ABORT] detection: the skill emits "[ABORT] " when it + // cannot complete the program. Kill the SDK query immediately — + // the prompt doesn't need to cooperate with "and exit" because the + // abort is enforced here. The reason is surfaced via the returned + // AgentErrorType.ABORT so the runner can render a custom screen. + if ( + abortCases.length > 0 && + !abortReason && + message.type === 'assistant' + ) { + const content = message.message?.content; + if (Array.isArray(content)) { + for (const block of content) { + if (block.type === 'text' && typeof block.text === 'string') { + const match = block.text.match(/\[ABORT\]\s*(.+?)(?:\n|$)/); + if (match) { + abortReason = match[1].trim(); + logToFile(`Agent emitted [ABORT]: ${abortReason}`); + abortController.abort(); + signalDone(); + break; + } + } } } } - } - } - // 401: show auth error screen and exit immediately - if (message.type === 'assistant' && signals.hasApiErrorStatus(401)) { - signalDone!(); - spinner.stop('Authentication failed'); - // Re-check at error time: a settings conflict can be the *real* cause - // of a 401, distinct from bad PAT / wrong region / expired key. - // Only the conflict case warrants telling the user to log out of - // Claude Code. - const authError = buildAuthErrorContext( - options.installDir, - agentConfig.gatewayAuth.gatewayUrl, - os.homedir(), - signals.apiKeySource, - ); - // A refresh that already failed on a dead grant explains this 401 - // outright; without it the screen falls through to generic key-type - // and scope advice that cannot apply. - const sessionExpired = isGrantRevoked(); - const authCode = classifyAuthFailure({ - hasSettingsConflict: authError.hasSettingsConflict, - usingManagedLogin: authError.usingManagedLogin, - sessionExpired, - apiKey: options.apiKey, - gatewayRegion: authError.region, - sessionRegion: options.cloudRegion, - }); - logToFile('Agent error: 401, showing auth error screen', { - ...authError, - sessionExpired, - }); - getUI().showAuthError({ - hasSettingsConflict: authError.hasSettingsConflict, - conflicts: authError.conflicts, - usingManagedLogin: authError.usingManagedLogin, - credentialPlaces: authError.credentialPlaces, - sessionExpired, - logFilePath: getLogFilePath(), - }); - await wizardAbort({ - code: authCode, - message: 'Authentication failed (401)', - error: new WizardError( - 'Authentication failed', - { + // 401 on a bearer past its refresh instant: it aged out, so re-mint + // once and resume. Any other 401 is a bad credential: show the auth + // error screen and exit. + if (message.type === 'assistant' && signals.hasApiErrorStatus(401)) { + signalDone(); + if ( + agentConfig.refreshGatewayAuth && + !reminted && + isPastRefresh(agentConfig.gatewayAuth) + ) { + logToFile( + 'Agent error: 401 on an aged gateway bearer; re-minting', + ); + remintRequested = true; + abortController.abort(); + break; + } + spinner.stop('Authentication failed'); + // Re-check at error time: a settings conflict can be the *real* cause + // of a 401, distinct from bad PAT / wrong region / expired key. + // Only the conflict case warrants telling the user to log out of + // Claude Code. + const authError = buildAuthErrorContext( + options.installDir, + agentConfig.gatewayAuth.gatewayUrl, + os.homedir(), + signals.apiKeySource, + ); + // A refresh that already failed on a dead grant explains this 401 + // outright; without it the screen falls through to generic key-type + // and scope advice that cannot apply. + const sessionExpired = isGrantRevoked(); + const authCode = classifyAuthFailure({ hasSettingsConflict: authError.hasSettingsConflict, - conflictSources: authError.conflictSources, - conflictKeys: authError.conflictKeys, - gatewayUrl: authError.gatewayUrl, - region: authError.region, usingManagedLogin: authError.usingManagedLogin, - apiKeySource: authError.apiKeySource, - }, - authCode, - ), - }); - } + sessionExpired, + apiKey: options.apiKey, + gatewayRegion: authError.region, + sessionRegion: options.cloudRegion, + }); + logToFile('Agent error: 401, showing auth error screen', { + ...authError, + sessionExpired, + }); + getUI().showAuthError({ + hasSettingsConflict: authError.hasSettingsConflict, + conflicts: authError.conflicts, + usingManagedLogin: authError.usingManagedLogin, + credentialPlaces: authError.credentialPlaces, + sessionExpired, + logFilePath: getLogFilePath(), + }); + await wizardAbort({ + code: authCode, + message: 'Authentication failed (401)', + error: new WizardError( + 'Authentication failed', + { + hasSettingsConflict: authError.hasSettingsConflict, + conflictSources: authError.conflictSources, + conflictKeys: authError.conflictKeys, + gatewayUrl: authError.gatewayUrl, + region: authError.region, + usingManagedLogin: authError.usingManagedLogin, + apiKeySource: authError.apiKeySource, + }, + authCode, + ), + }); + } - try { - middleware?.onMessage(message); - } catch (e) { - logToFile(`${AgentSignals.BENCHMARK} Middleware onMessage error:`, e); - } + try { + middleware?.onMessage(message); + } catch (e) { + logToFile( + `${AgentSignals.BENCHMARK} Middleware onMessage error:`, + e, + ); + } - // Signal completion when result received - if (message.type === 'result') { - // Track successful results before any potential cleanup errors - // The SDK may emit a second error result during cleanup due to a race condition - if (message.subtype === 'success' && !message.is_error) { - receivedSuccessResult = true; - lastResultMessage = message; + // Signal completion when result received + if (message.type === 'result') { + // Track successful results before any potential cleanup errors + // The SDK may emit a second error result during cleanup due to a race condition + if (message.subtype === 'success' && !message.is_error) { + receivedSuccessResult = true; + lastResultMessage = message; + } + signalDone(); + } } - signalDone!(); + } catch (error) { + // The abort we asked for; anything else belongs to the outer catch. + if (remintRequested) return 'remint'; + throw error; } + return remintRequested ? 'remint' : 'done'; + }; + + const refreshGatewayAuth = agentConfig.refreshGatewayAuth; + if ((await runQuery()) === 'remint' && refreshGatewayAuth) { + // The subprocess froze the dead bearer in its env at spawn, so it cannot + // be handed a new one: mint, then resume the session in a new one. + reminted = true; + remintRequested = false; + abortController = new AbortController(); + signals.forgetApiErrors(); + spinner.message('Renewing the gateway token...'); + const stale = agentConfig.gatewayAuth; + // A refusal or failure here ends the run with its own message. + agentConfig.gatewayAuth = await refreshGatewayAuth(); + logToFile( + `Gateway token renewed after a 401 (${Math.round( + (Date.now() - stale.refreshAtMs) / 1000, + )}s past refresh); resuming session ${ + sessionId ?? '(none: fresh session)' + }`, + ); + analytics.wizardCapture('gateway token reminted', { + resumed: sessionId !== undefined, + }); + spinner.message(spinnerMessage); + await runQuery(sessionId); } // A YARA hook detected a terminal violation and aborted the run. @@ -1294,7 +1372,7 @@ export async function runAgent( return completeWithSuccess(); } catch (error) { // Signal done to unblock the async generator - signalDone!(); + signalDone(); // A YARA hook aborted the run (the SDK throws AbortError once the hook // calls abortController.abort()). Surface it before anything else so it is diff --git a/src/lib/agent/output-signals.ts b/src/lib/agent/output-signals.ts index 184b7f7d0..f5e543c75 100644 --- a/src/lib/agent/output-signals.ts +++ b/src/lib/agent/output-signals.ts @@ -41,6 +41,18 @@ export class AgentOutputSignals { if (SIGNAL_NEEDLES.some((n) => text.includes(n))) this.lines.push(text); } + /** + * Drop the retained API-error lines and keep every other signal. Used after + * a re-mint so the run is judged on the resumed session, not on the 401 + * that ended the first one. + */ + forgetApiErrors(): void { + const kept = this.lines.filter( + (line) => !line.includes(OUTPUT_SIGNALS.API_ERROR), + ); + this.lines.splice(0, this.lines.length, ...kept); + } + /** * Record the SDK's `apiKeySource` from its `init` message (e.g. * `"/login managed key"`, `"ANTHROPIC_API_KEY"`). Used to triage a 401: diff --git a/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts b/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts index 09c4d1fff..a0cb47201 100644 --- a/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts +++ b/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts @@ -1,4 +1,11 @@ -import { buildGatewayProvider, buildGatewayHeaders } from '../gateway'; +import { + buildGatewayProvider, + buildGatewayHeaders, + isGatewayAuthRejection, + withGatewayRemint, + GATEWAY_PROVIDER, +} from '../gateway'; +import type { GatewayAuth } from '@lib/gateway-session'; describe('buildGatewayProvider effort', () => { const base = { @@ -78,3 +85,132 @@ describe('buildGatewayHeaders', () => { expect(headers['X-POSTHOG-PROPERTY-run_id']).toBeUndefined(); }); }); + +describe('isGatewayAuthRejection', () => { + it.each([ + 'OpenAI API error (401): token expired', + '401 {"type":"error","error":{"type":"authentication_error"}}', + 'Unauthorized', + ])('recognises %s', (message) => { + expect(isGatewayAuthRejection(message)).toBe(true); + }); + + it.each([ + 'OpenAI API error (429): rate limit', + 'connection reset', + undefined, + ])('ignores %s', (message) => { + expect(isGatewayAuthRejection(message)).toBe(false); + }); +}); + +describe('withGatewayRemint', () => { + const HOUR = 3600_000; + const rejected = { + role: 'assistant', + stopReason: 'error', + errorMessage: 'OpenAI API error (401): token expired', + }; + const fine = { role: 'assistant', stopReason: 'stop' }; + const gatewayAuth = (token: string, refreshAtMs: number): GatewayAuth => ({ + gatewayUrl: 'https://ai-gateway.us.posthog.com', + token, + teamId: 42, + refreshAtMs, + }); + + // A fake session: each prompt ends on the next scripted turn. The re-minted + // bearer is fresh unless a test ages it to pin the once-per-session rule. + function harness( + auth: GatewayAuth, + turns: unknown[], + remintedAt = Date.now() + HOUR, + ) { + const prompts: string[] = []; + const session = { + prompt: vi.fn((text: string) => { + prompts.push(text); + wrapped.noteAssistantTurn(turns.shift() ?? fine); + return Promise.resolve(); + }), + }; + const registry = { registerProvider: vi.fn() }; + const refreshAuth = vi + .fn() + .mockResolvedValue(gatewayAuth('phe_new', remintedAt)); + const wrapped = withGatewayRemint({ + session, + registry, + auth, + refreshAuth, + providerInputs: (a) => ({ + gatewayUrl: a.gatewayUrl, + accessToken: a.token, + teamId: a.teamId, + wizardMetadata: {}, + wizardFlags: {}, + modelId: 'openai/gpt-5.6-terra', + }), + continueText: 'continue', + }); + return { wrapped, registry, refreshAuth, prompts }; + } + + it('re-mints once and continues when a turn ends on a 401 from an aged bearer', async () => { + const { wrapped, registry, refreshAuth, prompts } = harness( + gatewayAuth('phe_old', Date.now() - 1), + [rejected, fine], + ); + + await wrapped.prompt('do it'); + + expect(refreshAuth).toHaveBeenCalledTimes(1); + // pi resolves the apiKey per request, so the re-registered provider is + // what the continued turn signs with. + expect(registry.registerProvider).toHaveBeenCalledWith( + GATEWAY_PROVIDER, + expect.objectContaining({ apiKey: 'phe_new' }), + ); + expect(prompts).toEqual(['do it', 'continue']); + }); + + it('leaves a 401 on a fresh bearer to the harness', async () => { + const { wrapped, refreshAuth, prompts } = harness( + gatewayAuth('phe_fresh', Date.now() + HOUR), + [rejected], + ); + + await wrapped.prompt('do it'); + + expect(refreshAuth).not.toHaveBeenCalled(); + expect(prompts).toEqual(['do it']); + }); + + it('does not re-mint a second time', async () => { + // Even with the re-minted bearer already past refresh, one mint per + // session is the rule. + const { wrapped, refreshAuth, prompts } = harness( + gatewayAuth('phe_old', Date.now() - 1), + [rejected, rejected, rejected], + Date.now() - 1, + ); + + await wrapped.prompt('do it'); + await wrapped.prompt('again'); + + expect(refreshAuth).toHaveBeenCalledTimes(1); + expect(prompts).toEqual(['do it', 'continue', 'again']); + }); + + it('ignores a turn that ended without an auth error', async () => { + const { wrapped, refreshAuth, prompts } = harness( + gatewayAuth('phe_old', Date.now() - 1), + [fine], + ); + + await wrapped.prompt('do it'); + + expect(refreshAuth).not.toHaveBeenCalled(); + expect(prompts).toEqual(['do it']); + }); +}); diff --git a/src/lib/agent/runner/harness/pi/gateway.ts b/src/lib/agent/runner/harness/pi/gateway.ts index d1c76ae0e..3b6b5dfea 100644 --- a/src/lib/agent/runner/harness/pi/gateway.ts +++ b/src/lib/agent/runner/harness/pi/gateway.ts @@ -6,7 +6,11 @@ * (lazily imported, properly typed) pi ModelRegistry. */ -import { buildWizardPropertiesBlob } from '@lib/gateway-session'; +import { + buildWizardPropertiesBlob, + isPastRefresh, + type GatewayAuth, +} from '@lib/gateway-session'; import { modelCapabilities, type ThinkingLevel, @@ -135,3 +139,66 @@ export function buildGatewayProvider(inputs: GatewayProviderInputs): { }; return { provider, api, caps, gatewayUrl, baseUrl: model.baseUrl }; } + +/** + * Whether a turn's error is the gateway rejecting the bearer. pi-ai keeps the + * HTTP status in its error text; Anthropic's SDK also names the error type. + */ +export function isGatewayAuthRejection( + errorMessage: string | undefined, +): boolean { + return /\b401\b|authentication_error|unauthorized/i.test(errorMessage ?? ''); +} + +export interface GatewayRemintOptions { + session: { prompt(text: string): Promise }; + registry: { registerProvider(providerName: string, config: never): void }; + auth: GatewayAuth; + /** The cache: the same token while fresh, a new mint past the refresh point. */ + refreshAuth: () => Promise; + providerInputs: (auth: GatewayAuth) => GatewayProviderInputs; + /** The prompt that resumes the work after a re-mint. */ + continueText: string | (() => string); + onRemint?: () => void; +} + +/** + * Wraps a pi session's prompt(): when a turn ends on a 401 from a bearer past + * its refresh instant, mint once, re-register the provider with the new + * bearer (pi resolves the apiKey per request), and continue. A 401 on a fresh + * bearer, or a second one, is left to the harness's normal failure path. + */ +export function withGatewayRemint(opts: GatewayRemintOptions): { + prompt(text: string): Promise; + /** Feed every assistant `message_end`; the last turn decides. */ + noteAssistantTurn(message: unknown): void; +} { + let auth = opts.auth; + let rejected = false; + let reminted = false; + return { + noteAssistantTurn(message) { + const turn = message as + | { stopReason?: string; errorMessage?: string } + | undefined; + rejected = + turn?.stopReason === 'error' && + isGatewayAuthRejection(turn.errorMessage); + }, + async prompt(text) { + rejected = false; + await opts.session.prompt(text); + if (!rejected || reminted || !isPastRefresh(auth)) return; + reminted = true; + auth = await opts.refreshAuth(); + opts.registry.registerProvider( + GATEWAY_PROVIDER, + buildGatewayProvider(opts.providerInputs(auth)).provider as never, + ); + opts.onRemint?.(); + rejected = false; + const next = opts.continueText; + await opts.session.prompt(typeof next === 'function' ? next() : next); + }, + }; +} diff --git a/src/lib/agent/runner/harness/pi/index.ts b/src/lib/agent/runner/harness/pi/index.ts index 1ee59d1a0..a6028ff0e 100644 --- a/src/lib/agent/runner/harness/pi/index.ts +++ b/src/lib/agent/runner/harness/pi/index.ts @@ -27,8 +27,12 @@ import { AgentErrorType } from '@lib/agent/agent-interface'; import { AgentSignals, REMARK_INSTRUCTION } from '@lib/agent/signals'; import { AgentOutputSignals } from '@lib/agent/output-signals'; import { assembleCommandments } from '../../switchboard/commandments'; -import { gatewayAuth } from '@lib/gateway-session'; -import { buildGatewayProvider, GATEWAY_PROVIDER } from './gateway'; +import { gatewayAuth, type GatewayAuth } from '@lib/gateway-session'; +import { + buildGatewayProvider, + GATEWAY_PROVIDER, + withGatewayRemint, +} from './gateway'; import { createAioCapture } from '@lib/agent/aio-capture'; import type { AgentResult, @@ -248,20 +252,23 @@ export const piBackend: AgentHarness = { // the claude-agent-sdk path. The provider spec is shared with the // orchestrator's per-task sessions (gateway.ts). gatewayAuth mints the // run's scoped token. - const auth = await gatewayAuth( - boot.credentials.host, - boot.credentials.accessToken, - boot.programId, - ); - const { provider, caps } = buildGatewayProvider({ - gatewayUrl: auth.gatewayUrl, - accessToken: auth.token, - teamId: auth.teamId, + const refreshAuth = () => + gatewayAuth( + boot.credentials.host, + boot.credentials.accessToken, + boot.programId, + ); + const auth = await refreshAuth(); + const providerInputs = (current: GatewayAuth) => ({ + gatewayUrl: current.gatewayUrl, + accessToken: current.token, + teamId: current.teamId, wizardMetadata: boot.wizardMetadata, wizardFlags: boot.wizardFlags, modelId, effort: inputs.thinkingLevel, }); + const { provider, caps } = buildGatewayProvider(providerInputs(auth)); const registry = ModelRegistry.inMemory(AuthStorage.create()); registry.registerProvider(GATEWAY_PROVIDER, provider as never); @@ -454,6 +461,22 @@ export const piBackend: AgentHarness = { // event; without this its tools report "MCP not initialized". await agentSession.bindExtensions({}); + // A turn that ends on a 401 from an aged bearer re-mints once and + // continues; pi resolves the provider's apiKey per request, so + // re-registering is enough. + const turns = withGatewayRemint({ + session: agentSession, + registry, + auth, + refreshAuth, + providerInputs, + continueText: CONTINUE_INSTRUCTION, + onRemint: () => { + logToFile('[pi] gateway token renewed after a 401; continuing'); + analytics.wizardCapture('gateway token reminted', { harness: 'pi' }); + }, + }); + // Map pi events onto the run spinner + the log file, mirroring the // anthropic path's log shape (assistant turns + tool I/O) and driving the // single run spinner with one stable status at a time (no overlap). @@ -470,6 +493,7 @@ export const piBackend: AgentHarness = { break; } assistantTurns += 1; + turns.noteAssistantTurn(event.message); const assistant = extractText(event.message).trim(); if (assistant) { logToFile(`[pi] assistant: ${assistant.slice(0, 1000)}`); @@ -526,7 +550,7 @@ export const piBackend: AgentHarness = { try { // Non-streaming: resolves when the agent run completes. Throws if no // model/api key, or on a transport error. - await agentSession.prompt(prompt); + await turns.prompt(prompt); // Completion guard: pi's prompt() resolves the moment the model returns // a turn with no tool call (e.g. a lone [STATUS] line), even mid-plan. @@ -541,7 +565,7 @@ export const piBackend: AgentHarness = { logToFile( `[pi] completion guard: tasks still open, nudge ${continueNudges}/${MAX_CONTINUE_NUDGES}`, ); - await agentSession.prompt(CONTINUE_INSTRUCTION); + await turns.prompt(CONTINUE_INSTRUCTION); } // Best-effort remark ask — a failed turn never fails a successful run. diff --git a/src/lib/agent/runner/harness/pi/task.ts b/src/lib/agent/runner/harness/pi/task.ts index 091fccfff..a6e575def 100644 --- a/src/lib/agent/runner/harness/pi/task.ts +++ b/src/lib/agent/runner/harness/pi/task.ts @@ -36,8 +36,12 @@ import { AgentOutputSignals } from '@lib/agent/output-signals'; import { TaskStatus } from '../../sequence/orchestrator/queue'; import type { OrchestratorToolsContext } from '../../sequence/orchestrator/queue-tools'; import type { AgentResult, TaskRunInputs } from '../types'; -import { gatewayAuth } from '@lib/gateway-session'; -import { buildGatewayProvider, GATEWAY_PROVIDER } from './gateway'; +import { gatewayAuth, type GatewayAuth } from '@lib/gateway-session'; +import { + buildGatewayProvider, + GATEWAY_PROVIDER, + withGatewayRemint, +} from './gateway'; import { assembleCommandments } from '../../switchboard/commandments'; import { applyOutroMarkers, @@ -215,15 +219,17 @@ export async function runPiTask(inputs: TaskRunInputs): Promise { createWriteToolDefinition, } = sdk; - const auth = await gatewayAuth( - boot.credentials.host, - boot.credentials.accessToken, - boot.programId, - ); - const { provider, caps } = buildGatewayProvider({ - gatewayUrl: auth.gatewayUrl, - accessToken: auth.token, - teamId: auth.teamId, + const refreshAuth = () => + gatewayAuth( + boot.credentials.host, + boot.credentials.accessToken, + boot.programId, + ); + const auth = await refreshAuth(); + const providerInputs = (current: GatewayAuth) => ({ + gatewayUrl: current.gatewayUrl, + accessToken: current.token, + teamId: current.teamId, wizardMetadata: boot.wizardMetadata, wizardFlags: boot.wizardFlags, modelId, @@ -231,6 +237,7 @@ export async function runPiTask(inputs: TaskRunInputs): Promise { // back to the model table. effort, }); + const { provider, caps } = buildGatewayProvider(providerInputs(auth)); const registry = ModelRegistry.inMemory(AuthStorage.create()); registry.registerProvider(GATEWAY_PROVIDER, provider as never); const model = registry.find(GATEWAY_PROVIDER, modelId); @@ -369,6 +376,22 @@ export async function runPiTask(inputs: TaskRunInputs): Promise { }); await agentSession.bindExtensions({}); + // A turn that ends on a 401 from an aged bearer re-mints once and + // continues with the nudge the task would get anyway. + const turns = withGatewayRemint({ + session: agentSession, + registry, + auth, + refreshAuth, + providerInputs, + continueText: () => + orchestrator.currentTaskId ? TASK_NUDGE : SEED_NUDGE, + onRemint: () => { + logToFile('[pi-task] gateway token renewed after a 401; continuing'); + analytics.wizardCapture('gateway token reminted', { harness: 'pi' }); + }, + }); + // The one complete list: exactly the tools registered on this session, in // the names the agent will call them by. posthog_exec binds as an extension. const toolNames = [ @@ -390,6 +413,7 @@ export async function runPiTask(inputs: TaskRunInputs): Promise { break; } assistantTurns += 1; + turns.noteAssistantTurn(event.message); const assistant = extractText(event.message).trim(); if (assistant) { logToFile(`[pi-task] assistant: ${assistant.slice(0, 1000)}`); @@ -429,7 +453,7 @@ export async function runPiTask(inputs: TaskRunInputs): Promise { capture.setInitialPrompt(taskPrompt); try { - await agentSession.prompt(taskPrompt); + await turns.prompt(taskPrompt); // pi's prompt() resolves the moment a turn carries no tool call — which // an agent mid-plan does emit. While the work has not reached its @@ -444,7 +468,7 @@ export async function runPiTask(inputs: TaskRunInputs): Promise { logToFile( `[pi-task] completion guard: not settled, nudge ${nudges}/${MAX_TASK_NUDGES}`, ); - await agentSession.prompt( + await turns.prompt( orchestrator.currentTaskId ? TASK_NUDGE : SEED_NUDGE, ); } diff --git a/src/lib/agent/runner/shared/bootstrap.ts b/src/lib/agent/runner/shared/bootstrap.ts index 152e48952..6c2bd794e 100644 --- a/src/lib/agent/runner/shared/bootstrap.ts +++ b/src/lib/agent/runner/shared/bootstrap.ts @@ -307,12 +307,12 @@ export async function bootstrapProgram( // set them — so downstream readers get a non-null type without asserting. const credentials = session.credentials!; - // Mint the run's scoped gateway token once for the boot. - const auth = await gatewayAuth( - credentials.host, - credentials.accessToken, - programConfig.id, - ); + // Mint now so a refusal fails the boot before any agent starts. Later + // readers re-resolve through the cache, which re-mints past the refresh + // point. + const currentGatewayAuth = () => + gatewayAuth(credentials.host, credentials.accessToken, programConfig.id); + await currentGatewayAuth(); return { skillsBaseUrl, @@ -327,14 +327,17 @@ export async function bootstrapProgram( // Resolved once, here: the only place holding both the switchboard inputs // and the gateway auth. Every skill install downstream reads it off boot. triageProvider: createTriageLLMProvider( - { - baseURL: auth.gatewayUrl, - authToken: auth.token, - teamId: auth.teamId, - // `call_type` splits scan spend out of the program's agent cost — - // same tag the in-run triage provider carries. - wizardMetadata: { ...wizardMetadata, call_type: CallType.yaraTriage }, - wizardFlags, + async () => { + const auth = await currentGatewayAuth(); + return { + baseURL: auth.gatewayUrl, + authToken: auth.token, + teamId: auth.teamId, + // `call_type` splits scan spend out of the program's agent cost, + // the same tag the in-run triage provider carries. + wizardMetadata: { ...wizardMetadata, call_type: CallType.yaraTriage }, + wizardFlags, + }; }, resolveHarness({ program: programConfig.id, diff --git a/src/lib/agent/signals.ts b/src/lib/agent/signals.ts index f66f3ee36..dd087e1fa 100644 --- a/src/lib/agent/signals.ts +++ b/src/lib/agent/signals.ts @@ -53,6 +53,13 @@ export type AgentSignal = (typeof AgentSignals)[keyof typeof AgentSignals]; */ export const REMARK_INSTRUCTION = `Reply with a single line that starts with ${AgentSignals.WIZARD_REMARK} and no other lines. In that line, state briefly what information or guidance would have been useful to have in the integration prompt or documentation for this run — specifically anything that would have prevented tool failures, erroneous edits, or other wasted turns.`; +/** + * First prompt of a session resumed after a mid-run re-mint. The transcript + * carries the work so far; the model only needs to pick it up. + */ +export const RESUME_INSTRUCTION = + 'Your previous request failed with a transient gateway authentication error that has since been fixed. Continue the task from where you left off.'; + /** * Error types that can be returned from agent execution. * These correspond to the error signals that the agent emits. diff --git a/src/lib/agent/triage-provider.ts b/src/lib/agent/triage-provider.ts index 293c38262..dfb03861d 100644 --- a/src/lib/agent/triage-provider.ts +++ b/src/lib/agent/triage-provider.ts @@ -7,7 +7,10 @@ import { Harness } from '@lib/constants'; import { logToFile } from '@utils/debug'; -import { buildGatewayModel } from '@lib/agent/runner/harness/pi/gateway'; +import { + buildGatewayModel, + gatewayApiFor, +} from '@lib/agent/runner/harness/pi/gateway'; import { modelCapabilities, triageModelFor, @@ -36,31 +39,38 @@ export interface TriageGatewayAuth { * Triage provider for a harness. Auth is always explicit: every caller already * holds the gateway url and the run's token, and reading them back out of * ANTHROPIC_* made an unauthed provider silent — it returned undefined, the - * caller failed closed, and a clean first-party skill got deleted. + * caller failed closed, and a clean first-party skill got deleted. A resolver + * is re-read on every call, so a run that re-mints mid-way scans with the + * current bearer rather than the one it started with. */ export function createTriageLLMProvider( - auth: TriageGatewayAuth, + auth: TriageGatewayAuth | (() => Promise), harness: Harness, ): LLMProvider { - const { baseURL, authToken } = auth; + const resolveAuth = + typeof auth === 'function' ? auth : () => Promise.resolve(auth); const modelId = triageModelFor(harness); - const model = buildGatewayModel({ - gatewayUrl: baseURL, - accessToken: authToken, - teamId: auth.teamId, - wizardMetadata: auth?.wizardMetadata ?? {}, - wizardFlags: auth?.wizardFlags ?? {}, - modelId, - }); const { reasoning, thinkingLevel } = modelCapabilities(modelId); logToFile( - `[YARA] triage provider ready (model: ${modelId}, api: ${model.api})`, + `[YARA] triage provider ready (model: ${modelId}, api: ${gatewayApiFor( + modelId, + )})`, ); return async (prompt: string): Promise => { // Lazy: pi-ai is a 5MB ESM tree, and this module is in the static graph of // every command. Same constraint as the pi harness's SDK imports. const { completeSimple } = await import('@earendil-works/pi-ai'); + const current = await resolveAuth(); + const authToken = current.authToken; + const model = buildGatewayModel({ + gatewayUrl: current.baseURL, + accessToken: authToken, + teamId: current.teamId, + wizardMetadata: current.wizardMetadata ?? {}, + wizardFlags: current.wizardFlags ?? {}, + modelId, + }); const controller = new AbortController(); const timer = setTimeout(() => controller.abort(), TRIAGE_TIMEOUT_MS); try { diff --git a/src/lib/gateway-session.ts b/src/lib/gateway-session.ts index 647492a26..b26da0238 100644 --- a/src/lib/gateway-session.ts +++ b/src/lib/gateway-session.ts @@ -19,6 +19,12 @@ export interface GatewayAuth { token: string; /** The team the mint verified; rides the blob so dashboards keep a breakdown. */ teamId?: number; + /** + * Instant past which a 401 on this bearer is age rather than a bad + * credential: the cache re-mints past it, and a session still holding the + * old bearer may re-mint once. Before it the mint has to be trusted. + */ + refreshAtMs: number; } interface CachedAuth { @@ -36,8 +42,8 @@ let cached: CachedAuth | null = null; let inFlight: { key: string; promise: Promise } | null = null; /** - * Adoption floor. The anthropic subprocess holds its credential for the whole - * session, so a token below this 401s mid-run. + * Adoption floor. The anthropic subprocess holds its credential until a 401 + * forces a re-mint, so a token below this would churn mints. */ const MIN_USABLE_TTL_MS = 2 * 60 * 1000; /** Re-resolve at this fraction of the token's life, leaving a usable remainder. */ @@ -111,6 +117,7 @@ async function resolveGatewayAuth( gatewayUrl: minted.gatewayUrl, token: minted.token, teamId: minted.teamId, + refreshAtMs: staleAtMs, }; cached = { key, auth, staleAtMs }; return auth; @@ -122,6 +129,11 @@ export function resetGatewaySession(): void { inFlight = null; } +/** Whether a 401 on this bearer may be age (past its refresh instant) rather than a bad credential. */ +export function isPastRefresh(auth: GatewayAuth, now = Date.now()): boolean { + return now >= auth.refreshAtMs; +} + /** * Whether a server-supplied origin may receive a bearer and prompt content: * https (loopback excepted), and either a posthog.com host or the one the run From 223bfb1419d1ab835a777012f8529b07189863db Mon Sep 17 00:00:00 2001 From: Edwin Lim Date: Tue, 8 Sep 2026 15:18:09 -0700 Subject: [PATCH 2/4] add aws context mill fallback to anthropic --- src/lib/agent/agent-interface.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/lib/agent/agent-interface.ts b/src/lib/agent/agent-interface.ts index 8d1098767..815498a25 100644 --- a/src/lib/agent/agent-interface.ts +++ b/src/lib/agent/agent-interface.ts @@ -21,6 +21,7 @@ import { WIZARD_REMARK_EVENT_NAME, wizardUserAgentForProgram, DEFAULT_AGENT_MODEL, + AWS_SKILLS_BASE_URL, } from '@lib/constants'; import { type AdditionalFeature, @@ -1011,6 +1012,10 @@ export async function runAgent( 'raw.githubusercontent.com', 'release-assets.githubusercontent.com', 'objects.githubusercontent.com', + // The AWS mirror GitHub downloads fail over to + // (fetch-retry.ts); without it the failover dies in the + // sandbox exactly when GitHub is down. + new URL(AWS_SKILLS_BASE_URL).hostname, ], }, }, From e4b081e57793addd46647238ace8ec86aa687f58 Mon Sep 17 00:00:00 2001 From: Brandon Leung Date: Tue, 8 Sep 2026 19:17:11 -0400 Subject: [PATCH 3/4] fix(gateway): decide a re-mint on pi's diagnostic code, not its error prose Co-Authored-By: Claude Opus 5 (1M context) --- .../harness/pi/__tests__/gateway.test.ts | 24 +++++++++++++++ src/lib/agent/runner/harness/pi/gateway.ts | 30 ++++++++++++++----- 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts b/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts index a0cb47201..38a78436b 100644 --- a/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts +++ b/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts @@ -102,6 +102,30 @@ describe('isGatewayAuthRejection', () => { ])('ignores %s', (message) => { expect(isGatewayAuthRejection(message)).toBe(false); }); + + it("reads pi's diagnostic code rather than the message text", () => { + expect( + isGatewayAuthRejection({ + errorMessage: 'the model is unhappy', + diagnostics: [{ error: { code: 401 } }], + }), + ).toBe(true); + expect( + isGatewayAuthRejection({ + errorMessage: 'the model is unhappy', + diagnostics: [{ error: { name: 'AuthenticationError' } }], + }), + ).toBe(true); + }); + + it('does not re-mint on a diagnostic that is not an auth rejection', () => { + expect( + isGatewayAuthRejection({ + errorMessage: 'rate limited', + diagnostics: [{ error: { code: 429, name: 'RateLimitError' } }], + }), + ).toBe(false); + }); }); describe('withGatewayRemint', () => { diff --git a/src/lib/agent/runner/harness/pi/gateway.ts b/src/lib/agent/runner/harness/pi/gateway.ts index 3b6b5dfea..c827bcfcb 100644 --- a/src/lib/agent/runner/harness/pi/gateway.ts +++ b/src/lib/agent/runner/harness/pi/gateway.ts @@ -140,13 +140,31 @@ export function buildGatewayProvider(inputs: GatewayProviderInputs): { return { provider, api, caps, gatewayUrl, baseUrl: model.baseUrl }; } +/** The part of a pi assistant turn that says why it failed. */ +export interface GatewayTurnError { + errorMessage?: string; + diagnostics?: { error?: { name?: string; code?: string | number } }[]; +} + /** - * Whether a turn's error is the gateway rejecting the bearer. pi-ai keeps the - * HTTP status in its error text; Anthropic's SDK also names the error type. + * Whether a turn's error is the gateway rejecting the bearer. pi attaches the + * SDK's own error to `diagnostics`, so its code decides when one is present. + * The message match is the fallback for a turn that failed before pi built a + * diagnostic, where the status survives only as prose. */ export function isGatewayAuthRejection( - errorMessage: string | undefined, + turn: GatewayTurnError | string | undefined, ): boolean { + const { errorMessage, diagnostics } = + typeof turn === 'string' + ? { errorMessage: turn, diagnostics: undefined } + : turn ?? {}; + for (const diagnostic of diagnostics ?? []) { + const code = diagnostic.error?.code; + if (code === 401 || code === '401') return true; + if (/^authentication_?error$/i.test(diagnostic.error?.name ?? '')) + return true; + } return /\b401\b|authentication_error|unauthorized/i.test(errorMessage ?? ''); } @@ -179,11 +197,9 @@ export function withGatewayRemint(opts: GatewayRemintOptions): { return { noteAssistantTurn(message) { const turn = message as - | { stopReason?: string; errorMessage?: string } + | ({ stopReason?: string } & GatewayTurnError) | undefined; - rejected = - turn?.stopReason === 'error' && - isGatewayAuthRejection(turn.errorMessage); + rejected = turn?.stopReason === 'error' && isGatewayAuthRejection(turn); }, async prompt(text) { rejected = false; From c80e97c6a55e404233ecf1393e4aa334939b48ab Mon Sep 17 00:00:00 2001 From: "Vincent (Wen Yu) Ge" <29069505+gewenyu99@users.noreply.github.com> Date: Tue, 8 Sep 2026 23:02:32 -0400 Subject: [PATCH 4/4] fix(models): use supported wizard models (#1229) --- src/lib/__tests__/agent-interface.test.ts | 4 ++-- .../__tests__/agent-prompt-loader.test.ts | 14 ++++++------- src/lib/agent/__tests__/token-pricing.test.ts | 6 ++++-- .../runner/__tests__/switchboard.test.ts | 21 ++++++++++++------- .../agent/runner/harness/anthropic/README.md | 2 +- .../harness/pi/__tests__/gateway.test.ts | 2 +- src/lib/agent/runner/sequence/README.md | 2 +- .../agent/runner/switchboard/flags/schemes.ts | 4 ++-- src/lib/agent/runner/switchboard/models.ts | 7 +------ src/lib/agent/token-pricing.ts | 20 ++++++++---------- src/lib/constants.ts | 20 ++++-------------- src/ui/tui/__tests__/store.test.ts | 2 +- src/wizard.ts | 2 +- 13 files changed, 48 insertions(+), 58 deletions(-) diff --git a/src/lib/__tests__/agent-interface.test.ts b/src/lib/__tests__/agent-interface.test.ts index f87979395..7baafa7e7 100644 --- a/src/lib/__tests__/agent-interface.test.ts +++ b/src/lib/__tests__/agent-interface.test.ts @@ -631,7 +631,7 @@ describe('subprocess gateway credentials', () => { const config = { workingDirectory: '/test/dir', mcpServers: {}, - model: 'claude-sonnet-4-6', + model: 'claude-sonnet-5', // Deliberately different from the gateway bearer below: identical values // would let either source pass. posthogApiKey: 'phx_user_oauth_token', @@ -716,7 +716,7 @@ describe('gateway re-mint on 401', () => { ) => ({ workingDirectory: '/test/dir', mcpServers: {}, - model: 'claude-sonnet-4-6', + model: 'claude-sonnet-5', posthogApiKey: 'phx_user_oauth_token', sequence: Sequence.linear, triageProvider: () => Promise.resolve('false_positive'), diff --git a/src/lib/agent/__tests__/agent-prompt-loader.test.ts b/src/lib/agent/__tests__/agent-prompt-loader.test.ts index 9ece1795c..ab23e6a49 100644 --- a/src/lib/agent/__tests__/agent-prompt-loader.test.ts +++ b/src/lib/agent/__tests__/agent-prompt-loader.test.ts @@ -36,7 +36,7 @@ describe('parseAgentPrompt', () => { type: instrument-events model_pi: openai/gpt-5.6-terra # per-profile model targets effort_pi: medium -model_sdk: claude-sonnet-4-6 +model_sdk: claude-sonnet-5 skills: [instrument-events] allowedTools: [Read, Edit, Grep, Glob, Bash] disallowedTools: [enqueue_task] @@ -52,7 +52,7 @@ Add at least one capture call. expect(p.type).toBe('instrument-events'); expect(p.modelPi).toBe('openai/gpt-5.6-terra'); expect(p.effortPi).toBe('medium'); - expect(p.modelSdk).toBe('claude-sonnet-4-6'); + expect(p.modelSdk).toBe('claude-sonnet-5'); expect(p.skills).toEqual(['instrument-events']); expect(p.allowedTools).toEqual(['Read', 'Edit', 'Grep', 'Glob', 'Bash']); expect(p.disallowedTools).toEqual(['enqueue_task']); @@ -86,7 +86,7 @@ Connect the sources. effort: 'medium', }); expect(promptModelFor(p, 'anthropic')).toEqual({ - model: 'claude-sonnet-4-6', + model: 'claude-sonnet-5', effort: undefined, }); }); @@ -281,7 +281,7 @@ describe('buildRegistry', () => { flow: 'f', modelPi: 'openai/gpt-5.6-terra', effortPi: 'medium', - modelSdk: 'claude-sonnet-4-6', + modelSdk: 'claude-sonnet-5', }), prompt({ type: 'install', flow: 'f', modelPi: 'openai/gpt-5.6-luna' }), ]; @@ -294,7 +294,7 @@ describe('buildRegistry', () => { expect(registry.get('review')).toMatchObject({ modelPi: 'openai/gpt-5.6-sol', effortPi: 'medium', - modelSdk: 'claude-sonnet-4-6', + modelSdk: 'claude-sonnet-5', }); expect(registry.seed).toMatchObject({ modelPi: 'openai/gpt-5.6-terra', @@ -324,7 +324,7 @@ describe('resolveTask', () => { runnerSeeded: false, modelPi: 'openai/gpt-5.6-luna', effortPi: 'low', - modelSdk: 'claude-haiku-4-5-20251001', + modelSdk: 'claude-haiku-4-5', skills: ['instrument-events'], allowedTools: ['Read', 'Edit'], disallowedTools: ['enqueue_task'], @@ -356,7 +356,7 @@ describe('resolveTask', () => { effort: 'low', }); expect(taskModelSpec(registry, task, 'anthropic').model).toBe( - 'claude-haiku-4-5-20251001', + 'claude-haiku-4-5', ); }); diff --git a/src/lib/agent/__tests__/token-pricing.test.ts b/src/lib/agent/__tests__/token-pricing.test.ts index 4c69fa8c2..c0c22a519 100644 --- a/src/lib/agent/__tests__/token-pricing.test.ts +++ b/src/lib/agent/__tests__/token-pricing.test.ts @@ -7,7 +7,9 @@ import { describe('pricePerMtokForModel', () => { it('defaults to Sonnet pricing (DEFAULT_AGENT_MODEL) when no model is given', () => { - expect(pricePerMtokForModel(undefined)).toEqual({ + expect( + pricePerMtokForModel(undefined, new Date('2026-09-01T00:00:00Z')), + ).toEqual({ input: 3, output: 15, cacheRead: 0.3, @@ -17,7 +19,7 @@ describe('pricePerMtokForModel', () => { }); it('strips a dated release suffix to match the undated table entry', () => { - // HAIKU_MODEL is exactly this string. + // Resumed sessions can still report a dated model id. expect(pricePerMtokForModel('claude-haiku-4-5-20251001')?.input).toBe(1); }); diff --git a/src/lib/agent/runner/__tests__/switchboard.test.ts b/src/lib/agent/runner/__tests__/switchboard.test.ts index aa925c6dc..e160674c7 100644 --- a/src/lib/agent/runner/__tests__/switchboard.test.ts +++ b/src/lib/agent/runner/__tests__/switchboard.test.ts @@ -13,6 +13,7 @@ import { describe, it, expect } from 'vitest'; import { PROGRAM_REGISTRY } from '@lib/programs/program-registry'; import { DEFAULT_AGENT_MODEL, + HAIKU_MODEL, GPT5_6_LUNA_MODEL, GPT5_6_SOL_MODEL, GPT5_6_TERRA_MODEL, @@ -268,9 +269,8 @@ describe('switchboard composed clamp', () => { describe('switchboard modelCapabilities (stage 2: effective effort)', () => { it('marks the known reasoning models as reasoning', () => { for (const m of [ - 'claude-sonnet-4-6', - 'claude-opus-4-8', - 'claude-haiku-4-5-20251001', + 'claude-sonnet-5', + 'claude-haiku-4-5', 'openai/gpt-5.6-terra', ]) { expect(modelCapabilities(m).reasoning).toBe(true); @@ -315,18 +315,25 @@ describe('switchboard modelCapabilities (stage 2: effective effort)', () => { }); describe('switchboard model allow-list', () => { - it('allow-lists the sonnets and the gpt-5.6 line, nothing older', () => { + it('allow-lists Sonnet 5, Haiku 4.5 and the gpt-5.6 line', () => { for (const m of [ - DEFAULT_AGENT_MODEL, SONNET_5_MODEL, + HAIKU_MODEL, GPT5_6_LUNA_MODEL, GPT5_6_TERRA_MODEL, GPT5_6_SOL_MODEL, ]) { expect(isValidModel(m)).toBe(true); } - // The retired openai ids are gone — no longer valid to dispatch on. - for (const m of ['openai/gpt-5', 'openai/gpt-5.4', 'openai/gpt-5.5']) { + // Retired model ids must not reach the gateway. + for (const m of [ + 'claude-sonnet-4-6', + 'claude-opus-4-8', + 'claude-haiku-4-5-20251001', + 'openai/gpt-5', + 'openai/gpt-5.4', + 'openai/gpt-5.5', + ]) { expect(isValidModel(m)).toBe(false); } expect(isValidModel('')).toBe(false); diff --git a/src/lib/agent/runner/harness/anthropic/README.md b/src/lib/agent/runner/harness/anthropic/README.md index 6d939d37e..b94ece27c 100644 --- a/src/lib/agent/runner/harness/anthropic/README.md +++ b/src/lib/agent/runner/harness/anthropic/README.md @@ -28,7 +28,7 @@ Both entry points are implemented: - **Custom headers:** wizard flags (`X-POSTHOG-FLAG-*`) and metadata (`X-POSTHOG-PROPERTY-*`) piggyback on every gateway request for tracing. - **Model routing:** `AgentConfig.modelOverride` accepts any gateway model id - (`DEFAULT_AGENT_MODEL`, `HAIKU_MODEL`, `OPUS_MODEL`, `GPT5_6_TERRA_MODEL`), so + (`DEFAULT_AGENT_MODEL`, `HAIKU_MODEL`, `GPT5_6_TERRA_MODEL`), so mechanical work (repo classification, source-map detection) can route to `HAIKU_MODEL` while integration work stays on Sonnet. diff --git a/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts b/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts index 38a78436b..c68420193 100644 --- a/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts +++ b/src/lib/agent/runner/harness/pi/__tests__/gateway.test.ts @@ -61,7 +61,7 @@ describe('buildGatewayProvider transport', () => { it('routes anthropic models over anthropic-messages without /v1', () => { const { api, baseUrl } = buildGatewayProvider({ ...base, - modelId: 'claude-sonnet-4-6', + modelId: 'claude-sonnet-5', }); expect(api).toBe('anthropic-messages'); expect(baseUrl).toBe('https://ai-gateway.us.posthog.com'); diff --git a/src/lib/agent/runner/sequence/README.md b/src/lib/agent/runner/sequence/README.md index 5f57544fd..78e8592b4 100644 --- a/src/lib/agent/runner/sequence/README.md +++ b/src/lib/agent/runner/sequence/README.md @@ -63,7 +63,7 @@ Each step is one markdown file whose frontmatter declares its shape: type: dashboard flow: integration-v2 label: Create a starter dashboard -model: claude-sonnet-4-6 +model: claude-sonnet-5 skills: [basic-integration-dashboard] allowedTools: [Read, Glob, Grep] disallowedTools: [Write, Edit, Bash, enqueue_task] diff --git a/src/lib/agent/runner/switchboard/flags/schemes.ts b/src/lib/agent/runner/switchboard/flags/schemes.ts index 9f6bf2b57..c30629699 100644 --- a/src/lib/agent/runner/switchboard/flags/schemes.ts +++ b/src/lib/agent/runner/switchboard/flags/schemes.ts @@ -6,7 +6,6 @@ */ import { z } from 'zod'; import { - DEFAULT_AGENT_MODEL, GPT5_6_LUNA_MODEL, GPT5_6_SOL_MODEL, GPT5_6_TERRA_MODEL, @@ -25,7 +24,8 @@ const MODEL_FLAG_VARIANTS: Record = { 'gpt-5-6-luna': GPT5_6_LUNA_MODEL, 'gpt-5-6-terra': GPT5_6_TERRA_MODEL, 'gpt-5-6-sol': GPT5_6_SOL_MODEL, - 'sonnet-4-6': DEFAULT_AGENT_MODEL, + // Existing flag payloads keep working while dispatching the supported Sonnet. + 'sonnet-4-6': SONNET_5_MODEL, 'sonnet-5': SONNET_5_MODEL, }; diff --git a/src/lib/agent/runner/switchboard/models.ts b/src/lib/agent/runner/switchboard/models.ts index 108cbb304..50f68dcc6 100644 --- a/src/lib/agent/runner/switchboard/models.ts +++ b/src/lib/agent/runner/switchboard/models.ts @@ -10,9 +10,7 @@ * are silent when wrong, so they live here as one configurable table. */ import { - DEFAULT_AGENT_MODEL, SONNET_5_MODEL, - OPUS_MODEL, HAIKU_MODEL, GPT5_6_LUNA_MODEL, GPT5_6_SOL_MODEL, @@ -52,9 +50,7 @@ export interface ModelCapabilities { /** Explicit per-model traits. Anything absent falls back to `defaultCaps`. */ export const MODEL_CAPABILITIES: Record = { - [DEFAULT_AGENT_MODEL]: { reasoning: true }, // claude-sonnet-4-6 [SONNET_5_MODEL]: { reasoning: true }, - [OPUS_MODEL]: { reasoning: true }, [HAIKU_MODEL]: { reasoning: true }, // The openai 5.6 line; all reasoning models, so they must opt in past the // openai-completions default (reasoning off). Luna stays low for cheap, @@ -100,8 +96,7 @@ function defaultCaps(modelId: string): ModelCapabilities { /** * Scan-triage classifier per harness: the cheapest tier of the line that harness * already speaks. Undated ids on purpose — triage is a boolean classifier, so it - * should follow the current release rather than pin one, and these are not - * dispatchable agent models (absent from MODEL_CAPABILITIES by design). + * should follow the current release rather than pin one. */ export const TRIAGE_MODELS: Record = { [Harness.anthropic]: HAIKU_TRIAGE_MODEL, diff --git a/src/lib/agent/token-pricing.ts b/src/lib/agent/token-pricing.ts index 2a489ac92..ccbfbd0c2 100644 --- a/src/lib/agent/token-pricing.ts +++ b/src/lib/agent/token-pricing.ts @@ -64,11 +64,9 @@ function pricingFromInput(input: number, output: number): PricePerMtok { /** * Exact-match price table, keyed by the *undated* model id prefix — * `pricePerMtokForModel` strips a dated suffix before falling back to this - * table, so e.g. `HAIKU_MODEL` (`claude-haiku-4-5-20251001`) resolves via - * the `'claude-haiku-4-5'` entry without needing its own duplicate key. - * `DEFAULT_AGENT_MODEL` (`claude-sonnet-4-6`) has no date suffix, so it's - * keyed directly. The rest are additional real Anthropic model ids kept - * priced correctly in case a subagent or a future default ever reports one + * table. Historical dated Haiku ids resolve via the `'claude-haiku-4-5'` + * entry without needing duplicate keys. Historical model ids stay priced + * correctly in case a resumed session or subagent reports one * (a turn on an id NOT in this table contributes $0 to the live estimate * rather than guessing — see `pricePerMtokForModel`). * @@ -77,8 +75,8 @@ function pricingFromInput(input: number, output: number): PricePerMtok { * `pricePerMtokForModel` via `SONNET_5_PRICE`. */ const PRICE_TABLE: Record = { - [DEFAULT_AGENT_MODEL]: pricingFromInput(3, 15), - 'claude-haiku-4-5': pricingFromInput(1, 5), // HAIKU_MODEL + a date suffix + 'claude-sonnet-4-6': pricingFromInput(3, 15), + 'claude-haiku-4-5': pricingFromInput(1, 5), 'claude-opus-4-5': pricingFromInput(5, 25), }; @@ -88,7 +86,7 @@ const PRICE_TABLE: Record = { * — not a coincidence, Anthropic prices it identically to 4.6 once the * promo ends. Re-verify against https://models.dev/api.json if this ever * looks off, and delete this special case once the promo window has - * passed (its `PRICE_TABLE[DEFAULT_AGENT_MODEL]` price is the same anyway). + * passed. */ const SONNET_5_PROMO_ENDS_UTC = new Date('2026-09-01T00:00:00Z'); @@ -123,10 +121,10 @@ export function pricePerMtokForModel( model?: string, now: Date = new Date(), ): PricePerMtok | undefined { - if (!model) return PRICE_TABLE[DEFAULT_AGENT_MODEL]; - const undated = stripDateSuffix(model); + const resolvedModel = model || DEFAULT_AGENT_MODEL; + const undated = stripDateSuffix(resolvedModel); if (undated === 'claude-sonnet-5') return sonnet5Price(now); - return PRICE_TABLE[model] ?? PRICE_TABLE[undated]; + return PRICE_TABLE[resolvedModel] ?? PRICE_TABLE[undated]; } /** diff --git a/src/lib/constants.ts b/src/lib/constants.ts index 13494b960..5ddfd9a1c 100644 --- a/src/lib/constants.ts +++ b/src/lib/constants.ts @@ -6,29 +6,17 @@ import { VERSION } from './version'; // ── Models ────────────────────────────────────────────────────────── -/** - * Default model for agent runs. Bare model IDs (no `anthropic/` prefix) so the - * LLM gateway's Bedrock fallback can match map_to_bedrock_model(). - */ -export const DEFAULT_AGENT_MODEL = 'claude-sonnet-4-6'; - -/** Next sonnet generation (a `MODEL_FLAG_VARIANTS` key in the switchboard). */ +/** Bare model IDs let the gateway match its provider routing aliases. */ export const SONNET_5_MODEL = 'claude-sonnet-5'; +export const DEFAULT_AGENT_MODEL = SONNET_5_MODEL; /** * Cheaper, faster model for mechanical agent work (e.g. repo classification * during source-map detection). Passed via AgentConfig.modelOverride. */ -export const HAIKU_MODEL = 'claude-haiku-4-5-20251001'; +export const HAIKU_MODEL = 'claude-haiku-4-5'; -/** Undated haiku, for scan triage — the alias tracks the current 4.5 release rather than pinning one. */ -export const HAIKU_TRIAGE_MODEL = 'claude-haiku-4-5'; - -/** - * Larger model for planning / hard work. Named the switchboard could route to - * from `PROGRAM_BINDINGS[id].model` or `contextMillOverride`. - */ -export const OPUS_MODEL = 'claude-opus-4-8'; +export const HAIKU_TRIAGE_MODEL = HAIKU_MODEL; // The only openai models the wizard runs. export const GPT5_6_LUNA_MODEL = 'openai/gpt-5.6-luna'; diff --git a/src/ui/tui/__tests__/store.test.ts b/src/ui/tui/__tests__/store.test.ts index a58923606..0201a6320 100644 --- a/src/ui/tui/__tests__/store.test.ts +++ b/src/ui/tui/__tests__/store.test.ts @@ -808,7 +808,7 @@ describe('WizardStore', () => { cacheCreationTokens: 0, cacheCreation5m: 0, cacheCreation1h: 0, - model: 'claude-haiku-4-5-20251001', + model: 'claude-haiku-4-5', }); store.addTokenUsage({ inputTokens: 1_000_000, diff --git a/src/wizard.ts b/src/wizard.ts index 0c441bd2e..e8e497320 100644 --- a/src/wizard.ts +++ b/src/wizard.ts @@ -122,7 +122,7 @@ export class Wizard { }) .option('model', { describe: - 'Override the agent model (gateway id, e.g. claude-sonnet-4-6 | openai/gpt-5). Wins over the binding default.\nenv: POSTHOG_WIZARD_MODEL', + 'Override the agent model (gateway id, e.g. claude-sonnet-5 | openai/gpt-5.6-terra). Wins over the binding default.\nenv: POSTHOG_WIZARD_MODEL', type: 'string', hidden: true, })