From 61dc49ca4da0be5463d7896c94a9f6ddd926d0de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BA=8C=E4=B8=AB=E8=AE=B2=E6=A2=B5?= Date: Sun, 2 Aug 2026 09:01:22 +0800 Subject: [PATCH] fix(ai): preserve Dameng database and schema scopes --- .../src/components/editor/AiAssistant.vue | 85 ++++-- .../__tests__/useDatabaseOptions.spec.ts | 34 +++ .../src/lib/__tests__/ai/aiContext.spec.ts | 79 +++++- apps/desktop/src/lib/ai/ai.ts | 50 +++- apps/desktop/src/lib/backend/http.ts | 4 + apps/desktop/src/lib/backend/tauri.ts | 4 + crates/dbx-core/src/agent_loop.rs | 49 +++- crates/dbx-core/src/agent_tools.rs | 264 ++++++++++++++++-- crates/dbx-core/src/ai_claude_code_cli.rs | 1 + crates/dbx-core/src/ai_cli_agent.rs | 13 +- crates/dbx-core/src/ai_codex_cli.rs | 1 + crates/dbx-core/src/ai_pi_agent_cli.rs | 2 + crates/dbx-mcp/src/backend.rs | 12 +- crates/dbx-mcp/src/server.rs | 74 ++++- crates/dbx-web/src/routes/ai.rs | 7 + src-tauri/src/commands/ai.rs | 31 ++ 16 files changed, 634 insertions(+), 76 deletions(-) diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index 19f8852d1e..be6ae1ff1d 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -55,7 +55,7 @@ import ConnectionGroupBadge from "@/components/connection/ConnectionGroupBadge.v import { useQueryStore } from "@/stores/queryStore"; import { useToast } from "@/composables/useToast"; import { useNavigationTargets } from "@/composables/useNavigationTargets"; -import { buildAiContext, runAgentStream, isVectorDbType, isValidActionForMode, defaultActionForMode, type AiAction, type AiAssistantMode, type AiSqlFileContext, type CustomPromptContext } from "@/lib/ai/ai"; +import { buildAiContext, resolveAiDatabaseTarget, resolveAiNamespaceSelection, resolveDefaultAiSchema, runAgentStream, isVectorDbType, isValidActionForMode, defaultActionForMode, type AiAction, type AiAssistantMode, type AiSqlFileContext, type CustomPromptContext } from "@/lib/ai/ai"; import { isAiConfigModelCandidate } from "@/lib/ai/aiConfigCandidates"; import { orderAiConfigsForDisplay } from "@/lib/ai/aiConfigOrdering"; import { effortSelectionEquals, runtimeEffortFromPreference } from "@/lib/ai/aiEffortPreference"; @@ -75,7 +75,7 @@ import { aiCancelStream, saveAiConversation, loadAiConversations, deleteAiConver import type { AiMessage } from "@/lib/backend/api"; import type { AiConfigItem, AiEffortCapability, AiEffortOption, AiEffortSelection } from "@/types/ai"; import type { ConnectionConfig, QueryTab, SavedSqlFile, TableInfo } from "@/types/database"; -import { useDatabaseOptions } from "@/composables/useDatabaseOptions"; +import { fetchNamespaceOptionsForConnection, useDatabaseOptions } from "@/composables/useDatabaseOptions"; import { decodeSelectableDatabaseValue, encodeSelectableDatabaseValue, formatDatabaseLabel, resolveDefaultDatabase } from "@/lib/database/defaultDatabase"; import { normalizeSqliteNamespace } from "@/lib/database/sqliteNamespace"; import { isSchemaAware } from "@/lib/database/databaseCapabilities"; @@ -730,6 +730,7 @@ let confirmedWriteSqlText: string | undefined = undefined; * to prevent a database change between confirmation and execution. */ let confirmedConnectionId: string | undefined = undefined; let confirmedDatabase: string | undefined = undefined; +let confirmedSchema: string | undefined = undefined; /** Clear all pending write-confirmation state. Call on every early-return * and failure path so a stale grant cannot leak into a subsequent send(). */ @@ -738,9 +739,13 @@ function clearPendingWriteGrant() { confirmedWriteSqlText = undefined; confirmedConnectionId = undefined; confirmedDatabase = undefined; + confirmedSchema = undefined; } -const productionContext = computed(() => productionContextForDatabase(props.connection, props.tab?.database)); +const productionContext = computed(() => { + const target = props.connection && props.tab ? resolveAiDatabaseTarget(props.tab, props.connection) : undefined; + return productionContextForDatabase(props.connection, target?.database); +}); function sendProposalReply(positive: boolean) { // Disable while a stream is in flight or no proposal is currently active. @@ -762,7 +767,11 @@ function sendProposalReply(positive: boolean) { if (confirmedWriteSqlText) { allowWriteSqlForNextRun = true; confirmedConnectionId = props.connection?.id; - confirmedDatabase = props.tab?.database || ""; + if (props.tab && props.connection) { + const target = resolveAiDatabaseTarget(props.tab, props.connection); + confirmedDatabase = target.database; + confirmedSchema = target.schema; + } } // When no SQL code block is found in the proposal, treat the // confirmation as rejected — we cannot bind the agent to a @@ -804,11 +813,17 @@ function selectModeActionItem(action: AiAction) { modeActionOpen.value = false; } -const { databaseOptions: allDbOptions, loadDatabaseOptions } = useDatabaseOptions(); +const { databaseOptions, loadDatabaseOptions } = useDatabaseOptions(); + +// Dameng presents schemas as its top-level namespace, unlike the other +// connection types that rely on the shared database-options loader. +const aiDatabaseOptions = ref>({}); const dbOptions = computed(() => { - if (!props.connection) return []; - return allDbOptions.value[props.connection.id] || []; + const connection = props.connection; + if (!connection) return []; + if (connection.db_type === "dameng") return aiDatabaseOptions.value[connection.id] || []; + return databaseOptions.value[connection.id] || []; }); const dbSelectOptions = computed(() => { @@ -824,20 +839,29 @@ const dbSelectOptions = computed(() => { })); }); -const selectedDatabaseSelectValue = computed(() => (props.connection ? encodeSelectableDatabaseValue(props.connection.db_type, props.tab?.database || "") : "")); +const selectedNamespace = computed(() => (props.connection && props.tab ? resolveAiNamespaceSelection(props.tab, props.connection).value : "")); + +const selectedDatabaseSelectValue = computed(() => (props.connection ? encodeSelectableDatabaseValue(props.connection.db_type, selectedNamespace.value) : "")); const selectedDatabaseLabel = computed(() => { if (!props.connection) return t("editor.selectDatabase"); if (!props.tab) return t("editor.selectDatabase"); - return formatDatabaseLabel(props.connection, props.tab.database || "", { + return formatDatabaseLabel(props.connection, selectedNamespace.value, { defaultDatabase: t("editor.defaultDatabase"), noDatabase: t("editor.noDatabase"), }); }); -async function loadDatabases() { - if (!props.connection) return; - await loadDatabaseOptions(props.connection.id); +async function loadDatabases(connection = props.connection): Promise { + if (!connection) return []; + if (connection.db_type !== "dameng") { + await loadDatabaseOptions(connection.id); + return databaseOptions.value[connection.id] || []; + } + await connectionStore.ensureConnected(connection.id); + const options = await fetchNamespaceOptionsForConnection(connection.id, connection); + aiDatabaseOptions.value[connection.id] = options; + return options; } async function changeConnection(connectionId: string) { @@ -845,16 +869,16 @@ async function changeConnection(connectionId: string) { if (!conn) return; connectionStore.activeConnectionId = connectionId; const tab = props.tab; + const tabId = tab ? tab.id : queryStore.createTab(connectionId, resolveDefaultDatabase(conn, [])); if (tab) { queryStore.updateConnection(tab.id, connectionId, resolveDefaultDatabase(conn, [])); - } else { - queryStore.createTab(connectionId, resolveDefaultDatabase(conn, [])); } try { - await loadDatabaseOptions(connectionId); - const database = resolveDefaultDatabase(conn, allDbOptions.value[connectionId] || []); - if (tab) { - queryStore.updateDatabase(tab.id, database); + const options = await loadDatabases(conn); + if (conn.db_type === "dameng") { + queryStore.updateSchema(tabId, resolveDefaultAiSchema(conn, options)); + } else { + queryStore.updateDatabase(tabId, resolveDefaultDatabase(conn, options)); } } catch (e: unknown) { const message = e instanceof Error ? e.message : String(e); @@ -862,11 +886,16 @@ async function changeConnection(connectionId: string) { } } -function changeDatabase(value: string) { +function changeNamespace(value: string) { const tab = props.tab; const connection = props.connection; if (!tab || !connection) return; - queryStore.updateDatabase(tab.id, decodeSelectableDatabaseValue(connection.db_type, value)); + const namespace = decodeSelectableDatabaseValue(connection.db_type, value); + if (resolveAiNamespaceSelection(tab, connection).kind === "schema") { + queryStore.updateSchema(tab.id, namespace || undefined); + } else { + queryStore.updateDatabase(tab.id, namespace); + } } function flushAssistantDeltas() { @@ -1717,7 +1746,9 @@ async function send() { if (msg.role === "assistant" && msg.content) { confirmedWriteSqlText = extractSingleSqlCodeBlock(msg.content); confirmedConnectionId = connection.id; - confirmedDatabase = tab.database || ""; + const target = resolveAiDatabaseTarget(tab, connection); + confirmedDatabase = target.database; + confirmedSchema = target.schema; break; } if (msg.role === "user") break; @@ -1727,11 +1758,12 @@ async function send() { } } } - // Verify the connection/database haven't changed since the user confirmed - // the write operation. If the user switched connections or databases between + // Verify the connection/database/schema haven't changed since the user confirmed + // the write operation. If the user switched connections or namespaces between // confirmation and execution, the grant is void. if (allowWriteSqlForNextRun && confirmedWriteSqlText) { - if (confirmedConnectionId !== connection.id || confirmedDatabase !== (tab.database || "")) { + const target = resolveAiDatabaseTarget(tab, connection); + if (confirmedConnectionId !== connection.id || confirmedDatabase !== target.database || confirmedSchema !== target.schema) { allowWriteSqlForNextRun = false; confirmedWriteSqlText = undefined; } @@ -1743,10 +1775,12 @@ async function send() { // state, so the values survive to be passed through to the backend. const confirmedTargetConnId = allowWriteSql ? confirmedConnectionId : undefined; const confirmedTargetDb = allowWriteSql ? confirmedDatabase : undefined; + const confirmedTargetSchema = allowWriteSql ? confirmedSchema : undefined; allowWriteSqlForNextRun = false; confirmedWriteSqlText = undefined; confirmedConnectionId = undefined; confirmedDatabase = undefined; + confirmedSchema = undefined; messages.value.push({ role: "assistant", content: "", sourceConnectionName: connection.name }); const assistantIdx = messages.value.length - 1; const sessionId = uuid(); @@ -1770,6 +1804,7 @@ async function send() { confirmedWriteSql, confirmedConnectionId: confirmedTargetConnId, confirmedDatabase: confirmedTargetDb, + confirmedSchema: confirmedTargetSchema, }, history, (event: AgentEvent) => { @@ -2398,7 +2433,7 @@ async function openExternalUrl(url: string) { :model-value="selectedDatabaseSelectValue" @update:model-value=" (v) => { - if (typeof v === 'string') changeDatabase(v); + if (typeof v === 'string') changeNamespace(v); } " @update:open=" diff --git a/apps/desktop/src/composables/__tests__/useDatabaseOptions.spec.ts b/apps/desktop/src/composables/__tests__/useDatabaseOptions.spec.ts index b73721fadc..28b2a44861 100644 --- a/apps/desktop/src/composables/__tests__/useDatabaseOptions.spec.ts +++ b/apps/desktop/src/composables/__tests__/useDatabaseOptions.spec.ts @@ -17,6 +17,8 @@ const mocks = vi.hoisted(() => ({ getConfig: vi.fn(), listDatabases: vi.fn(), listSchemas: vi.fn(), + redisListDatabases: vi.fn(), + mongoListDatabases: vi.fn(), listDorisCatalogs: vi.fn(), listDorisCatalogDatabases: vi.fn(), })); @@ -24,6 +26,8 @@ const mocks = vi.hoisted(() => ({ vi.mock("@/lib/backend/api", () => ({ listDatabases: mocks.listDatabases, listSchemas: mocks.listSchemas, + redisListDatabases: mocks.redisListDatabases, + mongoListDatabases: mocks.mongoListDatabases, listDorisCatalogs: mocks.listDorisCatalogs, listDorisCatalogDatabases: mocks.listDorisCatalogDatabases, })); @@ -185,3 +189,33 @@ describe("namespace options", () => { expect(mocks.listSchemas).not.toHaveBeenCalled(); }); }); + +describe("database options loader", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("uses the Redis database API", async () => { + mocks.getConfig.mockReturnValue({ db_type: "redis" }); + mocks.redisListDatabases.mockResolvedValue([{ db: 0 }, { db: 1 }]); + const options = useDatabaseOptions(); + + await options.loadDatabaseOptions("connection-1"); + + expect(options.databaseOptions.value["connection-1"]).toEqual(["0", "1"]); + expect(mocks.redisListDatabases).toHaveBeenCalledWith("connection-1"); + expect(mocks.listDatabases).not.toHaveBeenCalled(); + }); + + it("uses the MongoDB database API", async () => { + mocks.getConfig.mockReturnValue({ db_type: "mongodb" }); + mocks.mongoListDatabases.mockResolvedValue(["app", "analytics"]); + const options = useDatabaseOptions(); + + await options.loadDatabaseOptions("connection-1"); + + expect(options.databaseOptions.value["connection-1"]).toEqual(["app", "analytics"]); + expect(mocks.mongoListDatabases).toHaveBeenCalledWith("connection-1"); + expect(mocks.listDatabases).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/desktop/src/lib/__tests__/ai/aiContext.spec.ts b/apps/desktop/src/lib/__tests__/ai/aiContext.spec.ts index f3b3a73a4a..de365e14dd 100644 --- a/apps/desktop/src/lib/__tests__/ai/aiContext.spec.ts +++ b/apps/desktop/src/lib/__tests__/ai/aiContext.spec.ts @@ -1,5 +1,6 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; -import { buildAiContext } from "@/lib/ai/ai"; +import { buildAiContext, resolveAiDatabaseTarget, resolveAiNamespaceSelection, resolveDefaultAiSchema, runAgentStream } from "@/lib/ai/ai"; +import type { AiConfig } from "@/types/ai"; import type { ConnectionConfig, QueryTab } from "@/types/database"; const apiMock = vi.hoisted(() => ({ @@ -7,6 +8,7 @@ const apiMock = vi.hoisted(() => ({ getColumns: vi.fn(), listIndexes: vi.fn(), listForeignKeys: vi.fn(), + aiAgentStream: vi.fn(), })); vi.mock("@/lib/backend/api", () => apiMock); @@ -24,12 +26,13 @@ function sqliteConnection(database?: string): ConnectionConfig { }; } -function queryTab(database: string): QueryTab { +function queryTab(database: string, schema?: string): QueryTab { return { id: "tab-1", title: "Query", connectionId: "sqlite-1", database, + schema, sql: "", isExecuting: false, isCancelling: false, @@ -38,6 +41,30 @@ function queryTab(database: string): QueryTab { }; } +function damengConnection(database = "APPDB"): ConnectionConfig { + return { + id: "dameng-1", + name: "Dameng", + db_type: "dameng", + host: "127.0.0.1", + port: 5236, + username: "APP_USER", + password: "", + database, + }; +} + +function aiConfig(): AiConfig { + return { + provider: "openai", + apiKey: "", + authMethod: "api-key", + endpoint: "", + model: "test-model", + apiStyle: "responses", + }; +} + describe("SQLite AI context routing", () => { beforeEach(() => { vi.clearAllMocks(); @@ -61,3 +88,51 @@ describe("SQLite AI context routing", () => { expect(apiMock.listTables).toHaveBeenCalledWith("sqlite-1", "analytics", "analytics"); }); }); + +describe("Dameng AI context routing", () => { + beforeEach(() => { + vi.clearAllMocks(); + apiMock.listTables.mockResolvedValue([]); + apiMock.getColumns.mockResolvedValue([]); + apiMock.listIndexes.mockResolvedValue([]); + apiMock.listForeignKeys.mockResolvedValue([]); + apiMock.aiAgentStream.mockResolvedValue("done"); + }); + + it("keeps the connection database while routing metadata through the selected schema", async () => { + const tab = { ...queryTab("APPDB", "REPORTING"), connectionId: "dameng-1" }; + const context = await buildAiContext(tab, damengConnection()); + + expect(context.database).toBe("APPDB"); + expect(context.schema).toBe("REPORTING"); + expect(apiMock.listTables).toHaveBeenCalledWith("dameng-1", "APPDB", "REPORTING"); + + await runAgentStream( + { + config: aiConfig(), + action: "general", + mode: "agent", + instruction: "inspect orders", + context, + }, + [], + vi.fn(), + "session-1", + ); + expect(apiMock.aiAgentStream).toHaveBeenCalledWith("session-1", expect.any(Object), "dameng-1", "APPDB", "REPORTING", "dameng", expect.any(Function), "agent", false, undefined, undefined, undefined, undefined); + }); + + it("models the AI selector as a schema and chooses the connected user by default", () => { + const connection = damengConnection(); + expect(resolveAiNamespaceSelection(queryTab("APPDB", "REPORTING"), connection)).toEqual({ + kind: "schema", + value: "REPORTING", + }); + expect(resolveDefaultAiSchema(connection, ["ARCHIVE", "APP_USER", "REPORTING"])).toBe("APP_USER"); + expect(resolveDefaultAiSchema(connection, ["REPORTING", "ARCHIVE"])).toBe("REPORTING"); + }); + + it("does not change non-Dameng namespace behavior", () => { + expect(resolveAiDatabaseTarget(queryTab("analytics"), sqliteConnection())).toEqual({ database: "analytics" }); + }); +}); diff --git a/apps/desktop/src/lib/ai/ai.ts b/apps/desktop/src/lib/ai/ai.ts index 2130edb59f..9d32f15197 100644 --- a/apps/desktop/src/lib/ai/ai.ts +++ b/apps/desktop/src/lib/ai/ai.ts @@ -80,6 +80,8 @@ export interface AiContext { connectionName: string; databaseType: DatabaseType; database: string; + /** Selected schema when it is distinct from the connection database (for example Dameng). */ + schema?: string; currentSql: string; lastError?: string; lastResultPreview?: string; @@ -98,9 +100,15 @@ export interface AiRequestInput { allowWriteSql?: boolean; /** When allowWriteSql is true, the specific write SQL the user confirmed. */ confirmedWriteSql?: string; - /** Connection/database snapshot at confirmation time; verified at backend. */ + /** Connection/database/schema snapshot at confirmation time; verified at backend. */ confirmedConnectionId?: string; confirmedDatabase?: string; + confirmedSchema?: string; +} + +export interface AiNamespaceSelection { + kind: "database" | "schema"; + value: string; } export interface CustomPromptContext { @@ -189,6 +197,7 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa }, input.context.connectionId, input.context.database, + input.context.schema, input.context.databaseType, onEvent, input.mode || "ask", @@ -196,6 +205,7 @@ export async function runAgentStream(input: AiRequestInput, history: api.AiMessa input.confirmedWriteSql, input.confirmedConnectionId, input.confirmedDatabase, + input.confirmedSchema, ); } @@ -264,6 +274,7 @@ export function buildSystemPrompt(action: AiAction, context: AiContext, mode: Ai `Database type: ${context.databaseType}`, `Connection: ${context.connectionName}`, `Database: ${context.database}`, + context.schema ? `Selected schema: ${context.schema}` : "", schemaCoverageLine(context, isZh), "", `Current SQL:\n${context.currentSql.trim() || "(empty)"}`, @@ -435,7 +446,7 @@ export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig const maxIndexesPerTable = options.maxIndexesPerTable ?? 10; const maxFksPerTable = options.maxFksPerTable ?? 10; const databaseType = aiDatabaseTypeForConnection(connection); - const database = aiDatabaseNamespace(tab, connection); + const { database, schema } = resolveAiDatabaseTarget(tab, connection); const tables: AiSchemaTable[] = []; const tableKeys = new Set(); let truncated = false; @@ -551,6 +562,7 @@ export async function buildAiContext(tab: QueryTab, connection: ConnectionConfig connectionName: connection.name, databaseType, database, + schema, currentSql: currentCollectionName ?? tab.sql, lastError: extractLastError(tab.result), lastResultPreview: formatResultPreview(tab.result), @@ -604,7 +616,8 @@ async function resolveMentionedTableSchema(tab: QueryTab, connection: Connection } async function loadCandidateSchemas(tab: QueryTab, connection: ConnectionConfig): Promise { - const database = aiDatabaseNamespace(tab, connection); + const { database, schema } = resolveAiDatabaseTarget(tab, connection); + if (schema) return [schema]; if (isSchemaAware(aiDatabaseTypeForConnection(connection))) { const schemas = await api.listSchemas(tab.connectionId, database); return prioritizeSchemas(schemas); @@ -617,8 +630,37 @@ function aiDatabaseTypeForConnection(connection: ConnectionConfig): DatabaseType } function aiDatabaseNamespace(tab: QueryTab, connection: ConnectionConfig): string { + return resolveAiDatabaseTarget(tab, connection).database; +} + +export function resolveAiNamespaceSelection(tab: QueryTab, connection: ConnectionConfig): AiNamespaceSelection { + if (connection.db_type === "dameng") { + return { kind: "schema", value: tab.schema?.trim() || "" }; + } + return { kind: "database", value: tab.database || "" }; +} + +export function resolveDefaultAiSchema(connection: ConnectionConfig, schemaOptions: string[]): string | undefined { + if (connection.db_type !== "dameng") return undefined; + const username = connection.username.trim().toLowerCase(); + return schemaOptions.find((schema) => schema.trim().toLowerCase() === username) ?? schemaOptions[0]; +} + +/** + * Resolve the namespace used by an AI request without treating a Dameng schema + * selection as a connection database override. Dameng connections stay bound to + * their configured database while the query tab's selection scopes metadata and + * SQL execution through the schema parameter. + */ +export function resolveAiDatabaseTarget(tab: QueryTab, connection: ConnectionConfig): { database: string; schema?: string } { const database = tab.database || connection.database || "main"; - return connection.db_type === "sqlite" ? normalizeSqliteNamespace(database, connection) : database; + if (connection.db_type === "dameng") { + return { + database, + schema: resolveAiNamespaceSelection(tab, connection).value || undefined, + }; + } + return { database: connection.db_type === "sqlite" ? normalizeSqliteNamespace(database, connection) : database }; } function prioritizeSchemas(schemas: string[]): string[] { diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 601cb92486..971a9fcd39 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -1336,6 +1336,7 @@ export async function aiAgentStream( request: AiCompletionRequest, connectionId: string, database: string, + schema: string | undefined, dbType: string, onEvent: (event: import("@/lib/backend/tauri").AgentEvent) => void, mode?: string, @@ -1343,6 +1344,7 @@ export async function aiAgentStream( confirmedWriteSql?: string, confirmedConnectionId?: string, confirmedDatabase?: string, + confirmedSchema?: string, signal?: AbortSignal, ): Promise { const res = await fetch(apiUrl("/api/ai/agent-stream"), { @@ -1353,12 +1355,14 @@ export async function aiAgentStream( request, connectionId, database, + schema, dbType, mode: mode || "ask", allowWriteSql, confirmedWriteSql, confirmedConnectionId, confirmedDatabase, + confirmedSchema, }), signal, }); diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 804320448d..da66bdde27 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -431,6 +431,7 @@ export async function aiAgentStream( request: AiCompletionRequest, connectionId: string, database: string, + schema: string | undefined, dbType: string, onEvent: (event: AgentEvent) => void, mode?: string, @@ -438,6 +439,7 @@ export async function aiAgentStream( confirmedWriteSql?: string, confirmedConnectionId?: string, confirmedDatabase?: string, + confirmedSchema?: string, _signal?: AbortSignal, ): Promise { const unlisten: UnlistenFn = await listen("ai-agent-event", (event) => { @@ -452,12 +454,14 @@ export async function aiAgentStream( request, connectionId, database, + schema, dbType, mode, allowWriteSql, confirmedWriteSql, confirmedConnectionId, confirmedDatabase, + confirmedSchema, }); } catch (e) { unlisten(); diff --git a/crates/dbx-core/src/agent_loop.rs b/crates/dbx-core/src/agent_loop.rs index 7460346bfb..1e52dfacf7 100644 --- a/crates/dbx-core/src/agent_loop.rs +++ b/crates/dbx-core/src/agent_loop.rs @@ -76,6 +76,8 @@ pub struct AgentLoopContext { pub state: Arc, pub connection_id: String, pub database: String, + /// Selected schema that scopes Agent metadata and SQL execution. + pub schema: Option, pub db_type: DatabaseType, pub cli_mcp_server_command: Option, pub sql_permissions: agent_tools::AgentSqlPermissions, @@ -149,6 +151,7 @@ pub async fn run_agent_loop( connection_id: agent_ctx.connection_id.clone(), connection_name, database: agent_ctx.database.clone(), + schema: agent_ctx.schema.clone(), agent_mode: is_agent_mode, allow_writes: agent_ctx.sql_permissions.allow_writes, allow_dangerous: agent_ctx.sql_permissions.allow_dangerous, @@ -435,6 +438,7 @@ pub async fn run_agent_loop( let state2 = Arc::clone(&agent_ctx.state); let conn2 = agent_ctx.connection_id.clone(); let db2 = agent_ctx.database.clone(); + let schema2 = agent_ctx.schema.clone(); let db_type = agent_ctx.db_type; let sql_permissions = agent_ctx.sql_permissions.clone(); @@ -452,25 +456,39 @@ pub async fn run_agent_loop( }; // Run parallel group - let parallel_futures: Vec<_> = parallel_indices - .iter() - .map(|&i| { - let tc = make_tc(&collected_tool_calls[i]); - let state = Arc::clone(&state2); - let conn = conn2.clone(); - let db = db2.clone(); - let perms = sql_permissions.clone(); - async move { agent_tools::execute_tool(&tc, &state, &conn, &db, &db_type, perms).await } - }) - .collect(); + let parallel_futures: Vec<_> = + parallel_indices + .iter() + .map(|&i| { + let tc = make_tc(&collected_tool_calls[i]); + let state = Arc::clone(&state2); + let conn = conn2.clone(); + let db = db2.clone(); + let schema = schema2.clone(); + let perms = sql_permissions.clone(); + async move { + agent_tools::execute_tool(&tc, &state, &conn, &db, schema.as_deref(), &db_type, perms).await + } + }) + .collect(); let parallel_results = join_all(parallel_futures).await; // Run sequential group one-by-one let mut sequential_results = Vec::with_capacity(sequential_indices.len()); for &i in &sequential_indices { let tc = make_tc(&collected_tool_calls[i]); - sequential_results - .push(agent_tools::execute_tool(&tc, &state2, &conn2, &db2, &db_type, sql_permissions.clone()).await); + sequential_results.push( + agent_tools::execute_tool( + &tc, + &state2, + &conn2, + &db2, + schema2.as_deref(), + &db_type, + sql_permissions.clone(), + ) + .await, + ); } // Merge results back into original order @@ -866,7 +884,7 @@ async fn build_schema_prompt(agent_ctx: &AgentLoopContext, system_prompt: &str) &agent_ctx.state, &agent_ctx.connection_id, &agent_ctx.database, - "", + agent_ctx.schema.as_deref().unwrap_or(""), None, Some(50), // smaller limit for prompt injection None, @@ -879,6 +897,9 @@ async fn build_schema_prompt(agent_ctx: &AgentLoopContext, system_prompt: &str) Ok(tables) if !tables.is_empty() => { enriched.push_str("\n\n## Database Schema (for context — no tools available)\n"); enriched.push_str(&format!("Database: {}\n", agent_ctx.database)); + if let Some(schema) = agent_ctx.schema.as_deref() { + enriched.push_str(&format!("Schema: {schema}\n")); + } enriched.push_str("Tables:\n"); for t in &tables { enriched.push_str(&format!(" - {} ({})", t.name, t.table_type)); diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index 601fd3004e..721154721e 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -66,22 +66,27 @@ pub fn verify_confirmed_target( confirmed_write_sql: Option, confirmed_connection_id: Option, confirmed_database: Option, + confirmed_schema: Option, actual_connection_id: &str, actual_database: &str, + actual_schema: Option<&str>, ) -> (Option, Option) { let Some(ref confirmed_sql) = confirmed_write_sql else { return (allow_write_sql, confirmed_write_sql); }; // Only verify when a write SQL was actually confirmed. let target_mismatch = confirmed_connection_id.as_deref() != Some(actual_connection_id) - || confirmed_database.as_deref() != Some(actual_database); + || confirmed_database.as_deref() != Some(actual_database) + || confirmed_schema.as_deref() != actual_schema; if target_mismatch { log::warn!( - "Write-SQL grant voided: confirmed target (conn={:?}, db={:?}) does not match actual (conn={}, db={}).", + "Write-SQL grant voided: confirmed target (conn={:?}, db={:?}, schema={:?}) does not match actual (conn={}, db={}, schema={:?}).", confirmed_connection_id, confirmed_database, + confirmed_schema, actual_connection_id, actual_database, + actual_schema, ); // SQL can contain literals or credentials. Keep diagnostic visibility // behind the shared debug-only redaction boundary. @@ -313,21 +318,25 @@ pub async fn execute_tool( state: &Arc, connection_id: &str, database: &str, + default_schema: Option<&str>, db_type: &DatabaseType, sql_permissions: AgentSqlPermissions, ) -> ToolResult { let result = match tool_call.name.as_str() { - "list_tables" => execute_list_tables(tool_call, state, connection_id, database, db_type).await, - "get_columns" => execute_get_columns(tool_call, state, connection_id, database, db_type).await, + "list_tables" => execute_list_tables(tool_call, state, connection_id, database, default_schema, db_type).await, + "get_columns" => execute_get_columns(tool_call, state, connection_id, database, default_schema, db_type).await, "execute_query" => { - execute_execute_query(tool_call, state, connection_id, database, db_type, sql_permissions).await + execute_execute_query(tool_call, state, connection_id, database, default_schema, db_type, sql_permissions) + .await + } + "get_sample_data" => { + execute_get_sample_data(tool_call, state, connection_id, database, default_schema, db_type).await } - "get_sample_data" => execute_get_sample_data(tool_call, state, connection_id, database, db_type).await, "list_collections" => execute_list_collections(tool_call, state, connection_id, database, db_type).await, "browse_collection" => execute_browse_collection(tool_call, state, connection_id, database, db_type).await, "explain_query" => { let (text_result, explain_data) = - execute_explain_query(tool_call, state, connection_id, database, db_type).await; + execute_explain_query(tool_call, state, connection_id, database, default_schema, db_type).await; match text_result { Ok(content) => { return ToolResult { @@ -375,9 +384,10 @@ async fn execute_list_tables( state: &Arc, connection_id: &str, database: &str, + default_schema: Option<&str>, _db_type: &DatabaseType, ) -> Result { - let schema = tool_call.arguments.get("schema").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let schema = effective_schema(tool_call, default_schema).unwrap_or_default(); // Request one extra to detect whether more tables exist beyond the limit. let tables = crate::schema::list_tables_core( @@ -426,6 +436,7 @@ async fn execute_get_columns( state: &Arc, connection_id: &str, database: &str, + default_schema: Option<&str>, _db_type: &DatabaseType, ) -> Result { let table = tool_call @@ -447,7 +458,7 @@ async fn execute_get_columns( return Err(format!("Table name contains invalid characters: '{}'", table)); } - let schema = tool_call.arguments.get("schema").and_then(|v| v.as_str()).unwrap_or("").to_string(); + let schema = effective_schema(tool_call, default_schema).unwrap_or_default(); let columns = crate::schema::get_columns_core(state, connection_id, database, &schema, &table) .await @@ -529,6 +540,7 @@ async fn execute_execute_query( state: &Arc, connection_id: &str, database: &str, + default_schema: Option<&str>, db_type: &DatabaseType, sql_permissions: AgentSqlPermissions, ) -> Result { @@ -588,9 +600,16 @@ async fn execute_execute_query( client_session_id: client_session_id.map(str::to_string), ..Default::default() }; - let result = - crate::query::execute_sql_statement_with_options(state, connection_id, database, sql, None, None, options) - .await?; + let result = crate::query::execute_sql_statement_with_options( + state, + connection_id, + database, + sql, + default_schema, + None, + options, + ) + .await?; format_query_result_as_text(&result, limit) } @@ -647,6 +666,7 @@ async fn execute_get_sample_data( state: &Arc, connection_id: &str, database: &str, + default_schema: Option<&str>, db_type: &DatabaseType, ) -> Result { let table = @@ -659,7 +679,7 @@ async fn execute_get_sample_data( return Err(format!("Table name contains invalid characters: '{}'", table)); } - let schema = tool_call.arguments.get("schema").and_then(|v| v.as_str()).map(str::trim).filter(|s| !s.is_empty()); + let schema = effective_schema(tool_call, default_schema); let limit = tool_call .arguments .get("limit") @@ -669,7 +689,7 @@ async fn execute_get_sample_data( // Reuse the table-data builder so identifier quoting and row limiting follow // the active database instead of assuming PostgreSQL syntax. - let sql = build_sample_data_sql(db_type, schema, table, limit); + let sql = build_sample_data_sql(db_type, schema.as_deref(), table, limit); // Delegate to execute_execute_query with a synthetic tool call let synthetic_call = ToolCall { @@ -678,8 +698,16 @@ async fn execute_get_sample_data( arguments: serde_json::json!({ "sql": sql, "limit": limit }), provider_payload: None, }; - execute_execute_query(&synthetic_call, state, connection_id, database, db_type, AgentSqlPermissions::default()) - .await + execute_execute_query( + &synthetic_call, + state, + connection_id, + database, + schema.as_deref(), + db_type, + AgentSqlPermissions::default(), + ) + .await } fn build_sample_data_sql(db_type: &DatabaseType, schema: Option<&str>, table: &str, limit: usize) -> String { @@ -699,6 +727,7 @@ async fn execute_explain_query( state: &Arc, connection_id: &str, database: &str, + default_schema: Option<&str>, db_type: &DatabaseType, ) -> (Result, Option) { let sql = match tool_call.arguments.get("sql").and_then(|v| v.as_str()) { @@ -733,7 +762,7 @@ async fn execute_explain_query( state, connection_id, Some(database), - None, + default_schema, sql, Some("explain"), ) @@ -764,7 +793,7 @@ async fn execute_explain_query( connection_id, database, &explain_sql, - None, + default_schema, None, options, ) @@ -784,6 +813,20 @@ async fn execute_explain_query( (Ok(text), explain_data) } +/// A selected context schema is authoritative for an Agent run. If none was +/// selected, keep the existing per-tool schema parameter behavior. +fn effective_schema(tool_call: &ToolCall, default_schema: Option<&str>) -> Option { + default_schema.map(str::trim).filter(|schema| !schema.is_empty()).map(ToOwned::to_owned).or_else(|| { + tool_call + .arguments + .get("schema") + .and_then(|value| value.as_str()) + .map(str::trim) + .filter(|schema| !schema.is_empty()) + .map(ToOwned::to_owned) + }) +} + /// Execute list_collections tool (vector databases). async fn execute_list_collections( _tool_call: &ToolCall, @@ -933,6 +976,114 @@ async fn resolve_chroma_collection_uuid( #[cfg(test)] mod tests { use super::*; + #[cfg(unix)] + use crate::connection::PoolKind; + #[cfg(unix)] + use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec}; + #[cfg(unix)] + use crate::models::connection::{default_redis_key_separator, ConnectionConfig}; + #[cfg(unix)] + use crate::storage::Storage; + + #[cfg(unix)] + async fn spawn_recording_agent(record_path: &std::path::Path) -> (AgentDriverClient, tempfile::NamedTempFile) { + use std::io::Write; + + let mut script = tempfile::NamedTempFile::new().unwrap(); + write!( + script, + r#"import json +import sys + +record_path = sys.argv[1] +print(json.dumps({{"ready": True}}), flush=True) +for line in sys.stdin: + request = json.loads(line) + with open(record_path, "a", encoding="utf-8") as record: + record.write(json.dumps(request) + "\n") + result = {{ + "columns": [], + "column_types": [], + "column_sortables": [], + "rows": [], + "affected_rows": 1, + "execution_time_ms": 0, + "truncated": False, + "session_id": None, + "has_more": False + }} + print(json.dumps({{"jsonrpc": "2.0", "id": request["id"], "result": result}}), flush=True) +"# + ) + .unwrap(); + script.flush().unwrap(); + + let client = AgentDriverClient::spawn( + AgentLaunchSpec::new("python3") + .with_args([script.path().to_string_lossy().to_string(), record_path.to_string_lossy().to_string()]), + ) + .await + .unwrap(); + (client, script) + } + + #[cfg(unix)] + fn dameng_test_connection() -> ConnectionConfig { + ConnectionConfig { + id: "dameng-1".to_string(), + name: "Dameng".to_string(), + note: String::new(), + db_type: DatabaseType::Dameng, + driver_profile: None, + driver_label: None, + url_params: None, + agent_java_options: Vec::new(), + host: "localhost".to_string(), + port: 5236, + username: "APP_USER".to_string(), + password: String::new(), + database: Some("APPDB".to_string()), + visible_databases: None, + visible_schemas: None, + show_system_schemas: false, + attached_databases: Vec::new(), + init_script: None, + color: None, + transport_layers: Vec::new(), + connect_timeout_secs: 10, + query_timeout_secs: 30, + idle_timeout_secs: 60, + keepalive_interval_secs: 30, + ssl: false, + ca_cert_path: String::new(), + client_cert_path: String::new(), + client_key_path: String::new(), + sysdba: false, + oracle_connection_type: None, + connection_string: None, + redis_connection_mode: None, + redis_sentinel_master: String::new(), + redis_sentinel_nodes: String::new(), + redis_sentinel_username: String::new(), + redis_sentinel_password: String::new(), + redis_sentinel_tls: false, + redis_cluster_nodes: String::new(), + redis_key_separator: default_redis_key_separator(), + redis_scan_page_size: None, + redis_database_aliases: Default::default(), + etcd_endpoints: String::new(), + gbase_server: String::new(), + informix_server: String::new(), + external_config: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + one_time: false, + read_only: false, + is_production: false, + production_databases: vec![], + database_info: None, + } + } #[test] fn vector_read_only_tools_do_not_include_collection_browsing() { @@ -1004,6 +1155,83 @@ mod tests { ); } + #[test] + fn selected_schema_overrides_the_tool_schema_argument() { + let call = ToolCall { + id: "call-1".to_string(), + name: "list_tables".to_string(), + arguments: serde_json::json!({ "schema": "OTHER" }), + provider_payload: None, + }; + + assert_eq!(effective_schema(&call, Some("REPORTING")).as_deref(), Some("REPORTING")); + assert_eq!(effective_schema(&call, None).as_deref(), Some("OTHER")); + } + + #[cfg(unix)] + #[tokio::test] + async fn dameng_agent_queries_and_confirmed_writes_use_selected_schema() { + let temp_dir = tempfile::tempdir().unwrap(); + let record_path = temp_dir.path().join("agent-requests.jsonl"); + let (client, _script) = spawn_recording_agent(&record_path).await; + let storage = Storage::open(&temp_dir.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let connection = dameng_test_connection(); + state.configs.write().await.insert(connection.id.clone(), connection); + state.connections.write().await.insert("dameng-1:APPDB".to_string(), PoolKind::agent(client)); + + let read = ToolCall { + id: "read".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": "SELECT * FROM orders" }), + provider_payload: None, + }; + let read_result = execute_tool( + &read, + &state, + "dameng-1", + "APPDB", + Some("REPORTING"), + &DatabaseType::Dameng, + AgentSqlPermissions::default(), + ) + .await; + assert!(!read_result.is_error, "{}", read_result.content); + + let confirmed_sql = "DELETE FROM orders WHERE id = 1"; + let write = ToolCall { + id: "write".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": confirmed_sql }), + provider_payload: None, + }; + let write_result = execute_tool( + &write, + &state, + "dameng-1", + "APPDB", + Some("REPORTING"), + &DatabaseType::Dameng, + confirmed_write_sql_permissions(false, true, Some(confirmed_sql.to_string())), + ) + .await; + assert!(!write_result.is_error, "{}", write_result.content); + + let requests = std::fs::read_to_string(&record_path) + .unwrap() + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .filter(|request| request["method"] == "execute_query") + .collect::>(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0]["params"]["sql"], "SELECT * FROM orders"); + assert_eq!(requests[1]["params"]["sql"], confirmed_sql); + for request in requests { + assert_eq!(request["params"]["database"], "APPDB"); + assert_eq!(request["params"]["schema"], "REPORTING"); + } + } + #[test] fn build_browse_query_qdrant() { let q = build_browse_query(&DatabaseType::Qdrant, "articles", "", 10).unwrap(); diff --git a/crates/dbx-core/src/ai_claude_code_cli.rs b/crates/dbx-core/src/ai_claude_code_cli.rs index 5234e2e2ef..6b9019d388 100644 --- a/crates/dbx-core/src/ai_claude_code_cli.rs +++ b/crates/dbx-core/src/ai_claude_code_cli.rs @@ -588,6 +588,7 @@ mod tests { connection_id: "conn-1".to_string(), connection_name: "local".to_string(), database: "demo".to_string(), + schema: None, agent_mode: true, allow_writes: false, allow_dangerous: false, diff --git a/crates/dbx-core/src/ai_cli_agent.rs b/crates/dbx-core/src/ai_cli_agent.rs index 302dd6283d..e483ba9c76 100644 --- a/crates/dbx-core/src/ai_cli_agent.rs +++ b/crates/dbx-core/src/ai_cli_agent.rs @@ -14,6 +14,7 @@ pub struct CliAgentRunOptions { pub connection_id: String, pub connection_name: String, pub database: String, + pub schema: Option, pub agent_mode: bool, pub allow_writes: bool, pub allow_dangerous: bool, @@ -71,6 +72,9 @@ pub fn dbx_mcp_scope_env(options: &CliAgentRunOptions) -> Vec<(&'static str, Str ("DBX_MCP_SCOPE_CONNECTION_NAME", options.connection_name.clone()), ("DBX_MCP_SCOPE_DATABASE", options.database.clone()), ]; + if let Some(schema) = options.schema.as_deref().filter(|schema| !schema.trim().is_empty()) { + env.push(("DBX_MCP_SCOPE_SCHEMA", schema.to_string())); + } if let Some(ref sql) = options.confirmed_write_sql { env.push(("DBX_MCP_CONFIRMED_WRITE_SQL", sql.clone())); } @@ -84,9 +88,10 @@ mod scope_env_tests { #[test] fn confirmed_write_sql_is_passed_to_the_scoped_mcp_subprocess() { let options = CliAgentRunOptions { - connection_id: "connection-1".to_string(), - connection_name: "staging".to_string(), - database: "app".to_string(), + connection_id: "dameng-1".to_string(), + connection_name: "Dameng".to_string(), + database: "APPDB".to_string(), + schema: Some("REPORTING".to_string()), agent_mode: true, allow_writes: true, allow_dangerous: true, @@ -98,6 +103,8 @@ mod scope_env_tests { assert!(env.contains(&("DBX_MCP_CONFIRMED_WRITE_SQL", "DELETE FROM sessions WHERE id = 7".to_string()))); assert!(env.contains(&("DBX_MCP_ALLOW_WRITES", "1".to_string()))); assert!(env.contains(&("DBX_MCP_ALLOW_DANGEROUS_SQL", "1".to_string()))); + assert!(env.contains(&("DBX_MCP_SCOPE_DATABASE", "APPDB".to_string()))); + assert!(env.contains(&("DBX_MCP_SCOPE_SCHEMA", "REPORTING".to_string()))); } } diff --git a/crates/dbx-core/src/ai_codex_cli.rs b/crates/dbx-core/src/ai_codex_cli.rs index 29df8d7517..f3fbe95058 100644 --- a/crates/dbx-core/src/ai_codex_cli.rs +++ b/crates/dbx-core/src/ai_codex_cli.rs @@ -1045,6 +1045,7 @@ mod tests { connection_id: "conn-1".to_string(), connection_name: "local".to_string(), database: "demo".to_string(), + schema: None, agent_mode: true, allow_writes: false, allow_dangerous: false, diff --git a/crates/dbx-core/src/ai_pi_agent_cli.rs b/crates/dbx-core/src/ai_pi_agent_cli.rs index 28e4697331..8ea32688f1 100644 --- a/crates/dbx-core/src/ai_pi_agent_cli.rs +++ b/crates/dbx-core/src/ai_pi_agent_cli.rs @@ -825,6 +825,7 @@ mod tests { connection_id: "connection-1".to_string(), connection_name: "Test connection".to_string(), database: "dbx_test".to_string(), + schema: Some("reporting".to_string()), agent_mode: true, allow_writes: true, allow_dangerous: false, @@ -852,6 +853,7 @@ mod tests { assert_eq!(env.get("DBX_MCP_SCOPE_CONNECTION_ID").map(String::as_str), Some("connection-1")); assert_eq!(env.get("DBX_MCP_SCOPE_CONNECTION_NAME").map(String::as_str), Some("Test connection")); assert_eq!(env.get("DBX_MCP_SCOPE_DATABASE").map(String::as_str), Some("dbx_test")); + assert_eq!(env.get("DBX_MCP_SCOPE_SCHEMA").map(String::as_str), Some("reporting")); let enabled_tools = serde_json::from_str::>(env.get("DBX_PI_ENABLED_TOOLS").unwrap()).unwrap(); assert!(enabled_tools.iter().any(|tool| tool == "dbx_execute_query")); diff --git a/crates/dbx-mcp/src/backend.rs b/crates/dbx-mcp/src/backend.rs index b11fc4d1f8..fcd31693ae 100644 --- a/crates/dbx-mcp/src/backend.rs +++ b/crates/dbx-mcp/src/backend.rs @@ -322,9 +322,19 @@ impl DbxBackend for LocalBackend { arguments: Value, permissions: AgentSqlPermissions, ) -> ToolResult { + let schema = arguments.get("schema").and_then(|value| value.as_str()).map(ToOwned::to_owned); let call = ToolCall { id: format!("mcp-{tool_name}"), name: tool_name.to_string(), arguments, provider_payload: None }; - agent_tools::execute_tool(&call, &self.state, &connection.id, database, &connection.db_type, permissions).await + agent_tools::execute_tool( + &call, + &self.state, + &connection.id, + database, + schema.as_deref(), + &connection.db_type, + permissions, + ) + .await } async fn execute_query( diff --git a/crates/dbx-mcp/src/server.rs b/crates/dbx-mcp/src/server.rs index 6626be6d58..200d0f75d5 100644 --- a/crates/dbx-mcp/src/server.rs +++ b/crates/dbx-mcp/src/server.rs @@ -160,6 +160,7 @@ pub struct McpScope { pub connection_ids: Vec, pub connection_name: Option, pub database: Option, + pub schema: Option, } struct ResolvedConnection { @@ -179,11 +180,12 @@ impl McpScope { connection_ids, connection_name: non_empty_env("DBX_MCP_SCOPE_CONNECTION_NAME"), database: non_empty_env("DBX_MCP_SCOPE_DATABASE"), + schema: non_empty_env("DBX_MCP_SCOPE_SCHEMA"), } } fn enabled(&self) -> bool { - self.connection_scope_enabled() || self.database.is_some() + self.connection_scope_enabled() || self.database.is_some() || self.schema.is_some() } fn connection_scope_enabled(&self) -> bool { @@ -257,7 +259,11 @@ impl DbxMcpServer { Ok(database) => database, Err(error) => return error, }; - match self.backend.list_tables(&resolved.connection, &database, &request.schema.unwrap_or_default()).await { + let schema = match self.resolve_schema(request.schema) { + Ok(schema) => schema, + Err(error) => return error, + }; + match self.backend.list_tables(&resolved.connection, &database, &schema).await { Ok(tables) if tables.is_empty() => text("No tables found."), Ok(tables) => text( tables @@ -287,11 +293,11 @@ impl DbxMcpServer { Ok(database) => database, Err(error) => return error, }; - match self - .backend - .get_columns(&resolved.connection, &database, &request.schema.unwrap_or_default(), &request.table) - .await - { + let schema = match self.resolve_schema(request.schema) { + Ok(schema) => schema, + Err(error) => return error, + }; + match self.backend.get_columns(&resolved.connection, &database, &schema, &request.table).await { Ok(columns) if columns.is_empty() => text("No columns found."), Ok(columns) => text(format_columns(&columns)), Err(error) => tool_error("TABLE_DESCRIPTION_ERROR", error), @@ -375,6 +381,9 @@ impl DbxMcpServer { Err(error) => return error, }; let mut arguments = json!({ "sql": request.sql, "limit": 100 }); + if let Some(schema) = self.scope.schema.as_deref() { + arguments["schema"] = json!(schema); + } if let Some(session) = &session { arguments["client_session_id"] = json!(session.client_session_id); } @@ -523,7 +532,10 @@ impl DbxMcpServer { Ok(database) => database, Err(error) => return error, }; - let schema = request.schema.unwrap_or_default(); + let schema = match self.resolve_schema(request.schema) { + Ok(schema) => schema, + Err(error) => return error, + }; let max_tables = request.max_tables.unwrap_or(8).clamp(1, 20); let available = match self.backend.list_tables(connection, &database, &schema).await { Ok(tables) => tables, @@ -658,6 +670,10 @@ impl DbxMcpServer { Ok(database) => database, Err(error) => return error, }; + let schema = match self.resolve_schema(request.schema) { + Ok(schema) => schema, + Err(error) => return error, + }; match self .backend .bridge_request( @@ -667,7 +683,7 @@ impl DbxMcpServer { "connection_name": connection.name, "table": request.table, "database": database, - "schema": request.schema, + "schema": schema, }), ) .await @@ -765,6 +781,25 @@ impl DbxMcpServer { Ok(requested.or_else(|| connection.database.clone()).unwrap_or_default()) } + /// Resolve the schema for scoped CLI agents. A selected schema is a hard + /// bound, matching the existing database scope behavior. + #[allow(clippy::result_large_err)] + fn resolve_schema(&self, requested: Option) -> Result { + let requested = requested.map(|schema| schema.trim().to_string()).filter(|schema| !schema.is_empty()); + if let Some(scoped) = self.scope.schema.as_deref() { + if let Some(requested) = requested.as_deref() { + if requested != scoped { + return Err(tool_error( + "SCHEMA_OUT_OF_SCOPE", + format!("Schema \"{requested}\" is outside the scoped schema \"{scoped}\"."), + )); + } + } + return Ok(scoped.to_string()); + } + Ok(requested.unwrap_or_default()) + } + // CallToolResult is the rmcp wire response type; keeping it unboxed avoids conversions at every tool boundary. #[allow(clippy::result_large_err)] fn resolve_redis_database( @@ -1351,6 +1386,7 @@ mod tests { connection_ids: vec!["first".to_string()], connection_name: Some("scope-name".to_string()), database: None, + schema: None, }; assert!(scope.matches(&first)); @@ -1377,6 +1413,26 @@ mod tests { assert!(!names.iter().any(|name| name == "dbx_execute_and_show")); } + #[test] + fn schema_scope_is_a_hard_bound() { + let dameng = connection("dameng-1", "Dameng", "dameng", "APPDB"); + let server = DbxMcpServer::with_runtime_options( + Arc::new(FakeBackend::default()), + McpScope { + database: Some("APPDB".to_string()), + schema: Some("REPORTING".to_string()), + ..Default::default() + }, + false, + ); + + assert_eq!(server.resolve_database(None, &dameng).unwrap(), "APPDB"); + assert_eq!(server.resolve_schema(None).unwrap(), "REPORTING"); + assert_eq!(server.resolve_schema(Some("REPORTING".to_string())).unwrap(), "REPORTING"); + let error = server.resolve_schema(Some("APP_USER".to_string())).unwrap_err(); + assert!(result_text(&error).contains("SCHEMA_OUT_OF_SCOPE")); + } + #[test] fn redis_database_scope_fails_closed_and_cannot_be_overridden() { let redis = connection("redis", "redis", "redis", "1"); diff --git a/crates/dbx-web/src/routes/ai.rs b/crates/dbx-web/src/routes/ai.rs index 975dd6412c..78315ed00c 100644 --- a/crates/dbx-web/src/routes/ai.rs +++ b/crates/dbx-web/src/routes/ai.rs @@ -92,6 +92,8 @@ pub struct AiAgentStreamRequest { pub request: AiCompletionRequest, pub connection_id: String, pub database: String, + #[serde(default)] + pub schema: Option, pub db_type: String, /// Agent mode: "ask" (read-only tools) or "agent" (all tools including execute_query). /// Defaults to "ask" if not provided. @@ -107,6 +109,8 @@ pub struct AiAgentStreamRequest { pub confirmed_connection_id: Option, #[serde(default)] pub confirmed_database: Option, + #[serde(default)] + pub confirmed_schema: Option, } fn default_agent_mode() -> String { @@ -424,8 +428,10 @@ pub async fn ai_agent_stream( body.confirmed_write_sql, body.confirmed_connection_id, body.confirmed_database, + body.confirmed_schema, &body.connection_id, &body.database, + body.schema.as_deref(), ); // Writes are only allowed when a specific SQL statement was confirmed — // an empty confirmed_write_sql is treated as "no confirmation" so the @@ -439,6 +445,7 @@ pub async fn ai_agent_stream( state: state.app.clone(), connection_id: body.connection_id, database: body.database, + schema: body.schema, db_type: parsed_db_type, cli_mcp_server_command: None, sql_permissions, diff --git a/src-tauri/src/commands/ai.rs b/src-tauri/src/commands/ai.rs index 8d74a34efb..3697177510 100644 --- a/src-tauri/src/commands/ai.rs +++ b/src-tauri/src/commands/ai.rs @@ -137,12 +137,14 @@ pub async fn ai_agent_stream( request: AiCompletionRequest, connection_id: String, database: String, + schema: Option, db_type: String, mode: Option, allow_write_sql: Option, confirmed_write_sql: Option, confirmed_connection_id: Option, confirmed_database: Option, + confirmed_schema: Option, ) -> Result { let mut request = resolve_cli_provider_request(request); merge_global_max_retries( @@ -179,8 +181,10 @@ pub async fn ai_agent_stream( confirmed_write_sql, confirmed_connection_id, confirmed_database, + confirmed_schema, &connection_id, &database, + schema.as_deref(), ); // Explicit confirmation grants write access only to this agent run, never to // production. Writes are only allowed when a specific SQL statement was @@ -195,6 +199,7 @@ pub async fn ai_agent_stream( state: state.inner().clone(), connection_id, database, + schema, db_type: parsed_db_type, cli_mcp_server_command, sql_permissions, @@ -329,8 +334,10 @@ mod tests { Some("DELETE FROM users WHERE id = 1".to_string()), Some("conn-1".to_string()), Some("app".to_string()), + None, "conn-1", "app", + None, ); assert_eq!(allow, Some(true)); assert_eq!(confirmed, Some("DELETE FROM users WHERE id = 1".to_string())); @@ -343,8 +350,10 @@ mod tests { Some("DELETE FROM users WHERE id = 1".to_string()), Some("conn-staging".to_string()), Some("app".to_string()), + None, "conn-production", "app", + None, ); assert_eq!(allow, Some(false)); assert_eq!(confirmed, None); @@ -357,8 +366,10 @@ mod tests { Some("DELETE FROM users WHERE id = 1".to_string()), Some("conn-1".to_string()), Some("staging".to_string()), + None, "conn-1", "production", + None, ); assert_eq!(allow, Some(false)); assert_eq!(confirmed, None); @@ -373,8 +384,10 @@ mod tests { None, Some("conn-staging".to_string()), Some("staging".to_string()), + None, "conn-production", "production", + None, ); assert_eq!(allow, Some(false)); assert_eq!(confirmed, None); @@ -390,8 +403,26 @@ mod tests { Some("DELETE FROM users WHERE id = 1".to_string()), None, None, + None, "conn-1", "app", + None, + ); + assert_eq!(allow, Some(false)); + assert_eq!(confirmed, None); + } + + #[test] + fn verify_confirmed_target_rejects_mismatched_schema() { + let (allow, confirmed) = dbx_core::agent_tools::verify_confirmed_target( + Some(true), + Some("DELETE FROM users WHERE id = 1".to_string()), + Some("conn-1".to_string()), + Some("APPDB".to_string()), + Some("APP_USER".to_string()), + "conn-1", + "APPDB", + Some("REPORTING"), ); assert_eq!(allow, Some(false)); assert_eq!(confirmed, None);