improvement(ai): stamp context attachments

[Light2Dark]

📝 Summary

Follow-up to #9985. That PR introduced two parallel paths for resolving
@-context out of a prompt — resolveChatContext (chat) and
getAICompletionBodyWithAttachments

Only the chat path stamped its context-derived file attachments with
providerMetadata.marimo.source = "context". So isContextAttachment — the
predicate used to dedupe context attachments on message edit — silently
returned false for completion-path attachments. That's a latent footgun and a
drift risk between the two code paths.

Why stamp attachments at all?

Some @-mentions resolve to file attachments (e.g. a cell's image output) that
get appended to the user message alongside files the user uploaded by hand.
Once in the message the two are indistinguishable. The stamp lets
handleMessageEdit re-resolve fresh context on edit while dropping only stale
context attachments and preserving the user's own uploads.

📋 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.

[vercel]

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

Project	Deployment	Actions	Updated (UTC)
marimo-docs	Ready	Preview, Comment	Jun 26, 2026 2:57am

[cubic-dev-ai]

No issues found across 2 files

Architecture diagram

sequenceDiagram
    participant UI as Chat UI / Completion Trigger
    participant Resolve as resolveChatContext
    participant CompResolve as getAICompletionBodyWithAttachments
    participant Shared as resolveContextAttachments (NEW)
    participant Stamp as stampContextAttachment (NEW)
    participant Registry as AI Context Registry
    participant Store as Jotai Store

    Note over UI,Store: Context attachment resolution (both paths)

    UI->>Resolve: User edits message with @-mentions
    UI->>CompResolve: "Generate with AI" triggered

    Resolve->>Shared: resolveContextAttachments(input)
    CompResolve->>Shared: resolveContextAttachments(input)

    alt Input contains "@"
        Shared->>Store: getAIContextRegistry(store)
        Store-->>Shared: registry instance
        Shared->>Registry: parseAllContextIds(input)
        Registry-->>Shared: contextLocatorIds[]
        alt Has context IDs
            Shared->>Registry: formatContextForAI(contextIds)
            Registry-->>Shared: plain text context
            Shared->>Registry: getAttachmentsForContext(contextIds)
            alt Success
                Registry-->>Shared: FileUIPart[]
                Shared->>Stamp: stampContextAttachment(attachment) (for each)
                Stamp->>Stamp: Set providerMetadata.marimo.source = "context"
                Stamp-->>Shared: Stamped FileUIPart
            else Error
                Registry-->>Shared: Error
                Shared->>Shared: Log error, return no attachments
            end
        else No context IDs
            Shared-->>Resolve: empty result
            Shared-->>CompResolve: empty result
        end
    else No "@" in input
        Shared-->>Resolve: empty result
        Shared-->>CompResolve: empty result
    end

    Shared-->>Resolve: { plainText, contextIds, attachments }
    Shared-->>CompResolve: { plainText, attachments }

    Note over Resolve: Build MarimoContextUIPart
    Resolve-->>UI: ResolvedChatContext

    Note over CompResolve: Build AICompletionBodyWithAttachments
    CompResolve-->>UI: body with context.plainText + attachments

Loading

_{Re-trigger cubic}