Compare commits

...
Author SHA1 Message Date
Lawrence Chen 3fad5abf1d Merge pull request #10424 from manaflow-ai/feat/coderouter-status-session-counts
Expose per-account session counts and cooldowns in account listing
2026-08-19 04:13:14 -07:00
Lawrence Chen 34f0e1eeca Expose per-account session counts and cooldowns in account listing
The status endpoint now reports, per account, how many sessions are
bound with traffic inside the routing load window, plus the account's
cooldown deadline. The coderouter CLI uses both to make sticky routing
and rate-limit cooldowns visible. The session-count read is display
only and fails to zero so it can never take the status endpoint down.
2026-08-19 03:59:50 -07:00
Lawrence Chen 882ab10e75 Merge pull request #10402 from manaflow-ai/feat/coderouter-free-tier
Make hosted coderouter free for up to 3 accounts
2026-08-18 23:44:59 -07:00
Lawrence Chen 9ad361ed37 Make hosted coderouter free for up to 3 accounts
Pricing change: a team may connect and route up to 3 provider
accounts (subscriptions) for free. More than 3 requires an active
cmux Pro or Team subscription.

- New services/coderouter/entitlement.ts: free tier checked first
  (one indexed count read); the Stripe read runs only for teams over
  the limit. Both checks fail closed.
- Session issuance: replaces the flat Pro gate. Over-limit teams
  without a subscription get 402 with the count in the message.
- Account add: connecting an account beyond the limit without a
  subscription returns 402 before anything is stored. Re-importing
  an account the team already has is always allowed, so a broken
  account stays repairable on the free tier. The accounts POST
  handler is now factory-built for dependency-injected tests.
- Billing-lapse token revocation stays unchanged: after a lapse the
  CLI renews and the free tier re-qualifies teams with <= 3 accounts.
- New analytics: entitlement_basis on route_session_issued and a
  coderouter_account_limit_reached event, both schema-whitelisted.
- CODEROUTER_HOSTED_PRO_REQUIRED=0 still disables all gating.

Boundary: exactly 3 accounts is free; the 4th needs Pro/Team.
2026-08-18 23:21:33 -07:00
Lawrence Chen ae7d761e03 Merge pull request #10387 from manaflow-ai/fix/coderouter-codex-session-spread
Spread Codex placements and pin sessions to accounts
2026-08-18 21:57:20 -07:00
Austin Wang c15d665f65 Merge pull request #10363 from manaflow-ai/issue-10103-copy-mode-selection-bleed
Fix terminal Copy Mode selection scope across split panes
2026-08-18 21:46:09 -07:00
Lawrence Chen 425f2d7efe Keep sticky sessions through in-flight credential refreshes
A bound account in state 'refreshing' is healthy; its refresh lease
resolves in seconds. Treat it as usable for sticky reuse, and give a
sticky session a short bounded wait (4 x 500ms) when the refresh lease
is busy, instead of instantly moving the session and discarding its
prompt cache. Non-sticky requests keep the fail-fast move.
2026-08-18 21:43:42 -07:00
Lawrence Chen 7d79a858f6 Degrade to legacy routing while the session table migration is pending
If the code deploys before the additive migration is applied, the
session-stickiness statements hit undefined_table (42P01). Treat that
as: no binding found, claim without the session-load ordering term, and
skip the pin. Routing behaves exactly like the pre-change code until
the migration lands, instead of erroring every Codex request. Covered
by a DB behavior test that renames the table away and back.
2026-08-18 21:38:20 -07:00
Lawrence Chen 18ed976a1b Spread Codex placements and pin sessions to accounts
Port of subrouter PR #228 to the coderouter TypeScript data plane.

Problems in the current routing:
- selectAccountForRequest picks an account per request with three
  independent statements (sweep, select, update). Nothing spans the
  read-pick-write sequence, so parallel session starts read the same
  snapshot and herd onto one account. That account burns quota first,
  then every session reroutes to the next account at once.
- No session stickiness exists. Each request rotates to the
  least-recently-used account, so the provider prompt cache is
  discarded on nearly every turn and the whole prefix is re-billed
  as uncached input across all accounts.

Fix:
- New coderouter_session_accounts table pins one agent session (the
  Codex CLI session_id header) to one account.
- selectAccountForSession honors a usable binding first (sticky), and
  moves a session only when its account is broken, cooling down,
  removed, or already attempted in this request.
- New placements claim an account atomically in one statement with
  FOR UPDATE SKIP LOCKED, ordered by fewest recently-active bound
  sessions, then least-recently-used. Concurrent claims take
  different accounts instead of the same snapshot argmax. When every
  candidate row is locked, a blocking fallback claim accepts a
  collision instead of failing the request.
- selectAccountForRequest (models, opencode) now uses the same atomic
  claim, which closes its read-pick-write race too.

Tests: unit tests for the selector and the responses proxy, and
CMUX_DB_TEST-gated behavior tests that prove sequential spread,
concurrent non-herding, stickiness, move-on-cooldown, exclusion,
and last-write-wins binding against real Postgres.
2026-08-18 21:34:14 -07:00
austinpower1258 6e56e6ff14 fix: cancel copy mode when pane deactivates 2026-08-18 14:16:33 -07:00
austinpower1258 b0948c0341 fix: cancel copy mode on responder loss 2026-08-18 14:11:59 -07:00
austinpower1258 d3e3abd367 fix: scope terminal copy mode to focused pane 2026-08-18 14:01:01 -07:00
austinpower1258 9f556f118e test: require copy mode to end on focus loss 2026-08-18 14:00:14 -07:00
20 changed files with 1588 additions and 89 deletions
@@ -27,6 +27,9 @@ public protocol TerminalSurfaceNativeViewing: NSView, TerminalSurfaceHosting {
@discardableResult
func toggleKeyboardCopyMode() -> Bool
/// Ends keyboard copy mode and clears its selection when focus leaves the view.
func cancelKeyboardCopyMode()
/// Re-applies the window background for the active surface.
func applyWindowBackgroundIfActive()
@@ -21,6 +21,9 @@ extension TerminalSurface {
/// Applies a focus state to the runtime surface (deduplicated).
@MainActor
public func setFocus(_ focused: Bool, force: Bool = false) {
if !focused {
surfaceView.cancelKeyboardCopyMode()
}
// Only send focus events when the state changes to avoid redundant
// prompt redraws with zsh themes like Powerlevel10k.
guard force || focused != desiredFocusState else { return }
@@ -9,7 +9,8 @@ final class FakeTerminalSurfaceNativeView: NSView {
weak var attachedController: (any TerminalSurfaceControlling)?
var attachedSurfaceController: (any TerminalSurfaceControlling)? { attachedController }
var currentKeyStateIndicatorText: String? { nil }
var isKeyboardCopyModeActive: Bool { false }
var isKeyboardCopyModeActive = false
private(set) var keyboardCopyModeCancellationCount = 0
var shouldDeferRuntimeInput = false
var runtimeInputDeferralResponses: [Bool] = []
var runtimeInputDeferralCallCount = 0
@@ -18,6 +19,10 @@ final class FakeTerminalSurfaceNativeView: NSView {
var mobileMouseButtonEvents: [String] = []
func toggleKeyboardCopyMode() -> Bool { false }
func cancelKeyboardCopyMode() {
keyboardCopyModeCancellationCount += 1
isKeyboardCopyModeActive = false
}
func applyWindowBackgroundIfActive() {}
func forceRefreshSurface() -> Bool { true }
func runtimeSurfaceDidBecomeReady() {}
@@ -184,6 +184,18 @@ struct TerminalSurfaceExplicitInputTests {
#expect(fixture.paneHost.explicitInputCount == 1)
}
@Test func losingFocusCancelsKeyboardCopyModeOnTheSurface() {
let fixture = makeFixture()
defer { fixture.surface.releaseSurfaceForTesting() }
fixture.nativeView.isKeyboardCopyModeActive = true
fixture.surface.setFocus(true)
fixture.surface.setFocus(false)
#expect(fixture.nativeView.keyboardCopyModeCancellationCount == 1)
#expect(!fixture.nativeView.isKeyboardCopyModeActive)
}
@Test func mobileGesturesNotifyPaneHost() {
let fixture = makeFixture()
defer { fixture.surface.releaseSurfaceForTesting() }
+12
View File
@@ -4521,6 +4521,14 @@ class GhosttyNSView: NSView, NSUserInterfaceValidations {
return true
}
func cancelKeyboardCopyMode() {
guard keyboardCopyModeActive else { return }
if let surface {
_ = GhosttyRuntimeCInterop.clearSelection(surface)
}
setKeyboardCopyModeActive(false)
}
private func setKeyboardCopyModeActive(_ active: Bool) {
keyboardCopyModeInputState.reset()
keyboardCopyModeSelectionKind = nil
@@ -5341,6 +5349,7 @@ class GhosttyNSView: NSView, NSUserInterfaceValidations {
imeConsumedKeyUps.removeAll()
manualNamedKeyConsumedKeyUps.removeAll()
desiredFocus = false
cancelKeyboardCopyMode()
terminalSurface?.hostedView.cancelSuppressedFirstResponderFocusReapply()
terminalSurface?.recordExternalFocusState(false)
}
@@ -10105,6 +10114,9 @@ final class GhosttySurfaceScrollView: NSView {
func setActive(_ active: Bool) {
let wasActive = isActive
if !active {
surfaceView.cancelKeyboardCopyMode()
}
isActive = active
#if DEBUG
if wasActive != active {
+81 -3
View File
@@ -1,7 +1,12 @@
import { env } from "../../../env";
import {
addAccount,
parseCredential,
} from "../../../../services/coderouter/accounts";
import {
CODEROUTER_FREE_ACCOUNT_LIMIT,
accountAdditionAllowed,
} from "../../../../services/coderouter/entitlement";
import {
resolveCoderouterUsageTeam,
resolveCodeRouterRequestContext,
@@ -70,8 +75,27 @@ export async function GET(request: Request): Promise<Response> {
});
}
export async function POST(request: Request): Promise<Response> {
const resolved = await resolveCodeRouterRequestContext(request, "manage");
type AccountsPostDependencies = {
readonly resolveContext: typeof resolveCodeRouterRequestContext;
readonly additionAllowed: typeof accountAdditionAllowed;
readonly add: typeof addAccount;
readonly hostedProRequired: () => boolean;
};
const defaultAccountsPostDependencies: AccountsPostDependencies = {
resolveContext: resolveCodeRouterRequestContext,
additionAllowed: accountAdditionAllowed,
add: addAccount,
hostedProRequired: () => env.CODEROUTER_HOSTED_PRO_REQUIRED === "1",
};
export const POST = makeCoderouterAccountsPostHandler();
export function makeCoderouterAccountsPostHandler(
dependencies: AccountsPostDependencies = defaultAccountsPostDependencies,
) {
return async function POST(request: Request): Promise<Response> {
const resolved = await dependencies.resolveContext(request, "manage");
if (!resolved.ok) return resolved.response;
const length = Number(request.headers.get("content-length") ?? "0");
if (Number.isFinite(length) && length > MAX_BODY_BYTES) {
@@ -91,8 +115,61 @@ export async function POST(request: Request): Promise<Response> {
if (!credential) {
return Response.json({ error: "invalid_request" }, { status: 400 });
}
if (dependencies.hostedProRequired()) {
let decision;
try {
decision = await dependencies.additionAllowed({
stackUserId: resolved.value.user.id,
teamId: resolved.value.team.teamId,
provider: credential.provider,
providerAccountId: credential.accountId,
});
} catch (error) {
reportCoderouterFailure("rds", error, {
operation: "account_addition_gate",
});
return Response.json(
{
error: "entitlement_unavailable",
message:
"coderouter could not verify your plan. Nothing was changed; retry shortly.",
retryable: true,
},
{
status: 503,
headers: { "cache-control": "no-store", "retry-after": "5" },
},
);
}
if (!decision.allowed) {
captureCoderouterEvent({
event: "coderouter_account_limit_reached",
userId: resolved.value.user.id,
teamId: resolved.value.team.teamId,
properties: {
provider: credential.provider,
account_count: decision.accountCount,
free_limit: CODEROUTER_FREE_ACCOUNT_LIMIT,
},
});
return Response.json(
{
error: "pro_required",
message:
`Free hosted coderouter covers up to ${CODEROUTER_FREE_ACCOUNT_LIMIT} connected accounts; ` +
`this team already has ${decision.accountCount}. ` +
"Upgrade to cmux Pro or Team to connect more, or remove an account first.",
retryable: false,
},
{
status: 402,
headers: { "cache-control": "no-store" },
},
);
}
}
try {
const result = await addAccount(resolved.value.team.teamId, credential);
const result = await dependencies.add(resolved.value.team.teamId, credential);
captureCoderouterEvent({
event: "coderouter_account_added",
userId: resolved.value.user.id,
@@ -129,6 +206,7 @@ export async function POST(request: Request): Promise<Response> {
},
);
}
};
}
function timing(name: string, duration: number): string {
+20 -11
View File
@@ -1,5 +1,8 @@
import { env } from "../../../env";
import { hasActiveCoderouterSubscription } from "../../../../services/billing/pro";
import {
CODEROUTER_FREE_ACCOUNT_LIMIT,
coderouterEntitlement,
} from "../../../../services/coderouter/entitlement";
import {
authenticateRouteToken,
issueRouteToken,
@@ -16,14 +19,14 @@ import {
type SessionDependencies = {
readonly resolveContext: typeof resolveCodeRouterRequestContext;
readonly hasActiveEntitlement: typeof hasActiveCoderouterSubscription;
readonly entitlement: typeof coderouterEntitlement;
readonly issueToken: typeof issueRouteToken;
readonly hostedProRequired: () => boolean;
};
const defaultDependencies: SessionDependencies = {
resolveContext: resolveCodeRouterRequestContext,
hasActiveEntitlement: hasActiveCoderouterSubscription,
entitlement: coderouterEntitlement,
issueToken: issueRouteToken,
hostedProRequired: () => env.CODEROUTER_HOSTED_PRO_REQUIRED === "1",
};
@@ -77,19 +80,22 @@ export function makeCoderouterSessionPostHandler(
const resolved = await dependencies.resolveContext(request, "use");
if (!resolved.ok) return resolved.response;
const userId = resolved.value.user.id;
let entitlementBasis = "ungated";
if (dependencies.hostedProRequired()) {
try {
if (
!(await dependencies.hasActiveEntitlement(
userId,
resolved.value.team.teamId,
))
) {
const entitlement = await dependencies.entitlement(
userId,
resolved.value.team.teamId,
);
entitlementBasis = entitlement.basis;
if (!entitlement.allowed) {
return Response.json(
{
error: "pro_required",
message:
"Hosted coderouter requires cmux Pro or Team. Upgrade or connect a self-hosted server.",
`Free hosted coderouter covers up to ${CODEROUTER_FREE_ACCOUNT_LIMIT} connected accounts; ` +
`this team has ${entitlement.accountCount}. ` +
"Upgrade to cmux Pro or Team, remove accounts, or connect a self-hosted server.",
retryable: false,
},
{
@@ -144,7 +150,10 @@ export function makeCoderouterSessionPostHandler(
event: "coderouter_route_session_issued",
userId,
teamId: resolved.value.team.teamId,
properties: { hosted_pro_required: dependencies.hostedProRequired() },
properties: {
hosted_pro_required: dependencies.hostedProRequired(),
entitlement_basis: entitlementBasis,
},
});
addCoderouterBreadcrumb("session", "Route session issued");
return Response.json(
@@ -0,0 +1,19 @@
CREATE TABLE "coderouter_session_accounts" (
"team_id" text NOT NULL,
"provider" text NOT NULL,
"session_key" text NOT NULL,
"account_id" uuid NOT NULL
REFERENCES "coderouter_accounts" ("id") ON DELETE CASCADE,
"created_at" timestamp with time zone DEFAULT now() NOT NULL,
"last_seen_at" timestamp with time zone DEFAULT now() NOT NULL,
CONSTRAINT "coderouter_session_accounts_pkey"
PRIMARY KEY ("team_id", "provider", "session_key"),
CONSTRAINT "coderouter_session_accounts_provider_check"
CHECK ("provider" IN ('codex', 'opencode-go'))
);
CREATE INDEX "coderouter_session_accounts_account_idx"
ON "coderouter_session_accounts" ("account_id");
CREATE INDEX "coderouter_session_accounts_last_seen_idx"
ON "coderouter_session_accounts" ("last_seen_at");
+33
View File
@@ -8,6 +8,7 @@ import {
jsonb,
pgEnum,
pgTable,
primaryKey,
text,
timestamp,
uniqueIndex,
@@ -581,6 +582,38 @@ export const coderouterVaultLeases = pgTable(
],
);
/**
* Session -> account stickiness for coderouter routing.
*
* Providers cache prompt prefixes per account, so moving a live session to a
* different account re-bills its whole prompt prefix as uncached input. A row
* here pins one agent session (the Codex CLI `session_id` header) to one
* account. Placement of a new session spreads across the least-loaded usable
* accounts under FOR UPDATE SKIP LOCKED, so concurrent session starts cannot
* herd onto a single account (port of subrouter PR #228).
*/
export const coderouterSessionAccounts = pgTable(
"coderouter_session_accounts",
{
teamId: text("team_id").notNull(),
provider: text("provider").$type<"codex" | "opencode-go">().notNull(),
sessionKey: text("session_key").notNull(),
accountId: uuid("account_id")
.notNull()
.references(() => coderouterAccounts.id, { onDelete: "cascade" }),
createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
lastSeenAt: timestamp("last_seen_at", { withTimezone: true }).notNull().defaultNow(),
},
(table) => [
primaryKey({
name: "coderouter_session_accounts_pkey",
columns: [table.teamId, table.provider, table.sessionKey],
}),
index("coderouter_session_accounts_account_idx").on(table.accountId),
index("coderouter_session_accounts_last_seen_idx").on(table.lastSeenAt),
],
);
export const stripeCustomers = pgTable(
"stripe_customers",
{
+25 -4
View File
@@ -16,6 +16,7 @@ import {
export type CoderouterAnalyticsEvent =
| "coderouter_account_added"
| "coderouter_account_limit_reached"
| "coderouter_account_removed"
| "coderouter_account_status_viewed"
| "coderouter_auth_rejected"
@@ -170,6 +171,7 @@ async function deliver(
function eventNeedsUserScope(event: CoderouterAnalyticsEvent): boolean {
return event === "coderouter_account_added" ||
event === "coderouter_account_limit_reached" ||
event === "coderouter_account_removed" ||
event === "coderouter_route_session_issued" ||
event === "coderouter_route_session_revoked";
@@ -188,6 +190,15 @@ function eventProperties(
}
return { provider, source, already_exists: input.already_exists };
}
case "coderouter_account_limit_reached": {
const provider = accountProvider(input.provider);
if (!provider) return null;
return {
provider,
account_count_bucket: countBucket(input.account_count),
free_limit: typeof input.free_limit === "number" ? input.free_limit : 0,
};
}
case "coderouter_account_removed": {
const source = lifecycleSource(input.source);
if (!source) return null;
@@ -212,10 +223,20 @@ function eventProperties(
const reason = authReason(input.reason);
return surface && reason ? { surface, reason } : null;
}
case "coderouter_route_session_issued":
return typeof input.hosted_pro_required === "boolean"
? { hosted_pro_required: input.hosted_pro_required }
: null;
case "coderouter_route_session_issued": {
if (typeof input.hosted_pro_required !== "boolean") return null;
const output: Record<string, AnalyticsScalar> = {
hosted_pro_required: input.hosted_pro_required,
};
const basis = enumValue(input.entitlement_basis, [
"free_tier",
"subscription",
"pro_required",
"ungated",
]);
if (basis) output.entitlement_basis = basis;
return output;
}
case "coderouter_route_session_revoked":
return {};
case "coderouter_organization_catalog_viewed":
+85 -14
View File
@@ -2,6 +2,7 @@ import {
authenticateRouteToken,
markAccountCooldown,
selectAccountForRequest,
selectAccountForSession,
} from "./repository";
import { freshCredential } from "./refresh";
import { fetchProviderRead } from "./providerFetch";
@@ -24,7 +25,70 @@ const ALLOWED_REQUEST_HEADERS = [
"user-agent",
] as const;
export async function proxyCodexRequest(request: Request): Promise<Response> {
type CodexResponsesDependencies = {
readonly authenticate: typeof authenticateRouteToken;
readonly select: typeof selectAccountForSession;
readonly credential: typeof freshCredential;
readonly cooldown: typeof markAccountCooldown;
};
/**
* The Codex CLI sends a stable `session_id` header for every request of one
* agent session. That key pins the session to one account so the provider's
* prompt cache stays warm across turns.
*/
function sessionKeyFromRequest(request: Request): string | null {
const raw = request.headers.get("session_id")?.trim();
if (!raw || raw.length > 512) return null;
return raw;
}
const STICKY_REFRESH_RETRIES = 4;
const STICKY_REFRESH_RETRY_DELAY_MS = 500;
/**
* A sticky session that hits a refresh already in flight should wait for the
* winner's fresh credential rather than move to another account: a move
* discards the session's prompt cache and re-bills its whole prefix, while
* the in-flight refresh completes within seconds. Non-sticky requests keep
* the fail-fast behavior.
*/
async function credentialWithStickyPatience(
dependencies: Pick<CodexResponsesDependencies, "credential">,
input: { teamId: string; accountId: string; expectedRevision: number },
sticky: boolean,
): Promise<Awaited<ReturnType<CodexResponsesDependencies["credential"]>>> {
for (let attempt = 0; ; attempt++) {
try {
return await dependencies.credential(input);
} catch (error) {
const busy = error && typeof error === "object" && "_tag" in error &&
(error as { _tag: string })._tag === "CodeRouterRefreshBusy";
if (!busy || !sticky || attempt >= STICKY_REFRESH_RETRIES) throw error;
await new Promise((resolve) =>
setTimeout(resolve, STICKY_REFRESH_RETRY_DELAY_MS)
);
}
}
}
export function createCodexResponsesProxy(
dependencies: CodexResponsesDependencies,
): (request: Request) => Promise<Response> {
return async (request) => proxyCodexRequestWith(dependencies, request);
}
export const proxyCodexRequest = createCodexResponsesProxy({
authenticate: authenticateRouteToken,
select: selectAccountForSession,
credential: freshCredential,
cooldown: markAccountCooldown,
});
async function proxyCodexRequestWith(
dependencies: CodexResponsesDependencies,
request: Request,
): Promise<Response> {
const startedAt = performance.now();
const token = bearerToken(request);
if (!token) {
@@ -51,7 +115,7 @@ export async function proxyCodexRequest(request: Request): Promise<Response> {
false,
);
}
const identity = await authenticateRouteToken(token);
const identity = await dependencies.authenticate(token);
if (!identity) {
addCoderouterBreadcrumb("auth", "Route token rejected", {}, "warning");
captureCoderouterEvent({
@@ -85,30 +149,37 @@ export async function proxyCodexRequest(request: Request): Promise<Response> {
const value = request.headers.get(name);
if (value) forwardedHeaders.set(name, value);
}
const sessionKey = sessionKeyFromRequest(request);
const attempted: string[] = [];
let refreshRetries = 0;
let failureStage: "account_selection" | "credential_refresh" | "upstream_transport" =
"account_selection";
let upstream: Response | null = null;
for (let attempt = 0; attempt < 8; attempt++) {
const account = await selectAccountForRequest(
identity.teamId,
"codex",
attempted,
);
const account = await dependencies.select({
teamId: identity.teamId,
provider: "codex",
sessionKey,
excludedAccountIds: attempted,
});
if (!account) break;
attempted.push(account.id);
addCoderouterBreadcrumb("routing", "Selected provider account", {
provider: "codex",
attempt: attempt + 1,
sticky: account.sticky,
});
let credential;
try {
credential = await freshCredential({
teamId: identity.teamId,
accountId: account.id,
expectedRevision: account.vaultRevision,
});
credential = await credentialWithStickyPatience(
dependencies,
{
teamId: identity.teamId,
accountId: account.id,
expectedRevision: account.vaultRevision,
},
account.sticky,
);
} catch (error) {
failureStage = "credential_refresh";
if (error && typeof error === "object" && "_tag" in error) {
@@ -141,7 +212,7 @@ export async function proxyCodexRequest(request: Request): Promise<Response> {
"warning",
);
try {
const refreshed = await freshCredential({
const refreshed = await dependencies.credential({
teamId: identity.teamId,
accountId: account.id,
expectedRevision: account.vaultRevision,
@@ -172,7 +243,7 @@ export async function proxyCodexRequest(request: Request): Promise<Response> {
status: 429,
},
);
await markAccountCooldown(account.id, rateLimitDelay(upstream.headers));
await dependencies.cooldown(account.id, rateLimitDelay(upstream.headers));
continue;
}
break;
+101
View File
@@ -0,0 +1,101 @@
import { hasActiveCoderouterSubscription } from "../billing/pro";
import { countAccountsForTeam, findAccountByProviderIdentity } from "./repository";
import type { CodeRouterProvider } from "./types";
/**
* Hosted coderouter pricing: a team may connect and route up to this many
* provider accounts (subscriptions) for free. More than this requires an
* active cmux Pro (user) or Team (team) subscription.
*/
export const CODEROUTER_FREE_ACCOUNT_LIMIT = 3;
export type CoderouterEntitlement = {
readonly allowed: boolean;
/** What granted (or would grant) access. */
readonly basis: "free_tier" | "subscription" | "pro_required";
readonly accountCount: number;
};
export type CoderouterEntitlementDependencies = {
readonly countAccounts: typeof countAccountsForTeam;
readonly hasActiveSubscription: typeof hasActiveCoderouterSubscription;
};
/**
* Free tier first: the account count is one cheap indexed RDS read and covers
* most teams, so the Stripe-subscription read only runs for teams over the
* limit. Failures propagate to the caller, which must fail closed.
*/
export function createCoderouterEntitlementCheck(
dependencies: CoderouterEntitlementDependencies,
): (stackUserId: string, teamId: string) => Promise<CoderouterEntitlement> {
return async (stackUserId, teamId) => {
const accountCount = await dependencies.countAccounts(teamId);
if (accountCount <= CODEROUTER_FREE_ACCOUNT_LIMIT) {
return { allowed: true, basis: "free_tier", accountCount };
}
const subscribed = await dependencies.hasActiveSubscription(
stackUserId,
teamId,
);
return subscribed
? { allowed: true, basis: "subscription", accountCount }
: { allowed: false, basis: "pro_required", accountCount };
};
}
export const coderouterEntitlement = createCoderouterEntitlementCheck({
countAccounts: countAccountsForTeam,
hasActiveSubscription: hasActiveCoderouterSubscription,
});
export type AccountAdditionGateDependencies =
& CoderouterEntitlementDependencies
& {
readonly findExisting: typeof findAccountByProviderIdentity;
};
export type AccountAdditionDecision = {
readonly allowed: boolean;
readonly accountCount: number;
};
/**
* Gate for connecting one more provider account. Re-importing a credential
* for an account the team already has never increases the count, so it is
* always allowed — a broken account must stay repairable on the free tier.
* A genuinely new account is allowed while the team stays at or under the
* free limit after adding it, or when a subscription covers the team.
*/
export function createAccountAdditionGate(
dependencies: AccountAdditionGateDependencies,
): (input: {
stackUserId: string;
teamId: string;
provider: CodeRouterProvider;
providerAccountId: string;
}) => Promise<AccountAdditionDecision> {
return async (input) => {
const accountCount = await dependencies.countAccounts(input.teamId);
const existing = await dependencies.findExisting(
input.teamId,
input.provider,
input.providerAccountId,
);
if (existing) return { allowed: true, accountCount };
if (accountCount < CODEROUTER_FREE_ACCOUNT_LIMIT) {
return { allowed: true, accountCount };
}
const subscribed = await dependencies.hasActiveSubscription(
input.stackUserId,
input.teamId,
);
return { allowed: subscribed, accountCount };
};
}
export const accountAdditionAllowed = createAccountAdditionGate({
countAccounts: countAccountsForTeam,
hasActiveSubscription: hasActiveCoderouterSubscription,
findExisting: findAccountByProviderIdentity,
});
+381 -46
View File
@@ -1,10 +1,11 @@
import { createHash, randomBytes, randomUUID } from "node:crypto";
import { and, eq, gt, isNotNull, isNull, lt, lte, notInArray, or, sql } from "drizzle-orm";
import { and, eq, gt, isNotNull, isNull, lt, lte, or, sql } from "drizzle-orm";
import { cloudDb } from "../../db/client";
import {
coderouterAccounts,
coderouterCredentials,
coderouterRouteTokens,
coderouterSessionAccounts,
coderouterVaultLeases,
} from "../../db/schema";
import type { EncryptedCredential } from "./encryption";
@@ -150,22 +151,56 @@ export async function deleteAccount(input: {
export async function listAccounts(
teamId: string,
): Promise<readonly CodeRouterAccountSummary[]> {
return await cloudDb()
.select({
id: coderouterAccounts.id,
provider: coderouterAccounts.provider,
providerAccountId: coderouterAccounts.providerAccountId,
label: coderouterAccounts.label,
state: coderouterAccounts.state,
credentialExpiresAt: coderouterAccounts.credentialExpiresAt,
lastFailureCode: coderouterAccounts.lastFailureCode,
})
.from(coderouterAccounts)
.where(eq(coderouterAccounts.teamId, teamId))
.then((rows) => rows.map((row) => ({
...row,
credentialExpiresAt: row.credentialExpiresAt?.toISOString() ?? null,
})));
const [rows, sessionCounts] = await Promise.all([
cloudDb()
.select({
id: coderouterAccounts.id,
provider: coderouterAccounts.provider,
providerAccountId: coderouterAccounts.providerAccountId,
label: coderouterAccounts.label,
state: coderouterAccounts.state,
credentialExpiresAt: coderouterAccounts.credentialExpiresAt,
lastFailureCode: coderouterAccounts.lastFailureCode,
cooldownUntil: coderouterAccounts.cooldownUntil,
})
.from(coderouterAccounts)
.where(eq(coderouterAccounts.teamId, teamId)),
countActiveSessionsByAccount(teamId),
]);
return rows.map((row) => ({
...row,
credentialExpiresAt: row.credentialExpiresAt?.toISOString() ?? null,
cooldownUntil: row.cooldownUntil?.toISOString() ?? null,
activeSessions: sessionCounts.get(row.id) ?? 0,
}));
}
/**
* Sessions bound per account with traffic inside the load window. Display
* only: a failure here must never take the status endpoint down.
*/
async function countActiveSessionsByAccount(
teamId: string,
): Promise<ReadonlyMap<string, number>> {
try {
const rows = await cloudDb()
.select({
accountId: coderouterSessionAccounts.accountId,
sessions: sql<number>`count(*)::int`,
})
.from(coderouterSessionAccounts)
.where(and(
eq(coderouterSessionAccounts.teamId, teamId),
gt(
coderouterSessionAccounts.lastSeenAt,
sql`now() - interval '${sql.raw(SESSION_BINDING_LOAD_WINDOW)}'`,
),
))
.groupBy(coderouterSessionAccounts.accountId);
return new Map(rows.map((row) => [row.accountId, Number(row.sessions)]));
} catch {
return new Map();
}
}
export async function listCoderouterTeamIds(): Promise<readonly string[]> {
@@ -386,6 +421,19 @@ export async function upsertAccountMetadata(input: {
});
}
/**
* Counts every provider account the team has connected, in any state.
* Broken and cooling-down accounts still occupy a slot: the team controls
* them and can remove them; only removal frees the slot.
*/
export async function countAccountsForTeam(teamId: string): Promise<number> {
const [row] = await cloudDb()
.select({ count: sql<number>`count(*)::int` })
.from(coderouterAccounts)
.where(eq(coderouterAccounts.teamId, teamId));
return Number(row?.count ?? 0);
}
export async function findAccountByProviderIdentity(
teamId: string,
provider: CodeRouterProvider,
@@ -407,15 +455,23 @@ export async function findAccountByProviderIdentity(
return row ?? null;
}
export async function selectAccountForRequest(
teamId: string,
provider: CodeRouterProvider,
excludedAccountIds: readonly string[] = [],
): Promise<{
export type RoutedAccount = {
id: string;
vaultRevision: number;
credentialExpiresAt: Date | null;
} | null> {
};
export type StickyRoutedAccount = RoutedAccount & {
/** True when the session's existing account binding was honored. */
sticky: boolean;
};
/** Bindings older than this stop counting toward an account's session load. */
const SESSION_BINDING_LOAD_WINDOW = "6 hours";
/** Bindings idle longer than this are pruned opportunistically. */
const SESSION_BINDING_RETENTION = "7 days";
async function sweepExpiredRefreshLeases(teamId: string): Promise<void> {
const now = new Date();
await cloudDb()
.update(coderouterAccounts)
@@ -430,33 +486,312 @@ export async function selectAccountForRequest(
eq(coderouterAccounts.state, "refreshing"),
lte(coderouterAccounts.refreshLeaseExpiresAt, now),
));
const [row] = await cloudDb()
.select({
id: coderouterAccounts.id,
vaultRevision: coderouterAccounts.vaultRevision,
credentialExpiresAt: coderouterAccounts.credentialExpiresAt,
})
.from(coderouterAccounts)
.where(and(
eq(coderouterAccounts.teamId, teamId),
eq(coderouterAccounts.provider, provider),
eq(coderouterAccounts.state, "active"),
or(
isNull(coderouterAccounts.cooldownUntil),
lte(coderouterAccounts.cooldownUntil, now),
),
excludedAccountIds.length === 0
? sql`true`
: notInArray(coderouterAccounts.id, [...excludedAccountIds]),
))
.orderBy(sql`${coderouterAccounts.lastUsedAt} asc nulls first`, coderouterAccounts.createdAt)
.limit(1);
}
/**
* Returns the session's bound account when that account is still usable.
* Bumps the binding's last-seen time and the account's last-used time so
* new-session placement steers away from accounts with live traffic.
*/
export async function findSessionAccount(
teamId: string,
provider: CodeRouterProvider,
sessionKey: string,
excludedAccountIds: readonly string[] = [],
): Promise<RoutedAccount | null> {
let result: unknown;
try {
result = await findSessionAccountStatement(
teamId,
provider,
sessionKey,
excludedAccountIds,
);
} catch (error) {
// The session table's migration has not been applied yet. Route without
// stickiness rather than failing the request.
if (isMissingSessionTableError(error)) return null;
throw error;
}
const [row] = databaseRows(result);
if (!row) return null;
await cloudDb()
.update(coderouterAccounts)
.set({ lastUsedAt: new Date() })
.where(eq(coderouterAccounts.id, row.id));
return row;
.where(eq(coderouterAccounts.id, String(row.id)));
return routedAccountRow(row);
}
async function findSessionAccountStatement(
teamId: string,
provider: CodeRouterProvider,
sessionKey: string,
excludedAccountIds: readonly string[],
): Promise<unknown> {
return await cloudDb().execute(sql`
update "coderouter_session_accounts" as binding
set "last_seen_at" = now()
from "coderouter_accounts" as account
where binding."team_id" = ${teamId}
and binding."provider" = ${provider}
and binding."session_key" = ${sessionKey}
and account."id" = binding."account_id"
-- 'refreshing' is a healthy account with a credential refresh in
-- flight (seconds). Moving the session would discard its prompt
-- cache for no reason, so the binding stays usable.
and account."state" in ('active', 'refreshing')
and (account."cooldown_until" is null or account."cooldown_until" <= now())
${accountExclusion(sql`account."id"`, excludedAccountIds)}
returning
account."id" as "id",
account."vault_revision" as "vaultRevision",
account."credential_expires_at" as "credentialExpiresAt"
`);
}
/**
* Atomically claims the best placement candidate for a new session or an
* unbound request. One statement performs read, pick, and write:
* FOR UPDATE SKIP LOCKED makes concurrent claims take different accounts
* instead of all reading the same snapshot and herding onto one account
* (the TypeScript port of subrouter PR #228's placement spread). Ordering
* prefers the fewest recently-active bound sessions, then least-recently-used.
*/
export async function claimAccountForPlacement(
teamId: string,
provider: CodeRouterProvider,
excludedAccountIds: readonly string[] = [],
): Promise<RoutedAccount | null> {
try {
return await claimWithOrdering(teamId, provider, excludedAccountIds, true);
} catch (error) {
// The session table's migration has not been applied yet. Claim without
// the session-load ordering term rather than failing the request.
if (!isMissingSessionTableError(error)) throw error;
return await claimWithOrdering(teamId, provider, excludedAccountIds, false);
}
}
async function claimWithOrdering(
teamId: string,
provider: CodeRouterProvider,
excludedAccountIds: readonly string[],
withSessionLoad: boolean,
): Promise<RoutedAccount | null> {
// First pass skips rows other placements hold locked, so overlapping claims
// fan out across different accounts instead of herding onto one.
const spread = await claimStatement(
teamId,
provider,
excludedAccountIds,
true,
withSessionLoad,
);
if (spread) return spread;
// Every usable account was locked by a concurrent claim (or none exists).
// Fall back to a blocking claim: colliding with another placement is far
// better than telling the caller no account is available.
return await claimStatement(
teamId,
provider,
excludedAccountIds,
false,
withSessionLoad,
);
}
async function claimStatement(
teamId: string,
provider: CodeRouterProvider,
excludedAccountIds: readonly string[],
skipLocked: boolean,
withSessionLoad: boolean,
): Promise<RoutedAccount | null> {
const result = await cloudDb().execute(sql`
with candidate as (
select account."id"
from "coderouter_accounts" as account
where account."team_id" = ${teamId}
and account."provider" = ${provider}
and account."state" = 'active'
and (account."cooldown_until" is null or account."cooldown_until" <= now())
${accountExclusion(sql`account."id"`, excludedAccountIds)}
order by
${withSessionLoad
? sql`(
select count(*)
from "coderouter_session_accounts" as binding
where binding."account_id" = account."id"
and binding."last_seen_at" > now() - interval '${sql.raw(SESSION_BINDING_LOAD_WINDOW)}'
) asc,`
: sql``}
account."last_used_at" asc nulls first,
account."created_at" asc
limit 1
${skipLocked ? sql.raw("for update of account skip locked") : sql``}
)
update "coderouter_accounts" as claimed
set "last_used_at" = now(), "updated_at" = now()
from candidate
where claimed."id" = candidate."id"
returning
claimed."id" as "id",
claimed."vault_revision" as "vaultRevision",
claimed."credential_expires_at" as "credentialExpiresAt"
`);
const [row] = databaseRows(result);
return row ? routedAccountRow(row) : null;
}
/** Pins a session to an account. Last write wins on a same-session race. */
export async function bindSessionAccount(
teamId: string,
provider: CodeRouterProvider,
sessionKey: string,
accountId: string,
): Promise<void> {
const db = cloudDb();
try {
await db
.insert(coderouterSessionAccounts)
.values({ teamId, provider, sessionKey, accountId })
.onConflictDoUpdate({
target: [
coderouterSessionAccounts.teamId,
coderouterSessionAccounts.provider,
coderouterSessionAccounts.sessionKey,
],
set: { accountId, lastSeenAt: new Date() },
});
await db.execute(sql`
delete from "coderouter_session_accounts"
where "team_id" = ${teamId}
and "last_seen_at" < now() - interval '${sql.raw(SESSION_BINDING_RETENTION)}'
`);
} catch (error) {
// The session table's migration has not been applied yet. Skip the pin
// rather than failing the request; routing degrades to the legacy
// per-request behavior until the migration lands.
if (isMissingSessionTableError(error)) return;
throw error;
}
}
/** Postgres undefined_table (42P01), possibly wrapped by the ORM. */
function isMissingSessionTableError(error: unknown): boolean {
let current: unknown = error;
for (let depth = 0; depth < 5 && current; depth++) {
if (
typeof current === "object" &&
"code" in current &&
(current as { code?: unknown }).code === "42P01"
) {
return true;
}
current = typeof current === "object" && current !== null && "cause" in current
? (current as { cause?: unknown }).cause
: undefined;
}
return false;
}
export type SessionAccountSelectorDependencies = {
readonly sweepLeases: typeof sweepExpiredRefreshLeases;
readonly findBound: typeof findSessionAccount;
readonly claim: typeof claimAccountForPlacement;
readonly bind: typeof bindSessionAccount;
};
/**
* Session-sticky account selection.
*
* A bound session keeps riding its account while that account is usable, so
* the provider's prompt cache stays warm. A session moves only when its
* account is broken, cooling down, removed, or already attempted this request
* (the move-worthiness gate from subrouter PR #228, reduced to the binary
* usability signal this schema has). Placement of a new or moving session
* spreads across the least-loaded usable accounts.
*/
export function createSessionAccountSelector(
dependencies: SessionAccountSelectorDependencies,
): (input: {
teamId: string;
provider: CodeRouterProvider;
sessionKey: string | null;
excludedAccountIds?: readonly string[];
}) => Promise<StickyRoutedAccount | null> {
return async (input) => {
const excluded = input.excludedAccountIds ?? [];
await dependencies.sweepLeases(input.teamId);
if (input.sessionKey) {
const bound = await dependencies.findBound(
input.teamId,
input.provider,
input.sessionKey,
excluded,
);
if (bound) return { ...bound, sticky: true };
}
const placed = await dependencies.claim(
input.teamId,
input.provider,
excluded,
);
if (!placed) return null;
if (input.sessionKey) {
await dependencies.bind(
input.teamId,
input.provider,
input.sessionKey,
placed.id,
);
}
return { ...placed, sticky: false };
};
}
export const selectAccountForSession = createSessionAccountSelector({
sweepLeases: sweepExpiredRefreshLeases,
findBound: findSessionAccount,
claim: claimAccountForPlacement,
bind: bindSessionAccount,
});
export async function selectAccountForRequest(
teamId: string,
provider: CodeRouterProvider,
excludedAccountIds: readonly string[] = [],
): Promise<RoutedAccount | null> {
await sweepExpiredRefreshLeases(teamId);
return await claimAccountForPlacement(teamId, provider, excludedAccountIds);
}
function accountExclusion(
column: ReturnType<typeof sql>,
excludedAccountIds: readonly string[],
) {
if (excludedAccountIds.length === 0) return sql``;
return sql` and ${column} not in (${
sql.join(excludedAccountIds.map((id) => sql`${id}`), sql`, `)
})`;
}
function routedAccountRow(row: Record<string, unknown>): RoutedAccount {
return {
id: String(row.id),
vaultRevision: Number(row.vaultRevision),
credentialExpiresAt: row.credentialExpiresAt instanceof Date
? row.credentialExpiresAt
: row.credentialExpiresAt
? new Date(String(row.credentialExpiresAt))
: null,
};
}
function databaseRows(result: unknown): readonly Record<string, unknown>[] {
if (Array.isArray(result)) return result as readonly Record<string, unknown>[];
const rows = (result as { readonly rows?: unknown } | null)?.rows;
return Array.isArray(rows) ? rows as readonly Record<string, unknown>[] : [];
}
export async function markAccountCooldown(
+3
View File
@@ -41,4 +41,7 @@ export type CodeRouterAccountSummary = {
readonly state: "active" | "refreshing" | "expired" | "broken";
readonly credentialExpiresAt: string | null;
readonly lastFailureCode: string | null;
readonly cooldownUntil: string | null;
/** Sessions bound to this account with traffic in the recent window. */
readonly activeSessions: number;
};
@@ -0,0 +1,98 @@
import { describe, expect, mock, test } from "bun:test";
import { makeCoderouterAccountsPostHandler } from "../app/api/coderouter/accounts/route";
const context = {
ok: true as const,
value: {
user: { id: "user_1" },
team: {
teamId: "team_1",
teamName: "Team",
use: true,
manageAccounts: true,
},
},
};
const credentialBody = JSON.stringify({
provider: "codex",
accessToken: "access",
refreshToken: "refresh",
idToken: "id",
accountId: "acct-openai-1",
email: "[email protected]",
expiresAt: Date.now() + 60_000,
});
function addRequest(): Request {
return new Request("https://coderouter.dev/api/coderouter/accounts", {
method: "POST",
headers: { "content-type": "application/json" },
body: credentialBody,
});
}
describe("coderouter account addition limit", () => {
test("blocks connecting a fourth account without a subscription", async () => {
const add = mock(async () => ({ accountId: "new", alreadyExists: false }));
const POST = makeCoderouterAccountsPostHandler({
resolveContext: mock(async () => context) as never,
additionAllowed: async () => ({ allowed: false, accountCount: 3 }),
add,
hostedProRequired: () => true,
});
const response = await POST(addRequest());
expect(response.status).toBe(402);
await expect(response.json()).resolves.toMatchObject({
error: "pro_required",
});
expect(add).not.toHaveBeenCalled();
});
test("stores an account the gate allows", async () => {
const POST = makeCoderouterAccountsPostHandler({
resolveContext: mock(async () => context) as never,
additionAllowed: async () => ({ allowed: true, accountCount: 1 }),
add: async () => ({ accountId: "new", alreadyExists: false }),
hostedProRequired: () => true,
});
const response = await POST(addRequest());
expect(response.status).toBe(201);
});
test("fails closed with a retryable error when the gate is unavailable", async () => {
const add = mock(async () => ({ accountId: "new", alreadyExists: false }));
const POST = makeCoderouterAccountsPostHandler({
resolveContext: mock(async () => context) as never,
additionAllowed: async () => {
throw new Error("database unavailable");
},
add,
hostedProRequired: () => true,
});
const response = await POST(addRequest());
expect(response.status).toBe(503);
await expect(response.json()).resolves.toMatchObject({
error: "entitlement_unavailable",
retryable: true,
});
expect(add).not.toHaveBeenCalled();
});
test("skips the gate entirely when hosted billing is off", async () => {
const additionAllowed = mock(async () => ({
allowed: false,
accountCount: 9,
}));
const POST = makeCoderouterAccountsPostHandler({
resolveContext: mock(async () => context) as never,
additionAllowed,
add: async () => ({ accountId: "new", alreadyExists: false }),
hostedProRequired: () => false,
});
const response = await POST(addRequest());
expect(response.status).toBe(201);
expect(additionAllowed).not.toHaveBeenCalled();
});
});
+115
View File
@@ -0,0 +1,115 @@
import { describe, expect, mock, test } from "bun:test";
import {
CODEROUTER_FREE_ACCOUNT_LIMIT,
createAccountAdditionGate,
createCoderouterEntitlementCheck,
} from "../services/coderouter/entitlement";
describe("coderouter entitlement", () => {
test("free limit is 3 connected accounts", () => {
expect(CODEROUTER_FREE_ACCOUNT_LIMIT).toBe(3);
});
test("a team at or under the free limit routes without a subscription", async () => {
for (const count of [0, 1, 2, 3]) {
const hasActiveSubscription = mock(async () => false);
const check = createCoderouterEntitlementCheck({
countAccounts: async () => count,
hasActiveSubscription,
});
const result = await check("user_1", "team_1");
expect(result).toEqual({
allowed: true,
basis: "free_tier",
accountCount: count,
});
// The Stripe read must not run for free-tier teams.
expect(hasActiveSubscription).not.toHaveBeenCalled();
}
});
test("a team over the free limit requires a subscription", async () => {
const check = createCoderouterEntitlementCheck({
countAccounts: async () => 4,
hasActiveSubscription: async () => false,
});
await expect(check("user_1", "team_1")).resolves.toEqual({
allowed: false,
basis: "pro_required",
accountCount: 4,
});
});
test("a subscription covers a team over the free limit", async () => {
const check = createCoderouterEntitlementCheck({
countAccounts: async () => 7,
hasActiveSubscription: async (userId, teamId) =>
userId === "user_1" && teamId === "team_1",
});
await expect(check("user_1", "team_1")).resolves.toEqual({
allowed: true,
basis: "subscription",
accountCount: 7,
});
});
});
describe("coderouter account addition gate", () => {
const input = {
stackUserId: "user_1",
teamId: "team_1",
provider: "codex" as const,
providerAccountId: "acct-openai-1",
};
test("allows a new account while the team stays within the free limit", async () => {
const gate = createAccountAdditionGate({
countAccounts: async () => 2,
hasActiveSubscription: async () => false,
findExisting: async () => null,
});
await expect(gate(input)).resolves.toEqual({
allowed: true,
accountCount: 2,
});
});
test("blocks the fourth account without a subscription", async () => {
const gate = createAccountAdditionGate({
countAccounts: async () => 3,
hasActiveSubscription: async () => false,
findExisting: async () => null,
});
await expect(gate(input)).resolves.toEqual({
allowed: false,
accountCount: 3,
});
});
test("allows the fourth account with a subscription", async () => {
const gate = createAccountAdditionGate({
countAccounts: async () => 3,
hasActiveSubscription: async () => true,
findExisting: async () => null,
});
await expect(gate(input)).resolves.toEqual({
allowed: true,
accountCount: 3,
});
});
test("always allows re-importing an existing account", async () => {
const hasActiveSubscription = mock(async () => false);
const gate = createAccountAdditionGate({
countAccounts: async () => 5,
hasActiveSubscription,
findExisting: async () => ({ id: "acct-1", state: "broken", vaultRevision: 2 }),
});
await expect(gate(input)).resolves.toEqual({
allowed: true,
accountCount: 5,
});
expect(hasActiveSubscription).not.toHaveBeenCalled();
});
});
@@ -0,0 +1,161 @@
import { afterAll, beforeAll, beforeEach, describe, expect, mock, test } from "bun:test";
type SelectInput = {
teamId: string;
provider: string;
sessionKey: string | null;
excludedAccountIds?: readonly string[];
};
let selectInputs: SelectInput[] = [];
let accountsToServe: { id: string; sticky: boolean }[] = [];
let cooldowns: string[] = [];
let upstreamStatuses: number[] = [];
let credentialBusyBudgets = new Map<string, number>();
let credentialCalls: string[] = [];
const originalFetch = globalThis.fetch;
beforeAll(() => {
globalThis.fetch = mock(async () => {
const status = upstreamStatuses.shift() ?? 200;
return new Response("data: done\n\n", {
status,
headers: { "content-type": "text/event-stream" },
});
}) as typeof fetch;
});
afterAll(() => {
globalThis.fetch = originalFetch;
});
const { createCodexResponsesProxy } = await import("../services/coderouter/codexProxy");
const proxy = createCodexResponsesProxy({
authenticate: async () => ({ teamId: "team-1", stackUserId: "stack-user-1" }),
select: async (input) => {
selectInputs.push({
...(input as SelectInput),
excludedAccountIds: [...(input.excludedAccountIds ?? [])],
});
const next = accountsToServe.shift();
return next
? {
id: next.id,
vaultRevision: 1,
credentialExpiresAt: null,
sticky: next.sticky,
}
: null;
},
credential: async ({ accountId }) => {
if (credentialBusyBudgets.get(accountId)) {
credentialBusyBudgets.set(
accountId,
(credentialBusyBudgets.get(accountId) ?? 1) - 1,
);
throw Object.assign(new Error("busy"), { _tag: "CodeRouterRefreshBusy" });
}
credentialCalls.push(accountId);
return {
provider: "codex",
accessToken: `access-${accountId}`,
refreshToken: "refresh",
idToken: "id",
accountId: "chatgpt-account",
email: "[email protected]",
expiresAt: Date.now() + 60_000,
};
},
cooldown: async (accountId) => {
cooldowns.push(accountId);
},
});
beforeEach(() => {
selectInputs = [];
accountsToServe = [];
cooldowns = [];
upstreamStatuses = [];
credentialBusyBudgets = new Map();
credentialCalls = [];
});
function responsesRequest(headers: Record<string, string> = {}): Request {
return new Request("https://coderouter.dev/v1/responses", {
method: "POST",
headers: {
authorization: "Bearer crt_token",
"content-type": "application/json",
...headers,
},
body: JSON.stringify({ model: "gpt-test", input: [] }),
});
}
describe("codex responses proxy session routing", () => {
test("passes the session_id header to account selection", async () => {
accountsToServe = [{ id: "acct-1", sticky: true }];
const response = await proxy(responsesRequest({ session_id: "session-abc" }));
expect(response.status).toBe(200);
expect(selectInputs).toHaveLength(1);
expect(selectInputs[0]?.sessionKey).toBe("session-abc");
expect(selectInputs[0]?.teamId).toBe("team-1");
expect(selectInputs[0]?.provider).toBe("codex");
});
test("selects without a session key when the header is missing", async () => {
accountsToServe = [{ id: "acct-1", sticky: false }];
const response = await proxy(responsesRequest());
expect(response.status).toBe(200);
expect(selectInputs[0]?.sessionKey).toBeNull();
});
test("ignores oversized session ids", async () => {
accountsToServe = [{ id: "acct-1", sticky: false }];
await proxy(responsesRequest({ session_id: "x".repeat(600) }));
expect(selectInputs[0]?.sessionKey).toBeNull();
});
test("cools down a rate-limited account and retries excluding it", async () => {
accountsToServe = [
{ id: "acct-1", sticky: true },
{ id: "acct-2", sticky: false },
];
upstreamStatuses = [429, 200];
const response = await proxy(responsesRequest({ session_id: "session-move" }));
expect(response.status).toBe(200);
expect(cooldowns).toEqual(["acct-1"]);
expect(selectInputs).toHaveLength(2);
expect(selectInputs[1]?.excludedAccountIds).toEqual(["acct-1"]);
expect(selectInputs[1]?.sessionKey).toBe("session-move");
});
test("a sticky session waits out an in-flight refresh instead of moving", async () => {
accountsToServe = [{ id: "acct-1", sticky: true }];
credentialBusyBudgets.set("acct-1", 2);
const response = await proxy(responsesRequest({ session_id: "session-wait" }));
expect(response.status).toBe(200);
expect(selectInputs).toHaveLength(1);
expect(credentialCalls).toEqual(["acct-1"]);
});
test("a non-sticky request moves immediately on refresh-busy", async () => {
accountsToServe = [
{ id: "acct-1", sticky: false },
{ id: "acct-2", sticky: false },
];
credentialBusyBudgets.set("acct-1", 1);
const response = await proxy(responsesRequest());
expect(response.status).toBe(200);
expect(credentialCalls).toEqual(["acct-2"]);
expect(selectInputs).toHaveLength(2);
});
test("returns no_usable_account when selection is exhausted", async () => {
accountsToServe = [];
const response = await proxy(responsesRequest({ session_id: "session-dry" }));
expect(response.status).toBe(503);
const body = await response.json() as { error: string };
expect(body.error).toBe("no_usable_account");
});
});
@@ -0,0 +1,268 @@
import { afterAll, beforeAll, beforeEach, describe, expect, test } from "bun:test";
import { randomUUID } from "node:crypto";
import postgres, { type Sql } from "postgres";
import { closeCloudDbForTests } from "../db/client";
import { listAccounts } from "../services/coderouter/repository";
import {
bindSessionAccount,
claimAccountForPlacement,
findSessionAccount,
markAccountCooldown,
selectAccountForSession,
} from "../services/coderouter/repository";
const runDbTests = process.env.CMUX_DB_TEST === "1";
const dbTest = runDbTests ? test : test.skip;
const TEAM = "team-routing-test";
let sql: Sql | null = null;
beforeAll(() => {
if (!runDbTests) return;
const databaseURL = process.env.DIRECT_DATABASE_URL ?? process.env.DATABASE_URL;
if (!databaseURL) throw new Error("DATABASE_URL is required when CMUX_DB_TEST=1");
sql = postgres(databaseURL, { max: 4 });
});
afterAll(async () => {
await closeCloudDbForTests();
await sql?.end({ timeout: 5 });
});
beforeEach(async () => {
if (!sql) return;
await sql`truncate coderouter_session_accounts, coderouter_accounts cascade`;
});
async function insertAccounts(count: number): Promise<string[]> {
if (!sql) throw new Error("no sql client");
const ids: string[] = [];
for (let index = 0; index < count; index++) {
const id = randomUUID();
ids.push(id);
await sql`
insert into coderouter_accounts
(id, team_id, provider, provider_account_id, label, state)
values
(${id}, ${TEAM}, 'codex', ${`provider-${id}`}, ${`Account ${index}`}, 'active')
`;
}
return ids;
}
async function bindingCounts(): Promise<Map<string, number>> {
if (!sql) throw new Error("no sql client");
const rows = await sql`
select account_id, count(*)::int as sessions
from coderouter_session_accounts
where team_id = ${TEAM}
group by account_id
`;
return new Map(rows.map((row) => [String(row.account_id), Number(row.sessions)]));
}
describe("coderouter routing db behavior", () => {
dbTest("sequential new sessions spread evenly across accounts", async () => {
const accounts = await insertAccounts(4);
for (let index = 0; index < 8; index++) {
const placed = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: `session-${index}`,
});
expect(placed).not.toBeNull();
expect(placed?.sticky).toBe(false);
}
const counts = await bindingCounts();
expect(counts.size).toBe(4);
for (const accountId of accounts) {
expect(counts.get(accountId)).toBe(2);
}
});
dbTest("concurrent new sessions do not all land on one account", async () => {
await insertAccounts(4);
const placements = await Promise.all(
Array.from({ length: 12 }, (_, index) =>
selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: `burst-${index}`,
})),
);
const ids = placements.map((placed) => {
expect(placed).not.toBeNull();
return placed?.id ?? "";
});
const concentration = new Map<string, number>();
for (const id of ids) {
concentration.set(id, (concentration.get(id) ?? 0) + 1);
}
// The old read-pick-write race let all 12 pick the same account.
expect(concentration.size).toBeGreaterThanOrEqual(2);
const max = Math.max(...concentration.values());
expect(max).toBeLessThan(12);
});
dbTest("a session sticks to its account across requests", async () => {
await insertAccounts(3);
const first = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "sticky-session",
});
expect(first?.sticky).toBe(false);
for (let index = 0; index < 5; index++) {
const again = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "sticky-session",
});
expect(again?.id).toBe(first?.id ?? "");
expect(again?.sticky).toBe(true);
}
// Other sessions in between must not steal the binding.
await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "other-session",
});
const after = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "sticky-session",
});
expect(after?.id).toBe(first?.id ?? "");
});
dbTest("a session moves once when its account cools down, then resticks", async () => {
await insertAccounts(3);
const first = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "moving-session",
});
expect(first).not.toBeNull();
await markAccountCooldown(first?.id ?? "", 60_000);
const moved = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "moving-session",
});
expect(moved).not.toBeNull();
expect(moved?.id).not.toBe(first?.id ?? "");
expect(moved?.sticky).toBe(false);
const stuck = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "moving-session",
});
expect(stuck?.id).toBe(moved?.id ?? "");
expect(stuck?.sticky).toBe(true);
});
dbTest("claim skips excluded accounts and drains to null", async () => {
const accounts = await insertAccounts(2);
const first = await claimAccountForPlacement(TEAM, "codex", []);
expect(first).not.toBeNull();
const second = await claimAccountForPlacement(TEAM, "codex", [first?.id ?? ""]);
expect(second).not.toBeNull();
expect(second?.id).not.toBe(first?.id ?? "");
const drained = await claimAccountForPlacement(TEAM, "codex", accounts);
expect(drained).toBeNull();
});
dbTest("bindings honor exclusion and report null when excluded", async () => {
await insertAccounts(2);
const placed = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "excluded-session",
});
expect(placed).not.toBeNull();
const bound = await findSessionAccount(
TEAM,
"codex",
"excluded-session",
[placed?.id ?? ""],
);
expect(bound).toBeNull();
});
dbTest("a binding survives its account being mid-refresh", async () => {
if (!sql) throw new Error("no sql client");
await insertAccounts(2);
const first = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "refreshing-session",
});
expect(first).not.toBeNull();
await sql`
update coderouter_accounts
set state = 'refreshing',
refresh_lease_id = ${randomUUID()},
refresh_lease_expires_at = now() + interval '30 seconds'
where id = ${first?.id ?? ""}
`;
const during = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "refreshing-session",
});
expect(during?.id).toBe(first?.id ?? "");
expect(during?.sticky).toBe(true);
});
dbTest("account listing reports active session counts and cooldowns", async () => {
if (!sql) throw new Error("no sql client");
const accounts = await insertAccounts(2);
await bindSessionAccount(TEAM, "codex", "fresh-session-1", accounts[0] ?? "");
await bindSessionAccount(TEAM, "codex", "fresh-session-2", accounts[0] ?? "");
await bindSessionAccount(TEAM, "codex", "stale-session", accounts[1] ?? "");
await sql`
update coderouter_session_accounts
set last_seen_at = now() - interval '7 hours'
where session_key = 'stale-session'
`;
await markAccountCooldown(accounts[1] ?? "", 60_000);
const listed = await listAccounts(TEAM);
const byId = new Map(listed.map((account) => [account.id, account]));
expect(byId.get(accounts[0] ?? "")?.activeSessions).toBe(2);
expect(byId.get(accounts[0] ?? "")?.cooldownUntil).toBeNull();
// A binding idle beyond the load window no longer counts.
expect(byId.get(accounts[1] ?? "")?.activeSessions).toBe(0);
expect(byId.get(accounts[1] ?? "")?.cooldownUntil).not.toBeNull();
});
dbTest("routes without stickiness while the session table is missing", async () => {
if (!sql) throw new Error("no sql client");
await insertAccounts(2);
await sql`alter table coderouter_session_accounts rename to coderouter_session_accounts_gone`;
try {
const placed = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "pre-migration-session",
});
expect(placed).not.toBeNull();
expect(placed?.sticky).toBe(false);
const again = await selectAccountForSession({
teamId: TEAM,
provider: "codex",
sessionKey: "pre-migration-session",
});
expect(again).not.toBeNull();
} finally {
await sql`alter table coderouter_session_accounts_gone rename to coderouter_session_accounts`;
}
});
dbTest("bind is last-write-wins for a session key", async () => {
const accounts = await insertAccounts(2);
await bindSessionAccount(TEAM, "codex", "raced-session", accounts[0] ?? "");
await bindSessionAccount(TEAM, "codex", "raced-session", accounts[1] ?? "");
const bound = await findSessionAccount(TEAM, "codex", "raced-session", []);
expect(bound?.id).toBe(accounts[1] ?? "");
});
});
+46 -10
View File
@@ -46,7 +46,11 @@ describe("coderouter hosted entitlement", () => {
}));
const POST = makeCoderouterSessionPostHandler({
resolveContext: mock(async () => context) as never,
hasActiveEntitlement: mock(async () => false),
entitlement: mock(async () => ({
allowed: false as const,
basis: "pro_required" as const,
accountCount: 5,
})),
issueToken,
hostedProRequired: () => true,
});
@@ -64,15 +68,45 @@ describe("coderouter hosted entitlement", () => {
expect(issueToken).not.toHaveBeenCalled();
});
test("issues a hosted route token on the free tier without a subscription", async () => {
const issueToken = mock(async () => ({
token: "crt_free",
expiresAt: new Date("2026-09-01T00:00:00Z"),
}));
const POST = makeCoderouterSessionPostHandler({
resolveContext: mock(async () => context) as never,
entitlement: mock(async () => ({
allowed: true as const,
basis: "free_tier" as const,
accountCount: 3,
})),
issueToken,
hostedProRequired: () => true,
});
const response = await POST(
new Request("https://coderouter.dev/api/coderouter/session", {
method: "POST",
}),
);
expect(response.status).toBe(200);
expect(issueToken).toHaveBeenCalledWith("team_1", "user_1");
});
test("keeps self-hosted servers independent from hosted billing", async () => {
const issueToken = mock(async () => ({
token: "crt_test",
expiresAt: new Date("2026-09-01T00:00:00Z"),
}));
const hasActiveEntitlement = mock(async () => false);
const entitlement = mock(async () => ({
allowed: false as const,
basis: "pro_required" as const,
accountCount: 5,
}));
const POST = makeCoderouterSessionPostHandler({
resolveContext: mock(async () => context) as never,
hasActiveEntitlement,
entitlement,
issueToken,
hostedProRequired: () => false,
});
@@ -88,14 +122,14 @@ describe("coderouter hosted entitlement", () => {
token: "crt_test",
openaiBaseUrl: "https://router.example.com/v1",
});
expect(hasActiveEntitlement).not.toHaveBeenCalled();
expect(entitlement).not.toHaveBeenCalled();
expect(issueToken).toHaveBeenCalledWith("team_1", "user_1");
});
test("fails closed when hosted entitlement storage is unavailable", async () => {
const POST = makeCoderouterSessionPostHandler({
resolveContext: mock(async () => context) as never,
hasActiveEntitlement: mock(async () => {
entitlement: mock(async () => {
throw new Error("database unavailable");
}),
issueToken: mock(async () => {
@@ -119,12 +153,14 @@ describe("coderouter hosted entitlement", () => {
token: "crt_team",
expiresAt: new Date("2026-09-01T00:00:00Z"),
}));
const hasActiveEntitlement = mock(async (...args: unknown[]) =>
args[0] === "user_1" && args[1] === "team_1"
);
const entitlement = mock(async (...args: unknown[]) => ({
allowed: args[0] === "user_1" && args[1] === "team_1",
basis: "subscription" as const,
accountCount: 5,
}));
const POST = makeCoderouterSessionPostHandler({
resolveContext: mock(async () => context) as never,
hasActiveEntitlement: hasActiveEntitlement as never,
entitlement: entitlement as never,
issueToken,
hostedProRequired: () => true,
});
@@ -136,7 +172,7 @@ describe("coderouter hosted entitlement", () => {
);
expect(response.status).toBe(200);
expect(hasActiveEntitlement).toHaveBeenCalledWith("user_1", "team_1");
expect(entitlement).toHaveBeenCalledWith("user_1", "team_1");
expect(issueToken).toHaveBeenCalledWith("team_1", "user_1");
});
});
@@ -0,0 +1,116 @@
import { describe, expect, test } from "bun:test";
import { createSessionAccountSelector } from "../services/coderouter/repository";
type Call = { fn: string; args: unknown[] };
function makeDependencies(options: {
bound?: { id: string; vaultRevision: number; credentialExpiresAt: Date | null } | null;
placed?: { id: string; vaultRevision: number; credentialExpiresAt: Date | null } | null;
}) {
const calls: Call[] = [];
return {
calls,
dependencies: {
sweepLeases: async (...args: unknown[]) => {
calls.push({ fn: "sweepLeases", args });
},
findBound: async (...args: unknown[]) => {
calls.push({ fn: "findBound", args });
return options.bound ?? null;
},
claim: async (...args: unknown[]) => {
calls.push({ fn: "claim", args });
return options.placed ?? null;
},
bind: async (...args: unknown[]) => {
calls.push({ fn: "bind", args });
},
},
};
}
const account = (id: string) => ({
id,
vaultRevision: 3,
credentialExpiresAt: null,
});
describe("coderouter session account selector", () => {
test("honors an existing usable binding and does not place", async () => {
const { calls, dependencies } = makeDependencies({ bound: account("acct-1") });
const select = createSessionAccountSelector(dependencies);
const result = await select({
teamId: "team-1",
provider: "codex",
sessionKey: "session-a",
});
expect(result).toEqual({ ...account("acct-1"), sticky: true });
expect(calls.map((c) => c.fn)).toEqual(["sweepLeases", "findBound"]);
});
test("places and binds a new session when no binding exists", async () => {
const { calls, dependencies } = makeDependencies({
bound: null,
placed: account("acct-2"),
});
const select = createSessionAccountSelector(dependencies);
const result = await select({
teamId: "team-1",
provider: "codex",
sessionKey: "session-b",
});
expect(result).toEqual({ ...account("acct-2"), sticky: false });
expect(calls.map((c) => c.fn)).toEqual([
"sweepLeases",
"findBound",
"claim",
"bind",
]);
expect(calls[3]?.args).toEqual(["team-1", "codex", "session-b", "acct-2"]);
});
test("rebinds when the bound account is no longer usable", async () => {
const { calls, dependencies } = makeDependencies({
bound: null,
placed: account("acct-3"),
});
const select = createSessionAccountSelector(dependencies);
const result = await select({
teamId: "team-1",
provider: "codex",
sessionKey: "session-c",
excludedAccountIds: ["acct-1"],
});
expect(result?.sticky).toBe(false);
expect(result?.id).toBe("acct-3");
const findBound = calls.find((c) => c.fn === "findBound");
expect(findBound?.args).toEqual(["team-1", "codex", "session-c", ["acct-1"]]);
const claim = calls.find((c) => c.fn === "claim");
expect(claim?.args).toEqual(["team-1", "codex", ["acct-1"]]);
});
test("skips stickiness entirely without a session key", async () => {
const { calls, dependencies } = makeDependencies({ placed: account("acct-4") });
const select = createSessionAccountSelector(dependencies);
const result = await select({
teamId: "team-1",
provider: "codex",
sessionKey: null,
});
expect(result).toEqual({ ...account("acct-4"), sticky: false });
expect(calls.map((c) => c.fn)).toEqual(["sweepLeases", "claim"]);
});
test("returns null when no account is usable", async () => {
const { calls, dependencies } = makeDependencies({ bound: null, placed: null });
const select = createSessionAccountSelector(dependencies);
const result = await select({
teamId: "team-1",
provider: "codex",
sessionKey: "session-d",
});
expect(result).toBeNull();
expect(calls.some((c) => c.fn === "bind")).toBe(false);
});
});