From 12da615daafa8c5bd157c2a3182ad71bb1a3064a Mon Sep 17 00:00:00 2001 From: Abeautifulsnow Date: Fri, 21 Aug 2026 10:36:17 +0800 Subject: [PATCH] fix(ai): require exact SQL confirmation for agent writes --- .../src/components/editor/AiAssistant.vue | 70 ++++- ...ssistant.writeConfirmationProposal.spec.ts | 47 ++++ apps/desktop/src/i18n/locales/en.ts | 9 + apps/desktop/src/i18n/locales/es.ts | 9 + apps/desktop/src/i18n/locales/it.ts | 9 + apps/desktop/src/i18n/locales/ja.ts | 9 + apps/desktop/src/i18n/locales/ko.ts | 9 + apps/desktop/src/i18n/locales/pt-BR.ts | 9 + apps/desktop/src/i18n/locales/zh-CN.ts | 9 + apps/desktop/src/i18n/locales/zh-TW.ts | 9 + .../ai/aiWriteSqlConfirmation.spec.ts | 53 +++- apps/desktop/src/lib/ai/aiProposalDetect.ts | 35 ++- apps/desktop/src/lib/backend/tauri.ts | 4 +- crates/dbx-core/src/agent_events.rs | 6 + crates/dbx-core/src/agent_loop.rs | 257 +++++++++++++++++- crates/dbx-core/src/agent_tools.rs | 112 +++++++- 16 files changed, 626 insertions(+), 30 deletions(-) create mode 100644 apps/desktop/src/components/editor/__tests__/AiAssistant.writeConfirmationProposal.spec.ts diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue index 2c73182571..505930773a 100644 --- a/apps/desktop/src/components/editor/AiAssistant.vue +++ b/apps/desktop/src/components/editor/AiAssistant.vue @@ -132,7 +132,7 @@ import { parseExplainResult, parseOracleExplainText, type ParsedExplainPlan } fr import { copyToClipboard } from "@/lib/common/clipboard"; import { AI_TABLE_MENTION_CANDIDATE_LIMIT, AI_TABLE_MENTION_SCHEMA_LIMIT, filterAiTableMentionCandidates, formatAiTableMention, parseAiTableMentions, type AiTableMention } from "@/lib/ai/aiTableMentions"; import { isAiPromptImeCompositionEvent, shouldSubmitAiPromptOnKeydown } from "@/lib/ai/aiPromptKeyboard"; -import { looksLikeActionProposal, containsChinese, looksLikeWriteSqlProposal, shouldGrantWriteSqlOnShortAffirmative } from "@/lib/ai/aiProposalDetect"; +import { isActionableWriteProposalMessage, isActionableWriteSqlProposal, looksLikeActionProposal, looksLikeWriteSqlProposal, shouldGrantWriteSqlOnShortAffirmative } from "@/lib/ai/aiProposalDetect"; import { visibleToActualIndex } from "@/lib/ai/aiMessageEdit"; import { shouldShowReasoningCharCount, reasoningCharCountClass } from "@/lib/ai/aiReasoningPresentation"; import { saveTextFile } from "@/lib/export/saveTextFile"; @@ -199,8 +199,11 @@ interface ChatMessage { reasoning?: string; isThinking?: boolean; agentSteps?: AiAgentStepItem[]; - /** Hidden system-generated context summary; not rendered in chat UI but included in LLM history. */ - kind?: "contextSummary"; + /** Hidden system-generated context summary; not rendered in chat UI but included in LLM history. + * `writeSqlConfirmation` / `productionWriteBlocked` mark backend-generated, + * localized confirmation/block messages so the UI does not re-detect them + * from text phrasing. */ + kind?: "contextSummary" | "writeSqlConfirmation" | "productionWriteBlocked"; /** Per-message token stats from the last agent run; ephemeral, not persisted. */ tokens?: { input: number; output: number }; } @@ -906,7 +909,15 @@ const proposalConfirmMessage = computed(() => { if (msg.kind === "contextSummary") continue; if (msg.role !== "assistant") return null; if (!msg.content) return null; - return looksLikeActionProposal(msg.content) ? msg : null; + // Backend-generated write confirmations are localized, so the English/Chinese + // phrase detectors cannot recognize them; the `kind` marker plus one exact + // SQL block is the structural proof of actionability. + if (msg.kind === "writeSqlConfirmation") return extractSingleSqlCodeBlock(msg.content) ? msg : null; + if (!looksLikeActionProposal(msg.content)) return null; + // A generic write question cannot authorize a later, unseen tool call. + // Hide its action bar until the assistant displays one exact SQL block. + if (looksLikeWriteSqlProposal(msg.content) && !isActionableWriteSqlProposal(msg.content)) return null; + return msg; } return null; }); @@ -940,17 +951,17 @@ function sendProposalReply(positive: boolean) { if (isGenerating.value) return; const target = proposalConfirmMessage.value; if (!target) return; - if (positive && productionContext.value.active && looksLikeWriteSqlProposal(target.content)) { + const isWriteConfirmation = isActionableWriteProposalMessage(target); + if (positive && productionContext.value.active && (target.kind === "writeSqlConfirmation" || looksLikeWriteSqlProposal(target.content))) { const sql = extractFirstSqlCodeBlock(target.content); if (sql) emit("replaceSql", sql); toast(t("production.aiReviewRequired"), 5000); return; } - const isZh = containsChinese(target.content || ""); - const replyZh = positive ? "请执行上面你刚提议的操作，不要再反问确认。" : "不用执行上面提到的操作，继续当前对话。"; - const replyEn = positive ? "Execute the action you just proposed above; do not ask for confirmation again." : "Do not execute the action mentioned above; continue the current conversation."; - prompt.value = isZh ? replyZh : replyEn; - if (positive && assistantMode.value === "agent" && looksLikeWriteSqlProposal(target.content)) { + // Write confirmations carry the exact-SQL reply; other action proposals keep + // the generic wording so the model does not receive SQL-specific instructions. + prompt.value = positive ? (isWriteConfirmation ? t("ai.writeSqlConfirmationReplyYes") : t("ai.proposalConfirmReplyYes")) : isWriteConfirmation ? t("ai.writeSqlConfirmationReplyNo") : t("ai.proposalConfirmReplyNo"); + if (positive && assistantMode.value === "agent" && isWriteConfirmation) { confirmedWriteSqlText = extractSingleSqlCodeBlock(target.content); if (confirmedWriteSqlText) { allowWriteSqlForNextRun = true; @@ -1141,6 +1152,31 @@ function appendAssistantDelta(assistantIdx: number, delta: string) { scheduleAssistantDeltaFlush(assistantIdx); } +function replaceAssistantText(assistantIdx: number, content: string) { + // A model can stream prose or a code block before returning a write tool call. + // Discard that partial output so the confirmation detector sees exactly one SQL block. + if (assistantDeltaFrame !== null) { + cancelAnimationFrame(assistantDeltaFrame); + assistantDeltaFrame = null; + } + pendingAssistantDelta = ""; + pendingAssistantReasoning = ""; + pendingAssistantIndex = -1; + const msg = messages.value[assistantIdx]; + if (!msg) return; + msg.content = content; + msg.reasoning = undefined; + msg.isThinking = false; +} + +function writeSqlConfirmationText(sql: string): string { + return `${t("ai.writeSqlConfirmationRequired")}\n\n\`\`\`sql\n${sql.trim()}\n\`\`\`\n\n${t("ai.writeSqlConfirmationQuestion")}`; +} + +function productionWriteBlockedText(sql: string): string { + return `${t("ai.productionWriteBlocked")}\n\n\`\`\`sql\n${sql.trim()}\n\`\`\``; +} + function appendAssistantReasoning(assistantIdx: number, delta: string) { pendingAssistantReasoning += delta; scheduleAssistantDeltaFlush(assistantIdx); @@ -2449,6 +2485,16 @@ async function send() { if (event.type === "text_delta" && event.delta) { appendAssistantDelta(assistantIdx, event.delta); } + if (event.type === "write_sql_confirmation_required") { + replaceAssistantText(assistantIdx, writeSqlConfirmationText(event.sql)); + const msg = messages.value[assistantIdx]; + if (msg) msg.kind = "writeSqlConfirmation"; + } + if (event.type === "production_write_blocked") { + replaceAssistantText(assistantIdx, productionWriteBlockedText(event.sql)); + const msg = messages.value[assistantIdx]; + if (msg) msg.kind = "productionWriteBlocked"; + } if (event.type === "reasoning_delta" && event.delta) { appendAssistantReasoning(assistantIdx, event.delta); } @@ -3293,11 +3339,11 @@ async function openExternalUrl(url: string) {
diff --git a/apps/desktop/src/components/editor/__tests__/AiAssistant.writeConfirmationProposal.spec.ts b/apps/desktop/src/components/editor/__tests__/AiAssistant.writeConfirmationProposal.spec.ts new file mode 100644 index 0000000000..adc2a69999 --- /dev/null +++ b/apps/desktop/src/components/editor/__tests__/AiAssistant.writeConfirmationProposal.spec.ts @@ -0,0 +1,47 @@ +import { readFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; + +const source = readFileSync(new URL("../AiAssistant.vue", import.meta.url), "utf8"); +const localePaths = ["en", "es", "it", "ja", "ko", "pt-BR", "zh-CN", "zh-TW"]; + +describe("AI assistant write confirmation replacement", () => { + it("renders a localized confirmation from the backend's semantic event", () => { + expect(source).toContain("function replaceAssistantText(assistantIdx: number, content: string)"); + expect(source).toContain('pendingAssistantDelta = "";'); + expect(source).toContain("msg.content = content;"); + expect(source).toContain("function writeSqlConfirmationText(sql: string)"); + expect(source).toContain('event.type === "write_sql_confirmation_required"'); + expect(source).toContain("replaceAssistantText(assistantIdx, writeSqlConfirmationText(event.sql));"); + expect(source).toContain('event.type === "production_write_blocked"'); + }); + + it("does not render an executable write confirmation without exact SQL", () => { + expect(source).toContain("looksLikeWriteSqlProposal(msg.content) && !isActionableWriteSqlProposal(msg.content)"); + expect(source).toContain('if (positive && assistantMode.value === "agent" && isWriteConfirmation)'); + }); + + it("routes generic action replies separately from write-SQL replies", () => { + expect(source).toContain('isWriteConfirmation ? t("ai.writeSqlConfirmationReplyYes") : t("ai.proposalConfirmReplyYes")'); + expect(source).toContain('isWriteConfirmation ? t("ai.writeSqlConfirmationReplyNo") : t("ai.proposalConfirmReplyNo")'); + }); + + it("marks backend-generated confirmations as structurally actionable", () => { + expect(source).toContain('msg.kind = "writeSqlConfirmation";'); + expect(source).toContain('msg.kind = "productionWriteBlocked";'); + expect(source).toContain('if (msg.kind === "writeSqlConfirmation") return extractSingleSqlCodeBlock(msg.content) ? msg : null;'); + expect(source).toContain("isActionableWriteProposalMessage(msg)"); + }); + + it("provides the write-confirmation wording in every supported locale", () => { + for (const locale of localePaths) { + const messages = readFileSync(new URL(`../../../i18n/locales/${locale}.ts`, import.meta.url), "utf8"); + expect(messages).toContain("writeSqlConfirmationRequired:"); + expect(messages).toContain("writeSqlConfirmationQuestion:"); + expect(messages).toContain("writeSqlConfirmYes:"); + expect(messages).toContain("writeSqlConfirmNo:"); + expect(messages).toContain("productionWriteBlocked:"); + expect(messages).toContain("proposalConfirmReplyYes:"); + expect(messages).toContain("proposalConfirmReplyNo:"); + } + }); +}); diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 658ec7c709..cc827e2a49 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -2487,6 +2487,15 @@ export default { truncated: "Context truncated", contextSummary: "{database} · {tables} tables", autoSqlBlocked: "The AI-generated SQL looked too risky to auto-execute. Review it manually before running.", + writeSqlConfirmationRequired: "This SQL changes data or schema. Review it carefully before execution.", + writeSqlConfirmationQuestion: "Execute this SQL once?", + writeSqlConfirmYes: "Confirm execution", + writeSqlConfirmNo: "Cancel", + writeSqlConfirmationReplyYes: "Execute the exact SQL shown above once. Do not ask for confirmation again.", + writeSqlConfirmationReplyNo: "Do not execute the SQL shown above. Continue the conversation.", + productionWriteBlocked: "AI cannot execute data or schema changes on a production database. Review this SQL and run it manually in DBX.", + proposalConfirmReplyYes: "Execute the action you just proposed above; do not ask for confirmation again.", + proposalConfirmReplyNo: "Do not execute the action mentioned above; continue the current conversation.", proposalConfirmYes: "Yes, proceed", proposalConfirmNo: "No thanks", proposalConfirmTitle: "Action confirmation required", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index fcc837dc29..123412a5cf 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -2526,6 +2526,15 @@ export default withEnglishFallback({ truncated: "Contexto truncado", contextSummary: "{database} · {tables} tablas", autoSqlBlocked: "El SQL generado por la IA parecía demasiado riesgoso para ejecutarse automáticamente. Revísalo manualmente antes de ejecutar.", + writeSqlConfirmationRequired: "Este SQL modifica datos o el esquema. Revísalo cuidadosamente antes de ejecutarlo.", + writeSqlConfirmationQuestion: "¿Ejecutar este SQL una vez?", + writeSqlConfirmYes: "Confirmar ejecución", + writeSqlConfirmNo: "Cancelar", + writeSqlConfirmationReplyYes: "Ejecuta una vez el SQL exacto mostrado arriba. No vuelvas a pedir confirmación.", + writeSqlConfirmationReplyNo: "No ejecutes el SQL mostrado arriba. Continúa la conversación.", + productionWriteBlocked: "La IA no puede ejecutar cambios de datos o esquema en una base de datos de producción. Revisa este SQL y ejecútalo manualmente en DBX.", + proposalConfirmReplyYes: "Ejecuta la acción que acabas de proponer arriba; no vuelvas a pedir confirmación.", + proposalConfirmReplyNo: "No ejecutes la acción mencionada arriba; continúa la conversación.", proposalConfirmYes: "Sí, continuar", proposalConfirmNo: "No, gracias", proposalConfirmTitle: "Confirmación de acción requerida", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index d63f8a1370..27a776b398 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -2420,6 +2420,15 @@ export default withEnglishFallback({ truncated: "Contesto troncato", contextSummary: "{database} · {tables} tabelle", autoSqlBlocked: "L'SQL generato dall'AI è sembrato troppo rischioso per l'esecuzione automatica. Rivedilo manualmente prima di eseguirlo.", + writeSqlConfirmationRequired: "Questo SQL modifica dati o schema. Rivedilo attentamente prima dell'esecuzione.", + writeSqlConfirmationQuestion: "Eseguire questo SQL una volta?", + writeSqlConfirmYes: "Conferma esecuzione", + writeSqlConfirmNo: "Annulla", + writeSqlConfirmationReplyYes: "Esegui una sola volta l'SQL esatto mostrato sopra. Non chiedere di nuovo conferma.", + writeSqlConfirmationReplyNo: "Non eseguire l'SQL mostrato sopra. Continua la conversazione.", + productionWriteBlocked: "L'AI non può eseguire modifiche a dati o schema su un database di produzione. Rivedi questo SQL ed eseguilo manualmente in DBX.", + proposalConfirmReplyYes: "Esegui l'azione che hai appena proposto sopra; non chiedere di nuovo conferma.", + proposalConfirmReplyNo: "Non eseguire l'azione menzionata sopra; continua la conversazione.", proposalConfirmYes: "Sì, procedi", proposalConfirmNo: "No, grazie", proposalConfirmTitle: "Conferma azione richiesta", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index 93a70a9933..84e571ff36 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -2560,6 +2560,15 @@ export default withEnglishFallback({ truncated: "コンテキストが切り詰められました", contextSummary: "{database} · {tables}テーブル", autoSqlBlocked: "AIが生成したSQLは自動実行するにはリスクが高すぎると判断されました。実行前に手動で確認してください。", + writeSqlConfirmationRequired: "この SQL はデータまたはスキーマを変更します。実行前に内容をよく確認してください。", + writeSqlConfirmationQuestion: "この SQL を一度実行しますか？", + writeSqlConfirmYes: "実行を確認", + writeSqlConfirmNo: "キャンセル", + writeSqlConfirmationReplyYes: "上に表示された正確な SQL を一度だけ実行してください。確認を再度求めないでください。", + writeSqlConfirmationReplyNo: "上に表示された SQL を実行しないでください。会話を続けてください。", + productionWriteBlocked: "AI は本番データベースでデータまたはスキーマの変更を実行できません。この SQL を確認し、DBX で手動実行してください。", + proposalConfirmReplyYes: "上で提案した操作を実行してください。再度確認を求めないでください。", + proposalConfirmReplyNo: "上で述べた操作は実行しないでください。会話を続けてください。", proposalConfirmYes: "はい、続行", proposalConfirmNo: "いいえ", proposalConfirmTitle: "操作の確認", diff --git a/apps/desktop/src/i18n/locales/ko.ts b/apps/desktop/src/i18n/locales/ko.ts index fb056afd50..d5c662034e 100644 --- a/apps/desktop/src/i18n/locales/ko.ts +++ b/apps/desktop/src/i18n/locales/ko.ts @@ -2337,6 +2337,15 @@ export default withEnglishFallback({ truncated: "컨텍스트가 잘렸습니다", contextSummary: "{database} · 테이블 {tables}개", autoSqlBlocked: "AI가 생성한 SQL이 자동 실행하기에는 너무 위험해 보입니다. 실행 전에 수동으로 검토하세요.", + writeSqlConfirmationRequired: "이 SQL은 데이터 또는 스키마를 변경합니다. 실행 전에 신중하게 검토하세요.", + writeSqlConfirmationQuestion: "이 SQL을 한 번 실행할까요?", + writeSqlConfirmYes: "실행 확인", + writeSqlConfirmNo: "취소", + writeSqlConfirmationReplyYes: "위에 표시된 정확한 SQL을 한 번만 실행하세요. 다시 확인하지 마세요.", + writeSqlConfirmationReplyNo: "위에 표시된 SQL을 실행하지 말고 대화를 계속하세요.", + productionWriteBlocked: "AI는 운영 데이터베이스에서 데이터 또는 스키마 변경을 실행할 수 없습니다. 이 SQL을 검토한 후 DBX에서 수동으로 실행하세요.", + proposalConfirmReplyYes: "위에서 제안한 작업을 실행하세요. 다시 확인을 요청하지 마세요.", + proposalConfirmReplyNo: "위에서 언급한 작업을 실행하지 말고 대화를 계속하세요.", proposalConfirmYes: "예, 진행", proposalConfirmNo: "아니오", proposalConfirmTitle: "작업 확인이 필요합니다", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index c2c157b78c..dd84ef300c 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -2528,6 +2528,15 @@ export default withEnglishFallback({ truncated: "Contexto truncado", contextSummary: "{database} · {tables} tabelas", autoSqlBlocked: "O SQL gerado pela AI pareceu arriscado demais para execução automática. Revise-o manualmente antes de executar.", + writeSqlConfirmationRequired: "Este SQL altera dados ou o esquema. Revise-o cuidadosamente antes de executar.", + writeSqlConfirmationQuestion: "Executar este SQL uma vez?", + writeSqlConfirmYes: "Confirmar execução", + writeSqlConfirmNo: "Cancelar", + writeSqlConfirmationReplyYes: "Execute uma vez o SQL exato mostrado acima. Não peça confirmação novamente.", + writeSqlConfirmationReplyNo: "Não execute o SQL mostrado acima. Continue a conversa.", + productionWriteBlocked: "A IA não pode executar alterações de dados ou esquema em um banco de dados de produção. Revise este SQL e execute-o manualmente no DBX.", + proposalConfirmReplyYes: "Execute a ação que você acabou de propor acima; não peça confirmação novamente.", + proposalConfirmReplyNo: "Não execute a ação mencionada acima; continue a conversa.", proposalConfirmYes: "Sim, continuar", proposalConfirmNo: "Não, obrigado", proposalConfirmTitle: "Confirmação de ação necessária", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index 6442b2eccf..172bc79e27 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -2411,6 +2411,15 @@ export default withEnglishFallback({ truncated: "上下文已截断", contextSummary: "{database} · {tables} 张表", autoSqlBlocked: "AI 生成的 SQL 风险较高，已阻止自动执行，请手动检查后再运行。", + writeSqlConfirmationRequired: "此 SQL 会修改数据或结构。执行前请仔细核对。", + writeSqlConfirmationQuestion: "确认执行此 SQL 一次吗？", + writeSqlConfirmYes: "确认执行", + writeSqlConfirmNo: "取消", + writeSqlConfirmationReplyYes: "仅执行上方展示的精确 SQL 一次，不要再次询问确认。", + writeSqlConfirmationReplyNo: "不要执行上方展示的 SQL，继续当前对话。", + productionWriteBlocked: "AI 不能在生产数据库执行数据或结构变更。请核对此 SQL 后在 DBX 中手动执行。", + proposalConfirmReplyYes: "请执行上面你刚提议的操作，不要再反问确认。", + proposalConfirmReplyNo: "不用执行上面提到的操作，继续当前对话。", proposalConfirmYes: "是，请继续", proposalConfirmNo: "不用了", proposalConfirmTitle: "需要确认是否执行以下操作", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index a48e8a19d5..9e63309eae 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -2421,6 +2421,15 @@ export default withEnglishFallback({ truncated: "上下文已截斷", contextSummary: "{database} · {tables} 張資料表", autoSqlBlocked: "AI 產生的 SQL 風險較高，已阻止自動執行，請手動檢查後再執行。", + writeSqlConfirmationRequired: "此 SQL 會修改資料或結構。執行前請仔細核對。", + writeSqlConfirmationQuestion: "確認執行此 SQL 一次嗎？", + writeSqlConfirmYes: "確認執行", + writeSqlConfirmNo: "取消", + writeSqlConfirmationReplyYes: "僅執行上方顯示的精確 SQL 一次，不要再次詢問確認。", + writeSqlConfirmationReplyNo: "不要執行上方顯示的 SQL，繼續目前對話。", + productionWriteBlocked: "AI 無法在生產資料庫執行資料或結構變更。請核對此 SQL 後在 DBX 中手動執行。", + proposalConfirmReplyYes: "請執行上方你剛剛提議的操作，不要再反問確認。", + proposalConfirmReplyNo: "不要執行上方提到的操作，繼續目前對話。", proposalConfirmYes: "是，請繼續", proposalConfirmNo: "不用了", proposalConfirmTitle: "需要確認是否執行以下操作", diff --git a/apps/desktop/src/lib/__tests__/ai/aiWriteSqlConfirmation.spec.ts b/apps/desktop/src/lib/__tests__/ai/aiWriteSqlConfirmation.spec.ts index b359be4427..25dd93f3a4 100644 --- a/apps/desktop/src/lib/__tests__/ai/aiWriteSqlConfirmation.spec.ts +++ b/apps/desktop/src/lib/__tests__/ai/aiWriteSqlConfirmation.spec.ts @@ -1,11 +1,18 @@ import { describe, expect, it } from "vitest"; -import { containsWriteSql, looksLikeActionProposal, looksLikeWriteSqlProposal, shouldGrantWriteSqlOnShortAffirmative, type WriteSqlGrantParams } from "@/lib/ai/aiProposalDetect"; +import { containsWriteSql, isActionableWriteProposalMessage, isActionableWriteSqlProposal, looksLikeActionProposal, looksLikeWriteSqlProposal, shouldGrantWriteSqlOnShortAffirmative, type WriteSqlGrantParams } from "@/lib/ai/aiProposalDetect"; import { aiSkillForAction } from "@/lib/ai/aiSkills"; import { extractFirstSqlCodeBlock, extractSingleSqlCodeBlock, countSqlCodeBlocks } from "@/lib/ai/aiSqlExecutionPolicy"; // ── looksLikeWriteSqlProposal ────────────────────────────────────────────── describe("looksLikeWriteSqlProposal", () => { + it("detects the backend-generated write confirmation proposal", () => { + const proposal = ["This database change requires your explicit confirmation before DBX can execute it.", "", "```sql", "CREATE TABLE users (id INT);", "```", "", "Should I execute this SQL statement?"].join("\n"); + expect(looksLikeActionProposal(proposal)).toBe(true); + expect(looksLikeWriteSqlProposal(proposal)).toBe(true); + expect(extractSingleSqlCodeBlock(proposal)).toBe("CREATE TABLE users (id INT);"); + }); + it("detects zh write-SQL proposal (CREATE in last line)", () => { expect(looksLikeWriteSqlProposal("需要我执行 CREATE TABLE users 吗？")).toBe(true); }); @@ -102,15 +109,15 @@ describe("looksLikeWriteSqlProposal", () => { describe("shouldGrantWriteSqlOnShortAffirmative (manual-typing confirmation)", () => { const writeProposal = { role: "assistant" as const, - content: "已经分析了表结构。\n需要我执行这条 CREATE TABLE users 语句吗？", + content: "已经分析了表结构。\n```sql\nCREATE TABLE users (id INT);\n```\n需要我执行这条 CREATE TABLE users 语句吗？", }; const insertProposal = { role: "assistant" as const, - content: "I understand the schema.\nShould I run this INSERT for you?", + content: "I understand the schema.\n```sql\nINSERT INTO users (id) VALUES (1);\n```\nShould I run this INSERT for you?", }; const writeProposalEn = { role: "assistant" as const, - content: "Should I execute this CREATE TABLE users?", + content: "```sql\nCREATE TABLE users (id INT);\n```\nShould I execute this CREATE TABLE users?", }; // Prompt-mirrored generic confirmation with SQL code block. const genericZhProposal = { @@ -145,6 +152,44 @@ describe("shouldGrantWriteSqlOnShortAffirmative (manual-typing confirmation)", ( expect(shouldGrantWriteSqlOnShortAffirmative(params({ messages: [{ role: "user", content: "创建表" }, genericZhProposal] }))).toBe(true); }); + it("does NOT grant for a generic write question without reviewable SQL", () => { + const genericWriteQuestion = "需要我执行这条 INSERT 语句吗？"; + expect(looksLikeWriteSqlProposal(genericWriteQuestion)).toBe(true); + expect(isActionableWriteSqlProposal(genericWriteQuestion)).toBe(false); + expect( + shouldGrantWriteSqlOnShortAffirmative( + params({ + messages: [ + { role: "user", content: "插入数据" }, + { role: "assistant", content: genericWriteQuestion }, + ], + }), + ), + ).toBe(false); + }); + + it("grants for a backend-generated localized confirmation via its structural kind", () => { + // The exact wording DBX renders from i18n (en): it does NOT contain an + // English/Chinese ask phrase, so the text detectors reject it. + const localizedEn = "This SQL changes data or schema. Review it carefully before execution.\n\n```sql\nCREATE TABLE users (id INT);\n```\n\nExecute this SQL once?"; + expect(looksLikeActionProposal(localizedEn)).toBe(false); + expect(isActionableWriteSqlProposal(localizedEn)).toBe(false); + // The backend marks the message, making it actionable without phrase detection. + expect(isActionableWriteProposalMessage({ role: "assistant", content: localizedEn, kind: "writeSqlConfirmation" })).toBe(true); + expect( + shouldGrantWriteSqlOnShortAffirmative( + params({ + messages: [ + { role: "user", content: "创建表" }, + { role: "assistant", content: localizedEn, kind: "writeSqlConfirmation" }, + ], + }), + ), + ).toBe(true); + // A production-block message must never be treated as a confirmable write. + expect(isActionableWriteProposalMessage({ role: "assistant", content: localizedEn, kind: "productionWriteBlocked" })).toBe(false); + }); + it("skips contextSummary messages to find the assistant message", () => { expect( shouldGrantWriteSqlOnShortAffirmative( diff --git a/apps/desktop/src/lib/ai/aiProposalDetect.ts b/apps/desktop/src/lib/ai/aiProposalDetect.ts index f37067b7c1..b8c6c52006 100644 --- a/apps/desktop/src/lib/ai/aiProposalDetect.ts +++ b/apps/desktop/src/lib/ai/aiProposalDetect.ts @@ -1,3 +1,5 @@ +import { extractSingleSqlCodeBlock } from "@/lib/ai/aiSqlExecutionPolicy"; + /** * Pure detectors for recognizing AI assistant action proposals and user * short replies (affirmative / negative). Used by the chat UI to render @@ -235,6 +237,35 @@ export function looksLikeWriteSqlProposal(content: string): boolean { return lastLineMentionsWriteSql(lastLine, isZh, content); } +/** + * A write proposal is actionable only when the user can review exactly one + * SQL statement. A generic “execute this INSERT?” prompt cannot safely bind + * the next agent run, so the UI must not present it as an executable action. + */ +export function isActionableWriteSqlProposal(content: string): boolean { + return looksLikeWriteSqlProposal(content) && extractSingleSqlCodeBlock(content) !== undefined; +} + +/** A chat message that the write-confirmation gate inspects. */ +export interface WriteSqlGrantMessage { + role: "user" | "assistant"; + content: string; + kind?: "contextSummary" | "writeSqlConfirmation" | "productionWriteBlocked"; +} + +/** + * True when a chat message can bind the next agent run to one exact SQL + * statement. Backend-generated confirmations are marked with a structural + * `kind` and carry exactly one SQL block, so they are actionable regardless of + * the locale their localized wording is in; model-authored proposals still + * have to pass the English/Chinese text detectors. + */ +export function isActionableWriteProposalMessage(msg: WriteSqlGrantMessage): boolean { + if (msg.role !== "assistant" || !msg.content) return false; + if (msg.kind === "writeSqlConfirmation") return extractSingleSqlCodeBlock(msg.content) !== undefined; + return isActionableWriteSqlProposal(msg.content); +} + /** * Returns true when the proposal line explicitly references a write/DDL * SQL operation (not a read-only action that happens to coexist with a @@ -279,7 +310,7 @@ export interface WriteSqlGrantParams { * The function scans backward from the last message, skipping * contextSummary entries, and stops at the first user message. */ - messages: Array<{ role: "user" | "assistant"; content: string; kind?: "contextSummary" }>; + messages: WriteSqlGrantMessage[]; } /** @@ -296,7 +327,7 @@ export function shouldGrantWriteSqlOnShortAffirmative(params: WriteSqlGrantParam for (let i = params.messages.length - 1; i >= 0; i--) { const msg = params.messages[i]; if (msg.kind === "contextSummary") continue; - if (msg.role === "assistant" && msg.content && looksLikeWriteSqlProposal(msg.content)) return true; + if (isActionableWriteProposalMessage(msg)) return true; if (msg.role === "user") return false; } return false; diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index e090830209..8b6553c68f 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -438,6 +438,8 @@ export async function aiStream(sessionId: string, request: AiCompletionRequest, export type AgentEvent = | { type: "turn_start"; turn: number } | { type: "text_delta"; delta: string } + | { type: "write_sql_confirmation_required"; sql: string } + | { type: "production_write_blocked"; sql: string } | { type: "reasoning_delta"; delta: string } | { type: "tool_call_start"; @@ -850,7 +852,7 @@ export interface AiChatMessage { content: string; mentions?: unknown[]; reasoning?: string; - kind?: "contextSummary"; + kind?: "contextSummary" | "writeSqlConfirmation" | "productionWriteBlocked"; } export interface AiConversation { diff --git a/crates/dbx-core/src/agent_events.rs b/crates/dbx-core/src/agent_events.rs index ecdef4e139..41b9b3aa08 100644 --- a/crates/dbx-core/src/agent_events.rs +++ b/crates/dbx-core/src/agent_events.rs @@ -8,6 +8,12 @@ pub enum AgentEvent { TurnStart { turn: u32 }, /// Text delta from the LLM response. TextDelta { delta: String }, + /// Request a localized, exact-SQL confirmation after an unconfirmed write + /// tool call. The desktop client owns the user-facing wording. + WriteSqlConfirmationRequired { sql: String }, + /// Tell the desktop client to show localized manual-review guidance after + /// an AI write attempt against a production target. + ProductionWriteBlocked { sql: String }, /// Reasoning/thinking delta (for models that support it). ReasoningDelta { delta: String }, /// The LLM wants to call a tool. diff --git a/crates/dbx-core/src/agent_loop.rs b/crates/dbx-core/src/agent_loop.rs index c2c1555abc..e63919335f 100644 --- a/crates/dbx-core/src/agent_loop.rs +++ b/crates/dbx-core/src/agent_loop.rs @@ -248,8 +248,9 @@ pub async fn run_agent_loop( ) .await; } + let mut sql_permissions = agent_ctx.sql_permissions.clone(); let tools = if is_agent_mode { - agent_tools::all_tools(agent_ctx.db_type, agent_ctx.sql_permissions.clone()) + agent_tools::all_tools(agent_ctx.db_type, sql_permissions.clone()) } else { agent_tools::read_only_tools(agent_ctx.db_type) }; @@ -469,6 +470,33 @@ pub async fn run_agent_loop( break; } + // Some models call execute_query for a write despite the prompt requiring + // a proposal first. Stop before dispatch so DBX can return an exact SQL + // proposal for non-production targets, or a non-confirmable production + // block rather than an impossible confirmation loop. + if let Some(sql) = unconfirmed_write_sql(&collected_tool_calls, agent_ctx.db_type, &sql_permissions) { + let targets_production = { + let configs = agent_ctx.state.configs.read().await; + configs.get(&agent_ctx.connection_id).is_some_and(|config| { + crate::production_safety::targets_production_database(config, &agent_ctx.database, sql) + }) + }; + let response = write_attempt_response(sql, targets_production); + on_event(match &response { + WriteAttemptResponse::ConfirmationRequired { sql } => { + AgentEvent::WriteSqlConfirmationRequired { sql: sql.clone() } + } + WriteAttemptResponse::ProductionBlocked { sql } => { + AgentEvent::ProductionWriteBlocked { sql: sql.clone() } + } + }); + // The frontend localizes this semantic event before persisting it in + // chat history. Keep a non-user-facing result for API consumers. + final_text = response.result_text(); + loop_exit = LoopExit::Completed; + break; + } + // Execute each tool call // Emit all ToolCallStart events first for tc in &collected_tool_calls { @@ -485,7 +513,7 @@ pub async fn run_agent_loop( 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(); + let parallel_sql_permissions = sql_permissions.clone(); // Split by index into parallel and sequential groups using tool metadata let tool_parallel_map: std::collections::HashMap<&str, bool> = @@ -510,7 +538,7 @@ pub async fn run_agent_loop( let conn = conn2.clone(); let db = db2.clone(); let schema = schema2.clone(); - let perms = sql_permissions.clone(); + let perms = parallel_sql_permissions.clone(); async move { agent_tools::execute_tool(&tc, &state, &conn, &db, schema.as_deref(), &db_type, perms).await } @@ -522,6 +550,7 @@ pub async fn run_agent_loop( let mut sequential_results = Vec::with_capacity(sequential_indices.len()); for &i in &sequential_indices { let tc = make_tc(&collected_tool_calls[i]); + let execution_permissions = sequential_tool_permissions(&tc, db_type, &mut sql_permissions); sequential_results.push( agent_tools::execute_tool( &tc, @@ -530,7 +559,7 @@ pub async fn run_agent_loop( &db2, schema2.as_deref(), &db_type, - sql_permissions.clone(), + execution_permissions, ) .await, ); @@ -1184,6 +1213,72 @@ async fn maybe_compact( CompactResult::Compacted } +/// Semantic write-attempt outcomes. Keep user-visible language in the desktop +/// i18n catalog rather than hard-coding English into a backend event. +enum WriteAttemptResponse { + ConfirmationRequired { sql: String }, + ProductionBlocked { sql: String }, +} + +impl WriteAttemptResponse { + fn result_text(&self) -> String { + match self { + Self::ConfirmationRequired { sql } | Self::ProductionBlocked { sql } => sql.clone(), + } + } +} + +fn write_attempt_response(sql: &str, targets_production: bool) -> WriteAttemptResponse { + let sql = sql.trim().to_string(); + if targets_production { + WriteAttemptResponse::ProductionBlocked { sql } + } else { + WriteAttemptResponse::ConfirmationRequired { sql } + } +} + +fn sequential_tool_permissions( + tool_call: &ToolCall, + db_type: DatabaseType, + permissions: &mut agent_tools::AgentSqlPermissions, +) -> agent_tools::AgentSqlPermissions { + if tool_call.name == "execute_query" { + if let Some(sql) = tool_call.arguments.get("sql").and_then(|value| value.as_str()) { + return agent_tools::take_sql_permissions_for_execution(sql, db_type, permissions); + } + } + permissions.clone() +} + +/// Returns a deterministic confirmation proposal when the current turn contains +/// an unconfirmed write/DDL tool call. The caller must stop before dispatching +/// any tools so the database never sees the attempted SQL. +fn unconfirmed_write_sql<'a>( + tool_calls: &'a [ToolCall], + db_type: DatabaseType, + permissions: &agent_tools::AgentSqlPermissions, +) -> Option<&'a str> { + // Mongo execute_query is a read-only shell-command tool; mutating commands + // are rejected by the tool itself, so a confirmation proposal could never + // be executed. And because Mongo commands are not SQL, the risk classifier + // would flag even reads (db.collection.find(...)) as writes, turning every + // Mongo agent query into an impossible confirmation. Preserve the previous + // behavior: Mongo writes keep failing with a clear read-only error. + if db_type == DatabaseType::MongoDb { + return None; + } + tool_calls.iter().find_map(|tool_call| { + if tool_call.name != "execute_query" { + return None; + } + let sql = tool_call.arguments.get("sql").and_then(|value| value.as_str())?; + agent_tools::write_requires_confirmation(sql, db_type, permissions) + .ok() + .filter(|required| *required) + .map(|_| sql) + }) +} + fn tool_result_for_followup_context(tool_name: &str, content: &str) -> String { let result = compact_tool_result_for_context(tool_name, content); format!( @@ -1440,6 +1535,160 @@ mod tests { } } + #[test] + fn unconfirmed_write_attempt_requests_confirmation_with_exact_sql() { + let sql = "CREATE TABLE users (id INT);"; + let response = write_attempt_response(sql, false); + + assert!( + matches!(response, WriteAttemptResponse::ConfirmationRequired { sql: response_sql } if response_sql == sql) + ); + } + + #[test] + fn production_write_attempt_is_not_an_executable_confirmation() { + let response = write_attempt_response("DELETE FROM users WHERE id = 7;", true); + + assert!( + matches!(response, WriteAttemptResponse::ProductionBlocked { sql } if sql == "DELETE FROM users WHERE id = 7;") + ); + } + + #[test] + fn direct_unconfirmed_ddl_tool_call_becomes_a_confirmation_proposal() { + let tool_calls = vec![ToolCall { + id: "create-table".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": "CREATE TABLE users (id INT);" }), + provider_payload: None, + }]; + + let sql = + unconfirmed_write_sql(&tool_calls, DatabaseType::Postgres, &agent_tools::AgentSqlPermissions::default()) + .expect("unconfirmed DDL must be intercepted before tool dispatch"); + assert!( + matches!(write_attempt_response(sql, false), WriteAttemptResponse::ConfirmationRequired { sql } if sql == "CREATE TABLE users (id INT);") + ); + } + + #[test] + fn confirmed_or_read_only_tool_calls_do_not_become_confirmation_proposals() { + let create = ToolCall { + id: "create-table".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": "CREATE TABLE users (id INT);" }), + provider_payload: None, + }; + let select = ToolCall { + id: "select-users".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": "SELECT * FROM users" }), + provider_payload: None, + }; + let confirmed = + agent_tools::confirmed_write_sql_permissions(false, true, Some("CREATE TABLE users (id INT);".to_string())); + + assert!(unconfirmed_write_sql(&[create], DatabaseType::Postgres, &confirmed).is_none()); + assert!(unconfirmed_write_sql(&[select], DatabaseType::Postgres, &agent_tools::AgentSqlPermissions::default()) + .is_none()); + } + + #[test] + fn duplicate_confirmed_writes_in_one_batch_receive_one_grant() { + let sql = "DELETE FROM users WHERE id = 7;"; + let call = ToolCall { + id: "delete-user".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": sql }), + provider_payload: None, + }; + let mut permissions = agent_tools::confirmed_write_sql_permissions(false, true, Some(sql.to_string())); + + let first = sequential_tool_permissions(&call, DatabaseType::Postgres, &mut permissions); + let second = sequential_tool_permissions(&call, DatabaseType::Postgres, &mut permissions); + + assert!(first.allow_writes); + assert_eq!(first.confirmed_write_sql.as_deref(), Some(sql)); + assert!(!second.allow_writes); + assert_eq!(second.confirmed_write_sql, None); + assert_eq!(permissions, agent_tools::AgentSqlPermissions::default()); + } + + #[test] + fn confirmed_write_grant_stays_consumed_across_turns() { + let sql = "DELETE FROM users WHERE id = 7;"; + let call = ToolCall { + id: "delete-user".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": sql }), + provider_payload: None, + }; + let mut permissions = agent_tools::confirmed_write_sql_permissions(false, true, Some(sql.to_string())); + + let _ = sequential_tool_permissions(&call, DatabaseType::Postgres, &mut permissions); + + assert_eq!(unconfirmed_write_sql(std::slice::from_ref(&call), DatabaseType::Postgres, &permissions), Some(sql)); + } + + #[test] + fn read_before_confirmed_write_does_not_consume_the_grant() { + let confirmed_sql = "DELETE FROM users WHERE id = 7;"; + let read = ToolCall { + id: "read-user".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": "SELECT * FROM users WHERE id = 7" }), + provider_payload: None, + }; + let write = ToolCall { + id: "delete-user".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": confirmed_sql }), + provider_payload: None, + }; + let mut permissions = + agent_tools::confirmed_write_sql_permissions(false, true, Some(confirmed_sql.to_string())); + + let read_permissions = sequential_tool_permissions(&read, DatabaseType::Postgres, &mut permissions); + let write_permissions = sequential_tool_permissions(&write, DatabaseType::Postgres, &mut permissions); + + assert!(read_permissions.allow_writes); + assert!(write_permissions.allow_writes); + assert_eq!(write_permissions.confirmed_write_sql.as_deref(), Some(confirmed_sql)); + assert_eq!(permissions, agent_tools::AgentSqlPermissions::default()); + } + + #[test] + fn mongo_agent_commands_are_never_intercepted_as_write_confirmations() { + // Mongo execute_query is read-only at the tool level, and its commands are + // not SQL — the risk classifier would flag even reads as writes. Neither a + // read nor a mutating command may become an impossible confirmation. + let read = ToolCall { + id: "mongo-read".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": "db.collection.find({})" }), + provider_payload: None, + }; + let write = ToolCall { + id: "mongo-write".to_string(), + name: "execute_query".to_string(), + arguments: json!({ "sql": "db.collection.insertOne({ name: \"x\" })" }), + provider_payload: None, + }; + let permissions = agent_tools::AgentSqlPermissions::default(); + // Without the Mongo exclusion this read would be flagged as an unconfirmed + // write (Mongo commands are not SQL), producing a confirmation for every + // agent query. Assert the misclassification so the exclusion's necessity + // stays explicit. + assert!(agent_tools::write_requires_confirmation( + "db.collection.find({})", + DatabaseType::MongoDb, + &permissions + ) + .unwrap()); + assert!(unconfirmed_write_sql(&[read], DatabaseType::MongoDb, &permissions).is_none()); + assert!(unconfirmed_write_sql(&[write], DatabaseType::MongoDb, &permissions).is_none()); + } + #[test] fn generate_contract_rejects_schema_summary_without_sql() { let contract = generate_contract("帮我生成统计 2026年1月2日新注册会员数量的 sql", "ask"); diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index 94d0d4e71d..07d32381e2 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -152,6 +152,44 @@ fn sql_risk_allowed(risk: SqlRisk, permissions: AgentSqlPermissions) -> bool { } } +/// Returns true when an Agent attempted a write or DDL call before DBX has a +/// user-confirmed SQL binding for the current run. The caller must turn that +/// attempt into a confirmation proposal instead of sending it to the database. +pub fn write_requires_confirmation( + sql: &str, + db_type: DatabaseType, + permissions: &AgentSqlPermissions, +) -> Result { + if permissions.allow_writes || permissions.allow_dangerous || permissions.confirmed_write_sql.is_some() { + return Ok(false); + } + let risk = crate::sql_risk::classify_sql_risk_for_database(sql, db_type)?; + Ok(matches!(risk, SqlRisk::Write | SqlRisk::Ddl)) +} + +/// Snapshot the permissions for one execute_query call and consume an exact +/// confirmed write/DDL grant before dispatch. Later calls in the same agent run +/// receive the cleared permissions and must request a new confirmation. +pub(crate) fn take_sql_permissions_for_execution( + sql: &str, + db_type: DatabaseType, + permissions: &mut AgentSqlPermissions, +) -> AgentSqlPermissions { + let execution_permissions = permissions.clone(); + let consumes_confirmation = match crate::sql_risk::classify_sql_risk_for_database(sql, db_type) { + Ok(risk @ (SqlRisk::Write | SqlRisk::Ddl)) => { + sql_risk_allowed(risk, execution_permissions.clone()) + && sql_matches_confirmed_write(sql, &execution_permissions.confirmed_write_sql) + && execution_permissions.confirmed_write_sql.is_some() + } + _ => false, + }; + if consumes_confirmation { + *permissions = AgentSqlPermissions::default(); + } + execution_permissions +} + /// Returns true for vector database types (Qdrant, Milvus, Weaviate, ChromaDb). /// If modifying this, also update VECTOR_DB_TYPES in apps/desktop/src/lib/ai.ts. pub fn is_vector_db(db_type: DatabaseType) -> bool { @@ -328,7 +366,7 @@ fn execute_query_tool(sql_permissions: AgentSqlPermissions) -> ToolDefinition { } else if sql_permissions.allow_writes { "Execute SQL after the user explicitly confirmed this operation. Read queries and non-DDL writes are allowed for this run." } else { - "Execute a read-only SQL query and return results (max 50 rows). Only SELECT, WITH, SHOW, DESCRIBE, EXPLAIN statements are allowed. Write operations (INSERT/UPDATE/DELETE/DDL) are blocked." + "Execute a read-only SQL query and return results (max 50 rows). This run cannot execute writes or DDL because no specific SQL has been confirmed yet; this does not mean the database itself is read-only. When the user requests a write, first propose the exact SQL in one ```sql code block and ask for confirmation. After confirmation, DBX starts a new run that can execute only that exact SQL. Only SELECT, WITH, SHOW, DESCRIBE, EXPLAIN statements may be executed in this run." }; ToolDefinition { name: "execute_query", @@ -712,12 +750,26 @@ async fn execute_execute_query( // Classify SQL risk using the concrete database dialect. let risk = crate::sql_risk::classify_sql_risk_for_database(sql, *db_type)?; let connection_config = state.configs.read().await.get(connection_id).cloned(); - if let Some(config) = connection_config { - if risk != SqlRisk::ReadOnly && crate::production_safety::targets_production_database(&config, database, sql) { - return Err("Blocked: AI agents cannot execute writes or DDL on a production database. Return the SQL for the user to review and execute manually in DBX.".to_string()); - } + let targets_production = connection_config.as_ref().is_some_and(|config| { + risk != SqlRisk::ReadOnly && crate::production_safety::targets_production_database(config, database, sql) + }); + let risk_allowed = sql_risk_allowed(risk, sql_permissions.clone()); + let confirmed_sql_matches = + risk == SqlRisk::ReadOnly || sql_matches_confirmed_write(sql, &sql_permissions.confirmed_write_sql); + + // Keep permission diagnostics useful without logging SQL, connection details, or result data. + log::debug!( + "[agent:execute_query:permission] risk={risk:?} targets_production={targets_production} allow_writes={} allow_dangerous={} has_confirmed_sql={} confirmed_sql_matches={confirmed_sql_matches} will_execute={}", + sql_permissions.allow_writes, + sql_permissions.allow_dangerous, + sql_permissions.confirmed_write_sql.is_some(), + !targets_production && risk_allowed && confirmed_sql_matches, + ); + + if targets_production { + return Err("Blocked: AI agents cannot execute writes or DDL on a production database. Return the SQL for the user to review and execute manually in DBX.".to_string()); } - if !sql_risk_allowed(risk, sql_permissions.clone()) { + if !risk_allowed { if risk == SqlRisk::Transaction { return Err("Blocked: transaction control statements are not available to the AI agent.".to_string()); } @@ -729,7 +781,7 @@ async fn execute_execute_query( // When the user confirmed a specific write SQL, the agent must // execute only that SQL — not an arbitrary different statement. - if risk != SqlRisk::ReadOnly && !sql_matches_confirmed_write(sql, &sql_permissions.confirmed_write_sql) { + if !confirmed_sql_matches { let confirmed = sql_permissions.confirmed_write_sql.as_deref().unwrap_or(""); return Err(format!( "Blocked: the executed SQL does not match the user-confirmed SQL.\n\ @@ -1438,6 +1490,52 @@ for line in sys.stdin: )); } + #[test] + fn unconfirmed_sql_tool_describes_the_confirmation_flow() { + let execute_query = all_tools(DatabaseType::Postgres, AgentSqlPermissions::default()) + .into_iter() + .find(|tool| tool.name == "execute_query") + .expect("PostgreSQL Agent mode should expose execute_query"); + + assert!(execute_query.description.contains("no specific SQL has been confirmed yet")); + assert!(execute_query.description.contains("does not mean the database itself is read-only")); + assert!(execute_query.description.contains("propose the exact SQL in one ```sql code block")); + assert!(!sql_risk_allowed(SqlRisk::Write, AgentSqlPermissions::default())); + assert!(!sql_risk_allowed(SqlRisk::Ddl, AgentSqlPermissions::default())); + } + + #[test] + fn write_confirmation_is_required_only_for_unconfirmed_write_or_ddl() { + let permissions = AgentSqlPermissions::default(); + assert!(write_requires_confirmation("INSERT INTO users (id) VALUES (1)", DatabaseType::Postgres, &permissions) + .unwrap()); + assert!(write_requires_confirmation("CREATE TABLE users (id INT)", DatabaseType::Mysql, &permissions).unwrap()); + assert!(!write_requires_confirmation("SELECT * FROM users", DatabaseType::Postgres, &permissions).unwrap()); + + let confirmed = confirmed_write_sql_permissions(false, true, Some("CREATE TABLE users (id INT)".to_string())); + assert!( + !write_requires_confirmation("CREATE TABLE users (id INT)", DatabaseType::Postgres, &confirmed).unwrap() + ); + } + + #[test] + fn confirmed_postgres_article_tags_insert_is_write_enabled_for_nonproduction_agent() { + let sql = "INSERT INTO \"public\".\"article_tags\" (article_id, tag_id)\nVALUES (1, 4), (2, 4), (3, 8), (4, 5), (4, 6);"; + let permissions = confirmed_write_sql_permissions(false, true, Some(sql.to_string())); + + assert!(permissions.allow_writes); + assert!(permissions.allow_dangerous); + assert!(sql_risk_allowed(SqlRisk::Write, permissions.clone())); + assert!(sql_matches_confirmed_write(sql, &permissions.confirmed_write_sql)); + + let execute_query = all_tools(DatabaseType::Postgres, permissions) + .into_iter() + .find(|tool| tool.name == "execute_query") + .expect("PostgreSQL Agent mode should expose execute_query"); + assert!(execute_query.description.contains("writes")); + assert!(!execute_query.description.contains("Write operations (INSERT/UPDATE/DELETE/DDL) are blocked")); + } + #[test] fn oracle_agent_tools_include_explain_query() { let tools = all_tools(DatabaseType::Oracle, AgentSqlPermissions::default());