diag(ai-gateway): measure whether pool slots are leaked

apps/web/src/app/api/internal/usage/record/route.ts

@@ -14,6 +14,7 @@ import {
 import {
   createPhaseTimer,
   emitUsageRecordTiming,
+  noteRequestStart,
   readPoolGauges,
   shouldEmitUsageRecordTiming,
 } from '@/lib/ai-gateway/usage-record-diagnostics';
@@ -48,7 +49,11 @@ export async function POST(request: NextRequest): Promise<NextResponse> {
   // pairing them with the in-process pool gauges distinguishes waiting for a pool
   // connection from waiting for PostgreSQL from a stalled event loop.
   const timer = createPhaseTimer();
+  // Read before this request acquires anything: on an instance that has been quiet
+  // longer than `idleTimeoutMillis`, anything still checked out belongs to nobody.
+  const msSinceLastRequest = noteRequestStart();
   const poolBefore = readPoolGauges();
+  const checkedOutAtEntry = poolBefore.total - poolBefore.idle;
   let poolWaitingPeak = poolBefore.waiting;
   const samplePool = () => {
     const gauges = readPoolGauges();
@@ -57,7 +62,7 @@ export async function POST(request: NextRequest): Promise<NextResponse> {
   };
   const reportTiming = (usageId: string, outcome: UsageRecordResponse['status']) => {
     const totalMs = timer.totalMs();
-    if (!shouldEmitUsageRecordTiming(totalMs)) return;
+    if (!shouldEmitUsageRecordTiming(totalMs, msSinceLastRequest)) return;
     emitUsageRecordTiming({
       usageId,
       outcome,
@@ -66,6 +71,8 @@ export async function POST(request: NextRequest): Promise<NextResponse> {
       poolBefore,
       poolAfter: samplePool(),
       poolWaitingPeak,
+      msSinceLastRequest,
+      checkedOutAtEntry,
     });
   };
 

apps/web/src/lib/ai-gateway/processUsage.ts

@@ -6,6 +6,7 @@ import {
   isUsageRowConflict,
   stackFramesUnderHeader,
 } from './usage-record-diagnostics';
+import { isTransactionBeginFailure, recordTransactionBeginFailure } from '@/lib/db-pool-leak-probe';
 import type { MicrodollarUsage } from '@kilocode/db/schema';
 import { microdollar_usage } from '@kilocode/db/schema';
 import { createTimer } from '@/lib/timer';
@@ -689,6 +690,12 @@ export async function insertUsageRecord(
             // retrying — `id` is fixed for the delivery — so stop immediately and
             // let the outer handler recover the committed row's identity. Retrying
             // it burned three attempts and rebuilt the statement each time.
+            // Count a failed `BEGIN` before deciding what to do with it. Under
+            // drizzle 0.45.2 the `BEGIN` is issued outside the `try`/`finally`, so
+            // this failure never reaches `release()` and burns a pool slot for the
+            // life of the process. The count is what makes the leak measurable
+            // against the pool's low-water mark.
+            if (isTransactionBeginFailure(error)) recordTransactionBeginFailure();
             if (attempt >= 2 || isUsageRowConflict(error)) throw error;
             // Never log the raw error: its message is the interpolated statement,
             // roughly 30KB including prompt prefixes and the client IP.

apps/web/src/lib/ai-gateway/usage-record-diagnostics.test.ts

@@ -2,15 +2,17 @@ jest.mock('@/lib/drizzle', () => ({
   pool: { totalCount: 4, idleCount: 1, waitingCount: 0, options: { max: 10 } },
 }));
 
-import { describe, expect, test } from '@jest/globals';
+import { beforeEach, describe, expect, test } from '@jest/globals';
 
 import {
   createPhaseTimer,
   describeDatabaseError,
   isUsageRowConflict,
+  noteRequestStart,
   readPoolGauges,
   shouldEmitUsageRecordTiming,
   stackFramesUnderHeader,
+  __resetRequestClockForTest,
 } from './usage-record-diagnostics';
 
 /** Shape of the drizzle wrapper: a message carrying the statement, plus a cause. */
@@ -254,12 +256,41 @@ describe('readPoolGauges', () => {
 
 describe('shouldEmitUsageRecordTiming', () => {
   test('always emits at or above the slow threshold', () => {
-    expect(shouldEmitUsageRecordTiming(1_000, () => 1)).toBe(true);
-    expect(shouldEmitUsageRecordTiming(50_000, () => 1)).toBe(true);
+    expect(shouldEmitUsageRecordTiming(1_000, null, () => 1)).toBe(true);
+    expect(shouldEmitUsageRecordTiming(50_000, null, () => 1)).toBe(true);
   });
 
   test('samples a small fraction of fast requests for a baseline', () => {
-    expect(shouldEmitUsageRecordTiming(5, () => 0.005)).toBe(true);
-    expect(shouldEmitUsageRecordTiming(5, () => 0.5)).toBe(false);
+    expect(shouldEmitUsageRecordTiming(5, null, () => 0.005)).toBe(true);
+    expect(shouldEmitUsageRecordTiming(5, null, () => 0.5)).toBe(false);
+  });
+
+  // The leak observation: a request landing on a quiet instance is fast, so the
+  // duration threshold would never surface it, yet it is the only moment when a
+  // leaked slot is distinguishable from a busy one.
+  test('always emits the first request after the instance has been quiet', () => {
+    expect(shouldEmitUsageRecordTiming(5, 5_000, () => 1)).toBe(true);
+    expect(shouldEmitUsageRecordTiming(5, 60_000, () => 1)).toBe(true);
+  });
+
+  test('does not treat a busy instance as an observation point', () => {
+    expect(shouldEmitUsageRecordTiming(5, 4_999, () => 1)).toBe(false);
+    expect(shouldEmitUsageRecordTiming(5, 0, () => 1)).toBe(false);
+  });
+});
+
+describe('noteRequestStart', () => {
+  beforeEach(() => {
+    __resetRequestClockForTest();
+  });
+
+  test('reports null for the first request on an instance', () => {
+    expect(noteRequestStart(1_000)).toBeNull();
+  });
+
+  test('reports the gap since the previous request', () => {
+    noteRequestStart(1_000);
+    expect(noteRequestStart(7_500)).toBe(6_500);
+    expect(noteRequestStart(7_600)).toBe(100);
   });
 });

apps/web/src/lib/ai-gateway/usage-record-diagnostics.ts

@@ -2,6 +2,7 @@ import 'server-only';
 
 import { monitorEventLoopDelay } from 'node:perf_hooks';
 import { pool } from '@/lib/drizzle';
+import { readPoolLeakStats } from '@/lib/db-pool-leak-probe';
 
 /**
  * Diagnostics for the Frankfurt-local usage write.
@@ -95,13 +96,54 @@ export type UsageRecordTiming = {
   poolAfter: PoolGauges;
   /** Highest `waitingCount` sampled during the request. */
   poolWaitingPeak: number;
+  /**
+   * Milliseconds since the previous request on this instance, `null` for the
+   * first. Above `QUIET_INSTANCE_MS` the pool should have drained, so a non-zero
+   * `checked_out` in the same line is a leaked slot.
+   */
+  msSinceLastRequest: number | null;
+  /**
+   * Checked-out connections observed at handler entry, before this request
+   * acquired anything. On a quiet instance this should be zero.
+   */
+  checkedOutAtEntry: number;
 };
 
+/**
+ * A request arriving after this much instance quiet is treated as an observation
+ * point regardless of how fast it was. `idleTimeoutMillis` is 5s, so by then a
+ * healthy pool has closed its idle clients and `checked_out` should be zero —
+ * which is precisely when a leaked slot is visible. These requests are fast, so
+ * the duration threshold would never surface them.
+ */
+const QUIET_INSTANCE_MS = 5_000;
+
 export function shouldEmitUsageRecordTiming(
   totalMs: number,
+  msSinceLastRequest: number | null = null,
   random: () => number = Math.random
 ): boolean {
-  return totalMs >= SLOW_EMIT_THRESHOLD_MS || random() < BASELINE_SAMPLE_RATE;
+  if (totalMs >= SLOW_EMIT_THRESHOLD_MS) return true;
+  if (msSinceLastRequest !== null && msSinceLastRequest >= QUIET_INSTANCE_MS) return true;
+  return random() < BASELINE_SAMPLE_RATE;
+}
+
+/**
+ * Milliseconds since the previous request on this instance, or `null` for the
+ * first one. Tracked here rather than in the route so the notion of "quiet" and
+ * its emission threshold stay together.
+ */
+let lastRequestAtMs: number | null = null;
+
+export function noteRequestStart(now: number = Date.now()): number | null {
+  const since = lastRequestAtMs === null ? null : now - lastRequestAtMs;
+  lastRequestAtMs = now;
+  return since;
+}
+
+/** Test-only reset so quiet-period assertions do not depend on suite ordering. */
+export function __resetRequestClockForTest(): void {
+  lastRequestAtMs = null;
 }
 
 /**
@@ -120,6 +162,9 @@ export function emitUsageRecordTiming(timing: UsageRecordTiming): void {
       pool_waiting_peak: timing.poolWaitingPeak,
       pool_max: poolMax(),
       event_loop_lag: eventLoopLagMs(),
+      ms_since_last_request: timing.msSinceLastRequest,
+      checked_out_at_entry: timing.checkedOutAtEntry,
+      pool_leak: readPoolLeakStats(),
     })
   );
 }

apps/web/src/lib/db-pool-leak-probe.test.ts

@@ -0,0 +1,102 @@
+// Counters attach to the real pool only outside the test environment, so the pool
+// is stubbed here and the gauges are driven directly. The stub is built inside the
+// factory because `jest.mock` is hoisted above this file's own declarations, so
+// referencing an outer const here throws "Cannot access before initialization".
+jest.mock('@/lib/drizzle', () => ({
+  pool: { totalCount: 0, idleCount: 0, waitingCount: 0, on: jest.fn() },
+}));
+
+import { beforeEach, describe, expect, test } from '@jest/globals';
+
+import {
+  isTransactionBeginFailure,
+  readPoolLeakStats,
+  recordTransactionBeginFailure,
+  __resetPoolLeakCountersForTest,
+} from './db-pool-leak-probe';
+
+const { pool: mockPool } = jest.requireMock<{
+  pool: { totalCount: number; idleCount: number };
+}>('@/lib/drizzle');
+
+function drizzleError(message: string, cause?: unknown) {
+  const error = new Error(message);
+  error.name = 'DrizzleQueryError';
+  if (cause !== undefined) (error as unknown as { cause: unknown }).cause = cause;
+  return error;
+}
+
+beforeEach(() => {
+  __resetPoolLeakCountersForTest();
+  mockPool.totalCount = 0;
+  mockPool.idleCount = 0;
+});
+
+describe('isTransactionBeginFailure', () => {
+  // BEGIN takes no parameters, which is what makes matching the message safe here
+  // and unsafe for every other statement in this path.
+  test.each([
+    'Failed query: begin',
+    'Failed query: begin\nparams: ',
+    'Failed query: BEGIN',
+    '  Failed query: begin isolation level serializable',
+  ])('recognises %j', message => {
+    expect(isTransactionBeginFailure(drizzleError(message))).toBe(true);
+  });
+
+  test('recognises a begin failure nested in the cause chain', () => {
+    expect(
+      isTransactionBeginFailure(drizzleError('wrapper', drizzleError('Failed query: begin')))
+    ).toBe(true);
+  });
+
+  // These must not be counted as leaks: the statement failed inside the
+  // try/finally, so drizzle released the client.
+  test.each([
+    'Failed query: WITH microdollar_usage_ins AS (...) params: 3f5826e7,You are Kilo',
+    'Failed query: commit',
+    'Failed query: rollback',
+    'Failed query: select 1',
+    'Connection terminated unexpectedly',
+  ])('does not match %j', message => {
+    expect(isTransactionBeginFailure(drizzleError(message))).toBe(false);
+  });
+
+  test('does not match a statement that merely mentions begin', () => {
+    expect(isTransactionBeginFailure(drizzleError('Failed query: select * from beginnings'))).toBe(
+      false
+    );
+  });
+
+  test('does not throw on null or a primitive', () => {
+    expect(isTransactionBeginFailure(null)).toBe(false);
+    expect(isTransactionBeginFailure('nope')).toBe(false);
+  });
+});
+
+describe('readPoolLeakStats', () => {
+  test('derives checked_out from the pool gauges', () => {
+    mockPool.totalCount = 10;
+    mockPool.idleCount = 3;
+    expect(readPoolLeakStats().checked_out).toBe(7);
+  });
+
+  test('reports min_checked_out as null until a pool event is seen', () => {
+    expect(readPoolLeakStats().min_checked_out).toBeNull();
+  });
+
+  test('counts begin failures, the per-leak unit', () => {
+    recordTransactionBeginFailure();
+    recordTransactionBeginFailure();
+    expect(readPoolLeakStats().begin_failures).toBe(2);
+  });
+
+  test('reports outstanding as acquires minus releases', () => {
+    const stats = readPoolLeakStats();
+    expect(stats.outstanding).toBe(stats.acquires - stats.releases);
+  });
+
+  test('exposes uptime so a low-water mark can be read against instance age', () => {
+    expect(readPoolLeakStats().instance_uptime_ms).toBeGreaterThanOrEqual(0);
+  });
+});