fix(runtime): close Worker database pools per request

src/lib/adapters/cloudflare-runtime.ts

@@ -117,6 +117,7 @@ type CloudflareRuntimeEnv = Record<string, unknown> & {
 type CloudflareRuntimeContext = {
   cache: Map<symbol, unknown>;
   cacheStorage?: CloudflareCacheStorage;
+  cleanups: Set<() => Promise<void> | void>;
   env?: CloudflareRuntimeEnv;
   request?: CloudflareRequestContext;
   scheduleTask?: CloudflareTaskScheduler;
@@ -177,6 +178,61 @@ function getCurrentCloudflareRuntimeEnv() {
   return globalForCloudflareRuntime.__lifeUstcCloudflareRuntimeEnv;
 }
 
+async function cleanupCloudflareRuntimeContext(
+  context: CloudflareRuntimeContext,
+) {
+  const cleanupResults = await Promise.allSettled(
+    [...context.cleanups].map((cleanup) => Promise.resolve().then(cleanup)),
+  );
+  context.cache.clear();
+  context.cleanups.clear();
+  const failures = cleanupResults
+    .filter(
+      (cleanupResult): cleanupResult is PromiseRejectedResult =>
+        cleanupResult.status === "rejected",
+    )
+    .map((cleanupResult) => cleanupResult.reason);
+  if (failures.length === 1) throw failures[0];
+  if (failures.length > 1) {
+    throw new AggregateError(failures, "Cloudflare runtime cleanup failed");
+  }
+}
+
+function responseWithRuntimeCleanup(
+  response: Response,
+  cleanup: () => Promise<void>,
+) {
+  if (!response.body) return response;
+  const reader = response.body.getReader();
+  const body = new ReadableStream<Uint8Array>(
+    {
+      async pull(controller) {
+        try {
+          const chunk = await reader.read();
+          if (!chunk.done) {
+            controller.enqueue(chunk.value);
+            return;
+          }
+          await cleanup();
+          controller.close();
+        } catch (error) {
+          await cleanup().catch(() => undefined);
+          controller.error(error);
+        }
+      },
+      async cancel(reason) {
+        try {
+          await reader.cancel(reason);
+        } finally {
+          await cleanup();
+        }
+      },
+    },
+    { highWaterMark: 0 },
+  );
+  return new Response(body, response);
+}
+
 export function runWithCloudflareRuntimeEnv<T>(
   env: unknown,
   callback: () => T | Promise<T>,
@@ -195,20 +251,43 @@ export function runWithCloudflareRuntimeEnv<T>(
   const context: CloudflareRuntimeContext = {
     cache: new Map(),
     cacheStorage: normalizeCloudflareCacheStorage(),
+    cleanups: new Set(),
     env: normalizeCloudflareRuntimeEnv(env),
     scheduleTask: normalizeCloudflareTaskScheduler(executionContext),
     tracing,
   };
 
   return cloudflareRuntimeStorage.run(context, async () => {
+    let cleanupPromise: Promise<void> | undefined;
+    const cleanup = () => {
+      cleanupPromise ??= cleanupCloudflareRuntimeContext(context);
+      return cleanupPromise;
+    };
+    let result: T;
     try {
-      return await callback();
-    } finally {
-      context.cache.clear();
+      result = await callback();
+    } catch (error) {
+      await cleanup().catch(() => undefined);
+      throw error;
+    }
+    if (
+      result instanceof Response &&
+      result.body &&
+      context.cleanups.size > 0
+    ) {
+      return responseWithRuntimeCleanup(result, cleanup) as T;
     }
+    await cleanup();
+    return result;
   });
 }
 
+export function registerCloudflareRuntimeCleanup(
+  cleanup: () => Promise<void> | void,
+) {
+  cloudflareRuntimeStorage.getStore()?.cleanups.add(cleanup);
+}
+
 export function runCloudflareTraceSpan<T>(
   name: string,
   attributes: Record<string, boolean | number | string | undefined>,

src/lib/db/auth-prisma.ts

@@ -3,6 +3,7 @@ import type { PrismaClient } from "@/generated/prisma/client";
 import {
   getCloudflareRuntimeContext,
   hasCloudflareRuntimeEnv,
+  registerCloudflareRuntimeCleanup,
 } from "@/lib/adapters/cloudflare-runtime";
 import { createBasePrisma, logPrismaQuery } from "@/lib/db/prisma-query-events";
 import { shouldEnablePrismaQueryLogging } from "@/lib/db/prisma-query-logging";
@@ -31,6 +32,7 @@ function getBaseAuthPrisma() {
         | undefined;
       if (cached) return cached;
       const client = createAuthPrismaClient();
+      registerCloudflareRuntimeCleanup(() => client.$disconnect());
       cache.set(cloudflareAuthPrismaCacheKey, client);
       return client;
     }

src/lib/db/maintenance-prisma.ts

@@ -3,6 +3,7 @@ import type { PrismaClient } from "@/generated/prisma/client";
 import {
   getCloudflareRuntimeContext,
   hasCloudflareRuntimeEnv,
+  registerCloudflareRuntimeCleanup,
 } from "@/lib/adapters/cloudflare-runtime";
 import { createBasePrisma, logPrismaQuery } from "@/lib/db/prisma-query-events";
 import { shouldEnablePrismaQueryLogging } from "@/lib/db/prisma-query-logging";
@@ -33,6 +34,7 @@ function getBaseMaintenancePrisma() {
         | undefined;
       if (cached) return cached;
       const client = createMaintenancePrismaClient();
+      registerCloudflareRuntimeCleanup(() => client.$disconnect());
       cache.set(cloudflareMaintenancePrismaCacheKey, client);
       return client;
     }

src/lib/db/prisma-adapter.ts

@@ -79,14 +79,13 @@ export function createPrismaAdapter(
     {
       connectionString: resolvedConnectionString,
       // On Workers every request builds a fresh pool (pg sockets cannot be
-      // reused across requests) and each new connection pays full SCRAM
-      // deriveBits CPU. Concurrent queries (Promise.all, RLS tx + session
-      // lookup) otherwise open up to pg's default of 10 connections per
-      // request; cap the pool so a single request can never open more than 3.
+      // reused across requests), which the runtime context disconnects before
+      // the request completes. Concurrent queries (Promise.all, RLS tx +
+      // session lookup) otherwise open up to pg's default of 10 connections;
+      // cap the pool so a single request can never open more than 3.
       max: 3,
-      // Idle connections from a finished request are never reusable, so close
-      // them quickly instead of holding sockets (and server slots) for pg's
-      // 10s default.
+      // Keep a short idle timeout as a safety net for clients created outside
+      // the managed request context.
       idleTimeoutMillis: 5_000,
     },
     {

src/lib/db/prisma.ts

@@ -3,6 +3,7 @@ import type { PrismaClient } from "@/generated/prisma/client";
 import {
   getCloudflareRuntimeContext,
   hasCloudflareRuntimeEnv,
+  registerCloudflareRuntimeCleanup,
 } from "@/lib/adapters/cloudflare-runtime";
 import { localizedNamesExtension } from "@/lib/db/prisma-localized-names";
 import { createBasePrisma, logPrismaQuery } from "@/lib/db/prisma-query-events";
@@ -67,11 +68,16 @@ function getBasePrisma() {
   if (hasCloudflareRuntimeEnv()) {
     const cache = getCloudflarePrismaCache();
     if (cache) {
-      cache.base ??= createPrismaClient(cache);
+      if (!cache.base) {
+        cache.base = createPrismaClient(cache);
+        registerCloudflareRuntimeCleanup(() => cache.base?.$disconnect());
+      }
       return cache.base;
     }
 
-    return createPrismaClient();
+    const client = createPrismaClient();
+    registerCloudflareRuntimeCleanup(() => client.$disconnect());
+    return client;
   }
 
   const cached = globalForPrisma.prisma ?? basePrisma;

tests/unit/auth-prisma-boundary.test.ts

@@ -1,15 +1,21 @@
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import {
   getCloudflareAuthHyperdriveConnectionString,
   runWithCloudflareRuntimeEnv,
 } from "@/lib/adapters/cloudflare-runtime";
 
 const { appClient, createBasePrismaMock, firstClient, secondClient } =
   vi.hoisted(() => ({
-    appClient: { user: { boundary: "app" } },
+    appClient: { $disconnect: vi.fn(), user: { boundary: "app" } },
     createBasePrismaMock: vi.fn(),
-    firstClient: { user: { boundary: "first-auth" } },
-    secondClient: { user: { boundary: "second-auth" } },
+    firstClient: {
+      $disconnect: vi.fn(),
+      user: { boundary: "first-auth" },
+    },
+    secondClient: {
+      $disconnect: vi.fn(),
+      user: { boundary: "second-auth" },
+    },
   }));
 
 vi.mock("@/lib/db/prisma-query-events", () => ({
@@ -22,6 +28,12 @@ vi.mock("@/lib/db/prisma-query-logging", () => ({
 }));
 
 describe("auth Prisma boundary", () => {
+  beforeEach(() => {
+    appClient.$disconnect.mockClear();
+    firstClient.$disconnect.mockClear();
+    secondClient.$disconnect.mockClear();
+  });
+
   it("keeps overlapping Cloudflare requests on their own auth clients", async () => {
     createBasePrismaMock.mockReset().mockImplementation((_url, database) => {
       if (database !== "auth") return appClient;
@@ -65,6 +77,8 @@ describe("auth Prisma boundary", () => {
 
     expect(createBasePrismaMock).toHaveBeenCalledTimes(2);
     expect(createBasePrismaMock).toHaveBeenCalledWith(undefined, "auth");
+    expect(firstClient.$disconnect).toHaveBeenCalledOnce();
+    expect(secondClient.$disconnect).toHaveBeenCalledOnce();
   });
 
   it("keeps app and auth clients distinct inside one request", async () => {
@@ -94,5 +108,7 @@ describe("auth Prisma boundary", () => {
     expect(createBasePrismaMock).toHaveBeenCalledTimes(2);
     expect(createBasePrismaMock).toHaveBeenCalledWith();
     expect(createBasePrismaMock).toHaveBeenCalledWith(undefined, "auth");
+    expect(appClient.$disconnect).toHaveBeenCalledOnce();
+    expect(firstClient.$disconnect).toHaveBeenCalledOnce();
   });
 });

tests/unit/cloudflare-runtime-tracing.test.ts

@@ -2,6 +2,7 @@ import { afterEach, describe, expect, it, vi } from "vitest";
 import {
   getCloudflareNamedCache,
   getCloudflareRuntimeTaskScheduler,
+  registerCloudflareRuntimeCleanup,
   runCloudflareTraceSpan,
   runWithCloudflareRuntimeEnv,
 } from "@/lib/adapters/cloudflare-runtime";
@@ -148,4 +149,58 @@ describe("Cloudflare runtime tracing", () => {
     expect(getCloudflareNamedCache("outside-request")).toBeUndefined();
     expect(getCloudflareRuntimeTaskScheduler()).toBeUndefined();
   });
+
+  it("awaits request-scoped cleanup before resolving", async () => {
+    const events: string[] = [];
+
+    await runWithCloudflareRuntimeEnv({}, async () => {
+      registerCloudflareRuntimeCleanup(async () => {
+        await Promise.resolve();
+        events.push("cleanup");
+      });
+      events.push("callback");
+    });
+
+    expect(events).toEqual(["callback", "cleanup"]);
+  });
+
+  it("defers cleanup until a response body finishes streaming", async () => {
+    const cleanup = vi.fn();
+
+    const response = await runWithCloudflareRuntimeEnv({}, () => {
+      registerCloudflareRuntimeCleanup(cleanup);
+      return new Response("streamed");
+    });
+
+    expect(cleanup).not.toHaveBeenCalled();
+    await expect(response.text()).resolves.toBe("streamed");
+    expect(cleanup).toHaveBeenCalledOnce();
+  });
+
+  it("cleans up when a response body is canceled", async () => {
+    const cleanup = vi.fn();
+    const cancel = vi.fn();
+
+    const response = await runWithCloudflareRuntimeEnv({}, () => {
+      registerCloudflareRuntimeCleanup(cleanup);
+      return new Response(new ReadableStream({ cancel }));
+    });
+
+    await response.body?.cancel("client disconnected");
+    expect(cancel).toHaveBeenCalledWith("client disconnected");
+    expect(cleanup).toHaveBeenCalledOnce();
+  });
+
+  it("preserves callback errors when cleanup also fails", async () => {
+    const callbackFailure = new Error("callback failed");
+
+    await expect(
+      runWithCloudflareRuntimeEnv({}, () => {
+        registerCloudflareRuntimeCleanup(() => {
+          throw new Error("cleanup failed");
+        });
+        throw callbackFailure;
+      }),
+    ).rejects.toBe(callbackFailure);
+  });
 });

tests/unit/prisma-rls-context.test.ts

@@ -17,6 +17,7 @@ const { baseClient, extendedClient, todoFindManyMock } = vi.hoisted(() => {
   };
   return {
     baseClient: {
+      $disconnect: vi.fn(),
       $extends: vi.fn(() => extended),
     },
     extendedClient: extended,
@@ -92,6 +93,7 @@ describe("localized Prisma clients in RLS context", () => {
         );
       });
     });
+    expect(baseClient.$disconnect).toHaveBeenCalledOnce();
   });
 
   it("blocks saved localized clients, delegates, and methods inside RLS context", async () => {