fix(ai): close auth broker account pool gaps
This commit is contained in:
@@ -168,7 +168,7 @@ SDK hosts can supply the same provider-to-identity mapping as `accountPool` in `
|
||||
- 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.
|
||||
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.
|
||||
|
||||
|
||||
@@ -5,7 +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.
|
||||
- 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
|
||||
|
||||
|
||||
@@ -114,7 +114,7 @@ async function readConfigYaml(agentDir: string): Promise<ConfigSnapshot> {
|
||||
return {};
|
||||
}
|
||||
|
||||
async function readAuthBrokerAccountPool(): Promise<AuthBrokerAccountPool | undefined> {
|
||||
export async function loadAuthBrokerAccountPool(): Promise<AuthBrokerAccountPool | undefined> {
|
||||
const filePath = process.env.OMP_AUTH_BROKER_ACCOUNT_POOL_FILE?.trim();
|
||||
if (!filePath) return undefined;
|
||||
|
||||
@@ -132,9 +132,15 @@ async function readAuthBrokerAccountPool(): Promise<AuthBrokerAccountPool | unde
|
||||
|
||||
const accountPool = new Map<string, ReadonlySet<string>>();
|
||||
for (const [provider, value] of Object.entries(parsed)) {
|
||||
if (provider.trim().length === 0) {
|
||||
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`,
|
||||
@@ -225,7 +231,7 @@ export async function discoverAuthStorage(options: DiscoverAuthStorageOptions =
|
||||
});
|
||||
|
||||
if (brokerConfig) {
|
||||
const accountPool = options.accountPool ?? (await readAuthBrokerAccountPool());
|
||||
const accountPool = options.accountPool ?? (await loadAuthBrokerAccountPool());
|
||||
const client = new AuthBrokerClient({ url: brokerConfig.url, token: brokerConfig.token });
|
||||
const cachePath = options.cachePath ?? getAuthBrokerSnapshotCachePath();
|
||||
const ttlMs = resolveSnapshotTtlMs();
|
||||
|
||||
@@ -945,13 +945,12 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
|
||||
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);
|
||||
});
|
||||
return this.#snapshot.credentials.some(
|
||||
entry =>
|
||||
entry.provider === report.provider &&
|
||||
entry.credential.type === "oauth" &&
|
||||
usageReportMatchesCredential(report, entry.credential),
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -76,6 +76,7 @@ describe("resolveAuthBrokerConfig config discovery", () => {
|
||||
["[]", "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"],
|
||||
] as const;
|
||||
for (const [content, expectedError] of invalidFiles) {
|
||||
await Bun.write(poolPath, content);
|
||||
|
||||
@@ -1159,7 +1159,7 @@ describe("RemoteAuthCredentialStore + AuthStorage integration", () => {
|
||||
}
|
||||
});
|
||||
|
||||
test("account pool preserves usage for a provider with a visible API key", async () => {
|
||||
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 = {
|
||||
@@ -1214,7 +1214,7 @@ describe("RemoteAuthCredentialStore + AuthStorage integration", () => {
|
||||
},
|
||||
});
|
||||
try {
|
||||
expect(await remoteStore.fetchUsageReports()).toEqual(reports);
|
||||
expect(await remoteStore.fetchUsageReports()).toEqual([reports[0]]);
|
||||
} finally {
|
||||
remoteStore.close();
|
||||
}
|
||||
|
||||
@@ -29,6 +29,7 @@
|
||||
|
||||
### Fixed
|
||||
|
||||
- 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 `error.notify` raising a "Stopped with error" toast for provider failures while an auto-retry or async-delivery continuation was pending; the toast now waits for the true terminal settle.
|
||||
- Fixed terminal `yield` results racing post-turn maintenance, which could trigger an unnecessary automatic handoff or compaction.
|
||||
- Fixed credential-shaped tokens (GitHub/GitLab/OpenAI/Anthropic key patterns) being redacted from outbound provider requests even with `secrets.enabled` off; the pattern redaction now follows the `secrets.enabled` ("Hide Secrets") setting like the secret obfuscator.
|
||||
|
||||
@@ -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<void> {
|
||||
|
||||
// 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<void> {
|
||||
);
|
||||
}
|
||||
|
||||
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();
|
||||
|
||||
@@ -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"]);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user