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.
This commit is contained in:
can1357
@@ -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 `<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. |
|
||||
|
||||
Resolution order in `resolveAuthBrokerConfig()`:
|
||||
|
||||
|
||||
@@ -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 `<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. |
|
||||
|
||||
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`.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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 () => {
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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<OAuthAccess | undefined>;
|
||||
getOAuthAccess(
|
||||
provider: string,
|
||||
sessionId?: string,
|
||||
options?: DryBalanceAuthOptions,
|
||||
): Promise<OAuthAccess | undefined>;
|
||||
}
|
||||
|
||||
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<Api>,
|
||||
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<T, R>(items: T[], concurrency: number, fn: (item: T
|
||||
return results;
|
||||
}
|
||||
|
||||
|
||||
function sortedStats(map: Map<string, number>, samples: number): Array<{ label: string; count: number; percent: number }> {
|
||||
function sortedStats(
|
||||
map: Map<string, number>,
|
||||
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<DryBalanceSummary> {
|
||||
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;
|
||||
|
||||
@@ -106,8 +106,8 @@ import {
|
||||
AuthBrokerClient,
|
||||
AuthStorage,
|
||||
DEFAULT_SNAPSHOT_CACHE_TTL_MS,
|
||||
readAuthBrokerSnapshotCache,
|
||||
RemoteAuthCredentialStore,
|
||||
readAuthBrokerSnapshotCache,
|
||||
type SnapshotResponse,
|
||||
writeAuthBrokerSnapshotCache,
|
||||
} from "./session/auth-storage";
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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<Record<(typeof ENV_KEYS)[number], string | undefined>> = {};
|
||||
|
||||
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<boolean>, timeoutMs = 2_000): Promise<void> {
|
||||
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();
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user