From 940ffa2ece7fee34ba0c39d6619db254bf25c151 Mon Sep 17 00:00:00 2001 From: FullerStackDev <263060202+fuller-stack-dev@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:27:30 -0600 Subject: [PATCH] fix(gateway): keep context cache warmup responsive --- src/agents/context-cache.ts | 22 +- src/agents/context.lookup.test.ts | 181 ++++++++- .../context.prewarm.integration.test.ts | 157 ++++++++ src/agents/context.ts | 365 +++++++++++++----- .../server-startup-context-cache-prewarm.ts | 9 +- .../server-startup-post-attach.test.ts | 34 +- 6 files changed, 655 insertions(+), 113 deletions(-) create mode 100644 src/agents/context.prewarm.integration.test.ts diff --git a/src/agents/context-cache.ts b/src/agents/context-cache.ts index 6204d975794b..9c5984d7749a 100644 --- a/src/agents/context-cache.ts +++ b/src/agents/context-cache.ts @@ -1,7 +1,23 @@ /** Process-local model context window cache keyed by model id. */ -export const MODEL_CONTEXT_TOKEN_CACHE = new Map(); -export const MODEL_CONFIGURED_CONTEXT_TOKEN_CACHE = new Map(); -export const MODEL_CONTEXT_WINDOW_CACHE = new Map(); +export let MODEL_CONTEXT_TOKEN_CACHE = new Map(); +export let MODEL_CONFIGURED_CONTEXT_TOKEN_CACHE = new Map(); +export let MODEL_CONTEXT_WINDOW_CACHE = new Map(); + +/** Publish one complete cache generation without an O(N) main-loop copy. */ +export function replaceContextWindowCaches(params: { + configuredTokenCache: Map; + discoveredTokenCache: Map; + contextWindowCache: Map; +}): void { + MODEL_CONFIGURED_CONTEXT_TOKEN_CACHE = params.configuredTokenCache; + MODEL_CONTEXT_TOKEN_CACHE = params.discoveredTokenCache; + MODEL_CONTEXT_WINDOW_CACHE = params.contextWindowCache; +} + +/** Publish one complete discovered-metadata generation. */ +export function replaceDiscoveredContextTokenCache(cache: Map): void { + MODEL_CONTEXT_TOKEN_CACHE = cache; +} const PROVIDER_CONTEXT_TOKEN_CACHE_PREFIX = "\0provider:"; diff --git a/src/agents/context.lookup.test.ts b/src/agents/context.lookup.test.ts index 6624debd317e..a7c54e7ab1c8 100644 --- a/src/agents/context.lookup.test.ts +++ b/src/agents/context.lookup.test.ts @@ -3,7 +3,10 @@ import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import type { OpenClawConfig } from "../config/types.openclaw.js"; import { ANTHROPIC_CONTEXT_1M_TOKENS } from "./context-resolution.js"; -import { CONTEXT_WINDOW_RUNTIME_STATE } from "./context-runtime-state.js"; +import { + beginContextWindowCacheRefresh, + CONTEXT_WINDOW_RUNTIME_STATE, +} from "./context-runtime-state.js"; type DiscoveredModel = { id: string; @@ -21,6 +24,17 @@ const contextTestState = vi.hoisted(() => { runtimeConfigSnapshot: null as OpenClawConfig | null, runtimeConfigSourceSnapshot: null as OpenClawConfig | null, loadModelCatalogOwnerSnapshot: vi.fn(async (_params: unknown) => ({ + config: state.loadConfigImpl() as OpenClawConfig, + agentDir: "/tmp/context-catalog-agent", + modelCatalog: { + entries: state.discoveredModels, + routeVariants: [], + staticEntries: state.staticCatalogModels, + }, + })), + loadPublishedModelCatalogOwnerSnapshot: vi.fn(async (_params: unknown) => ({ + config: state.loadConfigImpl() as OpenClawConfig, + agentDir: "/tmp/context-catalog-agent", modelCatalog: { entries: state.discoveredModels, routeVariants: [], @@ -44,6 +58,8 @@ vi.mock("../config/runtime-source-projection.js", () => ({ vi.mock("./prepared-model-catalog.js", () => ({ loadPreparedModelCatalogOwnerSnapshot: contextTestState.loadModelCatalogOwnerSnapshot, + loadPublishedPreparedModelCatalogOwnerSnapshot: + contextTestState.loadPublishedModelCatalogOwnerSnapshot, })); function mockContextDeps(params: { @@ -88,6 +104,28 @@ function createContextOverrideConfig( }; } +function createLargeContextOverrideConfig(prefix: string, modelCount: number): OpenClawConfig { + return { + models: { + providers: { + synthetic: { + baseUrl: "https://example.invalid", + api: "openai-completions", + models: Array.from({ length: modelCount }, (_, index) => ({ + id: `${prefix}-${index}`, + name: `${prefix} ${index}`, + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128_000, + maxTokens: 8_192, + })), + }, + }, + }, + }; +} + async function flushAsyncWarmup() { // Warmup may run via timers or microtasks depending on the import path; flush // both so assertions observe stable cache state. @@ -134,6 +172,18 @@ describe("lookupContextTokens", () => { contextTestState.runtimeConfigSourceSnapshot = null; contextTestState.loadModelCatalogOwnerSnapshot.mockClear(); contextTestState.loadModelCatalogOwnerSnapshot.mockImplementation(async () => ({ + config: contextTestState.loadConfigImpl() as OpenClawConfig, + agentDir: "/tmp/context-catalog-agent", + modelCatalog: { + entries: contextTestState.discoveredModels, + routeVariants: [], + staticEntries: contextTestState.staticCatalogModels, + }, + })); + contextTestState.loadPublishedModelCatalogOwnerSnapshot.mockClear(); + contextTestState.loadPublishedModelCatalogOwnerSnapshot.mockImplementation(async () => ({ + config: contextTestState.loadConfigImpl() as OpenClawConfig, + agentDir: "/tmp/context-catalog-agent", modelCatalog: { entries: contextTestState.discoveredModels, routeVariants: [], @@ -379,6 +429,31 @@ describe("lookupContextTokens", () => { ).toBe(ANTHROPIC_CONTEXT_1M_TOKENS); }); + it("cooperatively projects a large catalog when request-time loading starts first", async () => { + const modelCount = 1_025; + const config = createContextOverrideConfig("synthetic", "configured", 128_000); + contextTestState.discoveredModels = Array.from({ length: modelCount }, (_, index) => ({ + id: `request-discovered-${index}`, + provider: "synthetic", + contextWindow: 64_000, + })); + const immediateSpy = vi.spyOn(globalThis, "setImmediate"); + try { + await contextModule.ensureContextWindowCacheLoaded(config); + + expect(contextTestState.loadModelCatalogOwnerSnapshot).toHaveBeenCalledOnce(); + expect(immediateSpy).toHaveBeenCalled(); + expect( + contextModule.lookupContextTokens("request-discovered-1024", { + allowAsyncLoad: false, + skipRuntimeConfigLoad: true, + }), + ).toBe(64_000); + } finally { + immediateSpy.mockRestore(); + } + }); + it("warms fresh caches instead of reusing a pre-generation load promise", async () => { const legacyLoadPromise = Promise.resolve(); CONTEXT_WINDOW_RUNTIME_STATE.loadPromise = legacyLoadPromise; @@ -425,6 +500,110 @@ describe("lookupContextTokens", () => { } }); + it("cooperatively warms the published owner without exact-config matching", async () => { + const modelCount = 1_025; + const config = createLargeContextOverrideConfig("configured", modelCount); + contextTestState.loadConfigImpl = () => config; + contextTestState.discoveredModels = Array.from({ length: modelCount }, (_, index) => ({ + id: `discovered-${index}`, + provider: "synthetic", + contextWindow: 64_000, + })); + const immediateSpy = vi.spyOn(globalThis, "setImmediate"); + try { + await contextModule.prewarmContextWindowCacheAfterReady({ config }); + + expect(contextTestState.loadPublishedModelCatalogOwnerSnapshot).toHaveBeenCalledWith( + expect.objectContaining({ config, readOnly: true }), + ); + expect(contextTestState.loadModelCatalogOwnerSnapshot).not.toHaveBeenCalled(); + expect(immediateSpy).toHaveBeenCalled(); + expect( + contextModule.lookupContextTokens("discovered-1024", { + allowAsyncLoad: false, + skipRuntimeConfigLoad: true, + }), + ).toBe(64_000); + expect( + contextModule.lookupContextTokens("configured-1024", { + allowAsyncLoad: false, + skipRuntimeConfigLoad: true, + }), + ).toBe(128_000); + } finally { + immediateSpy.mockRestore(); + } + }); + + it("does not yield while warming a small published catalog", async () => { + const config = createContextOverrideConfig("synthetic", "configured-small", 128_000); + contextTestState.loadConfigImpl = () => config; + contextTestState.discoveredModels = [ + { id: "discovered-small", provider: "synthetic", contextWindow: 64_000 }, + ]; + const immediateSpy = vi.spyOn(globalThis, "setImmediate"); + try { + await contextModule.prewarmContextWindowCacheAfterReady({ config }); + expect(immediateSpy).not.toHaveBeenCalled(); + } finally { + immediateSpy.mockRestore(); + } + }); + + it("does not publish a cooperatively prepared stale generation", async () => { + const config = createLargeContextOverrideConfig("stale", 1_025); + contextTestState.loadConfigImpl = () => config; + const originalSetImmediate = globalThis.setImmediate; + const immediateSpy = vi + .spyOn(globalThis, "setImmediate") + .mockImplementationOnce((callback, ...args) => + originalSetImmediate(() => { + beginContextWindowCacheRefresh(); + callback(...args); + }), + ); + try { + await contextModule.prewarmContextWindowCacheAfterReady({ config }); + expect( + contextModule.lookupContextTokens("stale-0", { + allowAsyncLoad: false, + skipRuntimeConfigLoad: true, + }), + ).toBeUndefined(); + } finally { + immediateSpy.mockRestore(); + } + }); + + it("stops an in-flight optional warmup without publishing partial caches", async () => { + const config = createLargeContextOverrideConfig("cancelled", 1_025); + contextTestState.loadConfigImpl = () => config; + let cancelled = false; + const originalSetImmediate = globalThis.setImmediate; + const immediateSpy = vi + .spyOn(globalThis, "setImmediate") + .mockImplementationOnce((callback, ...args) => + originalSetImmediate(() => { + cancelled = true; + callback(...args); + }), + ); + try { + await contextModule.prewarmContextWindowCacheAfterReady({ + config, + isCancelled: () => cancelled, + }); + expect( + contextModule.lookupContextTokens("cancelled-0", { + allowAsyncLoad: false, + skipRuntimeConfigLoad: true, + }), + ).toBeUndefined(); + } finally { + immediateSpy.mockRestore(); + } + }); + it("warms context metadata from bundled provider static catalogs", async () => { contextTestState.staticCatalogModels = [ { diff --git a/src/agents/context.prewarm.integration.test.ts b/src/agents/context.prewarm.integration.test.ts new file mode 100644 index 000000000000..ef762531c4c5 --- /dev/null +++ b/src/agents/context.prewarm.integration.test.ts @@ -0,0 +1,157 @@ +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { monitorEventLoopDelay, performance } from "node:perf_hooks"; +import { afterAll, describe, expect, it, vi } from "vitest"; +import type { OpenClawConfig } from "../config/types.openclaw.js"; +import { resetContextWindowCacheForTest } from "./context-runtime-state.js"; +import { resetPreparedModelRuntimeSnapshotsForTest } from "./prepared-model-runtime.test-support.js"; + +const originalHome = process.env.HOME; +const originalOpenClawHome = process.env.OPENCLAW_HOME; +const originalStateDir = process.env.OPENCLAW_STATE_DIR; +const root = fs.mkdtempSync(path.join(os.tmpdir(), "openclaw-context-prewarm-")); +const workspaceDir = path.join(root, "workspace"); +fs.mkdirSync(workspaceDir, { recursive: true }); +process.env.HOME = root; +process.env.OPENCLAW_HOME = root; +process.env.OPENCLAW_STATE_DIR = path.join(root, "state"); + +afterAll(() => { + resetContextWindowCacheForTest(); + resetPreparedModelRuntimeSnapshotsForTest(); + process.env.HOME = originalHome; + if (originalOpenClawHome === undefined) { + delete process.env.OPENCLAW_HOME; + } else { + process.env.OPENCLAW_HOME = originalOpenClawHome; + } + if (originalStateDir === undefined) { + delete process.env.OPENCLAW_STATE_DIR; + } else { + process.env.OPENCLAW_STATE_DIR = originalStateDir; + } + fs.rmSync(root, { recursive: true, force: true }); +}); + +describe("post-ready context cache prewarm", () => { + it("yields through the real prepared catalog lifecycle and converges atomically", async () => { + const modelCount = Number.parseInt(process.env.SYNTHETIC_MODEL_COUNT ?? "20000", 10); + const makeModels = (provider: string, baseWindow: number) => + Array.from({ length: modelCount }, (_, index) => ({ + id: index === 0 ? "shared-model" : `${provider}-model-${index}`, + name: `${provider} model ${index}`, + reasoning: false, + input: ["text" as const], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: baseWindow + (index % 17), + maxTokens: 8_192, + })); + const config = { + agents: { + defaults: { + workspace: workspaceDir, + model: { primary: "synthetic-a/shared-model" }, + models: { + "synthetic-a/shared-model": { agentRuntime: { id: "openclaw" } }, + "synthetic-b/shared-model": { agentRuntime: { id: "openclaw" } }, + }, + }, + }, + models: { + providers: { + "synthetic-a": { + baseUrl: "http://127.0.0.1:1/v1", + api: "openai-completions" as const, + models: makeModels("synthetic-a", 128_000), + }, + "synthetic-b": { + baseUrl: "http://127.0.0.1:2/v1", + api: "openai-completions" as const, + models: makeModels("synthetic-b", 64_000), + }, + }, + }, + } satisfies OpenClawConfig; + + const runtimeModule = await import("./prepared-model-runtime.js"); + const contextModule = await import("./context.js"); + await runtimeModule.refreshPreparedModelRuntimeSnapshots(config, { + gatewayLifecycle: true, + catalogMode: "static", + defaultWorkspaceDir: workspaceDir, + }); + contextModule.resetContextWindowCacheForTest(); + + const heartbeatGaps: number[] = []; + let lastHeartbeatAt = performance.now(); + const heartbeat = setInterval(() => { + const now = performance.now(); + heartbeatGaps.push(now - lastHeartbeatAt); + lastHeartbeatAt = now; + }, 20); + const delayMonitor = monitorEventLoopDelay({ resolution: 20 }); + delayMonitor.enable(); + const eluStart = performance.eventLoopUtilization(); + const immediateSpy = vi.spyOn(globalThis, "setImmediate"); + let immediateCallCount: number; + const warmStartedAt = performance.now(); + let warmCompletedAt = warmStartedAt; + try { + await contextModule.prewarmContextWindowCacheAfterReady({ + config: process.env.SYNTHETIC_CLONE_CONFIG === "0" ? config : structuredClone(config), + }); + warmCompletedAt = performance.now(); + await new Promise((resolve) => { + setTimeout(resolve, 40); + }); + } finally { + immediateCallCount = immediateSpy.mock.calls.length; + immediateSpy.mockRestore(); + delayMonitor.disable(); + clearInterval(heartbeat); + } + const warmMs = warmCompletedAt - warmStartedAt; + const elu = performance.eventLoopUtilization(performance.eventLoopUtilization(), eluStart); + const maxHeartbeatGapMs = Math.max(...heartbeatGaps); + const delayMaxMs = delayMonitor.max / 1_000_000; + console.log( + `CONTEXT_PREWARM_GREEN ${JSON.stringify({ + modelCount, + warmMs, + maxHeartbeatGapMs, + delayMaxMs, + delayP99Ms: delayMonitor.percentile(99) / 1_000_000, + eventLoopUtilization: elu.utilization, + immediateCallCount, + })}`, + ); + + expect(immediateCallCount).toBeGreaterThan(0); + expect(heartbeatGaps.length).toBeGreaterThan(0); + expect(maxHeartbeatGapMs).toBeLessThan(500); + expect(delayMaxMs).toBeLessThan(500); + expect( + contextModule.lookupContextTokens("shared-model", { + allowAsyncLoad: false, + skipRuntimeConfigLoad: true, + }), + ).toBe(64_000); + expect( + contextModule.resolveContextTokensForModel({ + cfg: config, + provider: "synthetic-a", + model: "shared-model", + allowAsyncLoad: false, + }), + ).toBe(128_000); + expect( + contextModule.resolveContextTokensForModel({ + cfg: config, + provider: "synthetic-b", + model: "shared-model", + allowAsyncLoad: false, + }), + ).toBe(64_000); + }, 30_000); +}); diff --git a/src/agents/context.ts b/src/agents/context.ts index 30661cd5b4af..a85dca179382 100644 --- a/src/agents/context.ts +++ b/src/agents/context.ts @@ -12,9 +12,10 @@ import { lookupCachedContextWindow, minPositiveContextTokens, MODEL_CONFIGURED_CONTEXT_TOKEN_CACHE, - MODEL_CONTEXT_TOKEN_CACHE, MODEL_CONTEXT_WINDOW_CACHE, providerContextTokenCacheKey, + replaceContextWindowCaches, + replaceDiscoveredContextTokenCache, } from "./context-cache.js"; import { type ContextTokenResolutionParams, @@ -44,57 +45,112 @@ type ModelEntry = { contextWindow?: number; contextTokens?: number; }; +type ConfiguredProviderEntry = NonNullable[string]>; +type ConfiguredModelEntry = NonNullable[number]; const CONFIG_LOAD_RETRY_POLICY: BackoffPolicy = { initialMs: 1_000, maxMs: 60_000, factor: 2, jitter: 0, }; +const CONTEXT_CACHE_PREWARM_BATCH_SIZE = 512; const loadPreparedModelCatalogRuntime = () => import("./prepared-model-catalog.js"); +function cacheMinimum(cache: Map, key: string, contextTokens: number): void { + const existing = cache.get(key); + if (existing === undefined || contextTokens < existing) { + cache.set(key, contextTokens); + } +} + +function applyDiscoveredContextWindow(cache: Map, model: ModelEntry): void { + if (!model?.id) { + return; + } + const discoveredContextTokens = + typeof model.contextTokens === "number" + ? Math.trunc(model.contextTokens) + : typeof model.contextWindow === "number" + ? Math.trunc(model.contextWindow) + : undefined; + const contextTokens = + resolveDiscoveredAnthropicFixedContextWindow(model) ?? discoveredContextTokens; + if (!contextTokens || contextTokens <= 0) { + return; + } + // Cache the most conservative effective limit. Provider/runtime callers that + // know the active provider prefer the provider-owned entry below. + cacheMinimum(cache, model.id, contextTokens); + if (typeof model.provider !== "string") { + return; + } + const provider = normalizeProviderId(model.provider); + if (!provider) { + return; + } + cacheMinimum(cache, providerContextTokenCacheKey(provider, model.id), contextTokens); + const slash = model.id.indexOf("/"); + const prefixedProvider = slash > 0 ? normalizeProviderId(model.id.slice(0, slash)) : ""; + const bareModelId = slash > 0 ? model.id.slice(slash + 1).trim() : ""; + // Some registries preserve a self-prefixed id alongside provider ownership. + // Cache its bare form without stripping cross-provider ids such as OpenRouter rows. + if (prefixedProvider === provider && bareModelId) { + cacheMinimum(cache, providerContextTokenCacheKey(provider, bareModelId), contextTokens); + } +} + export function applyDiscoveredContextWindows(params: { cache: Map; models: ModelEntry[]; }) { - const cacheMinimum = (key: string, contextTokens: number) => { - const existing = params.cache.get(key); - if (existing === undefined || contextTokens < existing) { - params.cache.set(key, contextTokens); - } - }; - for (const model of params.models) { - if (!model?.id) { - continue; - } - const discoveredContextTokens = - typeof model.contextTokens === "number" - ? Math.trunc(model.contextTokens) - : typeof model.contextWindow === "number" - ? Math.trunc(model.contextWindow) - : undefined; - const contextTokens = - resolveDiscoveredAnthropicFixedContextWindow(model) ?? discoveredContextTokens; - if (!contextTokens || contextTokens <= 0) { - continue; - } - // Cache the most conservative effective limit. Provider/runtime callers that - // know the active provider prefer the provider-owned entry below. - cacheMinimum(model.id, contextTokens); - if (typeof model.provider === "string") { - const provider = normalizeProviderId(model.provider); - if (provider) { - cacheMinimum(providerContextTokenCacheKey(provider, model.id), contextTokens); - const slash = model.id.indexOf("/"); - const prefixedProvider = slash > 0 ? normalizeProviderId(model.id.slice(0, slash)) : ""; - const bareModelId = slash > 0 ? model.id.slice(slash + 1).trim() : ""; - // Some registries preserve a self-prefixed id alongside provider ownership. - // Cache its bare form without stripping cross-provider ids such as OpenRouter rows. - if (prefixedProvider === provider && bareModelId) { - cacheMinimum(providerContextTokenCacheKey(provider, bareModelId), contextTokens); - } - } - } + applyDiscoveredContextWindow(params.cache, model); + } +} + +function applyConfiguredContextWindow(params: { + cache: Map; + windowCache: Map; + providerId: string; + provider: ConfiguredProviderEntry; + model: ConfiguredModelEntry; +}): void { + const modelId = typeof params.model?.id === "string" ? params.model.id : undefined; + const contextTokens = + typeof params.model?.contextTokens === "number" + ? params.model.contextTokens + : typeof params.provider?.contextTokens === "number" + ? params.provider.contextTokens + : undefined; + const contextWindow = + typeof params.model?.contextWindow === "number" + ? params.model.contextWindow + : typeof params.provider?.contextWindow === "number" + ? params.provider.contextWindow + : undefined; + const configuredValue = + contextTokens && contextTokens > 0 + ? { cache: params.cache, value: contextTokens } + : contextWindow && contextWindow > 0 + ? { cache: params.windowCache, value: contextWindow } + : undefined; + if (!modelId || !configuredValue) { + return; + } + configuredValue.cache.set(modelId, configuredValue.value); + configuredValue.cache.set( + providerContextTokenCacheKey(normalizeProviderId(params.providerId), modelId), + configuredValue.value, + ); + const normalizedProvider = normalizeProviderId(params.providerId); + const slash = modelId.indexOf("/"); + const prefixedProvider = slash > 0 ? normalizeProviderId(modelId.slice(0, slash)) : ""; + const bareModelId = slash > 0 ? modelId.slice(slash + 1).trim() : ""; + if (normalizedProvider && prefixedProvider === normalizedProvider && bareModelId) { + configuredValue.cache.set( + providerContextTokenCacheKey(normalizedProvider, bareModelId), + configuredValue.value, + ); } } @@ -112,47 +168,178 @@ export function applyConfiguredContextWindows(params: { continue; } for (const model of provider.models) { - const modelId = typeof model?.id === "string" ? model.id : undefined; - const contextTokens = - typeof model?.contextTokens === "number" - ? model.contextTokens - : typeof provider?.contextTokens === "number" - ? provider.contextTokens - : undefined; - const contextWindow = - typeof model?.contextWindow === "number" - ? model.contextWindow - : typeof provider?.contextWindow === "number" - ? provider.contextWindow - : undefined; - const configuredValue = - contextTokens && contextTokens > 0 - ? { cache: params.cache, value: contextTokens } - : contextWindow && contextWindow > 0 - ? { cache: params.windowCache, value: contextWindow } - : undefined; - if (!modelId || !configuredValue) { - continue; - } - configuredValue.cache.set(modelId, configuredValue.value); - configuredValue.cache.set( - providerContextTokenCacheKey(normalizeProviderId(providerId), modelId), - configuredValue.value, - ); - const normalizedProvider = normalizeProviderId(providerId); - const slash = modelId.indexOf("/"); - const prefixedProvider = slash > 0 ? normalizeProviderId(modelId.slice(0, slash)) : ""; - const bareModelId = slash > 0 ? modelId.slice(slash + 1).trim() : ""; - if (normalizedProvider && prefixedProvider === normalizedProvider && bareModelId) { - configuredValue.cache.set( - providerContextTokenCacheKey(normalizedProvider, bareModelId), - configuredValue.value, - ); - } + applyConfiguredContextWindow({ + cache: params.cache, + windowCache: params.windowCache, + providerId, + provider, + model, + }); } } } +function yieldToEventLoop(): Promise { + return new Promise((resolve) => { + setImmediate(resolve); + }); +} + +async function applyConfiguredContextWindowsCooperatively(params: { + cache: Map; + windowCache: Map; + modelsConfig: ModelsConfig | undefined; + shouldStop: () => boolean; +}): Promise { + const providers = params.modelsConfig?.providers; + if (!providers || typeof providers !== "object") { + return !params.shouldStop(); + } + let processed = 0; + for (const [providerId, provider] of Object.entries(providers)) { + if (!Array.isArray(provider?.models)) { + continue; + } + for (const model of provider.models) { + if (params.shouldStop()) { + return false; + } + if (processed > 0 && processed % CONTEXT_CACHE_PREWARM_BATCH_SIZE === 0) { + await yieldToEventLoop(); + if (params.shouldStop()) { + return false; + } + } + applyConfiguredContextWindow({ + cache: params.cache, + windowCache: params.windowCache, + providerId, + provider, + model, + }); + processed += 1; + } + } + return !params.shouldStop(); +} + +async function applyDiscoveredContextWindowsCooperatively(params: { + cache: Map; + modelGroups: readonly (readonly ModelEntry[])[]; + shouldStop: () => boolean; +}): Promise { + let processed = 0; + for (const models of params.modelGroups) { + for (const model of models) { + if (params.shouldStop()) { + return false; + } + if (processed > 0 && processed % CONTEXT_CACHE_PREWARM_BATCH_SIZE === 0) { + await yieldToEventLoop(); + if (params.shouldStop()) { + return false; + } + } + applyDiscoveredContextWindow(params.cache, model); + processed += 1; + } + } + return !params.shouldStop(); +} + +/** + * Warm the process cache from the Gateway's currently published catalog owner + * without letting optional post-ready projection monopolize the main loop. + */ +export function prewarmContextWindowCacheAfterReady(params: { + config: OpenClawConfig; + isCancelled?: () => boolean; +}): Promise { + const generation = CONTEXT_WINDOW_RUNTIME_STATE.generation; + if ( + CONTEXT_WINDOW_RUNTIME_STATE.loadPromise && + CONTEXT_WINDOW_RUNTIME_STATE.loadGeneration === generation + ) { + return CONTEXT_WINDOW_RUNTIME_STATE.loadPromise; + } + + const shouldStop = () => + CONTEXT_WINDOW_RUNTIME_STATE.generation !== generation || params.isCancelled?.() === true; + const loadPromise = Promise.resolve() + .then(async () => { + if (shouldStop()) { + return; + } + let owner: + | Awaited< + ReturnType< + typeof import("./prepared-model-catalog.js").loadPublishedPreparedModelCatalogOwnerSnapshot + > + > + | undefined; + try { + const { loadPublishedPreparedModelCatalogOwnerSnapshot } = + await loadPreparedModelCatalogRuntime(); + const defaultAgentId = resolveDefaultAgentId(params.config); + owner = await loadPublishedPreparedModelCatalogOwnerSnapshot({ + config: params.config, + agentId: defaultAgentId, + agentDir: resolveAgentDir(params.config, defaultAgentId), + readOnly: true, + }); + } catch { + // Config-backed overrides still converge when the prepared owner is unavailable. + } + if (shouldStop()) { + return; + } + + const sourceConfig = owner?.config ?? params.config; + const stagedConfiguredTokenCache = new Map(); + const stagedContextWindowCache = new Map(); + const stagedDiscoveredTokenCache = new Map(); + if ( + !(await applyConfiguredContextWindowsCooperatively({ + cache: stagedConfiguredTokenCache, + windowCache: stagedContextWindowCache, + modelsConfig: sourceConfig.models as ModelsConfig | undefined, + shouldStop, + })) + ) { + return; + } + if ( + !(await applyDiscoveredContextWindowsCooperatively({ + cache: stagedDiscoveredTokenCache, + modelGroups: [owner?.modelCatalog.entries ?? [], owner?.modelCatalog.staticEntries ?? []], + shouldStop, + })) + ) { + return; + } + if (shouldStop()) { + return; + } + + // Publish one complete generation so yielded preparation never exposes a + // mix of old and new configured, static, or discovered metadata. + replaceContextWindowCaches({ + configuredTokenCache: stagedConfiguredTokenCache, + contextWindowCache: stagedContextWindowCache, + discoveredTokenCache: stagedDiscoveredTokenCache, + }); + CONTEXT_WINDOW_RUNTIME_STATE.configuredConfig = sourceConfig; + CONTEXT_WINDOW_RUNTIME_STATE.configLoadFailures = 0; + CONTEXT_WINDOW_RUNTIME_STATE.nextConfigLoadAttemptAtMs = 0; + }) + .catch(() => { + // Keep optional Gateway warmup best-effort. + }); + CONTEXT_WINDOW_RUNTIME_STATE.loadPromise = loadPromise; + CONTEXT_WINDOW_RUNTIME_STATE.loadGeneration = generation; + return loadPromise; +} + function primeConfiguredContextWindowsFromConfig(cfg: OpenClawConfig): OpenClawConfig { applyConfiguredContextWindows({ cache: MODEL_CONFIGURED_CONTEXT_TOKEN_CACHE, @@ -223,16 +410,17 @@ export function ensureContextWindowCacheLoaded(cfgOverride?: OpenClawConfig): Pr if (CONTEXT_WINDOW_RUNTIME_STATE.generation !== generation) { return; } - const models = - catalogResult.status === "fulfilled" ? catalogResult.value.modelCatalog.entries : []; - const providerStaticModels = - catalogResult.status === "fulfilled" - ? (catalogResult.value.modelCatalog.staticEntries ?? []) - : []; - applyDiscoveredContextWindows({ - cache: stagedTokenCache, - models: [...models, ...providerStaticModels], - }); + const modelCatalog = + catalogResult.status === "fulfilled" ? catalogResult.value.modelCatalog : undefined; + if ( + !(await applyDiscoveredContextWindowsCooperatively({ + cache: stagedTokenCache, + modelGroups: [modelCatalog?.entries ?? [], modelCatalog?.staticEntries ?? []], + shouldStop: () => CONTEXT_WINDOW_RUNTIME_STATE.generation !== generation, + })) + ) { + return; + } } catch { // Static and discovered rows belong to one atomic generation. If its owner fails, keep // config overrides only instead of mixing in independently rediscovered static metadata. @@ -241,10 +429,7 @@ export function ensureContextWindowCacheLoaded(cfgOverride?: OpenClawConfig): Pr if (CONTEXT_WINDOW_RUNTIME_STATE.generation !== generation) { return; } - MODEL_CONTEXT_TOKEN_CACHE.clear(); - for (const [key, value] of stagedTokenCache) { - MODEL_CONTEXT_TOKEN_CACHE.set(key, value); - } + replaceDiscoveredContextTokenCache(stagedTokenCache); }) .catch(() => { // Keep lookup best-effort. diff --git a/src/gateway/server-startup-context-cache-prewarm.ts b/src/gateway/server-startup-context-cache-prewarm.ts index 6cb29242158e..ab7d416857ba 100644 --- a/src/gateway/server-startup-context-cache-prewarm.ts +++ b/src/gateway/server-startup-context-cache-prewarm.ts @@ -14,7 +14,7 @@ type ContextCachePrewarmHandle = { }; export function scheduleContextCachePrewarm(params: { - cfgAtStart: OpenClawConfig; + getConfig: () => OpenClawConfig; startupTrace?: StartupTrace; log: { warn: (msg: string) => void }; }): ContextCachePrewarmHandle { @@ -23,9 +23,12 @@ export function scheduleContextCachePrewarm(params: { if (stopped) { return; } - const { ensureContextWindowCacheLoaded } = await import("../agents/context.js"); + const { prewarmContextWindowCacheAfterReady } = await import("../agents/context.js"); if (!stopped) { - await ensureContextWindowCacheLoaded(params.cfgAtStart); + await prewarmContextWindowCacheAfterReady({ + config: params.getConfig(), + isCancelled: () => stopped, + }); } }; diff --git a/src/gateway/server-startup-post-attach.test.ts b/src/gateway/server-startup-post-attach.test.ts index cb76b56ad8ad..9944480fa0a8 100644 --- a/src/gateway/server-startup-post-attach.test.ts +++ b/src/gateway/server-startup-post-attach.test.ts @@ -67,7 +67,7 @@ const hoisted = vi.hoisted(() => { async (_cfg?: unknown, _options?: unknown) => {}, ); const loadAgentRuntimePluginRegistryHandle = vi.fn(); - const ensureContextWindowCacheLoaded = vi.fn(async () => {}); + const prewarmContextWindowCacheAfterReady = vi.fn(async () => {}); const scheduleGatewayHandlerPrewarm = vi.fn(() => ({ stop: vi.fn() })); const clearCurrentProviderAuthState = vi.fn(); const warmCurrentProviderAuthStateOffMainThread = vi.fn( @@ -107,7 +107,7 @@ const hoisted = vi.hoisted(() => { prepareModelRuntimeSnapshot, refreshPreparedModelRuntimeSnapshots, loadAgentRuntimePluginRegistryHandle, - ensureContextWindowCacheLoaded, + prewarmContextWindowCacheAfterReady, scheduleGatewayHandlerPrewarm, clearCurrentProviderAuthState, warmCurrentProviderAuthStateOffMainThread, @@ -217,7 +217,7 @@ vi.mock("../agents/runtime-plugins.js", () => ({ })); vi.mock("../agents/context.js", () => ({ - ensureContextWindowCacheLoaded: hoisted.ensureContextWindowCacheLoaded, + prewarmContextWindowCacheAfterReady: hoisted.prewarmContextWindowCacheAfterReady, })); vi.mock("./server-startup-handler-prewarm.js", () => ({ @@ -368,8 +368,8 @@ describe("startGatewayPostAttachRuntime", () => { hoisted.refreshPreparedModelRuntimeSnapshots.mockReset(); hoisted.refreshPreparedModelRuntimeSnapshots.mockResolvedValue(undefined); hoisted.loadAgentRuntimePluginRegistryHandle.mockReset(); - hoisted.ensureContextWindowCacheLoaded.mockReset(); - hoisted.ensureContextWindowCacheLoaded.mockResolvedValue(undefined); + hoisted.prewarmContextWindowCacheAfterReady.mockReset(); + hoisted.prewarmContextWindowCacheAfterReady.mockResolvedValue(undefined); hoisted.scheduleGatewayHandlerPrewarm.mockClear(); hoisted.clearCurrentProviderAuthState.mockClear(); hoisted.warmCurrentProviderAuthStateOffMainThread.mockReset(); @@ -1258,32 +1258,34 @@ describe("startGatewayPostAttachRuntime", () => { it("defers context-window cache prewarm to a post-ready sidecar", async () => { vi.useFakeTimers(); - const cfg = { agents: { defaults: { model: "openai/gpt-5.5" } } } as never; + const cfg = { agents: { defaults: { model: "openai/gpt-5.5" } } }; + const currentConfig = { ...cfg }; const admission = tryBeginGatewayRootWorkAdmission(); if (!admission) { throw new Error("Expected request work admission"); } const sidecar = scheduleContextCachePrewarm({ - cfgAtStart: cfg, + getConfig: () => currentConfig, log: { warn: vi.fn() }, }); try { - // Earlier gateway lifetimes may finish during the fake-clock window; - // this sidecar's captured config identifies its own prewarm precisely. - expect(hoisted.ensureContextWindowCacheLoaded).not.toHaveBeenCalledWith(cfg); + expect(hoisted.prewarmContextWindowCacheAfterReady).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(4_999); - expect(hoisted.ensureContextWindowCacheLoaded).not.toHaveBeenCalledWith(cfg); + expect(hoisted.prewarmContextWindowCacheAfterReady).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(1); - expect(hoisted.ensureContextWindowCacheLoaded).not.toHaveBeenCalledWith(cfg); + expect(hoisted.prewarmContextWindowCacheAfterReady).not.toHaveBeenCalled(); admission.release(); await vi.advanceTimersByTimeAsync(249); - expect(hoisted.ensureContextWindowCacheLoaded).not.toHaveBeenCalledWith(cfg); + expect(hoisted.prewarmContextWindowCacheAfterReady).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(1); await vi.dynamicImportSettled(); await waitForGatewayTestState(() => { - expect(hoisted.ensureContextWindowCacheLoaded).toHaveBeenCalledWith(cfg); + expect(hoisted.prewarmContextWindowCacheAfterReady).toHaveBeenCalledWith({ + config: currentConfig, + isCancelled: expect.any(Function), + }); }); } finally { admission.release(); @@ -1294,13 +1296,13 @@ describe("startGatewayPostAttachRuntime", () => { it("cancels context-window cache prewarm when the gateway stops first", async () => { vi.useFakeTimers(); const sidecar = scheduleContextCachePrewarm({ - cfgAtStart: {} as never, + getConfig: () => ({}) as never, log: { warn: vi.fn() }, }); await sidecar.stop(); await vi.runAllTimersAsync(); - expect(hoisted.ensureContextWindowCacheLoaded).not.toHaveBeenCalled(); + expect(hoisted.prewarmContextWindowCacheAfterReady).not.toHaveBeenCalled(); }); it("keeps provider auth prewarm alive when Gmail post-ready sidecars stop", async () => {