From 0dbfa1f6be128a998f383730eec8f6487be080db Mon Sep 17 00:00:00 2001 From: Peter Lindsey Date: Wed, 3 Jun 2026 17:10:57 +0800 Subject: [PATCH] fix(hooks): make usageState limits credential-aware + document best-effort delivery Addresses review feedback on #89629. 1) Provider-limit resolution no longer defaults to OAuth. getProviderUsageLimits and getProviderUsageLimitsCached resolved with `credentialType ?? "oauth"`, and the agent-runner snapshot call passed no credential type, so an api-key OpenAI turn could borrow cached OAuth/ChatGPT usage windows. Drop the "oauth" default (missing credential type => no OpenAI usage provider) and thread the turn's authMode through at the call site. Adds provider-usage.limits.test.ts covering api-key/no-credential (no fetch), oauth/token (resolves), and non-OpenAI. 2) usageState is documented as best-effort, present only on live dispatcher delivery. Routed durable and recovered queue replays re-run this hook as a stateless transform over the original payload (see QueuedDeliveryPayload); a point-in-time usage snapshot is not stateless and would replay stale after a restart, so it is intentionally omitted there. Consumers must treat the field as optional. Co-Authored-By: Claude Opus 4.8 --- src/auto-reply/reply/agent-runner.ts | 8 ++- src/infra/provider-usage.limits.test.ts | 74 +++++++++++++++++++++++++ src/infra/provider-usage.limits.ts | 12 +++- src/plugins/hook-types.ts | 9 +++ 4 files changed, 99 insertions(+), 4 deletions(-) create mode 100644 src/infra/provider-usage.limits.test.ts diff --git a/src/auto-reply/reply/agent-runner.ts b/src/auto-reply/reply/agent-runner.ts index c4951ded7e69..f296cbd99b7d 100644 --- a/src/auto-reply/reply/agent-runner.ts +++ b/src/auto-reply/reply/agent-runner.ts @@ -1802,8 +1802,12 @@ export async function runReplyAgent(params: { // Provider subscription/limit windows for the 📊 readout. Non-blocking // (stale-while-revalidate): returns cached windows or undefined on a // cold cache and refreshes in the background, so it never delays the - // reply. Undefined for api-key / unmapped providers. - limits: getProviderUsageLimitsCached(providerUsed), + // reply. Credential-aware: the turn's authMode is threaded through so an + // api-key OpenAI turn resolves to no provider (undefined limits) instead + // of borrowing OAuth/ChatGPT windows. Undefined for api-key / unmapped. + limits: getProviderUsageLimitsCached(providerUsed, { + credentialType: runResult.meta?.requestShaping?.authMode ?? undefined, + }), }, ); } diff --git a/src/infra/provider-usage.limits.test.ts b/src/infra/provider-usage.limits.test.ts new file mode 100644 index 000000000000..af48eee95967 --- /dev/null +++ b/src/infra/provider-usage.limits.test.ts @@ -0,0 +1,74 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +// Mock the network-backed loader so we can assert WHETHER a usage fetch is even +// attempted for a given credential type (the regression is about defaulting to +// OAuth and borrowing ChatGPT windows for an api-key turn). +const loadMock = vi.hoisted(() => vi.fn()); +vi.mock("./provider-usage.load.js", () => ({ + loadProviderUsageSummary: loadMock, +})); + +const { getProviderUsageLimits, clearProviderUsageLimitsCacheForTest } = await import( + "./provider-usage.limits.js" +); + +afterEach(() => { + loadMock.mockReset(); + clearProviderUsageLimitsCacheForTest(); +}); + +describe("getProviderUsageLimits credential awareness", () => { + it("returns undefined for an api-key OpenAI turn and never fetches", async () => { + const out = await getProviderUsageLimits("openai", { credentialType: "api-key" }); + expect(out).toBeUndefined(); + expect(loadMock).not.toHaveBeenCalled(); + }); + + it("returns undefined for OpenAI with no credential type (no implicit oauth default)", async () => { + const out = await getProviderUsageLimits("openai"); + expect(out).toBeUndefined(); + expect(loadMock).not.toHaveBeenCalled(); + }); + + it("resolves OpenAI limits for an oauth turn", async () => { + loadMock.mockResolvedValue({ + updatedAt: 0, + providers: [ + { + provider: "openai", + displayName: "OpenAI", + windows: [{ label: "5h", usedPercent: 40 }], + }, + ], + }); + const out = await getProviderUsageLimits("openai", { credentialType: "oauth" }); + expect(loadMock).toHaveBeenCalledTimes(1); + expect(out?.available).toBe(true); + expect(out?.windows[0]).toMatchObject({ label: "5h", used_pct: 40, pct_left: 60 }); + }); + + it("resolves OpenAI limits for a token turn", async () => { + loadMock.mockResolvedValue({ + updatedAt: 0, + providers: [{ provider: "openai", displayName: "OpenAI", windows: [] }], + }); + await getProviderUsageLimits("openai", { credentialType: "token" }); + expect(loadMock).toHaveBeenCalledTimes(1); + }); + + it("resolves non-OpenAI providers regardless of credential type", async () => { + loadMock.mockResolvedValue({ + updatedAt: 0, + providers: [ + { + provider: "anthropic", + displayName: "Anthropic", + windows: [{ label: "week", usedPercent: 10 }], + }, + ], + }); + const out = await getProviderUsageLimits("anthropic"); + expect(loadMock).toHaveBeenCalledTimes(1); + expect(out?.available).toBe(true); + }); +}); diff --git a/src/infra/provider-usage.limits.ts b/src/infra/provider-usage.limits.ts index 58b2ad7177b9..2881ac194e58 100644 --- a/src/infra/provider-usage.limits.ts +++ b/src/infra/provider-usage.limits.ts @@ -26,8 +26,12 @@ export async function getProviderUsageLimits( provider: string | undefined | null, options?: { credentialType?: string | null; timeoutMs?: number; now?: number }, ): Promise { + // Pass the turn's real credential type through unchanged. Do NOT default to + // "oauth": resolveUsageProviderId only returns the OpenAI usage provider for + // oauth/token, so defaulting would attach OAuth/ChatGPT windows to an API-key + // turn. A missing credential type ⇒ no OpenAI limits (correct, not OAuth). const usageId = resolveUsageProviderId(provider, { - credentialType: options?.credentialType ?? "oauth", + credentialType: options?.credentialType ?? null, }); if (!usageId) { return undefined; @@ -103,8 +107,12 @@ export function getProviderUsageLimitsCached( provider: string | undefined | null, options?: { credentialType?: string | null; timeoutMs?: number }, ): ReplyUsageLimits | undefined { + // Pass the turn's real credential type through unchanged. Do NOT default to + // "oauth": resolveUsageProviderId only returns the OpenAI usage provider for + // oauth/token, so defaulting would attach OAuth/ChatGPT windows to an API-key + // turn. A missing credential type ⇒ no OpenAI limits (correct, not OAuth). const usageId = resolveUsageProviderId(provider, { - credentialType: options?.credentialType ?? "oauth", + credentialType: options?.credentialType ?? null, }); if (!usageId) { return undefined; diff --git a/src/plugins/hook-types.ts b/src/plugins/hook-types.ts index eae98fc76ea6..9c20604ba924 100644 --- a/src/plugins/hook-types.ts +++ b/src/plugins/hook-types.ts @@ -560,6 +560,15 @@ export type PluginHookReplyPayloadSendingEvent = { channel?: string; sessionKey?: string; runId?: string; + /** + * Per-turn usage snapshot — **best-effort, present only on live dispatcher + * delivery.** It is intentionally absent on routed durable deliveries and on + * recovered queue replays: those re-run this hook as a stateless transform over + * the original payload (see `QueuedDeliveryPayload`), and a point-in-time usage + * snapshot is not stateless — replaying it after a restart would surface stale + * numbers. Consumers (e.g. a usage footer) must treat this as optional and + * degrade gracefully when it is undefined. + */ usageState?: PluginHookReplyUsageState; };