Compare commits
13
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3fad5abf1d | ||
|
|
34f0e1eeca | ||
|
|
882ab10e75 | ||
|
|
9ad361ed37 | ||
|
|
ae7d761e03 | ||
|
|
c15d665f65 | ||
|
|
425f2d7efe | ||
|
|
7d79a858f6 | ||
|
|
18ed976a1b | ||
|
|
6e56e6ff14 | ||
|
|
b0948c0341 | ||
|
|
d3e3abd367 | ||
|
|
9f556f118e |
+3
@@ -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()
|
||||
|
||||
|
||||
+3
@@ -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 }
|
||||
|
||||
+6
-1
@@ -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() {}
|
||||
|
||||
+12
@@ -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() }
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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");
|
||||
@@ -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",
|
||||
{
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
@@ -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(
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -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,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);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user