Give a LangGraph Agent a Python Sandbox

python3 -m venv .venv && . .venv/bin/activate
pip install langgraph langchain-core sandboxer-sdk
export SANDBOXER_API_KEY=sbx_live_…

import os
import threading

from langchain_core.messages import AIMessage, ToolMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition
from sandboxer import Sandboxer

client = Sandboxer()  # reads SANDBOXER_API_KEY
PREFIX = os.environ.get("SANDBOX_PREFIX", "langgraph-")
_sandboxes = {}  # thread_id -> Sandbox
_lock = threading.Lock()


def sandbox_for(thread_id: str):
    """One sandbox per LangGraph thread, created on first use."""
    with _lock:
        if thread_id not in _sandboxes:
            _sandboxes[thread_id] = client.create(
                type="small",
                ttl_seconds=1800,  # stops itself if this process dies
                name=f"{PREFIX}{thread_id}",
            )
        return _sandboxes[thread_id]


def close_all():
    for thread_id, sandbox in _sandboxes.items():
        sandbox.delete()  # or sandbox.stop() to keep files for later
        print(f"deleted {sandbox.id} ({thread_id})")
    _sandboxes.clear()


@tool
def run_python(code: str, config: RunnableConfig) -> str:
    """Run a Python 3 program in this conversation's cloud sandbox and return its output.
    Files and installed packages persist between calls in the same conversation."""
    sandbox = sandbox_for(config["configurable"]["thread_id"])
    sandbox.files.write("main.py", code)
    r = sandbox.exec("python3 main.py", timeout_seconds=int(os.environ.get("RUN_TIMEOUT", "120")))
    status = "timed out" if r.timed_out else f"exit code {r.exit_code}"
    out = (r.stdout + r.stderr)[-4000:]  # keep the model's context small
    return f"[{sandbox.id}] {status}\n{out}".strip()


# A stand-in for the model: it always asks for the same tool call, then reports the result.
# Replace this node with a real chat model bound to [run_python].
PROGRAM = """
import json
n = 1_000_000
sieve = bytearray([1]) * n
sieve[0:2] = b"\\x00\\x00"
for i in range(2, int(n ** 0.5) + 1):
    if sieve[i]:
        sieve[i * i::i] = bytearray(len(range(i * i, n, i)))
primes = [i for i, p in enumerate(sieve) if p]
gaps = [b - a for a, b in zip(primes, primes[1:])]
json.dump(primes, open("primes.json", "w"))
print("primes below 1,000,000:", len(primes))
print("largest gap:", max(gaps), "after", primes[gaps.index(max(gaps))])
"""


def fake_model(state: MessagesState):
    last = state["messages"][-1]
    if isinstance(last, ToolMessage):
        return {"messages": [AIMessage(content=f"Tool said:\n{last.content}")]}
    code = PROGRAM if last.content.startswith("count the primes") else last.content
    return {"messages": [AIMessage(content="", tool_calls=[
        {"name": "run_python", "args": {"code": code}, "id": f"call_{len(state['messages'])}"}
    ])]}


builder = StateGraph(MessagesState)
builder.add_node("agent", fake_model)
builder.add_node("tools", ToolNode([run_python]))
builder.add_edge(START, "agent")
builder.add_conditional_edges("agent", tools_condition)
builder.add_edge("tools", "agent")
graph = builder.compile(checkpointer=InMemorySaver())


def ask(thread_id: str, text: str):
    config = {"configurable": {"thread_id": thread_id}}
    result = graph.invoke({"messages": [{"role": "user", "content": text}]}, config)
    print(f"--- thread {thread_id}: {result['messages'][-1].content}\n")


if __name__ == "__main__":
    try:
        # 1. The tool on its own, no graph
        print(run_python.invoke(
            {"code": "import platform, os; print(platform.python_version(), os.getcwd())"},
            config={"configurable": {"thread_id": "1"}},
        ), "\n")
        # 2. Graph -> tool -> sandbox, same thread, so the same sandbox
        ask("1", "count the primes below 1,000,000")
        ask("1", "import json; print(len(json.load(open('primes.json'))), 'primes still on disk')")
        # 3. A second thread gets its own sandbox
        ask("2", "import os; print(os.path.exists('primes.json'))")
        # 4. A runaway program hits the timeout
        os.environ["RUN_TIMEOUT"] = "5"
        ask("2", "import time\nwhile True: time.sleep(1)")
    finally:
        close_all()

python3 codeact.py

[sb_g6jl6xmzhc] exit code 0
3.12.3 /home/user

--- thread 1: Tool said:
[sb_g6jl6xmzhc] exit code 0
primes below 1,000,000: 78498
largest gap: 114 after 492113

--- thread 1: Tool said:
[sb_g6jl6xmzhc] exit code 0
78498 primes still on disk

--- thread 2: Tool said:
[sb_lqkqv4gsj5] exit code 0
False

--- thread 2: Tool said:
[sb_lqkqv4gsj5] timed out

deleted sb_g6jl6xmzhc (1)
deleted sb_lqkqv4gsj5 (2)