address comments

Author: icecrasher321

apps/sim/lib/execution/payloads/large-value-metadata.test.ts

@@ -27,6 +27,7 @@ const {
   mockValues,
   mockWhere,
   mockTxWhere,
+  mockNotInArray,
 } = vi.hoisted(() => {
   const mockOnConflictDoNothing = vi.fn(async () => undefined)
   const mockValues = vi.fn(() => ({ onConflictDoNothing: mockOnConflictDoNothing }))
@@ -54,6 +55,7 @@ const {
     mockEq: vi.fn((...args: unknown[]) => ({ op: 'eq', args })),
     mockExecute: vi.fn(async () => [{ count: 0 }]),
     mockInsert,
+    mockNotInArray: vi.fn((...args: unknown[]) => ({ op: 'notInArray', args })),
     mockOnConflictDoNothing,
     mockSelect,
     mockSelectFrom,
@@ -127,6 +129,7 @@ vi.mock('drizzle-orm', () => ({
   and: mockAnd,
   eq: mockEq,
   inArray: vi.fn((...args: unknown[]) => ({ op: 'inArray', args })),
+  notInArray: mockNotInArray,
   sql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })),
 }))
 
@@ -286,6 +289,38 @@ describe('large value metadata', () => {
     expect(mockTxSelectLimit).toHaveBeenCalledWith(MAX_LARGE_VALUE_REFERENCES_PER_SCOPE)
   })
 
+  it('filters known dependency children before applying the remaining reference budget', async () => {
+    const directKeys = Array.from({ length: MAX_LARGE_VALUE_REFERENCES_PER_SCOPE }, (_, index) =>
+      largeValueKey(`e${index.toString(36).padStart(11, '0')}`)
+    )
+    const knownChildKey = directKeys[1]
+    const unseenChildKey = largeValueKey('unseenchild1', 'source-execution')
+    mockTxSelectLimit.mockImplementationOnce(async () => {
+      const filtersKnownChildren = mockNotInArray.mock.calls.some(
+        ([field, values]) =>
+          field === 'executionLargeValueDependencies.childKey' &&
+          Array.isArray(values) &&
+          values.includes(knownChildKey)
+      )
+      return [{ childKey: filtersKnownChildren ? unseenChildKey : knownChildKey }]
+    })
+
+    await expect(
+      registerLargeValueOwner(
+        {
+          key: largeValueKey('zyxwvutsrqpo', 'execution-1'),
+          workspaceId: 'workspace-1',
+          workflowId: 'workflow-1',
+          executionId: 'execution-1',
+          size: 123,
+        },
+        directKeys
+      )
+    ).rejects.toThrow('Large value dependency closure exceeds the limit')
+
+    expect(mockTxSelectLimit).toHaveBeenCalledWith(1)
+  })
+
   it('replaces an execution reference set with same-workspace unique keys', async () => {
     const matchingKey = largeValueKey('abcdefghijkl')
     const otherWorkspaceKey =

apps/sim/lib/execution/payloads/large-value-metadata.ts

@@ -7,7 +7,7 @@ import {
   workflowExecutionLogs,
 } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
-import { and, eq, inArray, sql } from 'drizzle-orm'
+import { and, eq, inArray, notInArray, sql } from 'drizzle-orm'
 import { collectLargeValueKeys } from '@/lib/execution/payloads/large-execution-value'
 
 const logger = createLogger('LargeValueMetadata')
@@ -101,6 +101,16 @@ function getCount(rows: unknown): number {
   return Number((row as { count: unknown }).count) || 0
 }
 
+export function collectLargeValueReferenceKeys(value: unknown, workspaceId?: string): string[] {
+  return getBoundedUniqueKeys(
+    collectLargeValueKeys(value).filter((key) => {
+      const parsed = parseLargeValueStorageKey(key)
+      return workspaceId ? parsed?.workspaceId === workspaceId : Boolean(parsed)
+    }),
+    'Large value reference set'
+  )
+}
+
 async function getDependencyClosure(
   client: LargeValueMetadataClient,
   ownerKey: string,
@@ -133,7 +143,8 @@ async function getDependencyClosure(
         .where(
           and(
             eq(executionLargeValueDependencies.workspaceId, workspaceId),
-            inArray(executionLargeValueDependencies.parentKey, keyChunk)
+            inArray(executionLargeValueDependencies.parentKey, keyChunk),
+            notInArray(executionLargeValueDependencies.childKey, Array.from(closureKeys))
           )
         )
         .limit(remainingBudget + 1)
@@ -225,14 +236,30 @@ export async function replaceLargeValueReferencesWithClient(
   client: LargeValueMetadataClient,
   scope: LargeValueReferenceScope,
   value: unknown
+): Promise<void> {
+  if (!scope.workspaceId || !scope.executionId) {
+    return
+  }
+
+  await replaceLargeValueReferenceKeysWithClient(
+    client,
+    scope,
+    collectLargeValueReferenceKeys(value, scope.workspaceId)
+  )
+}
+
+export async function replaceLargeValueReferenceKeysWithClient(
+  client: LargeValueMetadataClient,
+  scope: LargeValueReferenceScope,
+  referenceKeys: string[]
 ): Promise<void> {
   const { workspaceId, workflowId, executionId, source } = scope
   if (!workspaceId || !executionId) {
     return
   }
 
   const keys = getBoundedUniqueKeys(
-    collectLargeValueKeys(value).filter((key) => {
+    referenceKeys.filter((key) => {
       const parsed = parseLargeValueStorageKey(key)
       return parsed?.workspaceId === workspaceId
     }),
@@ -340,8 +367,11 @@ export async function replaceLargeValueReferences(
   scope: LargeValueReferenceScope,
   value: unknown
 ): Promise<void> {
+  const referenceKeys = scope.workspaceId
+    ? collectLargeValueReferenceKeys(value, scope.workspaceId)
+    : []
   await db.transaction(async (tx) => {
-    await replaceLargeValueReferencesWithClient(tx, scope, value)
+    await replaceLargeValueReferenceKeysWithClient(tx, scope, referenceKeys)
   })
 }
 

apps/sim/lib/logs/execution/logger.ts

@@ -22,7 +22,10 @@ import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing'
 import { isBillingEnabled } from '@/lib/core/config/feature-flags'
 import { redactApiKeys } from '@/lib/core/security/redaction'
 import { filterForDisplay } from '@/lib/core/utils/display-filters'
-import { replaceLargeValueReferencesWithClient } from '@/lib/execution/payloads/large-value-metadata'
+import {
+  collectLargeValueReferenceKeys,
+  replaceLargeValueReferenceKeysWithClient,
+} from '@/lib/execution/payloads/large-value-metadata'
 import { emitWorkflowExecutionCompleted } from '@/lib/logs/events'
 import { snapshotService } from '@/lib/logs/execution/snapshot/service'
 import type {
@@ -740,6 +743,7 @@ export class ExecutionLogger implements IExecutionLoggerService {
       },
       executionId
     )
+    const completedExecutionLargeValueKeys = collectLargeValueReferenceKeys(completedExecutionData)
 
     const updatedLog = await db.transaction(async (tx) => {
       await setExecutionLogWriteTimeouts(tx)
@@ -762,15 +766,15 @@ export class ExecutionLogger implements IExecutionLoggerService {
         throw new Error(`Workflow log not found for execution ${executionId}`)
       }
 
-      await replaceLargeValueReferencesWithClient(
+      await replaceLargeValueReferenceKeysWithClient(
         tx,
         {
           workspaceId: log.workspaceId,
           workflowId: log.workflowId,
           executionId,
           source: 'execution_log',
         },
-        completedExecutionData
+        completedExecutionLargeValueKeys
       )
 
       return log

apps/sim/lib/workflows/executor/human-in-the-loop-manager.ts

@@ -13,7 +13,10 @@ import {
   resetExecutionStreamBuffer,
   type TerminalExecutionStreamStatus,
 } from '@/lib/execution/event-buffer'
-import { replaceLargeValueReferencesWithClient } from '@/lib/execution/payloads/large-value-metadata'
+import {
+  collectLargeValueReferenceKeys,
+  replaceLargeValueReferenceKeysWithClient,
+} from '@/lib/execution/payloads/large-value-metadata'
 import { compactBlockLogs, compactExecutionPayload } from '@/lib/execution/payloads/serializer'
 import { preprocessExecution } from '@/lib/execution/preprocessing'
 import { LoggingSession } from '@/lib/logs/execution/logging-session'
@@ -189,6 +192,9 @@ export class PauseResumeManager {
     const snapshotWorkspaceId = getSnapshotWorkspaceId(
       isRecord(snapshotReferenceValue) ? snapshotReferenceValue.snapshot : undefined
     )
+    const snapshotReferenceKeys = snapshotWorkspaceId
+      ? collectLargeValueReferenceKeys(snapshotReferenceValue, snapshotWorkspaceId)
+      : []
 
     const pausePointsRecord = pausePoints.reduce<Record<string, any>>((acc, point) => {
       acc[point.contextId] = {
@@ -241,15 +247,15 @@ export class PauseResumeManager {
           nextResumeAt,
         })
         if (snapshotWorkspaceId) {
-          await replaceLargeValueReferencesWithClient(
+          await replaceLargeValueReferenceKeysWithClient(
             tx,
             {
               workspaceId: snapshotWorkspaceId,
               workflowId,
               executionId,
               source: 'paused_snapshot',
             },
-            snapshotReferenceValue
+            snapshotReferenceKeys
           )
         }
         return
@@ -298,15 +304,15 @@ export class PauseResumeManager {
         .where(eq(pausedExecutions.id, existing.id))
 
       if (snapshotWorkspaceId) {
-        await replaceLargeValueReferencesWithClient(
+        await replaceLargeValueReferenceKeysWithClient(
           tx,
           {
             workspaceId: snapshotWorkspaceId,
             workflowId,
             executionId,
             source: 'paused_snapshot',
           },
-          snapshotReferenceValue
+          snapshotReferenceKeys
         )
       }
     })
@@ -1614,6 +1620,10 @@ export class PauseResumeManager {
       triggerIds: currentSnapshot.triggerIds,
     }
     const snapshotWorkspaceId = getSnapshotWorkspaceId(snapshotData)
+    const snapshotReferenceValue = { ...updatedSnapshot, snapshot: snapshotData }
+    const snapshotReferenceKeys = snapshotWorkspaceId
+      ? collectLargeValueReferenceKeys(snapshotReferenceValue, snapshotWorkspaceId)
+      : []
 
     await db.transaction(async (tx) => {
       await tx
@@ -1625,15 +1635,15 @@ export class PauseResumeManager {
         .where(eq(pausedExecutions.id, pausedExecutionId))
 
       if (snapshotWorkspaceId) {
-        await replaceLargeValueReferencesWithClient(
+        await replaceLargeValueReferenceKeysWithClient(
           tx,
           {
             workspaceId: snapshotWorkspaceId,
             workflowId: pausedExecution.workflowId,
             executionId: pausedExecution.executionId,
             source: 'paused_snapshot',
           },
-          { ...updatedSnapshot, snapshot: snapshotData }
+          snapshotReferenceKeys
         )
       }
     })

packages/testing/src/mocks/database.mock.ts

@@ -234,20 +234,25 @@ export const dbChainMock = {
  * Creates a mock database connection.
  */
 export function createMockDb() {
+  const fromBuilder = () => ({
+    where: vi.fn(() => ({
+      limit: vi.fn(() => Promise.resolve([])),
+      orderBy: vi.fn(() => Promise.resolve([])),
+    })),
+    leftJoin: vi.fn(() => ({
+      where: vi.fn(() => Promise.resolve([])),
+    })),
+    innerJoin: vi.fn(() => ({
+      where: vi.fn(() => Promise.resolve([])),
+    })),
+  })
+
   return {
     select: vi.fn(() => ({
-      from: vi.fn(() => ({
-        where: vi.fn(() => ({
-          limit: vi.fn(() => Promise.resolve([])),
-          orderBy: vi.fn(() => Promise.resolve([])),
-        })),
-        leftJoin: vi.fn(() => ({
-          where: vi.fn(() => Promise.resolve([])),
-        })),
-        innerJoin: vi.fn(() => ({
-          where: vi.fn(() => Promise.resolve([])),
-        })),
-      })),
+      from: vi.fn(fromBuilder),
+    })),
+    selectDistinct: vi.fn(() => ({
+      from: vi.fn(fromBuilder),
     })),
     insert: vi.fn(() => ({
       values: vi.fn(() => ({

packages/testing/src/mocks/schema.mock.ts

@@ -157,6 +157,29 @@ export const schemaMock = {
     files: 'files',
     createdAt: 'createdAt',
   },
+  executionLargeValues: {
+    key: 'key',
+    workspaceId: 'workspaceId',
+    workflowId: 'workflowId',
+    ownerExecutionId: 'ownerExecutionId',
+    size: 'size',
+    createdAt: 'createdAt',
+    deletedAt: 'deletedAt',
+  },
+  executionLargeValueReferences: {
+    key: 'key',
+    executionId: 'executionId',
+    source: 'source',
+    workspaceId: 'workspaceId',
+    workflowId: 'workflowId',
+    createdAt: 'createdAt',
+  },
+  executionLargeValueDependencies: {
+    parentKey: 'parentKey',
+    childKey: 'childKey',
+    workspaceId: 'workspaceId',
+    createdAt: 'createdAt',
+  },
   pausedExecutions: {
     id: 'id',
     workflowId: 'workflowId',