From 8920439569c39eb03ff8eccdc11e6aca97f7fee4 Mon Sep 17 00:00:00 2001 From: Charis <26616127+charislam@users.noreply.github.com> Date: Fri, 21 Aug 2026 14:58:32 -0400 Subject: [PATCH] Expose previous notebook content in update_notebook (#49401) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Plumb pre-update notebook snapshot through `update_notebook` tool response as `previous_content` - Add sanitizers in `tool-sanitizer.ts` to strip snapshot before model sees it - Add client-side stripping in `prepareMessagesForAPI` to avoid re-uploading snapshot on subsequent turns - This is PR 2 of 3 fixing Linear issue FE-4243 (notebook update proposal shows 'unapplyable' error for already-completed updates) - Ships no visible behavior change on its own; enables PR 3 to restore diff preview for completed updates ## Test plan - [x] Unit tests: 80/80 passing across notebook-tools.test.ts, tool-sanitizer.test.ts, generate-assistant-response.utils.test.ts, message-utils.test.ts, and mock-tools.test.ts - [x] Typecheck: clean for all changed files - [x] ESLint: zero errors, lint:ratchet passes (exit 0) - [x] Integration: previous_content is correctly populated with pre-update notebook, stripped before model context, and stripped on client-side re-upload ## Summary by CodeRabbit - **Bug Fixes** - Notebook updates now retain previous content for recovery and history. - AI responses expose only the notebook’s ID and name, keeping previous content out of model-visible data. - **Tests** - Added coverage for notebook update results, content sanitization, and message preparation, including cases where previous content is absent or preserved. --- .../generate-assistant-response.utils.test.ts | 24 ++++++++++++ apps/studio/lib/ai/message-utils.test.ts | 37 ++++++++++++++++++- apps/studio/lib/ai/message-utils.ts | 14 +++++++ apps/studio/lib/ai/test-fixtures.ts | 23 ++++++++++++ .../lib/ai/tools/notebook-tools.test.ts | 32 +++++++++++++++- apps/studio/lib/ai/tools/notebook-tools.ts | 9 ++++- .../lib/ai/tools/tool-sanitizer.test.ts | 20 ++++++++++ apps/studio/lib/ai/tools/tool-sanitizer.ts | 16 ++++++++ 8 files changed, 172 insertions(+), 3 deletions(-) diff --git a/apps/studio/lib/ai/generate-assistant-response.utils.test.ts b/apps/studio/lib/ai/generate-assistant-response.utils.test.ts index 7e5d23429df..7bce339e1a4 100644 --- a/apps/studio/lib/ai/generate-assistant-response.utils.test.ts +++ b/apps/studio/lib/ai/generate-assistant-response.utils.test.ts @@ -115,6 +115,30 @@ describe('prepareMessagesForModel', () => { expect(result[0].parts).toEqual([]) }) + it('strips update_notebook previous_content before it reaches the model on history replay', () => { + const messages = [ + assistantMessage([ + toolPart({ + state: 'output-available', + output: { + id: 'notebook-1', + name: 'Signup funnel', + previous_content: { schema_version: 1, cells: [] }, + }, + }), + ]), + ] + + const result = prepareMessagesForModel(messages, 'schema') + + expect(result[0].parts).toEqual([ + toolPart({ + state: 'output-available', + output: { id: 'notebook-1', name: 'Signup funnel' }, + }), + ]) + }) + it('keeps non-tool parts untouched', () => { const messages = [assistantMessage([{ type: 'text', text: 'hello' }])] diff --git a/apps/studio/lib/ai/message-utils.test.ts b/apps/studio/lib/ai/message-utils.test.ts index d5acac0b328..a0be44767f2 100644 --- a/apps/studio/lib/ai/message-utils.test.ts +++ b/apps/studio/lib/ai/message-utils.test.ts @@ -1,4 +1,4 @@ -import type { DynamicToolUIPart, UIMessage } from 'ai' +import type { DynamicToolUIPart, ToolUIPart, UIMessage } from 'ai' import { describe, expect, it } from 'vitest' import { @@ -6,6 +6,7 @@ import { isManualApprovalRequested, prepareMessagesForAPI, } from './message-utils' +import { createAssistantMessageWithUpdateNotebookTool } from './test-fixtures' const makeApprovalPart = (id: string, isAutomatic = false): DynamicToolUIPart => ({ @@ -262,4 +263,38 @@ describe('prepareMessagesForAPI', () => { expect(result[2]).toEqual(messages[2]) expect(result[3]).not.toHaveProperty('results') }) + + it('strips update_notebook previous_content before re-uploading to the API', () => { + const messages = [createAssistantMessageWithUpdateNotebookTool()] + + const result = prepareMessagesForAPI(messages) + + expect((result[0].parts[0] as ToolUIPart).output).toEqual({ + id: 'notebook-1', + name: 'Signup funnel', + }) + }) + + it('does not mutate the original message parts when stripping previous_content', () => { + const messages = [createAssistantMessageWithUpdateNotebookTool()] + const originalParts = messages[0].parts + + prepareMessagesForAPI(messages) + + expect(messages[0].parts).toBe(originalParts) + expect((originalParts[0] as ToolUIPart).output).toHaveProperty('previous_content') + }) + + it('leaves an update_notebook output without previous_content unchanged', () => { + const messages = [ + createAssistantMessageWithUpdateNotebookTool({ id: 'notebook-1', name: 'Signup funnel' }), + ] + + const result = prepareMessagesForAPI(messages) + + expect((result[0].parts[0] as ToolUIPart).output).toEqual({ + id: 'notebook-1', + name: 'Signup funnel', + }) + }) }) diff --git a/apps/studio/lib/ai/message-utils.ts b/apps/studio/lib/ai/message-utils.ts index bb00d524a2c..4d86fad7a87 100644 --- a/apps/studio/lib/ai/message-utils.ts +++ b/apps/studio/lib/ai/message-utils.ts @@ -8,6 +8,15 @@ import { type UIPart = UIMessagePart +/** Strips `update_notebook`'s `previous_content` snapshot — display-only, never re-uploaded. */ +function stripNotebookSnapshot(part: UIPart): UIPart { + if (!isToolUIPart(part) || part.type !== 'tool-update_notebook') return part + if (!part.output || typeof part.output !== 'object') return part + + const { previous_content, ...sanitizedOutput } = part.output as Record + return { ...part, output: sanitizedOutput } as UIPart +} + /** * Prepares messages for API transmission by cleaning and limiting history */ @@ -26,6 +35,11 @@ export function prepareMessagesForAPI(messages: UIMessage[]): UIMessage[] { if (message.role === 'assistant' && message.results) { delete cleanedMessage.results } + // Map into a new array rather than mutating in place — `parts` is shared by reference + // with the locally persisted message the assistant panel reads from. + if (cleanedMessage.parts) { + cleanedMessage.parts = cleanedMessage.parts.map(stripNotebookSnapshot) + } return cleanedMessage as UIMessage }) diff --git a/apps/studio/lib/ai/test-fixtures.ts b/apps/studio/lib/ai/test-fixtures.ts index 1270e795763..83bcfbc3043 100644 --- a/apps/studio/lib/ai/test-fixtures.ts +++ b/apps/studio/lib/ai/test-fixtures.ts @@ -50,6 +50,29 @@ export function createAssistantMessageWithExecuteSqlTool( } } +export function createAssistantMessageWithUpdateNotebookTool( + output: Record = { + id: 'notebook-1', + name: 'Signup funnel', + previous_content: { schema_version: 1, cells: [] }, + }, + id = 'assistant-notebook-msg-1' +): UIMessage { + return { + id, + role: 'assistant', + parts: [ + { + type: 'tool-update_notebook', + state: 'output-available', + toolCallId: 'call-notebook-1', + input: { id: 'notebook-1', expected_updated_at: '2026-01-01T00:00:00.000Z' }, + output, + } satisfies ToolUIPart, + ], + } +} + export function createAssistantMessageWithMultipleTools( id = 'assistant-multi-tool-msg-1' ): UIMessage { diff --git a/apps/studio/lib/ai/tools/notebook-tools.test.ts b/apps/studio/lib/ai/tools/notebook-tools.test.ts index 7df9b8b72a6..2ab12beb1a4 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.test.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.test.ts @@ -632,7 +632,37 @@ describe('ai/tools/notebook-tools', () => { 'cell-2', ]) expect(content.cells[2].sql).toBe('select * from auth.users limit 100') - expect(result).toEqual({ id: 'notebook-1', name: 'Signup funnel' }) + // previous_content is the pre-update notebook (unaffected by this update's + // operations), not the post-update `sentBody` asserted above — it lets the client + // re-derive the diff at the time of approval + expect(result).toEqual({ + id: 'notebook-1', + name: 'Signup funnel', + previous_content: { + schema_version: 1, + cells: [ + { _tag: 'markdown_cell', _id: 'cell-1', text: '# Signup funnel' }, + { + _tag: 'database_cell', + _id: 'cell-2', + sql: 'select * from auth.users limit 100', + row_limit: 100, + view: 'table', + }, + { + _tag: 'log_cell', + _id: 'cell-3', + sql: 'select timestamp, event_message from edge_logs limit 10', + time_range: { _tag: 'relative_time_range', unit: 'hour', amount: 1 }, + view: 'table', + }, + ], + }, + }) + expect((tools.update_notebook as any).toModelOutput({ output: result })).toEqual({ + type: 'json', + value: { id: 'notebook-1', name: 'Signup funnel' }, + }) }) it('should throw a descriptive, assistant-exposable error instead of PUTting when an operation targets an unknown cell id', async () => { diff --git a/apps/studio/lib/ai/tools/notebook-tools.ts b/apps/studio/lib/ai/tools/notebook-tools.ts index 0b7f18ba2df..558bc562a56 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.ts @@ -314,8 +314,15 @@ export const getNotebookTools = (ctx: NotebookToolsContext = {}) => { authHeaders ) - return { id, name: notebook.name } + // `previous_content` lets the client re-derive the diff it already showed for + // approval, without re-fetching a notebook that's now post-update. Never reaches + // the model — see toModelOutput below. + return { id, name: notebook.name, previous_content: wireNotebook } }, + toModelOutput: ({ output }) => ({ + type: 'json', + value: { id: output.id, name: output.name }, + }), }), } } diff --git a/apps/studio/lib/ai/tools/tool-sanitizer.test.ts b/apps/studio/lib/ai/tools/tool-sanitizer.test.ts index 43ecfcc78f5..14705228f16 100644 --- a/apps/studio/lib/ai/tools/tool-sanitizer.test.ts +++ b/apps/studio/lib/ai/tools/tool-sanitizer.test.ts @@ -7,6 +7,7 @@ import { prepareMessagesForAPI } from '../message-utils' import { createAssistantMessageWithExecuteSqlTool, createAssistantMessageWithMultipleTools, + createAssistantMessageWithUpdateNotebookTool, createLongConversation, } from '../test-fixtures' import { NO_DATA_PERMISSIONS, sanitizeMessagePart } from './tool-sanitizer' @@ -173,4 +174,23 @@ describe('messages are sanitized based on opt-in level', () => { } }) }) + + test('update_notebook previous_content is stripped from the tool output regardless of opt-in level', () => { + const message = createAssistantMessageWithUpdateNotebookTool() + + const sanitized = sanitizeMessagePart(message.parts[0], 'schema') as ToolUIPart + + expect(sanitized.output).toEqual({ id: 'notebook-1', name: 'Signup funnel' }) + }) + + test('an update_notebook output that never had previous_content passes through unchanged', () => { + const message = createAssistantMessageWithUpdateNotebookTool({ + id: 'notebook-1', + name: 'Signup funnel', + }) + + const sanitized = sanitizeMessagePart(message.parts[0], 'schema') as ToolUIPart + + expect(sanitized.output).toEqual({ id: 'notebook-1', name: 'Signup funnel' }) + }) }) diff --git a/apps/studio/lib/ai/tools/tool-sanitizer.ts b/apps/studio/lib/ai/tools/tool-sanitizer.ts index 1f277dc49a3..766bdd0dbe1 100644 --- a/apps/studio/lib/ai/tools/tool-sanitizer.ts +++ b/apps/studio/lib/ai/tools/tool-sanitizer.ts @@ -32,8 +32,24 @@ const executeSqlSanitizer: ToolSanitizer = { }, } +// `previous_content` is UI-only and must never reach the model — `toModelOutput` on the +// tool covers the same turn, but history replay skips it, so it's stripped here too. +const updateNotebookSanitizer: ToolSanitizer = { + toolName: 'update_notebook', + sanitize: (tool) => { + if (!tool.output || typeof tool.output !== 'object') return tool + + const { previous_content, ...sanitizedOutput } = tool.output as Record + return { + ...tool, + output: sanitizedOutput, + } + }, +} + export const ALL_TOOL_SANITIZERS = { [executeSqlSanitizer.toolName]: executeSqlSanitizer, + [updateNotebookSanitizer.toolName]: updateNotebookSanitizer, } export function sanitizeMessagePart(