feat(ai): add process-scoped OAuth account pools

This commit is contained in:
Alex TYRODE
2026-07-23 17:54:35 +00:00
parent c64e7146e0
commit faabe089e8
8 changed files with 508 additions and 15 deletions
+25
View File
@@ -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. 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 `<config-dir>/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()`:
+1
View File
@@ -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 `<config-dir>/auth-broker.token` | Resolution: this env → `auth.broker.token` (`$ENV_NAME` indirection supported) → `<config-dir>/auth-broker.token` (mode `0600`). `<config-dir>` 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 `<config-dir>/auth-gateway.token` and is managed via `omp auth-gateway token`.
+1
View File
@@ -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 while leaving API keys and the shared encrypted snapshot cache unrestricted.
### Fixed
+45 -1
View File
@@ -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<string | undefined>;
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,46 @@ async function readConfigYaml(agentDir: string): Promise<ConfigSnapshot> {
return {};
}
async function readAuthBrokerAccountPool(): Promise<AuthBrokerAccountPool | undefined> {
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<string, ReadonlySet<string>>();
for (const [provider, value] of Object.entries(parsed)) {
if (provider.trim().length === 0) {
throw new AIError.ConfigurationError("OMP_AUTH_BROKER_ACCOUNT_POOL_FILE contains an empty provider id");
}
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<string>();
for (const identity of value) {
if (typeof identity !== "string" || identity.trim().length === 0) {
throw new AIError.ConfigurationError(
`OMP_AUTH_BROKER_ACCOUNT_POOL_FILE entry for ${provider} contains an invalid identity key`,
);
}
identities.add(identity.trim());
}
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 +225,7 @@ export async function discoverAuthStorage(options: DiscoverAuthStorageOptions =
});
if (brokerConfig) {
const accountPool = options.accountPool ?? (await readAuthBrokerAccountPool());
const client = new AuthBrokerClient({ url: brokerConfig.url, token: brokerConfig.token });
const cachePath = options.cachePath ?? getAuthBrokerSnapshotCachePath();
const ttlMs = resolveSnapshotTtlMs();
@@ -226,6 +269,7 @@ export async function discoverAuthStorage(options: DiscoverAuthStorageOptions =
client,
initialSnapshot,
onSnapshot: persist,
accountPool,
});
const storage = new AuthStorage(store, {
configValueResolver: options.configValueResolver,
+94 -13
View File
@@ -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<string, ReadonlySet<string>>;
function isCredentialInAccountPool(
entry: Pick<SnapshotEntry, "provider" | "credential" | "identityKey">,
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<OAuthCredentials> {
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<UsageReport[] | null> {
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,28 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
signal?: AbortSignal,
): Promise<UsageReport | null> {
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;
// Aggregate reports do not identify their source credential type. A
// visible API key therefore makes the provider's whole report set visible.
return this.#snapshot.credentials.some(entry => {
if (entry.provider !== report.provider) return false;
if (entry.credential.type !== "oauth") return true;
return usageReportMatchesCredential(report, entry.credential);
});
});
}
ingestUsageReport(provider: Provider, credential: OAuthCredential, report: UsageReport): boolean {
const key = usageOverlayKey(provider, credential);
if (!key) return false;
@@ -1059,6 +1125,21 @@ function matchUsageReport(reports: UsageReport[], provider: Provider, credential
return null;
}
function usageReportMatchesCredential(report: UsageReport, credential: OAuthCredential): boolean {
const metadata = (report.metadata ?? {}) as Record<string, unknown>;
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 }))
@@ -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,30 @@ 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"],
] as const;
for (const [content, expectedError] of invalidFiles) {
await Bun.write(poolPath, content);
await expect(discoverAuthStorage({ agentDir })).rejects.toThrow(expectedError);
}
});
});
});
@@ -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<Record<(typeof ANTHROPIC_ENV)[number], string | undefined>> = {};
@@ -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();
}
},
);
});
});
+182
View File
@@ -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 preserves usage for a provider 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);
} 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();
}
});
});