Skip to content

feat(table): allow hosts to size-gate downloads#9510

Merged
kirangadhave merged 5 commits into
mainfrom
kg/table-download-size-gate
May 12, 2026
Merged

feat(table): allow hosts to size-gate downloads#9510
kirangadhave merged 5 commits into
mainfrom
kg/table-download-size-gate

Conversation

@kirangadhave
Copy link
Copy Markdown
Member

@kirangadhave kirangadhave commented May 11, 2026

This is step 1 toward fixing marimo-team/marimo-lsp#542; the companion change will live in marimo-lsp, which will set the atom to a concrete VS Code–appropriate limit at renderer init.

Summary

Adds size_bytes to DownloadAsResponse and a new downloadSizeLimitAtom (jotai) that gates mo.ui.table exports and clipboard copies when the serialized payload exceeds a host-supplied limit. The default policy is null (gate dormant); hosts opt in by setting the atom at init time.

This is a bandaid until streaming exports lands — the goal is to give VS Code's NotebookRenderer iframe (and any other sandboxed host) a way to refuse downloads it can't fulfill, rather than failing silently.

  • Backend: download_as now returns (url, filename, size_bytes); the size is computed once from the already-serialized bytes per format (csv/json/parquet). Threaded through both ui.table._download_as and ui.dataframe._download_as.
  • Frontend: ExportMenu reads downloadSizeLimitAtom; when set and `size_bytes > limitBytes`, a danger toast fires and the download/copy is skipped. Fail-open when `size_bytes` is absent (older backend) or the atom is unset.
  • Schema: DownloadAsSchema zod output mirrors the new optional field.

This is blocking marimo-team/marimo-lsp#554

Test plan

  • `make py-check` clean
  • `make fe-check` clean
  • Backend tests for `download_as` size correctness (csv/json/parquet) and `_download_as` populating `size_bytes`
  • Manual: in browser, with atom unset, big-table export still works (no regression)
  • Manual: with the atom temporarily defaulted to `{ limitBytes: 1, ... }`, any non-empty table export and clipboard copy triggers the toast and skips the action

@vercel
Copy link
Copy Markdown

vercel Bot commented May 11, 2026

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

Project Deployment Actions Updated (UTC)
marimo-docs Ready Ready Preview, Comment May 12, 2026 7:23pm

Request Review

@kirangadhave
Copy link
Copy Markdown
Member Author

@cubic-dev-ai

@cubic-dev-ai
Copy link
Copy Markdown
Contributor

cubic-dev-ai Bot commented May 11, 2026

@cubic-dev-ai

@kirangadhave I have started the AI code review. It will take a few minutes to complete.

Copy link
Copy Markdown
Contributor

@cubic-dev-ai cubic-dev-ai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No issues found across 8 files

Architecture diagram
sequenceDiagram
    participant FE as Frontend (ExportMenu)
    participant Atom as downloadSizeLimitAtom
    participant BE as Backend (download_as)
    participant DF as mo.dataframe/table (manager)
    participant DB as Data

    Note over FE,DB: Export / Clipboard flow with optional size-gating

    FE->>BE: downloadAs({ format })
    BE->>DF: to_csv/to_json/to_parquet()
    DF->>DB: read data
    DF-->>BE: payload bytes
    BE->>BE: len(payload) – NEW: compute size_bytes
    BE->>DF: mo_data.csv/json/parquet(url)
    DF-->>BE: vfile with url
    BE-->>FE: { url, filename, size_bytes }

    alt Host has set downloadSizeLimitAtom
        FE->>Atom: read policy (limitBytes, unavailableMessage)
        alt size_bytes != null AND size_bytes > limitBytes
            FE->>FE: toast("Data too large to download/copy")
            Note over FE: Skip download/clipboard action
        else size_bytes absent (older backend) or within limit
            FE->>FE: proceed with download or clipboard copy
        end
    else Atom is null (default)
        Note over FE: No gate – always proceed
        FE->>FE: download or clipboard as normal
    end
Loading

@kirangadhave kirangadhave requested a review from manzt May 11, 2026 23:05
@kirangadhave kirangadhave marked this pull request as ready for review May 11, 2026 23:05
@kirangadhave kirangadhave added the bug Something isn't working label May 11, 2026
Comment thread frontend/src/components/data-table/export-actions.tsx Outdated
Expose a per-table JSON-serialized size in the render payload and a new `downloadSizeLimitAtom` so hosts (e.g. marimo-lsp in VS Code) can disable the Export button up-front when the data would not fit through their transport.

JSON size is cached on the UI element via `memoize_last_value` keyed on manager identity, so it is recomputed only when filters/search produce a new manager. When the host has not seeded the atom (browser, islands), the gate is dormant and the button stays enabled regardless of size.
@kirangadhave
Copy link
Copy Markdown
Member Author

@cubic-dev-ai

@cubic-dev-ai
Copy link
Copy Markdown
Contributor

cubic-dev-ai Bot commented May 12, 2026

@cubic-dev-ai

@kirangadhave I have started the AI code review. It will take a few minutes to complete.

@kirangadhave kirangadhave marked this pull request as ready for review May 12, 2026 18:42
Copy link
Copy Markdown
Contributor

@cubic-dev-ai cubic-dev-ai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2 issues found across 9 files

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/plugins/impl/DataTablePlugin.tsx">

<violation number="1" location="frontend/src/plugins/impl/DataTablePlugin.tsx:556">
P2: `props.sizeBytes` is used inside `useAsyncData` but not tracked as a dependency, which can leave `data.sizeBytes` stale and cause incorrect size-gating decisions.</violation>
</file>

<file name="marimo/_plugins/ui/_impl/table.py">

<violation number="1" location="marimo/_plugins/ui/_impl/table.py:800">
P2: Avoid computing `size-bytes` during lazy initialization; this eagerly serializes the full table and negates the lazy-loading path.</violation>
</file>
Architecture diagram
sequenceDiagram
    participant UI as Browser/VS Code Renderer
    participant ExportMenu as ExportMenu Component
    participant Atom as downloadSizeLimitAtom
    participant Plugin as DataTablePlugin
    participant Backend as Marimo Python Backend
    participant Table as mo.ui.table
    participant DataFrame as mo.ui.dataframe
    participant Manager as TableManager

    Note over UI,Manager: Download Size Gate Flow

    UI->>Plugin: Render table with sizeBytes
    Plugin->>ExportMenu: Pass sizeBytes prop
    ExportMenu->>Atom: Read downloadSizeLimitAtom

    alt Host sets limit (e.g., VS Code)
        Atom-->>ExportMenu: { limitBytes, unavailableMessage }
        
        alt sizeBytes > limitBytes
            ExportMenu-->>UI: Disable Export button
            ExportMenu-->>UI: Show tooltip with unavailableMessage
            ExportMenu-->>ExportMenu: Handle export → skip download/copy
        else sizeBytes ≤ limitBytes
            ExportMenu-->>ExportMenu: Normal export flow continues
        end
    else Host doesn't set limit (default null)
        Atom-->>ExportMenu: null
        ExportMenu-->>ExportMenu: Gate disabled, normal flow continues
    end

    Note over UI,Backend: Download Request (when gate passes)

    UI->>ExportMenu: user clicks export format
    ExportMenu->>Plugin: downloadAs({ format })
    Plugin->>Backend: RPC call to _download_as()
    
    Backend->>Table: _download_as(args)
    Table->>Manager: serialize data (csv/json/parquet)
    Manager-->>Table: payload bytes
    
    Table->>Table: Compute size_bytes = len(payload)
    Table-->>Backend: DownloadAsResponse(url, filename, size_bytes)
    
    Backend-->>Plugin: Response with size_bytes
    Plugin-->>ExportMenu: Return url + filename
    ExportMenu-->>UI: Trigger download / clipboard copy

    Note over UI,Backend: Search/Filter Recompute

    Plugin->>Backend: search({ query, ... })
    Backend->>Table: _search(args)
    Table->>Manager: Filter data
    Manager-->>Table: Filtered result
    
    Table->>Table: _get_json_size_bytes(result) (memoized)
    Table-->>Backend: SearchTableResponse(data, total_rows, size_bytes)
    Backend-->>Plugin: Response with size_bytes
    Plugin-->>UI: Update sizeBytes prop on DataTable

    Note over Table,DataFrame: DataFrame Plugin also participates

    Backend->>DataFrame: init with size_bytes
    DataFrame->>Plugin: Pass size_bytes prop
    Plugin->>ExportMenu: Forward sizeBytes to ExportMenu

    Note over Backend,Manager: Conservative size estimate
    Note over Backend,Manager: JSON is largest format → if JSON fits, CSV/Parquet also fit
Loading

Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.

Comment thread frontend/src/plugins/impl/DataTablePlugin.tsx
Comment thread marimo/_plugins/ui/_impl/table.py Outdated
The size-bytes JSON computation runs on every table init; in lazy mode no
search has happened yet so this serialized the full manager unnecessarily.
Also adds props.sizeBytes to useAsyncData deps so the gate updates when the
prop changes.
The dropdown gated open on a disabled flag while keeping aria-disabled only.
Switch to a native disabled attribute on the Button and DropdownMenuTrigger,
wrapped in a span so the Tooltip keeps a live hover target.
export const ExportMenu: React.FC<ExportActionProps> = (props) => {
const { locale } = useLocale();
const [open, setOpen] = React.useState(false);
const policy = useAtomValue(downloadSizeLimitAtom);
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice

Copy link
Copy Markdown
Collaborator

@manzt manzt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks great. nice solution

@kirangadhave kirangadhave merged commit 8cb7deb into main May 12, 2026
46 checks passed
@kirangadhave kirangadhave deleted the kg/table-download-size-gate branch May 12, 2026 19:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants