From ddb3e2c44298559304bed6bc02b6d42b413e021d Mon Sep 17 00:00:00 2001 From: Charis <26616127+charislam@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:54:49 -0400 Subject: [PATCH] feat(studio): create_notebook AI tool (#48938) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Adds a `create_notebook` AI assistant tool (`needsApproval: true`) that lets the assistant create a new notebook after explicit user approval. - Cell SQL is promoted from untrusted to safe via `acceptUntrustedSql`/`acceptUntrustedLogsSql` inside `execute`, using the approval gate as the confirming user gesture (same pattern as `execute_sql`). - Input is validated against the existing agent-writable notebook schema, which rejects any agent-supplied cell `id` at the schema level. - Threads an optional auth-headers param through `upsertContent`/`createNotebook`/`updateNotebook` so the tool can pass its own bearer token server-side. - Registers the tool in the tool-filter (`SCHEMA` category, alongside `list_notebooks`/`get_notebook`) and adds a `## Notebooks` prompt section guiding the assistant on when to use `create_notebook` vs. one-off `execute_sql`. Resolves FE-4082 ## Test plan - [x] `notebook-tools.test.ts` covers: tool registration, `needsApproval`, cell-id rejection, valid input, PUT body shape, and the returned id — all passing - [x] Typecheck clean - [x] Lint clean (no new warnings) ## Summary by CodeRabbit * **New Features** * Added AI-assisted notebook creation for saving multi-step investigations. * Added support for database and log SQL cells in newly created notebooks. * Notebook creation requires approval before saving and returns the notebook’s name and identifier. * Added support for custom request headers during notebook and content operations. * Added guidance for choosing between one-time SQL execution and reusable notebooks when Explorer is enabled. * **Improvements** * Improved validation and normalization of notebook content before saving. --- .../data/content/content-upsert-mutation.ts | 8 +- .../notebooks/notebook-upsert-mutation.ts | 28 +++-- .../lib/ai/generate-assistant-response.ts | 15 ++- apps/studio/lib/ai/prompts.ts | 10 ++ apps/studio/lib/ai/tool-filter.ts | 2 + .../lib/ai/tools/notebook-tools.test.ts | 119 +++++++++++++++++- apps/studio/lib/ai/tools/notebook-tools.ts | 49 ++++++++ apps/studio/pages/api/ai/sql/generate-v4.ts | 1 + 8 files changed, 216 insertions(+), 16 deletions(-) diff --git a/apps/studio/data/content/content-upsert-mutation.ts b/apps/studio/data/content/content-upsert-mutation.ts index 8ce336b7684..ce91d62bed5 100644 --- a/apps/studio/data/content/content-upsert-mutation.ts +++ b/apps/studio/data/content/content-upsert-mutation.ts @@ -22,12 +22,16 @@ export type UpsertContentVariables = { export async function upsertContent( { projectRef, payload }: UpsertContentVariables, - signal?: AbortSignal + signal?: AbortSignal, + headersInit?: HeadersInit ): Promise { + const headers = new Headers(headersInit) + headers.set('Version', '2') + const { data, error } = await put('/platform/projects/{ref}/content', { params: { path: { ref: projectRef } }, body: unmapSqlContentField(payload), - headers: { Version: '2' }, + headers, signal, }) if (error) handleError(error) diff --git a/apps/studio/data/content/notebooks/notebook-upsert-mutation.ts b/apps/studio/data/content/notebooks/notebook-upsert-mutation.ts index f585dd4e4a5..39155ee8d49 100644 --- a/apps/studio/data/content/notebooks/notebook-upsert-mutation.ts +++ b/apps/studio/data/content/notebooks/notebook-upsert-mutation.ts @@ -47,16 +47,19 @@ export type CreateNotebookVariables = { export async function createNotebook( { projectRef, name, description, content }: CreateNotebookVariables, - signal?: AbortSignal + signal?: AbortSignal, + headersInit?: HeadersInit ) { - const payload = buildNotebookUpsertPayload({ - id: crypto.randomUUID(), - name, - description, - content, - }) + const id = crypto.randomUUID() + const payload = buildNotebookUpsertPayload({ id, name, description, content }) - return upsertContent({ projectRef, payload: payload as unknown as UpsertContentPayload }, signal) + await upsertContent( + { projectRef, payload: payload as unknown as UpsertContentPayload }, + signal, + headersInit + ) + + return { id } } export type CreateNotebookData = Awaited> @@ -65,11 +68,16 @@ export type UpdateNotebookVariables = CreateNotebookVariables & { id: string } export async function updateNotebook( { projectRef, id, name, description, content }: UpdateNotebookVariables, - signal?: AbortSignal + signal?: AbortSignal, + headersInit?: HeadersInit ) { const payload = buildNotebookUpsertPayload({ id, name, description, content }) - return upsertContent({ projectRef, payload: payload as unknown as UpsertContentPayload }, signal) + return upsertContent( + { projectRef, payload: payload as unknown as UpsertContentPayload }, + signal, + headersInit + ) } export type UpdateNotebookData = Awaited> diff --git a/apps/studio/lib/ai/generate-assistant-response.ts b/apps/studio/lib/ai/generate-assistant-response.ts index f131119b549..7d949003a22 100644 --- a/apps/studio/lib/ai/generate-assistant-response.ts +++ b/apps/studio/lib/ai/generate-assistant-response.ts @@ -16,7 +16,13 @@ import type { AssistantEvalInput } from '@/evals/scorer' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' import { buildAssistantContextMessages, NO_SCHEMA_ACCESS_MESSAGE } from '@/lib/ai/assistant-context' import { IS_TRACING_ENABLED } from '@/lib/ai/braintrust-logger' -import { CHAT_PROMPT, GENERAL_PROMPT, LIMITATIONS_PROMPT, SECURITY_PROMPT } from '@/lib/ai/prompts' +import { + CHAT_PROMPT, + GENERAL_PROMPT, + LIMITATIONS_PROMPT, + NOTEBOOKS_PROMPT, + SECURITY_PROMPT, +} from '@/lib/ai/prompts' import { sanitizeMessagePart } from '@/lib/ai/tools/tool-sanitizer' const { streamText: tracedStreamText } = wrapAISDK(ai) @@ -36,6 +42,7 @@ export async function generateAssistantResponse({ orgId, planId, includesLogsSnippets, + isExplorerEnabled, systemProviderOptions, providerOptions, requestedModel, @@ -57,6 +64,7 @@ export async function generateAssistantResponse({ planId?: string /** Whether any user message in the conversation attached a logs (ClickHouse) query. */ includesLogsSnippets?: boolean + isExplorerEnabled?: boolean requestedModel?: string systemProviderOptions?: Record providerOptions?: Record @@ -104,10 +112,13 @@ export async function generateAssistantResponse({ : await getSchemas() : NO_SCHEMA_ACCESS_MESSAGE - // Important: do not use dynamic content in the system prompt or Bedrock will not cache it + // Important: do not use per-request dynamic content in the system prompt or Bedrock will + // not cache it. isExplorerEnabled is a per-user flag, not per-request, so it only produces + // two prompt variants (on/off) rather than defeating caching. const system = source` ${GENERAL_PROMPT} ${CHAT_PROMPT} + ${isExplorerEnabled ? NOTEBOOKS_PROMPT : ''} ${SECURITY_PROMPT} ${LIMITATIONS_PROMPT} diff --git a/apps/studio/lib/ai/prompts.ts b/apps/studio/lib/ai/prompts.ts index 81c6f0435f6..3ffa11c6828 100644 --- a/apps/studio/lib/ai/prompts.ts +++ b/apps/studio/lib/ai/prompts.ts @@ -757,6 +757,16 @@ When asked about restoring/recovering deleted data: DO NOT start searching for recovery docs before checking deletion docs ` +// Notebooks haven't shipped yet — gated behind the Explorer feature flag, same as the +// notebook AI tools (see lib/ai/is-explorer-enabled.ts). Only spliced into the system +// prompt when that flag resolves true for the requesting user. +export const NOTEBOOKS_PROMPT = ` +## Notebooks +- Use \`create_notebook\` for a saved, shareable, multi-step investigation or dashboard the user will revisit — e.g. "build me a signup funnel notebook" or "create a notebook to track auth errors". +- 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\` directly; the tool handles user approval. +` + export const OUTPUT_ONLY_PROMPT = ` # Output-Only Mode diff --git a/apps/studio/lib/ai/tool-filter.ts b/apps/studio/lib/ai/tool-filter.ts index e74977ab37e..e1af605f4ac 100644 --- a/apps/studio/lib/ai/tool-filter.ts +++ b/apps/studio/lib/ai/tool-filter.ts @@ -39,6 +39,7 @@ export const toolSetValidationSchema = z.record( 'get_report', 'list_notebooks', 'get_notebook', + 'create_notebook', // Fallback tools for self-hosted 'getSchemaTables', @@ -92,6 +93,7 @@ export const TOOL_CATEGORY_MAP: Record = { get_report: TOOL_CATEGORIES.SCHEMA, list_notebooks: TOOL_CATEGORIES.SCHEMA, get_notebook: TOOL_CATEGORIES.SCHEMA, + create_notebook: TOOL_CATEGORIES.SCHEMA, getSchemaTables: TOOL_CATEGORIES.SCHEMA, getRlsKnowledge: TOOL_CATEGORIES.SCHEMA, getFunctions: TOOL_CATEGORIES.SCHEMA, diff --git a/apps/studio/lib/ai/tools/notebook-tools.test.ts b/apps/studio/lib/ai/tools/notebook-tools.test.ts index eadd9f10305..371b7543963 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.test.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.test.ts @@ -3,8 +3,26 @@ import { HttpResponse } from 'msw' import { describe, expect, it } from 'vitest' import { getNotebookTools } from './notebook-tools' +import type { AgentNotebook } from '@/data/content/notebooks/notebook-schema' import { addAPIMock, type APIErrorBody } from '@/tests/lib/msw' +const VALID_AGENT_CONTENT: AgentNotebook = { + schema_version: 1, + cells: [ + { _tag: 'markdown_cell', text: '# Signup funnel' }, + { + _tag: 'database_cell', + sql: 'select * from auth.users limit 100', + row_limit: 100, + }, + { + _tag: 'log_cell', + sql: "select timestamp, event_message from edge_logs where source = 'edge_logs' limit 10", + time_range: { _tag: 'relative_time_range', unit: 'hour', amount: 1 }, + }, + ], +} + type GetUserContentByIdResponse = components['schemas']['GetUserContentByIdResponse'] type GetUserContentResponse = components['schemas']['GetUserContentResponse'] @@ -29,10 +47,10 @@ const NOTEBOOK_CONTENT = { describe('ai/tools/notebook-tools', () => { describe('getNotebookTools', () => { - it('should return list_notebooks and get_notebook tools', () => { + it('should return list_notebooks, get_notebook, and create_notebook tools', () => { const tools = getNotebookTools() - expect(Object.keys(tools)).toEqual(['list_notebooks', 'get_notebook']) + expect(Object.keys(tools)).toEqual(['list_notebooks', 'get_notebook', 'create_notebook']) }) it('should not require approval to read notebooks', () => { @@ -41,6 +59,12 @@ describe('ai/tools/notebook-tools', () => { expect(tools.list_notebooks.needsApproval).toBeUndefined() expect(tools.get_notebook.needsApproval).toBeUndefined() }) + + it('should require approval to create a notebook', () => { + const tools = getNotebookTools() + + expect(tools.create_notebook.needsApproval).toBe(true) + }) }) describe('list_notebooks', () => { @@ -207,4 +231,95 @@ describe('ai/tools/notebook-tools', () => { ).rejects.toThrow() }) }) + + describe('create_notebook', () => { + it('should reject content whose cells carry an agent-supplied id', () => { + const tools = getNotebookTools() + const schema = tools.create_notebook.inputSchema + + if (!('safeParse' in schema)) throw new Error('inputSchema has no safeParse') + + const result = schema.safeParse({ + name: 'Signup funnel', + content: { + schema_version: 1, + cells: [{ _tag: 'markdown_cell', id: 'cell-1', text: '# Signup funnel' }], + }, + }) + + expect(result.success).toBe(false) + }) + + it('should accept a valid id-less notebook with all three cell types', () => { + const tools = getNotebookTools() + const schema = tools.create_notebook.inputSchema + + if (!('safeParse' in schema)) throw new Error('inputSchema has no safeParse') + + const result = schema.safeParse({ + name: 'Signup funnel', + description: 'Tracks signups over time', + content: VALID_AGENT_CONTENT, + }) + + expect(result.success).toBe(true) + }) + + it('should PUT a notebook with promoted SQL and no cell ids', async () => { + let sentBody: Record | undefined + addAPIMock({ + method: 'put', + path: '/platform/projects/:ref/content', + response: async ({ request }) => { + sentBody = (await request.json()) as Record + return new HttpResponse(null) + }, + }) + + const tools = getNotebookTools({ projectRef: 'test-project' }) + if (!tools.create_notebook.execute) throw new Error('execute is undefined') + + await tools.create_notebook.execute( + { name: 'Signup funnel', content: VALID_AGENT_CONTENT }, + { toolCallId: 'test', messages: [] } + ) + + expect(sentBody?.type).toBe('notebook') + expect(sentBody?.visibility).toBe('project') + + const content = sentBody?.content as { cells: Array> } + for (const cell of content.cells) { + expect(cell).not.toHaveProperty('id') + } + + const [, databaseCell, logCell] = content.cells + expect(databaseCell.sql).toBe('select * from auth.users limit 100') + expect(logCell.sql).toBe( + "select timestamp, event_message from edge_logs where source = 'edge_logs' limit 10" + ) + }) + + it('should return the id it generated and sent, since a successful create response has no body', async () => { + let sentBody: Record | undefined + addAPIMock({ + method: 'put', + path: '/platform/projects/:ref/content', + response: async ({ request }) => { + sentBody = (await request.json()) as Record + return new HttpResponse(null) + }, + }) + + const tools = getNotebookTools({ projectRef: 'test-project' }) + if (!tools.create_notebook.execute) throw new Error('execute is undefined') + + const result = await tools.create_notebook.execute( + { name: 'Signup funnel', content: VALID_AGENT_CONTENT }, + { toolCallId: 'test', messages: [] } + ) + + expect(typeof sentBody?.id).toBe('string') + expect(result).toEqual({ id: sentBody?.id, name: 'Signup funnel' }) + }) + }) }) diff --git a/apps/studio/lib/ai/tools/notebook-tools.ts b/apps/studio/lib/ai/tools/notebook-tools.ts index 1e5b53c62c7..88556cf5f12 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.ts @@ -1,8 +1,16 @@ +import { acceptUntrustedSql, untrustedSql } from '@supabase/pg-meta' import { tool } from 'ai' import { z } from 'zod' import { getContent } from '@/data/content/content-infinite-query' import { getNotebook } from '@/data/content/notebooks/notebook-query' +import { + agentNotebookSchema, + type WritableCell, + type WritableNotebook, +} from '@/data/content/notebooks/notebook-schema' +import { createNotebook } from '@/data/content/notebooks/notebook-upsert-mutation' +import { acceptUntrustedLogsSql, untrustedLogSql } from '@/data/logs/safe-analytics-sql' import type { Notebooks } from '@/types' export type NotebookToolsContext = { @@ -81,5 +89,46 @@ export const getNotebookTools = (ctx: NotebookToolsContext = {}) => { } }, }), + create_notebook: tool({ + description: + 'Asks the user to create a new notebook with the given cells. Requires user approval before creating.', + inputSchema: z.object({ + name: z.string().describe('A short, descriptive name for the notebook.'), + description: z + .string() + .optional() + .describe('A short description of what the notebook is for.'), + content: agentNotebookSchema.describe( + 'The notebook content: a schema version and an ordered list of cells (markdown, database, or log). Cells must not include an id — one is assigned when the notebook is saved.' + ), + }), + needsApproval: true, + execute: async ({ name, description, content }) => { + const cells: WritableNotebook['cells'] = content.cells.map((cell): WritableCell => { + switch (cell._tag) { + case 'markdown_cell': + return cell + case 'database_cell': + // The `needsApproval: true` gate above is the user gesture that promotes this SQL from untrusted to safe. + return { ...cell, sql: acceptUntrustedSql(untrustedSql(cell.sql)) } + case 'log_cell': + return { ...cell, sql: acceptUntrustedLogsSql(untrustedLogSql(cell.sql)) } + } + }) + + const result = await createNotebook( + { + projectRef: projectRef ?? '', + name, + description, + content: { schema_version: content.schema_version, cells }, + }, + undefined, + authHeaders + ) + + return { id: result.id, name } + }, + }), } } diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index ed76f31dd4f..b4a3a9dc5c2 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -235,6 +235,7 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: Jw orgId, planId, includesLogsSnippets, + isExplorerEnabled: explorerEnabled, requestedModel, systemProviderOptions, abortSignal: abortController.signal,