fix: injection protections and concurrency improvements

tests/sandbox/test_jupyter_gateway_isolation.py

@@ -145,3 +145,102 @@ def test_prime_creates_workspace_mode_0700(tmp_path, monkeypatch):
     assert kernel.initialized
     assert stat.S_IMODE(os.stat(root).st_mode) == 0o700
     assert stat.S_IMODE(os.stat(workspace).st_mode) == 0o700
+
+
+# -- Output cap: rich outputs count against the byte budget --------------------
+
+
+def test_rich_payload_bytes_sums_all_data():
+    content = {"data": {"text/html": "x" * 100, "text/plain": "y" * 50, "application/json": {"a": 1}}}
+    n = JupyterKernelGatewaySandbox._rich_payload_bytes(content)
+    assert n >= 150  # html + plain + serialized json all counted
+
+
+def _frame(msg_id, msg_type, content):
+    return json.dumps({"parent_header": {"msg_id": msg_id}, "msg_type": msg_type, "content": content})
+
+
+class _FakeWS:
+    def __init__(self, frames):
+        self._frames = list(frames)
+
+    def settimeout(self, _t):
+        pass
+
+    def recv(self):
+        if self._frames:
+            return self._frames.pop(0)
+        import websocket
+
+        raise websocket.WebSocketConnectionClosedException()
+
+
+def test_collect_caps_oversize_rich_output(monkeypatch):
+    # A huge execute_result must NOT be buffered: once it would exceed the byte
+    # budget the bundle is dropped and the result is marked truncated.
+    sb = JupyterKernelGatewaySandbox(gateway_url="http://unused", max_output_bytes=1000)
+    monkeypatch.setattr(sb, "_interrupt_and_drain", lambda *a, **k: None)
+    msg_id = "m1"
+    frames = [_frame(msg_id, "execute_result", {"data": {"text/html": "H" * 5000}})]
+    result = sb._collect(_FakeWS(frames), msg_id, timeout=5, kernel_id="k1")
+    assert result.results == []  # over-budget bundle dropped, not materialized
+    assert "[output truncated" in result.stderr
+
+
+def test_collect_keeps_small_rich_output(monkeypatch):
+    sb = JupyterKernelGatewaySandbox(gateway_url="http://unused", max_output_bytes=10000)
+    monkeypatch.setattr(sb, "_interrupt_and_drain", lambda *a, **k: None)
+    msg_id = "m2"
+    frames = [
+        _frame(msg_id, "execute_result", {"data": {"text/plain": "small"}}),
+        _frame(msg_id, "execute_reply", {"status": "ok", "execution_count": 1}),
+        _frame(msg_id, "status", {"execution_state": "idle"}),
+    ]
+    result = sb._collect(_FakeWS(frames), msg_id, timeout=5, kernel_id="k2")
+    assert len(result.results) == 1
+    assert "[output truncated" not in (result.stderr or "")
+
+
+# -- open() is idempotent under concurrency -----------------------------------
+
+
+def test_concurrent_open_creates_one_kernel(monkeypatch):
+    # Two threads opening the same session must POST exactly one kernel; a second
+    # would orphan on the gateway. The CV guard serializes per-session creation.
+    import threading
+
+    sb = JupyterKernelGatewaySandbox(gateway_url="http://unused")
+    posts = {"n": 0}
+    post_lock = threading.Lock()
+    start = threading.Event()
+
+    class _Resp:
+        def raise_for_status(self):
+            pass
+
+        def json(self):
+            return {"id": f"kernel-{posts['n']}"}
+
+    def _fake_post(url, **kwargs):
+        with post_lock:
+            posts["n"] += 1
+        start.wait(timeout=2)  # hold both threads at the POST to force the race
+        return _Resp()
+
+    monkeypatch.setattr(jupyter_gateway.requests, "post", _fake_post)
+    monkeypatch.setattr(sb, "_prime", lambda kernel: None)
+
+    results = {}
+
+    def _open(i):
+        results[i] = sb.open("session-x")
+
+    threads = [threading.Thread(target=_open, args=(i,)) for i in range(2)]
+    for t in threads:
+        t.start()
+    start.set()
+    for t in threads:
+        t.join(timeout=5)
+
+    assert posts["n"] == 1, "concurrent open() created more than one kernel"
+    assert results[0] == results[1]  # both callers got the same kernel id

tests/sandbox/test_sandbox_manager.py

@@ -1,6 +1,7 @@
 """Unit tests for SandboxManager and SandboxCreator using an in-memory backend."""
 
 import threading
+import time
 from typing import Dict, List
 
 import pytest
@@ -256,6 +257,57 @@ def test_reuse_existing_session_never_evicts_at_cap(backend):
     assert mgr.session_count() == 1
 
 
+def test_concurrent_open_same_id_does_not_evict_innocent():
+    """A 2nd open() of an id mid cold-open reuses its placeholder, never evicts another session.
+
+    The session_id is derived from conversation/run id, so two concurrent requests on
+    one conversation both call open(same_id). The second must NOT see the first's
+    not-yet-ready placeholder as a reason to make room (evict an innocent LRU-idle
+    session or raise SandboxCapacityError) -- overwriting the same key adds no slot.
+    """
+    state = {"entries": 0}
+    both_in_backend = threading.Event()
+    release = threading.Event()
+    guard = threading.Lock()
+
+    class _BlockingSameId(FakeBackend):
+        def open(self, session_id: str) -> str:
+            if session_id == "slow":
+                with guard:
+                    state["entries"] += 1
+                    if state["entries"] >= 2:
+                        both_in_backend.set()
+                assert release.wait(timeout=5)
+            return super().open(session_id)
+
+    backend = _BlockingSameId()
+    mgr = SandboxManager(backend, max_ttl=600, max_sessions=2)
+    mgr.open("keep")  # a ready, idle session occupying one of the two slots
+
+    t1 = threading.Thread(target=lambda: mgr.open("slow"))
+    t1.start()
+    # Wait until t1 has registered the "slow" placeholder and is blocked in backend.open
+    # (registry is now {keep, slow} = at cap).
+    deadline = time.monotonic() + 5
+    while mgr.session_count() < 2 and time.monotonic() < deadline:
+        time.sleep(0.01)
+    assert mgr.session_count() == 2
+
+    t2 = threading.Thread(target=lambda: mgr.open("slow"))
+    t2.start()
+    try:
+        # Both opens have passed the lock section into backend.open. With the bug, t2
+        # would have evicted "keep" in its lock section before reaching here.
+        assert both_in_backend.wait(timeout=5)
+        assert mgr.has_session("keep"), "concurrent same-id open evicted an innocent session"
+        assert "keep" not in backend.torn_down
+    finally:
+        release.set()
+        t1.join(timeout=5)
+        t2.join(timeout=5)
+    assert mgr.has_session("slow") and mgr.has_session("keep")
+
+
 # ---------------------------------------------------------------------------
 # Idle reaper
 # ---------------------------------------------------------------------------