diff --git a/apps/studio/lib/ai/prompts.ts b/apps/studio/lib/ai/prompts.ts index b0c637c4d1f80..f1ababc6acb93 100644 --- a/apps/studio/lib/ai/prompts.ts +++ b/apps/studio/lib/ai/prompts.ts @@ -766,7 +766,7 @@ export const NOTEBOOKS_PROMPT = ` - Use \`update_notebook\` to edit an existing notebook — insert, replace, delete, or move cells — instead of recreating it from scratch. - Use \`execute_sql\` for a single ad-hoc question with no need to persist it. - When the request clearly calls for a notebook, call \`create_notebook\` or \`update_notebook\` directly; both tools handle user approval. -- \`update_notebook\` re-fetches the notebook right before applying edits, so the latest save always wins — it cannot detect edits made by someone else in between. +- \`update_notebook\` requires \`expected_updated_at\`, the \`updated_at\` you got from \`get_notebook\`. If the notebook changed since, the call is rejected — call \`get_notebook\` again and reissue \`update_notebook\` against the current content. - When describing an existing notebook, report each query cell's configuration that changes what it returns — a log cell's time range, a database cell's row limit — and don't count markdown cells as queries. ` diff --git a/apps/studio/lib/ai/tools/mock-tools.test.ts b/apps/studio/lib/ai/tools/mock-tools.test.ts index a9b808b703f21..bbd437a5add38 100644 --- a/apps/studio/lib/ai/tools/mock-tools.test.ts +++ b/apps/studio/lib/ai/tools/mock-tools.test.ts @@ -182,6 +182,7 @@ describe('ai/tools/mock-tools getMockTools', () => { const result = await mockTools.update_notebook.execute( { id: AUTH_HEALTH_NOTEBOOK_ID, + expected_updated_at: before.updated_at, operations: [ { _tag: 'insert_cell', @@ -218,10 +219,16 @@ describe('ai/tools/mock-tools getMockTools', () => { if (!mockTools.get_notebook.execute) throw new Error('execute is undefined') if (!mockTools.update_notebook.execute) throw new Error('execute is undefined') + const before = await mockTools.get_notebook.execute( + { id: EDGE_FUNCTION_NOTEBOOK_ID }, + { toolCallId: 'test', messages: [] } + ) + await expect( mockTools.update_notebook.execute( { id: EDGE_FUNCTION_NOTEBOOK_ID, + expected_updated_at: before.updated_at, operations: [{ _tag: 'delete_cell', cell_id: 'does-not-exist' }], }, { toolCallId: 'test', messages: [] } diff --git a/apps/studio/lib/ai/tools/mock-tools.ts b/apps/studio/lib/ai/tools/mock-tools.ts index a5fc459e8142c..0c06860d0f712 100644 --- a/apps/studio/lib/ai/tools/mock-tools.ts +++ b/apps/studio/lib/ai/tools/mock-tools.ts @@ -451,6 +451,7 @@ function createMockNotebookTools(store: MockNotebookStore) { name: notebook.name, description: notebook.description, visibility: notebook.visibility, + updated_at: notebook.updated_at, cells: notebook.content.cells, } }, @@ -484,8 +485,14 @@ function createMockNotebookTools(store: MockNotebookStore) { ...update_notebook, // Same reasoning as create_notebook's override above. needsApproval: false, + // expected_updated_at is validated by the real inputSchema (spread above) but not + // checked here: the in-memory store has no concurrent writers for the eval harness + // to race against. execute: async ( - { id, operations }: { id: string; operations: NotebookOperation[] }, + { + id, + operations, + }: { id: string; expected_updated_at: string; operations: NotebookOperation[] }, _options: ToolCallOptions ) => { const notebook = store.get(id) diff --git a/apps/studio/lib/ai/tools/notebook-tools.test.ts b/apps/studio/lib/ai/tools/notebook-tools.test.ts index c6889a74a0677..3185c8315a994 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.test.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.test.ts @@ -202,6 +202,7 @@ describe('ai/tools/notebook-tools', () => { name: 'Signup funnel', description: undefined, visibility: 'project', + updated_at: '2026-01-01T00:00:00.000Z', cells: [ { _tag: 'markdown_cell', id: 'cell-1', text: '# Signup funnel' }, { @@ -399,6 +400,7 @@ describe('ai/tools/notebook-tools', () => { const result = await tools.update_notebook.execute( { id: 'notebook-1', + expected_updated_at: '2026-01-01T00:00:00.000Z', operations: [ { _tag: 'delete_cell', cell_id: 'cell-3' }, { @@ -432,10 +434,32 @@ describe('ai/tools/notebook-tools', () => { await expect( tools.update_notebook.execute( - { id: 'notebook-1', operations: [{ _tag: 'delete_cell', cell_id: 'missing-cell' }] }, + { + id: 'notebook-1', + expected_updated_at: '2026-01-01T00:00:00.000Z', + operations: [{ _tag: 'delete_cell', cell_id: 'missing-cell' }], + }, { toolCallId: 'test', messages: [] } ) ).rejects.toThrow('No cell with id "missing-cell"') }) + + it('should throw instead of PUTting when the notebook changed since expected_updated_at', async () => { + mockGetNotebook() + + const tools = getNotebookTools({ projectRef: 'test-project' }) + if (!tools.update_notebook.execute) throw new Error('execute is undefined') + + await expect( + tools.update_notebook.execute( + { + id: 'notebook-1', + expected_updated_at: '2025-12-31T00:00:00.000Z', + operations: [{ _tag: 'delete_cell', cell_id: 'cell-3' }], + }, + { toolCallId: 'test', messages: [] } + ) + ).rejects.toThrow(/changed since expected_updated_at/) + }) }) }) diff --git a/apps/studio/lib/ai/tools/notebook-tools.ts b/apps/studio/lib/ai/tools/notebook-tools.ts index c5e95f3d077a4..a0d62f1adc800 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.ts @@ -85,6 +85,7 @@ export const getNotebookTools = (ctx: NotebookToolsContext = {}) => { name: notebook.name, description: notebook.description, visibility: notebook.visibility, + updated_at: notebook.updated_at, // Inlined rather than a shared helper: this discards the `unchecked_sql` brand for // display purposes only — the result is returned to the agent, never written back. cells: notebook.content.cells.map((cell) => { @@ -149,17 +150,28 @@ export const getNotebookTools = (ctx: NotebookToolsContext = {}) => { }), update_notebook: tool({ description: - 'Asks the user to apply an ordered list of cell operations (insert, replace, delete, move) to an existing notebook. Requires user approval before updating. Re-fetches the notebook right before applying the operations; concurrent edits are last-write-wins.', + 'Asks the user to apply an ordered list of cell operations (insert, replace, delete, move) to an existing notebook. Requires user approval before updating. Re-fetches the notebook right before applying the operations and rejects the update if it changed since expected_updated_at.', inputSchema: z.object({ id: z.string().describe('The id of the notebook to update.'), + expected_updated_at: z + .string() + .describe( + 'The `updated_at` you received from `get_notebook`. The update is rejected if the notebook changed since.' + ), operations: notebookOperationsSchema.describe( 'An ordered list of operations to apply to the notebook, addressing existing cells by id.' ), }), needsApproval: true, - execute: async ({ id, operations }) => { + execute: async ({ id, expected_updated_at, operations }) => { const notebook = await getNotebook({ projectRef, id }, undefined, authHeaders) + if (notebook.updated_at !== expected_updated_at) { + throw new Error( + `Notebook "${id}" changed since expected_updated_at (${expected_updated_at}); it is now ${notebook.updated_at}. Call get_notebook again and reissue update_notebook against the current content.` + ) + } + // Inlined rather than a shared helper, right beside this tool's own // `needsApproval: true`: this discards each cell's `unchecked_sql` brand so // applyNotebookOperations can splice cells as plain data. The result is never