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

package/src/auth-gateway/server.ts

@@ -131,15 +131,37 @@ function deriveSessionId(modelId: string, context: Context): string {
 }
 
 /**
- * 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.
+ * The client's own session key, or `undefined` when it sent none. A blank key
+ * counts as none: 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 normalizeClientSessionKey(clientKey: string | undefined): string | undefined {
+	return clientKey !== undefined && clientKey.trim().length > 0 ? clientKey : undefined;
+}
+
+/**
+ * Stable identity of the account a request's credential belongs to.
+ *
+ * `markUsageLimitReached` and the auth-retry resolver switch a session to a
+ * sibling credential, so the provider state retained for that session can
+ * outlive the account that taught it. OAuth rows expose an account id / email
+ * that survives token refresh — fingerprinting the bearer instead would look
+ * like a rotation every time a token refreshes and discard the retained
+ * lessons for nothing. Key-based rows fall back to a hash of the key, never
+ * the key itself: this value is held for the lifetime of the entry.
+ */
+function resolveGatewayAccount(storage: AuthStorage, provider: string, sessionId: string, apiKey: string): string {
+	const identity = storage.getOAuthAccountIdentity(provider, sessionId);
+	if (identity) {
+		return `oauth:${JSON.stringify([
+			identity.accountId ?? "",
+			identity.email ?? "",
+			identity.projectId ?? "",
+			identity.orgId ?? "",
+		])}`;
+	}
+	return `key:${Bun.hash(apiKey).toString(36)}`;
 }
 
 function buildStreamOptions(parsed: ParsedFormatRequest, api: Api, signal: AbortSignal): SimpleStreamOptions {
@@ -182,7 +204,8 @@ 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 = resolveSessionId(options.promptCacheKey, parsed.modelId, parsed.context);
+	const promptCacheKey =
+		normalizeClientSessionKey(options.promptCacheKey) ?? deriveSessionId(parsed.modelId, parsed.context);
 	opts.promptCacheKey = promptCacheKey;
 	opts.sessionId = promptCacheKey;
 	if (options.thinkingBudgets) {
@@ -314,6 +337,7 @@ function buildGatewayApiKeyResolver(
 	requestSignal: AbortSignal,
 	format: string,
 	peer: string,
+	onResolvedKey: (apiKey: string) => void,
 ): ApiKeyResolver {
 	let lastKey = initialKey;
 	return async ({ lastChance, error, signal }) => {
@@ -329,6 +353,7 @@ function buildGatewayApiKeyResolver(
 				forceRefresh: true,
 			});
 			lastKey = refreshed ?? lastKey;
+			if (refreshed) onResolvedKey(refreshed);
 			return refreshed;
 		}
 		const next = await refreshGatewayApiKeyAfterAuthError(
@@ -343,6 +368,7 @@ function buildGatewayApiKeyResolver(
 			peer,
 		);
 		lastKey = next ?? lastKey;
+		if (next) onResolvedKey(next);
 		return next;
 	};
 }
@@ -476,7 +502,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 = resolveSessionId(parsed.options.promptCacheKey, parsed.modelId, parsed.context);
+	const clientKey = normalizeClientSessionKey(parsed.options.promptCacheKey);
+	const sessionId = clientKey ?? deriveSessionId(parsed.modelId, parsed.context);
 	parsed.options.promptCacheKey = sessionId;
 
 	// pi-ai's stream() does NOT consult AuthStorage — the caller (us) is
@@ -505,6 +532,19 @@ async function handleFormatEndpoint(
 	}
 
 	const streamOpts = buildStreamOptions(parsed, model.api, controller.signal);
+	// 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. The lease
+	// keeps the entry out of reach of eviction until this request is done with
+	// it, so it MUST be released on every exit path.
+	const lease = sessionStates.acquire({
+		clientKey,
+		model,
+		context: parsed.context,
+		account: resolveGatewayAccount(bootOpts.storage, model.provider, sessionId, apiKey),
+	});
+	streamOpts.providerSessionState = lease.states;
 	streamOpts.apiKey = buildGatewayApiKeyResolver(
 		bootOpts.storage,
 		model,
@@ -513,12 +553,9 @@ async function handleFormatEndpoint(
 		controller.signal,
 		route.label,
 		peer,
+		resolvedKey =>
+			lease.updateAccount(resolveGatewayAccount(bootOpts.storage, model.provider, sessionId, resolvedKey)),
 	);
-	// 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,
@@ -565,45 +602,59 @@ async function handleFormatEndpoint(
 				peer,
 			});
 			return route.module.formatError(classified.status, classified.type, classified.message);
+		} finally {
+			// Every non-streaming outcome — answered, upstream error, thrown,
+			// client gone — is done with the provider state here.
+			lease.release();
 		}
 	}
 
-	let events: AssistantMessageEventStream;
+	// A streamed turn outlives this function, so the lease travels with the
+	// event stream and is released when the turn settles. Until that handoff
+	// happens, the `finally` below owns it.
+	let streamOwnsLease = false;
 	try {
+		let events: AssistantMessageEventStream;
+		try {
+			if (controller.signal.aborted) return clientClosedResponse(route);
+			events = streamSimple(model, parsed.context, streamOpts);
+		} catch (error) {
+			const classified = classifyGatewayError(error);
+			logger.warn("auth-gateway streamSimple threw", { format: route.label, error: classified.message, peer });
+			return route.module.formatError(classified.status, classified.type, classified.message);
+		}
 		if (controller.signal.aborted) return clientClosedResponse(route);
-		events = streamSimple(model, parsed.context, streamOpts);
-	} catch (error) {
-		const classified = classifyGatewayError(error);
-		logger.warn("auth-gateway streamSimple threw", { format: route.label, error: classified.message, peer });
-		return route.module.formatError(classified.status, classified.type, classified.message);
-	}
-	if (controller.signal.aborted) return clientClosedResponse(route);
-	void events
-		.result()
-		.then(message => recordGatewayUsage(bootOpts.storage, model, client, message))
-		.catch(() => {});
+		void events
+			.result()
+			.then(message => recordGatewayUsage(bootOpts.storage, model, client, message))
+			.catch(() => {})
+			.finally(() => lease.release());
+		streamOwnsLease = true;
 
-	const sseStream = route.module.encodeStream(events, parsed.modelId, parsed.options, {
-		signal: controller.signal,
-		onCancel: reason => {
-			if (!controller.signal.aborted) {
-				controller.abort(reason instanceof Error ? reason : new Error("client closed request"));
-			}
-		},
-	});
-	return new Response(sseStream, {
-		status: 200,
-		headers: {
-			...gatewayResponseHeaders(model, { requestId }),
-			"Content-Type": "text/event-stream; charset=utf-8",
-			"Cache-Control": "no-cache",
-			Connection: "keep-alive",
-			// Disable proxy buffering (nginx and ingress controllers honor this).
-			// Without it the SSE stream gets held until the buffer flushes, which
-			// stalls the long-thinking-budget calls we exist to support.
-			"X-Accel-Buffering": "no",
-		},
-	});
+		const sseStream = route.module.encodeStream(events, parsed.modelId, parsed.options, {
+			signal: controller.signal,
+			onCancel: reason => {
+				if (!controller.signal.aborted) {
+					controller.abort(reason instanceof Error ? reason : new Error("client closed request"));
+				}
+			},
+		});
+		return new Response(sseStream, {
+			status: 200,
+			headers: {
+				...gatewayResponseHeaders(model, { requestId }),
+				"Content-Type": "text/event-stream; charset=utf-8",
+				"Cache-Control": "no-cache",
+				Connection: "keep-alive",
+				// Disable proxy buffering (nginx and ingress controllers honor this).
+				// Without it the SSE stream gets held until the buffer flushes, which
+				// stalls the long-thinking-budget calls we exist to support.
+				"X-Accel-Buffering": "no",
+			},
+		});
+	} finally {
+		if (!streamOwnsLease) lease.release();
+	}
 }
 
 /**
@@ -660,7 +711,8 @@ async function handlePiNative(
 	// 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 = resolveSessionId(parsed.options.sessionId, parsed.modelId, parsed.context);
+	const clientKey = normalizeClientSessionKey(parsed.options.sessionId);
+	const sessionId = clientKey ?? deriveSessionId(parsed.modelId, parsed.context);
 	parsed.options.sessionId = sessionId;
 
 	let apiKey: string | undefined;
@@ -684,6 +736,17 @@ async function handlePiNative(
 		);
 	}
 
+	// 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. The lease keeps the entry out of reach of eviction until this
+	// request is done with it, so it MUST be released on every exit path.
+	const lease = sessionStates.acquire({
+		clientKey,
+		model,
+		context: parsed.context,
+		account: resolveGatewayAccount(bootOpts.storage, model.provider, sessionId, apiKey),
+	});
 	// Build the SimpleStreamOptions actually handed to `streamSimple`. We
 	// trust the client's options (already allow-listed by `parseRequest`) and
 	// only inject server-controlled fields. The codex sampling strip mirrors
@@ -693,11 +756,7 @@ async function handlePiNative(
 		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),
+		providerSessionState: lease.states,
 	};
 	streamOpts.apiKey = buildGatewayApiKeyResolver(
 		bootOpts.storage,
@@ -707,6 +766,8 @@ async function handlePiNative(
 		controller.signal,
 		"pi-native",
 		peer,
+		resolvedKey =>
+			lease.updateAccount(resolveGatewayAccount(bootOpts.storage, model.provider, sessionId, resolvedKey)),
 	);
 	if (model.api === "openai-codex-responses") {
 		delete streamOpts.temperature;
@@ -761,42 +822,56 @@ async function handlePiNative(
 			const classified = classifyGatewayError(error);
 			logger.warn("auth-gateway non-streaming aborted", { format: "pi-native", error: classified.message, peer });
 			return piNative.formatError(classified.status, classified.type, classified.message);
+		} finally {
+			// Every non-streaming outcome — answered, upstream error, thrown,
+			// client gone — is done with the provider state here.
+			lease.release();
 		}
 	}
 
-	let events: AssistantMessageEventStream;
+	// A streamed turn outlives this function, so the lease travels with the
+	// event stream and is released when the turn settles. Until that handoff
+	// happens, the `finally` below owns it.
+	let streamOwnsLease = false;
 	try {
+		let events: AssistantMessageEventStream;
+		try {
+			if (controller.signal.aborted) return aborted();
+			events = streamSimple(model, parsed.context, streamOpts);
+		} catch (error) {
+			const classified = classifyGatewayError(error);
+			logger.warn("auth-gateway streamSimple threw", { format: "pi-native", error: classified.message, peer });
+			return piNative.formatError(classified.status, classified.type, classified.message);
+		}
 		if (controller.signal.aborted) return aborted();
-		events = streamSimple(model, parsed.context, streamOpts);
-	} catch (error) {
-		const classified = classifyGatewayError(error);
-		logger.warn("auth-gateway streamSimple threw", { format: "pi-native", error: classified.message, peer });
-		return piNative.formatError(classified.status, classified.type, classified.message);
-	}
-	if (controller.signal.aborted) return aborted();
-	void events
-		.result()
-		.then(message => recordGatewayUsage(bootOpts.storage, model, client, message))
-		.catch(() => {});
+		void events
+			.result()
+			.then(message => recordGatewayUsage(bootOpts.storage, model, client, message))
+			.catch(() => {})
+			.finally(() => lease.release());
+		streamOwnsLease = true;
 
-	const sseStream = piNative.encodeStream(events, parsed.modelId, parsed.options, {
-		signal: controller.signal,
-		onCancel: reason => {
-			if (!controller.signal.aborted) {
-				controller.abort(reason instanceof Error ? reason : new Error("client closed request"));
-			}
-		},
-	});
-	return new Response(sseStream, {
-		status: 200,
-		headers: {
-			...gatewayResponseHeaders(model, { requestId }),
-			"Content-Type": "text/event-stream; charset=utf-8",
-			"Cache-Control": "no-cache",
-			Connection: "keep-alive",
-			"X-Accel-Buffering": "no",
-		},
-	});
+		const sseStream = piNative.encodeStream(events, parsed.modelId, parsed.options, {
+			signal: controller.signal,
+			onCancel: reason => {
+				if (!controller.signal.aborted) {
+					controller.abort(reason instanceof Error ? reason : new Error("client closed request"));
+				}
+			},
+		});
+		return new Response(sseStream, {
+			status: 200,
+			headers: {
+				...gatewayResponseHeaders(model, { requestId }),
+				"Content-Type": "text/event-stream; charset=utf-8",
+				"Cache-Control": "no-cache",
+				Connection: "keep-alive",
+				"X-Accel-Buffering": "no",
+			},
+		});
+	} finally {
+		if (!streamOwnsLease) lease.release();
+	}
 }
 
 /**

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

@@ -19,12 +19,14 @@
  * 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.
+ * and on everything it still holds at shutdown — but it only ever drops an
+ * entry no request is holding, because `close()` on a live entry tears down
+ * state an in-flight stream is still streaming through.
  */
 
 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";
+import { resetAccountScopedProviderSessionState } from "../provider-session-state";
+import type { Api, Context, Model, ProviderSessionState } from "../types";
 
 /**
  * Retained logical sessions. Each entry is a handful of small provider records
@@ -35,6 +37,61 @@ import type { Api, Model, ProviderSessionState } from "../types";
  */
 export const AUTH_GATEWAY_MAX_SESSION_STATES = 256;
 
+/** Why an entry's provider records were closed. Logged on teardown failure. */
+type SessionDisposeReason = "evict" | "shutdown";
+
+/**
+ * One request's claim on a retained session.
+ *
+ * `release()` is what makes the entry evictable again, so it MUST run for every
+ * outcome of the request — a `finally` at the call site for the synchronous
+ * paths, stream completion for the streaming ones. It is idempotent, so the
+ * two can overlap.
+ */
+export interface AuthGatewaySessionStateLease {
+	/** The map to hand to `streamSimple` as `providerSessionState`. */
+	readonly states: Map<string, ProviderSessionState>;
+	/** Reset account-scoped records if an in-request auth retry switches accounts. */
+	updateAccount(account: string): void;
+	/** Give up this request's claim. Idempotent. */
+	release(): void;
+}
+
+/** Everything the store needs to place one request on a retained session. */
+export interface AuthGatewaySessionStateRequest {
+	/**
+	 * The client's own session key (`prompt_cache_key` / `sessionId`), or
+	 * `undefined` when it sent none — blank counts as none. A supplied key is
+	 * authoritative: the client is telling us which conversation this is.
+	 */
+	clientKey: string | undefined;
+	model: Model<Api>;
+	/**
+	 * System prompt, tools and message history of this request. Used only when
+	 * `clientKey` is absent, to place the request on the conversation it
+	 * continues.
+	 */
+	context: Context;
+	/**
+	 * Stable identity of the account this request's credential resolved to.
+	 * A change means the gateway switched the session to a sibling credential,
+	 * so the account-dependent lessons in the retained map are re-probed. The
+	 * comparison happens on acquire and whenever an in-request auth retry
+	 * resolves a sibling credential.
+	 */
+	account: string;
+}
+
+interface RetainedSession {
+	/** Current index key. Advances as a keyless conversation grows. */
+	key: string;
+	states: Map<string, ProviderSessionState>;
+	/** Account identity of the most recent request placed on this entry. */
+	account: string;
+	/** Requests currently holding this entry. Eviction never takes one of these. */
+	leases: number;
+}
+
 /**
  * Close every provider record an evicted (or drained) session held.
  *
@@ -42,15 +99,18 @@ export const AUTH_GATEWAY_MAX_SESSION_STATES = 256;
  * 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 {
+function closeSessionState(
+	states: Map<string, ProviderSessionState>,
+	sessionKey: string,
+	reason: SessionDisposeReason,
+): 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.
+			// propagates out of the eviction into whichever request happened to
+			// trigger it, or abandons the remainder of the shutdown drain.
 			logger.warn("auth-gateway provider session state close failed", {
 				sessionKey,
 				providerKey,
@@ -62,17 +122,82 @@ function closeSessionState(states: Map<string, ProviderSessionState>, sessionKey
 	states.clear();
 }
 
+/**
+ * Index keys this request may be placed on, most specific first.
+ *
+ * With a client key there is exactly one: the client named its conversation, so
+ * provider + model + that key is the identity.
+ *
+ * Without one the gateway has to infer the conversation, and the request's
+ * message history is the only thing that can distinguish two of them. The
+ * derived `sessionId` used for prefix caching and credential stickiness hashes
+ * the model, system prompt, tools and *first* message, which is deliberately
+ * prefix-shaped — two chats that open the same way share a cache bucket, which
+ * is a cache hit rather than a leak, and share a sticky account, which is a
+ * load-balancing hint. Retained provider state is neither: sharing it means one
+ * chat's rejection silences another chat's request, and one chat's Codex
+ * transport session answers another chat's turn. So provider state gets its own
+ * key, and only provider state: `deriveSessionId` keeps its two other jobs.
+ *
+ * The key is therefore a running hash over (provider, model, system, tools) and
+ * then every message, one key per message — the last of which identifies the
+ * exact history this request presented. Turn N+1 of a conversation extends turn
+ * N's history, so turn N's key is one of the earlier entries in turn N+1's
+ * chain: the store finds the nearest ancestor and moves that entry forward onto
+ * the new key. Two conversations that share an opening therefore share an entry
+ * only until they diverge; after that the first branch to arrive keeps the
+ * ancestor and the other starts clean. That is the most a stateless wire can
+ * tell us — before divergence the two requests are byte-identical.
+ *
+ * System prompt and tools sit in the root rather than per-message because they
+ * are not history: a client that re-stamps its system prompt every turn (a
+ * date, a cwd) starts a new lineage, exactly as it already starts a new derived
+ * `sessionId` today.
+ */
+function sessionKeys(request: AuthGatewaySessionStateRequest): string[] {
+	const { model } = request;
+	const scope = `${model.provider}\u0000${model.id}`;
+	if (request.clientKey !== undefined) return [`c\u0000${scope}\u0000${request.clientKey}`];
+	const { context } = request;
+	// NUL separates the components so none of them can forge the boundary.
+	let hash = Bun.hash(
+		`${scope}\u0000${context.systemPrompt?.join("\n\n") ?? ""}\u0000${context.tools ? JSON.stringify(context.tools) : ""}`,
+	);
+	const keys: string[] = [];
+	for (const message of context.messages) {
+		// Role + content only: omp re-stamps `timestamp` and provider metadata on
+		// every parsed message, so hashing those would break the chain on turn
+		// two of every conversation.
+		hash = Bun.hash(JSON.stringify({ role: message.role, content: message.content }), hash);
+		keys.push(`h\u0000${scope}\u0000${hash.toString(36)}`);
+	}
+	// A request with no messages has no history to place; its root is the key.
+	if (keys.length === 0) return [`h\u0000${scope}\u0000${hash.toString(36)}`];
+	keys.reverse();
+	return keys;
+}
+
 /**
  * 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.
+ *
+ * The recency order is this class's own (a `Map` iterates in insertion order,
+ * and every acquire re-inserts) rather than `LRUCache`'s, because the policy
+ * needs two things a general cache cannot express: an entry that a request is
+ * still holding must be skipped when picking a victim, and an entry must be
+ * able to change key — `LRUCache` disposes on every removal, which is precisely
+ * the `close()` we must not run here.
  */
 export class AuthGatewaySessionStateStore {
-	readonly #sessions: LRUCache<string, Map<string, ProviderSessionState>>;
+	/** Least recently acquired first — insertion order is the LRU order. */
+	readonly #sessions = new Map<string, RetainedSession>();
+	readonly #max: number;
 
 	constructor(max: number = AUTH_GATEWAY_MAX_SESSION_STATES) {
-		this.#sessions = new LRUCache({ max, dispose: closeSessionState });
+		if (!Number.isInteger(max) || max < 1) throw new TypeError("max must be a positive integer");
+		this.#max = max;
 	}
 
 	/** Retained logical sessions. */
@@ -81,34 +206,107 @@ export class AuthGatewaySessionStateStore {
 	}
 
 	/**
-	 * 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.
+	 * Claim the provider-session map for one request, 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
+	 * Keyed by provider + model + conversation (see {@link sessionKeys}). 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.
+	 * would only fragment the map. The credential is absent for the same reason
+	 * — most of what is retained is true of the endpoint whoever calls it, and
+	 * Codex already sub-keys its transport by account and bearer — so a
+	 * credential switch resets the account-dependent subset instead of
+	 * splitting the entry (see `resetAccountScopedProviderSessionState`).
+	 *
+	 * The returned lease MUST be released; until then the entry cannot be
+	 * evicted.
 	 */
-	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;
+	acquire(request: AuthGatewaySessionStateRequest): AuthGatewaySessionStateLease {
+		const session = this.#claim(sessionKeys(request), request.account);
+		let released = false;
+		return {
+			states: session.states,
+			updateAccount: (account: string): void => {
+				if (session.account === account) return;
+				resetAccountScopedProviderSessionState(session.states);
+				session.account = account;
+			},
+			release: (): void => {
+				if (released) return;
+				released = true;
+				session.leases--;
+				// This entry may be the victim the bound has been waiting for.
+				if (session.leases === 0) this.#evict();
+			},
+		};
 	}
 
 	/** Close and drop every retained state. Called when the gateway shuts down. */
 	close(): void {
+		// The only place a leased entry is torn down: the listener is already
+		// down, every in-flight stream is being cancelled with it, and the
+		// process cannot settle while a Codex WebSocket or Duo workflow is open.
+		for (const session of this.#sessions.values()) closeSessionState(session.states, session.key, "shutdown");
 		this.#sessions.clear();
 	}
+
+	/**
+	 * Resolve `keys` to an entry — reusing the nearest ancestor when a keyless
+	 * conversation has grown — mark it most recently used, and hand the caller
+	 * the claim. The claim is taken before the ceiling is enforced: a brand-new
+	 * entry belongs to the request that just created it, and is not a candidate
+	 * for making room for itself.
+	 */
+	#claim(keys: readonly string[], account: string): RetainedSession {
+		const key = keys[0] ?? "";
+		for (const candidate of keys) {
+			const session = this.#sessions.get(candidate);
+			if (session === undefined) continue;
+			// Re-insert at the tail for recency, under this request's own key so
+			// the next turn of this conversation finds it as its ancestor. A
+			// sibling branch of the same ancestor no longer matches, which is the
+			// point: it gets an entry of its own.
+			this.#sessions.delete(candidate);
+			session.key = key;
+			this.#sessions.set(key, session);
+			session.leases++;
+			if (session.account !== account) {
+				resetAccountScopedProviderSessionState(session.states);
+				session.account = account;
+			}
+			return session;
+		}
+		const created: RetainedSession = { key, states: new Map(), account, leases: 1 };
+		this.#sessions.set(key, created);
+		this.#evict();
+		return created;
+	}
+
+	/**
+	 * Enforce the ceiling against entries no request is holding.
+	 *
+	 * A long-running stream is exactly the entry LRU order would pick — it was
+	 * acquired when the stream opened and not touched since — so blind eviction
+	 * would `close()` the sockets and flags that stream is still using. Live
+	 * entries are skipped instead, and the store sits above its bound until
+	 * their requests release; the excess is therefore capped by the number of
+	 * concurrent requests, each of which holds a client connection.
+	 */
+	#evict(): void {
+		if (this.#sessions.size <= this.#max) return;
+		// Deleting during Map iteration is well-defined: the current and later
+		// keys stay consistent, so this walks least-recently-acquired first.
+		for (const session of this.#sessions.values()) {
+			if (this.#sessions.size <= this.#max) return;
+			if (session.leases > 0) continue;
+			this.#sessions.delete(session.key);
+			closeSessionState(session.states, session.key, "evict");
+		}
+	}
 }