feat(ai): defer marimo-pair references and improve code mode prompts

frontend/src/components/chat/chat-panel.tsx

@@ -67,6 +67,8 @@ import { PromptInput } from "../editor/ai/add-cell-with-ai";
 import {
   addContextCompletion,
   CONTEXT_TRIGGER,
+  isContextAttachment,
+  resolveChatContext,
 } from "../editor/ai/completion-utils";
 import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
 import { CopyClipboardIcon } from "../icons/copy-icon";
@@ -83,7 +85,6 @@ import {
 import { renderUIMessage } from "./chat-display";
 import { ChatHistoryPopover } from "./chat-history-popover";
 import {
-  buildCompletionRequestBody,
   convertToFileUIPart,
   generateChatTitle,
   handleToolCall,
@@ -92,6 +93,7 @@ import {
   PROVIDERS_THAT_SUPPORT_ATTACHMENTS,
   useFileState,
 } from "./chat-utils";
+import { getCodes } from "@/core/codemirror/copilot/getCodes";
 
 // Default mode for the AI
 const DEFAULT_MODE = "manual";
@@ -278,16 +280,13 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
         subtitle: "AI with access to read and write tools",
         Icon: HatGlasses,
       },
-    ];
-
-    if (import.meta.env.DEV) {
-      modeOptions.push({
+      {
         value: "code_mode",
         label: "Code Mode (experimental)",
         subtitle: "AI with access to the notebook's kernel. Use with caution.",
         Icon: CodeIcon,
-      });
-    }
+      },
+    ];
 
     const isAttachmentSupported =
       PROVIDERS_THAT_SUPPORT_ATTACHMENTS.has(currentProvider);
@@ -535,9 +534,10 @@ const ChatPanelBody = () => {
           );
         }
 
-        const completionBody = await buildCompletionRequestBody(
-          options.messages,
-        );
+        const completionBody = {
+          uiMessages: options.messages,
+          includeOtherCode: getCodes(""),
+        };
 
         // Call this here to ensure the value is not stale
         const chatMode = store.get(aiAtom)?.mode || DEFAULT_MODE;
@@ -626,6 +626,8 @@ const ChatPanelBody = () => {
       initialAttachments && initialAttachments.length > 0
         ? await convertToFileUIPart(initialAttachments)
         : undefined;
+    const { contextPart, attachments } =
+      await resolveChatContext(initialMessage);
 
     // Trigger AI conversation with append
     sendMessage({
@@ -635,7 +637,9 @@ const ChatPanelBody = () => {
           type: "text" as const,
           text: initialMessage,
         },
+        ...(contextPart ? [contextPart] : []),
         ...(fileParts ?? []),
+        ...attachments,
       ],
     });
     clearFiles();
@@ -653,17 +657,31 @@ const ChatPanelBody = () => {
     openModal(<PairWithAgentModal onClose={closeModal} />);
   });
 
-  const handleMessageEdit = useEvent((index: number, newValue: string) => {
-    const editedMessage = messages[index];
-    const fileParts = editedMessage.parts?.filter((p) => p.type === "file");
-
-    const messageId = editedMessage.id;
-    sendMessage({
-      messageId: messageId, // replace the message
-      role: "user",
-      parts: [{ type: "text", text: newValue }, ...fileParts],
-    });
-  });
+  const handleMessageEdit = useEvent(
+    async (index: number, newValue: string) => {
+      const editedMessage = messages[index];
+      // Keep the user's own uploaded files, but drop the previous @-context
+      // snapshot (data part + its attachments) so we can re-resolve a fresh,
+      // point-in-time snapshot from the edited text below.
+      const userFileParts =
+        editedMessage.parts?.filter(
+          (p) => p.type === "file" && !isContextAttachment(p),
+        ) ?? [];
+      const { contextPart, attachments } = await resolveChatContext(newValue);
+
+      const messageId = editedMessage.id;
+      sendMessage({
+        messageId: messageId, // replace the message
+        role: "user",
+        parts: [
+          { type: "text", text: newValue },
+          ...(contextPart ? [contextPart] : []),
+          ...userFileParts,
+          ...attachments,
+        ],
+      });
+    },
+  );
 
   const handleChatInputSubmit = useEvent(
     async (e: KeyboardEvent | undefined, newValue: string): Promise<void> => {
@@ -674,11 +692,17 @@ const ChatPanelBody = () => {
         storePrompt(newMessageInputRef.current.view);
       }
       const fileParts = files ? await convertToFileUIPart(files) : undefined;
+      const { contextPart, attachments } = await resolveChatContext(newValue);
 
       e?.preventDefault();
       sendMessage({
-        text: newValue,
-        files: fileParts,
+        role: "user",
+        parts: [
+          { type: "text", text: newValue },
+          ...(contextPart ? [contextPart] : []),
+          ...(fileParts ?? []),
+          ...attachments,
+        ],
       });
       setInput("");
       clearFiles();

frontend/src/components/editor/ai/__tests__/completion-utils.test.ts

@@ -8,7 +8,14 @@ import { datasetsAtom } from "@/core/datasets/state";
 import type { DatasetsState } from "@/core/datasets/types";
 import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
-import { codeToCells, getAICompletionBody } from "../completion-utils";
+import type { UIMessage } from "ai";
+import {
+  codeToCells,
+  getAICompletionBody,
+  isContextAttachment,
+  MARIMO_CONTEXT_PART_TYPE,
+  resolveChatContext,
+} from "../completion-utils";
 
 // Mock getCodes function
 vi.mock("@/core/codemirror/copilot/getCodes", () => ({
@@ -350,6 +357,89 @@ describe("getAICompletionBody", () => {
   });
 });
 
+describe("resolveChatContext", () => {
+  beforeEach(() => {
+    store.set(datasetsAtom, {
+      tables: [],
+    } as unknown as DatasetsState);
+    store.set(dataSourceConnectionsAtom, {
+      latestEngineSelected: DUCKDB_ENGINE,
+      connectionsMap: new Map(),
+    });
+    store.set(variablesAtom, {});
+  });
+
+  it("returns no context when the input has no @-mentions", async () => {
+    const result = await resolveChatContext("just a plain question");
+    expect(result).toEqual({ contextPart: null, attachments: [] });
+  });
+
+  it("returns no context part when @-mentions resolve to nothing", async () => {
+    const result = await resolveChatContext("look at @variable://ghost");
+    expect(result.contextPart).toBeNull();
+    expect(result.attachments).toEqual([]);
+  });
+
+  it("captures resolved @-context into a data part", async () => {
+    store.set(variablesAtom, {
+      [variableName("var1")]: {
+        name: variableName("var1"),
+        value: "string value",
+        dataType: "string",
+        declaredBy: [],
+        usedBy: [],
+      },
+    });
+
+    const result = await resolveChatContext("inspect @variable://var1");
+
+    expect(result.contextPart?.type).toBe(MARIMO_CONTEXT_PART_TYPE);
+    expect(result.contextPart?.data.contextIds).toEqual(["variable://var1"]);
+    expect(result.contextPart?.data.plainText).toMatchInlineSnapshot(
+      `"<variable name="var1" dataType="string">"string value"</variable>"`,
+    );
+  });
+});
+
+describe("isContextAttachment", () => {
+  type Part = UIMessage["parts"][number];
+
+  it("is true for a file part tagged as context", () => {
+    const part = {
+      type: "file",
+      mediaType: "image/png",
+      url: "data:image/png;base64,abc",
+      providerMetadata: { marimo: { source: "context" } },
+    } as Part;
+    expect(isContextAttachment(part)).toBe(true);
+  });
+
+  it("is false for a user-uploaded file part (no marker)", () => {
+    const part = {
+      type: "file",
+      mediaType: "image/png",
+      url: "data:image/png;base64,abc",
+    } as Part;
+    expect(isContextAttachment(part)).toBe(false);
+  });
+
+  it("is false for a file part with unrelated provider metadata", () => {
+    const part = {
+      type: "file",
+      mediaType: "image/png",
+      url: "data:image/png;base64,abc",
+      providerMetadata: { openai: { foo: "bar" } },
+    } as Part;
+    expect(isContextAttachment(part)).toBe(false);
+  });
+
+  it("is false for non-file parts", () => {
+    expect(isContextAttachment({ type: "text", text: "hi" } as Part)).toBe(
+      false,
+    );
+  });
+});
+
 describe("codeToCells", () => {
   it("should return empty array for empty string", () => {
     const code = "";

frontend/src/components/editor/ai/completion-utils.ts

@@ -7,7 +7,7 @@ import {
   startCompletion,
 } from "@codemirror/autocomplete";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
-import type { FileUIPart } from "ai";
+import type { DataUIPart, FileUIPart, UIMessage } from "ai";
 import { getAIContextRegistry } from "@/core/ai/context/context";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
@@ -50,6 +50,91 @@ export function getAICompletionBody({
   };
 }
 
+export interface MarimoContextData {
+  plainText: string;
+  contextIds: string[];
+}
+
+export type MarimoContextUIPart = DataUIPart<{
+  "marimo-context": MarimoContextData;
+}>;
+
+/**
+ * Wire `type` of the @-context data part. Must match
+ * `MARIMO_CONTEXT_PART_TYPE` on the backend.
+ */
+export const MARIMO_CONTEXT_PART_TYPE =
+  "data-marimo-context" as const satisfies MarimoContextUIPart["type"];
+
+export interface ResolvedChatContext {
+  contextPart: MarimoContextUIPart | null;
+  attachments: FileUIPart[];
+}
+
+/**
+ * Marker stamped onto attachments derived from @-context (as opposed to files
+ * the user uploaded directly).
+ */
+const CONTEXT_ATTACHMENT_METADATA = {
+  marimo: { source: "context" },
+} as const;
+
+/** Whether a part is an attachment that was derived from @-context. */
+export function isContextAttachment(part: UIMessage["parts"][number]): boolean {
+  return (
+    part.type === "file" &&
+    part.providerMetadata?.marimo?.source ===
+      CONTEXT_ATTACHMENT_METADATA.marimo.source
+  );
+}
+
+/**
+ * Resolve @-context for messages. They represent referenced
+ * datasets, variables, or other context from the user's prompt.
+ */
+export async function resolveChatContext(
+  input: string,
+): Promise<ResolvedChatContext> {
+  if (!input.includes(CONTEXT_TRIGGER)) {
+    return { contextPart: null, attachments: [] };
+  }
+
+  const registry = getAIContextRegistry(store);
+  const contextIds = registry.parseAllContextIds(input);
+  if (contextIds.length === 0) {
+    return { contextPart: null, attachments: [] };
+  }
+
+  const plainText = registry.formatContextForAI(contextIds);
+
+  let attachments: FileUIPart[] = [];
+  try {
+    const resolved = await registry.getAttachmentsForContext(contextIds);
+    attachments = resolved.map((attachment) => ({
+      ...attachment,
+      providerMetadata: {
+        ...attachment.providerMetadata,
+        marimo: {
+          ...attachment.providerMetadata?.marimo,
+          ...CONTEXT_ATTACHMENT_METADATA.marimo,
+        },
+      },
+    }));
+  } catch (error) {
+    Logger.error("Error getting attachments:", error);
+  }
+
+  let contextPart: MarimoContextUIPart | null = null;
+  if (plainText.trim()) {
+    contextPart = {
+      type: MARIMO_CONTEXT_PART_TYPE,
+      data: { plainText, contextIds: contextIds.map(String) },
+    };
+  }
+
+  return { contextPart, attachments };
+}
+
 /**
  * Gets the request body and attachments for the AI completion API.
  */