From 8f7f8f2c35eae5457d293b739c3b2a8adbde9491 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 18 Oct 2026 14:51:52 -0700 Subject: [PATCH 1/2] fix(memory): read memory records and their provenance in one statement --- apps/sim/executor/handlers/agent/memory.ts | 13 +- .../handlers/agent/memory.use-cases.test.ts | 8 +- .../lib/memory/application/use-cases.test.ts | 1 + apps/sim/lib/memory/application/use-cases.ts | 199 ++++++------------ apps/sim/lib/memory/conversation-store.ts | 5 +- .../memory/message-provenance.integration.ts | 80 ++++++- .../sim/lib/memory/message-provenance.test.ts | 4 +- apps/sim/lib/memory/secret-provenance.ts | 12 ++ 8 files changed, 169 insertions(+), 153 deletions(-) diff --git a/apps/sim/executor/handlers/agent/memory.ts b/apps/sim/executor/handlers/agent/memory.ts index 5123addab5a..3c8f1fbe5a7 100644 --- a/apps/sim/executor/handlers/agent/memory.ts +++ b/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] diff --git a/apps/sim/executor/handlers/agent/memory.use-cases.test.ts b/apps/sim/executor/handlers/agent/memory.use-cases.test.ts index 507745bc8fc..8dfd6805a90 100644 --- a/apps/sim/executor/handlers/agent/memory.use-cases.test.ts +++ b/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( diff --git a/apps/sim/lib/memory/application/use-cases.test.ts b/apps/sim/lib/memory/application/use-cases.test.ts index 10f25d50523..2df8dac9ca8 100644 --- a/apps/sim/lib/memory/application/use-cases.test.ts +++ b/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(), })) diff --git a/apps/sim/lib/memory/application/use-cases.ts b/apps/sim/lib/memory/application/use-cases.ts index fcf14cf05f5..6711d5d987e 100644 --- a/apps/sim/lib/memory/application/use-cases.ts +++ b/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>[number] + export interface MemoryRecord { id: string key: string @@ -103,112 +117,45 @@ function memoryMessageError(data: unknown): string | null { return null } -async function loadReadProvenance( - records: MemoryRecord[], - signal?: AbortSignal -): Promise { - if (records.length === 0) return [] - - const recordsById = new Map() - 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 + const records: MemoryRecord[] = [] + const readProvenance: MemoryReadProvenance[] = [] const tailBudget = new PlainMemoryReadBudget() - const tailProvenance: Array = [] - 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) { + 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 +175,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 +192,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 +261,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, diff --git a/apps/sim/lib/memory/conversation-store.ts b/apps/sim/lib/memory/conversation-store.ts index 1a2c50330f1..c2f0af67472 100644 --- a/apps/sim/lib/memory/conversation-store.ts +++ b/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)) diff --git a/apps/sim/lib/memory/message-provenance.integration.ts b/apps/sim/lib/memory/message-provenance.integration.ts index a706f483f58..00ff4c2d659 100644 --- a/apps/sim/lib/memory/message-provenance.integration.ts +++ b/apps/sim/lib/memory/message-provenance.integration.ts @@ -52,8 +52,13 @@ vi.mock('@/lib/workspaces/application/context', () => ({ }), })) +import type { BillingAttributionSnapshot } from '@/lib/billing/core/attribution' import { hashDurableSecretProvenanceValue } from '@/lib/execution/durable-secret-provenance' -import { appendMemoryUseCase } from '@/lib/memory/application/use-cases' +import { + appendMemoryUseCase, + listMemoriesUseCase, + readMemoryUseCase, +} from '@/lib/memory/application/use-cases' import { Memory } from '@/executor/handlers/agent/memory' import type { AgentInputs } from '@/executor/handlers/agent/types' import type { ExecutionContext } from '@/executor/types' @@ -244,6 +249,79 @@ describe('memory provenance in PostgreSQL', () => { expect(execution.resolvedSecretTraceRegistry?.isComplete()).toBe(true) }) + it.each(['list', 'read', 'append'] as const)( + 'binds a %s to the revision it read while another writer appends', + async (mode) => { + const key = `concurrent-${mode}-read` + await toolAppend(key, 'A') + // An append's readback includes its own message; every other read sees only the first. + const expected = mode === 'append' ? ['TOKEN_A', 'TOKEN_C'] : ['TOKEN_A'] + let appended = false + const now = Date.now() + const input = { + workspaceId: SCOPE.workspaceId, + includePersistedSecretProvenance: true, + // Lands after the read is taken, where a separate sidecar read once paired the old + // record with this append's provenance. + resolveBillingAttribution: async (): Promise => { + await toolAppend(key, 'B') + appended = true + return { + actorUserId: SCOPE.userId, + workspaceId: SCOPE.workspaceId, + billedAccountUserId: SCOPE.userId, + organizationId: null, + billingEntity: { type: 'user', id: SCOPE.userId }, + payerSubscription: null, + billingPeriod: { + start: new Date(now - 86_400_000).toISOString(), + end: new Date(now + 86_400_000).toISOString(), + }, + } + }, + } + const result = + mode === 'list' + ? await listMemoriesUseCase.execute({ + principal: principal(), + input: { ...input, query: key, limit: 10 }, + }) + : mode === 'read' + ? await readMemoryUseCase.execute({ + principal: principal(), + input: { ...input, key }, + }) + : await appendMemoryUseCase.execute({ + principal: principal(), + input: { + ...input, + key, + data: { role: 'user', content: 'secret-C' }, + writeProvenance: { + status: 'exact', + entries: [ + { + encryptedValue: 'cipher-C', + name: 'TOKEN_C', + sourceUserId: SCOPE.userId, + sourceWorkspaceId: SCOPE.workspaceId, + }, + ], + }, + }, + }) + expect(appended).toBe(true) + expect(result.readProvenance).toHaveLength(1) + const [read] = result.readProvenance ?? [] + expect(read.data).toHaveLength(expected.length) + expect(read.provenance.status).toBe('exact') + expect( + read.provenance.status === 'exact' && + read.provenance.entries.map(({ name }) => name).sort() + ).toEqual(expected) + } + ) + it.each(['tool-tool', 'tool-native', 'native-native'] as const)( 'preserves both first appends and secret bindings for %s', async (mode) => { diff --git a/apps/sim/lib/memory/message-provenance.test.ts b/apps/sim/lib/memory/message-provenance.test.ts index 13970743255..3d3d84c7c18 100644 --- a/apps/sim/lib/memory/message-provenance.test.ts +++ b/apps/sim/lib/memory/message-provenance.test.ts @@ -73,8 +73,8 @@ function queueStoredMemory(data: Message[], entries: readonly DurableSecretProve data, secretProvenanceVersion: 1, provenanceContentHash: hashDurableSecretProvenanceValue(data), - provenanceStatus: 'exact', - provenanceEntries: entries, + status: 'exact', + entries, }, ]) } diff --git a/apps/sim/lib/memory/secret-provenance.ts b/apps/sim/lib/memory/secret-provenance.ts index 22d200ca4bf..e4ebef664f7 100644 --- a/apps/sim/lib/memory/secret-provenance.ts +++ b/apps/sim/lib/memory/secret-provenance.ts @@ -162,6 +162,18 @@ export async function bindMemorySecretProvenanceToMessages( return normalized ? { status: 'exact', entries: normalized } : { status: 'unknown' } } +/** + * A memory row's sidecar, selected beside it with + * `leftJoin(memorySecretProvenance, eq(memorySecretProvenance.memoryId, memory.id))`. The writer + * binds both in one transaction, so a reader takes both in one statement: a second read can pair + * the row with a later writer's sidecar, and that hash mismatch reads as unknown. + */ +export const MEMORY_SECRET_PROVENANCE_COLUMNS = { + provenanceContentHash: memorySecretProvenance.contentHash, + status: memorySecretProvenance.status, + entries: memorySecretProvenance.entries, +} + interface MemorySecretProvenanceRow { secretProvenanceVersion: number | null data: unknown From 075c4b603329caaa6611db4a89727d5bba965ff5 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sun, 18 Oct 2026 15:02:42 -0700 Subject: [PATCH 2/2] fix(memory): stop cancelled memory reads before tails and provenance --- .../lib/memory/application/use-cases.test.ts | 30 +++++++++++++++++++ apps/sim/lib/memory/application/use-cases.ts | 2 ++ 2 files changed, 32 insertions(+) diff --git a/apps/sim/lib/memory/application/use-cases.test.ts b/apps/sim/lib/memory/application/use-cases.test.ts index 2df8dac9ca8..82b7d800abd 100644 --- a/apps/sim/lib/memory/application/use-cases.test.ts +++ b/apps/sim/lib/memory/application/use-cases.test.ts @@ -157,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, diff --git a/apps/sim/lib/memory/application/use-cases.ts b/apps/sim/lib/memory/application/use-cases.ts index 6711d5d987e..eb34201c131 100644 --- a/apps/sim/lib/memory/application/use-cases.ts +++ b/apps/sim/lib/memory/application/use-cases.ts @@ -137,6 +137,7 @@ async function readMemoryResults( input.resolveBillingAttribution ))) : undefined + input.signal?.throwIfAborted() const records: MemoryRecord[] = [] const readProvenance: MemoryReadProvenance[] = [] const tailBudget = new PlainMemoryReadBudget() @@ -145,6 +146,7 @@ async function readMemoryResults( ? 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]),