From acf0fffc40a6b04cf23e65d7aaeacb90f6870717 Mon Sep 17 00:00:00 2001 From: Sawyer Cutler Date: Mon, 28 Sep 2026 08:02:22 -0700 Subject: [PATCH] feat(subagent): refresh OAuth credential before continuation infers Continuations run after tool results land, possibly long after the attempt-start refresh, so an expired token 401s the next infer. --- src/auth/refresh-source-credential.ts | 16 +- src/subagent/refresh-inference-source.test.ts | 185 ++++++++++++ src/subagent/refresh-inference-source.ts | 148 +++++++++- src/subagent/run-continuation-refresh.test.ts | 268 ++++++++++++++++++ src/subagent/run.ts | 16 +- src/tui/runner/exit.test.ts | 14 +- 6 files changed, 636 insertions(+), 11 deletions(-) create mode 100644 src/subagent/run-continuation-refresh.test.ts diff --git a/src/auth/refresh-source-credential.ts b/src/auth/refresh-source-credential.ts index 0238e1974..5a0a0eb76 100644 --- a/src/auth/refresh-source-credential.ts +++ b/src/auth/refresh-source-credential.ts @@ -5,6 +5,7 @@ import { findSourceCredentialRecord, rotateSourceCredentialMaterialIfCurrent, type SourceCredentialProvenance, + type SourceCredentialRecord, } from "../config/source-credentials.js"; import { getValidCodexToken } from "./codex/session.js"; import { getValidXaiToken } from "./xai/session.js"; @@ -39,7 +40,20 @@ export async function refreshSourceCredentialByProvenance( credentialId: string, ): Promise { const record = findSourceCredentialRecord(credentialId); - if (record?.provenance.kind !== "oauth") return false; + if (record === undefined) return false; + return refreshSourceCredentialFromRecord(credentialId, record); +} + +/** + * Refreshes from a caller-held record snapshot. The snapshot must be taken + * before the caller's first await so a deferred refresh still compares + * against the registration it started from and cannot overwrite a newer one. + */ +export async function refreshSourceCredentialFromRecord( + credentialId: string, + record: SourceCredentialRecord, +): Promise { + if (record.provenance.kind !== "oauth") return false; const material = await resolveOAuthCredentialMaterial(record.provenance); return rotateSourceCredentialMaterialIfCurrent( credentialId, diff --git a/src/subagent/refresh-inference-source.test.ts b/src/subagent/refresh-inference-source.test.ts index f0e3b5a93..cb948b6fb 100644 --- a/src/subagent/refresh-inference-source.test.ts +++ b/src/subagent/refresh-inference-source.test.ts @@ -1,6 +1,7 @@ import { afterEach, describe, expect, spyOn, test } from "bun:test"; import * as codexSession from "../auth/codex/session.js"; import * as xaiSession from "../auth/xai/session.js"; +import * as oauthStores from "../config/oauth-stores.js"; import type { InferenceSource } from "@intx/types/runtime"; import { @@ -159,4 +160,188 @@ describe("refresh-inference-source", () => { ]); expect(out).toBe(source); }); + + describe("expiry-aware ensure", () => { + afterEach(() => { + spyOn(oauthStores, "loadCodexProfile").mockRestore(); + spyOn(oauthStores, "loadXaiProfile").mockRestore(); + }); + + function stageCodexProfile(tokens: { + access: string; + refresh: string; + expiresAt: number; + }) { + spyOn(oauthStores, "loadCodexProfile").mockResolvedValue({ + name: "default", + createdAt: 0, + tokens, + }); + } + + test("fresh staged Codex tokens skip the token session", async () => { + const getValid = spyOn(codexSession, "getValidCodexToken"); + stageCodexProfile({ + access: "live-token", + refresh: "refresh-1", + expiresAt: Date.now() + 3_600_000, + }); + const { ensureFreshInferenceSource } = + await import("./refresh-inference-source.js"); + const source = baseSource("codex/fresh"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: "live-token" }, + }); + const out = await ensureFreshInferenceSource(source, []); + expect(out).toBe(source); + expect(getValid).not.toHaveBeenCalled(); + expect(peekSourceCredentialSecret(source.credentialId)).toBe( + "live-token", + ); + }); + + test("expiring staged Codex tokens refresh exactly once", async () => { + const getValid = spyOn( + codexSession, + "getValidCodexToken", + ).mockResolvedValue({ access: "fresh-token" }); + stageCodexProfile({ + access: "stale-token", + refresh: "refresh-1", + expiresAt: Date.now() - 1_000, + }); + const { ensureFreshInferenceSource } = + await import("./refresh-inference-source.js"); + const source = baseSource("codex/expiring"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: "stale-token" }, + }); + const out = await ensureFreshInferenceSource(source, []); + expect(out).toBe(source); + expect(getValid).toHaveBeenCalledTimes(1); + expect(peekSourceCredentialSecret(source.credentialId)).toBe( + "fresh-token", + ); + }); + + test("missing staged Codex profile falls back to refresh", async () => { + const getValid = spyOn( + codexSession, + "getValidCodexToken", + ).mockResolvedValue({ access: "fresh-token" }); + spyOn(oauthStores, "loadCodexProfile").mockResolvedValue(undefined); + const { ensureFreshInferenceSource } = + await import("./refresh-inference-source.js"); + const source = baseSource("codex/missing"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: "stale-token" }, + }); + await ensureFreshInferenceSource(source, []); + expect(getValid).toHaveBeenCalledTimes(1); + expect(peekSourceCredentialSecret(source.credentialId)).toBe( + "fresh-token", + ); + }); + + test("unreadable staged Codex store falls back to refresh", async () => { + const getValid = spyOn( + codexSession, + "getValidCodexToken", + ).mockResolvedValue({ access: "fresh-token" }); + spyOn(oauthStores, "loadCodexProfile").mockRejectedValue( + new Error("disk gone"), + ); + const { ensureFreshInferenceSource } = + await import("./refresh-inference-source.js"); + const source = baseSource("codex/unreadable"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: "stale-token" }, + }); + await ensureFreshInferenceSource(source, []); + expect(getValid).toHaveBeenCalledTimes(1); + }); + + test("fresh staged tokens with a rotated cell secret refresh to heal", async () => { + const getValid = spyOn( + codexSession, + "getValidCodexToken", + ).mockResolvedValue({ access: "rotated-token" }); + stageCodexProfile({ + access: "staged-token", + refresh: "refresh-1", + expiresAt: Date.now() + 3_600_000, + }); + const { ensureFreshInferenceSource } = + await import("./refresh-inference-source.js"); + const source = baseSource("codex/rotated"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: "cell-token" }, + }); + await ensureFreshInferenceSource(source, []); + expect(getValid).toHaveBeenCalledTimes(1); + expect(peekSourceCredentialSecret(source.credentialId)).toBe( + "rotated-token", + ); + }); + + test("concurrent ensures share one in-flight refresh", async () => { + let resolveRefresh!: (value: { access: string }) => void; + const getValid = spyOn( + codexSession, + "getValidCodexToken", + ).mockReturnValue( + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + spyOn(oauthStores, "loadCodexProfile").mockResolvedValue(undefined); + const { ensureFreshInferenceSource } = + await import("./refresh-inference-source.js"); + const source = baseSource("codex/concurrent"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: "stale-token" }, + }); + const pending = [ + ensureFreshInferenceSource(source, []), + ensureFreshInferenceSource(source, []), + ensureFreshInferenceSource(source, []), + ]; + resolveRefresh({ access: "fresh-token" }); + const outs = await Promise.all(pending); + expect(getValid).toHaveBeenCalledTimes(1); + expect(outs.every((out) => out === source)).toBe(true); + expect(peekSourceCredentialSecret(source.credentialId)).toBe( + "fresh-token", + ); + }); + + test("fresh staged xAI tokens skip the token session", async () => { + const getValid = spyOn(xaiSession, "getValidXaiToken"); + spyOn(oauthStores, "loadXaiProfile").mockResolvedValue({ + name: "default", + createdAt: 0, + tokens: { + access: "live-token", + refresh: "refresh-1", + expiresAt: Date.now() + 3_600_000, + }, + }); + const { ensureFreshInferenceSource } = + await import("./refresh-inference-source.js"); + const source = baseSource("xai/fresh"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "xai", profile: "default" }, + material: { secret: "live-token" }, + }); + const out = await ensureFreshInferenceSource(source, []); + expect(out).toBe(source); + expect(getValid).not.toHaveBeenCalled(); + }); + }); }); diff --git a/src/subagent/refresh-inference-source.ts b/src/subagent/refresh-inference-source.ts index c0ed60043..3d7021f2b 100644 --- a/src/subagent/refresh-inference-source.ts +++ b/src/subagent/refresh-inference-source.ts @@ -1,16 +1,104 @@ -import type { InferenceSource } from "@intx/types/runtime"; +import type { + InferenceSource, + ReactorAction, + ReactorDirector, +} from "@intx/types/runtime"; -import { refreshSourceCredentialByProvenance } from "../auth/refresh-source-credential.js"; +import { isCodexTokenExpired } from "../auth/codex/session.js"; +import { isXaiTokenExpired } from "../auth/xai/session.js"; +import type { OAuthCredentialProvenance } from "../auth/refresh-source-credential.js"; +import { refreshSourceCredentialFromRecord } from "../auth/refresh-source-credential.js"; +import { loadCodexProfile, loadXaiProfile } from "../config/oauth-stores.js"; +import { + findSourceCredentialRecord, + type SourceCredentialRecord, +} from "../config/source-credentials.js"; import type { ProviderCatalogEntry } from "../config/index.js"; +interface ContinuationRefreshAgent { + setSources(sources: InferenceSource[], defaultSource: string): void; +} + +export interface ContinuationRefreshOptions { + getAgent: () => ContinuationRefreshAgent | null; + sources: readonly InferenceSource[]; + defaultSource: string; + catalog: readonly ProviderCatalogEntry[] | undefined; +} + +// Concurrent continuation ensures share one in-flight refresh per credential +// so a burst of post-tool infer decisions issues a single token grant. +const inFlightRefreshes = new Map>(); + +// Directors already wrapped by withContinuationOAuthRefresh, mapped back to +// their proxy. The run.ts factory applies the wrapper once, but the cache +// keeps a second application a no-op instead of stacking duplicate refreshes. +const wrappedDirectors = new WeakMap(); + +async function stagedOAuthTokensFresh( + provenance: OAuthCredentialProvenance, + secret: string, +): Promise { + try { + if (provenance.provider === "codex") { + const tokens = (await loadCodexProfile(provenance.profile))?.tokens; + if (tokens === undefined || isCodexTokenExpired(tokens, Date.now())) + return false; + return tokens.access === secret; + } + const tokens = (await loadXaiProfile(provenance.profile))?.tokens; + if (tokens === undefined || isXaiTokenExpired(tokens, Date.now())) + return false; + return tokens.access === secret; + } catch { + // An unreadable credential store says nothing about expiry. Fall through + // to the unconditional refresh path, which surfaces the real auth error. + return false; + } +} + +async function ensureFreshOAuthSource( + source: InferenceSource, + record: SourceCredentialRecord, +): Promise { + const provenance = record.provenance; + if (provenance.kind !== "oauth") return source; + if (await stagedOAuthTokensFresh(provenance, record.material.secret)) + return source; + await refreshSourceCredentialFromRecord(source.credentialId, record); + return source; +} + +/** + * Ensures the inference source's OAuth credential is valid before a + * continuation infer runs. Fresh staged tokens return the source untouched — + * no token-session call, no lock, no latency. Expiring, missing, or + * unreadable staged tokens fall through to the unconditional refresh path. + */ export async function ensureFreshInferenceSource( source: InferenceSource, _catalog: readonly ProviderCatalogEntry[] | undefined, ): Promise { - await refreshSourceCredentialByProvenance(source.credentialId); - return source; + const record = findSourceCredentialRecord(source.credentialId); + const provenance = record?.provenance; + if (record === undefined || provenance?.kind !== "oauth") return source; + const inFlight = inFlightRefreshes.get(source.credentialId); + if (inFlight !== undefined) return inFlight; + const pending = ensureFreshOAuthSource(source, record); + inFlightRefreshes.set(source.credentialId, pending); + try { + return await pending; + } finally { + if (inFlightRefreshes.get(source.credentialId) === pending) + inFlightRefreshes.delete(source.credentialId); + } } +/** + * Refreshes every inference source in the bundle whose OAuth credential is + * expiring. Sources with fresh staged tokens (or no OAuth provenance) pass + * through untouched. + */ export async function refreshInferenceSourceBundle( sources: readonly InferenceSource[], defaultSource: string, @@ -21,3 +109,55 @@ export async function refreshInferenceSourceBundle( ); return { sources: refreshed, defaultSource }; } + +function returnsInferAction( + result: ReactorAction | readonly ReactorAction[], +): boolean { + const list = Array.isArray(result) ? result : [result]; + return list.some((action) => action?.type === "infer"); +} + +/** + * Host-layer wrapper: after the inner director returns an infer action, the + * OAuth credential backing the continuation is ensured fresh and pushed to + * the live agent before the reactor executes the infer. Non-infer decisions + * pass through untouched, directors stay pure, and the inner policy stamping + * runs exactly once. + * + * Implemented as a Proxy so the factory keeps returning the director it + * built: instanceof checks, observe* hooks, and every other member forward + * to the inner director; only decide() is intercepted. + */ +export function withContinuationOAuthRefresh( + inner: TDirector, + options: ContinuationRefreshOptions, +): TDirector { + const cached = wrappedDirectors.get(inner); + if (cached !== undefined) return cached as TDirector; + const decide = async ( + event: Parameters[0], + state: Parameters[1], + capabilities: Parameters[2], + ) => { + const result = await inner.decide(event, state, capabilities); + if (!returnsInferAction(result)) return result; + const agent = options.getAgent(); + if (agent === null) return result; + const fresh = await refreshInferenceSourceBundle( + options.sources, + options.defaultSource, + options.catalog, + ); + agent.setSources(fresh.sources, fresh.defaultSource); + return result; + }; + const proxy = new Proxy(inner, { + get(target, property, _receiver) { + if (property === "decide") return decide; + return Reflect.get(target, property); + }, + }); + wrappedDirectors.set(inner, proxy); + wrappedDirectors.set(proxy, proxy); + return proxy; +} diff --git a/src/subagent/run-continuation-refresh.test.ts b/src/subagent/run-continuation-refresh.test.ts new file mode 100644 index 000000000..cffec981e --- /dev/null +++ b/src/subagent/run-continuation-refresh.test.ts @@ -0,0 +1,268 @@ +import { afterEach, describe, expect, spyOn, test } from "bun:test"; +import * as codexSession from "../auth/codex/session.js"; +import * as oauthStores from "../config/oauth-stores.js"; + +import type { + InferenceSource, + ReactorAction, + ReactorCapabilities, + ReactorDirector, + ReactorInboundEvent, + ReactorState, +} from "@intx/types/runtime"; +import { + clearSourceCredentials, + peekSourceCredentialSecret, + registerSourceCredentialRecord, +} from "../config/source-credentials.js"; +import { createTestCapabilities } from "./director-test-harness.js"; +import { SubAgentDirector } from "./nudge-director.js"; +import type { ContinuationRefreshOptions } from "./refresh-inference-source.js"; +import { withContinuationOAuthRefresh } from "./refresh-inference-source.js"; + +const state = { turns: [] } as unknown as ReactorState; +const retryPolicyStub = { id: "test-retry-policy" } as never; + +function inferenceDone(callIds: string[]): ReactorInboundEvent { + return { + type: "inference.done", + turn: { + role: "assistant", + model: "test", + timestamp: 0, + content: callIds.map((id) => ({ + type: "tool_call", + id, + name: "read_file", + arguments: { path: `${id}.ts` }, + })), + }, + usage: { input: 0, output: 1, cacheRead: 0, cacheWrite: 0, thinking: 0 }, + source: { model: "test-model" }, + } as unknown as ReactorInboundEvent; +} + +function toolDone(callId: string): ReactorInboundEvent { + return { + type: "tool.done", + result: { callId, content: "ok", isError: false }, + } as unknown as ReactorInboundEvent; +} + +function actions(result: ReactorAction | ReactorAction[]): ReactorAction[] { + return Array.isArray(result) ? result : [result]; +} + +function inferOptions(result: ReactorAction | ReactorAction[]) { + const infer = actions(result).find((action) => action.type === "infer"); + if (infer?.type !== "infer") throw new Error("expected infer action"); + return infer.options as Record | undefined; +} + +const bundleSource = (id: string): InferenceSource => ({ + id, + provider: "openai", + baseURL: "https://api.openai.com/v1", + credentialId: id, + model: "test", +}); + +describe("run continuation OAuth refresh", () => { + afterEach(() => { + clearSourceCredentials(); + spyOn(codexSession, "getValidCodexToken").mockRestore(); + spyOn(oauthStores, "loadCodexProfile").mockRestore(); + }); + + async function setup(expiring: boolean) { + const now = Date.now(); + spyOn(oauthStores, "loadCodexProfile").mockResolvedValue({ + name: "default", + createdAt: 0, + tokens: { + access: "staged-token", + refresh: "refresh-1", + expiresAt: expiring ? now - 1_000 : now + 3_600_000, + }, + }); + const getValid = spyOn( + codexSession, + "getValidCodexToken", + ).mockResolvedValue({ access: "fresh-token" }); + const source = bundleSource("codex/default"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: expiring ? "stale-token" : "staged-token" }, + }); + const setSourcesCalls: { + sources: InferenceSource[]; + defaultSource: string; + }[] = []; + const agent = { + setSources(sources: InferenceSource[], defaultSource: string) { + setSourcesCalls.push({ sources, defaultSource }); + }, + }; + const options: ContinuationRefreshOptions = { + getAgent: () => agent, + sources: [source], + defaultSource: source.id, + catalog: [], + }; + const inner = new SubAgentDirector( + "system", + [], + undefined, + 30, + Date.now, + false, + false, + retryPolicyStub, + ); + const wrapped = withContinuationOAuthRefresh(inner, options); + const caps = createTestCapabilities(); + return { getValid, setSourcesCalls, wrapped, caps, source }; + } + + test("post-tool continuation refreshes when expiring", async () => { + const { getValid, setSourcesCalls, wrapped, caps, source } = + await setup(true); + const first = await wrapped.decide(inferenceDone(["call-1"]), state, caps); + expect(actions(first).some((action) => action.type === "infer")).toBe( + false, + ); + expect(getValid).not.toHaveBeenCalled(); + expect(setSourcesCalls).toHaveLength(0); + + const result = await wrapped.decide(toolDone("call-1"), state, caps); + const options = inferOptions(result); + expect(getValid).toHaveBeenCalledTimes(1); + expect(peekSourceCredentialSecret(source.credentialId)).toBe("fresh-token"); + expect(setSourcesCalls).toHaveLength(1); + expect(setSourcesCalls[0]?.defaultSource).toBe(source.id); + expect(setSourcesCalls[0]?.sources.map((entry) => entry.id)).toEqual([ + source.id, + ]); + // The inner policy stamping composes untouched through the wrapper. + expect(options?.retryPolicy).toBe(retryPolicyStub); + expect(options?.systemPrompt).toBeDefined(); + }); + + test("post-tool continuation skips the session when fresh", async () => { + const { getValid, setSourcesCalls, wrapped, caps } = await setup(false); + await wrapped.decide(inferenceDone(["call-1"]), state, caps); + const result = await wrapped.decide(toolDone("call-1"), state, caps); + expect(inferOptions(result)?.retryPolicy).toBe(retryPolicyStub); + expect(getValid).not.toHaveBeenCalled(); + expect(setSourcesCalls).toHaveLength(1); + }); + + test("infer passes through untouched when the agent is not running yet", async () => { + const getValid = spyOn(codexSession, "getValidCodexToken"); + const inner = new SubAgentDirector( + "system", + [], + undefined, + 30, + Date.now, + false, + false, + retryPolicyStub, + ); + const wrapped = withContinuationOAuthRefresh(inner, { + getAgent: () => null, + sources: [bundleSource("codex/default")], + defaultSource: "codex/default", + catalog: [], + }); + const caps = createTestCapabilities(); + await wrapped.decide(inferenceDone(["call-1"]), state, caps); + const result = await wrapped.decide(toolDone("call-1"), state, caps); + expect(actions(result).some((action) => action.type === "infer")).toBe( + true, + ); + expect(getValid).not.toHaveBeenCalled(); + }); + + test("wrapping an already-wrapped director is a no-op", async () => { + const getValid = spyOn( + codexSession, + "getValidCodexToken", + ).mockResolvedValue({ access: "fresh-token" }); + spyOn(oauthStores, "loadCodexProfile").mockResolvedValue(undefined); + const source = bundleSource("codex/double"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: "stale-token" }, + }); + let setSources = 0; + const options: ContinuationRefreshOptions = { + getAgent: () => ({ + setSources: () => { + setSources += 1; + }, + }), + sources: [source], + defaultSource: source.id, + catalog: [], + }; + const fakeInner: ReactorDirector = { + decide: async ( + _event: ReactorInboundEvent, + _state: ReactorState, + capabilities: ReactorCapabilities, + ) => capabilities.infer(), + }; + const wrapped = withContinuationOAuthRefresh(fakeInner, options); + expect(withContinuationOAuthRefresh(wrapped, options)).toBe(wrapped); + await wrapped.decide(toolDone("call-1"), state, createTestCapabilities()); + expect(getValid).toHaveBeenCalledTimes(1); + expect(setSources).toBe(1); + }); + + test("concurrent continuation decides share one refresh", async () => { + let resolveRefresh!: (value: { access: string }) => void; + const getValid = spyOn(codexSession, "getValidCodexToken").mockReturnValue( + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + spyOn(oauthStores, "loadCodexProfile").mockResolvedValue(undefined); + const source = bundleSource("codex/race"); + registerSourceCredentialRecord(source.credentialId, { + provenance: { kind: "oauth", provider: "codex", profile: "default" }, + material: { secret: "stale-token" }, + }); + const setSourcesCalls: { sources: unknown; defaultSource: string }[] = []; + const options: ContinuationRefreshOptions = { + getAgent: () => ({ + setSources: (sources: InferenceSource[], defaultSource: string) => { + setSourcesCalls.push({ sources, defaultSource }); + }, + }), + sources: [source], + defaultSource: source.id, + catalog: [], + }; + const fakeInner: ReactorDirector = { + decide: async ( + _event: ReactorInboundEvent, + _state: ReactorState, + capabilities: ReactorCapabilities, + ) => capabilities.infer(), + }; + const caps = createTestCapabilities(); + const first = withContinuationOAuthRefresh(fakeInner, options); + const second = withContinuationOAuthRefresh(fakeInner, options); + const pending = [ + first.decide(toolDone("call-1"), state, caps), + second.decide(toolDone("call-1"), state, caps), + ]; + resolveRefresh({ access: "fresh-token" }); + const results = await Promise.all(pending); + expect(results).toHaveLength(2); + expect(getValid).toHaveBeenCalledTimes(1); + expect(setSourcesCalls).toHaveLength(2); + expect(peekSourceCredentialSecret(source.credentialId)).toBe("fresh-token"); + }); +}); diff --git a/src/subagent/run.ts b/src/subagent/run.ts index e35b395b5..2f7bbd83d 100644 --- a/src/subagent/run.ts +++ b/src/subagent/run.ts @@ -89,7 +89,10 @@ import { generateSessionId } from "../session/index.js"; import { consumeStream } from "../session/stream-consumer.js"; import { createCycleTextRecorder } from "../session/stream-journal.js"; import { onTurnBoundary } from "../agent/reactor-events.js"; -import { refreshInferenceSourceBundle } from "./refresh-inference-source.js"; +import { + refreshInferenceSourceBundle, + withContinuationOAuthRefresh, +} from "./refresh-inference-source.js"; import { createResolvedProviderFailureError, isResolvedProviderFailureError, @@ -1091,7 +1094,16 @@ async function runSubAgentInner( director.observeInterventions((event) => { interventions(event); }); - return director; + // Continuation infers run after tool results land, possibly long after + // the attempt-start refresh. The wrapper ensures the OAuth credential + // is fresh and pushes it to the live agent before the reactor executes + // the infer; decide() itself stays pure. + return withContinuationOAuthRefresh(director, { + getAgent: () => agent, + sources: bundle.sources, + defaultSource: bundle.defaultSource, + catalog: params.catalog, + }); }, }); diff --git a/src/tui/runner/exit.test.ts b/src/tui/runner/exit.test.ts index 171829c02..1f6f1ae49 100644 --- a/src/tui/runner/exit.test.ts +++ b/src/tui/runner/exit.test.ts @@ -383,6 +383,7 @@ function stubSendLifecycle(agent: Agent): { function hangCodexRefresh(): { settle: (value: { access: string }) => void; spy: ReturnType; + waitForHang: () => Promise; } { let settle: ((value: { access: string }) => void) | undefined; const spy = spyOn(codexSession, "getValidCodexToken").mockImplementation( @@ -394,6 +395,13 @@ function hangCodexRefresh(): { return { spy, settle: (value) => defined(settle, "settleRefresh")(value), + // The refresh reaches the token session after an async staged-expiry + // probe, so wait for the hang instead of counting microtask ticks. + waitForHang: async () => { + for (let i = 0; i < 1000 && settle === undefined; i++) { + await new Promise((resolve) => setTimeout(resolve, 0)); + } + }, }; } @@ -408,8 +416,7 @@ describe("agentProxy.send vs /clear", () => { try { const { agentProxy } = await createRunLifecycle(state, services); const pending = agentProxy.send("keep me out of the new session"); - await Promise.resolve(); - await Promise.resolve(); + await hung.waitForHang(); resetSessionForRotation(state, services); state.currentAgent = newAgent; @@ -430,8 +437,7 @@ describe("agentProxy.send vs /clear", () => { try { const { agentProxy } = await createRunLifecycle(state, services); const pending = agentProxy.send("deliver this"); - await Promise.resolve(); - await Promise.resolve(); + await hung.waitForHang(); hung.settle({ access: "fresh-token" }); await pending; expect(sends).toEqual(["deliver this"]);