From 9ae8432dc71af9dbefcd3500feaab01bba91ebe5 Mon Sep 17 00:00:00 2001 From: Charis <26616127+charislam@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:16:58 -0400 Subject: [PATCH] refactor(studio): migrate create, delete, and run notebook tools to tool contracts (#51754) ## Summary Migrates the create, delete, and run notebook assistant tools to the tool contract + server pattern used by `update_notebook`. Each tool now defines its schema once in a shared contract, with server-only execution split into a `.server.ts` module. This removes duplicated client-side schemas and gives the three tools one shared notebook mutation output schema. ## Changes - Add `create-notebook`, `delete-notebook`, and `run-notebook` tool contracts, each with a `.server.ts` implementation. - Remove the duplicated client-side schemas for these tools from `notebook-tools.ts` and the registry. - Add `notebook-mutation-output.ts` as the single output schema shared by the notebook mutation tools. - Add `lib/ai/opt-in-level.ts` with the opt-in level helpers, extracted from the existing module to break an import cycle. - Add optional `ToolContract.outputOnInvalid` so `run_notebook` history replay can reuse a stored output when input fails validation. - Update the AI assistant panel renderers (proposal, run, opt-in modal and request) and the tool filter, sanitizer, and cache invalidation to use the new contracts. ## Deviations worth reviewing - **`outputOnInvalid` on `ToolContract`.** This is a new optional hook. It exists so `run_notebook` history replay can return a valid output when stored input no longer passes validation. Reviewers should confirm this is the right level for the hook rather than handling it inside the tool. - **`opt-in-level.ts` extraction.** The helpers were moved into their own module to break a circular import between the tool modules and the opt-in components. This adds one file, but no behavior change is intended. - **`create-notebook` contract imports the data-layer `notebook-schema`.** This matches `update-notebook`. It should be moved behind the contract before the planned ESLint rule step, which restricts contract modules from importing data-layer code. ## Testing - `vitest` for `lib/ai`, `components/ui/AIAssistantPanel`, and `state`: 920 tests pass across 71 files. - `tsc`: one error remains in `packages/ui-patterns` (`InstructionBlocks.tsx`), which is unrelated to this change. I have not verified whether it also fails on the base branch. - ESLint reports no errors on the changed files (9 `no-explicit-any` warnings in `tool-filter.ts` and its test); Prettier is clean. - Not run: knip and the lint ratchet. --- .../ui/AIAssistantPanel/AIOptInModal.tsx | 2 +- .../ui/AIAssistantPanel/Message.utils.ts | 19 -- .../NotebookProposalRenderer.tsx | 23 +- .../AIAssistantPanel/NotebookRunRenderer.tsx | 7 +- .../ui/AIAssistantPanel/OptInRequest.tsx | 2 +- apps/studio/eslint.config.cjs | 10 +- apps/studio/evals/scorer-wasm.ts | 4 +- .../lib/ai/notebook-cache-invalidation.ts | 17 +- apps/studio/lib/ai/opt-in-level.ts | 21 ++ apps/studio/lib/ai/tool-filter.test.ts | 37 ++- apps/studio/lib/ai/tool-filter.ts | 22 +- .../lib/ai/tools/create-notebook.server.ts | 61 +++++ apps/studio/lib/ai/tools/create-notebook.ts | 24 ++ .../lib/ai/tools/delete-notebook.server.ts | 19 ++ apps/studio/lib/ai/tools/delete-notebook.ts | 18 ++ .../lib/ai/tools/notebook-mutation-output.ts | 4 + .../lib/ai/tools/notebook-run-output.ts | 2 +- apps/studio/lib/ai/tools/notebook-tools.ts | 246 +----------------- apps/studio/lib/ai/tools/registry.test.ts | 57 ++++ apps/studio/lib/ai/tools/registry.ts | 10 +- .../lib/ai/tools/run-notebook.server.ts | 153 +++++++++++ apps/studio/lib/ai/tools/run-notebook.ts | 33 +++ apps/studio/lib/ai/tools/studio-tools.ts | 3 +- apps/studio/lib/ai/tools/tool-sanitizer.ts | 13 +- apps/studio/lib/ai/tools/types.ts | 5 + apps/studio/lib/ai/tools/update-notebook.ts | 5 +- 26 files changed, 487 insertions(+), 330 deletions(-) create mode 100644 apps/studio/lib/ai/opt-in-level.ts create mode 100644 apps/studio/lib/ai/tools/create-notebook.server.ts create mode 100644 apps/studio/lib/ai/tools/create-notebook.ts create mode 100644 apps/studio/lib/ai/tools/delete-notebook.server.ts create mode 100644 apps/studio/lib/ai/tools/delete-notebook.ts create mode 100644 apps/studio/lib/ai/tools/notebook-mutation-output.ts create mode 100644 apps/studio/lib/ai/tools/run-notebook.server.ts create mode 100644 apps/studio/lib/ai/tools/run-notebook.ts diff --git a/apps/studio/components/ui/AIAssistantPanel/AIOptInModal.tsx b/apps/studio/components/ui/AIAssistantPanel/AIOptInModal.tsx index 51726ba0dc8..5aba4b048d7 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AIOptInModal.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AIOptInModal.tsx @@ -19,7 +19,7 @@ import { useAIOptInForm } from '@/hooks/forms/useAIOptInForm' import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization' -import { isOptInLevelAtLeast } from '@/lib/ai/tool-filter' +import { isOptInLevelAtLeast } from '@/lib/ai/opt-in-level' interface AIOptInModalProps { visible: boolean diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.utils.ts b/apps/studio/components/ui/AIAssistantPanel/Message.utils.ts index 01b078911ce..932864198f7 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.utils.ts +++ b/apps/studio/components/ui/AIAssistantPanel/Message.utils.ts @@ -1,8 +1,6 @@ import { untrustedSql } from '@supabase/pg-meta' import { z, type SafeParseReturnType } from 'zod' -import { agentNotebookSchema } from '@/data/content/notebooks/notebook-schema' - // Splits markdown into alternating [plain, code, plain, code, ...] segments. // Odd-indexed segments are already inside code spans/fences and should be left alone. const CODE_SEGMENT_REGEX = /(```[\s\S]*?```|`[^`]*`)/g @@ -128,23 +126,6 @@ export const deployEdgeFunctionOutputSchema = z .object({ success: z.boolean().optional() }) .passthrough() -export const createNotebookInputSchema = z.object({ - name: z.string(), - description: z.string().optional(), - content: agentNotebookSchema, -}) - -export const deleteNotebookInputSchema = z.object({ - id: z.string(), -}) - -export const runNotebookInputSchema = z.object({ - id: z.string(), - expected_updated_at: z.string(), -}) - -export const notebookToolOutputSchema = z.object({ id: z.string(), name: z.string() }) - export const rateMessageResponseSchema = z.object({ category: z.enum([ 'sql_generation', diff --git a/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.tsx b/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.tsx index 089bc787f75..cbc82b07fd5 100644 --- a/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/NotebookProposalRenderer.tsx @@ -9,11 +9,6 @@ import { CodeBlock } from 'ui-patterns/CodeBlock' import { AssistantNotebookPreview } from './AssistantNotebookPreview' import { Confirm } from './Confirm' import { type ConfirmFooterApprovalState } from './Confirm.utils' -import { - createNotebookInputSchema, - deleteNotebookInputSchema, - notebookToolOutputSchema, -} from './Message.utils' import { AlertError } from '@/components/ui/AlertError' import { hasDiscardableChanges } from '@/data/content/notebooks/notebook-cache' import { @@ -23,6 +18,8 @@ import { } from '@/data/content/notebooks/notebook-operations' import { useNotebookQuery } from '@/data/content/notebooks/notebook-query' import { toWireNotebook } from '@/data/content/notebooks/notebook-schema' +import { contract as createNotebookContract } from '@/lib/ai/tools/create-notebook' +import { contract as deleteNotebookContract } from '@/lib/ai/tools/delete-notebook' import { contract as updateNotebookContract } from '@/lib/ai/tools/update-notebook' import { useNotebooksStateSnapshot } from '@/state/notebooks/notebooks-state' @@ -96,6 +93,12 @@ const NOTEBOOK_ACTION_NOUN: Record = { * fetch, and the `expected_updated_at` version check. Wraps the preview in `Confirm` the * same way `AssistantQueryCell` wraps `QueryEditor`. */ +const MODE_CONTRACTS = { + create: createNotebookContract, + update: updateNotebookContract, + delete: deleteNotebookContract, +} as const satisfies Record + export const NotebookProposalRenderer = (props: NotebookProposalRendererProps) => { const { ref } = useParams() const { mode, state, input, output, confirmState, onApprove, onDeny, denyWithReason } = props @@ -107,7 +110,7 @@ export const NotebookProposalRenderer = (props: NotebookProposalRendererProps) = return } - const parsedOutput = notebookToolOutputSchema.safeParse(output) + const parsedOutput = MODE_CONTRACTS[mode].outputSchema.safeParse(output) // A deleted notebook no longer exists to open, so this action only applies to create/update. const canOpenNotebook = mode !== 'delete' && state === 'output-available' const footerAction = @@ -163,7 +166,7 @@ function UpdateNotebookTerminalSummary({ output, }: Pick) { const { ref } = useParams() - const parsedOutput = notebookToolOutputSchema.safeParse(output) + const parsedOutput = updateNotebookContract.outputSchema.safeParse(output) const label = state === 'output-available' ? parsedOutput.success @@ -271,7 +274,7 @@ function CreateNotebookProposal({ onApprove, onDeny, }: NotebookProposalStepProps) { - const parsedInput = createNotebookInputSchema.safeParse(input) + const parsedInput = createNotebookContract.inputSchema.safeParse(input) if (!parsedInput.success) { return ( @@ -523,7 +526,7 @@ function DeleteNotebookProposal({ }: NotebookProposalStepPropsWithOutput) { const { ref } = useParams() const notebooksSnap = useNotebooksStateSnapshot() - const parsedInput = deleteNotebookInputSchema.safeParse(input) + const parsedInput = deleteNotebookContract.inputSchema.safeParse(input) const isCompleted = state === 'output-available' const hasUnsavedLocalChanges = confirmState === 'approval-requested' && @@ -552,7 +555,7 @@ function DeleteNotebookProposal({ } if (isCompleted) { - const parsedOutput = notebookToolOutputSchema.safeParse(output) + const parsedOutput = deleteNotebookContract.outputSchema.safeParse(output) const notebookName = parsedOutput.success ? parsedOutput.data.name : undefined return ( diff --git a/apps/studio/components/ui/AIAssistantPanel/NotebookRunRenderer.tsx b/apps/studio/components/ui/AIAssistantPanel/NotebookRunRenderer.tsx index 28f763095de..b01bc252f30 100644 --- a/apps/studio/components/ui/AIAssistantPanel/NotebookRunRenderer.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/NotebookRunRenderer.tsx @@ -8,11 +8,10 @@ import { AssistantNotebookPreview } from './AssistantNotebookPreview' import { toAssistantQueryResult } from './AssistantQueryCell.utils' import { Confirm } from './Confirm' import type { ConfirmFooterApprovalState } from './Confirm.utils' -import { runNotebookInputSchema } from './Message.utils' import { AlertError } from '@/components/ui/AlertError' import { useNotebookQuery } from '@/data/content/notebooks/notebook-query' import { toWireNotebook } from '@/data/content/notebooks/notebook-schema' -import { notebookRunOutputSchema } from '@/lib/ai/tools/notebook-run-output' +import { contract as runNotebookContract } from '@/lib/ai/tools/run-notebook' export type NotebookRunState = | 'input-available' @@ -40,7 +39,7 @@ export const NotebookRunRenderer = ({ onDeny, }: NotebookRunRendererProps) => { const { ref } = useParams() - const parsedInput = runNotebookInputSchema.safeParse(input) + const parsedInput = runNotebookContract.inputSchema.safeParse(input) const { data: notebook, isLoading, @@ -139,7 +138,7 @@ export const NotebookRunRenderer = ({ _tag: 'unchanged' as const, cell, })) - const parsedOutput = notebookRunOutputSchema.safeParse(output) + const parsedOutput = runNotebookContract.outputSchema.safeParse(output) const isHistoricalRun = state === 'output-available' || state === 'output-error' const referencedUpdatedAt = state === 'output-available' && parsedOutput.success diff --git a/apps/studio/components/ui/AIAssistantPanel/OptInRequest.tsx b/apps/studio/components/ui/AIAssistantPanel/OptInRequest.tsx index ce476900dba..3e8107edf9a 100644 --- a/apps/studio/components/ui/AIAssistantPanel/OptInRequest.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/OptInRequest.tsx @@ -9,7 +9,7 @@ import type { ConfirmFooterApprovalState } from './Confirm.utils' import { NoPermission } from '@/components/ui/NoPermission' import { useAsyncCheckPermissions } from '@/hooks/misc/useCheckPermissions' import { useOrgAiOptInLevel, type AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' -import { optInLevelSchema } from '@/lib/ai/tool-filter' +import { optInLevelSchema } from '@/lib/ai/opt-in-level' const updateOptInLevelOutputSchema = z.object({ levelAfterReview: optInLevelSchema, diff --git a/apps/studio/eslint.config.cjs b/apps/studio/eslint.config.cjs index 667637022f1..35d48bdf5ab 100644 --- a/apps/studio/eslint.config.cjs +++ b/apps/studio/eslint.config.cjs @@ -69,6 +69,14 @@ const TOOL_MODULE_RESTRICTED_IMPORT_PATHS = [ message: APPROVAL_GATE_MESSAGE, }, ] +// The gate itself is also off limits: only `approvalGatedToolFromContract` may hand it out. +const TOOL_MODULE_RESTRICTED_IMPORT_PATTERNS = [ + { + group: ['**/approval-gate', './approval-gate'], + importNames: ['approvalGate'], + message: APPROVAL_GATE_MESSAGE, + }, +] module.exports = defineConfig([ { files: ['**/*.ts', '**/*.tsx'] }, @@ -145,7 +153,7 @@ module.exports = defineConfig([ 'error', { paths: [...BASE_RESTRICTED_IMPORT_PATHS, ...TOOL_MODULE_RESTRICTED_IMPORT_PATHS], - patterns: LUCIDE_ICON_IMPORT_PATTERNS, + patterns: [...LUCIDE_ICON_IMPORT_PATTERNS, ...TOOL_MODULE_RESTRICTED_IMPORT_PATTERNS], }, ], }, diff --git a/apps/studio/evals/scorer-wasm.ts b/apps/studio/evals/scorer-wasm.ts index 114ef234d60..447c866adbb 100644 --- a/apps/studio/evals/scorer-wasm.ts +++ b/apps/studio/evals/scorer-wasm.ts @@ -3,7 +3,7 @@ import { parse } from 'libpg-query' import { AssistantEvalScorer } from './scorer' import { getParsedToolSpans } from './trace-utils' -import { createNotebookInputSchema } from '@/components/ui/AIAssistantPanel/Message.utils' +import { contract as createNotebookContract } from '@/lib/ai/tools/create-notebook' import { executeSqlInputSchema } from '@/lib/ai/tools/studio-tools' import { extractIdentifiers, isQuotedInSql, needsQuoting } from '@/lib/sql-identifier-quoting' @@ -17,7 +17,7 @@ async function getSqlQueries(trace: Trace): Promise { inputSchema: executeSqlInputSchema, }) const createNotebookSpans = await getParsedToolSpans(trace, 'create_notebook', { - inputSchema: createNotebookInputSchema, + inputSchema: createNotebookContract.inputSchema, }) const notebookCellSql = createNotebookSpans.flatMap((s) => diff --git a/apps/studio/lib/ai/notebook-cache-invalidation.ts b/apps/studio/lib/ai/notebook-cache-invalidation.ts index fee7fbb4a69..e0a08ed4e5a 100644 --- a/apps/studio/lib/ai/notebook-cache-invalidation.ts +++ b/apps/studio/lib/ai/notebook-cache-invalidation.ts @@ -1,10 +1,9 @@ import type { QueryClient } from '@tanstack/react-query' import type { ToolUIPart, UIMessage } from 'ai' -import { notebookToolOutputSchema } from '@/components/ui/AIAssistantPanel/Message.utils' import { contentKeys } from '@/data/content/keys' import { evictNotebookFromCaches } from '@/data/content/notebooks/notebook-cache' -import type { ToolName } from '@/lib/ai/tool-filter' +import { TOOL_CONTRACTS, type ContractToolName } from '@/lib/ai/tools/registry' import { notebooksState } from '@/state/notebooks/notebooks-state' export type NotebookCacheEffect = @@ -15,13 +14,21 @@ const NOTEBOOK_MUTATION_TOOLS = [ 'create_notebook', 'update_notebook', 'delete_notebook', -] as const satisfies readonly ToolName[] +] as const satisfies readonly ContractToolName[] + +const NOTEBOOK_MUTATION_CONTRACTS = { + 'tool-create_notebook': TOOL_CONTRACTS.create_notebook, + 'tool-update_notebook': TOOL_CONTRACTS.update_notebook, + 'tool-delete_notebook': TOOL_CONTRACTS.delete_notebook, +} as const const NOTEBOOK_MUTATION_TOOL_TYPES: ReadonlySet = new Set( NOTEBOOK_MUTATION_TOOLS.map((name) => `tool-${name}`) ) -function isNotebookMutationPart(part: UIMessage['parts'][number]): part is ToolUIPart { +function isNotebookMutationPart( + part: UIMessage['parts'][number] +): part is ToolUIPart & { type: keyof typeof NOTEBOOK_MUTATION_CONTRACTS } { return NOTEBOOK_MUTATION_TOOL_TYPES.has(part.type) } @@ -39,7 +46,7 @@ export function collectNotebookCacheEffects( if (part.state !== 'output-available') continue if (processed.has(part.toolCallId)) continue - const result = notebookToolOutputSchema.safeParse(part.output) + const result = NOTEBOOK_MUTATION_CONTRACTS[part.type].outputSchema.safeParse(part.output) if (!result.success) continue effects.push({ diff --git a/apps/studio/lib/ai/opt-in-level.ts b/apps/studio/lib/ai/opt-in-level.ts new file mode 100644 index 00000000000..f2a27c28176 --- /dev/null +++ b/apps/studio/lib/ai/opt-in-level.ts @@ -0,0 +1,21 @@ +import { z } from 'zod' + +import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' + +/** Ordered from least to most data shared with the model. */ +export const OPT_IN_LEVELS = [ + 'disabled', + 'schema', + 'schema_and_log', + 'schema_and_log_and_data', +] as const + +export const optInLevelSchema = z.enum(OPT_IN_LEVELS) + +export function isOptInLevelAtLeast(level: AiOptInLevel, minimum: AiOptInLevel): boolean { + return OPT_IN_LEVELS.indexOf(level) >= OPT_IN_LEVELS.indexOf(minimum) +} + +export function lowerOptInLevel(a: AiOptInLevel, b: AiOptInLevel): AiOptInLevel { + return isOptInLevelAtLeast(a, b) ? b : a +} diff --git a/apps/studio/lib/ai/tool-filter.test.ts b/apps/studio/lib/ai/tool-filter.test.ts index 68e0386125e..9792c9dadbd 100644 --- a/apps/studio/lib/ai/tool-filter.test.ts +++ b/apps/studio/lib/ai/tool-filter.test.ts @@ -9,31 +9,44 @@ import { TOOL_CATEGORY_MAP, toolSetValidationSchema, } from './tool-filter' +import { contract as createNotebookContract } from './tools/create-notebook' +import { contract as deleteNotebookContract } from './tools/delete-notebook' +import { contract as runNotebookContract } from './tools/run-notebook' import { contract as updateNotebookContract } from './tools/update-notebook' describe('TOOL_CATEGORY_MAP', () => { it('should categorize tools correctly', () => { expect(TOOL_CATEGORY_MAP['execute_sql']).toBe(TOOL_CATEGORIES.UI) expect(TOOL_CATEGORY_MAP['list_tables']).toBe(TOOL_CATEGORIES.SCHEMA) - expect(TOOL_CATEGORY_MAP['run_notebook']).toBe(TOOL_CATEGORIES.SCHEMA) + expect(TOOL_CATEGORY_MAP['list_notebooks']).toBe(TOOL_CATEGORIES.SCHEMA) }) }) describe('tools with a contract', () => { - const tools = { - update_notebook: { description: 'Update a notebook', execute: vitest.fn() }, - } as unknown as ToolSet + const contracts = { + create_notebook: createNotebookContract, + delete_notebook: deleteNotebookContract, + run_notebook: runNotebookContract, + update_notebook: updateNotebookContract, + } - it('is gated by the contract minOptInLevel, not the category map', () => { - expect(updateNotebookContract.minOptInLevel).toBe('schema') - expect(Object.hasOwn(TOOL_CATEGORY_MAP, 'update_notebook')).toBe(false) + it.each(Object.entries(contracts))( + '%s is gated by the contract minOptInLevel, not the category map', + (name, contract) => { + const tools = { + [name]: { description: 'A notebook tool', execute: vitest.fn() }, + } as unknown as ToolSet - const blocked = filterToolsByOptInLevel(tools, 'disabled').update_notebook - const allowed = filterToolsByOptInLevel(tools, 'schema').update_notebook + expect(contract.minOptInLevel).toBe('schema') + expect(Object.hasOwn(TOOL_CATEGORY_MAP, name)).toBe(false) - expect(blocked.description).toContain('Requires opting in') - expect(allowed.description).toBe('Update a notebook') - }) + const blocked = filterToolsByOptInLevel(tools, 'disabled')[name] + const allowed = filterToolsByOptInLevel(tools, 'schema')[name] + + expect(blocked.description).toContain('Requires opting in') + expect(allowed.description).toBe('A notebook tool') + } + ) }) describe('tool allowance by opt-in level', () => { diff --git a/apps/studio/lib/ai/tool-filter.ts b/apps/studio/lib/ai/tool-filter.ts index 4ea1cef33a5..77d2b289f57 100644 --- a/apps/studio/lib/ai/tool-filter.ts +++ b/apps/studio/lib/ai/tool-filter.ts @@ -1,6 +1,7 @@ import type { Tool, ToolSet } from 'ai' import { z } from 'zod' +import { isOptInLevelAtLeast, optInLevelSchema } from './opt-in-level' import { getToolContract, type ContractToolName } from './tools/registry' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' @@ -60,16 +61,6 @@ export const toolSetValidationSchema = z.record( ) export type ToolName = keyof z.infer -/** Ordered from least to most data shared with the model. */ -export const OPT_IN_LEVELS = [ - 'disabled', - 'schema', - 'schema_and_log', - 'schema_and_log_and_data', -] as const - -export const optInLevelSchema = z.enum(OPT_IN_LEVELS) - export const updateOptInLevelInputSchema = z.object({ requiredLevel: z .enum(['schema', 'schema_and_log', 'schema_and_log_and_data']) @@ -84,14 +75,6 @@ export const storedUpdateOptInLevelInputSchema = updateOptInLevelInputSchema.ext levelWhenAsked: optInLevelSchema.optional(), }) -export function isOptInLevelAtLeast(level: AiOptInLevel, minimum: AiOptInLevel): boolean { - return OPT_IN_LEVELS.indexOf(level) >= OPT_IN_LEVELS.indexOf(minimum) -} - -export function lowerOptInLevel(a: AiOptInLevel, b: AiOptInLevel): AiOptInLevel { - return isOptInLevelAtLeast(a, b) ? b : a -} - /** * Tool categories based on the data they access */ @@ -135,9 +118,6 @@ export const TOOL_CATEGORY_MAP: Record, Tool list_databases: TOOL_CATEGORIES.SCHEMA, list_notebooks: TOOL_CATEGORIES.SCHEMA, get_notebook: TOOL_CATEGORIES.SCHEMA, - run_notebook: TOOL_CATEGORIES.SCHEMA, - create_notebook: TOOL_CATEGORIES.SCHEMA, - delete_notebook: TOOL_CATEGORIES.SCHEMA, getSchemaTables: TOOL_CATEGORIES.SCHEMA, getRlsKnowledge: TOOL_CATEGORIES.SCHEMA, getFunctions: TOOL_CATEGORIES.SCHEMA, diff --git a/apps/studio/lib/ai/tools/create-notebook.server.ts b/apps/studio/lib/ai/tools/create-notebook.server.ts new file mode 100644 index 00000000000..780e50f052c --- /dev/null +++ b/apps/studio/lib/ai/tools/create-notebook.server.ts @@ -0,0 +1,61 @@ +import { untrustedSql } from '@supabase/pg-meta' + +import { getAuthHeaders } from './auth-headers' +import { contract } from './create-notebook' +import { assertValidDatabaseIdentifiers } from './notebook-tool-error' +import { approvalGatedToolFromContract } from './tool-from-contract' +import type { ApprovalGatedToolFactory } from './types' +import type { WritableCell, WritableNotebook } from '@/data/content/notebooks/notebook-schema' +import { createNotebook } from '@/data/content/notebooks/notebook-upsert-mutation' +import { untrustedLogSql } from '@/data/logs/safe-analytics-sql' +import { getReadReplicas } from '@/data/read-replicas/replicas-query' + +export const createTool: ApprovalGatedToolFactory = (ctx) => { + const { projectRef } = ctx + const authHeaders = getAuthHeaders(ctx.authorization) + + return approvalGatedToolFromContract( + contract, + ctx, + async ({ name, description, content }, { abortSignal }, gate) => { + if ( + content.cells.some( + (cell) => cell._tag === 'database_cell' && cell.database_identifier !== undefined + ) + ) { + const databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders) + assertValidDatabaseIdentifiers( + content.cells, + new Set((databases ?? []).map((database) => database.identifier)) + ) + } + + // `gate` is only handed to an `execute` that runs after the user approves, so that + // approval is the user gesture that promotes each cell's SQL from untrusted to safe. + // Keep the promotion here, not in a shared helper, so it's auditable directly. + const cells: WritableNotebook['cells'] = content.cells.map((cell): WritableCell => { + switch (cell._tag) { + case 'markdown_cell': + return cell + case 'database_cell': + return { ...cell, sql: gate.acceptUntrustedSql(untrustedSql(cell.sql)) } + case 'log_cell': + return { ...cell, sql: gate.acceptUntrustedLogsSql(untrustedLogSql(cell.sql)) } + } + }) + + const result = await createNotebook( + { + projectRef: projectRef ?? '', + name, + description, + content: { schema_version: content.schema_version, cells }, + }, + abortSignal, + authHeaders + ) + + return { id: result.id, name } + } + ) +} diff --git a/apps/studio/lib/ai/tools/create-notebook.ts b/apps/studio/lib/ai/tools/create-notebook.ts new file mode 100644 index 00000000000..eed15604c78 --- /dev/null +++ b/apps/studio/lib/ai/tools/create-notebook.ts @@ -0,0 +1,24 @@ +import { z } from 'zod' + +import { notebookMutationOutputSchema } from './notebook-mutation-output' +import type { ToolContract } from './types' +import { agentNotebookSchema } from '@/data/content/notebooks/notebook-schema' + +const 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.' + ), +}) + +export const contract = { + name: 'create_notebook', + description: + 'Asks the user to create a new notebook with the given cells. Requires user approval before creating.', + inputSchema, + outputSchema: notebookMutationOutputSchema, + minOptInLevel: 'schema', + needsApproval: true, + toModelOutput: ({ id, name }) => ({ id, name }), +} as const satisfies ToolContract diff --git a/apps/studio/lib/ai/tools/delete-notebook.server.ts b/apps/studio/lib/ai/tools/delete-notebook.server.ts new file mode 100644 index 00000000000..fc2cf5c1352 --- /dev/null +++ b/apps/studio/lib/ai/tools/delete-notebook.server.ts @@ -0,0 +1,19 @@ +import { getAuthHeaders } from './auth-headers' +import { contract } from './delete-notebook' +import { toolFromContract } from './tool-from-contract' +import type { ToolFactory } from './types' +import { deleteContents } from '@/data/content/content-delete-mutation' +import { getNotebook } from '@/data/content/notebooks/notebook-query' + +export const createTool: ToolFactory = (ctx) => { + const { projectRef } = ctx + const authHeaders = getAuthHeaders(ctx.authorization) + + return toolFromContract(contract, ctx, async ({ id }, { abortSignal }) => { + const notebook = await getNotebook({ projectRef, id }, abortSignal, authHeaders) + + await deleteContents({ projectRef: projectRef ?? '', ids: [id] }, abortSignal, authHeaders) + + return { id, name: notebook.name } + }) +} diff --git a/apps/studio/lib/ai/tools/delete-notebook.ts b/apps/studio/lib/ai/tools/delete-notebook.ts new file mode 100644 index 00000000000..44c9940c838 --- /dev/null +++ b/apps/studio/lib/ai/tools/delete-notebook.ts @@ -0,0 +1,18 @@ +import { z } from 'zod' + +import { notebookMutationOutputSchema } from './notebook-mutation-output' +import type { ToolContract } from './types' + +const inputSchema = z.object({ + id: z.string().describe('The id of the notebook to delete.'), +}) + +export const contract = { + name: 'delete_notebook', + description: 'Asks the user to delete a notebook. Requires user approval before deleting.', + inputSchema, + outputSchema: notebookMutationOutputSchema, + minOptInLevel: 'schema', + needsApproval: true, + toModelOutput: ({ id, name }) => ({ id, name }), +} as const satisfies ToolContract diff --git a/apps/studio/lib/ai/tools/notebook-mutation-output.ts b/apps/studio/lib/ai/tools/notebook-mutation-output.ts new file mode 100644 index 00000000000..53842241173 --- /dev/null +++ b/apps/studio/lib/ai/tools/notebook-mutation-output.ts @@ -0,0 +1,4 @@ +import { z } from 'zod' + +/** What a notebook mutation (create, update, delete) reports back: the notebook it touched. */ +export const notebookMutationOutputSchema = z.object({ id: z.string(), name: z.string() }) diff --git a/apps/studio/lib/ai/tools/notebook-run-output.ts b/apps/studio/lib/ai/tools/notebook-run-output.ts index 8b2bc214f4e..2e1585439b1 100644 --- a/apps/studio/lib/ai/tools/notebook-run-output.ts +++ b/apps/studio/lib/ai/tools/notebook-run-output.ts @@ -1,7 +1,7 @@ import type { JSONValue } from 'ai' import { z } from 'zod' -import { lowerOptInLevel, optInLevelSchema } from '../tool-filter' +import { lowerOptInLevel, optInLevelSchema } from '../opt-in-level' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' const jsonValueSchema: z.ZodType = z.lazy(() => diff --git a/apps/studio/lib/ai/tools/notebook-tools.ts b/apps/studio/lib/ai/tools/notebook-tools.ts index 2229c2c02ed..07913af05c9 100644 --- a/apps/studio/lib/ai/tools/notebook-tools.ts +++ b/apps/studio/lib/ai/tools/notebook-tools.ts @@ -1,46 +1,20 @@ -import { acceptUntrustedSql, untrustedSql } from '@supabase/pg-meta' -import { tool, type JSONValue } from 'ai' +import { tool } from 'ai' import { z } from 'zod' import { getAuthHeaders } from './auth-headers' -import { - sanitizeNotebookRunOutput, - type NotebookRunCellOutput, - type NotebookRunOutput, -} from './notebook-run-output' -import { assertValidDatabaseIdentifiers, NotebookToolError } from './notebook-tool-error' +import { createTool as createCreateNotebookTool } from './create-notebook.server' +import { createTool as createDeleteNotebookTool } from './delete-notebook.server' +import { createTool as createRunNotebookTool } from './run-notebook.server' import type { ToolContext } from './types' import { createTool as createUpdateNotebookTool } from './update-notebook.server' -import { resolveLogTimeRange } from '@/components/interfaces/QuerySources/LogTimeRange.utils' -import { deleteContents } from '@/data/content/content-delete-mutation' import { getContent } from '@/data/content/content-infinite-query' import { getNotebook } from '@/data/content/notebooks/notebook-query' -import { - agentNotebookSchema, - toWireNotebook, - type WritableCell, - type WritableNotebook, -} from '@/data/content/notebooks/notebook-schema' -import { createNotebook } from '@/data/content/notebooks/notebook-upsert-mutation' -import { executeLogsSql } from '@/data/logs/execute-logs-sql-mutation' -import { acceptUntrustedLogsSql, untrustedLogSql } from '@/data/logs/safe-analytics-sql' -import { QUERY_SOURCE_REGISTRY } from '@/data/query-sources/query-source-registry' +import { toWireNotebook } from '@/data/content/notebooks/notebook-schema' import { getReadReplicas } from '@/data/read-replicas/replicas-query' -import { executeSql } from '@/data/sql/execute-sql-mutation' -import { applyAutoLimit } from '@/data/sql/utils' import type { Notebooks } from '@/types' -function getErrorMessage(error: unknown): string { - if (error instanceof Error) return error.message - if (typeof error === 'object' && error !== null && 'message' in error) { - const message = (error as { message?: unknown }).message - if (typeof message === 'string') return message - } - return 'An unexpected error occurred while running this query.' -} - export const getNotebookTools = (ctx: ToolContext = {}) => { - const { projectRef, connectionString, authorization, aiOptInLevel = 'schema' } = ctx + const { projectRef, authorization } = ctx const authHeaders = getAuthHeaders(authorization) return { @@ -123,211 +97,9 @@ export const getNotebookTools = (ctx: ToolContext = {}) => { } }, }), - run_notebook: tool({ - description: - 'Asks the user to run every database and logs query cell in a notebook, in notebook order, and returns each result according to the organization data-sharing permissions. Requires one user approval before any cell executes.', - inputSchema: z.object({ - id: z.string().describe('The id of the notebook to run.'), - expected_updated_at: z - .string() - .describe( - 'The `updated_at` received from `get_notebook`. The run is rejected if the notebook changed since.' - ), - }), - needsApproval: true, - execute: async ({ id, expected_updated_at }, { abortSignal }) => { - const notebook = await getNotebook({ projectRef, id }, abortSignal, authHeaders) - - if (notebook.updated_at !== expected_updated_at) { - throw new NotebookToolError( - `Notebook "${id}" changed since expected_updated_at (${expected_updated_at}); it is now ${notebook.updated_at}. Call get_notebook again and reissue run_notebook against the current content.`, - { exposeToAssistant: true } - ) - } - - const queryCells = notebook.content.cells.filter( - (cell) => cell._tag === 'database_cell' || cell._tag === 'log_cell' - ) - const shouldLookupReplica = queryCells.some( - (cell) => - cell._tag === 'database_cell' && - cell.database_identifier !== undefined && - cell.database_identifier !== projectRef - ) - let databases: Awaited> = [] - let replicaLookupFailure: { error: unknown } | undefined - if (shouldLookupReplica) { - try { - databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders) - } catch (error) { - replicaLookupFailure = { error } - } - } - - const cells: NotebookRunCellOutput[] = [] - - // This execute function only runs after the AI SDK approval gate above resolves true. - // That approval is the explicit user gesture which promotes every persisted cell's - // untrusted SQL immediately before execution. Run sequentially to preserve notebook - // order and avoid racing cells that may depend on earlier writes. - for (const cell of queryCells) { - const title = cell.title?.trim() || 'Untitled query' - - try { - if (cell._tag === 'log_cell') { - const result = await executeLogsSql({ - projectRef: projectRef ?? '', - sql: acceptUntrustedLogsSql(cell.unchecked_sql), - range: resolveLogTimeRange(cell.time_range), - endpoint: QUERY_SOURCE_REGISTRY.logs.endpoint, - signal: abortSignal, - headers: authHeaders, - }) - if (result.error) throw result.error - cells.push({ - cell_id: cell._id, - title, - source: 'logs', - status: 'success', - // The analytics response crossed a JSON fetch boundary, although its shared - // data-layer type remains `unknown[]` because individual endpoints vary. - rows: result.rows as JSONValue[], - }) - continue - } - - const isReplicaTarget = - cell.database_identifier !== undefined && cell.database_identifier !== projectRef - - if (isReplicaTarget && replicaLookupFailure !== undefined) { - throw replicaLookupFailure.error - } - - const selectedConnectionString = !isReplicaTarget - ? connectionString - : databases?.find((database) => database.identifier === cell.database_identifier) - ?.connectionString - - if (!selectedConnectionString) { - throw new Error( - `Unable to run query: connection string is missing for database "${cell.database_identifier ?? projectRef ?? 'primary'}"` - ) - } - - const limitedSql = applyAutoLimit( - acceptUntrustedSql(cell.unchecked_sql), - cell.row_limit - ) - const { result } = await executeSql( - { - projectRef, - connectionString: selectedConnectionString, - sql: limitedSql.sql, - isStatementTimeoutDisabled: true, - }, - abortSignal, - authHeaders - ) - cells.push({ - cell_id: cell._id, - title, - source: 'database', - status: 'success', - rows: Array.isArray(result) ? result : [], - }) - } catch (error) { - cells.push({ - cell_id: cell._id, - title, - source: cell._tag === 'log_cell' ? 'logs' : 'database', - status: 'error', - error: { message: getErrorMessage(error) }, - }) - } - } - - return { - id: notebook.id, - name: notebook.name, - updated_at: notebook.updated_at, - optInLevel: aiOptInLevel, - cells, - } satisfies NotebookRunOutput - }, - toModelOutput: ({ output }) => ({ - type: 'json', - value: sanitizeNotebookRunOutput(output, aiOptInLevel), - }), - }), - 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 }, { abortSignal }) => { - if ( - content.cells.some( - (cell) => cell._tag === 'database_cell' && cell.database_identifier !== undefined - ) - ) { - const databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders) - assertValidDatabaseIdentifiers( - content.cells, - new Set((databases ?? []).map((database) => database.identifier)) - ) - } - - // This approval gate is the user gesture that promotes each cell's SQL from - // untrusted to safe — keep the promotion here, not in a shared helper, so it's - // auditable directly alongside the `needsApproval: true` above. - const cells: WritableNotebook['cells'] = content.cells.map((cell): WritableCell => { - switch (cell._tag) { - case 'markdown_cell': - return cell - case 'database_cell': - 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 }, - }, - abortSignal, - authHeaders - ) - - return { id: result.id, name } - }, - }), + run_notebook: createRunNotebookTool(ctx), + create_notebook: createCreateNotebookTool(ctx), update_notebook: createUpdateNotebookTool(ctx), - delete_notebook: tool({ - description: 'Asks the user to delete a notebook. Requires user approval before deleting.', - inputSchema: z.object({ - id: z.string().describe('The id of the notebook to delete.'), - }), - needsApproval: true, - execute: async ({ id }, { abortSignal }) => { - const notebook = await getNotebook({ projectRef, id }, abortSignal, authHeaders) - - await deleteContents({ projectRef: projectRef ?? '', ids: [id] }, abortSignal, authHeaders) - - return { id, name: notebook.name } - }, - }), + delete_notebook: createDeleteNotebookTool(ctx), } } diff --git a/apps/studio/lib/ai/tools/registry.test.ts b/apps/studio/lib/ai/tools/registry.test.ts index 1a83a862b88..f69ae096bb8 100644 --- a/apps/studio/lib/ai/tools/registry.test.ts +++ b/apps/studio/lib/ai/tools/registry.test.ts @@ -1,6 +1,7 @@ import type { ToolUIPart } from 'ai' import { describe, expect, expectTypeOf, it } from 'vitest' +import { INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE } from './notebook-run-output' import { sanitizeContractPart, type StudioUIMessage } from './registry' describe('TOOL_CONTRACTS', () => { @@ -46,3 +47,59 @@ describe('sanitizeContractPart', () => { ).toBe(undefined) }) }) + +describe('sanitizeContractPart for run_notebook', () => { + const output = { + id: 'notebook-1', + name: 'Notebook', + updated_at: '2026-01-01T00:00:00.000Z', + optInLevel: 'schema_and_log_and_data', + cells: [ + { + cell_id: 'db', + title: 'Database', + source: 'database', + status: 'success', + rows: [{ secret: 'value' }], + }, + ], + } + const part = { + type: 'tool-run_notebook', + toolCallId: 'call-2', + state: 'output-available', + input: {}, + output, + } as ToolUIPart + + it('hides rows below the data level and keeps them at the data level', () => { + const hidden = sanitizeContractPart(part, 'schema')?.output as typeof output + expect(hidden.cells[0]).not.toHaveProperty('rows') + expect(sanitizeContractPart(part, 'schema_and_log_and_data')?.output).toEqual(output) + }) + + it('replaces an invalid output with an explanation instead of dropping it', () => { + expect(sanitizeContractPart({ ...part, output: 'oops' } as ToolUIPart, 'schema')?.output).toBe( + INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE + ) + }) +}) + +describe('notebook mutation contracts', () => { + it.each(['create_notebook', 'delete_notebook'] as const)( + '%s output is limited to id and name on replay', + (name) => { + const part = { + type: `tool-${name}`, + toolCallId: 'call-3', + state: 'output-available', + input: {}, + output: { id: 'notebook-1', name: 'Notebook', extra: 'ignored' }, + } as ToolUIPart + expect(sanitizeContractPart(part, 'schema')?.output).toEqual({ + id: 'notebook-1', + name: 'Notebook', + }) + } + ) +}) diff --git a/apps/studio/lib/ai/tools/registry.ts b/apps/studio/lib/ai/tools/registry.ts index d126e26bf34..3a65a417013 100644 --- a/apps/studio/lib/ai/tools/registry.ts +++ b/apps/studio/lib/ai/tools/registry.ts @@ -1,6 +1,9 @@ import type { InferUITools, Tool, ToolUIPart, UIDataTypes, UIMessage } from 'ai' import type { ToolName } from '../tool-filter' +import { contract as createNotebook } from './create-notebook' +import { contract as deleteNotebook } from './delete-notebook' +import { contract as runNotebook } from './run-notebook' import type { ToolContract, ToolContractInput, ToolContractOutput } from './types' import { contract as updateNotebook } from './update-notebook' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' @@ -10,6 +13,9 @@ import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' * by hand in the places the contract replaces. */ export const TOOL_CONTRACTS = { + create_notebook: createNotebook, + delete_notebook: deleteNotebook, + run_notebook: runNotebook, update_notebook: updateNotebook, } as const satisfies Partial> @@ -53,6 +59,8 @@ export function sanitizeContractPart( const parsed = contract.outputSchema.safeParse(part.output) return { ...part, - output: parsed.success ? contract.toModelOutput(parsed.data, optInLevel) : undefined, + output: parsed.success + ? contract.toModelOutput(parsed.data, optInLevel) + : contract.outputOnInvalid, } } diff --git a/apps/studio/lib/ai/tools/run-notebook.server.ts b/apps/studio/lib/ai/tools/run-notebook.server.ts new file mode 100644 index 00000000000..7b7bdc1fdb3 --- /dev/null +++ b/apps/studio/lib/ai/tools/run-notebook.server.ts @@ -0,0 +1,153 @@ +import type { JSONValue } from 'ai' + +import { getAuthHeaders } from './auth-headers' +import type { NotebookRunCellOutput, NotebookRunOutput } from './notebook-run-output' +import { NotebookToolError } from './notebook-tool-error' +import { contract } from './run-notebook' +import { approvalGatedToolFromContract } from './tool-from-contract' +import type { ApprovalGatedToolFactory } from './types' +import { resolveLogTimeRange } from '@/components/interfaces/QuerySources/LogTimeRange.utils' +import { getNotebook } from '@/data/content/notebooks/notebook-query' +import { executeLogsSql } from '@/data/logs/execute-logs-sql-mutation' +import { QUERY_SOURCE_REGISTRY } from '@/data/query-sources/query-source-registry' +import { getReadReplicas } from '@/data/read-replicas/replicas-query' +import { executeSql } from '@/data/sql/execute-sql-mutation' +import { applyAutoLimit } from '@/data/sql/utils' + +function getErrorMessage(error: unknown): string { + if (error instanceof Error) return error.message + if (typeof error === 'object' && error !== null && 'message' in error) { + const message = (error as { message?: unknown }).message + if (typeof message === 'string') return message + } + return 'An unexpected error occurred while running this query.' +} + +export const createTool: ApprovalGatedToolFactory = (ctx) => { + const { projectRef, connectionString, aiOptInLevel = 'schema' } = ctx + const authHeaders = getAuthHeaders(ctx.authorization) + + return approvalGatedToolFromContract( + contract, + ctx, + async ({ id, expected_updated_at }, { abortSignal }, gate) => { + const notebook = await getNotebook({ projectRef, id }, abortSignal, authHeaders) + + if (notebook.updated_at !== expected_updated_at) { + throw new NotebookToolError( + `Notebook "${id}" changed since expected_updated_at (${expected_updated_at}); it is now ${notebook.updated_at}. Call get_notebook again and reissue run_notebook against the current content.`, + { exposeToAssistant: true } + ) + } + + const queryCells = notebook.content.cells.filter( + (cell) => cell._tag === 'database_cell' || cell._tag === 'log_cell' + ) + const shouldLookupReplica = queryCells.some( + (cell) => + cell._tag === 'database_cell' && + cell.database_identifier !== undefined && + cell.database_identifier !== projectRef + ) + let databases: Awaited> = [] + let replicaLookupFailure: { error: unknown } | undefined + if (shouldLookupReplica) { + try { + databases = await getReadReplicas({ projectRef }, abortSignal, authHeaders) + } catch (error) { + replicaLookupFailure = { error } + } + } + + const cells: NotebookRunCellOutput[] = [] + + // This execute function only runs after the AI SDK approval gate resolves true. That + // approval is the explicit user gesture which promotes every persisted cell's + // untrusted SQL immediately before execution. Run sequentially to preserve notebook + // order and avoid racing cells that may depend on earlier writes. + for (const cell of queryCells) { + const title = cell.title?.trim() || 'Untitled query' + + try { + if (cell._tag === 'log_cell') { + const result = await executeLogsSql({ + projectRef: projectRef ?? '', + sql: gate.acceptUntrustedLogsSql(cell.unchecked_sql), + range: resolveLogTimeRange(cell.time_range), + endpoint: QUERY_SOURCE_REGISTRY.logs.endpoint, + signal: abortSignal, + headers: authHeaders, + }) + if (result.error) throw result.error + cells.push({ + cell_id: cell._id, + title, + source: 'logs', + status: 'success', + // The analytics response crossed a JSON fetch boundary, although its shared + // data-layer type remains `unknown[]` because individual endpoints vary. + rows: result.rows as JSONValue[], + }) + continue + } + + const isReplicaTarget = + cell.database_identifier !== undefined && cell.database_identifier !== projectRef + + if (isReplicaTarget && replicaLookupFailure !== undefined) { + throw replicaLookupFailure.error + } + + const selectedConnectionString = !isReplicaTarget + ? connectionString + : databases?.find((database) => database.identifier === cell.database_identifier) + ?.connectionString + + if (!selectedConnectionString) { + throw new Error( + `Unable to run query: connection string is missing for database "${cell.database_identifier ?? projectRef ?? 'primary'}"` + ) + } + + const limitedSql = applyAutoLimit( + gate.acceptUntrustedSql(cell.unchecked_sql), + cell.row_limit + ) + const { result } = await executeSql( + { + projectRef, + connectionString: selectedConnectionString, + sql: limitedSql.sql, + isStatementTimeoutDisabled: true, + }, + abortSignal, + authHeaders + ) + cells.push({ + cell_id: cell._id, + title, + source: 'database', + status: 'success', + rows: Array.isArray(result) ? result : [], + }) + } catch (error) { + cells.push({ + cell_id: cell._id, + title, + source: cell._tag === 'log_cell' ? 'logs' : 'database', + status: 'error', + error: { message: getErrorMessage(error) }, + }) + } + } + + return { + id: notebook.id, + name: notebook.name, + updated_at: notebook.updated_at, + optInLevel: aiOptInLevel, + cells, + } satisfies NotebookRunOutput + } + ) +} diff --git a/apps/studio/lib/ai/tools/run-notebook.ts b/apps/studio/lib/ai/tools/run-notebook.ts new file mode 100644 index 00000000000..fba5b1a0cc6 --- /dev/null +++ b/apps/studio/lib/ai/tools/run-notebook.ts @@ -0,0 +1,33 @@ +import { z } from 'zod' + +import { + INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE, + notebookRunOutputSchema, + sanitizeNotebookRunOutput, + type NotebookRunOutput, +} from './notebook-run-output' +import type { ToolContract } from './types' + +const inputSchema = z.object({ + id: z.string().describe('The id of the notebook to run.'), + expected_updated_at: z + .string() + .describe( + 'The `updated_at` received from `get_notebook`. The run is rejected if the notebook changed since.' + ), +}) + +export const contract = { + name: 'run_notebook', + description: + 'Asks the user to run every database and logs query cell in a notebook, in notebook order, and returns each result according to the organization data-sharing permissions. Requires one user approval before any cell executes.', + inputSchema, + outputSchema: notebookRunOutputSchema, + minOptInLevel: 'schema', + needsApproval: true, + toModelOutput: (output: NotebookRunOutput, optInLevel) => + sanitizeNotebookRunOutput(output, optInLevel), + // Persisted output comes from the client, so a malformed run is replaced with an + // explanation rather than dropped. + outputOnInvalid: INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE, +} as const satisfies ToolContract diff --git a/apps/studio/lib/ai/tools/studio-tools.ts b/apps/studio/lib/ai/tools/studio-tools.ts index 7f2ca4f325e..28b93d772bb 100644 --- a/apps/studio/lib/ai/tools/studio-tools.ts +++ b/apps/studio/lib/ai/tools/studio-tools.ts @@ -9,6 +9,7 @@ import { DEFAULT_CELL_ROW_LIMIT } from '@/components/interfaces/Explorer/QueryCe import { deployEdgeFunction } from '@/data/edge-functions/edge-functions-deploy-mutation' import { executeSql } from '@/data/sql/execute-sql-mutation' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' +import { isOptInLevelAtLeast } from '@/lib/ai/opt-in-level' import { EDGE_FUNCTION_PROMPT, LOGS_PROMPT, @@ -17,7 +18,7 @@ import { RLS_PROMPT, STORAGE_PROMPT, } from '@/lib/ai/prompts' -import { isOptInLevelAtLeast, updateOptInLevelInputSchema } from '@/lib/ai/tool-filter' +import { updateOptInLevelInputSchema } from '@/lib/ai/tool-filter' import { NO_DATA_PERMISSIONS, toModelRows } from '@/lib/ai/tools/tool-sanitizer' import { fixSqlBackslashEscapes } from '@/lib/ai/util' diff --git a/apps/studio/lib/ai/tools/tool-sanitizer.ts b/apps/studio/lib/ai/tools/tool-sanitizer.ts index 34dce217f03..deb8625eada 100644 --- a/apps/studio/lib/ai/tools/tool-sanitizer.ts +++ b/apps/studio/lib/ai/tools/tool-sanitizer.ts @@ -1,8 +1,8 @@ import type { ToolUIPart, UIMessage } from 'ai' import { z } from 'zod' -import { lowerOptInLevel, optInLevelSchema, type ToolName } from '../tool-filter' -import { sanitizeNotebookRunOutput } from './notebook-run-output' +import { lowerOptInLevel, optInLevelSchema } from '../opt-in-level' +import type { ToolName } from '../tool-filter' import { sanitizeContractPart } from './registry' import { DEFAULT_CELL_ROW_LIMIT } from '@/components/interfaces/Explorer/QueryCell/QueryCell.utils' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' @@ -68,17 +68,8 @@ const executeSqlSanitizer: ToolSanitizer = { }, } -const runNotebookSanitizer: ToolSanitizer = { - toolName: 'run_notebook', - sanitize: (tool, optInLevel) => ({ - ...tool, - output: sanitizeNotebookRunOutput(tool.output, optInLevel), - }), -} - export const ALL_TOOL_SANITIZERS = { [executeSqlSanitizer.toolName]: executeSqlSanitizer, - [runNotebookSanitizer.toolName]: runNotebookSanitizer, } export function sanitizeMessagePart( diff --git a/apps/studio/lib/ai/tools/types.ts b/apps/studio/lib/ai/tools/types.ts index 54dbae60a4e..e5832cceacd 100644 --- a/apps/studio/lib/ai/tools/types.ts +++ b/apps/studio/lib/ai/tools/types.ts @@ -27,6 +27,11 @@ export type ToolContract = { * its own result. Return the output unchanged only if all of it is safe at every level. */ toModelOutput(output: TOutput, optInLevel: AiOptInLevel): JSONValue + /** + * What history replay shows the model when a persisted output doesn't match `outputSchema`. + * Omit it to drop the output. + */ + outputOnInvalid?: JSONValue } export type ToolContractInput = z.infer diff --git a/apps/studio/lib/ai/tools/update-notebook.ts b/apps/studio/lib/ai/tools/update-notebook.ts index 40cfefea35b..22076d86a97 100644 --- a/apps/studio/lib/ai/tools/update-notebook.ts +++ b/apps/studio/lib/ai/tools/update-notebook.ts @@ -1,5 +1,6 @@ import { z } from 'zod' +import { notebookMutationOutputSchema } from './notebook-mutation-output' import type { ToolContract } from './types' import { notebookOperationsSchema } from '@/data/content/notebooks/notebook-operations' import { notebookSchema } from '@/data/content/notebooks/notebook-schema' @@ -20,9 +21,7 @@ const inputSchema = z.object({ * `previous_content` lets the client re-derive the diff it already showed for approval, * without re-fetching a notebook that's now post-update. It is UI-only: `toModelOutput` drops it. */ -const outputSchema = z.object({ - id: z.string(), - name: z.string(), +const outputSchema = notebookMutationOutputSchema.extend({ previous_content: notebookSchema.optional(), })