@oh-my-pi/pi-ai 16.3.4 → 16.3.6

package/src/auth-broker/remote-store.ts

@@ -16,13 +16,20 @@ import {
 	type OAuthCredential,
 	REMOTE_REFRESH_SENTINEL,
 	type StoredAuthCredential,
+	type StoredCredentialBlock,
 } from "../auth-storage";
 import * as AIError from "../error";
 import type { OAuthCredentials } from "../registry/oauth/types";
 import type { Provider } from "../types";
 import type { UsageReport } from "../usage";
 import { type AuthBrokerClient, AuthBrokerStreamUnsupportedError } from "./client";
-import type { RefresherSchedule, SnapshotEntry, SnapshotResponse, SnapshotStreamEvent } from "./types";
+import type {
+	CredentialBlockSnapshot,
+	RefresherSchedule,
+	SnapshotEntry,
+	SnapshotResponse,
+	SnapshotStreamEvent,
+} from "./types";
 
 /**
  * Client-side TTL for the aggregate `/v1/usage` response. The broker dedups
@@ -38,6 +45,31 @@ const BACKGROUND_WAIT_MS = 30_000;
 const BACKGROUND_BACKOFF_INITIAL_MS = 500;
 const BACKGROUND_BACKOFF_MAX_MS = 30_000;
 
+function compareCredentialBlockSnapshots(a: CredentialBlockSnapshot, b: CredentialBlockSnapshot): number {
+	const provider = a.providerKey.localeCompare(b.providerKey);
+	if (provider !== 0) return provider;
+	const scope = a.blockScope.localeCompare(b.blockScope);
+	if (scope !== 0) return scope;
+	return a.blockedUntilMs - b.blockedUntilMs;
+}
+
+function toCredentialBlockSnapshot(block: StoredCredentialBlock): CredentialBlockSnapshot {
+	return {
+		providerKey: block.providerKey,
+		blockScope: block.blockScope,
+		blockedUntilMs: block.blockedUntilMs,
+	};
+}
+
+function credentialEntryWithBlocks(
+	entry: AuthCredentialSnapshotEntry,
+	blocks: readonly CredentialBlockSnapshot[] | undefined,
+): SnapshotEntry {
+	const incoming: SnapshotEntry = { ...entry, rotatesInMs: null };
+	if (blocks && blocks.length > 0) incoming.blocks = [...blocks].sort(compareCredentialBlockSnapshots);
+	return incoming;
+}
+
 function emptySnapshot(): SnapshotResponse {
 	return {
 		generation: 0,
@@ -169,13 +201,17 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 	}
 
 	#applySnapshot(snapshot: SnapshotResponse, generation: number): void {
-		this.#snapshot = snapshot;
+		const nowMs = Date.now();
+		this.#snapshot = {
+			...snapshot,
+			credentials: snapshot.credentials.map(entry => this.#normalizeSnapshotEntryBlocks(entry, nowMs)),
+		};
 		this.#generation = generation;
-		this.#snapshotReceivedAt = Date.now();
+		this.#snapshotReceivedAt = nowMs;
 		const onSnapshot = this.#onSnapshot;
 		if (!onSnapshot) return;
 		try {
-			onSnapshot(snapshot, generation);
+			onSnapshot(this.#snapshot, generation);
 		} catch (error) {
 			logger.debug("auth-broker snapshot callback failed", { error: String(error) });
 		}
@@ -266,11 +302,12 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		generation: number,
 		serverNowMs: number,
 	): void {
-		const index = this.#snapshot.credentials.findIndex(candidate => candidate.id === entry.id);
+		const incoming = this.#normalizeSnapshotEntryBlocks(entry, Date.now());
+		const index = this.#snapshot.credentials.findIndex(candidate => candidate.id === incoming.id);
 		const credentials =
 			index === -1
-				? [...this.#snapshot.credentials, entry]
-				: this.#snapshot.credentials.map((candidate, i) => (i === index ? entry : candidate));
+				? [...this.#snapshot.credentials, incoming]
+				: this.#snapshot.credentials.map((candidate, i) => (i === index ? incoming : candidate));
 		this.#snapshot = { ...this.#snapshot, generation, serverNowMs, refresher, credentials };
 		this.#generation = generation;
 		this.#snapshotReceivedAt = Date.now();
@@ -304,6 +341,76 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		return out;
 	}
 
+	getCredentialBlock(credentialId: number, providerKey: string, blockScope: string): number | undefined {
+		const nowMs = Date.now();
+		this.cleanExpiredCredentialBlocks(nowMs);
+		const entry = this.#snapshot.credentials.find(candidate => candidate.id === credentialId);
+		if (!entry?.blocks) return undefined;
+		const block = entry.blocks.find(
+			candidate => candidate.providerKey === providerKey && candidate.blockScope === blockScope,
+		);
+		if (!block || block.blockedUntilMs <= nowMs) return undefined;
+		return block.blockedUntilMs;
+	}
+
+	listCredentialBlocks(credentialIds: readonly number[]): StoredCredentialBlock[] {
+		const nowMs = Date.now();
+		this.cleanExpiredCredentialBlocks(nowMs);
+		const ids = new Set(credentialIds);
+		const blocks: StoredCredentialBlock[] = [];
+		for (const entry of this.#snapshot.credentials) {
+			if (!ids.has(entry.id) || !entry.blocks) continue;
+			for (const block of entry.blocks) {
+				if (block.blockedUntilMs <= nowMs) continue;
+				blocks.push({
+					credentialId: entry.id,
+					providerKey: block.providerKey,
+					blockScope: block.blockScope,
+					blockedUntilMs: block.blockedUntilMs,
+				});
+			}
+		}
+		blocks.sort((a, b) => a.credentialId - b.credentialId || compareCredentialBlockSnapshots(a, b));
+		return blocks;
+	}
+
+	upsertCredentialBlock(block: StoredCredentialBlock): void {
+		this.#upsertSnapshotBlock(block);
+		const body = toCredentialBlockSnapshot(block);
+		void this.#client
+			.upsertCredentialBlock(block.credentialId, body)
+			.then(() => {
+				this.#maybeRefreshSnapshot("credential block");
+			})
+			.catch(error => {
+				logger.warn("auth-broker credential block propagation failed", {
+					id: block.credentialId,
+					providerKey: block.providerKey,
+					blockScope: block.blockScope,
+					error: String(error),
+				});
+			});
+	}
+
+	deleteCredentialBlocks(credentialId: number): void {
+		this.#deleteSnapshotBlocks(credentialId);
+		void this.#client
+			.deleteCredentialBlocks(credentialId)
+			.then(() => {
+				this.#maybeRefreshSnapshot("credential blocks delete");
+			})
+			.catch(error => {
+				logger.warn("auth-broker credential blocks delete propagation failed", {
+					id: credentialId,
+					error: String(error),
+				});
+			});
+	}
+
+	cleanExpiredCredentialBlocks(nowMs: number): void {
+		this.#pruneExpiredCredentialBlocks(nowMs);
+	}
+
 	/**
 	 * In-memory update from a successful refresh through the broker. AuthStorage
 	 * calls this after `#replaceCredentialAt`; the broker already persisted the
@@ -459,13 +566,19 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		// `entries` is the broker's authoritative post-upsert list of rows for
 		// `provider`. Drop our existing rows for the same provider and splice in
 		// the fresh set — preserving every other provider's rows in place.
+		const existingBlocks = new Map(
+			this.#snapshot.credentials
+				.filter(entry => entry.provider === provider && entry.blocks !== undefined)
+				.map(entry => [entry.id, entry.blocks] as const),
+		);
 		const others = this.#snapshot.credentials.filter(entry => entry.provider !== provider);
-		const incoming = entries.map(entry => ({ ...entry, rotatesInMs: null }));
+		const incoming = entries.map(entry => credentialEntryWithBlocks(entry, existingBlocks.get(entry.id)));
 		this.#snapshot = { ...this.#snapshot, credentials: [...others, ...incoming] };
 	}
 	#applyCredentialEntry(entry: AuthCredentialSnapshotEntry): void {
-		const incoming = { ...entry, rotatesInMs: null };
 		const index = this.#snapshot.credentials.findIndex(candidate => candidate.id === entry.id);
+		const existingBlocks = index === -1 ? undefined : this.#snapshot.credentials[index]?.blocks;
+		const incoming = credentialEntryWithBlocks(entry, existingBlocks);
 		if (index === -1) {
 			this.#snapshot = { ...this.#snapshot, credentials: [...this.#snapshot.credentials, incoming] };
 			return;
@@ -485,6 +598,73 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		this.#snapshot = { ...this.#snapshot, credentials: next };
 	}
 
+	#normalizeSnapshotEntryBlocks(entry: SnapshotEntry, nowMs: number): SnapshotEntry {
+		if (!entry.blocks || entry.blocks.length === 0) return entry;
+		const blocks = entry.blocks
+			.filter(block => block.blockedUntilMs > nowMs)
+			.map(block => ({
+				providerKey: block.providerKey,
+				blockScope: block.blockScope,
+				blockedUntilMs: block.blockedUntilMs,
+			}))
+			.sort(compareCredentialBlockSnapshots);
+		if (blocks.length > 0) return { ...entry, blocks };
+		const next: SnapshotEntry = { ...entry };
+		delete next.blocks;
+		return next;
+	}
+
+	#upsertSnapshotBlock(block: StoredCredentialBlock): void {
+		const index = this.#snapshot.credentials.findIndex(entry => entry.id === block.credentialId);
+		if (index === -1) return;
+		const entry = this.#snapshot.credentials[index]!;
+		const incoming = toCredentialBlockSnapshot(block);
+		const blocks = entry.blocks ? [...entry.blocks] : [];
+		const blockIndex = blocks.findIndex(
+			candidate => candidate.providerKey === incoming.providerKey && candidate.blockScope === incoming.blockScope,
+		);
+		if (blockIndex === -1) {
+			blocks.push(incoming);
+		} else {
+			const existing = blocks[blockIndex]!;
+			blocks[blockIndex] = {
+				...existing,
+				blockedUntilMs: Math.max(existing.blockedUntilMs, incoming.blockedUntilMs),
+			};
+		}
+		blocks.sort(compareCredentialBlockSnapshots);
+		const credentials = [...this.#snapshot.credentials];
+		credentials[index] = { ...entry, blocks };
+		this.#snapshot = { ...this.#snapshot, credentials };
+	}
+
+	#deleteSnapshotBlocks(credentialId: number): void {
+		const index = this.#snapshot.credentials.findIndex(entry => entry.id === credentialId);
+		if (index === -1) return;
+		const entry = this.#snapshot.credentials[index]!;
+		if (!entry.blocks || entry.blocks.length === 0) return;
+		const next: SnapshotEntry = { ...entry };
+		delete next.blocks;
+		const credentials = [...this.#snapshot.credentials];
+		credentials[index] = next;
+		this.#snapshot = { ...this.#snapshot, credentials };
+	}
+
+	#pruneExpiredCredentialBlocks(nowMs: number): void {
+		let changed = false;
+		const credentials = this.#snapshot.credentials.map(entry => {
+			if (!entry.blocks || entry.blocks.length === 0) return entry;
+			const blocks = entry.blocks.filter(block => block.blockedUntilMs > nowMs);
+			if (blocks.length === entry.blocks.length) return entry;
+			changed = true;
+			if (blocks.length > 0) return { ...entry, blocks };
+			const next: SnapshotEntry = { ...entry };
+			delete next.blocks;
+			return next;
+		});
+		if (changed) this.#snapshot = { ...this.#snapshot, credentials };
+	}
+
 	/**
 	 * Fire-and-forget `refreshSnapshot()` after a write. When the SSE stream is
 	 * active the broker will deliver the new generation push, so the extra GET

package/src/auth-broker/server.ts

@@ -11,10 +11,13 @@
  */
 import { logger } from "@oh-my-pi/pi-utils";
 import { type Type, type } from "arktype";
-import type { AuthStorage } from "../auth-storage";
+import type { AuthStorage, StoredCredentialBlock } from "../auth-storage";
 import { parseBind } from "../utils/parse-bind";
 import { AuthBrokerRefresher, type AuthBrokerRefresherSchedule } from "./refresher";
 import type {
+	CredentialBlockResponse,
+	CredentialBlockSnapshot,
+	CredentialBlocksDeleteResponse,
 	CredentialDisableResponse,
 	CredentialRefreshResponse,
 	CredentialUploadResponse,
@@ -33,7 +36,11 @@ import {
 	DEFAULT_SERVER_IDLE_TIMEOUT_S,
 	DEFAULT_STREAM_KEEPALIVE_MS,
 } from "./types";
-import { credentialDisableRequestSchema, credentialUploadRequestSchema } from "./wire-schemas";
+import {
+	credentialBlockRequestSchema,
+	credentialDisableRequestSchema,
+	credentialUploadRequestSchema,
+} from "./wire-schemas";
 
 export interface AuthBrokerServerOptions {
 	/** Underlying credential storage (wraps the local SQLite store on the broker). */
@@ -120,6 +127,8 @@ async function parseBody<t>(
 
 const REFRESH_ROUTE = /^\/v1\/credential\/(\d+)\/refresh$/;
 const DISABLE_ROUTE = /^\/v1\/credential\/(\d+)\/disable$/;
+const BLOCK_ROUTE = /^\/v1\/credential\/(\d+)\/block$/;
+const BLOCKS_ROUTE = /^\/v1\/credential\/(\d+)\/blocks$/;
 
 const MAX_SNAPSHOT_WAIT_MS = 30_000;
 const DISABLED_NEXT_SWEEP_IN_MS = Number.MAX_SAFE_INTEGER;
@@ -262,14 +271,48 @@ function computeRotatesInMs(
 	return Math.max(0, rotatesAt - serverNowMs);
 }
 
+function compareCredentialBlockSnapshots(a: CredentialBlockSnapshot, b: CredentialBlockSnapshot): number {
+	const provider = a.providerKey.localeCompare(b.providerKey);
+	if (provider !== 0) return provider;
+	const scope = a.blockScope.localeCompare(b.blockScope);
+	if (scope !== 0) return scope;
+	return a.blockedUntilMs - b.blockedUntilMs;
+}
+
+function buildCredentialBlockGroups(
+	blocks: readonly StoredCredentialBlock[],
+	serverNowMs: number,
+): Map<number, CredentialBlockSnapshot[]> {
+	const byCredentialId = new Map<number, CredentialBlockSnapshot[]>();
+	for (const block of blocks) {
+		if (block.blockedUntilMs <= serverNowMs) continue;
+		const snapshotBlock: CredentialBlockSnapshot = {
+			providerKey: block.providerKey,
+			blockScope: block.blockScope,
+			blockedUntilMs: block.blockedUntilMs,
+		};
+		const existing = byCredentialId.get(block.credentialId);
+		if (existing) {
+			existing.push(snapshotBlock);
+		} else {
+			byCredentialId.set(block.credentialId, [snapshotBlock]);
+		}
+	}
+	for (const credentialBlocks of byCredentialId.values()) credentialBlocks.sort(compareCredentialBlockSnapshots);
+	return byCredentialId;
+}
+
 function buildSnapshot(storage: AuthStorage, refresher: AuthBrokerRefresher | undefined): SnapshotResponse {
 	const serverNowMs = Date.now();
 	const base = storage.exportSnapshot();
 	const { wire, nextSweepAt } = resolveRefresherSchedule(refresher, serverNowMs);
-	const credentials: SnapshotEntry[] = base.credentials.map(entry => ({
-		...entry,
-		rotatesInMs: computeRotatesInMs(entry, wire, nextSweepAt, serverNowMs),
-	}));
+	const credentialIds = base.credentials.map(entry => entry.id);
+	const blocksByCredentialId = buildCredentialBlockGroups(storage.listCredentialBlocks(credentialIds), serverNowMs);
+	const credentials: SnapshotEntry[] = base.credentials.map(entry => {
+		const blocks = blocksByCredentialId.get(entry.id);
+		const rotatesInMs = computeRotatesInMs(entry, wire, nextSweepAt, serverNowMs);
+		return blocks && blocks.length > 0 ? { ...entry, rotatesInMs, blocks } : { ...entry, rotatesInMs };
+	});
 	return {
 		generation: base.generation,
 		generatedAt: base.generatedAt,
@@ -336,7 +379,14 @@ async function serveSnapshot(
  * keep the stale projection.
  */
 function fingerprintEntry(entry: SnapshotEntry): string {
-	return JSON.stringify([entry.id, entry.provider, entry.identityKey, entry.rotatesInMs, entry.credential]);
+	return JSON.stringify([
+		entry.id,
+		entry.provider,
+		entry.identityKey,
+		entry.rotatesInMs,
+		entry.credential,
+		entry.blocks ?? [],
+	]);
 }
 
 function sseEvent(event: string, body: unknown): string {
@@ -593,6 +643,58 @@ export function startAuthBroker(opts: AuthBrokerServerOptions): AuthBrokerServer
 					const response: CredentialDisableResponse = { ok: true };
 					return json(200, response);
 				}
+				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);
+					if (!parsed.ok) return parsed.response;
+					const block: StoredCredentialBlock = {
+						credentialId: id,
+						providerKey: parsed.data.providerKey,
+						blockScope: parsed.data.blockScope,
+						blockedUntilMs: parsed.data.blockedUntilMs,
+					};
+					if (!opts.storage.exportSnapshot().credentials.some(entry => entry.id === id)) {
+						logger.info("auth-broker credential block miss", { id, peer });
+						return json(404, { error: `No credential with id=${id}` });
+					}
+					try {
+						opts.storage.upsertCredentialBlock(block);
+						const response: CredentialBlockResponse = { ok: true };
+						logger.info("auth-broker credential block upserted", {
+							id,
+							peer,
+							providerKey: block.providerKey,
+							blockScope: block.blockScope,
+							blockedUntilMs: block.blockedUntilMs,
+						});
+						return json(200, response);
+					} catch (error) {
+						const message = error instanceof Error ? error.message : String(error);
+						logger.warn("auth-broker credential block upsert failed", { id, peer, error: message });
+						const status = message.includes("No credential with id") ? 404 : 500;
+						return json(status, { error: message });
+					}
+				}
+				const blocksDeleteMatch = req.method === "DELETE" ? pathname.match(BLOCKS_ROUTE) : null;
+				if (blocksDeleteMatch) {
+					const id = Number.parseInt(blocksDeleteMatch[1], 10);
+					if (!opts.storage.exportSnapshot().credentials.some(entry => entry.id === id)) {
+						logger.info("auth-broker credential blocks delete miss", { id, peer });
+						return json(404, { error: `No credential with id=${id}` });
+					}
+					try {
+						opts.storage.deleteCredentialBlocks(id);
+						const response: CredentialBlocksDeleteResponse = { ok: true };
+						logger.info("auth-broker credential blocks deleted", { id, peer });
+						return json(200, response);
+					} catch (error) {
+						const message = error instanceof Error ? error.message : String(error);
+						logger.warn("auth-broker credential blocks delete failed", { id, peer, error: message });
+						const status = message.includes("No credential with id") ? 404 : 500;
+						return json(status, { error: message });
+					}
+				}
 				if (req.method === "POST" && pathname === "/v1/credential") {
 					const parsed = await parseBody(req, credentialUploadRequestSchema);
 					if (!parsed.ok) return parsed.response;

package/src/auth-broker/types.ts

@@ -6,7 +6,12 @@
  * credential expires or a 401 surfaces on a supposedly-fresh credential.
  */
 
-import type { AuthCredential, AuthCredentialSnapshot, AuthCredentialSnapshotEntry } from "../auth-storage";
+import type {
+	AuthCredential,
+	AuthCredentialSnapshot,
+	AuthCredentialSnapshotEntry,
+	StoredCredentialBlock,
+} from "../auth-storage";
 import type { UsageReport } from "../usage";
 
 /** GET /v1/healthz response body. */
@@ -22,8 +27,11 @@ export interface RefresherSchedule {
 	nextSweepInMs: number;
 }
 
+export type CredentialBlockSnapshot = Omit<StoredCredentialBlock, "credentialId">;
+
 export type SnapshotEntry = AuthCredentialSnapshotEntry & {
 	rotatesInMs: number | null;
+	blocks?: CredentialBlockSnapshot[];
 };
 
 /** GET /v1/snapshot response body. */
@@ -54,6 +62,19 @@ export interface CredentialDisableResponse {
 	ok: boolean;
 }
 
+/** POST /v1/credential/:id/block request body. */
+export type CredentialBlockRequest = CredentialBlockSnapshot;
+
+/** POST /v1/credential/:id/block response body. */
+export interface CredentialBlockResponse {
+	ok: boolean;
+}
+
+/** DELETE /v1/credential/:id/blocks response body. */
+export interface CredentialBlocksDeleteResponse {
+	ok: boolean;
+}
+
 /**
  * POST /v1/credential request body. The OAuth `refresh` must be the *real*
  * refresh token (not the sentinel) — the broker is the canonical writer.

package/src/auth-broker/wire-schemas.ts

@@ -69,6 +69,13 @@ export const credentialSnapshotEntrySchema = type({
 	identityKey: "string | null",
 });
 
+export const credentialBlockSnapshotSchema = type({
+	"+": "reject",
+	providerKey: type("string").atLeastLength(1),
+	blockScope: "string",
+	blockedUntilMs: "number",
+});
+
 export const snapshotEntrySchema = type({
 	"+": "reject",
 	id: "number.integer",
@@ -76,6 +83,7 @@ export const snapshotEntrySchema = type({
 	credential: snapshotCredentialSchema,
 	identityKey: "string | null",
 	rotatesInMs: "number | null",
+	"blocks?": credentialBlockSnapshotSchema.array(),
 });
 
 export const refresherScheduleSchema = type({
@@ -235,6 +243,20 @@ export const credentialDisableResponseSchema = type({
 	ok: "boolean",
 });
 
+// ─── Credential blocks ──────────────────────────────────────────────────────
+
+export const credentialBlockRequestSchema = credentialBlockSnapshotSchema;
+
+export const credentialBlockResponseSchema = type({
+	"+": "reject",
+	ok: "boolean",
+});
+
+export const credentialBlocksDeleteResponseSchema = type({
+	"+": "reject",
+	ok: "boolean",
+});
+
 // ─── Upload ────────────────────────────────────────────────────────────────
 
 export const credentialUploadRequestSchema = type({