perf(gateway): prepare provider usage runtime (#125520)

* perf(gateway): prepare provider usage runtime

* fix(gateway): keep provider usage internals private
This commit is contained in:
Peter Steinberger
2026-08-17 20:47:29 -07:00
committed by GitHub
parent 7242074ecc
commit ba3599627c
9 changed files with 425 additions and 165 deletions
@@ -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<string, ProviderUsageCacheEntry>();
const usageRefreshByAgentId = new Map<string, ProviderUsageRefresh>();
let cacheGeneration = 0;
function sortedRecordEntries<T>(value: Record<string, T> | undefined) {
return Object.entries(value ?? {}).toSorted(([left], [right]) => left.localeCompare(right));
}
export function fingerprintProviderUsageCredentials(params: {
cfg: OpenClawConfig;
directApiKeys: ReadonlyMap<string, { source: "config" | "env"; envVar?: string } | undefined>;
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<ReturnType<typeof loadProviderUsageSumm
return usageByProvider;
}
function retainLastGoodOnTimeout(
summary: UsageSummary,
lastGood: UsageSummary | undefined,
): UsageSummary {
if (!lastGood) {
return summary;
}
const lastGoodByProvider = new Map(
lastGood.providers
.filter((provider) => 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<UsageSummary> {
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<UsageSummary> {
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(),
});
}
@@ -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<UsageSummary> => 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);
@@ -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,
@@ -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<string, ResolvedDirectApiKey>;
providerIds: UsageProviderId[];
store: AuthProfileStore;
};
type ProviderUsageRuntimeGeneration = ProviderUsageRuntimeSnapshot & {
authStoreGeneration: number;
pluginRegistryGeneration: number;
};
let current: ProviderUsageRuntimeGeneration | undefined;
function sortedRecordEntries<T>(value: Record<string, T> | undefined) {
return Object.entries(value ?? {}).toSorted(([left], [right]) => left.localeCompare(right));
}
function fingerprintProviderUsageCredentials(params: {
cfg: OpenClawConfig;
directApiKeys: ReadonlyMap<string, ResolvedDirectApiKey>;
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<string, ResolvedDirectApiKey> {
const directApiKeys = new Map<string, ResolvedDirectApiKey>();
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;
}
@@ -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<typeof import("../../agents/auth-profiles.js")>(
@@ -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 }> }>;
};
+10 -3
View File
@@ -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)
+48
View File
@@ -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<ProviderUsageSnapshot>(() => {});
}
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") {
+52 -35
View File
@@ -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;
+1 -1
View File
@@ -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 = {