From db4030e8915219693edf521a4a6fc28ef1e704e6 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Mon, 14 Sep 2026 13:37:40 -0700 Subject: [PATCH] fix(agent): retain conversation memory attachments --- .github/workflows/test-build.yml | 4 +- .../content/docs/academy/agents/memory.mdx | 2 + apps/sim/executor/constants.ts | 1 + .../handlers/agent/agent-handler.test.ts | 166 ++++++ .../executor/handlers/agent/agent-handler.ts | 90 ++- .../agent/memory-harness.postgres.test.ts | 547 ++++++++++++++++++ .../executor/handlers/agent/memory.test.ts | 185 +++++- apps/sim/executor/handlers/agent/memory.ts | 65 ++- apps/sim/executor/handlers/agent/types.ts | 6 + .../sim/lib/memory/message-provenance.test.ts | 100 +++- 10 files changed, 1099 insertions(+), 67 deletions(-) create mode 100644 apps/sim/executor/handlers/agent/memory-harness.postgres.test.ts diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index 6ad195ff86d..ab28f35e8a5 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -170,15 +170,17 @@ jobs: if-no-files-found: ignore retention-days: 7 - - name: Verify durable provenance bindings and concurrent memory writes + - name: Verify durable provenance, concurrent memory writes, and attachment replay working-directory: apps/sim env: TABLE_PROVENANCE_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_auth_scim MEMORY_PROVENANCE_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_auth_scim + AGENT_MEMORY_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_auth_scim run: >- bunx vitest run lib/table/rows/secret-provenance.postgres.test.ts lib/memory/message-provenance.postgres.test.ts + executor/handlers/agent/memory-harness.postgres.test.ts - name: Verify Search progress, pagination, and outbox scheduling in PostgreSQL working-directory: apps/sim diff --git a/apps/docs/content/docs/academy/agents/memory.mdx b/apps/docs/content/docs/academy/agents/memory.mdx index 483a70c404d..48a11d441d7 100644 --- a/apps/docs/content/docs/academy/agents/memory.mdx +++ b/apps/docs/content/docs/academy/agents/memory.mdx @@ -18,6 +18,8 @@ import { AV_MEMORY_WORKFLOW } from '@/components/workflow-preview/academy-video- By default, an agent keeps nothing between runs: every conversation starts completely fresh. The Memory setting changes that: choose Conversation, give it a conversation ID, and everything said under that key is kept and loaded back before the model runs. +Uploaded attachments stay linked to the message that included them. Memory stores file references; each later run reads the accessible files again and prepares them for the selected provider. Attachments follow the selected memory window and the source file's storage retention. A replay can include up to 20 attachment references; use a smaller memory window for longer file-heavy conversations. Files omitted by older versions of memory need to be attached again. + { }) }) + describe('conversation attachment replay', () => { + beforeEach(() => { + dbChainMockFns.returning.mockResolvedValue([{ id: 'memory-1' }]) + }) + + afterEach(() => { + vi.restoreAllMocks() + }) + + const file = { + id: 'file-1', + name: 'example.png', + key: 'execution/test-workspace/test-workflow/exec-1/example.png', + url: 'https://storage.example.com/expired', + size: 8, + type: 'image/png', + context: 'execution', + base64: 'iVBORw0KGgo=', + } + + it.each(['files', 'messages', 'userPrompt'] as const)( + 'replays a previous turn from %s with a fresh provider attachment', + async (source) => { + mockGetProviderFromModel.mockReturnValue('openai') + const hydrate = vi + .spyOn(userFileBase64, 'hydrateUserFilesWithBase64') + .mockImplementation(async (value) => { + const files = value as (typeof file)[] + return files.map((attachment) => ({ + ...attachment, + base64: file.base64, + })) as typeof value + }) + const inputs: AgentInputs = { + model: 'gpt-4o', + memoryType: 'conversation', + conversationId: 'conversation-1', + ...(source === 'userPrompt' + ? { userPrompt: 'Analyze this file', files: [file] } + : { + messages: [ + { + role: 'user', + content: 'Analyze this file', + ...(source === 'messages' ? { files: [file] } : {}), + }, + ], + ...(source === 'files' ? { files: [file] } : {}), + }), + } + const original = structuredClone(inputs) + await handler.execute({ ...mockContext, executionId: 'exec-1' }, mockBlock, inputs) + const stored = dbChainMockFns.values.mock.calls + .map(([row]) => row) + .find((row) => Array.isArray(row.data) && row.data[0]?.role === 'user')?.data as Message[] + expect(stored).toBeDefined() + expect(stored[0].files).toEqual([ + { + id: file.id, + name: file.name, + key: file.key, + url: '', + size: file.size, + type: file.type, + context: file.context, + }, + ]) + expect(inputs).toEqual(original) + + queueTableRows(schemaMock.memory, [ + { data: [...stored, { role: 'assistant', content: 'First answer' }] }, + ]) + mockGetProviderFromModel.mockReturnValue('anthropic') + const nextContext = { ...mockContext, executionId: 'exec-2' } + await handler.execute(nextContext, mockBlock, { + model: 'claude-sonnet-4-5', + memoryType: 'conversation', + conversationId: 'conversation-1', + messages: [{ role: 'user', content: 'What is in that file?' }], + }) + const request = mockExecuteProviderRequest.mock.calls.at(-1)?.[1] + expect(request.messages[0]).toMatchObject({ + role: 'user', + content: 'Analyze this file', + files: [{ key: file.key, base64: file.base64 }], + }) + expect(request.messages.at(-1)).toMatchObject({ + role: 'user', + content: 'What is in that file?', + }) + expect(request.messages.at(-1).files).toBeUndefined() + expect(hydrate.mock.calls.at(-1)?.[0]).toEqual(stored[0].files) + expect(hydrate.mock.calls.at(-1)?.[1]).toMatchObject({ + executionId: 'exec-2', + fileKeys: [file.key], + }) + hydrate.mockRestore() + } + ) + + it('does not duplicate an attachment when the same execution revisits the agent', async () => { + mockGetProviderFromModel.mockReturnValue('openai') + queueTableRows(schemaMock.memory, [ + { + data: [ + { role: 'user', content: 'Analyze this file', executionId: 'exec-1', files: [file] }, + ], + }, + ]) + await handler.execute({ ...mockContext, executionId: 'exec-1' }, mockBlock, { + model: 'gpt-4o', + memoryType: 'conversation', + conversationId: 'conversation-1', + messages: [{ role: 'user', content: 'Analyze this file' }], + files: [file], + }) + expect(mockExecuteProviderRequest.mock.calls[0][1].messages[0].files).toHaveLength(1) + expect(dbChainMockFns.values.mock.calls.some(([row]) => row.data?.[0]?.role === 'user')).toBe( + false + ) + }) + + it('saves a new attachment appended to an existing conversation', async () => { + mockGetProviderFromModel.mockReturnValue('openai') + queueTableRows(schemaMock.memory, [{ data: [{ role: 'assistant', content: 'Hello' }] }]) + await handler.execute({ ...mockContext, executionId: 'exec-2' }, mockBlock, { + model: 'gpt-4o', + memoryType: 'conversation', + conversationId: 'conversation-1', + messages: [{ role: 'user', content: 'Analyze this file' }], + files: [file], + }) + const stored = dbChainMockFns.values.mock.calls + .map(([row]) => row) + .find((row) => row.data?.[0]?.role === 'user')?.data as Message[] + expect(stored[0].files?.[0]).toMatchObject({ key: file.key, url: '' }) + expect(stored[0].files?.[0].base64).toBeUndefined() + }) + + it('does not hydrate attachments excluded by the conversation window', async () => { + mockGetProviderFromModel.mockReturnValue('openai') + const hydrate = vi.spyOn(userFileBase64, 'hydrateUserFilesWithBase64') + queueTableRows(schemaMock.memory, [ + { + data: [ + { role: 'user', content: 'Old file', files: [file] }, + { role: 'assistant', content: 'Recent answer' }, + ], + }, + ]) + const context = { ...mockContext, executionId: 'exec-2' } + await handler.execute(context, mockBlock, { + model: 'gpt-4o', + memoryType: 'sliding_window', + slidingWindowSize: '1', + conversationId: 'conversation-1', + messages: [{ role: 'user', content: 'Hello' }], + }) + expect(hydrate).not.toHaveBeenCalled() + expect(context.fileKeys).toBeUndefined() + hydrate.mockRestore() + }) + }) + describe('execute', () => { it('should execute a basic agent block request', async () => { const inputs = { diff --git a/apps/sim/executor/handlers/agent/agent-handler.ts b/apps/sim/executor/handlers/agent/agent-handler.ts index a78105e8aba..0ccc9c52f80 100644 --- a/apps/sim/executor/handlers/agent/agent-handler.ts +++ b/apps/sim/executor/handlers/agent/agent-handler.ts @@ -53,6 +53,7 @@ import { } from '@/executor/handlers/agent/skills-resolver' import type { AgentInputs, + FileNameProjection, Message, StreamingConfig, ToolInput, @@ -374,14 +375,12 @@ export class AgentBlockHandler implements BlockHandler { } const streamingConfig = this.getStreamingConfig(ctx, block) - const messages = await this.buildMessages(ctx, filteredInputs, modelInputs, skillMetadata) - const messagesWithInputFiles = this.attachFilesToLastUserMessage( + const messagesWithInputFiles = await this.buildMessages( ctx, - messages, - filteredInputs.files, - fileProjection.projectedFiles, - fileProjection.projectedNameByFile, - fileProjection.directNameInputPaths + filteredInputs, + modelInputs, + skillMetadata, + fileProjection ) const messagesWithFiles = await this.hydrateMessageFilesForProvider( ctx, @@ -1212,10 +1211,13 @@ export class AgentBlockHandler implements BlockHandler { ctx: ExecutionContext, inputs: AgentInputs, modelInputs: AgentInputs, - skillMetadata: Array<{ name: string; description: string }> = [] + skillMetadata: Array<{ name: string; description: string }>, + fileProjection: ReturnType ): Promise { const messages: Message[] = [] const memoryEnabled = inputs.memoryType && inputs.memoryType !== 'none' + const pendingMemoryMessages: Array<{ raw: Message; model: Message }> = [] + let seedMessageCount = 0 // 1. Extract and validate messages from messages-input subblock const inputMessages = this.extractValidMessages(inputs.messages) @@ -1226,7 +1228,11 @@ export class AgentBlockHandler implements BlockHandler { // 2. Handle native memory: seed on first run, then fetch and append new user input if (memoryEnabled && ctx.workspaceId) { - const memoryMessages = await memoryService.fetchMemoryMessages(ctx, inputs) + const memoryMessages = await memoryService.fetchMemoryMessages( + ctx, + inputs, + fileProjection.projectedNameByFile + ) const hasExisting = memoryMessages.length > 0 if (!hasExisting && conversationMessages.length > 0) { @@ -1236,7 +1242,13 @@ export class AgentBlockHandler implements BlockHandler { const rawTaggedMessages = rawConversationMessages.map((m) => m.role === 'user' ? { ...m, executionId: ctx.executionId } : m ) - await memoryService.seedMemory(ctx, inputs, rawTaggedMessages) + for (let index = 0; index < taggedMessages.length; index++) { + pendingMemoryMessages.push({ + raw: rawTaggedMessages[index], + model: taggedMessages[index], + }) + } + seedMessageCount = taggedMessages.length messages.push(...taggedMessages) } else { messages.push(...memoryMessages) @@ -1261,9 +1273,9 @@ export class AgentBlockHandler implements BlockHandler { if (!userMessageInThisRun) { const taggedMessage = { ...latestUserFromInput, executionId: ctx.executionId } messages.push(taggedMessage) - await memoryService.appendToMemory(ctx, inputs, { - ...latestRawUserFromInput, - executionId: ctx.executionId, + pendingMemoryMessages.push({ + raw: { ...latestRawUserFromInput, executionId: ctx.executionId }, + model: taggedMessage, }) } } @@ -1300,9 +1312,9 @@ export class AgentBlockHandler implements BlockHandler { const userMessages = messages.filter((m) => m.role === 'user') const lastUserMessage = userMessages[userMessages.length - 1] if (lastUserMessage) { - await memoryService.appendToMemory(ctx, inputs, { - ...lastUserMessage, - content: this.formatUserPrompt(inputs.userPrompt), + pendingMemoryMessages.push({ + raw: { ...lastUserMessage, content: this.formatUserPrompt(inputs.userPrompt) }, + model: lastUserMessage, }) } } @@ -1328,7 +1340,33 @@ export class AgentBlockHandler implements BlockHandler { } } - return messages.length > 0 ? messages : undefined + const messagesWithFiles = this.attachFilesToLastUserMessage( + ctx, + messages.length > 0 ? messages : undefined, + inputs.files, + fileProjection.projectedFiles, + fileProjection.projectedNameByFile, + fileProjection.directNameInputPaths + ) + + /** Persist the complete turn before provider hydration adds bytes or transient handles. */ + const lastUserMessage = messages.filter((message) => message.role === 'user').at(-1) + const attachedUserMessage = messagesWithFiles + ?.filter((message) => message.role === 'user') + .at(-1) + const messagesToStore = pendingMemoryMessages.map(({ raw, model }) => + model === lastUserMessage && attachedUserMessage?.files + ? { ...raw, files: attachedUserMessage.files } + : raw + ) + if (seedMessageCount > 0) { + await memoryService.seedMemory(ctx, inputs, messagesToStore.slice(0, seedMessageCount)) + } + for (const message of messagesToStore.slice(seedMessageCount)) { + await memoryService.appendToMemory(ctx, inputs, message) + } + + return messagesWithFiles } private attachFilesToLastUserMessage( @@ -1336,7 +1374,7 @@ export class AgentBlockHandler implements BlockHandler { messages: Message[] | undefined, filesInput: unknown, projectedFilesInput: unknown, - projectedNameByFile: WeakMap, + projectedNameByFile: WeakMap, directNameInputPaths: readonly ResolvedSecretInputPath[] ): Message[] | undefined { const normalizedFiles = normalizeFileInput(filesInput) @@ -1397,10 +1435,13 @@ export class AgentBlockHandler implements BlockHandler { } const lastUserMessage = messages[lastUserMessageIndex] + const filesByKey = new Map( + [...(lastUserMessage.files ?? []), ...userFiles].map((file) => [file.key || file.id, file]) + ) const nextMessages = [...messages] nextMessages[lastUserMessageIndex] = { ...lastUserMessage, - files: [...(lastUserMessage.files ?? []), ...userFiles], + files: Array.from(filesByKey.values()), } return nextMessages @@ -1410,7 +1451,7 @@ export class AgentBlockHandler implements BlockHandler { ctx: ExecutionContext, messages: Message[] | undefined, providerId: string, - projectedNameByFile: WeakMap, + projectedNameByFile: WeakMap, modelBoundInputPaths: ResolvedSecretInputPath[] ): Promise { if (!messages?.some((message) => message.files?.length)) { @@ -1478,7 +1519,7 @@ export class AgentBlockHandler implements BlockHandler { return [file] } - modelBoundInputPaths.push(nameProjection.inputPath) + if (nameProjection.inputPath) modelBoundInputPaths.push(nameProjection.inputPath) const extension = getFileExtension(file.name) const suffix = extension ? `.${extension}` : '' const keepsSuffix = @@ -2004,7 +2045,7 @@ export class AgentBlockHandler implements BlockHandler { inputs: AgentInputs ): { projectedFiles: unknown - projectedNameByFile: WeakMap + projectedNameByFile: WeakMap directNameInputPaths: ResolvedSecretInputPath[] modelBoundInputPaths: ResolvedSecretInputPath[] } { @@ -2104,10 +2145,7 @@ export class AgentBlockHandler implements BlockHandler { } } - const projectedNameByFile = new WeakMap< - object, - { name: string; inputPath: ResolvedSecretInputPath } - >() + const projectedNameByFile = new WeakMap() const projectedMessages = Array.isArray(projection.value.messages) ? projection.value.messages : [] diff --git a/apps/sim/executor/handlers/agent/memory-harness.postgres.test.ts b/apps/sim/executor/handlers/agent/memory-harness.postgres.test.ts new file mode 100644 index 00000000000..7c9ebf1a3f5 --- /dev/null +++ b/apps/sim/executor/handlers/agent/memory-harness.postgres.test.ts @@ -0,0 +1,547 @@ +/** + * @vitest-environment node + * + * Opt-in harness: AGENT_MEMORY_TEST_DATABASE_URL must point to disposable local PostgreSQL. + * Run from apps/sim with `bun run test executor/handlers/agent/memory-harness.postgres.test.ts`. + * AGENT_MEMORY_TEST_LIVE=1 uses configured OpenAI/Anthropic credentials and synthetic PDFs. + * Otherwise only provider HTTP responses are simulated; storage, SQL, hydration, dispatch, + * SDK request construction, response parsing, and memory persistence execute real code. + * Account policy/key lookup use fixtures, Redis uses the in-memory fallback, and no API route runs. + * AGENT_MEMORY_TEST_REPORT optionally writes SQL snapshots and outgoing attachment hashes. + */ +import { createHash } from 'node:crypto' +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { generateId, generateShortId } from '@sim/utils/id' +import { sql } from 'drizzle-orm' +import { getTableConfig, PgDialect, type PgTable } from 'drizzle-orm/pg-core' +import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js' +import { PDFDocument, StandardFonts } from 'pdf-lib' +import postgres from 'postgres' +import { fetch as networkFetch } from 'undici' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' + +const fixture = vi.hoisted(() => ({ + database: undefined as PostgresJsDatabase | undefined, + uploads: '', +})) + +vi.unmock('drizzle-orm') +vi.unmock('@sim/db/schema') +vi.mock('@sim/db', () => ({ + db: new Proxy( + {}, + { + get(_target, property) { + if (!fixture.database) throw new Error('Harness database is not initialized') + const value = Reflect.get(fixture.database, property) + return typeof value === 'function' ? value.bind(fixture.database) : value + }, + } + ), +})) +vi.mock('@/lib/uploads/core/setup.server', () => ({ + get UPLOAD_DIR_SERVER() { + return fixture.uploads + }, +})) +/** Fixture principals and explicit keys replace account configuration, not file authorization. */ +vi.mock('@/ee/access-control/utils/permission-check', () => ({ + assertPermissionsAllowed: async () => {}, + validateModelProvider: async () => {}, + validateBlockType: async () => {}, +})) +vi.mock('@/lib/api-key/byok', () => ({ + getApiKeyWithBYOK: async ( + _provider: string, + _model: string, + _workspace: string, + apiKey: string + ) => ({ apiKey, isBYOK: true }), +})) + +import { + memory, + memorySecretProvenance, + workspaceFileSecretProvenance, + workspaceFiles, +} from '@sim/db/schema' +import { hashDurableSecretProvenanceValue } from '@/lib/execution/durable-secret-provenance' +import { uploadExecutionFile } from '@/lib/uploads/contexts/execution/execution-file-manager' +import { EXACT_EMPTY_WORKSPACE_FILE_SECRET_PROVENANCE } from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance' +import { deleteFile } from '@/lib/uploads/core/storage-service' +import { AgentBlockHandler } from '@/executor/handlers/agent/agent-handler' +import type { AgentInputs, Message } from '@/executor/handlers/agent/types' +import type { ExecutionContext, StreamingExecution, UserFile } from '@/executor/types' +import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +import { createAgentStreamPump } from '@/providers/stream-pump' +import type { SerializedBlock } from '@/serializer/types' + +const databaseUrl = process.env.AGENT_MEMORY_TEST_DATABASE_URL +if (databaseUrl && !['localhost', '127.0.0.1', '[::1]'].includes(new URL(databaseUrl).hostname)) { + throw new Error('The Agent memory harness requires a disposable local database') +} +const live = process.env.AGENT_MEMORY_TEST_LIVE === '1' +const schemaName = `agent_memory_${generateId().replaceAll('-', '')}` +const connection = databaseUrl + ? postgres(databaseUrl, { + max: 3, + connection: { search_path: schemaName }, + onnotice: () => {}, + }) + : undefined +const scope = { workspaceId: generateId(), workflowId: generateId(), userId: generateId() } +const block = { + id: generateId(), + metadata: { id: 'agent', name: 'Memory harness' }, + position: { x: 0, y: 0 }, + config: { tool: '', params: {} }, + inputs: {}, + outputs: {}, + enabled: true, +} as SerializedBlock +const models = { + openai: process.env.AGENT_MEMORY_TEST_OPENAI_MODEL || 'gpt-4.1-mini', + anthropic: process.env.AGENT_MEMORY_TEST_ANTHROPIC_MODEL || 'claude-haiku-4-5', +} as const +type Provider = keyof typeof models +interface WireRequest { + host: string + body: Record +} +interface StoredConversation { + data: Message[] + secret_provenance_version: number + content_hash: string + status: string + entries: unknown[] +} +const report: Array> = [] +let outbound: WireRequest[] = [] +let transportReply = 'READY' + +function apiKey(provider: Provider): string { + if (!live) return 'synthetic-harness-key' + const value = + provider === 'openai' + ? process.env.OPENAI_API_KEY + : process.env.ANTHROPIC_API_KEY || process.env.ANTHROPIC_API_KEY_1 + if (!value) throw new Error(`Live harness requires a configured ${provider} API key`) + return value +} + +function context(streaming = false): ExecutionContext { + return { + ...scope, + executionId: generateId(), + principal: { + kind: 'system', + serviceId: 'schedule', + workspaceId: scope.workspaceId, + workflowId: scope.workflowId, + }, + stream: streaming, + selectedOutputs: [block.id], + blockStates: new Map(), + blockLogs: [], + metadata: { startTime: new Date().toISOString(), duration: 0 }, + environmentVariables: {}, + decisions: { router: new Map(), condition: new Map() }, + loopExecutions: new Map(), + completedLoops: new Set(), + executedBlocks: new Set(), + activeExecutionPath: new Set(), + workflow: { blocks: [], connections: [], loops: {}, version: '1' }, + resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry([], scope), + } as ExecutionContext +} + +/** Use the executor's stream pump and completion hook to persist streamed assistant turns. */ +async function executeTurn(ctx: ExecutionContext, inputs: AgentInputs): Promise { + const result = await new AgentBlockHandler().execute(ctx, block, inputs) + if (!ctx.stream) return String(result.content) + expect(result).toHaveProperty('stream') + const streamingResult = result as StreamingExecution + expect(streamingResult.onFullContent).toBeTypeOf('function') + const pump = createAgentStreamPump({ + source: streamingResult.stream, + streamFormat: streamingResult.streamFormat ?? 'text', + sinkMode: true, + }) + const drained = await pump.run() + expect(drained.fullyDrained).toBe(true) + expect(drained.cancelled).toBe(false) + await streamingResult.onFullContent!(drained.answerText) + return drained.answerText +} + +/** Generate only the four production tables exercised here; unrelated application FKs are omitted. */ +async function createTable(table: PgTable): Promise { + if (!connection) throw new Error('Missing harness database') + const dialect = new PgDialect() + const config = getTableConfig(table) + const columns = config.columns.map((column) => { + const defaultValue = + column.dataType === 'json' && column.default !== undefined + ? sql`${JSON.stringify(column.default)}::jsonb` + : sql`${column.default}` + const defaultSql = + column.default === undefined + ? '' + : ` DEFAULT ${dialect.sqlToQuery(defaultValue.inlineParams()).sql}` + return `"${column.name}" ${column.getSQLType()}${column.primary ? ' PRIMARY KEY' : ''}${column.notNull ? ' NOT NULL' : ''}${defaultSql}` + }) + await connection.unsafe(`CREATE TABLE "${config.name}" (${columns.join(', ')})`) +} + +async function readConversation(key: string): Promise { + if (!connection) throw new Error('Missing harness database') + const [row] = await connection` + SELECT m.data, m.secret_provenance_version, p.content_hash, p.status, p.entries + FROM memory m LEFT JOIN memory_secret_provenance p ON p.memory_id = m.id + WHERE m.workspace_id = ${scope.workspaceId} AND m.key = ${key}` + if (!row) throw new Error('Conversation was not persisted') + expect(row.secret_provenance_version).toBe(1) + expect(row.status).toBe('exact') + expect(row.content_hash).toBe(hashDurableSecretProvenanceValue(row.data)) + expect(row.entries).toEqual([]) + return row +} + +function requestFiles(request: WireRequest): string[] { + const messages = request.host === 'api.openai.com' ? request.body.input : request.body.messages + if (!Array.isArray(messages)) throw new Error('Provider did not send a message array') + return messages.flatMap((message) => { + if (!Array.isArray(message.content)) return [] + return message.content.flatMap((part: Record) => { + if (part.type === 'input_file' && typeof part.file_data === 'string') + return [part.file_data.split(',')[1]] + if ( + part.type === 'document' && + part.source && + typeof part.source === 'object' && + 'data' in part.source + ) { + return [String(part.source.data)] + } + return [] + }) + }) +} + +async function interceptFetch(input: RequestInfo | URL, init?: RequestInit): Promise { + const request = new Request(input, init) + const url = new URL(request.url) + if (!['api.openai.com', 'api.anthropic.com'].includes(url.hostname)) { + throw new Error(`Unexpected harness network destination: ${url.hostname}`) + } + const text = await request.text() + const body = JSON.parse(text) as Record + outbound.push({ host: url.hostname, body }) + if (live) { + const response = await networkFetch(url, { + method: request.method, + headers: Object.fromEntries(request.headers), + body: text, + signal: AbortSignal.timeout(60_000), + }) + return new Response(await response.arrayBuffer(), { + status: response.status, + headers: Object.fromEntries(response.headers), + }) + } + const content = transportReply + const response = + url.hostname === 'api.openai.com' + ? { + id: 'resp_harness', + object: 'response', + status: 'completed', + model: body.model, + output: [ + { + id: 'msg_harness', + type: 'message', + role: 'assistant', + status: 'completed', + content: [{ type: 'output_text', text: content, annotations: [] }], + }, + ], + usage: { input_tokens: 100, output_tokens: 8, total_tokens: 108 }, + } + : { + id: 'msg_harness', + type: 'message', + role: 'assistant', + model: body.model, + content: [{ type: 'text', text: content }], + stop_reason: 'end_turn', + stop_sequence: null, + usage: { input_tokens: 100, output_tokens: 8 }, + } + if (!body.stream) return Response.json(response) + const events = + url.hostname === 'api.openai.com' + ? [ + { type: 'response.output_text.delta', delta: content }, + { type: 'response.completed', response }, + ] + : [ + { type: 'message_start', message: { ...response, content: [] } }, + { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } }, + { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: content } }, + { type: 'content_block_stop', index: 0 }, + { + type: 'message_delta', + delta: { stop_reason: 'end_turn', stop_sequence: null }, + usage: { output_tokens: 8 }, + }, + { type: 'message_stop' }, + ] + return new Response( + events.map((event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`).join(''), + { headers: { 'Content-Type': 'text/event-stream' } } + ) +} + +describe.skipIf(!databaseUrl)( + 'Agent memory through PostgreSQL, storage, and provider transports', + () => { + beforeAll(async () => { + if (!connection) throw new Error('Missing harness database') + fixture.uploads = await mkdtemp(join(tmpdir(), 'sim-memory-files-')) + await connection`CREATE SCHEMA ${connection(schemaName)}` + fixture.database = drizzle(connection) + for (const table of [ + memory, + memorySecretProvenance, + workspaceFiles, + workspaceFileSecretProvenance, + ]) + await createTable(table) + await connection.unsafe( + `CREATE UNIQUE INDEX memory_workspace_key_idx ON memory(workspace_id, key)` + ) + await connection.unsafe( + `CREATE UNIQUE INDEX workspace_files_key_active_unique ON workspace_files(key) WHERE deleted_at IS NULL` + ) + await connection.unsafe(` + CREATE FUNCTION demote_memory() RETURNS trigger LANGUAGE plpgsql AS $body$ + BEGIN NEW.secret_provenance_version := NULL; RETURN NEW; END; $body$; + CREATE TRIGGER memory_demote BEFORE UPDATE OF data ON memory FOR EACH ROW + WHEN(OLD.data IS DISTINCT FROM NEW.data) EXECUTE FUNCTION demote_memory(); + `) + }) + + afterAll(async () => { + try { + if (process.env.AGENT_MEMORY_TEST_REPORT) { + await writeFile( + process.env.AGENT_MEMORY_TEST_REPORT, + JSON.stringify({ live, cases: report }, null, 2) + ) + } + if (connection) await connection`DROP SCHEMA IF EXISTS ${connection(schemaName)} CASCADE` + } finally { + fixture.database = undefined + if (connection) await connection.end() + if (fixture.uploads) await rm(fixture.uploads, { recursive: true, force: true }) + vi.unstubAllGlobals() + } + }) + + it.each( + ( + [ + { first: 'openai', second: 'openai', source: 'files' }, + { first: 'anthropic', second: 'anthropic', source: 'messages' }, + { first: 'openai', second: 'anthropic', source: 'userPrompt' }, + { first: 'anthropic', second: 'openai', source: 'files' }, + ] as const + ).flatMap((scenario) => [false, true].map((streaming) => ({ ...scenario, streaming }))) + )( + '$first → $second, using $source, streaming=$streaming', + async ({ first, second, source, streaming }) => { + vi.stubGlobal('fetch', interceptFetch) + outbound = [] + const conversationId = generateId() + const firstContext = context(streaming) + const marker = `PROBE-${generateShortId()}` + const pdf = await PDFDocument.create() + const font = await pdf.embedFont(StandardFonts.Helvetica) + pdf.addPage().drawText(`memory_probe = ${marker}`, { x: 40, y: 700, size: 18, font }) + const buffer = Buffer.from(await pdf.save()) + const file = await uploadExecutionFile( + { + workspaceId: scope.workspaceId, + workflowId: scope.workflowId, + executionId: firstContext.executionId!, + }, + buffer, + 'memory-probe.pdf', + 'application/pdf', + scope.userId, + EXACT_EMPTY_WORKSPACE_FILE_SECRET_PROVENANCE + ) + expect(file.key).toContain(firstContext.executionId) + if (!connection) throw new Error('Missing harness database') + /** Production provenance compares JavaScript Dates, so compare versions at millisecond precision. */ + const [storedFile] = await connection` + SELECT f.id, f.key, f.workspace_id, f.context, f.original_name, f.content_type, + f.size_bytes::integer AS size_bytes, f.secret_provenance_version, + p.status, p.entries, + date_trunc('milliseconds', f.content_updated_at) = p.content_updated_at AS provenance_bound + FROM workspace_files f + LEFT JOIN workspace_file_secret_provenance p ON p.file_id = f.id + WHERE f.key = ${file.key!}` + expect(storedFile).toEqual({ + id: file.id, + key: file.key, + workspace_id: scope.workspaceId, + context: 'execution', + original_name: 'memory-probe.pdf', + content_type: 'application/pdf', + size_bytes: buffer.length, + secret_provenance_version: 1, + status: 'exact', + entries: [], + provenance_bound: true, + }) + const firstPrompt = + 'Read the attached PDF. Reply exactly READY. Do not quote or mention anything from the file.' + const firstInputs: AgentInputs = { + model: models[first], + apiKey: apiKey(first), + maxTokens: '128', + memoryType: 'conversation', + conversationId, + ...(source === 'userPrompt' + ? { userPrompt: firstPrompt, files: [file] } + : { + messages: [ + { + role: 'user', + content: firstPrompt, + ...(source === 'messages' ? { files: [file] } : {}), + }, + ], + ...(source === 'files' ? { files: [file] } : {}), + }), + } + transportReply = 'READY' + expect(await executeTurn(firstContext, firstInputs)).toBe('READY') + const firstStored = await readConversation(conversationId) + expect(firstStored.data.map((message) => message.role)).toEqual(['user', 'assistant']) + const reference: UserFile = { + id: file.id, + name: file.name, + key: file.key, + url: '', + size: buffer.length, + type: 'application/pdf', + context: 'execution', + } + expect(firstStored.data[0].files).toEqual([reference]) + expect(JSON.stringify(firstStored)).not.toContain(marker) + expect(JSON.stringify(firstStored)).not.toContain('base64') + expect(JSON.stringify(firstStored)).not.toContain('providerFile') + expect(requestFiles(outbound[0])).toEqual([buffer.toString('base64')]) + + /** Fresh handler, execution, and registry force a DB read and a new execution-scoped byte cache. */ + const secondContext = context(streaming) + transportReply = marker + const followupInputs: AgentInputs = { + model: models[second], + apiKey: apiKey(second), + maxTokens: '128', + memoryType: 'conversation', + conversationId, + messages: [ + { + role: 'user', + content: + 'What is the memory_probe value from the PDF I attached earlier? Reply exactly the value. If no PDF is available, reply NO_FILE.', + }, + ], + } + expect(await executeTurn(secondContext, followupInputs)).toBe(marker) + expect(secondContext.fileKeys).toEqual([file.key]) + expect(outbound).toHaveLength(2) + expect(outbound.every((request) => Boolean(request.body.stream) === streaming)).toBe(true) + expect(requestFiles(outbound[1])).toEqual([buffer.toString('base64')]) + const secondStored = await readConversation(conversationId) + expect(secondStored.data.map((message) => message.role)).toEqual([ + 'user', + 'assistant', + 'user', + 'assistant', + ]) + expect(secondStored.data[0]).toEqual(firstStored.data[0]) + expect(secondStored.data[2].files).toBeUndefined() + expect(secondStored.data[3].content).toBe(marker) + + transportReply = 'NO_FILE' + const unrelatedConversationId = generateId() + const unrelatedContext = context(streaming) + expect( + await executeTurn(unrelatedContext, { + ...followupInputs, + conversationId: unrelatedConversationId, + }) + ).toBe('NO_FILE') + expect(unrelatedContext.fileKeys ?? []).toEqual([]) + expect(outbound).toHaveLength(3) + expect(requestFiles(outbound[2])).toEqual([]) + const unrelatedStored = await readConversation(unrelatedConversationId) + expect(unrelatedStored.data).toHaveLength(2) + expect(unrelatedStored.data.every((message) => !message.files)).toBe(true) + + const otherWorkflow = context(streaming) + otherWorkflow.workflowId = generateId() + otherWorkflow.principal = { + kind: 'system', + serviceId: 'schedule', + workspaceId: scope.workspaceId, + workflowId: otherWorkflow.workflowId, + } + await expect(executeTurn(otherWorkflow, followupInputs)).rejects.toThrow( + 'could not be read' + ) + expect(outbound).toHaveLength(3) + + await deleteFile({ key: file.key!, context: 'execution' }) + await expect(executeTurn(context(streaming), followupInputs)).rejects.toThrow( + 'could not be read' + ) + expect(outbound).toHaveLength(3) + const afterDeletion = await readConversation(conversationId) + expect(afterDeletion.data[0].files).toEqual([reference]) + report.push({ + first, + second, + source, + streaming, + marker, + storedFile, + firstStored, + secondStored, + controls: { + unrelatedConversation: 'passed', + differentWorkflow: 'blocked before HTTP', + missingSource: 'blocked before HTTP', + }, + transports: outbound.map((request) => ({ + host: request.host, + fileCount: requestFiles(request).length, + fileSha256: requestFiles(request).map((base64) => + createHash('sha256').update(Buffer.from(base64, 'base64')).digest('hex') + ), + })), + passed: true, + }) + }, + 150_000 + ) + } +) diff --git a/apps/sim/executor/handlers/agent/memory.test.ts b/apps/sim/executor/handlers/agent/memory.test.ts index 465dec67f15..96f629e62ba 100644 --- a/apps/sim/executor/handlers/agent/memory.test.ts +++ b/apps/sim/executor/handlers/agent/memory.test.ts @@ -1,4 +1,4 @@ -import { loggerMock } from '@sim/testing' +import { loggerMock, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing' import { beforeEach, describe, expect, it, vi } from 'vitest' const { mockDecryptSecret, mockRedactObjectStrings, mockIsEnforced, mockReportUnrecorded } = @@ -24,10 +24,21 @@ vi.mock('@/lib/logs/execution/pii-redaction', () => ({ })) import { hashDurableSecretProvenanceValue } from '@/lib/execution/durable-secret-provenance' +import { assertUserFileContentAccess } from '@/lib/execution/payloads/materialization.server' import { MEMORY } from '@/executor/constants' import { Memory } from '@/executor/handlers/agent/memory' import type { Message } from '@/executor/handlers/agent/types' +import type { ExecutionContext, UserFile } from '@/executor/types' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +import { + buildAnthropicMessageContent, + buildBedrockMessageContent, + buildGeminiMessageParts, + buildOpenAICompatibleChatContent, + buildOpenAIMessageContent, + buildOpenRouterMessageContent, + prepareProviderAttachments, +} from '@/providers/attachments' const mockMemoryLogger = vi.mocked(loggerMock.createLogger).mock.results[ vi.mocked(loggerMock.createLogger).mock.calls.findIndex(([name]) => name === 'Memory') @@ -44,6 +55,7 @@ describe('Memory', () => { beforeEach(() => { vi.clearAllMocks() + resetDbChainMock() mockIsEnforced.mockReturnValue(false) mockDecryptSecret.mockImplementation(async (encryptedValue: string) => ({ decrypted: `decrypted:${encryptedValue}`, @@ -214,7 +226,7 @@ describe('Memory', () => { }) describe('sanitizeMessageForStorage', () => { - it('should strip file payloads but preserve tool-call fields before memory persistence', () => { + it('preserves storage references and tool calls without file payloads or provider handles', () => { const message: Message = { role: 'user', content: 'Analyze this file', @@ -228,6 +240,9 @@ describe('Memory', () => { size: 128, type: 'image/png', base64: 'iVBORw0KGgo=', + providerFileId: 'expired-provider-file', + providerFileUri: 'expired-provider-uri', + remoteUrl: 'https://storage.example.com/expired', }, ], tool_calls: [{ id: 'call-1' }], @@ -237,11 +252,177 @@ describe('Memory', () => { role: 'user', content: 'Analyze this file', executionId: 'exec-1', + files: [ + { + id: 'file-1', + key: 'workspace/ws-1/example.png', + name: 'example.png', + url: '', + size: 128, + type: 'image/png', + }, + ], tool_calls: [{ id: 'call-1' }], }) }) }) + describe('provider-independent file references', () => { + const storedFile: UserFile = { + id: 'file-1', + key: 'workspace/workspace-1/image.png', + name: 'image.png', + url: '', + type: 'image/png', + size: 8, + context: 'workspace', + } + const bytes = 'iVBORw0KGgo=' + const renderers: Array<{ + providers: string[] + render: (content: string, files: UserFile[], provider: string) => unknown + }> = [ + { providers: ['openai', 'azure-openai'], render: buildOpenAIMessageContent }, + { providers: ['anthropic', 'azure-anthropic'], render: buildAnthropicMessageContent }, + { providers: ['google', 'vertex'], render: buildGeminiMessageParts }, + { providers: ['bedrock'], render: buildBedrockMessageContent }, + { providers: ['openrouter'], render: buildOpenRouterMessageContent }, + { + providers: [ + 'mistral', + 'groq', + 'fireworks', + 'together', + 'baseten', + 'ollama', + 'ollama-cloud', + 'vllm', + 'litellm', + 'xai', + 'kimi', + ], + render: buildOpenAICompatibleChatContent, + }, + ] + const providers = renderers.flatMap(({ providers, render }) => + providers.map((provider) => ({ provider, render })) + ) + + it.each(providers)( + 'keeps the same attachment wire content for $provider', + async ({ provider, render }) => { + queueTableRows(schemaMock.memory, [ + { + data: [ + { + role: 'user', + content: 'Describe the image', + files: [ + { + ...storedFile, + url: 'https://expired.example/file', + base64: 'stale-bytes', + providerFileId: 'stale-id', + providerFileUri: 'stale-uri', + remoteUrl: 'https://expired.example/provider-file', + }, + ], + }, + ], + }, + ]) + const [message] = await memoryService.fetchMemoryMessages( + { workspaceId: 'workspace-1' } as ExecutionContext, + { memoryType: 'conversation', conversationId: 'conversation-1' } + ) + expect(message.files).toEqual([storedFile]) + const hydrated = message.files!.map((file) => ({ ...file, base64: bytes })) + expect(render(message.content, hydrated, provider)).toEqual( + render('Describe the image', [{ ...storedFile, base64: bytes }], provider) + ) + } + ) + + it.each(['deepseek', 'cerebras', 'sakana', 'nvidia', 'meta', 'zai'])( + 'keeps the explicit unsupported-attachment error for %s', + (provider) => { + expect(() => + prepareProviderAttachments([{ ...storedFile, base64: bytes }], provider) + ).toThrow('File attachments are not supported') + } + ) + + it('admits only the remembered execution file and preserves workspace and workflow scope', async () => { + const file = { + ...storedFile, + key: 'execution/workspace-1/workflow-1/exec-1/image.png', + context: 'execution', + } + queueTableRows(schemaMock.memory, [ + { data: [{ role: 'user', content: 'File', files: [file] }] }, + ]) + const context = { + workspaceId: 'workspace-1', + workflowId: 'workflow-1', + executionId: 'exec-2', + } as ExecutionContext + await memoryService.fetchMemoryMessages(context, { + memoryType: 'conversation', + conversationId: 'conversation-1', + }) + await expect(assertUserFileContentAccess(file, context)).resolves.toBeUndefined() + await expect( + assertUserFileContentAccess( + { ...file, key: 'execution/workspace-1/workflow-1/exec-1/other.png' }, + context + ) + ).rejects.toThrow('File is not available') + await expect( + assertUserFileContentAccess(file, { ...context, workspaceId: 'workspace-2' }) + ).rejects.toThrow('File is not available') + await expect( + assertUserFileContentAccess(file, { ...context, workflowId: 'workflow-2' }) + ).rejects.toThrow('File is not available') + }) + + it('bounds historical file loading even when the messages contain no text', async () => { + queueTableRows(schemaMock.memory, [ + { + data: Array.from({ length: MEMORY.MAX_REPLAY_FILE_REFERENCES + 1 }, () => ({ + role: 'user', + content: '', + files: [storedFile], + })), + }, + ]) + await expect( + memoryService.fetchMemoryMessages({ workspaceId: 'workspace-1' } as ExecutionContext, { + memoryType: 'conversation', + conversationId: 'conversation-1', + }) + ).rejects.toThrow('Use a smaller memory window') + }) + + it('does not carry inline-only or malformed file objects into a later turn', async () => { + queueTableRows(schemaMock.memory, [ + { + data: [ + { + role: 'user', + content: '', + files: [null, { name: 'invalid.png' }, { ...storedFile, key: '', base64: bytes }], + }, + ], + }, + ]) + const messages = await memoryService.fetchMemoryMessages( + { workspaceId: 'workspace-1' } as ExecutionContext, + { memoryType: 'conversation', conversationId: 'conversation-1' } + ) + expect(messages).toEqual([{ role: 'user', content: '' }]) + }) + }) + describe('secret projection', () => { function createContext(registry: ResolvedSecretTraceRegistry) { return { diff --git a/apps/sim/executor/handlers/agent/memory.ts b/apps/sim/executor/handlers/agent/memory.ts index cb12dddfca6..2da2bbec8d9 100644 --- a/apps/sim/executor/handlers/agent/memory.ts +++ b/apps/sim/executor/handlers/agent/memory.ts @@ -4,6 +4,7 @@ import { createLogger } from '@sim/logger' import { generateId } from '@sim/utils/id' import { isPlainRecord } from '@sim/utils/object' import { and, eq, sql } from 'drizzle-orm' +import { isUserFileWithMetadata } from '@/lib/core/utils/user-file' import { bindDurableSecretProvenanceToValue, durableSecretProvenanceFromRegistry, @@ -14,6 +15,7 @@ import { isDurableSecretProvenanceEnforced, reportUnrecordedDurableProvenance, } from '@/lib/execution/durable-secret-provenance-enforcement' +import { mergeFileKeys } from '@/lib/execution/payloads/access-keys' import { redactObjectStrings } from '@/lib/logs/execution/pii-redaction' import { lockMemoryConversationInTx } from '@/lib/memory/locks' import { @@ -23,7 +25,7 @@ import { } from '@/lib/memory/secret-provenance' import { getAccurateTokenCount } from '@/lib/tokenization/accurate' import { MEMORY } from '@/executor/constants' -import type { AgentInputs, Message } from '@/executor/handlers/agent/types' +import type { AgentInputs, FileNameProjection, Message } from '@/executor/handlers/agent/types' import type { ExecutionContext } from '@/executor/types' import { projectResolvedSecretModelContent, @@ -38,7 +40,11 @@ const logger = createLogger('Memory') const MEMORY_CONTENT_REFUSAL = 'Memory content could not be safely projected' export class Memory { - async fetchMemoryMessages(ctx: ExecutionContext, inputs: AgentInputs): Promise { + async fetchMemoryMessages( + ctx: ExecutionContext, + inputs: AgentInputs, + projectedNameByFile?: WeakMap + ): Promise { if (!inputs.memoryType || inputs.memoryType === 'none') { return [] } @@ -76,6 +82,14 @@ export class Memory { messages = stored.messages } + /** Bound historical downloads independently of text-token windows and per-file byte caps. */ + const fileCount = messages.reduce((count, message) => count + (message.files?.length ?? 0), 0) + if (fileCount > MEMORY.MAX_REPLAY_FILE_REFERENCES) { + throw new Error( + `Conversation memory exceeds ${MEMORY.MAX_REPLAY_FILE_REFERENCES} file attachments. Use a smaller memory window.` + ) + } + const selection = await createMemorySecretProvenanceSelector( stored.provenance, stored.messages, @@ -166,7 +180,7 @@ export class Memory { }) } - return Promise.all( + const projectedMessages = await Promise.all( messages.map(async (message) => { const messageProvenance = selectProvenance([message]) const modelRegistry = new ResolvedSecretTraceRegistry( @@ -188,9 +202,15 @@ export class Memory { inputPath: 'messages', }) } - return this.projectMessageForModel(modelRegistry, message) + return this.projectMessageForModel(modelRegistry, message, projectedNameByFile) }) ) + /** Saved references admit only these files; materialization still enforces their scope. */ + mergeFileKeys( + ctx, + projectedMessages.flatMap((message) => message.files?.map((file) => file.key) ?? []) + ) + return projectedMessages } private captureMessagesProvenance( @@ -305,7 +325,25 @@ export class Memory { } } - private projectMessageForModel(registry: ResolvedSecretTraceRegistry, message: Message): Message { + private projectMessageForModel( + registry: ResolvedSecretTraceRegistry, + message: Message, + projectedNameByFile?: WeakMap + ): Message { + for (const file of message.files ?? []) { + const projection = projectResolvedSecretModelContent(file.name, registry) + if (!projection.safe || typeof projection.value !== 'string') { + refuseResolvedSecretProjection({ + site: 'memory.fileNameProjection', + message: MEMORY_CONTENT_REFUSAL, + registry, + inputPath: 'files.name', + }) + } + if (projection.value !== file.name) { + projectedNameByFile?.set(file, { name: projection.value }) + } + } const functionArguments = this.readFunctionCallArguments( message.function_call, registry, @@ -451,9 +489,24 @@ export class Memory { return messages.slice(-limit) } + /** Storage keys survive turns; inline bytes, signed URLs, and provider handles do not. */ private sanitizeMessageForStorage(message: Message): Message { const { files: _files, ...messageWithoutFiles } = message - return messageWithoutFiles + const files = Array.isArray(message.files) + ? message.files + .filter(isUserFileWithMetadata) + .filter((file) => file.key) + .map((file) => ({ + id: file.id, + name: file.name, + key: file.key, + url: '', + size: file.size, + type: file.type, + ...(typeof file.context === 'string' ? { context: file.context } : {}), + })) + : [] + return files.length > 0 ? { ...messageWithoutFiles, files } : messageWithoutFiles } private applyTokenWindow(messages: Message[], maxTokens: number, model?: string): Message[] { diff --git a/apps/sim/executor/handlers/agent/types.ts b/apps/sim/executor/handlers/agent/types.ts index 6ba968f9904..514ecad4c58 100644 --- a/apps/sim/executor/handlers/agent/types.ts +++ b/apps/sim/executor/handlers/agent/types.ts @@ -1,5 +1,11 @@ import type { McpOperationPolicy } from '@/lib/mcp/operation-policy' import type { UserFile } from '@/executor/types' +import type { ResolvedSecretInputPath } from '@/executor/utils/resolved-secret-trace-registry' + +export interface FileNameProjection { + name: string + inputPath?: ResolvedSecretInputPath +} export interface SkillInput { skillId: string diff --git a/apps/sim/lib/memory/message-provenance.test.ts b/apps/sim/lib/memory/message-provenance.test.ts index 0104bb00e1a..af5d95e202a 100644 --- a/apps/sim/lib/memory/message-provenance.test.ts +++ b/apps/sim/lib/memory/message-provenance.test.ts @@ -48,7 +48,7 @@ import { createMemorySecretProvenanceSelector, } from '@/lib/memory/secret-provenance' import { Memory } from '@/executor/handlers/agent/memory' -import type { AgentInputs, Message } from '@/executor/handlers/agent/types' +import type { AgentInputs, FileNameProjection, Message } from '@/executor/handlers/agent/types' import type { ExecutionContext } from '@/executor/types' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' import { memoryAddTool } from '@/tools/memory/add' @@ -196,38 +196,74 @@ describe.each([false, true])('memory message provenance with enforcement %s', (e expect(mocks.logger.error).not.toHaveBeenCalled() }) - it.each(['append', 'seed'] as const)('binds %s messages after removing files', async (mode) => { - const service = new Memory() - const writes = service as unknown as MemoryWrites - const append = vi.spyOn(writes, 'appendMessage').mockResolvedValue(undefined) - const seed = vi.spyOn(writes, 'seedMemoryRecord').mockResolvedValue(undefined) - const registry = new ResolvedSecretTraceRegistry( - [{ name: 'TOKEN', plaintext: SECRET, encryptedValue: 'ciphertext' }], - SCOPE - ) - registry.recordResolved('TOKEN', SECRET) - const message = { - role: 'user', - content: SECRET, - files: [{ id: 'file-1', name: 'document.txt' }], - } as Message - if (mode === 'append') await service.appendToMemory(executionContext(registry), INPUTS, message) - else await service.seedMemory(executionContext(registry), INPUTS, [message]) + it.each(['append', 'seed'] as const)( + 'binds %s messages after removing transient attachment fields', + async (mode) => { + const service = new Memory() + const writes = service as unknown as MemoryWrites + const append = vi.spyOn(writes, 'appendMessage').mockResolvedValue(undefined) + const seed = vi.spyOn(writes, 'seedMemoryRecord').mockResolvedValue(undefined) + const registry = new ResolvedSecretTraceRegistry( + [{ name: 'TOKEN', plaintext: SECRET, encryptedValue: 'ciphertext' }], + SCOPE + ) + registry.recordResolved('TOKEN', SECRET) + const message = { + role: 'user', + content: SECRET, + files: [ + { + id: 'file-1', + name: `${SECRET}.txt`, + key: 'workspace/workspace-1/document.txt', + url: 'https://storage.example.com/signed', + size: 8, + type: 'text/plain', + base64: 'cGF5bG9hZA==', + providerFileId: 'expired-file', + }, + ], + } as Message + if (mode === 'append') + await service.appendToMemory(executionContext(registry), INPUTS, message) + else await service.seedMemory(executionContext(registry), INPUTS, [message]) - const stored = mode === 'append' ? [append.mock.calls[0][2]] : seed.mock.calls[0][2] - const provenance = mode === 'append' ? append.mock.calls[0][3] : seed.mock.calls[0][3] - expect(stored).toEqual([{ role: 'user', content: SECRET }]) - expect(provenance).toMatchObject({ - status: 'exact', - entries: [{ sourceValueHash: hashDurableSecretProvenanceValue(stored[0]) }], - }) - if (provenance?.status !== 'exact') throw new Error('Expected exact provenance') - queueStoredMemory(stored, provenance.entries) - expect((await service.fetchMemoryMessages(executionContext(), INPUTS))[0].content).toBe( - '{{TOKEN}}' - ) - expect(mocks.logger.error).not.toHaveBeenCalled() - }) + const stored = mode === 'append' ? [append.mock.calls[0][2]] : seed.mock.calls[0][2] + const provenance = mode === 'append' ? append.mock.calls[0][3] : seed.mock.calls[0][3] + expect(stored).toEqual([ + { + role: 'user', + content: SECRET, + files: [ + { + id: 'file-1', + name: `${SECRET}.txt`, + key: 'workspace/workspace-1/document.txt', + url: '', + size: 8, + type: 'text/plain', + }, + ], + }, + ]) + expect(provenance).toMatchObject({ + status: 'exact', + entries: [{ sourceValueHash: hashDurableSecretProvenanceValue(stored[0]) }], + }) + if (provenance?.status !== 'exact') throw new Error('Expected exact provenance') + queueStoredMemory(stored, provenance.entries) + const projectedNames = new WeakMap() + const [replayed] = await service.fetchMemoryMessages( + executionContext(), + INPUTS, + projectedNames + ) + expect(replayed.content).toBe('{{TOKEN}}') + expect(replayed.files?.[0].name).toBe(`${SECRET}.txt`) + expect(projectedNames.get(replayed.files![0])).toEqual({ name: '{{TOKEN}}.txt' }) + expect(mocks.logger.error).not.toHaveBeenCalled() + } + ) it.each(['unbound', 'before-file-sanitization'] as const)( 'redacts historical %s entries without refusing the run or exposing telemetry values',