Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 59 additions & 0 deletions src/subagent/refresh-inference-source.cache.test.ts
Original file line number Diff line number Diff line change
@@ -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);
});
});
34 changes: 25 additions & 9 deletions src/subagent/refresh-inference-source.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
Expand All @@ -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;
}
Loading