## 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.
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.tsto 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:
- Implement it in the appropriate tool module and include it in
getTools. - Update the allowlist/validation in
tool-filter.ts. Declare minimum access in its contract or the existing category map, as applicable. - 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.
- Update
Message.Parts.tsxand its helpers if the tool needs a UI renderer. - 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.
Related AI features
The chat Assistant shares code with other Studio AI features, but those have separate endpoints:
- pages/api/ai/sql/ contains policy generation, SQL title, filter, and cron helpers alongside the chat endpoint.
- pages/api/ai/code/complete.ts handles code completions.
- pages/api/ai/onboarding/design.ts handles onboarding design.
- data/ai/ contains client queries and mutations for AI helpers and feedback. The chat transport itself lives in
ai-assistant-state.tsx. - routes/api/ai/ contains the corresponding TanStack API adapters.
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.