fix(claude): persist image cell outputs to files instead of inlining base64

[herikwebb]

Problem

In Claude Code mode _resolve_supports_vision returns True, so image output bundles (matplotlib plots etc.) were rendered as markdown data-URLs and joined into the CLI query as plain text — an enormous base64 string the model pays tokens for but cannot render as an image.

Change

• New image_saver hook on format_output_context: takes (mime_type, cleaned_base64), persists the image, returns its path (or None to fall back to text rendering). It takes precedence over supports_vision.
• In Claude Code mode, image bundles are decoded and written under NBI's upload dir (same lifecycle + retention sweeps as pasted images) and the message references the saved path, which the agent's Read tool can open natively.
• Claude-mode output context is now budgeted by the formatted message's real token count instead of the client-side base64 estimate (the message is small once images become paths).
• Non-Claude providers keep the existing data-URL / text rendering.
• 6 tests, including an end-to-end test verifying the decoded PNG lands on disk and no base64 reaches chat history.

Testing

Full Python suite passes (1298 tests).

[github-actions]

Automated PR Review

Model: gpt-5.5
Verdict: CHANGES_REQUESTED

Changed files:

notebook_intelligence/cell_output.py
notebook_intelligence/extension.py
tests/test_cell_output.py
tests/test_websocket_handler_integration.py

Review result:

Automated review findings

High — notebook_intelligence/extension.py — Image files are written before token-budget acceptance

In Claude Code mode, the handler now calls:

context_message = _format_output_context(
    output_context,
    supports_vision=False,
    image_saver=_save_output_context_image,
)
estimated_tokens = _token_count(context_message)
if estimated_tokens > remaining_token_budget:
    continue

Because _format_output_context(..., image_saver=...) has the side effect of decoding and writing image bundles to disk, rejected contexts can still create files. A crafted authenticated websocket request can include many additionalContext entries with large cell sources that exceed the token budget, causing each context to be skipped while still persisting image files. Since additionalContext is not bounded here, this can become a disk/CPU exhaustion vector.

Recommended fix: perform a side-effect-free budget check before saving images, or split formatting into “estimate” and “persist/render” phases so files are only written after the context is accepted. Add a regression test where an over-budget output context with an image does not create a file.

Medium — notebook_intelligence/extension.py — Saved output-context images bypass upload retention sweeping

_save_output_context_image() stores files under the shared upload directory, but unlike the upload handler path, it does not trigger _sweep_upload_dir(...). As a result, in long-lived servers that use cell-output image contexts but do not perform regular file uploads, these generated image directories can accumulate until process exit.

Recommended fix: invoke the same retention cleanup path after successful image persistence, or centralize temp-file creation/retention so uploads and output-context images use the same lifecycle. Add coverage that saved output images participate in retention cleanup.

[github-actions]

Automated PR Review

Model: gpt-5.5
Verdict: APPROVE

Changed files:

notebook_intelligence/cell_output.py
notebook_intelligence/extension.py
tests/test_cell_output.py
tests/test_websocket_handler_integration.py

Review result:

No automated findings were found in the changed code.

Residual risks to keep in mind:

• Persisting cell-output images is still best-effort: disk-full or permission failures may degrade the Claude Code experience, though this is not a new security issue.
• Upload/temp-dir retention remains cleanup-based rather than quota-based, so sustained authenticated usage could still consume disk until sweeps/exit cleanup run.