Harden scoped Code Mode execution

agents/restricted/src/hyperforge_restricted/agent.py

@@ -341,10 +341,13 @@ async def inner_from_config(
 
     def _runner(self, memory: QuestionMemory, manager: Manager):
         if sandbox_settings.sandbox_socket is not None:
-            # Remote mode, use a socket to communicate with the sandbox server
+            # Remote mode, use a socket to communicate with the sandbox server.
+            # Published functions may return arbitrary Pydantic models, so keep
+            # converting callback results to their legacy dictionary form.
             return SandboxRunner.remote(
                 sandbox_settings.sandbox_socket,
                 functools.partial(self.handle_queue_item, manager, memory),
+                legacy_callback_results=True,
             )
         # Process pool mode, run in a separate process for isolation
         return SandboxRunner.with_pool(

agents/restricted/tests/test_sandbox.py

@@ -1,6 +1,9 @@
 import asyncio
+import functools
 import os
 import tempfile
+from concurrent.futures import ProcessPoolExecutor
+from contextlib import asynccontextmanager
 
 import pytest
 from hyperforge.codemode import sandbox
@@ -10,13 +13,41 @@
     WorkerExecutionRequest,
 )
 from hyperforge.definition import FunctionDefinition
+from pydantic import BaseModel
 
 from hyperforge_restricted.agent import PythonAgent
 from hyperforge_restricted.config import PythonAgentConfig
 
 
 @pytest.fixture
-async def sandbox_conn(monkeypatch):
+async def long_timeout(monkeypatch):
+    monkeypatch.setenv("SANDBOX_CALLBACK_WAIT_SECONDS", "15")
+    yield
+
+
+@asynccontextmanager
+async def configured_sandbox(monkeypatch, **environment: str):
+    with tempfile.TemporaryDirectory() as path:
+        socket = f"{path}/sandbox.sock"
+        monkeypatch.setenv("SANDBOX_SOCKET", socket)
+        monkeypatch.setenv("SANDBOX_VERIFY", "false")
+        monkeypatch.setenv("SANDBOX_TOKEN", "test-token")
+        for name, value in environment.items():
+            monkeypatch.setenv(name, value)
+        task = asyncio.create_task(sandbox.run_sandbox_server())
+        for _ in range(100):
+            if os.path.exists(socket):
+                break
+            await asyncio.sleep(0.01)
+        try:
+            yield socket
+        finally:
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
+@pytest.fixture
+async def sandbox_conn(monkeypatch, long_timeout):
     with tempfile.TemporaryDirectory() as path:
         socket = f"{path}/sandbox.sock"
         monkeypatch.setenv("SANDBOX_SOCKET", socket)
@@ -36,7 +67,7 @@ async def sandbox_conn(monkeypatch):
 
 
 @pytest.fixture
-async def tokenless_sandbox_conn(monkeypatch):
+async def tokenless_sandbox_conn(monkeypatch, long_timeout):
     with tempfile.TemporaryDirectory() as path:
         socket = f"{path}/sandbox.sock"
         monkeypatch.setenv("SANDBOX_SOCKET", socket)
@@ -53,12 +84,6 @@ async def tokenless_sandbox_conn(monkeypatch):
         await asyncio.gather(task, return_exceptions=True)
 
 
-@pytest.fixture
-async def long_timeout(monkeypatch):
-    monkeypatch.setenv("SANDBOX_CALLBACK_WAIT_SECONDS", "15")
-    yield
-
-
 async def test_sandbox_no_code(sandbox_conn, long_timeout):
     request = WorkerExecutionRequest(
         code="", question="Q?", local_vars={}, global_vars={}, function_names={}
@@ -84,7 +109,11 @@ async def test_sandbox_token_is_optional(tokenless_sandbox_conn, long_timeout):
 async def test_sandbox_bad_code(sandbox_conn, long_timeout):
     code = 'return "patata"'  # This does not compile (return not supported)
     request = WorkerExecutionRequest(
-        code=code, question="Q?", local_vars={}, global_vars={}, function_names={}
+        code=code,
+        question="Q?",
+        local_vars={},
+        global_vars={},
+        function_names={},
     )
 
     (rx, tx) = sandbox_conn
@@ -97,14 +126,19 @@ async def test_sandbox_bad_code(sandbox_conn, long_timeout):
 async def test_sandbox_infinite_loop(sandbox_conn):
     code = "while True: pass"  # An infinite loop and should be interrupted
     request = WorkerExecutionRequest(
-        code=code, question="Q?", local_vars={}, global_vars={}, function_names={}
+        code=code,
+        question="Q?",
+        local_vars={},
+        global_vars={},
+        function_names={},
+        max_runtime_seconds=0.1,
     )
 
     (rx, tx) = sandbox_conn
     await tx.write_message(SandboxMessage.Run(request, token="test-token"))
     message = await rx.read_message()
     assert isinstance(message, SandboxMessage.Error)
-    assert "timeout" in message.error
+    assert "timed out" in message.error
 
 
 async def test_sandbox_runtime_limit(sandbox_conn, long_timeout):
@@ -124,6 +158,34 @@ async def test_sandbox_runtime_limit(sandbox_conn, long_timeout):
     assert "timed out" in message.error
 
 
+async def test_sandbox_enforces_server_runtime_ceiling(monkeypatch):
+    request = WorkerExecutionRequest(
+        code="while True: pass",
+        question="Q?",
+        local_vars={},
+        global_vars={},
+        function_names={},
+    )
+    async with configured_sandbox(
+        monkeypatch,
+        SANDBOX_CALLBACK_WAIT_SECONDS="15",
+        SANDBOX_MAX_SESSION_RUNTIME_SECONDS="0.1",
+    ) as socket:
+        rx_raw, tx_raw = await asyncio.open_unix_connection(socket)
+        rx = sandbox.SandboxReader(rx_raw)
+        tx = sandbox.SandboxWriter(tx_raw)
+        try:
+            await tx.write_message(SandboxMessage.Run(request, token="test-token"))
+
+            message = await rx.read_message()
+        finally:
+            tx.close()
+            await tx.wait_closed()
+
+    assert isinstance(message, SandboxMessage.Error)
+    assert "timed out" in message.error
+
+
 async def test_sandbox_function_call(sandbox_conn, long_timeout):
     request = WorkerExecutionRequest(
         code="show(msg=foo())",
@@ -176,6 +238,15 @@ def test_restricted_agent_config_cannot_enable_debug_mode():
     assert "debug" not in config.model_dump()
 
 
+def test_restricted_agent_remote_runner_uses_legacy_callback_results(monkeypatch):
+    agent = PythonAgent.__new__(PythonAgent)
+    monkeypatch.setattr(sandbox.settings, "sandbox_socket", "/tmp/sandbox.sock")
+
+    runner = agent._runner(None, None)
+
+    assert runner.legacy_callback_results is True
+
+
 async def test_restricted_agent_rejects_undeclared_callback():
     agent = PythonAgent.__new__(PythonAgent)
     agent.function_names = {"child": {}}
@@ -188,3 +259,187 @@ async def test_restricted_agent_rejects_undeclared_callback():
 
     with pytest.raises(ValueError, match="not authorized"):
         await agent.handle_queue_item(None, None, task)
+
+
+async def test_remote_restricted_agent_preserves_arbitrary_model_callback_results(
+    monkeypatch, long_timeout
+):
+    class ApplicationResult(BaseModel):
+        value: int
+
+    received = []
+
+    class Child:
+        async def lookup(self, **_kwargs):
+            return ApplicationResult(value=7)
+
+        async def show(self, result, **_kwargs):
+            received.append(result)
+            return None
+
+    agent = PythonAgent.__new__(PythonAgent)
+    agent.function_names = {
+        "child": {
+            function: FunctionDefinition(name=function, description="", parameters={})
+            for function in ("lookup", "show")
+        }
+    }
+    agent.agents = {"child": Child()}
+    agent.flow_id = None
+
+    async with configured_sandbox(monkeypatch) as socket:
+        runner = sandbox.SandboxRunner.remote(
+            socket,
+            functools.partial(agent.handle_queue_item, None, None),
+            legacy_callback_results=True,
+        )
+        await runner.run(
+            WorkerExecutionRequest(
+                code="show(result=lookup())",
+                question="Q?",
+                local_vars={},
+                global_vars={},
+                function_names=agent.function_names,
+            )
+        )
+
+    assert received == [{"__model__": "ApplicationResult", "value": 7}]
+
+
+async def test_sandbox_rejects_wrong_token(sandbox_conn):
+    request = WorkerExecutionRequest(
+        code="", question="Q?", local_vars={}, global_vars={}, function_names={}
+    )
+    (rx, tx) = sandbox_conn
+
+    await tx.write_message(SandboxMessage.Run(request, token="wrong-token"))
+
+    with pytest.raises(asyncio.IncompleteReadError):
+        await rx.read_message()
+
+
+async def test_remote_runner_allows_missing_token(monkeypatch):
+    with tempfile.TemporaryDirectory() as path:
+        socket = f"{path}/sandbox.sock"
+        monkeypatch.setenv("SANDBOX_SOCKET", socket)
+        monkeypatch.setenv("SANDBOX_VERIFY", "false")
+        monkeypatch.setenv("SANDBOX_CALLBACK_WAIT_SECONDS", "15")
+        monkeypatch.delenv("SANDBOX_TOKEN", raising=False)
+        task = asyncio.create_task(sandbox.run_sandbox_server())
+        for _ in range(100):
+            if os.path.exists(socket):
+                break
+            await asyncio.sleep(0.01)
+
+        async def callback(_task):
+            return None
+
+        runner = sandbox.SandboxRunner.remote(socket, callback)
+        request = WorkerExecutionRequest(
+            code="", question="Q?", local_vars={}, global_vars={}, function_names={}
+        )
+        try:
+            await runner.run(request)
+        finally:
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
+async def test_local_runner_preserves_legacy_python_values():
+    received = []
+
+    async def callback(task):
+        received.append(task.keyword_args["value"])
+        return task.keyword_args["value"]
+
+    with ProcessPoolExecutor(max_workers=1) as pool:
+        runner = sandbox.SandboxRunner.with_pool(pool, callback)
+        await runner.run(
+            WorkerExecutionRequest(
+                code="passthrough(value=value)",
+                question="Q?",
+                local_vars={"value": b"legacy-pickle-value"},
+                global_vars={},
+                function_names={
+                    "self": {
+                        "passthrough": FunctionDefinition(
+                            name="passthrough", description="", parameters={}
+                        )
+                    }
+                },
+            )
+        )
+
+    assert received == [b"legacy-pickle-value"]
+
+
+async def test_sandbox_rejects_concurrent_session(monkeypatch):
+    request = WorkerExecutionRequest(
+        code="foo()",
+        question="Q?",
+        local_vars={},
+        global_vars={},
+        function_names={
+            "self": {
+                "foo": FunctionDefinition(name="foo", description="", parameters={})
+            }
+        },
+    )
+    async with configured_sandbox(
+        monkeypatch,
+        SANDBOX_CALLBACK_WAIT_SECONDS="15",
+        SANDBOX_MAX_CONCURRENT_SESSIONS="1",
+    ) as socket:
+        first_rx_raw, first_tx_raw = await asyncio.open_unix_connection(socket)
+        first_rx = sandbox.SandboxReader(first_rx_raw)
+        first_tx = sandbox.SandboxWriter(first_tx_raw)
+        try:
+            await first_tx.write_message(
+                SandboxMessage.Run(request, token="test-token")
+            )
+            assert isinstance(await first_rx.read_message(), SandboxMessage.Request)
+
+            second_rx_raw, second_tx_raw = await asyncio.open_unix_connection(socket)
+            second_rx = sandbox.SandboxReader(second_rx_raw)
+            second_tx = sandbox.SandboxWriter(second_tx_raw)
+            try:
+                await second_tx.write_message(
+                    SandboxMessage.Run(request, token="test-token")
+                )
+                response = await second_rx.read_message()
+                assert isinstance(response, SandboxMessage.Error)
+                assert "concurrency limit" in response.error
+            finally:
+                second_tx.close()
+                await second_tx.wait_closed()
+
+            await first_tx.write_message(SandboxMessage.Response(result=None))
+            assert isinstance(await first_rx.read_message(), SandboxMessage.Done)
+        finally:
+            first_tx.close()
+            await first_tx.wait_closed()
+
+
+def test_sandbox_settings_preserve_legacy_metrics_defaults() -> None:
+    configured = sandbox.SandboxSettings()
+
+    assert configured.sandbox_metrics_enabled is True
+    assert configured.sandbox_metrics_host == "0.0.0.0"
+    assert configured.sandbox_socket_mode == "0600"
+    assert configured.sandbox_max_concurrent_sessions is None
+    assert configured.sandbox_max_session_runtime_seconds is None
+    assert configured.sandbox_max_session_memory_bytes is None
+
+
+def test_sandbox_settings_accept_shared_group_socket_mode() -> None:
+    configured = sandbox.SandboxSettings(
+        sandbox_socket_mode="0660", sandbox_socket_group="sandbox"
+    )
+
+    assert configured.sandbox_socket_mode == "0660"
+    assert configured.sandbox_socket_group == "sandbox"
+
+
+def test_sandbox_settings_reject_invalid_socket_mode() -> None:
+    with pytest.raises(ValueError, match="octal mode"):
+        sandbox.SandboxSettings(sandbox_socket_mode="invalid")