Files
supabase/apps/studio/lib/ai
Charis 9ae8432dc7 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.
2026-10-15 11:16:58 -04:00
..

Supabase Assistant code guide

Start here when working on the Assistant in Studio. This guide maps the implementation for contributors and coding agents. Follow Studio's agent instructions and load the skills that apply to your change.

Request flow

flowchart TD
  UI[AssistantChat] --> State[ai-assistant-state: chat transport]
  State --> API[POST /api/ai/sql/generate-v4]
  API --> Access[Resolve organization access and model]
  Access --> Tools[getTools]
  Tools --> Response[generateAssistantResponse]
  Response --> Stream[Stream messages and tool calls to the UI]
  Stream --> Parts[Message.Parts: render text and tool cards]
  Parts --> Approval[User approves a tool when required]
  Approval --> State

The sidebar provider mounts AIAssistant. AssistantChat manages the active conversation UI and supplies project and organization context to ai-assistant-state.

The state module owns the AI SDK chat instances, request transport, chat persistence, and approval continuation. Its transport sends messages to generate-v4, which validates and authenticates the request, resolves access and the model, builds tools, and streams the response. The TanStack route adapts this same handler. Keep both runtimes in mind; see the migration guide.

generateAssistantResponse combines prompts, context, prepared conversation history, and tools for the model. Message.Parts dispatches streamed message parts to their UI renderers. Approval responses travel back through the chat transport so the server can continue tool execution.

Rules to keep

  • Enforce access in code. Resolve organization/project access on the server and keep the restrictive fallback for an unverified pairing. Prompts and UI visibility are not authorization checks.
  • Keep execution approval separate from data sharing. Full data opt-in does not approve AI-generated SQL execution. Promote untrusted SQL only after explicit user approval.
  • Protect results on every model path. Apply output filtering to both live tool results and replayed history. Raising opt-in must not expose SQL results collected under a lower level.
  • Keep the tool allowlist. Register tools explicitly and declare their minimum opt-in level. Remote MCP tools must pass the filter; the remote server's read-only setting does not replace it.
  • Keep request context out of the shared system prompt. Put project, chat, clock, and logs context in assistant-context.ts to preserve provider caching.
  • Keep tool contracts safe to import in the browser. Shared contracts contain schemas and pure output transformations; database calls and other server execution belong in the server implementation.
  • Preserve both runtimes. Keep the TanStack adapter and Next.js handler aligned. The chat route currently delegates to the shared handler; do not remove that handler while the adapter uses it.

Where to make changes

Chat and rendering

Change Start here
Panel layout, chat selection, and header AIAssistant.tsx, AIAssistantHeader.tsx
Composer, attachments, and sending messages AssistantChat.tsx, AssistantChatForm.tsx
Chat state, persistence, transport, and selected model ai-assistant-state.tsx
Message rendering and tool-specific cards Message.Parts.tsx, Message.Parts.utils.ts
SQL results and charts AssistantQueryCell.tsx, AssistantQueryCell.utils.ts
Approval UI Confirm.tsx, Confirm.utils.ts

Model requests and context

Change Start here
Chat endpoint and request validation generate-v4.ts
Model streaming and response orchestration generate-assistant-response.ts
System instructions and topic knowledge prompts.ts; load_knowledge in studio-tools.ts
Project, chat, clock, and logs context assistant-context.ts, assistant-message-metadata.ts
History preparation and output sanitization generate-assistant-response.utils.ts, tool-sanitizer.ts
Model catalog, defaults, and provider configuration model.utils.ts, model.ts, bedrock.ts; entitlement and throttling decisions in generate-v4.ts
Tracing and feedback braintrust-logger.ts, feedback API, rate-message-mutation.ts

Keep per-request context in assistant-context.ts. The shared system prompt is deliberately static for provider caching.

Tools

Change Start here
Tool assembly and platform/self-hosted selection tools/index.ts
Allowlist, minimum opt-in levels, and blocked-tool responses tool-filter.ts
SQL execution, Edge Function deployment, chat renaming, knowledge, and opt-in requests studio-tools.ts
Remote MCP connection and tool selection supabase-mcp.ts, mcp-tools.ts
Schema and policy tools schema-tools.ts
Reports, notebooks, incidents, and support report-tools.ts, notebook-tools.ts, incident-tools.ts, support-tools.ts
Self-hosted tools fallback-tools.ts
Shared tool contracts types.ts, registry.ts, tool-from-contract.ts

Some tools use shared contracts; others declare their schemas and output handling directly. update-notebook.ts and update-notebook.server.ts show the contract/server split. Contracts define input/output schemas, minimum opt-in level, approval requirements, and what the model may see. The registry also applies contract output handling when replaying history.

For a new or changed tool:

  1. Implement it in the appropriate tool module and include it in getTools.
  2. Update the allowlist/validation in tool-filter.ts. Declare minimum access in its contract or the existing category map, as applicable.
  3. Decide whether execution needs user approval and what output may reach the model, including on history replay. Read the SQL safety skill before changing SQL execution.
  4. Update Message.Parts.tsx and its helpers if the tool needs a UI renderer.
  5. Update relevant tests, mock tools, and evaluation cases. See local development and validation.

Access and data sharing

These checks have different responsibilities:

Responsibility Implementation
Resolve the saved organization opt-in tags for the UI useOrgOptedIntoAi.ts
Resolve server-side access for an organization/project pair ai-details.ts
Check permission to change organization settings useCheckPermissions.ts, used by OptInRequest.tsx and AIOptInModal.tsx
Save organization opt-in settings useAIOptInForm.ts, DataPrivacyForm.tsx
Gate tool availability by opt-in level tool-filter.ts
Approve SQL execution and call the database execute_sql in studio-tools.ts, using execute-sql-mutation.ts
Control which SQL results reach the model toModelOutput in studio-tools.ts and history sanitization in tool-sanitizer.ts

The opt-in levels are disabled, schema, schema_and_log, and schema_and_log_and_data. execute_sql is available at every level and requires user approval; sharing its result rows with the model requires schema_and_log_and_data. History sanitization also considers the level stamped on the original SQL output. Permission to change opt-in settings, user approval to run SQL, and the database's authorization checks are separate from result sharing.

The model can call update_opt_in_level with requiredLevel to request higher access, or without requiredLevel when the user asks to review or change the setting. Every call requires approval through the opt-in request card, even when the organization already has the requested level. If the requested level is already satisfied, the input transform removes requiredLevel and treats the call as a review request; levelWhenAsked records the organization's level when the call was made.

On platform, the endpoint resolves organization access and getTools includes platform tools when an access token is available. It also registers the opt-in request tool. Self-hosted requests use the schema level and fallback tools instead of the platform tool set. Check getTools and generate-v4 when changing either environment.

The chat Assistant shares code with other Studio AI features, but those have separate endpoints:

Local development and validation

Use the Studio setup instructions for the application and backend configuration. From the repository root, pnpm dev:studio starts Studio on port 8082 by default. The chat endpoint currently selects the OpenAI provider, which requires a server-side OPENAI_API_KEY; see model.ts. A self-hosted session exercises the fallback tool set, so use a platform development environment when working on platform-only tools or organization settings.

Tests are colocated with the implementation under lib/ai/, components/ui/AIAssistantPanel/, and state/. For example, from apps/studio:

pnpm exec vitest run lib/ai/tool-filter.test.ts lib/ai/tools/studio-tools.test.ts

Use the Assistant evaluation guide for behavioral datasets, local evaluation commands, CI labels, and scoring. dataset.ts contains scenarios and mock-tools.ts supplies tool behavior for evaluations. The browser test is assistant.spec.ts; read the E2E setup guide before running it.

Update this map when moving entry points or changing the request/tool lifecycle.