Feature: Jupyter Support, Traces & Persistent REPL (Specs 00-02)

[petroslamb]

Summary

This PR implements three major specifications:

1. Persistent Multi-turn REPL (Spec 00): Implements logic for persistent sessions and multi-turn interactions.
2. Trace Artifacts (Spec 01): Adds durable notebook-based traces (.trace.ipynb) and markdown traces providing visibility into RLM execution steps.
3. Jupyter Support (Spec 02): Introduces a jupyter environment allowing RLM code to execute directly within a Jupyter kernel with captured output and variable syncing.

Changes

• Implemented persistent multi-turn REPL logic defined in Spec 00.
• Added trace_notebook.py and trace_markdown.py for trace generation.
• Implemented JupyterREPL environment.
• Updated RLM core to support trace accumulation and trace_markdown on completion objects.
• Added visualizer/ updates to support new log formats.
• Added extensive tests for notebook and multi-turn integration.

Testing

• Verified with pytest tests/repl/test_jupyter_repl.py.
• Verified system-wide with uv run pytest.
• Manual verification in Jupyter Lab.

[alexzhang13]

Cool, this is a pretty big PR so may take some time for me to thoroughly test. A lot of nice features here though, I wanted to add Jupyter support for a while!

[petroslamb]

Cool, this is a pretty big PR so may take some time for me to thoroughly test. A lot of nice features here though, I wanted to add Jupyter support for a while!

i know, it got. big, tried to break it down to three commits each following a spec. if you need futher help or breaking down let me know.

[petroslamb]

Actually i can break this down to 3 prs that build on each other, if you see is fit, with a more descriptive description for each.

Before I do that work, wanted to check if you're aligned with the direction.

Proposed PR Sequence

PR 1: Session vs Completion Interface + Context Tracking

Formalizes the distinction between RLMSession.chat() (multi-turn conversations with history) and RLM.completion() (one-off calls). Introduces SupportsPersistence protocol so environments can expose structured context variables (session_context_0, session_context_1, etc.) to the REPL.

Why: Separates two usage patterns that were previously conflated, and gives models structured access to prior prompts/contexts.

Spec: docs/specs/00-persistent-multi-turn-repl.md | Files: ~15

PR 2: Trace Artifacts (Notebook + Markdown)

Auto-generates two trace artifacts per run:

• .trace.ipynb notebook (executable, replayable from JSONL logs)
• result.trace_markdown (lightweight inspection in notebooks)

Why: Better observability into RLM trajectories and nested llm_query calls for debugging.

Spec: docs/specs/01-trace-artifacts.md | Files: ~12 | Depends on: PR 1

PR 3: Jupyter Support

Adds native jupyter environment that executes code blocks directly in the notebook kernel (no Docker/Modal needed). Includes sync_to_user_ns=True to expose variables to user namespace.

Why: Reduces friction for researchers who want to experiment locally in notebooks.

Spec: docs/specs/02-jupyter-support.md | Files: ~8 | Depends on: PR 1, 2

Questions

1. Do these features align with RLM's direction? (Especially the new session vs completion interface)
2. Prefer 3 sequential PRs or keep the single large PR?
3. Any architectural concerns? Happy to discuss alternatives.
   If you're on board, I'll split this into 3 PRs with review guides. If any don't fit your vision, let me know and I can drop them. Thanks!

There is a notebook that showcases this new implementation 00-RLM-Notebook-Integration.ipynb, along with the specs. If you find it useful i can help you out to make things easier to merge.