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 <noreply@anthropic.com>
This commit is contained in:
Peter Lindsey
2026-06-03 17:10:57 +08:00
committed by Ayaan Zaidi
parent f06f2f17c2
commit 0dbfa1f6be
4 changed files with 99 additions and 4 deletions
+6 -2
View File
@@ -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,
}),
},
);
}
+74
View File
@@ -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);
});
});
+10 -2
View File
@@ -26,8 +26,12 @@ export async function getProviderUsageLimits(
provider: string | undefined | null,
options?: { credentialType?: string | null; timeoutMs?: number; now?: number },
): Promise<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;
@@ -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;
+9
View File
@@ -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;
};