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

[Light2Dark]

📝 Summary

Code mode had no way to load deeper reference docs without bloating the system prompt. I split the reference material into deferred pydantic-ai capabilities so the model pulls in gotchas, notebook-improvements, and rich-representations only when needed.

• Bundle three reference markdown files and expose them as deferred Capability objects wired into the code-mode agent harness
• Update SKILL.md and reference docs to tell the model to use load_capability instead of reading files from disk
• Streamline the code-mode system prompt: drop per-language rules and injected cell code (the model can inspect the notebook via execute_code), and add a nudge to inspect notebook state first
• Expose code mode in the chat mode picker outside dev builds

📋 Pre-Review Checklist

• For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on Discord, or the community discussions (Please provide a link if applicable).
• Any AI generated code has been reviewed line-by-line by the human PR author, who stands by it.
• Video or media evidence is provided for any visual changes (optional).

✅ Merge Checklist

• I have read the contributor guidelines.
• Documentation has been updated where applicable, including docstrings for API changes.
• Tests have been added for the changes made.

Made with Cursor

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
marimo-docs	Ready	Preview, Comment	Jun 25, 2026 3:50pm

[cubic-dev-ai]

No issues found across 13 files

Architecture diagram

sequenceDiagram
    participant Client as Chat Client (Frontend)
    participant API as API Route (stream_completion)
    participant Provider as AI Provider (stream_completion_harness)
    participant Agent as pydantic-ai Agent
    participant Capability as Deferred Capability
    participant Tool as execute_code Tool
    participant Session as Kernel Session
    participant SkillLoader as Skill/Reference Loader

    Note over Client,SkillLoader: Code Mode Chat Flow - NEW: Chat Mode Picker Unrestricted

    Client->>API: POST /api/ai/completion (mode=code_mode)
    
    Note over API,Session: NEW: Code mode no longer receives cell code injection

    Provider->>Provider: Build system prompt (no language rules, no other cell code)

    Provider->>SkillLoader: Load SKILL.md (marimo-pair)
    SkillLoader-->>Provider: SKILL.md content (with load_capability instructions)

    Provider->>Provider: Create pydantic-ai model

    Provider->>Capability: NEW: references_capability()
    Capability->>SkillLoader: Load gotchas.md (deferred)
    Capability->>SkillLoader: Load notebook-improvements.md (deferred)
    Capability->>SkillLoader: Load rich-representations.md (deferred)
    SkillLoader-->>Capability: Reference content (cached, not inlined)

    Provider->>Agent: Create Agent(toolsets=[build_execute_code_toolset], capabilities=3 deferred refs, instructions=system_prompt)

    Provider->>Agent: Submit message

    alt Model needs reference material
        Agent->>Capability: load_capability("gotchas")
        Capability-->>Agent: Gotchas content (instructions)
        Agent->>Capability: load_capability("rich-representations")
        Capability-->>Agent: Rich representations content
    end

    alt Model needs to inspect/run code
        Agent->>Tool: execute_code(code)
        Tool->>Session: Run scratchpad code with credentials
        Session-->>Tool: Execution result
        Tool-->>Agent: CodeExecutionResult
    end

    Agent-->>Provider: Stream response tokens
    Provider->>API: StreamingResponse
    API-->>Client: Stream chat response

    Note over Client,Agent: NEW: Reference docs loaded on-demand, not inline in prompt

Loading

_{Re-trigger cubic}

[mscolnick]

i wonder if you can make this dynamic from reading the references file. not necessary now, but maybe in future if we vendor this from marimo-pair

[cubic-dev-ai]

2 issues found across 12 files (changes from recent commits).

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="frontend/src/components/chat/chat-panel.tsx">

<violation number="1" location="frontend/src/components/chat/chat-panel.tsx:537">
P2: Potential loss of legacy context resolution for old persisted chats after removing `buildCompletionRequestBody`. The old helper re-scanned all historical message text for `@` references on every completion request (regenerate/retry). The new code only preserves context for messages that already contain the new `data-marimo-context` part. Old persisted chats without that part will silently lose `@` context on regeneration.</violation>
</file>

<file name="frontend/src/components/editor/ai/completion-utils.ts">

<violation number="1" location="frontend/src/components/editor/ai/completion-utils.ts:95">
P2: New `resolveChatContext` duplicates core parsing/attachment logic from `getAICompletionBodyWithAttachments`, creating drift risk. Only the new path stamps `providerMetadata.marimo.source = "context"`, so `isContextAttachment`-based filtering behaves inconsistently across chat and completion attachments. Extract a shared helper for the common registry parsing/attachment fetching logic.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

P2: Potential loss of legacy context resolution for old persisted chats after removing buildCompletionRequestBody. The old helper re-scanned all historical message text for @ references on every completion request (regenerate/retry). The new code only preserves context for messages that already contain the new data-marimo-context part. Old persisted chats without that part will silently lose @ context on regeneration.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At frontend/src/components/chat/chat-panel.tsx, line 537:

<comment>Potential loss of legacy context resolution for old persisted chats after removing `buildCompletionRequestBody`. The old helper re-scanned all historical message text for `@` references on every completion request (regenerate/retry). The new code only preserves context for messages that already contain the new `data-marimo-context` part. Old persisted chats without that part will silently lose `@` context on regeneration.</comment>

<file context>
@@ -532,9 +534,10 @@ const ChatPanelBody = () => {
-        const completionBody = await buildCompletionRequestBody(
-          options.messages,
-        );
+        const completionBody = {
+          uiMessages: options.messages,
+          includeOtherCode: getCodes(""),
</file context>

[cubic-dev-ai]

P2: New resolveChatContext duplicates core parsing/attachment logic from getAICompletionBodyWithAttachments, creating drift risk. Only the new path stamps providerMetadata.marimo.source = "context", so isContextAttachment-based filtering behaves inconsistently across chat and completion attachments. Extract a shared helper for the common registry parsing/attachment fetching logic.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At frontend/src/components/editor/ai/completion-utils.ts, line 95:

<comment>New `resolveChatContext` duplicates core parsing/attachment logic from `getAICompletionBodyWithAttachments`, creating drift risk. Only the new path stamps `providerMetadata.marimo.source = "context"`, so `isContextAttachment`-based filtering behaves inconsistently across chat and completion attachments. Extract a shared helper for the common registry parsing/attachment fetching logic.</comment>

<file context>
@@ -50,6 +50,91 @@ export function getAICompletionBody({
+ * 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> {
</file context>

[github-actions]

🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.23.12-dev3