sidebar code-mode: add run_scratchpad utility functions

[Light2Dark]

📝 Summary

This is in preparation for adding code_mode to the chat sidebar. Refactors some functions, and creates a toolset function for pydantic-ai which we will eventually use.

📋 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	May 21, 2026 8:16am

[Copilot]

Pull request overview

Refactors scratchpad execution utilities in preparation for adding code_mode support to the chat sidebar, centralizing screenshot credential wiring and introducing a reusable runner + toolset hook.

Changes:

• Added get_code_mode_credentials() to consistently derive (server_url, auth_token) and avoid relying on spoofable request host headers.
• Added run_scratchpad_code() to drive scratchpad execution for code-mode (AI tool usage), including timeout handling and result extraction.
• Introduced shared screenshot meta keys and updated server/runtime code + tests to use them.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/_server/test_scratchpad.py	Adds regression tests for run_scratchpad_code, including screenshot meta propagation, locking, timeout behavior, and child-cell error plumbing.
tests/_server/api/test_api_utils.py	Adds tests ensuring get_code_mode_credentials() uses configured host/port, handles base_url, and propagates scheme.
tests/_server/api/endpoints/test_execution.py	Updates assertions to use screenshot meta constants and checks request meta presence.
marimo/_server/scratchpad.py	Introduces run_scratchpad_code() and wires screenshot credentials into HTTPRequest.meta.
marimo/_server/api/utils.py	Adds get_code_mode_credentials() helper to build trusted callback URL + auth token.
marimo/_server/api/endpoints/execution.py	Refactors /execute to use get_code_mode_credentials() and shared screenshot meta keys.
marimo/_server/ai/tools/code_mode.py	Adds a pydantic_ai FunctionToolset builder exposing an execute_code tool bound to session/request.
marimo/_code_mode/screenshot_meta.py	Defines shared meta keys for screenshot server URL + auth token.
marimo/_code_mode/_context.py	Switches screenshot meta access to shared constants and adjusts missing-credentials error text.

[cubic-dev-ai]

4 issues found across 9 files

Architecture diagram

sequenceDiagram
    participant AIAgent as AI Agent (Pydantic AI)
    participant CodeTool as execute_code Tool
    participant ExecEndpoint as /api/kernel/execute SSE
    participant Utils as get_code_mode_credentials()
    participant Scratchpad as run_scratchpad_code()
    participant Session as Session
    participant Ctx as code-mode ctx (kernel)
    participant Browser as Playwright Browser

    Note over AIAgent,Browser: Code-mode AI tool calling back into the running notebook server

    AIAgent->>CodeTool: Call execute_code(code="...")
    CodeTool->>Utils: get_code_mode_credentials(AppState, Request)
    Utils->>Utils: Build server_url from configured host/port + auth_token
    Utils-->>CodeTool: Return (server_url, auth_token)
    CodeTool->>Scratchpad: run_scratchpad_code(session, request, code, server_url, auth_token)

    Scratchpad->>Scratchpad: HTTPRequest.from_request(request)
    Scratchpad->>Scratchpad: Stamp server_url + auth_token into request.meta
    Scratchpad->>Session: instantiate(... auto_run=False)
    Scratchpad->>Session: put_control_request(ExecuteScratchpadCommand(code, request, run_id))
    activate Session
    Note over Session: Kernel executes scratchpad code

    rect rgb(200, 220, 240)
        Note over Ctx,Browser: Inside scratchpad cell execution
        Ctx->>Ctx: ctx.screenshot() called from code
        Ctx->>Ctx: Read screenshot_server_url from request.meta
        Ctx->>Ctx: Read screenshot_auth_token from request.meta
        Ctx->>Browser: Connect Playwright to server_url
        Ctx->>ExecEndpoint: POST to screenshot endpoint with auth_token
        ExecEndpoint-->>Ctx: Screenshot taken
        Browser-->>Ctx: Return screenshot data
    end

    Session-->>Scratchpad: CompletedRunNotification
    deactivate Session
    Scratchpad->>Scratchpad: extract_result(session, listener)

    alt Normal completion
        Scratchpad-->>CodeTool: CodeExecutionResult(success=True, ...)
    else Timeout (no completion within timeout)
        Scratchpad->>Session: try_interrupt()
        Scratchpad-->>CodeTool: CodeExecutionResult(success=False, errors=[...])
    end

    CodeTool-->>AIAgent: Return CodeExecutionResult

Loading

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

Re-trigger cubic</sub>