From b95f808505fe7a5356cfb8a2e30425eb49a49ff4 Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Wed, 16 Sep 2026 09:04:15 -0400 Subject: [PATCH 1/5] Harden scoped Code Mode execution --- agents/restricted/tests/test_sandbox.py | 221 ++- docs/harness-sdk.md | 275 ++- hyperforge/CHANGELOG.md | 13 + .../src/hyperforge/codemode/__init__.py | 12 + hyperforge/src/hyperforge/codemode/model.py | 566 +++++- hyperforge/src/hyperforge/codemode/sandbox.py | 615 ++++++- hyperforge/src/hyperforge/codemode/worker.py | 105 +- .../src/hyperforge/harness_sdk/__init__.py | 4 + .../src/hyperforge/harness_sdk/harness.py | 15 +- .../src/hyperforge/harness_sdk/models.py | 1 + .../hyperforge/harness_sdk/tools/codemode.py | 674 ++++++- .../src/hyperforge/harness_sdk/tools/utils.py | 46 +- hyperforge/src/hyperforge/server/sandbox.py | 21 +- hyperforge/tests/codemode/test_model.py | 453 +++++ hyperforge/tests/codemode/test_protocol.py | 881 ++++++++- hyperforge/tests/harness_sdk/test_codemode.py | 1593 ++++++++++++++++- hyperforge/tests/harness_sdk/test_harness.py | 14 + 17 files changed, 5220 insertions(+), 289 deletions(-) diff --git a/agents/restricted/tests/test_sandbox.py b/agents/restricted/tests/test_sandbox.py index 25b550cb..43e142a6 100644 --- a/agents/restricted/tests/test_sandbox.py +++ b/agents/restricted/tests/test_sandbox.py @@ -1,6 +1,8 @@ import asyncio import os import tempfile +from concurrent.futures import ProcessPoolExecutor +from contextlib import asynccontextmanager import pytest from hyperforge.codemode import sandbox @@ -16,7 +18,34 @@ @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 +65,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 +82,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 +107,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 +124,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 +156,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())", @@ -188,3 +248,142 @@ 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_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") diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index a1b06f18..7f445a73 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -191,26 +191,80 @@ Tools can attach typed context to their result by setting `context_type` and, if needed, registering a schema and formatter with `register_context()`. Tools are inherited by spawned sub-agents by default. Set -`inheritance=ToolInheritancePolicy.DO_NOT_INHERIT` when a tool must remain on -the current agent. +`inheritance=ToolInheritancePolicy.DO_NOT_INHERIT` on `@tool` or `HarnessTool` +when a tool must remain on the current agent. Scoped Code Mode tools use this +policy by default. ## Scoped Code Mode -`create_codemode_tool()` exposes generated Python as one model tool while making -only an explicit, immutable tuple of capabilities callable from that Python. -The factory never discovers tools from the harness, so capabilities can remain -hidden from the model's top-level tool list. +Use `create_codemode_tool()` when generated Python should orchestrate a small, +explicit capability set without exposing those capabilities as top-level model +tools. The capability list must be an immutable tuple. Scoped Code Mode never +calls `harness.iter_tools()` and never includes core or external tools unless the +caller explicitly passes them. ```python +from pydantic import BaseModel, Field + from hyperforge.harness_sdk import ( AgentHarness, CodeModeCapability, + CodeModeExecutionLimiter, + CodeModeLimits, + HarnessTool, + ToolCallContext, create_codemode_tool, + tool, ) + +class SearchInput(BaseModel): + query: str = Field(description="Read-only catalog search query") + + +class SearchOutput(BaseModel): + matches: list[dict[str, str]] + internal_cursor: str | None = None + + +@tool(description="Search the approved catalog without modifying it.") +async def search_catalog( + context: ToolCallContext, + input_value: SearchInput, +) -> SearchOutput: + return SearchOutput( + matches=[{"id": "item-1", "title": input_value.query}], + internal_cursor="do-not-expose", + ) + + +def project_search_result( + capability: HarnessTool, + output: BaseModel, +) -> dict[str, object]: + del capability + validated = SearchOutput.model_validate(output) + return {"matches": validated.matches} + + code_mode = create_codemode_tool( - capabilities=(CodeModeCapability(weather),), + capabilities=( + CodeModeCapability( + search_catalog, + result_adapter=project_search_result, + ), + ), + limits=CodeModeLimits( + max_source_bytes=64 * 1024, + max_result_bytes=256 * 1024, + max_cumulative_result_bytes=1024 * 1024, + max_output_bytes=256 * 1024, + max_nested_calls=10, + ), + execution_limiter=CodeModeExecutionLimiter(max_concurrent_executions=4), + remote_required=True, ) + agent = AgentHarness( model="your-model", model_client=model_client, @@ -218,40 +272,162 @@ agent = AgentHarness( ) ``` -The generated tool description includes each capability's description and -argument JSON Schema. Capability calls use normal `HarnessTool` input and output -validation and count toward ordinary harness tool-call usage. The optional -`question` Code Mode input is available to generated code as `question`. - -By default, capability results are projected through the tool's formatted -model-facing context. Supply a custom `result_adapter` for another safe -projection, or opt into the complete JSON-mode Pydantic output with -`raw_codemode_result_adapter`. Every adapter result is normalized through a -strict JSON round trip before generated code receives it, so tuples become -lists and non-JSON or non-finite values are rejected consistently across -runners. Adapters must not return Pydantic models, including models nested in -containers, or dictionaries containing the reserved `__model__` key. -Values passed to scoped `output(...)` use the same normalization and rejection -rules. Generated code must call `output(value)` exactly once; a missing or -repeated call fails the invocation. - -Capability names must be unique ASCII public Python identifiers, cannot use -reserved worker names, and cannot expose an `agent_id` input field. Scoped Code -Mode tools default to `ToolInheritancePolicy.DO_NOT_INHERIT`. If the outer Code -Mode tool is explicitly configured with `INHERIT`, all of its capabilities must -also be inheritable; the factory rejects combinations that would bypass a -capability's `DO_NOT_INHERIT` policy. - -For tests, `runner=` accepts a `CodeModeRunner`; its `run()` method receives the -worker request and capability dispatch callback. Without an injected runner, -the factory uses the existing `SandboxRunner.remote` when `SANDBOX_SOCKET` is -configured and the existing isolated process otherwise. Runtime and memory -behavior continues to use the existing sandbox implementation and -`UsageLimits`; this foundation adds no separate sandbox limits. - -The existing exported `codemode` tool remains available for compatibility and -continues to discover registered harness tools. Prefer the scoped factory when -the generated program should have an explicit capability boundary. +Only `code_mode` is registered with the harness in this example. +`search_catalog` is callable from generated code but is not advertised as a +top-level model tool. The Code Mode tool description includes each capability's +description and argument schema so the model can write valid calls. Nested +arguments still pass through `HarnessTool.execute()`, including JSON Schema, +Pydantic input, and output validation. + +Every `CodeModeCapability` has a result adapter. The safe default returns the +tool's formatted model-facing context as a string. Prefer an application adapter, +as above, when generated code needs selected structured fields. Projected values +must contain only JSON worker values; the SDK rejects non-serializable values, +non-finite numbers, Pydantic models at any nesting depth, and the reserved +`__model__` transport key. Use +`raw_codemode_result_adapter` only after explicitly deciding that the complete +JSON-mode Pydantic output is safe for generated code. A raw adapter does not +bypass serialization or result-size validation. + +Capability names must be unique ASCII public, non-keyword Python identifiers and must not +conflict with worker names such as `codemode`, `output`, `save`, `question`, +`agent_id`, `dataclass`, `Chunk`, `Context`, `List`, `Any`, or `Dict`. + +`CodeModeLimits` applies source, per-call projected-result, cumulative +projected-result, final-output, and nested-call limits to each invocation. +Values are measured as UTF-8 JSON bytes where applicable, and the cumulative +cap must be at least the per-call cap. `CodeModeExecutionLimiter` provides +fail-fast admission control. The default instance is process-wide; pass one +shared application-owned instance to limit a particular run or group of tools. +Projected results are also preflighted against the fixed JSON packet ceiling, +using the raw callback value for isolated execution and the complete response +envelope for remote execution, before a completion event is emitted. +Existing runtime and memory limits remain configured through `UsageLimits`: + +```python +from hyperforge.harness_sdk import UsageLimits + +usage_limits = UsageLimits( + max_tool_calls=20, + max_codemode_runtime_seconds=30, + max_codemode_memory_bytes=512 * 1024 * 1024, +) +``` + +The normal `max_tool_calls` count includes the outer Code Mode call and every +nested capability call exactly once. `max_nested_calls` independently bounds one +generated program. Generated code must call `output(value)` with exactly one +value, exactly once; a missing, empty, or repeated `output` call fails the +invocation. Authorization, nested/global call limits, projected-result limits, +and invalid or repeated `output` attempts remain terminal even if generated code +catches the immediate callback error. An output attempt rejected by worker +transport serialization also remains terminal. The optional `question` input is +exposed to generated code as the worker's `question` variable. + +Pass `runner=` to inject a custom `CodeModeRunner` for deterministic tests; the +default runner uses the remote sandbox, or the isolated local process when +`remote_required=False`. An injected runner bypasses the socket and token +fail-closed checks and owns its callback lifecycle. `remote_required=False` +always selects the isolated process even if ambient sandbox socket and token +settings are present. + +Each nested call emits `TOOL_REQUESTED`, followed by `TOOL_COMPLETED` or +`TOOL_FAILED`, with a stable call ID. Event payloads have `codemode=true`, +`nested=true`, and the outer call's `parent_call_id`; `HarnessEvent` also stores +`parent_call_id` as a first-class field. Arguments and projected results are +sanitized. Arbitrary dictionaries use an ordered entry-list representation whose +keys contain only type and UTF-8 length metadata, so attacker-controlled keys are +not persisted and sensitive key names still cause recursive value redaction. +Sanitization emits at most 64 value nodes, 16 items per container, and six nested +container levels. A bounded `truncated` object records the reason and omitted item +count. All strings, including JSON-looking strings, are persisted only as type and +UTF-8 byte-count metadata and are never parsed. Failure payloads contain only the +exception type, and normal event fields preserve conversation, turn, agent, and +parent-agent context. Scalar actor, tenant, and user identifiers from +`execution_context` are copied into the nested payload. Marker strings and decimal +integers are limited to 1 KiB each, all marker values share a 4 KiB aggregate +budget, and marker floats must be finite; invalid context fails before generated +execution or nested event emission. Nested wrapper and tool-authored events retain +the turn that started the Code Mode invocation, +including events from callbacks that finish during a later harness turn. A nested +event persistence or publication failure is reported to generated code only as +`Code Mode capability failed` and remains invocation-terminal. + +The outer Code Mode `TOOL_REQUESTED` event follows the standard harness event +contract and persists its raw `code` and `question` arguments. Applications must +not embed secrets in either field. The nested events described above are the +sanitized audit surface for capability calls and results. + +The existing exported `codemode` tool remains available for compatibility. It +discovers registered tools and returns raw JSON-mode outputs. New applications +that require capability isolation should use `create_codemode_tool()`. + +Remote and isolated-process Code Mode execution move values over the bounded +JSON worker protocol. Values crossing those worker boundaries must be JSON +worker values; types such as `bytes`, sets, and dataclasses are rejected before +transport. Arbitrary Pydantic models, including root models and application +models with aliases, are also rejected; only the protocol's explicit model types +are encoded. Explicitly supported marked models such as `Context`, `RemiResponse`, +`WorkerError`, and `RestrictedPythonTask` retain their markers when nested in run +request local or global variables; the run envelope itself remains unmarked. +Plain dictionaries containing `__model__` are escaped unambiguously and remain +dictionaries after decoding, including dictionaries that resemble the escape +wrapper itself. Convert binary data to base64 strings in tool results. The legacy +local pool used by the restricted agent retains its existing pickle transport for +backward compatibility. Scoped requests redact generated-code exceptions to a +fixed generic error before they leave the worker. Generic legacy requests retain +detailed syntax and runtime errors by default. + +### Code Mode Security + +RestrictedPython reduces the available Python language surface. It is not a +security boundary. Production generated-code execution requires a separately +isolated OS process or container with a dedicated non-root identity, no network, +a read-only or minimal filesystem, dropped capabilities, no privilege +escalation, seccomp or an equivalent syscall policy, and CPU, memory, PID, +file-descriptor, and wall-clock limits. Keep the API process outside that +boundary. + +Scoped Code Mode defaults to `remote_required=True`. It fails closed when +`SANDBOX_SOCKET` or `SANDBOX_TOKEN` is absent and never silently falls back to +local execution. The generic sandbox service can run without a token for +deployments that rely solely on Unix socket ACLs, but scoped Code Mode still +requires one; the socket ACL is an additional control, not a replacement for +token authentication. `remote_required=False` enables the isolated local +process and is intended only for deterministic tests and explicitly trusted +development environments. + +Sandbox deployment controls and compatibility defaults are: + +- `SANDBOX_MAX_CONCURRENT_SESSIONS` bounds authenticated worker sessions and + concurrent authentication handshakes when configured. It is unset by default. +- `SANDBOX_MAX_SESSION_RUNTIME_SECONDS` and + `SANDBOX_MAX_SESSION_MEMORY_BYTES` impose server-owned ceilings when + configured. Both are unset by default for compatibility. Scoped + Code Mode supplies per-request runtime and memory limits. The remote client + bounds connections that have a requested or server runtime plus + `SANDBOX_TIMEOUT_SLACK_SECONDS=10`. +- `SANDBOX_SOCKET_MODE=0600` restricts the socket to its owner. +- `SANDBOX_SOCKET_GROUP` optionally changes group ownership. Use an explicitly + provisioned shared group with `SANDBOX_SOCKET_MODE=0660` when the API and + sandbox run as different non-root users. +- `SANDBOX_METRICS_ENABLED=true`, `SANDBOX_METRICS_HOST=0.0.0.0`, and + `SANDBOX_METRICS_PORT=8091` preserve the existing metrics endpoint. Set + `SANDBOX_METRICS_ENABLED=false` to avoid creating an IP listener, or bind the + endpoint to `127.0.0.1` when it does not need to be remotely scraped. +- Local process IPC and remote length-prefixed JSON messages are bounded; + oversized run requests, callbacks, responses, and frames are rejected. + +The generic sandbox server accepts tokenless sessions unless `SANDBOX_TOKEN` or +an application `token_verifier` is configured. Server-owned runtime, memory, and +concurrency enforcement likewise requires the deployment settings above. These +compatibility defaults do not weaken the scoped client factory's fail-closed +requirement for a socket and token when `remote_required=True`. + +Place the socket in a dedicated directory writable only by the sandbox identity +and, when configured, the shared group. Provision the directory and group before +startup; do not make the socket or parent directory world-writable. Rotate +`SANDBOX_TOKEN` as a secret and restart both peers together during rotation. ## Usage Limits @@ -364,8 +540,12 @@ publication never waits for the consumer. If the consumer falls behind by `event_queue_size` events, additional `TEXT_DELTA` and `REASONING_DELTA` events are dropped; persisted lifecycle events remain lossless. Events from child agents may interleave, but `turn_id`, `agent_id`, and `parent_agent_id` identify their -origin. Use `async with agent` when a caller may stop consuming early so the -active turn and descendants are cleaned up. +origin. Events emitted through `ToolCallContext.emit()` carry its +`parent_call_id`; each call receives its own context, so use `context.emit()` +when an event should be attributed to the current call. Scoped Code Mode contexts +also preserve their originating turn for delayed callback events. Use +`async with agent` when a caller may stop consuming early so the active turn and +descendants are cleaned up. At most four agents run concurrently in a conversation by default, including the root agent. Configure `max_concurrent_agents` to change this limit. When the @@ -439,10 +619,11 @@ implementation is responsible for enforcing tenant and user isolation. The model can call `spawn_agent` with a self-contained prompt, then call `wait_agent` with the returned internal ID. It can use `send_message` to steer a -running child. Children inherit the model client, model, external tools, -execution context, storage, and usage limits. They start with isolated history -unless `include_history=true` is requested. Included history contains the current -system and user messages but never an incomplete spawn tool exchange. +running child. Children inherit the model client, model, external tools whose +inheritance policy is `INHERIT`, execution context, storage, and usage limits. +They start with isolated history unless `include_history=true` is requested. +Included history contains the current system and user messages but never an +incomplete spawn tool exchange. Delegation is limited to depth one by default. Set `max_spawn_depth` on the root harness to change the limit; an agent at the limit does not advertise the diff --git a/hyperforge/CHANGELOG.md b/hyperforge/CHANGELOG.md index b2496f6d..c4e9ce68 100644 --- a/hyperforge/CHANGELOG.md +++ b/hyperforge/CHANGELOG.md @@ -3,6 +3,19 @@ ## Unreleased - Added explicit scoped Code Mode capabilities and child-tool inheritance policy. +- Hardened scoped Code Mode with fail-closed remote execution, bounded JSON + transport, execution limits, atomic callback cleanup, immutable sandbox + admission/configuration snapshots, and memory-bounded, key-safe sanitized + nested-call events. Scoped worker errors are redacted before transport and + remain redacted if transport fails; scoped source/result prechecks stop at their + configured bounds. Nested-call and output violations are terminal, and late + nested events retain the turn that originated the invocation. Event backend, + global call, projected-result, cumulative-result, worker output serialization, + and execution-marker failures now fail closed across caught retries. Supported + marked protocol models remain available in run request local and global values, + while plain marker-like dictionaries round-trip without retyping. Integer event + markers and exact projected-result response envelopes are now bounded before + nested execution or completion emission. ## 1.0.0 diff --git a/hyperforge/src/hyperforge/codemode/__init__.py b/hyperforge/src/hyperforge/codemode/__init__.py index c2a775ee..8f131d6d 100644 --- a/hyperforge/src/hyperforge/codemode/__init__.py +++ b/hyperforge/src/hyperforge/codemode/__init__.py @@ -1,11 +1,17 @@ from .model import ( + MAX_PROTOCOL_BYTES, RestrictedPythonTask, SandboxMessage, WorkerError, WorkerExecutionRequest, WorkerTypes, + decode_json_value, + decode_protocol_value, deserialize, + encode_protocol_value, + encode_sandbox_message, serialize, + validate_protocol_value, ) from .sandbox import ( InsecureSandbox, @@ -21,6 +27,7 @@ __all__ = [ "InsecureSandbox", + "MAX_PROTOCOL_BYTES", "PythonAgentWorker", "RestrictedPythonTask", "SandboxMessage", @@ -32,8 +39,13 @@ "WorkerExecutionRequest", "WorkerError", "WorkerTypes", + "decode_json_value", + "decode_protocol_value", + "encode_sandbox_message", "deserialize", + "encode_protocol_value", "run_sandbox_server", "serialize", "settings", + "validate_protocol_value", ] diff --git a/hyperforge/src/hyperforge/codemode/model.py b/hyperforge/src/hyperforge/codemode/model.py index 4e2f34f6..a2572096 100644 --- a/hyperforge/src/hyperforge/codemode/model.py +++ b/hyperforge/src/hyperforge/codemode/model.py @@ -1,11 +1,47 @@ -from dataclasses import dataclass -from typing import Any, Sequence, assert_never, cast +import json +import math +from dataclasses import dataclass, is_dataclass +from typing import Any, Iterator, Sequence, assert_never, cast -from nuclia_models.predict.remi import RemiResponse -from pydantic import BaseModel +from nuclia_models.common.consumption import Consumption, TokensDetail +from nuclia_models.predict.remi import AnswerRelevance, RemiResponse +from pydantic import BaseModel, Field from hyperforge.definition import FunctionDefinition from hyperforge.memory import Context +from hyperforge.models import Chunk, Image, JSONObject, Prompt + +MAX_PROTOCOL_BYTES = 16 * 1024 * 1024 +MAX_PROTOCOL_DEPTH = 256 +_JSON_STRING_CHUNK_CHARS = 4096 +_PLAIN_DICT_MODEL = "__hyperforge_plain_dict__" + + +class _ProtocolValueTooLarge(ValueError): + pass + + +class _StrictWorkerModelValue(TypeError): + pass + + +class _ReservedWorkerKey(ValueError): + pass + + +@dataclass(frozen=True) +class _PlainModel: + value: BaseModel + + +@dataclass(frozen=True) +class _MarkedProtocolValue: + value: Any + + +@dataclass(frozen=True) +class _RawPlainDict: + value: dict[Any, Any] class WorkerExecutionRequest(BaseModel): @@ -14,8 +50,9 @@ class WorkerExecutionRequest(BaseModel): local_vars: dict[str, Any] global_vars: dict[str, Any] function_names: dict[str, dict[str, FunctionDefinition]] - max_runtime_seconds: float | None = None - max_memory_bytes: int | None = None + max_runtime_seconds: float | None = Field(default=None, gt=0, allow_inf_nan=False) + max_memory_bytes: int | None = Field(default=None, gt=0) + redact_errors: bool = False class RestrictedPythonTask(BaseModel): @@ -47,27 +84,60 @@ class WorkerError(BaseModel): "RemiResponse": RemiResponse, "WorkerError": WorkerError, } +_MARKED_PROTOCOL_MODELS = frozenset(WorkerModels.values()) +_UNMARKED_PROTOCOL_MODELS = frozenset( + { + AnswerRelevance, + Chunk, + Consumption, + FunctionDefinition, + Image, + JSONObject, + Prompt, + TokensDetail, + } +) -def deserialize(msg: Any) -> WorkerTypes: +def deserialize(msg: Any, *, _depth: int = 0) -> WorkerTypes: + if _depth > MAX_PROTOCOL_DEPTH: + raise ValueError("Protocol value is too deeply nested") if isinstance(msg, dict) and "__model__" in msg: model_name = msg["__model__"] + if not isinstance(model_name, str): + raise ValueError("Invalid protocol model marker") + if model_name == _PLAIN_DICT_MODEL: + value = msg.get("value") + if set(msg) != {"__model__", "value"} or not isinstance(value, dict): + raise ValueError("Invalid escaped protocol dictionary") + return { + key: deserialize(item, _depth=_depth + 1) for key, item in value.items() + } if model_name == "RestrictedPythonTask": + args = deserialize(msg.get("args"), _depth=_depth + 1) + keyword_args = deserialize(msg.get("keyword_args"), _depth=_depth + 1) + if not isinstance(args, (list, tuple)) or not isinstance( + keyword_args, dict + ): + raise ValueError("Invalid RestrictedPythonTask payload") return RestrictedPythonTask( function=msg["function"], agent=msg["agent"], - args=tuple(deserialize(a) for a in msg["args"]), - keyword_args={ - k: deserialize(v) for k, v in msg["keyword_args"].items() - }, + args=tuple(args), + keyword_args=cast(dict[str, Any], keyword_args), ) model = WorkerModels.get(model_name) if model is not None: - return cast(WorkerTypes, model.model_validate(msg)) + payload = { + key: deserialize(value, _depth=_depth + 1) + for key, value in msg.items() + if key != "__model__" + } + return cast(WorkerTypes, model.model_validate(payload)) if isinstance(msg, list): - return [deserialize(m) for m in msg] + return [deserialize(m, _depth=_depth + 1) for m in msg] if isinstance(msg, dict): - return {k: deserialize(v) for k, v in msg.items()} + return {k: deserialize(v, _depth=_depth + 1) for k, v in msg.items()} return msg @@ -77,20 +147,438 @@ def serialize(message: Any) -> Any: "__model__": "RestrictedPythonTask", "function": message.function, "agent": message.agent, - "args": [serialize(a) for a in message.args], - "keyword_args": {k: serialize(v) for k, v in message.keyword_args.items()}, + "args": serialize(message.args), + "keyword_args": serialize(message.keyword_args), } if isinstance(message, BaseModel): - converted = message.model_dump() + if type(message) not in _MARKED_PROTOCOL_MODELS: + raise TypeError( + f"Pydantic model {type(message).__name__} is not supported by the " + "sandbox protocol" + ) + converted = { + key: serialize(value) for key, value in message.model_dump().items() + } converted["__model__"] = message.__class__.__name__ return converted if isinstance(message, (list, tuple)): return [serialize(m) for m in message] if isinstance(message, dict): - return {k: serialize(v) for k, v in message.items()} + converted = {k: serialize(v) for k, v in message.items()} + if "__model__" in message: + return {"__model__": _PLAIN_DICT_MODEL, "value": converted} + return converted return message +def validate_protocol_value( + value: Any, label: str, *, max_bytes: int = MAX_PROTOCOL_BYTES +) -> None: + encode_protocol_value(value, label, max_bytes=max_bytes) + + +def encode_protocol_value( + value: Any, + label: str, + *, + max_bytes: int = MAX_PROTOCOL_BYTES, + strict_worker_value: bool = False, +) -> bytes: + try: + encoded = bytearray() + for chunk in _iter_json_bytes( + value, + depth=0, + markers=set(), + max_bytes=max_bytes, + include_model_markers=True, + strict_worker_value=strict_worker_value, + ): + if len(encoded) + len(chunk) > max_bytes: + raise _ProtocolValueTooLarge + encoded.extend(chunk) + except _ProtocolValueTooLarge as exc: + raise ValueError(f"{label} exceeds maximum size") from exc + except _StrictWorkerModelValue as exc: + raise TypeError(f"{label} must not contain Pydantic model values") from exc + except _ReservedWorkerKey as exc: + raise ValueError( + f"{label} cannot contain the reserved '__model__' key" + ) from exc + except (RecursionError, TypeError, UnicodeError, ValueError) as exc: + raise ValueError(f"{label} is not serializable") from exc + return bytes(encoded) + + +def _iter_json_bytes( + value: Any, + *, + depth: int, + markers: set[int], + max_bytes: int, + include_model_markers: bool, + strict_worker_value: bool, +) -> Iterator[bytes]: + if depth > MAX_PROTOCOL_DEPTH: + raise ValueError("Protocol value is too deeply nested") + if strict_worker_value and isinstance(value, BaseModel): + raise _StrictWorkerModelValue + if is_dataclass(value) and not isinstance( + value, (_PlainModel, _MarkedProtocolValue, _RawPlainDict) + ): + raise TypeError("Dataclass values are not supported by the sandbox protocol") + if isinstance(value, _PlainModel): + model = value.value + if type(model) is not WorkerExecutionRequest: + raise TypeError( + f"Pydantic model {type(model).__name__} is not supported by the " + "sandbox protocol" + ) + yield from _iter_json_object_bytes( + _worker_request_items(cast(WorkerExecutionRequest, model)), + marker=id(model), + depth=depth, + markers=markers, + max_bytes=max_bytes, + include_model_markers=False, + strict_worker_value=strict_worker_value, + ) + return + if isinstance(value, _MarkedProtocolValue): + yield from _iter_json_bytes( + value.value, + depth=depth, + markers=markers, + max_bytes=max_bytes, + include_model_markers=True, + strict_worker_value=strict_worker_value, + ) + return + if isinstance(value, _RawPlainDict): + yield b"{" + for index, (key, item) in enumerate(value.value.items()): + if index: + yield b"," + yield from _iter_json_string_bytes(_json_key(key, max_bytes)) + yield b":" + yield from _iter_json_bytes( + item, + depth=depth + 1, + markers=markers, + max_bytes=max_bytes, + include_model_markers=include_model_markers, + strict_worker_value=strict_worker_value, + ) + yield b"}" + return + if isinstance(value, RestrictedPythonTask): + yield from _iter_json_object_bytes( + ( + ("__model__", "RestrictedPythonTask"), + ("function", value.function), + ("agent", value.agent), + ("args", value.args), + ("keyword_args", value.keyword_args), + ), + marker=id(value), + depth=depth, + markers=markers, + max_bytes=max_bytes, + include_model_markers=True, + strict_worker_value=strict_worker_value, + ) + return + elif isinstance(value, BaseModel): + supported_models = ( + _MARKED_PROTOCOL_MODELS + if include_model_markers + else _UNMARKED_PROTOCOL_MODELS + ) + if type(value) not in supported_models: + raise TypeError( + f"Pydantic model {type(value).__name__} is not supported by the " + "sandbox protocol" + ) + yield from _iter_json_object_bytes( + _model_items(value, include_marker=include_model_markers), + marker=id(value), + depth=depth, + markers=markers, + max_bytes=max_bytes, + include_model_markers=False, + strict_worker_value=strict_worker_value, + ) + return + if value is None: + yield b"null" + return + if value is True: + yield b"true" + return + if value is False: + yield b"false" + return + if isinstance(value, str): + yield from _iter_json_string_bytes(value) + return + if isinstance(value, int): + if int.bit_length(value) > max_bytes * 4: + raise _ProtocolValueTooLarge + yield int.__repr__(value).encode("ascii") + return + if isinstance(value, float): + if not math.isfinite(value): + raise ValueError("Out of range float values are not JSON compliant") + yield json.dumps(value, allow_nan=False).encode("ascii") + return + if isinstance(value, (list, tuple)): + marker = id(value) + if marker in markers: + raise ValueError("Circular reference detected") + markers.add(marker) + try: + yield b"[" + for index, item in enumerate(value): + if index: + yield b"," + yield from _iter_json_bytes( + item, + depth=depth + 1, + markers=markers, + max_bytes=max_bytes, + include_model_markers=include_model_markers, + strict_worker_value=strict_worker_value, + ) + yield b"]" + finally: + markers.remove(marker) + return + if isinstance(value, dict): + if strict_worker_value and "__model__" in value: + raise _ReservedWorkerKey + if "__model__" in value: + yield from _iter_json_object_bytes( + ( + ("__model__", _PLAIN_DICT_MODEL), + ("value", _RawPlainDict(value)), + ), + marker=id(value), + depth=depth, + markers=markers, + max_bytes=max_bytes, + include_model_markers=include_model_markers, + strict_worker_value=strict_worker_value, + ) + return + marker = id(value) + if marker in markers: + raise ValueError("Circular reference detected") + markers.add(marker) + try: + yield b"{" + for index, (key, item) in enumerate(value.items()): + if index: + yield b"," + yield from _iter_json_string_bytes(_json_key(key, max_bytes)) + yield b":" + yield from _iter_json_bytes( + item, + depth=depth + 1, + markers=markers, + max_bytes=max_bytes, + include_model_markers=include_model_markers, + strict_worker_value=strict_worker_value, + ) + yield b"}" + finally: + markers.remove(marker) + return + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") + + +def _iter_json_object_bytes( + items: Iterator[tuple[str, Any]] | Sequence[tuple[str, Any]], + *, + marker: int, + depth: int, + markers: set[int], + max_bytes: int, + include_model_markers: bool, + strict_worker_value: bool, +) -> Iterator[bytes]: + if marker in markers: + raise ValueError("Circular reference detected") + markers.add(marker) + try: + yield b"{" + for index, (key, item) in enumerate(items): + if index: + yield b"," + yield from _iter_json_string_bytes(key) + yield b":" + yield from _iter_json_bytes( + item, + depth=depth + 1, + markers=markers, + max_bytes=max_bytes, + include_model_markers=include_model_markers, + strict_worker_value=strict_worker_value, + ) + yield b"}" + finally: + markers.remove(marker) + + +def _model_items( + value: BaseModel, *, include_marker: bool +) -> Iterator[tuple[str, Any]]: + decorators = value.__pydantic_decorators__ + if ( + decorators.field_serializers + or decorators.model_serializers + or _schema_has_serialization(value.__class__.__pydantic_core_schema__) + ): + raise TypeError( + "Pydantic models with custom serializers are not supported by the " + "sandbox protocol" + ) + field_names: set[str] = set() + for name, field in value.__class__.model_fields.items(): + if field.exclude is True: + continue + item = getattr(value, name) + exclude_if = getattr(field, "exclude_if", None) + if exclude_if is not None and exclude_if(item): + continue + field_names.add(name) + yield name, item + for name, item in (value.__pydantic_extra__ or {}).items(): + if name not in field_names: + yield name, item + for name in value.__class__.model_computed_fields: + if name not in field_names: + yield name, getattr(value, name) + if include_marker: + yield "__model__", value.__class__.__name__ + + +def _worker_request_items( + value: WorkerExecutionRequest, +) -> Iterator[tuple[str, Any]]: + for name, item in _model_items(value, include_marker=False): + if name in {"local_vars", "global_vars"}: + item = _MarkedProtocolValue(item) + yield name, item + + +def _schema_has_serialization(schema: Any) -> bool: + pending = [iter((schema,))] + visited: set[int] = set() + while pending: + values = pending[-1] + try: + value = next(values) + except StopIteration: + pending.pop() + continue + if isinstance(value, dict): + if "serialization" in value: + return True + marker = id(value) + if marker not in visited: + visited.add(marker) + pending.append(iter(value.values())) + elif isinstance(value, (list, tuple)): + marker = id(value) + if marker not in visited: + visited.add(marker) + pending.append(iter(value)) + return False + + +def _iter_json_string_bytes(value: str) -> Iterator[bytes]: + yield b'"' + for start in range(0, len(value), _JSON_STRING_CHUNK_CHARS): + encoded = json.dumps( + value[start : start + _JSON_STRING_CHUNK_CHARS], ensure_ascii=False + ) + yield encoded[1:-1].encode("utf-8") + yield b'"' + + +def _json_key(value: Any, max_bytes: int) -> str: + if isinstance(value, str): + return value + if value is None: + return "null" + if value is True: + return "true" + if value is False: + return "false" + if isinstance(value, int): + if int.bit_length(value) > max_bytes * 4: + raise _ProtocolValueTooLarge + return int.__repr__(value) + if isinstance(value, float) and math.isfinite(value): + return json.dumps(value, allow_nan=False) + raise TypeError( + f"keys must be str, int, float, bool or None, not {type(value).__name__}" + ) + + +def decode_protocol_value( + encoded: bytes, label: str, *, max_bytes: int = MAX_PROTOCOL_BYTES +) -> WorkerTypes: + value = decode_json_value(encoded, label, max_bytes=max_bytes) + try: + return deserialize(value) + except (AttributeError, KeyError, RecursionError, TypeError, ValueError) as exc: + raise ValueError(f"{label} is invalid") from exc + + +def decode_json_value( + encoded: bytes, label: str, *, max_bytes: int = MAX_PROTOCOL_BYTES +) -> Any: + if len(encoded) > max_bytes: + raise ValueError(f"{label} exceeds maximum size") + try: + value = json.loads( + encoded.decode("utf-8"), parse_constant=_reject_json_constant + ) + _validate_json_depth(value) + return value + except ( + RecursionError, + UnicodeDecodeError, + json.JSONDecodeError, + ValueError, + ) as exc: + raise ValueError(f"{label} is not valid JSON") from exc + + +def _reject_json_constant(value: str) -> None: + raise ValueError(f"Invalid JSON constant: {value}") + + +def _validate_json_depth(value: Any) -> None: + pending = [(iter((value,)), 0)] + while pending: + values, depth = pending[-1] + try: + current = next(values) + except StopIteration: + pending.pop() + continue + if depth > MAX_PROTOCOL_DEPTH: + raise ValueError("Protocol value is too deeply nested") + if isinstance(current, dict): + pending.append((iter(current.values()), depth + 1)) + elif isinstance(current, list): + pending.append((iter(current), depth + 1)) + elif isinstance(current, float) and not math.isfinite(current): + raise ValueError("Protocol value contains a non-finite number") + + class SandboxMessage: @dataclass class Run: @@ -119,13 +607,20 @@ class Error: def parse(cls, data: dict[str, Any]) -> AnyMessage: match data["_"]: case "run": + run = dict(data["run"]) + for name in ("local_vars", "global_vars"): + values = deserialize(run.get(name)) + if not isinstance(values, dict): + raise ValueError(f"Invalid run request {name}") + run[name] = values return cls.Run( - run=WorkerExecutionRequest.model_validate(data["run"]), + run=WorkerExecutionRequest.model_validate(run), token=data.get("token"), ) case "request": task = deserialize(data["task"]) - assert isinstance(task, RestrictedPythonTask) + if not isinstance(task, RestrictedPythonTask): + raise ValueError("Invalid sandbox request task") return cls.Request(task=task) case "response": return cls.Response(result=deserialize(data["result"])) @@ -140,9 +635,12 @@ def parse(cls, data: dict[str, Any]) -> AnyMessage: def serialize(cls, message: AnyMessage) -> dict[str, Any]: match message: case cls.Run(): + run = message.run.model_dump() + run["local_vars"] = serialize(message.run.local_vars) + run["global_vars"] = serialize(message.run.global_vars) return { "_": "run", - "run": message.run.model_dump(), + "run": run, "token": message.token, } case cls.Request(): @@ -155,3 +653,29 @@ def serialize(cls, message: AnyMessage) -> dict[str, Any]: return {"_": "error", "error": message.error} case _: assert_never(message) + + +def encode_sandbox_message( + message: SandboxMessage.AnyMessage, + label: str, + *, + max_bytes: int = MAX_PROTOCOL_BYTES, +) -> bytes: + match message: + case SandboxMessage.Run(): + value = { + "_": "run", + "run": _PlainModel(message.run), + "token": message.token, + } + case SandboxMessage.Request(): + value = {"_": "request", "task": message.task} + case SandboxMessage.Response(): + value = {"_": "response", "result": message.result} + case SandboxMessage.Done(): + value = {"_": "done"} + case SandboxMessage.Error(): + value = {"_": "error", "error": message.error} + case _: + assert_never(message) + return encode_protocol_value(value, label, max_bytes=max_bytes) diff --git a/hyperforge/src/hyperforge/codemode/sandbox.py b/hyperforge/src/hyperforge/codemode/sandbox.py index e448bf14..3c1c6e72 100644 --- a/hyperforge/src/hyperforge/codemode/sandbox.py +++ b/hyperforge/src/hyperforge/codemode/sandbox.py @@ -1,11 +1,13 @@ import asyncio import contextvars import hmac -import json import logging import os +import shutil +import threading from concurrent.futures import Executor, ThreadPoolExecutor from concurrent.futures import Future as ConcurrentFuture +from dataclasses import dataclass from multiprocessing import Pipe, Process from multiprocessing.connection import Connection from typing import Any, Awaitable, Callable, Coroutine @@ -13,31 +15,75 @@ import httpx import nucliadb_telemetry.metrics import prometheus_client +from pydantic import Field, field_validator from pydantic_settings import BaseSettings from .model import ( + MAX_PROTOCOL_BYTES, RestrictedPythonTask, SandboxMessage, WorkerError, WorkerExecutionRequest, WorkerTypes, + decode_json_value, + decode_protocol_value, + encode_protocol_value, + encode_sandbox_message, ) from .worker import PythonAgentWorker logger = logging.getLogger("hyperforge_codemode_sandbox") WORKER_EXIT_TIMEOUT = 1 -MAX_PACKET_BYTES = 16 * 1024 * 1024 +CALLBACK_CANCEL_TIMEOUT = 1 +MAX_PACKET_BYTES = MAX_PROTOCOL_BYTES class SandboxSettings(BaseSettings): sandbox_verify: bool = True sandbox_socket: str | None = None + sandbox_socket_mode: str = "0600" + sandbox_socket_group: str | None = None sandbox_metrics_port: int = 8091 + sandbox_metrics_host: str = "0.0.0.0" + sandbox_metrics_enabled: bool = True sandbox_token: str | None = None + sandbox_max_concurrent_sessions: int | None = Field(default=None, gt=0) + sandbox_max_session_runtime_seconds: float | None = Field( + default=None, gt=0, allow_inf_nan=False + ) + sandbox_max_session_memory_bytes: int | None = Field(default=None, gt=0) + sandbox_timeout_slack_seconds: float = Field( + default=10.0, ge=0, allow_inf_nan=False + ) sandbox_callback_wait_seconds: float = 1.0 + @field_validator("sandbox_socket_mode") + @classmethod + def validate_socket_mode(cls, value: str) -> str: + try: + mode = int(value, 8) + except ValueError as exc: + raise ValueError("SANDBOX_SOCKET_MODE must be an octal mode") from exc + if mode < 0 or mode > 0o777: + raise ValueError("SANDBOX_SOCKET_MODE must be between 0000 and 0777") + return value + settings = SandboxSettings() +_remote_admission_lock = threading.Lock() + + +@dataclass +class _RemoteAdmissionCounter: + active_sessions: int = 0 + + +@dataclass(frozen=True) +class _RemoteSessionTicket: + counter: _RemoteAdmissionCounter | None + + +_remote_admission_counter = _RemoteAdmissionCounter() class InsecureSandbox(Exception): @@ -94,32 +140,78 @@ def __init__( self.debug = debug self.isolated = False self.token = token + self._callback_lock = threading.Lock() + self._callbacks_closed = True + self._run_active = False + self._pending_callback_futures: set[ConcurrentFuture[WorkerTypes]] = set() + self._callback_tasks: set[asyncio.Task[WorkerTypes]] = set() + self._callback_futures: dict[ + asyncio.Task[WorkerTypes], ConcurrentFuture[WorkerTypes] + ] = {} + self._orphaned_callback_tasks: set[asyncio.Task[WorkerTypes]] = set() async def run(self, request: WorkerExecutionRequest): - if self.pool is not None: - return await self._run_in_pool(request) - if self.socket is not None: + self._begin_callback_run() + try: + if self.pool is not None: + return await self._run_in_pool(request) + request = _canonicalize_run_request(request) + if self.socket is not None: + runtime = request.max_runtime_seconds + if runtime is None: + runtime = settings.sandbox_max_session_runtime_seconds + timeout = ( + runtime + settings.sandbox_timeout_slack_seconds + if runtime is not None + else None + ) + try: + async with asyncio.timeout(timeout): + return await self._run_with_remote_admission(request) + except TimeoutError as exc: + raise RuntimeError("Codemode execution timed out") from exc + if self.isolated: + return await self._run_isolated(request) + raise ValueError( + "SandboxRunner must be initialized with a pool or a socket" + ) + finally: + try: + await self._cancel_callbacks() + finally: + self._finish_callback_run() + + async def _run_with_remote_admission(self, request: WorkerExecutionRequest): + ticket = await _acquire_remote_session(settings.sandbox_max_concurrent_sessions) + try: return await self._run_remotely(request) - if self.isolated: - return await self._run_isolated(request) - raise ValueError("SandboxRunner must be initialized with a pool or a socket") + finally: + _release_remote_session(ticket) async def _run_in_pool(self, request: WorkerExecutionRequest): loop = asyncio.get_running_loop() pipe_worker, pipe_restricted = Pipe() - worker = PythonAgentWorker(pipe_worker, debug=self.debug) - controller_task = asyncio.create_task(self.background_task(pipe_restricted)) - await loop.run_in_executor( - self.pool, - worker._process_question_context_sync, - request.code, - request.question, - request.local_vars, - request.global_vars, - request.function_names, - request.max_memory_bytes, + worker = PythonAgentWorker(pipe_worker, debug=self.debug, json_protocol=False) + controller_task = asyncio.create_task( + self.background_task(pipe_restricted, json_protocol=False) ) - await controller_task + try: + await loop.run_in_executor( + self.pool, + worker._process_question_context_sync, + request.code, + request.question, + request.local_vars, + request.global_vars, + request.function_names, + request.max_memory_bytes, + request.redact_errors, + ) + await controller_task + finally: + if not controller_task.done(): + controller_task.cancel() + await asyncio.gather(controller_task, return_exceptions=True) async def _run_isolated(self, request: WorkerExecutionRequest): controller_task, process = self.run_in_process(request) @@ -156,33 +248,87 @@ async def _run_remotely(self, request: WorkerExecutionRequest): token = token_source() rx, tx = await asyncio.open_unix_connection(self.socket) reader, writer = SandboxReader(rx), SandboxWriter(tx) + incoming: asyncio.Task[SandboxMessage.AnyMessage] | None = None try: await writer.write_message(SandboxMessage.Run(run=request, token=token)) while True: + if incoming is None: + incoming = asyncio.create_task(reader.read_message()) try: - msg = await reader.read_message() + msg = await incoming except asyncio.IncompleteReadError as exc: raise RuntimeError( "Sandbox connection closed unexpectedly" ) from exc + incoming = None if isinstance(msg, SandboxMessage.Done): break if isinstance(msg, SandboxMessage.Error): raise RuntimeError(f"Python agent error: {msg.error}") - assert isinstance(msg, SandboxMessage.Request) - try: - response = await self.callback(msg.task) - except Exception as exc: - response = WorkerError(error=str(exc)) - await writer.write_message(SandboxMessage.Response(result=response)) + if not isinstance(msg, SandboxMessage.Request): + raise RuntimeError("Unexpected sandbox protocol message") + incoming = await self._run_remote_callback(reader, writer, msg.task) finally: + if incoming is not None: + incoming.cancel() + await asyncio.gather(incoming, return_exceptions=True) writer.close() await writer.wait_closed() + async def _run_remote_callback( + self, + reader: "SandboxReader", + writer: "SandboxWriter", + task: RestrictedPythonTask, + ) -> asyncio.Task[SandboxMessage.AnyMessage]: + with self._callback_lock: + if self._callbacks_closed: + raise RuntimeError("Sandbox callback lifecycle is closed") + callback_task = asyncio.create_task(self.callback(task)) + self._callback_tasks.add(callback_task) + incoming = asyncio.create_task(reader.read_message()) + try: + done, _ = await asyncio.wait( + {callback_task, incoming}, return_when=asyncio.FIRST_COMPLETED + ) + if incoming in done: + message = incoming.result() + if isinstance(message, SandboxMessage.Error): + raise RuntimeError(f"Python agent error: {message.error}") + raise RuntimeError( + "Unexpected sandbox protocol message during callback" + ) + try: + response = callback_task.result() + except Exception as exc: + response = WorkerError(error=str(exc)) + await writer.write_message(SandboxMessage.Response(result=response)) + except BaseException: + if incoming.done(): + _consume_task_result(incoming) + else: + incoming.cancel() + await asyncio.gather(incoming, return_exceptions=True) + raise + finally: + if not callback_task.done(): + callback_task.cancel() + callback_done, still_pending = await asyncio.wait( + {callback_task}, timeout=CALLBACK_CANCEL_TIMEOUT + ) + for completed in callback_done: + await asyncio.gather(completed, return_exceptions=True) + if still_pending: + logger.warning("Sandbox callback ignored cancellation") + with self._callback_lock: + self._callback_tasks.discard(callback_task) + for remaining in still_pending: + self._track_orphaned_callback(remaining) + return incoming + def run_in_process(self, request: WorkerExecutionRequest): pipe_worker, pipe_restricted = Pipe() worker = PythonAgentWorker(pipe_worker, debug=self.debug) - controller_task = asyncio.create_task(self.background_task(pipe_restricted)) process = Process( target=worker._process_question_context_sync, args=( @@ -192,13 +338,22 @@ def run_in_process(self, request: WorkerExecutionRequest): request.global_vars, request.function_names, request.max_memory_bytes, + request.redact_errors, ), ) - process.start() + try: + process.start() + except BaseException: + pipe_worker.close() + pipe_restricted.close() + raise pipe_worker.close() + controller_task = asyncio.create_task(self.background_task(pipe_restricted)) return controller_task, process - async def background_task(self, pipe: Connection) -> bool: + async def background_task( + self, pipe: Connection, *, json_protocol: bool = True + ) -> bool: loop = asyncio.get_running_loop() context = contextvars.copy_context() try: @@ -208,42 +363,202 @@ async def background_task(self, pipe: Connection) -> bool: pipe, loop, context, + json_protocol, ) finally: pipe.close() + def _begin_callback_run(self) -> None: + with self._callback_lock: + if self._run_active: + raise RuntimeError("SandboxRunner does not support concurrent runs") + unsettled = any( + not future.done() for future in self._pending_callback_futures + ) or any( + not task.done() + for task in self._callback_tasks | self._orphaned_callback_tasks + ) + if unsettled: + raise RuntimeError( + "SandboxRunner cannot be reused until callbacks have completed" + ) + self._pending_callback_futures.clear() + self._callback_tasks.clear() + self._callback_futures.clear() + self._orphaned_callback_tasks.clear() + self._callbacks_closed = False + self._run_active = True + + def _finish_callback_run(self) -> None: + with self._callback_lock: + self._callbacks_closed = True + self._run_active = False + + def _register_callback_intent(self, future: ConcurrentFuture[WorkerTypes]) -> bool: + with self._callback_lock: + if self._callbacks_closed: + future.set_exception( + RuntimeError("Sandbox callback lifecycle is closed") + ) + return False + self._pending_callback_futures.add(future) + return True + + def _schedule_callback( + self, + future: ConcurrentFuture[WorkerTypes], + item: RestrictedPythonTask, + loop: asyncio.AbstractEventLoop, + context: contextvars.Context, + ) -> None: + with self._callback_lock: + if future not in self._pending_callback_futures: + return + self._pending_callback_futures.remove(future) + if self._callbacks_closed: + if not future.done(): + future.set_exception( + RuntimeError("Sandbox callback lifecycle is closed") + ) + return + task = loop.create_task(self.callback(item), context=context.copy()) + self._callback_tasks.add(task) + self._callback_futures[task] = future + + def copy_result(completed: asyncio.Task[WorkerTypes]) -> None: + with self._callback_lock: + self._callback_tasks.discard(completed) + controller_future = self._callback_futures.pop(completed, None) + if controller_future is None or controller_future.done(): + return + try: + controller_future.set_result(completed.result()) + except BaseException as exc: + controller_future.set_exception(exc) + + task.add_done_callback(copy_result) + + async def _cancel_callbacks(self) -> None: + with self._callback_lock: + self._callbacks_closed = True + intents = list(self._pending_callback_futures) + self._pending_callback_futures.clear() + tasks = list(self._callback_tasks) + task_futures = [ + future + for task in tasks + if (future := self._callback_futures.pop(task, None)) is not None + ] + cancellation = RuntimeError("Sandbox callback cancelled") + for future in [*intents, *task_futures]: + if not future.done(): + future.set_exception(cancellation) + for task in tasks: + task.cancel() + if tasks: + done, pending = await asyncio.wait(tasks, timeout=CALLBACK_CANCEL_TIMEOUT) + for completed in done: + await asyncio.gather(completed, return_exceptions=True) + if pending: + logger.warning("Sandbox callbacks ignored cancellation") + with self._callback_lock: + self._callback_tasks.difference_update(tasks) + for task in pending: + self._track_orphaned_callback(task) + + def _track_orphaned_callback(self, task: asyncio.Task[WorkerTypes]) -> None: + with self._callback_lock: + self._orphaned_callback_tasks.add(task) + + def completed(done: asyncio.Task[WorkerTypes]) -> None: + with self._callback_lock: + self._orphaned_callback_tasks.discard(done) + _consume_task_result(done) + + task.add_done_callback(completed) + + def run_when_callbacks_complete(self, callback: Callable[[], None]) -> bool: + with self._callback_lock: + pending: set[Any] = { + future for future in self._pending_callback_futures if not future.done() + } + pending.update( + task + for task in self._callback_tasks | self._orphaned_callback_tasks + if not task.done() + ) + if not pending: + return False + + completion_lock = threading.Lock() + remaining = len(pending) + registration_complete = False + + def completed(_done: Any) -> None: + nonlocal remaining + with completion_lock: + remaining -= 1 + release = registration_complete and remaining == 0 + if release: + callback() + + for future in pending: + future.add_done_callback(completed) + with completion_lock: + registration_complete = True + release = remaining == 0 + if release: + callback() + return True + def read_write_queue( self, pipe: Connection, loop: asyncio.AbstractEventLoop, context: contextvars.Context, + json_protocol: bool = True, ) -> bool: try: while True: - item = pipe.recv() + if json_protocol: + item = decode_protocol_value( + pipe.recv_bytes(MAX_PACKET_BYTES), + "Local sandbox request", + max_bytes=MAX_PACKET_BYTES, + ) + else: + item = pipe.recv() + if not isinstance(item, RestrictedPythonTask): + raise ValueError("Invalid local sandbox request") if item.function == "_error": raise RuntimeError(item.args[0]) if item.function == "_close": return True future: ConcurrentFuture[WorkerTypes] = ConcurrentFuture() - - def schedule_callback() -> None: - task = loop.create_task(self.callback(item), context=context.copy()) - - def copy_result(completed: asyncio.Task[WorkerTypes]) -> None: - try: - future.set_result(completed.result()) - except Exception as exc: - future.set_exception(exc) - - task.add_done_callback(copy_result) - - loop.call_soon_threadsafe(schedule_callback) + if self._register_callback_intent(future): + loop.call_soon_threadsafe( + self._schedule_callback, future, item, loop, context + ) try: result = future.result() except Exception as exc: result = WorkerError(error=str(exc)) - pipe.send(result) + if json_protocol: + try: + encoded = encode_protocol_value( + result, + "Local sandbox response", + max_bytes=MAX_PACKET_BYTES, + ) + except ValueError as exc: + encoded = encode_protocol_value( + WorkerError(error=str(exc)), + "Local sandbox response error", + max_bytes=MAX_PACKET_BYTES, + ) + pipe.send_bytes(encoded) + else: + pipe.send(result) except EOFError: return False @@ -257,6 +572,25 @@ def _url_reachable(url: str) -> bool: return False +def _canonicalize_run_request( + request: WorkerExecutionRequest, +) -> WorkerExecutionRequest: + encoded_request = encode_sandbox_message( + SandboxMessage.Run(run=request), + "Sandbox run request", + max_bytes=MAX_PACKET_BYTES, + ) + encoded_value = decode_json_value( + encoded_request, "Sandbox run request", max_bytes=MAX_PACKET_BYTES + ) + if not isinstance(encoded_value, dict): + raise ValueError("Invalid sandbox run request") + message = SandboxMessage.parse(encoded_value) + if not isinstance(message, SandboxMessage.Run): + raise ValueError("Invalid sandbox run request") + return message.run + + def _verify_connectivity(): kubernetes_host = os.environ.get("KUBERNETES_SERVICE_HOST") if kubernetes_host and _url_reachable(f"https://{kubernetes_host}"): @@ -269,16 +603,16 @@ class SandboxWriter: def __init__(self, writer: asyncio.StreamWriter): self.writer = writer - async def _write_packet(self, msg: dict[str, Any]): - msg_bytes = json.dumps(msg).encode("utf-8") - if len(msg_bytes) > MAX_PACKET_BYTES: - raise ValueError("Sandbox message exceeds maximum size") + async def _write_packet(self, msg: SandboxMessage.AnyMessage): + msg_bytes = encode_sandbox_message( + msg, "Sandbox message", max_bytes=MAX_PACKET_BYTES + ) self.writer.write(len(msg_bytes).to_bytes(4, "little")) self.writer.write(msg_bytes) await self.writer.drain() async def write_message(self, message: SandboxMessage.AnyMessage): - await self._write_packet(SandboxMessage.serialize(message)) + await self._write_packet(message) def close(self): self.writer.close() @@ -297,10 +631,25 @@ async def _read_packet(self) -> dict[str, Any]: if msg_len > MAX_PACKET_BYTES: raise ValueError("Sandbox message exceeds maximum size") msg_bytes = await self.reader.readexactly(msg_len) - return json.loads(msg_bytes.decode("utf-8")) + value = decode_json_value( + msg_bytes, "Sandbox message", max_bytes=MAX_PACKET_BYTES + ) + if not isinstance(value, dict): + raise ValueError("Sandbox message must be a JSON object") + return value async def read_message(self) -> SandboxMessage.AnyMessage: - return SandboxMessage.parse(await self._read_packet()) + value = await self._read_packet() + try: + return SandboxMessage.parse(value) + except ( + AttributeError, + KeyError, + RecursionError, + TypeError, + ValueError, + ) as exc: + raise ValueError("Invalid sandbox message") from exc class SandboxSession: @@ -309,32 +658,47 @@ def __init__( reader: SandboxReader, writer: SandboxWriter, callback_wait_seconds: float = 1.0, + *, + settings_snapshot: SandboxSettings | None = None, ): self.reader = reader self.writer = writer self.callback_wait_seconds = callback_wait_seconds + self.settings = settings_snapshot or settings.model_copy(deep=True) self.runner = SandboxRunner( pool=None, socket=None, callback=self._callback, debug=False ) + self.task: asyncio.Task[bool] | None = None + self.process: Process | None = None self.waiting_for_agent: asyncio.Task[None] | None = None self.timed_out = False async def run(self, request: WorkerExecutionRequest): - self.task, self.process = self.runner.run_in_process(request) + self.runner._begin_callback_run() try: + memory = self.settings.sandbox_max_session_memory_bytes + if memory is None: + memory = request.max_memory_bytes + elif request.max_memory_bytes is not None: + memory = min(memory, request.max_memory_bytes) + request = request.model_copy(update={"max_memory_bytes": memory}) + task, process = self.runner.run_in_process(request) + self.task, self.process = task, process await self._start_timeout(self.callback_wait_seconds) - if request.max_runtime_seconds is None: - clean_exit = await self.task - else: - async with asyncio.timeout(request.max_runtime_seconds): - clean_exit = await self.task + runtime = self.settings.sandbox_max_session_runtime_seconds + if runtime is None: + runtime = request.max_runtime_seconds + elif request.max_runtime_seconds is not None: + runtime = min(runtime, request.max_runtime_seconds) + async with asyncio.timeout(runtime): + clean_exit = await task await self._stop_timeout() if self.timed_out: raise RuntimeError("Python agent timeout") if not clean_exit: raise RuntimeError("Codemode worker exited unexpectedly") - await asyncio.to_thread(self.process.join, WORKER_EXIT_TIMEOUT) - if self.process.exitcode != 0: + await asyncio.to_thread(process.join, WORKER_EXIT_TIMEOUT) + if process.exitcode != 0: raise RuntimeError( "Codemode worker exceeded its memory limit or exited unexpectedly" ) @@ -344,27 +708,41 @@ async def run(self, request: WorkerExecutionRequest): except Exception as exc: await self._fail(exc) finally: - await self._stop_timeout() - if self.process.is_alive(): - self.process.kill() - await asyncio.to_thread(self.process.join) - if not self.task.done(): - self.task.cancel() - await asyncio.gather(self.task, return_exceptions=True) + try: + await self._stop_timeout() + await self._stop_worker() + finally: + try: + await self.runner._cancel_callbacks() + finally: + self.runner._finish_callback_run() async def _fail(self, exc: Exception): logger.error("Error in sandbox session: %s", exc) - if self.process.is_alive(): - self.process.kill() - await asyncio.to_thread(self.process.join) - if not self.task.done(): - self.task.cancel() - await asyncio.gather(self.task, return_exceptions=True) + await self._stop_worker() try: await self.writer.write_message(SandboxMessage.Error(error=str(exc))) except Exception: pass + async def _stop_worker(self) -> None: + process = self.process + if process is not None: + try: + if process.is_alive(): + process.kill() + await asyncio.to_thread(process.join) + except Exception: + logger.exception("Failed to clean up sandbox process") + task = self.task + if task is not None: + try: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + except Exception: + logger.exception("Failed to clean up sandbox controller task") + async def _callback(self, item: RestrictedPythonTask) -> WorkerTypes: await self._stop_timeout() await self.writer.write_message(SandboxMessage.Request(task=item)) @@ -372,11 +750,14 @@ async def _callback(self, item: RestrictedPythonTask) -> WorkerTypes: response = await self.reader.read_message() except asyncio.IncompleteReadError: logger.warning("Got disconnected from agent") - self.process.kill() - self.process.join() - self.task.cancel() + if self.process is not None: + self.process.kill() + self.process.join() + if self.task is not None: + self.task.cancel() return None - assert isinstance(response, SandboxMessage.Response) + if not isinstance(response, SandboxMessage.Response): + raise RuntimeError("Unexpected sandbox protocol response") await self._start_timeout(self.callback_wait_seconds) return response.result @@ -391,7 +772,7 @@ async def _stop_timeout(self): async def _timeout(self, timeout: float): await asyncio.sleep(timeout) self.timed_out = True - if self.process.is_alive(): + if self.process is not None and self.process.is_alive(): self.process.kill() @@ -405,18 +786,34 @@ async def run_sandbox_server( *, token_verifier: Callable[[str], Awaitable[bool]] | None = None ): server_settings = SandboxSettings() - assert server_settings.sandbox_socket is not None + if server_settings.sandbox_socket is None: + raise RuntimeError("SANDBOX_SOCKET is required for the sandbox server") sandbox_token = server_settings.sandbox_token if server_settings.sandbox_verify: _verify_connectivity() + active_sessions = 0 + admission_lock = asyncio.Lock() + max_sessions = server_settings.sandbox_max_concurrent_sessions + handshake_limit = ( + asyncio.Semaphore(max_sessions) if max_sessions is not None else None + ) + async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter): + nonlocal active_sessions reader = SandboxReader(rx) writer = SandboxWriter(tx) + if handshake_limit is not None and handshake_limit.locked(): + writer.close() + await writer.wait_closed() + return + if handshake_limit is not None: + await handshake_limit.acquire() try: async with asyncio.timeout(1): msg = await reader.read_message() - assert isinstance(msg, SandboxMessage.Run) + if not isinstance(msg, SandboxMessage.Run): + raise ValueError("Expected a sandbox run request") token = msg.token or "" valid_token = ( await token_verifier(token) @@ -434,25 +831,79 @@ async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter): writer.close() await writer.wait_closed() return + finally: + if handshake_limit is not None: + handshake_limit.release() + + async with admission_lock: + if max_sessions is not None and active_sessions >= max_sessions: + await writer.write_message( + SandboxMessage.Error(error="Sandbox concurrency limit reached") + ) + writer.close() + await writer.wait_closed() + return + active_sessions += 1 agents_running.inc() observation = execution_observer() observation.start() try: session = SandboxSession( - reader, writer, SandboxSettings().sandbox_callback_wait_seconds + reader, + writer, + server_settings.sandbox_callback_wait_seconds, + settings_snapshot=server_settings, ) await session.run(msg.run) except Exception as exc: observation.set_status("error") await writer.write_message(SandboxMessage.Error(error=str(exc))) finally: + async with admission_lock: + active_sessions -= 1 agents_running.dec() observation.end() writer.close() await writer.wait_closed() server = await asyncio.start_unix_server(handler, server_settings.sandbox_socket) - os.chmod(server_settings.sandbox_socket, 0o600) + if server_settings.sandbox_socket_group is not None: + group: str | int = server_settings.sandbox_socket_group + if server_settings.sandbox_socket_group.isdecimal(): + group = int(server_settings.sandbox_socket_group) + shutil.chown(server_settings.sandbox_socket, group=group) + os.chmod( + server_settings.sandbox_socket, int(server_settings.sandbox_socket_mode, 8) + ) async with server: await server.serve_forever() + + +def _consume_task_result(task: asyncio.Future[Any]) -> None: + try: + task.exception() + except asyncio.CancelledError: + pass + + +async def _acquire_remote_session( + max_concurrent_sessions: int | None, +) -> _RemoteSessionTicket: + if max_concurrent_sessions is None: + return _RemoteSessionTicket(counter=None) + while True: + with _remote_admission_lock: + if _remote_admission_counter.active_sessions < max_concurrent_sessions: + _remote_admission_counter.active_sessions += 1 + return _RemoteSessionTicket(counter=_remote_admission_counter) + await asyncio.sleep(0.01) + + +def _release_remote_session(ticket: _RemoteSessionTicket) -> None: + if ticket.counter is None: + return + with _remote_admission_lock: + if ticket.counter.active_sessions <= 0: + raise RuntimeError("Remote sandbox admission counter underflow") + ticket.counter.active_sessions -= 1 diff --git a/hyperforge/src/hyperforge/codemode/worker.py b/hyperforge/src/hyperforge/codemode/worker.py index 89b134bc..0adda404 100644 --- a/hyperforge/src/hyperforge/codemode/worker.py +++ b/hyperforge/src/hyperforge/codemode/worker.py @@ -31,7 +31,13 @@ from hyperforge.memory import Chunk from hyperforge.memory.memory import Context -from .model import RestrictedPythonTask, WorkerError +from .model import ( + MAX_PROTOCOL_BYTES, + RestrictedPythonTask, + WorkerError, + decode_protocol_value, + encode_protocol_value, +) BLOCKED_EXCEPTIONS = { "BaseException", @@ -39,6 +45,7 @@ "KeyboardInterrupt", "SystemExit", } +REDACTED_EXECUTION_ERROR = "Generated code execution failed" def _set_memory_limit(max_memory_bytes: int | None) -> None: @@ -106,10 +113,16 @@ def guarded_inplace(operator: str, left: Any, right: Any) -> Any: class PythonAgentWorker: context: Optional[Context] = None - def __init__(self, pipe: Connection, debug: bool = False): + def __init__( + self, pipe: Connection, debug: bool = False, json_protocol: bool = True + ): self.pipe = pipe self.functions_agent_id: Dict[str, List[str]] = {} self.debug = debug + self.json_protocol = json_protocol + self._redact_errors = False + self._output_attempted = False + self._output_send_failed = False def _process_question_context_sync( self, @@ -119,7 +132,11 @@ def _process_question_context_sync( global_vars: Dict[str, Any], function_names: Dict[str, Dict[str, FunctionDefinition]], max_memory_bytes: int | None = None, + redact_errors: bool = False, ): + self._redact_errors = redact_errors + self._output_attempted = False + self._output_send_failed = False try: _harden_process(max_memory_bytes) for agent, functions in function_names.items(): @@ -161,24 +178,40 @@ def _process_question_context_sync( if self.debug: global_vars["pdb"] = __import__("pdb") exec(byte_code, global_vars, local_vars) + if self._output_attempted and self._output_send_failed: + raise RuntimeError(REDACTED_EXECUTION_ERROR) except BaseException as exc: - self.pipe.send( - RestrictedPythonTask( - function="_error", - agent="_controller", - args=(str(exc),), - keyword_args={}, - ) + error = REDACTED_EXECUTION_ERROR if redact_errors else str(exc) + task = RestrictedPythonTask( + function="_error", + agent="_controller", + args=(error,), + keyword_args={}, ) + if redact_errors: + self._send_or_close(task) + else: + self._send(task) finally: - self.pipe.send( - RestrictedPythonTask( - function="_close", - agent="_controller", - args=(None,), - keyword_args={}, - ) + task = RestrictedPythonTask( + function="_close", + agent="_controller", + args=(None,), + keyword_args={}, ) + if redact_errors: + self._send_or_close(task) + else: + self._send(task) + + def _send_or_close(self, task: RestrictedPythonTask) -> None: + try: + self._send(task) + except BaseException: + try: + self.pipe.close() + except BaseException: + pass def save( self, @@ -188,7 +221,7 @@ def save( question: str = "", final_answer: Optional[str] = None, ): - self.pipe.send( + self._send( RestrictedPythonTask( function="save", agent="_controller", @@ -224,18 +257,40 @@ def valid_functions_call( ) agent_id = agents[0] - self.pipe.send( - RestrictedPythonTask( - function=function_name, - agent=agent_id, - args=args, - keyword_args=kwargs, - ) + task = RestrictedPythonTask( + function=function_name, + agent=agent_id, + args=args, + keyword_args=kwargs, ) + output_attempt = self._redact_errors and function_name == "output" + if output_attempt: + self._output_attempted = True + try: + self._send(task) + except BaseException: + if output_attempt: + self._output_send_failed = True + raise return self._receive() + def _send(self, task: RestrictedPythonTask) -> None: + if self.json_protocol: + self.pipe.send_bytes( + encode_protocol_value( + task, "Local sandbox request", max_bytes=MAX_PROTOCOL_BYTES + ) + ) + else: + self.pipe.send(task) + def _receive(self) -> Any: - result = self.pipe.recv() + if self.json_protocol: + result = decode_protocol_value( + self.pipe.recv_bytes(MAX_PROTOCOL_BYTES), "Local sandbox response" + ) + else: + result = self.pipe.recv() if isinstance(result, WorkerError): raise RuntimeError(result.error) return result diff --git a/hyperforge/src/hyperforge/harness_sdk/__init__.py b/hyperforge/src/hyperforge/harness_sdk/__init__.py index 4d7f2539..d38cdeec 100644 --- a/hyperforge/src/hyperforge/harness_sdk/__init__.py +++ b/hyperforge/src/hyperforge/harness_sdk/__init__.py @@ -43,7 +43,9 @@ from .tools.codemode import ( CodeModeCapability, CodeModeDispatch, + CodeModeExecutionLimiter, CodemodeInput, + CodeModeLimits, CodemodeOutput, CodeModeResultAdapter, CodeModeRunner, @@ -69,6 +71,8 @@ "ChatCompletionUsage", "CodeModeCapability", "CodeModeDispatch", + "CodeModeExecutionLimiter", + "CodeModeLimits", "CodeModeResultAdapter", "CodeModeRunner", "CodemodeInput", diff --git a/hyperforge/src/hyperforge/harness_sdk/harness.py b/hyperforge/src/hyperforge/harness_sdk/harness.py index 9dd06a9a..9964c9d6 100644 --- a/hyperforge/src/hyperforge/harness_sdk/harness.py +++ b/hyperforge/src/hyperforge/harness_sdk/harness.py @@ -42,6 +42,13 @@ MAX_EMPTY_RETRIES = 2 +class _CurrentTurn: + pass + + +_CURRENT_TURN = _CurrentTurn() + + @dataclass class AgentResult: text: str @@ -587,13 +594,19 @@ async def emit( payload: dict[str, Any], *, persist: bool = True, + parent_call_id: str | None = None, + turn_id: str | None | _CurrentTurn = _CURRENT_TURN, ) -> HarnessEvent: + resolved_turn_id = ( + self._turn_id if isinstance(turn_id, _CurrentTurn) else turn_id + ) event = HarnessEvent( id=uuid.uuid4().hex, conversation_id=self.conversation_id, - turn_id=self._turn_id, + turn_id=resolved_turn_id, agent_id=self.agent_id, parent_agent_id=self.parent_agent_id, + parent_call_id=parent_call_id, category=self.category, tags=self.tags, metadata=self._persisted_metadata(), diff --git a/hyperforge/src/hyperforge/harness_sdk/models.py b/hyperforge/src/hyperforge/harness_sdk/models.py index b452960f..bf052025 100644 --- a/hyperforge/src/hyperforge/harness_sdk/models.py +++ b/hyperforge/src/hyperforge/harness_sdk/models.py @@ -84,6 +84,7 @@ class HarnessEvent(BaseModel): turn_id: str | None = None agent_id: str | None = None parent_agent_id: str | None = None + parent_call_id: str | None = None category: str = "assistant" tags: list[str] = Field(default_factory=list) metadata: dict[str, Any] = Field(default_factory=dict) diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py index 1c8352e5..9d31d688 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -1,25 +1,40 @@ +import asyncio import inspect import json import keyword +import math +import threading +import time import uuid from dataclasses import dataclass -from typing import Any, Awaitable, Callable, Protocol +from typing import Any, Awaitable, Callable, NoReturn, Protocol, cast from pydantic import BaseModel from hyperforge.codemode import ( RestrictedPythonTask, + SandboxMessage, SandboxRunner, WorkerExecutionRequest, + encode_protocol_value, + encode_sandbox_message, ) -from hyperforge.codemode.sandbox import settings +from hyperforge.codemode.sandbox import MAX_PACKET_BYTES, settings from hyperforge.definition import FunctionDefinition from ..context import format_context +from ..models import HarnessEventType +from ..usage import UsageLimitExceeded from . import HarnessTool, ToolCallContext, ToolInheritancePolicy, tool CODEMODE_TOOL_NAME = "codemode" OUTPUT_FUNCTION_NAME = "output" +DEFAULT_MAX_SOURCE_BYTES = 64 * 1024 +DEFAULT_MAX_RESULT_BYTES = 1024 * 1024 +DEFAULT_MAX_CUMULATIVE_RESULT_BYTES = 4 * 1024 * 1024 +DEFAULT_MAX_OUTPUT_BYTES = 1024 * 1024 +DEFAULT_MAX_NESTED_CALLS = 20 +DEFAULT_MAX_CONCURRENT_EXECUTIONS = 4 RESERVED_CAPABILITY_NAMES = frozenset( { "Any", @@ -44,12 +59,39 @@ "save", } ) +_SENSITIVE_FIELD_PARTS = ( + "api_key", + "apikey", + "authorization", + "cookie", + "credential", + "passwd", + "password", + "private_key", + "privatekey", + "secret", + "token", +) +_EVENT_SANITIZE_MAX_NODES = 64 +_EVENT_SANITIZE_MAX_ITEMS = 16 +_EVENT_SANITIZE_MAX_DEPTH = 6 +_EVENT_STRING_CHUNK_CHARS = 4096 +_EVENT_CONTEXT_MAX_STRING_BYTES = 1024 +_EVENT_CONTEXT_MAX_INTEGER_BYTES = 1024 +_EVENT_CONTEXT_MAX_TOTAL_BYTES = 4096 +_SOURCE_SIZE_CHUNK_CHARS = 4096 +_CAPABILITY_EVENT_FAILURE = "Code Mode capability failed" +_POLICY_CHECK_FAILURE = "Code Mode policy check failed" type CodeModeResultAdapter[OutputT: BaseModel] = Callable[ [HarnessTool[Any, OutputT], OutputT], Any | Awaitable[Any] ] type CodeModeDispatch = Callable[[RestrictedPythonTask], Awaitable[Any]] +class _CapabilityEventFailure(RuntimeError): + pass + + class CodemodeInput(BaseModel): code: str question: str = "" @@ -74,6 +116,13 @@ def raw_codemode_result_adapter[OutputT: BaseModel]( class CodeModeRunner(Protocol): + """Execute a request and route worker callbacks through dispatch. + + Injected runners are explicitly trusted and own their callback lifecycle. + The production default uses the remote sandbox, or an isolated local process + only when ``remote_required=False``. + """ + async def run( self, request: WorkerExecutionRequest, dispatch: CodeModeDispatch ) -> Any: ... @@ -85,6 +134,59 @@ class CodeModeCapability[OutputT: BaseModel]: result_adapter: CodeModeResultAdapter[OutputT] = context_codemode_result_adapter +@dataclass(frozen=True) +class CodeModeLimits: + max_source_bytes: int = DEFAULT_MAX_SOURCE_BYTES + max_result_bytes: int = DEFAULT_MAX_RESULT_BYTES + max_cumulative_result_bytes: int = DEFAULT_MAX_CUMULATIVE_RESULT_BYTES + max_output_bytes: int = DEFAULT_MAX_OUTPUT_BYTES + max_nested_calls: int = DEFAULT_MAX_NESTED_CALLS + + def __post_init__(self) -> None: + for name, value in vars(self).items(): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + if self.max_cumulative_result_bytes < self.max_result_bytes: + raise ValueError("max_cumulative_result_bytes must be >= max_result_bytes") + + +class CodeModeExecutionLimiter: + """Fail-fast admission control shared by one or more Code Mode tools.""" + + def __init__( + self, max_concurrent_executions: int = DEFAULT_MAX_CONCURRENT_EXECUTIONS + ) -> None: + if ( + not isinstance(max_concurrent_executions, int) + or isinstance(max_concurrent_executions, bool) + or max_concurrent_executions <= 0 + ): + raise ValueError("max_concurrent_executions must be a positive integer") + self._max_concurrent_executions = max_concurrent_executions + self._active = 0 + self._lock = threading.Lock() + + @property + def max_concurrent_executions(self) -> int: + return self._max_concurrent_executions + + def acquire(self) -> None: + with self._lock: + if self._active >= self.max_concurrent_executions: + raise RuntimeError( + "Code Mode concurrency limit reached: " + f"{self._active} >= {self.max_concurrent_executions}" + ) + self._active += 1 + + def release(self) -> None: + with self._lock: + self._active -= 1 + + +_process_execution_limiter = CodeModeExecutionLimiter() + + @tool( name=CODEMODE_TOOL_NAME, description=( @@ -158,12 +260,20 @@ async def dispatch(task: RestrictedPythonTask) -> Any: def create_codemode_tool( *, capabilities: tuple[CodeModeCapability[Any], ...], + limits: CodeModeLimits = CodeModeLimits(), + execution_limiter: CodeModeExecutionLimiter = _process_execution_limiter, + remote_required: bool = True, runner: CodeModeRunner | None = None, inheritance: ToolInheritancePolicy = ToolInheritancePolicy.DO_NOT_INHERIT, name: str = CODEMODE_TOOL_NAME, description: str | None = None, ) -> HarnessTool[CodemodeInput, CodemodeOutput]: - """Create a Code Mode tool from an explicit immutable capability set.""" + """Create a Code Mode tool with only the explicitly supplied capabilities. + + ``runner`` defaults to the remote sandbox (or the isolated local process + when ``remote_required=False``); inject a ``CodeModeRunner`` for + deterministic tests without any sandbox environment. + """ if not isinstance(capabilities, tuple): raise TypeError("capabilities must be an immutable tuple") capability_map: dict[str, CodeModeCapability[Any]] = {} @@ -198,6 +308,7 @@ def create_codemode_tool( "'agent_id' input field" ) capability_map[capability_name] = capability + scoped_capabilities = tuple(capabilities) if inheritance == ToolInheritancePolicy.INHERIT: non_inherited = sorted( capability.tool.name @@ -214,37 +325,144 @@ async def execute( context: ToolCallContext, input_value: CodemodeInput ) -> CodemodeOutput: harness = context.harness - result = CodemodeOutput() - output_calls = 0 - output_succeeded = False - - async def dispatch(task: RestrictedPythonTask) -> Any: - nonlocal output_calls, output_succeeded - if task.function == OUTPUT_FUNCTION_NAME: - output_calls += 1 - if output_calls > 1: - raise ValueError("output(value) may only be called once") - if task.args and task.keyword_args: - raise ValueError( - "output accepts either a positional or keyword value" - ) - if len(task.args) > 1: - raise ValueError("output accepts one value") - if not task.args and set(task.keyword_args) != {"value"}: - raise ValueError("output requires exactly one 'value' argument") - value = task.args[0] if task.args else task.keyword_args["value"] - result.value = _normalize_json_value( - value, - label="Code Mode output", + turn_id = harness.turn_id + if _utf8_exceeds_limit(input_value.code, limits.max_source_bytes): + raise ValueError( + "Code Mode source exceeds maximum size: " + f"> {limits.max_source_bytes} bytes" + ) + socket = settings.sandbox_socket if runner is None and remote_required else None + sandbox_token = ( + settings.sandbox_token if runner is None and remote_required else None + ) + if runner is None and remote_required: + if socket is None: + raise RuntimeError( + "Remote Code Mode execution is required but SANDBOX_SOCKET " + "is absent" + ) + if not sandbox_token: + raise RuntimeError( + "Remote Code Mode execution is required but SANDBOX_TOKEN is absent" ) - output_succeeded = True - return None + execution_limiter.acquire() + return await _execute_scoped_codemode( + harness, + input_value, + scoped_capabilities, + limits, + socket, + sandbox_token, + execution_limiter, + runner, + remote_required, + parent_call_id=context.id, + turn_id=turn_id, + ) - capability = capability_map.get(task.function) - if capability is None: - raise ValueError(f"Unknown Code Mode capability: {task.function}") - harness.usage.tool_calls += 1 + execute.__name__ = name + tool_description = _scoped_description(scoped_capabilities, description) + return HarnessTool( + name=name, + handler=execute, + description=tool_description, + inheritance=inheritance, + ) + + +async def _execute_scoped_codemode( + harness: Any, + input_value: CodemodeInput, + capabilities: tuple[CodeModeCapability[Any], ...], + limits: CodeModeLimits, + socket: str | None, + sandbox_token: str | None, + execution_limiter: CodeModeExecutionLimiter, + runner: CodeModeRunner | None = None, + remote_required: bool = True, + *, + parent_call_id: str | None = None, + turn_id: str | None = None, +) -> CodemodeOutput: + capability_map = {capability.tool.name: capability for capability in capabilities} + result = CodemodeOutput() + invocation_state = _InvocationState() + output_state = _OutputState(invocation_state) + nested_calls = 0 + cumulative_result_bytes = 0 + try: + event_context = _nested_event_context(harness) + except (TypeError, ValueError) as exc: + invocation_state.latch(exc) + event_context = {} + + def event_failure() -> Exception: + error = _CapabilityEventFailure(_CAPABILITY_EVENT_FAILURE) + invocation_state.latch(error) + return error + + async def emit_nested( + event_type: HarnessEventType, payload: dict[str, Any] + ) -> None: + try: + await harness.emit( + event_type, + payload, + parent_call_id=parent_call_id, + turn_id=turn_id, + ) + except asyncio.CancelledError: + raise + except Exception: + raise event_failure() from None + + async def dispatch(task: RestrictedPythonTask) -> Any: + nonlocal nested_calls, cumulative_result_bytes + if task.function == OUTPUT_FUNCTION_NAME: + result.value = output_state.record(task, limits.max_output_bytes) + return None + + invocation_state.raise_if_failed() + capability = capability_map.get(task.function) + if capability is None: + invocation_state.fail( + ValueError(f"Unknown Code Mode capability: {task.function}") + ) + if nested_calls >= limits.max_nested_calls: + invocation_state.fail( + RuntimeError( + "Code Mode nested call limit exceeded: " + f"{nested_calls + 1} > {limits.max_nested_calls}" + ) + ) + nested_calls += 1 + harness.usage.tool_calls += 1 + try: harness._check_limit("max_tool_calls", harness.usage.tool_calls) + except UsageLimitExceeded as exc: + invocation_state.fail(exc) + except Exception: + invocation_state.fail(RuntimeError(_POLICY_CHECK_FAILURE)) + call_id = uuid.uuid4().hex + marker = { + "codemode": True, + "nested": True, + "parent_call_id": parent_call_id, + **event_context, + } + await emit_nested( + HarnessEventType.TOOL_REQUESTED, + { + "call": { + "id": call_id, + "name": capability.tool.name, + "arguments": _sanitize_task_arguments(capability.tool, task), + }, + **marker, + }, + ) + started = time.perf_counter() + try: arguments = _scoped_tool_arguments( capability.tool, task.args, task.keyword_args ) @@ -252,18 +470,97 @@ async def dispatch(task: RestrictedPythonTask) -> Any: ToolCallContext( harness=harness, name=capability.tool.name, - id=uuid.uuid4().hex, + id=call_id, + turn_id=turn_id, + _emit_failure=event_failure, ), arguments, ) + invocation_state.raise_if_failed() projected = capability.result_adapter(capability.tool, output) if inspect.isawaitable(projected): projected = await projected - return _normalize_json_value( - projected, - label="Code Mode result adapter", + try: + normalized, result_bytes = _normalize_worker_value( + projected, + limits.max_result_bytes, + f"Code Mode result from {capability.tool.name}", + ) + except (TypeError, ValueError) as exc: + invocation_state.latch(exc) + raise + try: + transport_label = ( + f"Code Mode result transport from {capability.tool.name}" + ) + if remote_required: + encode_sandbox_message( + SandboxMessage.Response(result=normalized), + transport_label, + max_bytes=MAX_PACKET_BYTES, + ) + else: + encode_protocol_value( + normalized, + transport_label, + max_bytes=MAX_PACKET_BYTES, + ) + except (TypeError, ValueError) as exc: + invocation_state.latch(exc) + raise + projected_cumulative_bytes = cumulative_result_bytes + result_bytes + if projected_cumulative_bytes > limits.max_cumulative_result_bytes: + error = RuntimeError( + "Code Mode cumulative result limit exceeded: " + f"{projected_cumulative_bytes} > " + f"{limits.max_cumulative_result_bytes} bytes" + ) + invocation_state.latch(error) + raise error + cumulative_result_bytes = projected_cumulative_bytes + sanitized_result = _sanitize_event_value(normalized) + except BaseException as exc: + await emit_nested( + HarnessEventType.TOOL_FAILED, + { + "call_id": call_id, + "tool": capability.tool.name, + "result": {"error": type(exc).__name__}, + "duration_ms": round((time.perf_counter() - started) * 1000, 3), + **marker, + }, + ) + if isinstance(exc, asyncio.CancelledError): + raise + if isinstance(exc, _CapabilityEventFailure): + raise RuntimeError(_CAPABILITY_EVENT_FAILURE) from None + raise RuntimeError( + f"Code Mode capability {capability.tool.name!r} failed " + f"({type(exc).__name__})" + ) from None + else: + await emit_nested( + HarnessEventType.TOOL_COMPLETED, + { + "call_id": call_id, + "tool": capability.tool.name, + "result": sanitized_result, + "duration_ms": round((time.perf_counter() - started) * 1000, 3), + "result_bytes": result_bytes, + **marker, + }, ) + return normalized + sandbox_runner: SandboxRunner | None = None + try: + invocation_state.raise_if_failed() + if runner is None: + sandbox_runner = ( + SandboxRunner.remote(cast(str, socket), dispatch, token=sandbox_token) + if remote_required + else SandboxRunner.isolated_process(dispatch) + ) definitions = { capability.tool.name: FunctionDefinition( name=capability.tool.name, @@ -277,7 +574,7 @@ async def dispatch(task: RestrictedPythonTask) -> Any: description="Set the value returned by Code Mode.", parameters={"value": {}}, ) - request = WorkerExecutionRequest( + worker_request = WorkerExecutionRequest( code=input_value.code, question=input_value.question, local_vars={}, @@ -285,55 +582,281 @@ async def dispatch(task: RestrictedPythonTask) -> Any: function_names={"harness": definitions}, max_runtime_seconds=harness.usage_limits.max_codemode_runtime_seconds, max_memory_bytes=harness.usage_limits.max_codemode_memory_bytes, + redact_errors=True, ) if runner is not None: - await runner.run(request, dispatch) - else: - sandbox_runner = ( - SandboxRunner.remote(settings.sandbox_socket, dispatch) - if settings.sandbox_socket is not None - else SandboxRunner.isolated_process(dispatch) + await runner.run(worker_request, dispatch) + elif sandbox_runner is not None: + await sandbox_runner.run(worker_request) + finally: + hold = ( + sandbox_runner.run_when_callbacks_complete + if sandbox_runner is not None + else None + ) + if hold is None or not hold(execution_limiter.release): + execution_limiter.release() + invocation_state.raise_if_failed() + if not output_state.called: + raise ValueError("Code Mode must call output(value) exactly once") + return result + + +class _InvocationState: + def __init__(self) -> None: + self.violation: tuple[type[Exception], str] | None = None + + def latch(self, exc: Exception) -> None: + if self.violation is None: + self.violation = (type(exc), str(exc)) + + def fail(self, exc: Exception) -> NoReturn: + self.latch(exc) + if self.violation is None: + raise AssertionError("Code Mode violation was not latched") + exception_type, message = self.violation + raise exception_type(message) from None + + def raise_if_failed(self) -> None: + if self.violation is not None: + exception_type, message = self.violation + raise exception_type(message) from None + + +class _OutputState: + """Validate output(value) calls: exactly one call, exactly one value.""" + + def __init__(self, invocation_state: _InvocationState) -> None: + self.invocation_state = invocation_state + self.attempts = 0 + self.called = False + + def record(self, task: RestrictedPythonTask, max_output_bytes: int) -> Any: + self.attempts += 1 + self.invocation_state.raise_if_failed() + if self.attempts > 1: + error = ValueError("output(value) may only be called once") + self.invocation_state.latch(error) + raise error + try: + value = _output_value(task) + normalized, _ = _normalize_worker_value( + value, max_output_bytes, "Code Mode output" + ) + except (TypeError, ValueError) as exc: + self.invocation_state.latch(exc) + raise + self.called = True + return normalized + + +def _output_value(task: RestrictedPythonTask) -> Any: + if task.args and task.keyword_args: + raise ValueError("output accepts either a positional or keyword value") + if len(task.args) > 1: + raise ValueError("output accepts one value") + if task.keyword_args.keys() - {"value"}: + raise ValueError("output accepts only the 'value' keyword") + if not task.args and "value" not in task.keyword_args: + raise ValueError("output requires a value") + return task.args[0] if task.args else task.keyword_args["value"] + + +def _normalize_worker_value(value: Any, max_bytes: int, label: str) -> tuple[Any, int]: + try: + encoded = encode_protocol_value( + value, + label, + max_bytes=max_bytes, + strict_worker_value=True, + ) + return json.loads(encoded), len(encoded) + except TypeError: + raise + except (RecursionError, ValueError) as exc: + if "exceeds maximum size" in str(exc): + raise ValueError(str(exc)) from exc + if "reserved '__model__' key" in str(exc): + raise ValueError(str(exc)) from exc + raise TypeError(f"{label} is not a serializable worker value") from exc + + +def _utf8_exceeds_limit(value: str, max_bytes: int) -> bool: + observed = 0 + for start in range(0, len(value), _SOURCE_SIZE_CHUNK_CHARS): + observed += len(value[start : start + _SOURCE_SIZE_CHUNK_CHARS].encode("utf-8")) + if observed > max_bytes: + return True + return False + + +class _EventSanitizer: + def __init__(self) -> None: + self.remaining_nodes = _EVENT_SANITIZE_MAX_NODES + + def sanitize(self, value: Any, *, depth: int = 0) -> Any: + self.remaining_nodes -= 1 + if isinstance(value, dict): + return self._sanitize_dict(value, depth) + if isinstance(value, (list, tuple)): + return self._sanitize_sequence(value, depth) + if isinstance(value, str): + return {"type": "string", "utf8_bytes": _utf8_size(value)} + if isinstance(value, float) and not math.isfinite(value): + return "" + if isinstance(value, (bool, int, float)) or value is None: + return value + return {"type": "other"} + + def _sanitize_dict(self, value: dict[Any, Any], depth: int) -> dict[str, Any]: + entries: list[dict[str, Any]] = [] + result: dict[str, Any] = {"type": "object", "entries": entries} + total = len(value) + if depth >= _EVENT_SANITIZE_MAX_DEPTH: + if total: + result["truncated"] = { + "reason": "depth_limit", + "omitted_entries": total, + } + return result + + reason = "item_limit" + for key, nested in value.items(): + if len(entries) >= _EVENT_SANITIZE_MAX_ITEMS: + break + if self.remaining_nodes <= 0: + reason = "node_budget" + break + key_text = _event_key_text(key) + sensitive = any( + part in key_text.casefold() for part in _SENSITIVE_FIELD_PARTS + ) + if sensitive: + self.remaining_nodes -= 1 + sanitized = "[REDACTED]" + else: + sanitized = self.sanitize(nested, depth=depth + 1) + entries.append( + { + "key": ( + {"type": "string", "utf8_bytes": _utf8_size(key_text)} + if isinstance(key, str) + else {"type": "other"} + ), + "value": sanitized, + } ) - await sandbox_runner.run(request) - if output_calls != 1 or not output_succeeded: - raise ValueError("Code Mode must call output(value) exactly once") + if len(entries) < total: + result["truncated"] = { + "reason": reason, + "omitted_entries": total - len(entries), + } return result - execute.__name__ = name - return HarnessTool( - name=name, - handler=execute, - description=_scoped_description(capabilities, description), - inheritance=inheritance, + def _sanitize_sequence( + self, value: list[Any] | tuple[Any, ...], depth: int + ) -> dict[str, Any]: + items: list[Any] = [] + result: dict[str, Any] = {"type": "array", "items": items} + total = len(value) + if depth >= _EVENT_SANITIZE_MAX_DEPTH: + if total: + result["truncated"] = { + "reason": "depth_limit", + "omitted_items": total, + } + return result + + reason = "item_limit" + for nested in value: + if len(items) >= _EVENT_SANITIZE_MAX_ITEMS: + break + if self.remaining_nodes <= 0: + reason = "node_budget" + break + items.append(self.sanitize(nested, depth=depth + 1)) + if len(items) < total: + result["truncated"] = { + "reason": reason, + "omitted_items": total - len(items), + } + return result + + +def _sanitize_event_value(value: Any) -> Any: + return _EventSanitizer().sanitize(value) + + +def _event_key_text(value: Any) -> str: + if isinstance(value, str): + return value + if value is None or isinstance(value, (bool, int, float)): + return str(value) + return "" + + +def _utf8_size(value: str) -> int: + return sum( + len(value[start : start + _EVENT_STRING_CHUNK_CHARS].encode("utf-8")) + for start in range(0, len(value), _EVENT_STRING_CHUNK_CHARS) ) -def _normalize_json_value(value: Any, *, label: str) -> Any: - _reject_reserved_json_values(value, label=label) +def _sanitize_task_arguments( + capability: HarnessTool[Any, Any], task: RestrictedPythonTask +) -> Any: try: - return json.loads(json.dumps(value, allow_nan=False)) - except TypeError as exc: - raise TypeError(f"{label} must be JSON-serializable") from exc - except ValueError as exc: - raise ValueError(f"{label} must be a strict JSON value: {exc}") from exc - - -def _reject_reserved_json_values(value: Any, *, label: str) -> None: - if isinstance(value, BaseModel): - raise TypeError(f"{label} must not contain Pydantic model values") - if isinstance(value, dict): - if "__model__" in value: - raise ValueError(f"{label} cannot contain the reserved '__model__' key") - for key, nested in value.items(): - _reject_reserved_json_values(key, label=label) - _reject_reserved_json_values(nested, label=label) - elif isinstance(value, (list, tuple)): - for nested in value: - _reject_reserved_json_values(nested, label=label) + arguments = _scoped_tool_arguments(capability, task.args, task.keyword_args) + except ValueError: + return {"invalid": True} + return _sanitize_event_value(arguments) + + +def _nested_event_context(harness: Any) -> dict[str, Any]: + context: dict[str, str | int | float | bool] = {} + total_value_bytes = 0 + for key in ("actor", "actor_id", "tenant", "tenant_id", "user_id"): + value = harness.execution_context.get(key) + if isinstance(value, str): + if _utf8_exceeds_limit(value, _EVENT_CONTEXT_MAX_STRING_BYTES): + raise ValueError( + f"Code Mode execution context marker {key!r} exceeds maximum size" + ) + value_bytes = _utf8_size(value) + context[key] = value + elif isinstance(value, bool): + value_bytes = len(str(value).lower()) + context[key] = value + elif isinstance(value, int): + if int.bit_length(value) > _EVENT_CONTEXT_MAX_INTEGER_BYTES * 4: + raise ValueError( + f"Code Mode execution context marker {key!r} exceeds maximum size" + ) + value_bytes = len(int.__repr__(value)) + if value_bytes > _EVENT_CONTEXT_MAX_INTEGER_BYTES: + raise ValueError( + f"Code Mode execution context marker {key!r} exceeds maximum size" + ) + context[key] = value + elif isinstance(value, float): + if not math.isfinite(value): + raise ValueError( + "Code Mode execution context markers must contain finite numbers" + ) + value_bytes = len(json.dumps(value, allow_nan=False)) + context[key] = value + else: + continue + total_value_bytes += value_bytes + if total_value_bytes > _EVENT_CONTEXT_MAX_TOTAL_BYTES: + raise ValueError( + "Code Mode execution context markers exceed maximum total size" + ) + return {"execution_context": context} if context else {} def _scoped_description( - capabilities: tuple[CodeModeCapability[Any], ...], description: str | None + capabilities: tuple[CodeModeCapability[Any], ...], description: str | None = None ) -> str: introduction = description or ( "Execute restricted Python code using only the scoped capabilities below; " @@ -350,8 +873,7 @@ def _scoped_description( sort_keys=True, ) definitions.append( - f"{capability.tool.name}: " - f"{capability.tool.description or 'No description.'} " + f"{capability.tool.name}: {capability.tool.description or 'No description.'} " f"Arguments JSON Schema: {schema}" ) return f"{introduction}\nCapability functions:\n" + "\n".join(definitions) @@ -398,6 +920,8 @@ def _scoped_tool_arguments( __all__ = [ "CodeModeCapability", "CodeModeDispatch", + "CodeModeExecutionLimiter", + "CodeModeLimits", "CodeModeResultAdapter", "CodeModeRunner", "CodemodeInput", diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/utils.py b/hyperforge/src/hyperforge/harness_sdk/tools/utils.py index 7c6740bd..9e0b5951 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/utils.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/utils.py @@ -10,7 +10,12 @@ from pydantic import BaseModel, ValidationError from ..context import make_context -from ..models import HarnessContextReference, HarnessContextType +from ..models import ( + HarnessContextReference, + HarnessContextType, + HarnessEvent, + HarnessEventType, +) from ..schema import flatten_json_schema if TYPE_CHECKING: @@ -22,6 +27,13 @@ type ContextFactory[OutputT: BaseModel] = Callable[[OutputT], HarnessContextReference] +class _CurrentTurn: + pass + + +_CURRENT_TURN = _CurrentTurn() + + class ToolInheritancePolicy(StrEnum): INHERIT = "inherit" DO_NOT_INHERIT = "do_not_inherit" @@ -38,6 +50,38 @@ class ToolCallContext: harness: "AgentHarness" name: str id: str | None = None + turn_id: str | None | _CurrentTurn = _CURRENT_TURN + _emit_failure: Callable[[], Exception] | None = field( + default=None, repr=False, compare=False + ) + + async def emit( + self, + event_type: HarnessEventType, + payload: dict[str, Any], + *, + persist: bool = True, + ) -> "HarnessEvent": + """Emit an event attributed to this tool call.""" + try: + if not isinstance(self.turn_id, _CurrentTurn): + return await self.harness.emit( + event_type, + payload, + persist=persist, + parent_call_id=self.id, + turn_id=self.turn_id, + ) + return await self.harness.emit( + event_type, + payload, + persist=persist, + parent_call_id=self.id, + ) + except Exception: + if self._emit_failure is None: + raise + raise self._emit_failure() from None @dataclass(frozen=True) diff --git a/hyperforge/src/hyperforge/server/sandbox.py b/hyperforge/src/hyperforge/server/sandbox.py index 36b2b2ba..31e4a0af 100644 --- a/hyperforge/src/hyperforge/server/sandbox.py +++ b/hyperforge/src/hyperforge/server/sandbox.py @@ -22,20 +22,21 @@ def set_sentry(zone: str, environment: str, sentry_url: str): sentry_sdk.set_tag("zone", zone) -async def run_metrics_server(port: int): +async def run_metrics_server(host: str, port: int): import uvicorn - config = uvicorn.Config( - application_metrics, host="0.0.0.0", port=port, log_level="info" - ) + config = uvicorn.Config(application_metrics, host=host, port=port, log_level="info") server = uvicorn.Server(config) await server.serve() -async def run_servers(metrics_port: int) -> None: +async def run_servers( + *, metrics_enabled: bool, metrics_host: str, metrics_port: int +) -> None: async with asyncio.TaskGroup() as tasks: - tasks.create_task(run_metrics_server(metrics_port)) tasks.create_task(sandbox.run_sandbox_server()) + if metrics_enabled: + tasks.create_task(run_metrics_server(metrics_host, metrics_port)) def run(): # pragma: no cover @@ -57,4 +58,10 @@ def run(): # pragma: no cover if sandbox_settings.sandbox_socket is None: raise RuntimeError("SANDBOX_SOCKET is required for the sandbox server") - asyncio.run(run_servers(sandbox_settings.sandbox_metrics_port)) + asyncio.run( + run_servers( + metrics_enabled=sandbox_settings.sandbox_metrics_enabled, + metrics_host=sandbox_settings.sandbox_metrics_host, + metrics_port=sandbox_settings.sandbox_metrics_port, + ) + ) diff --git a/hyperforge/tests/codemode/test_model.py b/hyperforge/tests/codemode/test_model.py index 92f86389..f437ba41 100644 --- a/hyperforge/tests/codemode/test_model.py +++ b/hyperforge/tests/codemode/test_model.py @@ -1,9 +1,33 @@ +import json +import tracemalloc +from dataclasses import dataclass +from typing import Annotated + +import pytest +from nuclia_models.predict.remi import RemiResponse +from pydantic import ( + BaseModel, + Field, + PlainSerializer, + RootModel, + ValidationError, + field_serializer, +) + from hyperforge.codemode import ( + RestrictedPythonTask, SandboxMessage, + WorkerError, WorkerExecutionRequest, + decode_protocol_value, deserialize, + encode_protocol_value, + encode_sandbox_message, serialize, ) +from hyperforge.codemode import model as model_module +from hyperforge.memory import Context +from hyperforge.models import Chunk def test_nested_json_values_round_trip() -> None: @@ -27,3 +51,432 @@ def test_sandbox_run_token_round_trip() -> None: assert isinstance(parsed, SandboxMessage.Run) assert parsed.token == "secret" + + +def test_known_protocol_models_keep_exact_wire_encoding() -> None: + error = WorkerError(error="failed") + task = RestrictedPythonTask( + function="lookup", agent="agent", args=(1,), keyword_args={"value": 2} + ) + + assert encode_protocol_value(error, "Worker error") == ( + b'{"error":"failed","__model__":"WorkerError"}' + ) + assert encode_protocol_value(task, "Worker task") == ( + b'{"__model__":"RestrictedPythonTask","function":"lookup",' + b'"agent":"agent","args":[1],"keyword_args":{"value":2}}' + ) + assert ( + decode_protocol_value( + encode_protocol_value(error, "Worker error"), "Worker error" + ) + == error + ) + assert ( + decode_protocol_value(encode_protocol_value(task, "Worker task"), "Worker task") + == task + ) + + +def test_context_protocol_model_keeps_existing_wire_encoding() -> None: + context = Context( + original_question_uuid=None, + actual_question_uuid=None, + question="question", + chunks=[Chunk(chunk_id="chunk", text="content")], + source="source", + agent="agent", + ) + expected = json.dumps( + serialize(context), ensure_ascii=False, separators=(",", ":") + ).encode() + + encoded = encode_protocol_value(context, "Context") + + assert encoded == expected + assert decode_protocol_value(encoded, "Context") == context + + +def test_protocol_encoder_rejects_arbitrary_aliased_pydantic_model() -> None: + class AliasedModel(BaseModel): + value: str = Field(alias="wireValue") + + with pytest.raises(ValueError, match="not serializable"): + encode_protocol_value(AliasedModel(wireValue="value"), "Aliased model") + + +def test_protocol_encoder_rejects_arbitrary_root_model() -> None: + class Values(RootModel[list[int]]): + pass + + with pytest.raises(ValueError, match="not serializable"): + encode_protocol_value(Values([1, 2]), "Root model") + + +def test_strict_worker_encoder_rejects_nested_models_and_reserved_keys() -> None: + class Value(BaseModel): + item: int + + with pytest.raises(TypeError, match="must not contain Pydantic model values"): + encode_protocol_value( + {"nested": [Value(item=1)]}, + "Worker value", + strict_worker_value=True, + ) + with pytest.raises(ValueError, match="reserved '__model__' key"): + encode_protocol_value( + {"nested": {"__model__": "Context"}}, + "Worker value", + strict_worker_value=True, + ) + + +def test_worker_request_preserves_detailed_errors_by_default() -> None: + request = WorkerExecutionRequest( + code="", + local_vars={}, + global_vars={}, + function_names={}, + ) + + assert request.redact_errors is False + + +def test_bounded_sandbox_encoding_matches_existing_wire_format() -> None: + message = SandboxMessage.Run( + run=WorkerExecutionRequest( + code="output(value)", + local_vars={"value": (1, 2)}, + global_vars={}, + function_names={}, + ), + token="secret", + ) + + encoded = encode_sandbox_message(message, "Sandbox message") + + expected = json.loads(json.dumps(SandboxMessage.serialize(message))) + assert json.loads(encoded) == expected + + +def test_sandbox_run_round_trips_supported_marked_variable_models() -> None: + context = Context( + original_question_uuid=None, + actual_question_uuid=None, + question="question", + chunks=[Chunk(chunk_id="chunk", text="content")], + source="source", + agent="agent", + ) + remi = RemiResponse(time=0.25) + error = WorkerError(error="failed") + task = RestrictedPythonTask( + function="lookup", agent="agent", args=(), keyword_args={} + ) + message = SandboxMessage.Run( + run=WorkerExecutionRequest( + code="", + local_vars={"context": context, "remi": remi, "task": task}, + global_vars={"error": error}, + function_names={}, + ) + ) + + encoded = encode_sandbox_message(message, "Sandbox message") + wire = json.loads(encoded) + parsed = SandboxMessage.parse(wire) + + assert "__model__" not in wire["run"] + assert wire["run"]["local_vars"]["context"]["__model__"] == "Context" + assert wire["run"]["local_vars"]["remi"]["__model__"] == "RemiResponse" + assert wire["run"]["local_vars"]["task"]["__model__"] == "RestrictedPythonTask" + assert wire["run"]["global_vars"]["error"]["__model__"] == "WorkerError" + assert isinstance(parsed, SandboxMessage.Run) + assert parsed.run.local_vars == {"context": context, "remi": remi, "task": task} + assert parsed.run.global_vars == {"error": error} + + +def test_sandbox_run_round_trips_marker_bearing_variable_containers() -> None: + context = Context( + original_question_uuid=None, + actual_question_uuid=None, + question="question", + chunks=[Chunk(chunk_id="chunk", text="content")], + source="source", + agent="agent", + ) + error = WorkerError(error="actual error") + local_vars = { + "__model__": "WorkerError", + "error": "plain container", + "actual_error": error, + } + global_vars = { + "__model__": model_module._PLAIN_DICT_MODEL, + "value": {"actual_context": context}, + } + message = SandboxMessage.Run( + run=WorkerExecutionRequest( + code="", + local_vars=local_vars, + global_vars=global_vars, + function_names={}, + ) + ) + + parsed_streaming = SandboxMessage.parse( + json.loads(encode_sandbox_message(message, "Sandbox message")) + ) + parsed_compatibility = SandboxMessage.parse(SandboxMessage.serialize(message)) + + assert isinstance(parsed_streaming, SandboxMessage.Run) + assert isinstance(parsed_streaming.run.local_vars, dict) + assert isinstance(parsed_streaming.run.global_vars, dict) + assert parsed_streaming.run.local_vars == local_vars + assert parsed_streaming.run.global_vars == global_vars + assert isinstance(parsed_compatibility, SandboxMessage.Run) + assert parsed_compatibility.run.local_vars == local_vars + assert parsed_compatibility.run.global_vars == global_vars + + +@pytest.mark.parametrize("marker_kind", ["model", "malformed", "escape"]) +def test_restricted_task_round_trips_marker_bearing_keyword_args( + marker_kind: str, +) -> None: + context = Context( + original_question_uuid=None, + actual_question_uuid=None, + question="question", + chunks=[Chunk(chunk_id="chunk", text="content")], + source="source", + agent="agent", + ) + error = WorkerError(error="actual error") + if marker_kind == "model": + keyword_args = { + "__model__": "WorkerError", + "error": "plain container", + "actual": error, + } + elif marker_kind == "malformed": + keyword_args = {"__model__": 42, "actual": error} + else: + keyword_args = { + "__model__": model_module._PLAIN_DICT_MODEL, + "value": {"actual": error}, + } + marker_argument = {"__model__": "Context", "value": "plain argument"} + task = RestrictedPythonTask( + function="lookup", + agent="agent", + args=(marker_argument, context), + keyword_args=keyword_args, + ) + + parsed_streaming = decode_protocol_value( + encode_protocol_value(task, "Worker task"), "Worker task" + ) + parsed_compatibility = deserialize(serialize(task)) + + assert isinstance(parsed_streaming, RestrictedPythonTask) + assert parsed_streaming.args == (marker_argument, context) + assert parsed_streaming.keyword_args == keyword_args + assert isinstance(parsed_compatibility, RestrictedPythonTask) + assert parsed_compatibility.args == (marker_argument, context) + assert parsed_compatibility.keyword_args == keyword_args + + +@pytest.mark.parametrize( + "value", + [ + {"__model__": "WorkerError", "error": "plain dictionary"}, + {"__model__": 42, "value": "malformed marker"}, + { + "__model__": model_module._PLAIN_DICT_MODEL, + "value": {"plain": True}, + }, + ], +) +def test_plain_marker_dictionaries_round_trip_without_retyping( + value: dict[str, object], +) -> None: + message = SandboxMessage.Run( + run=WorkerExecutionRequest( + code="", + local_vars={"value": value}, + global_vars={"nested": [value]}, + function_names={}, + ) + ) + + encoded_value = decode_protocol_value( + encode_protocol_value(value, "Plain marker dictionary"), + "Plain marker dictionary", + ) + parsed_streaming = SandboxMessage.parse( + json.loads(encode_sandbox_message(message, "Sandbox message")) + ) + parsed_compatibility = SandboxMessage.parse(SandboxMessage.serialize(message)) + + assert encoded_value == value + assert isinstance(encoded_value, dict) + assert deserialize(serialize(value)) == value + assert isinstance(parsed_streaming, SandboxMessage.Run) + assert parsed_streaming.run.local_vars["value"] == value + assert parsed_streaming.run.global_vars["nested"] == [value] + assert isinstance(parsed_compatibility, SandboxMessage.Run) + assert parsed_compatibility.run.local_vars["value"] == value + assert parsed_compatibility.run.global_vars["nested"] == [value] + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("max_runtime_seconds", float("inf")), + ("max_runtime_seconds", 0), + ("max_memory_bytes", 0), + ], +) +def test_worker_request_rejects_invalid_resource_limits( + field: str, value: float | int +) -> None: + with pytest.raises(ValidationError): + WorkerExecutionRequest( + code="", + local_vars={}, + global_vars={}, + function_names={}, + **{field: value}, + ) + + +def test_protocol_encoder_streams_large_strings( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_dumps = model_module.json.dumps + encoded_string_lengths = [] + + def recording_dumps(value, *args, **kwargs): + if isinstance(value, str): + encoded_string_lengths.append(len(value)) + return original_dumps(value, *args, **kwargs) + + monkeypatch.setattr(model_module.json, "dumps", recording_dumps) + + encoded = encode_protocol_value("x" * 10_000, "Test value", max_bytes=10_002) + + assert len(encoded) == 10_002 + assert max(encoded_string_lengths) <= model_module._JSON_STRING_CHUNK_CHARS + + +def test_protocol_encoder_stops_at_byte_limit() -> None: + with pytest.raises(ValueError, match="maximum size"): + encode_protocol_value("x" * 1_000_000, "Test value", max_bytes=512) + + +def test_protocol_encoder_uses_base_integer_representation() -> None: + class MisleadingInt(int): + def __str__(self) -> str: + return "not-json" + + def __repr__(self) -> str: + return "also-not-json" + + value = MisleadingInt(7) + + assert encode_protocol_value(value, "Test value") == b"7" + assert encode_protocol_value({value: 1}, "Test value") == b'{"7":1}' + + +def test_protocol_decoder_rejects_deep_json_as_value_error() -> None: + encoded = b"[" * 1000 + b"0" + b"]" * 1000 + + with pytest.raises(ValueError, match="not valid JSON"): + decode_protocol_value(encoded, "Test value") + + +def test_protocol_decoder_rejects_malformed_model_marker() -> None: + encoded = ( + b'{"__model__":"RestrictedPythonTask","function":"test","agent":"agent",' + b'"args":[],"keyword_args":[]}' + ) + + with pytest.raises(ValueError, match="Test value is invalid"): + decode_protocol_value(encoded, "Test value") + + +@pytest.mark.parametrize("encoded", [b"1e9999", b'{"value":1e9999}']) +def test_protocol_decoder_rejects_non_finite_numbers(encoded: bytes) -> None: + with pytest.raises(ValueError, match="not valid JSON"): + decode_protocol_value(encoded, "Test value") + + +def test_sandbox_encoder_rejects_custom_pydantic_serializers() -> None: + class SecretModel(BaseModel): + secret: str + + @field_serializer("secret") + def redact_secret(self, _value: str) -> str: + return "[REDACTED]" + + message = SandboxMessage.Run( + run=WorkerExecutionRequest( + code="", + local_vars={"value": SecretModel(secret="must-not-leak")}, + global_vars={}, + function_names={}, + ) + ) + + with pytest.raises(ValueError, match="not serializable"): + encode_sandbox_message(message, "Sandbox message") + + +def test_sandbox_encoder_rejects_schema_level_serializers() -> None: + class SecretModel(BaseModel): + secret: Annotated[ + str, + PlainSerializer(lambda _value: "[REDACTED]", return_type=str), + ] + + message = SandboxMessage.Run( + run=WorkerExecutionRequest( + code="", + local_vars={"value": SecretModel(secret="must-not-leak")}, + global_vars={}, + function_names={}, + ) + ) + + with pytest.raises(ValueError, match="not serializable"): + encode_sandbox_message(message, "Sandbox message") + + +def test_sandbox_encoder_rejects_dataclass_values() -> None: + @dataclass + class Value: + item: int + + message = SandboxMessage.Run( + run=WorkerExecutionRequest( + code="", + local_vars={"value": Value(item=1)}, + global_vars={}, + function_names={}, + ) + ) + + with pytest.raises(ValueError, match="not serializable"): + encode_sandbox_message(message, "Sandbox message") + + +def test_protocol_depth_validation_uses_bounded_auxiliary_memory() -> None: + wide_value = [0] * 100_000 + tracemalloc.start() + try: + model_module._validate_json_depth(wide_value) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert peak < 1_000_000 diff --git a/hyperforge/tests/codemode/test_protocol.py b/hyperforge/tests/codemode/test_protocol.py index e452bc6a..de748375 100644 --- a/hyperforge/tests/codemode/test_protocol.py +++ b/hyperforge/tests/codemode/test_protocol.py @@ -1,18 +1,40 @@ import asyncio +import contextvars +import io import tempfile +import threading +import traceback +from concurrent.futures import Future as ConcurrentFuture +from contextlib import redirect_stderr from pathlib import Path import pytest +from nuclia_models.predict.remi import RemiResponse -from hyperforge.codemode.model import WorkerExecutionRequest, deserialize +from hyperforge.codemode import sandbox +from hyperforge.codemode import worker as worker_module +from hyperforge.codemode.model import ( + RestrictedPythonTask, + SandboxMessage, + WorkerError, + WorkerExecutionRequest, + decode_protocol_value, + deserialize, + encode_protocol_value, +) from hyperforge.codemode.sandbox import ( MAX_PACKET_BYTES, SandboxReader, SandboxRunner, + SandboxSession, SandboxSettings, + SandboxWriter, run_sandbox_server, ) +from hyperforge.codemode.worker import PythonAgentWorker from hyperforge.definition import FunctionDefinition +from hyperforge.memory import Context +from hyperforge.models import Chunk def _empty_request() -> WorkerExecutionRequest: @@ -30,6 +52,154 @@ def _function_names(*names: str) -> dict[str, dict[str, FunctionDefinition]]: } +def test_redacted_worker_transport_failure_does_not_expose_original_exception( + monkeypatch: pytest.MonkeyPatch, +) -> None: + secret = "worker-stderr-secret" + + class ClosedPipe: + closed = False + + def send_bytes(self, _payload: bytes) -> None: + raise BrokenPipeError("pipe closed") + + def close(self) -> None: + self.closed = True + + pipe = ClosedPipe() + worker = PythonAgentWorker(pipe) # type: ignore[arg-type] + monkeypatch.setattr(worker_module, "_harden_process", lambda _limit: None) + stderr = io.StringIO() + + with redirect_stderr(stderr): + try: + worker._process_question_context_sync( + f"raise RuntimeError('{secret}')", + "", + {}, + {}, + {}, + redact_errors=True, + ) + except BaseException: + traceback.print_exc() + + assert pipe.closed + assert secret not in stderr.getvalue() + + +def test_generic_worker_transport_failure_still_propagates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class ClosedPipe: + def send_bytes(self, _payload: bytes) -> None: + raise BrokenPipeError("pipe closed") + + worker = PythonAgentWorker(ClosedPipe()) # type: ignore[arg-type] + monkeypatch.setattr(worker_module, "_harden_process", lambda _limit: None) + + with pytest.raises(BrokenPipeError, match="pipe closed"): + worker._process_question_context_sync("1 + 1", "", {}, {}, {}) + + +@pytest.mark.parametrize( + "first_output", + ["output({1, 2})", "output('x' * 512)"], +) +def test_redacted_worker_latches_caught_output_serialization_failure( + monkeypatch: pytest.MonkeyPatch, first_output: str +) -> None: + class Pipe: + def __init__(self) -> None: + self.sent: list[bytes] = [] + + def send_bytes(self, payload: bytes) -> None: + self.sent.append(payload) + + def recv_bytes(self, _max_bytes: int) -> bytes: + return encode_protocol_value(None, "Worker response", max_bytes=256) + + pipe = Pipe() + worker = PythonAgentWorker(pipe) # type: ignore[arg-type] + monkeypatch.setattr(worker_module, "_harden_process", lambda _limit: None) + monkeypatch.setattr(worker_module, "MAX_PROTOCOL_BYTES", 256) + code = f""" +try: + {first_output} +except Exception: + pass +output('valid') +""" + + worker._process_question_context_sync( + code, + "", + {}, + {}, + _function_names("output"), + redact_errors=True, + ) + + tasks = [ + decode_protocol_value(payload, "Worker task", max_bytes=256) + for payload in pipe.sent + ] + assert [ + task.function for task in tasks if isinstance(task, RestrictedPythonTask) + ] == [ + "output", + "_error", + "_close", + ] + error = tasks[1] + assert isinstance(error, RestrictedPythonTask) + assert error.args == ("Generated code execution failed",) + + +def test_generic_worker_can_recover_from_output_serialization_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Pipe: + def __init__(self) -> None: + self.sent: list[bytes] = [] + + def send_bytes(self, payload: bytes) -> None: + self.sent.append(payload) + + def recv_bytes(self, _max_bytes: int) -> bytes: + return encode_protocol_value(None, "Worker response", max_bytes=256) + + pipe = Pipe() + worker = PythonAgentWorker(pipe) # type: ignore[arg-type] + monkeypatch.setattr(worker_module, "_harden_process", lambda _limit: None) + monkeypatch.setattr(worker_module, "MAX_PROTOCOL_BYTES", 256) + + worker._process_question_context_sync( + """ +try: + output({1, 2}) +except Exception: + pass +output('valid') +""", + "", + {}, + {}, + _function_names("output"), + ) + + tasks = [ + decode_protocol_value(payload, "Worker task", max_bytes=256) + for payload in pipe.sent + ] + assert [ + task.function for task in tasks if isinstance(task, RestrictedPythonTask) + ] == [ + "output", + "_close", + ] + + @pytest.fixture def socket_path(): with tempfile.TemporaryDirectory(prefix="sbx", dir="/tmp") as directory: @@ -51,6 +221,678 @@ async def test_reader_rejects_oversized_packet_before_payload() -> None: await SandboxReader(reader)._read_packet() +@pytest.mark.asyncio +async def test_reader_rejects_deep_packet_as_value_error() -> None: + payload = b'{"_":"response","result":' + b"[" * 1000 + b"0" + b"]" * 1000 + b"}" + reader = asyncio.StreamReader() + reader.feed_data(len(payload).to_bytes(4, "little") + payload) + + with pytest.raises(ValueError, match="not valid JSON"): + await SandboxReader(reader).read_message() + + +@pytest.mark.asyncio +async def test_writer_rejects_oversized_packet() -> None: + class Writer: + def write(self, data: bytes) -> None: + del data + + async def drain(self) -> None: + pass + + with pytest.raises(ValueError, match="maximum size"): + await SandboxWriter(Writer()).write_message( # type: ignore[arg-type] + SandboxMessage.Response(result="x" * MAX_PACKET_BYTES) + ) + + +@pytest.mark.asyncio +async def test_writer_rejects_deep_packet_as_value_error() -> None: + class Writer: + def write(self, data: bytes) -> None: + del data + + async def drain(self) -> None: + pass + + value: object = 0 + for _ in range(1000): + value = [value] + + with pytest.raises(ValueError, match="not serializable"): + await SandboxWriter(Writer()).write_message( # type: ignore[arg-type] + SandboxMessage.Response(result=value) # type: ignore[arg-type] + ) + + +@pytest.mark.asyncio +async def test_reader_normalizes_malformed_model_marker() -> None: + payload = ( + b'{"_":"response","result":{"__model__":"RestrictedPythonTask",' + b'"function":"test","agent":"agent","args":[],"keyword_args":[]}}' + ) + reader = asyncio.StreamReader() + reader.feed_data(len(payload).to_bytes(4, "little") + payload) + + with pytest.raises(ValueError, match="Invalid sandbox message"): + await SandboxReader(reader).read_message() + + +@pytest.mark.asyncio +async def test_local_protocol_rejects_oversized_callback_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("hyperforge.codemode.sandbox.MAX_PACKET_BYTES", 512) + + async def callback(_task: RestrictedPythonTask): + return "x" * 512 + + runner = SandboxRunner.isolated_process(callback) + + with pytest.raises(RuntimeError, match="maximum size"): + await runner.run( + WorkerExecutionRequest( + code="capability()", + local_vars={}, + global_vars={}, + function_names={ + "test": { + "capability": FunctionDefinition( + name="capability", description="", parameters={} + ) + } + }, + max_runtime_seconds=30, + ) + ) + + +@pytest.mark.asyncio +async def test_local_runner_preserves_python_request_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def callback(_task: RestrictedPythonTask): + return None + + runner = SandboxRunner.with_pool(object(), callback) # type: ignore[arg-type] + captured = None + + async def run_in_pool(request: WorkerExecutionRequest) -> None: + nonlocal captured + captured = request.local_vars["value"] + + monkeypatch.setattr(runner, "_run_in_pool", run_in_pool) + + await runner.run( + WorkerExecutionRequest( + code="", + local_vars={"value": b"not-json"}, + global_vars={}, + function_names={}, + ) + ) + + assert captured == b"not-json" + assert isinstance(captured, bytes) + + +@pytest.mark.asyncio +async def test_remote_callback_cancellation_is_bounded( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started = asyncio.Event() + release = asyncio.Event() + callback_done = asyncio.Event() + + async def callback(_task: RestrictedPythonTask): + started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + await release.wait() + finally: + callback_done.set() + + class Reader: + async def read_message(self): + await asyncio.Future() + + class Writer: + async def write_message(self, _message): + pytest.fail("cancelled callback must not send a response") + + monkeypatch.setattr("hyperforge.codemode.sandbox.CALLBACK_CANCEL_TIMEOUT", 0.01) + runner = SandboxRunner.isolated_process(callback) + runner._begin_callback_run() + execution = asyncio.create_task( + runner._run_remote_callback( + Reader(), # type: ignore[arg-type] + Writer(), # type: ignore[arg-type] + RestrictedPythonTask( + function="capability", agent="test", args=(), keyword_args={} + ), + ) + ) + await started.wait() + + execution.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(execution, timeout=0.2) + + release.set() + await asyncio.wait_for(callback_done.wait(), timeout=0.2) + await runner._cancel_callbacks() + runner._finish_callback_run() + + +@pytest.mark.asyncio +async def test_remote_callback_keeps_frame_arriving_during_response_write() -> None: + write_started = asyncio.Event() + frame_read = asyncio.Event() + + class Reader: + async def read_message(self): + await write_started.wait() + frame_read.set() + return SandboxMessage.Done() + + class Writer: + async def write_message(self, _message): + write_started.set() + await frame_read.wait() + + async def callback(_task: RestrictedPythonTask): + return None + + runner = SandboxRunner.isolated_process(callback) + runner._begin_callback_run() + incoming = await runner._run_remote_callback( + Reader(), # type: ignore[arg-type] + Writer(), # type: ignore[arg-type] + RestrictedPythonTask( + function="capability", agent="test", args=(), keyword_args={} + ), + ) + + message = await asyncio.wait_for(incoming, timeout=0.2) + + assert isinstance(message, SandboxMessage.Done) + await runner._cancel_callbacks() + runner._finish_callback_run() + + +@pytest.mark.asyncio +async def test_callback_scheduled_after_cleanup_cannot_start() -> None: + called = False + intent_registered = threading.Event() + allow_schedule = threading.Event() + registration_result = False + + async def callback(_task: RestrictedPythonTask): + nonlocal called + called = True + return None + + runner = SandboxRunner.isolated_process(callback) + runner._begin_callback_run() + future: ConcurrentFuture = ConcurrentFuture() + loop = asyncio.get_running_loop() + + def schedule_from_executor_thread() -> None: + nonlocal registration_result + registration_result = runner._register_callback_intent(future) + intent_registered.set() + allow_schedule.wait() + loop.call_soon_threadsafe( + runner._schedule_callback, + future, + RestrictedPythonTask( + function="capability", agent="test", args=(), keyword_args={} + ), + loop, + contextvars.copy_context(), + ) + + executor_thread = threading.Thread(target=schedule_from_executor_thread) + executor_thread.start() + await asyncio.to_thread(intent_registered.wait) + assert registration_result + + await runner._cancel_callbacks() + runner._finish_callback_run() + allow_schedule.set() + await asyncio.to_thread(executor_thread.join) + await asyncio.sleep(0) + + assert not called + with pytest.raises(RuntimeError, match="cancelled"): + future.result() + assert not runner.run_when_callbacks_complete(lambda: None) + + +@pytest.mark.asyncio +async def test_runner_rejects_reuse_until_orphaned_callback_finishes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started = asyncio.Event() + release = asyncio.Event() + + async def callback(_task: RestrictedPythonTask): + started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + await release.wait() + + monkeypatch.setattr(sandbox, "CALLBACK_CANCEL_TIMEOUT", 0.01) + runner = SandboxRunner.isolated_process(callback) + runner._begin_callback_run() + future: ConcurrentFuture = ConcurrentFuture() + assert runner._register_callback_intent(future) + runner._schedule_callback( + future, + RestrictedPythonTask( + function="capability", agent="test", args=(), keyword_args={} + ), + asyncio.get_running_loop(), + contextvars.copy_context(), + ) + callback_task = next(iter(runner._callback_tasks)) + await started.wait() + + await runner._cancel_callbacks() + runner._finish_callback_run() + + with pytest.raises(RuntimeError, match="callbacks have completed"): + await runner.run(_empty_request()) + + release.set() + await callback_task + with pytest.raises(RuntimeError, match="cancelled"): + future.result() + + async def complete(_request: WorkerExecutionRequest) -> None: + return None + + monkeypatch.setattr(runner, "_run_isolated", complete) + await runner.run(_empty_request()) + + +@pytest.mark.asyncio +async def test_isolated_runner_rejects_non_json_request_before_execution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def callback(_task: RestrictedPythonTask): + return None + + runner = SandboxRunner.isolated_process(callback) + + def unexpected_run(_request: WorkerExecutionRequest): + pytest.fail("invalid request reached the isolated worker") + + monkeypatch.setattr(runner, "run_in_process", unexpected_run) + + with pytest.raises(ValueError, match="not serializable"): + await runner.run( + WorkerExecutionRequest( + code="", + local_vars={"value": b"not-json"}, + global_vars={}, + function_names={}, + ) + ) + + +@pytest.mark.asyncio +async def test_isolated_runner_rejects_oversized_request_before_execution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def callback(_task: RestrictedPythonTask): + return None + + runner = SandboxRunner.isolated_process(callback) + + def unexpected_run(_request: WorkerExecutionRequest): + pytest.fail("oversized request reached the isolated worker") + + monkeypatch.setattr(runner, "run_in_process", unexpected_run) + monkeypatch.setattr(sandbox, "MAX_PACKET_BYTES", 256) + + with pytest.raises(ValueError, match="maximum size"): + await runner.run( + WorkerExecutionRequest( + code="x" * 512, + local_vars={}, + global_vars={}, + function_names={}, + ) + ) + + +@pytest.mark.asyncio +async def test_isolated_runner_executes_canonical_decoded_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def callback(_task: RestrictedPythonTask): + return None + + captured: WorkerExecutionRequest | None = None + runner = SandboxRunner.isolated_process(callback) + + async def capture(request: WorkerExecutionRequest) -> None: + nonlocal captured + captured = request + + monkeypatch.setattr(runner, "_run_isolated", capture) + + await runner.run( + WorkerExecutionRequest( + code="", + local_vars={"value": (1, 2)}, + global_vars={}, + function_names={}, + ) + ) + + assert captured is not None + assert captured.local_vars["value"] == [1, 2] + + +@pytest.mark.asyncio +async def test_isolated_runner_preserves_supported_marked_request_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def callback(_task: RestrictedPythonTask): + return None + + context = Context( + original_question_uuid=None, + actual_question_uuid=None, + question="question", + chunks=[Chunk(chunk_id="chunk", text="content")], + source="source", + agent="agent", + ) + remi = RemiResponse(time=0.25) + error = WorkerError(error="failed") + task = RestrictedPythonTask( + function="lookup", agent="agent", args=(), keyword_args={} + ) + captured: WorkerExecutionRequest | None = None + runner = SandboxRunner.isolated_process(callback) + + async def capture(request: WorkerExecutionRequest) -> None: + nonlocal captured + captured = request + + monkeypatch.setattr(runner, "_run_isolated", capture) + + await runner.run( + WorkerExecutionRequest( + code="", + local_vars={"context": context, "remi": remi, "task": task}, + global_vars={"error": error}, + function_names={}, + ) + ) + + assert captured is not None + assert captured.local_vars == {"context": context, "remi": remi, "task": task} + assert captured.global_vars == {"error": error} + + +@pytest.mark.asyncio +async def test_generic_isolated_runner_preserves_generated_runtime_error() -> None: + async def callback(_task: RestrictedPythonTask): + return None + + runner = SandboxRunner.isolated_process(callback) + request = WorkerExecutionRequest( + code="raise RuntimeError('legacy runtime detail')", + local_vars={}, + global_vars={}, + function_names={}, + ) + + with pytest.raises(RuntimeError, match="legacy runtime detail"): + await runner.run(request) + + +@pytest.mark.asyncio +async def test_sandbox_session_reports_process_start_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Writer: + def __init__(self) -> None: + self.messages: list[SandboxMessage.AnyMessage] = [] + + async def write_message(self, message: SandboxMessage.AnyMessage) -> None: + self.messages.append(message) + + writer = Writer() + session = SandboxSession( + object(), # type: ignore[arg-type] + writer, # type: ignore[arg-type] + ) + + def fail_to_start(_request: WorkerExecutionRequest): + raise OSError("process start failed") + + monkeypatch.setattr(session.runner, "run_in_process", fail_to_start) + + await session.run(_empty_request()) + + assert len(writer.messages) == 1 + response = writer.messages[0] + assert isinstance(response, SandboxMessage.Error) + assert response.error == "process start failed" + assert session.task is None + assert session.process is None + assert session.runner._callbacks_closed + assert not session.runner._run_active + + +@pytest.mark.asyncio +async def test_remote_run_consumes_frame_arriving_during_response_write( + monkeypatch: pytest.MonkeyPatch, socket_path: str +) -> None: + done_frame_read = asyncio.Event() + original_read = SandboxReader.read_message + original_write = SandboxWriter.write_message + + async def tracked_read(self): + message = await original_read(self) + if isinstance(message, SandboxMessage.Done): + done_frame_read.set() + return message + + async def stalled_write(self, message): + await original_write(self, message) + if isinstance(message, SandboxMessage.Response): + await done_frame_read.wait() + + monkeypatch.setattr(SandboxReader, "read_message", tracked_read) + monkeypatch.setattr(SandboxWriter, "write_message", stalled_write) + + async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter) -> None: + reader, writer = SandboxReader(rx), SandboxWriter(tx) + run_message = await reader.read_message() + assert isinstance(run_message, SandboxMessage.Run) + await writer.write_message( + SandboxMessage.Request( + task=RestrictedPythonTask( + function="capability", agent="test", args=(), keyword_args={} + ) + ) + ) + response = await reader.read_message() + assert isinstance(response, SandboxMessage.Response) + await writer.write_message(SandboxMessage.Done()) + writer.close() + await writer.wait_closed() + + server = await asyncio.start_unix_server(handler, socket_path) + + async def callback(_task: RestrictedPythonTask): + return None + + runner = SandboxRunner.remote(socket_path, callback, token="server-secret") + request = WorkerExecutionRequest( + code="", local_vars={}, global_vars={}, function_names={} + ) + try: + await asyncio.wait_for(runner.run(request), timeout=2) + finally: + server.close() + await server.wait_closed() + + +@pytest.mark.asyncio +async def test_remote_runners_share_client_session_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + active = 0 + maximum_active = 0 + + async def run_remotely( + _runner: SandboxRunner, _request: WorkerExecutionRequest + ) -> None: + nonlocal active, maximum_active + active += 1 + maximum_active = max(maximum_active, active) + await asyncio.sleep(0.01) + active -= 1 + + async def callback(_task: RestrictedPythonTask): + return None + + monkeypatch.setattr(sandbox.settings, "sandbox_max_concurrent_sessions", 2) + monkeypatch.setattr(SandboxRunner, "_run_remotely", run_remotely) + request = WorkerExecutionRequest( + code="", local_vars={}, global_vars={}, function_names={} + ) + runners = [SandboxRunner.remote("/unused.sock", callback) for _ in range(5)] + + await asyncio.gather(*(runner.run(request) for runner in runners)) + + assert maximum_active == 2 + + +@pytest.mark.asyncio +async def test_remote_session_limit_is_shared_across_event_loops( + monkeypatch: pytest.MonkeyPatch, +) -> None: + active = 0 + maximum_active = 0 + counter_lock = threading.Lock() + + async def run_remotely( + _runner: SandboxRunner, _request: WorkerExecutionRequest + ) -> None: + nonlocal active, maximum_active + with counter_lock: + active += 1 + maximum_active = max(maximum_active, active) + try: + await asyncio.sleep(0.02) + finally: + with counter_lock: + active -= 1 + + async def callback(_task: RestrictedPythonTask): + return None + + request = WorkerExecutionRequest( + code="", local_vars={}, global_vars={}, function_names={} + ) + + def run_in_new_loop() -> None: + runner = SandboxRunner.remote("/unused.sock", callback) + asyncio.run(runner.run(request)) + + monkeypatch.setattr(sandbox.settings, "sandbox_max_concurrent_sessions", 2) + monkeypatch.setattr(SandboxRunner, "_run_remotely", run_remotely) + + await asyncio.gather(*(asyncio.to_thread(run_in_new_loop) for _ in range(6))) + + assert maximum_active == 2 + + +@pytest.mark.asyncio +async def test_remote_admission_release_uses_acquisition_ticket( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def callback(_task: RestrictedPythonTask): + return None + + async def change_limit( + _runner: SandboxRunner, _request: WorkerExecutionRequest + ) -> None: + monkeypatch.setattr(sandbox.settings, "sandbox_max_concurrent_sessions", None) + + monkeypatch.setattr(sandbox.settings, "sandbox_max_concurrent_sessions", 1) + monkeypatch.setattr(SandboxRunner, "_run_remotely", change_limit) + runner = SandboxRunner.remote("/unused.sock", callback) + + await runner.run(_empty_request()) + + assert sandbox._remote_admission_counter.active_sessions == 0 + + +@pytest.mark.asyncio +async def test_remote_run_timeout_adds_slack_to_requested_runtime( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(sandbox.settings, "sandbox_timeout_slack_seconds", 5.0) + completed = [] + + async def slow_admission(_runner: SandboxRunner, _request: WorkerExecutionRequest): + await asyncio.sleep(0.05) + completed.append(True) + + async def callback(_task: RestrictedPythonTask): + return None + + monkeypatch.setattr(SandboxRunner, "_run_with_remote_admission", slow_admission) + runner = SandboxRunner.remote("/unused.sock", callback) + request = WorkerExecutionRequest( + code="", + local_vars={}, + global_vars={}, + function_names={}, + max_runtime_seconds=0.01, + ) + + await runner.run(request) + + assert completed == [True] + + +@pytest.mark.asyncio +async def test_remote_run_without_requested_runtime_uses_server_ceiling( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(sandbox.settings, "sandbox_timeout_slack_seconds", 0.0) + monkeypatch.setattr(sandbox.settings, "sandbox_max_session_runtime_seconds", 0.05) + + async def hanging_admission( + _runner: SandboxRunner, _request: WorkerExecutionRequest + ): + await asyncio.Future() + + async def callback(_task: RestrictedPythonTask): + return None + + monkeypatch.setattr(SandboxRunner, "_run_with_remote_admission", hanging_admission) + runner = SandboxRunner.remote("/unused.sock", callback) + request = WorkerExecutionRequest( + code="", local_vars={}, global_vars={}, function_names={} + ) + + with pytest.raises(RuntimeError, match="timed out"): + await runner.run(request) + + async def _serve(monkeypatch, socket_path: str, **kwargs) -> asyncio.Task: monkeypatch.setenv("SANDBOX_SOCKET", socket_path) server_task = asyncio.create_task(run_sandbox_server(**kwargs)) @@ -100,6 +942,43 @@ async def test_sandbox_server_accepts_programmatic_client_token( await asyncio.gather(server_task, return_exceptions=True) +@pytest.mark.asyncio +async def test_sandbox_server_uses_one_settings_snapshot( + monkeypatch: pytest.MonkeyPatch, socket_path: str +) -> None: + snapshot = SandboxSettings( + sandbox_socket=socket_path, + sandbox_verify=False, + sandbox_token=None, + sandbox_callback_wait_seconds=15, + sandbox_max_concurrent_sessions=1, + sandbox_max_session_runtime_seconds=30, + sandbox_max_session_memory_bytes=256 * 1024 * 1024, + ) + calls = 0 + + def settings_factory() -> SandboxSettings: + nonlocal calls + calls += 1 + if calls > 1: + raise AssertionError("sandbox server reloaded settings") + return snapshot + + monkeypatch.setattr(sandbox, "SandboxSettings", settings_factory) + server_task = asyncio.create_task(run_sandbox_server()) + for _ in range(200): + if Path(socket_path).exists(): + break + await asyncio.sleep(0.01) + try: + await _run_client(socket_path, "value = 1", token="client-token") + finally: + server_task.cancel() + await asyncio.gather(server_task, return_exceptions=True) + + assert calls == 1 + + @pytest.mark.asyncio async def test_sandbox_server_rejects_invalid_client_token( monkeypatch, socket_path: str diff --git a/hyperforge/tests/harness_sdk/test_codemode.py b/hyperforge/tests/harness_sdk/test_codemode.py index f44699ed..cc9e6050 100644 --- a/hyperforge/tests/harness_sdk/test_codemode.py +++ b/hyperforge/tests/harness_sdk/test_codemode.py @@ -1,19 +1,28 @@ +import asyncio +import logging import sys +import tempfile +import tracemalloc from collections.abc import AsyncIterator from contextvars import ContextVar +from pathlib import Path from typing import Any import pytest -from pydantic import BaseModel, Field, ValidationError +from pydantic import BaseModel, Field from hyperforge.codemode import RestrictedPythonTask, WorkerExecutionRequest +from hyperforge.codemode import sandbox as sandbox_module from hyperforge.harness_sdk import ( AgentHarness, CodeModeCapability, CodeModeDispatch, + CodeModeExecutionLimiter, CodemodeInput, + CodeModeLimits, HarnessContextReference, HarnessContextType, + HarnessEventType, HarnessTool, HarnessToolCall, ModelDelta, @@ -28,6 +37,10 @@ from hyperforge.harness_sdk.tools import codemode as codemode_module +def _ctx(harness: Any) -> ToolCallContext: + return ToolCallContext(harness=harness, name="test") + + class UpperInput(BaseModel): value: str @@ -77,6 +90,24 @@ async def run( self.results.append((task.function, await dispatch(task))) +class CatchingRunner: + def __init__(self, tasks: list[RestrictedPythonTask]) -> None: + self.tasks = tasks + self.errors: list[str] = [] + self.results: list[tuple[str, Any]] = [] + + async def run( + self, _request: WorkerExecutionRequest, dispatch: CodeModeDispatch + ) -> None: + for task in self.tasks: + try: + result = await dispatch(task) + except Exception as exc: + self.errors.append(str(exc)) + else: + self.results.append((task.function, result)) + + def worker_task(function: str, *args: Any, **kwargs: Any) -> RestrictedPythonTask: return RestrictedPythonTask( function=function, @@ -86,6 +117,31 @@ def worker_task(function: str, *args: Any, **kwargs: Any) -> RestrictedPythonTas ) +def local_codemode( + *capabilities: CodeModeCapability[Any], + limits: CodeModeLimits | None = None, + execution_limiter: CodeModeExecutionLimiter | None = None, +) -> HarnessTool[Any, Any]: + if execution_limiter is None: + return create_codemode_tool( + capabilities=capabilities, + limits=limits or CodeModeLimits(), + remote_required=False, + ) + return create_codemode_tool( + capabilities=capabilities, + limits=limits or CodeModeLimits(), + execution_limiter=execution_limiter, + remote_required=False, + ) + + +def value_adapter( + _tool: HarnessTool[Any, UpperOutput], output: UpperOutput +) -> dict[str, str]: + return {"value": UpperOutput.model_validate(output).value} + + @pytest.mark.asyncio async def test_codemode_calls_registered_tools_and_returns_output() -> None: harness = AgentHarness( @@ -298,14 +354,18 @@ async def lower(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: } assert runner.requests[0].code == "hidden_upper('hello'); output('done')" assert runner.requests[0].question == "Why?" + assert runner.requests[0].redact_errors is True assert runner.results[0] == ("hidden_upper", {"value": "HELLO"}) @pytest.mark.asyncio -async def test_scoped_codemode_runs_hidden_capability_in_isolated_process( +async def test_remote_required_false_forces_isolated_process_with_ambient_remote( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setattr(codemode_module.settings, "sandbox_socket", None) + monkeypatch.setattr( + codemode_module.settings, "sandbox_socket", "/ambient-sandbox.sock" + ) + monkeypatch.setattr(codemode_module.settings, "sandbox_token", "ambient-token") @tool() async def lower(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: @@ -317,7 +377,8 @@ async def lower(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: upper, result_adapter=raw_codemode_result_adapter, ), - ) + ), + remote_required=False, ) harness = AgentHarness( model="test", @@ -334,7 +395,7 @@ async def lower(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: assert result.value == "HELLO" assert upper.name not in harness._tools - with pytest.raises(RuntimeError, match="lower.*not defined"): + with pytest.raises(RuntimeError, match="Generated code execution failed"): await scoped.execute( ToolCallContext(harness=harness, name=scoped.name), {"code": "lower(value='HELLO')"}, @@ -344,7 +405,7 @@ async def lower(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: @pytest.mark.asyncio async def test_scoped_codemode_default_projection_uses_formatted_context() -> None: runner = RecordingRunner( - [worker_task("lookup", value="hello"), worker_task("output", "done")] + [worker_task("lookup", value="hello"), worker_task("output", None)] ) scoped = create_codemode_tool( capabilities=(CodeModeCapability(lookup),), @@ -363,7 +424,7 @@ async def test_scoped_codemode_default_projection_uses_formatted_context() -> No @pytest.mark.asyncio async def test_scoped_codemode_raw_projection_must_be_explicit() -> None: runner = RecordingRunner( - [worker_task("lookup", value="hello"), worker_task("output", "done")] + [worker_task("lookup", value="hello"), worker_task("output", None)] ) scoped = create_codemode_tool( capabilities=( @@ -394,20 +455,37 @@ def wrap_model( ) -> dict[str, BaseModel]: return {"wrapped": output} - runner = RecordingRunner([worker_task("upper", value="hello")]) + runner = RecordingRunner( + [worker_task("upper", value="hello"), worker_task("output", None)] + ) scoped = create_codemode_tool( capabilities=(CodeModeCapability(upper, result_adapter=wrap_model),), runner=runner, ) harness = AgentHarness(model="test", model_client=UnusedModel()) - with pytest.raises(TypeError, match="must not contain Pydantic model values"): + with pytest.raises(RuntimeError, match="Code Mode capability 'upper' failed"): await scoped.execute( ToolCallContext(harness=harness, name=scoped.name), {"code": "upper(value='hello')"}, ) +@pytest.mark.asyncio +async def test_scoped_codemode_rejects_nested_models_in_output() -> None: + runner = RecordingRunner( + [worker_task("output", {"items": [UpperOutput(value="unsafe")]})] + ) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(TypeError, match="must not contain Pydantic model values"): + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "output({'items': []})"}, + ) + + @pytest.mark.asyncio async def test_scoped_codemode_normalizes_adapter_results_as_strict_json() -> None: def project( @@ -416,7 +494,7 @@ def project( return {"value": output.value, "positions": (1, 2)} runner = RecordingRunner( - [worker_task("upper", value="hello"), worker_task("output", "done")] + [worker_task("upper", value="hello"), worker_task("output", None)] ) scoped = create_codemode_tool( capabilities=(CodeModeCapability(upper, result_adapter=project),), @@ -449,7 +527,7 @@ def project( ) harness = AgentHarness(model="test", model_client=UnusedModel()) - with pytest.raises(ValueError, match="reserved '__model__' key"): + with pytest.raises(RuntimeError, match="Code Mode capability 'upper' failed"): await scoped.execute( ToolCallContext(harness=harness, name=scoped.name), {"code": "upper(value='hello')"}, @@ -470,7 +548,7 @@ def project( ) harness = AgentHarness(model="test", model_client=UnusedModel()) - with pytest.raises(ValueError, match="strict JSON value"): + with pytest.raises(RuntimeError, match="Code Mode capability 'upper' failed"): await scoped.execute( ToolCallContext(harness=harness, name=scoped.name), {"code": "upper(value='hello')"}, @@ -547,7 +625,7 @@ async def test_scoped_codemode_uses_normal_tool_validation() -> None: ) harness = AgentHarness(model="test", model_client=UnusedModel()) - with pytest.raises(ValueError, match="Invalid upper arguments"): + with pytest.raises(RuntimeError, match=r"capability 'upper' failed \(ValueError\)"): await input_scoped.execute( ToolCallContext(harness=harness, name=input_scoped.name), {"code": "upper(missing='value')"}, @@ -566,7 +644,9 @@ async def invalid_output( capabilities=(CodeModeCapability(invalid_output),), runner=invalid_output_runner, ) - with pytest.raises(ValidationError): + with pytest.raises( + RuntimeError, match=r"capability 'invalid_output' failed \(ValidationError\)" + ): await output_scoped.execute( ToolCallContext(harness=harness, name=output_scoped.name), {"code": "invalid_output(value='hello')"}, @@ -585,7 +665,7 @@ async def aliased( return UpperOutput(value=input_value.value.upper()) runner = RecordingRunner( - [worker_task("aliased", "hello"), worker_task("output", "done")] + [worker_task("aliased", "hello"), worker_task("output", None)] ) scoped = create_codemode_tool( capabilities=( @@ -618,7 +698,7 @@ async def test_scoped_codemode_counts_nested_calls_as_ordinary_usage() -> None: [ worker_task("upper", value="one"), worker_task("upper", value="two"), - worker_task("output", "done"), + worker_task("output", None), ] ) scoped = create_codemode_tool( @@ -662,7 +742,7 @@ async def record_id( [ worker_task("record_id", value="one"), worker_task("record_id", value="two"), - worker_task("output", "done"), + worker_task("output", None), ] ) scoped = create_codemode_tool( @@ -682,7 +762,24 @@ async def record_id( assert "outer-call" not in call_ids -@pytest.mark.parametrize("name", ["codemode", "output", "question", "agent_id"]) +@pytest.mark.parametrize( + "name", + [ + "codemode", + "output", + "print", + "printed", + "save", + "question", + "agent_id", + "dataclass", + "Chunk", + "Context", + "List", + "Any", + "Dict", + ], +) def test_scoped_codemode_rejects_reserved_capability_names(name: str) -> None: reserved = HarnessTool(name, upper.handler) @@ -760,3 +857,1463 @@ def test_scoped_codemode_description_includes_capability_schema() -> None: assert '"required":["value"]' in scoped.description assert "Custom orchestration instructions." in custom.description assert '"required":["value"]' in custom.description + + +@pytest.mark.asyncio +async def test_scoped_codemode_default_projection_uses_model_context() -> None: + @tool( + context_factory=lambda output: HarnessContextReference( + type=HarnessContextType.STRUCTURED, + content={"value": output.value}, + ) + ) + async def lookup(_: AgentHarness, input_value: UpperInput) -> SensitiveOutput: + return SensitiveOutput(value=input_value.value.upper(), secret="internal") + + scoped = local_codemode(CodeModeCapability(lookup)) + harness = AgentHarness(model="test", model_client=UnusedModel(), tools=[scoped]) + + result = await scoped.execute( + _ctx(harness), + {"code": "result = lookup(value='hello')\noutput(result)"}, + ) + + assert result.value == '{"value":"HELLO"}' + + +@pytest.mark.asyncio +async def test_scoped_codemode_sanitizes_formatted_context_events() -> None: + @tool() + async def lookup(_: AgentHarness, input_value: UpperInput) -> SensitiveOutput: + return SensitiveOutput(value=input_value.value, secret="internal-secret") + + scoped = local_codemode(CodeModeCapability(lookup)) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute(_ctx(harness), {"code": "lookup(value='safe'); output(None)"}) + + events = [event async for event in harness.history()] + completed = next( + event + for event in events + if event.type == HarnessEventType.TOOL_COMPLETED and event.payload.get("nested") + ) + assert completed.payload["result"]["type"] == "string" + assert completed.payload["result"]["utf8_bytes"] > 0 + assert "internal-secret" not in str(completed.payload) + + +@pytest.mark.asyncio +async def test_scoped_codemode_does_not_persist_opaque_string_results() -> None: + def project(_tool: HarnessTool, _output: BaseModel) -> str: + return "authorization: Bearer projected-secret" + + scoped = local_codemode(CodeModeCapability(upper, result_adapter=project)) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute(_ctx(harness), {"code": "upper(value='safe'); output(None)"}) + + events = [event async for event in harness.history()] + completed = next( + event + for event in events + if event.type == HarnessEventType.TOOL_COMPLETED and event.payload.get("nested") + ) + assert completed.payload["result"] == {"type": "string", "utf8_bytes": 38} + assert "projected-secret" not in str(completed.payload) + + +@pytest.mark.asyncio +async def test_scoped_codemode_does_not_persist_nested_opaque_strings() -> None: + def project(_tool: HarnessTool, _output: BaseModel) -> dict[str, object]: + return { + "value": "authorization: Bearer nested-secret", + "items": ["another-secret"], + } + + scoped = local_codemode(CodeModeCapability(upper, result_adapter=project)) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute(_ctx(harness), {"code": "upper(value='safe'); output(None)"}) + + events = [event async for event in harness.history()] + completed = next( + event + for event in events + if event.type == HarnessEventType.TOOL_COMPLETED and event.payload.get("nested") + ) + assert completed.payload["result"] == { + "type": "object", + "entries": [ + { + "key": {"type": "string", "utf8_bytes": 5}, + "value": {"type": "string", "utf8_bytes": 35}, + }, + { + "key": {"type": "string", "utf8_bytes": 5}, + "value": { + "type": "array", + "items": [{"type": "string", "utf8_bytes": 14}], + }, + }, + ], + } + assert "secret" not in str(completed.payload) + + +@pytest.mark.asyncio +async def test_scoped_codemode_treats_json_looking_strings_as_opaque() -> None: + def project(_tool: HarnessTool, _output: BaseModel) -> str: + return '{"score":1e9999}' + + scoped = local_codemode(CodeModeCapability(upper, result_adapter=project)) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute(_ctx(harness), {"code": "upper(value='safe'); output(None)"}) + + events = [event async for event in harness.history()] + completed = next( + event + for event in events + if event.type == HarnessEventType.TOOL_COMPLETED and event.payload.get("nested") + ) + assert completed.payload["result"] == { + "type": "string", + "utf8_bytes": len('{"score":1e9999}'.encode()), + } + + +def test_event_sanitizer_bounds_large_list() -> None: + sanitized = codemode_module._sanitize_event_value(list(range(100_000))) + + assert sanitized["type"] == "array" + assert len(sanitized["items"]) == codemode_module._EVENT_SANITIZE_MAX_ITEMS + assert sanitized["truncated"] == { + "reason": "item_limit", + "omitted_items": 100_000 - codemode_module._EVENT_SANITIZE_MAX_ITEMS, + } + + +def test_event_sanitizer_bounds_large_dictionary() -> None: + sanitized = codemode_module._sanitize_event_value( + {f"untrusted-key-{index}": index for index in range(100_000)} + ) + + assert sanitized["type"] == "object" + assert len(sanitized["entries"]) == codemode_module._EVENT_SANITIZE_MAX_ITEMS + assert sanitized["truncated"] == { + "reason": "item_limit", + "omitted_entries": 100_000 - codemode_module._EVENT_SANITIZE_MAX_ITEMS, + } + assert "untrusted-key" not in str(sanitized) + + +def test_event_sanitizer_enforces_global_node_budget() -> None: + sanitized = codemode_module._sanitize_event_value( + [[0] * codemode_module._EVENT_SANITIZE_MAX_ITEMS for _ in range(10_000)] + ) + + assert sanitized["truncated"]["reason"] == "node_budget" + assert len(sanitized["items"]) < codemode_module._EVENT_SANITIZE_MAX_ITEMS + assert sanitized["items"][-1]["truncated"]["reason"] == "node_budget" + + +def test_event_sanitizer_enforces_depth_limit() -> None: + value: Any = 0 + for _ in range(100): + value = [value] + + sanitized = codemode_module._sanitize_event_value(value) + for _ in range(codemode_module._EVENT_SANITIZE_MAX_DEPTH): + sanitized = sanitized["items"][0] + + assert sanitized == { + "type": "array", + "items": [], + "truncated": {"reason": "depth_limit", "omitted_items": 1}, + } + + +def test_scoped_worker_encoding_stops_before_traversing_wide_value() -> None: + class TrackingList(list[int]): + visited = 0 + + def __iter__(self): + for item in super().__iter__(): + self.visited += 1 + if self.visited >= 100: + pytest.fail("strict worker encoding traversed the full container") + yield item + + value = TrackingList(range(100_000)) + + with pytest.raises(ValueError, match="exceeds maximum size"): + codemode_module._normalize_worker_value(value, 32, "Wide value") + + assert value.visited < 100 + + +def test_scoped_worker_encoding_stops_before_traversing_wide_dictionary() -> None: + class TrackingDict(dict[str, int]): + visited = 0 + + def items(self): + for item in super().items(): + self.visited += 1 + if self.visited >= 100: + pytest.fail("strict worker encoding traversed the full dictionary") + yield item + + value = TrackingDict({str(index): index for index in range(100_000)}) + + with pytest.raises(ValueError, match="exceeds maximum size"): + codemode_module._normalize_worker_value(value, 32, "Wide value") + + assert value.visited < 100 + + +def test_scoped_source_sizing_stops_after_crossing_limit() -> None: + class GuardedSource(str): + def encode(self, *_args, **_kwargs): + pytest.fail("source sizing encoded the complete source") + + source = GuardedSource("é" * 1_000_000) + tracemalloc.start() + try: + assert codemode_module._utf8_exceeds_limit(source, 32) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert peak < 100_000 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "projection", + [ + lambda _tool, _output: object(), + lambda _tool, output: output, + lambda _tool, _output: {"__model__": "Context", "value": "unsafe"}, + ], +) +async def test_scoped_codemode_rejects_unsafe_projected_results( + projection: Any, +) -> None: + scoped = local_codemode(CodeModeCapability(upper, result_adapter=projection)) + harness = AgentHarness(model="test", model_client=UnusedModel(), tools=[scoped]) + + with pytest.raises(RuntimeError, match="Generated code execution failed"): + await scoped.execute(_ctx(harness), {"code": "upper(value='hello')"}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_enforces_source_result_and_output_byte_limits() -> None: + harness = AgentHarness(model="test", model_client=UnusedModel()) + source_limited = local_codemode(limits=CodeModeLimits(max_source_bytes=1)) + with pytest.raises(ValueError, match="source exceeds maximum size"): + await source_limited.execute(_ctx(harness), {"code": "é"}) + + result_limited = local_codemode( + CodeModeCapability(upper, result_adapter=value_adapter), + limits=CodeModeLimits(max_result_bytes=4), + ) + with pytest.raises(RuntimeError, match="Generated code execution failed"): + await result_limited.execute(_ctx(harness), {"code": "upper(value='hello')"}) + + output_limited = local_codemode(limits=CodeModeLimits(max_output_bytes=4)) + with pytest.raises(RuntimeError, match="Generated code execution failed"): + await output_limited.execute(_ctx(harness), {"code": "output('hello')"}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_result_policy_failure_remains_terminal_when_caught() -> ( + None +): + calls: list[str] = [] + + @tool() + async def count(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + calls.append(input_value.value) + return UpperOutput(value=input_value.value) + + scoped = local_codemode( + CodeModeCapability(count, result_adapter=value_adapter), + limits=CodeModeLimits(max_result_bytes=8), + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + code = """ +try: + count(value='one') +except Exception: + pass +try: + count(value='retry') +except Exception: + pass +try: + output('apparently-successful') +except Exception: + pass +""" + + with pytest.raises(ValueError, match="result from count exceeds maximum size"): + await scoped.execute(_ctx(harness), {"code": code}) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert calls == ["one"] + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + ] + + +@pytest.mark.asyncio +async def test_scoped_codemode_preflights_isolated_result_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + @tool() + async def transport_result( + _: ToolCallContext, input_value: UpperInput + ) -> UpperOutput: + calls.append(input_value.value) + return UpperOutput(value=input_value.value) + + scoped = local_codemode( + CodeModeCapability( + transport_result, + result_adapter=lambda _tool, _output: "x" * 127, + ), + limits=CodeModeLimits(max_result_bytes=256), + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + monkeypatch.setattr(codemode_module, "MAX_PACKET_BYTES", 128) + code = """ +try: + transport_result(value='one') +except Exception: + pass +try: + transport_result(value='retry') +except Exception: + pass +try: + output('apparently-successful') +except Exception: + pass +""" + + with pytest.raises(ValueError, match="result transport.*exceeds maximum size"): + await scoped.execute(_ctx(harness), {"code": code}) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert calls == ["one"] + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + ] + assert all(event.type != HarnessEventType.TOOL_COMPLETED for event in nested) + + +@pytest.mark.asyncio +async def test_scoped_codemode_preflights_remote_result_envelope( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + + @tool() + async def transport_result( + _: ToolCallContext, input_value: UpperInput + ) -> UpperOutput: + calls.append(input_value.value) + return UpperOutput(value=input_value.value) + + runner = CatchingRunner( + [ + worker_task("transport_result", value="one"), + worker_task("transport_result", value="retry"), + worker_task("output", "apparently-successful"), + ] + ) + scoped = create_codemode_tool( + capabilities=( + CodeModeCapability( + transport_result, + result_adapter=lambda _tool, _output: "x" * 110, + ), + ), + limits=CodeModeLimits(max_result_bytes=128), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + monkeypatch.setattr(codemode_module, "MAX_PACKET_BYTES", 128) + + with pytest.raises(ValueError, match="result transport.*exceeds maximum size"): + await scoped.execute(_ctx(harness), {"code": "unused"}) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert calls == ["one"] + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + ] + assert all(event.type != HarnessEventType.TOOL_COMPLETED for event in nested) + + +@pytest.mark.asyncio +async def test_scoped_codemode_counts_nested_calls_once_and_enforces_local_limit() -> ( + None +): + scoped = local_codemode( + CodeModeCapability(upper, result_adapter=value_adapter), + limits=CodeModeLimits(max_nested_calls=1), + ) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + tools=[scoped], + usage_limits=UsageLimits(max_tool_calls=10), + ) + + await harness._execute_tool_calls( + [ + HarnessToolCall( + id="outer-call", + name="codemode", + arguments={"code": "upper(value='one')\nupper(value='two')"}, + ) + ] + ) + + assert harness.usage.tool_calls == 2 + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_COMPLETED, + ] + + +@pytest.mark.asyncio +async def test_scoped_codemode_nested_limit_remains_terminal_when_caught() -> None: + calls: list[str] = [] + + @tool() + async def count(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + calls.append(input_value.value) + return UpperOutput(value=input_value.value) + + scoped = local_codemode( + CodeModeCapability(count, result_adapter=value_adapter), + limits=CodeModeLimits(max_nested_calls=1), + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + code = """ +count(value='one') +try: + count(value='two') +except Exception: + pass +try: + count(value='three') +except Exception: + pass +try: + output('apparently-successful') +except Exception: + pass +""" + + with pytest.raises(RuntimeError, match="nested call limit exceeded: 2 > 1"): + await scoped.execute(_ctx(harness), {"code": code}) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert calls == ["one"] + assert harness.usage.tool_calls == 1 + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_COMPLETED, + ] + + +@pytest.mark.asyncio +async def test_scoped_codemode_global_tool_limit_remains_terminal_when_caught() -> None: + calls: list[str] = [] + + @tool() + async def count(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + calls.append(input_value.value) + return UpperOutput(value=input_value.value) + + scoped = local_codemode(CodeModeCapability(count, result_adapter=value_adapter)) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + usage_limits=UsageLimits(max_tool_calls=1), + ) + code = """ +count(value='one') +try: + count(value='two') +except Exception: + pass +try: + count(value='three') +except Exception: + pass +try: + output('apparently-successful') +except Exception: + pass +""" + + with pytest.raises(RuntimeError, match="max_tool_calls limit exceeded: 2 > 1"): + await scoped.execute(_ctx(harness), {"code": code}) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert calls == ["one"] + assert harness.usage.tool_calls == 2 + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_COMPLETED, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure_point", + [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_COMPLETED, + HarnessEventType.TOOL_FAILED, + ], +) +async def test_scoped_codemode_event_failures_are_redacted_and_terminal( + monkeypatch: pytest.MonkeyPatch, failure_point: HarnessEventType +) -> None: + backend_secret = f"{failure_point.value}-backend-secret" + fixed_failure = "Code Mode capability failed" + calls: list[str] = [] + + @tool() + async def event_target(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + calls.append(input_value.value) + if failure_point == HarnessEventType.TOOL_FAILED: + raise RuntimeError("handler-secret") + return UpperOutput(value=input_value.value) + + runner = CatchingRunner( + [ + worker_task("event_target", value="one"), + worker_task("event_target", value="two"), + worker_task("output", "apparently-successful"), + ] + ) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(event_target, result_adapter=value_adapter),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + original_emit = harness.emit + + async def failing_emit( + event_type: HarnessEventType, payload: dict[str, Any], **kwargs: Any + ): + if event_type == failure_point: + raise RuntimeError(backend_secret) + return await original_emit(event_type, payload, **kwargs) + + monkeypatch.setattr(harness, "emit", failing_emit) + + with pytest.raises(RuntimeError, match=f"^{fixed_failure}$") as exc: + await scoped.execute(_ctx(harness), {"code": "unused"}) + + events = [event async for event in harness.history()] + assert runner.errors == [fixed_failure] * 3 + assert calls == ( + [] if failure_point == HarnessEventType.TOOL_REQUESTED else ["one"] + ) + assert [event.type for event in events] == ( + [] + if failure_point == HarnessEventType.TOOL_REQUESTED + else [HarnessEventType.TOOL_REQUESTED] + ) + exposed = str([exc.value, runner.errors, events]) + assert backend_secret not in exposed + assert "handler-secret" not in exposed + + +@pytest.mark.asyncio +async def test_scoped_codemode_tool_authored_event_failure_is_terminal( + monkeypatch: pytest.MonkeyPatch, +) -> None: + backend_secret = "tool-authored-event-backend-secret" + calls: list[str] = [] + + @tool() + async def event_target( + context: ToolCallContext, input_value: UpperInput + ) -> UpperOutput: + calls.append(input_value.value) + try: + await context.emit(HarnessEventType.MESSAGE_ADDED, {"authored": True}) + except RuntimeError: + pass + return UpperOutput(value=input_value.value) + + runner = CatchingRunner( + [ + worker_task("event_target", value="one"), + worker_task("event_target", value="two"), + worker_task("output", "apparently-successful"), + ] + ) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(event_target, result_adapter=value_adapter),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + original_emit = harness.emit + + async def failing_emit( + event_type: HarnessEventType, payload: dict[str, Any], **kwargs: Any + ): + if event_type == HarnessEventType.MESSAGE_ADDED: + raise RuntimeError(backend_secret) + return await original_emit(event_type, payload, **kwargs) + + monkeypatch.setattr(harness, "emit", failing_emit) + + with pytest.raises(RuntimeError, match="^Code Mode capability failed$") as exc: + await scoped.execute(_ctx(harness), {"code": "unused"}) + + events = [event async for event in harness.history()] + assert calls == ["one"] + assert runner.errors == ["Code Mode capability failed"] * 3 + assert [event.type for event in events] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + ] + assert backend_secret not in str([exc.value, runner.errors, events]) + + +@pytest.mark.asyncio +async def test_scoped_codemode_can_recover_from_ordinary_handler_failure() -> None: + calls: list[str] = [] + + @tool() + async def recoverable(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + calls.append(input_value.value) + if input_value.value == "fail": + raise RuntimeError("handler-secret") + return UpperOutput(value=input_value.value) + + runner = CatchingRunner( + [ + worker_task("recoverable", value="fail"), + worker_task("recoverable", value="ok"), + worker_task("output", "done"), + ] + ) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(recoverable, result_adapter=value_adapter),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + result = await scoped.execute(_ctx(harness), {"code": "unused"}) + + events = [event async for event in harness.history()] + assert result.value == "done" + assert calls == ["fail", "ok"] + assert runner.errors == ["Code Mode capability 'recoverable' failed (RuntimeError)"] + assert [event.type for event in events] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_COMPLETED, + ] + assert "handler-secret" not in str([runner.errors, events]) + + +@pytest.mark.asyncio +async def test_scoped_codemode_unknown_capability_is_terminal() -> None: + calls: list[str] = [] + + @tool() + async def allowed(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + calls.append(input_value.value) + return UpperOutput(value=input_value.value) + + runner = CatchingRunner( + [ + worker_task("unknown"), + worker_task("allowed", value="retry"), + worker_task("output", "apparently-successful"), + ] + ) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(allowed, result_adapter=value_adapter),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match="Unknown Code Mode capability: unknown"): + await scoped.execute(_ctx(harness), {"code": "unused"}) + + assert calls == [] + assert runner.errors == ["Unknown Code Mode capability: unknown"] * 3 + assert [event async for event in harness.history()] == [] + + +@pytest.mark.asyncio +async def test_scoped_codemode_emits_sanitized_nested_events_with_parent_id() -> None: + @tool() + async def inspect_value( + _: AgentHarness, input_value: UpperInput + ) -> SensitiveOutput: + return SensitiveOutput(value=input_value.value, secret="result-secret") + + def project(_tool: HarnessTool, output: BaseModel) -> dict[str, str]: + value = SensitiveOutput.model_validate(output) + return {"value": value.value, "token": "projected-secret"} + + scoped = local_codemode(CodeModeCapability(inspect_value, result_adapter=project)) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + tools=[scoped], + execution_context={"actor_id": "actor-1", "tenant_id": "tenant-1"}, + ) + + await harness._execute_tool_calls( + [ + HarnessToolCall( + id="outer-call", + name="codemode", + arguments={ + "code": "inspect_value(value='safe', token='argument-secret')" + }, + ) + ] + ) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_COMPLETED, + ] + nested_id = nested[0].payload["call"]["id"] + assert nested_id == nested[1].payload["call_id"] + assert all(event.parent_call_id == "outer-call" for event in nested) + assert all(event.payload["parent_call_id"] == "outer-call" for event in nested) + assert nested[0].payload["call"]["arguments"] == { + "type": "object", + "entries": [ + { + "key": {"type": "string", "utf8_bytes": 5}, + "value": {"type": "string", "utf8_bytes": 4}, + }, + { + "key": {"type": "string", "utf8_bytes": 5}, + "value": "[REDACTED]", + }, + ], + } + assert nested[1].payload["result"] == { + "type": "object", + "entries": [ + { + "key": {"type": "string", "utf8_bytes": 5}, + "value": {"type": "string", "utf8_bytes": 4}, + }, + { + "key": {"type": "string", "utf8_bytes": 5}, + "value": "[REDACTED]", + }, + ], + } + assert nested[0].payload["execution_context"] == { + "actor_id": "actor-1", + "tenant_id": "tenant-1", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("marker", "expected"), + [ + ( + "x" * (codemode_module._EVENT_CONTEXT_MAX_STRING_BYTES + 1), + "exceeds maximum size", + ), + (float("nan"), "must contain finite numbers"), + (float("inf"), "must contain finite numbers"), + (float("-inf"), "must contain finite numbers"), + ( + 1 << (codemode_module._EVENT_CONTEXT_MAX_INTEGER_BYTES * 4 + 1), + "exceeds maximum size", + ), + ( + -(1 << (codemode_module._EVENT_CONTEXT_MAX_INTEGER_BYTES * 4 + 1)), + "exceeds maximum size", + ), + ], +) +async def test_scoped_codemode_rejects_unsafe_event_context_before_execution( + marker: Any, expected: str +) -> None: + runner = RecordingRunner([worker_task("output", "not-run")]) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + execution_context={"actor_id": marker}, + ) + + with pytest.raises(ValueError, match=expected): + await scoped.execute(_ctx(harness), {"code": "unused"}) + + assert runner.requests == [] + assert [event async for event in harness.history()] == [] + + +@pytest.mark.asyncio +async def test_scoped_codemode_enforces_total_event_context_budget() -> None: + runner = RecordingRunner([worker_task("output", "not-run")]) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + execution_context={ + "actor": "x" * codemode_module._EVENT_CONTEXT_MAX_STRING_BYTES, + "actor_id": "x" * codemode_module._EVENT_CONTEXT_MAX_STRING_BYTES, + "tenant": "x" * codemode_module._EVENT_CONTEXT_MAX_STRING_BYTES, + "tenant_id": "x" * codemode_module._EVENT_CONTEXT_MAX_STRING_BYTES, + "user_id": 1, + }, + ) + + with pytest.raises(ValueError, match="exceed maximum total size"): + await scoped.execute(_ctx(harness), {"code": "unused"}) + + assert runner.requests == [] + assert [event async for event in harness.history()] == [] + + +def test_scoped_codemode_event_context_keeps_booleans_distinct() -> None: + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + execution_context={"actor": True, "tenant": False}, + ) + + assert codemode_module._nested_event_context(harness) == { + "execution_context": {"actor": True, "tenant": False} + } + + +@pytest.mark.asyncio +async def test_scoped_codemode_emits_projected_completion_event() -> None: + scoped = local_codemode(CodeModeCapability(upper, result_adapter=value_adapter)) + harness = AgentHarness(model="test", model_client=UnusedModel(), tools=[scoped]) + + await harness._execute_tool_calls( + [ + HarnessToolCall( + id="outer-call", + name="codemode", + arguments={"code": "upper(value='safe'); output(None)"}, + ) + ] + ) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_COMPLETED, + ] + assert nested[1].payload["result"] == { + "type": "object", + "entries": [ + { + "key": {"type": "string", "utf8_bytes": 5}, + "value": {"type": "string", "utf8_bytes": 4}, + } + ], + } + assert nested[1].payload["result_bytes"] == len('{"value":"SAFE"}') + assert nested[1].payload["duration_ms"] >= 0 + assert nested[0].payload["call"]["id"] == nested[1].payload["call_id"] + + +@pytest.mark.asyncio +async def test_scoped_codemode_never_persists_raw_event_dictionary_keys() -> None: + class MappingInput(BaseModel): + payload: dict[str, str] + + @tool() + async def inspect_mapping( + _: AgentHarness, _input_value: MappingInput + ) -> UpperOutput: + return UpperOutput(value="safe") + + def project(_tool: HarnessTool, _output: BaseModel) -> dict[str, str]: + return { + "result-secret-in-key": "visible", + "password": "sensitive-result-value", + } + + scoped = local_codemode(CodeModeCapability(inspect_mapping, result_adapter=project)) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute( + _ctx(harness), + { + "code": ( + "inspect_mapping(payload={'argument-secret-in-key': 'visible', " + "'token': 'sensitive-argument-value'}); output(None)" + ) + }, + ) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + requested_arguments = nested[0].payload["call"]["arguments"] + completed_result = nested[1].payload["result"] + persisted = str([requested_arguments, completed_result]) + + assert ( + requested_arguments["entries"][0]["value"]["entries"][1]["value"] + == "[REDACTED]" + ) + assert completed_result["entries"][1]["value"] == "[REDACTED]" + assert "argument-secret-in-key" not in persisted + assert "result-secret-in-key" not in persisted + assert "token" not in persisted + assert "password" not in persisted + assert "sensitive-argument-value" not in persisted + assert "sensitive-result-value" not in persisted + + +@pytest.mark.asyncio +async def test_scoped_codemode_records_nested_failure_before_propagating() -> None: + @tool() + async def fail(_: AgentHarness, _input_value: UpperInput) -> UpperOutput: + raise ValueError("secret failure detail") + + scoped = local_codemode(CodeModeCapability(fail)) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(RuntimeError, match="Generated code execution failed"): + await scoped.execute(_ctx(harness), {"code": "fail(value='safe')"}) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + ] + assert nested[1].payload["result"] == {"error": "ValueError"} + assert "secret failure detail" not in str(nested) + + +def _secret_capability(secret: str) -> CodeModeCapability[Any]: + @tool(name="secret_value") + async def secret_value( + _context: ToolCallContext, _input_value: UpperInput + ) -> UpperOutput: + return UpperOutput(value="internal") + + def project(_tool: HarnessTool, _output: BaseModel) -> str: + return secret + + return CodeModeCapability(secret_value, result_adapter=project) + + +async def _assert_reraised_secret_is_redacted( + scoped: HarnessTool[Any, Any], harness: AgentHarness, secret: str, caplog +) -> None: + message = await harness._execute_tool_call( + HarnessToolCall( + id="outer-call", + name=scoped.name, + arguments={ + "code": ( + "projected = secret_value(value='x')\nraise RuntimeError(projected)" + ) + }, + ) + ) + events = [event async for event in harness.history()] + persisted = str([event.payload for event in events]) + outer_failure = next( + event + for event in events + if event.type == HarnessEventType.TOOL_FAILED + and not event.payload.get("nested") + ) + + assert "Generated code execution failed" in message.content + assert "Generated code execution failed" in str(outer_failure.payload) + assert secret not in message.content + assert secret not in persisted + assert secret not in caplog.text + + +@pytest.mark.asyncio +async def test_isolated_scoped_codemode_redacts_reraised_projected_secret( + caplog, +) -> None: + secret = "isolated-projected-secret" + scoped = local_codemode(_secret_capability(secret)) + harness = AgentHarness(model="test", model_client=UnusedModel(), tools=[scoped]) + caplog.set_level(logging.INFO) + + await _assert_reraised_secret_is_redacted(scoped, harness, secret, caplog) + + +@pytest.mark.asyncio +async def test_remote_scoped_codemode_redacts_reraised_projected_secret( + monkeypatch: pytest.MonkeyPatch, caplog +) -> None: + secret = "remote-projected-secret" + caplog.set_level(logging.INFO) + with tempfile.TemporaryDirectory(prefix="sbx", dir="/tmp") as directory: + socket = str(Path(directory) / "sandbox.sock") + monkeypatch.setenv("SANDBOX_SOCKET", socket) + monkeypatch.setenv("SANDBOX_VERIFY", "false") + monkeypatch.setenv("SANDBOX_TOKEN", "server-token") + monkeypatch.setenv("SANDBOX_CALLBACK_WAIT_SECONDS", "15") + monkeypatch.setattr(codemode_module.settings, "sandbox_socket", socket) + monkeypatch.setattr(codemode_module.settings, "sandbox_token", "server-token") + server_task = asyncio.create_task(sandbox_module.run_sandbox_server()) + for _ in range(200): + if Path(socket).exists(): + break + await asyncio.sleep(0.01) + + scoped = create_codemode_tool(capabilities=(_secret_capability(secret),)) + harness = AgentHarness(model="test", model_client=UnusedModel(), tools=[scoped]) + try: + await _assert_reraised_secret_is_redacted(scoped, harness, secret, caplog) + assert any( + record.name == "hyperforge_codemode_sandbox" + and "Generated code execution failed" in record.getMessage() + for record in caplog.records + ) + finally: + server_task.cancel() + await asyncio.gather(server_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_scoped_codemode_cancels_pending_nested_callback_on_timeout() -> None: + started = asyncio.Event() + cancelled = asyncio.Event() + + @tool() + async def wait(_: AgentHarness, _input_value: UpperInput) -> UpperOutput: + started.set() + try: + await asyncio.Future() + finally: + cancelled.set() + + scoped = local_codemode(CodeModeCapability(wait)) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + usage_limits=UsageLimits(max_codemode_runtime_seconds=2), + ) + + with pytest.raises(RuntimeError, match="timed out"): + await scoped.execute(_ctx(harness), {"code": "wait(value='x')"}) + + assert started.is_set() + assert cancelled.is_set() + + +@pytest.mark.asyncio +async def test_scoped_codemode_propagates_external_cancellation() -> None: + started = asyncio.Event() + cancelled = asyncio.Event() + + @tool() + async def wait(_: AgentHarness, _input_value: UpperInput) -> UpperOutput: + started.set() + try: + await asyncio.Future() + finally: + cancelled.set() + + scoped = local_codemode(CodeModeCapability(wait)) + harness = AgentHarness(model="test", model_client=UnusedModel()) + execution = asyncio.create_task( + scoped.execute(_ctx(harness), {"code": "wait(value='x')"}) + ) + await started.wait() + + execution.cancel() + with pytest.raises(asyncio.CancelledError): + await execution + + assert cancelled.is_set() + + +@pytest.mark.asyncio +async def test_scoped_codemode_fails_closed_without_remote_socket( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(codemode_module.settings, "sandbox_socket", None) + scoped = create_codemode_tool(capabilities=()) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(RuntimeError, match="SANDBOX_SOCKET is absent"): + await scoped.execute(_ctx(harness), {"code": "output(None)"}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_fails_closed_without_remote_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(codemode_module.settings, "sandbox_socket", "/unused.sock") + monkeypatch.setattr(codemode_module.settings, "sandbox_token", None) + scoped = create_codemode_tool(capabilities=()) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(RuntimeError, match="SANDBOX_TOKEN is absent"): + await scoped.execute(_ctx(harness), {"code": "output(None)"}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_passes_captured_token_when_settings_drift( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, str | None] = {} + + class RemoteRunner: + def __init__(self, dispatch: CodeModeDispatch) -> None: + self.dispatch = dispatch + + async def run(self, _request: WorkerExecutionRequest) -> None: + monkeypatch.setattr( + codemode_module.settings, "sandbox_token", "replacement-token" + ) + await self.dispatch(worker_task("output", "ok")) + + def run_when_callbacks_complete(self, _callback: Any) -> bool: + return False + + def remote( + _socket: str, + dispatch: CodeModeDispatch, + _debug: bool = False, + *, + token: str | None = None, + ) -> RemoteRunner: + captured["token"] = token + return RemoteRunner(dispatch) + + monkeypatch.setattr(codemode_module.settings, "sandbox_socket", "/sandbox.sock") + monkeypatch.setattr(codemode_module.settings, "sandbox_token", "checked-token") + monkeypatch.setattr(codemode_module.SandboxRunner, "remote", remote) + scoped = create_codemode_tool(capabilities=()) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + result = await scoped.execute(_ctx(harness), {"code": "output('ok')"}) + + assert result.value == "ok" + assert captured == {"token": "checked-token"} + assert codemode_module.settings.sandbox_token == "replacement-token" + + +@pytest.mark.asyncio +async def test_scoped_codemode_uses_injected_runner_without_sandbox( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(codemode_module.settings, "sandbox_socket", None) + monkeypatch.setattr(codemode_module.settings, "sandbox_token", None) + runner = RecordingRunner( + [ + worker_task("upper", "hello"), + worker_task("output", {"done": True}), + ] + ) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(upper, result_adapter=value_adapter),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + result = await scoped.execute( + _ctx(harness), + {"code": "upper('hello'); output({'done': True})", "question": "Q?"}, + ) + + assert result.value == {"done": True} + request = runner.requests[0] + assert request.code == "upper('hello'); output({'done': True})" + assert request.question == "Q?" + assert set(request.function_names["harness"]) == {"upper", "output"} + assert ("upper", {"value": "HELLO"}) in runner.results + assert harness.usage.tool_calls == 1 + + +@pytest.mark.asyncio +async def test_scoped_codemode_output_must_be_called_exactly_once() -> None: + runner = RecordingRunner( + [worker_task("output", 1), worker_task("output", 2)], + ) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match="may only be called once"): + await scoped.execute(_ctx(harness), {"code": "output(1); output(2)"}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_repeated_output_remains_terminal_when_caught() -> None: + scoped = local_codemode() + harness = AgentHarness(model="test", model_client=UnusedModel()) + code = """ +output(1) +try: + output(2) +except Exception: + pass +""" + + with pytest.raises(ValueError, match="may only be called once"): + await scoped.execute(_ctx(harness), {"code": code}) + + +@pytest.mark.asyncio +async def test_scoped_worker_output_serialization_failure_is_terminal() -> None: + scoped = local_codemode() + harness = AgentHarness(model="test", model_client=UnusedModel()) + code = """ +try: + output({1, 2}) +except Exception: + pass +output('apparently-successful') +""" + + with pytest.raises(RuntimeError, match="Generated code execution failed"): + await scoped.execute(_ctx(harness), {"code": code}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first_attempt", "expected"), + [ + ("output()", "output requires a value"), + ("output('too-large')", "exceeds maximum size"), + ( + "output({'__model__': 'attacker-controlled'})", + "reserved '__model__' key", + ), + ], +) +async def test_scoped_codemode_invalid_first_output_attempt_remains_terminal( + first_attempt: str, expected: str +) -> None: + scoped = local_codemode(limits=CodeModeLimits(max_output_bytes=8)) + harness = AgentHarness(model="test", model_client=UnusedModel()) + code = f""" +try: + {first_attempt} +except Exception: + pass +try: + output(None) +except Exception: + pass +""" + + with pytest.raises((TypeError, ValueError), match=expected): + await scoped.execute(_ctx(harness), {"code": code}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_output_requires_a_value() -> None: + runner = RecordingRunner([worker_task("output")]) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match="output requires a value"): + await scoped.execute(_ctx(harness), {"code": "output()"}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_requires_output_call() -> None: + runner = RecordingRunner([]) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match="must call output\\(value\\) exactly once"): + await scoped.execute(_ctx(harness), {"code": "1 + 1"}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_enforces_cumulative_result_limit() -> None: + calls: list[str] = [] + + @tool() + async def count(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + calls.append(input_value.value) + return UpperOutput(value=input_value.value) + + scoped = local_codemode( + CodeModeCapability(count, result_adapter=value_adapter), + limits=CodeModeLimits(max_result_bytes=20, max_cumulative_result_bytes=20), + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + code = """ +count(value='hi') +try: + count(value='ok') +except Exception: + pass +try: + count(value='retry') +except Exception: + pass +try: + output('apparently-successful') +except Exception: + pass +""" + + with pytest.raises(RuntimeError, match="cumulative result limit exceeded"): + await scoped.execute(_ctx(harness), {"code": code}) + + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + assert calls == ["hi", "ok"] + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_COMPLETED, + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + ] + assert nested[-1].payload["result"] == {"error": "RuntimeError"} + + +@pytest.mark.parametrize( + "kwargs", + [ + {"max_cumulative_result_bytes": 0}, + {"max_result_bytes": 1024, "max_cumulative_result_bytes": 512}, + ], +) +def test_scoped_codemode_rejects_inconsistent_result_limits(kwargs: dict): + with pytest.raises(ValueError, match="max_cumulative_result_bytes"): + CodeModeLimits(**kwargs) + + +@pytest.mark.parametrize("value", [0, -1, True, 1.5, float("nan")]) +def test_scoped_codemode_rejects_invalid_limits(value: Any) -> None: + with pytest.raises(ValueError, match="positive integer"): + CodeModeLimits(max_source_bytes=value) # type: ignore[arg-type] + + +def test_codemode_execution_limiter_rejects_invalid_limit() -> None: + with pytest.raises(ValueError, match="positive integer"): + CodeModeExecutionLimiter(0) + + +@pytest.mark.asyncio +async def test_scoped_codemode_concurrency_saturation() -> None: + started = asyncio.Event() + + @tool() + async def wait(_: AgentHarness, _input_value: UpperInput) -> UpperOutput: + started.set() + await asyncio.Future() + + scoped = local_codemode( + CodeModeCapability(wait), + execution_limiter=CodeModeExecutionLimiter(1), + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + first = asyncio.create_task( + scoped.execute(_ctx(harness), {"code": "wait(value='x')"}) + ) + await started.wait() + + with pytest.raises(RuntimeError, match="concurrency limit reached"): + await scoped.execute(_ctx(harness), {"code": "output(None)"}) + + first.cancel() + await asyncio.gather(first, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_scoped_codemode_holds_slot_for_non_cooperative_callback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started = asyncio.Event() + cancellation_seen = asyncio.Event() + release = asyncio.Event() + callback_done = asyncio.Event() + + @tool() + async def wait(context: ToolCallContext, _input_value: UpperInput) -> UpperOutput: + started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_seen.set() + await release.wait() + await context.emit(HarnessEventType.MESSAGE_ADDED, {"tool_authored": True}) + raise + finally: + callback_done.set() + + monkeypatch.setattr("hyperforge.codemode.sandbox.CALLBACK_CANCEL_TIMEOUT", 0.01) + limiter = CodeModeExecutionLimiter(1) + scoped = local_codemode(CodeModeCapability(wait), execution_limiter=limiter) + harness = AgentHarness(model="test", model_client=UnusedModel()) + harness._turn_id = "originating-turn" + execution = asyncio.create_task( + scoped.execute(_ctx(harness), {"code": "wait(value='x')"}) + ) + await started.wait() + + execution.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(execution, timeout=0.5) + await cancellation_seen.wait() + with pytest.raises(RuntimeError, match="concurrency limit reached"): + await scoped.execute(_ctx(harness), {"code": "output(None)"}) + + harness._turn_id = "later-turn" + release.set() + await asyncio.wait_for(callback_done.wait(), timeout=0.5) + await asyncio.sleep(0) + result = await scoped.execute(_ctx(harness), {"code": "output('released')"}) + events = [event async for event in harness.history()] + nested = [event for event in events if event.payload.get("nested")] + tool_authored = [event for event in events if event.payload.get("tool_authored")] + + assert result.value == "released" + assert [event.type for event in nested] == [ + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + ] + assert len(tool_authored) == 1 + assert tool_authored[0].parent_call_id == nested[0].payload["call"]["id"] + assert {event.turn_id for event in [*nested, *tool_authored]} == { + "originating-turn" + } diff --git a/hyperforge/tests/harness_sdk/test_harness.py b/hyperforge/tests/harness_sdk/test_harness.py index 80c0f621..e2e369ad 100644 --- a/hyperforge/tests/harness_sdk/test_harness.py +++ b/hyperforge/tests/harness_sdk/test_harness.py @@ -320,6 +320,20 @@ async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: assert events[-1].payload == {"text": "answer-2"} +@pytest.mark.asyncio +async def test_emit_distinguishes_explicit_no_turn_from_current_turn() -> None: + harness = AgentHarness(model="test-model", model_client=Model()) + harness._turn_id = "current-turn" + + current = await harness.emit(HarnessEventType.TURN_STARTED, {}, persist=False) + no_turn = await harness.emit( + HarnessEventType.TURN_STARTED, {}, persist=False, turn_id=None + ) + + assert current.turn_id == "current-turn" + assert no_turn.turn_id is None + + @pytest.mark.asyncio async def test_unconsumed_turn_does_not_buffer_live_events() -> None: harness = AgentHarness(model="test-model", model_client=Model()) From 9132784e57759e31be8fd6f09b4699fd0f783a82 Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Thu, 17 Sep 2026 09:29:47 -0400 Subject: [PATCH 2/5] Preserve restricted callback compatibility --- .../src/hyperforge_restricted/agent.py | 5 +- agents/restricted/tests/test_sandbox.py | 56 +++++++ .../src/hyperforge/codemode/__init__.py | 2 + hyperforge/src/hyperforge/codemode/model.py | 35 +++++ hyperforge/src/hyperforge/codemode/sandbox.py | 7 + hyperforge/tests/codemode/test_protocol.py | 143 ++++++++++++++++++ 6 files changed, 247 insertions(+), 1 deletion(-) diff --git a/agents/restricted/src/hyperforge_restricted/agent.py b/agents/restricted/src/hyperforge_restricted/agent.py index 8f0ece8f..a0fcda0f 100644 --- a/agents/restricted/src/hyperforge_restricted/agent.py +++ b/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( diff --git a/agents/restricted/tests/test_sandbox.py b/agents/restricted/tests/test_sandbox.py index 43e142a6..0cce7e0c 100644 --- a/agents/restricted/tests/test_sandbox.py +++ b/agents/restricted/tests/test_sandbox.py @@ -1,4 +1,5 @@ import asyncio +import functools import os import tempfile from concurrent.futures import ProcessPoolExecutor @@ -12,6 +13,7 @@ WorkerExecutionRequest, ) from hyperforge.definition import FunctionDefinition +from pydantic import BaseModel from hyperforge_restricted.agent import PythonAgent from hyperforge_restricted.config import PythonAgentConfig @@ -236,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": {}} @@ -250,6 +261,51 @@ async def test_restricted_agent_rejects_undeclared_callback(): 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={} diff --git a/hyperforge/src/hyperforge/codemode/__init__.py b/hyperforge/src/hyperforge/codemode/__init__.py index 8f131d6d..a8418181 100644 --- a/hyperforge/src/hyperforge/codemode/__init__.py +++ b/hyperforge/src/hyperforge/codemode/__init__.py @@ -11,6 +11,7 @@ encode_protocol_value, encode_sandbox_message, serialize, + serialize_legacy_callback_result, validate_protocol_value, ) from .sandbox import ( @@ -46,6 +47,7 @@ "encode_protocol_value", "run_sandbox_server", "serialize", + "serialize_legacy_callback_result", "settings", "validate_protocol_value", ] diff --git a/hyperforge/src/hyperforge/codemode/model.py b/hyperforge/src/hyperforge/codemode/model.py index a2572096..3508fa0f 100644 --- a/hyperforge/src/hyperforge/codemode/model.py +++ b/hyperforge/src/hyperforge/codemode/model.py @@ -171,6 +171,41 @@ def serialize(message: Any) -> Any: return message +def serialize_legacy_callback_result(value: Any) -> Any: + """Serialize restricted-agent callback results with the pre-hardening behavior. + + Arbitrary Pydantic models are converted to their plain dictionary + representation with an ``__model__`` marker so callback results keep + crossing the remote sandbox boundary. Protocol models retain their native + wire representation, while unknown markers deserialize back to plain + dictionaries on the worker side. + """ + if isinstance(value, RestrictedPythonTask): + return { + "__model__": "RestrictedPythonTask", + "function": value.function, + "agent": value.agent, + "args": serialize_legacy_callback_result(value.args), + "keyword_args": serialize_legacy_callback_result(value.keyword_args), + } + if isinstance(value, BaseModel): + if type(value) in _MARKED_PROTOCOL_MODELS: + return value + converted = { + key: serialize_legacy_callback_result(item) + for key, item in value.model_dump().items() + } + converted["__model__"] = type(value).__name__ + return converted + if isinstance(value, (list, tuple)): + return [serialize_legacy_callback_result(item) for item in value] + if isinstance(value, dict): + return { + key: serialize_legacy_callback_result(item) for key, item in value.items() + } + return value + + def validate_protocol_value( value: Any, label: str, *, max_bytes: int = MAX_PROTOCOL_BYTES ) -> None: diff --git a/hyperforge/src/hyperforge/codemode/sandbox.py b/hyperforge/src/hyperforge/codemode/sandbox.py index 3c1c6e72..4921fccb 100644 --- a/hyperforge/src/hyperforge/codemode/sandbox.py +++ b/hyperforge/src/hyperforge/codemode/sandbox.py @@ -29,6 +29,7 @@ decode_protocol_value, encode_protocol_value, encode_sandbox_message, + serialize_legacy_callback_result, ) from .worker import PythonAgentWorker @@ -108,6 +109,7 @@ def remote( debug: bool = False, *, token: str | Callable[[], str] | None = None, + legacy_callback_results: bool = False, ) -> "SandboxRunner": return cls( pool=None, @@ -115,6 +117,7 @@ def remote( callback=callback, debug=debug, token=token, + legacy_callback_results=legacy_callback_results, ) @classmethod @@ -133,6 +136,7 @@ def __init__( callback: Callable[[RestrictedPythonTask], Coroutine[Any, Any, WorkerTypes]], debug: bool, token: str | Callable[[], str] | None = None, + legacy_callback_results: bool = False, ): self.pool = pool self.socket = socket @@ -140,6 +144,7 @@ def __init__( self.debug = debug self.isolated = False self.token = token + self.legacy_callback_results = legacy_callback_results self._callback_lock = threading.Lock() self._callbacks_closed = True self._run_active = False @@ -302,6 +307,8 @@ async def _run_remote_callback( response = callback_task.result() except Exception as exc: response = WorkerError(error=str(exc)) + if self.legacy_callback_results: + response = serialize_legacy_callback_result(response) await writer.write_message(SandboxMessage.Response(result=response)) except BaseException: if incoming.done(): diff --git a/hyperforge/tests/codemode/test_protocol.py b/hyperforge/tests/codemode/test_protocol.py index de748375..0930f2d9 100644 --- a/hyperforge/tests/codemode/test_protocol.py +++ b/hyperforge/tests/codemode/test_protocol.py @@ -10,6 +10,7 @@ import pytest from nuclia_models.predict.remi import RemiResponse +from pydantic import BaseModel from hyperforge.codemode import sandbox from hyperforge.codemode import worker as worker_module @@ -18,9 +19,12 @@ SandboxMessage, WorkerError, WorkerExecutionRequest, + decode_json_value, decode_protocol_value, deserialize, encode_protocol_value, + encode_sandbox_message, + serialize_legacy_callback_result, ) from hyperforge.codemode.sandbox import ( MAX_PACKET_BYTES, @@ -421,6 +425,145 @@ async def callback(_task: RestrictedPythonTask): runner._finish_callback_run() +@pytest.mark.asyncio +async def test_remote_runner_legacy_callback_results_round_trip() -> None: + class ApplicationResult(BaseModel): + value: int + + write_finished = asyncio.Event() + + class Reader: + async def read_message(self): + await write_finished.wait() + return SandboxMessage.Done() + + encoded_messages: list[bytes] = [] + + class Writer: + async def write_message(self, message): + encoded_messages.append(encode_sandbox_message(message, "Sandbox message")) + write_finished.set() + + async def callback(_task: RestrictedPythonTask): + return ApplicationResult(value=7) + + runner = SandboxRunner.remote( + "unused-socket", callback, legacy_callback_results=True + ) + runner._begin_callback_run() + incoming = await runner._run_remote_callback( + Reader(), # type: ignore[arg-type] + Writer(), # type: ignore[arg-type] + RestrictedPythonTask( + function="capability", agent="test", args=(), keyword_args={} + ), + ) + message = await asyncio.wait_for(incoming, timeout=0.2) + assert isinstance(message, SandboxMessage.Done) + await runner._cancel_callbacks() + runner._finish_callback_run() + + message = SandboxMessage.parse( + decode_json_value(encoded_messages[0], "Sandbox message") + ) + assert isinstance(message, SandboxMessage.Response) + assert message.result == {"__model__": "ApplicationResult", "value": 7} + + +def test_remote_runner_legacy_callback_results_cover_subclasses() -> None: + class CustomContext(Context): + extra: str = "detail" + + context = CustomContext( + original_question_uuid=None, + actual_question_uuid=None, + question="question", + chunks=[], + source="source", + agent="agent", + ) + + serialized = serialize_legacy_callback_result([context]) + + assert serialized == [{**context.model_dump(), "__model__": "CustomContext"}] + + +@pytest.mark.parametrize( + "value", + [ + Context( + original_question_uuid=None, + actual_question_uuid=None, + question="question", + chunks=[], + source="source", + agent="agent", + ), + RemiResponse(time=0.25), + WorkerError(error="failed"), + ], +) +def test_remote_runner_legacy_callback_results_preserve_protocol_models( + value: BaseModel, +) -> None: + serialized = serialize_legacy_callback_result(value) + encoded = encode_sandbox_message( + SandboxMessage.Response(result=serialized), "Sandbox message" + ) + message = SandboxMessage.parse(decode_json_value(encoded, "Sandbox message")) + + assert serialized is value + assert isinstance(message, SandboxMessage.Response) + assert type(message.result) is type(value) + assert message.result == value + + +def test_legacy_callback_worker_error_still_raises_in_worker() -> None: + class Pipe: + def recv_bytes(self, _max_bytes: int) -> bytes: + result = serialize_legacy_callback_result(WorkerError(error="denied")) + return encode_protocol_value(result, "Local sandbox response") + + worker = PythonAgentWorker(Pipe()) # type: ignore[arg-type] + + with pytest.raises(RuntimeError, match="denied"): + worker._receive() + + +@pytest.mark.asyncio +async def test_remote_runner_strict_callback_results_reject_arbitrary_models() -> None: + class ApplicationResult(BaseModel): + value: int + + write_finished = asyncio.Event() + + class Reader: + async def read_message(self): + await write_finished.wait() + return SandboxMessage.Done() + + class Writer: + async def write_message(self, message): + encode_sandbox_message(message, "Sandbox message") + write_finished.set() + + async def callback(_task: RestrictedPythonTask): + return ApplicationResult(value=7) + + runner = SandboxRunner.remote("unused-socket", callback) + runner._begin_callback_run() + with pytest.raises(ValueError, match="not serializable"): + await runner._run_remote_callback( + Reader(), # type: ignore[arg-type] + Writer(), # type: ignore[arg-type] + RestrictedPythonTask( + function="capability", agent="test", args=(), keyword_args={} + ), + ) + await runner._cancel_callbacks() + runner._finish_callback_run() + + @pytest.mark.asyncio async def test_callback_scheduled_after_cleanup_cannot_start() -> None: called = False From 221006f8643b82beccf7339d1f53d4e6a499b1f8 Mon Sep 17 00:00:00 2001 From: Nathan Van Gheem <524413+vangheem@users.noreply.github.com> Date: Mon, 21 Sep 2026 18:03:49 -0400 Subject: [PATCH 3/5] better guidance --- docs/harness-sdk.md | 4 +++- .../src/hyperforge/harness_sdk/tools/codemode.py | 13 ++++++++----- hyperforge/tests/harness_sdk/test_codemode.py | 9 +++++++++ 3 files changed, 20 insertions(+), 6 deletions(-) diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index 8c57b47e..199c620a 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -318,7 +318,9 @@ The normal `max_tool_calls` count includes the outer Code Mode call and every nested capability call exactly once. `max_nested_calls` independently bounds one generated program. Generated code must call `output(value)` with exactly one value, exactly once; a missing, empty, or repeated `output` call fails the -invocation. Authorization, nested/global call limits, projected-result limits, +invocation. `print()` is unavailable and must not be used to return a result. +Assign the final value to a variable and finish with `output(result)`. +Authorization, nested/global call limits, projected-result limits, and invalid or repeated `output` attempts remain terminal even if generated code catches the immediate callback error. An output attempt rejected by worker transport serialization also remains terminal. The optional `question` input is diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py index 9d31d688..db455b70 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -29,6 +29,11 @@ CODEMODE_TOOL_NAME = "codemode" OUTPUT_FUNCTION_NAME = "output" +CODEMODE_USAGE_GUIDANCE = ( + "Return the final result by calling output(value) exactly once. " + "Do not use print(); it is unavailable. " + "For example, assign the result to a variable and finish with output(result)." +) DEFAULT_MAX_SOURCE_BYTES = 64 * 1024 DEFAULT_MAX_RESULT_BYTES = 1024 * 1024 DEFAULT_MAX_CUMULATIVE_RESULT_BYTES = 4 * 1024 * 1024 @@ -191,7 +196,7 @@ def release(self) -> None: name=CODEMODE_TOOL_NAME, description=( "Execute restricted Python code. Registered tools are available as functions; " - "call output(value) to return a result." + f"{CODEMODE_USAGE_GUIDANCE}" ), ) async def codemode( @@ -858,10 +863,8 @@ def _nested_event_context(harness: Any) -> dict[str, Any]: def _scoped_description( capabilities: tuple[CodeModeCapability[Any], ...], description: str | None = None ) -> str: - introduction = description or ( - "Execute restricted Python code using only the scoped capabilities below; " - "call output(value) exactly once to return a result." - ) + introduction = description or "Execute restricted Python code using only the scoped capabilities below." + introduction = f"{introduction} {CODEMODE_USAGE_GUIDANCE}" if not capabilities: return introduction definitions = [] diff --git a/hyperforge/tests/harness_sdk/test_codemode.py b/hyperforge/tests/harness_sdk/test_codemode.py index cc9e6050..b30c994f 100644 --- a/hyperforge/tests/harness_sdk/test_codemode.py +++ b/hyperforge/tests/harness_sdk/test_codemode.py @@ -855,8 +855,17 @@ def test_scoped_codemode_description_includes_capability_schema() -> None: assert "Uppercase a value" in scoped.description assert '"required":["value"]' in scoped.description + assert "Do not use print(); it is unavailable." in scoped.description + assert "finish with output(result)" in scoped.description assert "Custom orchestration instructions." in custom.description assert '"required":["value"]' in custom.description + assert "Do not use print(); it is unavailable." in custom.description + + +def test_compatibility_codemode_description_explains_output_contract() -> None: + assert "output(value) exactly once" in codemode.description + assert "Do not use print(); it is unavailable." in codemode.description + assert "finish with output(result)" in codemode.description @pytest.mark.asyncio From 167a0cd91ec6469cd94f4b08b8ebf6b5c769015a Mon Sep 17 00:00:00 2001 From: Nathan Van Gheem <524413+vangheem@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:34:59 -0400 Subject: [PATCH 4/5] . --- docs/harness-sdk.md | 5 +++++ hyperforge/src/hyperforge/harness_sdk/tools/codemode.py | 4 ++++ hyperforge/tests/harness_sdk/test_codemode.py | 2 ++ 3 files changed, 11 insertions(+) diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index 5adaf733..59e16a23 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -322,6 +322,11 @@ generated program. Generated code must call `output(value)` with exactly one value, exactly once; a missing, empty, or repeated `output` call fails the invocation. `print()` is unavailable and must not be used to return a result. Assign the final value to a variable and finish with `output(result)`. +Generated code can use the provided capability functions and these common +built-ins: `abs`, `bool`, `bytes`, `chr`, `complex`, `divmod`, `float`, `hash`, +`hex`, `id`, `int`, `isinstance`, `issubclass`, `len`, `oct`, `ord`, `pow`, +`range`, `repr`, `round`, `slice`, `sorted`, `str`, `sum`, `tuple`, and `zip`. +Other helpers, including `all`, `any`, `min`, and `max`, are unavailable. Authorization, nested/global call limits, projected-result limits, and invalid or repeated `output` attempts remain terminal even if generated code catches the immediate callback error. An output attempt rejected by worker diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py index db455b70..21274f89 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -32,6 +32,10 @@ CODEMODE_USAGE_GUIDANCE = ( "Return the final result by calling output(value) exactly once. " "Do not use print(); it is unavailable. " + "Available common built-ins are abs, bool, bytes, chr, complex, divmod, float, " + "hash, hex, id, int, isinstance, issubclass, len, oct, ord, pow, range, repr, " + "round, slice, sorted, str, sum, tuple, and zip; helpers such as all, any, min, " + "and max are unavailable. " "For example, assign the result to a variable and finish with output(result)." ) DEFAULT_MAX_SOURCE_BYTES = 64 * 1024 diff --git a/hyperforge/tests/harness_sdk/test_codemode.py b/hyperforge/tests/harness_sdk/test_codemode.py index 342e120a..93b01993 100644 --- a/hyperforge/tests/harness_sdk/test_codemode.py +++ b/hyperforge/tests/harness_sdk/test_codemode.py @@ -857,6 +857,7 @@ def test_scoped_codemode_description_includes_capability_schema() -> None: assert "Uppercase a value" in scoped.description assert '"required":["value"]' in scoped.description assert "Do not use print(); it is unavailable." in scoped.description + assert "helpers such as all, any, min, and max are unavailable" in scoped.description assert "finish with output(result)" in scoped.description assert "Custom orchestration instructions." in custom.description assert '"required":["value"]' in custom.description @@ -866,6 +867,7 @@ def test_scoped_codemode_description_includes_capability_schema() -> None: def test_compatibility_codemode_description_explains_output_contract() -> None: assert "output(value) exactly once" in codemode.description assert "Do not use print(); it is unavailable." in codemode.description + assert "helpers such as all, any, min, and max are unavailable" in codemode.description assert "finish with output(result)" in codemode.description From 4149c884a8b619755c0f7ce1396cc8fdc67a739d Mon Sep 17 00:00:00 2001 From: Nathan Van Gheem <524413+vangheem@users.noreply.github.com> Date: Tue, 22 Sep 2026 16:21:54 -0400 Subject: [PATCH 5/5] fix --- hyperforge/src/hyperforge/harness_sdk/tools/codemode.py | 5 ++++- hyperforge/tests/harness_sdk/test_codemode.py | 8 ++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py index 21274f89..6e943d8c 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -867,7 +867,10 @@ def _nested_event_context(harness: Any) -> dict[str, Any]: def _scoped_description( capabilities: tuple[CodeModeCapability[Any], ...], description: str | None = None ) -> str: - introduction = description or "Execute restricted Python code using only the scoped capabilities below." + introduction = ( + description + or "Execute restricted Python code using only the scoped capabilities below." + ) introduction = f"{introduction} {CODEMODE_USAGE_GUIDANCE}" if not capabilities: return introduction diff --git a/hyperforge/tests/harness_sdk/test_codemode.py b/hyperforge/tests/harness_sdk/test_codemode.py index 93b01993..eb85dce8 100644 --- a/hyperforge/tests/harness_sdk/test_codemode.py +++ b/hyperforge/tests/harness_sdk/test_codemode.py @@ -857,7 +857,9 @@ def test_scoped_codemode_description_includes_capability_schema() -> None: assert "Uppercase a value" in scoped.description assert '"required":["value"]' in scoped.description assert "Do not use print(); it is unavailable." in scoped.description - assert "helpers such as all, any, min, and max are unavailable" in scoped.description + assert ( + "helpers such as all, any, min, and max are unavailable" in scoped.description + ) assert "finish with output(result)" in scoped.description assert "Custom orchestration instructions." in custom.description assert '"required":["value"]' in custom.description @@ -867,7 +869,9 @@ def test_scoped_codemode_description_includes_capability_schema() -> None: def test_compatibility_codemode_description_explains_output_contract() -> None: assert "output(value) exactly once" in codemode.description assert "Do not use print(); it is unavailable." in codemode.description - assert "helpers such as all, any, min, and max are unavailable" in codemode.description + assert ( + "helpers such as all, any, min, and max are unavailable" in codemode.description + ) assert "finish with output(result)" in codemode.description