fix(auth): skipped local usage probes in health checks
- Added a UsageProvider flag for providers whose reports do not validate upstream credentials.\n- Kept OpenCode Go quota reporting local-only while leaving credential health unknown unless a strict completion probe runs.\n- Covered the OpenCode Go health-check regression.\n\nFixes #2942
This commit is contained in:
@@ -2729,6 +2729,8 @@ export class AuthStorage {
|
||||
base.reason = `no usage probe configured for provider ${row.provider}`;
|
||||
} else if (providerImpl.supports && !providerImpl.supports(initialRequest)) {
|
||||
base.reason = `usage probe does not support ${cred.type} credentials for ${row.provider}`;
|
||||
} else if (providerImpl.validatesCredentials === false) {
|
||||
base.reason = `usage probe for ${row.provider} does not validate credentials`;
|
||||
} else {
|
||||
try {
|
||||
const report = await providerImpl.fetchUsage(params, ctx);
|
||||
|
||||
@@ -257,6 +257,8 @@ export interface UsageProvider {
|
||||
/** Parse provider rate-limit response headers (lowercased keys) into a usage report, if supported. */
|
||||
parseRateLimitHeaders?(headers: Record<string, string>, now?: number): UsageReport | null;
|
||||
supports?(params: UsageFetchParams): boolean;
|
||||
/** True when fetchUsage contacts upstream and can authenticate the credential for health checks. */
|
||||
validatesCredentials?: boolean;
|
||||
}
|
||||
|
||||
/** Request context used when ranking usage for a specific model. */
|
||||
|
||||
@@ -69,6 +69,7 @@ function buildWindowLimit(
|
||||
export const opencodeGoUsageProvider: UsageProvider = {
|
||||
id: OPENCODE_GO_PROVIDER,
|
||||
supports: params => params.provider === OPENCODE_GO_PROVIDER && params.credential.type === "api_key",
|
||||
validatesCredentials: false,
|
||||
async fetchUsage(params, ctx) {
|
||||
if (params.provider !== OPENCODE_GO_PROVIDER || params.credential.type !== "api_key") return null;
|
||||
const nowMs = Date.now();
|
||||
|
||||
@@ -15,10 +15,12 @@
|
||||
* 5. Providers with no registered `UsageProvider` report `ok: null` with
|
||||
* "no usage probe configured" — the credential's status is unknown,
|
||||
* not failed.
|
||||
* 6. When a `completionProbe` is supplied, it receives the post-refresh
|
||||
* 6. Local-only usage providers opt out of health validation and leave
|
||||
* `ok: null` unless a separate completion probe is supplied.
|
||||
* 7. When a `completionProbe` is supplied, it receives the post-refresh
|
||||
* bearer for every row, runs independently of the usage probe (i.e. it
|
||||
* still runs for providers without a `UsageProvider`), but is skipped
|
||||
* when OAuth refresh fails.
|
||||
* still runs for providers without a validating `UsageProvider`), but is
|
||||
* skipped when OAuth refresh fails.
|
||||
*/
|
||||
import { afterEach, describe, expect, it, vi } from "bun:test";
|
||||
import {
|
||||
@@ -32,6 +34,7 @@ import {
|
||||
} from "@oh-my-pi/pi-ai/auth-storage";
|
||||
import type { UsageProvider } from "@oh-my-pi/pi-ai/usage";
|
||||
import * as claudeUsage from "@oh-my-pi/pi-ai/usage/claude";
|
||||
import { opencodeGoUsageProvider } from "@oh-my-pi/pi-ai/usage/opencode-go";
|
||||
|
||||
function oauthRow(id: number, email: string, opts?: { expired?: boolean }): StoredAuthCredential {
|
||||
const credential: AuthCredential = {
|
||||
@@ -396,6 +399,33 @@ describe("AuthStorage.checkCredentials", () => {
|
||||
}
|
||||
});
|
||||
|
||||
it("does not mark local-only usage providers healthy without upstream validation", async () => {
|
||||
const apiKeyRow: StoredAuthCredential = {
|
||||
id: 12,
|
||||
provider: "opencode-go",
|
||||
credential: { type: "api_key", key: "sk-opencode-go" },
|
||||
disabledCause: null,
|
||||
};
|
||||
const store = makeStore([apiKeyRow]);
|
||||
const storage = new AuthStorage(store, {
|
||||
usageProviderResolver: provider => (provider === "opencode-go" ? opencodeGoUsageProvider : undefined),
|
||||
});
|
||||
await storage.reload();
|
||||
|
||||
const probe = vi.fn<CompletionProbe>().mockResolvedValue({ ok: false, reason: "401 invalid_api_key" });
|
||||
|
||||
try {
|
||||
const [result] = await storage.checkCredentials({ completionProbe: probe });
|
||||
expect(result.ok).toBeNull();
|
||||
expect(result.reason).toMatch(/does not validate credentials/);
|
||||
expect(result.report).toBeUndefined();
|
||||
expect(probe).toHaveBeenCalledTimes(1);
|
||||
expect(result.completion).toEqual({ ok: false, reason: "401 invalid_api_key" });
|
||||
} finally {
|
||||
storage.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("skips the completionProbe when OAuth refresh fails", async () => {
|
||||
const refreshSpy = vi
|
||||
.fn<NonNullable<AuthCredentialStore["refreshOAuthCredential"]>>()
|
||||
|
||||
Reference in New Issue
Block a user