1import { describe, expect, it } from "@effect/vitest";
2import * as NodeServices from "@effect/platform-node/NodeServices";
3import * as Context from "effect/Context";
4import * as Effect from "effect/Effect";
5import * as FileSystem from "effect/FileSystem";
6import * as Layer from "effect/Layer";
7import * as Path from "effect/Path";
8import * as Semaphore from "effect/Semaphore";
9import { PGlite } from "@electric-sql/pglite";
10import { drizzle } from "drizzle-orm/pglite";
11import { ne } from "drizzle-orm";
12import { PgDialect } from "drizzle-orm/pg-core";
13import { DEFAULT_MANAGED_TUNNEL_LIMIT, type RelayT3ChatPlan } from "@t3tools/contracts/relay";
14
15import * as RelayDb from "../db.ts";
16import { relayManagedEndpointAllocations } from "../persistence/schema.ts";
17import * as T3ChatAccounts from "../t3chat/T3ChatAccounts.ts";
18import * as ManagedTunnelLimits from "./ManagedTunnelLimits.ts";
19
20const BACKFILL_MIGRATION = "20261028062605_backfill_withdrawn_tunnel_allocations";
21
22class FakeTransaction extends Context.Service<FakeTransaction, { holdsLock: boolean }>()(
23 "t3code-relay/environments/ManagedTunnelLimits.test/FakeTransaction",
24) {}
25
26interface FakeDbInput {
27 readonly otherSlots: () => number;
28 readonly ownSlot?: boolean;
29 readonly override?: number;
30}
31
32// Models Postgres just enough: the advisory lock only works inside a transaction and is held until it ends.
33function makeFakeDb(input: FakeDbInput) {
34 const userLock = Semaphore.makeUnsafe(1);
35 const dialect = new PgDialect();
36 return {
37 $client: {
38 withTransaction: <A, E, R>(effect: Effect.Effect<A, E, R>) => {
39 const transaction = { holdsLock: false };
40 return effect.pipe(
41 Effect.provideService(FakeTransaction, transaction),
42 Effect.ensuring(
43 Effect.suspend(() => (transaction.holdsLock ? userLock.release(1) : Effect.void)),
44 ),
45 );
46 },
47 },
48 execute: () =>
49 Effect.gen(function* () {
50 const transaction = yield* Effect.serviceOption(FakeTransaction);
51 if (transaction._tag === "None") {
52 return yield* Effect.die("advisory lock taken outside a transaction");
53 }
54 yield* userLock.take(1);
55 transaction.value.holdsLock = true;
56 return [];
57 }),
58 select: () => ({
59 from: (table: unknown) => ({
60 where: (condition: unknown) => {
61 if (table !== relayManagedEndpointAllocations) {
62 return {
63 limit: () =>
64 Effect.succeed(
65 input.override === undefined ? [] : [{ maxTunnels: input.override }],
66 ),
67 };
68 }
69 const countsOthers = dialect
70 .sqlToQuery(condition as never)
71 .sql.includes('"relay_managed_endpoint_allocations"."environment_id" <>');
72 return Effect.yieldNow.pipe(
73 Effect.andThen(
74 Effect.sync(() => [
75 { slots: countsOthers ? input.otherSlots() : input.ownSlot ? 1 : 0 },
76 ]),
77 ),
78 );
79 },
80 }),
81 }),
82 } as unknown as RelayDb.RelayDb["Service"];
83}
84
85const t3ChatLayer = (plan: RelayT3ChatPlan | "lookup-fails", lookups?: Array<string>) =>
86 Layer.mock(T3ChatAccounts.T3ChatAccounts, {
87 available: true,
88 getLink: ({ userId }) => {
89 lookups?.push(userId);
90 return plan === "lookup-fails"
91 ? Effect.fail(
92 new T3ChatAccounts.T3ChatRequestFailed({
93 operation: "get-link",
94 cause: new Error("down"),
95 }),
96 )
97 : Effect.succeed({ email: "theo@example.com", plan });
98 },
99 });
100
101const reserve = (
102 input: Omit<FakeDbInput, "otherSlots"> & {
103 readonly activeTunnels: number;
104 readonly plan: RelayT3ChatPlan | "lookup-fails";
105 readonly lookups?: Array<string>;
106 },
107) =>
108 Effect.gen(function* () {
109 const limits = yield* ManagedTunnelLimits.ManagedTunnelLimits;
110 return yield* limits.withCapacity({ userId: "user-1", environmentId: "env-new" }, Effect.void);
111 }).pipe(
112 Effect.provide(
113 ManagedTunnelLimits.layer.pipe(
114 Layer.provide(
115 Layer.succeed(
116 RelayDb.RelayDb,
117 makeFakeDb({ ...input, otherSlots: () => input.activeTunnels }),
118 ),
119 ),
120 Layer.provide(t3ChatLayer(input.plan, input.lookups)),
121 ),
122 ),
123 );
124
125const getOverride = (override?: number) =>
126 Effect.gen(function* () {
127 const limits = yield* ManagedTunnelLimits.ManagedTunnelLimits;
128 return yield* limits.getOverride({ userId: "user-1" });
129 }).pipe(
130 Effect.provide(
131 ManagedTunnelLimits.layer.pipe(
132 Layer.provide(
133 Layer.succeed(
134 RelayDb.RelayDb,
135 makeFakeDb({
136 otherSlots: () => 0,
137 ...(override === undefined ? {} : { override }),
138 }),
139 ),
140 ),
141 Layer.provide(t3ChatLayer("free")),
142 ),
143 ),
144 );
145
146describe("ManagedTunnelLimits", () => {
147 it.effect("rejects concurrent reservations past the limit", () => {
148 const reserved: Array<string> = [];
149 const attempts = DEFAULT_MANAGED_TUNNEL_LIMIT + 2;
150
151 return Effect.gen(function* () {
152 const limits = yield* ManagedTunnelLimits.ManagedTunnelLimits;
153 const results = yield* Effect.forEach(
154 Array.from({ length: attempts }, (_, index) => `environment-${index}`),
155 (environmentId) =>
156 limits
157 .withCapacity(
158 { userId: "user-1", environmentId },
159 Effect.yieldNow.pipe(Effect.andThen(Effect.sync(() => reserved.push(environmentId)))),
160 )
161 .pipe(Effect.result),
162 { concurrency: "unbounded" },
163 );
164
165 expect(reserved).toHaveLength(DEFAULT_MANAGED_TUNNEL_LIMIT);
166 expect(
167 results.filter(
168 (result) =>
169 result._tag === "Failure" && result.failure._tag === "ManagedTunnelLimitExceeded",
170 ),
171 ).toHaveLength(2);
172 }).pipe(
173 Effect.provide(
174 ManagedTunnelLimits.layer.pipe(
175 Layer.provide(
176 Layer.succeed(RelayDb.RelayDb, makeFakeDb({ otherSlots: () => reserved.length })),
177 ),
178 Layer.provide(t3ChatLayer("free")),
179 ),
180 ),
181 );
182 });
183
184 it.effect("counts each allocation by its tunnel and link state", () =>
185 Effect.gen(function* () {
186 const fs = yield* FileSystem.FileSystem;
187 const pg = yield* Effect.acquireRelease(
188 Effect.promise(() => PGlite.create()),
189 (client) => Effect.promise(() => client.close()),
190 );
191 const path = yield* Path.Path;
192 const migrationsDir = yield* path.fromFileUrl(
193 new URL("../../migrations/postgres/", import.meta.url),
194 );
195 const migrations = (yield* fs.readDirectory(migrationsDir)).toSorted();
196 const applyMigrations = (names: ReadonlyArray<string>) =>
197 Effect.forEach(names, (name) =>
198 fs
199 .readFileString(`${migrationsDir}${name}/migration.sql`)
200 .pipe(Effect.flatMap((migration) => Effect.promise(() => pg.exec(migration)))),
201 );
202 const backfill = migrations.indexOf(BACKFILL_MIGRATION);
203 expect(backfill).toBeGreaterThan(-1);
204 yield* applyMigrations(migrations.slice(0, backfill));
205
206 const at = (minutes: number) => `2026-10-27T12:${String(minutes).padStart(2, "0")}:00.000Z`;
207 const cases: ReadonlyArray<{
208 readonly environmentId: string;
209 readonly counted: boolean;
210 readonly allocation: {
211 readonly tunnelId?: string;
212 readonly releasedAt?: number;
213 readonly updatedAt: number;
214 };
215 readonly link?: {
216 readonly revoked?: boolean;
217 readonly publishOnly?: boolean;
218 readonly updatedAt: number;
219 };
220 }> = [
221 { environmentId: "pending", counted: true, allocation: { updatedAt: 50 } },
222 { environmentId: "abandoned-reservation", counted: false, allocation: { updatedAt: 10 } },
223 {
224 environmentId: "live",
225 counted: true,
226 allocation: { tunnelId: "t", updatedAt: 10 },
227 link: { updatedAt: 12 },
228 },
229 {
230 environmentId: "released",
231 counted: false,
232 allocation: { tunnelId: "t", releasedAt: 20, updatedAt: 20 },
233 },
234 {
235 environmentId: "re-reserved-after-release",
236 counted: true,
237 allocation: { tunnelId: "t", releasedAt: 20, updatedAt: 50 },
238 },
239 {
240 environmentId: "revoked",
241 counted: false,
242 allocation: { tunnelId: "t", updatedAt: 10 },
243 link: { revoked: true, updatedAt: 20 },
244 },
245 {
246 environmentId: "revoked-pending",
247 counted: false,
248 allocation: { updatedAt: 50 },
249 link: { revoked: true, updatedAt: 55 },
250 },
251 {
252 environmentId: "publish-only",
253 counted: false,
254 allocation: { tunnelId: "t", updatedAt: 10 },
255 link: { publishOnly: true, updatedAt: 20 },
256 },
257 // Reprovisioned after a revoke, but the Worker died before rewriting the link.
258 {
259 environmentId: "reprovision-before-link",
260 counted: true,
261 allocation: { tunnelId: "t", updatedAt: 20 },
262 link: { revoked: true, updatedAt: 10 },
263 },
264 {
265 environmentId: "reserved-before-link-rewrite",
266 counted: true,
267 allocation: { updatedAt: 50 },
268 link: { revoked: true, updatedAt: 48 },
269 },
270 // Tunnel deleted but the row survived; deprovision marked it released, then a relink stalled.
271 {
272 environmentId: "released-after-failed-removal-then-reserved",
273 counted: false,
274 allocation: { tunnelId: "t", releasedAt: 20, updatedAt: 30 },
275 link: { revoked: true, updatedAt: 15 },
276 },
277 {
278 environmentId: "withdrawn-at-provisioning-write",
279 counted: false,
280 allocation: { tunnelId: "t", updatedAt: 20 },
281 link: { revoked: true, updatedAt: 20 },
282 },
283 { environmentId: "self", counted: false, allocation: { tunnelId: "t", updatedAt: 10 } },
284 ];
285 // Written before the deploy: claimDeprovision bumped updatedAt past the revoke, then cleanup failed.
286 const legacyCases: typeof cases = [
287 {
288 environmentId: "legacy-failed-deprovision",
289 counted: false,
290 allocation: { tunnelId: "t", updatedAt: 30 },
291 link: { revoked: true, updatedAt: 20 },
292 },
293 {
294 environmentId: "legacy-failed-publish-only-deprovision",
295 counted: false,
296 allocation: { tunnelId: "t", updatedAt: 30 },
297 link: { publishOnly: true, updatedAt: 20 },
298 },
299 {
300 environmentId: "legacy-live",
301 counted: true,
302 allocation: { tunnelId: "t", updatedAt: 30 },
303 link: { updatedAt: 20 },
304 },
305 ];
306 const insert = (entries: typeof cases) =>
307 Effect.forEach(entries, (entry) =>
308 Effect.promise(async () => {
309 await pg.query(
310 `insert into relay_managed_endpoint_allocations
311 (user_id, environment_id, hostname, tunnel_id, tunnel_name, tunnel_released_at, created_at, updated_at)
312 values ('user-1', $1::text, $1::text, $2, $1::text, $3, $4, $4)`,
313 [
314 entry.environmentId,
315 entry.allocation.tunnelId ?? null,
316 entry.allocation.releasedAt === undefined ? null : at(entry.allocation.releasedAt),
317 at(entry.allocation.updatedAt),
318 ],
319 );
320 const link = entry.link;
321 if (link !== undefined) {
322 await pg.query(
323 `insert into relay_environment_links
324 (user_id, environment_id, environment_public_key, endpoint_http_base_url,
325 endpoint_ws_base_url, endpoint_provider_kind, revoked_at, managed_tunnels_enabled,
326 created_at, updated_at)
327 values ('user-1', $1, 'key', 'https://e', 'wss://e', 'cloudflare_tunnel', $2, $3, $4, $4)`,
328 [
329 entry.environmentId,
330 link.revoked === true ? at(link.updatedAt) : null,
331 link.publishOnly !== true,
332 at(link.updatedAt),
333 ],
334 );
335 }
336 }),
337 );
338 yield* insert(legacyCases);
339 yield* applyMigrations(migrations.slice(backfill));
340 yield* insert(cases);
341 yield* Effect.promise(() =>
342 pg.query(
343 `insert into relay_managed_endpoint_allocations
344 (user_id, environment_id, hostname, tunnel_id, tunnel_name, created_at, updated_at)
345 values ('user-2', 'other-user', 'other-user', 't', 'other-user', $1, $1)`,
346 [at(10)],
347 ),
348 );
349
350 const counted = yield* Effect.promise(() =>
351 drizzle({ client: pg })
352 .select({ environmentId: relayManagedEndpointAllocations.environmentId })
353 .from(relayManagedEndpointAllocations)
354 .where(
355 ManagedTunnelLimits.countedAllocations({
356 userId: "user-1",
357 environmentId: "self",
358 pendingCutoff: at(45),
359 environment: ne,
360 }),
361 ),
362 );
363
364 expect(counted.map((row) => row.environmentId).toSorted()).toEqual(
365 [...legacyCases, ...cases]
366 .filter((entry) => entry.counted)
367 .map((entry) => entry.environmentId)
368 .toSorted(),
369 );
370 }).pipe(Effect.provide(NodeServices.layer)),
371 );
372
373 it.effect("does not ask T3 Chat while the account is under its cap", () =>
374 Effect.gen(function* () {
375 const lookups: Array<string> = [];
376 yield* reserve({ activeTunnels: 2, plan: "free", lookups });
377 expect(lookups).toEqual([]);
378 }),
379 );
380
381 it.effect("raises the cap to the linked T3 Chat plan's limit", () =>
382 Effect.gen(function* () {
383 for (const [plan, limit] of [
384 ["pro", 10],
385 ["premier", 100],
386 ] as const) {
387 yield* reserve({ activeTunnels: limit - 1, plan });
388 const error = yield* Effect.flip(reserve({ activeTunnels: limit, plan }));
389 expect(error).toMatchObject({ _tag: "ManagedTunnelLimitExceeded", maxTunnels: limit });
390 }
391 }),
392 );
393
394 it.effect("keeps the default cap for free plans and failed lookups", () =>
395 Effect.gen(function* () {
396 for (const plan of ["free", "lookup-fails"] as const) {
397 const error = yield* Effect.flip(reserve({ activeTunnels: 3, plan }));
398 expect(error).toMatchObject({ _tag: "ManagedTunnelLimitExceeded", maxTunnels: 3 });
399 }
400 }),
401 );
402
403 it.effect("treats a manual override as the exact cap", () =>
404 Effect.gen(function* () {
405 const lookups: Array<string> = [];
406 for (const override of [1, 200]) {
407 const error = yield* Effect.flip(
408 reserve({ override, activeTunnels: override, plan: "premier", lookups }),
409 );
410 expect(error).toMatchObject({ _tag: "ManagedTunnelLimitExceeded", maxTunnels: override });
411 expect(yield* getOverride(override)).toBe(override);
412 }
413 expect(yield* getOverride()).toBeNull();
414 expect(lookups).toEqual([]);
415 }),
416 );
417
418 it.effect("lets an environment holding a slot relink without asking T3 Chat", () =>
419 Effect.gen(function* () {
420 const lookups: Array<string> = [];
421 yield* reserve({ activeTunnels: 10, plan: "lookup-fails", ownSlot: true, lookups });
422 expect(lookups).toEqual([]);
423 }),
424 );
425});