CL-8988: refresh OAuth credential before continuation infers

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<boolean> {
   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<boolean> {
+  if (record.provenance.kind !== "oauth") return false;
   const material = await resolveOAuthCredentialMaterial(record.provenance);
   return rotateSourceCredentialMaterialIfCurrent(
     credentialId,

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();
+    });
+  });
 });

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<string, Promise<InferenceSource>>();
+
+// 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<object, ReactorDirector>();
+
+async function stagedOAuthTokensFresh(
+  provenance: OAuthCredentialProvenance,
+  secret: string,
+): Promise<boolean> {
+  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<InferenceSource> {
+  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<InferenceSource> {
-  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<TDirector extends ReactorDirector>(
+  inner: TDirector,
+  options: ContinuationRefreshOptions,
+): TDirector {
+  const cached = wrappedDirectors.get(inner);
+  if (cached !== undefined) return cached as TDirector;
+  const decide = async (
+    event: Parameters<TDirector["decide"]>[0],
+    state: Parameters<TDirector["decide"]>[1],
+    capabilities: Parameters<TDirector["decide"]>[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;
+}