diff --git a/apps/web/src/cloud/managedRelayState.ts b/apps/web/src/cloud/managedRelayState.ts index c5f7536305337..679f362fdf5a0 100644 --- a/apps/web/src/cloud/managedRelayState.ts +++ b/apps/web/src/cloud/managedRelayState.ts @@ -150,11 +150,14 @@ export function useManagedRelayDevices() { /** How this account can raise its tunnel cap: link T3 Chat, or upgrade the linked plan. */ export type ManagedTunnelUpgrade = "link" | "pro" | "premier" | null; -// Mirrors the relay's cap without manual overrides, which the relay does not expose. +// Mirrors the relay's cap: a manual override is exact, otherwise the T3 Chat plan raises the default. function managedTunnelCap(t3Chat: RelayT3ChatLinkResponse): { readonly limit: number; readonly upgrade: ManagedTunnelUpgrade; } { + if (t3Chat.maxTunnelsOverride !== undefined) { + return { limit: t3Chat.maxTunnelsOverride, upgrade: null }; + } if (!t3Chat.available) return { limit: DEFAULT_MANAGED_TUNNEL_LIMIT, upgrade: null }; switch (t3Chat.link?.plan) { case "premier": diff --git a/infra/relay/src/environments/ManagedEndpointProvider.test.ts b/infra/relay/src/environments/ManagedEndpointProvider.test.ts index 01230171be9ba..96661fad42a61 100644 --- a/infra/relay/src/environments/ManagedEndpointProvider.test.ts +++ b/infra/relay/src/environments/ManagedEndpointProvider.test.ts @@ -343,6 +343,7 @@ function makeTunnelLimits( } return yield* reserve; }), + getOverride: () => Effect.die("unused"), }); } diff --git a/infra/relay/src/environments/ManagedTunnelLimits.test.ts b/infra/relay/src/environments/ManagedTunnelLimits.test.ts index 6aca885b85b87..6b8dd817ee01c 100644 --- a/infra/relay/src/environments/ManagedTunnelLimits.test.ts +++ b/infra/relay/src/environments/ManagedTunnelLimits.test.ts @@ -121,6 +121,27 @@ const reserve = ( ), ); +const getOverride = (override?: number) => + Effect.gen(function* () { + const limits = yield* ManagedTunnelLimits.ManagedTunnelLimits; + return yield* limits.getOverride({ userId: "user-1" }); + }).pipe( + Effect.provide( + ManagedTunnelLimits.layer.pipe( + Layer.provide( + Layer.succeed( + RelayDb.RelayDb, + makeFakeDb({ + otherSlots: () => 0, + ...(override === undefined ? {} : { override }), + }), + ), + ), + Layer.provide(t3ChatLayer("free")), + ), + ), + ); + describe("ManagedTunnelLimits", () => { it.effect("rejects concurrent reservations past the limit", () => { const reserved: Array = []; @@ -383,7 +404,9 @@ describe("ManagedTunnelLimits", () => { reserve({ override, activeTunnels: override, plan: "premier", lookups }), ); expect(error).toMatchObject({ _tag: "ManagedTunnelLimitExceeded", maxTunnels: override }); + expect(yield* getOverride(override)).toBe(override); } + expect(yield* getOverride()).toBeNull(); expect(lookups).toEqual([]); }), ); diff --git a/infra/relay/src/environments/ManagedTunnelLimits.ts b/infra/relay/src/environments/ManagedTunnelLimits.ts index aa373ad27ffb2..3af90f86c4163 100644 --- a/infra/relay/src/environments/ManagedTunnelLimits.ts +++ b/infra/relay/src/environments/ManagedTunnelLimits.ts @@ -106,6 +106,10 @@ export class ManagedTunnelLimits extends Context.Service< input: { readonly userId: string; readonly environmentId: string }, reserve: Effect.Effect, ) => Effect.Effect; + /** The user's manual cap, which replaces the default and any T3 Chat plan cap. */ + readonly getOverride: (input: { + readonly userId: string; + }) => Effect.Effect; } >()("t3code-relay/environments/ManagedTunnelLimits") {} @@ -126,7 +130,24 @@ const make = Effect.gen(function* () { ), ); + const loadOverride = (userId: string) => + db + .select({ maxTunnels: relayManagedTunnelLimits.maxTunnels }) + .from(relayManagedTunnelLimits) + .where(eq(relayManagedTunnelLimits.userId, userId)) + .limit(1) + .pipe( + Effect.map((rows) => rows[0]?.maxTunnels ?? null), + Effect.mapError( + (cause) => + new ManagedTunnelLimitPersistenceError({ operation: "load-limit", userId, cause }), + ), + ); + return ManagedTunnelLimits.of({ + getOverride: Effect.fn("relay.managed_tunnel_limits.get_override")(function* (input) { + return yield* loadOverride(input.userId); + }), withCapacity: Effect.fn("relay.managed_tunnel_limits.with_capacity")(function* ( input: { readonly userId: string; readonly environmentId: string }, reserve: Effect.Effect, @@ -148,22 +169,8 @@ const make = Effect.gen(function* () { ), ); - const overrides = yield* db - .select({ maxTunnels: relayManagedTunnelLimits.maxTunnels }) - .from(relayManagedTunnelLimits) - .where(eq(relayManagedTunnelLimits.userId, input.userId)) - .limit(1) - .pipe( - Effect.mapError( - (cause) => - new ManagedTunnelLimitPersistenceError({ - operation: "load-limit", - userId: input.userId, - cause, - }), - ), - ); - const baseMaxTunnels = overrides[0]?.maxTunnels ?? DEFAULT_MANAGED_TUNNEL_LIMIT; + const override = yield* loadOverride(input.userId); + const baseMaxTunnels = override ?? DEFAULT_MANAGED_TUNNEL_LIMIT; const pendingCutoff = DateTime.formatIso( DateTime.subtract(yield* DateTime.now, { @@ -196,7 +203,7 @@ const make = Effect.gen(function* () { // A manual override is the exact cap; a T3 Chat plan raises only the default. // Runs under the user lock, bounded by the T3 Chat request timeout. const maxTunnels = - t3Chat.available && overrides.length === 0 + t3Chat.available && override === null ? Math.max(baseMaxTunnels, yield* t3ChatTunnelLimit(input.userId)) : baseMaxTunnels; if (activeTunnels >= maxTunnels) { diff --git a/infra/relay/src/http/Api.test.ts b/infra/relay/src/http/Api.test.ts index a20989cd3c025..e6db7fc8b7aa0 100644 --- a/infra/relay/src/http/Api.test.ts +++ b/infra/relay/src/http/Api.test.ts @@ -64,6 +64,7 @@ import * as HookInbox from "../hooks/HookInbox.ts"; import * as EnvironmentLinks from "../environments/EnvironmentLinks.ts"; import * as ManagedEndpointAllocations from "../environments/ManagedEndpointAllocations.ts"; import * as ManagedEndpointProvider from "../environments/ManagedEndpointProvider.ts"; +import * as ManagedTunnelLimits from "../environments/ManagedTunnelLimits.ts"; import * as AgentActivityPublisher from "../agentActivity/AgentActivityPublisher.ts"; import * as EnvironmentPublishSignatures from "../environments/EnvironmentPublishSignatures.ts"; import * as T3ChatAccounts from "../t3chat/T3ChatAccounts.ts"; @@ -130,6 +131,7 @@ describe("managed endpoint provisioning errors", () => { Layer.mock(ManagedEndpointProvider.ManagedEndpointProvider, {}), Layer.mock(Devices.Devices, {}), Layer.mock(T3ChatAccounts.T3ChatAccounts, { available: false }), + Layer.mock(ManagedTunnelLimits.ManagedTunnelLimits, {}), ), ), Layer.provideMerge( @@ -241,6 +243,7 @@ describe("device listing compatibility", () => { }, }), Layer.mock(T3ChatAccounts.T3ChatAccounts, { available: false }), + Layer.mock(ManagedTunnelLimits.ManagedTunnelLimits, {}), ), ), Layer.provideMerge( @@ -322,6 +325,7 @@ describe("Cloudflare rate-limit retry hint", () => { Layer.mock(ManagedEndpointProvider.ManagedEndpointProvider, {}), Layer.mock(Devices.Devices, {}), Layer.mock(T3ChatAccounts.T3ChatAccounts, { available: false }), + Layer.mock(ManagedTunnelLimits.ManagedTunnelLimits, {}), ), ), Layer.provideMerge( diff --git a/infra/relay/src/http/Api.ts b/infra/relay/src/http/Api.ts index 5a4d77c28d3c2..390f3c18f785f 100644 --- a/infra/relay/src/http/Api.ts +++ b/infra/relay/src/http/Api.ts @@ -74,6 +74,7 @@ import * as AgentActivityPublisher from "../agentActivity/AgentActivityPublisher import * as EnvironmentConnector from "../environments/EnvironmentConnector.ts"; import * as EnvironmentLinker from "../environments/EnvironmentLinker.ts"; import * as ManagedEndpointProvider from "../environments/ManagedEndpointProvider.ts"; +import * as ManagedTunnelLimits from "../environments/ManagedTunnelLimits.ts"; import * as ManagedEndpointAllocations from "../environments/ManagedEndpointAllocations.ts"; import * as EnvironmentPublishSignatures from "../environments/EnvironmentPublishSignatures.ts"; import * as MobileRegistrations from "../agentActivity/MobileRegistrations.ts"; @@ -825,6 +826,7 @@ export const layerClientApi = HttpApiBuilder.group( const managedEndpointProvider = yield* ManagedEndpointProvider.ManagedEndpointProvider; const devices = yield* Devices.Devices; const t3Chat = yield* T3ChatAccounts.T3ChatAccounts; + const tunnelLimits = yield* ManagedTunnelLimits.ManagedTunnelLimits; return handlers .handle( "listEnvironments", @@ -996,12 +998,20 @@ export const layerClientApi = HttpApiBuilder.group( "getT3ChatLink", Effect.fn("relay.api.client.getT3ChatLink")(function* () { yield* appendRelayCredentialResponseHeaders; - if (!t3Chat.available) return { available: false, link: null }; const { userId } = yield* RelayClientPrincipal; - const link = yield* t3Chat - .getLink({ userId }) - .pipe(Effect.catch(() => relayInternalErrorResponse("upstream_unavailable"))); - return { available: true, link }; + const link = t3Chat.available + ? yield* t3Chat + .getLink({ userId }) + .pipe(Effect.catch(() => relayInternalErrorResponse("upstream_unavailable"))) + : null; + const maxTunnelsOverride = yield* tunnelLimits + .getOverride({ userId }) + .pipe(Effect.catch(() => relayInternalErrorResponse("internal_error"))); + return { + available: t3Chat.available, + link, + ...(maxTunnelsOverride === null ? {} : { maxTunnelsOverride }), + }; }, mapRelayCommonApiErrors("not_authorized")), ) .handle( diff --git a/packages/contracts/src/relay.ts b/packages/contracts/src/relay.ts index 4dbb6a199f472..8ebc9df23a0bb 100644 --- a/packages/contracts/src/relay.ts +++ b/packages/contracts/src/relay.ts @@ -981,6 +981,8 @@ export const RelayT3ChatLinkResponse = Schema.Struct({ // False when this relay has no T3 Chat secret, so clients hide the feature. available: Schema.Boolean, link: Schema.NullOr(RelayT3ChatLink), + // A manual cap that replaces the plan's; T3 Chat upgrades do not raise it. + maxTunnelsOverride: Schema.optionalKey(Schema.Number), }); export type RelayT3ChatLinkResponse = typeof RelayT3ChatLinkResponse.Type;