From 09b326558f9ce261378e3fb8586beb4c5bc32238 Mon Sep 17 00:00:00 2001
From: Allen <maybe.zhuo@qq.com>
Date: Fri, 21 Aug 2026 13:03:14 +0800
Subject: [PATCH] feat(ai): drag scoped table references into chat

---
 .../src/components/editor/AiAssistant.vue     | 24 ++++++-
 .../AiAssistant.tableReferenceDrop.spec.ts    | 53 +++++++++++++++
 .../src/components/sidebar/TreeItem.vue       |  3 +-
 .../__tests__/ai/aiTableReferenceDrop.spec.ts | 67 +++++++++++++++++++
 .../src/lib/ai/aiTableReferenceDrop.ts        | 43 ++++++++++++
 .../src/lib/editor/queryEditorTableDrop.ts    | 11 ++-
 6 files changed, 197 insertions(+), 4 deletions(-)
 create mode 100644 apps/desktop/src/components/editor/__tests__/AiAssistant.tableReferenceDrop.spec.ts
 create mode 100644 apps/desktop/src/lib/__tests__/ai/aiTableReferenceDrop.spec.ts
 create mode 100644 apps/desktop/src/lib/ai/aiTableReferenceDrop.ts

diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue
index 505930773a..e8f5ca1739 100644
--- a/apps/desktop/src/components/editor/AiAssistant.vue
+++ b/apps/desktop/src/components/editor/AiAssistant.vue
@@ -1,4 +1,4 @@
-<script setup lang="ts">
+<script setup lang="ts">
 import { computed, nextTick, onMounted, onUnmounted, ref, toRaw, watch, type Component } from "vue";
 import { uuid } from "@/lib/common/utils";
 import { useI18n } from "vue-i18n";
@@ -131,6 +131,8 @@ import ExplainPlanViewer from "@/components/explain/ExplainPlanViewer.vue";
 import { parseExplainResult, parseOracleExplainText, type ParsedExplainPlan } from "@/lib/diagram/explainPlan";
 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 { handleAiTableReferenceDropEvent } from "@/lib/ai/aiTableReferenceDrop";
+import { DBX_TABLE_REFERENCE_DROP_EVENT, clearActiveTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
 import { isAiPromptImeCompositionEvent, shouldSubmitAiPromptOnKeydown } from "@/lib/ai/aiPromptKeyboard";
 import { isActionableWriteProposalMessage, isActionableWriteSqlProposal, looksLikeActionProposal, looksLikeWriteSqlProposal, shouldGrantWriteSqlOnShortAffirmative } from "@/lib/ai/aiProposalDetect";
 import { visibleToActualIndex } from "@/lib/ai/aiMessageEdit";
@@ -2286,6 +2288,22 @@ function onTauriFileDrop(event: Event) {
   void addDroppedAttachmentPaths(payload.paths);
 }
 
+function onTableReferenceDropEvent(event: Event) {
+  handleAiTableReferenceDropEvent(event, {
+    context: {
+      connectionId: props.tab?.connectionId || props.connection?.id,
+      database: props.tab?.database || props.connection?.database || "",
+    },
+    assistantRoot: assistantRootRef.value,
+    elementFromPoint: (x, y) => document.elementFromPoint(x, y),
+    onMention: (mention, payload) => {
+      addSelectedMention({ kind: "table", schema: mention.schema, name: mention.table, tableType: "table" });
+      clearActiveTableReferencePayload(payload);
+      nextTick(() => promptTextareaRef.value?.focus());
+    },
+  });
+}
+
 async function send() {
   const text = prompt.value.trim();
   if ((!text && !selectedMentions.value.length && !selectedSqlFileMentions.value.length && !selectedCsvAttachments.value.length && !selectedImageAttachments.value.length) || isGenerating.value) return;
@@ -2891,6 +2909,7 @@ onMounted(async () => {
 
   window.addEventListener("resize", handlePanelResize);
   document.addEventListener("dbx:tauri-file-drop", onTauriFileDrop as EventListener);
+  window.addEventListener(DBX_TABLE_REFERENCE_DROP_EVENT, onTableReferenceDropEvent);
   if (typeof ResizeObserver !== "undefined" && assistantRootRef.value) {
     promptPanelResizeObserver = new ResizeObserver(handlePanelResize);
     promptPanelResizeObserver.observe(assistantRootRef.value);
@@ -2969,6 +2988,7 @@ onUnmounted(() => {
   document.body.style.cursor = "";
   window.removeEventListener("resize", handlePanelResize);
   document.removeEventListener("dbx:tauri-file-drop", onTauriFileDrop as EventListener);
+  window.removeEventListener(DBX_TABLE_REFERENCE_DROP_EVENT, onTableReferenceDropEvent);
   promptPanelResizeObserver?.disconnect();
 });
 
@@ -3042,7 +3062,7 @@ async function openExternalUrl(url: string) {
 </script>
 
 <template>
-  <div ref="assistantRootRef" class="flex h-full min-h-0 flex-col overflow-hidden" @dragenter="onAttachmentDragEnter" @dragover="onAttachmentDragOver" @dragleave="onAttachmentDragLeave" @drop="onAttachmentDrop">
+  <div ref="assistantRootRef" data-ai-assistant-root class="flex h-full min-h-0 flex-col overflow-hidden" @dragenter="onAttachmentDragEnter" @dragover="onAttachmentDragOver" @dragleave="onAttachmentDragLeave" @drop="onAttachmentDrop">
     <div class="flex items-center gap-2 border-b px-3 shrink-0" :class="settings.editorSettings.appLayout === 'classic' ? 'h-9' : 'h-10'">
       <span class="flex flex-1 self-stretch items-center truncate text-xs font-medium" data-tauri-drag-region>
         {{ chatTitle }}
diff --git a/apps/desktop/src/components/editor/__tests__/AiAssistant.tableReferenceDrop.spec.ts b/apps/desktop/src/components/editor/__tests__/AiAssistant.tableReferenceDrop.spec.ts
new file mode 100644
index 0000000000..63bc6b5f9e
--- /dev/null
+++ b/apps/desktop/src/components/editor/__tests__/AiAssistant.tableReferenceDrop.spec.ts
@@ -0,0 +1,53 @@
+// @vitest-environment happy-dom
+
+import { afterEach, describe, expect, it } from "vitest";
+import { handleAiTableReferenceDropEvent, type AiTableReferenceDropContext } from "@/lib/ai/aiTableReferenceDrop";
+import { createTableReferenceDropEvent, createTableReferencePayload, DBX_TABLE_REFERENCE_DROP_EVENT, type QueryEditorTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
+
+function dispatchTableReferenceDrop(payload: QueryEditorTableReferencePayload, context: AiTableReferenceDropContext) {
+  const assistantRoot = document.createElement("div");
+  const target = document.createElement("span");
+  assistantRoot.append(target);
+  document.body.append(assistantRoot);
+  const mentions: string[] = [];
+  const listener = (event: Event) => {
+    handleAiTableReferenceDropEvent(event, {
+      context,
+      assistantRoot,
+      elementFromPoint: () => target,
+      onMention: (mention) => mentions.push(mention.raw),
+    });
+  };
+  window.addEventListener(DBX_TABLE_REFERENCE_DROP_EVENT, listener);
+  window.dispatchEvent(createTableReferenceDropEvent({ payload, clientX: 12, clientY: 24 }));
+  window.removeEventListener(DBX_TABLE_REFERENCE_DROP_EVENT, listener);
+  return mentions;
+}
+
+function tablePayload(overrides: Partial<QueryEditorTableReferencePayload> = {}) {
+  return createTableReferencePayload({
+    connectionId: overrides.connectionId ?? "conn-1",
+    database: overrides.database ?? "app-db",
+    schema: "public",
+    tableName: "users",
+    databaseType: "postgres",
+  })!;
+}
+
+afterEach(() => {
+  document.body.innerHTML = "";
+});
+
+describe("AI assistant table reference drop", () => {
+  it("accepts a table dropped from the active connection and database", () => {
+    expect(dispatchTableReferenceDrop(tablePayload(), { connectionId: "conn-1", database: "app-db" })).toEqual(["@public.users"]);
+  });
+
+  it("rejects a table dropped from another connection", () => {
+    expect(dispatchTableReferenceDrop(tablePayload({ connectionId: "conn-2" }), { connectionId: "conn-1", database: "app-db" })).toEqual([]);
+  });
+
+  it("rejects a table dropped from another database", () => {
+    expect(dispatchTableReferenceDrop(tablePayload({ database: "analytics" }), { connectionId: "conn-1", database: "app-db" })).toEqual([]);
+  });
+});
diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue
index 6428ccd756..ca48aa454e 100644
--- a/apps/desktop/src/components/sidebar/TreeItem.vue
+++ b/apps/desktop/src/components/sidebar/TreeItem.vue
@@ -54,6 +54,7 @@ import LightTooltip from "@/components/ui/LightTooltip.vue";
 import type { ColumnInfo, ConnectionConfig, CustomTypeTreeMemberMeta, DatabaseType, TreeNode, TriggerInfo } from "@/types/database";
 import { alignedCommentLeadingWidth, canTreeNodePin, canTreeNodeShowExpander, sidebarTreeNodeComment, trailingCommentAvailableWidth, trailingCommentGapPx, treeItemPaddingLeft, treeLabelWidthClass, usesFullWidthTreeLabel } from "@/lib/sidebar/sidebarTreeItemLayout";
 import { clearActiveTableReferencePayload, createTableReferencePayload, createTableReferenceDropEvent, setActiveTableReferencePayload, type QueryEditorTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
+import { AI_ASSISTANT_TABLE_DROP_ROOT_SELECTOR } from "@/lib/ai/aiTableReferenceDrop";
 import { formatSidebarObjectStorage } from "@/lib/sidebar/sidebarDatabaseStorage";
 import { dataTabOpenModeFromTreeClick } from "@/lib/sidebar/dataTabOpenPolicy";
 import { effectiveDatabaseTypeForConnection } from "@/lib/database/jdbcDialect";
@@ -1189,7 +1190,7 @@ function onTableReferenceMouseUp(event: MouseEvent) {
   if (payload) {
     suppressNextTableReferenceClick = true;
     const target = document.elementFromPoint(event.clientX, event.clientY);
-    if (target instanceof Element && target.closest("[data-query-editor-root]")) {
+    if (target instanceof Element && target.closest(`[data-query-editor-root], ${AI_ASSISTANT_TABLE_DROP_ROOT_SELECTOR}`)) {
       window.dispatchEvent(
         createTableReferenceDropEvent({
           payload,
diff --git a/apps/desktop/src/lib/__tests__/ai/aiTableReferenceDrop.spec.ts b/apps/desktop/src/lib/__tests__/ai/aiTableReferenceDrop.spec.ts
new file mode 100644
index 0000000000..bd26d54585
--- /dev/null
+++ b/apps/desktop/src/lib/__tests__/ai/aiTableReferenceDrop.spec.ts
@@ -0,0 +1,67 @@
+import { describe, expect, it } from "vitest";
+import { aiTableMentionFromTableReference } from "@/lib/ai/aiTableReferenceDrop";
+import { createTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
+
+describe("ai table reference drop", () => {
+  const context = { connectionId: "conn-1", database: "app-db" };
+
+  it("maps a table payload to a table mention", () => {
+    const payload = createTableReferencePayload({
+      connectionId: "conn-1",
+      database: "app-db",
+      schema: "public",
+      tableName: "users",
+      databaseType: "postgres",
+    });
+
+    expect(aiTableMentionFromTableReference(payload, context)).toEqual({
+      raw: "@public.users",
+      schema: "public",
+      table: "users",
+    });
+  });
+
+  it("quotes mention parts that are not simple identifiers", () => {
+    const payload = createTableReferencePayload({
+      connectionId: "conn-1",
+      database: "app-db",
+      tableName: "order items",
+      databaseType: "mysql",
+    });
+
+    expect(aiTableMentionFromTableReference(payload, context)).toEqual({
+      raw: '@"order items"',
+      schema: undefined,
+      table: "order items",
+    });
+  });
+
+  it("ignores database references", () => {
+    const payload = createTableReferencePayload({
+      connectionId: "conn-1",
+      database: "app-db",
+      referenceType: "database",
+      databaseType: "mysql",
+    });
+
+    expect(aiTableMentionFromTableReference(payload, context)).toBeNull();
+  });
+
+  it("ignores column references", () => {
+    const payload = createTableReferencePayload({
+      connectionId: "conn-1",
+      database: "app-db",
+      schema: "public",
+      tableName: "users",
+      columnName: "email",
+      databaseType: "postgres",
+    });
+
+    expect(aiTableMentionFromTableReference(payload, context)).toBeNull();
+  });
+
+  it("ignores empty payloads", () => {
+    expect(aiTableMentionFromTableReference(null, context)).toBeNull();
+    expect(aiTableMentionFromTableReference(undefined, context)).toBeNull();
+  });
+});
diff --git a/apps/desktop/src/lib/ai/aiTableReferenceDrop.ts b/apps/desktop/src/lib/ai/aiTableReferenceDrop.ts
new file mode 100644
index 0000000000..2a9fb59af0
--- /dev/null
+++ b/apps/desktop/src/lib/ai/aiTableReferenceDrop.ts
@@ -0,0 +1,43 @@
+import { formatAiTableMention, type AiTableMention } from "@/lib/ai/aiTableMentions";
+import type { QueryEditorTableReferenceDropDetail, QueryEditorTableReferencePayload } from "@/lib/editor/queryEditorTableDrop";
+
+/** Selector marking the AI assistant panel root as a table-reference drop target. */
+export const AI_ASSISTANT_TABLE_DROP_ROOT_SELECTOR = "[data-ai-assistant-root]";
+
+export interface AiTableReferenceDropContext {
+  connectionId?: string;
+  database?: string;
+}
+
+export interface AiTableReferenceDropHandlerOptions {
+  context: AiTableReferenceDropContext;
+  assistantRoot: Element | null | undefined;
+  elementFromPoint: (x: number, y: number) => Element | null;
+  onMention: (mention: AiTableMention, payload: QueryEditorTableReferencePayload) => void;
+}
+
+/**
+ * Maps a sidebar table-reference drag payload to an AI table mention chip.
+ * Only plain table/view references become mentions; database and column
+ * references are not representable as table mentions and return null.
+ */
+export function aiTableMentionFromTableReference(payload: QueryEditorTableReferencePayload | null | undefined, context: AiTableReferenceDropContext): AiTableMention | null {
+  if (!payload || payload.referenceType === "database" || payload.columnName) return null;
+  if (!context.connectionId || context.database == null || payload.connectionId !== context.connectionId || payload.database !== context.database) return null;
+  const table = payload.tableName;
+  if (!table) return null;
+  const schema = payload.schema;
+  return { raw: formatAiTableMention(schema, table), schema, table };
+}
+
+export function handleAiTableReferenceDropEvent(event: Event, options: AiTableReferenceDropHandlerOptions): boolean {
+  if (!(event instanceof CustomEvent)) return false;
+  const detail = event.detail as QueryEditorTableReferenceDropDetail | undefined;
+  if (!detail?.payload) return false;
+  const target = options.elementFromPoint(detail.clientX, detail.clientY);
+  if (!target || !options.assistantRoot?.contains(target)) return false;
+  const mention = aiTableMentionFromTableReference(detail.payload, options.context);
+  if (!mention) return false;
+  options.onMention(mention, detail.payload);
+  return true;
+}
diff --git a/apps/desktop/src/lib/editor/queryEditorTableDrop.ts b/apps/desktop/src/lib/editor/queryEditorTableDrop.ts
index b284e36c96..c11205092a 100644
--- a/apps/desktop/src/lib/editor/queryEditorTableDrop.ts
+++ b/apps/desktop/src/lib/editor/queryEditorTableDrop.ts
@@ -24,7 +24,16 @@ export interface QueryEditorTableReferenceDropDetail {
 
 let activeTableReferencePayload: QueryEditorTableReferencePayload | null = null;
 
-export function createTableReferencePayload(options: { connectionId?: string; database?: string; schema?: string; tableName?: string; columnName?: string; referenceType?: "database" | "table" | "column"; databaseType?: DatabaseType; driverProfile?: string }): QueryEditorTableReferencePayload | null {
+export function createTableReferencePayload(options: {
+  connectionId?: string;
+  database?: string;
+  schema?: string;
+  tableName?: string;
+  columnName?: string;
+  referenceType?: "database" | "table" | "column";
+  databaseType?: DatabaseType;
+  driverProfile?: string;
+}): QueryEditorTableReferencePayload | null {
   if (!options.connectionId || options.database == null) return null;
   const referenceType = options.referenceType ?? (options.columnName ? "column" : "table");
   if (referenceType !== "database" && !options.tableName) return null;