From b8a602ac2a770b3d3a145efa81485a7046329a13 Mon Sep 17 00:00:00 2001 From: can1357 Date: Fri, 5 Jun 2026 11:13:44 +0200 Subject: [PATCH] test(auth): switched OAuth ranking tests to weighted selection - Replaced top-rank assertions with weighted-preference distribution checks. - Added cases for equal-priority balancing and 2x best-bucket cap. - Added coding-agent snapshot-cache boot and seed tests. --- docs/auth-broker-gateway.md | 10 ++ docs/environment-variables.md | 2 + packages/ai/CHANGELOG.md | 8 + packages/ai/src/auth-broker/index.ts | 2 +- .../ai/test/auth-broker-remote-store.test.ts | 1 - .../test/auth-broker-snapshot-cache.test.ts | 6 +- .../test/auth-storage-codex-selection.test.ts | 95 +++++++++--- packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/cli/dry-balance-cli.ts | 40 +++-- packages/coding-agent/src/sdk.ts | 2 +- .../coding-agent/src/session/auth-storage.ts | 2 +- .../test/auth-broker-snapshot-cache.test.ts | 140 ++++++++++++++++++ packages/utils/CHANGELOG.md | 4 + 13 files changed, 273 insertions(+), 40 deletions(-) create mode 100644 packages/coding-agent/test/auth-broker-snapshot-cache.test.ts diff --git a/docs/auth-broker-gateway.md b/docs/auth-broker-gateway.md index ec7948d1b..1fccca164 100644 --- a/docs/auth-broker-gateway.md +++ b/docs/auth-broker-gateway.md @@ -140,6 +140,14 @@ When the gateway (or any other broker client) calls `fetchUsageReports()` / `get The 15 s client window deliberately sits below the broker’s 5 min server cache, so almost every client poll is served from the broker’s already-cached value; the client cache exists to absorb the parallel fan-out generated by `AuthStorage.#rankOAuthSelections` into a single broker round-trip. +## Client snapshot cache + +`discoverAuthStorage()` persists the broker snapshot to `~/.omp/cache/auth-broker-snapshot.enc` after the initial `/v1/snapshot` fetch and after later broker-sourced full snapshots. The file is AES-256-GCM encrypted with `SHA-256(OMP_AUTH_BROKER_TOKEN)` and authenticated with the broker URL as additional data, so changing either the token or URL makes the cache unreadable. The file is written atomically with mode `0600`. + +Freshness is anchored to the broker-stamped `snapshot.generatedAt`, not local write time. Default TTL is 1 h (`OMP_AUTH_BROKER_SNAPSHOT_TTL_MS`); `0` disables the cache and restores the old always-fetch boot path. When the cached snapshot is still fresh, `omp` boots from it and skips the blocking `/v1/snapshot` query. `RemoteAuthCredentialStore` still starts its normal SSE / long-poll background sync immediately, so deleted or rotated credentials reconcile after startup, and expired OAuth access tokens still refresh through `POST /v1/credential/:id/refresh`. + +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. + ## 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. @@ -150,6 +158,8 @@ The broker is **off** unless `OMP_AUTH_BROKER_URL` (or `auth.broker.url` in `con | ----------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------- | | `OMP_AUTH_BROKER_URL` | Base URL of the remote auth-broker (e.g. `https://broker.tailnet:8765`). Selecting this puts the client in broker mode — local SQLite is bypassed. | Any time the omp client should resolve credentials through a broker (and required by `omp auth-gateway serve`). | | `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. | Resolution order in `resolveAuthBrokerConfig()`: diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 131d9170c..e556301e8 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -97,6 +97,8 @@ When the broker is enabled, the local SQLite credential store is bypassed and al | ----------------------- | -------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `OMP_AUTH_BROKER_URL` | Base URL of the remote auth-broker (e.g. `https://broker.tailnet:8765`); selects broker mode | Resolving credentials through a broker; also required by `omp auth-gateway serve` (the gateway is itself a broker client) | Wins over `auth.broker.url` in `config.yml`. When set with no resolvable token, `resolveAuthBrokerConfig()` hard-errors instead of falling back to local SQLite. | | `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. | 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 2789f64c8..59f2033dd 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,14 @@ ## [Unreleased] +### Added + +- Added an AES-256-GCM auth-broker snapshot cache module and `RemoteAuthCredentialStoreOptions.onSnapshot` so broker clients can persist broker-sourced full snapshots without blocking startup on every run. + +### Changed + +- Changed usage-ranked OAuth credential selection to pick deterministic session-sticky weighted buckets instead of always choosing the top-ranked account, capping the best account at 2x the baseline session likelihood while keeping equal-priority accounts evenly balanced. + ## [15.9.1] - 2026-06-04 ### Added diff --git a/packages/ai/src/auth-broker/index.ts b/packages/ai/src/auth-broker/index.ts index a2f31df85..189d377c5 100644 --- a/packages/ai/src/auth-broker/index.ts +++ b/packages/ai/src/auth-broker/index.ts @@ -1,6 +1,6 @@ export * from "./client"; export * from "./refresher"; export * from "./remote-store"; -export * from "./snapshot-cache"; export * from "./server"; +export * from "./snapshot-cache"; export * from "./types"; diff --git a/packages/ai/test/auth-broker-remote-store.test.ts b/packages/ai/test/auth-broker-remote-store.test.ts index ee288a0bb..50be18938 100644 --- a/packages/ai/test/auth-broker-remote-store.test.ts +++ b/packages/ai/test/auth-broker-remote-store.test.ts @@ -126,7 +126,6 @@ describe("RemoteAuthCredentialStore SSE integration", () => { }); expect(callbacks).toHaveLength(0); - storage!.upsertCredential("anthropic", mintOAuthCredential("callback", Date.now() + 120_000)); const refreshed = await remote.refreshSnapshot(); expect(callbacks).toHaveLength(1); diff --git a/packages/ai/test/auth-broker-snapshot-cache.test.ts b/packages/ai/test/auth-broker-snapshot-cache.test.ts index 58aa20b72..9480472c6 100644 --- a/packages/ai/test/auth-broker-snapshot-cache.test.ts +++ b/packages/ai/test/auth-broker-snapshot-cache.test.ts @@ -2,11 +2,7 @@ import { 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 { - readAuthBrokerSnapshotCache, - type SnapshotResponse, - writeAuthBrokerSnapshotCache, -} from "../src"; +import { readAuthBrokerSnapshotCache, type SnapshotResponse, writeAuthBrokerSnapshotCache } from "../src"; const TOKEN = "broker-cache-token"; const URL = "http://127.0.0.1:8765"; diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index b1eb5b63b..7e23abae9 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -173,7 +173,7 @@ describe("AuthStorage codex oauth ranking", () => { } }); - test("prefers near-reset weekly account over lower-used far-reset account", async () => { + test("weights near-reset weekly account over lower-used far-reset account", async () => { if (!authStorage) throw new Error("test setup failed"); await authStorage.set("openai-codex", [ @@ -198,11 +198,11 @@ describe("AuthStorage codex oauth ranking", () => { }), ); - const apiKey = await authStorage.getApiKey("openai-codex", "session-weekly-reset"); - expect(apiKey).toBe("api-acct-near"); + const counts = await countApiKeySelections(authStorage, "openai-codex", "weighted-codex-near"); + expectWeightedPreference(counts, "api-acct-near", "api-acct-far"); }); - test("prioritizes fresh 5h ticker account at 0% usage", async () => { + test("weights fresh 5h ticker account at 0% usage", async () => { if (!authStorage) throw new Error("test setup failed"); await authStorage.set("openai-codex", [ @@ -235,8 +235,8 @@ describe("AuthStorage codex oauth ranking", () => { }), ); - const apiKey = await authStorage.getApiKey("openai-codex", "session-five-hour-start"); - expect(apiKey).toBe("api-acct-zero"); + const counts = await countApiKeySelections(authStorage, "openai-codex", "weighted-codex-zero"); + expectWeightedPreference(counts, "api-acct-zero", "api-acct-progress"); }); test("skips exhausted weekly account even when reset is near", async () => { if (!authStorage) throw new Error("test setup failed"); @@ -426,7 +426,7 @@ describe("AuthStorage codex oauth ranking", () => { expect(elapsedMs).toBeLessThan(1_000); }); - test("sorts 3 accounts by weekly drain rate", async () => { + test("weights 3 accounts by weekly drain rate", async () => { if (!authStorage) throw new Error("test setup failed"); await authStorage.set("openai-codex", [ @@ -460,8 +460,9 @@ describe("AuthStorage codex oauth ranking", () => { }), ); - const apiKey = await authStorage.getApiKey("openai-codex", "session-three-accounts"); - expect(apiKey).toBe("api-acct-slow"); + const counts = await countApiKeySelections(authStorage, "openai-codex", "weighted-codex-three"); + expect(countFor(counts, "api-acct-slow")).toBeGreaterThan(countFor(counts, "api-acct-medium")); + expect(countFor(counts, "api-acct-slow")).toBeGreaterThan(countFor(counts, "api-acct-fast")); }); test("handles usage fetch failure gracefully (null report)", async () => { @@ -482,8 +483,8 @@ describe("AuthStorage codex oauth ranking", () => { }), ); - const apiKey = await authStorage.getApiKey("openai-codex", "session-null-usage"); - expect(apiKey).toBe("api-acct-known"); + const counts = await countApiKeySelections(authStorage, "openai-codex", "weighted-codex-known", 300); + expectWeightedPreference(counts, "api-acct-known", "api-acct-null"); }); test("refreshes expired oauth candidates in parallel before selection", async () => { if (!authStorage) throw new Error("test setup failed"); @@ -650,7 +651,7 @@ describe("AuthStorage claude oauth ranking", () => { } }); - test("prefers lower secondary drain rate account", async () => { + test("weights lower secondary drain rate account", async () => { if (!authStorage) throw new Error("test setup failed"); await authStorage.set("anthropic", [ @@ -675,8 +676,67 @@ describe("AuthStorage claude oauth ranking", () => { }), ); - const apiKey = await authStorage.getApiKey("anthropic", "session-claude-drain"); - expect(apiKey).toBe("api-acct-near"); + const counts = await countApiKeySelections(authStorage, "anthropic", "weighted-claude-near"); + expectWeightedPreference(counts, "api-acct-near", "api-acct-far"); + }); + + test("balances equal-priority accounts evenly", async () => { + if (!authStorage) throw new Error("test setup failed"); + + await authStorage.set("anthropic", [ + { type: "oauth", ...createCredential("acct-a", "a@example.com") }, + { type: "oauth", ...createCredential("acct-b", "b@example.com") }, + ]); + + for (const accountId of ["acct-a", "acct-b"]) { + usageByAccount.set( + accountId, + createClaudeUsageReport({ + accountId, + primary: { usedFraction: 0.25, resetInMs: 4 * HOUR_MS }, + secondary: { usedFraction: 0.25, resetInMs: 4 * 24 * HOUR_MS }, + }), + ); + } + + const counts = await countApiKeySelections(authStorage, "anthropic", "weighted-claude-equal", 200); + expect(Math.abs(countFor(counts, "api-acct-a") - countFor(counts, "api-acct-b"))).toBeLessThanOrEqual(25); + }); + + test("caps the strongest priority bucket at about 2x baseline weight", async () => { + if (!authStorage) throw new Error("test setup failed"); + + await authStorage.set("anthropic", [ + { type: "oauth", ...createCredential("acct-best", "best@example.com") }, + { type: "oauth", ...createCredential("acct-base-a", "base-a@example.com") }, + { type: "oauth", ...createCredential("acct-base-b", "base-b@example.com") }, + ]); + + usageByAccount.set( + "acct-best", + createClaudeUsageReport({ + accountId: "acct-best", + primary: { usedFraction: 0.05, resetInMs: 4 * HOUR_MS }, + secondary: { usedFraction: 0.05, resetInMs: 6 * 24 * HOUR_MS }, + }), + ); + for (const accountId of ["acct-base-a", "acct-base-b"]) { + usageByAccount.set( + accountId, + createClaudeUsageReport({ + accountId, + primary: { usedFraction: 0.7, resetInMs: 2 * HOUR_MS }, + secondary: { usedFraction: 0.7, resetInMs: 2 * 24 * HOUR_MS }, + }), + ); + } + + const counts = await countApiKeySelections(authStorage, "anthropic", "claude-cap", 300); + expectWeightedPreference(counts, "api-acct-best", "api-acct-base-a"); + expectWeightedPreference(counts, "api-acct-best", "api-acct-base-b"); + expect(Math.abs(countFor(counts, "api-acct-base-a") - countFor(counts, "api-acct-base-b"))).toBeLessThanOrEqual( + 15, + ); }); test("skips exhausted account and picks healthy", async () => { @@ -737,7 +797,7 @@ describe("AuthStorage claude oauth ranking", () => { expect(apiKey).toBe("api-acct-soon"); }); - test("sorts 3 accounts by secondary drain rate", async () => { + test("weights 3 accounts by secondary drain rate", async () => { if (!authStorage) throw new Error("test setup failed"); await authStorage.set("anthropic", [ @@ -771,8 +831,9 @@ describe("AuthStorage claude oauth ranking", () => { }), ); - const apiKey = await authStorage.getApiKey("anthropic", "session-claude-three"); - expect(apiKey).toBe("api-acct-slow"); + const counts = await countApiKeySelections(authStorage, "anthropic", "weighted-claude-three"); + expect(countFor(counts, "api-acct-slow")).toBeGreaterThan(countFor(counts, "api-acct-medium")); + expect(countFor(counts, "api-acct-slow")).toBeGreaterThan(countFor(counts, "api-acct-fast")); }); test("single credential works without ranking", async () => { diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 18f2b7ddc..792cb455e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -3,6 +3,7 @@ ## [Unreleased] ### Added +- Added an encrypted local auth-broker snapshot cache for `discoverAuthStorage`, with `OMP_AUTH_BROKER_SNAPSHOT_TTL_MS` and `OMP_AUTH_BROKER_SNAPSHOT_CACHE`, so fresh cached broker credentials can boot without a blocking `/v1/snapshot` fetch and survive broker-down startup windows. - Added `dry-balance` CLI command to perform a dry-run OAuth account balancing check across configurable random session IDs, with sample and concurrency options, JSON output, and success/failure summary reporting - Added `--json` output mode and machine-readable result format to `omp dry-balance` for automated use diff --git a/packages/coding-agent/src/cli/dry-balance-cli.ts b/packages/coding-agent/src/cli/dry-balance-cli.ts index 03249c705..1e16d5237 100644 --- a/packages/coding-agent/src/cli/dry-balance-cli.ts +++ b/packages/coding-agent/src/cli/dry-balance-cli.ts @@ -4,10 +4,10 @@ import chalk from "chalk"; import { ModelRegistry } from "../config/model-registry"; import { formatModelString, + type ModelMatchPreferences, resolveAllowedModels, resolveCliModel, resolveModelRoleValue, - type ModelMatchPreferences, } from "../config/model-resolver"; import { Settings } from "../config/settings"; import { discoverAuthStorage } from "../sdk"; @@ -32,7 +32,11 @@ export interface DryBalanceAuthOptions { } export interface DryBalanceAuthStorage { - getOAuthAccess(provider: string, sessionId?: string, options?: DryBalanceAuthOptions): Promise; + getOAuthAccess( + provider: string, + sessionId?: string, + options?: DryBalanceAuthOptions, + ): Promise; } export interface DryBalanceModelRegistry { @@ -142,7 +146,9 @@ async function resolveDryBalanceModel( const allowedModels = await resolveAllowedModels(modelRegistry, settings, preferences); if (allowedModels.length === 0) { - throw new Error("No models available. Use --model to select a model or configure enabledModels/default model settings."); + throw new Error( + "No models available. Use --model to select a model or configure enabledModels/default model settings.", + ); } const defaultRoleSpec = resolveModelRoleValue(settings?.getModelRole("default"), allowedModels, { @@ -161,12 +167,11 @@ async function resolveDryBalanceModel( return { model: allowedModels[0], - warning: "No allowed model had usable credentials during default resolution; dry-balance will report OAuth failures for the first allowed model.", + warning: + "No allowed model had usable credentials during default resolution; dry-balance will report OAuth failures for the first allowed model.", }; } - - async function runOneAttempt( model: Model, modelRegistry: DryBalanceModelRegistry, @@ -181,7 +186,8 @@ async function runOneAttempt( modelId: model.id, }); if (!access) return { ok: false, reason: "no OAuth access resolved" }; - const account = access.email ?? access.accountId ?? access.projectId ?? access.enterpriseUrl ?? "(unknown oauth account)"; + const account = + access.email ?? access.accountId ?? access.projectId ?? access.enterpriseUrl ?? "(unknown oauth account)"; return { ok: true, account }; } catch (error) { return { ok: false, reason: error instanceof Error ? error.message : String(error) }; @@ -205,8 +211,10 @@ async function mapConcurrent(items: T[], concurrency: number, fn: (item: T return results; } - -function sortedStats(map: Map, samples: number): Array<{ label: string; count: number; percent: number }> { +function sortedStats( + map: Map, + samples: number, +): Array<{ label: string; count: number; percent: number }> { return [...map.entries()] .map(([label, count]) => ({ label, count, percent: (count / samples) * 100 })) .sort((left, right) => right.count - left.count || left.label.localeCompare(right.label)); @@ -253,7 +261,6 @@ function summarizeResults( }; } - function formatRows(rows: Array<{ count: number; percent: number; label: string }>): string[] { if (rows.length === 0) return [` ${chalk.dim("(none)")}`]; const maxCountWidth = Math.max(...rows.map(row => row.count.toString().length)); @@ -296,13 +303,18 @@ export async function runDryBalanceCommand( deps: DryBalanceDependencies = {}, ): Promise { const samples = normalizePositiveInteger("count", command.flags.count, DEFAULT_SAMPLE_COUNT); - const concurrency = Math.min(samples, normalizePositiveInteger("concurrency", command.flags.concurrency, DEFAULT_CONCURRENCY)); + const concurrency = Math.min( + samples, + normalizePositiveInteger("concurrency", command.flags.concurrency, DEFAULT_CONCURRENCY), + ); const randomSessionId = deps.randomSessionId ?? (() => Bun.randomUUIDv7()); const writeStdout = deps.writeStdout ?? ((text: string) => process.stdout.write(text)); const writeStderr = deps.writeStderr ?? ((text: string) => process.stderr.write(text)); - const setExitCode = deps.setExitCode ?? ((code: number) => { - process.exitCode = code; - }); + const setExitCode = + deps.setExitCode ?? + ((code: number) => { + process.exitCode = code; + }); const runtime = await (deps.createRuntime ?? createDefaultRuntime)(); try { const modelSelector = command.flags.model ?? command.model; diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 9ae14fd17..63a82447f 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -106,8 +106,8 @@ import { AuthBrokerClient, AuthStorage, DEFAULT_SNAPSHOT_CACHE_TTL_MS, - readAuthBrokerSnapshotCache, RemoteAuthCredentialStore, + readAuthBrokerSnapshotCache, type SnapshotResponse, writeAuthBrokerSnapshotCache, } from "./session/auth-storage"; diff --git a/packages/coding-agent/src/session/auth-storage.ts b/packages/coding-agent/src/session/auth-storage.ts index 8a9b14fad..d3486d7e2 100644 --- a/packages/coding-agent/src/session/auth-storage.ts +++ b/packages/coding-agent/src/session/auth-storage.ts @@ -19,9 +19,9 @@ export { AuthBrokerClient, AuthStorage, DEFAULT_SNAPSHOT_CACHE_TTL_MS, - readAuthBrokerSnapshotCache, REMOTE_REFRESH_SENTINEL, RemoteAuthCredentialStore, + readAuthBrokerSnapshotCache, SqliteAuthCredentialStore, writeAuthBrokerSnapshotCache, } from "@oh-my-pi/pi-ai"; diff --git a/packages/coding-agent/test/auth-broker-snapshot-cache.test.ts b/packages/coding-agent/test/auth-broker-snapshot-cache.test.ts new file mode 100644 index 000000000..7dba4027c --- /dev/null +++ b/packages/coding-agent/test/auth-broker-snapshot-cache.test.ts @@ -0,0 +1,140 @@ +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 { type AuthBrokerServerHandle, AuthStorage, SqliteAuthCredentialStore, startAuthBroker } from "@oh-my-pi/pi-ai"; +import { discoverAuthStorage } from "../src/sdk"; +import { + readAuthBrokerSnapshotCache, + type SnapshotResponse, + writeAuthBrokerSnapshotCache, +} from "../src/session/auth-storage"; + +const ENV_KEYS = [ + "OMP_AUTH_BROKER_URL", + "OMP_AUTH_BROKER_TOKEN", + "OMP_AUTH_BROKER_SNAPSHOT_CACHE", + "OMP_AUTH_BROKER_SNAPSHOT_TTL_MS", +] as const; +const PROVIDER = "unit-auth-broker-cache"; +const TOKEN = "coding-agent-cache-token"; + +const savedEnv: Partial> = {}; + +function makeSnapshot(urlTime: number): SnapshotResponse { + return { + generation: 11, + generatedAt: urlTime, + serverNowMs: urlTime, + refresher: { + enabled: false, + intervalMs: 60_000, + skewMs: 300_000, + nextSweepInMs: Number.MAX_SAFE_INTEGER, + }, + credentials: [ + { + id: 1, + provider: PROVIDER, + credential: { type: "api_key", key: "cached-api-key" }, + identityKey: null, + rotatesInMs: null, + }, + ], + }; +} + +async function waitUntil(predicate: () => boolean | Promise, timeoutMs = 2_000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (await predicate()) return; + await Bun.sleep(10); + } + if (!(await predicate())) throw new Error("waitUntil timeout"); +} + +describe("discoverAuthStorage auth-broker snapshot cache", () => { + let tempDir = ""; + + beforeEach(async () => { + for (const key of ENV_KEYS) savedEnv[key] = process.env[key]; + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "coding-agent-auth-broker-cache-")); + }); + + afterEach(async () => { + for (const key of ENV_KEYS) { + if (savedEnv[key] === undefined) delete process.env[key]; + else process.env[key] = savedEnv[key]; + } + await fs.rm(tempDir, { recursive: true, force: true }); + }); + + test("boots from a fresh encrypted cache when the broker is down", async () => { + const cachePath = path.join(tempDir, "snapshot.enc"); + const downUrl = "http://127.0.0.1:1"; + process.env.OMP_AUTH_BROKER_URL = downUrl; + process.env.OMP_AUTH_BROKER_TOKEN = TOKEN; + process.env.OMP_AUTH_BROKER_SNAPSHOT_CACHE = cachePath; + process.env.OMP_AUTH_BROKER_SNAPSHOT_TTL_MS = "3600000"; + await writeAuthBrokerSnapshotCache({ + path: cachePath, + token: TOKEN, + url: downUrl, + snapshot: makeSnapshot(Date.now()), + }); + + const storage = await discoverAuthStorage(tempDir); + try { + expect(await storage.getApiKey(PROVIDER)).toBe("cached-api-key"); + } finally { + storage.close(); + } + }); + + test("seeds the encrypted cache after an initial broker fetch", async () => { + const cachePath = path.join(tempDir, "snapshot.enc"); + const brokerStore = await SqliteAuthCredentialStore.open(path.join(tempDir, "broker.db")); + brokerStore.saveApiKey(PROVIDER, "broker-api-key"); + const brokerStorage = new AuthStorage(brokerStore); + await brokerStorage.reload(); + let handle: AuthBrokerServerHandle | undefined; + let storage: AuthStorage | undefined; + try { + handle = startAuthBroker({ + storage: brokerStorage, + bind: "127.0.0.1:0", + bearerTokens: [TOKEN], + disableRefresher: true, + }); + process.env.OMP_AUTH_BROKER_URL = handle.url; + process.env.OMP_AUTH_BROKER_TOKEN = TOKEN; + process.env.OMP_AUTH_BROKER_SNAPSHOT_CACHE = cachePath; + process.env.OMP_AUTH_BROKER_SNAPSHOT_TTL_MS = "3600000"; + + storage = await discoverAuthStorage(tempDir); + expect(await storage.getApiKey(PROVIDER)).toBe("broker-api-key"); + await waitUntil(async () => { + const cached = await readAuthBrokerSnapshotCache({ + path: cachePath, + token: TOKEN, + url: handle!.url, + ttlMs: 3_600_000, + }); + return cached?.credentials.some(entry => entry.provider === PROVIDER) ?? false; + }); + const cached = await readAuthBrokerSnapshotCache({ + path: cachePath, + token: TOKEN, + url: handle.url, + ttlMs: 3_600_000, + }); + const entry = cached?.credentials.find(candidate => candidate.provider === PROVIDER); + expect(entry?.credential).toEqual({ type: "api_key", key: "broker-api-key" }); + } finally { + storage?.close(); + await handle?.close(); + brokerStorage.close(); + brokerStore.close(); + } + }); +}); diff --git a/packages/utils/CHANGELOG.md b/packages/utils/CHANGELOG.md index a33a5cd87..dc16a8eeb 100644 --- a/packages/utils/CHANGELOG.md +++ b/packages/utils/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `getAuthBrokerSnapshotCachePath()` with `OMP_AUTH_BROKER_SNAPSHOT_CACHE` override support for isolating the encrypted broker snapshot cache. + ## [15.9.1] - 2026-06-04 ### Fixed