From 20bad972fb5fe2d56d9abb02fa0648aedf59c8df Mon Sep 17 00:00:00 2001 From: Michiel Westerbeek Date: Tue, 1 Sep 2026 21:01:34 +0000 Subject: [PATCH] Add subscription gateway account pool core --- bun.lock | 6 + package.json | 1 + packages/core/subscription-gateway/README.md | 5 + .../core/subscription-gateway/package.json | 18 ++ .../core/subscription-gateway/src/index.ts | 2 + .../src/pool/account-walk.test.ts | 121 +++++++++++ .../src/pool/account-walk.ts | 110 ++++++++++ .../src/pool/cooldowns.test.ts | 107 +++++++++ .../src/pool/cooldowns.ts | 144 ++++++++++++ .../subscription-gateway/src/pool/index.ts | 5 + .../src/pool/router.test.ts | 205 ++++++++++++++++++ .../subscription-gateway/src/pool/router.ts | 163 ++++++++++++++ .../src/pool/selectors.ts | 60 +++++ .../subscription-gateway/src/pool/types.ts | 71 ++++++ .../src/storage/cooldown-repository.ts | 104 +++++++++ .../subscription-gateway/src/storage/index.ts | 1 + 16 files changed, 1123 insertions(+) create mode 100644 packages/core/subscription-gateway/README.md create mode 100644 packages/core/subscription-gateway/package.json create mode 100644 packages/core/subscription-gateway/src/index.ts create mode 100644 packages/core/subscription-gateway/src/pool/account-walk.test.ts create mode 100644 packages/core/subscription-gateway/src/pool/account-walk.ts create mode 100644 packages/core/subscription-gateway/src/pool/cooldowns.test.ts create mode 100644 packages/core/subscription-gateway/src/pool/cooldowns.ts create mode 100644 packages/core/subscription-gateway/src/pool/index.ts create mode 100644 packages/core/subscription-gateway/src/pool/router.test.ts create mode 100644 packages/core/subscription-gateway/src/pool/router.ts create mode 100644 packages/core/subscription-gateway/src/pool/selectors.ts create mode 100644 packages/core/subscription-gateway/src/pool/types.ts create mode 100644 packages/core/subscription-gateway/src/storage/cooldown-repository.ts create mode 100644 packages/core/subscription-gateway/src/storage/index.ts diff --git a/bun.lock b/bun.lock index e87712f37f..ade15dc1db 100644 --- a/bun.lock +++ b/bun.lock @@ -139,6 +139,10 @@ "name": "@tellahq/opensession-protocol", "version": "0.1.0", }, + "packages/core/subscription-gateway": { + "name": "@tellahq/subscription-gateway", + "version": "0.0.0", + }, "packages/integrations/apple-mobile": { "name": "@tellahq/opensession-apple-mobile", "version": "0.1.0", @@ -813,6 +817,8 @@ "@tellahq/opensession-website": ["@tellahq/opensession-website@workspace:packages/clients/website"], + "@tellahq/subscription-gateway": ["@tellahq/subscription-gateway@workspace:packages/core/subscription-gateway"], + "@types/bun": ["@types/bun@1.3.14", "", { "dependencies": { "bun-types": "1.3.14" } }, "sha512-h1hFqFVcvAvD9j9K7ZW7vd82aSA+rTdznZa+5bwvCwqSB1jmmfLcbIWhOLx1/+boy/xmjgCs/OMUL8hRJSmnPw=="], "@types/d3": ["@types/d3@7.4.3", "", { "dependencies": { "@types/d3-array": "*", "@types/d3-axis": "*", "@types/d3-brush": "*", "@types/d3-chord": "*", "@types/d3-color": "*", "@types/d3-contour": "*", "@types/d3-delaunay": "*", "@types/d3-dispatch": "*", "@types/d3-drag": "*", "@types/d3-dsv": "*", "@types/d3-ease": "*", "@types/d3-fetch": "*", "@types/d3-force": "*", "@types/d3-format": "*", "@types/d3-geo": "*", "@types/d3-hierarchy": "*", "@types/d3-interpolate": "*", "@types/d3-path": "*", "@types/d3-polygon": "*", "@types/d3-quadtree": "*", "@types/d3-random": "*", "@types/d3-scale": "*", "@types/d3-scale-chromatic": "*", "@types/d3-selection": "*", "@types/d3-shape": "*", "@types/d3-time": "*", "@types/d3-time-format": "*", "@types/d3-timer": "*", "@types/d3-transition": "*", "@types/d3-zoom": "*" } }, "sha512-lZXZ9ckh5R8uiFVt8ogUNf+pIrK4EsWrx2Np75WvF/eTpJ0FMHNhjXk8CKEx/+gpHbNQyJWehbFaTvqmHWB3ww=="], diff --git a/package.json b/package.json index e2fd4fde1e..eb0e68853c 100644 --- a/package.json +++ b/package.json @@ -16,6 +16,7 @@ "packages/clients/website", "packages/core/opensession-server", "packages/core/protocol", + "packages/core/subscription-gateway", "packages/integrations/apple-mobile" ], "scripts": { diff --git a/packages/core/subscription-gateway/README.md b/packages/core/subscription-gateway/README.md new file mode 100644 index 0000000000..d1866ac9fa --- /dev/null +++ b/packages/core/subscription-gateway/README.md @@ -0,0 +1,5 @@ +# Subscription Gateway + +Subscription Gateway will expose Claude and ChatGPT subscription accounts through an OpenAI-compatible local API. It is private while the account pool, provider adapters, and HTTP contract are being extracted and tested. + +The first package layer contains no provider SDK or Open Session runtime code. It owns account routing, durable cooldowns, and retries that stop once an attempt emits client-visible output. diff --git a/packages/core/subscription-gateway/package.json b/packages/core/subscription-gateway/package.json new file mode 100644 index 0000000000..5a901e75e4 --- /dev/null +++ b/packages/core/subscription-gateway/package.json @@ -0,0 +1,18 @@ +{ + "name": "@tellahq/subscription-gateway", + "version": "0.0.0", + "private": true, + "description": "Account routing and OpenAI-compatible access for AI subscription providers", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/tellahq/opensession.git", + "directory": "packages/core/subscription-gateway" + }, + "type": "module", + "exports": { + ".": "./src/index.ts", + "./pool": "./src/pool/index.ts", + "./storage": "./src/storage/index.ts" + } +} diff --git a/packages/core/subscription-gateway/src/index.ts b/packages/core/subscription-gateway/src/index.ts new file mode 100644 index 0000000000..26a2b47fec --- /dev/null +++ b/packages/core/subscription-gateway/src/index.ts @@ -0,0 +1,2 @@ +export * from "./pool"; +export * from "./storage"; diff --git a/packages/core/subscription-gateway/src/pool/account-walk.test.ts b/packages/core/subscription-gateway/src/pool/account-walk.test.ts new file mode 100644 index 0000000000..0147372657 --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/account-walk.test.ts @@ -0,0 +1,121 @@ +import { describe, expect, test } from "bun:test"; +import { walkAccounts } from "./account-walk"; +import type { AccountResolution, RoutableAccount } from "./types"; + +const shared = (index: number): RoutableAccount => ({ + id: `account-${index}`, + name: `Account ${index}`, + access: { kind: "shared" }, +}); + +describe("walkAccounts", () => { + test("walks every account without a fixed retry limit", async () => { + const accounts = Array.from({ length: 25 }, (_, index) => shared(index)); + const attempted: string[] = []; + const result = await walkAccounts({ + acquire: (excluded): AccountResolution => { + const account = accounts.find( + (candidate) => !excluded.has(candidate.id), + ); + return account + ? { kind: "selected", account, reason: "pool" } + : { kind: "refused", refusal: { kind: "pool-dry" } }; + }, + attempt: ({ account }) => { + attempted.push(account.id); + return account.id === "account-24" + ? { kind: "succeeded", value: "done" } + : { kind: "retry-account", error: new Error("limited") }; + }, + onEvent: () => undefined, + }); + + expect(result).toMatchObject({ + kind: "succeeded", + account: { id: "account-24" }, + value: "done", + }); + expect(attempted).toHaveLength(25); + }); + + test("retries after replay-safe events", async () => { + const events: string[] = []; + const attempts: string[] = []; + const result = await walkAccounts({ + acquire: (excluded): AccountResolution => { + const account = [shared(0), shared(1)].find( + (candidate) => !excluded.has(candidate.id), + ); + return account + ? { kind: "selected", account, reason: "pool" } + : { kind: "refused", refusal: { kind: "pool-dry" } }; + }, + attempt: ({ account, emit }) => { + attempts.push(account.id); + emit({ kind: "replay-safe", value: `usage:${account.id}` }); + return account.id === "account-0" + ? { kind: "retry-account", error: new Error("limited") } + : { kind: "succeeded", value: "done" }; + }, + onEvent: (event: string) => events.push(event), + }); + expect(result.kind).toBe("succeeded"); + expect(attempts).toEqual(["account-0", "account-1"]); + expect(events).toEqual(["usage:account-0", "usage:account-1"]); + }); + + test("never replays after client-visible output", async () => { + const attempts: string[] = []; + const result = await walkAccounts({ + acquire: (excluded): AccountResolution => { + const account = [shared(0), shared(1)].find( + (candidate) => !excluded.has(candidate.id), + ); + return account + ? { kind: "selected", account, reason: "pool" } + : { kind: "refused", refusal: { kind: "pool-dry" } }; + }, + attempt: ({ account, emit }) => { + attempts.push(account.id); + emit({ kind: "client-visible", value: "partial answer" }); + return { kind: "retry-account", error: new Error("limited") }; + }, + onEvent: () => undefined, + }); + expect(result).toMatchObject({ + kind: "failed", + account: { id: "account-0" }, + replayBlocked: true, + }); + expect(attempts).toEqual(["account-0"]); + }); + + test("a strict pin refuses rather than widening", async () => { + const pinned = shared(0); + const fallback = shared(1); + const result = await walkAccounts({ + acquire: (excluded): AccountResolution => { + if (excluded.has(pinned.id)) { + return { + kind: "refused", + refusal: { + kind: "pin-unusable", + pinnedId: pinned.id, + pinName: pinned.name, + }, + }; + } + return { kind: "selected", account: pinned, reason: "pinned" }; + }, + attempt: () => ({ + kind: "retry-account", + error: new Error(`do not use ${fallback.id}`), + }), + onEvent: () => undefined, + }); + expect(result).toMatchObject({ + kind: "refused", + refusal: { kind: "pin-unusable", pinnedId: "account-0" }, + }); + }); +}); diff --git a/packages/core/subscription-gateway/src/pool/account-walk.ts b/packages/core/subscription-gateway/src/pool/account-walk.ts new file mode 100644 index 0000000000..3af677d9bb --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/account-walk.ts @@ -0,0 +1,110 @@ +import type { + AccountRefusal, + AccountResolution, + RoutableAccount, +} from "./types"; + +export type AttemptEvent = + | { readonly kind: "replay-safe"; readonly value: TEvent } + | { readonly kind: "client-visible"; readonly value: TEvent }; + +export type AttemptOutcome = + | { readonly kind: "succeeded"; readonly value: TResult } + | { readonly kind: "retry-account"; readonly error: unknown } + | { readonly kind: "failed"; readonly error: unknown }; + +export interface AccountAttempt { + readonly account: TAccount; + readonly signal?: AbortSignal; + readonly emit: (event: AttemptEvent) => void; +} + +export type AccountWalkResult = + | { + readonly kind: "succeeded"; + readonly account: TAccount; + readonly value: TResult; + } + | { readonly kind: "refused"; readonly refusal: AccountRefusal } + | { + readonly kind: "failed"; + readonly account: TAccount; + readonly error: unknown; + readonly replayBlocked: boolean; + } + | { readonly kind: "aborted" } + | { readonly kind: "invalid-selection"; readonly accountId: string }; + +export interface WalkAccountsOptions< + TAccount extends RoutableAccount, + TEvent, + TResult, +> { + readonly acquire: ( + excludedIds: ReadonlySet, + ) => AccountResolution | Promise>; + readonly attempt: ( + input: AccountAttempt, + ) => AttemptOutcome | Promise>; + readonly onEvent: (event: TEvent) => void; + readonly onRetry?: ( + account: TAccount, + error: unknown, + ) => void | Promise; + readonly signal?: AbortSignal; +} + +/** + * Tries accounts until one succeeds or the router refuses another pick. + * A retry is allowed only before an attempt emits client-visible output. + */ +export async function walkAccounts< + TAccount extends RoutableAccount, + TEvent, + TResult, +>( + options: WalkAccountsOptions, +): Promise> { + const excludedIds = new Set(); + + while (!options.signal?.aborted) { + const resolution = await options.acquire(excludedIds); + if (resolution.kind === "refused") return resolution; + const { account } = resolution; + if (excludedIds.has(account.id)) { + return { kind: "invalid-selection", accountId: account.id }; + } + + let replayBlocked = false; + let outcome: AttemptOutcome; + try { + outcome = await options.attempt({ + account, + signal: options.signal, + emit: (event) => { + if (event.kind === "client-visible") replayBlocked = true; + options.onEvent(event.value); + }, + }); + } catch (error) { + outcome = { kind: "failed", error }; + } + + if (outcome.kind === "succeeded") { + return { kind: "succeeded", account, value: outcome.value }; + } + if (outcome.kind === "failed" || replayBlocked) { + return { + kind: "failed", + account, + error: outcome.error, + replayBlocked, + }; + } + + excludedIds.add(account.id); + await options.onRetry?.(account, outcome.error); + } + + return { kind: "aborted" }; +} diff --git a/packages/core/subscription-gateway/src/pool/cooldowns.test.ts b/packages/core/subscription-gateway/src/pool/cooldowns.test.ts new file mode 100644 index 0000000000..4de3c07792 --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/cooldowns.test.ts @@ -0,0 +1,107 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import { mkdtemp, readFile, rm, stat } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { JsonCooldownRepository, MemoryCooldownRepository } from "../storage"; +import { CooldownRegistry } from "./cooldowns"; + +const temporaryDirectories: string[] = []; + +afterEach(async () => { + await Promise.all( + temporaryDirectories + .splice(0) + .map((path) => rm(path, { recursive: true, force: true })), + ); +}); + +describe("CooldownRegistry", () => { + test("separates model exhaustion from account-wide exhaustion", async () => { + let now = 1_000; + const registry = await CooldownRegistry.open( + new MemoryCooldownRepository(), + () => now, + ); + await registry.markExhausted({ + accountId: "account-a", + model: "model-a", + until: 2_000, + }); + expect(registry.isActive("account-a", "model-a")).toBe(true); + expect(registry.isActive("account-a", "model-b")).toBe(false); + + await registry.markExhausted({ accountId: "account-a", until: 3_000 }); + expect(registry.isActive("account-a", "model-b")).toBe(true); + now = 3_001; + expect(registry.isActive("account-a", "model-a")).toBe(false); + }); + + test("does not let a wedge shorten exhaustion", async () => { + const registry = await CooldownRegistry.open( + new MemoryCooldownRepository(), + () => 1_000, + ); + await registry.markExhausted({ accountId: "account-a", until: 10_000 }); + expect(await registry.markWedged("account-a", 5_000)).toBeUndefined(); + expect(registry.isActive("account-a")).toBe(true); + }); + + test("clears only the wedge represented by its token", async () => { + const registry = await CooldownRegistry.open( + new MemoryCooldownRepository(), + () => 1_000, + ); + const token = await registry.markWedged("account-a", 5_000); + expect(token).toBeDefined(); + if (!token) return; + await registry.markExhausted({ accountId: "account-a", until: 4_000 }); + expect(await registry.clearWedge(token)).toBe(false); + expect(registry.isActive("account-a")).toBe(true); + }); + + test("restores a shorter cooldown when a wedge is rolled back", async () => { + const registry = await CooldownRegistry.open( + new MemoryCooldownRepository(), + () => 1_000, + ); + await registry.markExhausted({ accountId: "account-a", until: 2_000 }); + const token = await registry.markWedged("account-a", 5_000); + expect(token).toBeDefined(); + if (!token) return; + expect(await registry.clearWedge(token)).toBe(true); + expect(registry.activeRecords()).toEqual([ + { + scope: "account", + accountId: "account-a", + reason: "exhausted", + until: 2_000, + }, + ]); + }); + + test("hydrates durable cooldowns after restart", async () => { + const directory = await mkdtemp(join(tmpdir(), "subscription-gateway-")); + temporaryDirectories.push(directory); + const path = join(directory, "state", "cooldowns.json"); + const repository = new JsonCooldownRepository(path); + const first = await CooldownRegistry.open(repository, () => 1_000); + await first.markExhausted({ + accountId: "account-a", + model: "model-a", + until: 10_000, + }); + + const second = await CooldownRegistry.open(repository, () => 2_000); + expect(second.isActive("account-a", "model-a")).toBe(true); + expect(JSON.parse(await readFile(path, "utf8"))).toEqual([ + { + scope: "model", + accountId: "account-a", + model: "model-a", + reason: "exhausted", + until: 10_000, + }, + ]); + expect((await stat(path)).mode & 0o777).toBe(0o600); + }); +}); diff --git a/packages/core/subscription-gateway/src/pool/cooldowns.ts b/packages/core/subscription-gateway/src/pool/cooldowns.ts new file mode 100644 index 0000000000..1c69f75a19 --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/cooldowns.ts @@ -0,0 +1,144 @@ +import type { + CooldownRecord, + CooldownRepository, +} from "../storage/cooldown-repository"; + +export interface WedgeCooldownToken { + readonly accountId: string; + readonly until: number; + readonly previous?: Extract; +} + +function recordKey(record: CooldownRecord): string { + return record.scope === "account" + ? `account\0${record.accountId}` + : `model\0${record.accountId}\0${record.model}`; +} + +export class CooldownRegistry { + readonly #repository: CooldownRepository; + readonly #clock: () => number; + readonly #records = new Map(); + #saveTail: Promise = Promise.resolve(); + + private constructor( + repository: CooldownRepository, + clock: () => number, + records: readonly CooldownRecord[], + ) { + this.#repository = repository; + this.#clock = clock; + const now = clock(); + for (const record of records) { + if (record.until > now) this.#records.set(recordKey(record), record); + } + } + + static async open( + repository: CooldownRepository, + clock: () => number = Date.now, + ): Promise { + return new CooldownRegistry(repository, clock, await repository.load()); + } + + isActive(accountId: string, model?: string): boolean { + return ( + this.#activeRecord(`account\0${accountId}`) !== undefined || + (model + ? this.#activeRecord(`model\0${accountId}\0${model}`) !== undefined + : false) + ); + } + + async markExhausted(input: { + readonly accountId: string; + readonly until: number; + readonly model?: string; + }): Promise { + const record: CooldownRecord = input.model + ? { + scope: "model", + accountId: input.accountId, + model: input.model, + reason: "exhausted", + until: input.until, + } + : { + scope: "account", + accountId: input.accountId, + reason: "exhausted", + until: input.until, + }; + const key = recordKey(record); + const current = this.#activeRecord(key); + if (current && current.until >= record.until) { + if (current.reason === "exhausted") return; + this.#records.set(key, { ...record, until: current.until }); + } else { + this.#records.set(key, record); + } + await this.#persist(); + } + + async markWedged( + accountId: string, + durationMs: number, + ): Promise { + const until = this.#clock() + durationMs; + const key = `account\0${accountId}`; + const current = this.#activeRecord(key); + if (current && current.until >= until) return undefined; + this.#records.set(key, { + scope: "account", + accountId, + reason: "wedged", + until, + }); + await this.#persist(); + return { + accountId, + until, + ...(current?.scope === "account" ? { previous: current } : {}), + }; + } + + async clearWedge(token: WedgeCooldownToken): Promise { + const key = `account\0${token.accountId}`; + const current = this.#activeRecord(key); + if ( + !current || + current.scope !== "account" || + current.reason !== "wedged" || + current.until !== token.until + ) { + return false; + } + if (token.previous && token.previous.until > this.#clock()) { + this.#records.set(key, token.previous); + } else { + this.#records.delete(key); + } + await this.#persist(); + return true; + } + + activeRecords(): readonly CooldownRecord[] { + for (const key of this.#records.keys()) this.#activeRecord(key); + return [...this.#records.values()]; + } + + #activeRecord(key: string): CooldownRecord | undefined { + const record = this.#records.get(key); + if (!record) return undefined; + if (record.until > this.#clock()) return record; + this.#records.delete(key); + return undefined; + } + + async #persist(): Promise { + const snapshot = [...this.#records.values()]; + const save = this.#saveTail.then(() => this.#repository.save(snapshot)); + this.#saveTail = save.catch(() => undefined); + await save; + } +} diff --git a/packages/core/subscription-gateway/src/pool/index.ts b/packages/core/subscription-gateway/src/pool/index.ts new file mode 100644 index 0000000000..b12a38ced7 --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/index.ts @@ -0,0 +1,5 @@ +export * from "./account-walk"; +export * from "./cooldowns"; +export * from "./router"; +export * from "./selectors"; +export * from "./types"; diff --git a/packages/core/subscription-gateway/src/pool/router.test.ts b/packages/core/subscription-gateway/src/pool/router.test.ts new file mode 100644 index 0000000000..3959d18336 --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/router.test.ts @@ -0,0 +1,205 @@ +import { describe, expect, test } from "bun:test"; +import { AccountRouter } from "./router"; +import { hrwScore, leastRecentlyPicked, rendezvousAffinity } from "./selectors"; +import type { AccountAvailability, RoutableAccount } from "./types"; + +interface TestAccount extends RoutableAccount { + readonly capacity: "subscription" | "paid"; +} + +const accounts: TestAccount[] = [ + { + id: "shared-a", + name: "Shared A", + access: { kind: "shared" }, + capacity: "subscription", + }, + { + id: "shared-b", + name: "Shared B", + access: { kind: "shared" }, + capacity: "paid", + }, + { + id: "shared-c", + name: "Shared C", + access: { kind: "shared" }, + capacity: "subscription", + }, + { + id: "alex", + name: "Alex", + access: { kind: "personal", owner: "Alex" }, + capacity: "subscription", + }, + { + id: "grant", + name: "Grant", + access: { kind: "personal", owner: "Grant" }, + capacity: "subscription", + }, +]; + +function router(input?: { + readonly unavailable?: ReadonlySet; + readonly affinity?: boolean; +}): AccountRouter { + let now = 0; + return new AccountRouter({ + accounts: () => accounts, + availability: (account, request): AccountAvailability => { + if (input?.unavailable?.has(account.id)) return { kind: "unavailable" }; + if (account.capacity === "paid" && !request.allowPaidUsage) { + return { kind: "unavailable" }; + } + return { + kind: "available", + priority: account.capacity === "subscription" ? 0 : 100, + }; + }, + select: input?.affinity + ? rendezvousAffinity() + : leastRecentlyPicked(), + ownerMatches: (principalId, owner) => + principalId.toLowerCase() === owner.toLowerCase(), + clock: () => { + now += 1; + return now; + }, + }); +} + +function selectedId( + resolution: ReturnType["resolve"]>, +) { + expect(resolution.kind).toBe("selected"); + return resolution.kind === "selected" ? resolution.account.id : undefined; +} + +describe("AccountRouter", () => { + test("applies the owner gate to pins, designations, and pool picks", () => { + const subject = router(); + const softPin = subject.resolve({ principalId: "Robin", pinnedId: "alex" }); + expect(selectedId(softPin)).toBe("shared-a"); + + expect( + subject.resolve({ + principalId: "Robin", + pinnedId: "alex", + strictPin: true, + }), + ).toEqual({ + kind: "refused", + refusal: { + kind: "pin-unusable", + pinnedId: "alex", + pinName: "Alex", + }, + }); + + expect( + subject.resolve({ principalId: "Robin", designatedIds: ["alex"] }), + ).toEqual({ + kind: "refused", + refusal: { kind: "designated-dry", tried: "Alex" }, + }); + expect(selectedId(subject.resolve({ principalId: "Alex" }))).toBe("alex"); + expect(selectedId(subject.resolve({}))).toBe("shared-c"); + }); + + test("honors pin, sticky, and designation order", () => { + const subject = router(); + expect(selectedId(subject.resolve({ pinnedId: "shared-a" }))).toBe( + "shared-a", + ); + expect(selectedId(subject.resolve({ stickyId: "shared-a" }))).toBe( + "shared-a", + ); + expect( + selectedId( + subject.resolve({ + designatedIds: ["shared-b", "shared-a"], + allowPaidUsage: true, + }), + ), + ).toBe("shared-b"); + expect( + subject.resolve({ + pinnedId: "shared-a", + strictPin: true, + designatedIds: ["shared-b"], + allowPaidUsage: true, + }), + ).toMatchObject({ + kind: "refused", + refusal: { kind: "pin-not-designated" }, + }); + }); + + test("uses subscription capacity before opted-in paid capacity", () => { + const subject = router(); + expect(selectedId(subject.resolve({ allowPaidUsage: true }))).toBe( + "shared-a", + ); + const withoutSubscription = router({ + unavailable: new Set(["shared-a", "shared-c"]), + }); + expect( + selectedId(withoutSubscription.resolve({ allowPaidUsage: true })), + ).toBe("shared-b"); + expect(withoutSubscription.resolve({})).toMatchObject({ + kind: "refused", + refusal: { kind: "pool-dry" }, + }); + }); + + test("does not consume an LRU turn while peeking", () => { + const subject = router(); + expect(selectedId(subject.resolve({ allowPaidUsage: true }))).toBe( + "shared-a", + ); + const firstPeek = selectedId( + subject.resolve({ allowPaidUsage: true, recordPick: false }), + ); + const secondPeek = selectedId( + subject.resolve({ allowPaidUsage: true, recordPick: false }), + ); + expect(firstPeek).toBe("shared-c"); + expect(secondPeek).toBe(firstPeek); + }); + + test("keeps rendezvous affinity stable while excluding failed accounts", () => { + const subject = router({ affinity: true }); + const request = { affinityKey: "bks-test-session", allowPaidUsage: true }; + const first = selectedId(subject.resolve(request)); + const second = selectedId(subject.resolve(request)); + expect(second).toBe(first); + const fallback = selectedId( + subject.resolve({ ...request, excludeIds: new Set([first ?? ""]) }), + ); + expect(fallback).not.toBe(first); + }); +}); + +describe("hrwScore", () => { + test("keeps the existing affinity vectors pinned", () => { + expect( + hrwScore( + "bks-019f7182-a597-7000-96b0-50fdc06f8694", + "eae22618-bd72-45ab-8307-4949b5e409cd", + ), + ).toBe(1742935766); + expect( + hrwScore( + "bks-019f7182-a597-7000-96b0-50fdc06f8694", + "13fde4f9-e1f2-486c-8e04-1d0f322b7636", + ), + ).toBe(3956256899); + expect( + hrwScore("bks-test-session", "eae22618-bd72-45ab-8307-4949b5e409cd"), + ).toBe(3693026164); + expect( + hrwScore("bks-test-session", "13fde4f9-e1f2-486c-8e04-1d0f322b7636"), + ).toBe(1275860373); + }); +}); diff --git a/packages/core/subscription-gateway/src/pool/router.ts b/packages/core/subscription-gateway/src/pool/router.ts new file mode 100644 index 0000000000..afb2cefac4 --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/router.ts @@ -0,0 +1,163 @@ +import type { + AccountAvailability, + AccountRequest, + AccountResolution, + AccountSelector, + RoutableAccount, + SelectionCandidate, +} from "./types"; + +export interface AccountRouterOptions< + TAccount extends RoutableAccount, + TModel, +> { + readonly accounts: () => readonly TAccount[]; + readonly availability: ( + account: TAccount, + request: Pick, "model" | "allowPaidUsage">, + ) => AccountAvailability; + readonly select: AccountSelector; + readonly ownerMatches?: (principalId: string, owner: string) => boolean; + readonly clock?: () => number; +} + +export class AccountRouter { + readonly #options: AccountRouterOptions; + readonly #lastPickedAt = new Map(); + + constructor(options: AccountRouterOptions) { + this.#options = options; + } + + resolve(request: AccountRequest): AccountResolution { + const accounts = this.#options.accounts(); + const byId = new Map(accounts.map((account) => [account.id, account])); + const designatedIds = request.designatedIds?.length + ? request.designatedIds + : undefined; + const isDesignated = (id: string): boolean => + !designatedIds || designatedIds.includes(id); + const nameOf = (id: string): string => byId.get(id)?.name ?? id; + const usable = (id: string): SelectionCandidate | undefined => { + if (request.excludeIds?.has(id)) return undefined; + const account = byId.get(id); + if (!account || !this.#canUse(account, request.principalId)) { + return undefined; + } + const availability = this.#options.availability(account, request); + return availability.kind === "available" + ? { account, priority: availability.priority } + : undefined; + }; + + if (request.pinnedId) { + if (!isDesignated(request.pinnedId)) { + if (request.strictPin) { + return { + kind: "refused", + refusal: { + kind: "pin-not-designated", + pinnedId: request.pinnedId, + pinName: nameOf(request.pinnedId), + }, + }; + } + } else { + const pinned = usable(request.pinnedId); + if (pinned) { + return { + kind: "selected", + account: pinned.account, + reason: "pinned", + }; + } + if (request.strictPin) { + return { + kind: "refused", + refusal: { + kind: "pin-unusable", + pinnedId: request.pinnedId, + pinName: nameOf(request.pinnedId), + }, + }; + } + } + } + + if (request.stickyId && isDesignated(request.stickyId)) { + const sticky = usable(request.stickyId); + if (sticky) { + return { + kind: "selected", + account: sticky.account, + reason: "sticky", + }; + } + } + + if (designatedIds) { + for (const id of designatedIds) { + const designated = usable(id); + if (designated) { + return { + kind: "selected", + account: designated.account, + reason: "designated", + }; + } + } + return { + kind: "refused", + refusal: { + kind: "designated-dry", + tried: designatedIds.map(nameOf).join(", "), + }, + }; + } + + const candidates = accounts.flatMap((account) => { + const candidate = usable(account.id); + return candidate ? [candidate] : []; + }); + const personal = candidates.filter( + ({ account }) => account.access.kind === "personal", + ); + const pool = personal.length + ? personal + : candidates.filter(({ account }) => account.access.kind === "shared"); + if (!pool.length) { + return { + kind: "refused", + refusal: accounts.length + ? { kind: "pool-dry" } + : { kind: "none-configured" }, + }; + } + + const account = this.#options.select(pool, { + affinityKey: request.affinityKey, + lastPickedAt: (accountId) => this.#lastPickedAt.get(accountId) ?? 0, + }); + if (request.recordPick ?? true) { + this.#lastPickedAt.set(account.id, this.#options.clock?.() ?? Date.now()); + } + return { + kind: "selected", + account, + reason: account.access.kind === "personal" ? "personal" : "pool", + }; + } + + forget(accountId: string): void { + this.#lastPickedAt.delete(accountId); + } + + #canUse(account: TAccount, principalId: string | undefined): boolean { + if (account.access.kind === "shared") return true; + if (!principalId) return false; + return ( + this.#options.ownerMatches?.(principalId, account.access.owner) ?? + principalId === account.access.owner + ); + } +} diff --git a/packages/core/subscription-gateway/src/pool/selectors.ts b/packages/core/subscription-gateway/src/pool/selectors.ts new file mode 100644 index 0000000000..031df4b456 --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/selectors.ts @@ -0,0 +1,60 @@ +import type { + AccountSelector, + RoutableAccount, + SelectionCandidate, +} from "./types"; + +function compareByPriorityAndLastPick( + left: SelectionCandidate, + right: SelectionCandidate, + lastPickedAt: (accountId: string) => number, +): number { + return ( + left.priority - right.priority || + lastPickedAt(left.account.id) - lastPickedAt(right.account.id) + ); +} + +export function leastRecentlyPicked< + TAccount extends RoutableAccount, +>(): AccountSelector { + return (candidates, context) => { + const picked = candidates.toSorted((left, right) => + compareByPriorityAndLastPick(left, right, context.lastPickedAt), + )[0]; + if (!picked) throw new Error("Account selector received no candidates"); + return picked.account; + }; +} + +/** + * FNV-1a rendezvous score shared with Open Session's current Codex picker. + * Changing it moves existing affinity keys to different accounts. + */ +export function hrwScore(affinityKey: string, accountId: string): number { + let hash = 0x811c9dc5; + const value = `${affinityKey}\0${accountId}`; + for (let index = 0; index < value.length; index += 1) { + hash ^= value.charCodeAt(index); + hash = Math.imul(hash, 0x01000193) >>> 0; + } + return hash >>> 0; +} + +export function rendezvousAffinity< + TAccount extends RoutableAccount, +>(): AccountSelector { + const lru = leastRecentlyPicked(); + return (candidates, context) => { + if (!context.affinityKey) return lru(candidates, context); + const picked = candidates.toSorted( + (left, right) => + left.priority - right.priority || + hrwScore(context.affinityKey ?? "", right.account.id) - + hrwScore(context.affinityKey ?? "", left.account.id) || + left.account.id.localeCompare(right.account.id), + )[0]; + if (!picked) throw new Error("Account selector received no candidates"); + return picked.account; + }; +} diff --git a/packages/core/subscription-gateway/src/pool/types.ts b/packages/core/subscription-gateway/src/pool/types.ts new file mode 100644 index 0000000000..decc9e07e6 --- /dev/null +++ b/packages/core/subscription-gateway/src/pool/types.ts @@ -0,0 +1,71 @@ +export type AccountAccess = + | { readonly kind: "shared" } + | { readonly kind: "personal"; readonly owner: string }; + +export interface RoutableAccount { + readonly id: string; + readonly name: string; + readonly access: AccountAccess; +} + +export type AccountAvailability = + | { readonly kind: "available"; readonly priority: number } + | { readonly kind: "unavailable" }; + +export type PickReason = + | "pinned" + | "sticky" + | "designated" + | "personal" + | "pool"; + +export type AccountRefusal = + | { readonly kind: "none-configured" } + | { readonly kind: "pool-dry" } + | { + readonly kind: "pin-unusable"; + readonly pinnedId: string; + readonly pinName: string; + } + | { + readonly kind: "pin-not-designated"; + readonly pinnedId: string; + readonly pinName: string; + } + | { readonly kind: "designated-dry"; readonly tried: string }; + +export type AccountResolution = + | { + readonly kind: "selected"; + readonly account: TAccount; + readonly reason: PickReason; + } + | { readonly kind: "refused"; readonly refusal: AccountRefusal }; + +export interface AccountRequest { + readonly principalId?: string; + readonly model?: TModel; + readonly pinnedId?: string; + readonly strictPin?: boolean; + readonly stickyId?: string; + readonly designatedIds?: readonly string[]; + readonly affinityKey?: string; + readonly excludeIds?: ReadonlySet; + readonly allowPaidUsage?: boolean; + readonly recordPick?: boolean; +} + +export interface SelectionCandidate { + readonly account: TAccount; + readonly priority: number; +} + +export interface SelectionContext { + readonly affinityKey?: string; + readonly lastPickedAt: (accountId: string) => number; +} + +export type AccountSelector = ( + candidates: readonly SelectionCandidate[], + context: SelectionContext, +) => TAccount; diff --git a/packages/core/subscription-gateway/src/storage/cooldown-repository.ts b/packages/core/subscription-gateway/src/storage/cooldown-repository.ts new file mode 100644 index 0000000000..b8b09cc53e --- /dev/null +++ b/packages/core/subscription-gateway/src/storage/cooldown-repository.ts @@ -0,0 +1,104 @@ +import { mkdir, open, readFile, rename, rm } from "node:fs/promises"; +import { dirname } from "node:path"; + +export type CooldownRecord = + | { + readonly scope: "account"; + readonly accountId: string; + readonly reason: "exhausted" | "wedged"; + readonly until: number; + } + | { + readonly scope: "model"; + readonly accountId: string; + readonly model: string; + readonly reason: "exhausted"; + readonly until: number; + }; + +export interface CooldownRepository { + load(): Promise; + save(records: readonly CooldownRecord[]): Promise; +} + +export class MemoryCooldownRepository implements CooldownRepository { + #records: CooldownRecord[]; + + constructor(records: readonly CooldownRecord[] = []) { + this.#records = [...records]; + } + + async load(): Promise { + return [...this.#records]; + } + + async save(records: readonly CooldownRecord[]): Promise { + this.#records = [...records]; + } +} + +function isCooldownRecord(value: unknown): value is CooldownRecord { + if (!value || typeof value !== "object") return false; + if (!("scope" in value) || !("accountId" in value)) return false; + if (!("reason" in value) || !("until" in value)) return false; + if (typeof value.accountId !== "string" || !value.accountId) return false; + if (typeof value.until !== "number" || !Number.isFinite(value.until)) { + return false; + } + if (value.scope === "account") { + return value.reason === "exhausted" || value.reason === "wedged"; + } + return ( + value.scope === "model" && + value.reason === "exhausted" && + "model" in value && + typeof value.model === "string" && + !!value.model + ); +} + +export class JsonCooldownRepository implements CooldownRepository { + readonly #path: string; + + constructor(path: string) { + this.#path = path; + } + + async load(): Promise { + let text: string; + try { + text = await readFile(this.#path, "utf8"); + } catch (error) { + if ( + error instanceof Error && + "code" in error && + error.code === "ENOENT" + ) { + return []; + } + throw error; + } + const parsed: unknown = JSON.parse(text); + if (!Array.isArray(parsed) || !parsed.every(isCooldownRecord)) { + throw new Error(`Invalid cooldown store at ${this.#path}`); + } + return parsed; + } + + async save(records: readonly CooldownRecord[]): Promise { + const directory = dirname(this.#path); + await mkdir(directory, { recursive: true, mode: 0o700 }); + const temporaryPath = `${this.#path}.${crypto.randomUUID()}.tmp`; + const file = await open(temporaryPath, "wx", 0o600); + try { + await file.writeFile(`${JSON.stringify(records, null, 2)}\n`, "utf8"); + await file.sync(); + await file.close(); + await rename(temporaryPath, this.#path); + } catch (error) { + await file.close().catch(() => undefined); + await rm(temporaryPath, { force: true }).catch(() => undefined); + throw error; + } + } +} diff --git a/packages/core/subscription-gateway/src/storage/index.ts b/packages/core/subscription-gateway/src/storage/index.ts new file mode 100644 index 0000000000..1cc9c95a8d --- /dev/null +++ b/packages/core/subscription-gateway/src/storage/index.ts @@ -0,0 +1 @@ +export * from "./cooldown-repository";