diff --git a/docs/auth-broker-gateway.md b/docs/auth-broker-gateway.md index 1fccca164..c77aaf1e6 100644 --- a/docs/auth-broker-gateway.md +++ b/docs/auth-broker-gateway.md @@ -148,6 +148,30 @@ Freshness is anchored to the broker-stamped `snapshot.generatedAt`, not local wr If the broker is down at boot and a fresh cache exists, startup now succeeds from the cached snapshot. If the cache is missing, expired, corrupt, written for a different URL, or encrypted with a different token, startup falls back to the live fetch and fails the same way it did before if the broker is unreachable. +## Client account pools (routing, not authorization) + +Broker clients can restrict their visible OAuth accounts by setting `OMP_AUTH_BROKER_ACCOUNT_POOL_FILE` to a JSON file. The file maps provider IDs to exact `identityKey` values from the broker snapshot protocol: + +```json +{ + "anthropic": ["email:alice@example.com|org:org-team"], + "openai-codex": [] +} +``` + +`identityKey` is the token-free identity field already carried by each authenticated `/v1/snapshot` credential entry. Operator tooling should project only `provider` and `identityKey`; it must not retain or print the accompanying credential payload. A dedicated account-listing CLI is intentionally outside this routing feature's scope. + +SDK hosts can supply the same provider-to-identity mapping as `accountPool` in `discoverAuthStorage()` or `RemoteAuthCredentialStore`. An explicit programmatic pool takes precedence over the environment file. + +- A missing provider is unrestricted. +- An empty array hides every OAuth credential for that provider. +- A non-empty array exposes only exact identity matches, including organization/workspace qualifiers. +- API-key credentials remain visible; the pool applies only to OAuth accounts. + +The file is parsed once when broker-backed auth storage starts. An unreadable file, malformed JSON, or invalid provider entry aborts initialization rather than silently broadening the pool. Full snapshots, SSE updates, refresh responses, and aggregate usage are filtered consistently. For a provider named in the pool, aggregate reports are returned only when they can be attributed to a visible OAuth identity; reports attributable only to an API key or lacking matching identity metadata fail closed. The encrypted snapshot cache remains a raw broker snapshot so trusted processes sharing that cache can apply different pools. + +This is a **trusted-client routing policy, not an authorization boundary**. The client still holds a broker bearer token, receives raw broker responses before applying its local view, and can call broker endpoints directly. Use server-side authorization—not account pools—when clients must be prevented from retrieving other credentials. + ## Operator opt-in The broker is **off** unless `OMP_AUTH_BROKER_URL` (or `auth.broker.url` in `config.yml`) is set. When set, `discoverAuthStorage` in `packages/coding-agent/src/sdk.ts` swaps the local SQLite credential store for `RemoteAuthCredentialStore` and every API call resolves credentials through the broker. @@ -160,6 +184,7 @@ The broker is **off** unless `OMP_AUTH_BROKER_URL` (or `auth.broker.url` in `con | `OMP_AUTH_BROKER_TOKEN` | Bearer token used for every broker endpoint except `/v1/healthz`. | When `OMP_AUTH_BROKER_URL` is set and no token is available from `auth.broker.token` or `/auth-broker.token`. | | `OMP_AUTH_BROKER_SNAPSHOT_TTL_MS` | Freshness window for the encrypted local snapshot cache. Default `3600000` (1 h); `0` disables cache reads and writes. | Optional in broker mode. | | `OMP_AUTH_BROKER_SNAPSHOT_CACHE` | Path override for the encrypted local snapshot cache. Default `~/.omp/cache/auth-broker-snapshot.enc` (or XDG cache equivalent). | Optional in broker mode. | +| `OMP_AUTH_BROKER_ACCOUNT_POOL_FILE` | JSON file mapping provider IDs to OAuth `identityKey` values visible to this trusted client. Parsed once; invalid files abort initialization. API keys are unaffected. | Optional in broker mode. | Resolution order in `resolveAuthBrokerConfig()`: diff --git a/docs/environment-variables.md b/docs/environment-variables.md index d146a97f1..038366d9b 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -104,6 +104,7 @@ When the broker is enabled, the local SQLite credential store is bypassed and al | `OMP_AUTH_BROKER_TOKEN` | Bearer token sent on every broker endpoint except `/v1/healthz` | `OMP_AUTH_BROKER_URL` is set and no token is available from `auth.broker.token` or `/auth-broker.token` | Resolution: this env → `auth.broker.token` (`$ENV_NAME` indirection supported) → `/auth-broker.token` (mode `0600`). `` is `~/.omp/` (respecting `PI_CONFIG_DIR`). | | `OMP_AUTH_BROKER_SNAPSHOT_TTL_MS` | Freshness window for the encrypted local broker snapshot cache | Optional in broker mode | Default `3600000` (1 h). Freshness is based on broker `snapshot.generatedAt`; `0` disables cache reads/writes and forces the old blocking fetch every startup. | | `OMP_AUTH_BROKER_SNAPSHOT_CACHE` | Path to the encrypted local broker snapshot cache | Optional in broker mode | Defaults to `~/.omp/cache/auth-broker-snapshot.enc` (or XDG cache equivalent). Useful for tests, ephemeral hosts, or relocating the `0600` cache file. | +| `OMP_AUTH_BROKER_ACCOUNT_POOL_FILE` | Process-scoped OAuth account routing for a trusted broker client | Optional in broker mode | Path to a JSON object mapping provider IDs to exact broker `identityKey` arrays. Missing providers are unrestricted; `[]` hides that provider's OAuth accounts; API keys remain visible. Parsed once at startup and fails closed on invalid input. This is not server authorization. | The gateway has no dedicated env vars — it inherits `OMP_AUTH_BROKER_*`. Its own inbound bearer token lives at `/auth-gateway.token` and is managed via `omp auth-gateway token`. diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 117bb9072..0292f3d13 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -5,6 +5,7 @@ ### Added - Added Anthropic extra-usage reporting across `omp usage`, interactive `/usage`, and ACP `/usage`: the OAuth usage endpoint's authoritative `spend` payload (or legacy `extra_usage` fallback when absent) is normalized into a `Claude Extra Usage` USD row; capped accounts show limit/remaining/fractions and status, while uncapped spend exposes only its absolute used amount—rendered as `$… used` in CLI/TUI and `123.45 usd used` in ACP—without a fabricated cap, percentage, or status. ([#5575](https://github.com/can1357/oh-my-pi/issues/5575)) +- Added process-scoped OAuth account pools for trusted auth-broker clients via `OMP_AUTH_BROKER_ACCOUNT_POOL_FILE`, consistently filtering snapshots, streaming updates, refreshes, and usage reports to selected OAuth identities while leaving API-key credentials and the shared encrypted snapshot cache unrestricted. ### Fixed diff --git a/packages/ai/src/auth-broker/discover.ts b/packages/ai/src/auth-broker/discover.ts index 7cb08cb26..2198a81c2 100644 --- a/packages/ai/src/auth-broker/discover.ts +++ b/packages/ai/src/auth-broker/discover.ts @@ -18,7 +18,7 @@ import { YAML } from "bun"; import { AuthStorage } from "../auth-storage"; import * as AIError from "../error"; import { AuthBrokerClient } from "./client"; -import { RemoteAuthCredentialStore } from "./remote-store"; +import { type AuthBrokerAccountPool, RemoteAuthCredentialStore } from "./remote-store"; import { readAuthBrokerSnapshotCache, writeAuthBrokerSnapshotCache } from "./snapshot-cache"; import { DEFAULT_SNAPSHOT_CACHE_TTL_MS, type SnapshotResponse } from "./types"; @@ -37,6 +37,8 @@ export interface DiscoverAuthStorageOptions { configValueResolver?: (config: string) => Promise; cachePath?: string; sourceLabel?: string; + /** Programmatic pool for SDK hosts. Takes precedence over the environment file. */ + accountPool?: AuthBrokerAccountPool; } /** Path to the local bearer token file. Created by `omp auth-broker token`. */ @@ -112,6 +114,57 @@ async function readConfigYaml(agentDir: string): Promise { return {}; } +export async function loadAuthBrokerAccountPool(): Promise { + const filePath = process.env.OMP_AUTH_BROKER_ACCOUNT_POOL_FILE?.trim(); + if (!filePath) return undefined; + + let parsed: unknown; + try { + parsed = await Bun.file(filePath).json(); + } catch (error) { + throw new AIError.ConfigurationError(`Unable to read OMP_AUTH_BROKER_ACCOUNT_POOL_FILE at ${filePath}`, { + cause: error, + }); + } + if (parsed === null || typeof parsed !== "object" || Array.isArray(parsed)) { + throw new AIError.ConfigurationError("OMP_AUTH_BROKER_ACCOUNT_POOL_FILE must contain a JSON object"); + } + + const accountPool = new Map>(); + for (const [provider, value] of Object.entries(parsed)) { + const normalizedProvider = provider.trim(); + if (normalizedProvider.length === 0) { + throw new AIError.ConfigurationError("OMP_AUTH_BROKER_ACCOUNT_POOL_FILE contains an empty provider id"); + } + if (provider !== normalizedProvider) { + throw new AIError.ConfigurationError( + "OMP_AUTH_BROKER_ACCOUNT_POOL_FILE contains a provider id with surrounding whitespace", + ); + } + if (!Array.isArray(value)) { + throw new AIError.ConfigurationError( + `OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} must be an array of identity keys`, + ); + } + const identities = new Set(); + for (const identity of value) { + if (typeof identity !== "string" || identity.length === 0) { + throw new AIError.ConfigurationError( + `OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} contains an invalid identity key`, + ); + } + if (identity !== identity.trim()) { + throw new AIError.ConfigurationError( + `OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} contains an identity key with surrounding whitespace`, + ); + } + identities.add(identity); + } + accountPool.set(provider, identities); + } + return accountPool; +} + function resolveSnapshotTtlMs(): number { const raw = process.env.OMP_AUTH_BROKER_SNAPSHOT_TTL_MS; if (raw === undefined) return DEFAULT_SNAPSHOT_CACHE_TTL_MS; @@ -183,6 +236,7 @@ export async function discoverAuthStorage(options: DiscoverAuthStorageOptions = }); if (brokerConfig) { + const accountPool = options.accountPool ?? (await loadAuthBrokerAccountPool()); const client = new AuthBrokerClient({ url: brokerConfig.url, token: brokerConfig.token }); const cachePath = options.cachePath ?? getAuthBrokerSnapshotCachePath(); const ttlMs = resolveSnapshotTtlMs(); @@ -226,6 +280,7 @@ export async function discoverAuthStorage(options: DiscoverAuthStorageOptions = client, initialSnapshot, onSnapshot: persist, + accountPool, }); const storage = new AuthStorage(store, { configValueResolver: options.configValueResolver, diff --git a/packages/ai/src/auth-broker/remote-store.ts b/packages/ai/src/auth-broker/remote-store.ts index 20e074a86..7eb249df1 100644 --- a/packages/ai/src/auth-broker/remote-store.ts +++ b/packages/ai/src/auth-broker/remote-store.ts @@ -31,6 +31,23 @@ import type { SnapshotStreamEvent, } from "./types"; +/** + * Per-provider OAuth identities visible to this trusted broker client. + * Missing providers are unrestricted; an empty set excludes that provider's + * OAuth credentials. API keys are never filtered. + */ +export type AuthBrokerAccountPool = ReadonlyMap>; + +function isCredentialInAccountPool( + entry: Pick, + accountPool: AuthBrokerAccountPool | undefined, +): boolean { + if (entry.credential.type !== "oauth") return true; + const identities = accountPool?.get(entry.provider); + if (identities === undefined) return true; + return entry.identityKey !== null && identities.has(entry.identityKey); +} + /** * Client-side TTL for the aggregate `/v1/usage` response. The broker dedups * upstream `/usage` hits via AuthStorage's 5-minute per-credential cache plus @@ -205,16 +222,23 @@ export interface RemoteAuthCredentialStoreOptions { */ streamSnapshots?: boolean; /** - * Called after broker-sourced full snapshots are applied. The constructor's - * initial snapshot intentionally does not trigger this hook. + * Called with each broker-sourced raw full snapshot after the filtered + * public view is applied. The constructor's initial snapshot intentionally + * does not trigger this hook. */ onSnapshot?: (snapshot: SnapshotResponse, generation: number) => void; + /** + * OAuth identities visible through this store. This is a trusted-client + * routing policy, not broker authorization. + */ + accountPool?: AuthBrokerAccountPool; } export class RemoteAuthCredentialStore implements AuthCredentialStore { readonly #client: AuthBrokerClient; readonly #streamSnapshots: boolean; readonly #onSnapshot?: (snapshot: SnapshotResponse, generation: number) => void; + readonly #accountPool?: AuthBrokerAccountPool; #snapshot: SnapshotResponse = emptySnapshot(); #snapshotReceivedAt = Date.now(); #generation = 0; @@ -239,6 +263,9 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { constructor(opts: RemoteAuthCredentialStoreOptions) { this.#client = opts.client; this.#streamSnapshots = opts.streamSnapshots ?? true; + this.#accountPool = opts.accountPool + ? new Map([...opts.accountPool].map(([provider, identities]) => [provider, new Set(identities)])) + : undefined; this.#applySnapshot(opts.initialSnapshot ?? emptySnapshot(), opts.initialSnapshot?.generation ?? 0); this.#onSnapshot = opts.onSnapshot; void this.#runBackground(); @@ -255,7 +282,9 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { #applySnapshot(snapshot: SnapshotResponse, generation: number, protectNewBlocks = true): void { const nowMs = Date.now(); const previousCredentials = this.#snapshot.credentials; - const credentials = snapshot.credentials.map(entry => this.#normalizeSnapshotEntryBlocks(entry, nowMs)); + const credentials = snapshot.credentials + .filter(entry => isCredentialInAccountPool(entry, this.#accountPool)) + .map(entry => this.#normalizeSnapshotEntryBlocks(entry, nowMs)); if (snapshotBlocksChanged(previousCredentials, credentials)) this.#invalidateUsageCache(); if (protectNewBlocks) this.#protectNewSnapshotBlocks(previousCredentials, credentials, nowMs); this.#snapshot = { ...snapshot, credentials }; @@ -264,7 +293,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { const onSnapshot = this.#onSnapshot; if (!onSnapshot) return; try { - onSnapshot(this.#snapshot, generation); + onSnapshot(snapshot, generation); } catch (error) { logger.debug("auth-broker snapshot callback failed", { error: String(error) }); } @@ -383,6 +412,10 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { generation: number, serverNowMs: number, ): void { + if (!isCredentialInAccountPool(entry, this.#accountPool)) { + this.#removeStreamCredential(entry.id, refresher, generation, serverNowMs); + return; + } const incoming = this.#normalizeSnapshotEntryBlocks(entry, Date.now()); const index = this.#snapshot.credentials.findIndex(candidate => candidate.id === incoming.id); const previousBlocks = index === -1 ? undefined : this.#snapshot.credentials[index]?.blocks; @@ -577,7 +610,11 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { if (entry.credential.type !== "oauth") { throw new AIError.AuthBrokerError(`Broker returned non-OAuth credential for id=${credentialId}`); } - this.#applyCredentialEntry(entry); + if (!this.#applyCredentialEntry(entry)) { + throw new AIError.AuthBrokerError( + `Broker refreshed credential id=${credentialId} outside the configured account pool`, + ); + } this.#maybeRefreshSnapshot("suspect credential refresh"); } @@ -677,20 +714,27 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { .map(entry => [entry.id, entry.blocks] as const), ); const others = this.#snapshot.credentials.filter(entry => entry.provider !== provider); - const incoming = entries.map(entry => credentialEntryWithBlocks(entry, existingBlocks.get(entry.id))); + const incoming = entries + .filter(entry => isCredentialInAccountPool(entry, this.#accountPool)) + .map(entry => credentialEntryWithBlocks(entry, existingBlocks.get(entry.id))); this.#snapshot = { ...this.#snapshot, credentials: [...others, ...incoming] }; } - #applyCredentialEntry(entry: AuthCredentialSnapshotEntry): void { + #applyCredentialEntry(entry: AuthCredentialSnapshotEntry): boolean { + if (!isCredentialInAccountPool(entry, this.#accountPool)) { + this.#removeCredentialById(entry.id); + return false; + } const index = this.#snapshot.credentials.findIndex(candidate => candidate.id === entry.id); const existingBlocks = index === -1 ? undefined : this.#snapshot.credentials[index]?.blocks; const incoming = credentialEntryWithBlocks(entry, existingBlocks); if (index === -1) { this.#snapshot = { ...this.#snapshot, credentials: [...this.#snapshot.credentials, incoming] }; - return; + return true; } const credentials = [...this.#snapshot.credentials]; credentials[index] = incoming; this.#snapshot = { ...this.#snapshot, credentials }; + return true; } #removeProviderEntries(provider: string): void { @@ -837,14 +881,19 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { signal?: AbortSignal, ): Promise { const { entry } = await this.#client.refreshCredential(credentialId, signal); + if (entry.credential.type !== "oauth") { + throw new AIError.AuthBrokerError(`Broker returned non-OAuth credential for id=${credentialId}`); + } + if (!this.#applyCredentialEntry(entry)) { + throw new AIError.AuthBrokerError( + `Broker refreshed credential id=${credentialId} outside the configured account pool`, + ); + } if (!this.#streamingActive) { await this.refreshSnapshot().catch(error => { logger.debug("auth-broker snapshot refresh after credential refresh failed", { error: String(error) }); }); } - if (entry.credential.type !== "oauth") { - throw new AIError.AuthBrokerError(`Broker returned non-OAuth credential for id=${credentialId}`); - } const refreshed = entry.credential; return { access: refreshed.access, @@ -865,7 +914,8 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { */ async fetchUsageReports(signal?: AbortSignal): Promise { const reports = await this.#raceWithSignal(this.#loadUsageReports(), signal); - return reports ? this.#applyUsageOverlays(reports) : null; + if (!reports) return null; + return this.#filterUsageReports(this.#applyUsageOverlays(reports)); } /** @@ -883,12 +933,27 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { signal?: AbortSignal, ): Promise { const reports = await this.#raceWithSignal(this.#loadUsageReports(), signal); - const matched = reports ? matchUsageReport(reports, provider, credential) : null; + const visibleReports = reports ? this.#filterUsageReports(reports) : null; + const matched = visibleReports ? matchUsageReport(visibleReports, provider, credential) : null; const overlay = this.#getActiveUsageOverlay(provider, credential); if (matched && overlay) return mergeUsageReports(matched, overlay); return overlay ?? matched; } + #filterUsageReports(reports: UsageReport[]): UsageReport[] { + const accountPool = this.#accountPool; + if (!accountPool) return reports; + return reports.filter(report => { + if (!accountPool.has(report.provider)) return true; + return this.#snapshot.credentials.some( + entry => + entry.provider === report.provider && + entry.credential.type === "oauth" && + usageReportMatchesCredential(report, entry.credential), + ); + }); + } + ingestUsageReport(provider: Provider, credential: OAuthCredential, report: UsageReport): boolean { const key = usageOverlayKey(provider, credential); if (!key) return false; @@ -1059,6 +1124,21 @@ function matchUsageReport(reports: UsageReport[], provider: Provider, credential return null; } +function usageReportMatchesCredential(report: UsageReport, credential: OAuthCredential): boolean { + const metadata = (report.metadata ?? {}) as Record; + const credentialOrg = credential.orgId?.trim().toLowerCase(); + const reportOrg = readMetadataString(metadata, "orgId")?.toLowerCase(); + if (credentialOrg !== reportOrg) return false; + + const accountId = credential.accountId?.trim().toLowerCase(); + const email = credential.email?.trim().toLowerCase(); + const projectId = credential.projectId?.trim().toLowerCase(); + if (accountId || email || projectId) { + return reportMatchesIdentity(report, accountId, email, projectId); + } + return credentialOrg !== undefined; +} + function findMatchingReportIndex(reports: UsageReport[], overlay: UsageReport): number { const all = reports .map((report, index) => ({ report, index })) diff --git a/packages/ai/test/auth-broker-config-discovery.test.ts b/packages/ai/test/auth-broker-config-discovery.test.ts index 67d4033aa..752a5b7a2 100644 --- a/packages/ai/test/auth-broker-config-discovery.test.ts +++ b/packages/ai/test/auth-broker-config-discovery.test.ts @@ -2,13 +2,14 @@ import { afterEach, beforeEach, describe, expect, test } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as path from "node:path"; -import { resolveAuthBrokerConfig } from "@oh-my-pi/pi-ai/auth-broker"; +import { discoverAuthStorage, resolveAuthBrokerConfig } from "@oh-my-pi/pi-ai/auth-broker"; import { removeWithRetries } from "../../utils/src/temp"; import { withEnv } from "./helpers"; const SUPPRESS_AUTH_BROKER_ENV = { OMP_AUTH_BROKER_URL: undefined, OMP_AUTH_BROKER_TOKEN: undefined, + OMP_AUTH_BROKER_ACCOUNT_POOL_FILE: undefined, } as const; describe("resolveAuthBrokerConfig config discovery", () => { @@ -56,4 +57,32 @@ describe("resolveAuthBrokerConfig config discovery", () => { }); }); }); + + test("rejects unreadable or malformed account-pool files before connecting", async () => { + const poolPath = path.join(agentDir, "account-pool.json"); + const brokerEnv = { + ...SUPPRESS_AUTH_BROKER_ENV, + OMP_AUTH_BROKER_URL: "http://127.0.0.1:1", + OMP_AUTH_BROKER_TOKEN: "test-token", + OMP_AUTH_BROKER_ACCOUNT_POOL_FILE: poolPath, + } as const; + + await withEnv(brokerEnv, async () => { + await expect(discoverAuthStorage({ agentDir })).rejects.toThrow( + "Unable to read OMP_AUTH_BROKER_ACCOUNT_POOL_FILE", + ); + + const invalidFiles = [ + ["[]", "must contain a JSON object"], + ['{"anthropic":"email:a@example.com"}', "must be an array of identity keys"], + ['{"anthropic":[42]}', "contains an invalid identity key"], + ['{" anthropic":["email:a@example.com"]}', "provider id with surrounding whitespace"], + ['{"anthropic":[" email:a@example.com"]}', "identity key with surrounding whitespace"], + ] as const; + for (const [content, expectedError] of invalidFiles) { + await Bun.write(poolPath, content); + await expect(discoverAuthStorage({ agentDir })).rejects.toThrow(expectedError); + } + }); + }); }); diff --git a/packages/ai/test/auth-broker-remote-store.test.ts b/packages/ai/test/auth-broker-remote-store.test.ts index 589f0a160..ae7b11502 100644 --- a/packages/ai/test/auth-broker-remote-store.test.ts +++ b/packages/ai/test/auth-broker-remote-store.test.ts @@ -6,11 +6,13 @@ import { AuthStorage, REMOTE_REFRESH_SENTINEL, SqliteAuthCredentialStore } from import { AuthBrokerClient, type AuthBrokerServerHandle, + discoverAuthStorage, RemoteAuthCredentialStore, type SnapshotResponse, startAuthBroker, } from "@oh-my-pi/pi-ai/auth-broker"; import { removeWithRetries } from "../../utils/src/temp"; +import { withEnv } from "./helpers"; const ANTHROPIC_ENV = ["ANTHROPIC_API_KEY", "ANTHROPIC_OAUTH_TOKEN"] as const; const savedEnv: Partial> = {}; @@ -131,4 +133,134 @@ describe("RemoteAuthCredentialStore SSE integration", () => { expect(callbacks[0].generation).toBe(refreshed.generation); expect(callbacks[0].snapshot).toEqual(refreshed); }); + + test("filters configured OAuth identities while preserving API keys and raw snapshot callbacks", async () => { + storage!.upsertCredential("anthropic", mintOAuthCredential("b", Date.now() + 120_000)); + storage!.upsertCredential("anthropic", { type: "api_key", key: "visible-api-key" }); + const client = new AuthBrokerClient({ url: handle!.url, token }); + const initialResult = await client.fetchSnapshot(); + if (initialResult.status !== 200) throw new Error("expected initial snapshot"); + const allowed = initialResult.snapshot.credentials.find(entry => entry.identityKey?.includes("a@example.com")); + const excluded = initialResult.snapshot.credentials.find(entry => entry.identityKey?.includes("b@example.com")); + if (!allowed?.identityKey || !excluded?.identityKey) throw new Error("expected OAuth identity keys"); + const identities = new Set([allowed.identityKey]); + const callbacks: SnapshotResponse[] = []; + remote = new RemoteAuthCredentialStore({ + client, + initialSnapshot: initialResult.snapshot, + streamSnapshots: false, + accountPool: new Map([["anthropic", identities]]), + onSnapshot: snapshot => { + callbacks.push(snapshot); + }, + }); + + identities.add(excluded.identityKey); + expect( + remote + .listAuthCredentials("anthropic") + .map(entry => entry.credential.type) + .sort(), + ).toEqual(["api_key", "oauth"]); + const refreshed = await remote.refreshSnapshot(); + expect(refreshed.credentials.filter(entry => entry.credential.type === "oauth")).toHaveLength(1); + expect(callbacks.at(-1)?.credentials).toHaveLength(3); + }); + + test("advances the SSE generation without exposing an excluded entry", async () => { + const client = new AuthBrokerClient({ url: handle!.url, token }); + const initialResult = await client.fetchSnapshot(); + if (initialResult.status !== 200) throw new Error("expected initial snapshot"); + const allowed = initialResult.snapshot.credentials[0]; + if (!allowed?.identityKey) throw new Error("expected OAuth identity key"); + remote = new RemoteAuthCredentialStore({ + client, + initialSnapshot: initialResult.snapshot, + accountPool: new Map([["anthropic", new Set([allowed.identityKey])]]), + }); + const initialGeneration = remote.snapshot.generation; + + storage!.upsertCredential("anthropic", mintOAuthCredential("b", Date.now() + 120_000)); + await waitUntil(() => remote!.snapshot.generation > initialGeneration); + + expect(remote.snapshot.credentials).toHaveLength(1); + expect(remote.snapshot.credentials[0]?.identityKey).toBe(allowed.identityKey); + }); + + test("treats a missing provider as unrestricted and an empty provider pool as OAuth-disabled", async () => { + storage!.upsertCredential("openai-codex", mintOAuthCredential("codex", Date.now() + 120_000)); + const client = new AuthBrokerClient({ url: handle!.url, token }); + const initialResult = await client.fetchSnapshot(); + if (initialResult.status !== 200) throw new Error("expected initial snapshot"); + remote = new RemoteAuthCredentialStore({ + client, + initialSnapshot: initialResult.snapshot, + streamSnapshots: false, + accountPool: new Map([["anthropic", new Set()]]), + }); + + expect(remote.listAuthCredentials("anthropic")).toEqual([]); + expect(remote.listAuthCredentials("openai-codex")).toHaveLength(1); + }); + + test("loads the account pool once for broker-backed discovery", async () => { + storage!.upsertCredential("anthropic", mintOAuthCredential("b", Date.now() + 120_000)); + const client = new AuthBrokerClient({ url: handle!.url, token }); + const initialResult = await client.fetchSnapshot(); + if (initialResult.status !== 200) throw new Error("expected initial snapshot"); + const allowed = initialResult.snapshot.credentials.find(entry => entry.identityKey?.includes("a@example.com")); + const excluded = initialResult.snapshot.credentials.find(entry => entry.identityKey?.includes("b@example.com")); + if (!allowed?.identityKey || !excluded?.identityKey) throw new Error("expected OAuth identity keys"); + const poolPath = path.join(tempDir, "account-pool.json"); + await Bun.write(poolPath, JSON.stringify({ anthropic: [allowed.identityKey] })); + + await withEnv( + { + OMP_AUTH_BROKER_URL: handle!.url, + OMP_AUTH_BROKER_TOKEN: token, + OMP_AUTH_BROKER_ACCOUNT_POOL_FILE: poolPath, + }, + async () => { + const discovered = await discoverAuthStorage({ + agentDir: tempDir, + cachePath: path.join(tempDir, "snapshot-cache.enc"), + }); + try { + expect(discovered.listOAuthAccounts("anthropic").map(account => account.email)).toEqual([ + "a@example.com", + ]); + + await Bun.write(poolPath, JSON.stringify({ anthropic: [allowed.identityKey, excluded.identityKey] })); + await discovered.reload(); + expect(discovered.listOAuthAccounts("anthropic").map(account => account.email)).toEqual([ + "a@example.com", + ]); + } finally { + discovered.close(); + } + }, + ); + }); + + test("prefers a programmatic SDK account pool over the environment file", async () => { + await withEnv( + { + OMP_AUTH_BROKER_URL: handle!.url, + OMP_AUTH_BROKER_TOKEN: token, + OMP_AUTH_BROKER_ACCOUNT_POOL_FILE: path.join(tempDir, "missing-account-pool.json"), + }, + async () => { + const discovered = await discoverAuthStorage({ + agentDir: tempDir, + cachePath: path.join(tempDir, "sdk-snapshot-cache.enc"), + accountPool: new Map([["anthropic", new Set()]]), + }); + try { + expect(discovered.listOAuthAccounts("anthropic")).toEqual([]); + } finally { + discovered.close(); + } + }, + ); + }); }); diff --git a/packages/ai/test/remote-auth-store.test.ts b/packages/ai/test/remote-auth-store.test.ts index 9192b5209..c33909b94 100644 --- a/packages/ai/test/remote-auth-store.test.ts +++ b/packages/ai/test/remote-auth-store.test.ts @@ -1092,4 +1092,186 @@ describe("RemoteAuthCredentialStore + AuthStorage integration", () => { expect(serverInvalidateSpy).toHaveBeenCalled(); clientStorage.close(); }); + + test("account pool exposes only qualified usage reports for visible OAuth identities", async () => { + const brokerClient = new AuthBrokerClient({ url: "http://127.0.0.1:9", token: "unused" }); + const now = Date.now(); + const makeCredential = (orgId: string) => ({ + type: "oauth" as const, + access: `access-${orgId}`, + refresh: REMOTE_REFRESH_SENTINEL, + expires: now + 120_000, + accountId: "account-shared", + email: "shared@example.com", + orgId, + }); + const reports: UsageReport[] = [ + { + provider: "anthropic", + fetchedAt: now, + limits: [], + metadata: { accountId: "account-shared", email: "shared@example.com", orgId: "org-team" }, + }, + { + provider: "anthropic", + fetchedAt: now, + limits: [], + metadata: { accountId: "account-shared", email: "shared@example.com", orgId: "org-max" }, + }, + { provider: "anthropic", fetchedAt: now, limits: [] }, + ]; + vi.spyOn(brokerClient, "fetchUsage").mockResolvedValue({ generatedAt: now, reports }); + const teamIdentity = "email:shared@example.com|org:org-team"; + const remoteStore = new RemoteAuthCredentialStore({ + client: brokerClient, + streamSnapshots: false, + accountPool: new Map([["anthropic", new Set([teamIdentity])]]), + initialSnapshot: { + generation: 1, + generatedAt: now, + serverNowMs: now, + refresher: { enabled: false, intervalMs: 0, skewMs: 0, nextSweepInMs: Number.MAX_SAFE_INTEGER }, + credentials: [ + { + id: 1, + provider: "anthropic", + credential: makeCredential("org-team"), + identityKey: teamIdentity, + rotatesInMs: null, + }, + { + id: 2, + provider: "anthropic", + credential: makeCredential("org-max"), + identityKey: "email:shared@example.com|org:org-max", + rotatesInMs: null, + }, + ], + }, + }); + try { + expect(remoteStore.snapshot.credentials.map(entry => entry.identityKey)).toEqual([teamIdentity]); + const visibleReports = await remoteStore.fetchUsageReports(); + expect(visibleReports?.map(report => report.metadata?.orgId)).toEqual(["org-team"]); + expect(await remoteStore.getUsageReport("anthropic", makeCredential("org-max"))).toBeNull(); + } finally { + remoteStore.close(); + } + }); + + test("account pool hides unattributable usage even with a visible API key", async () => { + const brokerClient = new AuthBrokerClient({ url: "http://127.0.0.1:9", token: "unused" }); + const now = Date.now(); + const oauthCredential = { + type: "oauth" as const, + access: "oauth-access", + refresh: REMOTE_REFRESH_SENTINEL, + expires: now + 120_000, + accountId: "oauth-account", + email: "oauth@example.com", + }; + const reports: UsageReport[] = [ + { + provider: "anthropic", + fetchedAt: now, + limits: [], + metadata: { accountId: "oauth-account", email: "oauth@example.com" }, + }, + { + provider: "anthropic", + fetchedAt: now, + limits: [], + metadata: { accountId: "api-key-account" }, + }, + ]; + vi.spyOn(brokerClient, "fetchUsage").mockResolvedValue({ generatedAt: now, reports }); + const oauthIdentity = "email:oauth@example.com"; + const remoteStore = new RemoteAuthCredentialStore({ + client: brokerClient, + streamSnapshots: false, + accountPool: new Map([["anthropic", new Set([oauthIdentity])]]), + initialSnapshot: { + generation: 1, + generatedAt: now, + serverNowMs: now, + refresher: { enabled: false, intervalMs: 0, skewMs: 0, nextSweepInMs: Number.MAX_SAFE_INTEGER }, + credentials: [ + { + id: 1, + provider: "anthropic", + credential: oauthCredential, + identityKey: oauthIdentity, + rotatesInMs: null, + }, + { + id: 2, + provider: "anthropic", + credential: { type: "api_key", key: "visible-api-key" }, + identityKey: null, + rotatesInMs: null, + }, + ], + }, + }); + try { + expect(await remoteStore.fetchUsageReports()).toEqual([reports[0]]); + } finally { + remoteStore.close(); + } + }); + + test("rejects a refreshed credential whose identity leaves the account pool", async () => { + const brokerClient = new AuthBrokerClient({ url: "http://127.0.0.1:9", token: "unused" }); + const now = Date.now(); + const allowedCredential = { + type: "oauth" as const, + access: "allowed-access", + refresh: REMOTE_REFRESH_SENTINEL, + expires: now + 120_000, + accountId: "account-allowed", + email: "allowed@example.com", + }; + const allowedIdentity = "email:allowed@example.com"; + vi.spyOn(brokerClient, "refreshCredential").mockResolvedValue({ + entry: { + id: 1, + provider: "anthropic", + credential: { + ...allowedCredential, + access: "excluded-access", + accountId: "account-excluded", + email: "excluded@example.com", + }, + identityKey: "email:excluded@example.com", + }, + }); + const remoteStore = new RemoteAuthCredentialStore({ + client: brokerClient, + streamSnapshots: false, + accountPool: new Map([["anthropic", new Set([allowedIdentity])]]), + initialSnapshot: { + generation: 1, + generatedAt: now, + serverNowMs: now, + refresher: { enabled: false, intervalMs: 0, skewMs: 0, nextSweepInMs: Number.MAX_SAFE_INTEGER }, + credentials: [ + { + id: 1, + provider: "anthropic", + credential: allowedCredential, + identityKey: allowedIdentity, + rotatesInMs: null, + }, + ], + }, + }); + try { + await expect(remoteStore.refreshOAuthCredential("anthropic", 1, allowedCredential)).rejects.toThrow( + "outside the configured account pool", + ); + expect(remoteStore.snapshot.credentials).toEqual([]); + } finally { + remoteStore.close(); + } + }); }); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index f8e3d7aa9..6d0abcc25 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -39,6 +39,7 @@ - Pinned displaceable transcript snapshots (`hub` waiting polls and live `todo` lists) to the viewport like the `vibe_wait` wall: when the transcript outgrows the terminal, their still-mutating rows are no longer committed to native scrollback on every spinner tick, which previously spammed hundreds of duplicated "waiting on N jobs" rows and force-sealed the poll so follow-up polls stacked instead of replacing it. - Fixed legacy pi extensions failing to load on npm/source-link installs when their module graph contained a transitive CommonJS dependency (`Missing graph-owned CommonJS definition`). The `@(scope)/pi-coding-agent` root shim is served from `src/` on source-link installs, so an extension's import evaluated a second instance of `legacy-pi-compat.ts` whose top-level global registration clobbered the host bundle's populated CommonJS graph bridge with empty state; the registration is now first-wins so the host bridge survives ([#6449](https://github.com/can1357/oh-my-pi/issues/6449)). +- Fixed `omp auth-gateway serve` and `omp auth-gateway check` bypassing the process-scoped OAuth account pool configured by `OMP_AUTH_BROKER_ACCOUNT_POOL_FILE`. - Fixed a first-use race in `ArtifactManager` where two concurrent `allocatePath`/`save` callers on a fresh instance both re-seeded `#nextId` across the directory-scan yield and allocated the same artifact id, silently overwriting the first artifact (same tool type) or making `artifact://` resolution ambiguous (different tool types). The initial scan is now memoized as a single in-flight promise so all concurrent callers share one initialization and receive distinct ids ([#4091](https://github.com/can1357/oh-my-pi/issues/4091)). - Fixed blob reference resolution passing unvalidated `blob:sha256:` suffixes into `path.join`, allowing a crafted ref (e.g. `blob:sha256:../../../etc/passwd`) in a persisted/shared session to escape the blob directory and read arbitrary files into resolved image history; `parseBlobRef` now rejects any suffix that is not a canonical 64-char lowercase hex hash, gating every resolution path ([#4088](https://github.com/can1357/oh-my-pi/issues/4088)). - Fixed pressing Esc in the `/omfg` amendment input discarding the whole generated-rule flow; it now returns to the save destination selector with the candidate preserved. diff --git a/packages/coding-agent/src/cli/auth-gateway-cli.ts b/packages/coding-agent/src/cli/auth-gateway-cli.ts index 73776916c..99947afa0 100644 --- a/packages/coding-agent/src/cli/auth-gateway-cli.ts +++ b/packages/coding-agent/src/cli/auth-gateway-cli.ts @@ -24,7 +24,12 @@ import { completeSimple, type Model, } from "@oh-my-pi/pi-ai"; -import { AuthBrokerClient, RemoteAuthCredentialStore, type SnapshotResponse } from "@oh-my-pi/pi-ai/auth-broker"; +import { + AuthBrokerClient, + loadAuthBrokerAccountPool, + RemoteAuthCredentialStore, + type SnapshotResponse, +} from "@oh-my-pi/pi-ai/auth-broker"; import { DEFAULT_AUTH_GATEWAY_BIND, startAuthGateway } from "@oh-my-pi/pi-ai/auth-gateway"; import { type GeneratedProvider, getBundledModels, getBundledProviders } from "@oh-my-pi/pi-catalog/models"; import { getConfigRootDir, isEnoent, VERSION } from "@oh-my-pi/pi-utils"; @@ -147,9 +152,14 @@ async function runServe(flags: AuthGatewayCommandArgs["flags"]): Promise { // Build a broker-backed AuthStorage — same pattern as discoverAuthStorage() // in sdk.ts. The gateway never touches local SQLite. + const accountPool = await loadAuthBrokerAccountPool(); const client = createBrokerClient(brokerConfig); const initialSnapshot = await fetchBrokerSnapshot(client); - const store = new RemoteAuthCredentialStore({ client, initialSnapshot }); + const store = new RemoteAuthCredentialStore({ + client, + initialSnapshot, + accountPool, + }); // Refresh + usage both flow through the store's broker hooks automatically — // `RemoteAuthCredentialStore.refreshOAuthCredential` and `.fetchUsageReports`. // AuthStorage discovers them when no explicit option overrides them, so the @@ -538,9 +548,14 @@ async function runCheck(flags: AuthGatewayCommandArgs["flags"]): Promise { ); } + const accountPool = await loadAuthBrokerAccountPool(); const client = createBrokerClient(brokerConfig); const initialSnapshot = await fetchBrokerSnapshot(client); - const store = new RemoteAuthCredentialStore({ client, initialSnapshot }); + const store = new RemoteAuthCredentialStore({ + client, + initialSnapshot, + accountPool, + }); const storage = new AuthStorage(store, { sourceLabel: `broker ${brokerConfig.url}` }); try { await storage.reload(); diff --git a/packages/coding-agent/test/auth-gateway-account-pool.test.ts b/packages/coding-agent/test/auth-gateway-account-pool.test.ts new file mode 100644 index 000000000..084809767 --- /dev/null +++ b/packages/coding-agent/test/auth-gateway-account-pool.test.ts @@ -0,0 +1,76 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; +import { AuthStorage, SqliteAuthCredentialStore } from "@oh-my-pi/pi-ai"; +import { type AuthBrokerServerHandle, startAuthBroker } from "@oh-my-pi/pi-ai/auth-broker"; +import { runAuthGatewayCommand } from "@oh-my-pi/pi-coding-agent/cli/auth-gateway-cli"; +import { removeWithRetries } from "@oh-my-pi/pi-utils"; + +const BROKER_TOKEN = "gateway-account-pool-token"; +const ENV_KEYS = ["OMP_AUTH_BROKER_URL", "OMP_AUTH_BROKER_TOKEN", "OMP_AUTH_BROKER_ACCOUNT_POOL_FILE"] as const; + +describe("auth-gateway account pool", () => { + let tempDir = ""; + let brokerStore: SqliteAuthCredentialStore | undefined; + let brokerStorage: AuthStorage | undefined; + let handle: AuthBrokerServerHandle | undefined; + let savedEnv: Record<(typeof ENV_KEYS)[number], string | undefined>; + + beforeEach(async () => { + savedEnv = Object.fromEntries(ENV_KEYS.map(key => [key, process.env[key]])) as typeof savedEnv; + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "omp-auth-gateway-pool-")); + brokerStore = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); + brokerStore.saveOAuth("anthropic", { + access: "allowed-access", + refresh: "allowed-refresh", + expires: Date.now() + 120_000, + email: "allowed@example.com", + }); + brokerStore.saveOAuth("anthropic", { + access: "excluded-access", + refresh: "excluded-refresh", + expires: Date.now() + 120_000, + email: "excluded@example.com", + }); + brokerStorage = new AuthStorage(brokerStore); + await brokerStorage.reload(); + handle = startAuthBroker({ + storage: brokerStorage, + bind: "127.0.0.1:0", + bearerTokens: [BROKER_TOKEN], + disableRefresher: true, + }); + const poolPath = path.join(tempDir, "account-pool.json"); + await Bun.write(poolPath, JSON.stringify({ anthropic: ["email:allowed@example.com"] })); + process.env.OMP_AUTH_BROKER_URL = handle.url; + process.env.OMP_AUTH_BROKER_TOKEN = BROKER_TOKEN; + process.env.OMP_AUTH_BROKER_ACCOUNT_POOL_FILE = poolPath; + }); + + afterEach(async () => { + vi.restoreAllMocks(); + await handle?.close(); + brokerStorage?.close(); + brokerStore?.close(); + if (tempDir) await removeWithRetries(tempDir); + for (const key of ENV_KEYS) { + const value = savedEnv[key]; + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } + }); + + test("check probes only credentials selected by the environment pool", async () => { + let output = ""; + vi.spyOn(process.stdout, "write").mockImplementation(chunk => { + output += typeof chunk === "string" ? chunk : new TextDecoder().decode(chunk); + return true; + }); + + await runAuthGatewayCommand({ action: "check", flags: { json: true } }); + + const result = JSON.parse(output) as { credentials: Array<{ email?: string }> }; + expect(result.credentials.map(credential => credential.email)).toEqual(["allowed@example.com"]); + }); +});