Skip to content

Commit 8454d5e

Browse files
TheGreatAxiosTheGreatAxios
authored
feat(subagent): refresh OAuth credential before continuation infers (#1208)
Continuations run after tool results land, possibly long after the attempt-start refresh, so an expired token 401s the next infer.
1 parent b1841c8 commit 8454d5e

6 files changed

Lines changed: 636 additions & 11 deletions

File tree

‎src/auth/refresh-source-credential.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import {
55
findSourceCredentialRecord,
66
rotateSourceCredentialMaterialIfCurrent,
77
type SourceCredentialProvenance,
8+
type SourceCredentialRecord,
89
} from "../config/source-credentials.js";
910
import { getValidCodexToken } from "./codex/session.js";
1011
import { getValidXaiToken } from "./xai/session.js";
@@ -39,7 +40,20 @@ export async function refreshSourceCredentialByProvenance(
3940
credentialId: string,
4041
): Promise<boolean> {
4142
const record = findSourceCredentialRecord(credentialId);
42-
if (record?.provenance.kind !== "oauth") return false;
43+
if (record === undefined) return false;
44+
return refreshSourceCredentialFromRecord(credentialId, record);
45+
}
46+
47+
/**
48+
* Refreshes from a caller-held record snapshot. The snapshot must be taken
49+
* before the caller's first await so a deferred refresh still compares
50+
* against the registration it started from and cannot overwrite a newer one.
51+
*/
52+
export async function refreshSourceCredentialFromRecord(
53+
credentialId: string,
54+
record: SourceCredentialRecord,
55+
): Promise<boolean> {
56+
if (record.provenance.kind !== "oauth") return false;
4357
const material = await resolveOAuthCredentialMaterial(record.provenance);
4458
return rotateSourceCredentialMaterialIfCurrent(
4559
credentialId,

‎src/subagent/refresh-inference-source.test.ts‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { afterEach, describe, expect, spyOn, test } from "bun:test";
22
import * as codexSession from "../auth/codex/session.js";
33
import * as xaiSession from "../auth/xai/session.js";
4+
import * as oauthStores from "../config/oauth-stores.js";
45

56
import type { InferenceSource } from "@intx/types/runtime";
67
import {
@@ -159,4 +160,188 @@ describe("refresh-inference-source", () => {
159160
]);
160161
expect(out).toBe(source);
161162
});
163+
164+
describe("expiry-aware ensure", () => {
165+
afterEach(() => {
166+
spyOn(oauthStores, "loadCodexProfile").mockRestore();
167+
spyOn(oauthStores, "loadXaiProfile").mockRestore();
168+
});
169+
170+
function stageCodexProfile(tokens: {
171+
access: string;
172+
refresh: string;
173+
expiresAt: number;
174+
}) {
175+
spyOn(oauthStores, "loadCodexProfile").mockResolvedValue({
176+
name: "default",
177+
createdAt: 0,
178+
tokens,
179+
});
180+
}
181+
182+
test("fresh staged Codex tokens skip the token session", async () => {
183+
const getValid = spyOn(codexSession, "getValidCodexToken");
184+
stageCodexProfile({
185+
access: "live-token",
186+
refresh: "refresh-1",
187+
expiresAt: Date.now() + 3_600_000,
188+
});
189+
const { ensureFreshInferenceSource } =
190+
await import("./refresh-inference-source.js");
191+
const source = baseSource("codex/fresh");
192+
registerSourceCredentialRecord(source.credentialId, {
193+
provenance: { kind: "oauth", provider: "codex", profile: "default" },
194+
material: { secret: "live-token" },
195+
});
196+
const out = await ensureFreshInferenceSource(source, []);
197+
expect(out).toBe(source);
198+
expect(getValid).not.toHaveBeenCalled();
199+
expect(peekSourceCredentialSecret(source.credentialId)).toBe(
200+
"live-token",
201+
);
202+
});
203+
204+
test("expiring staged Codex tokens refresh exactly once", async () => {
205+
const getValid = spyOn(
206+
codexSession,
207+
"getValidCodexToken",
208+
).mockResolvedValue({ access: "fresh-token" });
209+
stageCodexProfile({
210+
access: "stale-token",
211+
refresh: "refresh-1",
212+
expiresAt: Date.now() - 1_000,
213+
});
214+
const { ensureFreshInferenceSource } =
215+
await import("./refresh-inference-source.js");
216+
const source = baseSource("codex/expiring");
217+
registerSourceCredentialRecord(source.credentialId, {
218+
provenance: { kind: "oauth", provider: "codex", profile: "default" },
219+
material: { secret: "stale-token" },
220+
});
221+
const out = await ensureFreshInferenceSource(source, []);
222+
expect(out).toBe(source);
223+
expect(getValid).toHaveBeenCalledTimes(1);
224+
expect(peekSourceCredentialSecret(source.credentialId)).toBe(
225+
"fresh-token",
226+
);
227+
});
228+
229+
test("missing staged Codex profile falls back to refresh", async () => {
230+
const getValid = spyOn(
231+
codexSession,
232+
"getValidCodexToken",
233+
).mockResolvedValue({ access: "fresh-token" });
234+
spyOn(oauthStores, "loadCodexProfile").mockResolvedValue(undefined);
235+
const { ensureFreshInferenceSource } =
236+
await import("./refresh-inference-source.js");
237+
const source = baseSource("codex/missing");
238+
registerSourceCredentialRecord(source.credentialId, {
239+
provenance: { kind: "oauth", provider: "codex", profile: "default" },
240+
material: { secret: "stale-token" },
241+
});
242+
await ensureFreshInferenceSource(source, []);
243+
expect(getValid).toHaveBeenCalledTimes(1);
244+
expect(peekSourceCredentialSecret(source.credentialId)).toBe(
245+
"fresh-token",
246+
);
247+
});
248+
249+
test("unreadable staged Codex store falls back to refresh", async () => {
250+
const getValid = spyOn(
251+
codexSession,
252+
"getValidCodexToken",
253+
).mockResolvedValue({ access: "fresh-token" });
254+
spyOn(oauthStores, "loadCodexProfile").mockRejectedValue(
255+
new Error("disk gone"),
256+
);
257+
const { ensureFreshInferenceSource } =
258+
await import("./refresh-inference-source.js");
259+
const source = baseSource("codex/unreadable");
260+
registerSourceCredentialRecord(source.credentialId, {
261+
provenance: { kind: "oauth", provider: "codex", profile: "default" },
262+
material: { secret: "stale-token" },
263+
});
264+
await ensureFreshInferenceSource(source, []);
265+
expect(getValid).toHaveBeenCalledTimes(1);
266+
});
267+
268+
test("fresh staged tokens with a rotated cell secret refresh to heal", async () => {
269+
const getValid = spyOn(
270+
codexSession,
271+
"getValidCodexToken",
272+
).mockResolvedValue({ access: "rotated-token" });
273+
stageCodexProfile({
274+
access: "staged-token",
275+
refresh: "refresh-1",
276+
expiresAt: Date.now() + 3_600_000,
277+
});
278+
const { ensureFreshInferenceSource } =
279+
await import("./refresh-inference-source.js");
280+
const source = baseSource("codex/rotated");
281+
registerSourceCredentialRecord(source.credentialId, {
282+
provenance: { kind: "oauth", provider: "codex", profile: "default" },
283+
material: { secret: "cell-token" },
284+
});
285+
await ensureFreshInferenceSource(source, []);
286+
expect(getValid).toHaveBeenCalledTimes(1);
287+
expect(peekSourceCredentialSecret(source.credentialId)).toBe(
288+
"rotated-token",
289+
);
290+
});
291+
292+
test("concurrent ensures share one in-flight refresh", async () => {
293+
let resolveRefresh!: (value: { access: string }) => void;
294+
const getValid = spyOn(
295+
codexSession,
296+
"getValidCodexToken",
297+
).mockReturnValue(
298+
new Promise((resolve) => {
299+
resolveRefresh = resolve;
300+
}),
301+
);
302+
spyOn(oauthStores, "loadCodexProfile").mockResolvedValue(undefined);
303+
const { ensureFreshInferenceSource } =
304+
await import("./refresh-inference-source.js");
305+
const source = baseSource("codex/concurrent");
306+
registerSourceCredentialRecord(source.credentialId, {
307+
provenance: { kind: "oauth", provider: "codex", profile: "default" },
308+
material: { secret: "stale-token" },
309+
});
310+
const pending = [
311+
ensureFreshInferenceSource(source, []),
312+
ensureFreshInferenceSource(source, []),
313+
ensureFreshInferenceSource(source, []),
314+
];
315+
resolveRefresh({ access: "fresh-token" });
316+
const outs = await Promise.all(pending);
317+
expect(getValid).toHaveBeenCalledTimes(1);
318+
expect(outs.every((out) => out === source)).toBe(true);
319+
expect(peekSourceCredentialSecret(source.credentialId)).toBe(
320+
"fresh-token",
321+
);
322+
});
323+
324+
test("fresh staged xAI tokens skip the token session", async () => {
325+
const getValid = spyOn(xaiSession, "getValidXaiToken");
326+
spyOn(oauthStores, "loadXaiProfile").mockResolvedValue({
327+
name: "default",
328+
createdAt: 0,
329+
tokens: {
330+
access: "live-token",
331+
refresh: "refresh-1",
332+
expiresAt: Date.now() + 3_600_000,
333+
},
334+
});
335+
const { ensureFreshInferenceSource } =
336+
await import("./refresh-inference-source.js");
337+
const source = baseSource("xai/fresh");
338+
registerSourceCredentialRecord(source.credentialId, {
339+
provenance: { kind: "oauth", provider: "xai", profile: "default" },
340+
material: { secret: "live-token" },
341+
});
342+
const out = await ensureFreshInferenceSource(source, []);
343+
expect(out).toBe(source);
344+
expect(getValid).not.toHaveBeenCalled();
345+
});
346+
});
162347
});

‎src/subagent/refresh-inference-source.ts‎

Lines changed: 144 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,104 @@
1-
import type { InferenceSource } from "@intx/types/runtime";
1+
import type {
2+
InferenceSource,
3+
ReactorAction,
4+
ReactorDirector,
5+
} from "@intx/types/runtime";
26

3-
import { refreshSourceCredentialByProvenance } from "../auth/refresh-source-credential.js";
7+
import { isCodexTokenExpired } from "../auth/codex/session.js";
8+
import { isXaiTokenExpired } from "../auth/xai/session.js";
9+
import type { OAuthCredentialProvenance } from "../auth/refresh-source-credential.js";
10+
import { refreshSourceCredentialFromRecord } from "../auth/refresh-source-credential.js";
11+
import { loadCodexProfile, loadXaiProfile } from "../config/oauth-stores.js";
12+
import {
13+
findSourceCredentialRecord,
14+
type SourceCredentialRecord,
15+
} from "../config/source-credentials.js";
416
import type { ProviderCatalogEntry } from "../config/index.js";
517

18+
interface ContinuationRefreshAgent {
19+
setSources(sources: InferenceSource[], defaultSource: string): void;
20+
}
21+
22+
export interface ContinuationRefreshOptions {
23+
getAgent: () => ContinuationRefreshAgent | null;
24+
sources: readonly InferenceSource[];
25+
defaultSource: string;
26+
catalog: readonly ProviderCatalogEntry[] | undefined;
27+
}
28+
29+
// Concurrent continuation ensures share one in-flight refresh per credential
30+
// so a burst of post-tool infer decisions issues a single token grant.
31+
const inFlightRefreshes = new Map<string, Promise<InferenceSource>>();
32+
33+
// Directors already wrapped by withContinuationOAuthRefresh, mapped back to
34+
// their proxy. The run.ts factory applies the wrapper once, but the cache
35+
// keeps a second application a no-op instead of stacking duplicate refreshes.
36+
const wrappedDirectors = new WeakMap<object, ReactorDirector>();
37+
38+
async function stagedOAuthTokensFresh(
39+
provenance: OAuthCredentialProvenance,
40+
secret: string,
41+
): Promise<boolean> {
42+
try {
43+
if (provenance.provider === "codex") {
44+
const tokens = (await loadCodexProfile(provenance.profile))?.tokens;
45+
if (tokens === undefined || isCodexTokenExpired(tokens, Date.now()))
46+
return false;
47+
return tokens.access === secret;
48+
}
49+
const tokens = (await loadXaiProfile(provenance.profile))?.tokens;
50+
if (tokens === undefined || isXaiTokenExpired(tokens, Date.now()))
51+
return false;
52+
return tokens.access === secret;
53+
} catch {
54+
// An unreadable credential store says nothing about expiry. Fall through
55+
// to the unconditional refresh path, which surfaces the real auth error.
56+
return false;
57+
}
58+
}
59+
60+
async function ensureFreshOAuthSource(
61+
source: InferenceSource,
62+
record: SourceCredentialRecord,
63+
): Promise<InferenceSource> {
64+
const provenance = record.provenance;
65+
if (provenance.kind !== "oauth") return source;
66+
if (await stagedOAuthTokensFresh(provenance, record.material.secret))
67+
return source;
68+
await refreshSourceCredentialFromRecord(source.credentialId, record);
69+
return source;
70+
}
71+
72+
/**
73+
* Ensures the inference source's OAuth credential is valid before a
74+
* continuation infer runs. Fresh staged tokens return the source untouched —
75+
* no token-session call, no lock, no latency. Expiring, missing, or
76+
* unreadable staged tokens fall through to the unconditional refresh path.
77+
*/
678
export async function ensureFreshInferenceSource(
779
source: InferenceSource,
880
_catalog: readonly ProviderCatalogEntry[] | undefined,
981
): Promise<InferenceSource> {
10-
await refreshSourceCredentialByProvenance(source.credentialId);
11-
return source;
82+
const record = findSourceCredentialRecord(source.credentialId);
83+
const provenance = record?.provenance;
84+
if (record === undefined || provenance?.kind !== "oauth") return source;
85+
const inFlight = inFlightRefreshes.get(source.credentialId);
86+
if (inFlight !== undefined) return inFlight;
87+
const pending = ensureFreshOAuthSource(source, record);
88+
inFlightRefreshes.set(source.credentialId, pending);
89+
try {
90+
return await pending;
91+
} finally {
92+
if (inFlightRefreshes.get(source.credentialId) === pending)
93+
inFlightRefreshes.delete(source.credentialId);
94+
}
1295
}
1396

97+
/**
98+
* Refreshes every inference source in the bundle whose OAuth credential is
99+
* expiring. Sources with fresh staged tokens (or no OAuth provenance) pass
100+
* through untouched.
101+
*/
14102
export async function refreshInferenceSourceBundle(
15103
sources: readonly InferenceSource[],
16104
defaultSource: string,
@@ -21,3 +109,55 @@ export async function refreshInferenceSourceBundle(
21109
);
22110
return { sources: refreshed, defaultSource };
23111
}
112+
113+
function returnsInferAction(
114+
result: ReactorAction | readonly ReactorAction[],
115+
): boolean {
116+
const list = Array.isArray(result) ? result : [result];
117+
return list.some((action) => action?.type === "infer");
118+
}
119+
120+
/**
121+
* Host-layer wrapper: after the inner director returns an infer action, the
122+
* OAuth credential backing the continuation is ensured fresh and pushed to
123+
* the live agent before the reactor executes the infer. Non-infer decisions
124+
* pass through untouched, directors stay pure, and the inner policy stamping
125+
* runs exactly once.
126+
*
127+
* Implemented as a Proxy so the factory keeps returning the director it
128+
* built: instanceof checks, observe* hooks, and every other member forward
129+
* to the inner director; only decide() is intercepted.
130+
*/
131+
export function withContinuationOAuthRefresh<TDirector extends ReactorDirector>(
132+
inner: TDirector,
133+
options: ContinuationRefreshOptions,
134+
): TDirector {
135+
const cached = wrappedDirectors.get(inner);
136+
if (cached !== undefined) return cached as TDirector;
137+
const decide = async (
138+
event: Parameters<TDirector["decide"]>[0],
139+
state: Parameters<TDirector["decide"]>[1],
140+
capabilities: Parameters<TDirector["decide"]>[2],
141+
) => {
142+
const result = await inner.decide(event, state, capabilities);
143+
if (!returnsInferAction(result)) return result;
144+
const agent = options.getAgent();
145+
if (agent === null) return result;
146+
const fresh = await refreshInferenceSourceBundle(
147+
options.sources,
148+
options.defaultSource,
149+
options.catalog,
150+
);
151+
agent.setSources(fresh.sources, fresh.defaultSource);
152+
return result;
153+
};
154+
const proxy = new Proxy(inner, {
155+
get(target, property, _receiver) {
156+
if (property === "decide") return decide;
157+
return Reflect.get(target, property);
158+
},
159+
});
160+
wrappedDirectors.set(inner, proxy);
161+
wrappedDirectors.set(proxy, proxy);
162+
return proxy;
163+
}

0 commit comments

Comments
 (0)