diff --git a/apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.utils.ts b/apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.utils.ts index 53ae319964b..c78e1515e13 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.utils.ts +++ b/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 = { - 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) { diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.utils.ts b/apps/studio/components/ui/AIAssistantPanel/Message.utils.ts index 932864198f7..5a19856f84f 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.utils.ts +++ b/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', diff --git a/apps/studio/evals/scorer-wasm.ts b/apps/studio/evals/scorer-wasm.ts index 447c866adbb..0fd62a27fe9 100644 --- a/apps/studio/evals/scorer-wasm.ts +++ b/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 { const executeSqlSpans = await getParsedToolSpans(trace, 'execute_sql', { - inputSchema: executeSqlInputSchema, + inputSchema: executeSqlContract.inputSchema, }) const createNotebookSpans = await getParsedToolSpans(trace, 'create_notebook', { inputSchema: createNotebookContract.inputSchema, diff --git a/apps/studio/evals/scorer.ts b/apps/studio/evals/scorer.ts index f13c2d06362..15ea0b11ee2 100644 --- a/apps/studio/evals/scorer.ts +++ b/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) diff --git a/apps/studio/lib/ai/generate-assistant-response.ts b/apps/studio/lib/ai/generate-assistant-response.ts index b9ac4eadd9c..e5e97bcb6b7 100644 --- a/apps/studio/lib/ai/generate-assistant-response.ts +++ b/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 metadata.chatName = newName } } diff --git a/apps/studio/lib/ai/tool-filter.test.ts b/apps/studio/lib/ai/tool-filter.test.ts index 9792c9dadbd..244c0858328 100644 --- a/apps/studio/lib/ai/tool-filter.test.ts +++ b/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) }) diff --git a/apps/studio/lib/ai/tool-filter.ts b/apps/studio/lib/ai/tool-filter.ts index 77d2b289f57..e49c107b720 100644 --- a/apps/studio/lib/ai/tool-filter.ts +++ b/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, 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, diff --git a/apps/studio/lib/ai/tools/deploy-edge-function.server.ts b/apps/studio/lib/ai/tools/deploy-edge-function.server.ts new file mode 100644 index 00000000000..b89a4aecdfb --- /dev/null +++ b/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 = (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 } + }) diff --git a/apps/studio/lib/ai/tools/deploy-edge-function.ts b/apps/studio/lib/ai/tools/deploy-edge-function.ts new file mode 100644 index 00000000000..bf92f8071c7 --- /dev/null +++ b/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 diff --git a/apps/studio/lib/ai/tools/execute-sql-output.ts b/apps/studio/lib/ai/tools/execute-sql-output.ts new file mode 100644 index 00000000000..0c582e19bd9 --- /dev/null +++ b/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 = ( + 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 + +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 +} diff --git a/apps/studio/lib/ai/tools/execute-sql.server.ts b/apps/studio/lib/ai/tools/execute-sql.server.ts new file mode 100644 index 00000000000..0b26954341a --- /dev/null +++ b/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 = (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, + } + } + ) +} diff --git a/apps/studio/lib/ai/tools/execute-sql.ts b/apps/studio/lib/ai/tools/execute-sql.ts new file mode 100644 index 00000000000..07e17cc1adf --- /dev/null +++ b/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 diff --git a/apps/studio/lib/ai/tools/knowledge-topics.ts b/apps/studio/lib/ai/tools/knowledge-topics.ts new file mode 100644 index 00000000000..8173dc61e4a --- /dev/null +++ b/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 diff --git a/apps/studio/lib/ai/tools/load-knowledge.server.ts b/apps/studio/lib/ai/tools/load-knowledge.server.ts new file mode 100644 index 00000000000..33993b89ff0 --- /dev/null +++ b/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 + +export const createTool: ToolFactory = (ctx) => + toolFromContract(contract, ctx, ({ name }) => KNOWLEDGE[name]) diff --git a/apps/studio/lib/ai/tools/load-knowledge.ts b/apps/studio/lib/ai/tools/load-knowledge.ts new file mode 100644 index 00000000000..ce6e2938f4c --- /dev/null +++ b/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 diff --git a/apps/studio/lib/ai/tools/registry.test.ts b/apps/studio/lib/ai/tools/registry.test.ts index f69ae096bb8..c823ee22cd8 100644 --- a/apps/studio/lib/ai/tools/registry.test.ts +++ b/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) }) }) diff --git a/apps/studio/lib/ai/tools/registry.ts b/apps/studio/lib/ai/tools/registry.ts index 3a65a417013..25c2ccd964c 100644 --- a/apps/studio/lib/ai/tools/registry.ts +++ b/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> diff --git a/apps/studio/lib/ai/tools/rename-chat.server.ts b/apps/studio/lib/ai/tools/rename-chat.server.ts new file mode 100644 index 00000000000..449181db451 --- /dev/null +++ b/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 = (ctx) => + toolFromContract(contract, ctx, async () => ({ status: 'Chat request sent to client' })) diff --git a/apps/studio/lib/ai/tools/rename-chat.ts b/apps/studio/lib/ai/tools/rename-chat.ts new file mode 100644 index 00000000000..31a7d3047c5 --- /dev/null +++ b/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 diff --git a/apps/studio/lib/ai/tools/studio-tools.test.ts b/apps/studio/lib/ai/tools/studio-tools.test.ts index c5e1a65f307..c035b54ec88 100644 --- a/apps/studio/lib/ai/tools/studio-tools.test.ts +++ b/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, }) }) diff --git a/apps/studio/lib/ai/tools/studio-tools.ts b/apps/studio/lib/ai/tools/studio-tools.ts index 28b93d772bb..41ee0e028dc 100644 --- a/apps/studio/lib/ai/tools/studio-tools.ts +++ b/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) => diff --git a/apps/studio/lib/ai/tools/tool-sanitizer.test.ts b/apps/studio/lib/ai/tools/tool-sanitizer.test.ts index a9c2e84eccd..396f71b92bf 100644 --- a/apps/studio/lib/ai/tools/tool-sanitizer.test.ts +++ b/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', diff --git a/apps/studio/lib/ai/tools/tool-sanitizer.ts b/apps/studio/lib/ai/tools/tool-sanitizer.ts index deb8625eada..d925d78c747 100644 --- a/apps/studio/lib/ai/tools/tool-sanitizer.ts +++ b/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: 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 = ( - 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 diff --git a/apps/studio/state/ai-assistant-state.tsx b/apps/studio/state/ai-assistant-state.tsx index 2e418b93ad2..73445ed4e31 100644 --- a/apps/studio/state/ai-assistant-state.tsx +++ b/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 if (options.id && newName?.trim()) { state.renameChat(options.id, newName.trim())