getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-17 10:02:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

WIP: Remove global state from pi-ai OAuth/API key handling

Browse source 
- Remove setApiKey, resolveApiKey, and global apiKeys Map from stream.ts
- Rename getApiKey to getApiKeyFromEnv (only checks env vars)
- Remove OAuth storage layer (storage.ts deleted)
- OAuth login/refresh functions now return credentials instead of saving
- getOAuthApiKey/refreshOAuthToken now take credentials as params
- Add test/oauth.ts helper for ai package tests
- Simplify root npm run check (single biome + tsgo pass)
- Remove redundant check scripts from most packages
- Add web-ui and coding-agent examples to biome/tsgo includes

coding-agent still has compile errors - needs refactoring for new API
This commit is contained in:
Mario Zechner 2025-12-25 01:01:03 +01:00
parent d93cbf8c32
commit 030788140a
51 changed files with 646 additions and 570 deletions
Download patch file Download diff file Expand all files Collapse all files

11
packages/ai/CHANGELOG.md
View file

@ -2,9 +2,14 @@
## [Unreleased]
## Breaking Changes
- **setApiKey, resolveApiKey**: removed. You need to create your own api key storage/resolution
- **getApiKey**: renamed to `getApiKeyFromEnv`. Given a provider, checks for the known env variable holding the API key.
### Breaking Changes
- **setApiKey, resolveApiKey**: Removed. Callers must manage their own API key storage/resolution.
- **getApiKey**: Renamed to `getApiKeyFromEnv`. Only checks environment variables for known providers.
- **OAuth storage removed**: All storage functions (`loadOAuthCredentials`, `saveOAuthCredentials`, `setOAuthStorage`, etc.) removed. Callers are responsible for storing credentials.
- **OAuth login functions**: `loginAnthropic`, `loginGitHubCopilot`, `loginGeminiCli`, `loginAntigravity` now return `OAuthCredentials` instead of saving to disk.
- **refreshOAuthToken**: Now takes `(provider, credentials)` and returns new `OAuthCredentials` instead of saving.
- **getOAuthApiKey**: Now takes `(provider, credentials)` and returns `{ newCredentials, apiKey }` or null.
- **OAuthCredentials type**: No longer includes `type: "oauth"` discriminator. Callers add discriminator when storing.
## [0.27.7] - 2025-12-24

1
packages/ai/package.json
View file

@ -15,7 +15,6 @@
"build": "npm run generate-models && tsgo -p tsconfig.build.json",
"dev": "tsgo -p tsconfig.build.json --watch --preserveWatchOutput",
"dev:tsc": "tsgo -p tsconfig.build.json --watch --preserveWatchOutput",
"check": "biome check --write . && tsgo --noEmit",
"test": "vitest --run",
"prepublishOnly": "npm run clean && npm run build"
},

5
packages/ai/src/stream.ts
View file

@ -18,8 +18,9 @@ import type {
} from "./types.js";
/**
* Get API key from environment variables (sync).
* Does NOT check OAuth credentials - use getApiKeyAsync for that.
* Get API key for provider from known environment variables, e.g. OPENAI_API_KEY.
*
* Will not return API keys for providers that require OAuth tokens.
*/
export function getApiKeyFromEnv(provider: KnownProvider): string | undefined;
export function getApiKeyFromEnv(provider: string): string | undefined;

10
packages/ai/src/utils/oauth/anthropic.ts
View file

@ -3,7 +3,7 @@
*/
import { createHash, randomBytes } from "crypto";
import { type OAuthCredentials, saveOAuthCredentials } from "./storage.js";
import type { OAuthCredentials } from "./types.js";
const decode = (s: string) => Buffer.from(s, "base64").toString();
const CLIENT_ID = decode("OWQxYzI1MGEtZTYxYi00NGQ5LTg4ZWQtNTk0NGQxOTYyZjVl");
@ -30,7 +30,7 @@ function generatePKCE(): { verifier: string; challenge: string } {
export async function loginAnthropic(
onAuthUrl: (url: string) => void,
onPromptCode: () => Promise<string>,
): Promise<void> {
): Promise<OAuthCredentials> {
const { verifier, challenge } = generatePKCE();
// Build authorization URL
@ -87,14 +87,11 @@ export async function loginAnthropic(
const expiresAt = Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000;
// Save credentials
const credentials: OAuthCredentials = {
type: "oauth",
return {
refresh: tokenData.refresh_token,
access: tokenData.access_token,
expires: expiresAt,
};
saveOAuthCredentials("anthropic", credentials);
}
/**
@ -123,7 +120,6 @@ export async function refreshAnthropicToken(refreshToken: string): Promise<OAuth
};
return {
type: "oauth",
refresh: data.refresh_token,
access: data.access_token,
expires: Date.now() + data.expires_in * 1000 - 5 * 60 * 1000,

17
packages/ai/src/utils/oauth/github-copilot.ts
View file

@ -3,7 +3,7 @@
*/
import { getModels } from "../../models.js";
import { type OAuthCredentials, saveOAuthCredentials } from "./storage.js";
import type { OAuthCredentials } from "./types.js";
const decode = (s: string) => Buffer.from(s, "base64").toString();
const CLIENT_ID = decode("SXYxLmI1MDdhMDhjODdlY2ZlOTg=");
@ -63,7 +63,7 @@ function getUrls(domain: string): {
* Token format: tid=...;exp=...;proxy-ep=proxy.individual.githubcopilot.com;...
* Returns API URL like https://api.individual.githubcopilot.com
*/
export function getBaseUrlFromToken(token: string): string | null {
function getBaseUrlFromToken(token: string): string | null {
const match = token.match(/proxy-ep=([^;]+)/);
if (!match) return null;
const proxyHost = match[1];
@ -217,7 +217,6 @@ export async function refreshGitHubCopilotToken(
}
return {
type: "oauth",
refresh: refreshToken,
access: token,
expires: expiresAt * 1000 - 5 * 60 * 1000,
@ -229,11 +228,7 @@ export async function refreshGitHubCopilotToken(
* Enable a model for the user's GitHub Copilot account.
* This is required for some models (like Claude, Grok) before they can be used.
*/
export async function enableGitHubCopilotModel(
token: string,
modelId: string,
enterpriseDomain?: string,
): Promise<boolean> {
async function enableGitHubCopilotModel(token: string, modelId: string, enterpriseDomain?: string): Promise<boolean> {
const baseUrl = getGitHubCopilotBaseUrl(token, enterpriseDomain);
const url = `${baseUrl}/models/${modelId}/policy`;
@ -259,7 +254,7 @@ export async function enableGitHubCopilotModel(
* Enable all known GitHub Copilot models that may require policy acceptance.
* Called after successful login to ensure all models are available.
*/
export async function enableAllGitHubCopilotModels(
async function enableAllGitHubCopilotModels(
token: string,
enterpriseDomain?: string,
onProgress?: (model: string, success: boolean) => void,
@ -312,9 +307,5 @@ export async function loginGitHubCopilot(options: {
// Enable all models after successful login
options.onProgress?.("Enabling models...");
await enableAllGitHubCopilotModels(credentials.access, enterpriseDomain ?? undefined);
// Save credentials
saveOAuthCredentials("github-copilot", credentials);
return credentials;
}

15
packages/ai/src/utils/oauth/google-antigravity.ts
View file

@ -5,7 +5,7 @@
import { createHash, randomBytes } from "crypto";
import { createServer, type Server } from "http";
import { type OAuthCredentials, saveOAuthCredentials } from "./storage.js";
import type { OAuthCredentials } from "./types.js";
// Antigravity OAuth credentials (different from Gemini CLI)
const decode = (s: string) => Buffer.from(s, "base64").toString();
@ -30,11 +30,6 @@ const TOKEN_URL = "https://oauth2.googleapis.com/token";
// Fallback project ID when discovery fails
const DEFAULT_PROJECT_ID = "rising-fact-p41fc";
export interface AntigravityCredentials extends OAuthCredentials {
projectId: string;
email?: string;
}
/**
* Generate PKCE code verifier and challenge
*/
@ -220,7 +215,6 @@ export async function refreshAntigravityToken(refreshToken: string, projectId: s
};
return {
type: "oauth",
refresh: data.refresh_token || refreshToken,
access: data.access_token,
expires: Date.now() + data.expires_in * 1000 - 5 * 60 * 1000,
@ -237,7 +231,7 @@ export async function refreshAntigravityToken(refreshToken: string, projectId: s
export async function loginAntigravity(
onAuth: (info: { url: string; instructions?: string }) => void,
onProgress?: (message: string) => void,
): Promise<AntigravityCredentials> {
): Promise<OAuthCredentials> {
const { verifier, challenge } = generatePKCE();
// Start local server for callback
@ -317,8 +311,7 @@ export async function loginAntigravity(
// Calculate expiry time (current time + expires_in seconds - 5 min buffer)
const expiresAt = Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000;
const credentials: AntigravityCredentials = {
type: "oauth",
const credentials: OAuthCredentials = {
refresh: tokenData.refresh_token,
access: tokenData.access_token,
expires: expiresAt,
@ -326,8 +319,6 @@ export async function loginAntigravity(
email,
};
saveOAuthCredentials("google-antigravity", credentials);
return credentials;
} finally {
server.close();

15
packages/ai/src/utils/oauth/google-gemini-cli.ts
View file

@ -5,7 +5,7 @@
import { createHash, randomBytes } from "crypto";
import { createServer, type Server } from "http";
import { type OAuthCredentials, saveOAuthCredentials } from "./storage.js";
import type { OAuthCredentials } from "./types.js";
const decode = (s: string) => Buffer.from(s, "base64").toString();
const CLIENT_ID = decode(
@ -22,11 +22,6 @@ const AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth";
const TOKEN_URL = "https://oauth2.googleapis.com/token";
const CODE_ASSIST_ENDPOINT = "https://cloudcode-pa.googleapis.com";
export interface GoogleCloudCredentials extends OAuthCredentials {
projectId: string;
email?: string;
}
/**
* Generate PKCE code verifier and challenge
*/
@ -251,7 +246,6 @@ export async function refreshGoogleCloudToken(refreshToken: string, projectId: s
};
return {
type: "oauth",
refresh: data.refresh_token || refreshToken,
access: data.access_token,
expires: Date.now() + data.expires_in * 1000 - 5 * 60 * 1000,
@ -268,7 +262,7 @@ export async function refreshGoogleCloudToken(refreshToken: string, projectId: s
export async function loginGeminiCli(
onAuth: (info: { url: string; instructions?: string }) => void,
onProgress?: (message: string) => void,
): Promise<GoogleCloudCredentials> {
): Promise<OAuthCredentials> {
const { verifier, challenge } = generatePKCE();
// Start local server for callback
@ -348,8 +342,7 @@ export async function loginGeminiCli(
// Calculate expiry time (current time + expires_in seconds - 5 min buffer)
const expiresAt = Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000;
const credentials: GoogleCloudCredentials = {
type: "oauth",
const credentials: OAuthCredentials = {
refresh: tokenData.refresh_token,
access: tokenData.access_token,
expires: expiresAt,
@ -357,8 +350,6 @@ export async function loginGeminiCli(
email,
};
saveOAuthCredentials("google-gemini-cli", credentials);
return credentials;
} finally {
server.close();

116
packages/ai/src/utils/oauth/index.ts
View file

@ -13,9 +13,6 @@
export { loginAnthropic, refreshAnthropicToken } from "./anthropic.js";
// GitHub Copilot
export {
enableAllGitHubCopilotModels,
enableGitHubCopilotModel,
getBaseUrlFromToken,
getGitHubCopilotBaseUrl,
loginGitHubCopilot,
normalizeDomain,
@ -23,32 +20,16 @@ export {
} from "./github-copilot.js";
// Google Antigravity
export {
type AntigravityCredentials,
loginAntigravity,
refreshAntigravityToken,
} from "./google-antigravity.js";
// Google Gemini CLI
export {
type GoogleCloudCredentials,
loginGeminiCli,
refreshGoogleCloudToken,
} from "./google-gemini-cli.js";
// Storage
export {
getOAuthPath,
hasOAuthCredentials,
listOAuthProviders,
loadOAuthCredentials,
loadOAuthStorage,
type OAuthCredentials,
type OAuthProvider,
type OAuthStorage,
type OAuthStorageBackend,
removeOAuthCredentials,
resetOAuthStorage,
saveOAuthCredentials,
setOAuthStorage,
} from "./storage.js";
export * from "./types.js";
// ============================================================================
// High-level API
@ -58,15 +39,16 @@ import { refreshAnthropicToken } from "./anthropic.js";
import { refreshGitHubCopilotToken } from "./github-copilot.js";
import { refreshAntigravityToken } from "./google-antigravity.js";
import { refreshGoogleCloudToken } from "./google-gemini-cli.js";
import type { OAuthCredentials, OAuthProvider } from "./storage.js";
import { loadOAuthCredentials, removeOAuthCredentials, saveOAuthCredentials } from "./storage.js";
import type { OAuthCredentials, OAuthProvider, OAuthProviderInfo } from "./types.js";
/**
* Refresh token for any OAuth provider.
* Saves the new credentials and returns the new access token.
*/
export async function refreshToken(provider: OAuthProvider): Promise<string> {
const credentials = loadOAuthCredentials(provider);
export async function refreshOAuthToken(
provider: OAuthProvider,
credentials: OAuthCredentials,
): Promise<OAuthCredentials> {
if (!credentials) {
throw new Error(`No OAuth credentials found for ${provider}`);
}
@ -96,8 +78,7 @@ export async function refreshToken(provider: OAuthProvider): Promise<string> {
throw new Error(`Unknown OAuth provider: ${provider}`);
}
saveOAuthCredentials(provider, newCredentials);
return newCredentials.access;
return newCredentials;
}
/**
@ -107,81 +88,30 @@ export async function refreshToken(provider: OAuthProvider): Promise<string> {
* For google-gemini-cli and antigravity, returns JSON-encoded { token, projectId }
*
* @returns API key string, or null if no credentials
* @throws Error if refresh fails
*/
export async function getOAuthApiKey(provider: OAuthProvider): Promise<string | null> {
const credentials = loadOAuthCredentials(provider);
if (!credentials) {
export async function getOAuthApiKey(
provider: OAuthProvider,
credentials: Record<string, OAuthCredentials>,
): Promise<{ newCredentials: OAuthCredentials; apiKey: string } | null> {
let creds = credentials[provider];
if (!creds) {
return null;
}
// Providers that need projectId in the API key
const needsProjectId = provider === "google-gemini-cli" || provider === "google-antigravity";
// Check if expired
if (Date.now() >= credentials.expires) {
// Refresh if expired
if (Date.now() >= creds.expires) {
try {
const newToken = await refreshToken(provider);
// For providers that need projectId, return JSON
if (needsProjectId) {
const refreshedCreds = loadOAuthCredentials(provider);
if (refreshedCreds?.projectId) {
return JSON.stringify({ token: newToken, projectId: refreshedCreds.projectId });
}
}
return newToken;
} catch (error) {
console.error(`Failed to refresh OAuth token for ${provider}:`, error);
removeOAuthCredentials(provider);
return null;
creds = await refreshOAuthToken(provider, creds);
} catch (_error) {
throw new Error(`Failed to refresh OAuth token for ${provider}`);
}
}
// For providers that need projectId, return JSON
if (needsProjectId) {
if (!credentials.projectId) {
return null;
}
return JSON.stringify({ token: credentials.access, projectId: credentials.projectId });
}
return credentials.access;
}
/**
* Map model provider to OAuth provider.
* Returns undefined if the provider doesn't use OAuth.
*/
export function getOAuthProviderForModelProvider(modelProvider: string): OAuthProvider | undefined {
const mapping: Record<string, OAuthProvider> = {
anthropic: "anthropic",
"github-copilot": "github-copilot",
"google-gemini-cli": "google-gemini-cli",
"google-antigravity": "google-antigravity",
};
return mapping[modelProvider];
}
// ============================================================================
// Login/Logout types for convenience
// ============================================================================
export type OAuthPrompt = {
message: string;
placeholder?: string;
allowEmpty?: boolean;
};
export type OAuthAuthInfo = {
url: string;
instructions?: string;
};
export interface OAuthProviderInfo {
id: OAuthProvider;
name: string;
available: boolean;
const needsProjectId = provider === "google-gemini-cli" || provider === "google-antigravity";
const apiKey = needsProjectId ? JSON.stringify({ token: creds.access, projectId: creds.projectId }) : creds.access;
return { newCredentials: creds, apiKey };
}
/**

161
packages/ai/src/utils/oauth/storage.ts
View file

@ -1,161 +0,0 @@
/**
* OAuth credential storage with configurable backend.
*
* Default: ~/.pi/agent/oauth.json
* Override with setOAuthStorage() for custom storage locations or backends.
*/
import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs";
import { homedir } from "os";
import { dirname, join } from "path";
export interface OAuthCredentials {
type: "oauth";
refresh: string;
access: string;
expires: number;
enterpriseUrl?: string;
projectId?: string;
email?: string;
}
export interface OAuthStorage {
[provider: string]: OAuthCredentials;
}
export type OAuthProvider = "anthropic" | "github-copilot" | "google-gemini-cli" | "google-antigravity";
/**
* Storage backend interface.
* Implement this to use a custom storage location or backend.
*/
export interface OAuthStorageBackend {
/** Load all OAuth credentials. Return empty object if none exist. */
load(): OAuthStorage;
/** Save all OAuth credentials. */
save(storage: OAuthStorage): void;
}
// ============================================================================
// Default filesystem backend
// ============================================================================
const DEFAULT_PATH = join(homedir(), ".pi", "agent", "oauth.json");
function defaultLoad(): OAuthStorage {
if (!existsSync(DEFAULT_PATH)) {
return {};
}
try {
const content = readFileSync(DEFAULT_PATH, "utf-8");
return JSON.parse(content);
} catch {
return {};
}
}
function defaultSave(storage: OAuthStorage): void {
const configDir = dirname(DEFAULT_PATH);
if (!existsSync(configDir)) {
mkdirSync(configDir, { recursive: true, mode: 0o700 });
}
writeFileSync(DEFAULT_PATH, JSON.stringify(storage, null, 2), "utf-8");
chmodSync(DEFAULT_PATH, 0o600);
}
// ============================================================================
// Configurable backend
// ============================================================================
let currentBackend: OAuthStorageBackend = {
load: defaultLoad,
save: defaultSave,
};
/**
* Configure the OAuth storage backend.
*
* @example
* // Custom file path
* setOAuthStorage({
* load: () => JSON.parse(readFileSync('/custom/path/oauth.json', 'utf-8')),
* save: (storage) => writeFileSync('/custom/path/oauth.json', JSON.stringify(storage))
* });
*
* @example
* // In-memory storage (for testing)
* let memoryStorage = {};
* setOAuthStorage({
* load: () => memoryStorage,
* save: (storage) => { memoryStorage = storage; }
* });
*/
export function setOAuthStorage(backend: OAuthStorageBackend): void {
currentBackend = backend;
}
/**
* Reset to default filesystem storage (~/.pi/agent/oauth.json)
*/
export function resetOAuthStorage(): void {
currentBackend = { load: defaultLoad, save: defaultSave };
}
/**
* Get the default OAuth path (for reference, may not be used if custom backend is set)
*/
export function getOAuthPath(): string {
return DEFAULT_PATH;
}
// ============================================================================
// Public API (uses current backend)
// ============================================================================
/**
* Load all OAuth credentials
*/
export function loadOAuthStorage(): OAuthStorage {
return currentBackend.load();
}
/**
* Load OAuth credentials for a specific provider
*/
export function loadOAuthCredentials(provider: string): OAuthCredentials | null {
const storage = currentBackend.load();
return storage[provider] || null;
}
/**
* Save OAuth credentials for a specific provider
*/
export function saveOAuthCredentials(provider: string, creds: OAuthCredentials): void {
const storage = currentBackend.load();
storage[provider] = creds;
currentBackend.save(storage);
}
/**
* Remove OAuth credentials for a specific provider
*/
export function removeOAuthCredentials(provider: string): void {
const storage = currentBackend.load();
delete storage[provider];
currentBackend.save(storage);
}
/**
* Check if OAuth credentials exist for a provider
*/
export function hasOAuthCredentials(provider: string): boolean {
return loadOAuthCredentials(provider) !== null;
}
/**
* List all providers with OAuth credentials
*/
export function listOAuthProviders(): string[] {
const storage = currentBackend.load();
return Object.keys(storage);
}

27
packages/ai/src/utils/oauth/types.ts Normal file
View file

@ -0,0 +1,27 @@
export type OAuthCredentials = {
refresh: string;
access: string;
expires: number;
enterpriseUrl?: string;
projectId?: string;
email?: string;
};
export type OAuthProvider = "anthropic" | "github-copilot" | "google-gemini-cli" | "google-antigravity";
export type OAuthPrompt = {
message: string;
placeholder?: string;
allowEmpty?: boolean;
};
export type OAuthAuthInfo = {
url: string;
instructions?: string;
};
export interface OAuthProviderInfo {
id: OAuthProvider;
name: string;
available: boolean;
}

3
packages/ai/test/abort.test.ts
View file

@ -1,7 +1,8 @@
import { describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { complete, resolveApiKey, stream } from "../src/stream.js";
import { complete, stream } from "../src/stream.js";
import type { Api, Context, Model, OptionsForApi } from "../src/types.js";
import { resolveApiKey } from "./oauth.js";
// Resolve OAuth tokens at module level (async, runs before tests)
const geminiCliToken = await resolveApiKey("google-gemini-cli");

2
packages/ai/test/agent.test.ts
View file

@ -3,7 +3,6 @@ import { agentLoop, agentLoopContinue } from "../src/agent/agent-loop.js";
import { calculateTool } from "../src/agent/tools/calculate.js";
import type { AgentContext, AgentEvent, AgentLoopConfig } from "../src/agent/types.js";
import { getModel } from "../src/models.js";
import { resolveApiKey } from "../src/stream.js";
import type {
Api,
AssistantMessage,
@ -13,6 +12,7 @@ import type {
ToolResultMessage,
UserMessage,
} from "../src/types.js";
import { resolveApiKey } from "./oauth.js";
// Resolve OAuth tokens at module level (async, runs before tests)
const oauthTokens = await Promise.all([

3
packages/ai/test/context-overflow.test.ts
View file

@ -15,9 +15,10 @@ import type { ChildProcess } from "child_process";
import { execSync, spawn } from "child_process";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { complete, resolveApiKey } from "../src/stream.js";
import { complete } from "../src/stream.js";
import type { AssistantMessage, Context, Model, Usage } from "../src/types.js";
import { isContextOverflow } from "../src/utils/overflow.js";
import { resolveApiKey } from "./oauth.js";
// Resolve OAuth tokens at module level (async, runs before tests)
const oauthTokens = await Promise.all([

3
packages/ai/test/empty.test.ts
View file

@ -1,7 +1,8 @@
import { describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { complete, resolveApiKey } from "../src/stream.js";
import { complete } from "../src/stream.js";
import type { Api, AssistantMessage, Context, Model, OptionsForApi, UserMessage } from "../src/types.js";
import { resolveApiKey } from "./oauth.js";
// Resolve OAuth tokens at module level (async, runs before tests)
const oauthTokens = await Promise.all([

3
packages/ai/test/image-tool-result.test.ts
View file

@ -3,8 +3,9 @@ import { join } from "node:path";
import { Type } from "@sinclair/typebox";
import { describe, expect, it } from "vitest";
import type { Api, Context, Model, Tool, ToolResultMessage } from "../src/index.js";
import { complete, getModel, resolveApiKey } from "../src/index.js";
import { complete, getModel } from "../src/index.js";
import type { OptionsForApi } from "../src/types.js";
import { resolveApiKey } from "./oauth.js";
// Resolve OAuth tokens at module level (async, runs before tests)
const oauthTokens = await Promise.all([

89
packages/ai/test/oauth.ts Normal file
View file

@ -0,0 +1,89 @@
/**
* Test helper for resolving API keys from ~/.pi/agent/auth.json
*
* Supports both API key and OAuth credentials.
* OAuth tokens are automatically refreshed if expired and saved back to auth.json.
*/
import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs";
import { homedir } from "os";
import { dirname, join } from "path";
import { getOAuthApiKey } from "../src/utils/oauth/index.js";
import type { OAuthCredentials, OAuthProvider } from "../src/utils/oauth/types.js";
const AUTH_PATH = join(homedir(), ".pi", "agent", "auth.json");
type ApiKeyCredential = {
type: "api_key";
key: string;
};
type OAuthCredentialEntry = {
type: "oauth";
} & OAuthCredentials;
type AuthCredential = ApiKeyCredential | OAuthCredentialEntry;
type AuthStorage = Record<string, AuthCredential>;
function loadAuthStorage(): AuthStorage {
if (!existsSync(AUTH_PATH)) {
return {};
}
try {
const content = readFileSync(AUTH_PATH, "utf-8");
return JSON.parse(content);
} catch {
return {};
}
}
function saveAuthStorage(storage: AuthStorage): void {
const configDir = dirname(AUTH_PATH);
if (!existsSync(configDir)) {
mkdirSync(configDir, { recursive: true, mode: 0o700 });
}
writeFileSync(AUTH_PATH, JSON.stringify(storage, null, 2), "utf-8");
chmodSync(AUTH_PATH, 0o600);
}
/**
* Resolve API key for a provider from ~/.pi/agent/auth.json
*
* For API key credentials, returns the key directly.
* For OAuth credentials, returns the access token (refreshing if expired and saving back).
*
* For google-gemini-cli and google-antigravity, returns JSON-encoded { token, projectId }
*/
export async function resolveApiKey(provider: string): Promise<string | undefined> {
const storage = loadAuthStorage();
const entry = storage[provider];
if (!entry) return undefined;
if (entry.type === "api_key") {
return entry.key;
}
if (entry.type === "oauth") {
// Build OAuthCredentials record for getOAuthApiKey
const oauthCredentials: Record<string, OAuthCredentials> = {};
for (const [key, value] of Object.entries(storage)) {
if (value.type === "oauth") {
const { type: _, ...creds } = value;
oauthCredentials[key] = creds;
}
}
const result = await getOAuthApiKey(provider as OAuthProvider, oauthCredentials);
if (!result) return undefined;
// Save refreshed credentials back to auth.json
storage[provider] = { type: "oauth", ...result.newCredentials };
saveAuthStorage(storage);
return result.apiKey;
}
return undefined;
}

3
packages/ai/test/stream.test.ts
View file

@ -5,9 +5,10 @@ import { dirname, join } from "path";
import { fileURLToPath } from "url";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { complete, resolveApiKey, stream } from "../src/stream.js";
import { complete, stream } from "../src/stream.js";
import type { Api, Context, ImageContent, Model, OptionsForApi, Tool, ToolResultMessage } from "../src/types.js";
import { StringEnum } from "../src/utils/typebox-helpers.js";
import { resolveApiKey } from "./oauth.js";
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

3
packages/ai/test/tokens.test.ts
View file

@ -1,7 +1,8 @@
import { describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { resolveApiKey, stream } from "../src/stream.js";
import { stream } from "../src/stream.js";
import type { Api, Context, Model, OptionsForApi } from "../src/types.js";
import { resolveApiKey } from "./oauth.js";
// Resolve OAuth tokens at module level (async, runs before tests)
const oauthTokens = await Promise.all([

3
packages/ai/test/tool-call-without-result.test.ts
View file

@ -1,8 +1,9 @@
import { Type } from "@sinclair/typebox";
import { describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { complete, resolveApiKey } from "../src/stream.js";
import { complete } from "../src/stream.js";
import type { Api, Context, Model, OptionsForApi, Tool } from "../src/types.js";
import { resolveApiKey } from "./oauth.js";
// Resolve OAuth tokens at module level (async, runs before tests)
const oauthTokens = await Promise.all([

3
packages/ai/test/total-tokens.test.ts
View file

@ -14,8 +14,9 @@
import { describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { complete, resolveApiKey } from "../src/stream.js";
import { complete } from "../src/stream.js";
import type { Api, Context, Model, OptionsForApi, Usage } from "../src/types.js";
import { resolveApiKey } from "./oauth.js";
// Resolve OAuth tokens at module level (async, runs before tests)
const oauthTokens = await Promise.all([

3
packages/ai/test/unicode-surrogate.test.ts
View file

@ -1,8 +1,9 @@
import { Type } from "@sinclair/typebox";
import { describe, expect, it } from "vitest";
import { getModel } from "../src/models.js";
import { complete, resolveApiKey } from "../src/stream.js";
import { complete } from "../src/stream.js";
import type { Api, Context, Model, OptionsForApi, ToolResultMessage } from "../src/types.js";
import { resolveApiKey } from "./oauth.js";
// Empty schema for test tools - must be proper OBJECT type for Cloud Code Assist
const emptySchema = Type.Object({});