fix(subagent): apply the latest wrap's sources on continuation OAuth refresh

Author: TheGreatAxios

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

@@ -0,0 +1,59 @@
+import { describe, expect, test } from "bun:test";
+import type {
+  InferenceSource,
+  ReactorAction,
+  ReactorDirector,
+} from "@intx/types/runtime";
+
+import { withContinuationOAuthRefresh } from "./refresh-inference-source.js";
+
+const baseSource = (id: string): InferenceSource => ({
+  id,
+  provider: "openai",
+  baseURL: "https://api.openai.com/v1",
+  credentialId: id,
+  model: "gpt-4o",
+});
+
+describe("withContinuationOAuthRefresh cache", () => {
+  test("a second wrap of the same inner director applies the second wrap's sources", async () => {
+    const seen: { sources: InferenceSource[]; defaultSource: string }[] = [];
+    const agent = {
+      setSources(sources: InferenceSource[], defaultSource: string) {
+        seen.push({ sources, defaultSource });
+      },
+    };
+
+    const inner: ReactorDirector = {
+      async decide(): Promise<ReactorAction[]> {
+        return [{ type: "infer" }];
+      },
+    };
+
+    const sourcesA = [baseSource("openai")];
+    const sourcesB = [baseSource("other")];
+
+    // First wrap with option set A.
+    withContinuationOAuthRefresh(inner, {
+      getAgent: () => agent,
+      sources: sourcesA,
+      defaultSource: "openai",
+      catalog: [],
+    });
+
+    // Second wrap of the SAME inner director with a different option set B.
+    const wrapped = withContinuationOAuthRefresh(inner, {
+      getAgent: () => agent,
+      sources: sourcesB,
+      defaultSource: "other",
+      catalog: [],
+    });
+
+    await wrapped.decide({} as never, {} as never, {} as never);
+
+    expect(seen).toHaveLength(1);
+    const call = seen[0];
+    expect(call?.defaultSource).toBe("other");
+    expect(call?.sources).toEqual(sourcesB);
+  });
+});

src/subagent/refresh-inference-source.ts

@@ -31,8 +31,15 @@ export interface ContinuationRefreshOptions {
 const inFlightRefreshes = new Map<string, Promise<InferenceSource>>();
 
 // Wrapped directors, mapped back to their proxy, so a second application is a
-// no-op instead of stacking duplicate refreshes.
-const wrappedDirectors = new WeakMap<object, ReactorDirector>();
+// no-op instead of stacking duplicate refreshes. The entry also holds the
+// latest options: a re-wrap with a new source set must update what decide()
+// reads, not be silently dropped by the cache.
+interface WrappedDirectorEntry {
+  proxy: ReactorDirector;
+  options: ContinuationRefreshOptions;
+}
+
+const wrappedDirectors = new WeakMap<object, WrappedDirectorEntry>();
 
 async function stagedOAuthTokensFresh(
   provenance: OAuthCredentialProvenance,
@@ -125,20 +132,28 @@ export function withContinuationOAuthRefresh<TDirector extends ReactorDirector>(
   options: ContinuationRefreshOptions,
 ): TDirector {
   const cached = wrappedDirectors.get(inner);
-  if (cached !== undefined) return cached as TDirector;
+  if (cached !== undefined) {
+    // Second application is a no-op for the proxy, but not for the options:
+    // the caller may re-wrap the same director with a different source set,
+    // and decide() below reads entry.options at call time, so the latest
+    // wrap's sources win instead of the first wrap's.
+    cached.options = options;
+    return cached.proxy as TDirector;
+  }
+  const entry: WrappedDirectorEntry = { proxy: inner, options };
   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();
+    const agent = entry.options.getAgent();
     if (agent === null) return result;
     const fresh = await refreshInferenceSourceBundle(
-      options.sources,
-      options.defaultSource,
-      options.catalog,
+      entry.options.sources,
+      entry.options.defaultSource,
+      entry.options.catalog,
     );
     agent.setSources(fresh.sources, fresh.defaultSource);
     return result;
@@ -149,7 +164,8 @@ export function withContinuationOAuthRefresh<TDirector extends ReactorDirector>(
       return Reflect.get(target, property);
     },
   });
-  wrappedDirectors.set(inner, proxy);
-  wrappedDirectors.set(proxy, proxy);
+  entry.proxy = proxy;
+  wrappedDirectors.set(inner, entry);
+  wrappedDirectors.set(proxy, entry);
   return proxy;
 }