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:
roboomp
2026-06-18 07:38:37 +00:00
parent 2b9047b79c
commit e5b35d38a7
4 changed files with 38 additions and 3 deletions
+2
View File
@@ -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);
+2
View File
@@ -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. */
+1
View File
@@ -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"]>>()