feat(ai): drag scoped table references into chat

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 }}

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([]);
+  });
+});

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,