chore: import upstream snapshot with attribution

apps/sim/lib/cleanup/batch-delete.ts

@@ -0,0 +1,264 @@
+import { db } from '@sim/db'
+import { createLogger } from '@sim/logger'
+import { and, inArray, isNotNull, lt, sql } from 'drizzle-orm'
+import type { PgColumn, PgTable } from 'drizzle-orm/pg-core'
+
+const logger = createLogger('BatchDelete')
+
+export const DEFAULT_BATCH_SIZE = 2000
+/** 50 × 2000 = 100K row cap per cleanup run; drains long-tail tenants in days, not weeks. */
+export const DEFAULT_MAX_BATCHES_PER_TABLE = 50
+/**
+ * Split workspaceIds into this-sized groups before running SELECT/DELETE. Large
+ * IN lists combined with `started_at < X` force Postgres to probe every
+ * workspace range in the composite index, which blows the 90s statement timeout
+ * at the scale of the full free tier.
+ */
+export const DEFAULT_WORKSPACE_CHUNK_SIZE = 50
+/** Bounds FK cascade trigger queue (per-statement in-memory) and bind-parameter count. */
+export const DEFAULT_DELETE_CHUNK_SIZE = 1000
+
+export function chunkArray<T>(arr: T[], size: number): T[][] {
+  const out: T[][] = []
+  for (let i = 0; i < arr.length; i += size) out.push(arr.slice(i, i + size))
+  return out
+}
+
+export interface SelectByIdChunksOptions {
+  /** Cap on rows returned across all chunks. Defaults to a full per-table cleanup budget. */
+  overallLimit?: number
+  chunkSize?: number
+}
+
+/**
+ * Run a SELECT query once per ID chunk and concatenate results up to
+ * `overallLimit`. Each chunk's query is passed the remaining row budget so the
+ * total never exceeds the cap. Use this when you need the selected row set
+ * (e.g. to drive S3 or copilot-backend cleanup alongside the DB delete).
+ *
+ * Works for any large ID set — workspace IDs, workflow IDs, etc. Avoids
+ * sending one massive `IN (...)` list that would blow Postgres's statement
+ * timeout.
+ */
+export async function selectRowsByIdChunks<T>(
+  ids: string[],
+  query: (chunkIds: string[], chunkLimit: number) => Promise<T[]>,
+  {
+    overallLimit = DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE,
+    chunkSize = DEFAULT_WORKSPACE_CHUNK_SIZE,
+  }: SelectByIdChunksOptions = {}
+): Promise<T[]> {
+  if (ids.length === 0) return []
+
+  const rows: T[] = []
+  for (const chunkIds of chunkArray(ids, chunkSize)) {
+    if (rows.length >= overallLimit) break
+    const remaining = overallLimit - rows.length
+    const chunkRows = await query(chunkIds, remaining)
+    rows.push(...chunkRows)
+  }
+  return rows
+}
+
+export interface TableCleanupResult {
+  table: string
+  deleted: number
+  failed: number
+}
+
+export interface ChunkedBatchDeleteOptions<TRow extends { id: string }> {
+  tableDef: PgTable
+  workspaceIds: string[]
+  tableName: string
+  /** SELECT eligible rows for one workspace chunk. The result must include `id`. */
+  selectChunk: (chunkIds: string[], limit: number) => Promise<TRow[]>
+  /** Runs between SELECT and DELETE; receives the just-selected rows. */
+  onBatch?: (rows: TRow[]) => Promise<void>
+  batchSize?: number
+  /** Max batches per workspace chunk. */
+  maxBatches?: number
+  /**
+   * Hard cap on rows processed (deleted + failed) across all chunks per call.
+   * Defaults to `DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE`. Cron
+   * runs frequently enough to catch up the backlog over multiple invocations.
+   */
+  totalRowLimit?: number
+  workspaceChunkSize?: number
+}
+
+/**
+ * Inner loop primitive for cleanup jobs.
+ *
+ * For each workspace chunk: SELECT a batch of eligible rows → run optional
+ * `onBatch` hook (e.g. to delete S3 files) → DELETE those rows by ID. Repeats
+ * until exhausted or `maxBatches` is hit, then moves to the next chunk. Stops
+ * the whole call once `totalRowLimit` rows have been processed.
+ *
+ * Workspace IDs are chunked before the SELECT — see
+ * `DEFAULT_WORKSPACE_CHUNK_SIZE` for why.
+ */
+export async function chunkedBatchDelete<TRow extends { id: string }>({
+  tableDef,
+  workspaceIds,
+  tableName,
+  selectChunk,
+  onBatch,
+  batchSize = DEFAULT_BATCH_SIZE,
+  maxBatches = DEFAULT_MAX_BATCHES_PER_TABLE,
+  totalRowLimit = DEFAULT_BATCH_SIZE * DEFAULT_MAX_BATCHES_PER_TABLE,
+  workspaceChunkSize = DEFAULT_WORKSPACE_CHUNK_SIZE,
+}: ChunkedBatchDeleteOptions<TRow>): Promise<TableCleanupResult> {
+  const result: TableCleanupResult = { table: tableName, deleted: 0, failed: 0 }
+
+  if (workspaceIds.length === 0) {
+    logger.info(`[${tableName}] Skipped — no workspaces in scope`)
+    return result
+  }
+
+  const chunks = chunkArray(workspaceIds, workspaceChunkSize)
+  let stoppedEarly = false
+  let attempted = 0
+
+  for (const [chunkIdx, chunkIds] of chunks.entries()) {
+    if (attempted >= totalRowLimit) {
+      stoppedEarly = true
+      break
+    }
+
+    let batchesProcessed = 0
+    let hasMore = true
+
+    while (hasMore && batchesProcessed < maxBatches && attempted < totalRowLimit) {
+      let rows: TRow[] = []
+      try {
+        const remainingLimit = totalRowLimit - attempted
+        const effectiveBatchSize = Math.min(batchSize, remainingLimit)
+        if (effectiveBatchSize <= 0) {
+          hasMore = false
+          break
+        }
+
+        rows = await selectChunk(chunkIds, effectiveBatchSize)
+
+        if (rows.length === 0) {
+          hasMore = false
+          break
+        }
+
+        attempted += rows.length
+        if (onBatch) await onBatch(rows)
+
+        const ids = rows.map((r) => r.id)
+        const deleted = await db
+          .delete(tableDef)
+          .where(inArray(sql`id`, ids))
+          .returning({ id: sql`id` })
+
+        result.deleted += deleted.length
+        result.failed += rows.length - deleted.length
+        hasMore = rows.length === effectiveBatchSize && attempted < totalRowLimit
+        batchesProcessed++
+      } catch (error) {
+        // Count rows we tried to delete; SELECT-stage errors leave rows=[].
+        result.failed += rows.length
+        logger.error(
+          `[${tableName}] Batch failed (chunk ${chunkIdx + 1}/${chunks.length}, ${rows.length} rows):`,
+          { error }
+        )
+        hasMore = false
+      }
+    }
+  }
+
+  logger.info(
+    `[${tableName}] Complete: ${result.deleted} deleted, ${result.failed} failed across ${chunks.length} chunks${stoppedEarly ? ' (row-limit reached, remaining chunks deferred to next run)' : ''}`
+  )
+
+  return result
+}
+
+export interface BatchDeleteOptions {
+  tableDef: PgTable
+  workspaceIdCol: PgColumn
+  timestampCol: PgColumn
+  workspaceIds: string[]
+  retentionDate: Date
+  tableName: string
+  /** When true, also requires `timestampCol IS NOT NULL` (soft-delete semantics). */
+  requireTimestampNotNull?: boolean
+  batchSize?: number
+  maxBatches?: number
+  workspaceChunkSize?: number
+}
+
+/**
+ * Convenience wrapper around `chunkedBatchDelete` for the common case: delete
+ * rows where `workspaceId IN (...) AND timestamp < retentionDate`. Use this
+ * when there's no per-row side effect (e.g. no S3 files to clean up alongside).
+ */
+export async function batchDeleteByWorkspaceAndTimestamp({
+  tableDef,
+  workspaceIdCol,
+  timestampCol,
+  workspaceIds,
+  retentionDate,
+  tableName,
+  requireTimestampNotNull = false,
+  ...rest
+}: BatchDeleteOptions): Promise<TableCleanupResult> {
+  return chunkedBatchDelete({
+    tableDef,
+    workspaceIds,
+    tableName,
+    selectChunk: (chunkIds, limit) => {
+      const predicates = [inArray(workspaceIdCol, chunkIds), lt(timestampCol, retentionDate)]
+      if (requireTimestampNotNull) predicates.push(isNotNull(timestampCol))
+      return db
+        .select({ id: sql<string>`id` })
+        .from(tableDef)
+        .where(and(...predicates))
+        .limit(limit)
+    },
+    ...rest,
+  })
+}
+
+/**
+ * Delete by explicit ID list, chunked so each statement is its own transaction.
+ * Partial progress survives chunk-level failures.
+ */
+export async function deleteRowsById(
+  tableDef: PgTable,
+  idCol: PgColumn,
+  ids: string[],
+  tableName: string,
+  chunkSize: number = DEFAULT_DELETE_CHUNK_SIZE
+): Promise<TableCleanupResult> {
+  const result: TableCleanupResult = { table: tableName, deleted: 0, failed: 0 }
+  if (ids.length === 0) return result
+
+  const chunks = chunkArray(ids, chunkSize)
+  for (const [chunkIdx, chunkIds] of chunks.entries()) {
+    try {
+      const deleted = await db
+        .delete(tableDef)
+        .where(inArray(idCol, chunkIds))
+        .returning({ id: idCol })
+      result.deleted += deleted.length
+    } catch (error) {
+      // Upper bound: Postgres rolls back the chunk on error, so actual deletes = 0,
+      // but we can't tell which IDs in the chunk would have matched. The next cron
+      // run picks up whatever's still expired, so this only inflates the metric.
+      result.failed += chunkIds.length
+      logger.error(
+        `[${tableName}] Delete chunk ${chunkIdx + 1}/${chunks.length} failed (up to ${chunkIds.length} rows):`,
+        { error }
+      )
+    }
+  }
+
+  logger.info(
+    `[${tableName}] Deleted ${result.deleted} rows across ${chunks.length} chunk(s)${result.failed > 0 ? `, ${result.failed} failed` : ''}`
+  )
+  return result
+}

apps/sim/lib/cleanup/chat-cleanup.ts

@@ -0,0 +1,215 @@
+import { db } from '@sim/db'
+import { copilotMessages, workspaceFiles } from '@sim/db/schema'
+import { createLogger } from '@sim/logger'
+import { and, inArray, isNull } from 'drizzle-orm'
+import { chunkArray } from '@/lib/cleanup/batch-delete'
+import { SIM_AGENT_API_URL } from '@/lib/copilot/constants'
+import { env } from '@/lib/core/config/env'
+import type { StorageContext } from '@/lib/uploads'
+import { isUsingCloudStorage, StorageService } from '@/lib/uploads'
+
+const logger = createLogger('ChatCleanup')
+
+const COPILOT_CLEANUP_BATCH_SIZE = 1000
+/** Bounds how many chats' `copilot_messages` rows are scanned per query. */
+const CHAT_FILE_COLLECT_CHUNK_SIZE = 500
+
+/**
+ * Only storage in these contexts is tied to chat/task lifecycle. Workspace
+ * files, execution logs, knowledge bases, profile pictures, etc. are owned by
+ * other subsystems and must never be touched by chat cleanup — even if a row
+ * somehow ends up with `chatId` set through a future flow.
+ */
+const CHAT_SCOPED_CONTEXTS = ['copilot', 'mothership'] as const satisfies readonly StorageContext[]
+type ChatScopedContext = (typeof CHAT_SCOPED_CONTEXTS)[number]
+
+interface FileRef {
+  key: string
+  context: ChatScopedContext
+}
+
+/**
+ * Collect all file storage keys for the given chat IDs from two sources:
+ * 1. workspaceFiles rows with chatId FK (chat-scoped contexts only)
+ * 2. fileAttachments[].key inside each copilot_messages.content
+ */
+export async function collectChatFiles(chatIds: string[]): Promise<FileRef[]> {
+  const files: FileRef[] = []
+  if (chatIds.length === 0) return files
+
+  const seen = new Set<string>()
+
+  for (const chunk of chunkArray(chatIds, CHAT_FILE_COLLECT_CHUNK_SIZE)) {
+    const [linkedFiles, messageRows] = await Promise.all([
+      db
+        .select({ key: workspaceFiles.key, context: workspaceFiles.context })
+        .from(workspaceFiles)
+        .where(
+          and(
+            inArray(workspaceFiles.chatId, chunk),
+            isNull(workspaceFiles.deletedAt),
+            inArray(workspaceFiles.context, [...CHAT_SCOPED_CONTEXTS])
+          )
+        ),
+      // Scan every message row for the chat (no deleted_at filter): this is a
+      // deletion path collecting blob keys, so attachments on any row count.
+      db
+        .select({ content: copilotMessages.content })
+        .from(copilotMessages)
+        .where(inArray(copilotMessages.chatId, chunk)),
+    ])
+
+    for (const f of linkedFiles) {
+      if (!seen.has(f.key)) {
+        seen.add(f.key)
+        files.push({ key: f.key, context: f.context as ChatScopedContext })
+      }
+    }
+
+    for (const row of messageRows) {
+      const msg = row.content
+      if (!msg || typeof msg !== 'object') continue
+      const attachments = (msg as Record<string, unknown>).fileAttachments
+      if (!Array.isArray(attachments)) continue
+      for (const attachment of attachments) {
+        if (
+          attachment &&
+          typeof attachment === 'object' &&
+          (attachment as Record<string, unknown>).key
+        ) {
+          const key = (attachment as Record<string, unknown>).key as string
+          if (!seen.has(key)) {
+            seen.add(key)
+            files.push({ key, context: 'copilot' })
+          }
+        }
+      }
+    }
+  }
+
+  return files
+}
+
+/** Groups files by storage context so each context can use one batch DELETE call. */
+export async function deleteStorageFiles(
+  files: FileRef[],
+  label: string
+): Promise<{ filesDeleted: number; filesFailed: number }> {
+  const stats = { filesDeleted: 0, filesFailed: 0 }
+  if (files.length === 0 || !isUsingCloudStorage()) return stats
+
+  const keysByContext = new Map<ChatScopedContext, string[]>()
+  for (const file of files) {
+    const bucket = keysByContext.get(file.context)
+    if (bucket) bucket.push(file.key)
+    else keysByContext.set(file.context, [file.key])
+  }
+
+  for (const [context, keys] of keysByContext) {
+    const result = await StorageService.deleteFiles(keys, context)
+    stats.filesDeleted += result.deleted
+    stats.filesFailed += result.failed.length
+    for (const { key, error } of result.failed) {
+      logger.error(`[${label}] Failed to delete storage file ${key} (context: ${context}):`, {
+        error,
+      })
+    }
+  }
+
+  return stats
+}
+
+/**
+ * Call the copilot backend to delete chat data (memory_files, checkpoints, task_chains, etc.)
+ * Chunked at 1000 per request.
+ */
+export async function cleanupCopilotBackend(
+  chatIds: string[],
+  label: string
+): Promise<{ deleted: number; failed: number }> {
+  const stats = { deleted: 0, failed: 0 }
+
+  if (chatIds.length === 0 || !env.COPILOT_API_KEY) {
+    if (!env.COPILOT_API_KEY) {
+      logger.warn(`[${label}] COPILOT_API_KEY not set, skipping copilot backend cleanup`)
+    }
+    return stats
+  }
+
+  for (let i = 0; i < chatIds.length; i += COPILOT_CLEANUP_BATCH_SIZE) {
+    const chunk = chatIds.slice(i, i + COPILOT_CLEANUP_BATCH_SIZE)
+    try {
+      const response = await fetch(`${SIM_AGENT_API_URL}/api/tasks/cleanup`, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          'x-api-key': env.COPILOT_API_KEY,
+        },
+        body: JSON.stringify({ chatIds: chunk }),
+      })
+
+      if (!response.ok) {
+        const errorBody = await response.text().catch(() => '')
+        logger.error(`[${label}] Copilot backend cleanup failed: ${response.status}`, {
+          errorBody,
+          chatCount: chunk.length,
+        })
+        stats.failed += chunk.length
+        continue
+      }
+
+      const result = await response.json()
+      stats.deleted += result.deleted ?? 0
+      logger.info(
+        `[${label}] Copilot backend cleanup: ${result.deleted} chats deleted (batch ${Math.floor(i / COPILOT_CLEANUP_BATCH_SIZE) + 1})`
+      )
+    } catch (error) {
+      stats.failed += chunk.length
+      logger.error(`[${label}] Copilot backend cleanup request failed:`, { error })
+    }
+  }
+
+  return stats
+}
+
+/**
+ * Full chat cleanup: collect file refs, then (after DB deletion by caller)
+ * call copilot backend and delete storage files.
+ *
+ * Usage:
+ *   const cleanup = await prepareChatCleanup(chatIds, label)
+ *   // ... delete DB rows ...
+ *   await cleanup.execute()
+ */
+export async function prepareChatCleanup(
+  chatIds: string[],
+  label: string
+): Promise<{ execute: () => Promise<void> }> {
+  // Collect file refs BEFORE DB deletion (keys + context are lost after cascade)
+  const files = await collectChatFiles(chatIds)
+  if (files.length > 0) {
+    logger.info(`[${label}] Collected ${files.length} files for cleanup`, {
+      files: files.map((f) => ({ key: f.key, context: f.context })),
+    })
+  }
+
+  return {
+    execute: async () => {
+      // Call copilot backend
+      if (chatIds.length > 0) {
+        const copilotResult = await cleanupCopilotBackend(chatIds, label)
+        logger.info(
+          `[${label}] Copilot backend: ${copilotResult.deleted} deleted, ${copilotResult.failed} failed`
+        )
+      }
+
+      // Delete storage files with correct context per file
+      if (files.length > 0) {
+        const fileStats = await deleteStorageFiles(files, label)
+        logger.info(
+          `[${label}] Storage cleanup: ${fileStats.filesDeleted} deleted, ${fileStats.filesFailed} failed`
+        )
+      }
+    },
+  }
+}