Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion apps/studio/lib/ai/prompts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
`

Expand Down
7 changes: 7 additions & 0 deletions apps/studio/lib/ai/tools/mock-tools.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down Expand Up @@ -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: [] }
Expand Down
9 changes: 8 additions & 1 deletion apps/studio/lib/ai/tools/mock-tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
},
Expand Down Expand Up @@ -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)
Expand Down
26 changes: 25 additions & 1 deletion apps/studio/lib/ai/tools/notebook-tools.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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' },
{
Expand Down Expand Up @@ -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' },
{
Expand Down Expand Up @@ -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/)
})
})
})
16 changes: 14 additions & 2 deletions apps/studio/lib/ai/tools/notebook-tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) => {
Expand Down Expand Up @@ -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
Expand Down
Loading