refactor(studio): migrate execute_sql, deploy_edge_function, rename_chat, and load_knowledge to tool contracts (#51803)

## Summary

Migrates four more assistant tools to the contract + `.server.ts`
pattern used by the notebook tools. Each tool now declares its schemas,
opt-in level, approval and model-visible output once.

## Changes

* Add contracts and server factories for `execute_sql`,
`deploy_edge_function`, `rename_chat`, `load_knowledge`;
`getStudioTools` now composes them.
* Add `execute-sql-output.ts` (output schema, `toModelRows`,
`NO_DATA_PERMISSIONS`) and `knowledge-topics.ts` (topic names + UI
labels, replacing the copy in `Message.Parts.utils.ts`).
* Delete `executeSqlSanitizer` and `ALL_TOOL_SANITIZERS`; history replay
goes through `sanitizeContractPart` only.
* Drop the four tools from `TOOL_CATEGORY_MAP`; the UI and evals import
schemas from the contracts; `rename_chat` input casts now use the
contract type.

## Deviations worth reviewing

* `execute_sql` replay: a bare-array legacy output is stamped `schema`
via a union transform (still hidden from the model, still shown in the
UI). A malformed non-array output is now dropped, where before it passed
through at the data level.
* The "no data permission" message now reaches the model as a JSON
string value instead of a text part. The text the model sees is
unchanged.
* `deploy_edge_function` output is now strictly `{ success: boolean }`
in the UI parse (was optional/passthrough).

apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.utils.ts

@@ -10,7 +10,7 @@ import {
   toQuerySourceBinding,
   type QuerySourceBinding,
 } from '@/data/query-sources/query-source-registry'
-import { executeSqlOutputSchema } from '@/lib/ai/tools/tool-sanitizer'
+import { contract as executeSqlContract } from '@/lib/ai/tools/execute-sql'
 
 export const DEFAULT_ASSISTANT_QUERY_TITLE = 'SQL query'
 export const DEFAULT_ASSISTANT_LOGS_QUERY_TITLE = 'Logs query'
@@ -109,7 +109,7 @@ export function inferAssistantChartDisplay(rows: readonly Record<string, unknown
 }
 
 export function toAssistantQueryResult(output: unknown): QueryResult | undefined {
-  const rows = Array.isArray(output) ? output : executeSqlOutputSchema.safeParse(output).data?.rows
+  const rows = executeSqlContract.outputSchema.safeParse(output).data?.rows
   return rows ? { rows: rows.filter(isPlainRow) } : undefined
 }
 

apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx

@@ -23,11 +23,7 @@ import {
   type CompactPart,
   type CompactPartStatus,
 } from './Message.Parts.utils'
-import {
-  deployEdgeFunctionInputSchema,
-  deployEdgeFunctionOutputSchema,
-  parseExecuteSqlChartResult,
-} from './Message.utils'
+import { deployEdgeFunctionInputSchema, parseExecuteSqlChartResult } from './Message.utils'
 import { MessageMarkdown } from './MessageMarkdown'
 import { MessagePartQueryLogs } from './MessagePartQueryLogs'
 import { NotebookProposalRenderer, type NotebookProposalMode } from './NotebookProposalRenderer'
@@ -36,6 +32,7 @@ import { OptInRequest } from './OptInRequest'
 import { parseSupportRequestMessage, SupportRequestMessage } from './SupportRequestMessage'
 import { useMinimumDisplayTime } from '@/hooks/misc/useMinimumDisplayTime'
 import { storedUpdateOptInLevelInputSchema } from '@/lib/ai/tool-filter'
+import { contract as deployEdgeFunctionContract } from '@/lib/ai/tools/deploy-edge-function'
 
 function MessagePartText({ textPart }: { textPart: TextUIPart }) {
   const { id, isLoading, readOnly, isUserMessage, state } = useMessageInfoContext()
@@ -182,7 +179,7 @@ function MessagePartDeployEdgeFunction({ toolPart }: { toolPart: ToolUIPart }) {
   const parsedInput = deployEdgeFunctionInputSchema.safeParse(input)
   if (!parsedInput.success) return null
 
-  const parsedOutput = deployEdgeFunctionOutputSchema.safeParse(output)
+  const parsedOutput = deployEdgeFunctionContract.outputSchema.safeParse(output)
   const isInitiallyDeployed =
     state === 'output-available' && parsedOutput.success && parsedOutput.data.success === true
 

apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.ts

@@ -8,6 +8,8 @@ import {
 } from 'ai'
 import isEqual from 'lodash/isEqual'
 
+import { KNOWLEDGE_TOPIC_LABELS, type KnowledgeName } from '@/lib/ai/tools/knowledge-topics'
+
 type MessagePart = UIMessage['parts'][number]
 type ToolPart = ToolUIPart | DynamicToolUIPart
 /** A one-line row that folds into a tool group: reasoning, or a lookup tool call. */
@@ -59,18 +61,11 @@ function forDocsQuery(input: unknown) {
   return query ? `for "${query}"` : undefined
 }
 
-const KNOWLEDGE_TOPICS: Record<string, string> = {
-  pg_best_practices: 'Postgres best practices',
-  rls: 'RLS',
-  storage: 'Storage',
-  edge_functions: 'Edge Functions',
-  realtime: 'Realtime',
-  logs: 'logs',
-}
-
 function onKnowledgeTopic(input: unknown) {
   const name = String(getField(input, 'name'))
-  return Object.hasOwn(KNOWLEDGE_TOPICS, name) ? `on ${KNOWLEDGE_TOPICS[name]}` : undefined
+  return Object.hasOwn(KNOWLEDGE_TOPIC_LABELS, name)
+    ? `on ${KNOWLEDGE_TOPIC_LABELS[name as KnowledgeName]}`
+    : undefined
 }
 
 function forAdvisorType(input: unknown) {

apps/studio/components/ui/AIAssistantPanel/Message.utils.ts

@@ -122,10 +122,6 @@ export const deployEdgeFunctionInputSchema = z
     }
   })
 
-export const deployEdgeFunctionOutputSchema = z
-  .object({ success: z.boolean().optional() })
-  .passthrough()
-
 export const rateMessageResponseSchema = z.object({
   category: z.enum([
     'sql_generation',

apps/studio/evals/scorer-wasm.ts

@@ -4,7 +4,7 @@ import { parse } from 'libpg-query'
 import { AssistantEvalScorer } from './scorer'
 import { getParsedToolSpans } from './trace-utils'
 import { contract as createNotebookContract } from '@/lib/ai/tools/create-notebook'
-import { executeSqlInputSchema } from '@/lib/ai/tools/studio-tools'
+import { contract as executeSqlContract } from '@/lib/ai/tools/execute-sql'
 import { extractIdentifiers, isQuotedInSql, needsQuoting } from '@/lib/sql-identifier-quoting'
 
 /**
@@ -14,7 +14,7 @@ import { extractIdentifiers, isQuotedInSql, needsQuoting } from '@/lib/sql-ident
  */
 async function getSqlQueries(trace: Trace): Promise<string[]> {
   const executeSqlSpans = await getParsedToolSpans(trace, 'execute_sql', {
-    inputSchema: executeSqlInputSchema,
+    inputSchema: executeSqlContract.inputSchema,
   })
   const createNotebookSpans = await getParsedToolSpans(trace, 'create_notebook', {
     inputSchema: createNotebookContract.inputSchema,

apps/studio/evals/scorer.ts

@@ -7,7 +7,7 @@ import { z } from 'zod'
 import { getParsedToolSpans, getThreadParts, getToolSpans } from './trace-utils'
 import type { Transcript } from './transcript'
 import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
-import { loadKnowledgeInputSchema } from '@/lib/ai/tools/studio-tools'
+import { contract as loadKnowledgeContract } from '@/lib/ai/tools/load-knowledge'
 import { extractUrls } from '@/lib/helpers'
 
 const LLM_AS_A_JUDGE_MODEL = 'gpt-5.2' // NOTE: `gpt-5.2-2025-12-11` snapshot not yet working with online scorers
@@ -166,7 +166,7 @@ export const knowledgeUsageScorer: AssistantEvalScorer = async ({ expected, trac
   if (!expected.requiredKnowledge || !trace) return null
 
   const knowledgeSpans = await getParsedToolSpans(trace, 'load_knowledge', {
-    inputSchema: loadKnowledgeInputSchema,
+    inputSchema: loadKnowledgeContract.inputSchema,
   })
   const loadedKnowledge: string[] = knowledgeSpans.map((s) => s.input.name)
 

apps/studio/lib/ai/generate-assistant-response.ts

@@ -24,6 +24,8 @@ import {
   NOTEBOOKS_PROMPT,
   SECURITY_PROMPT,
 } from '@/lib/ai/prompts'
+import type { contract as renameChatContract } from '@/lib/ai/tools/rename-chat'
+import type { ToolContractInput } from '@/lib/ai/tools/types'
 
 const { streamText: tracedStreamText } = wrapAISDK(ai)
 
@@ -159,7 +161,7 @@ export async function generateAssistantResponse({
           for (const step of steps) {
             for (const toolCall of step.toolCalls) {
               if (toolCall.toolName === 'rename_chat') {
-                const { newName } = toolCall.input as { newName: string }
+                const { newName } = toolCall.input as ToolContractInput<typeof renameChatContract>
                 metadata.chatName = newName
               }
             }

apps/studio/lib/ai/tool-filter.test.ts

@@ -16,7 +16,7 @@ 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['search_docs']).toBe(TOOL_CATEGORIES.UI)
     expect(TOOL_CATEGORY_MAP['list_tables']).toBe(TOOL_CATEGORIES.SCHEMA)
     expect(TOOL_CATEGORY_MAP['list_notebooks']).toBe(TOOL_CATEGORIES.SCHEMA)
   })

apps/studio/lib/ai/tool-filter.ts

@@ -97,15 +97,11 @@ type ToolCategory = (typeof TOOL_CATEGORIES)[keyof typeof TOOL_CATEGORIES]
  */
 export const TOOL_CATEGORY_MAP: Record<Exclude<ToolName, ContractToolName>, ToolCategory> = {
   // UI tools - always available
-  execute_sql: TOOL_CATEGORIES.UI,
-  deploy_edge_function: TOOL_CATEGORIES.UI,
-  rename_chat: TOOL_CATEGORIES.UI,
   update_opt_in_level: TOOL_CATEGORIES.UI,
   escalate_to_human: TOOL_CATEGORIES.UI,
   resolve_support_conversation: TOOL_CATEGORIES.UI,
   search_docs: TOOL_CATEGORIES.UI,
   get_active_incidents: TOOL_CATEGORIES.UI,
-  load_knowledge: TOOL_CATEGORIES.UI,
 
   // Schema tools - MCP
   list_tables: TOOL_CATEGORIES.SCHEMA,

apps/studio/lib/ai/tools/deploy-edge-function.server.ts

@@ -0,0 +1,20 @@
+import { contract } from './deploy-edge-function'
+import { toolFromContract } from './tool-from-contract'
+import type { ToolFactory } from './types'
+import { deployEdgeFunction } from '@/data/edge-functions/edge-functions-deploy-mutation'
+
+export const createTool: ToolFactory<typeof contract> = (ctx) =>
+  toolFromContract(contract, ctx, async ({ name, code }) => {
+    await deployEdgeFunction({
+      projectRef: ctx.projectRef ?? '',
+      slug: name,
+      metadata: {
+        entrypoint_path: 'index.ts',
+        name,
+        verify_jwt: true,
+      },
+      files: [{ name: 'index.ts', content: code }],
+      authorization: ctx.authorization,
+    })
+    return { success: true }
+  })

apps/studio/lib/ai/tools/deploy-edge-function.ts

@@ -0,0 +1,19 @@
+import { z } from 'zod'
+
+import type { ToolContract } from './types'
+
+const inputSchema = z.object({
+  name: z.string().describe('The URL-friendly name/slug of the Edge Function.'),
+  code: z.string().describe('The TypeScript code for the Edge Function.'),
+})
+
+export const contract = {
+  name: 'deploy_edge_function',
+  description:
+    'Asks the user to deploy a Supabase Edge Function from provided code. Requires user approval before deploying.',
+  inputSchema,
+  outputSchema: z.object({ success: z.boolean() }),
+  minOptInLevel: null,
+  needsApproval: true,
+  toModelOutput: ({ success }) => ({ success }),
+} as const satisfies ToolContract

apps/studio/lib/ai/tools/execute-sql-output.ts

@@ -0,0 +1,59 @@
+import type { JSONValue } from 'ai'
+import { z } from 'zod'
+
+import { lowerOptInLevel, optInLevelSchema } from '../opt-in-level'
+import { DEFAULT_CELL_ROW_LIMIT } from '@/components/interfaces/Explorer/QueryCell/QueryCell.utils'
+import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
+
+export const NO_DATA_PERMISSIONS =
+  'The query was executed and the user has viewed the results but decided not to share in the conversation due to permission levels. Continue with your plan unless instructed to interpret the result. If you need the rows to answer and `update_opt_in_level` is available, call it with `schema_and_log_and_data`, unless the user already skipped that request in this chat.'
+
+// the rows.length default covers chats saved before execute_sql sliced its rows
+export const toModelRows = <T>(
+  rows: T[],
+  {
+    totalRows = rows.length,
+    hasMoreRows = false,
+  }: { totalRows?: number; hasMoreRows?: boolean } = {}
+) => {
+  if (!hasMoreRows && totalRows <= DEFAULT_CELL_ROW_LIMIT) return rows
+
+  return {
+    rows: rows.slice(0, DEFAULT_CELL_ROW_LIMIT),
+    note: hasMoreRows
+      ? `The query matches more than ${DEFAULT_CELL_ROW_LIMIT} rows. Only the first ${DEFAULT_CELL_ROW_LIMIT} are included because results are capped to protect database performance, so the exact total is unknown. Tell the user it matched more than ${DEFAULT_CELL_ROW_LIMIT} rows and only the first ${DEFAULT_CELL_ROW_LIMIT} are shown. Do not state a specific total.`
+      : `Only the first ${DEFAULT_CELL_ROW_LIMIT} of ${totalRows} rows are included.`,
+  }
+}
+
+/**
+ * `optInLevel` is the level the query ran under. Rows are shared only at the lower of that
+ * and the current level, so rows from `schema` stay hidden after an upgrade to data.
+ */
+const stampedOutputSchema = z.object({
+  rows: z.array(z.unknown()),
+  totalRows: z.number().optional(),
+  hasMoreRows: z.boolean().optional(),
+  optInLevel: optInLevelSchema,
+})
+
+/**
+ * Output saved before the stamp existed is a bare array. It's stamped `schema` so its rows
+ * stay hidden from the model, and the UI can still show them.
+ */
+export const executeSqlOutputSchema = z.union([
+  stampedOutputSchema,
+  z.array(z.unknown()).transform((rows) => ({ rows, optInLevel: 'schema' as const })),
+])
+
+export type ExecuteSqlOutput = z.infer<typeof stampedOutputSchema>
+
+export function toModelExecuteSqlOutput(
+  output: ExecuteSqlOutput,
+  optInLevel: AiOptInLevel
+): JSONValue {
+  const level = lowerOptInLevel(output.optInLevel, optInLevel)
+  return level === 'schema_and_log_and_data'
+    ? (toModelRows(output.rows, output) as JSONValue)
+    : NO_DATA_PERMISSIONS
+}

apps/studio/lib/ai/tools/execute-sql.server.ts

@@ -0,0 +1,38 @@
+import { untrustedSql } from '@supabase/pg-meta'
+
+import { getAuthHeaders } from './auth-headers'
+import { capReadQuery } from './cap-read-query'
+import { contract } from './execute-sql'
+import { approvalGatedToolFromContract } from './tool-from-contract'
+import type { ApprovalGatedToolFactory } from './types'
+import { DEFAULT_CELL_ROW_LIMIT } from '@/components/interfaces/Explorer/QueryCell/QueryCell.utils'
+import { executeSql } from '@/data/sql/execute-sql-mutation'
+
+export const createTool: ApprovalGatedToolFactory<typeof contract> = (ctx) => {
+  const { projectRef, connectionString, aiOptInLevel = 'schema' } = ctx
+  const authHeaders = getAuthHeaders(ctx.authorization, { contentType: true })
+
+  return approvalGatedToolFromContract(
+    contract,
+    ctx,
+    async ({ sql, isWriteQuery }, { abortSignal }, gate) => {
+      const approvedSql = gate.acceptUntrustedSql(untrustedSql(sql))
+      const { sql: limitedSql, isCappedInDatabase } = isWriteQuery
+        ? { sql: approvedSql, isCappedInDatabase: false }
+        : capReadQuery(approvedSql, DEFAULT_CELL_ROW_LIMIT + 1)
+      const { result } = await executeSql(
+        { projectRef, connectionString, sql: limitedSql },
+        abortSignal,
+        authHeaders
+      )
+      const rows = Array.isArray(result) ? result : []
+      const isCapped = rows.length > DEFAULT_CELL_ROW_LIMIT
+      const hasMoreRows = isCappedInDatabase && isCapped
+      return {
+        rows: rows.slice(0, DEFAULT_CELL_ROW_LIMIT),
+        ...(hasMoreRows ? { hasMoreRows } : { totalRows: rows.length }),
+        optInLevel: aiOptInLevel,
+      }
+    }
+  )
+}

apps/studio/lib/ai/tools/execute-sql.ts

@@ -0,0 +1,36 @@
+import { z } from 'zod'
+
+import { executeSqlOutputSchema, toModelExecuteSqlOutput } from './execute-sql-output'
+import type { ToolContract } from './types'
+import { fixSqlBackslashEscapes } from '@/lib/ai/util'
+
+const inputSchema = z.object({
+  // Transform at parse time so the corrected SQL is what gets stored in
+  // toolCall.input — ensuring evals and logs reflect what actually runs.
+  sql: z.string().describe('The SQL statement to execute.').transform(fixSqlBackslashEscapes),
+  label: z.string().describe('A short 2-4 word label for the SQL statement.'),
+  chartConfig: z
+    .object({
+      view: z.enum(['table', 'chart']).describe('How to render the results after execution'),
+      xAxis: z.string().optional().describe('The column to use for the x-axis of the chart.'),
+      yAxis: z.string().optional().describe('The column to use for the y-axis of the chart.'),
+    })
+    .describe('Chart configuration for rendering the results'),
+  isWriteQuery: z
+    .boolean()
+    .default(false)
+    .describe(
+      'Whether the SQL statement performs a write operation or has side effects. Set true for INSERT/UPDATE/DELETE/DDL and for SELECT statements that call side-effecting functions, such as select cron.schedule(...), cron.unschedule(...), or functions that create, modify, schedule, enqueue, notify, or trigger work.'
+    ),
+})
+
+export const contract = {
+  name: 'execute_sql',
+  description:
+    'Asks the user to execute a SQL statement and return the results. Requires user approval before executing.',
+  inputSchema,
+  outputSchema: executeSqlOutputSchema,
+  minOptInLevel: null,
+  needsApproval: true,
+  toModelOutput: toModelExecuteSqlOutput,
+} as const satisfies ToolContract

apps/studio/lib/ai/tools/knowledge-topics.ts

@@ -0,0 +1,11 @@
+/** Topics `load_knowledge` can load, with how the UI refers to each. Prompt text lives server-side. */
+export const KNOWLEDGE_TOPIC_LABELS = {
+  pg_best_practices: 'Postgres best practices',
+  rls: 'RLS',
+  storage: 'Storage',
+  edge_functions: 'Edge Functions',
+  realtime: 'Realtime',
+  logs: 'logs',
+} as const
+
+export type KnowledgeName = keyof typeof KNOWLEDGE_TOPIC_LABELS

apps/studio/lib/ai/tools/load-knowledge.server.ts

@@ -0,0 +1,24 @@
+import type { KnowledgeName } from './knowledge-topics'
+import { contract } from './load-knowledge'
+import { toolFromContract } from './tool-from-contract'
+import type { ToolFactory } from './types'
+import {
+  EDGE_FUNCTION_PROMPT,
+  LOGS_PROMPT,
+  PG_BEST_PRACTICES,
+  REALTIME_PROMPT,
+  RLS_PROMPT,
+  STORAGE_PROMPT,
+} from '@/lib/ai/prompts'
+
+const KNOWLEDGE = {
+  pg_best_practices: PG_BEST_PRACTICES,
+  rls: RLS_PROMPT,
+  storage: STORAGE_PROMPT,
+  edge_functions: EDGE_FUNCTION_PROMPT,
+  realtime: REALTIME_PROMPT,
+  logs: LOGS_PROMPT,
+} satisfies Record<KnowledgeName, string>
+
+export const createTool: ToolFactory<typeof contract> = (ctx) =>
+  toolFromContract(contract, ctx, ({ name }) => KNOWLEDGE[name])

apps/studio/lib/ai/tools/load-knowledge.ts

@@ -0,0 +1,21 @@
+import { z } from 'zod'
+
+import { KNOWLEDGE_TOPIC_LABELS, type KnowledgeName } from './knowledge-topics'
+import type { ToolContract } from './types'
+
+const inputSchema = z.object({
+  name: z
+    .enum(Object.keys(KNOWLEDGE_TOPIC_LABELS) as [KnowledgeName, ...KnowledgeName[]])
+    .describe('The knowledge to load'),
+})
+
+export const contract = {
+  name: 'load_knowledge',
+  description:
+    'Load detailed knowledge about a Supabase topic before answering questions about it.',
+  inputSchema,
+  outputSchema: z.string(),
+  minOptInLevel: null,
+  needsApproval: false,
+  toModelOutput: (output: string) => output,
+} as const satisfies ToolContract

apps/studio/lib/ai/tools/registry.test.ts

@@ -43,7 +43,7 @@ describe('sanitizeContractPart', () => {
 
   it('returns undefined for tools without a contract', () => {
     expect(
-      sanitizeContractPart({ ...part, type: 'tool-execute_sql' } as ToolUIPart, 'schema')
+      sanitizeContractPart({ ...part, type: 'tool-list_tables' } as ToolUIPart, 'schema')
     ).toBe(undefined)
   })
 })

apps/studio/lib/ai/tools/registry.ts

@@ -3,6 +3,10 @@ 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 deployEdgeFunction } from './deploy-edge-function'
+import { contract as executeSql } from './execute-sql'
+import { contract as loadKnowledge } from './load-knowledge'
+import { contract as renameChat } from './rename-chat'
 import { contract as runNotebook } from './run-notebook'
 import type { ToolContract, ToolContractInput, ToolContractOutput } from './types'
 import { contract as updateNotebook } from './update-notebook'
@@ -15,6 +19,10 @@ import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
 export const TOOL_CONTRACTS = {
   create_notebook: createNotebook,
   delete_notebook: deleteNotebook,
+  deploy_edge_function: deployEdgeFunction,
+  execute_sql: executeSql,
+  load_knowledge: loadKnowledge,
+  rename_chat: renameChat,
   run_notebook: runNotebook,
   update_notebook: updateNotebook,
 } as const satisfies Partial<Record<ToolName, ToolContract>>

apps/studio/lib/ai/tools/rename-chat.server.ts

@@ -0,0 +1,7 @@
+import { contract } from './rename-chat'
+import { toolFromContract } from './tool-from-contract'
+import type { ToolFactory } from './types'
+
+// The client applies the rename in `onToolCall`; the server only acknowledges it.
+export const createTool: ToolFactory<typeof contract> = (ctx) =>
+  toolFromContract(contract, ctx, async () => ({ status: 'Chat request sent to client' }))

apps/studio/lib/ai/tools/rename-chat.ts

@@ -0,0 +1,17 @@
+import { z } from 'zod'
+
+import type { ToolContract } from './types'
+
+const inputSchema = z.object({
+  newName: z.string().describe('The new name for the chat session. Five words or less.'),
+})
+
+export const contract = {
+  name: 'rename_chat',
+  description: `Rename the current chat session when the current chat name doesn't describe the conversation topic.`,
+  inputSchema,
+  outputSchema: z.object({ status: z.string() }),
+  minOptInLevel: null,
+  needsApproval: false,
+  toModelOutput: ({ status }) => ({ status }),
+} as const satisfies ToolContract

apps/studio/lib/ai/tools/studio-tools.test.ts

@@ -3,7 +3,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 import { createUpdateOptInLevelInputSchema, getOptInTools, getStudioTools } from './studio-tools'
 import { executeSql } from '@/data/sql/execute-sql-mutation'
-import { NO_DATA_PERMISSIONS } from '@/lib/ai/tools/tool-sanitizer'
+import { NO_DATA_PERMISSIONS } from '@/lib/ai/tools/execute-sql-output'
 
 vi.mock('@/data/sql/execute-sql-mutation', () => ({
   executeSql: vi.fn(),
@@ -196,7 +196,7 @@ describe('ai/tools/studio-tools', () => {
       )
       expect(result).toEqual({ rows, totalRows: 1, optInLevel: 'schema' })
       expect((tools.execute_sql as any).toModelOutput({ output: result })).toEqual({
-        type: 'text',
+        type: 'json',
         value: NO_DATA_PERMISSIONS,
       })
     })

apps/studio/lib/ai/tools/studio-tools.ts

@@ -1,150 +1,21 @@
-import { acceptUntrustedSql, untrustedSql } from '@supabase/pg-meta'
 import { tool } from 'ai'
-import { z } from 'zod'
 
-import { getAuthHeaders } from './auth-headers'
-import { capReadQuery } from './cap-read-query'
+import { createTool as createDeployEdgeFunctionTool } from './deploy-edge-function.server'
+import { createTool as createExecuteSqlTool } from './execute-sql.server'
+import { createTool as createLoadKnowledgeTool } from './load-knowledge.server'
+import { createTool as createRenameChatTool } from './rename-chat.server'
 import type { ToolContext } from './types'
-import { DEFAULT_CELL_ROW_LIMIT } from '@/components/interfaces/Explorer/QueryCell/QueryCell.utils'
-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,
-  PG_BEST_PRACTICES,
-  REALTIME_PROMPT,
-  RLS_PROMPT,
-  STORAGE_PROMPT,
-} from '@/lib/ai/prompts'
 import { updateOptInLevelInputSchema } from '@/lib/ai/tool-filter'
-import { NO_DATA_PERMISSIONS, toModelRows } from '@/lib/ai/tools/tool-sanitizer'
-import { fixSqlBackslashEscapes } from '@/lib/ai/util'
 
-const KNOWLEDGE = {
-  pg_best_practices: PG_BEST_PRACTICES,
-  rls: RLS_PROMPT,
-  storage: STORAGE_PROMPT,
-  edge_functions: EDGE_FUNCTION_PROMPT,
-  realtime: REALTIME_PROMPT,
-  logs: LOGS_PROMPT,
-} as const
-
-type KnowledgeName = keyof typeof KNOWLEDGE
-
-export const executeSqlInputSchema = z.object({
-  // Transform at parse time so the corrected SQL is what gets stored in
-  // toolCall.input — ensuring evals and logs reflect what actually runs.
-  sql: z.string().describe('The SQL statement to execute.').transform(fixSqlBackslashEscapes),
-  label: z.string().describe('A short 2-4 word label for the SQL statement.'),
-  chartConfig: z
-    .object({
-      view: z.enum(['table', 'chart']).describe('How to render the results after execution'),
-      xAxis: z.string().optional().describe('The column to use for the x-axis of the chart.'),
-      yAxis: z.string().optional().describe('The column to use for the y-axis of the chart.'),
-    })
-    .describe('Chart configuration for rendering the results'),
-  isWriteQuery: z
-    .boolean()
-    .default(false)
-    .describe(
-      'Whether the SQL statement performs a write operation or has side effects. Set true for INSERT/UPDATE/DELETE/DDL and for SELECT statements that call side-effecting functions, such as select cron.schedule(...), cron.unschedule(...), or functions that create, modify, schedule, enqueue, notify, or trigger work.'
-    ),
+export const getStudioTools = (ctx: ToolContext = {}) => ({
+  execute_sql: createExecuteSqlTool(ctx),
+  deploy_edge_function: createDeployEdgeFunctionTool(ctx),
+  rename_chat: createRenameChatTool(ctx),
+  load_knowledge: createLoadKnowledgeTool(ctx),
 })
 
-export const loadKnowledgeInputSchema = z.object({
-  name: z
-    .enum(Object.keys(KNOWLEDGE) as [KnowledgeName, ...KnowledgeName[]])
-    .describe('The knowledge to load'),
-})
-
-export const getStudioTools = (ctx: ToolContext = {}) => {
-  const { projectRef, connectionString, authorization, aiOptInLevel = 'schema' } = ctx
-  const authHeaders = getAuthHeaders(authorization, { contentType: true })
-
-  return {
-    // Tools that wait on the user use `needsApproval` with a server-side `execute`, not a
-    // client tool resolved through `addToolResult`. Client results never land in a Braintrust
-    // tool span: https://github.com/supabase/supabase/pull/45654
-    execute_sql: tool({
-      description:
-        'Asks the user to execute a SQL statement and return the results. Requires user approval before executing.',
-      inputSchema: executeSqlInputSchema,
-      needsApproval: true,
-      execute: async ({ sql, isWriteQuery }, { abortSignal }) => {
-        // The `needsApproval: true` gate on this tool means the user has
-        // explicitly approved this AI-generated SQL before execute runs —
-        // that approval is the user gesture that promotes untrusted to safe.
-        const approvedSql = acceptUntrustedSql(untrustedSql(sql))
-        // Push the cap into the query so the database doesn't load rows we'd throw away. The
-        // extra row lets `totalRows > DEFAULT_CELL_ROW_LIMIT` still signal a truncated result.
-        const { sql: limitedSql, isCappedInDatabase } = isWriteQuery
-          ? { sql: approvedSql, isCappedInDatabase: false }
-          : capReadQuery(approvedSql, DEFAULT_CELL_ROW_LIMIT + 1)
-        const { result } = await executeSql(
-          { projectRef, connectionString, sql: limitedSql },
-          abortSignal,
-          authHeaders
-        )
-        const rows = Array.isArray(result) ? result : []
-        // same row cap as the query cell, unbounded results freeze the chat and overflow the model
-        const isCapped = rows.length > DEFAULT_CELL_ROW_LIMIT
-        // When the limit ran in the database, `rows.length` stops at the cap and isn't a real total
-        const hasMoreRows = isCappedInDatabase && isCapped
-        return {
-          rows: rows.slice(0, DEFAULT_CELL_ROW_LIMIT),
-          ...(hasMoreRows ? { hasMoreRows } : { totalRows: rows.length }),
-          optInLevel: aiOptInLevel,
-        }
-      },
-      toModelOutput: ({ output }) => {
-        return aiOptInLevel === 'schema_and_log_and_data'
-          ? { type: 'json', value: toModelRows(output.rows, output) }
-          : { type: 'text', value: NO_DATA_PERMISSIONS }
-      },
-    }),
-    deploy_edge_function: tool({
-      description:
-        'Asks the user to deploy a Supabase Edge Function from provided code. Requires user approval before deploying.',
-      inputSchema: z.object({
-        name: z.string().describe('The URL-friendly name/slug of the Edge Function.'),
-        code: z.string().describe('The TypeScript code for the Edge Function.'),
-      }),
-      needsApproval: true,
-      execute: async ({ name, code }) => {
-        await deployEdgeFunction({
-          projectRef: projectRef ?? '',
-          slug: name,
-          metadata: {
-            entrypoint_path: 'index.ts',
-            name,
-            verify_jwt: true,
-          },
-          files: [{ name: 'index.ts', content: code }],
-          authorization,
-        })
-        return { success: true }
-      },
-    }),
-    rename_chat: tool({
-      description: `Rename the current chat session when the current chat name doesn't describe the conversation topic.`,
-      inputSchema: z.object({
-        newName: z.string().describe('The new name for the chat session. Five words or less.'),
-      }),
-      execute: async () => {
-        return { status: 'Chat request sent to client' }
-      },
-    }),
-    load_knowledge: tool({
-      description:
-        'Load detailed knowledge about a Supabase topic before answering questions about it.',
-      inputSchema: loadKnowledgeInputSchema,
-      execute: ({ name }) => KNOWLEDGE[name],
-    }),
-  }
-}
-
 // Rewritten at parse time so the stored input is what the card shows. A level the org already
 // has is a review request, and `levelWhenAsked` outlives the user changing the setting.
 export const createUpdateOptInLevelInputSchema = (aiOptInLevel: AiOptInLevel) =>

apps/studio/lib/ai/tools/tool-sanitizer.test.ts

@@ -10,8 +10,9 @@ import {
   createAssistantMessageWithUpdateNotebookTool,
   createLongConversation,
 } from '../test-fixtures'
+import { NO_DATA_PERMISSIONS } from './execute-sql-output'
 import { INVALID_NOTEBOOK_RUN_OUTPUT_MESSAGE, type NotebookRunOutput } from './notebook-run-output'
-import { NO_DATA_PERMISSIONS, sanitizeMessagePart } from './tool-sanitizer'
+import { sanitizeMessagePart } from './tool-sanitizer'
 
 const notebookRunPart = {
   type: 'tool-run_notebook',

apps/studio/lib/ai/tools/tool-sanitizer.ts

@@ -1,90 +1,14 @@
 import type { ToolUIPart, UIMessage } from 'ai'
-import { z } from 'zod'
 
-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'
 
-interface ToolSanitizer {
-  toolName: ToolName
-  sanitize: <Tool extends ToolUIPart>(tool: Tool, optInLevel: AiOptInLevel) => Tool
-}
-
-export const NO_DATA_PERMISSIONS =
-  'The query was executed and the user has viewed the results but decided not to share in the conversation due to permission levels. Continue with your plan unless instructed to interpret the result. If you need the rows to answer and `update_opt_in_level` is available, call it with `schema_and_log_and_data`, unless the user already skipped that request in this chat.'
-
-// the rows.length default covers chats saved before execute_sql sliced its rows
-export const toModelRows = <T>(
-  rows: T[],
-  {
-    totalRows = rows.length,
-    hasMoreRows = false,
-  }: { totalRows?: number; hasMoreRows?: boolean } = {}
-) => {
-  if (!hasMoreRows && totalRows <= DEFAULT_CELL_ROW_LIMIT) return rows
-
-  return {
-    rows: rows.slice(0, DEFAULT_CELL_ROW_LIMIT),
-    note: hasMoreRows
-      ? `The query matches more than ${DEFAULT_CELL_ROW_LIMIT} rows. Only the first ${DEFAULT_CELL_ROW_LIMIT} are included because results are capped to protect database performance, so the exact total is unknown. Tell the user it matched more than ${DEFAULT_CELL_ROW_LIMIT} rows and only the first ${DEFAULT_CELL_ROW_LIMIT} are shown. Do not state a specific total.`
-      : `Only the first ${DEFAULT_CELL_ROW_LIMIT} of ${totalRows} rows are included.`,
-  }
-}
-
-/**
- * `optInLevel` is the level the query ran under. Rows are shared only at the lower of that
- * and the current level, so rows from `schema` stay hidden after an upgrade to data.
- * Output saved before the stamp existed is a bare array and stays hidden.
- */
-export const executeSqlOutputSchema = z.object({
-  rows: z.array(z.unknown()),
-  totalRows: z.number().optional(),
-  hasMoreRows: z.boolean().optional(),
-  optInLevel: optInLevelSchema,
-})
-
-const executeSqlSanitizer: ToolSanitizer = {
-  toolName: 'execute_sql',
-  sanitize: (tool, optInLevel) => {
-    const output = tool.output
-    const parsed = executeSqlOutputSchema.safeParse(output)
-
-    if (parsed.success) {
-      const level = lowerOptInLevel(parsed.data.optInLevel, optInLevel)
-      return {
-        ...tool,
-        output:
-          level === 'schema_and_log_and_data'
-            ? toModelRows(parsed.data.rows, parsed.data)
-            : NO_DATA_PERMISSIONS,
-      }
-    }
-
-    if (Array.isArray(output)) return { ...tool, output: NO_DATA_PERMISSIONS }
-
-    return { ...tool, output: optInLevel === 'schema_and_log_and_data' ? output : undefined }
-  },
-}
-
-export const ALL_TOOL_SANITIZERS = {
-  [executeSqlSanitizer.toolName]: executeSqlSanitizer,
-}
-
 export function sanitizeMessagePart(
   part: UIMessage['parts'][number],
   optInLevel: AiOptInLevel
 ): UIMessage['parts'][number] {
   if (part.type.startsWith('tool-')) {
-    const toolPart = part as ToolUIPart
-    const contractPart = sanitizeContractPart(toolPart, optInLevel)
-    if (contractPart) return contractPart
-    const toolName = toolPart.type.slice('tool-'.length)
-    const sanitizer = ALL_TOOL_SANITIZERS[toolName]
-    if (sanitizer) {
-      return sanitizer.sanitize(toolPart, optInLevel)
-    }
+    return sanitizeContractPart(part as ToolUIPart, optInLevel) ?? part
   }
 
   return part

apps/studio/state/ai-assistant-state.tsx

@@ -26,6 +26,8 @@ import {
   applyNotebookCacheEffects,
   collectNotebookCacheEffects,
 } from '@/lib/ai/notebook-cache-invalidation'
+import type { contract as renameChatContract } from '@/lib/ai/tools/rename-chat'
+import type { ToolContractInput } from '@/lib/ai/tools/types'
 import { BASE_PATH, IS_PLATFORM } from '@/lib/constants'
 
 type SuggestionsType = {
@@ -353,7 +355,7 @@ function createChatInstance(
       }
 
       if (toolCall.toolName === 'rename_chat') {
-        const { newName } = toolCall.input as { newName: string }
+        const { newName } = toolCall.input as ToolContractInput<typeof renameChatContract>
 
         if (options.id && newName?.trim()) {
           state.renameChat(options.id, newName.trim())