From 19fdbe2835e87a4c61349ae15087f43d8e10cbfc Mon Sep 17 00:00:00 2001 From: usr-bin-roygbiv Date: Mon, 27 Jul 2026 00:41:54 +0000 Subject: [PATCH] perf(ai): lazily construct auth broker wire schemas --- packages/ai/src/auth-broker/client.ts | 63 ++- packages/ai/src/auth-broker/server.ts | 17 +- .../src/auth-broker/wire-schema-resource.ts | 487 ++++++++++++++++++ packages/ai/src/auth-broker/wire-schemas.ts | 409 ++------------- ...auth-broker-wire-lazy-construction.test.ts | 37 ++ .../auth-broker-wire-schema-contract.test.ts | 267 ++++++++++ .../auth-broker-wire-construction-preload.ts | 15 + .../auth-broker-wire-construction-probe.ts | 55 ++ 8 files changed, 937 insertions(+), 413 deletions(-) create mode 100644 packages/ai/src/auth-broker/wire-schema-resource.ts create mode 100644 packages/ai/test/auth-broker-wire-lazy-construction.test.ts create mode 100644 packages/ai/test/auth-broker-wire-schema-contract.test.ts create mode 100644 packages/ai/test/fixtures/auth-broker-wire-construction-preload.ts create mode 100644 packages/ai/test/fixtures/auth-broker-wire-construction-probe.ts diff --git a/packages/ai/src/auth-broker/client.ts b/packages/ai/src/auth-broker/client.ts index 2b84ad26b..2e7b43347 100644 --- a/packages/ai/src/auth-broker/client.ts +++ b/packages/ai/src/auth-broker/client.ts @@ -28,22 +28,21 @@ import type { UsageResponse, UsageStaleResponse, } from "./types"; -import { - clientUsageReportResponseSchema, - clientUsageSummaryResponseSchema, - credentialBlockResponseSchema, - credentialBlocksDeleteResponseSchema, - credentialDisableResponseSchema, - credentialRefreshResponseSchema, - credentialUploadResponseSchema, - disabledCredentialsResponseSchema, - healthzResponseSchema, - snapshotResponseSchema, - snapshotStreamEventSchema, - usageHistoryResponseSchema, - usageResponseSchema, - usageStaleResponseSchema, -} from "./wire-schemas"; +import { getAuthBrokerWireSchemas } from "./wire-schema-resource"; + +type AuthBrokerResponseSchemaName = + | "clientUsageReportResponseSchema" + | "clientUsageSummaryResponseSchema" + | "credentialBlockResponseSchema" + | "credentialBlocksDeleteResponseSchema" + | "credentialDisableResponseSchema" + | "credentialRefreshResponseSchema" + | "credentialUploadResponseSchema" + | "disabledCredentialsResponseSchema" + | "healthzResponseSchema" + | "usageHistoryResponseSchema" + | "usageResponseSchema" + | "usageStaleResponseSchema"; export interface AuthBrokerClientOptions { /** Base URL (e.g. `https://broker.tailnet:8765`). Trailing slashes are trimmed. */ @@ -123,7 +122,7 @@ export class AuthBrokerClient { healthz(signal?: AbortSignal): Promise { return this.#request("GET", "/v1/healthz", { - schema: healthzResponseSchema, + schema: "healthzResponseSchema", auth: false, signal, }); @@ -152,7 +151,7 @@ export class AuthBrokerClient { } const text = await response.text(); const raw = this.#parseJson(text, response.status); - const validated = snapshotResponseSchema(raw); + const validated = getAuthBrokerWireSchemas().snapshotResponseSchema(raw); if (validated instanceof type.errors) { throw new AuthBrokerError("Auth broker response failed schema validation", { status: response.status, @@ -220,7 +219,7 @@ export class AuthBrokerClient { cause: err, }); } - const validated = snapshotStreamEventSchema(parsed); + const validated = getAuthBrokerWireSchemas().snapshotStreamEventSchema(parsed); if (validated instanceof type.errors) { throw new AuthBrokerError("Auth broker stream event failed schema validation", { body: validated.summary, @@ -250,7 +249,7 @@ export class AuthBrokerClient { // `metadata`) but leaves provider-specific extension fields permissive so // the broker can ship new shapes ahead of the client. `raw` is accepted // but normally stripped by the broker before send. - return this.#request("GET", "/v1/usage", { schema: usageResponseSchema, signal }); + return this.#request("GET", "/v1/usage", { schema: "usageResponseSchema", signal }); } /** Recorded usage-limit snapshots from the broker host, oldest first. */ @@ -262,14 +261,14 @@ export class AuthBrokerClient { if (query?.sinceMs !== undefined) params.set("sinceMs", String(query.sinceMs)); if (query?.provider) params.set("provider", query.provider); const path = `/v1/usage/history${params.size > 0 ? `?${params.toString()}` : ""}`; - return this.#request("GET", path, { schema: usageHistoryResponseSchema, signal }); + return this.#request("GET", path, { schema: "usageHistoryResponseSchema", signal }); } /** Report this client's batched observed request usage for per-install burn tracking. */ reportClientUsage(report: ClientUsageReportRequest, signal?: AbortSignal): Promise { return this.#request("POST", "/v1/usage/observed", { body: report, - schema: clientUsageReportResponseSchema, + schema: "clientUsageReportResponseSchema", signal, }); } @@ -280,21 +279,21 @@ export class AuthBrokerClient { if (query?.sinceMs !== undefined) params.set("sinceMs", String(query.sinceMs)); const path = `/v1/usage/clients${params.size > 0 ? `?${params.toString()}` : ""}`; return this.#request("GET", path, { - schema: clientUsageSummaryResponseSchema, + schema: "clientUsageSummaryResponseSchema", signal, }); } notifyUsageStale(signal?: AbortSignal): Promise { return this.#request("POST", "/v1/usage/stale", { - schema: usageStaleResponseSchema, + schema: "usageStaleResponseSchema", signal, }); } async refreshCredential(id: number, signal?: AbortSignal): Promise { return this.#request("POST", `/v1/credential/${id}/refresh`, { - schema: credentialRefreshResponseSchema, + schema: "credentialRefreshResponseSchema", signal, }); } @@ -303,7 +302,7 @@ export class AuthBrokerClient { const body: CredentialDisableRequest = { cause }; return this.#request("POST", `/v1/credential/${id}/disable`, { body, - schema: credentialDisableResponseSchema, + schema: "credentialDisableResponseSchema", signal, }); } @@ -319,7 +318,7 @@ export class AuthBrokerClient { const path = `/v1/credentials/disabled${params.size > 0 ? `?${params.toString()}` : ""}`; try { const response = await this.#request("GET", path, { - schema: disabledCredentialsResponseSchema, + schema: "disabledCredentialsResponseSchema", signal, }); return response.disabled; @@ -337,7 +336,7 @@ export class AuthBrokerClient { const body: CredentialUploadRequest = { provider, credential }; return this.#request("POST", "/v1/credential", { body, - schema: credentialUploadResponseSchema, + schema: "credentialUploadResponseSchema", signal, }); } @@ -350,14 +349,14 @@ export class AuthBrokerClient { const body: CredentialBlockRequest = block; return this.#request("POST", `/v1/credential/${id}/block`, { body, - schema: credentialBlockResponseSchema, + schema: "credentialBlockResponseSchema", signal, }); } async deleteCredentialBlocks(id: number, signal?: AbortSignal): Promise { return this.#request("DELETE", `/v1/credential/${id}/blocks`, { - schema: credentialBlocksDeleteResponseSchema, + schema: "credentialBlocksDeleteResponseSchema", signal, }); } @@ -365,12 +364,12 @@ export class AuthBrokerClient { async #request( method: "GET" | "POST" | "DELETE", path: string, - opts: { schema: (input: unknown) => unknown; auth?: boolean; body?: unknown; signal?: AbortSignal }, + opts: { schema: AuthBrokerResponseSchemaName; auth?: boolean; body?: unknown; signal?: AbortSignal }, ): Promise { const response = await this.#fetchRaw(method, path, opts); const text = await response.text(); const raw = this.#parseJson(text, response.status); - const validated = opts.schema(raw); + const validated = getAuthBrokerWireSchemas()[opts.schema](raw); if (validated instanceof type.errors) { throw new AuthBrokerError("Auth broker response failed schema validation", { status: response.status, diff --git a/packages/ai/src/auth-broker/server.ts b/packages/ai/src/auth-broker/server.ts index 56eeec91f..04a454f9c 100644 --- a/packages/ai/src/auth-broker/server.ts +++ b/packages/ai/src/auth-broker/server.ts @@ -38,12 +38,7 @@ import { DEFAULT_SERVER_IDLE_TIMEOUT_S, DEFAULT_STREAM_KEEPALIVE_MS, } from "./types"; -import { - clientUsageReportRequestSchema, - credentialBlockRequestSchema, - credentialDisableRequestSchema, - credentialUploadRequestSchema, -} from "./wire-schemas"; +import { getAuthBrokerWireSchemas } from "./wire-schema-resource"; export interface AuthBrokerServerOptions { /** Underlying credential storage (wraps the local SQLite store on the broker). */ @@ -622,7 +617,7 @@ export function startAuthBroker(opts: AuthBrokerServerOptions): AuthBrokerServer return json(200, { generatedAt: Date.now(), entries }); } if (req.method === "POST" && pathname === "/v1/usage/observed") { - const parsed = await parseBody(req, clientUsageReportRequestSchema); + const parsed = await parseBody(req, getAuthBrokerWireSchemas().clientUsageReportRequestSchema); if (!parsed.ok) return parsed.response; // Arktype's inferred union collides the `entries` field with // Array.prototype.entries; the schema already validated the shape. @@ -689,7 +684,9 @@ export function startAuthBroker(opts: AuthBrokerServerOptions): AuthBrokerServer const disableMatch = req.method === "POST" ? pathname.match(DISABLE_ROUTE) : null; if (disableMatch) { const id = Number.parseInt(disableMatch[1], 10); - const parsed = await parseBody(req, credentialDisableRequestSchema, { allowEmpty: true }); + const parsed = await parseBody(req, getAuthBrokerWireSchemas().credentialDisableRequestSchema, { + allowEmpty: true, + }); if (!parsed.ok) return parsed.response; const cause = parsed.data.cause && parsed.data.cause.length > 0 ? parsed.data.cause : "disabled via auth-broker"; @@ -705,7 +702,7 @@ export function startAuthBroker(opts: AuthBrokerServerOptions): AuthBrokerServer const blockMatch = req.method === "POST" ? pathname.match(BLOCK_ROUTE) : null; if (blockMatch) { const id = Number.parseInt(blockMatch[1], 10); - const parsed = await parseBody(req, credentialBlockRequestSchema); + const parsed = await parseBody(req, getAuthBrokerWireSchemas().credentialBlockRequestSchema); if (!parsed.ok) return parsed.response; const block: StoredCredentialBlock = { credentialId: id, @@ -755,7 +752,7 @@ export function startAuthBroker(opts: AuthBrokerServerOptions): AuthBrokerServer } } if (req.method === "POST" && pathname === "/v1/credential") { - const parsed = await parseBody(req, credentialUploadRequestSchema); + const parsed = await parseBody(req, getAuthBrokerWireSchemas().credentialUploadRequestSchema); if (!parsed.ok) return parsed.response; const { provider, credential } = parsed.data; try { diff --git a/packages/ai/src/auth-broker/wire-schema-resource.ts b/packages/ai/src/auth-broker/wire-schema-resource.ts new file mode 100644 index 000000000..d958c6f00 --- /dev/null +++ b/packages/ai/src/auth-broker/wire-schema-resource.ts @@ -0,0 +1,487 @@ +/** + * ArkType schemas for the auth-broker wire protocol. + * + * Shared between the server (validates inbound request bodies) and the client + * (validates responses from the broker). Schemas mirror the TypeScript types + * in `./types.ts` 1:1; the types remain the source of truth for static typing, + * and `Type` is asserted-compatible with them where possible. + * + * Envelope and fixed-shape schemas use `"+": "reject"` so unknown keys are + * rejected — the previous implementation used a hand-rolled `hasOnlyFields` + * allowlist for the same effect. The OAuth credential schema is the deliberate + * exception (standard type keeps extra keys): it preserves provider-specific extension fields so + * they round-trip through the broker instead of being dropped (see below). + */ +import { type Type, type } from "arktype"; +import { + type ApiKeyCredential, + type AuthCredential, + type AuthCredentialSnapshotEntry, + type DisabledCredentialSummary, + type OAuthCredential, + REMOTE_REFRESH_SENTINEL, + type RemoteOAuthCredential, + type SnapshotCredential, +} from "../auth-storage"; +import type { + ClientUsageReportRequest, + ClientUsageReportResponse, + ClientUsageSummaryResponse, + CredentialBlockRequest, + CredentialBlockResponse, + CredentialBlockSnapshot, + CredentialBlocksDeleteResponse, + CredentialDisableResponse, + CredentialRefreshResponse, + CredentialUploadRequest, + CredentialUploadResponse, + DisabledCredentialsResponse, + HealthzResponse, + RefresherSchedule, + SnapshotEntry, + SnapshotResponse, + SnapshotStreamEntryEvent, + SnapshotStreamEvent, + SnapshotStreamRemovedEvent, + SnapshotStreamSnapshotEvent, + UsageHistoryResponse, + UsageResponse, + UsageStaleResponse, +} from "./types"; + +export interface AuthBrokerWireSchemas { + readonly oauthCredentialSchema: Type; + readonly remoteOauthCredentialSchema: Type; + readonly apiKeyCredentialSchema: Type; + readonly writableAuthCredentialSchema: Type; + readonly snapshotCredentialSchema: Type; + readonly credentialSnapshotEntrySchema: Type; + readonly credentialBlockSnapshotSchema: Type; + readonly snapshotEntrySchema: Type; + readonly refresherScheduleSchema: Type; + readonly snapshotResponseSchema: Type; + readonly snapshotStreamSnapshotEventSchema: Type; + readonly snapshotStreamEntryEventSchema: Type; + readonly snapshotStreamRemovedEventSchema: Type; + readonly snapshotStreamEventSchema: Type; + readonly healthzResponseSchema: Type; + readonly usageResponseSchema: Type; + readonly usageHistoryResponseSchema: Type; + readonly clientUsageReportRequestSchema: Type; + readonly clientUsageReportResponseSchema: Type; + readonly clientUsageSummaryResponseSchema: Type; + readonly credentialRefreshResponseSchema: Type; + readonly credentialDisableRequestSchema: Type<{ cause?: string }>; + readonly credentialDisableResponseSchema: Type; + readonly disabledCredentialSummarySchema: Type; + readonly disabledCredentialsResponseSchema: Type; + readonly credentialBlockRequestSchema: Type; + readonly credentialBlockResponseSchema: Type; + readonly credentialBlocksDeleteResponseSchema: Type; + readonly usageStaleResponseSchema: Type; + readonly credentialUploadRequestSchema: Type; + readonly credentialUploadResponseSchema: Type; +} + +function buildAuthBrokerWireSchemas(): AuthBrokerWireSchemas { + // ─── Credential payloads ─────────────────────────────────────────────────── + + /** Real OAuth credential (broker-side) — refresh token is the actual upstream value. */ + const oauthCredentialSchema = type({ + "apiEndpoint?": "string", + type: "'oauth'", + refresh: type("string").narrow( + (value, ctx) => + value !== REMOTE_REFRESH_SENTINEL || + ctx.mustBe(`not equal to the remote sentinel (${REMOTE_REFRESH_SENTINEL})`), + ), + access: type("string").atLeastLength(1), + expires: "number", + "enterpriseUrl?": "string", + "projectId?": "string", + "email?": "string", + "accountId?": "string", + "orgId?": "string", + "orgName?": "string", + "authorizedAt?": "number", + }); + + /** OAuth credential as it appears in broker snapshots — refresh replaced with sentinel. */ + const remoteOauthCredentialSchema = type({ + "apiEndpoint?": "string", + type: "'oauth'", + refresh: type.enumerated(REMOTE_REFRESH_SENTINEL), + access: type("string").atLeastLength(1), + expires: "number", + "enterpriseUrl?": "string", + "projectId?": "string", + "email?": "string", + "accountId?": "string", + "orgId?": "string", + "orgName?": "string", + "authorizedAt?": "number", + }); + + const apiKeyCredentialSchema = type({ + "+": "reject", + type: "'api_key'", + key: type("string").atLeastLength(1), + "source?": "'login'", + }); + + /** Discriminated union accepted on POST /v1/credential (writes). */ + const writableAuthCredentialSchema = oauthCredentialSchema.or(apiKeyCredentialSchema); + + /** Discriminated union returned in snapshots (refresh is sentinel for OAuth). */ + const snapshotCredentialSchema = remoteOauthCredentialSchema.or(apiKeyCredentialSchema); + + // ─── Snapshot ────────────────────────────────────────────────────────────── + + const credentialSnapshotEntrySchema = type({ + "+": "reject", + id: "number.integer", + provider: type("string").atLeastLength(1), + credential: snapshotCredentialSchema, + identityKey: "string | null", + }); + + const credentialBlockSnapshotSchema = type({ + "+": "reject", + providerKey: type("string").atLeastLength(1), + blockScope: "string", + blockedUntilMs: "number", + "updatedAtMs?": "number", + }); + + const snapshotEntrySchema = type({ + "+": "reject", + id: "number.integer", + provider: type("string").atLeastLength(1), + credential: snapshotCredentialSchema, + identityKey: "string | null", + rotatesInMs: "number | null", + "blocks?": credentialBlockSnapshotSchema.array(), + }); + + const refresherScheduleSchema = type({ + "+": "reject", + enabled: "boolean", + intervalMs: "number", + skewMs: "number", + nextSweepInMs: "number", + }); + + const snapshotResponseSchema = type({ + "+": "reject", + generation: "number.integer", + generatedAt: "number", + serverNowMs: "number", + refresher: refresherScheduleSchema, + credentials: snapshotEntrySchema.array(), + }); + + // ─── Snapshot stream (SSE) ──────────────────────────────────────────────── + + /** First frame on connect — full snapshot embedded inline with a `kind` tag. */ + const snapshotStreamSnapshotEventSchema = type({ + "+": "reject", + generation: "number.integer", + generatedAt: "number", + serverNowMs: "number", + refresher: refresherScheduleSchema, + credentials: snapshotEntrySchema.array(), + kind: "'snapshot'", + }); + + /** Per-credential upsert/refresh delta. */ + const snapshotStreamEntryEventSchema = type({ + "+": "reject", + kind: "'entry'", + generation: "number.integer", + serverNowMs: "number", + refresher: refresherScheduleSchema, + entry: snapshotEntrySchema, + }); + + /** Per-credential delete delta. */ + const snapshotStreamRemovedEventSchema = type({ + "+": "reject", + kind: "'removed'", + generation: "number.integer", + serverNowMs: "number", + refresher: refresherScheduleSchema, + id: "number.integer", + }); + + /** Discriminated union over every event frame the snapshot stream emits. */ + const snapshotStreamEventSchema = snapshotStreamSnapshotEventSchema + .or(snapshotStreamEntryEventSchema) + .or(snapshotStreamRemovedEventSchema); + + // ─── Healthz ──────────────────────────────────────────────────────────────── + + const healthzResponseSchema = type({ + "+": "reject", + ok: "boolean", + "version?": "string", + }); + + // ─── Usage ───────────────────────────────────────────────────────────────── + + const usageUnitSchema = type("'percent' | 'tokens' | 'requests' | 'usd' | 'minutes' | 'bytes' | 'unknown'"); + const usageStatusSchema = type("'ok' | 'warning' | 'exhausted' | 'unknown'"); + + const usageWindowSchema = type({ + id: "string", + label: "string", + "durationMs?": "number", + "resetsAt?": "number", + }); + + const usageAmountSchema = type({ + "used?": "number", + "limit?": "number", + "remaining?": "number", + "usedFraction?": "number", + "remainingFraction?": "number", + unit: usageUnitSchema, + }); + + const usageScopeSchema = type({ + provider: "string", + "accountId?": "string", + "projectId?": "string", + "orgId?": "string", + "modelId?": "string", + "tier?": "string", + "windowId?": "string", + "shared?": "boolean", + }); + + const usageLimitSchema = type({ + id: "string", + label: "string", + scope: usageScopeSchema, + "window?": usageWindowSchema, + amount: usageAmountSchema, + "status?": usageStatusSchema, + "notes?": "string[]", + }); + + const usageResetCreditDetailSchema = type({ + "grantedAt?": "string", + "expiresAt?": "string", + "status?": "string", + }); + + const usageResetCreditsSchema = type({ + availableCount: "number", + "credits?": usageResetCreditDetailSchema.array(), + }); + + const arkUsageReportSchema = type({ + provider: "string", + fetchedAt: "number", + limits: usageLimitSchema.array(), + "resetCredits?": usageResetCreditsSchema, + "notes?": "string[]", + "metadata?": { "[string]": "unknown" }, + "raw?": "unknown", + }); + + /** + * Broker `/v1/usage` response. Reports are full {@link UsageReport}s minus the + * heavy provider-specific `raw` field (the server strips it before send) — we + * keep `raw` optional in the underlying schema so a misconfigured broker that + * forgot to strip still validates. + */ + const usageResponseSchema = type({ + "+": "reject", + generatedAt: "number", + reports: arkUsageReportSchema.array(), + }); + + const usageHistoryEntrySchema = type({ + recordedAt: "number", + provider: "string", + accountKey: "string", + "email?": "string", + "accountId?": "string", + limitId: "string", + label: "string", + "windowLabel?": "string", + "usedFraction?": "number", + "status?": "'ok' | 'warning' | 'exhausted' | 'unknown'", + "resetsAt?": "number", + }); + + /** Broker `/v1/usage/history` response — recorded usage-limit snapshots, oldest first. */ + const usageHistoryResponseSchema = type({ + "+": "reject", + generatedAt: "number", + entries: usageHistoryEntrySchema.array(), + }); + + const observedUsageEntrySchema = type({ + at: "number", + provider: "string", + model: "string", + requests: "number", + inputTokens: "number", + outputTokens: "number", + cacheReadTokens: "number", + cacheWriteTokens: "number", + costUsd: "number", + }); + + /** Broker `POST /v1/usage/observed` request — one client's batched observed usage. */ + const clientUsageReportRequestSchema = type({ + "+": "reject", + installId: "string", + "hostname?": "string", + entries: observedUsageEntrySchema.array(), + }); + + const clientUsageReportResponseSchema = type({ + "+": "reject", + ok: "boolean", + }); + + const clientProviderUsageSchema = type({ + provider: "string", + requests: "number", + inputTokens: "number", + outputTokens: "number", + cacheReadTokens: "number", + cacheWriteTokens: "number", + costUsd: "number", + }); + + const clientUsageClientSummarySchema = type({ + installId: "string", + "hostname?": "string", + firstSeen: "number", + lastSeen: "number", + providers: clientProviderUsageSchema.array(), + }); + + /** Broker `GET /v1/usage/clients` response — per-client token burn aggregates. */ + const clientUsageSummaryResponseSchema = type({ + "+": "reject", + generatedAt: "number", + clients: clientUsageClientSummarySchema.array(), + }); + + // ─── Refresh ─────────────────────────────────────────────────────────────── + + const credentialRefreshResponseSchema = type({ + "+": "reject", + entry: credentialSnapshotEntrySchema, + }); + + // ─── Disable ─────────────────────────────────────────────────────────────── + + const credentialDisableRequestSchema = type({ + "+": "reject", + "cause?": "string", + }); + + const credentialDisableResponseSchema = type({ + "+": "reject", + ok: "boolean", + }); + + /** One disabled-credential tombstone — identity + cause, never token material. */ + const disabledCredentialSummarySchema = type({ + "+": "reject", + id: "number.integer", + provider: type("string").atLeastLength(1), + type: "'oauth' | 'api_key'", + "email?": "string", + "accountId?": "string", + "orgId?": "string", + "orgName?": "string", + cause: "string", + "disabledAtMs?": "number", + }); + + /** Broker `GET /v1/credentials/disabled` response. */ + const disabledCredentialsResponseSchema = type({ + "+": "reject", + generatedAt: "number", + disabled: disabledCredentialSummarySchema.array(), + }); + + // ─── Credential blocks ────────────────────────────────────────────────────── + + const credentialBlockRequestSchema = credentialBlockSnapshotSchema; + + const credentialBlockResponseSchema = type({ + "+": "reject", + ok: "boolean", + }); + + const credentialBlocksDeleteResponseSchema = type({ + "+": "reject", + ok: "boolean", + }); + + const usageStaleResponseSchema = type({ + "+": "reject", + ok: "boolean", + }); + + // ─── Upload ──────────────────────────────────────────────────────────────── + + const credentialUploadRequestSchema = type({ + "+": "reject", + provider: type("string").atLeastLength(1), + credential: writableAuthCredentialSchema, + }); + + const credentialUploadResponseSchema = type({ + "+": "reject", + entries: credentialSnapshotEntrySchema.array(), + }); + + return { + oauthCredentialSchema, + remoteOauthCredentialSchema, + apiKeyCredentialSchema, + writableAuthCredentialSchema, + snapshotCredentialSchema, + credentialSnapshotEntrySchema, + credentialBlockSnapshotSchema, + snapshotEntrySchema, + refresherScheduleSchema, + snapshotResponseSchema, + snapshotStreamSnapshotEventSchema, + snapshotStreamEntryEventSchema, + snapshotStreamRemovedEventSchema, + snapshotStreamEventSchema, + healthzResponseSchema, + usageResponseSchema, + usageHistoryResponseSchema, + clientUsageReportRequestSchema, + clientUsageReportResponseSchema, + clientUsageSummaryResponseSchema, + credentialRefreshResponseSchema, + credentialDisableRequestSchema, + credentialDisableResponseSchema, + disabledCredentialSummarySchema, + disabledCredentialsResponseSchema, + credentialBlockRequestSchema, + credentialBlockResponseSchema, + credentialBlocksDeleteResponseSchema, + usageStaleResponseSchema, + credentialUploadRequestSchema, + credentialUploadResponseSchema, + }; +} + +let cachedAuthBrokerWireSchemas: AuthBrokerWireSchemas | undefined; + +export function getAuthBrokerWireSchemas(): AuthBrokerWireSchemas { + if (!cachedAuthBrokerWireSchemas) cachedAuthBrokerWireSchemas = buildAuthBrokerWireSchemas(); + return cachedAuthBrokerWireSchemas; +} diff --git a/packages/ai/src/auth-broker/wire-schemas.ts b/packages/ai/src/auth-broker/wire-schemas.ts index b3f4038e3..03af8b02a 100644 --- a/packages/ai/src/auth-broker/wire-schemas.ts +++ b/packages/ai/src/auth-broker/wire-schemas.ts @@ -1,376 +1,43 @@ /** - * ArkType schemas for the auth-broker wire protocol. + * Eager public facade for the auth-broker wire schemas. * - * Shared between the server (validates inbound request bodies) and the client - * (validates responses from the broker). Schemas mirror the TypeScript types - * in `./types.ts` 1:1; the types remain the source of truth for static typing, - * and `Type` is asserted-compatible with them where possible. - * - * Envelope and fixed-shape schemas use `"+": "reject"` so unknown keys are - * rejected — the previous implementation used a hand-rolled `hasOnlyFields` - * allowlist for the same effect. The OAuth credential schema is the deliberate - * exception (standard type keeps extra keys): it preserves provider-specific extension fields so - * they round-trip through the broker instead of being dropped (see below). + * Direct consumers of this module receive the shared real ArkType objects. + * Internal client/server startup imports the lazy resource instead. */ -import { type } from "arktype"; -import { REMOTE_REFRESH_SENTINEL } from "../auth-storage"; +import { getAuthBrokerWireSchemas } from "./wire-schema-resource"; -// ─── Credential payloads ─────────────────────────────────────────────────── +const wireSchemas = getAuthBrokerWireSchemas(); -/** Real OAuth credential (broker-side) — refresh token is the actual upstream value. */ -export const oauthCredentialSchema = type({ - "apiEndpoint?": "string", - type: "'oauth'", - refresh: type("string").narrow( - (value, ctx) => - value !== REMOTE_REFRESH_SENTINEL || - ctx.mustBe(`not equal to the remote sentinel (${REMOTE_REFRESH_SENTINEL})`), - ), - access: type("string").atLeastLength(1), - expires: "number", - "enterpriseUrl?": "string", - "projectId?": "string", - "email?": "string", - "accountId?": "string", - "orgId?": "string", - "orgName?": "string", - "authorizedAt?": "number", -}); - -/** OAuth credential as it appears in broker snapshots — refresh replaced with sentinel. */ -export const remoteOauthCredentialSchema = type({ - "apiEndpoint?": "string", - type: "'oauth'", - refresh: type.enumerated(REMOTE_REFRESH_SENTINEL), - access: type("string").atLeastLength(1), - expires: "number", - "enterpriseUrl?": "string", - "projectId?": "string", - "email?": "string", - "accountId?": "string", - "orgId?": "string", - "orgName?": "string", - "authorizedAt?": "number", -}); - -export const apiKeyCredentialSchema = type({ - "+": "reject", - type: "'api_key'", - key: type("string").atLeastLength(1), - "source?": "'login'", -}); - -/** Discriminated union accepted on POST /v1/credential (writes). */ -export const writableAuthCredentialSchema = oauthCredentialSchema.or(apiKeyCredentialSchema); - -/** Discriminated union returned in snapshots (refresh is sentinel for OAuth). */ -export const snapshotCredentialSchema = remoteOauthCredentialSchema.or(apiKeyCredentialSchema); - -// ─── Snapshot ────────────────────────────────────────────────────────────── - -export const credentialSnapshotEntrySchema = type({ - "+": "reject", - id: "number.integer", - provider: type("string").atLeastLength(1), - credential: snapshotCredentialSchema, - identityKey: "string | null", -}); - -export const credentialBlockSnapshotSchema = type({ - "+": "reject", - providerKey: type("string").atLeastLength(1), - blockScope: "string", - blockedUntilMs: "number", - "updatedAtMs?": "number", -}); - -export const snapshotEntrySchema = type({ - "+": "reject", - id: "number.integer", - provider: type("string").atLeastLength(1), - credential: snapshotCredentialSchema, - identityKey: "string | null", - rotatesInMs: "number | null", - "blocks?": credentialBlockSnapshotSchema.array(), -}); - -export const refresherScheduleSchema = type({ - "+": "reject", - enabled: "boolean", - intervalMs: "number", - skewMs: "number", - nextSweepInMs: "number", -}); - -export const snapshotResponseSchema = type({ - "+": "reject", - generation: "number.integer", - generatedAt: "number", - serverNowMs: "number", - refresher: refresherScheduleSchema, - credentials: snapshotEntrySchema.array(), -}); - -// ─── Snapshot stream (SSE) ──────────────────────────────────────────────── - -/** First frame on connect — full snapshot embedded inline with a `kind` tag. */ -export const snapshotStreamSnapshotEventSchema = type({ - "+": "reject", - generation: "number.integer", - generatedAt: "number", - serverNowMs: "number", - refresher: refresherScheduleSchema, - credentials: snapshotEntrySchema.array(), - kind: "'snapshot'", -}); - -/** Per-credential upsert/refresh delta. */ -export const snapshotStreamEntryEventSchema = type({ - "+": "reject", - kind: "'entry'", - generation: "number.integer", - serverNowMs: "number", - refresher: refresherScheduleSchema, - entry: snapshotEntrySchema, -}); - -/** Per-credential delete delta. */ -export const snapshotStreamRemovedEventSchema = type({ - "+": "reject", - kind: "'removed'", - generation: "number.integer", - serverNowMs: "number", - refresher: refresherScheduleSchema, - id: "number.integer", -}); - -/** Discriminated union over every event frame the snapshot stream emits. */ -export const snapshotStreamEventSchema = snapshotStreamSnapshotEventSchema - .or(snapshotStreamEntryEventSchema) - .or(snapshotStreamRemovedEventSchema); - -// ─── Healthz ──────────────────────────────────────────────────────────────── - -export const healthzResponseSchema = type({ - "+": "reject", - ok: "boolean", - "version?": "string", -}); - -// ─── Usage ───────────────────────────────────────────────────────────────── - -const usageUnitSchema = type("'percent' | 'tokens' | 'requests' | 'usd' | 'minutes' | 'bytes' | 'unknown'"); -const usageStatusSchema = type("'ok' | 'warning' | 'exhausted' | 'unknown'"); - -const usageWindowSchema = type({ - id: "string", - label: "string", - "durationMs?": "number", - "resetsAt?": "number", -}); - -const usageAmountSchema = type({ - "used?": "number", - "limit?": "number", - "remaining?": "number", - "usedFraction?": "number", - "remainingFraction?": "number", - unit: usageUnitSchema, -}); - -const usageScopeSchema = type({ - provider: "string", - "accountId?": "string", - "projectId?": "string", - "orgId?": "string", - "modelId?": "string", - "tier?": "string", - "windowId?": "string", - "shared?": "boolean", -}); - -const usageLimitSchema = type({ - id: "string", - label: "string", - scope: usageScopeSchema, - "window?": usageWindowSchema, - amount: usageAmountSchema, - "status?": usageStatusSchema, - "notes?": "string[]", -}); - -const usageResetCreditDetailSchema = type({ - "grantedAt?": "string", - "expiresAt?": "string", - "status?": "string", -}); - -const usageResetCreditsSchema = type({ - availableCount: "number", - "credits?": usageResetCreditDetailSchema.array(), -}); - -const arkUsageReportSchema = type({ - provider: "string", - fetchedAt: "number", - limits: usageLimitSchema.array(), - "resetCredits?": usageResetCreditsSchema, - "notes?": "string[]", - "metadata?": { "[string]": "unknown" }, - "raw?": "unknown", -}); - -/** - * Broker `/v1/usage` response. Reports are full {@link UsageReport}s minus the - * heavy provider-specific `raw` field (the server strips it before send) — we - * keep `raw` optional in the underlying schema so a misconfigured broker that - * forgot to strip still validates. - */ -export const usageResponseSchema = type({ - "+": "reject", - generatedAt: "number", - reports: arkUsageReportSchema.array(), -}); - -const usageHistoryEntrySchema = type({ - recordedAt: "number", - provider: "string", - accountKey: "string", - "email?": "string", - "accountId?": "string", - limitId: "string", - label: "string", - "windowLabel?": "string", - "usedFraction?": "number", - "status?": "'ok' | 'warning' | 'exhausted' | 'unknown'", - "resetsAt?": "number", -}); - -/** Broker `/v1/usage/history` response — recorded usage-limit snapshots, oldest first. */ -export const usageHistoryResponseSchema = type({ - "+": "reject", - generatedAt: "number", - entries: usageHistoryEntrySchema.array(), -}); - -const observedUsageEntrySchema = type({ - at: "number", - provider: "string", - model: "string", - requests: "number", - inputTokens: "number", - outputTokens: "number", - cacheReadTokens: "number", - cacheWriteTokens: "number", - costUsd: "number", -}); - -/** Broker `POST /v1/usage/observed` request — one client's batched observed usage. */ -export const clientUsageReportRequestSchema = type({ - "+": "reject", - installId: "string", - "hostname?": "string", - entries: observedUsageEntrySchema.array(), -}); - -export const clientUsageReportResponseSchema = type({ - "+": "reject", - ok: "boolean", -}); - -const clientProviderUsageSchema = type({ - provider: "string", - requests: "number", - inputTokens: "number", - outputTokens: "number", - cacheReadTokens: "number", - cacheWriteTokens: "number", - costUsd: "number", -}); - -const clientUsageClientSummarySchema = type({ - installId: "string", - "hostname?": "string", - firstSeen: "number", - lastSeen: "number", - providers: clientProviderUsageSchema.array(), -}); - -/** Broker `GET /v1/usage/clients` response — per-client token burn aggregates. */ -export const clientUsageSummaryResponseSchema = type({ - "+": "reject", - generatedAt: "number", - clients: clientUsageClientSummarySchema.array(), -}); - -// ─── Refresh ─────────────────────────────────────────────────────────────── - -export const credentialRefreshResponseSchema = type({ - "+": "reject", - entry: credentialSnapshotEntrySchema, -}); - -// ─── Disable ─────────────────────────────────────────────────────────────── - -export const credentialDisableRequestSchema = type({ - "+": "reject", - "cause?": "string", -}); - -export const credentialDisableResponseSchema = type({ - "+": "reject", - ok: "boolean", -}); - -/** One disabled-credential tombstone — identity + cause, never token material. */ -export const disabledCredentialSummarySchema = type({ - "+": "reject", - id: "number.integer", - provider: type("string").atLeastLength(1), - type: "'oauth' | 'api_key'", - "email?": "string", - "accountId?": "string", - "orgId?": "string", - "orgName?": "string", - cause: "string", - "disabledAtMs?": "number", -}); - -/** Broker `GET /v1/credentials/disabled` response. */ -export const disabledCredentialsResponseSchema = type({ - "+": "reject", - generatedAt: "number", - disabled: disabledCredentialSummarySchema.array(), -}); - -// ─── Credential blocks ────────────────────────────────────────────────────── - -export const credentialBlockRequestSchema = credentialBlockSnapshotSchema; - -export const credentialBlockResponseSchema = type({ - "+": "reject", - ok: "boolean", -}); - -export const credentialBlocksDeleteResponseSchema = type({ - "+": "reject", - ok: "boolean", -}); - -export const usageStaleResponseSchema = type({ - "+": "reject", - ok: "boolean", -}); - -// ─── Upload ──────────────────────────────────────────────────────────────── - -export const credentialUploadRequestSchema = type({ - "+": "reject", - provider: type("string").atLeastLength(1), - credential: writableAuthCredentialSchema, -}); - -export const credentialUploadResponseSchema = type({ - "+": "reject", - entries: credentialSnapshotEntrySchema.array(), -}); +export const { + oauthCredentialSchema, + remoteOauthCredentialSchema, + apiKeyCredentialSchema, + writableAuthCredentialSchema, + snapshotCredentialSchema, + credentialSnapshotEntrySchema, + credentialBlockSnapshotSchema, + snapshotEntrySchema, + refresherScheduleSchema, + snapshotResponseSchema, + snapshotStreamSnapshotEventSchema, + snapshotStreamEntryEventSchema, + snapshotStreamRemovedEventSchema, + snapshotStreamEventSchema, + healthzResponseSchema, + usageResponseSchema, + usageHistoryResponseSchema, + clientUsageReportRequestSchema, + clientUsageReportResponseSchema, + clientUsageSummaryResponseSchema, + credentialRefreshResponseSchema, + credentialDisableRequestSchema, + credentialDisableResponseSchema, + disabledCredentialSummarySchema, + disabledCredentialsResponseSchema, + credentialBlockRequestSchema, + credentialBlockResponseSchema, + credentialBlocksDeleteResponseSchema, + usageStaleResponseSchema, + credentialUploadRequestSchema, + credentialUploadResponseSchema, +} = wireSchemas; diff --git a/packages/ai/test/auth-broker-wire-lazy-construction.test.ts b/packages/ai/test/auth-broker-wire-lazy-construction.test.ts new file mode 100644 index 000000000..f3ada7990 --- /dev/null +++ b/packages/ai/test/auth-broker-wire-lazy-construction.test.ts @@ -0,0 +1,37 @@ +import { expect, test } from "bun:test"; +import * as path from "node:path"; +import { TempDir } from "@oh-my-pi/pi-utils"; + +const preloadPath = path.join(import.meta.dir, "fixtures", "auth-broker-wire-construction-preload.ts"); +const probePath = path.join(import.meta.dir, "fixtures", "auth-broker-wire-construction-probe.ts"); + +test("auth-broker wire schemas construct only on first validation", async () => { + const tempDir = TempDir.createSync("@auth-broker-wire-"); + try { + const proc = Bun.spawn([process.execPath, "--preload", preloadPath, probePath, tempDir.path()], { + cwd: path.join(import.meta.dir, "../../.."), + stdout: "pipe", + stderr: "pipe", + }); + const [stdout, stderr, exitCode] = await Promise.all([ + new Response(proc.stdout).text(), + new Response(proc.stderr).text(), + proc.exited, + ]); + + expect(exitCode, stderr).toBe(0); + expect(JSON.parse(stdout)).toEqual({ + counts: { + afterModuleImport: 0, + afterLocalDiscovery: 0, + afterConstruction: 0, + afterFirstHealth: 1, + afterSecondHealth: 1, + }, + firstHealth: { ok: true, version: "wire-lazy-probe" }, + secondHealth: { ok: true, version: "wire-lazy-probe" }, + }); + } finally { + await tempDir.remove().catch(() => {}); + } +}); diff --git a/packages/ai/test/auth-broker-wire-schema-contract.test.ts b/packages/ai/test/auth-broker-wire-schema-contract.test.ts new file mode 100644 index 000000000..428b42b58 --- /dev/null +++ b/packages/ai/test/auth-broker-wire-schema-contract.test.ts @@ -0,0 +1,267 @@ +import { describe, expect, test } from "bun:test"; +import * as wireSchemas from "@oh-my-pi/pi-ai/auth-broker/wire-schemas"; +import { Type, type } from "arktype"; + +const REFRESHER = { + enabled: false, + intervalMs: 60_000, + skewMs: 300_000, + nextSweepInMs: Number.MAX_SAFE_INTEGER, +}; +const BLOCK = { + providerKey: "anthropic:oauth", + blockScope: "tier:fable", + blockedUntilMs: 4_000, + updatedAtMs: 3_000, +}; +const REMOTE_OAUTH = { + type: "oauth", + access: "access", + refresh: "__remote__", + expires: 5_000, + tokenUrl: "https://example.test/token", + clientId: "provider-client", +}; +const REAL_OAUTH = { ...REMOTE_OAUTH, refresh: "real-refresh" }; +const API_KEY = { type: "api_key", key: "secret", source: "login" }; +const CREDENTIAL_ENTRY = { + id: 7, + provider: "anthropic", + credential: REMOTE_OAUTH, + identityKey: "account:test", +}; +const SNAPSHOT_ENTRY = { + ...CREDENTIAL_ENTRY, + rotatesInMs: null, + blocks: [BLOCK], +}; +const SNAPSHOT = { + generation: 2, + generatedAt: 1_000, + serverNowMs: 2_000, + refresher: REFRESHER, + credentials: [SNAPSHOT_ENTRY], +}; +const STREAM_SNAPSHOT = { kind: "snapshot", ...SNAPSHOT }; +const STREAM_ENTRY = { + kind: "entry", + generation: 3, + serverNowMs: 2_500, + refresher: REFRESHER, + entry: SNAPSHOT_ENTRY, +}; +const STREAM_REMOVED = { + kind: "removed", + generation: 4, + serverNowMs: 3_000, + refresher: REFRESHER, + id: 7, +}; +const USAGE_REPORT = { + provider: "anthropic", + fetchedAt: 2_000, + limits: [ + { + id: "rolling", + label: "Rolling window", + scope: { provider: "anthropic", windowId: "rolling", providerExtension: true }, + window: { id: "rolling", label: "5 hour", durationMs: 18_000_000 }, + amount: { used: 1, limit: 10, remaining: 9, unit: "tokens", providerExtension: "kept" }, + status: "ok", + notes: ["limit note"], + providerExtension: 1, + }, + ], + notes: ["report note"], + metadata: { plan: "max" }, + raw: { providerPayload: true }, + providerExtension: "kept", +}; +const OBSERVED_USAGE = { + at: 1_000, + provider: "anthropic", + model: "claude", + requests: 1, + inputTokens: 2, + outputTokens: 3, + cacheReadTokens: 4, + cacheWriteTokens: 5, + costUsd: 0.01, +}; + +const schemaNames = [ + "oauthCredentialSchema", + "remoteOauthCredentialSchema", + "apiKeyCredentialSchema", + "writableAuthCredentialSchema", + "snapshotCredentialSchema", + "credentialSnapshotEntrySchema", + "credentialBlockSnapshotSchema", + "snapshotEntrySchema", + "refresherScheduleSchema", + "snapshotResponseSchema", + "snapshotStreamSnapshotEventSchema", + "snapshotStreamEntryEventSchema", + "snapshotStreamRemovedEventSchema", + "snapshotStreamEventSchema", + "healthzResponseSchema", + "usageResponseSchema", + "usageHistoryResponseSchema", + "clientUsageReportRequestSchema", + "clientUsageReportResponseSchema", + "clientUsageSummaryResponseSchema", + "credentialRefreshResponseSchema", + "credentialDisableRequestSchema", + "credentialDisableResponseSchema", + "disabledCredentialSummarySchema", + "disabledCredentialsResponseSchema", + "credentialBlockRequestSchema", + "credentialBlockResponseSchema", + "credentialBlocksDeleteResponseSchema", + "usageStaleResponseSchema", + "credentialUploadRequestSchema", + "credentialUploadResponseSchema", +] as const; + +type SchemaName = (typeof schemaNames)[number]; +type CallableSchema = (input: unknown) => unknown; + +const validSamples: Record = { + oauthCredentialSchema: REAL_OAUTH, + remoteOauthCredentialSchema: REMOTE_OAUTH, + apiKeyCredentialSchema: API_KEY, + writableAuthCredentialSchema: REAL_OAUTH, + snapshotCredentialSchema: REMOTE_OAUTH, + credentialSnapshotEntrySchema: CREDENTIAL_ENTRY, + credentialBlockSnapshotSchema: BLOCK, + snapshotEntrySchema: SNAPSHOT_ENTRY, + refresherScheduleSchema: REFRESHER, + snapshotResponseSchema: SNAPSHOT, + snapshotStreamSnapshotEventSchema: STREAM_SNAPSHOT, + snapshotStreamEntryEventSchema: STREAM_ENTRY, + snapshotStreamRemovedEventSchema: STREAM_REMOVED, + snapshotStreamEventSchema: STREAM_ENTRY, + healthzResponseSchema: { ok: true, version: "contract" }, + usageResponseSchema: { generatedAt: 2_000, reports: [USAGE_REPORT] }, + usageHistoryResponseSchema: { + generatedAt: 2_000, + entries: [ + { + recordedAt: 1_000, + provider: "anthropic", + accountKey: "account:test", + limitId: "rolling", + label: "Rolling window", + usedFraction: 0.1, + status: "ok", + }, + ], + }, + clientUsageReportRequestSchema: { installId: "install", hostname: "host", entries: [OBSERVED_USAGE] }, + clientUsageReportResponseSchema: { ok: true }, + clientUsageSummaryResponseSchema: { + generatedAt: 2_000, + clients: [ + { + installId: "install", + hostname: "host", + firstSeen: 1_000, + lastSeen: 2_000, + providers: [{ ...OBSERVED_USAGE, firstSeen: undefined, at: undefined, model: undefined }], + }, + ], + }, + credentialRefreshResponseSchema: { entry: CREDENTIAL_ENTRY }, + credentialDisableRequestSchema: {}, + credentialDisableResponseSchema: { ok: true }, + disabledCredentialSummarySchema: { + id: 7, + provider: "anthropic", + type: "oauth", + email: "user@example.test", + cause: "revoked", + disabledAtMs: 2_000, + }, + disabledCredentialsResponseSchema: { + generatedAt: 2_000, + disabled: [{ id: 7, provider: "anthropic", type: "oauth", cause: "revoked" }], + }, + credentialBlockRequestSchema: BLOCK, + credentialBlockResponseSchema: { ok: true }, + credentialBlocksDeleteResponseSchema: { ok: true }, + usageStaleResponseSchema: { ok: true }, + credentialUploadRequestSchema: { provider: "anthropic", credential: REAL_OAUTH }, + credentialUploadResponseSchema: { entries: [CREDENTIAL_ENTRY] }, +}; + +function run(schema: unknown, input: unknown): unknown { + return (schema as CallableSchema)(input); +} + +function accept(schema: unknown, input: unknown): unknown { + const result = run(schema, input); + expect(result).not.toBeInstanceOf(type.errors); + if (result instanceof type.errors) throw new Error(`Expected schema acceptance: ${result.summary}`); + return result; +} + +function reject(schema: unknown, input: unknown): void { + expect(run(schema, input)).toBeInstanceOf(type.errors); +} + +describe("auth-broker public wire schemas", () => { + test("exports all 31 real callable ArkType values with canonical behavior", () => { + expect(Object.keys(wireSchemas).sort()).toEqual([...schemaNames].sort()); + for (const name of schemaNames) { + // biome-ignore lint/performance/noDynamicNamespaceImportAccess: this contract intentionally verifies the public namespace. + const schema = wireSchemas[name]; + expect(typeof schema).toBe("function"); + expect(schema).toBeInstanceOf(Type); + accept(schema, validSamples[name]); + } + }); + + test("preserves credential extension and sentinel boundaries", () => { + expect(accept(wireSchemas.oauthCredentialSchema, REAL_OAUTH)).toEqual(REAL_OAUTH); + expect(accept(wireSchemas.remoteOauthCredentialSchema, REMOTE_OAUTH)).toEqual(REMOTE_OAUTH); + reject(wireSchemas.oauthCredentialSchema, REMOTE_OAUTH); + reject(wireSchemas.remoteOauthCredentialSchema, REAL_OAUTH); + reject(wireSchemas.oauthCredentialSchema, { ...REAL_OAUTH, access: "" }); + reject(wireSchemas.apiKeyCredentialSchema, { ...API_KEY, extra: true }); + reject(wireSchemas.apiKeyCredentialSchema, { ...API_KEY, source: "environment" }); + reject(wireSchemas.credentialUploadRequestSchema, { provider: "", credential: REAL_OAUTH }); + }); + + test("preserves fixed envelopes, integer fields, discriminators, and block alias identity", () => { + expect(wireSchemas.credentialBlockRequestSchema).toBe(wireSchemas.credentialBlockSnapshotSchema); + accept(wireSchemas.credentialBlockRequestSchema, { + providerKey: BLOCK.providerKey, + blockScope: "", + blockedUntilMs: BLOCK.blockedUntilMs, + }); + accept(wireSchemas.credentialDisableRequestSchema, {}); + reject(wireSchemas.credentialDisableRequestSchema, { cause: 1 }); + reject(wireSchemas.credentialDisableRequestSchema, { extra: true }); + reject(wireSchemas.healthzResponseSchema, { ok: true, extra: true }); + reject(wireSchemas.snapshotResponseSchema, { ...SNAPSHOT, generation: 1.5 }); + reject(wireSchemas.snapshotResponseSchema, { ...SNAPSHOT, extra: true }); + reject(wireSchemas.snapshotEntrySchema, { ...SNAPSHOT_ENTRY, id: 1.5 }); + reject(wireSchemas.snapshotStreamEventSchema, { ...STREAM_ENTRY, kind: "snapshot" }); + reject(wireSchemas.snapshotStreamEventSchema, { ...STREAM_REMOVED, extra: true }); + reject(wireSchemas.snapshotStreamEventSchema, { kind: "unknown" }); + }); + + test("preserves usage extensions while rejecting envelope and enum violations", () => { + const response = { generatedAt: 2_000, reports: [USAGE_REPORT] }; + expect(accept(wireSchemas.usageResponseSchema, response)).toEqual(response); + reject(wireSchemas.usageResponseSchema, { ...response, extra: true }); + reject(wireSchemas.usageResponseSchema, { + ...response, + reports: [{ ...USAGE_REPORT, limits: [{ ...USAGE_REPORT.limits[0], amount: { unit: "seconds" } }] }], + }); + reject(wireSchemas.usageResponseSchema, { + ...response, + reports: [{ ...USAGE_REPORT, limits: [{ ...USAGE_REPORT.limits[0], status: "critical" }] }], + }); + }); +}); diff --git a/packages/ai/test/fixtures/auth-broker-wire-construction-preload.ts b/packages/ai/test/fixtures/auth-broker-wire-construction-preload.ts new file mode 100644 index 000000000..e3be08634 --- /dev/null +++ b/packages/ai/test/fixtures/auth-broker-wire-construction-preload.ts @@ -0,0 +1,15 @@ +import { spyOn } from "bun:test"; +import { type } from "arktype"; + +declare global { + var __authBrokerWireConstructionCount: number; +} + +globalThis.__authBrokerWireConstructionCount = 0; +const originalEnumerated = type.enumerated; +spyOn(type, "enumerated").mockImplementation((...values) => { + if (values.length === 1 && values[0] === "__remote__") { + globalThis.__authBrokerWireConstructionCount += 1; + } + return originalEnumerated(...values); +}); diff --git a/packages/ai/test/fixtures/auth-broker-wire-construction-probe.ts b/packages/ai/test/fixtures/auth-broker-wire-construction-probe.ts new file mode 100644 index 000000000..88a0f966c --- /dev/null +++ b/packages/ai/test/fixtures/auth-broker-wire-construction-probe.ts @@ -0,0 +1,55 @@ +import { + AuthBrokerClient, + type AuthBrokerServerHandle, + discoverAuthStorage, + startAuthBroker, +} from "@oh-my-pi/pi-ai/auth-broker"; +import type { AuthStorage } from "@oh-my-pi/pi-ai/auth-storage"; + +declare global { + var __authBrokerWireConstructionCount: number; +} + +const agentDir = process.argv[2]; +if (!agentDir) throw new Error("Expected an isolated agent directory"); + +const count = (): number => globalThis.__authBrokerWireConstructionCount; + +const counts = { + afterModuleImport: count(), + afterLocalDiscovery: -1, + afterConstruction: -1, + afterFirstHealth: -1, + afterSecondHealth: -1, +}; +let storage: AuthStorage | undefined; +let handle: AuthBrokerServerHandle | undefined; + +try { + delete process.env.OMP_AUTH_BROKER_URL; + delete process.env.OMP_AUTH_BROKER_TOKEN; + delete process.env.OMP_AUTH_BROKER_ACCOUNT_POOL_FILE; + + storage = await discoverAuthStorage({ agentDir }); + counts.afterLocalDiscovery = count(); + + handle = startAuthBroker({ + storage, + bind: "127.0.0.1:0", + bearerTokens: [], + version: "wire-lazy-probe", + disableRefresher: true, + }); + const client = new AuthBrokerClient({ url: handle.url, token: "unused", maxRetries: 0 }); + counts.afterConstruction = count(); + + const firstHealth = await client.healthz(); + counts.afterFirstHealth = count(); + const secondHealth = await client.healthz(); + counts.afterSecondHealth = count(); + + process.stdout.write(JSON.stringify({ counts, firstHealth, secondHealth })); +} finally { + await handle?.close(); + storage?.close(); +}