@@ -7,7 +7,7 @@ import { db } from '@sim/db'
77import { memory , memorySecretProvenance } from '@sim/db/schema'
88import { getPostgresErrorCode } from '@sim/utils/errors'
99import { generateId } from '@sim/utils/id'
10- import { and , eq , inArray , isNull , like } from 'drizzle-orm'
10+ import { and , eq , isNull , like , type SQL } from 'drizzle-orm'
1111import type { BillingAttributionSnapshot } from '@/lib/billing/core/attribution'
1212import { assertBillingAttributionSnapshot } from '@/lib/billing/core/attribution'
1313import { defineAuthorizedWorkspaceUseCase } from '@/lib/core/application'
@@ -21,10 +21,12 @@ import { memoryOperations } from '@/lib/memory/application/operations'
2121import { appendMemoryMessages , readPlainMemoryTail } from '@/lib/memory/conversation-store'
2222import { lockMemoryConversationInTx } from '@/lib/memory/locks'
2323import { PlainMemoryReadBudget } from '@/lib/memory/read-budget'
24- import { readBoundMemorySecretProvenance } from '@/lib/memory/secret-provenance'
24+ import {
25+ MEMORY_SECRET_PROVENANCE_COLUMNS ,
26+ readBoundMemorySecretProvenance ,
27+ } from '@/lib/memory/secret-provenance'
2528import { resolveActiveWorkspaceApplicationContext } from '@/lib/workspaces/application/context'
2629
27- const PRIVATE_MEMORY_QUERY_CHUNK_SIZE = 1_000
2830const MAX_MEMORY_LIST_LIMIT = 1_000
2931const MEMORY_READ_COLUMNS = {
3032 id : memory . id ,
@@ -38,6 +40,18 @@ const MEMORY_READ_COLUMNS = {
3840 deletedAt : memory . deletedAt ,
3941}
4042
43+ function selectMemoryRows ( where : SQL | undefined , limit : number ) {
44+ return db
45+ . select ( { ...MEMORY_READ_COLUMNS , ...MEMORY_SECRET_PROVENANCE_COLUMNS } )
46+ . from ( memory )
47+ . leftJoin ( memorySecretProvenance , eq ( memorySecretProvenance . memoryId , memory . id ) )
48+ . where ( where )
49+ . orderBy ( memory . createdAt )
50+ . limit ( limit )
51+ }
52+
53+ type MemoryReadRow = Awaited < ReturnType < typeof selectMemoryRows > > [ number ]
54+
4155export interface MemoryRecord {
4256 id : string
4357 key : string
@@ -103,112 +117,47 @@ function memoryMessageError(data: unknown): string | null {
103117 return null
104118}
105119
106- async function loadReadProvenance (
107- records : MemoryRecord [ ] ,
108- signal ?: AbortSignal
109- ) : Promise < MemoryReadProvenance [ ] > {
110- if ( records . length === 0 ) return [ ]
111-
112- const recordsById = new Map < string , MemoryRecord [ ] > ( )
113- for ( const record of records ) {
114- const matching = recordsById . get ( record . id ) ?? [ ]
115- matching . push ( record )
116- recordsById . set ( record . id , matching )
117- }
118-
119- const result : MemoryReadProvenance [ ] = [ ]
120- const ids = [ ...recordsById . keys ( ) ]
121-
122- for ( let index = 0 ; index < ids . length ; index += PRIVATE_MEMORY_QUERY_CHUNK_SIZE ) {
123- signal ?. throwIfAborted ( )
124- const pageIds = ids . slice ( index , index + PRIVATE_MEMORY_QUERY_CHUNK_SIZE )
125- const sidecars = await db
126- . select ( )
127- . from ( memorySecretProvenance )
128- . where ( inArray ( memorySecretProvenance . memoryId , pageIds ) )
129- const sidecarById = new Map ( sidecars . map ( ( sidecar ) => [ sidecar . memoryId , sidecar ] ) )
130-
131- for ( const memoryId of pageIds ) {
132- for ( const record of recordsById . get ( memoryId ) ?? [ ] ) {
133- const sidecar = sidecarById . get ( memoryId )
134- const provenance = readBoundMemorySecretProvenance ( {
135- secretProvenanceVersion : record . secretProvenanceVersion ,
136- data : record . data ,
137- provenanceContentHash : sidecar ?. contentHash ?? null ,
138- status : sidecar ?. status ?? null ,
139- entries : sidecar ?. entries ,
140- } )
141- result . push ( { data : record . data , provenance } )
142- }
143- }
144- }
145-
146- return result
147- }
148-
149- async function readResultProvenance (
150- records : MemoryRecord [ ] ,
120+ /** Keeps the legacy public projection separate from private Agent exchanges and checkpoints. */
121+ async function readMemoryResults (
122+ rows : MemoryReadRow [ ] ,
151123 principal : Principal ,
152124 workspaceId : string ,
153125 input : ReadProvenanceInput ,
154126 existingScope ?: MemoryLegacyProvenanceScope
155127) : Promise < {
128+ records : MemoryRecord [ ]
156129 readProvenance ?: MemoryReadProvenance [ ]
157130 provenanceScope ?: MemoryLegacyProvenanceScope
158131} > {
159- if ( ! input . includePersistedSecretProvenance ) return { }
160- const provenanceScope =
161- existingScope ??
162- ( await resolveMemoryLegacyProvenanceScope (
163- principal ,
164- workspaceId ,
165- input . resolveBillingAttribution
166- ) )
167- return {
168- readProvenance : await loadReadProvenance ( records , input . signal ) ,
169- provenanceScope,
170- }
171- }
172-
173- /** Keeps the legacy public projection separate from private Agent exchanges and checkpoints. */
174- async function readMemoryResults (
175- records : MemoryRecord [ ] ,
176- principal : Principal ,
177- workspaceId : string ,
178- input : ReadProvenanceInput ,
179- existingScope ?: MemoryLegacyProvenanceScope
180- ) {
181- const projected : MemoryRecord [ ] = [ ]
132+ const provenanceScope = input . includePersistedSecretProvenance
133+ ? ( existingScope ??
134+ ( await resolveMemoryLegacyProvenanceScope (
135+ principal ,
136+ workspaceId ,
137+ input . resolveBillingAttribution
138+ ) ) )
139+ : undefined
140+ input . signal ?. throwIfAborted ( )
141+ const records : MemoryRecord [ ] = [ ]
142+ const readProvenance : MemoryReadProvenance [ ] = [ ]
182143 const tailBudget = new PlainMemoryReadBudget ( )
183- const tailProvenance : Array < DurableSecretProvenance | undefined > = [ ]
184- for ( const record of records ) {
185- if ( record . storageVersion !== 2 ) {
186- projected . push ( record )
187- tailProvenance . push ( undefined )
188- continue
144+ for ( const { provenanceContentHash, status, entries, ...record } of rows ) {
145+ let provenance = provenanceScope
146+ ? readBoundMemorySecretProvenance ( { ...record , provenanceContentHash, status, entries } )
147+ : undefined
148+ if ( record . storageVersion === 2 ) {
149+ input . signal ?. throwIfAborted ( )
150+ const tail = await readPlainMemoryTail ( record . id , workspaceId , tailBudget )
151+ record . data = [
152+ ...( Array . isArray ( record . data ) ? record . data : [ record . data ] ) ,
153+ ...tail . messages ,
154+ ]
155+ if ( provenance ) provenance = mergeDurableSecretProvenance ( provenance , tail . provenance )
189156 }
190- const tail = await readPlainMemoryTail ( record . id , workspaceId , tailBudget )
191- projected . push ( {
192- ...record ,
193- data : [ ...( Array . isArray ( record . data ) ? record . data : [ record . data ] ) , ...tail . messages ] ,
194- } )
195- tailProvenance . push ( tail . provenance )
196- }
197- const result = await readResultProvenance ( records , principal , workspaceId , input , existingScope )
198- return {
199- records : projected ,
200- ...result ,
201- ...( result . readProvenance
202- ? {
203- readProvenance : result . readProvenance . map ( ( entry , index ) => ( {
204- data : projected [ index ] . data ,
205- provenance : tailProvenance [ index ]
206- ? mergeDurableSecretProvenance ( entry . provenance , tailProvenance [ index ] )
207- : entry . provenance ,
208- } ) ) ,
209- }
210- : { } ) ,
157+ records . push ( record )
158+ if ( provenance ) readProvenance . push ( { data : record . data , provenance } )
211159 }
160+ return provenanceScope ? { records, readProvenance, provenanceScope } : { records }
212161}
213162
214163export interface ListMemoriesInput extends WorkspaceInput , ReadProvenanceInput {
@@ -228,14 +177,9 @@ export const listMemoriesUseCase = defineAuthorizedWorkspaceUseCase({
228177 input . signal ?. throwIfAborted ( )
229178 const conditions = [ isNull ( memory . deletedAt ) , eq ( memory . workspaceId , context . workspaceId ) ]
230179 if ( input . query ) conditions . push ( like ( memory . key , `%${ input . query } %` ) )
231- const records = await db
232- . select ( MEMORY_READ_COLUMNS )
233- . from ( memory )
234- . where ( and ( ...conditions ) )
235- . orderBy ( memory . createdAt )
236- . limit ( input . limit )
180+ const rows = await selectMemoryRows ( and ( ...conditions ) , input . limit )
237181 input . signal ?. throwIfAborted ( )
238- return readMemoryResults ( records , principal , context . workspaceId , input )
182+ return readMemoryResults ( rows , principal , context . workspaceId , input )
239183 } ,
240184} )
241185
@@ -250,21 +194,17 @@ export const readMemoryUseCase = defineAuthorizedWorkspaceUseCase({
250194 authorizationOptions : { delegation : memoryDelegationPolicy } ,
251195 async execute ( { principal, input, context } ) {
252196 input . signal ?. throwIfAborted ( )
253- const records = await db
254- . select ( MEMORY_READ_COLUMNS )
255- . from ( memory )
256- . where (
257- and (
258- eq ( memory . key , input . key ) ,
259- eq ( memory . workspaceId , context . workspaceId ) ,
260- isNull ( memory . deletedAt )
261- )
262- )
263- . orderBy ( memory . createdAt )
264- . limit ( 1 )
197+ const rows = await selectMemoryRows (
198+ and (
199+ eq ( memory . key , input . key ) ,
200+ eq ( memory . workspaceId , context . workspaceId ) ,
201+ isNull ( memory . deletedAt )
202+ ) ,
203+ 1
204+ )
265205 input . signal ?. throwIfAborted ( )
266206 const { records : projected , ...provenance } = await readMemoryResults (
267- records ,
207+ rows ,
268208 principal ,
269209 context . workspaceId ,
270210 input
@@ -323,23 +263,18 @@ export const appendMemoryUseCase = defineAuthorizedWorkspaceUseCase({
323263 }
324264
325265 input . signal ?. throwIfAborted ( )
326- const records = await db
327- . select ( MEMORY_READ_COLUMNS )
328- . from ( memory )
329- . where (
330- and (
331- eq ( memory . key , input . key ) ,
332- eq ( memory . workspaceId , context . workspaceId ) ,
333- isNull ( memory . deletedAt )
334- )
335- )
336- . orderBy ( memory . createdAt )
337- . limit ( 1 )
338- const record = records [ 0 ]
339- if ( ! record ) throw new Error ( 'Failed to retrieve memory after creation/update' )
266+ const rows = await selectMemoryRows (
267+ and (
268+ eq ( memory . key , input . key ) ,
269+ eq ( memory . workspaceId , context . workspaceId ) ,
270+ isNull ( memory . deletedAt )
271+ ) ,
272+ 1
273+ )
274+ if ( ! rows [ 0 ] ) throw new Error ( 'Failed to retrieve memory after creation/update' )
340275 input . signal ?. throwIfAborted ( )
341276 const { records : projected , ...provenance } = await readMemoryResults (
342- records ,
277+ rows ,
343278 principal ,
344279 context . workspaceId ,
345280 input ,
0 commit comments