@oh-my-pi/pi-ai 18.2.0 → 18.2.1

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

@@ -255,6 +255,18 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 	#snapshot: SnapshotResponse = emptySnapshot();
 	#snapshotReceivedAt = Date.now();
 	#generation = 0;
+	/**
+	 * Content fingerprint of the credential set in {@link #snapshot} (id +
+	 * provider + credential material), recomputed after every snapshot mutation.
+	 * Drives {@link #credentialRevision} independently of the broker's numeric
+	 * generation, which is an in-memory counter that resets when the broker
+	 * process restarts and so cannot be trusted for change detection.
+	 */
+	#credentialFingerprint = "";
+	/** Monotonic local counter bumped whenever {@link #credentialFingerprint} changes. */
+	#credentialRevision = 0;
+	/** Revision last reported as "seen" by {@link pollExternalChanges}; seeded from the initial snapshot. */
+	#acknowledgedRevision = 0;
 	#usageOverlays: Map<string, UsageReport> = new Map();
 	#backgroundAbort = new AbortController();
 	readonly #backgroundIdleMs: number;
@@ -300,6 +312,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 			? new Map([...opts.accountPool].map(([provider, identities]) => [provider, new Set(identities)]))
 			: undefined;
 		this.#applySnapshot(opts.initialSnapshot ?? emptySnapshot(), opts.initialSnapshot?.generation ?? 0);
+		this.#acknowledgedRevision = this.#credentialRevision;
 		this.#onSnapshot = opts.onSnapshot;
 		void this.#runBackground();
 	}
@@ -325,6 +338,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		this.#snapshot = { ...snapshot, credentials };
 		this.#generation = generation;
 		this.#snapshotReceivedAt = nowMs;
+		this.#refreshCredentialRevision();
 		const onSnapshot = this.#onSnapshot;
 		if (!onSnapshot) return;
 		try {
@@ -333,6 +347,34 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 			logger.debug("auth-broker snapshot callback failed", { error: String(error) });
 		}
 	}
+
+	/**
+	 * Recompute the credential-content fingerprint and bump
+	 * {@link #credentialRevision} when it changes. Called after every snapshot
+	 * mutation so {@link pollExternalChanges} detects add/remove/replace even
+	 * when the broker's numeric generation repeats (e.g. after a broker
+	 * restart resets its in-memory counter).
+	 */
+	#refreshCredentialRevision(): void {
+		const fingerprint = this.#computeCredentialFingerprint();
+		if (fingerprint === this.#credentialFingerprint) return;
+		this.#credentialFingerprint = fingerprint;
+		this.#credentialRevision += 1;
+	}
+
+	/**
+	 * Order-independent digest of the routable credential material — exactly the
+	 * fields {@link listAuthCredentials} exposes (id, provider, credential). A
+	 * token rotation or an add/remove changes it; credential blocks and usage
+	 * overlays do not.
+	 */
+	#computeCredentialFingerprint(): string {
+		const parts = this.#snapshot.credentials.map(
+			entry => `${entry.id}\u0000${entry.provider}\u0000${JSON.stringify(entry.credential)}`,
+		);
+		parts.sort();
+		return parts.join("\u0001");
+	}
 	#protectNewSnapshotBlocks(previous: readonly SnapshotEntry[], next: readonly SnapshotEntry[], nowMs: number): void {
 		const previousBlocksByKey = new Map<string, string>();
 		for (const entry of previous) {
@@ -499,15 +541,13 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 	#applyStreamEvent(event: SnapshotStreamEvent): void {
 		switch (event.kind) {
 			case "snapshot": {
-				// Strip the discriminator so we store the wire-shape SnapshotResponse.
+				// The first frame of every SSE connection is a full authoritative
+				// snapshot. Always adopt it as the new generation baseline: the
+				// broker's in-memory generation counter resets on restart and may
+				// therefore be lower than the previous stream's last value.
+				// Subsequent entry/removal frames remain guarded against reordering
+				// relative to this new baseline below.
 				const { kind: _kind, ...snapshot } = event;
-				if (snapshot.generation < this.#generation) {
-					logger.debug("auth-broker stream snapshot older than local; ignoring", {
-						local: this.#generation,
-						incoming: snapshot.generation,
-					});
-					return;
-				}
 				this.#applySnapshot(snapshot, snapshot.generation);
 				return;
 			}
@@ -548,6 +588,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		this.#snapshot = { ...this.#snapshot, generation, serverNowMs, refresher, credentials };
 		this.#generation = generation;
 		this.#snapshotReceivedAt = Date.now();
+		this.#refreshCredentialRevision();
 	}
 
 	#removeStreamCredential(
@@ -564,6 +605,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		this.#snapshot = { ...this.#snapshot, generation, serverNowMs, refresher, credentials };
 		this.#generation = generation;
 		this.#snapshotReceivedAt = Date.now();
+		this.#refreshCredentialRevision();
 	}
 
 	/** Re-hydrate the in-memory snapshot from the broker. */
@@ -574,6 +616,29 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore {
 		return this.#snapshot;
 	}
 
+	/**
+	 * Stateful probe for broker-side credential changes, mirroring
+	 * {@link SqliteAuthCredentialStore.pollExternalChanges} so long-lived broker
+	 * clients (notably `auth-gateway serve`) pick up logins/logouts made by
+	 * another process without a restart.
+	 *
+	 * Compares a local content revision, not the broker's numeric generation:
+	 * generation is an in-memory counter that resets when the broker process
+	 * restarts, so a reconnecting stream can deliver a different credential set
+	 * under a repeated (or lower) generation. {@link #refreshCredentialRevision}
+	 * bumps the revision whenever the applied credential material actually
+	 * changes, catching those cases too. Records foreground activity first: a
+	 * low-traffic client's background sync parks after `#backgroundIdleMs`, and
+	 * without this wakeup it would never fetch the new snapshot to report in the
+	 * first place.
+	 */
+	pollExternalChanges(): boolean {
+		this.#noteActivity();
+		if (this.#credentialRevision === this.#acknowledgedRevision) return false;
+		this.#acknowledgedRevision = this.#credentialRevision;
+		return true;
+	}
+
 	listAuthCredentials(provider?: string): StoredAuthCredential[] {
 		this.#noteActivity();
 		const out: StoredAuthCredential[] = [];

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

@@ -220,6 +220,7 @@ const usageScopeSchema = type({
 	"tier?": "string",
 	"windowId?": "string",
 	"shared?": "boolean",
+	"sharedGroup?": "string",
 });
 
 const usageLimitSchema = type({

package/src/auth-gateway/index.ts

@@ -1,3 +1,4 @@
 export * from "./http";
+export * from "./session-state";
 export * from "./server";
 export * from "./types";

package/src/auth-gateway/server.ts

@@ -19,7 +19,7 @@
  */
 
 import { Effort } from "@oh-my-pi/pi-catalog/effort";
-import { extractHttpStatusFromError, extractRetryHint, logger } from "@oh-my-pi/pi-utils";
+import { extractHttpStatusFromError, logger } from "@oh-my-pi/pi-utils";
 import type { ApiKeyResolver } from "../auth-retry";
 import type { AuthStorage } from "../auth-storage";
 import * as AIError from "../error";
@@ -34,6 +34,7 @@ import type { Api, AssistantMessage, AssistantMessageEventStream, Context, Model
 import type { ClientUsageIdentity } from "../usage";
 import { deterministicUuid } from "../utils/deterministic-id";
 import { parseBind } from "../utils/parse-bind";
+import { extractProviderRetryHint } from "../utils/retry-after";
 import {
 	captureRequestHeaders,
 	corsHeaders,
@@ -44,6 +45,7 @@ import {
 	resolvePeer,
 	withCors,
 } from "./http";
+import { AuthGatewaySessionStateStore } from "./session-state";
 import type {
 	AuthGatewayServerHandle,
 	AuthGatewayServerOptions,
@@ -128,6 +130,18 @@ function deriveSessionId(modelId: string, context: Context): string {
 	return deterministicUuid(seed);
 }
 
+/**
+ * Resolve the logical session identity for one request. A client-supplied key
+ * wins so external session ids line up with the gateway's, but a blank one
+ * counts as absent: honouring it would collapse every caller that sends an
+ * empty key into one shared credential-sticky, prefix-cache and
+ * provider-session bucket.
+ */
+function resolveSessionId(clientKey: string | undefined, modelId: string, context: Context): string {
+	if (clientKey !== undefined && clientKey.trim().length > 0) return clientKey;
+	return deriveSessionId(modelId, context);
+}
+
 function buildStreamOptions(parsed: ParsedFormatRequest, api: Api, signal: AbortSignal): SimpleStreamOptions {
 	const opts: SimpleStreamOptions = { signal, cursorExternalToolExecutor: true };
 	const { options } = parsed;
@@ -168,7 +182,7 @@ function buildStreamOptions(parsed: ParsedFormatRequest, api: Api, signal: Abort
 	// Client-supplied `prompt_cache_key` wins; otherwise derive a stable
 	// key from the model + system + tools so prefix caching engages on
 	// Codex-class backends across turns of the same logical conversation.
-	const promptCacheKey = options.promptCacheKey ?? deriveSessionId(parsed.modelId, parsed.context);
+	const promptCacheKey = resolveSessionId(options.promptCacheKey, parsed.modelId, parsed.context);
 	opts.promptCacheKey = promptCacheKey;
 	opts.sessionId = promptCacheKey;
 	if (options.thinkingBudgets) {
@@ -248,7 +262,7 @@ async function refreshGatewayApiKeyAfterAuthError(
 	const message = error instanceof Error ? error.message : String(error);
 	const status = extractHttpStatusFromError(error);
 	if (AIError.isUsageLimit(error) || isUsageLimitOutcome(status, message)) {
-		const retryAfterMs = extractRetryHint(undefined, message);
+		const retryAfterMs = extractProviderRetryHint(provider, message);
 		const { switched, retryAtMs } = await storage.markUsageLimitReached(provider, sessionId, {
 			retryAfterMs,
 			providerTimed: retryAfterMs !== undefined,
@@ -379,6 +393,7 @@ async function handleFormatEndpoint(
 	bootOpts: AuthGatewayBootOptions,
 	req: Request,
 	peer: string,
+	sessionStates: AuthGatewaySessionStateStore,
 ): Promise<Response> {
 	const startedAt = performance.now();
 	const requestId = crypto.randomUUID();
@@ -461,8 +476,8 @@ async function handleFormatEndpoint(
 	// supplied (so external session ids align), otherwise derive from
 	// modelId + system + tools + first message. Mirrored into
 	// streamOpts.sessionId / promptCacheKey by `buildStreamOptions`.
-	const sessionId = parsed.options.promptCacheKey ?? deriveSessionId(parsed.modelId, parsed.context);
-	parsed.options.promptCacheKey ??= sessionId;
+	const sessionId = resolveSessionId(parsed.options.promptCacheKey, parsed.modelId, parsed.context);
+	parsed.options.promptCacheKey = sessionId;
 
 	// pi-ai's stream() does NOT consult AuthStorage — the caller (us) is
 	// expected to resolve the credential and pass it as `options.apiKey`.
@@ -499,6 +514,11 @@ async function handleFormatEndpoint(
 		route.label,
 		peer,
 	);
+	// Per-session provider learning (sticky strict-tools / fast-mode / thinking
+	// fallbacks, Codex transport sessions). Owned by this gateway instance: the
+	// map is non-serializable, so no client can supply it and every turn would
+	// otherwise re-learn each lesson from a fresh upstream rejection.
+	streamOpts.providerSessionState = sessionStates.acquire(sessionId, model);
 
 	logger.info("auth-gateway request", {
 		requestId,
@@ -600,7 +620,12 @@ async function handleFormatEndpoint(
  * `parseRequest`/`encodeResponse`/`encodeStream` differ from the format-endpoint
  * path.
  */
-async function handlePiNative(bootOpts: AuthGatewayBootOptions, req: Request, peer: string): Promise<Response> {
+async function handlePiNative(
+	bootOpts: AuthGatewayBootOptions,
+	req: Request,
+	peer: string,
+	sessionStates: AuthGatewaySessionStateStore,
+): Promise<Response> {
 	const startedAt = performance.now();
 	const requestId = crypto.randomUUID();
 	const controller = mirrorRequestAbort(req);
@@ -635,8 +660,8 @@ async function handlePiNative(bootOpts: AuthGatewayBootOptions, req: Request, pe
 	// up with cache-prefix stickiness — same identity used for both means
 	// the next turn of this conversation reuses the same credential until
 	// it hits a usage cap, then markUsageLimitReached can hand off.
-	const sessionId = parsed.options.sessionId ?? deriveSessionId(parsed.modelId, parsed.context);
-	parsed.options.sessionId ??= sessionId;
+	const sessionId = resolveSessionId(parsed.options.sessionId, parsed.modelId, parsed.context);
+	parsed.options.sessionId = sessionId;
 
 	let apiKey: string | undefined;
 	try {
@@ -668,6 +693,11 @@ async function handlePiNative(bootOpts: AuthGatewayBootOptions, req: Request, pe
 		apiKey,
 		signal: controller.signal,
 		cursorExternalToolExecutor: true,
+		// Per-session provider learning, owned by this gateway instance. The map
+		// is non-serializable, so `parseRequest` cannot accept one from the wire
+		// and every turn would otherwise re-learn each lesson from a fresh
+		// upstream rejection.
+		providerSessionState: sessionStates.acquire(sessionId, model),
 	};
 	streamOpts.apiKey = buildGatewayApiKeyResolver(
 		bootOpts.storage,
@@ -692,7 +722,7 @@ async function handlePiNative(bootOpts: AuthGatewayBootOptions, req: Request, pe
 	// headers — the client's values win when they collide.
 	const captured = captureRequestHeaders(req.headers);
 	streamOpts.headers = { ...captured, ...streamOpts.headers };
-	streamOpts.sessionId ??= sessionId;
+	streamOpts.sessionId = sessionId;
 
 	logger.info("auth-gateway request", {
 		requestId,
@@ -847,6 +877,9 @@ export function startAuthGateway(opts: AuthGatewayBootOptions): AuthGatewayServe
 	const bind = parseBind(opts.bind ?? DEFAULT_AUTH_GATEWAY_BIND);
 	const tokens = new Set<string>(opts.bearerTokens);
 	const version = opts.version;
+	// Owned by this server instance so two gateways in one process never share
+	// (or tear down) each other's provider state, and so `close()` can drain it.
+	const sessionStates = new AuthGatewaySessionStateStore();
 
 	const server = Bun.serve({
 		hostname: bind.hostname,
@@ -887,13 +920,13 @@ export function startAuthGateway(opts: AuthGatewayBootOptions): AuthGatewayServe
 				// Provider-format dispatch.
 				const formatRoute = FORMAT_ROUTES[pathname];
 				if (formatRoute && req.method === "POST") {
-					return withCors(await handleFormatEndpoint(formatRoute, opts, req, peer), req);
+					return withCors(await handleFormatEndpoint(formatRoute, opts, req, peer, sessionStates), req);
 				}
 
 				// Pi-native fast path. Same auth + provider plumbing as the
 				// foreign-wire routes, just without the wire-format translation.
 				if (req.method === "POST" && pathname === "/v1/pi/stream") {
-					return withCors(await handlePiNative(opts, req, peer), req);
+					return withCors(await handlePiNative(opts, req, peer, sessionStates), req);
 				}
 
 				// Model catalog.
@@ -927,6 +960,10 @@ export function startAuthGateway(opts: AuthGatewayBootOptions): AuthGatewayServe
 		hostname: boundHost,
 		close: async () => {
 			server.stop(true);
+			// Drain after the listener is down: the retained provider states own
+			// sockets and timers (Codex WebSockets, GitLab Duo workflows), so the
+			// process can't settle until each one is closed.
+			sessionStates.close();
 		},
 	};
 }

package/src/auth-gateway/session-state.ts

@@ -0,0 +1,114 @@
+/**
+ * Server-owned provider session state for the auth-gateway.
+ *
+ * `SimpleStreamOptions.providerSessionState` is how a provider keeps what it
+ * learned about an endpoint across turns of one conversation: Anthropic's
+ * sticky `strictToolsDisabled` / `fastModeDisabled` /
+ * `replayUnsignedThinkingDisabled` flags and dropped-thinking-prefix set,
+ * OpenAI's strict-tools and reasoning-effort fallbacks, Codex's WebSocket and
+ * turn-state sessions. An in-process omp session owns that `Map` for its whole
+ * lifetime, so a grammar-too-large 400 or a fast-mode rejection costs one
+ * wasted round-trip per session rather than one per turn.
+ *
+ * The map is deliberately non-serializable — `Set`/`Map` fields, live sockets,
+ * a `close()` method — so `pi-native-client` strips it from the wire and
+ * `pi-native-server` never accepts it. Gateway clients therefore cannot bring
+ * their own, and without a server-side owner every containerized / robomp turn
+ * re-learns every lesson from a fresh upstream rejection.
+ *
+ * A plain `Map<sessionId, …>` in a long-lived server process is a leak: nothing
+ * ever reclaims an entry, and the entries own timers and sockets. This store is
+ * an LRU with a hard entry ceiling that calls `close()` on everything it drops
+ * and on everything it still holds at shutdown.
+ */
+
+import { logger } from "@oh-my-pi/pi-utils";
+import { type DisposeReason, LRUCache } from "@oh-my-pi/pi-utils/lru";
+import type { Api, Model, ProviderSessionState } from "../types";
+
+/**
+ * Retained logical sessions. Each entry is a handful of small provider records
+ * plus, for Codex, a WebSocket session — cheap to keep, but not free, so the
+ * ceiling is what turns "one entry per session id forever" into a bounded cost.
+ * Eviction is least-recently-used, so the ceiling only ever drops sessions that
+ * have been quiet longer than the 256 most recent ones.
+ */
+export const AUTH_GATEWAY_MAX_SESSION_STATES = 256;
+
+/**
+ * Close every provider record an evicted (or drained) session held.
+ *
+ * Anthropic's `close()` resets its sticky flags, Codex's tears down WebSockets
+ * and GitLab Duo's stops the server-side workflow — so dropping an entry
+ * without closing it leaks exactly the resources the bound exists to cap.
+ */
+function closeSessionState(states: Map<string, ProviderSessionState>, sessionKey: string, reason: DisposeReason): void {
+	for (const [providerKey, state] of states) {
+		try {
+			state.close();
+		} catch (error) {
+			// One provider's teardown must not abort the rest: a throw here
+			// propagates out of `LRUCache.set` into whichever request happened to
+			// trigger the eviction, or abandons the remainder of the shutdown
+			// drain.
+			logger.warn("auth-gateway provider session state close failed", {
+				sessionKey,
+				providerKey,
+				reason,
+				error: String(error),
+			});
+		}
+	}
+	states.clear();
+}
+
+/**
+ * Bounded per-session provider state, owned by one gateway server instance.
+ *
+ * Two gateways in the same process get separate stores, so neither can hand a
+ * request another gateway's learned state or close it out from under one.
+ */
+export class AuthGatewaySessionStateStore {
+	readonly #sessions: LRUCache<string, Map<string, ProviderSessionState>>;
+
+	constructor(max: number = AUTH_GATEWAY_MAX_SESSION_STATES) {
+		this.#sessions = new LRUCache({ max, dispose: closeSessionState });
+	}
+
+	/** Retained logical sessions. */
+	get size(): number {
+		return this.#sessions.size;
+	}
+
+	/**
+	 * The provider-session map for one logical session on one model, created on
+	 * first use and returned by reference so provider mutations persist into the
+	 * next request.
+	 *
+	 * Keyed by session + provider + model id. The session is the identity that
+	 * matters — it is the same identity used for credential stickiness and
+	 * prefix-cache keying — but a client is free to reuse one session id across
+	 * models, and the coarsest provider entries do not separate models
+	 * themselves (`openai-responses` keys its strict-tools / history-replay
+	 * record by provider alone, Antigravity by a single constant), so the model
+	 * belongs in the key here. Endpoint is deliberately absent: every provider
+	 * whose learning is endpoint-specific already sub-keys it internally
+	 * (`anthropic-messages:${baseUrl}\0${modelId}`,
+	 * `openai-completions:${provider}:${baseUrl}:${modelId}`), and repeating it
+	 * would only fragment the map. NUL separates the components so none of them
+	 * can forge the boundary.
+	 */
+	acquire(sessionId: string, model: Model<Api>): Map<string, ProviderSessionState> {
+		const key = `${sessionId}\u0000${model.provider}\u0000${model.id}`;
+		const existing = this.#sessions.get(key);
+		if (existing) return existing;
+		const created = new Map<string, ProviderSessionState>();
+		this.#sessions.set(key, created);
+		return created;
+	}
+
+	/** Close and drop every retained state. Called when the gateway shuts down. */
+	close(): void {
+		this.#sessions.clear();
+	}
+}