diff --git a/src/subagent/refresh-inference-source.cache.test.ts b/src/subagent/refresh-inference-source.cache.test.ts new file mode 100644 index 000000000..2a6b38f9b --- /dev/null +++ b/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 { + 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); + }); +}); diff --git a/src/subagent/refresh-inference-source.ts b/src/subagent/refresh-inference-source.ts index eb0c8652f..6e5a6fd45 100644 --- a/src/subagent/refresh-inference-source.ts +++ b/src/subagent/refresh-inference-source.ts @@ -31,8 +31,15 @@ export interface ContinuationRefreshOptions { const inFlightRefreshes = new Map>(); // Wrapped directors, mapped back to their proxy, so a second application is a -// no-op instead of stacking duplicate refreshes. -const wrappedDirectors = new WeakMap(); +// 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(); async function stagedOAuthTokensFresh( provenance: OAuthCredentialProvenance, @@ -125,7 +132,15 @@ export function withContinuationOAuthRefresh( 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[0], state: Parameters[1], @@ -133,12 +148,12 @@ export function withContinuationOAuthRefresh( ) => { 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( 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; }