@oh-my-pi/pi-ai 18.8.0 → 18.8.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/CHANGELOG.md CHANGED
@@ -2,6 +2,17 @@
2
2
 
3
3
  ## [Unreleased]
4
4
 
5
+ ## [18.8.1] - 2026-10-07
6
+
7
+ ### Added
8
+
9
+ - Added session restrictions for OAuth account pools via `AuthStorage.sessions.restrict`, limiting selection, fallback, rotation, and authentication to specified accounts until the returned lease is released with `sessions.unrestrict`. API keys and other accounts are not used when a session is restricted.
10
+ - Exported `resolveCredentialIdentityKey` for determining the identity key used to match credentials with broker account pools and session restrictions.
11
+
12
+ ### Fixed
13
+
14
+ - Fixed Codex Fast (`priority`) and Ultrafast usage being recorded, billed, and reported as Standard when the backend echoed a default service tier; the requested tier is now preserved in usage and performance records.
15
+
5
16
  ## [18.8.0] - 2026-10-07
6
17
 
7
18
  ### Breaking Changes
@@ -1,4 +1,4 @@
1
- import type { AuthCredential, OAuthCredential, SessionsApi } from "./types.js";
1
+ import type { AuthCredential, OAuthCredential, SessionRestrictionLease, SessionsApi } from "./types.js";
2
2
  import type { AuthCredentialStore } from "./store.js";
3
3
  import type { CredentialPool } from "./pool.js";
4
4
  import type { KeyOverrides } from "./cascade.js";
@@ -13,10 +13,30 @@ export type SessionCredential = {
13
13
  /** Set only by the public user-facing pin API; automatic warm affinity leaves it absent. */
14
14
  explicit?: true;
15
15
  };
16
+ /** One installed restriction: the allowed OAuth identity keys and the lease that lifts it. */
17
+ type SessionRestriction = {
18
+ readonly allowed: ReadonlySet<string>;
19
+ readonly lease: SessionRestrictionLease;
20
+ };
21
+ /**
22
+ * Provider → session id → installed restriction. Owned by `AuthStorage` so
23
+ * restrictions outlive store replacement, which rebuilds every store-bound
24
+ * module (pins included) around the same identities.
25
+ */
26
+ export type SessionRestrictions = Map<string, Map<string, SessionRestriction>>;
16
27
  /** Session → credential affinity (pins), persisted in the store cache. */
17
28
  export declare class SessionAffinity implements SessionsApi {
18
29
  #private;
19
- constructor(store: AuthCredentialStore, pool: CredentialPool, overrides: KeyOverrides);
30
+ constructor(store: AuthCredentialStore, pool: CredentialPool, overrides: KeyOverrides, restrictions: SessionRestrictions);
31
+ restrict(provider: string, sessionId: string, identityKeys: readonly string[]): SessionRestrictionLease;
32
+ unrestrict(provider: string, sessionId: string, lease: SessionRestrictionLease): void;
33
+ /** True when {@link restrict} limits `sessionId` for `provider`. */
34
+ isRestricted(provider: string, sessionId: string | undefined): boolean;
35
+ /**
36
+ * True when `credential` may serve `sessionId`: always for an unrestricted
37
+ * session, otherwise only an OAuth credential whose identity key is allowed.
38
+ */
39
+ allows(provider: string, sessionId: string | undefined, credential: AuthCredential): boolean;
20
40
  /** Drop every pin for a provider, both in memory and in the persisted cache. */
21
41
  clearProvider(provider: string): void;
22
42
  /** Drop only this in-memory session pin after OAuth selection falls through. */
@@ -33,7 +53,11 @@ export declare class SessionAffinity implements SessionsApi {
33
53
  * rewrites were pure database churn.
34
54
  */
35
55
  record(provider: string, sessionId: string | undefined, type: AuthCredential["type"], index: number, lastUsedAtMs?: number, explicit?: boolean): void;
36
- /** Retrieves the last credential used by a session. */
56
+ /**
57
+ * Retrieves the last credential used by a session. A restricted session
58
+ * never sees a pin outside its allowlist — inherited, restored, or recorded
59
+ * before the restriction — so it re-ranks inside the allowlist instead.
60
+ */
37
61
  get(provider: string, sessionId: string | undefined): SessionCredential | undefined;
38
62
  /** Clears the last credential used by a session for a provider. */
39
63
  clear(provider: string, sessionId: string | undefined): void;
@@ -71,3 +95,4 @@ export declare class SessionAffinity implements SessionsApi {
71
95
  */
72
96
  release(provider: string, sessionId: string): boolean;
73
97
  }
98
+ export {};
@@ -8,6 +8,14 @@ export declare class KeyOverrides {
8
8
  #private;
9
9
  constructor(resolver?: (config: string) => Promise<string | undefined>);
10
10
  has(provider: string): boolean;
11
+ /**
12
+ * Whether an override replaces OAuth for a session. Any override does for an
13
+ * unrestricted session. A restricted session skips the runtime key, but a
14
+ * config key still blocks OAuth: it marks the provider's endpoint (often a
15
+ * proxy) as taking that key, so the session fails closed instead of sending
16
+ * a pooled OAuth token there (see `KeyCascade.get`).
17
+ */
18
+ suppressesOAuth(provider: string, restricted: boolean): boolean;
11
19
  runtimeKey(provider: string): string | undefined;
12
20
  configKey(provider: string): string | undefined;
13
21
  /** Config value consulted only after stored OAuth/login credentials. */
@@ -115,6 +123,10 @@ export declare class KeyCascade implements KeysApi {
115
123
  * Get API key for a provider.
116
124
  * Priority (first match wins): runtime override, config override, OAuth,
117
125
  * login API key, environment variable, then another stored API key.
126
+ * A session restricted by `sessions.restrict` resolves only its allowed
127
+ * OAuth accounts and throws {@link AIError.MissingApiKeyError} when none
128
+ * can serve or a config key (a models.yml `apiKey`, often for a proxy
129
+ * `baseUrl`) owns the provider, rather than using any other credential.
118
130
  */
119
131
  get(provider: string, sessionId?: string, options?: AuthApiKeyOptions, onCredentialId?: (id: number, identity?: OAuthRequestIdentity) => void): Promise<string | undefined>;
120
132
  /**
@@ -18,6 +18,12 @@ type SerializedCredentialRecord = {
18
18
  };
19
19
  export { isSqliteBusyError, isSqliteCorruptionError };
20
20
  export declare function serializeCredential(provider: string, credential: AuthCredential): SerializedCredentialRecord | null;
21
+ /**
22
+ * Identity key of a credential: the `identityKey` that broker snapshots carry
23
+ * and that broker account pools and `sessions.restrict` match, such as
24
+ * `email:<address>|org:<id>` for org-scoped providers. `null` for an API key or
25
+ * an OAuth credential without any account, email, project, or org identity.
26
+ */
21
27
  export declare function resolveCredentialIdentityKey(provider: string, credential: AuthCredential): string | null;
22
28
  /**
23
29
  * Default SQLite-backed implementation of {@link AuthCredentialStore}.
@@ -998,6 +998,11 @@ export interface OAuthApi {
998
998
  */
999
999
  refreshStored<T extends OAuthCredential = OAuthCredential>(provider: string, options: StoredOAuthRefreshOptions<T>): Promise<StoredOAuthRefreshResult<T>>;
1000
1000
  }
1001
+ /**
1002
+ * Handle for one {@link SessionsApi.restrict} call. {@link SessionsApi.unrestrict}
1003
+ * lifts a restriction only with the lease of the call that installed it.
1004
+ */
1005
+ export type SessionRestrictionLease = symbol;
1001
1006
  /** Session credential affinity operations. */
1002
1007
  export interface SessionsApi {
1003
1008
  /**
@@ -1025,6 +1030,27 @@ export interface SessionsApi {
1025
1030
  * source session.
1026
1031
  */
1027
1032
  inherit(sourceSessionId: string, targetSessionId: string): number;
1033
+ /**
1034
+ * Restrict one session's credentials for `provider` to the OAuth accounts
1035
+ * whose identity key (`email:<address>|org:<id>` for org-scoped providers;
1036
+ * otherwise `account:`, `email:`, or `project:` — the keys broker account
1037
+ * pools use) is listed. Selection, session pins, inherited affinity,
1038
+ * fallback passes, and rotation never leave the list, and runtime, config,
1039
+ * environment, and stored API keys are not used for that session; when no
1040
+ * listed account can serve, key resolution fails instead of borrowing
1041
+ * another account. An empty list allows no account. Replaces any earlier
1042
+ * restriction for the same provider and session; {@link inherit} does not
1043
+ * copy restrictions. Restrictions are never evicted, because that would
1044
+ * widen a live session; the owner passes the returned lease to
1045
+ * {@link unrestrict} when the session ends.
1046
+ */
1047
+ restrict(provider: string, sessionId: string, identityKeys: readonly string[]): SessionRestrictionLease;
1048
+ /**
1049
+ * Lift the restriction that `lease` installed. Once a later {@link restrict}
1050
+ * call has replaced it for the same provider and session, this does nothing,
1051
+ * so a stale owner never lifts the restriction a newer owner relies on.
1052
+ */
1053
+ unrestrict(provider: string, sessionId: string, lease: SessionRestrictionLease): void;
1028
1054
  /**
1029
1055
  * Release a session's sticky credential so its next {@link getApiKey} call
1030
1056
  * re-runs native pool ranking. This never blocks or penalizes the released
@@ -1,6 +1,6 @@
1
1
  import type { AuthCredentialStore } from "./auth/store.js";
2
2
  import type { AuthAccountPolicies, AuthApiKeyOptions, AuthStorageOptions, BlocksApi, CredentialsApi, HealthApi, KeysApi, LimitsApi, OAuthApi, ResetsApi, SessionsApi, UsageApi } from "./auth/types.js";
3
- export { isSqliteBusyError, isSqliteCorruptionError, SqliteAuthCredentialStore } from "./auth/sqlite-credential-store.js";
3
+ export { isSqliteBusyError, isSqliteCorruptionError, resolveCredentialIdentityKey, SqliteAuthCredentialStore, } from "./auth/sqlite-credential-store.js";
4
4
  export * from "./auth/store.js";
5
5
  export * from "./auth/types.js";
6
6
  /**
@@ -46,9 +46,9 @@ export declare class AuthStorage {
46
46
  * Swap the backing credential store in place (live `auth.broker.url` change).
47
47
  * Loads `store` into fresh store-bound state — pins, blocks, and usage caches are
48
48
  * keyed by the old store's row ids — then closes the previous store. Runtime key
49
- * overrides, account policies, usage-provider overrides, and credential event
50
- * subscribers carry over. On a load failure `store` is closed and the current
51
- * store stays active.
49
+ * overrides, account policies, session account restrictions, usage-provider
50
+ * overrides, and credential event subscribers carry over. On a load failure
51
+ * `store` is closed and the current store stays active.
52
52
  */
53
53
  replaceStore(store: AuthCredentialStore, options?: {
54
54
  sourceLabel?: string;
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@oh-my-pi/pi-ai",
3
- "version": "18.8.0",
3
+ "version": "18.8.1",
4
4
  "description": "Unified LLM API with automatic model discovery and provider configuration",
5
5
  "keywords": [
6
6
  "ai",
@@ -155,11 +155,11 @@
155
155
  "fmt": "oxfmt --no-error-on-unmatched-pattern 'src/**/*.{ts,tsx}' '{test,bench,examples,scripts}/**/*.ts' '*.ts'"
156
156
  },
157
157
  "dependencies": {
158
- "@oh-my-pi/omptype": "18.8.0",
159
- "@oh-my-pi/pi-catalog": "18.8.0",
160
- "@oh-my-pi/pi-natives": "18.8.0",
161
- "@oh-my-pi/pi-utils": "18.8.0",
162
- "@oh-my-pi/pi-wire": "18.8.0"
158
+ "@oh-my-pi/omptype": "18.8.1",
159
+ "@oh-my-pi/pi-catalog": "18.8.1",
160
+ "@oh-my-pi/pi-natives": "18.8.1",
161
+ "@oh-my-pi/pi-utils": "18.8.1",
162
+ "@oh-my-pi/pi-wire": "18.8.1"
163
163
  },
164
164
  "devDependencies": {
165
165
  "@types/bun": "^1.3.14"
@@ -1,10 +1,11 @@
1
1
  import { logger } from "@oh-my-pi/pi-utils";
2
2
  import { LRUCache } from "@oh-my-pi/pi-utils/lru";
3
3
  import { getEnvApiKey } from "../env-api-key";
4
- import type { AuthCredential, OAuthCredential, SessionsApi } from "./types";
4
+ import type { AuthCredential, OAuthCredential, SessionRestrictionLease, SessionsApi } from "./types";
5
5
  import type { AuthCredentialStore } from "./store";
6
6
  import type { CredentialPool } from "./pool";
7
7
  import type { KeyOverrides } from "./cascade";
8
+ import { resolveCredentialIdentityKey } from "./sqlite-credential-store";
8
9
 
9
10
  /** Prefix for persisted session-to-credential affinity. */
10
11
  export const SESSION_STICKY_CACHE_PREFIX = "session:sticky:";
@@ -39,6 +40,16 @@ export type SessionCredential = {
39
40
  explicit?: true;
40
41
  };
41
42
 
43
+ /** One installed restriction: the allowed OAuth identity keys and the lease that lifts it. */
44
+ type SessionRestriction = { readonly allowed: ReadonlySet<string>; readonly lease: SessionRestrictionLease };
45
+
46
+ /**
47
+ * Provider → session id → installed restriction. Owned by `AuthStorage` so
48
+ * restrictions outlive store replacement, which rebuilds every store-bound
49
+ * module (pins included) around the same identities.
50
+ */
51
+ export type SessionRestrictions = Map<string, Map<string, SessionRestriction>>;
52
+
42
53
  /** Session → credential affinity (pins), persisted in the store cache. */
43
54
  export class SessionAffinity implements SessionsApi {
44
55
  /** Tracks the last used credential per provider for a session (used for rate-limit switching). */
@@ -48,11 +59,58 @@ export class SessionAffinity implements SessionsApi {
48
59
  #store: AuthCredentialStore;
49
60
  #pool: CredentialPool;
50
61
  #overrides: KeyOverrides;
62
+ #restrictions: SessionRestrictions;
51
63
 
52
- constructor(store: AuthCredentialStore, pool: CredentialPool, overrides: KeyOverrides) {
64
+ constructor(
65
+ store: AuthCredentialStore,
66
+ pool: CredentialPool,
67
+ overrides: KeyOverrides,
68
+ restrictions: SessionRestrictions,
69
+ ) {
53
70
  this.#store = store;
54
71
  this.#pool = pool;
55
72
  this.#overrides = overrides;
73
+ this.#restrictions = restrictions;
74
+ }
75
+
76
+ restrict(provider: string, sessionId: string, identityKeys: readonly string[]): SessionRestrictionLease {
77
+ if (!provider || !sessionId) throw new Error("sessions.restrict requires a provider and a session id");
78
+ const lease = Symbol("sessions.restrict");
79
+ const sessions = this.#restrictions.get(provider) ?? new Map<string, SessionRestriction>();
80
+ sessions.set(sessionId, { allowed: new Set(identityKeys), lease });
81
+ this.#restrictions.set(provider, sessions);
82
+ return lease;
83
+ }
84
+
85
+ unrestrict(provider: string, sessionId: string, lease: SessionRestrictionLease): void {
86
+ const sessions = this.#restrictions.get(provider);
87
+ if (sessions?.get(sessionId)?.lease !== lease) return;
88
+ sessions.delete(sessionId);
89
+ if (sessions.size === 0) this.#restrictions.delete(provider);
90
+ }
91
+
92
+ /** True when {@link restrict} limits `sessionId` for `provider`. */
93
+ isRestricted(provider: string, sessionId: string | undefined): boolean {
94
+ return sessionId !== undefined && this.#restrictions.get(provider)?.has(sessionId) === true;
95
+ }
96
+
97
+ /**
98
+ * True when `credential` may serve `sessionId`: always for an unrestricted
99
+ * session, otherwise only an OAuth credential whose identity key is allowed.
100
+ */
101
+ allows(provider: string, sessionId: string | undefined, credential: AuthCredential): boolean {
102
+ const allowed = sessionId === undefined ? undefined : this.#restrictions.get(provider)?.get(sessionId)?.allowed;
103
+ if (allowed === undefined) return true;
104
+ if (credential.type !== "oauth") return false;
105
+ const identityKey = resolveCredentialIdentityKey(provider, credential);
106
+ return identityKey !== null && allowed.has(identityKey);
107
+ }
108
+
109
+ /** Whether the stored row at `index` may serve `sessionId` (see {@link allows}). */
110
+ #permits(provider: string, sessionId: string, index: number): boolean {
111
+ if (!this.isRestricted(provider, sessionId)) return true;
112
+ const credential = this.#pool.credentials(provider)[index];
113
+ return credential !== undefined && this.allows(provider, sessionId, credential);
56
114
  }
57
115
 
58
116
  /** Bounded per-provider session map, created on first use. */
@@ -101,7 +159,7 @@ export class SessionAffinity implements SessionsApi {
101
159
  lastUsedAtMs?: number,
102
160
  explicit = false,
103
161
  ): void {
104
- if (!sessionId) return;
162
+ if (!sessionId || !this.#permits(provider, sessionId, index)) return;
105
163
  const nowMs = lastUsedAtMs ?? Date.now();
106
164
  const credentialId = this.#pool.entries(provider)[index]?.id;
107
165
  const sessionMap = SessionAffinity.#sessionsFor(this.#sessionLastCredential, provider);
@@ -147,9 +205,18 @@ export class SessionAffinity implements SessionsApi {
147
205
  }
148
206
  }
149
207
 
150
- /** Retrieves the last credential used by a session. */
208
+ /**
209
+ * Retrieves the last credential used by a session. A restricted session
210
+ * never sees a pin outside its allowlist — inherited, restored, or recorded
211
+ * before the restriction — so it re-ranks inside the allowlist instead.
212
+ */
151
213
  get(provider: string, sessionId: string | undefined): SessionCredential | undefined {
152
214
  if (!sessionId) return undefined;
215
+ const credential = this.#lookup(provider, sessionId);
216
+ return credential && this.#permits(provider, sessionId, credential.index) ? credential : undefined;
217
+ }
218
+
219
+ #lookup(provider: string, sessionId: string): SessionCredential | undefined {
153
220
  const live = this.#sessionLastCredential.get(provider)?.get(sessionId);
154
221
  if (live) {
155
222
  // Another process can add or drop rows mid-session and the pool is an
@@ -233,12 +300,17 @@ export class SessionAffinity implements SessionsApi {
233
300
 
234
301
  activeOAuth(provider: string, sessionId?: string): OAuthCredential | undefined {
235
302
  const allCredentials = this.#pool.credentials(provider);
236
- const oauthCredentials = allCredentials.filter((c): c is OAuthCredential => c.type === "oauth");
303
+ const oauthCredentials = allCredentials.filter(
304
+ (c): c is OAuthCredential => c.type === "oauth" && this.allows(provider, sessionId, c),
305
+ );
237
306
  if (oauthCredentials.length === 0) return undefined;
238
307
 
239
308
  // Runtime / config overrides bypass OAuth account_uuid attribution — the
240
309
  // caller is authenticating with an explicit key, not the broker's OAuth.
241
- if (this.#overrides.has(provider)) return undefined;
310
+ // A restricted session skips a runtime key, and a config key fails it
311
+ // closed (see `KeyCascade.get`).
312
+ const restricted = this.isRestricted(provider, sessionId);
313
+ if (this.#overrides.suppressesOAuth(provider, restricted)) return undefined;
242
314
 
243
315
  // Prefer the session-sticky credential when available.
244
316
  const sessionPref = this.get(provider, sessionId);
@@ -251,7 +323,7 @@ export class SessionAffinity implements SessionsApi {
251
323
  // would misattribute traffic. Only apply this guard when sessionPref is absent; a
252
324
  // recorded OAuth sticky (sessionPref.type === "oauth") must NOT be blocked even if an
253
325
  // env key also happens to exist.
254
- if (!sessionPref && getEnvApiKey(provider)) return undefined;
326
+ if (!restricted && !sessionPref && getEnvApiKey(provider)) return undefined;
255
327
  // Resolve the sticky index against the full credential list — the index is
256
328
  // recorded against the unfiltered provider array (by record /
257
329
  // CredentialSelector.tryOAuth), not the OAuth-only subset, so dereferencing it into the
@@ -275,13 +347,13 @@ export class SessionAffinity implements SessionsApi {
275
347
  * account, a stale resume re-ranks.
276
348
  */
277
349
  pin(provider: string, sessionId: string, credentialId: number, options?: { restoredAtMs?: number }): boolean {
278
- if (!sessionId || this.#overrides.has(provider)) {
350
+ if (!sessionId || this.#overrides.suppressesOAuth(provider, this.isRestricted(provider, sessionId))) {
279
351
  return false;
280
352
  }
281
353
  const stored = this.#pool.entries(provider);
282
354
  const index = stored.findIndex(entry => entry.id === credentialId);
283
355
  const target = stored[index];
284
- if (target?.credential.type !== "oauth") return false;
356
+ if (target?.credential.type !== "oauth" || !this.allows(provider, sessionId, target.credential)) return false;
285
357
  const restoredAtMs = options?.restoredAtMs;
286
358
  this.record(provider, sessionId, "oauth", index, restoredAtMs, restoredAtMs === undefined);
287
359
  return true;
@@ -300,7 +372,7 @@ export class SessionAffinity implements SessionsApi {
300
372
  let inherited = 0;
301
373
  for (const provider of this.#pool.providers()) {
302
374
  const credential = this.get(provider, sourceSessionId);
303
- if (!credential) continue;
375
+ if (!credential || !this.#permits(provider, targetSessionId, credential.index)) continue;
304
376
  this.record(
305
377
  provider,
306
378
  targetSessionId,
@@ -42,6 +42,17 @@ export class KeyOverrides {
42
42
  return this.#runtimeOverrides.has(provider) || this.#configOverrides.has(provider);
43
43
  }
44
44
 
45
+ /**
46
+ * Whether an override replaces OAuth for a session. Any override does for an
47
+ * unrestricted session. A restricted session skips the runtime key, but a
48
+ * config key still blocks OAuth: it marks the provider's endpoint (often a
49
+ * proxy) as taking that key, so the session fails closed instead of sending
50
+ * a pooled OAuth token there (see `KeyCascade.get`).
51
+ */
52
+ suppressesOAuth(provider: string, restricted: boolean): boolean {
53
+ return restricted ? this.#configOverrides.has(provider) : this.has(provider);
54
+ }
55
+
45
56
  runtimeKey(provider: string): string | undefined {
46
57
  return this.#runtimeOverrides.get(provider);
47
58
  }
@@ -316,6 +327,10 @@ export class KeyCascade implements KeysApi {
316
327
  * Get API key for a provider.
317
328
  * Priority (first match wins): runtime override, config override, OAuth,
318
329
  * login API key, environment variable, then another stored API key.
330
+ * A session restricted by `sessions.restrict` resolves only its allowed
331
+ * OAuth accounts and throws {@link AIError.MissingApiKeyError} when none
332
+ * can serve or a config key (a models.yml `apiKey`, often for a proxy
333
+ * `baseUrl`) owns the provider, rather than using any other credential.
319
334
  */
320
335
  async get(
321
336
  provider: string,
@@ -323,8 +338,9 @@ export class KeyCascade implements KeysApi {
323
338
  options?: AuthApiKeyOptions,
324
339
  onCredentialId?: (id: number, identity?: OAuthRequestIdentity) => void,
325
340
  ): Promise<string | undefined> {
341
+ const restricted = this.#deps.affinity.isRestricted(provider, sessionId);
326
342
  // Runtime override takes highest priority
327
- const runtimeKey = this.#deps.overrides.runtimeKey(provider);
343
+ const runtimeKey = restricted ? undefined : this.#deps.overrides.runtimeKey(provider);
328
344
  if (runtimeKey) {
329
345
  return runtimeKey;
330
346
  }
@@ -334,7 +350,13 @@ export class KeyCascade implements KeysApi {
334
350
  // (e.g. an auth-gateway) and supplied the bearer for that endpoint —
335
351
  // honor it instead of forwarding an upstream OAuth token that the proxy
336
352
  // won't accept.
337
- const configKey = this.#deps.overrides.configKey(provider);
353
+ if (restricted && this.#deps.overrides.configKey(provider) !== undefined) {
354
+ throw new AIError.MissingApiKeyError(
355
+ provider,
356
+ `No API key for provider: ${provider} (session ${sessionId} is restricted to its OAuth account pool, but models.yml configures an apiKey for this provider, and pooled OAuth tokens are never sent past it)`,
357
+ );
358
+ }
359
+ const configKey = restricted ? undefined : this.#deps.overrides.configKey(provider);
338
360
  if (configKey !== undefined) {
339
361
  return this.#deps.overrides.resolve(configKey);
340
362
  }
@@ -349,6 +371,12 @@ export class KeyCascade implements KeysApi {
349
371
  }
350
372
  return oauthResolved.apiKey;
351
373
  }
374
+ if (restricted) {
375
+ throw new AIError.MissingApiKeyError(
376
+ provider,
377
+ `No API key for provider: ${provider} (session ${sessionId} is restricted to its OAuth account pool and none of those accounts is available)`,
378
+ );
379
+ }
352
380
  const loginApiKeySelection = await this.#deps.selector.selectApiKey(
353
381
  provider,
354
382
  sessionId,
package/src/auth/oauth.ts CHANGED
@@ -134,8 +134,9 @@ export class OAuthAccounts implements OAuthApi {
134
134
  async access(provider: string, sessionId?: string, options?: AuthApiKeyOptions): Promise<OAuthAccess | undefined> {
135
135
  // Runtime / config overrides intentionally short-circuit OAuth: when the
136
136
  // user has pinned an API key, they expect the OAuth identity to be
137
- // suppressed (same contract as account identity lookup).
138
- if (this.#deps.overrides.has(provider)) {
137
+ // suppressed (same contract as account identity lookup). A session
138
+ // restricted to an account pool never uses a runtime key.
139
+ if (this.#deps.overrides.suppressesOAuth(provider, this.#deps.affinity.isRestricted(provider, sessionId))) {
139
140
  return undefined;
140
141
  }
141
142
  const resolved = await this.#deps.selector.resolveOAuth(provider, sessionId, options);
@@ -236,7 +237,7 @@ export class OAuthAccounts implements OAuthApi {
236
237
  * credential.
237
238
  */
238
239
  accounts(provider: string, sessionId?: string): OAuthAccountSummary[] {
239
- if (this.#deps.overrides.has(provider)) {
240
+ if (this.#deps.overrides.suppressesOAuth(provider, this.#deps.affinity.isRestricted(provider, sessionId))) {
240
241
  return [];
241
242
  }
242
243
  const sessionCredential = this.#deps.affinity.get(provider, sessionId);
@@ -141,6 +141,7 @@ export class RateLimits implements LimitsApi {
141
141
  }
142
142
  #blockCredentialForRotation(
143
143
  provider: string,
144
+ sessionId: string | undefined,
144
145
  credentialType: AuthCredential["type"],
145
146
  targetIndex: number,
146
147
  blockedUntil: number,
@@ -188,7 +189,9 @@ export class RateLimits implements LimitsApi {
188
189
  .map((credential, index) => ({ credential, index }))
189
190
  .filter(
190
191
  (entry): entry is { credential: AuthCredential; index: number } =>
191
- entry.credential.type === credentialType && entry.index !== targetIndex,
192
+ entry.credential.type === credentialType &&
193
+ entry.index !== targetIndex &&
194
+ this.#deps.affinity.allows(provider, sessionId, entry.credential),
192
195
  );
193
196
 
194
197
  let retryAtMs: number | undefined;
@@ -292,6 +295,7 @@ export class RateLimits implements LimitsApi {
292
295
  .findIndex(entry => entry.id === targetCredentialId && entry.credential.type === credentialType);
293
296
  const rotation = this.#blockCredentialForRotation(
294
297
  provider,
298
+ sessionId,
295
299
  credentialType,
296
300
  targetIndex,
297
301
  blockedUntil,
@@ -456,6 +460,7 @@ export class RateLimits implements LimitsApi {
456
460
  }
457
461
  const mark = this.#blockCredentialForRotation(
458
462
  provider,
463
+ sessionId,
459
464
  sessionCredential.type,
460
465
  sessionCredential.index,
461
466
  Date.now() + DEFAULT_BLOCK_MS,
@@ -474,7 +479,8 @@ export class RateLimits implements LimitsApi {
474
479
  (credential, index) =>
475
480
  credential.type === sessionCredential.type &&
476
481
  index !== sessionCredential.index &&
477
- !this.#deps.blocks.isBlocked(provider, providerKey, index),
482
+ !this.#deps.blocks.isBlocked(provider, providerKey, index) &&
483
+ this.#deps.affinity.allows(provider, sessionId, credential),
478
484
  );
479
485
  const target = this.#deps.pool.entries(provider)[sessionCredential.index];
480
486
  const sticky = this.#deps.affinity.get(provider, sessionId);
@@ -516,14 +516,17 @@ export class CredentialSelector {
516
516
  options?: AuthApiKeyOptions,
517
517
  ): Promise<OAuthResolutionResult | undefined> {
518
518
  await this.#deps.pool.adoptExternalChanges();
519
- const credentials = this.#deps.pool
519
+ const stored = this.#deps.pool
520
520
  .credentials(provider)
521
521
  .map((credential, index) => ({ credential, index }))
522
522
  .filter((entry): entry is { credential: OAuthCredential; index: number } => entry.credential.type === "oauth");
523
523
  this.#deps.policies.validateFor(
524
524
  provider,
525
- credentials.map(entry => entry.credential),
525
+ stored.map(entry => entry.credential),
526
526
  );
527
+ // A session restriction drops every other account before ranking, pins,
528
+ // and the fallback passes below, so none of them can route back to it.
529
+ const credentials = stored.filter(entry => this.#deps.affinity.allows(provider, sessionId, entry.credential));
527
530
 
528
531
  if (credentials.length === 0) return undefined;
529
532
  this.#deps.policies.validateUsageCapability(provider, this.#deps.usage.canFetchOAuthUsage(provider));
@@ -217,6 +217,12 @@ function resolveProviderCredentialIdentityKey(provider: string, identifiers: str
217
217
  return null;
218
218
  }
219
219
 
220
+ /**
221
+ * Identity key of a credential: the `identityKey` that broker snapshots carry
222
+ * and that broker account pools and `sessions.restrict` match, such as
223
+ * `email:<address>|org:<id>` for org-scoped providers. `null` for an API key or
224
+ * an OAuth credential without any account, email, project, or org identity.
225
+ */
220
226
  export function resolveCredentialIdentityKey(provider: string, credential: AuthCredential): string | null {
221
227
  if (credential.type === "api_key") return null;
222
228
  return resolveProviderCredentialIdentityKey(provider, extractOAuthCredentialIdentifiers(credential));
package/src/auth/types.ts CHANGED
@@ -1088,6 +1088,12 @@ export interface OAuthApi {
1088
1088
  ): Promise<StoredOAuthRefreshResult<T>>;
1089
1089
  }
1090
1090
 
1091
+ /**
1092
+ * Handle for one {@link SessionsApi.restrict} call. {@link SessionsApi.unrestrict}
1093
+ * lifts a restriction only with the lease of the call that installed it.
1094
+ */
1095
+ export type SessionRestrictionLease = symbol;
1096
+
1091
1097
  /** Session credential affinity operations. */
1092
1098
  export interface SessionsApi {
1093
1099
  /**
@@ -1113,6 +1119,27 @@ export interface SessionsApi {
1113
1119
  * source session.
1114
1120
  */
1115
1121
  inherit(sourceSessionId: string, targetSessionId: string): number;
1122
+ /**
1123
+ * Restrict one session's credentials for `provider` to the OAuth accounts
1124
+ * whose identity key (`email:<address>|org:<id>` for org-scoped providers;
1125
+ * otherwise `account:`, `email:`, or `project:` — the keys broker account
1126
+ * pools use) is listed. Selection, session pins, inherited affinity,
1127
+ * fallback passes, and rotation never leave the list, and runtime, config,
1128
+ * environment, and stored API keys are not used for that session; when no
1129
+ * listed account can serve, key resolution fails instead of borrowing
1130
+ * another account. An empty list allows no account. Replaces any earlier
1131
+ * restriction for the same provider and session; {@link inherit} does not
1132
+ * copy restrictions. Restrictions are never evicted, because that would
1133
+ * widen a live session; the owner passes the returned lease to
1134
+ * {@link unrestrict} when the session ends.
1135
+ */
1136
+ restrict(provider: string, sessionId: string, identityKeys: readonly string[]): SessionRestrictionLease;
1137
+ /**
1138
+ * Lift the restriction that `lease` installed. Once a later {@link restrict}
1139
+ * call has replaced it for the same provider and session, this does nothing,
1140
+ * so a stale owner never lifts the restriction a newer owner relies on.
1141
+ */
1142
+ unrestrict(provider: string, sessionId: string, lease: SessionRestrictionLease): void;
1116
1143
  /**
1117
1144
  * Release a session's sticky credential so its next {@link getApiKey} call
1118
1145
  * re-runs native pool ranking. This never blocks or penalizes the released
@@ -19,7 +19,7 @@
19
19
  * const apiKey = await auth.keys.get("anthropic", sessionId, { modelId });
20
20
  */
21
21
  import { logger } from "@oh-my-pi/pi-utils";
22
- import { SessionAffinity } from "./auth/affinity";
22
+ import { SessionAffinity, type SessionRestrictions } from "./auth/affinity";
23
23
  import { BlockStoreHealth, CredentialBlocks } from "./auth/blocks";
24
24
  import { KeyCascade, KeyOverrides } from "./auth/cascade";
25
25
  import { CredentialHealth } from "./auth/health";
@@ -52,7 +52,12 @@ import { DEFAULT_USAGE_REQUEST_TIMEOUT_MS, UsageCache } from "./auth/usage-cache
52
52
  import type { UsageLogger } from "./usage";
53
53
  import { defaultRankingStrategy, defaultUsageProvider } from "./usage/registry";
54
54
 
55
- export { isSqliteBusyError, isSqliteCorruptionError, SqliteAuthCredentialStore } from "./auth/sqlite-credential-store";
55
+ export {
56
+ isSqliteBusyError,
57
+ isSqliteCorruptionError,
58
+ resolveCredentialIdentityKey,
59
+ SqliteAuthCredentialStore,
60
+ } from "./auth/sqlite-credential-store";
56
61
  export * from "./auth/store";
57
62
  export * from "./auth/types";
58
63
 
@@ -81,6 +86,7 @@ export class AuthStorage {
81
86
  readonly #options: AuthStorageOptions;
82
87
  readonly #overrides: KeyOverrides;
83
88
  readonly #policies: AccountPolicies;
89
+ readonly #restrictions: SessionRestrictions = new Map();
84
90
  #modules: AuthStorageModules;
85
91
 
86
92
  constructor(store: AuthCredentialStore, options: AuthStorageOptions = {}) {
@@ -145,9 +151,9 @@ export class AuthStorage {
145
151
  * Swap the backing credential store in place (live `auth.broker.url` change).
146
152
  * Loads `store` into fresh store-bound state — pins, blocks, and usage caches are
147
153
  * keyed by the old store's row ids — then closes the previous store. Runtime key
148
- * overrides, account policies, usage-provider overrides, and credential event
149
- * subscribers carry over. On a load failure `store` is closed and the current
150
- * store stays active.
154
+ * overrides, account policies, session account restrictions, usage-provider
155
+ * overrides, and credential event subscribers carry over. On a load failure
156
+ * `store` is closed and the current store stays active.
151
157
  */
152
158
  async replaceStore(store: AuthCredentialStore, options: { sourceLabel?: string } = {}): Promise<void> {
153
159
  const next = this.#compose(store, options.sourceLabel ?? this.#options.sourceLabel);
@@ -185,7 +191,7 @@ export class AuthStorage {
185
191
  // override's `cacheVersion` separates its rows from other processes sharing the store.
186
192
  const usageCache = new UsageCache(store, pool, provider => usage.providerFor(provider));
187
193
  const blocks = new CredentialBlocks({ store, pool, health: blockHealth, usageCache, strategies });
188
- const affinity = new SessionAffinity(store, pool, overrides);
194
+ const affinity = new SessionAffinity(store, pool, overrides, this.#restrictions);
189
195
  const usage = new UsageService({
190
196
  store,
191
197
  pool,
@@ -1348,14 +1348,19 @@ function getCodexServiceTierCostMultiplier(
1348
1348
  }
1349
1349
 
1350
1350
  /**
1351
- * The tier a Codex response was billed at. The response echo is authoritative
1352
- * whenever it reports a tier (the backend may serve a requested priority/flex
1353
- * turn as `default`); the requested tier is used only when the echo is absent.
1354
- * The tier's identity is preserved even when it has no pricing entry (`scale`),
1355
- * because the recorded tier also drives premium-request and speed accounting.
1351
+ * The tier a Codex turn was served at. The Codex backend echoes
1352
+ * `service_tier: "default"` on turns it serves at the sent tier (a `priority`
1353
+ * turn decodes at Fast speed yet reports `default`), so a `default` or `auto`
1354
+ * echo carries no information and, like an absent one, resolves to the sent
1355
+ * tier, or `default` when none was sent. Any other echo (`flex`, `priority`,
1356
+ * `ultrafast`, `scale`) is authoritative. The tier's identity is preserved even
1357
+ * when it has no pricing entry (`scale`), because the recorded tier also drives
1358
+ * premium-request and speed accounting. The echo cannot reveal a genuine
1359
+ * downgrade, so a downgraded turn is recorded at the sent tier.
1356
1360
  */
1357
1361
  function resolveCodexCostServiceTier(res: ServiceTier | undefined, req?: unknown): ServiceTier {
1358
- return res ?? parseServiceTier(req) ?? "default";
1362
+ if (res !== undefined && res !== "default" && res !== "auto") return res;
1363
+ return parseServiceTier(req) ?? "default";
1359
1364
  }
1360
1365
 
1361
1366
  function applyCodexServiceTierPricing(
@@ -1921,14 +1926,20 @@ async function openCodexWebSocketTransport(
1921
1926
  }
1922
1927
  recordCodexTurnRequestDiagnostics(websocketState, websocketRequest, "websocket", canAppendBeforeRequest);
1923
1928
  const requestBodyForState = cloneJsonTree(requestContext.transformedBody);
1924
- // `onPayload` may rewrite the outgoing frame (e.g. drop `stream_options`);
1925
- // recorded state must reflect what was actually sent — the sequential-cutoff
1926
- // summary decoder keys off it.
1929
+ // `onPayload` may rewrite the outgoing frame (e.g. drop `stream_options` or
1930
+ // change `service_tier`); recorded state must reflect what was actually
1931
+ // sent — the sequential-cutoff summary decoder and the served-tier
1932
+ // fallback key off it.
1927
1933
  if (websocketRequest.stream_options === undefined) {
1928
1934
  delete requestBodyForState.stream_options;
1929
1935
  } else {
1930
1936
  requestBodyForState.stream_options = websocketRequest.stream_options;
1931
1937
  }
1938
+ if (websocketRequest.service_tier === undefined) {
1939
+ delete requestBodyForState.service_tier;
1940
+ } else {
1941
+ requestBodyForState.service_tier = websocketRequest.service_tier;
1942
+ }
1932
1943
  requestContext.rawRequestDump.body = websocketRequest;
1933
1944
  CODEX_DEBUG &&
1934
1945
  logger.debug("[codex] codex websocket request", {
@@ -362,7 +362,7 @@ export function applyOpenAIServiceTier(
362
362
  * Standard OpenAI Responses service-tier cost multipliers. The non-Codex
363
363
  * Responses path bills the tier it was served (or requested): Flex processing is
364
364
  * half price; Priority (Fast mode) is a 2x premium. Codex bills the same tiers
365
- * with its own table (Priority is 2.5x on gpt-5.5) and applies that separately.
365
+ * with its own table (Fast is 2.5x on every model) and applies that separately.
366
366
  * `ultrafast` has no API-generic default — only models with a published
367
367
  * ultrafast price carry a `serviceTierCost.ultrafast` entry (Astra, 6x) and
368
368
  * everything else stays at 1x rather than an invented multiplier.