diff --git a/frontend/src/components/chat/chat-panel.tsx b/frontend/src/components/chat/chat-panel.tsx
index 70bbd86b304..8735fb985d7 100644
--- a/frontend/src/components/chat/chat-panel.tsx
+++ b/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();
diff --git a/frontend/src/components/editor/ai/__tests__/completion-utils.test.ts b/frontend/src/components/editor/ai/__tests__/completion-utils.test.ts
index 0e8d5814bd3..449c283ce27 100644
--- a/frontend/src/components/editor/ai/__tests__/completion-utils.test.ts
+++ b/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 = "";
diff --git a/frontend/src/components/editor/ai/completion-utils.ts b/frontend/src/components/editor/ai/completion-utils.ts
index ceb7f778806..ef408595d42 100644
--- a/frontend/src/components/editor/ai/completion-utils.ts
+++ b/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.
  */
diff --git a/marimo/_ai/_pydantic_ai_utils.py b/marimo/_ai/_pydantic_ai_utils.py
index 2b37f7e0e2c..c274dcf3847 100644
--- a/marimo/_ai/_pydantic_ai_utils.py
+++ b/marimo/_ai/_pydantic_ai_utils.py
@@ -23,6 +23,15 @@ def generate_id(prefix: str) -> str:
     return f"{prefix}_{uuid.uuid4().hex}"
 
 
+# Wire `type` of the @-context data part emitted by the frontend
+MARIMO_CONTEXT_PART_TYPE = "data-marimo-context"
+
+
+def format_inline_context(plain_text: str) -> str:
+    """Render @-context as a user-message text block."""
+    return f"<context>\n{plain_text.strip()}\n</context>"
+
+
 def form_toolsets(
     tools: list[ToolDefinition],
     tool_invoker: Callable[[str, dict[str, Any]], Any],
@@ -98,7 +107,7 @@ def safe_part_processor(
             or generate_id("message")
         )
         role = message.get("role", "assistant")
-        parts = [sanitize_part(part) for part in message.get("parts", [])]
+        parts = _prepare_parts(message.get("parts", []))
         metadata = message.get("metadata")
 
         ui_message = UIMessage(
@@ -117,6 +126,40 @@ def safe_part_processor(
     return pydantic_messages
 
 
+def _prepare_parts(raw_parts: list[Any]) -> list[Any]:
+    """Normalize a message's raw parts for pydantic-ai validation."""
+    parts: list[Any] = []
+    for part in raw_parts:
+        lowered = _expand_marimo_context_part(part)
+        if lowered is None:
+            continue  # empty context part, drop it
+        parts.append(sanitize_part(lowered))
+    return parts
+
+
+def _expand_marimo_context_part(part: Any) -> Any:
+    """Resolve a `data-marimo-context` part into a text part.
+
+    The @-context is shipped inside the user message as a data part because
+    pydantic-ai's VercelAIAdapter drops DataUIPart entirely.
+
+    Returns the part unchanged when it isn't a context part, a text part when
+    it carries non-empty context, or `None` to drop empty context parts.
+    """
+    if (
+        not isinstance(part, dict)
+        or part.get("type") != MARIMO_CONTEXT_PART_TYPE
+    ):
+        return part
+    data = part.get("data")
+    if not isinstance(data, dict):
+        return None
+    plain_text = (data.get("plainText") or "").strip()
+    if not plain_text:
+        return None
+    return {"type": "text", "text": format_inline_context(plain_text)}
+
+
 def create_simple_prompt(text: str) -> UIMessage:
     from pydantic_ai.ui.vercel_ai.request_types import TextUIPart, UIMessage
 
diff --git a/marimo/_server/ai/prompts.py b/marimo/_server/ai/prompts.py
index 0376060cb2b..eb66472ad10 100644
--- a/marimo/_server/ai/prompts.py
+++ b/marimo/_server/ai/prompts.py
@@ -311,8 +311,7 @@ def _multi_cell_language_rules() -> str:
 def _common_chat_sections(
     *,
     custom_rules: str | None,
-    include_other_code: str,
-    context: AiCompletionContext | None,
+    include_other_code: str | None,
 ) -> str:
     """Trailing sections shared by every chat mode."""
     out = ""
@@ -320,8 +319,6 @@ def _common_chat_sections(
         out += f"\n\n## Additional rules:\n{custom_rules}"
     if include_other_code:
         out += "\n\n" + _tag("code_from_other_cells", include_other_code)
-    if context:
-        out += format_context(context)
     return out
 
 
@@ -447,7 +444,6 @@ def _common_chat_sections(
 def get_chat_system_prompt(
     *,
     custom_rules: str | None,
-    context: AiCompletionContext | None,
     include_other_code: str,
     mode: CopilotMode,
     session_id: SessionId,
@@ -462,12 +458,12 @@ def get_chat_system_prompt(
             f"notebooks:\n\n{skill_md}"
         )
         system_prompt = f"{intro}\n\n{skill_section}"
-        system_prompt += _single_cell_language_rules()
-        return system_prompt + _common_chat_sections(
+        system_prompt += _common_chat_sections(
             custom_rules=custom_rules,
-            include_other_code=include_other_code,
-            context=context,
+            include_other_code=None,  # code mode can inspect code
         )
+        system_prompt += "\nIf you are not aware of the current notebook code, inspect it first before answering any questions."
+        return system_prompt
 
     system_prompt = (
         f"\n{_get_mode_intro_message(mode)}\n{_get_session_info(session_id)}"
@@ -487,7 +483,6 @@ def get_chat_system_prompt(
     return system_prompt + _common_chat_sections(
         custom_rules=custom_rules,
         include_other_code=include_other_code,
-        context=context,
     )
 
 
diff --git a/marimo/_server/ai/providers.py b/marimo/_server/ai/providers.py
index b55281182bd..290ce1620bd 100644
--- a/marimo/_server/ai/providers.py
+++ b/marimo/_server/ai/providers.py
@@ -253,14 +253,17 @@ async def stream_completion_harness(
 
         from marimo._server.ai.tools.code_mode import (
             build_execute_code_toolset,
+            references_capability,
         )
 
         model = self.create_model(max_tokens=max_tokens)
+        capabilities = references_capability()
         agent = Agent(
             model,
             model_settings=self._build_agent_settings(model),
             toolsets=[build_execute_code_toolset(session, request)],
             instructions=system_prompt,
+            capabilities=capabilities,
         )
 
         run_input = SubmitMessage(
@@ -268,7 +271,7 @@ async def stream_completion_harness(
             trigger="submit-message",
             messages=self.convert_messages(messages),
         )
-        stream_options.span_info.tool_count = 1
+        stream_options.span_info.tool_count = 1 + len(capabilities)
 
         adapter = VercelAIAdapter(
             agent=agent,
diff --git a/marimo/_server/ai/skills/marimo-pair/SKILL.md b/marimo/_server/ai/skills/marimo-pair/SKILL.md
index 2da0eac4f15..eb8e2780060 100644
--- a/marimo/_server/ai/skills/marimo-pair/SKILL.md
+++ b/marimo/_server/ai/skills/marimo-pair/SKILL.md
@@ -222,8 +222,8 @@ Submit the code that belongs in the cell.
   same-cell intermediates.
 - **Define each public name once** - a public name has one owning cell.
   Reassigning it in another cell fails with `Multiply-defined names`; edit the
-  owning cell or give the result a new name. See
-  [gotchas.md](references/gotchas.md).
+  owning cell or give the result a new name. Load the `gotchas` capability for
+  more traps.
 - **Run cells deliberately** - `create_cell` and `edit_cell` change structure
   only. Queue `ctx.run_cell(...)` when the cell should execute.
 
@@ -253,11 +253,14 @@ different paths.
 - **Set anywidget traitlets directly** - synced traitlets are Python
   attributes, for example `widget.value = 5`.
 
-For designing custom visual or interactive output, see
-[rich-representations.md](references/rich-representations.md).
+For designing custom visual or interactive output, load the
+`rich-representations` capability.
 
-## References
+## On-demand references
 
-- [gotchas.md](references/gotchas.md) — name redefinition, cached module proxies, and notebook traps
-- [rich-representations.md](references/rich-representations.md) — custom widgets and visualizations
-- [notebook-improvements.md](references/notebook-improvements.md) — improving existing notebooks
+Load these with the `load_capability` tool when you need deeper guidance. Do
+not read reference files from disk.
+
+- **`gotchas`** — name redefinition, cached module proxies, and notebook traps
+- **`rich-representations`** — custom widgets and visualizations
+- **`notebook-improvements`** — improving existing notebooks
diff --git a/marimo/_server/ai/skills/marimo-pair/references/gotchas.md b/marimo/_server/ai/skills/marimo-pair/references/gotchas.md
new file mode 100644
index 00000000000..3f55dd5d103
--- /dev/null
+++ b/marimo/_server/ai/skills/marimo-pair/references/gotchas.md
@@ -0,0 +1,91 @@
+# Gotchas
+
+Loaded on demand via the `gotchas` capability (`load_capability`). Do not read
+this file from disk.
+
+## Private variables are cell-scoped
+
+Variables with a `_` prefix are **private to the cell that defines them** in
+marimo. They cannot be referenced from other cells — you'll get a `NameError`.
+
+This matters when building notebooks programmatically. A common mistake:
+
+```python
+# Cell A
+_df = pd.DataFrame(results)   # _df is private to this cell
+
+# Cell B — FAILS
+mo.ui.table(_df)               # NameError: name '_df' is not defined
+```
+
+**Fix:** Either merge both into one cell, or use a non-private name (`df`).
+
+## Redefining a public name across cells
+
+Each public name has one owning cell. Defining it again in another cell fails
+with `Multiply-defined names`. This is easy to hit when building a notebook
+incrementally — a second cell reassigns `df`, `results`, `data`, etc.
+
+```python
+# Cell A
+df = pd.read_csv("data.csv")
+
+# Cell B — FAILS: df already defined in Cell A
+df = df.dropna()               # Multiply-defined names: df
+```
+
+**Fix — pick one:**
+
+- **Edit the owning cell** if the step belongs there (`ctx.edit_cell`).
+- **Use a new name** when later cells need the result (`clean = df.dropna()`).
+- **Use a private `_` name** for a throwaway intermediate (`_clean = df.dropna()`).
+
+`ctx.graph.cells[cid].defs` shows what a cell already owns.
+
+## Duplicate public imports across cells
+
+The same single-definition rule applies to imports: a public name (like `pd`)
+can only be defined in one cell. If two cells both `import pandas as pd`, you
+get a `Multiply-defined names` error at validation.
+
+**Fix:** Use a `_` prefix on the second import (`import pandas as _pd`) or
+consolidate imports into a shared cell.
+
+## `inspect.getsource()` on methods is indented
+
+`inspect.getsource()` on a class method preserves the original indentation.
+Passing this to `ast.parse()` fails with `IndentationError`.
+
+```python
+# FAILS
+src = inspect.getsource(SomeClass.some_method)
+tree = ast.parse(src)  # IndentationError: unexpected indent
+
+# FIX
+import textwrap
+src = textwrap.dedent(inspect.getsource(SomeClass.some_method))
+tree = ast.parse(src)
+```
+
+## Cached module availability
+
+Some libraries cache optional-dependency availability at import time. Installing
+a package mid-session via `ctx.packages.add()` won't update those caches.
+The user may need to restart the kernel — but try known workarounds first.
+
+### Polars + pyarrow
+
+`df.to_pandas()` fails with `ModuleNotFoundError: pa.Table requires 'pyarrow'`.
+
+**Workaround** — if this error occurs after installing pyarrow mid-session,
+run the following via `execute_code` (scratchpad), NOT in a cell. The patch
+mutates the cached module object in the running kernel, so it doesn't need to
+persist in the notebook.
+
+```python
+import pyarrow as _pa
+import polars.dataframe.frame as _frame_mod
+_frame_mod.pa = _pa
+```
+
+Then re-run the failing cell.
diff --git a/marimo/_server/ai/skills/marimo-pair/references/notebook-improvements.md b/marimo/_server/ai/skills/marimo-pair/references/notebook-improvements.md
new file mode 100644
index 00000000000..db8f2d6fe47
--- /dev/null
+++ b/marimo/_server/ai/skills/marimo-pair/references/notebook-improvements.md
@@ -0,0 +1,100 @@
+# Notebook Improvements
+
+Loaded on demand via the `notebook-improvements` capability (`load_capability`).
+Do not read this file from disk.
+
+When the user asks to improve, optimize, or clean up their notebook, scan the
+current cells for these opportunities. Use your judgment — don't over-apply,
+and if you're unsure whether a change is worthwhile, ask the user.
+
+## Cell names
+
+Low priority unless the user asks. `setup` and cells defining
+functions/classes are auto-named by marimo. Beyond that, naming is optional.
+Note that naming markdown cells clutters the UI by showing the cell header
+that's normally hidden.
+
+## Setup cell
+
+A setup cell is named `"setup"` and is guaranteed to run before all other
+cells. It's the place for module imports. Consolidating imports here keeps
+the notebook clean and ensures every cell can rely on those modules being
+available.
+
+**The setup cell cannot reference other cells' variables.** It runs first, so
+it must be self-contained: imports, constants, and definitions that depend only
+on each other. Reading a name defined elsewhere (e.g. `df`, a UI element) fails
+with `The setup cell cannot have references`.
+
+First check if the notebook already has a cell named `"setup"`. If not, create
+one and hoist scattered imports into it. `name="setup"` auto-positions the cell
+first — no `before`/`after` needed:
+
+```python
+cid = ctx.create_cell('''import polars as pl
+import marimo as mo
+import anywidget
+import traitlets''', name="setup")
+ctx.run_cell(cid)
+```
+
+If a setup cell already exists, `create_cell(name="setup")` raises `ValueError`;
+use `ctx.edit_cell("setup", code=...)` and `ctx.run_cell("setup")` instead.
+
+## Lift reusable functions into their own cells
+
+When a cell contains a single function or class that doesn't reference
+variables from other cells, marimo treats it specially — it can be written as
+a standalone definition and reused outside the notebook. These functions can
+use modules from the setup cell.
+
+Look for functions that **could belong in a library**: data loading, transforms,
+parsers, domain logic, custom widgets. If someone might reasonably `import` it
+from another module, it's a good candidate to lift into its own cell.
+
+Don't lift everything — notebook-specific wiring (UI layout, display logic,
+cell-level orchestration) should stay where it is. Use `_prefix` for
+cell-internal helpers that aren't meant to be reused.
+
+```python
+# before: useful logic buried in a larger cell
+objects = pl.read_csv("https://example.com/objects.csv")
+artists = pl.read_csv("https://example.com/artists.csv")
+
+def top_counts(df, col, n=5):
+    return df.group_by(col).len().sort("len", descending=True).head(n)
+
+result = top_counts(objects.join(artists, on="id"), "category")
+```
+
+```python
+# after: top_counts is general-purpose — give it its own cell
+
+# cell 1
+def top_counts(df, col, n=5):
+    return df.group_by(col).len().sort("len", descending=True).head(n)
+```
+
+```python
+# cell 2
+result = top_counts(df, "category")
+```
+
+## `mo.persistent_cache`
+
+`@mo.persistent_cache` caches a function's result to disk so it isn't
+recomputed on subsequent runs. The cache persists across kernel restarts.
+
+```python
+@mo.persistent_cache
+def load_data():
+    objects = pl.read_csv("https://example.com/objects.csv")
+    artists = pl.read_csv("https://example.com/artists.csv")
+    return objects.join(artists, on="id", how="left")
+
+df = load_data()
+```
+
+Good candidates: data loading, ETL, expensive computation that rarely changes.
+Don't over-optimize — if you're unsure, suggest it to the user rather than
+applying it.
diff --git a/marimo/_server/ai/skills/marimo-pair/references/rich-representations.md b/marimo/_server/ai/skills/marimo-pair/references/rich-representations.md
new file mode 100644
index 00000000000..913f1fe6a81
--- /dev/null
+++ b/marimo/_server/ai/skills/marimo-pair/references/rich-representations.md
@@ -0,0 +1,312 @@
+# Rich Representations
+
+Loaded on demand via the `rich-representations` capability (`load_capability`).
+Do not read this file from disk.
+
+Custom visual encodings for data that go beyond standard charts and tables.
+
+## Guiding principles
+
+**Visualization matters.** Helping users build custom visual representations
+is one of the highest-impact things the agent can do. A bespoke encoding
+tailored to the task — labeling, batch review, comparing variants — lets
+users _see_ their data in ways that tables and numbers never will. marimo
+is an environment where users create their own views, not just consume
+library charts. Help them imagine what's possible, then build it.
+
+**Use modern web APIs.** Models may default to older browser patterns; prefer
+modern HTML, CSS, and JavaScript that are supported in current browsers. Avoid
+build steps unless the task clearly needs them.
+
+**Prefer compact output.** marimo clips cell output at ~610px and scrolls.
+Avoid hitting that limit; if you need more space, manage your own scrolling
+inside a fixed-height container.
+
+**Keep it thin, make it compose.** A widget is a thin layer over data, not
+an application. One clear purpose, few traitlets, small `_esm`. Build small
+pieces that compose in the notebook — combine with other cells, UI elements,
+and views. Don't over-engineer.
+
+## Decision tree
+
+| Need                                           | Approach                                                                 |
+| ---------------------------------------------- | ------------------------------------------------------------------------ |
+| Custom output or interaction                   | **anywidget** — flexible enough to grow from display-only to interactive |
+| Tiny static HTML representation                | `_display_()` or `mo.Html`                                               |
+| Built-in control used as-is (slider, dropdown) | `mo.ui.*`                                                                |
+
+For custom representations, prefer anywidget unless the output is clearly a
+small static one-off.
+
+## anywidget
+
+[anywidget](https://anywidget.dev) bridges Python and JavaScript via
+traitlets. `.tag(sync=True)` makes a traitlet bidirectional — Python sets a
+value → JS sees it; JS calls `model.set()` + `model.save_changes()` →
+Python sees it. `_css` is optional global CSS.
+
+**marimo does not render traditional Jupyter widgets.** Libraries like jscatter,
+ipyvolume, etc. often have a top-level object whose default representation is a
+Jupyter widget (`application/vnd.jupyter.widget-view+json`). marimo cannot
+display these — you need to find the underlying **anywidget** instance, which
+marimo _does_ support.
+
+Common pattern: look for a `.widget` attribute on the library object:
+
+```python
+# jscatter example — Scatter is not renderable, but .widget is an anywidget
+scatter = jscatter.Scatter(data=df, x="x", y="y")
+scatter.widget  # <-- use this in the cell output
+```
+
+When unsure, check in the scratchpad:
+
+```python
+import anywidget
+obj = scatter.widget  # or whatever accessor the library provides
+print(isinstance(obj, anywidget.AnyWidget))  # True = marimo can render it
+```
+
+### `_esm` lifecycle
+
+**Render only** (most widgets):
+
+```js
+function render({ model, el }) {
+  /* ... */
+}
+export default { render };
+```
+
+**Initialize + render** (shared state across views, one-time setup):
+
+```js
+export default () => {
+  return {
+    initialize({ model }) {
+      // Once per widget instance — timers, connections, shared handlers
+      return () => {
+        /* cleanup */
+      };
+    },
+    render({ model, el }) {
+      // Once per view — display in 3 cells = 3 renders
+      return () => {
+        /* cleanup DOM listeners */
+      };
+    },
+  };
+};
+```
+
+- `model.on()` is auto-cleaned when a view is removed
+- DOM `addEventListener` is **not** — clean up with `AbortController`
+
+### Timer example (initialize + render)
+
+`initialize` owns one interval; each `render` view displays it.
+
+```python
+import anywidget
+import traitlets
+
+_TIMER_ESM = """
+export default () => {
+  return {
+    initialize({ model }) {
+      const id = setInterval(() => {
+        if (model.get("running")) {
+          model.set("seconds", model.get("seconds") + 1);
+          model.save_changes();
+        }
+      }, 1000);
+      return () => clearInterval(id);
+    },
+    render({ model, el }) {
+      const controller = new AbortController();
+      const { signal } = controller;
+
+      const span = document.createElement("span");
+      span.style.cssText = "font: 24px monospace;";
+
+      const btn = document.createElement("button");
+      btn.style.cssText = "margin-left: 8px; cursor: pointer;";
+
+      function update() {
+        const s = model.get("seconds");
+        const mm = String(Math.floor(s / 60)).padStart(2, "0");
+        const ss = String(s % 60).padStart(2, "0");
+        span.textContent = `${mm}:${ss}`;
+        btn.textContent = model.get("running") ? "⏸" : "▶";
+      }
+
+      model.on("change:seconds", update);
+      model.on("change:running", update);
+
+      btn.addEventListener("click", () => {
+        model.set("running", !model.get("running"));
+        model.save_changes();
+      }, { signal });
+
+      update();
+      el.append(span, btn);
+      return () => controller.abort();
+    }
+  };
+};
+"""
+
+class Timer(anywidget.AnyWidget):
+    seconds = traitlets.Int(0).tag(sync=True)
+    running = traitlets.Bool(True).tag(sync=True)
+    _esm = _TIMER_ESM
+```
+
+### Composing with the notebook
+
+Widgets become reactive notebook citizens when you bridge a traitlet to
+`mo.state`. This is a two-cell pattern — create the widget and wire up the
+observer in one cell, read the value in another:
+
+```python
+# Cell 1 — widget + observer
+timer = Timer()
+
+get_seconds, set_seconds = mo.state(timer.seconds)
+timer.observe(lambda _: set_seconds(timer.seconds), names=["seconds"])
+
+timer  # display the widget
+```
+
+```python
+# Cell 2 — reacts to changes
+seconds = get_seconds()
+mo.md(f"Timer is at **{seconds}s** — {'running' if seconds > 0 else 'stopped'}")
+```
+
+The common pattern is `mo.state(widget.trait)` for the initial value,
+`.observe()` on the specific trait name, and reading with the getter in a
+downstream cell. See [Reactive anywidgets](#reactive-anywidgets-in-marimo)
+for the details.
+
+### CDN dependencies
+
+Import JS libraries from [esm.sh](https://esm.sh) — no build step:
+
+```js
+import * as d3 from "https://esm.sh/d3@7";
+import { tableFromIPC } from "https://esm.sh/@uwdata/flechette@2";
+```
+
+### DataFrames and binary data
+
+**Prefer reducing data on the Python side.** Aggregate, filter, sample —
+send the widget only what it needs. Most widgets should receive a small,
+pre-processed payload via simple traitlets (lists, dicts). This keeps the
+widget simple and avoids extra dependencies.
+
+**For large tabular data (>2k rows)** where the widget genuinely needs
+row-level access, send Arrow IPC bytes instead of JSON. This adds
+complexity and dependencies, so only reach for it when the data volume
+justifies it.
+
+**Python — serialize:**
+
+```python
+# Polars (native, no pyarrow needed)
+_ipc=df.write_ipc(None).getvalue()
+
+# Any __arrow_c_stream__ source (pandas, narwhals, pyarrow, etc.)
+import io, pyarrow as pa, pyarrow.feather as feather
+
+def to_arrow_ipc(data) -> bytes:
+    table = pa.RecordBatchReader.from_stream(data).read_all()
+    sink = io.BytesIO()
+    feather.write_feather(table, sink, compression="uncompressed")
+    return sink.getvalue()
+```
+
+**JS — deserialize with `@uwdata/flechette`:**
+
+```js
+import { tableFromIPC } from "https://esm.sh/@uwdata/flechette@2";
+const table = tableFromIPC(new Uint8Array(model.get("_ipc").buffer));
+// table.numRows, table.numCols, table.get(i), table.getChild("col_name")
+```
+
+Use `traitlets.Any().tag(sync=True)` for the IPC bytes traitlet.
+
+## Reactive anywidgets in marimo
+
+When an anywidget trait (selection, value, zoom, etc.) should drive a
+downstream marimo cell, use `mo.state()` + `.observe()` on the **specific
+trait**. This is the preferred pattern:
+
+```python
+# In the cell that creates the widget:
+get_selection, set_selection = mo.state(widget.selection)
+widget.observe(
+    lambda _: set_selection(widget.selection),
+    names=["selection"],
+)
+
+# In a downstream cell — re-executes when selection changes:
+selection = get_selection()
+```
+
+Initialize `mo.state()` with the widget's current trait value — not a
+hardcoded default. Read the trait directly off the widget in the lambda.
+Do **not** use `change["new"]` or `allow_self_loops=True`.
+
+### `mo.state` + `.observe()` vs `mo.ui.anywidget()`
+
+Two strategies for reactive anywidgets. Choose one per widget — don't mix them.
+
+| Strategy                  | Reactivity                             | Best for                                               |
+| ------------------------- | -------------------------------------- | ------------------------------------------------------ |
+| `mo.state` + `.observe()` | Specific traits you pick               | Precision — only named traits trigger downstream cells |
+| `mo.ui.anywidget(widget)` | All synced traits as one `.value` dict | Convenience — observe everything at once               |
+
+### Programmatic widget control (scratchpad)
+
+Read widget state or set UI controls from the scratchpad — no clicking:
+
+```python
+print(timer.seconds)    # read
+timer.seconds = 0       # set — frontend updates automatically
+```
+
+`mo.ui.*` elements need `ctx.set_ui_value(...)` from code mode; anywidgets use
+direct assignment.
+
+## `_display_()` protocol
+
+Any object with a `_display_()` method renders richly in marimo. Return
+anything marimo can render — `mo.Html`, `mo.md()`, a chart, a string.
+
+Precedence: `_display_()` > built-in formatters > `_mime_()` > IPython
+`_repr_*_()` methods.
+
+```python
+from dataclasses import dataclass
+import marimo as mo
+
+@dataclass
+class ColorSwatch:
+    colors: list[str]
+
+    def _display_(self):
+        divs = "".join(
+            f'<div style="width:40px;height:40px;background:{c};border-radius:4px;"></div>'
+            for c in self.colors
+        )
+        return mo.Html(f'<div style="display:flex;gap:8px;">{divs}</div>')
+```
+
+For inline `<script>` tags, use `document.currentScript.previousElementSibling`
+to scope to the element — never hardcode IDs (breaks with multiple instances).
+
+## Minimize CLS (Cumulative Layout Shift)
+
+Use `min-height` or `aspect-ratio` on the outer container so the widget
+reserves space before content loads or when toggling between states.
diff --git a/marimo/_server/ai/skills/utils.py b/marimo/_server/ai/skills/utils.py
index c2d86808f9e..d3bfb906e4d 100644
--- a/marimo/_server/ai/skills/utils.py
+++ b/marimo/_server/ai/skills/utils.py
@@ -4,12 +4,33 @@
 from pathlib import Path
 from typing import Literal
 
-SKILL_NAMES: list[str] = ["marimo-pair"]
+SkillNameType = Literal["marimo-pair"]
+SKILL_NAMES: list[SkillNameType] = ["marimo-pair"]
+
+ReferenceNameType = Literal[
+    "gotchas", "rich-representations", "notebook-improvements"
+]
+REFERENCE_NAMES: list[ReferenceNameType] = [
+    "gotchas",
+    "rich-representations",
+    "notebook-improvements",
+]
 
 
 @functools.lru_cache(maxsize=1)
-def load_skill(skill_name: Literal["marimo-pair"]) -> str:
+def load_skill(skill_name: SkillNameType) -> str:
     """Read the bundled skill SKILL.md once and cache it."""
     return (Path(__file__).parent / skill_name / "SKILL.md").read_text(
         encoding="utf-8"
     )
+
+
+@functools.lru_cache(maxsize=3)
+def load_reference(reference_name: ReferenceNameType) -> str:
+    """Read the bundled reference file once and cache it."""
+    return (
+        Path(__file__).parent
+        / "marimo-pair"
+        / "references"
+        / f"{reference_name}.md"
+    ).read_text(encoding="utf-8")
diff --git a/marimo/_server/ai/tools/code_mode.py b/marimo/_server/ai/tools/code_mode.py
index 1609ba8e947..6ddceb670ea 100644
--- a/marimo/_server/ai/tools/code_mode.py
+++ b/marimo/_server/ai/tools/code_mode.py
@@ -4,12 +4,14 @@
 from typing import TYPE_CHECKING
 
 from marimo._ai._tools.types import CodeExecutionResult
+from marimo._server.ai.skills.utils import load_reference
 from marimo._server.api.deps import AppState
 from marimo._server.api.utils import get_code_mode_credentials
 from marimo._server.scratchpad import run_scratchpad_code
 
 if TYPE_CHECKING:
     from pydantic_ai import FunctionToolset
+    from pydantic_ai.capabilities import Capability
     from starlette.requests import Request
 
     from marimo._session.session import Session
@@ -53,3 +55,36 @@ async def execute_code(code: str) -> CodeExecutionResult:
         description=execute_code.__doc__,
     )
     return toolset
+
+
+def references_capability() -> list[Capability]:
+    from pydantic_ai.capabilities import Capability
+
+    gotchas_capability: Capability = Capability(
+        id="gotchas",
+        description=(
+            "Name redefinition, cached module proxies, and other notebook traps."
+        ),
+        instructions=load_reference("gotchas"),
+        defer_loading=True,
+    )
+
+    notebook_improvements_capability: Capability = Capability(
+        id="notebook-improvements",
+        description="Improving, optimizing, or cleaning up an existing notebook.",
+        instructions=load_reference("notebook-improvements"),
+        defer_loading=True,
+    )
+
+    rich_representations_capability: Capability = Capability(
+        id="rich-representations",
+        description="Custom widgets, visual encodings, and interactive output.",
+        instructions=load_reference("rich-representations"),
+        defer_loading=True,
+    )
+
+    return [
+        gotchas_capability,
+        notebook_improvements_capability,
+        rich_representations_capability,
+    ]
diff --git a/marimo/_server/api/endpoints/ai.py b/marimo/_server/api/endpoints/ai.py
index f93ba02c7b3..6f7f41c1bf2 100644
--- a/marimo/_server/api/endpoints/ai.py
+++ b/marimo/_server/api/endpoints/ai.py
@@ -209,7 +209,6 @@ async def ai_chat(
     mode: CopilotMode = ai_config.get("mode", "manual")
     system_prompt = get_chat_system_prompt(
         custom_rules=custom_rules,
-        context=body.context,
         include_other_code=body.include_other_code,
         mode=mode,
         session_id=session_id,
diff --git a/marimo/_server/models/completion.py b/marimo/_server/models/completion.py
index 355b2792f48..ce4bc57ae3e 100644
--- a/marimo/_server/models/completion.py
+++ b/marimo/_server/models/completion.py
@@ -66,7 +66,6 @@ class ChatRequest(msgspec.Struct, rename="camel"):
     See pydantic_ai.ui.vercel_ai.request_types.UIMessage or Vercel AI SDK documentation.
     """
 
-    context: AiCompletionContext
     include_other_code: str
     ui_messages: list[UIMessage]
     tools: list[ToolDefinition] | None = None
diff --git a/packages/openapi/api.yaml b/packages/openapi/api.yaml
index 47b24c0b1f2..c74348c7208 100644
--- a/packages/openapi/api.yaml
+++ b/packages/openapi/api.yaml
@@ -593,8 +593,6 @@ components:
 
         See pydantic_ai.ui.vercel_ai.request_types.UIMessage or Vercel AI SDK documentation.'
       properties:
-        context:
-          $ref: '#/components/schemas/AiCompletionContext'
         includeOtherCode:
           type: string
         model:
@@ -623,7 +621,6 @@ components:
           - type: 'null'
           default: null
       required:
-      - context
       - includeOtherCode
       - uiMessages
       title: ChatRequest
diff --git a/packages/openapi/src/api.ts b/packages/openapi/src/api.ts
index 035f4cb7ad7..65f0dc2cda9 100644
--- a/packages/openapi/src/api.ts
+++ b/packages/openapi/src/api.ts
@@ -3780,7 +3780,6 @@ export interface components {
      *     See pydantic_ai.ui.vercel_ai.request_types.UIMessage or Vercel AI SDK documentation.
      */
     ChatRequest: {
-      context: components["schemas"]["AiCompletionContext"];
       includeOtherCode: string;
       /** @default null */
       model?: string | null;
diff --git a/tests/_ai/test_pydantic_utils.py b/tests/_ai/test_pydantic_utils.py
index e3b79441611..6c428926719 100644
--- a/tests/_ai/test_pydantic_utils.py
+++ b/tests/_ai/test_pydantic_utils.py
@@ -304,6 +304,71 @@ def test_convert_message_without_metadata(self):
         result = convert_to_pydantic_messages(messages)
         assert result[0].metadata is None
 
+    def test_convert_expands_marimo_context_part(self):
+        from pydantic_ai.ui.vercel_ai.request_types import TextUIPart
+
+        messages = [
+            {
+                "id": "msg_ctx",
+                "role": "user",
+                "parts": [
+                    {"type": "text", "text": "What does df look like?"},
+                    {
+                        "type": "data-marimo-context",
+                        "data": {
+                            "plainText": "<variable name='df'>...</variable>",
+                            "contextIds": ["variable://df"],
+                        },
+                    },
+                ],
+            }
+        ]
+        result = convert_to_pydantic_messages(messages)
+        # The data part is lowered to a text part placed below the user text.
+        assert result[0].parts == [
+            TextUIPart(type="text", text="What does df look like?"),
+            TextUIPart(
+                type="text",
+                text="<context>\n<variable name='df'>...</variable>\n</context>",
+            ),
+        ]
+
+    def test_convert_drops_empty_marimo_context_part(self):
+        from pydantic_ai.ui.vercel_ai.request_types import TextUIPart
+
+        messages = [
+            {
+                "id": "msg_ctx",
+                "role": "user",
+                "parts": [
+                    {"type": "text", "text": "Hello"},
+                    {
+                        "type": "data-marimo-context",
+                        "data": {"plainText": "   ", "contextIds": []},
+                    },
+                ],
+            }
+        ]
+        result = convert_to_pydantic_messages(messages)
+        assert result[0].parts == [TextUIPart(type="text", text="Hello")]
+
+    def test_convert_drops_malformed_marimo_context_part(self):
+        from pydantic_ai.ui.vercel_ai.request_types import TextUIPart
+
+        # A malformed `data` payload (not an object) must not crash.
+        messages = [
+            {
+                "id": "msg_ctx",
+                "role": "user",
+                "parts": [
+                    {"type": "text", "text": "Hello"},
+                    {"type": "data-marimo-context", "data": "oops"},
+                ],
+            }
+        ]
+        result = convert_to_pydantic_messages(messages)
+        assert result[0].parts == [TextUIPart(type="text", text="Hello")]
+
     def test_convert_multiple_messages(self):
         messages = [
             {
diff --git a/tests/_server/ai/skills/test_skills_utils.py b/tests/_server/ai/skills/test_skills_utils.py
index 4464e5b46ee..f5777dfbe24 100644
--- a/tests/_server/ai/skills/test_skills_utils.py
+++ b/tests/_server/ai/skills/test_skills_utils.py
@@ -1,7 +1,12 @@
 # Copyright 2026 Marimo. All rights reserved.
 from __future__ import annotations
 
-from marimo._server.ai.skills.utils import SKILL_NAMES, load_skill
+from marimo._server.ai.skills.utils import (
+    REFERENCE_NAMES,
+    SKILL_NAMES,
+    load_reference,
+    load_skill,
+)
 
 
 def test_load_skill_returns_skill_contents() -> None:
@@ -19,3 +24,27 @@ def test_load_skill_is_cached() -> None:
 def test_all_declared_skills_are_loadable() -> None:
     for skill_name in SKILL_NAMES:
         assert load_skill(skill_name).strip()
+
+
+def test_load_reference_returns_reference_contents() -> None:
+    content = load_reference("gotchas")
+    assert content.strip()
+    assert content.startswith("# Gotchas")
+
+
+def test_load_reference_is_cached() -> None:
+    assert load_reference("gotchas") is load_reference("gotchas")
+
+
+def test_all_declared_references_are_loadable() -> None:
+    for reference_name in REFERENCE_NAMES:
+        assert load_reference(reference_name).strip()
+
+
+def test_reference_files_point_to_load_capability() -> None:
+    for reference_name in REFERENCE_NAMES:
+        content = load_reference(reference_name)
+        assert "load_capability" in content
+        assert f"`{reference_name}` capability" in content
+        assert "Do not read" in content
+        assert "from disk." in content
diff --git a/tests/_server/ai/snapshots/chat_system_prompts.txt b/tests/_server/ai/snapshots/chat_system_prompts.txt
index 50e10457797..94dc236ac10 100644
--- a/tests/_server/ai/snapshots/chat_system_prompts.txt
+++ b/tests/_server/ai/snapshots/chat_system_prompts.txt
@@ -297,476 +297,6 @@ chart
 ## Additional rules:
 Always be polite.
 
-==================== with variables ====================
-
-
-You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.
-Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.
-## Capabilities
-- Answer questions and provide guidance using only your internal knowledge and the notebook context provided by the user.
-
-## Limitations
-- You do NOT have access to any external tools, plugins, or APIs.
-- You may not perform any actions beyond generating text and code suggestions.
-
-Current notebook session ID: s_test. Use this session_id with tools that require it.
-
-Your goal is to do one of the following two things:
-
-1. Help users answer questions related to their notebook.
-2. Answer general-purpose questions unrelated to their particular notebook.
-
-It will be up to you to decide which of these you are doing based on what the user has told you. When unclear, ask clarifying questions to understand the user's intent before proceeding.
-
-The user may reference additional context in the form @kind://name. You can use this context to help you with the current task.
-
-You can respond with markdown, code, or a combination of both. You only work with two languages: Python and SQL.
-When responding in code, think of each block of code as a separate cell in the notebook.
-
-You have the following rules:
-
-- Do not import the same library twice.
-- Do not define a variable if it already exists. You may reference variables from other cells, but you may not define a variable if it already exists.
-
-# Marimo fundamentals
-
-Marimo is a reactive notebook that differs from traditional notebooks in key ways:
-- Cells execute automatically when their dependencies change
-- Variables cannot be redeclared across cells
-- The notebook forms a directed acyclic graph (DAG)
-- The last expression in a cell is automatically displayed
-- UI elements are reactive and update the notebook automatically
-
-Marimo's reactivity means:
-- When a variable changes, all cells that use that variable automatically re-execute
-- UI elements trigger updates when their values change without explicit callbacks
-- UI element values are accessed through `.value` attribute
-- You cannot access a UI element's value in the same cell where it's defined
-
-## Best Practices
-
-<data_handling>
-- Use polars for data manipulation
-- Implement proper data validation
-- Handle missing values appropriately
-- Use efficient data structures
-- A variable in the last expression of a cell is automatically displayed as a table
-</data_handling>
-
-<ui_elements>
-- Access UI element values with .value attribute (e.g., slider.value)
-- Create UI elements in one cell and reference them in later cells
-- Create intuitive layouts with mo.hstack(), mo.vstack(), and mo.tabs()
-- Prefer reactive updates over callbacks (marimo handles reactivity automatically)
-- Group related UI elements for better organization
-</ui_elements>
-
-## Available UI elements
-
-* `mo.ui.altair_chart(altair_chart)` - create a reactive Altair chart
-* `mo.ui.button(value=None, kind='primary')` - create a clickable button
-* `mo.ui.run_button(label=None, tooltip=None, kind='primary')` - create a button that runs code
-* `mo.ui.checkbox(label='', value=False)` - create a checkbox
-* `mo.ui.chat(placeholder='', value=None)` - create a chat interface
-* `mo.ui.date(value=None, label=None, full_width=False)` - create a date picker
-* `mo.ui.dropdown(options, value=None, label=None, full_width=False)` - create a dropdown menu
-* `mo.ui.file(label='', multiple=False, full_width=False)` - create a file upload element
-* `mo.ui.number(value=None, label=None, full_width=False)` - create a number input
-* `mo.ui.radio(options, value=None, label=None, full_width=False)` - create radio buttons
-* `mo.ui.refresh(options: List[str], default_interval: str)` - create a refresh control
-* `mo.ui.slider(start, stop, value=None, label=None, full_width=False, step=None)` - create a slider
-* `mo.ui.range_slider(start, stop, value=None, label=None, full_width=False, step=None)` - create a range slider
-* `mo.ui.table(data, columns=None, on_select=None, sortable=True, filterable=True)` - create an interactive table
-* `mo.ui.text(value='', label=None, full_width=False)` - create a text input
-* `mo.ui.text_area(value='', label=None, full_width=False)` - create a multi-line text input
-* `mo.ui.data_explorer(df)` - create an interactive dataframe explorer
-* `mo.ui.dataframe(df)` - display a dataframe with search, filter, and sort capabilities
-* `mo.ui.plotly(plotly_figure)` - create a reactive Plotly chart (supports scatter, treemap, and sunburst)
-* `mo.ui.tabs(elements: dict[str, mo.ui.Element])` - create a tabbed interface from a dictionary
-* `mo.ui.array(elements: list[mo.ui.Element])` - create an array of UI elements
-* `mo.ui.form(element: mo.ui.Element, label='', bordered=True)` - wrap an element in a form
-
-## Layout and utility functions
-
-* `mo.stop(predicate, output=None)` - stop execution conditionally
-* `mo.Html(html)` - display HTML
-* `mo.image(image)` - display an image
-* `mo.hstack(elements)` - stack elements horizontally
-* `mo.vstack(elements)` - stack elements vertically
-* `mo.tabs(elements)` - create a tabbed interface
-* `mo.mpl.interactive()` - make matplotlib plots interactive
-
-## Examples
-
-<example title="Basic UI with reactivity">
-import marimo as mo
-import altair as alt
-import polars as pl
-import numpy as np
-
-# Create a slider and display it
-n_points = mo.ui.slider(10, 100, value=50, label="Number of points")
-n_points  # Display the slider
-
-# Generate random data based on slider value
-# This cell automatically re-executes when n_points.value changes
-x = np.random.rand(n_points.value)
-y = np.random.rand(n_points.value)
-
-df = pl.DataFrame({"x": x, "y": y})
-
-chart = alt.Chart(df).mark_circle(opacity=0.7).encode(
-    x=alt.X('x', title='X axis'),
-    y=alt.Y('y', title='Y axis')
-).properties(
-    title=f"Scatter plot with {n_points.value} points",
-    width=400,
-    height=300
-)
-
-chart
-</example>
-
-## Rules for python:
-1. For matplotlib: use plt.gca() as the last expression instead of plt.show().
-2. For plotly: return the figure object directly.
-3. For altair: return the chart object directly. Add tooltips where appropriate. You can pass polars dataframes directly to altair (e.g., alt.Chart(df)).
-4. Include proper labels, titles, and color schemes.
-5. Make visualizations interactive where appropriate.
-6. If an import already exists, do not import it again.
-7. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
-
-## Rules for markdown:
-1. Use double dollar signs ($$) for ALL mathematical expressions (both inline and block) in markdown responses.
-2. Example: $$E=mc^2$$
-3. Do NOT use single dollar signs or square brackets for math.
-
-## Rules for sql:
-1. The SQL must use duckdb syntax.
-
-## Available variables from other cells:
-- variable: `var1`- variable: `var2`
-
-==================== with VariableContext objects ====================
-
-
-You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.
-Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.
-## Capabilities
-- Answer questions and provide guidance using only your internal knowledge and the notebook context provided by the user.
-
-## Limitations
-- You do NOT have access to any external tools, plugins, or APIs.
-- You may not perform any actions beyond generating text and code suggestions.
-
-Current notebook session ID: s_test. Use this session_id with tools that require it.
-
-Your goal is to do one of the following two things:
-
-1. Help users answer questions related to their notebook.
-2. Answer general-purpose questions unrelated to their particular notebook.
-
-It will be up to you to decide which of these you are doing based on what the user has told you. When unclear, ask clarifying questions to understand the user's intent before proceeding.
-
-The user may reference additional context in the form @kind://name. You can use this context to help you with the current task.
-
-You can respond with markdown, code, or a combination of both. You only work with two languages: Python and SQL.
-When responding in code, think of each block of code as a separate cell in the notebook.
-
-You have the following rules:
-
-- Do not import the same library twice.
-- Do not define a variable if it already exists. You may reference variables from other cells, but you may not define a variable if it already exists.
-
-# Marimo fundamentals
-
-Marimo is a reactive notebook that differs from traditional notebooks in key ways:
-- Cells execute automatically when their dependencies change
-- Variables cannot be redeclared across cells
-- The notebook forms a directed acyclic graph (DAG)
-- The last expression in a cell is automatically displayed
-- UI elements are reactive and update the notebook automatically
-
-Marimo's reactivity means:
-- When a variable changes, all cells that use that variable automatically re-execute
-- UI elements trigger updates when their values change without explicit callbacks
-- UI element values are accessed through `.value` attribute
-- You cannot access a UI element's value in the same cell where it's defined
-
-## Best Practices
-
-<data_handling>
-- Use polars for data manipulation
-- Implement proper data validation
-- Handle missing values appropriately
-- Use efficient data structures
-- A variable in the last expression of a cell is automatically displayed as a table
-</data_handling>
-
-<ui_elements>
-- Access UI element values with .value attribute (e.g., slider.value)
-- Create UI elements in one cell and reference them in later cells
-- Create intuitive layouts with mo.hstack(), mo.vstack(), and mo.tabs()
-- Prefer reactive updates over callbacks (marimo handles reactivity automatically)
-- Group related UI elements for better organization
-</ui_elements>
-
-## Available UI elements
-
-* `mo.ui.altair_chart(altair_chart)` - create a reactive Altair chart
-* `mo.ui.button(value=None, kind='primary')` - create a clickable button
-* `mo.ui.run_button(label=None, tooltip=None, kind='primary')` - create a button that runs code
-* `mo.ui.checkbox(label='', value=False)` - create a checkbox
-* `mo.ui.chat(placeholder='', value=None)` - create a chat interface
-* `mo.ui.date(value=None, label=None, full_width=False)` - create a date picker
-* `mo.ui.dropdown(options, value=None, label=None, full_width=False)` - create a dropdown menu
-* `mo.ui.file(label='', multiple=False, full_width=False)` - create a file upload element
-* `mo.ui.number(value=None, label=None, full_width=False)` - create a number input
-* `mo.ui.radio(options, value=None, label=None, full_width=False)` - create radio buttons
-* `mo.ui.refresh(options: List[str], default_interval: str)` - create a refresh control
-* `mo.ui.slider(start, stop, value=None, label=None, full_width=False, step=None)` - create a slider
-* `mo.ui.range_slider(start, stop, value=None, label=None, full_width=False, step=None)` - create a range slider
-* `mo.ui.table(data, columns=None, on_select=None, sortable=True, filterable=True)` - create an interactive table
-* `mo.ui.text(value='', label=None, full_width=False)` - create a text input
-* `mo.ui.text_area(value='', label=None, full_width=False)` - create a multi-line text input
-* `mo.ui.data_explorer(df)` - create an interactive dataframe explorer
-* `mo.ui.dataframe(df)` - display a dataframe with search, filter, and sort capabilities
-* `mo.ui.plotly(plotly_figure)` - create a reactive Plotly chart (supports scatter, treemap, and sunburst)
-* `mo.ui.tabs(elements: dict[str, mo.ui.Element])` - create a tabbed interface from a dictionary
-* `mo.ui.array(elements: list[mo.ui.Element])` - create an array of UI elements
-* `mo.ui.form(element: mo.ui.Element, label='', bordered=True)` - wrap an element in a form
-
-## Layout and utility functions
-
-* `mo.stop(predicate, output=None)` - stop execution conditionally
-* `mo.Html(html)` - display HTML
-* `mo.image(image)` - display an image
-* `mo.hstack(elements)` - stack elements horizontally
-* `mo.vstack(elements)` - stack elements vertically
-* `mo.tabs(elements)` - create a tabbed interface
-* `mo.mpl.interactive()` - make matplotlib plots interactive
-
-## Examples
-
-<example title="Basic UI with reactivity">
-import marimo as mo
-import altair as alt
-import polars as pl
-import numpy as np
-
-# Create a slider and display it
-n_points = mo.ui.slider(10, 100, value=50, label="Number of points")
-n_points  # Display the slider
-
-# Generate random data based on slider value
-# This cell automatically re-executes when n_points.value changes
-x = np.random.rand(n_points.value)
-y = np.random.rand(n_points.value)
-
-df = pl.DataFrame({"x": x, "y": y})
-
-chart = alt.Chart(df).mark_circle(opacity=0.7).encode(
-    x=alt.X('x', title='X axis'),
-    y=alt.Y('y', title='Y axis')
-).properties(
-    title=f"Scatter plot with {n_points.value} points",
-    width=400,
-    height=300
-)
-
-chart
-</example>
-
-## Rules for python:
-1. For matplotlib: use plt.gca() as the last expression instead of plt.show().
-2. For plotly: return the figure object directly.
-3. For altair: return the chart object directly. Add tooltips where appropriate. You can pass polars dataframes directly to altair (e.g., alt.Chart(df)).
-4. Include proper labels, titles, and color schemes.
-5. Make visualizations interactive where appropriate.
-6. If an import already exists, do not import it again.
-7. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
-
-## Rules for markdown:
-1. Use double dollar signs ($$) for ALL mathematical expressions (both inline and block) in markdown responses.
-2. Example: $$E=mc^2$$
-3. Do NOT use single dollar signs or square brackets for math.
-
-## Rules for sql:
-1. The SQL must use duckdb syntax.
-
-## Available variables from other cells:
-- variable: `df`
-  - value_type: DataFrame
-  - value_preview: <DataFrame with 100 rows and 5 columns>
-- variable: `model`
-  - value_type: Model
-  - value_preview: <Model object>
-
-
-==================== with context ====================
-
-
-You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.
-Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.
-## Capabilities
-- Answer questions and provide guidance using only your internal knowledge and the notebook context provided by the user.
-
-## Limitations
-- You do NOT have access to any external tools, plugins, or APIs.
-- You may not perform any actions beyond generating text and code suggestions.
-
-Current notebook session ID: s_test. Use this session_id with tools that require it.
-
-Your goal is to do one of the following two things:
-
-1. Help users answer questions related to their notebook.
-2. Answer general-purpose questions unrelated to their particular notebook.
-
-It will be up to you to decide which of these you are doing based on what the user has told you. When unclear, ask clarifying questions to understand the user's intent before proceeding.
-
-The user may reference additional context in the form @kind://name. You can use this context to help you with the current task.
-
-You can respond with markdown, code, or a combination of both. You only work with two languages: Python and SQL.
-When responding in code, think of each block of code as a separate cell in the notebook.
-
-You have the following rules:
-
-- Do not import the same library twice.
-- Do not define a variable if it already exists. You may reference variables from other cells, but you may not define a variable if it already exists.
-
-# Marimo fundamentals
-
-Marimo is a reactive notebook that differs from traditional notebooks in key ways:
-- Cells execute automatically when their dependencies change
-- Variables cannot be redeclared across cells
-- The notebook forms a directed acyclic graph (DAG)
-- The last expression in a cell is automatically displayed
-- UI elements are reactive and update the notebook automatically
-
-Marimo's reactivity means:
-- When a variable changes, all cells that use that variable automatically re-execute
-- UI elements trigger updates when their values change without explicit callbacks
-- UI element values are accessed through `.value` attribute
-- You cannot access a UI element's value in the same cell where it's defined
-
-## Best Practices
-
-<data_handling>
-- Use polars for data manipulation
-- Implement proper data validation
-- Handle missing values appropriately
-- Use efficient data structures
-- A variable in the last expression of a cell is automatically displayed as a table
-</data_handling>
-
-<ui_elements>
-- Access UI element values with .value attribute (e.g., slider.value)
-- Create UI elements in one cell and reference them in later cells
-- Create intuitive layouts with mo.hstack(), mo.vstack(), and mo.tabs()
-- Prefer reactive updates over callbacks (marimo handles reactivity automatically)
-- Group related UI elements for better organization
-</ui_elements>
-
-## Available UI elements
-
-* `mo.ui.altair_chart(altair_chart)` - create a reactive Altair chart
-* `mo.ui.button(value=None, kind='primary')` - create a clickable button
-* `mo.ui.run_button(label=None, tooltip=None, kind='primary')` - create a button that runs code
-* `mo.ui.checkbox(label='', value=False)` - create a checkbox
-* `mo.ui.chat(placeholder='', value=None)` - create a chat interface
-* `mo.ui.date(value=None, label=None, full_width=False)` - create a date picker
-* `mo.ui.dropdown(options, value=None, label=None, full_width=False)` - create a dropdown menu
-* `mo.ui.file(label='', multiple=False, full_width=False)` - create a file upload element
-* `mo.ui.number(value=None, label=None, full_width=False)` - create a number input
-* `mo.ui.radio(options, value=None, label=None, full_width=False)` - create radio buttons
-* `mo.ui.refresh(options: List[str], default_interval: str)` - create a refresh control
-* `mo.ui.slider(start, stop, value=None, label=None, full_width=False, step=None)` - create a slider
-* `mo.ui.range_slider(start, stop, value=None, label=None, full_width=False, step=None)` - create a range slider
-* `mo.ui.table(data, columns=None, on_select=None, sortable=True, filterable=True)` - create an interactive table
-* `mo.ui.text(value='', label=None, full_width=False)` - create a text input
-* `mo.ui.text_area(value='', label=None, full_width=False)` - create a multi-line text input
-* `mo.ui.data_explorer(df)` - create an interactive dataframe explorer
-* `mo.ui.dataframe(df)` - display a dataframe with search, filter, and sort capabilities
-* `mo.ui.plotly(plotly_figure)` - create a reactive Plotly chart (supports scatter, treemap, and sunburst)
-* `mo.ui.tabs(elements: dict[str, mo.ui.Element])` - create a tabbed interface from a dictionary
-* `mo.ui.array(elements: list[mo.ui.Element])` - create an array of UI elements
-* `mo.ui.form(element: mo.ui.Element, label='', bordered=True)` - wrap an element in a form
-
-## Layout and utility functions
-
-* `mo.stop(predicate, output=None)` - stop execution conditionally
-* `mo.Html(html)` - display HTML
-* `mo.image(image)` - display an image
-* `mo.hstack(elements)` - stack elements horizontally
-* `mo.vstack(elements)` - stack elements vertically
-* `mo.tabs(elements)` - create a tabbed interface
-* `mo.mpl.interactive()` - make matplotlib plots interactive
-
-## Examples
-
-<example title="Basic UI with reactivity">
-import marimo as mo
-import altair as alt
-import polars as pl
-import numpy as np
-
-# Create a slider and display it
-n_points = mo.ui.slider(10, 100, value=50, label="Number of points")
-n_points  # Display the slider
-
-# Generate random data based on slider value
-# This cell automatically re-executes when n_points.value changes
-x = np.random.rand(n_points.value)
-y = np.random.rand(n_points.value)
-
-df = pl.DataFrame({"x": x, "y": y})
-
-chart = alt.Chart(df).mark_circle(opacity=0.7).encode(
-    x=alt.X('x', title='X axis'),
-    y=alt.Y('y', title='Y axis')
-).properties(
-    title=f"Scatter plot with {n_points.value} points",
-    width=400,
-    height=300
-)
-
-chart
-</example>
-
-## Rules for python:
-1. For matplotlib: use plt.gca() as the last expression instead of plt.show().
-2. For plotly: return the figure object directly.
-3. For altair: return the chart object directly. Add tooltips where appropriate. You can pass polars dataframes directly to altair (e.g., alt.Chart(df)).
-4. Include proper labels, titles, and color schemes.
-5. Make visualizations interactive where appropriate.
-6. If an import already exists, do not import it again.
-7. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
-
-## Rules for markdown:
-1. Use double dollar signs ($$) for ALL mathematical expressions (both inline and block) in markdown responses.
-2. Example: $$E=mc^2$$
-3. Do NOT use single dollar signs or square brackets for math.
-
-## Rules for sql:
-1. The SQL must use duckdb syntax.
-
-## Available schema:
-- Table: df_1
-  - Column: age
-    - Type: int
-    - Sample values: 1, 2, 3
-  - Column: name
-    - Type: str
-    - Sample values: Alice, Bob, Charlie
-- Table: d2_2
-  - Column: a
-    - Type: int
-    - Sample values: 1, 2, 3
-  - Column: b
-    - Type: int
-    - Sample values: 4, 5, 6
-
-
 ==================== with other code ====================
 
 
@@ -1236,16 +766,4 @@ Always be polite.
 <code_from_other_cells>
 import pandas as pd
 import numpy as np
-</code_from_other_cells>
-
-## Available variables from other cells:
-- variable: `var1`- variable: `var2`
-
-## Available schema:
-- Table: df_1
-  - Column: age
-    - Type: int
-    - Sample values: 1, 2, 3
-  - Column: name
-    - Type: str
-    - Sample values: Alice, Bob, Charlie
+</code_from_other_cells>
\ No newline at end of file
diff --git a/tests/_server/ai/test_prompts.py b/tests/_server/ai/test_prompts.py
index 3be1bdd0431..ab974b2007d 100644
--- a/tests/_server/ai/test_prompts.py
+++ b/tests/_server/ai/test_prompts.py
@@ -226,7 +226,6 @@ def test_chat_system_prompts():
     result += get_chat_system_prompt(
         custom_rules=None,
         include_other_code="",
-        context=None,
         mode="manual",
         session_id=SessionId("s_test"),  # stable fake session id for snapshot
     )
@@ -235,76 +234,6 @@ def test_chat_system_prompts():
     result += get_chat_system_prompt(
         custom_rules="Always be polite.",
         include_other_code="",
-        context=None,
-        mode="manual",
-        session_id=SessionId("s_test"),
-    )
-
-    result += _header("with variables")
-    result += get_chat_system_prompt(
-        custom_rules=None,
-        include_other_code="",
-        context=AiCompletionContext(
-            variables=["var1", "var2"],
-        ),
-        mode="manual",
-        session_id=SessionId("s_test"),
-    )
-
-    result += _header("with VariableContext objects")
-    result += get_chat_system_prompt(
-        custom_rules=None,
-        include_other_code="",
-        context=AiCompletionContext(
-            variables=[
-                VariableContext(
-                    name="df",
-                    value_type="DataFrame",
-                    preview_value="<DataFrame with 100 rows and 5 columns>",
-                ),
-                VariableContext(
-                    name="model",
-                    value_type="Model",
-                    preview_value="<Model object>",
-                ),
-            ]
-        ),
-        mode="manual",
-        session_id=SessionId("s_test"),
-    )
-
-    result += _header("with context")
-    result += get_chat_system_prompt(
-        custom_rules=None,
-        include_other_code="",
-        context=AiCompletionContext(
-            schema=[
-                SchemaTable(
-                    name="df_1",
-                    columns=[
-                        SchemaColumn(
-                            "age", "int", sample_values=["1", "2", "3"]
-                        ),
-                        SchemaColumn(
-                            "name",
-                            "str",
-                            sample_values=["Alice", "Bob", "Charlie"],
-                        ),
-                    ],
-                ),
-                SchemaTable(
-                    name="d2_2",
-                    columns=[
-                        SchemaColumn(
-                            "a", "int", sample_values=["1", "2", "3"]
-                        ),
-                        SchemaColumn(
-                            "b", "int", sample_values=["4", "5", "6"]
-                        ),
-                    ],
-                ),
-            ],
-        ),
         mode="manual",
         session_id=SessionId("s_test"),
     )
@@ -313,7 +242,6 @@ def test_chat_system_prompts():
     result += get_chat_system_prompt(
         custom_rules=None,
         include_other_code="import pandas as pd\nimport numpy as np\n",
-        context=None,
         mode="manual",
         session_id=SessionId("s_test"),
     )
@@ -322,7 +250,6 @@ def test_chat_system_prompts():
     result += get_chat_system_prompt(
         custom_rules=None,
         include_other_code="",
-        context=None,
         mode="agent",
         session_id=SessionId("s_test"),
     )
@@ -331,24 +258,6 @@ def test_chat_system_prompts():
     result += get_chat_system_prompt(
         custom_rules="Always be polite.",
         include_other_code="import pandas as pd\nimport numpy as np\n",
-        context=AiCompletionContext(
-            variables=["var1", "var2"],
-            schema=[
-                SchemaTable(
-                    name="df_1",
-                    columns=[
-                        SchemaColumn(
-                            "age", "int", sample_values=["1", "2", "3"]
-                        ),
-                        SchemaColumn(
-                            "name",
-                            "str",
-                            sample_values=["Alice", "Bob", "Charlie"],
-                        ),
-                    ],
-                ),
-            ],
-        ),
         mode="manual",
         session_id=SessionId("s_test"),
     )
@@ -381,7 +290,6 @@ def test_markdown_rules_include_latex():
     chat_prompt = get_chat_system_prompt(
         custom_rules=None,
         include_other_code="",
-        context=None,
         mode="manual",
         session_id=SessionId("test"),
     )
@@ -483,17 +391,11 @@ def test_mode_intro_messages_share_base(mode: CopilotMode):
 
 def test_common_chat_sections_empty():
     assert (
-        _common_chat_sections(
-            custom_rules=None, include_other_code="", context=None
-        )
-        == ""
+        _common_chat_sections(custom_rules=None, include_other_code="") == ""
     )
     # Whitespace-only custom rules are treated as empty.
     assert (
-        _common_chat_sections(
-            custom_rules="   ", include_other_code="", context=None
-        )
-        == ""
+        _common_chat_sections(custom_rules="   ", include_other_code="") == ""
     )
 
 
@@ -501,19 +403,16 @@ def test_common_chat_sections_full():
     result = _common_chat_sections(
         custom_rules="Be concise.",
         include_other_code="import polars as pl",
-        context=AiCompletionContext(variables=["var1"]),
     )
     assert "## Additional rules:\nBe concise." in result
     assert "<code_from_other_cells>" in result
     assert "import polars as pl" in result
-    assert "## Available variables from other cells:" in result
 
 
 def test_chat_system_prompt_code_mode():
     prompt = get_chat_system_prompt(
         custom_rules=None,
         include_other_code="",
-        context=None,
         mode="code_mode",
         session_id=SessionId("s_test"),
     )
@@ -521,22 +420,22 @@ def test_chat_system_prompt_code_mode():
     assert _get_mode_intro_message("code_mode") in prompt
     assert "how to work with marimo" in prompt
     assert load_skill("marimo-pair") in prompt
-    # Code mode includes single-cell language rules.
-    assert "## Rules for python:" in prompt
+    assert "load_capability" in prompt
+    assert "`gotchas`" in prompt
+    assert "references/gotchas.md" not in prompt
 
 
 def test_chat_system_prompt_code_mode_includes_extras():
     prompt = get_chat_system_prompt(
         custom_rules="Always be polite.",
         include_other_code="import pandas as pd\n",
-        context=AiCompletionContext(variables=["var1", "var2"]),
         mode="code_mode",
         session_id=SessionId("s_test"),
     )
     assert "## Additional rules:\nAlways be polite." in prompt
-    assert "<code_from_other_cells>" in prompt
-    assert "import pandas as pd" in prompt
-    assert "## Available variables from other cells:" in prompt
+    # Code mode inspects code via tools, so other-cell code is not inlined.
+    assert "<code_from_other_cells>" not in prompt
+    assert "import pandas as pd" not in prompt
 
 
 def test_chat_system_prompt_non_code_mode_includes_session_info():
@@ -544,7 +443,6 @@ def test_chat_system_prompt_non_code_mode_includes_session_info():
         prompt = get_chat_system_prompt(
             custom_rules=None,
             include_other_code="",
-            context=None,
             mode=cast(CopilotMode, mode),
             session_id=SessionId("s_abc"),
         )
@@ -558,14 +456,12 @@ def test_chat_system_prompt_agent_mode_inserts_cell_rules():
     agent_prompt = get_chat_system_prompt(
         custom_rules=None,
         include_other_code="",
-        context=None,
         mode="agent",
         session_id=SessionId("s_test"),
     )
     manual_prompt = get_chat_system_prompt(
         custom_rules=None,
         include_other_code="",
-        context=None,
         mode="manual",
         session_id=SessionId("s_test"),
     )
diff --git a/tests/_server/ai/test_providers.py b/tests/_server/ai/test_providers.py
index 4a096f2127c..d1030698919 100644
--- a/tests/_server/ai/test_providers.py
+++ b/tests/_server/ai/test_providers.py
@@ -713,7 +713,7 @@ async def test_stream_completion_harness_wires_execute_code_toolset() -> None:
         )
 
     assert result is streaming_response
-    assert stream_options.span_info.tool_count == 1
+    assert stream_options.span_info.tool_count == 4
     # The toolset is bound to the caller's session and request.
     mock_build_toolset.assert_called_once_with(session, request)
 
@@ -722,3 +722,11 @@ async def test_stream_completion_harness_wires_execute_code_toolset() -> None:
     agent_kwargs = mock_agent.call_args.kwargs
     assert agent_kwargs["toolsets"] == [toolset]
     assert agent_kwargs["instructions"] == "SYSTEM PROMPT WITH SKILL"
+    capabilities = agent_kwargs["capabilities"]
+    assert len(capabilities) == 3
+    assert {capability.id for capability in capabilities} == {
+        "gotchas",
+        "notebook-improvements",
+        "rich-representations",
+    }
+    assert all(capability.defer_loading for capability in capabilities)
diff --git a/tests/_server/ai/tools/test_code_mode.py b/tests/_server/ai/tools/test_code_mode.py
index be1b3dd430a..82b68c6adf5 100644
--- a/tests/_server/ai/tools/test_code_mode.py
+++ b/tests/_server/ai/tools/test_code_mode.py
@@ -62,3 +62,27 @@ async def test_execute_code_tool_routes_to_scratchpad_with_credentials() -> (
         server_url="http://localhost:2718",
         auth_token="secret-token",
     )
+
+
+@pytest.mark.requires("pydantic_ai")
+def test_references_capability_exposes_deferred_reference_bundles() -> None:
+    from marimo._server.ai.skills.utils import load_reference
+    from marimo._server.ai.tools.code_mode import references_capability
+
+    capabilities = references_capability()
+    by_id = {capability.id: capability for capability in capabilities}
+
+    assert set(by_id) == {
+        "gotchas",
+        "notebook-improvements",
+        "rich-representations",
+    }
+    for capability in capabilities:
+        assert capability.defer_loading is True
+        assert capability.get_instructions() == [load_reference(capability.id)]
+
+    assert "Name redefinition" in by_id["gotchas"].description
+    assert (
+        "Improving, optimizing" in by_id["notebook-improvements"].description
+    )
+    assert "Custom widgets" in by_id["rich-representations"].description