diff --git a/src/gateway/server-methods/models-auth-status-usage-cache.ts b/src/gateway/server-methods/models-auth-status-usage-cache.ts index 2d8c8eefd897..ffc8232394ef 100644 --- a/src/gateway/server-methods/models-auth-status-usage-cache.ts +++ b/src/gateway/server-methods/models-auth-status-usage-cache.ts @@ -1,30 +1,19 @@ // Stale-while-revalidate cache for models.authStatus provider usage enrichment. -import { - ensureAuthProfileStore, - externalCliDiscoveryForConfigStatus, - type AuthProfileStore, -} from "../../agents/auth-profiles.js"; -import { - fingerprintAuthProfileCredential, - fingerprintAuthProfileOwnerShape, - fingerprintResolvedProviderAuth, -} from "../../agents/execution-auth-binding.js"; -import { - resolveLegacyInheritedAuthAgentId, - resolveLegacyInheritedAuthDir, -} from "../../agents/legacy-inherited-auth-dir.js"; -import { resolveEnvApiKey } from "../../agents/model-auth-env.js"; -import { resolveUsableCustomProviderApiKey } from "../../agents/model-auth.js"; +import type { AuthProfileStore } from "../../agents/auth-profiles.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; import { loadProviderUsageSummary } from "../../infra/provider-usage.load.js"; +import { PROVIDER_USAGE_TIMEOUT_MS } from "../../infra/provider-usage.shared.js"; import type { ProviderUsageSnapshot, UsageProviderId, UsageSummary, } from "../../infra/provider-usage.types.js"; import { createSubsystemLogger } from "../../logging/subsystem.js"; -import { listProviderUsagePluginDescriptors } from "../../plugins/provider-runtime.js"; import { formatForLog } from "../ws-log.js"; +import { + clearProviderUsageRuntimeSnapshot, + getProviderUsageRuntimeSnapshot, +} from "./provider-usage-runtime.js"; const log = createSubsystemLogger("provider-usage-cache"); const USAGE_CACHE_TTL_MS = 60_000; @@ -56,59 +45,11 @@ const usageCacheByAgentId = new Map(); const usageRefreshByAgentId = new Map(); let cacheGeneration = 0; -function sortedRecordEntries(value: Record | undefined) { - return Object.entries(value ?? {}).toSorted(([left], [right]) => left.localeCompare(right)); -} - -export function fingerprintProviderUsageCredentials(params: { - cfg: OpenClawConfig; - directApiKeys: ReadonlyMap; - store: AuthProfileStore; -}): string { - const profiles = Object.entries(params.store.profiles) - .toSorted(([left], [right]) => left.localeCompare(right)) - .map(([profileId, credential]) => { - const fingerprint = - fingerprintAuthProfileCredential({ profileId, credential }) ?? - fingerprintAuthProfileOwnerShape({ profileId, credential }); - return fingerprint ?? `${profileId}:${credential.type}:${credential.provider}`; - }); - const direct = [...params.directApiKeys] - .toSorted(([left], [right]) => left.localeCompare(right)) - .map(([provider, evidence]) => { - const configured = resolveUsableCustomProviderApiKey({ - cfg: params.cfg, - provider, - env: process.env, - }); - const envValue = evidence?.envVar ? process.env[evidence.envVar]?.trim() : undefined; - const resolved = - configured ?? - (envValue ? { apiKey: envValue, source: `env: ${evidence?.envVar}` } : undefined); - const fingerprint = resolved - ? fingerprintResolvedProviderAuth({ - apiKey: resolved.apiKey, - source: resolved.source, - mode: "api-key", - }) - : undefined; - return [provider, fingerprint ?? null]; - }); - // Profile selection can switch accounts without changing the profile set. - // Include every non-secret selector that resolveAuthProfileOrder consults. - return JSON.stringify({ - profiles, - direct, - order: sortedRecordEntries(params.store.order), - lastGood: sortedRecordEntries(params.store.lastGood), - usageStats: sortedRecordEntries(params.store.usageStats), - }); -} - export function clearModelAuthStatusUsageCache(): void { cacheGeneration += 1; usageCacheByAgentId.clear(); usageRefreshByAgentId.clear(); + clearProviderUsageRuntimeSnapshot(); } function providerUsageCacheKey(providerIds: readonly UsageProviderId[]): string { @@ -155,13 +96,37 @@ function mapProviderUsage(usage: Awaited provider.error === undefined) + .map((provider) => [provider.provider, provider]), + ); + return { + ...summary, + providers: summary.providers.map((provider) => + provider.error === "Timeout" + ? (lastGoodByProvider.get(provider.provider) ?? provider) + : provider, + ), + }; +} + function scheduleProviderUsageRefresh(params: { agentId: string; agentDir: string; + authStore?: AuthProfileStore; configRef: OpenClawConfig; credentialKey: string; providerIds: UsageProviderId[]; providerKey: string; + lastGood?: UsageSummary; }): Promise { const active = usageRefreshByAgentId.get(params.agentId); if ( @@ -176,10 +141,12 @@ function scheduleProviderUsageRefresh(params: { const promise = loadProviderUsageSummary({ providers: params.providerIds, agentDir: params.agentDir, + authStore: params.authStore, config: params.configRef, - timeoutMs: 3500, + timeoutMs: PROVIDER_USAGE_TIMEOUT_MS, }) - .then((usage) => { + .then((freshUsage) => { + const usage = retainLastGoodOnTimeout(freshUsage, params.lastGood); if ( publishGeneration === cacheGeneration && usageRefreshByAgentId.get(params.agentId) === refresh @@ -223,6 +190,7 @@ function scheduleProviderUsageRefresh(params: { type ProviderUsageCacheParams = { agentId: string; agentDir: string; + authStore?: AuthProfileStore; configRef: OpenClawConfig; credentialKey: string; forceRefresh?: boolean; @@ -264,10 +232,12 @@ export function readProviderUsageStaleWhileRevalidate( void scheduleProviderUsageRefresh({ agentId: params.agentId, agentDir: params.agentDir, + authStore: params.authStore, configRef: params.configRef, credentialKey, providerIds, providerKey, + lastGood: matching?.summary, }).catch(() => {}); } return matching?.usageByProvider ?? new Map(); @@ -289,10 +259,12 @@ async function loadProviderUsageSummaryStaleWhileRevalidate( const refresh = scheduleProviderUsageRefresh({ agentId: params.agentId, agentDir: params.agentDir, + authStore: params.authStore, configRef: params.configRef, credentialKey, providerIds, providerKey, + lastGood: matching?.summary, }); if (matching) { void refresh.catch(() => {}); @@ -306,42 +278,14 @@ export async function loadUsageStatusStaleWhileRevalidate(params: { config: OpenClawConfig; now?: number; }): Promise { - const agentId = resolveLegacyInheritedAuthAgentId(params.config); - const agentDir = resolveLegacyInheritedAuthDir(params.config); - const store = ensureAuthProfileStore(agentDir, { - externalCli: externalCliDiscoveryForConfigStatus({ cfg: params.config }), - }); - const providerIds = listProviderUsagePluginDescriptors({ - config: params.config, - env: process.env, - }).map((descriptor) => descriptor.provider); - const directApiKeys = new Map< - string, - { source: "config" | "env"; envVar?: string } | undefined - >(); - for (const provider of providerIds) { - const resolved = - resolveUsableCustomProviderApiKey({ - cfg: params.config, - provider, - env: process.env, - }) ?? resolveEnvApiKey(provider, process.env, { config: params.config }); - if (!resolved) { - continue; - } - const envVar = resolved.source.match(/^(?:shell env|env): ([A-Z][A-Z0-9_]*)$/u)?.[1]; - directApiKeys.set(provider, envVar ? { source: "env", envVar } : { source: "config" }); - } + const snapshot = getProviderUsageRuntimeSnapshot({ config: params.config }); return await loadProviderUsageSummaryStaleWhileRevalidate({ - agentId, - agentDir, - configRef: params.config, - credentialKey: fingerprintProviderUsageCredentials({ - cfg: params.config, - directApiKeys, - store, - }), - providerIds, + agentId: snapshot.agentId, + agentDir: snapshot.agentDir, + authStore: snapshot.store, + configRef: snapshot.configRef, + credentialKey: snapshot.credentialKey, + providerIds: snapshot.providerIds, now: params.now ?? Date.now(), }); } diff --git a/src/gateway/server-methods/models-auth-status.test.ts b/src/gateway/server-methods/models-auth-status.test.ts index 956386b5579f..1b2d834e5e3d 100644 --- a/src/gateway/server-methods/models-auth-status.test.ts +++ b/src/gateway/server-methods/models-auth-status.test.ts @@ -6,7 +6,10 @@ import { MAX_DATE_TIMESTAMP_MS } from "@openclaw/normalization-core/number-coerc import { createRequireRecord } from "openclaw/plugin-sdk/test-fixtures"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { AuthHealthSummary } from "../../agents/auth-health.js"; -import type { AuthProfileStore } from "../../agents/auth-profiles.js"; +import { + replaceRuntimeAuthProfileStoreSnapshots, + type AuthProfileStore, +} from "../../agents/auth-profiles.js"; import { NON_ENV_SECRETREF_MARKER } from "../../agents/model-auth-markers.js"; import type { UsageSummary } from "../../infra/provider-usage.types.js"; import { resolveInstalledPluginIndexPolicyHash } from "../../plugins/installed-plugin-index-policy.js"; @@ -54,6 +57,11 @@ const mocks = vi.hoisted(() => ({ (): AuthHealthSummary => ({ now: 0, warnAfterMs: 0, profiles: [], providers: [] }), ), loadProviderUsageSummary: vi.fn(async (): Promise => emptyUsageSummary()), + listProviderUsagePluginDescriptors: vi.fn(() => [ + { provider: "anthropic", displayName: "Claude" }, + { provider: "deepseek", displayName: "DeepSeek" }, + { provider: "openai", displayName: "OpenAI" }, + ]), })); vi.mock("../../config/config.js", () => ({ @@ -95,6 +103,10 @@ vi.mock("../../infra/provider-usage.load.js", () => ({ loadProviderUsageSummary: mocks.loadProviderUsageSummary, })); +vi.mock("../../plugins/provider-runtime.js", () => ({ + listProviderUsagePluginDescriptors: mocks.listProviderUsagePluginDescriptors, +})); + vi.mock("../../secrets/runtime.js", () => ({ refreshActiveProviderAuthRuntimeSnapshot: mocks.refreshActiveProviderAuthRuntimeSnapshot, })); @@ -201,6 +213,7 @@ let preparedMetadataSnapshot: unknown; function setPreparedAuthStore(store: AuthProfileStore): void { preparedAuthStore = store; + replaceRuntimeAuthProfileStoreSnapshots([{ agentDir: "/tmp/agent", store }]); } function setPreparedMetadataSnapshot(snapshot: unknown): void { @@ -1051,8 +1064,9 @@ describe("models.authStatus", () => { expect(mocks.loadProviderUsageSummary).toHaveBeenCalledWith({ providers: ["anthropic"], agentDir: "/tmp/agent", + authStore: preparedAuthStore, config: runtimeConfig, - timeoutMs: 3500, + timeoutMs: 5_000, }); let result: ModelAuthStatusResult | undefined; await waitForFast(async () => { @@ -1095,8 +1109,9 @@ describe("models.authStatus", () => { expect(mocks.loadProviderUsageSummary).toHaveBeenCalledWith({ providers: ["deepseek"], agentDir: "/tmp/agent", + authStore: preparedAuthStore, config: expect.any(Object), - timeoutMs: 3500, + timeoutMs: 5_000, }); let result: ModelAuthStatusResult | undefined; await waitForFast(async () => { @@ -1234,8 +1249,9 @@ describe("models.authStatus", () => { expect(mocks.loadProviderUsageSummary).toHaveBeenLastCalledWith({ providers: ["openai"], agentDir: "/tmp/rebound-agent", + authStore: preparedAuthStore, config: expect.any(Object), - timeoutMs: 3500, + timeoutMs: 5_000, }); }); @@ -1288,7 +1304,7 @@ describe("models.authStatus", () => { }); it("does not reuse usage after a direct provider key rotates", async () => { - const cfg = { + let cfg = { models: { providers: { deepseek: { apiKey: "first-direct-value" } } }, }; mocks.getRuntimeConfig.mockReturnValue(cfg); @@ -1316,7 +1332,10 @@ describe("models.authStatus", () => { expect(warmed.providers[0]?.usage?.summary).toBe("Balance 10"); }); - cfg.models.providers.deepseek.apiKey = "second-direct-value"; + cfg = { + models: { providers: { deepseek: { apiKey: "second-direct-value" } } }, + }; + mocks.getRuntimeConfig.mockReturnValue(cfg); const rotated = await readAuthStatus(); expect(rotated.providers[0]?.usage).toBeUndefined(); expect(mocks.loadProviderUsageSummary).toHaveBeenCalledTimes(2); diff --git a/src/gateway/server-methods/models-auth-status.ts b/src/gateway/server-methods/models-auth-status.ts index 785584996a55..4485152fa85e 100644 --- a/src/gateway/server-methods/models-auth-status.ts +++ b/src/gateway/server-methods/models-auth-status.ts @@ -59,7 +59,6 @@ import { modelAuthAgentScopeError, resolveModelAuthAgentScope } from "./model-au import { resolveModelProviderCapabilities } from "./model-provider-capabilities.js"; import { clearModelAuthStatusUsageCache, - fingerprintProviderUsageCredentials, type ProviderUsageStatus, readProviderUsageStaleWhileRevalidate, } from "./models-auth-status-usage-cache.js"; @@ -70,6 +69,7 @@ import type { ModelAuthStatusResult, ModelProviderCapability, } from "./models-auth-status.types.js"; +import { getProviderUsageRuntimeSnapshot } from "./provider-usage-runtime.js"; import type { GatewayRequestContext, GatewayRequestHandlers } from "./types.js"; export type { @@ -688,15 +688,18 @@ export const modelsAuthStatusHandlers: GatewayRequestHandlers = { ), ]; + const providerUsageRuntime = getProviderUsageRuntimeSnapshot({ + config: cfg, + agentId, + agentDir, + store, + }); const usageByProvider = readProviderUsageStaleWhileRevalidate({ agentId, agentDir, + authStore: providerUsageRuntime.store, configRef: cfg, - credentialKey: fingerprintProviderUsageCredentials({ - cfg, - directApiKeys: apiKeys, - store, - }), + credentialKey: providerUsageRuntime.credentialKey, forceRefresh: refreshRequested, providerIds: usageProviderIds, now, diff --git a/src/gateway/server-methods/provider-usage-runtime.ts b/src/gateway/server-methods/provider-usage-runtime.ts new file mode 100644 index 000000000000..b6691c2d8683 --- /dev/null +++ b/src/gateway/server-methods/provider-usage-runtime.ts @@ -0,0 +1,148 @@ +// Prepared provider-usage discovery and credential ownership for Gateway status RPCs. +import { + ensureAuthProfileStore, + externalCliDiscoveryForConfigStatus, + getRuntimeAuthProfileStoreSnapshotRevision, + type AuthProfileStore, +} from "../../agents/auth-profiles.js"; +import { + fingerprintAuthProfileCredential, + fingerprintAuthProfileOwnerShape, + fingerprintResolvedProviderAuth, +} from "../../agents/execution-auth-binding.js"; +import { + resolveLegacyInheritedAuthAgentId, + resolveLegacyInheritedAuthDir, +} from "../../agents/legacy-inherited-auth-dir.js"; +import { resolveEnvApiKey } from "../../agents/model-auth-env.js"; +import { resolveUsableCustomProviderApiKey } from "../../agents/model-auth.js"; +import type { OpenClawConfig } from "../../config/types.openclaw.js"; +import type { UsageProviderId } from "../../infra/provider-usage.types.js"; +import { + listProviderUsagePluginDescriptors, + type ProviderUsagePluginDescriptor, +} from "../../plugins/provider-runtime.js"; +import { getActivePluginRegistryVersion } from "../../plugins/runtime.js"; + +type ResolvedDirectApiKey = { apiKey: string; source: string }; + +type ProviderUsageRuntimeSnapshot = { + agentDir: string; + agentId: string; + configRef: OpenClawConfig; + credentialKey: string; + descriptors: ProviderUsagePluginDescriptor[]; + directApiKeys: ReadonlyMap; + providerIds: UsageProviderId[]; + store: AuthProfileStore; +}; + +type ProviderUsageRuntimeGeneration = ProviderUsageRuntimeSnapshot & { + authStoreGeneration: number; + pluginRegistryGeneration: number; +}; + +let current: ProviderUsageRuntimeGeneration | undefined; + +function sortedRecordEntries(value: Record | undefined) { + return Object.entries(value ?? {}).toSorted(([left], [right]) => left.localeCompare(right)); +} + +function fingerprintProviderUsageCredentials(params: { + cfg: OpenClawConfig; + directApiKeys: ReadonlyMap; + store: AuthProfileStore; +}): string { + const profiles = Object.entries(params.store.profiles) + .toSorted(([left], [right]) => left.localeCompare(right)) + .map(([profileId, credential]) => { + const fingerprint = + fingerprintAuthProfileCredential({ profileId, credential }) ?? + fingerprintAuthProfileOwnerShape({ profileId, credential }); + return fingerprint ?? `${profileId}:${credential.type}:${credential.provider}`; + }); + const direct = [...params.directApiKeys] + .toSorted(([left], [right]) => left.localeCompare(right)) + .map(([provider, resolved]) => [ + provider, + fingerprintResolvedProviderAuth({ ...resolved, mode: "api-key" }) ?? null, + ]); + // Profile selection can switch accounts without changing the profile set. + // Include every non-secret selector that resolveAuthProfileOrder consults. + return JSON.stringify({ + profiles, + direct, + order: sortedRecordEntries(params.store.order), + lastGood: sortedRecordEntries(params.store.lastGood), + usageStats: sortedRecordEntries(params.store.usageStats), + }); +} + +function resolveDirectApiKeys( + config: OpenClawConfig, + providerIds: readonly UsageProviderId[], +): Map { + const directApiKeys = new Map(); + for (const provider of providerIds) { + const resolved = + resolveUsableCustomProviderApiKey({ cfg: config, provider, env: process.env }) ?? + resolveEnvApiKey(provider, process.env, { config }); + if (!resolved) { + continue; + } + directApiKeys.set(provider, resolved); + } + return directApiKeys; +} + +export function clearProviderUsageRuntimeSnapshot(): void { + current = undefined; +} + +export function getProviderUsageRuntimeSnapshot(params: { + config: OpenClawConfig; + agentDir?: string; + agentId?: string; + store?: AuthProfileStore; +}): ProviderUsageRuntimeSnapshot { + const agentId = params.agentId ?? resolveLegacyInheritedAuthAgentId(params.config); + const agentDir = params.agentDir ?? resolveLegacyInheritedAuthDir(params.config); + // Config publication replaces the object, so identity is the exact mutation signal. + const configRef = params.config; + // Registry publication owns descriptor lifetime; request paths only compare its O(1) counter. + const pluginRegistryGeneration = getActivePluginRegistryVersion(); + // Auth writers and runtime overlay publishers advance this O(1) process generation. + const authStoreGeneration = getRuntimeAuthProfileStoreSnapshotRevision(agentDir); + if ( + current?.configRef === configRef && + current.agentDir === agentDir && + current.agentId === agentId && + current.pluginRegistryGeneration === pluginRegistryGeneration && + current.authStoreGeneration === authStoreGeneration + ) { + return current; + } + + const store = + params.store ?? + ensureAuthProfileStore(agentDir, { + externalCli: externalCliDiscoveryForConfigStatus({ cfg: configRef }), + }); + const descriptors = listProviderUsagePluginDescriptors({ config: configRef, env: process.env }); + const providerIds = descriptors.map((descriptor) => descriptor.provider); + const directApiKeys = resolveDirectApiKeys(configRef, providerIds); + current = { + agentDir, + agentId, + configRef, + credentialKey: fingerprintProviderUsageCredentials({ cfg: configRef, directApiKeys, store }), + descriptors, + directApiKeys, + providerIds, + store, + // Building can publish an external-auth overlay, so bind the finished snapshot to its result. + authStoreGeneration: getRuntimeAuthProfileStoreSnapshotRevision(agentDir), + pluginRegistryGeneration, + }; + return current; +} diff --git a/src/gateway/server-methods/usage.status-cache.test.ts b/src/gateway/server-methods/usage.status-cache.test.ts index ff8c1f7c054d..f8c8acdf5aea 100644 --- a/src/gateway/server-methods/usage.status-cache.test.ts +++ b/src/gateway/server-methods/usage.status-cache.test.ts @@ -1,14 +1,24 @@ import { expectDefined } from "@openclaw/normalization-core"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { useAutoCleanupTempDirTracker } from "../../../test/helpers/temp-dir.js"; import { resolveAgentDir, resolveDefaultAgentId } from "../../agents/agent-scope.js"; -import type { AuthProfileStore } from "../../agents/auth-profiles.js"; +import { + clearRuntimeAuthProfileStoreSnapshots, + replaceRuntimeAuthProfileStoreSnapshots, + saveAuthProfileStore, +} from "../../agents/auth-profiles.js"; import type { OpenClawConfig } from "../../config/types.openclaw.js"; +import type { UsageSummary } from "../../infra/provider-usage.types.js"; +import { createEmptyPluginRegistry } from "../../plugins/registry-empty.js"; +import { resetPluginRuntimeStateForTest, setActivePluginRegistry } from "../../plugins/runtime.js"; +import { closeOpenClawAgentDatabasesForTest } from "../../state/openclaw-agent-db.js"; const mocks = vi.hoisted(() => ({ ensureAuthProfileStore: vi.fn(), listProviderUsagePluginDescriptors: vi.fn(), loadProviderUsageSummary: vi.fn(), })); +const tempDirs = useAutoCleanupTempDirTracker(afterEach); vi.mock("../../agents/auth-profiles.js", async () => { const actual = await vi.importActual( @@ -37,9 +47,9 @@ vi.mock("../../infra/provider-usage.load.js", () => ({ import { clearModelAuthStatusUsageCache, - fingerprintProviderUsageCredentials, readProviderUsageStaleWhileRevalidate, } from "./models-auth-status-usage-cache.js"; +import { getProviderUsageRuntimeSnapshot } from "./provider-usage-runtime.js"; import { usageHandlers } from "./usage.js"; const config = { @@ -61,7 +71,7 @@ function createStore(access = "access-one") { }; } -async function runUsageStatus() { +async function runUsageStatus(runtimeConfig = config) { const respond = vi.fn(); await expectDefined( usageHandlers["usage.status"], @@ -69,7 +79,7 @@ async function runUsageStatus() { )({ respond, params: {}, - context: { getRuntimeConfig: () => config }, + context: { getRuntimeConfig: () => runtimeConfig }, } as unknown as Parameters<(typeof usageHandlers)["usage.status"]>[0]); expect(respond).toHaveBeenCalledTimes(1); expect(respond.mock.calls[0]?.[0]).toBe(true); @@ -110,6 +120,8 @@ describe("usage.status provider usage cache", () => { }); afterEach(() => { + clearRuntimeAuthProfileStoreSnapshots(); + resetPluginRuntimeStateForTest(); vi.unstubAllEnvs(); vi.restoreAllMocks(); }); @@ -137,11 +149,13 @@ describe("usage.status provider usage cache", () => { }); it("reuses byte-identical results within 60s and refreshes stale data in the background", async () => { - const first = await runUsageStatus(); + const first = (await runUsageStatus()) as UsageSummary; const repeated = await runUsageStatus(); expect(JSON.stringify(repeated)).toBe(JSON.stringify(first)); expect(mocks.loadProviderUsageSummary).toHaveBeenCalledTimes(1); + expect(mocks.ensureAuthProfileStore).toHaveBeenCalledTimes(1); + expect(mocks.listProviderUsagePluginDescriptors).toHaveBeenCalledTimes(1); now = 61_000; const stale = await runUsageStatus(); @@ -155,6 +169,69 @@ describe("usage.status provider usage cache", () => { expect(refreshed.providers[0]?.windows[0]?.usedPercent).toBe(20); }); expect(mocks.loadProviderUsageSummary).toHaveBeenCalledTimes(2); + expect(mocks.listProviderUsagePluginDescriptors).toHaveBeenCalledTimes(1); + }); + + it("rebuilds prepared usage facts once for each config and plugin generation", async () => { + await runUsageStatus(); + await runUsageStatus(); + + const nextConfig = { ...config }; + await runUsageStatus(nextConfig); + await runUsageStatus(nextConfig); + + setActivePluginRegistry(createEmptyPluginRegistry()); + await runUsageStatus(nextConfig); + await runUsageStatus(nextConfig); + + expect(mocks.listProviderUsagePluginDescriptors).toHaveBeenCalledTimes(3); + expect(mocks.ensureAuthProfileStore).toHaveBeenCalledTimes(3); + }); + + it("rebuilds prepared usage facts once after an auth-store write", async () => { + const writtenAgentDir = tempDirs.make("openclaw-usage-auth-"); + try { + replaceRuntimeAuthProfileStoreSnapshots([{ agentDir: writtenAgentDir, store }]); + + await runUsageStatus(); + await runUsageStatus(); + + store = createStore("access-two"); + saveAuthProfileStore(store, writtenAgentDir); + await runUsageStatus(); + await runUsageStatus(); + + expect(mocks.listProviderUsagePluginDescriptors).toHaveBeenCalledTimes(2); + expect(mocks.ensureAuthProfileStore).toHaveBeenCalledTimes(2); + } finally { + closeOpenClawAgentDatabasesForTest(); + } + }); + + it("keeps a provider's last-good snapshot when its refresh times out", async () => { + const first = (await runUsageStatus()) as UsageSummary; + now = 61_000; + mocks.loadProviderUsageSummary.mockResolvedValueOnce({ + updatedAt: now, + providers: [ + { + provider: "openai", + displayName: "OpenAI", + windows: [], + error: "Timeout", + }, + ], + }); + + const stale = await runUsageStatus(); + expect(JSON.stringify(stale)).toBe(JSON.stringify(first)); + await mocks.loadProviderUsageSummary.mock.results[1]?.value; + await vi.waitFor(async () => { + const retained = (await runUsageStatus()) as UsageSummary; + expect(retained.providers).toEqual(first.providers); + expect(retained.updatedAt).toBe(61_000); + expect(mocks.loadProviderUsageSummary).toHaveBeenCalledTimes(2); + }); }); it("shares the raw snapshot with models.authStatus and invalidates on credential rotation", async () => { @@ -165,18 +242,15 @@ describe("usage.status provider usage cache", () => { agentId, agentDir, configRef: config, - credentialKey: fingerprintProviderUsageCredentials({ - cfg: config, - directApiKeys: new Map(), - store: store as AuthProfileStore, - }), + credentialKey: getProviderUsageRuntimeSnapshot({ config }).credentialKey, providerIds: ["openai"], now, }); expect(usage.get("openai")?.windows[0]?.usedPercent).toBe(10); expect(mocks.loadProviderUsageSummary).toHaveBeenCalledTimes(1); - store.profiles["openai:default"].access = "access-two"; + store = createStore("access-two"); + replaceRuntimeAuthProfileStoreSnapshots([{ agentDir, store }]); const rotated = (await runUsageStatus()) as { providers: Array<{ windows: Array<{ usedPercent: number }> }>; }; diff --git a/src/infra/provider-usage.auth.ts b/src/infra/provider-usage.auth.ts index 0940924371d7..ada74cd1d128 100644 --- a/src/infra/provider-usage.auth.ts +++ b/src/infra/provider-usage.auth.ts @@ -46,13 +46,16 @@ type UsageAuthState = { env: NodeJS.ProcessEnv; agentDir?: string; allowAuthProfileStore: boolean; + getStore?: () => AuthStore; store?: AuthStore; }; function resolveUsageAuthStore(state: UsageAuthState): AuthStore { - state.store ??= ensureAuthProfileStore(state.agentDir, { - allowKeychainPrompt: false, - }); + state.store ??= + state.getStore?.() ?? + ensureAuthProfileStore(state.agentDir, { + allowKeychainPrompt: false, + }); return state.store; } @@ -477,6 +480,8 @@ function hasAuthProfileCredentialSource(params: { export async function resolveProviderAuths(params: { providers: UsageProviderId[]; auth?: ProviderAuth[]; + getStore?: () => AuthStore; + store?: AuthStore; agentDir?: string; config?: OpenClawConfig; env?: NodeJS.ProcessEnv; @@ -495,6 +500,8 @@ export async function resolveProviderAuths(params: { const authProfileSourceState: UsageAuthState = { ...stateBase, allowAuthProfileStore: true, + getStore: params.getStore, + store: params.store, }; const hasAuthProfileStoreSource = params.skipPluginAuthWithoutCredentialSource ? hasAnyAuthProfileStoreSource(params.agentDir) diff --git a/src/infra/provider-usage.load.test.ts b/src/infra/provider-usage.load.test.ts index 2cf7e8052646..e852b4d29821 100644 --- a/src/infra/provider-usage.load.test.ts +++ b/src/infra/provider-usage.load.test.ts @@ -216,6 +216,54 @@ describe("provider-usage.load", () => { ]); }); + it("returns live siblings when one provider never resolves before the deadline", async () => { + vi.useFakeTimers(); + try { + resolveProviderUsageSnapshotWithPluginMock.mockImplementation(async ({ provider }) => { + if (provider === "anthropic") { + return await new Promise(() => {}); + } + return { + provider, + displayName: "Codex", + windows: [{ label: "3h", usedPercent: 12 }], + }; + }); + const summaryPromise = loadProviderUsageSummary({ + auth: [ + { provider: "anthropic", token: "token-a" }, + { provider: "openai", token: "token-codex" }, + ], + config: {}, + env: {}, + timeoutMs: 5_000, + }); + let settled = false; + void summaryPromise.then(() => { + settled = true; + }); + + await vi.advanceTimersByTimeAsync(5_000); + const settledAtDeadline = settled; + if (!settledAtDeadline) { + await vi.advanceTimersByTimeAsync(1_000); + } + const summary = await summaryPromise; + + expect(settledAtDeadline).toBe(true); + expect(summary.providers).toEqual([ + { provider: "anthropic", displayName: "Claude", windows: [], error: "Timeout" }, + { + provider: "openai", + displayName: "Codex", + windows: [{ label: "3h", usedPercent: 12 }], + }, + ]); + } finally { + vi.useRealTimers(); + } + }); + it("keeps successful provider usage when a sibling auth hook rejects", async () => { resolveProviderUsageAuthWithPluginMock.mockImplementation(async ({ provider }) => { if (provider === "anthropic") { diff --git a/src/infra/provider-usage.load.ts b/src/infra/provider-usage.load.ts index 979134d78783..fedd6dcf07d0 100644 --- a/src/infra/provider-usage.load.ts +++ b/src/infra/provider-usage.load.ts @@ -1,10 +1,12 @@ // Loads provider usage snapshots from built-in and plugin providers. +import { ensureAuthProfileStore, type AuthProfileStore } from "../agents/auth-profiles.js"; import { getRuntimeConfig, type OpenClawConfig } from "../config/config.js"; import { listProviderUsagePluginDescriptors, resolveProviderUsageSnapshotWithPlugin, type ProviderUsagePluginDescriptor, } from "../plugins/provider-runtime.js"; +import { formatErrorMessage } from "./errors.js"; import { resolveFetch } from "./fetch.js"; import { resolveProxyFetchFromEnv } from "./net/proxy-fetch.js"; import { type ProviderAuth, resolveProviderAuths } from "./provider-usage.auth.js"; @@ -41,6 +43,7 @@ type UsageSummaryOptions = { timeoutMs?: number; providers?: UsageProviderId[]; auth?: ProviderAuth[]; + authStore?: AuthProfileStore; agentDir?: string; workspaceDir?: string; config?: OpenClawConfig; @@ -129,47 +132,61 @@ export async function loadProviderUsageSummary( windows: [], error, }); - const authFailures: ProviderUsageSnapshot[] = []; - const auths = await resolveProviderAuths({ - providers: descriptors.map((descriptor) => descriptor.provider), - auth: opts.auth, - agentDir: opts.agentDir, - config, - env, - skipPluginAuthWithoutCredentialSource: opts.skipPluginAuthWithoutCredentialSource, - onError: (provider, error) => { - const message = error instanceof Error ? error.message : String(error); - authFailures.push(failureSnapshot(provider, message.trim() || "Auth failed")); - }, - }); - if (auths.length === 0 && authFailures.length === 0) { - return { updatedAt: now, providers: [] }; - } - - const tasks = auths.map((auth) => { + let authStore = opts.authStore; + const getAuthStore = () => + (authStore ??= ensureAuthProfileStore(opts.agentDir, { allowKeychainPrompt: false })); + const tasks = descriptors.map(({ provider }) => { return raceUsageTimeout( - fetchProviderUsageSnapshot({ - auth, - config, - env, - agentDir: opts.agentDir, - workspaceDir: opts.workspaceDir, - timeoutMs, - fetchFn, - }), - timeoutMs + 1000, - failureSnapshot(auth.provider, "Timeout"), + (async () => { + let authError: unknown; + const auth = + opts.auth?.find((candidate) => candidate.provider === provider) ?? + ( + await resolveProviderAuths({ + providers: [provider], + agentDir: opts.agentDir, + config, + env, + getStore: getAuthStore, + store: opts.authStore, + skipPluginAuthWithoutCredentialSource: opts.skipPluginAuthWithoutCredentialSource, + onError: (_provider, error) => { + authError = error; + }, + }) + )[0]; + if (authError) { + const message = formatErrorMessage(authError); + return failureSnapshot(provider, message.trim() || "Auth failed"); + } + if (!auth) { + return undefined; + } + return await fetchProviderUsageSnapshot({ + auth, + config, + env, + agentDir: opts.agentDir, + workspaceDir: opts.workspaceDir, + timeoutMs, + fetchFn, + }); + })(), + timeoutMs, + failureSnapshot(provider, "Timeout"), ).catch((error: unknown) => { const message = error instanceof Error ? error.message : String(error); - return failureSnapshot(auth.provider, message.trim() || "Fetch failed"); + return failureSnapshot(provider, message.trim() || "Fetch failed"); }); }); - const snapshots = [...(await Promise.all(tasks)), ...authFailures].toSorted( - (left, right) => - (providerOrder.get(left.provider) ?? Number.MAX_SAFE_INTEGER) - - (providerOrder.get(right.provider) ?? Number.MAX_SAFE_INTEGER), - ); + const snapshots = (await Promise.all(tasks)) + .filter((snapshot): snapshot is ProviderUsageSnapshot => snapshot !== undefined) + .toSorted( + (left, right) => + (providerOrder.get(left.provider) ?? Number.MAX_SAFE_INTEGER) - + (providerOrder.get(right.provider) ?? Number.MAX_SAFE_INTEGER), + ); const providers = snapshots.filter((entry) => { if (entry.windows.length > 0) { return true; diff --git a/src/infra/provider-usage.shared.ts b/src/infra/provider-usage.shared.ts index 34544b8341b2..1eed852ceff9 100644 --- a/src/infra/provider-usage.shared.ts +++ b/src/infra/provider-usage.shared.ts @@ -3,7 +3,7 @@ import { normalizeProviderId } from "@openclaw/model-catalog-core/provider-id"; import { resolveTimerTimeoutMs } from "@openclaw/normalization-core/number-coercion"; import type { UsageProviderId } from "./provider-usage.types.js"; -/** Default timeout for provider usage collection. */ +/** One provider cannot hold the aggregate usage response beyond this deadline. */ export const PROVIDER_USAGE_TIMEOUT_MS = 5000; export const PROVIDER_LABELS = {