From c4c213ce3d6db0471ae009e2e0258097b387ab07 Mon Sep 17 00:00:00 2001 From: Pedro Rodrigues <44656907+Rodriguespn@users.noreply.github.com> Date: Tue, 7 Jul 2026 19:38:21 +0100 Subject: [PATCH] feat(studio): switch dashboard assistant to remote MCP server (#47479) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## I have read the [CONTRIBUTING.md]() file. YES ## What kind of change does this PR introduce? Feature / refactor. ## What is the current behavior? The dashboard assistant runs `@supabase/mcp-server-supabase` in-process over an in-memory transport (`lib/ai/supabase-mcp.ts`). ## What is the new behavior? The assistant connects to the **remote MCP server** over HTTP (`@ai-sdk/mcp`), forwarding the dashboard session token as a bearer. URL comes from `NEXT_PUBLIC_MCP_URL` with a local-dev fallback; platform-only, and Nimbus works via the same env var. * **Tool model unchanged:** UI-controlled `execute_sql` (with `needsApproval`) and `deploy_edge_function` still come from Studio; the allowlist (`TOOL_CATEGORY_MAP`) remains the gate keeping the remote's write tools away from the assistant (`read_only` is defense-in-depth). * **Attribution:** sends `x-source-name: supabase-studio` (+ `x-source-version`) → logged as `source_name`/`client_name`. * **Connection lifecycle:** the HTTP client is closed via the request's `AbortSignal` (tools execute later during streaming); `signal` is required on `getTools`/`getMcpTools`. * **Resilience:** a remote-MCP failure degrades to the remaining tools instead of failing the assistant. * **Drift protection:** relied-upon tools are typed against `keyof typeof supabaseMcpToolSchemas`, so a package bump that renames/removes one fails `pnpm typecheck`; a runtime check also warns if the deployed server returns fewer tools. * Adds unit tests for the above. ## Additional context * Verified end-to-end against a local remote MCP server with a dashboard token: `initialize` 200, tools listed, a tool executed, client closed cleanly. * The remote MCP (mgmt-api) already accepts dashboard session tokens (GoTrue-JWT auth path) — no backend change needed. `NEXT_PUBLIC_MCP_URL` must point at each env's `/mcp`. * `@supabase/mcp-server-supabase` is kept — still used by the self-hosted `/api/mcp` routes. Closes [AI-137](https://linear.app/supabase/issue/AI-137/switch-dashboard-assistant-to-remote-mcp) ## Rollout * **Rollout:** merges with `USE_REMOTE_MCP` off (in-process); flip it to `true` per environment (staging → prod → Nimbus) once each one's prerequisites land. * **Rollback:** unset `USE_REMOTE_MCP` and redeploy to fall back to the in-process client — no revert needed. ## Summary by CodeRabbit * **Bug Fixes** * Improved AI request handling so tool loading and generation clean up properly when a request is cancelled or the browser connection closes. * Added safer fallback behavior when remote tool loading fails, so AI features can continue with available tools instead of stopping entirely. * Updated remote tool access to use the current project reference and preserve the correct access headers. ## Summary by CodeRabbit * **New Features** * AI tools now connect more reliably to remote services and stop cleanly when requests end or are canceled. * Tool loading is more resilient, continuing with available tools if remote access is unavailable. * **Bug Fixes** * Improved cleanup to prevent lingering connections during SQL generation and policy workflows. * Added safer handling for remote tool changes and invalid responses. * **Tests** * Expanded automated coverage for remote tool setup, cancellation, and fallback behavior. --------- Co-authored-by: Claude Opus 4.8 (1M context) --- .github/workflows/braintrust-evals.yml | 38 ++++- apps/studio/evals/assistant.eval.ts | 36 ++-- apps/studio/evals/preflight.ts | 49 ++++++ apps/studio/lib/ai/supabase-mcp.test.ts | 139 ++++++++++++++++ apps/studio/lib/ai/supabase-mcp.ts | 113 ++++++++++++- apps/studio/lib/ai/tools/index.test.ts | 86 ++++++++++ apps/studio/lib/ai/tools/index.ts | 26 ++- apps/studio/lib/ai/tools/mcp-tools.test.ts | 172 ++++++++++++++++++++ apps/studio/lib/ai/tools/mcp-tools.ts | 135 ++++++++++++--- apps/studio/lib/ai/tools/mock-tools.test.ts | 64 ++++++++ apps/studio/lib/ai/tools/mock-tools.ts | 21 ++- apps/studio/lib/api/generate-v4.test.ts | 4 + apps/studio/package.json | 1 + apps/studio/pages/api/ai/sql/generate-v4.ts | 4 + apps/studio/pages/api/ai/sql/policy.ts | 92 ++++++----- apps/studio/turbo.jsonc | 3 + 16 files changed, 888 insertions(+), 95 deletions(-) create mode 100644 apps/studio/evals/preflight.ts create mode 100644 apps/studio/lib/ai/supabase-mcp.test.ts create mode 100644 apps/studio/lib/ai/tools/index.test.ts create mode 100644 apps/studio/lib/ai/tools/mcp-tools.test.ts create mode 100644 apps/studio/lib/ai/tools/mock-tools.test.ts diff --git a/.github/workflows/braintrust-evals.yml b/.github/workflows/braintrust-evals.yml index a2f9075f913..d05093acd61 100644 --- a/.github/workflows/braintrust-evals.yml +++ b/.github/workflows/braintrust-evals.yml @@ -15,8 +15,42 @@ permissions: contents: read jobs: + # Cheap, always-on gate: verifies the eval harness can reach the MCP server + # (its one real tool, search_docs). Runs on every push/PR — no OpenAI, no full + # eval suite — so a broken MCP connection is caught early with an actionable + # message. The eval job depends on this, so evals never run against a broken + # MCP connection. + preflight: + name: Eval MCP preflight + runs-on: ubuntu-latest + timeout-minutes: 10 + + steps: + - name: Checkout + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Install pnpm + uses: pnpm/action-setup@41ff72655975bd51cab0327fa583b6e92b6d3061 # v4.2.0 + with: + run_install: false + + - name: Use Node.js + uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0 + with: + node-version-file: '.nvmrc' + cache: 'pnpm' + + - name: Install Dependencies + run: pnpm install --frozen-lockfile + + - name: Preflight — MCP connectivity + run: cd apps/studio && pnpm evals:preflight + eval: name: Run evals + needs: preflight if: github.event_name == 'push' || (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'run-evals') && github.event.pull_request.head.repo.full_name == github.repository) runs-on: ubuntu-latest timeout-minutes: 20 @@ -43,8 +77,8 @@ jobs: - name: Use Node.js uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0 with: - node-version-file: ".nvmrc" - cache: "pnpm" + node-version-file: '.nvmrc' + cache: 'pnpm' - name: Install Dependencies run: pnpm install --frozen-lockfile diff --git a/apps/studio/evals/assistant.eval.ts b/apps/studio/evals/assistant.eval.ts index 207594cb24a..b540b630720 100644 --- a/apps/studio/evals/assistant.eval.ts +++ b/apps/studio/evals/assistant.eval.ts @@ -31,20 +31,30 @@ Eval('Assistant', { const modelResponse = await getModel({ provider: 'openai', modelEntry }) if (modelResponse.error) throw modelResponse.error - const result = await generateAssistantResponse({ - ...modelResponse.modelParams, - messages: [ - { - id: '1', - role: 'user', - parts: [{ type: 'text', text: input.prompt }], - }, - ], - tools: await getMockTools(input.mockTables ? { list_tables: input.mockTables } : undefined), - }) + // Owns the lifecycle of the remote MCP client opened inside getMockTools: + // aborting once generation is done closes that connection. + const toolsAbortController = new AbortController() + try { + const result = await generateAssistantResponse({ + ...modelResponse.modelParams, + messages: [ + { + id: '1', + role: 'user', + parts: [{ type: 'text', text: input.prompt }], + }, + ], + tools: await getMockTools( + input.mockTables ? { list_tables: input.mockTables } : undefined, + toolsAbortController.signal + ), + }) - const finishReason = await result.finishReason - return { finishReason } + const finishReason = await result.finishReason + return { finishReason } + } finally { + toolsAbortController.abort() + } }, scores: [ toolUsageScorer, diff --git a/apps/studio/evals/preflight.ts b/apps/studio/evals/preflight.ts new file mode 100644 index 00000000000..a8889540e2c --- /dev/null +++ b/apps/studio/evals/preflight.ts @@ -0,0 +1,49 @@ +/** + * Eval preflight — MCP connectivity check. + * + * The assistant eval harness (`getMockTools`) mocks every tool except + * `search_docs`, which it sources from a real MCP server. If that connection is + * broken (endpoint down, bad/expired token, contract drift, missing package), + * evals fail deep inside a Braintrust run with an opaque per-case error. + * + * This preflight exercises the exact same path and fails fast with an + * actionable message, so a broken MCP connection is caught up front when the + * eval job runs (e.g. on push). Keep it in lockstep with how `getMockTools` + * obtains `search_docs` — if that switches to the remote client (see AI-897), + * switch this too. + */ +import { createInProcessSupabaseMCPClient } from '@/lib/ai/supabase-mcp' + +async function runPreflight() { + let client: Awaited> | undefined + + try { + client = await createInProcessSupabaseMCPClient({ + accessToken: 'mock-access-token', + projectRef: 'mock-project-ref', + }) + + const tools = await client.tools() + + if (!tools || !('search_docs' in tools)) { + throw new Error( + 'Connected to the MCP server but `search_docs` was not returned. ' + + 'The tool contract may have drifted, or the server is misconfigured.' + ) + } + + console.log('✅ Eval MCP preflight OK — connected and `search_docs` is available.') + } finally { + await client?.close().catch(() => {}) + } +} + +runPreflight().catch((error) => { + console.error( + '❌ Eval MCP preflight failed — the eval harness cannot reach the MCP server, ' + + 'so evals would fail. Check NEXT_PUBLIC_MCP_URL, the access token, and the ' + + '@supabase/mcp-server-supabase dependency.' + ) + console.error(error instanceof Error ? error.message : error) + process.exit(1) +}) diff --git a/apps/studio/lib/ai/supabase-mcp.test.ts b/apps/studio/lib/ai/supabase-mcp.test.ts new file mode 100644 index 00000000000..eae88b4186e --- /dev/null +++ b/apps/studio/lib/ai/supabase-mcp.test.ts @@ -0,0 +1,139 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { createSupabaseMCPClient } from './supabase-mcp' + +const createMCPClientMock = vi.fn() + +vi.mock('@ai-sdk/mcp', () => ({ + createMCPClient: (...args: any[]) => createMCPClientMock(...args), +})) + +const ACCESS_TOKEN = 'test-access-token' +const PROJECT_REF = 'abcdefghijklmnopqrst' + +function getTransportConfig() { + expect(createMCPClientMock).toHaveBeenCalledTimes(1) + return createMCPClientMock.mock.calls[0][0] +} + +describe('createSupabaseMCPClient', () => { + const originalMcpUrl = process.env.NEXT_PUBLIC_MCP_URL + const originalSha = process.env.VERCEL_GIT_COMMIT_SHA + + beforeEach(() => { + vi.clearAllMocks() + createMCPClientMock.mockResolvedValue({ tools: vi.fn() }) + delete process.env.VERCEL_GIT_COMMIT_SHA + }) + + afterEach(() => { + if (originalMcpUrl === undefined) delete process.env.NEXT_PUBLIC_MCP_URL + else process.env.NEXT_PUBLIC_MCP_URL = originalMcpUrl + if (originalSha === undefined) delete process.env.VERCEL_GIT_COMMIT_SHA + else process.env.VERCEL_GIT_COMMIT_SHA = originalSha + }) + + it('connects over the HTTP transport with the supabase-studio client name', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + const config = getTransportConfig() + expect(config.name).toBe('supabase-studio') + expect(config.transport.type).toBe('http') + }) + + it('targets the URL from NEXT_PUBLIC_MCP_URL with project_ref and read_only', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + const url = new URL(getTransportConfig().transport.url) + expect(url.origin + url.pathname).toBe('https://mcp.supabase.com/mcp') + expect(url.searchParams.get('project_ref')).toBe(PROJECT_REF) + expect(url.searchParams.get('read_only')).toBe('true') + }) + + it('forwards the dashboard access token as a bearer header', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + expect(getTransportConfig().transport.headers.Authorization).toBe(`Bearer ${ACCESS_TOKEN}`) + }) + + it('identifies assistant traffic with the x-source-name header', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + expect(getTransportConfig().transport.headers['x-source-name']).toBe('supabase-studio') + }) + + it('sends x-source-version from VERCEL_GIT_COMMIT_SHA when available', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + process.env.VERCEL_GIT_COMMIT_SHA = 'abc1234' + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + expect(getTransportConfig().transport.headers['x-source-version']).toBe('abc1234') + }) + + it('omits x-source-version when VERCEL_GIT_COMMIT_SHA is not set', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + delete process.env.VERCEL_GIT_COMMIT_SHA + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + expect(getTransportConfig().transport.headers).not.toHaveProperty('x-source-version') + }) + + it('never leaks the access token into the URL', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + expect(getTransportConfig().transport.url).not.toContain(ACCESS_TOKEN) + }) + + it('always requests read-only mode', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + const url = new URL(getTransportConfig().transport.url) + expect(url.searchParams.get('read_only')).toBe('true') + }) + + it('falls back to the default local URL when NEXT_PUBLIC_MCP_URL is unset', async () => { + delete process.env.NEXT_PUBLIC_MCP_URL + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + const url = new URL(getTransportConfig().transport.url) + expect(url.origin + url.pathname).toBe('http://localhost:8080/mcp') + expect(url.searchParams.get('project_ref')).toBe(PROJECT_REF) + }) + + it('falls back to the default local URL when NEXT_PUBLIC_MCP_URL is an empty string', async () => { + process.env.NEXT_PUBLIC_MCP_URL = '' + + await createSupabaseMCPClient({ accessToken: ACCESS_TOKEN, projectRef: PROJECT_REF }) + + const url = new URL(getTransportConfig().transport.url) + expect(url.origin + url.pathname).toBe('http://localhost:8080/mcp') + }) + + it('returns the client created by createMCPClient', async () => { + process.env.NEXT_PUBLIC_MCP_URL = 'https://mcp.supabase.com/mcp' + const fakeClient = { tools: vi.fn() } + createMCPClientMock.mockResolvedValueOnce(fakeClient) + + const client = await createSupabaseMCPClient({ + accessToken: ACCESS_TOKEN, + projectRef: PROJECT_REF, + }) + + expect(client).toBe(fakeClient) + }) +}) diff --git a/apps/studio/lib/ai/supabase-mcp.ts b/apps/studio/lib/ai/supabase-mcp.ts index a508ad94a94..bbc65500f7b 100644 --- a/apps/studio/lib/ai/supabase-mcp.ts +++ b/apps/studio/lib/ai/supabase-mcp.ts @@ -1,17 +1,113 @@ import { createMCPClient } from '@ai-sdk/mcp' -import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js' // .js required for esbuild ESM resolution -import { createSupabaseMcpServer } from '@supabase/mcp-server-supabase' -import { createSupabaseApiPlatform } from '@supabase/mcp-server-supabase/platform/api' -import { API_URL } from '@/lib/constants' +/** + * Default MCP server URL used when `NEXT_PUBLIC_MCP_URL` is not configured (local + * development). Mirrors `DEFAULT_MCP_URL_PLATFORM` in `ui-patterns/McpUrlBuilder` + * so the assistant resolves the same endpoint as the Connect sheet. It's + * duplicated here (rather than imported) to keep + * `ui-patterns/McpUrlBuilder/constants` — which pulls in `next/image` and image + * assets — out of this server-side bundle. + */ +const DEFAULT_MCP_URL = 'http://localhost:8080/mcp' +/** + * Identifies assistant traffic to the remote MCP server. Sent both as the MCP + * client name (logged as `client_name`) and via the `x-source-name` header + * (logged as `source_name`) by the mgmt-api McpLogger, so assistant requests are + * attributable in the MCP server's logs. + */ +const SOURCE_NAME = 'supabase-studio' + +/** + * Builds the remote MCP endpoint URL for the dashboard assistant. + * + * Points at the remote MCP server configured via `NEXT_PUBLIC_MCP_URL` (e.g. + * https://mcp.supabase.com/mcp), falling back to a local default for development. + * The query parameters (`project_ref`, `read_only`) mirror `getMcpUrl` in + * `ui-patterns/McpUrlBuilder/utils/getMcpUrl` so the assistant and the Connect sheet stay in + * sync. The assistant only performs read operations, so `read_only` is always + * set. + * + * Note: the assistant only talks to the remote MCP server on the hosted platform + * (see `getTools` / `getMcpTools`), so no self-hosted branch is needed here. + */ +function getRemoteMcpUrl(projectRef: string) { + // `||` (not `??`) so an empty-string env var falls back instead of producing + // an invalid `new URL('')`. + const url = new URL(process.env.NEXT_PUBLIC_MCP_URL || DEFAULT_MCP_URL) + if (projectRef) { + url.searchParams.set('project_ref', projectRef) + } + url.searchParams.set('read_only', 'true') + + return url.toString() +} + +/** + * Creates an MCP client connected to the remote Supabase MCP server over HTTP. + * + * Previously the assistant instantiated the MCP server in-process and connected + * to it via an in-memory transport. It now connects to the remote MCP server so + * the dashboard assistant shares the same MCP surface as external clients. + * + * The dashboard session `accessToken` is forwarded as a bearer token. The remote + * MCP server is responsible for validating it and scoping access to the project. + */ export async function createSupabaseMCPClient({ accessToken, - projectId, + projectRef, }: { accessToken: string - projectId: string + projectRef: string }) { + // Identifies the deployed build in the MCP server's `source_version` log field. + const sourceVersion = process.env.VERCEL_GIT_COMMIT_SHA + + const client = await createMCPClient({ + name: SOURCE_NAME, + transport: { + type: 'http', + url: getRemoteMcpUrl(projectRef), + headers: { + Authorization: `Bearer ${accessToken}`, + // Identify assistant traffic in the remote MCP server's logs + 'x-source-name': SOURCE_NAME, + ...(sourceVersion ? { 'x-source-version': sourceVersion } : {}), + }, + }, + }) + + return client +} + +/** + * Legacy in-process MCP client — the pre-migration behavior, kept as a fallback + * behind the `USE_REMOTE_MCP` gate (see `tools/mcp-tools.ts`). + * + * Instantiates `@supabase/mcp-server-supabase` in-process and connects to it over + * an in-memory transport. The heavy server package is imported dynamically so it + * is code-split into its own chunk and stays out of the (default, post-migration) + * remote path's bundle. + * + * TODO(AI-897): remove in process mcp — delete this once every environment has + * been flipped to the remote MCP server and has been stable. Tracked alongside + * the `USE_REMOTE_MCP` rollout. + */ +export async function createInProcessSupabaseMCPClient({ + accessToken, + projectRef, +}: { + accessToken: string + projectRef: string +}) { + // Dynamic imports keep the in-process server + its transport out of the remote + // path's bundle (loaded only when this fallback is actually taken). + // `.js` is required for esbuild ESM resolution. + const { InMemoryTransport } = await import('@modelcontextprotocol/sdk/inMemory.js') + const { createSupabaseMcpServer } = await import('@supabase/mcp-server-supabase') + const { createSupabaseApiPlatform } = await import('@supabase/mcp-server-supabase/platform/api') + const { API_URL } = await import('@/lib/constants') + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair() // Instantiate the MCP server and connect to its transport @@ -22,14 +118,13 @@ export async function createSupabaseMCPClient({ apiUrl, }), contentApiUrl: process.env.NEXT_PUBLIC_CONTENT_API_URL, - projectId, + projectId: projectRef, readOnly: true, }) await server.connect(serverTransport) - // Create the MCP client and connect to its transport const client = await createMCPClient({ - name: 'supabase-studio', + name: SOURCE_NAME, transport: clientTransport, }) diff --git a/apps/studio/lib/ai/tools/index.test.ts b/apps/studio/lib/ai/tools/index.test.ts new file mode 100644 index 00000000000..fb927e1bb29 --- /dev/null +++ b/apps/studio/lib/ai/tools/index.test.ts @@ -0,0 +1,86 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { getTools } from './index' +import { getMcpTools } from './mcp-tools' + +vi.mock('common', () => ({ IS_PLATFORM: true })) + +vi.mock('./mcp-tools', () => ({ getMcpTools: vi.fn() })) +vi.mock('./studio-tools', () => ({ getStudioTools: vi.fn(() => ({ studio_tool: {} })) })) +vi.mock('./schema-tools', () => ({ getSchemaTools: vi.fn(() => ({ schema_tool: {} })) })) +vi.mock('./incident-tools', () => ({ getIncidentTools: vi.fn(() => ({ incident_tool: {} })) })) +vi.mock('./fallback-tools', () => ({ getFallbackTools: vi.fn(() => ({ fallback_tool: {} })) })) +// Identity filter so assertions can check the raw merged tool set +vi.mock('../tool-filter', () => ({ filterToolsByOptInLevel: vi.fn((tools) => tools) })) + +const BASE_PARAMS = { + projectRef: 'abcdefghijklmnopqrst', + connectionString: 'postgresql://localhost', + authorization: 'Bearer token', + aiOptInLevel: 'schema_and_log_and_data' as const, + accessToken: 'access-token', + baseUrl: 'https://supabase.com/dashboard', + signal: new AbortController().signal, +} + +describe('ai/tools getTools', () => { + beforeEach(async () => { + vi.clearAllMocks() + vi.mocked(getMcpTools).mockResolvedValue({ list_tables: {} } as any) + // Reset to platform each test; the self-hosted test overrides to false. + // Done here (not afterEach) so the spy can't leak across tests via order. + const common = await import('common') + vi.spyOn(common, 'IS_PLATFORM', 'get').mockReturnValue(true) + }) + + it('includes studio, MCP, schema and incident tools on platform', async () => { + const tools = await getTools(BASE_PARAMS) + + expect(getMcpTools).toHaveBeenCalledWith({ + accessToken: BASE_PARAMS.accessToken, + projectRef: BASE_PARAMS.projectRef, + aiOptInLevel: BASE_PARAMS.aiOptInLevel, + signal: BASE_PARAMS.signal, + }) + expect(tools).toHaveProperty('studio_tool') + expect(tools).toHaveProperty('list_tables') + expect(tools).toHaveProperty('schema_tool') + expect(tools).toHaveProperty('incident_tool') + }) + + it('degrades gracefully to the remaining tools when remote MCP fetch fails', async () => { + const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + vi.mocked(getMcpTools).mockRejectedValueOnce(new Error('remote MCP unreachable')) + + const tools = await getTools(BASE_PARAMS) + + // The assistant still works with the non-MCP tools instead of throwing + expect(tools).toHaveProperty('studio_tool') + expect(tools).toHaveProperty('schema_tool') + expect(tools).toHaveProperty('incident_tool') + expect(tools).not.toHaveProperty('list_tables') + expect(consoleSpy).toHaveBeenCalled() + + consoleSpy.mockRestore() + }) + + it('does not fetch MCP tools when no access token is provided', async () => { + const tools = await getTools({ ...BASE_PARAMS, accessToken: undefined }) + + expect(getMcpTools).not.toHaveBeenCalled() + expect(tools).toHaveProperty('studio_tool') + expect(tools).not.toHaveProperty('list_tables') + }) + + it('uses fallback tools and skips MCP when self-hosted', async () => { + const common = await import('common') + vi.spyOn(common, 'IS_PLATFORM', 'get').mockReturnValue(false) + + const tools = await getTools(BASE_PARAMS) + + expect(getMcpTools).not.toHaveBeenCalled() + expect(tools).toHaveProperty('studio_tool') + expect(tools).toHaveProperty('fallback_tool') + expect(tools).not.toHaveProperty('list_tables') + }) +}) diff --git a/apps/studio/lib/ai/tools/index.ts b/apps/studio/lib/ai/tools/index.ts index 56102becf92..295921d1c70 100644 --- a/apps/studio/lib/ai/tools/index.ts +++ b/apps/studio/lib/ai/tools/index.ts @@ -16,6 +16,7 @@ export const getTools = async ({ aiOptInLevel, accessToken, baseUrl, + signal, }: { projectRef: string connectionString: string @@ -23,6 +24,9 @@ export const getTools = async ({ aiOptInLevel: AiOptInLevel accessToken?: string baseUrl?: string + // Required: tools fetched from the remote MCP server hold an HTTP connection + // that is closed when this signal aborts (i.e. when the request ends). + signal: AbortSignal }) => { // Always include studio tools let tools: ToolSet = getStudioTools({ projectRef, connectionString, authorization, aiOptInLevel }) @@ -39,12 +43,22 @@ export const getTools = async ({ }), } } else if (accessToken) { - // If platform, fetch MCP and other platform specific tools - const mcpTools = await getMcpTools({ - accessToken, - projectRef, - aiOptInLevel, - }) + // If platform, fetch MCP and other platform specific tools. The MCP tools + // may be fetched from the remote MCP server over the network (see + // `USE_REMOTE_MCP`), so a failure there (outage, timeout, auth) should + // degrade gracefully to the remaining tools rather than break the entire + // assistant. + let mcpTools: ToolSet = {} + try { + mcpTools = await getMcpTools({ + accessToken, + projectRef, + aiOptInLevel, + signal, + }) + } catch (error) { + console.error('Failed to fetch MCP tools:', error) + } tools = { ...tools, diff --git a/apps/studio/lib/ai/tools/mcp-tools.test.ts b/apps/studio/lib/ai/tools/mcp-tools.test.ts new file mode 100644 index 00000000000..5290cd99b01 --- /dev/null +++ b/apps/studio/lib/ai/tools/mcp-tools.test.ts @@ -0,0 +1,172 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { createInProcessSupabaseMCPClient, createSupabaseMCPClient } from '../supabase-mcp' +import { getMcpTools } from './mcp-tools' + +vi.mock('../supabase-mcp', () => ({ + createSupabaseMCPClient: vi.fn(), + createInProcessSupabaseMCPClient: vi.fn(), +})) + +const BASE_PARAMS = { + accessToken: 'token', + projectRef: 'abcdefghijklmnopqrst', + aiOptInLevel: 'schema_and_log_and_data' as const, + // A fresh, non-aborted signal by default; lifecycle tests override it + signal: new AbortController().signal, +} + +// A realistic remote tool set: all expected read tools plus the UI-executed ones +const FULL_REMOTE_TOOLS = { + search_docs: { description: 'docs' }, + list_tables: { description: 'list tables' }, + list_extensions: { description: 'extensions' }, + list_edge_functions: { description: 'edge functions' }, + list_branches: { description: 'branches' }, + get_advisors: { description: 'advisors' }, + get_logs: { description: 'get logs' }, + execute_sql: { description: 'execute sql' }, + deploy_edge_function: { description: 'deploy' }, +} + +describe('ai/tools/mcp-tools getMcpTools', () => { + let close: ReturnType + let tools: ReturnType + let consoleErrorSpy: ReturnType + + beforeEach(() => { + vi.clearAllMocks() + // These tests exercise the remote transport; pin the migration gate to it. + process.env.USE_REMOTE_MCP = 'true' + consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + close = vi.fn().mockResolvedValue(undefined) + tools = vi.fn().mockResolvedValue({ ...FULL_REMOTE_TOOLS }) + vi.mocked(createSupabaseMCPClient).mockResolvedValue({ tools, close } as any) + }) + + afterEach(() => { + delete process.env.USE_REMOTE_MCP + consoleErrorSpy.mockRestore() + }) + + it('returns MCP tools and strips UI-executed tools handled locally', async () => { + const result = await getMcpTools(BASE_PARAMS) + + expect(result).toHaveProperty('list_tables') + expect(result).toHaveProperty('get_logs') + expect(result).not.toHaveProperty('execute_sql') + expect(result).not.toHaveProperty('deploy_edge_function') + }) + + it('warns when the remote server is missing an expected tool (contract drift)', async () => { + const { list_branches, ...withoutBranches } = FULL_REMOTE_TOOLS + tools.mockResolvedValueOnce(withoutBranches) + + await getMcpTools(BASE_PARAMS) + + expect(consoleErrorSpy).toHaveBeenCalledWith(expect.stringContaining('list_branches')) + }) + + it('does not warn about drift when all expected tools are present', async () => { + await getMcpTools(BASE_PARAMS) + + const driftWarnings = consoleErrorSpy.mock.calls.filter( + ([msg]: unknown[]) => typeof msg === 'string' && msg.includes('missing expected tools') + ) + expect(driftWarnings).toHaveLength(0) + }) + + it('keeps the connection open until the request signal aborts', async () => { + const controller = new AbortController() + + await getMcpTools({ ...BASE_PARAMS, signal: controller.signal }) + // Tools execute later during streaming, so the client must stay open + expect(close).not.toHaveBeenCalled() + + controller.abort() + await Promise.resolve() + expect(close).toHaveBeenCalledTimes(1) + }) + + it('closes the client and skips fetching tools when the signal is already aborted', async () => { + const controller = new AbortController() + controller.abort() + + const result = await getMcpTools({ ...BASE_PARAMS, signal: controller.signal }) + + expect(close).toHaveBeenCalledTimes(1) + expect(tools).not.toHaveBeenCalled() + expect(result).toEqual({}) + }) + + it('closes the client and rethrows when fetching tools fails, without double-closing on later abort', async () => { + tools.mockRejectedValueOnce(new Error('network unreachable')) + const controller = new AbortController() + + await expect(getMcpTools({ ...BASE_PARAMS, signal: controller.signal })).rejects.toThrow( + 'network unreachable' + ) + expect(close).toHaveBeenCalledTimes(1) + + // A subsequent abort must not close the (already closed) client again + controller.abort() + await Promise.resolve() + expect(close).toHaveBeenCalledTimes(1) + }) + + it('closes the client and rethrows when tool validation fails', async () => { + // A known tool name with a non-object value passes the opt-in filter but + // fails schema validation + tools.mockResolvedValueOnce({ ...FULL_REMOTE_TOOLS, list_tables: 'not-an-object' }) + + await expect(getMcpTools(BASE_PARAMS)).rejects.toThrow('MCP tools validation failed') + expect(close).toHaveBeenCalledTimes(1) + }) +}) + +describe('ai/tools/mcp-tools getMcpTools transport selection', () => { + let close: ReturnType + let tools: ReturnType + let consoleErrorSpy: ReturnType + + beforeEach(() => { + vi.clearAllMocks() + consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + close = vi.fn().mockResolvedValue(undefined) + tools = vi.fn().mockResolvedValue({ ...FULL_REMOTE_TOOLS }) + vi.mocked(createSupabaseMCPClient).mockResolvedValue({ tools, close } as any) + vi.mocked(createInProcessSupabaseMCPClient).mockResolvedValue({ tools, close } as any) + }) + + afterEach(() => { + delete process.env.USE_REMOTE_MCP + consoleErrorSpy.mockRestore() + }) + + it('uses the remote client when USE_REMOTE_MCP is "true"', async () => { + process.env.USE_REMOTE_MCP = 'true' + + await getMcpTools(BASE_PARAMS) + + expect(createSupabaseMCPClient).toHaveBeenCalledTimes(1) + expect(createInProcessSupabaseMCPClient).not.toHaveBeenCalled() + }) + + it('falls back to the in-process client when USE_REMOTE_MCP is unset (default)', async () => { + delete process.env.USE_REMOTE_MCP + + await getMcpTools(BASE_PARAMS) + + expect(createInProcessSupabaseMCPClient).toHaveBeenCalledTimes(1) + expect(createSupabaseMCPClient).not.toHaveBeenCalled() + }) + + it('uses the in-process client for any non-"true" value', async () => { + process.env.USE_REMOTE_MCP = 'false' + + await getMcpTools(BASE_PARAMS) + + expect(createInProcessSupabaseMCPClient).toHaveBeenCalledTimes(1) + expect(createSupabaseMCPClient).not.toHaveBeenCalled() + }) +}) diff --git a/apps/studio/lib/ai/tools/mcp-tools.ts b/apps/studio/lib/ai/tools/mcp-tools.ts index 3ccfcc9c992..f895f6d90b9 100644 --- a/apps/studio/lib/ai/tools/mcp-tools.ts +++ b/apps/studio/lib/ai/tools/mcp-tools.ts @@ -1,42 +1,135 @@ +// Type-only import (erased at build time — pulls no runtime code into this route). +import type * as SupabaseMcp from '@supabase/mcp-server-supabase' import type { ToolSet } from 'ai' -import { createSupabaseMCPClient } from '../supabase-mcp' +import { createInProcessSupabaseMCPClient, createSupabaseMCPClient } from '../supabase-mcp' import { filterToolsByOptInLevel, toolSetValidationSchema } from '../tool-filter' import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' -const UI_EXECUTED_TOOLS = ['execute_sql', 'deploy_edge_function'] +/** + * Union of the tool names exposed by the pinned `@supabase/mcp-server-supabase` + * version. Studio's dependency is bumped in lockstep with the remote MCP server + * (via automated bump PRs), so typing our tool-name lists against this makes an + * upstream rename/removal a **compile-time** failure (`pnpm typecheck`) in that + * PR, instead of a silent capability loss at runtime. + */ +type SupabaseMcpToolName = keyof typeof SupabaseMcp.supabaseMcpToolSchemas + +// UI-executed tools handled locally by Studio (see getStudioTools); the remote +// MCP server's versions are removed so the UI-controlled Studio versions win. +const UI_EXECUTED_TOOLS = [ + 'execute_sql', + 'deploy_edge_function', +] as const satisfies readonly SupabaseMcpToolName[] + +// Read-only tools the assistant relies on from the remote MCP server — the +// MCP-sourced subset of the allowlist in tool-filter.ts (TOOL_CATEGORY_MAP). +// `satisfies` gives the compile-time drift guard; the runtime check below also +// catches a deployed server that returns fewer tools (feature flags / version +// skew). The allowlist remains the source of truth for what is allowed. +const EXPECTED_MCP_TOOLS = [ + 'search_docs', + 'list_tables', + 'list_extensions', + 'list_edge_functions', + 'list_branches', + 'get_advisors', + 'get_logs', +] as const satisfies readonly SupabaseMcpToolName[] export const getMcpTools = async ({ accessToken, projectRef, aiOptInLevel, + signal, }: { accessToken: string projectRef: string aiOptInLevel: AiOptInLevel + // Required: the remote client holds an HTTP connection that must be torn down + // when the request ends. The caller owns that lifecycle via this signal. + signal: AbortSignal }) => { - // If platform, fetch MCP client and tools which replace old local tools - const mcpClient = await createSupabaseMCPClient({ + // Connect to the MCP server and fetch its tools, which replace the old local + // tools. `USE_REMOTE_MCP` gates the transport: the remote HTTP server (target + // state) or the legacy in-process server (fallback during the migration), + // defaulting to in-process until an environment opts in. Flip it per + // environment (staging → prod → Nimbus) once each one's prerequisites are met + // (dashboard-token support on the MCP API, remote MCP enabled for Nimbus); + // unset to roll back on the next deploy. Both transports expose the same tool + // surface, so the filtering, drift detection, and lifecycle handling below are + // transport-agnostic. + // + // TODO(AI-897): remove in process mcp — once every environment has been + // flipped and is stable, delete `createInProcessSupabaseMCPClient` and this + // fallback branch. + const useRemoteMcp = process.env.USE_REMOTE_MCP === 'true' + const createClient = useRemoteMcp ? createSupabaseMCPClient : createInProcessSupabaseMCPClient + const mcpClient = await createClient({ accessToken, - projectId: projectRef, + projectRef, }) - const availableMcpTools = (await mcpClient.tools()) as ToolSet - // Filter tools based on the (potentially modified) AI opt-in level - const allowedMcpTools = filterToolsByOptInLevel(availableMcpTools, aiOptInLevel) - - // Remove UI-executed tools handled locally - const filteredMcpTools: ToolSet = { ...allowedMcpTools } - UI_EXECUTED_TOOLS.forEach((toolName) => { - delete filteredMcpTools[toolName] - }) - - // Validate that only known tools are provided - const validation = toolSetValidationSchema.safeParse(filteredMcpTools) - if (!validation.success) { - console.error('MCP tools validation error:', validation.error) - throw new Error('Internal error: MCP tools validation failed') + // The remote client keeps an HTTP connection open. The tools' `execute` + // functions are invoked later, while the response is streaming, so the + // connection must stay open until the request ends. Close it exactly once when + // the request is done (normal completion or abort) to avoid leaking a + // connection per request. + let closed = false + const closeClient = () => { + if (closed) return + closed = true + void mcpClient.close().catch(() => {}) } - return validation.data + // The request already ended before we could fetch tools; don't bother. + if (signal.aborted) { + closeClient() + return {} as ToolSet + } + signal.addEventListener('abort', closeClient, { once: true }) + + try { + const availableMcpTools = (await mcpClient.tools()) as ToolSet + + // Runtime drift detection: `EXPECTED_MCP_TOOLS` is compile-time-checked + // against the pinned package (see its declaration), but the *deployed* remote + // server can still return fewer tools than the pinned types — feature flags, + // killswitches, or version skew during a bump. `filterToolsByOptInLevel` + // drops missing tools silently, so warn to make that observable. + const missingExpectedTools = EXPECTED_MCP_TOOLS.filter((name) => !(name in availableMcpTools)) + if (missingExpectedTools.length > 0) { + console.error( + `Remote MCP server is missing expected tools: ${missingExpectedTools.join(', ')}. ` + + 'The tool contract may have drifted; the assistant will operate without them.' + ) + } + + // Safety gate: `filterToolsByOptInLevel` keeps only tools in the allowlist + // (tool-filter.ts TOOL_CATEGORY_MAP) and drops everything else. This — not + // the `read_only` query param — is what prevents the remote server's + // write/destructive tools (apply_migration, create_branch, ...) from reaching + // the assistant. `read_only` is defense-in-depth (those tools throw at + // runtime). Do not remove this filter on the assumption `read_only` suffices. + const allowedMcpTools = filterToolsByOptInLevel(availableMcpTools, aiOptInLevel) + + // Remove UI-executed tools handled locally + const filteredMcpTools: ToolSet = { ...allowedMcpTools } + UI_EXECUTED_TOOLS.forEach((toolName) => { + delete filteredMcpTools[toolName] + }) + + // Validate that only known tools are provided + const validation = toolSetValidationSchema.safeParse(filteredMcpTools) + if (!validation.success) { + console.error('MCP tools validation error:', validation.error) + throw new Error('Internal error: MCP tools validation failed') + } + + return validation.data + } catch (error) { + // Don't leak the connection if fetching or validating tools fails + closeClient() + throw error + } } diff --git a/apps/studio/lib/ai/tools/mock-tools.test.ts b/apps/studio/lib/ai/tools/mock-tools.test.ts new file mode 100644 index 00000000000..d696d2d762d --- /dev/null +++ b/apps/studio/lib/ai/tools/mock-tools.test.ts @@ -0,0 +1,64 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { getMockTools } from './mock-tools' +import { createInProcessSupabaseMCPClient } from '@/lib/ai/supabase-mcp' + +// The one real tool in the eval harness (search_docs) is sourced from an +// in-process MCP client. Mock that client so this test stays hermetic and +// guards the wiring, not a live connection. +vi.mock('@/lib/ai/supabase-mcp', () => ({ + createInProcessSupabaseMCPClient: vi.fn(), +})) + +const SEARCH_DOCS = { description: 'search the docs' } + +describe('ai/tools/mock-tools getMockTools', () => { + let close: ReturnType + let tools: ReturnType + + beforeEach(() => { + vi.clearAllMocks() + close = vi.fn().mockResolvedValue(undefined) + tools = vi.fn().mockResolvedValue({ search_docs: SEARCH_DOCS }) + vi.mocked(createInProcessSupabaseMCPClient).mockResolvedValue({ tools, close } as any) + }) + + it('sources the real search_docs from the in-process MCP server alongside the deterministic mocks', async () => { + const result = await getMockTools(undefined, new AbortController().signal) + + expect(createInProcessSupabaseMCPClient).toHaveBeenCalledTimes(1) + // The real tool, wired through from the MCP client + expect(result).toHaveProperty('search_docs', SEARCH_DOCS) + // A couple of the deterministic mocks, to confirm the merge + expect(result).toHaveProperty('list_tables') + expect(result).toHaveProperty('get_logs') + }) + + // This is the regression guard: if the eval's MCP wiring breaks (contract + // drift, or a refactor that stops sourcing search_docs — e.g. the future + // AI-897 removal of the in-process client), fail loudly in normal CI instead + // of only surfacing during an opt-in Braintrust eval run. + it('throws a clear error when the MCP server does not expose search_docs', async () => { + tools.mockResolvedValueOnce({}) + + await expect(getMockTools(undefined, new AbortController().signal)).rejects.toThrow( + 'search_docs tool not available from MCP server' + ) + }) + + it('closes the MCP client when the caller aborts the signal', async () => { + const controller = new AbortController() + + await getMockTools(undefined, controller.signal) + // Connection stays open until generation ends (search_docs runs during it) + expect(close).not.toHaveBeenCalled() + + controller.abort() + await Promise.resolve() + expect(close).toHaveBeenCalledTimes(1) + }) + + afterEach(() => { + vi.clearAllMocks() + }) +}) diff --git a/apps/studio/lib/ai/tools/mock-tools.ts b/apps/studio/lib/ai/tools/mock-tools.ts index 8465459d3d3..f12d9b0e55e 100644 --- a/apps/studio/lib/ai/tools/mock-tools.ts +++ b/apps/studio/lib/ai/tools/mock-tools.ts @@ -1,9 +1,9 @@ import assert from 'node:assert' -import { tool } from 'ai' +import { tool, type ToolSet } from 'ai' import { z } from 'zod' import { getStudioTools } from '../tools/studio-tools' -import { getMcpTools } from '@/lib/ai/tools/mcp-tools' +import { createInProcessSupabaseMCPClient } from '@/lib/ai/supabase-mcp' const listTablesInputSchema = z.object({ schemas: z.array(z.string()).describe('The schema names to list.'), @@ -306,14 +306,25 @@ export type MockToolOverrides = { * * Note: search_docs uses the real implementation */ -export async function getMockTools(overrides?: MockToolOverrides) { +export async function getMockTools(overrides: MockToolOverrides | undefined, signal: AbortSignal) { const mockedStudioTools = createMockedStudioTools() - const { search_docs } = await getMcpTools({ + // Every tool here is a deterministic mock except `search_docs`, which uses the + // real implementation. We source it from an in-process MCP server directly + // (rather than `getMcpTools`) so the eval harness stays hermetic and decoupled + // from the assistant's transport gate (`USE_REMOTE_MCP`): the in-process server + // needs no live remote endpoint or real access token. See AI-897 for how to + // point evals at the remote MCP server instead. + const mcpClient = await createInProcessSupabaseMCPClient({ accessToken: 'mock-access-token', projectRef: 'mock-project-ref', - aiOptInLevel: 'schema_and_log_and_data', }) + // The caller owns this signal and aborts it once generation is done, which + // closes the client opened here (search_docs executes during generation, so + // the connection must stay open until then). + signal.addEventListener('abort', () => void mcpClient.close().catch(() => {}), { once: true }) + + const { search_docs } = (await mcpClient.tools()) as ToolSet assert(search_docs, 'search_docs tool not available from MCP server') diff --git a/apps/studio/lib/api/generate-v4.test.ts b/apps/studio/lib/api/generate-v4.test.ts index 7a064364246..5110d04fe8a 100644 --- a/apps/studio/lib/api/generate-v4.test.ts +++ b/apps/studio/lib/api/generate-v4.test.ts @@ -42,6 +42,7 @@ test('generateV4 calls the tool sanitizer', async () => { status: vi.fn(() => mockRes), json: vi.fn(() => mockRes), setHeader: vi.fn(() => mockRes), + on: vi.fn(), } vi.mock('@/lib/ai/ai-details', () => ({ @@ -83,4 +84,7 @@ test('generateV4 calls the tool sanitizer', async () => { await generateV4(mockReq as any, mockRes as any) expect(sanitizeMessagePart).toHaveBeenCalled() + // The response 'close' event must be wired up so the remote MCP connection + // opened in getTools is torn down when the stream finishes or the client drops + expect(mockRes.on).toHaveBeenCalledWith('close', expect.any(Function)) }) diff --git a/apps/studio/package.json b/apps/studio/package.json index 7da743e96fe..6eaf32965b7 100644 --- a/apps/studio/package.json +++ b/apps/studio/package.json @@ -33,6 +33,7 @@ "build:graphql-types": "tsx scripts/download-graphql-schema.mts && pnpm graphql-codegen --config scripts/codegen.ts", "build:graphql-types:watch": "pnpm graphql-codegen --config scripts/codegen.ts --watch", "evals:setup": "cp node_modules/libpg-query/wasm/libpg-query.wasm evals/libpg-query.wasm", + "evals:preflight": "tsx evals/preflight.ts", "evals:run": "braintrust eval --no-send-logs evals/assistant.eval.ts", "evals:upload": "braintrust eval evals/assistant.eval.ts", "scorers:deploy": "IS_BRAINTRUST_PUSH=true braintrust push evals/scorer-online.ts" diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index 457881205c7..1bbc78aa29b 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -165,6 +165,9 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: Jw const abortController = new AbortController() req.on('close', () => abortController.abort()) req.on('aborted', () => abortController.abort()) + // Fires when the response finishes streaming or the connection drops, which + // is what tears down the remote MCP connection opened in getTools. + res.on('close', () => abortController.abort()) const tools = await getTools({ projectRef, @@ -173,6 +176,7 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse, claims?: Jw aiOptInLevel, accessToken, baseUrl: getURL(), + signal: abortController.signal, }) // Get a list of all schemas to add to context diff --git a/apps/studio/pages/api/ai/sql/policy.ts b/apps/studio/pages/api/ai/sql/policy.ts index a6c30ebd76f..8cb27ed7061 100644 --- a/apps/studio/pages/api/ai/sql/policy.ts +++ b/apps/studio/pages/api/ai/sql/policy.ts @@ -100,53 +100,67 @@ export async function handlePost(req: NextApiRequest, res: NextApiResponse) { return res.status(500).json({ error: modelError.message }) } - const tools = await getTools({ - projectRef, - connectionString, - authorization, - aiOptInLevel, - accessToken, - }) + // Closes the remote MCP connection opened in getTools when generation is done, + // if anything below throws, or if the client disconnects mid-generation so the + // connection isn't held until generateText resolves on its own (mirrors the + // request-scoped cleanup in generate-v4.ts). + const toolsAbortController = new AbortController() + req.on('close', () => toolsAbortController.abort()) + req.on('aborted', () => toolsAbortController.abort()) + // Fires when the response finishes or the connection drops. + res.on('close', () => toolsAbortController.abort()) + try { + const tools = await getTools({ + projectRef, + connectionString, + authorization, + aiOptInLevel, + accessToken, + signal: toolsAbortController.signal, + }) - const { experimental_output } = await generateText({ - ...modelParams, - stopWhen: stepCountIs(5), - prompt: source` - You are a Postgres RLS (Row Level Security) expert. - Determine the most appropriate policies for the "${schema}"."${tableName}" table within a Supabase project. + const { experimental_output } = await generateText({ + ...modelParams, + stopWhen: stepCountIs(5), + prompt: source` + You are a Postgres RLS (Row Level Security) expert. + Determine the most appropriate policies for the "${schema}"."${tableName}" table within a Supabase project. - ${columns.length > 0 ? `Table columns: ${columns.join(', ')}` : 'No column metadata provided.'} + ${columns.length > 0 ? `Table columns: ${columns.join(', ')}` : 'No column metadata provided.'} - ${message ? `User request: ${message}` : ''} + ${message ? `User request: ${message}` : ''} - RLS Guide: ${RLS_PROMPT} + RLS Guide: ${RLS_PROMPT} - Requirements: - - Use the available planning and schema tools (like "list_policies" or "list_tables") to inspect the "${schema}" schema and existing policies before generating new ones. - - Ensure policies strictly adhere to the existing schema - - Return a curated list of recommended CREATE POLICY statements as JSON. - - Each policy must include: name, sql, command (SELECT/INSERT/UPDATE/DELETE/ALL), action (PERMISSIVE/RESTRICTIVE), roles (array of role names). - - Include "definition" (USING clause expression without the USING keyword) for SELECT, UPDATE, DELETE policies. - - Include "check" (WITH CHECK clause expression without the WITH CHECK keywords) for INSERT, UPDATE policies. - - Avoid duplicating existing policies and reference the public schema and typical Supabase best practices when deciding the coverage. - - Prefer PERMISSIVE policies unless a RESTRICTIVE policy is explicitly required - `, - tools, - experimental_output: Output.object({ - schema: z.object({ - policies: z.array(policySchema), + Requirements: + - Use the available planning and schema tools (like "list_policies" or "list_tables") to inspect the "${schema}" schema and existing policies before generating new ones. + - Ensure policies strictly adhere to the existing schema + - Return a curated list of recommended CREATE POLICY statements as JSON. + - Each policy must include: name, sql, command (SELECT/INSERT/UPDATE/DELETE/ALL), action (PERMISSIVE/RESTRICTIVE), roles (array of role names). + - Include "definition" (USING clause expression without the USING keyword) for SELECT, UPDATE, DELETE policies. + - Include "check" (WITH CHECK clause expression without the WITH CHECK keywords) for INSERT, UPDATE policies. + - Avoid duplicating existing policies and reference the public schema and typical Supabase best practices when deciding the coverage. + - Prefer PERMISSIVE policies unless a RESTRICTIVE policy is explicitly required + `, + tools, + experimental_output: Output.object({ + schema: z.object({ + policies: z.array(policySchema), + }), }), - }), - }) + }) - // Add table and schema to each policy from the request - const policies = (experimental_output?.policies ?? []).map((policy) => ({ - ...policy, - table: tableName, - schema, - })) + // Add table and schema to each policy from the request + const policies = (experimental_output?.policies ?? []).map((policy) => ({ + ...policy, + table: tableName, + schema, + })) - return res.json(policies) + return res.json(policies) + } finally { + toolsAbortController.abort() + } } catch (error) { if (error instanceof Error) { console.error(`AI policy generation failed: ${error.message}`) diff --git a/apps/studio/turbo.jsonc b/apps/studio/turbo.jsonc index d47f1beea95..08a66e00023 100644 --- a/apps/studio/turbo.jsonc +++ b/apps/studio/turbo.jsonc @@ -71,6 +71,9 @@ "OPENAI_API_KEY", "BRAINTRUST_API_KEY", "BRAINTRUST_PROJECT_ID", + // Gates the dashboard assistant between the remote MCP server and the + // legacy in-process one (see lib/ai/tools/mcp-tools.ts). + "USE_REMOTE_MCP", "AUTH_JWT_SECRET", "LOGFLARE_API_KEY", "LOGFLARE_PUBLIC_ACCESS_TOKEN",