1import { and, count, eq, exists, gt, gte, isNotNull, isNull, ne, not, or, sql } from "drizzle-orm";
2import { QueryBuilder } from "drizzle-orm/pg-core";
3import * as Context from "effect/Context";
4import * as DateTime from "effect/DateTime";
5import * as Effect from "effect/Effect";
6import * as Layer from "effect/Layer";
7import * as Schema from "effect/Schema";
8import * as SqlError from "effect/sql/SqlError";
9import { DEFAULT_MANAGED_TUNNEL_LIMIT, T3CHAT_PLAN_TUNNEL_LIMITS } from "@t3tools/contracts/relay";
10
11import * as RelayDb from "../db.ts";
12import {
13 relayEnvironmentLinks,
14 relayManagedEndpointAllocations,
15 relayManagedTunnelLimits,
16} from "../persistence/schema.ts";
17import * as T3ChatAccounts from "../t3chat/T3ChatAccounts.ts";
18
19// Reserve-to-tunnel takes two Cloudflare calls capped at 8 seconds each; a reservation older than this was abandoned.
20const MANAGED_TUNNEL_PENDING_WINDOW_MINUTES = 15;
21
22export class ManagedTunnelLimitPersistenceError extends Schema.TaggedError<ManagedTunnelLimitPersistenceError>()(
23 "ManagedTunnelLimitPersistenceError",
24 {
25 operation: Schema.Literals(["transaction", "lock-user", "load-limit", "count-tunnels"]),
26 userId: Schema.String,
27 cause: Schema.Defect(),
28 },
29) {
30 override get message(): string {
31 return `Managed tunnel limit '${this.operation}' failed for user '${this.userId}'`;
32 }
33}
34
35export class ManagedTunnelLimitExceeded extends Schema.TaggedError<ManagedTunnelLimitExceeded>()(
36 "ManagedTunnelLimitExceeded",
37 {
38 userId: Schema.String,
39 environmentId: Schema.String,
40 maxTunnels: Schema.Number,
41 activeTunnels: Schema.Number,
42 },
43) {
44 override get message(): string {
45 return `Managed tunnel limit reached for user '${this.userId}': ${this.activeTunnels} of ${this.maxTunnels} tunnels in use`;
46 }
47}
48
49const allocation = relayManagedEndpointAllocations;
50
51const withdrawnLink = (since: typeof allocation.updatedAt) =>
52 exists(
53 new QueryBuilder()
54 .select({ userId: relayEnvironmentLinks.userId })
55 .from(relayEnvironmentLinks)
56 .where(
57 and(
58 eq(relayEnvironmentLinks.userId, allocation.userId),
59 eq(relayEnvironmentLinks.environmentId, allocation.environmentId),
60 or(
61 isNotNull(relayEnvironmentLinks.revokedAt),
62 eq(relayEnvironmentLinks.managedTunnelsEnabled, false),
63 ),
64 gte(relayEnvironmentLinks.updatedAt, since),
65 ),
66 ),
67 );
68
69// Counts allocations, not links: a link is written only after its tunnel exists.
70// `environment` is `ne` for the user's other environments or `eq` for this one, which relinks without counting itself.
71// updatedAt moves only on provisioning writes and release, never on teardown claims, so a link
72// withdrawn at or after it means cleanup is owed; an older withdrawal means a relink that never finished.
73export const countedAllocations = (input: {
74 readonly userId: string;
75 readonly environmentId: string;
76 readonly pendingCutoff: string;
77 readonly environment: typeof eq;
78}) =>
79 and(
80 eq(allocation.userId, input.userId),
81 input.environment(allocation.environmentId, input.environmentId),
82 or(
83 // Reserved but not yet holding a tunnel. Release writes updatedAt = tunnelReleasedAt and
84 // a relink reserves before rewriting its link, so only a newer reservation counts here.
85 and(
86 gt(allocation.updatedAt, input.pendingCutoff),
87 or(
88 isNull(allocation.tunnelReleasedAt),
89 gt(allocation.updatedAt, allocation.tunnelReleasedAt),
90 ),
91 not(withdrawnLink(allocation.updatedAt)),
92 ),
93 and(
94 isNotNull(allocation.tunnelId),
95 isNull(allocation.tunnelReleasedAt),
96 not(withdrawnLink(allocation.updatedAt)),
97 ),
98 ),
99 );
100
101export class ManagedTunnelLimits extends Context.Service<
102 ManagedTunnelLimits,
103 {
104 /** Runs `reserve` only while the user is under their limit, serialized per user in one transaction. */
105 readonly withCapacity: <A, E, R>(
106 input: { readonly userId: string; readonly environmentId: string },
107 reserve: Effect.Effect<A, E, R>,
108 ) => Effect.Effect<A, E | ManagedTunnelLimitExceeded | ManagedTunnelLimitPersistenceError, R>;
109 /** The user's manual cap, which replaces the default and any T3 Chat plan cap. */
110 readonly getOverride: (input: {
111 readonly userId: string;
112 }) => Effect.Effect<number | null, ManagedTunnelLimitPersistenceError>;
113 }
114>()("t3code-relay/environments/ManagedTunnelLimits") {}
115
116const make = Effect.gen(function* () {
117 const db = yield* RelayDb.RelayDb;
118 const t3Chat = yield* T3ChatAccounts.T3ChatAccounts;
119
120 // Asked only at the cap, so most links never call T3 Chat. Any failure keeps the lower cap.
121 const t3ChatTunnelLimit = (userId: string) =>
122 t3Chat.getLink({ userId }).pipe(
123 Effect.map((link) =>
124 link && link.plan !== "free" ? T3CHAT_PLAN_TUNNEL_LIMITS[link.plan] : 0,
125 ),
126 Effect.catch((error) =>
127 Effect.logWarning("T3 Chat plan check failed", { userId, error: error.message }).pipe(
128 Effect.as(0),
129 ),
130 ),
131 );
132
133 const loadOverride = (userId: string) =>
134 db
135 .select({ maxTunnels: relayManagedTunnelLimits.maxTunnels })
136 .from(relayManagedTunnelLimits)
137 .where(eq(relayManagedTunnelLimits.userId, userId))
138 .limit(1)
139 .pipe(
140 Effect.map((rows) => rows[0]?.maxTunnels ?? null),
141 Effect.mapError(
142 (cause) =>
143 new ManagedTunnelLimitPersistenceError({ operation: "load-limit", userId, cause }),
144 ),
145 );
146
147 return ManagedTunnelLimits.of({
148 getOverride: Effect.fn("relay.managed_tunnel_limits.get_override")(function* (input) {
149 return yield* loadOverride(input.userId);
150 }),
151 withCapacity: Effect.fn("relay.managed_tunnel_limits.with_capacity")(function* <A, E, R>(
152 input: { readonly userId: string; readonly environmentId: string },
153 reserve: Effect.Effect<A, E, R>,
154 ): Effect.fn.Return<A, E | ManagedTunnelLimitExceeded | ManagedTunnelLimitPersistenceError, R> {
155 const checkAndReserve = Effect.gen(function* () {
156 // Held until commit, so the count below sees every reservation that took the lock first.
157 yield* db
158 .execute(
159 sql`select pg_advisory_xact_lock(hashtextextended(${`managed-tunnel-limit:${input.userId}`}, 0))`,
160 )
161 .pipe(
162 Effect.mapError(
163 (cause) =>
164 new ManagedTunnelLimitPersistenceError({
165 operation: "lock-user",
166 userId: input.userId,
167 cause,
168 }),
169 ),
170 );
171
172 const override = yield* loadOverride(input.userId);
173 const baseMaxTunnels = override ?? DEFAULT_MANAGED_TUNNEL_LIMIT;
174
175 const pendingCutoff = DateTime.formatIso(
176 DateTime.subtract(yield* DateTime.now, {
177 minutes: MANAGED_TUNNEL_PENDING_WINDOW_MINUTES,
178 }),
179 );
180 const countSlots = (environment: typeof eq) =>
181 db
182 .select({ slots: count() })
183 .from(allocation)
184 .where(countedAllocations({ ...input, pendingCutoff, environment }))
185 .pipe(
186 Effect.map((rows) => rows[0]?.slots ?? 0),
187 Effect.mapError(
188 (cause) =>
189 new ManagedTunnelLimitPersistenceError({
190 operation: "count-tunnels",
191 userId: input.userId,
192 cause,
193 }),
194 ),
195 );
196 const activeTunnels = yield* countSlots(ne);
197 if (activeTunnels < baseMaxTunnels) return yield* reserve;
198
199 // An environment that already holds a counted slot relinks without asking T3 Chat,
200 // so a downgrade or a T3 Chat outage only blocks new environments.
201 if ((yield* countSlots(eq)) > 0) return yield* reserve;
202
203 // A manual override is the exact cap; a T3 Chat plan raises only the default.
204 // Runs under the user lock, bounded by the T3 Chat request timeout.
205 const maxTunnels =
206 t3Chat.available && override === null
207 ? Math.max(baseMaxTunnels, yield* t3ChatTunnelLimit(input.userId))
208 : baseMaxTunnels;
209 if (activeTunnels >= maxTunnels) {
210 return yield* new ManagedTunnelLimitExceeded({
211 userId: input.userId,
212 environmentId: input.environmentId,
213 maxTunnels,
214 activeTunnels,
215 });
216 }
217 return yield* reserve;
218 });
219
220 // Statements inside map their own SqlErrors, so any left here come from opening or committing the transaction.
221 return yield* db.$client.withTransaction(checkAndReserve).pipe(
222 Effect.mapError((cause) =>
223 SqlError.isSqlError(cause)
224 ? new ManagedTunnelLimitPersistenceError({
225 operation: "transaction",
226 userId: input.userId,
227 cause,
228 })
229 : cause,
230 ),
231 );
232 }),
233 });
234});
235
236export const layer = Layer.effect(ManagedTunnelLimits, make);