fix(ai): preferred explicit env api key over stored static key

- Reordered `AuthStorage.getApiKey` and its async session variant to resolve stored OAuth, then an env var, then a stored static api_key — so an explicit `GEMINI_API_KEY`-style env var now wins over a stale broker-migrated key.
- Reworked `getCredentialOrigin` and `getApiKeySource` to mirror that precedence via a `describeStored` helper, reporting `oauth`/`env`/`api_key` in the new order.
- Updated `auth-storage-api-key-login` to neutralize ambient env keys with a `getEnvApiKey` spy, and `auth-storage-credential-origin` to assert OAuth outranks api_key and env outranks a stored api_key.
This commit is contained in:
can1357
2026-06-30 06:34:07 +02:00
parent 4e2acd3605
commit b87110c66b
4 changed files with 75 additions and 55 deletions
+1
View File
@@ -14,6 +14,7 @@
- Updated internal `coerceServiceTierByFamily` helper to facilitate migration from legacy settings
### Fixed
- Changed API-key resolution precedence so an explicit environment variable (e.g. `GEMINI_API_KEY`) overrides a stored/broker-migrated static API key; a deliberate OAuth login still takes precedence over the env var.
- Improved Vertex AI reliability by automatically falling back to global endpoints on 404 errors
+56 -52
View File
@@ -1736,8 +1736,8 @@ export class AuthStorage {
/**
* Classify where a provider's auth comes from, following the same precedence
* as {@link AuthStorage.getApiKey}: runtime override → config override →
* stored credential (api_key before oauth, matching getApiKey) → env var →
* fallback resolver. Returns undefined when no auth is configured.
* stored OAuth → env var → stored api_key → fallback resolver. Returns
* undefined when no auth is configured.
*
* Compact, structured counterpart to {@link describeCredentialSource}.
*/
@@ -1745,10 +1745,9 @@ export class AuthStorage {
if (this.#runtimeOverrides.has(provider)) return { kind: "runtime" };
if (this.#configOverrides.has(provider)) return { kind: "config" };
const stored = this.#getCredentialsForProvider(provider);
if (stored.length > 0) {
return { kind: stored.some(credential => credential.type === "api_key") ? "api_key" : "oauth" };
}
if (stored.some(credential => credential.type === "oauth")) return { kind: "oauth" };
if (getEnvApiKey(provider)) return { kind: "env", envVar: getEnvApiKeyName(provider) };
if (stored.some(credential => credential.type === "api_key")) return { kind: "api_key" };
if (this.#fallbackResolver?.(provider)) return { kind: "fallback" };
return undefined;
}
@@ -3760,12 +3759,8 @@ export class AuthStorage {
return configKey;
}
const apiKeySelection = this.#selectCredentialByType(provider, "api_key");
if (apiKeySelection) {
return this.#configValueResolver(apiKeySelection.credential.key);
}
// Return current OAuth access token only if it is not already expired.
// Precedence: a deliberate OAuth login wins, then an explicit env var, then a stored
// static api_key (which may be a stale broker-migrated copy) as a last resort.
const oauthSelection = this.#selectCredentialByType(provider, "oauth");
if (oauthSelection) {
const expiresAt = oauthSelection.credential.expires;
@@ -3784,17 +3779,22 @@ export class AuthStorage {
const envKey = getEnvApiKey(provider);
if (envKey) return envKey;
const apiKeySelection = this.#selectCredentialByType(provider, "api_key");
if (apiKeySelection) {
return this.#configValueResolver(apiKeySelection.credential.key);
}
return this.#fallbackResolver?.(provider) ?? undefined;
}
/**
* Get API key for a provider.
* Priority:
* Priority (first match wins):
* 1. Runtime override (CLI --api-key)
* 2. Config override (models.yml `providers.<name>.apiKey`)
* 3. API key from storage
* 4. OAuth token from storage (auto-refreshed)
* 5. Environment variable
* 3. OAuth token from storage (auto-refreshed)
* 4. Environment variable
* 5. Stored API key (e.g. a broker-migrated copy) — last resort, so an explicit env var wins
* 6. Fallback resolver (models.yml custom providers, last-resort)
*/
async getApiKey(provider: string, sessionId?: string, options?: AuthApiKeyOptions): Promise<string | undefined> {
@@ -3814,25 +3814,27 @@ export class AuthStorage {
return configKey;
}
const apiKeySelection = this.#selectCredentialByType(provider, "api_key", sessionId);
if (apiKeySelection) {
this.#recordSessionCredential(provider, sessionId, "api_key", apiKeySelection.index);
return this.#configValueResolver(apiKeySelection.credential.key);
}
// Precedence: a deliberate OAuth login wins, then an explicit env var, then a stored
// static api_key (which may be a stale broker-migrated copy) as a last resort.
const oauthResolved = await this.#resolveOAuthSelection(provider, sessionId, options);
if (oauthResolved) {
return oauthResolved.apiKey;
}
// Fall back to environment variable or custom resolver. If we reach here after
// an OAuth miss, the session sticky (if any) is stale — the request will
// authenticate via env/fallback, not OAuth, so clear the sticky now so that
// getOAuthAccountId() correctly suppresses account_uuid for this session.
// Past OAuth: the session sticky (if any) is stale — the request authenticates via
// env/api_key/fallback, not OAuth, so clear it now so getOAuthAccountId() correctly
// suppresses account_uuid for this session.
if (sessionId) this.#sessionLastCredential.get(provider)?.delete(sessionId);
const envKey = getEnvApiKey(provider);
if (envKey) return envKey;
const apiKeySelection = this.#selectCredentialByType(provider, "api_key", sessionId);
if (apiKeySelection) {
this.#recordSessionCredential(provider, sessionId, "api_key", apiKeySelection.index);
return this.#configValueResolver(apiKeySelection.credential.key);
}
// Fall back to custom resolver (e.g., models.json custom providers)
return this.#fallbackResolver?.(provider) ?? undefined;
}
@@ -4486,12 +4488,13 @@ export class AuthStorage {
/**
* Describe where the active credential for a provider came from.
*
* Surfaces four layers, highest precedence first:
* Mirrors {@link AuthStorage.getApiKey} precedence, highest first:
* 1. Runtime override (`--api-key`).
* 2. Config override (`models.yml` `providers.<name>.apiKey`).
* 3. Stored credential (the one this session is currently sticky to, or the
* one round-robin would pick next when no session id is supplied).
* 4. Env var / fallback resolver — when no stored credential exists.
* 3. Stored OAuth credential.
* 4. Env var — overrides a stored static api_key (e.g. a stale broker copy).
* 5. Stored api_key credential.
* 6. Fallback resolver.
*
* The string is purely informational; consumers must not parse it.
*/
@@ -4505,30 +4508,31 @@ export class AuthStorage {
const baseLabel = this.#sourceLabel ?? "local store";
const stored = this.#getStoredCredentials(provider);
if (stored.length === 0) {
if (getEnvApiKey(provider)) return `env ${baseLabel ? `(fallback over ${baseLabel})` : ""}`.trim();
if (this.#fallbackResolver?.(provider) !== undefined) return `fallback resolver`;
return undefined;
}
const session = sessionId ? this.#sessionLastCredential.get(provider)?.get(sessionId) : undefined;
// Same selection logic as #selectCredentialByType for "no session" lookups: prefer
// the type with stored credentials, lean OAuth before api_key. We don't run the
// full round-robin here because describing the source shouldn't advance the index.
const preferredType: AuthCredential["type"] =
session?.type ?? (stored.some(entry => entry.credential.type === "oauth") ? "oauth" : "api_key");
const typed = stored
.map((entry, index) => ({ entry, index }))
.filter(({ entry }) => entry.credential.type === preferredType);
if (typed.length === 0) return baseLabel;
const index = session?.index ?? typed[0].index;
const chosen = stored[index] ?? typed[0].entry;
const credential = chosen.credential;
const identity =
credential.type === "oauth"
? (credential.email ?? credential.accountId ?? credential.projectId ?? `cred ${chosen.id}`)
: `cred ${chosen.id}`;
return `${baseLabel} · ${preferredType} #${chosen.id} (${identity})`;
// Describe the stored credential of a given type, honoring the session sticky index.
const describeStored = (type: AuthCredential["type"]): string | undefined => {
const typed = stored
.map((entry, index) => ({ entry, index }))
.filter(({ entry }) => entry.credential.type === type);
if (typed.length === 0) return undefined;
const index = session?.type === type ? session.index : typed[0].index;
const chosen = stored[index] ?? typed[0].entry;
const credential = chosen.credential;
const identity =
credential.type === "oauth"
? (credential.email ?? credential.accountId ?? credential.projectId ?? `cred ${chosen.id}`)
: `cred ${chosen.id}`;
return `${baseLabel} · ${type} #${chosen.id} (${identity})`;
};
// A deliberate OAuth login wins; then an explicit env var; then a stored static api_key.
const oauthSource = describeStored("oauth");
if (oauthSource) return oauthSource;
if (getEnvApiKey(provider)) return `env (over ${baseLabel})`;
const apiKeySource = describeStored("api_key");
if (apiKeySource) return apiKeySource;
if (this.#fallbackResolver?.(provider) !== undefined) return "fallback resolver";
return undefined;
}
}
@@ -8,6 +8,7 @@ import { AuthStorage, SqliteAuthCredentialStore } from "@oh-my-pi/pi-ai/auth-sto
import * as deepseekModule from "@oh-my-pi/pi-ai/registry/deepseek";
import * as kagiModule from "@oh-my-pi/pi-ai/registry/kagi";
import * as ollamaCloudModule from "@oh-my-pi/pi-ai/registry/ollama-cloud";
import * as aiStream from "@oh-my-pi/pi-ai/stream";
import { removeWithRetries } from "../../utils/src/temp";
function countCredentialRows(dbPath: string, provider: string): number {
@@ -38,6 +39,9 @@ function countCredentialRowsByDisabledState(dbPath: string, provider: string, di
}
describe("AuthStorage api-key login upsert", () => {
// A live env var now (correctly) overrides a stored static api_key. These tests verify that a
// freshly stored api_key resolves through AuthStorage.getApiKey, so neutralize the env leg
// entirely — this ignores every provider's ambient env key, not just the few set locally.
let tempDir = "";
let dbPath = "";
let store: SqliteAuthCredentialStore | null = null;
@@ -47,6 +51,7 @@ describe("AuthStorage api-key login upsert", () => {
let loginOllamaCloudSpy: Mock<typeof ollamaCloudModule.loginOllamaCloud>;
beforeEach(async () => {
vi.spyOn(aiStream, "getEnvApiKey").mockReturnValue(undefined);
tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-api-key-login-"));
dbPath = path.join(tempDir, "agent.db");
store = await SqliteAuthCredentialStore.open(dbPath);
@@ -68,14 +68,24 @@ describe("AuthStorage.getCredentialOrigin", () => {
});
});
test("a stored api key reports api_key and outranks a co-stored OAuth credential", async () => {
test("a stored OAuth credential outranks a co-stored api key", async () => {
await withEnv(SUPPRESS_ENV, async () => {
// getApiKey() prefers api_key before oauth, so the origin must match.
// getApiKey() resolves stored OAuth before a stored api_key, so the origin must match.
await auth?.set("openai", [
{ type: "oauth", access: "a", refresh: "r", expires: Date.now() + 60_000 },
{ type: "api_key", key: "sk-stored" },
]);
expect(auth?.getCredentialOrigin("openai")).toEqual({ kind: "api_key" });
expect(auth?.getCredentialOrigin("openai")).toEqual({ kind: "oauth" });
});
});
test("an explicit env var outranks a stored api key", async () => {
// Regression: a live env var is the user's current choice and must win over a stored
// static api_key (e.g. a stale broker-migrated copy) so `GEMINI_API_KEY` etc. take effect.
await withEnv({ ...SUPPRESS_ENV, OPENAI_API_KEY: "sk-env" }, async () => {
await auth?.set("openai", [{ type: "api_key", key: "sk-stored" }]);
expect(auth?.getCredentialOrigin("openai")).toEqual({ kind: "env", envVar: "OPENAI_API_KEY" });
expect(await auth?.getApiKey("openai")).toBe("sk-env");
});
});