diff --git a/cloudflare/packages/core/src/account-store.ts b/cloudflare/packages/core/src/account-store.ts new file mode 100644 index 00000000..a1f5b521 --- /dev/null +++ b/cloudflare/packages/core/src/account-store.ts @@ -0,0 +1,162 @@ +import { Context, Effect, Layer } from "effect" + +/** + * SubrouterActor — multi-tenant rewrite of the Go subrouter, hosted on Rivet + * actors. v1 = functional parity with Go subrouter (sticky session → account + * routing). Aurora-backed account store, not local JSON. + * + * This file is the *interface and sticky-routing core*. The actual Rivet + * actor wrapper lives in packages/actors and forwards into here. The HTTP + * transport (the thing that *replaces* what Go subrouter listens for at + * 0.0.0.0:31415) lives in packages/proxy/src/gateway.ts since the AI + * integrations go through the same proxy as gh/registries/etc. + * + * Why: Go subrouter today requires a long-lived VM (Mac Mini or systemd box). + * Rivet actors let us host it with zero VM ops, multi-tenant from the start. + */ + +export interface Account { + readonly id: string + readonly orgId: string + readonly kind: "codex_oauth" | "anthropic_oauth" | "openai_apikey" | "anthropic_apikey" + readonly label: string + readonly enabled: boolean + readonly rateLimitRemaining?: number + readonly modelQuotas?: AccountModelQuotas + readonly lastUsedAt?: number +} + +export interface AccountModelQuota { + readonly remainingPercent: number + readonly resetsAt?: number + readonly protectedBelowPercent?: number +} + +export type AccountModelQuotas = Readonly> + +// Named model-family quota pools. A request whose model name contains one of +// these keywords draws from that pool instead of the account-wide "default": +// Codex Spark ("spark") and Claude Opus/Sonnet weekly limits ("opus"/"sonnet"). +// Keep this in sync with what the quota populator writes into Account.modelQuotas: +// a model keyed to a pool no account carries is ineligible everywhere (see +// accountHasQuotaForModel), which is also what isolates providers (a Claude +// "opus" model only matches accounts that carry an "opus" pool). Anthropic's +// opus/sonnet caps are sub-limits of the account-wide window, so the populator +// must set each to min(account-wide remaining, family remaining); the Codex +// Spark pool is independent of the account-wide window. +const MODEL_QUOTA_POOLS = ["spark", "opus", "sonnet"] as const + +export const quotaKeyForModel = (model: string | undefined): string => { + const normalized = model?.trim().toLowerCase() + if (!normalized) return "default" + for (const pool of MODEL_QUOTA_POOLS) { + if (normalized.includes(pool)) return pool + } + return "default" +} + +export const accountHasQuotaForModel = ( + account: Account, + quotaKey: string +): boolean => { + const quota = account.modelQuotas?.[quotaKey] + if (!quota) return quotaKey === "default" + return quota.remainingPercent > 0 +} + +export interface AccountStore { + readonly list: (orgId: string) => Effect.Effect> + readonly pick: (input: { + readonly orgId: string + readonly sessionId: string + readonly preferAccountId?: string + readonly model?: string + readonly quotaKey?: string + }) => Effect.Effect + readonly recordUse: (accountId: string) => Effect.Effect +} + +export class AccountStoreTag extends Context.Tag("AccountStore")< + AccountStoreTag, + AccountStore +>() {} + +/** + * Sticky session table. In production: a Postgres row per (orgId, sessionId). + * Here: in-memory Map. Stickiness ensures cached agent context stays useful + * (matches Go subrouter's `X-Subrouter-Session` semantics). + */ +export const makeInMemoryStickyStore = () => { + const map = new Map() // `${orgId}:${sessionId}` -> accountId + return { + get: (orgId: string, sessionId: string, quotaKey = "default"): string | null => + map.get(`${orgId}:${quotaKey}:${sessionId}`) ?? null, + set: ( + orgId: string, + sessionId: string, + accountId: string, + quotaKey = "default" + ): void => { + map.set(`${orgId}:${quotaKey}:${sessionId}`, accountId) + }, + clear: (): void => { + map.clear() + }, + } +} + +/** + * Reference in-memory AccountStore impl that round-robins enabled accounts and + * respects sticky sessions. Production swaps in a Postgres-backed implementation + * that reads from `accounts` and `subrouter_session_assignments`. + */ +export const makeInMemoryAccountStoreLayer = (initial: ReadonlyArray) => { + const accounts = [...initial] + const sticky = makeInMemoryStickyStore() + let cursor = 0 + + return Layer.succeed(AccountStoreTag, { + list: (orgId) => + Effect.succeed(accounts.filter((a) => a.orgId === orgId && a.enabled)), + + pick: ({ orgId, sessionId, preferAccountId, model, quotaKey }) => { + const resolvedQuotaKey = quotaKey ?? quotaKeyForModel(model) + const isEligible = (account: Account): boolean => + account.orgId === orgId && + account.enabled && + accountHasQuotaForModel(account, resolvedQuotaKey) + + // 1. sticky table first + const stickyId = sticky.get(orgId, sessionId, resolvedQuotaKey) + if (stickyId) { + const found = accounts.find((a) => a.id === stickyId && isEligible(a)) + if (found) return Effect.succeed(found) + } + // 2. explicit pin + if (preferAccountId) { + const found = accounts.find( + (a) => a.id === preferAccountId && isEligible(a) + ) + if (found) { + sticky.set(orgId, sessionId, found.id, resolvedQuotaKey) + return Effect.succeed(found) + } + } + // 3. round-robin over enabled accounts for the org + const eligible = accounts.filter(isEligible) + if (eligible.length === 0) return Effect.succeed(null) + const pick = eligible[cursor % eligible.length]! + cursor += 1 + sticky.set(orgId, sessionId, pick.id, resolvedQuotaKey) + return Effect.succeed(pick) + }, + + recordUse: (accountId) => + Effect.sync(() => { + const i = accounts.findIndex((a) => a.id === accountId) + if (i >= 0) { + accounts[i] = { ...accounts[i]!, lastUsedAt: Date.now() } + } + }), + } satisfies AccountStore) +} diff --git a/cloudflare/packages/core/src/index.ts b/cloudflare/packages/core/src/index.ts index 138bf977..4b7c8369 100644 --- a/cloudflare/packages/core/src/index.ts +++ b/cloudflare/packages/core/src/index.ts @@ -1,163 +1,19 @@ -import { Context, Effect, Layer } from "effect" -export * from "./service.ts" - -/** - * SubrouterActor — multi-tenant rewrite of the Go subrouter, hosted on Rivet - * actors. v1 = functional parity with Go subrouter (sticky session → account - * routing). Aurora-backed account store, not local JSON. - * - * This file is the *interface and sticky-routing core*. The actual Rivet - * actor wrapper lives in packages/actors and forwards into here. The HTTP - * transport (the thing that *replaces* what Go subrouter listens for at - * 0.0.0.0:31415) lives in packages/proxy/src/gateway.ts since the AI - * integrations go through the same proxy as gh/registries/etc. - * - * Why: Go subrouter today requires a long-lived VM (Mac Mini or systemd box). - * Rivet actors let us host it with zero VM ops, multi-tenant from the start. - */ - -export interface Account { - readonly id: string - readonly orgId: string - readonly kind: "codex_oauth" | "anthropic_oauth" | "openai_apikey" | "anthropic_apikey" - readonly label: string - readonly enabled: boolean - readonly rateLimitRemaining?: number - readonly modelQuotas?: AccountModelQuotas - readonly lastUsedAt?: number -} - -export interface AccountModelQuota { - readonly remainingPercent: number - readonly resetsAt?: number - readonly protectedBelowPercent?: number -} - -export type AccountModelQuotas = Readonly> - -// Named model-family quota pools. A request whose model name contains one of -// these keywords draws from that pool instead of the account-wide "default": -// Codex Spark ("spark") and Claude Opus/Sonnet weekly limits ("opus"/"sonnet"). -// Keep this in sync with what the quota populator writes into Account.modelQuotas: -// a model keyed to a pool no account carries is ineligible everywhere (see -// accountHasQuotaForModel), which is also what isolates providers (a Claude -// "opus" model only matches accounts that carry an "opus" pool). Anthropic's -// opus/sonnet caps are sub-limits of the account-wide window, so the populator -// must set each to min(account-wide remaining, family remaining); the Codex -// Spark pool is independent of the account-wide window. -const MODEL_QUOTA_POOLS = ["spark", "opus", "sonnet"] as const - -export const quotaKeyForModel = (model: string | undefined): string => { - const normalized = model?.trim().toLowerCase() - if (!normalized) return "default" - for (const pool of MODEL_QUOTA_POOLS) { - if (normalized.includes(pool)) return pool - } - return "default" -} - -export const accountHasQuotaForModel = ( - account: Account, - quotaKey: string -): boolean => { - const quota = account.modelQuotas?.[quotaKey] - if (!quota) return quotaKey === "default" - return quota.remainingPercent > 0 -} - -export interface AccountStore { - readonly list: (orgId: string) => Effect.Effect> - readonly pick: (input: { - readonly orgId: string - readonly sessionId: string - readonly preferAccountId?: string - readonly model?: string - readonly quotaKey?: string - }) => Effect.Effect - readonly recordUse: (accountId: string) => Effect.Effect -} - -export class AccountStoreTag extends Context.Tag("AccountStore")< +export { AccountStoreTag, - AccountStore ->() {} - -/** - * Sticky session table. In production: a Postgres row per (orgId, sessionId). - * Here: in-memory Map. Stickiness ensures cached agent context stays useful - * (matches Go subrouter's `X-Subrouter-Session` semantics). - */ -export const makeInMemoryStickyStore = () => { - const map = new Map() // `${orgId}:${sessionId}` -> accountId - return { - get: (orgId: string, sessionId: string, quotaKey = "default"): string | null => - map.get(`${orgId}:${quotaKey}:${sessionId}`) ?? null, - set: ( - orgId: string, - sessionId: string, - accountId: string, - quotaKey = "default" - ): void => { - map.set(`${orgId}:${quotaKey}:${sessionId}`, accountId) - }, - clear: (): void => { - map.clear() - }, - } -} - -/** - * Reference in-memory AccountStore impl that round-robins enabled accounts and - * respects sticky sessions. Production swaps in a Postgres-backed implementation - * that reads from `accounts` and `subrouter_session_assignments`. - */ -export const makeInMemoryAccountStoreLayer = (initial: ReadonlyArray) => { - const accounts = [...initial] - const sticky = makeInMemoryStickyStore() - let cursor = 0 - - return Layer.succeed(AccountStoreTag, { - list: (orgId) => - Effect.succeed(accounts.filter((a) => a.orgId === orgId && a.enabled)), - - pick: ({ orgId, sessionId, preferAccountId, model, quotaKey }) => { - const resolvedQuotaKey = quotaKey ?? quotaKeyForModel(model) - const isEligible = (account: Account): boolean => - account.orgId === orgId && - account.enabled && - accountHasQuotaForModel(account, resolvedQuotaKey) - - // 1. sticky table first - const stickyId = sticky.get(orgId, sessionId, resolvedQuotaKey) - if (stickyId) { - const found = accounts.find((a) => a.id === stickyId && isEligible(a)) - if (found) return Effect.succeed(found) - } - // 2. explicit pin - if (preferAccountId) { - const found = accounts.find( - (a) => a.id === preferAccountId && isEligible(a) - ) - if (found) { - sticky.set(orgId, sessionId, found.id, resolvedQuotaKey) - return Effect.succeed(found) - } - } - // 3. round-robin over enabled accounts for the org - const eligible = accounts.filter(isEligible) - if (eligible.length === 0) return Effect.succeed(null) - const pick = eligible[cursor % eligible.length]! - cursor += 1 - sticky.set(orgId, sessionId, pick.id, resolvedQuotaKey) - return Effect.succeed(pick) - }, - - recordUse: (accountId) => - Effect.sync(() => { - const i = accounts.findIndex((a) => a.id === accountId) - if (i >= 0) { - accounts[i] = { ...accounts[i]!, lastUsedAt: Date.now() } - } - }), - } satisfies AccountStore) -} + accountHasQuotaForModel, + makeInMemoryAccountStoreLayer, + makeInMemoryStickyStore, + quotaKeyForModel, +} from "./account-store.ts" +export type { + Account, + AccountModelQuota, + AccountModelQuotas, + AccountStore, +} from "./account-store.ts" +export { + NoEligibleAccount, + SubrouterService, + makeSubrouterServiceLayer, +} from "./service.ts" +export type { PickedRoute } from "./service.ts" diff --git a/cloudflare/packages/core/src/service.ts b/cloudflare/packages/core/src/service.ts index 4f2ae220..b3fe0e1f 100644 --- a/cloudflare/packages/core/src/service.ts +++ b/cloudflare/packages/core/src/service.ts @@ -14,7 +14,7 @@ import { Context, Effect, Layer } from "effect" import { AccountStoreTag, type Account, -} from "./index.ts" +} from "./account-store.ts" export class NoEligibleAccount extends Error { readonly _tag = "NoEligibleAccount" diff --git a/cloudflare/packages/core/test/public-api.test.ts b/cloudflare/packages/core/test/public-api.test.ts new file mode 100644 index 00000000..916855eb --- /dev/null +++ b/cloudflare/packages/core/test/public-api.test.ts @@ -0,0 +1,29 @@ +import { expect, test } from "bun:test" +import { Effect, Layer } from "effect" +import * as core from "../src/index.ts" +import * as service from "../src/service.ts" + +test("the core barrel exposes one coherent service and account-store API", async () => { + expect(core.SubrouterService).toBe(service.SubrouterService) + expect(core.NoEligibleAccount).toBe(service.NoEligibleAccount) + expect(typeof core.makeInMemoryAccountStoreLayer).toBe("function") + + const layer = core.makeInMemoryAccountStoreLayer([ + { + id: "account-1", + orgId: "org-1", + kind: "codex_oauth", + label: "primary", + enabled: true, + }, + ]) + const program = Effect.gen(function* () { + const subrouter = yield* core.SubrouterService + return yield* subrouter.route({ orgId: "org-1", sessionId: "session-1" }) + }).pipe(Effect.provide(core.makeSubrouterServiceLayer().pipe( + Layer.provideMerge(layer) + ))) + + const result = await Effect.runPromise(program) + expect(result.account.id).toBe("account-1") +}) diff --git a/cloudflare/packages/worker/src/core-routing.ts b/cloudflare/packages/worker/src/core-routing.ts new file mode 100644 index 00000000..58ebfe4f --- /dev/null +++ b/cloudflare/packages/worker/src/core-routing.ts @@ -0,0 +1,4 @@ +export { + accountHasQuotaForModel, + quotaKeyForModel, +} from "@subrouter/core" diff --git a/cloudflare/packages/worker/src/index.ts b/cloudflare/packages/worker/src/index.ts index c7148d4c..7bcae948 100644 --- a/cloudflare/packages/worker/src/index.ts +++ b/cloudflare/packages/worker/src/index.ts @@ -1,10 +1,6 @@ import { DurableObject } from "cloudflare:workers" -import { - accountHasQuotaForModel, - quotaKeyForModel, - type Account, - type AccountModelQuotas, -} from "@subrouter/core" +import type { Account, AccountModelQuotas } from "@subrouter/core" +import { accountHasQuotaForModel, quotaKeyForModel } from "./core-routing.ts" import { authModeForAccount, blockingRefreshFailure, diff --git a/cloudflare/packages/worker/test/proxy-streaming.test.ts b/cloudflare/packages/worker/test/proxy-streaming.test.ts index 7b5684c5..e84129a3 100644 --- a/cloudflare/packages/worker/test/proxy-streaming.test.ts +++ b/cloudflare/packages/worker/test/proxy-streaming.test.ts @@ -2,7 +2,7 @@ import { afterEach, describe, expect, mock, test } from "bun:test" import { Buffer } from "node:buffer" mock.module("cloudflare:workers", () => ({ DurableObject: class {} })) -mock.module("@subrouter/core", () => ({ +mock.module("../src/core-routing.ts", () => ({ accountHasQuotaForModel: () => true, quotaKeyForModel: (model: string | undefined) => model ?? "default", }))