diff --git a/src/gateway/server-methods/models-list-result.ts b/src/gateway/server-methods/models-list-result.ts index 959d2eb6cd86..27b1d8322512 100644 --- a/src/gateway/server-methods/models-list-result.ts +++ b/src/gateway/server-methods/models-list-result.ts @@ -56,6 +56,7 @@ import { resolveManifestProviderAuthChoices } from "../../plugins/provider-auth- import type { ProviderCatalogOutcome } from "../../plugins/provider-catalog.types.js"; import { normalizeAgentId } from "../../routing/session-key.js"; import type { GatewayAgentRuntime } from "../../shared/session-types.js"; +import { loadDeferredCatalog, resolveDeferredAuthStore } from "../server-model-catalog-auth.js"; import { resolveGatewayModelThinkingProfile } from "../session-utils-model.js"; import { createModelsListAuthResolver } from "./models-list-auth-resolver.js"; import type { GatewayRequestContext } from "./types.js"; @@ -547,10 +548,7 @@ export async function buildModelsListResult( if (params.preloadedOnly) { return { entries: [], routeVariants: [] }; } - loadedSnapshot = await params.context.loadGatewayModelCatalogSnapshot({ - agentId: initialAgentId, - readOnly: loadedReadOnly, - }); + loadedSnapshot = await loadDeferredCatalog(params.context, initialAgentId, loadedReadOnly); return loadedSnapshot; }, onTimeout: handleCatalogTimeout, @@ -572,10 +570,7 @@ export async function buildModelsListResult( agentId: escalationAgentId, view, loadCatalog: async ({ readOnly }) => { - fullSnapshot = await params.context.loadGatewayModelCatalogSnapshot({ - agentId: escalationAgentId, - readOnly, - }); + fullSnapshot = await loadDeferredCatalog(params.context, escalationAgentId, readOnly); return fullSnapshot; }, timeoutFullDiscovery: true, @@ -616,7 +611,8 @@ export async function buildModelsListResult( const outcomeProjection = providerOutcomes?.length ? { providerOutcomes } : {}; const preparedProjectionOwner = ownerSnapshot ?? params.catalogProjector; const metadataSnapshot = preparedProjectionOwner?.metadataSnapshot; - const preparedAuthStore = preparedProjectionOwner?.authStore; + const preparedAuthStore = + (await resolveDeferredAuthStore(ownerSnapshot)) ?? params.catalogProjector?.authStore; if (!metadataSnapshot || !preparedAuthStore) { throw new Error("Gateway model catalog owner omitted prepared metadata or auth state"); } diff --git a/src/gateway/server-methods/shared-types.ts b/src/gateway/server-methods/shared-types.ts index 7868457a215b..3320521b099f 100644 --- a/src/gateway/server-methods/shared-types.ts +++ b/src/gateway/server-methods/shared-types.ts @@ -224,6 +224,7 @@ type GatewayKernelContext = { loadGatewayModelCatalogSnapshot: (params?: { agentId?: string; agentDir?: string; + deferAuthRefresh?: boolean; readOnly?: boolean; workspaceDir?: string; }) => Promise; diff --git a/src/gateway/server-model-catalog-auth.ts b/src/gateway/server-model-catalog-auth.ts new file mode 100644 index 000000000000..53fa540f06af --- /dev/null +++ b/src/gateway/server-model-catalog-auth.ts @@ -0,0 +1,38 @@ +import type { AuthProfileStore } from "../agents/auth-profiles/types.js"; +import type { GatewayRequestContext } from "./server-methods/shared-types.js"; + +const pendingAuthStoreBySnapshot = new WeakMap>(); + +export function setPendingGatewayModelCatalogAuthStore( + snapshot: object, + pending: Promise, +): void { + pendingAuthStoreBySnapshot.set(snapshot, pending); + // A timed-out catalog read may abandon the snapshot before it reaches the auth projection. + // Observe rejection here while preserving it for a caller that does resolve this snapshot. + void pending.catch(() => undefined); +} + +export async function resolveDeferredAuthStore( + snapshot: + | { + authStore?: AuthProfileStore; + } + | undefined, +): Promise { + return snapshot + ? ((await pendingAuthStoreBySnapshot.get(snapshot)) ?? snapshot.authStore) + : undefined; +} + +export function loadDeferredCatalog( + context: Pick, + agentId: string, + readOnly: boolean, +) { + return context.loadGatewayModelCatalogSnapshot({ + agentId, + deferAuthRefresh: true, + readOnly, + }); +} diff --git a/src/gateway/server-model-catalog.ts b/src/gateway/server-model-catalog.ts index 53f69bb6a9f8..906b107ebb4f 100644 --- a/src/gateway/server-model-catalog.ts +++ b/src/gateway/server-model-catalog.ts @@ -6,6 +6,7 @@ import { } from "../agents/prepared-model-runtime-auth.js"; // Gateway catalog reads use the atomic prepared runtime generation. import { getRuntimeConfig } from "../config/io.js"; +import { setPendingGatewayModelCatalogAuthStore } from "./server-model-catalog-auth.js"; import type { GatewayModelCatalogOwnerSnapshot, GatewayModelCatalogSnapshot, @@ -25,6 +26,7 @@ type LoadPublishedPreparedModelCatalogOwnerSnapshot = (params: { type LoadGatewayModelCatalogParams = { agentId?: string; agentDir?: string; + deferAuthRefresh?: boolean; getConfig?: () => GatewayModelCatalogConfig; loadPublishedPreparedModelCatalogOwnerSnapshot?: LoadPublishedPreparedModelCatalogOwnerSnapshot; readOnly?: boolean; @@ -55,11 +57,12 @@ export async function resetPreparedModelCatalogStateForTest(): Promise { async function loadGatewayModelCatalogOwnerSnapshot( params?: LoadGatewayModelCatalogParams, -): Promise< - GatewayModelCatalogOwnerSnapshot & { +): Promise<{ + candidate: PublishedModelCatalogOwnerCandidate; + owner: GatewayModelCatalogOwnerSnapshot & { authMaterializations: GatewayModelCatalogSnapshot["authMaterializations"]; - } -> { + }; +}> { const loadOwner = await resolveLoader(params); const candidate = await loadOwner({ ...(params?.agentId ? { agentId: params.agentId } : {}), @@ -70,13 +73,11 @@ async function loadGatewayModelCatalogOwnerSnapshot( }); const owner = resolvePublishedModelCatalogOwner(candidate); return { - ...owner, - authStore: - (await loadPreparedModelRuntimeAuthStore( - candidate, - owner.modelCatalog.entries.map((entry) => entry.provider), - )) ?? owner.authStore, - authMaterializations: getPreparedModelRuntimeAuthMaterializations(candidate), + candidate, + owner: { + ...owner, + authMaterializations: getPreparedModelRuntimeAuthMaterializations(candidate), + }, }; } @@ -101,7 +102,20 @@ function projectGatewayModelCatalogSnapshot( export async function loadGatewayModelCatalogSnapshot( params?: LoadGatewayModelCatalogParams, ): Promise { - return projectGatewayModelCatalogSnapshot(await loadGatewayModelCatalogOwnerSnapshot(params)); + const { candidate, owner } = await loadGatewayModelCatalogOwnerSnapshot(params); + const pendingAuthStore = loadPreparedModelRuntimeAuthStore( + candidate, + owner.modelCatalog.entries.map((entry) => entry.provider), + ); + if (params?.deferAuthRefresh) { + const snapshot = projectGatewayModelCatalogSnapshot(owner); + setPendingGatewayModelCatalogAuthStore(snapshot, pendingAuthStore); + return snapshot; + } + return projectGatewayModelCatalogSnapshot({ + ...owner, + authStore: (await pendingAuthStore) ?? owner.authStore, + }); } export async function loadGatewayModelCatalog(