fix(knowledge): fix chunk_index race and storage-quota check/increment race (#5544)

* fix(knowledge): fix chunk_index race and storage-quota check/increment race

- createChunk now serializes concurrent writes to the same document via a
  transactional advisory lock before computing the next chunk_index, fixing
  the unique-constraint collision behind the ITSM 'Failed to create chunk'
  incident (two overlapping workflow runs appending to the same shared KB
  document).
- Document uploads now check and increment storage quota atomically in a
  single conditional UPDATE inside the insert transaction, instead of
  check-then-insert-then-increment-after-commit, closing the window where
  two concurrent uploads could both pass the quota check and over-commit
  storage.

* fix(knowledge): bound the chunk advisory-lock wait with lock_timeout

Address Greptile P1: the advisory lock in createChunk could wait
indefinitely on a stalled same-document transaction while holding a pooled
connection. Set a 5s lock_timeout before acquiring it, matching the
set_config + pg_advisory_xact_lock pattern already used by every other
advisory lock in this codebase (org membership, table schema/row-ordering,
BYOK keys, workspace env, usage-log flush, execution-log reconciliation).

apps/sim/lib/billing/storage/index.ts

@@ -1,6 +1,8 @@
 export { checkStorageQuota, getUserStorageLimit, getUserStorageUsage } from './limits'
 export {
+  checkAndIncrementStorageUsageInTx,
   decrementStorageUsage,
   decrementStorageUsageInTx,
   incrementStorageUsage,
+  maybeNotifyStorageLimit,
 } from './tracking'

apps/sim/lib/billing/storage/tracking.ts

@@ -7,11 +7,12 @@
 import { db } from '@sim/db'
 import { organization, userStats } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
-import { eq, sql } from 'drizzle-orm'
+import { generateId } from '@sim/utils/id'
+import { and, eq, sql } from 'drizzle-orm'
 import { maybeNotifyLimit } from '@/lib/billing/core/limit-notifications'
 import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
 import { getUserStorageLimit, getUserStorageUsage } from '@/lib/billing/storage/limits'
-import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
+import { getFreeTierLimit, isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
 import { isBillingEnabled } from '@/lib/core/config/env-flags'
 
 const logger = createLogger('StorageTracking')
@@ -33,7 +34,7 @@ function formatGb(bytes: number, decimals: number): string {
  * @param rearmOnly - True on decrements, so a shrink that leaves usage above a
  *   threshold re-arms but never sends (a drop is not a fresh crossing).
  */
-async function maybeNotifyStorageLimit(
+export async function maybeNotifyStorageLimit(
   userId: string,
   workspaceId: string,
   sub: HighestPrioritySubscription | null,
@@ -195,3 +196,94 @@ export async function decrementStorageUsageInTx(
       .where(eq(userStats.userId, userId))
   }
 }
+
+/**
+ * Atomically check quota and increment a user's (or their org's) storage
+ * counter inside an existing transaction, using a pre-resolved subscription.
+ * The check and the increment are a single conditional `UPDATE`, so two
+ * concurrent callers can no longer both read the same pre-increment usage,
+ * both pass the check, and both commit past the limit — the second caller's
+ * `UPDATE` re-evaluates the WHERE clause against the first caller's already
+ * -committed-within-the-same-DB-round-trip row and correctly fails. Replaces
+ * the old read-then-decide-then-increment-after-commit split (`checkStorageQuota`
+ * + a fire-and-forget `incrementStorageUsage` after the transaction), which left
+ * a window between the read and the increment.
+ *
+ * On success, callers should best-effort call {@link maybeNotifyStorageLimit}
+ * after the transaction commits (mirrors the existing post-increment threshold
+ * check) — this helper doesn't do it itself since it runs mid-transaction.
+ *
+ * For a personal (non-org-scoped) `userId`, this first upserts the
+ * `userStats` row on `tx` — a documented possibility for OAuth account
+ * linking (see `ensureUserStatsExists` in `lib/billing/core/usage.ts`, whose
+ * insert values this mirrors) — because the conditional `UPDATE` below
+ * matches 0 rows, and therefore reads as "quota exceeded", if that row
+ * doesn't exist yet. `ensureUserStatsExists` itself isn't reused here since
+ * it writes through the standalone `db` client, which would open a second
+ * pooled connection while this transaction's is held.
+ */
+export async function checkAndIncrementStorageUsageInTx(
+  tx: StorageTransaction,
+  sub: HighestPrioritySubscription | null,
+  userId: string,
+  bytes: number
+): Promise<{ allowed: boolean; currentUsage: number; limit: number; error?: string }> {
+  if (!isBillingEnabled) {
+    return { allowed: true, currentUsage: 0, limit: Number.MAX_SAFE_INTEGER }
+  }
+
+  const limit = await getUserStorageLimit(userId, sub)
+
+  if (bytes <= 0) {
+    return { allowed: true, currentUsage: await getUserStorageUsage(userId, sub), limit }
+  }
+
+  const orgScoped = isOrgScopedSubscription(sub, userId) && sub
+
+  if (!orgScoped) {
+    await tx
+      .insert(userStats)
+      .values({
+        id: generateId(),
+        userId,
+        currentUsageLimit: getFreeTierLimit().toString(),
+        usageLimitUpdatedAt: new Date(),
+      })
+      .onConflictDoNothing({ target: userStats.userId })
+  }
+
+  const [updated] = orgScoped
+    ? await tx
+        .update(organization)
+        .set({ storageUsedBytes: sql`${organization.storageUsedBytes} + ${bytes}` })
+        .where(
+          and(
+            eq(organization.id, sub.referenceId),
+            sql`${organization.storageUsedBytes} + ${bytes} <= ${limit}`
+          )
+        )
+        .returning({ storageUsedBytes: organization.storageUsedBytes })
+    : await tx
+        .update(userStats)
+        .set({ storageUsedBytes: sql`${userStats.storageUsedBytes} + ${bytes}` })
+        .where(
+          and(
+            eq(userStats.userId, userId),
+            sql`${userStats.storageUsedBytes} + ${bytes} <= ${limit}`
+          )
+        )
+        .returning({ storageUsedBytes: userStats.storageUsedBytes })
+
+  if (updated) {
+    return { allowed: true, currentUsage: updated.storageUsedBytes - bytes, limit }
+  }
+
+  const currentUsage = await getUserStorageUsage(userId, sub)
+  const newUsage = currentUsage + bytes
+  return {
+    allowed: false,
+    currentUsage,
+    limit,
+    error: `Storage limit exceeded. Used: ${(newUsage / (1024 * 1024 * 1024)).toFixed(2)}GB, Limit: ${(limit / (1024 * 1024 * 1024)).toFixed(0)}GB`,
+  }
+}

apps/sim/lib/knowledge/chunks/service.ts

@@ -17,6 +17,8 @@ import { estimateTokenCount } from '@/lib/tokenization/estimators'
 
 const logger = createLogger('ChunksService')
 
+const KB_CHUNK_LOCK_TIMEOUT_MS = 5_000
+
 /**
  * Query chunks for a document with filtering and pagination
  */
@@ -101,7 +103,23 @@ export async function queryChunks(
 }
 
 /**
- * Create a new chunk for a document
+ * Create a new chunk for a document.
+ *
+ * Assigns `chunkIndex` as `max(chunkIndex) + 1` under a transactional
+ * `pg_advisory_xact_lock` keyed on the document, so concurrent calls for the
+ * same document serialize instead of computing the same index and colliding
+ * on the `(document_id, chunk_index)` unique constraint. A `SELECT ... FOR
+ * UPDATE` on the current max row doesn't prevent that collision, since the
+ * row it would lock is unrelated to the not-yet-inserted next row. A row
+ * lock on `document` instead of an advisory lock would also work, but would
+ * invert the embedding-before-document lock order every other chunk
+ * mutation path uses (see lock-order.test.ts) — the advisory lock is a
+ * separate namespace, so it can't deadlock against that convention.
+ *
+ * `pg_advisory_xact_lock` auto-releases at transaction end, so there's no
+ * session lock to leak onto a pooled connection, and `lock_timeout` bounds
+ * the wait (it raises SQLSTATE 55P03 instead of hanging a pooled connection)
+ * if a same-document holder is stuck.
  */
 export async function createChunk(
   knowledgeBaseId: string,
@@ -135,8 +153,14 @@ export async function createChunk(
   const chunkId = generateId()
   const now = new Date()
 
-  // Use transaction to atomically get next index and insert chunk
   const newChunk = await db.transaction(async (tx) => {
+    await tx.execute(
+      sql`select set_config('lock_timeout', ${`${KB_CHUNK_LOCK_TIMEOUT_MS}ms`}, true)`
+    )
+    await tx.execute(
+      sql`select pg_advisory_xact_lock(hashtextextended(${`kb_chunk_seq:${documentId}`}, 0))`
+    )
+
     const activeDocument = await tx
       .select({ id: document.id })
       .from(document)
@@ -156,7 +180,6 @@ export async function createChunk(
       throw new Error('Document not found')
     }
 
-    // Get the next chunk index atomically within the transaction
     const lastChunk = await tx
       .select({ chunkIndex: embedding.chunkIndex })
       .from(embedding)

apps/sim/lib/knowledge/documents/service.ts

@@ -18,9 +18,9 @@ import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
 import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
 import { recordUsage } from '@/lib/billing/core/usage-log'
 import {
-  checkStorageQuota,
+  checkAndIncrementStorageUsageInTx,
   decrementStorageUsageInTx,
-  incrementStorageUsage,
+  maybeNotifyStorageLimit,
 } from '@/lib/billing/storage'
 import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing'
 import type { ChunkingStrategy, StrategyOptions } from '@/lib/chunkers/types'
@@ -864,6 +864,33 @@ export function isTriggerAvailable(): boolean {
   return Boolean(env.TRIGGER_SECRET_KEY) && isTriggerDevEnabled
 }
 
+/**
+ * Resolves the subscription to bill for a storage-metered upload, ahead of
+ * the FOR UPDATE transaction that will insert the document row.
+ *
+ * Must run before that transaction opens: `knowledgeBase.userId` is
+ * immutable once a KB is created, so resolving it here can't go stale
+ * relative to the transaction's own read of it, and resolving it inside the
+ * transaction would open a second pooled-database-connection checkout while
+ * the first is held.
+ */
+async function resolveQuotaSubscription(
+  knowledgeBaseId: string,
+  uploadedBy: string | null
+): Promise<HighestPrioritySubscription | null> {
+  const billedUserId =
+    uploadedBy ??
+    (
+      await db
+        .select({ userId: knowledgeBase.userId })
+        .from(knowledgeBase)
+        .where(eq(knowledgeBase.id, knowledgeBaseId))
+        .limit(1)
+    )[0]?.userId
+
+  return billedUserId ? getHighestPrioritySubscription(billedUserId) : null
+}
+
 export async function createDocumentRecords(
   documents: Array<{
     filename: string
@@ -883,7 +910,15 @@ export async function createDocumentRecords(
   requestId: string,
   uploadedBy: string | null = null
 ): Promise<DocumentData[]> {
-  let storageBilling: { userId: string; workspaceId: string | null; bytes: number } | null = null
+  let storageBilling: {
+    userId: string
+    workspaceId: string | null
+    bytes: number
+    sub: HighestPrioritySubscription | null
+  } | null = null
+
+  const totalBytes = documents.reduce((sum, docData) => sum + (docData.fileSize || 0), 0)
+  const sub = totalBytes > 0 ? await resolveQuotaSubscription(knowledgeBaseId, uploadedBy) : null
 
   const returnData = await db.transaction(async (tx) => {
     await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
@@ -912,13 +947,12 @@ export async function createDocumentRecords(
     )
 
     const billedUserId = uploadedBy ?? kb[0].userId
-    const totalBytes = documents.reduce((sum, docData) => sum + (docData.fileSize || 0), 0)
     if (totalBytes > 0) {
-      const quotaCheck = await checkStorageQuota(billedUserId, totalBytes)
+      const quotaCheck = await checkAndIncrementStorageUsageInTx(tx, sub, billedUserId, totalBytes)
       if (!quotaCheck.allowed) {
         throw new Error(quotaCheck.error || 'Storage limit exceeded')
       }
-      storageBilling = { userId: billedUserId, workspaceId: kbWorkspaceId, bytes: totalBytes }
+      storageBilling = { userId: billedUserId, workspaceId: kbWorkspaceId, bytes: totalBytes, sub }
     }
 
     // One load per batch (was N+1); skip entirely if no doc carries tags.
@@ -1011,11 +1045,14 @@ export async function createDocumentRecords(
   })
 
   if (storageBilling) {
-    const billing: { userId: string; workspaceId: string | null; bytes: number } = storageBilling
-    try {
-      await incrementStorageUsage(billing.userId, billing.bytes, billing.workspaceId ?? undefined)
-    } catch (storageError) {
-      logger.error(`[${requestId}] Failed to update storage tracking:`, storageError)
+    const billing: {
+      userId: string
+      workspaceId: string | null
+      bytes: number
+      sub: HighestPrioritySubscription | null
+    } = storageBilling
+    if (billing.workspaceId) {
+      void maybeNotifyStorageLimit(billing.userId, billing.workspaceId, billing.sub)
     }
   }
 
@@ -1338,7 +1375,15 @@ export async function createSingleDocument(
     ...processedTags,
   }
 
-  let storageBilling: { userId: string; workspaceId: string | null; bytes: number } | null = null
+  let storageBilling: {
+    userId: string
+    workspaceId: string | null
+    bytes: number
+    sub: HighestPrioritySubscription | null
+  } | null = null
+
+  const sub =
+    documentData.fileSize > 0 ? await resolveQuotaSubscription(knowledgeBaseId, uploadedBy) : null
 
   await db.transaction(async (tx) => {
     await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
@@ -1367,7 +1412,12 @@ export async function createSingleDocument(
 
     const billedUserId = uploadedBy ?? kb[0].userId
     if (documentData.fileSize > 0) {
-      const quotaCheck = await checkStorageQuota(billedUserId, documentData.fileSize)
+      const quotaCheck = await checkAndIncrementStorageUsageInTx(
+        tx,
+        sub,
+        billedUserId,
+        documentData.fileSize
+      )
       if (!quotaCheck.allowed) {
         throw new Error(quotaCheck.error || 'Storage limit exceeded')
       }
@@ -1375,6 +1425,7 @@ export async function createSingleDocument(
         userId: billedUserId,
         workspaceId: kb[0].workspaceId,
         bytes: documentData.fileSize,
+        sub,
       }
     }
 
@@ -1387,11 +1438,14 @@ export async function createSingleDocument(
   })
 
   if (storageBilling) {
-    const billing: { userId: string; workspaceId: string | null; bytes: number } = storageBilling
-    try {
-      await incrementStorageUsage(billing.userId, billing.bytes, billing.workspaceId ?? undefined)
-    } catch (storageError) {
-      logger.error(`[${requestId}] Failed to update storage tracking:`, storageError)
+    const billing: {
+      userId: string
+      workspaceId: string | null
+      bytes: number
+      sub: HighestPrioritySubscription | null
+    } = storageBilling
+    if (billing.workspaceId) {
+      void maybeNotifyStorageLimit(billing.userId, billing.workspaceId, billing.sub)
     }
   }