fix(memory): read memory records and their provenance in one statement

apps/sim/executor/handlers/agent/memory.ts

@@ -40,6 +40,7 @@ import { AGENT_MEMORY_RETRIEVAL_TOOL_ID } from '@/lib/memory/retrieval-tool-type
 import {
   bindMemorySecretProvenanceToMessages,
   createMemorySecretProvenanceSelector,
+  MEMORY_SECRET_PROVENANCE_COLUMNS,
   readBoundMemorySecretProvenance,
 } from '@/lib/memory/secret-provenance'
 import type { AgentInputs, FileNameProjection, Message } from '@/executor/handlers/agent/types'
@@ -674,9 +675,7 @@ export class Memory {
             storageVersion: memory.storageVersion,
             data: memory.data,
             secretProvenanceVersion: memory.secretProvenanceVersion,
-            provenanceContentHash: memorySecretProvenance.contentHash,
-            provenanceStatus: memorySecretProvenance.status,
-            provenanceEntries: memorySecretProvenance.entries,
+            ...MEMORY_SECRET_PROVENANCE_COLUMNS,
           })
           .from(memory)
           .leftJoin(memorySecretProvenance, eq(memorySecretProvenance.memoryId, memory.id))
@@ -689,13 +688,7 @@ export class Memory {
     }
 
     let data = row.data
-    let provenance = readBoundMemorySecretProvenance({
-      secretProvenanceVersion: row.secretProvenanceVersion,
-      data,
-      provenanceContentHash: row.provenanceContentHash,
-      status: row.provenanceStatus,
-      entries: row.provenanceEntries,
-    })
+    let provenance = readBoundMemorySecretProvenance(row)
     if (row.storageVersion === 2 && !options.richHistory) {
       const tail = await readPlainMemoryTail(row.id, workspaceId)
       data = [...(Array.isArray(data) ? data : []), ...tail.messages]

apps/sim/executor/handlers/agent/memory.use-cases.test.ts

@@ -51,8 +51,8 @@ describe('optional Agent memory durability failures', () => {
       data: prefix,
       secretProvenanceVersion: 1,
       provenanceContentHash: hashDurableSecretProvenanceValue(prefix),
-      provenanceStatus: 'exact',
-      provenanceEntries: [],
+      status: 'exact',
+      entries: [],
     })
     mocks.items.mockResolvedValue({ items: [] })
     mocks.append.mockResolvedValue(undefined)
@@ -69,8 +69,8 @@ describe('optional Agent memory durability failures', () => {
       data: history,
       secretProvenanceVersion: 1,
       provenanceContentHash: hashDurableSecretProvenanceValue(history),
-      provenanceStatus: 'exact',
-      provenanceEntries: [],
+      status: 'exact',
+      entries: [],
     })
     const memory = new Memory()
     await expect(

apps/sim/lib/memory/application/use-cases.test.ts

@@ -24,6 +24,7 @@ vi.mock('@sim/platform-authz/workspace', () => workspaceAuthzMock)
 vi.mock('@/lib/billing/core/attribution', () => billingAttributionMock)
 
 vi.mock('@/lib/memory/secret-provenance', () => ({
+  MEMORY_SECRET_PROVENANCE_COLUMNS: {},
   readBoundMemorySecretProvenance: hoisted.readBoundProvenance,
   replaceMemorySecretProvenanceInTx: vi.fn(),
 }))
@@ -156,6 +157,36 @@ describe('Memory application use cases', () => {
     )
     expect(mocks.readBoundProvenance).not.toHaveBeenCalled()
   })
+  it('stops a list cancelled during its billing lookup before reading tails or provenance', async () => {
+    queueTableRows(schemaMock.memory, [
+      {
+        id: 'memory-1',
+        key: 'conversation-1',
+        data: [],
+        storageVersion: 2,
+        secretProvenanceVersion: null,
+      },
+    ])
+    const controller = new AbortController()
+    await expect(
+      listMemoriesUseCase.execute({
+        principal: ACTORLESS_DEPLOYED_PRINCIPAL,
+        input: {
+          workspaceId: WORKSPACE_ID,
+          limit: 50,
+          includePersistedSecretProvenance: true,
+          signal: controller.signal,
+          resolveBillingAttribution: async () => {
+            controller.abort(new Error('cancelled'))
+            return BILLING_ATTRIBUTION
+          },
+        },
+      })
+    ).rejects.toThrow('cancelled')
+    expect(mocks.readPlainMemoryTail).not.toHaveBeenCalled()
+    expect(mocks.readBoundProvenance).not.toHaveBeenCalled()
+  })
+
   it('enforces one appended-history budget across the entire list response', async () => {
     queueTableRows(
       schemaMock.memory,

apps/sim/lib/memory/application/use-cases.ts

@@ -7,7 +7,7 @@ import { db } from '@sim/db'
 import { memory, memorySecretProvenance } from '@sim/db/schema'
 import { getPostgresErrorCode } from '@sim/utils/errors'
 import { generateId } from '@sim/utils/id'
-import { and, eq, inArray, isNull, like } from 'drizzle-orm'
+import { and, eq, isNull, like, type SQL } from 'drizzle-orm'
 import type { BillingAttributionSnapshot } from '@/lib/billing/core/attribution'
 import { assertBillingAttributionSnapshot } from '@/lib/billing/core/attribution'
 import { defineAuthorizedWorkspaceUseCase } from '@/lib/core/application'
@@ -21,10 +21,12 @@ import { memoryOperations } from '@/lib/memory/application/operations'
 import { appendMemoryMessages, readPlainMemoryTail } from '@/lib/memory/conversation-store'
 import { lockMemoryConversationInTx } from '@/lib/memory/locks'
 import { PlainMemoryReadBudget } from '@/lib/memory/read-budget'
-import { readBoundMemorySecretProvenance } from '@/lib/memory/secret-provenance'
+import {
+  MEMORY_SECRET_PROVENANCE_COLUMNS,
+  readBoundMemorySecretProvenance,
+} from '@/lib/memory/secret-provenance'
 import { resolveActiveWorkspaceApplicationContext } from '@/lib/workspaces/application/context'
 
-const PRIVATE_MEMORY_QUERY_CHUNK_SIZE = 1_000
 const MAX_MEMORY_LIST_LIMIT = 1_000
 const MEMORY_READ_COLUMNS = {
   id: memory.id,
@@ -38,6 +40,18 @@ const MEMORY_READ_COLUMNS = {
   deletedAt: memory.deletedAt,
 }
 
+function selectMemoryRows(where: SQL | undefined, limit: number) {
+  return db
+    .select({ ...MEMORY_READ_COLUMNS, ...MEMORY_SECRET_PROVENANCE_COLUMNS })
+    .from(memory)
+    .leftJoin(memorySecretProvenance, eq(memorySecretProvenance.memoryId, memory.id))
+    .where(where)
+    .orderBy(memory.createdAt)
+    .limit(limit)
+}
+
+type MemoryReadRow = Awaited<ReturnType<typeof selectMemoryRows>>[number]
+
 export interface MemoryRecord {
   id: string
   key: string
@@ -103,112 +117,47 @@ function memoryMessageError(data: unknown): string | null {
   return null
 }
 
-async function loadReadProvenance(
-  records: MemoryRecord[],
-  signal?: AbortSignal
-): Promise<MemoryReadProvenance[]> {
-  if (records.length === 0) return []
-
-  const recordsById = new Map<string, MemoryRecord[]>()
-  for (const record of records) {
-    const matching = recordsById.get(record.id) ?? []
-    matching.push(record)
-    recordsById.set(record.id, matching)
-  }
-
-  const result: MemoryReadProvenance[] = []
-  const ids = [...recordsById.keys()]
-
-  for (let index = 0; index < ids.length; index += PRIVATE_MEMORY_QUERY_CHUNK_SIZE) {
-    signal?.throwIfAborted()
-    const pageIds = ids.slice(index, index + PRIVATE_MEMORY_QUERY_CHUNK_SIZE)
-    const sidecars = await db
-      .select()
-      .from(memorySecretProvenance)
-      .where(inArray(memorySecretProvenance.memoryId, pageIds))
-    const sidecarById = new Map(sidecars.map((sidecar) => [sidecar.memoryId, sidecar]))
-
-    for (const memoryId of pageIds) {
-      for (const record of recordsById.get(memoryId) ?? []) {
-        const sidecar = sidecarById.get(memoryId)
-        const provenance = readBoundMemorySecretProvenance({
-          secretProvenanceVersion: record.secretProvenanceVersion,
-          data: record.data,
-          provenanceContentHash: sidecar?.contentHash ?? null,
-          status: sidecar?.status ?? null,
-          entries: sidecar?.entries,
-        })
-        result.push({ data: record.data, provenance })
-      }
-    }
-  }
-
-  return result
-}
-
-async function readResultProvenance(
-  records: MemoryRecord[],
+/** Keeps the legacy public projection separate from private Agent exchanges and checkpoints. */
+async function readMemoryResults(
+  rows: MemoryReadRow[],
   principal: Principal,
   workspaceId: string,
   input: ReadProvenanceInput,
   existingScope?: MemoryLegacyProvenanceScope
 ): Promise<{
+  records: MemoryRecord[]
   readProvenance?: MemoryReadProvenance[]
   provenanceScope?: MemoryLegacyProvenanceScope
 }> {
-  if (!input.includePersistedSecretProvenance) return {}
-  const provenanceScope =
-    existingScope ??
-    (await resolveMemoryLegacyProvenanceScope(
-      principal,
-      workspaceId,
-      input.resolveBillingAttribution
-    ))
-  return {
-    readProvenance: await loadReadProvenance(records, input.signal),
-    provenanceScope,
-  }
-}
-
-/** Keeps the legacy public projection separate from private Agent exchanges and checkpoints. */
-async function readMemoryResults(
-  records: MemoryRecord[],
-  principal: Principal,
-  workspaceId: string,
-  input: ReadProvenanceInput,
-  existingScope?: MemoryLegacyProvenanceScope
-) {
-  const projected: MemoryRecord[] = []
+  const provenanceScope = input.includePersistedSecretProvenance
+    ? (existingScope ??
+      (await resolveMemoryLegacyProvenanceScope(
+        principal,
+        workspaceId,
+        input.resolveBillingAttribution
+      )))
+    : undefined
+  input.signal?.throwIfAborted()
+  const records: MemoryRecord[] = []
+  const readProvenance: MemoryReadProvenance[] = []
   const tailBudget = new PlainMemoryReadBudget()
-  const tailProvenance: Array<DurableSecretProvenance | undefined> = []
-  for (const record of records) {
-    if (record.storageVersion !== 2) {
-      projected.push(record)
-      tailProvenance.push(undefined)
-      continue
+  for (const { provenanceContentHash, status, entries, ...record } of rows) {
+    let provenance = provenanceScope
+      ? readBoundMemorySecretProvenance({ ...record, provenanceContentHash, status, entries })
+      : undefined
+    if (record.storageVersion === 2) {
+      input.signal?.throwIfAborted()
+      const tail = await readPlainMemoryTail(record.id, workspaceId, tailBudget)
+      record.data = [
+        ...(Array.isArray(record.data) ? record.data : [record.data]),
+        ...tail.messages,
+      ]
+      if (provenance) provenance = mergeDurableSecretProvenance(provenance, tail.provenance)
     }
-    const tail = await readPlainMemoryTail(record.id, workspaceId, tailBudget)
-    projected.push({
-      ...record,
-      data: [...(Array.isArray(record.data) ? record.data : [record.data]), ...tail.messages],
-    })
-    tailProvenance.push(tail.provenance)
-  }
-  const result = await readResultProvenance(records, principal, workspaceId, input, existingScope)
-  return {
-    records: projected,
-    ...result,
-    ...(result.readProvenance
-      ? {
-          readProvenance: result.readProvenance.map((entry, index) => ({
-            data: projected[index].data,
-            provenance: tailProvenance[index]
-              ? mergeDurableSecretProvenance(entry.provenance, tailProvenance[index])
-              : entry.provenance,
-          })),
-        }
-      : {}),
+    records.push(record)
+    if (provenance) readProvenance.push({ data: record.data, provenance })
   }
+  return provenanceScope ? { records, readProvenance, provenanceScope } : { records }
 }
 
 export interface ListMemoriesInput extends WorkspaceInput, ReadProvenanceInput {
@@ -228,14 +177,9 @@ export const listMemoriesUseCase = defineAuthorizedWorkspaceUseCase({
     input.signal?.throwIfAborted()
     const conditions = [isNull(memory.deletedAt), eq(memory.workspaceId, context.workspaceId)]
     if (input.query) conditions.push(like(memory.key, `%${input.query}%`))
-    const records = await db
-      .select(MEMORY_READ_COLUMNS)
-      .from(memory)
-      .where(and(...conditions))
-      .orderBy(memory.createdAt)
-      .limit(input.limit)
+    const rows = await selectMemoryRows(and(...conditions), input.limit)
     input.signal?.throwIfAborted()
-    return readMemoryResults(records, principal, context.workspaceId, input)
+    return readMemoryResults(rows, principal, context.workspaceId, input)
   },
 })
 
@@ -250,21 +194,17 @@ export const readMemoryUseCase = defineAuthorizedWorkspaceUseCase({
   authorizationOptions: { delegation: memoryDelegationPolicy },
   async execute({ principal, input, context }) {
     input.signal?.throwIfAborted()
-    const records = await db
-      .select(MEMORY_READ_COLUMNS)
-      .from(memory)
-      .where(
-        and(
-          eq(memory.key, input.key),
-          eq(memory.workspaceId, context.workspaceId),
-          isNull(memory.deletedAt)
-        )
-      )
-      .orderBy(memory.createdAt)
-      .limit(1)
+    const rows = await selectMemoryRows(
+      and(
+        eq(memory.key, input.key),
+        eq(memory.workspaceId, context.workspaceId),
+        isNull(memory.deletedAt)
+      ),
+      1
+    )
     input.signal?.throwIfAborted()
     const { records: projected, ...provenance } = await readMemoryResults(
-      records,
+      rows,
       principal,
       context.workspaceId,
       input
@@ -323,23 +263,18 @@ export const appendMemoryUseCase = defineAuthorizedWorkspaceUseCase({
     }
 
     input.signal?.throwIfAborted()
-    const records = await db
-      .select(MEMORY_READ_COLUMNS)
-      .from(memory)
-      .where(
-        and(
-          eq(memory.key, input.key),
-          eq(memory.workspaceId, context.workspaceId),
-          isNull(memory.deletedAt)
-        )
-      )
-      .orderBy(memory.createdAt)
-      .limit(1)
-    const record = records[0]
-    if (!record) throw new Error('Failed to retrieve memory after creation/update')
+    const rows = await selectMemoryRows(
+      and(
+        eq(memory.key, input.key),
+        eq(memory.workspaceId, context.workspaceId),
+        isNull(memory.deletedAt)
+      ),
+      1
+    )
+    if (!rows[0]) throw new Error('Failed to retrieve memory after creation/update')
     input.signal?.throwIfAborted()
     const { records: projected, ...provenance } = await readMemoryResults(
-      records,
+      rows,
       principal,
       context.workspaceId,
       input,

apps/sim/lib/memory/conversation-store.ts

@@ -15,6 +15,7 @@ import { lockMemoryConversationInTx } from '@/lib/memory/locks'
 import { MAX_RICH_MEMORY_PAGE_BYTES, PlainMemoryReadBudget } from '@/lib/memory/read-budget'
 import {
   bindMemorySecretProvenanceToMessages,
+  MEMORY_SECRET_PROVENANCE_COLUMNS,
   readBoundMemorySecretProvenance,
   replaceMemorySecretProvenanceInTx,
 } from '@/lib/memory/secret-provenance'
@@ -234,9 +235,7 @@ export async function readConversationPrefix(input: ReadConversationPrefixInput)
       storageVersion: memory.storageVersion,
       data: memory.data,
       secretProvenanceVersion: memory.secretProvenanceVersion,
-      provenanceContentHash: memorySecretProvenance.contentHash,
-      provenanceStatus: memorySecretProvenance.status,
-      provenanceEntries: memorySecretProvenance.entries,
+      ...MEMORY_SECRET_PROVENANCE_COLUMNS,
     })
     .from(memory)
     .leftJoin(memorySecretProvenance, eq(memorySecretProvenance.memoryId, memory.id))