can1357
3676b5d97f
- Add usage history, reporting, and client summary endpoints to the auth broker. - Implement SQLite persistence, batching, and periodic flushing for observed client usage. - Integrate usage reporting into the coding agent session for completed assistant messages. - Update stats aggregator and dashboard to retrieve usage history snapshots from the broker.
512 lines
18 KiB
TypeScript
512 lines
18 KiB
TypeScript
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, REMOTE_REFRESH_SENTINEL, SqliteAuthCredentialStore } from "@oh-my-pi/pi-ai";
|
|
import {
|
|
AuthBrokerClient,
|
|
type AuthBrokerServerHandle,
|
|
AuthBrokerStreamUnsupportedError,
|
|
type SnapshotStreamEvent,
|
|
startAuthBroker,
|
|
} from "@oh-my-pi/pi-ai/auth-broker";
|
|
import * as oauthUtils from "@oh-my-pi/pi-ai/registry/oauth";
|
|
import { removeWithRetries } from "../../utils/src/temp";
|
|
|
|
const ANTHROPIC_ENV = ["ANTHROPIC_API_KEY", "ANTHROPIC_OAUTH_TOKEN"] as const;
|
|
const savedEnv: Partial<Record<(typeof ANTHROPIC_ENV)[number], string | undefined>> = {};
|
|
|
|
function mintOAuthCredential(suffix: string, expires: number) {
|
|
return {
|
|
type: "oauth" as const,
|
|
access: `access-${suffix}`,
|
|
refresh: `refresh-${suffix}`,
|
|
expires,
|
|
accountId: `account-${suffix}`,
|
|
email: `${suffix}@example.com`,
|
|
};
|
|
}
|
|
|
|
describe("auth-broker wire surface", () => {
|
|
let tempDir = "";
|
|
let store: SqliteAuthCredentialStore | undefined;
|
|
let storage: AuthStorage | undefined;
|
|
let handle: AuthBrokerServerHandle | undefined;
|
|
let token = "";
|
|
|
|
beforeEach(async () => {
|
|
for (const key of ANTHROPIC_ENV) {
|
|
savedEnv[key] = process.env[key];
|
|
delete process.env[key];
|
|
}
|
|
tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "auth-broker-wire-"));
|
|
store = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db"));
|
|
store.saveOAuth("anthropic", mintOAuthCredential("a", Date.now() + 60_000));
|
|
storage = new AuthStorage(store);
|
|
await storage.reload();
|
|
token = "test-bearer";
|
|
handle = startAuthBroker({
|
|
storage,
|
|
bind: "127.0.0.1:0",
|
|
bearerTokens: [token],
|
|
disableRefresher: true,
|
|
});
|
|
});
|
|
|
|
afterEach(async () => {
|
|
vi.restoreAllMocks();
|
|
await handle?.close();
|
|
storage?.close();
|
|
store?.close();
|
|
await removeWithRetries(tempDir);
|
|
for (const key of ANTHROPIC_ENV) {
|
|
if (savedEnv[key] === undefined) delete process.env[key];
|
|
else process.env[key] = savedEnv[key];
|
|
}
|
|
});
|
|
|
|
test("GET /v1/healthz returns ok without auth", async () => {
|
|
const res = await fetch(`${handle!.url}/v1/healthz`);
|
|
expect(res.status).toBe(200);
|
|
const body = (await res.json()) as { ok: boolean };
|
|
expect(body.ok).toBe(true);
|
|
});
|
|
|
|
test("GET /v1/snapshot requires bearer and redacts refresh tokens", async () => {
|
|
const unauthorized = await fetch(`${handle!.url}/v1/snapshot`);
|
|
expect(unauthorized.status).toBe(401);
|
|
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const snapshotResult = await client.fetchSnapshot();
|
|
if (snapshotResult.status !== 200) throw new Error("expected snapshot");
|
|
const snapshot = snapshotResult.snapshot;
|
|
expect(snapshot.credentials).toHaveLength(1);
|
|
const entry = snapshot.credentials[0];
|
|
expect(entry.provider).toBe("anthropic");
|
|
expect(entry.credential.type).toBe("oauth");
|
|
if (entry.credential.type === "oauth") {
|
|
expect(entry.credential.access).toBe("access-a");
|
|
// Refresh token is replaced with the wire sentinel — clients never see it.
|
|
expect(entry.credential.refresh).toBe(REMOTE_REFRESH_SENTINEL);
|
|
}
|
|
});
|
|
|
|
test("GET /v1/snapshot returns generation headers and 304 for unchanged long-poll", async () => {
|
|
const res = await fetch(`${handle!.url}/v1/snapshot`, {
|
|
headers: { Authorization: `Bearer ${token}` },
|
|
});
|
|
expect(res.status).toBe(200);
|
|
const body = (await res.json()) as { generation: number; serverNowMs: number; refresher: { enabled: boolean } };
|
|
expect(res.headers.get("etag")).toBe(`"${body.generation}"`);
|
|
expect(res.headers.get("cache-control")).toBe("no-store");
|
|
expect(body.generation).toBeGreaterThan(0);
|
|
expect(body.serverNowMs).toBeGreaterThan(0);
|
|
expect(body.refresher.enabled).toBe(false);
|
|
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const unchanged = await client.fetchSnapshot({ ifGenerationGt: body.generation, waitMs: 10 });
|
|
expect(unchanged.status).toBe(304);
|
|
expect(unchanged.generation).toBe(body.generation);
|
|
});
|
|
|
|
test("GET /v1/snapshot long-poll wakes when generation changes", async () => {
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const initial = await client.fetchSnapshot();
|
|
if (initial.status !== 200) throw new Error("expected snapshot");
|
|
|
|
const pending = client.fetchSnapshot({ ifGenerationGt: initial.generation, waitMs: 1000 });
|
|
setTimeout(() => {
|
|
storage!.upsertCredential("anthropic", mintOAuthCredential("b", Date.now() + 120_000));
|
|
}, 10);
|
|
|
|
const changed = await pending;
|
|
expect(changed.status).toBe(200);
|
|
if (changed.status !== 200) throw new Error("expected changed snapshot");
|
|
expect(changed.generation).toBeGreaterThan(initial.generation);
|
|
expect(
|
|
changed.snapshot.credentials.some(
|
|
entry => entry.credential.type === "oauth" && entry.credential.access === "access-b",
|
|
),
|
|
).toBe(true);
|
|
});
|
|
|
|
test("POST /v1/credential/:id/refresh forces a refresh and persists the new credential", async () => {
|
|
const refreshed = {
|
|
access: "access-rotated",
|
|
refresh: "refresh-rotated",
|
|
expires: Date.now() + 120_000,
|
|
accountId: "account-a",
|
|
email: "a@example.com",
|
|
};
|
|
vi.spyOn(oauthUtils, "refreshOAuthToken").mockResolvedValue(refreshed);
|
|
|
|
const initialResult = await new AuthBrokerClient({ url: handle!.url, token }).fetchSnapshot();
|
|
if (initialResult.status !== 200) throw new Error("expected snapshot");
|
|
const id = initialResult.snapshot.credentials[0].id;
|
|
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const result = await client.refreshCredential(id);
|
|
expect(result.entry.id).toBe(id);
|
|
if (result.entry.credential.type === "oauth") {
|
|
expect(result.entry.credential.access).toBe("access-rotated");
|
|
expect(result.entry.credential.refresh).toBe(REMOTE_REFRESH_SENTINEL);
|
|
}
|
|
|
|
// Underlying SQLite row was updated with the *real* refresh token (no sentinel).
|
|
const persisted = store!.getOAuth("anthropic");
|
|
expect(persisted?.access).toBe("access-rotated");
|
|
expect(persisted?.refresh).toBe("refresh-rotated");
|
|
});
|
|
|
|
test("POST /v1/credential/:id/disable soft-deletes the credential and surfaces 404 thereafter", async () => {
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const initialResult = await client.fetchSnapshot();
|
|
if (initialResult.status !== 200) throw new Error("expected snapshot");
|
|
const id = initialResult.snapshot.credentials[0].id;
|
|
|
|
const result = await client.disableCredential(id, "revoked by user");
|
|
expect(result.ok).toBe(true);
|
|
|
|
const afterResult = await client.fetchSnapshot();
|
|
if (afterResult.status !== 200) throw new Error("expected snapshot");
|
|
expect(afterResult.snapshot.credentials).toHaveLength(0);
|
|
|
|
await expect(client.refreshCredential(id)).rejects.toThrow();
|
|
});
|
|
|
|
test("GET /v1/usage/history serves recorded snapshots with sinceMs/provider filters and requires auth", async () => {
|
|
const hourMs = 60 * 60 * 1000;
|
|
const now = Date.now();
|
|
// Hours apart so the store's per-bucket dedupe keeps distinct rows.
|
|
store!.recordUsageSnapshots!([
|
|
{
|
|
recordedAt: now - 3 * hourMs,
|
|
provider: "anthropic",
|
|
accountKey: "acct-a",
|
|
email: "a@example.com",
|
|
limitId: "5h",
|
|
label: "Claude 5 Hour",
|
|
windowLabel: "5h",
|
|
usedFraction: 0.2,
|
|
status: "ok",
|
|
},
|
|
{
|
|
recordedAt: now - hourMs,
|
|
provider: "anthropic",
|
|
accountKey: "acct-a",
|
|
email: "a@example.com",
|
|
limitId: "5h",
|
|
label: "Claude 5 Hour",
|
|
windowLabel: "5h",
|
|
usedFraction: 0.9,
|
|
status: "warning",
|
|
},
|
|
{
|
|
recordedAt: now - hourMs,
|
|
provider: "openai-codex",
|
|
accountKey: "acct-b",
|
|
limitId: "weekly",
|
|
label: "Weekly",
|
|
usedFraction: 0.5,
|
|
},
|
|
]);
|
|
|
|
const unauthorized = await fetch(`${handle!.url}/v1/usage/history`);
|
|
expect(unauthorized.status).toBe(401);
|
|
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const all = await client.fetchUsageHistory();
|
|
expect(all.entries).toHaveLength(3);
|
|
expect(all.entries[0].recordedAt).toBeLessThanOrEqual(all.entries[1].recordedAt);
|
|
|
|
const recent = await client.fetchUsageHistory({ sinceMs: now - 2 * hourMs });
|
|
expect(recent.entries.map(e => e.usedFraction).sort()).toEqual([0.5, 0.9]);
|
|
|
|
const anthropicOnly = await client.fetchUsageHistory({ provider: "anthropic" });
|
|
expect(anthropicOnly.entries).toHaveLength(2);
|
|
expect(anthropicOnly.entries.every(e => e.provider === "anthropic")).toBe(true);
|
|
expect(anthropicOnly.entries[1]).toMatchObject({
|
|
accountKey: "acct-a",
|
|
email: "a@example.com",
|
|
windowLabel: "5h",
|
|
usedFraction: 0.9,
|
|
status: "warning",
|
|
});
|
|
});
|
|
|
|
test("POST /v1/usage/observed persists per-client usage served by GET /v1/usage/clients", async () => {
|
|
const unauthorized = await fetch(`${handle!.url}/v1/usage/observed`, { method: "POST" });
|
|
expect(unauthorized.status).toBe(401);
|
|
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const now = Date.now();
|
|
const entry = {
|
|
at: now,
|
|
provider: "anthropic",
|
|
model: "claude-x",
|
|
requests: 2,
|
|
inputTokens: 1000,
|
|
outputTokens: 400,
|
|
cacheReadTokens: 50,
|
|
cacheWriteTokens: 25,
|
|
costUsd: 3.25,
|
|
};
|
|
const ack = await client.reportClientUsage({ installId: "install-1", hostname: "mbp.local", entries: [entry] });
|
|
expect(ack.ok).toBe(true);
|
|
await client.reportClientUsage({
|
|
installId: "install-2",
|
|
entries: [{ ...entry, provider: "openai-codex", model: "gpt-y", requests: 1, costUsd: 0 }],
|
|
});
|
|
|
|
const summary = await client.fetchClientUsageSummary();
|
|
expect(summary.clients).toHaveLength(2);
|
|
const first = summary.clients.find(c => c.installId === "install-1");
|
|
expect(first).toMatchObject({ hostname: "mbp.local" });
|
|
expect(first?.lastSeen).toBeGreaterThan(0);
|
|
expect(first?.providers).toEqual([
|
|
{
|
|
provider: "anthropic",
|
|
requests: 2,
|
|
inputTokens: 1000,
|
|
outputTokens: 400,
|
|
cacheReadTokens: 50,
|
|
cacheWriteTokens: 25,
|
|
costUsd: 3.25,
|
|
},
|
|
]);
|
|
const second = summary.clients.find(c => c.installId === "install-2");
|
|
expect(second?.hostname).toBeUndefined();
|
|
expect(second?.providers[0]).toMatchObject({ provider: "openai-codex", requests: 1 });
|
|
|
|
// sinceMs beyond the recorded timestamps returns clients with no aggregates.
|
|
const future = await client.fetchClientUsageSummary({ sinceMs: now + 60_000 });
|
|
expect(future.clients.every(c => c.providers.length === 0)).toBe(true);
|
|
|
|
// Malformed body is rejected by schema validation.
|
|
const bad = await fetch(`${handle!.url}/v1/usage/observed`, {
|
|
method: "POST",
|
|
headers: { Authorization: `Bearer ${token}`, "Content-Type": "application/json" },
|
|
body: JSON.stringify({ installId: "x", entries: [{ at: "not-a-number" }] }),
|
|
});
|
|
expect(bad.status).toBe(400);
|
|
});
|
|
|
|
test("Unknown route returns 404", async () => {
|
|
const res = await fetch(`${handle!.url}/v1/nope`, {
|
|
headers: { Authorization: `Bearer ${token}` },
|
|
});
|
|
expect(res.status).toBe(404);
|
|
});
|
|
|
|
test("GET /v1/snapshot/stream requires bearer", async () => {
|
|
const res = await fetch(`${handle!.url}/v1/snapshot/stream`);
|
|
expect(res.status).toBe(401);
|
|
});
|
|
|
|
test("SSE stream emits initial snapshot then upsert delta", async () => {
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const controller = new AbortController();
|
|
const iter = client.openSnapshotStream({ signal: controller.signal });
|
|
try {
|
|
const first = await iter.next();
|
|
if (first.done) throw new Error("expected snapshot frame");
|
|
expect(first.value.kind).toBe("snapshot");
|
|
if (first.value.kind === "snapshot") {
|
|
expect(first.value.credentials).toHaveLength(1);
|
|
expect(first.value.credentials[0].provider).toBe("anthropic");
|
|
}
|
|
|
|
storage!.upsertCredential("anthropic", mintOAuthCredential("b", Date.now() + 120_000));
|
|
|
|
const next = await nextMatching(iter, event => event.kind === "entry");
|
|
if (next.kind !== "entry") throw new Error("expected entry frame");
|
|
expect(next.entry.provider).toBe("anthropic");
|
|
expect(next.entry.credential.type).toBe("oauth");
|
|
if (next.entry.credential.type === "oauth") {
|
|
expect(next.entry.credential.access).toBe("access-b");
|
|
expect(next.entry.credential.refresh).toBe(REMOTE_REFRESH_SENTINEL);
|
|
}
|
|
} finally {
|
|
controller.abort();
|
|
await iter.return(undefined).catch(() => {});
|
|
}
|
|
});
|
|
|
|
test("SSE stream pushes entry frame on refresh", async () => {
|
|
const refreshed = {
|
|
access: "access-rotated",
|
|
refresh: "refresh-rotated",
|
|
expires: Date.now() + 120_000,
|
|
accountId: "account-a",
|
|
email: "a@example.com",
|
|
};
|
|
vi.spyOn(oauthUtils, "refreshOAuthToken").mockResolvedValue(refreshed);
|
|
|
|
const initialSnapshot = await new AuthBrokerClient({ url: handle!.url, token }).fetchSnapshot();
|
|
if (initialSnapshot.status !== 200) throw new Error("expected snapshot");
|
|
const id = initialSnapshot.snapshot.credentials[0].id;
|
|
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const controller = new AbortController();
|
|
const iter = client.openSnapshotStream({ signal: controller.signal });
|
|
try {
|
|
const first = await iter.next();
|
|
if (first.done) throw new Error("expected snapshot frame");
|
|
|
|
await storage!.refreshCredentialById(id);
|
|
|
|
const next = await nextMatching(
|
|
iter,
|
|
event => event.kind === "entry" && event.entry.credential.type === "oauth" && event.entry.id === id,
|
|
);
|
|
if (next.kind !== "entry") throw new Error("expected entry frame");
|
|
if (next.entry.credential.type !== "oauth") throw new Error("expected oauth credential");
|
|
expect(next.entry.credential.access).toBe("access-rotated");
|
|
expect(next.entry.credential.refresh).toBe(REMOTE_REFRESH_SENTINEL);
|
|
} finally {
|
|
controller.abort();
|
|
await iter.return(undefined).catch(() => {});
|
|
}
|
|
});
|
|
|
|
test("SSE stream pushes removed frame on disable", async () => {
|
|
const initialSnapshot = await new AuthBrokerClient({ url: handle!.url, token }).fetchSnapshot();
|
|
if (initialSnapshot.status !== 200) throw new Error("expected snapshot");
|
|
const id = initialSnapshot.snapshot.credentials[0].id;
|
|
|
|
const client = new AuthBrokerClient({ url: handle!.url, token });
|
|
const controller = new AbortController();
|
|
const iter = client.openSnapshotStream({ signal: controller.signal });
|
|
try {
|
|
const first = await iter.next();
|
|
if (first.done) throw new Error("expected snapshot frame");
|
|
|
|
const disabled = storage!.disableCredentialById(id, "revoked by test");
|
|
expect(disabled).toBe(true);
|
|
|
|
const next = await nextMatching(iter, event => event.kind === "removed");
|
|
if (next.kind !== "removed") throw new Error("expected removed frame");
|
|
expect(next.id).toBe(id);
|
|
} finally {
|
|
controller.abort();
|
|
await iter.return(undefined).catch(() => {});
|
|
}
|
|
});
|
|
|
|
test("SSE stream keepalive comment arrives on cadence", async () => {
|
|
const localStore = await SqliteAuthCredentialStore.open(path.join(tempDir, "keepalive.db"));
|
|
localStore.saveOAuth("anthropic", mintOAuthCredential("k", Date.now() + 60_000));
|
|
const localStorage = new AuthStorage(localStore);
|
|
await localStorage.reload();
|
|
const localToken = "keepalive-bearer";
|
|
const localHandle = startAuthBroker({
|
|
storage: localStorage,
|
|
bind: "127.0.0.1:0",
|
|
bearerTokens: [localToken],
|
|
disableRefresher: true,
|
|
streamKeepaliveMs: 25,
|
|
});
|
|
const controller = new AbortController();
|
|
try {
|
|
const res = await fetch(`${localHandle.url}/v1/snapshot/stream`, {
|
|
headers: { Authorization: `Bearer ${localToken}`, Accept: "text/event-stream" },
|
|
signal: controller.signal,
|
|
});
|
|
expect(res.status).toBe(200);
|
|
expect(res.headers.get("content-type") ?? "").toContain("text/event-stream");
|
|
expect(res.body).not.toBeNull();
|
|
const reader = (res.body as ReadableStream<Uint8Array>).getReader();
|
|
const decoder = new TextDecoder();
|
|
const deadline = Date.now() + 1_000;
|
|
let seenKeepalive = false;
|
|
let buffer = "";
|
|
try {
|
|
while (Date.now() < deadline) {
|
|
const { value, done } = await reader.read();
|
|
if (done) break;
|
|
buffer += decoder.decode(value, { stream: true });
|
|
if (buffer.includes(": keepalive\n\n")) {
|
|
seenKeepalive = true;
|
|
break;
|
|
}
|
|
}
|
|
} finally {
|
|
await reader.cancel().catch(() => {});
|
|
}
|
|
expect(seenKeepalive).toBe(true);
|
|
} finally {
|
|
controller.abort();
|
|
await localHandle.close();
|
|
localStorage.close();
|
|
localStore.close();
|
|
}
|
|
});
|
|
|
|
test("openSnapshotStream throws AuthBrokerStreamUnsupportedError on 404", async () => {
|
|
const dummy = Bun.serve({
|
|
hostname: "127.0.0.1",
|
|
port: 0,
|
|
fetch: () => new Response("Not Found", { status: 404 }),
|
|
});
|
|
try {
|
|
const client = new AuthBrokerClient({ url: `http://${dummy.hostname}:${dummy.port}`, token });
|
|
const iter = client.openSnapshotStream();
|
|
await expect(iter.next()).rejects.toBeInstanceOf(AuthBrokerStreamUnsupportedError);
|
|
} finally {
|
|
dummy.stop(true);
|
|
}
|
|
});
|
|
|
|
test("openSnapshotStream rejects 200 responses that are not SSE", async () => {
|
|
const dummy = Bun.serve({
|
|
hostname: "127.0.0.1",
|
|
port: 0,
|
|
fetch: () => new Response("{}", { status: 200, headers: { "Content-Type": "application/json" } }),
|
|
});
|
|
try {
|
|
const client = new AuthBrokerClient({ url: `http://${dummy.hostname}:${dummy.port}`, token });
|
|
const iter = client.openSnapshotStream();
|
|
await expect(iter.next()).rejects.toThrow(/non-SSE/);
|
|
} finally {
|
|
dummy.stop(true);
|
|
}
|
|
});
|
|
|
|
test("openSnapshotStream rejects SSE responses without an initial snapshot", async () => {
|
|
const dummy = Bun.serve({
|
|
hostname: "127.0.0.1",
|
|
port: 0,
|
|
fetch: () =>
|
|
new Response(": keepalive\n\n", { status: 200, headers: { "Content-Type": "text/event-stream" } }),
|
|
});
|
|
try {
|
|
const client = new AuthBrokerClient({ url: `http://${dummy.hostname}:${dummy.port}`, token });
|
|
const iter = client.openSnapshotStream();
|
|
await expect(iter.next()).rejects.toThrow(/initial snapshot/);
|
|
} finally {
|
|
dummy.stop(true);
|
|
}
|
|
});
|
|
});
|
|
|
|
async function nextMatching(
|
|
iter: AsyncGenerator<SnapshotStreamEvent>,
|
|
predicate: (event: SnapshotStreamEvent) => boolean,
|
|
timeoutMs = 2_000,
|
|
): Promise<SnapshotStreamEvent> {
|
|
const deadline = Date.now() + timeoutMs;
|
|
for (;;) {
|
|
const remaining = deadline - Date.now();
|
|
if (remaining <= 0) throw new Error("nextMatching timeout");
|
|
const timer = Promise.withResolvers<never>();
|
|
const handle = setTimeout(() => timer.reject(new Error("nextMatching timeout")), remaining);
|
|
try {
|
|
const res = await Promise.race([iter.next(), timer.promise]);
|
|
if (res.done) throw new Error("stream ended before predicate satisfied");
|
|
if (predicate(res.value)) return res.value;
|
|
} finally {
|
|
clearTimeout(handle);
|
|
}
|
|
}
|
|
}
|