From 0da01d36d4cfe63fa8a9867f29b7dbc88c385649 Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Tue, 8 Sep 2026 17:37:34 -0400 Subject: [PATCH 1/7] Add scoped Code Mode execution (rebased onto main) Rebased onto main after #123 (pluggable sandbox auth and ToolCallContext) landed; the branch already absorbed that foundation, so conflicts were resolved in favor of the combined final state. - Capability-scoped Code Mode tool factory (create_codemode_tool) with result projection, source/result/output/nested-call limits, nested event ancestry, execution admission, and tool inheritance controls - Hardened local/remote sandbox execution with authenticated transport, bounded protocol handling, runtime/memory ceilings, cancellation cleanup, and configurable socket/metrics exposure - Fix test_sandbox_rejects_concurrent_session to read the socket path from SANDBOX_SOCKET env instead of the import-time settings singleton - Apply ruff format to codemode/sandbox.py, tools/codemode.py, test_codemode.py --- .../src/hyperforge_restricted/agent.py | 2 +- agents/restricted/tests/test_sandbox.py | 108 +++ docs/harness-sdk.md | 202 ++++- .../src/hyperforge/codemode/__init__.py | 12 + hyperforge/src/hyperforge/codemode/model.py | 393 +++++++++- hyperforge/src/hyperforge/codemode/sandbox.py | 339 +++++++- hyperforge/src/hyperforge/codemode/worker.py | 23 +- .../src/hyperforge/harness_sdk/__init__.py | 27 +- .../src/hyperforge/harness_sdk/execution.py | 24 + .../src/hyperforge/harness_sdk/harness.py | 20 +- .../src/hyperforge/harness_sdk/models.py | 1 + .../hyperforge/harness_sdk/tools/__init__.py | 10 +- .../hyperforge/harness_sdk/tools/codemode.py | 552 ++++++++++++- .../src/hyperforge/harness_sdk/tools/utils.py | 18 +- hyperforge/src/hyperforge/server/sandbox.py | 21 +- hyperforge/tests/codemode/test_model.py | 182 +++++ hyperforge/tests/codemode/test_protocol.py | 296 ++++++- hyperforge/tests/harness_sdk/test_codemode.py | 735 ++++++++++++++++++ hyperforge/tests/harness_sdk/test_harness.py | 54 ++ 19 files changed, 2937 insertions(+), 82 deletions(-) create mode 100644 hyperforge/src/hyperforge/harness_sdk/execution.py diff --git a/agents/restricted/src/hyperforge_restricted/agent.py b/agents/restricted/src/hyperforge_restricted/agent.py index 8f0ece8f..526ca19c 100644 --- a/agents/restricted/src/hyperforge_restricted/agent.py +++ b/agents/restricted/src/hyperforge_restricted/agent.py @@ -436,7 +436,7 @@ async def _get_question_context( local_vars=local_vars, global_vars=copy_of_global_vars, function_names=self.function_names, - ) + ), ) ) ) diff --git a/agents/restricted/tests/test_sandbox.py b/agents/restricted/tests/test_sandbox.py index 25b550cb..a3b53cd0 100644 --- a/agents/restricted/tests/test_sandbox.py +++ b/agents/restricted/tests/test_sandbox.py @@ -124,6 +124,26 @@ async def test_sandbox_runtime_limit(sandbox_conn, long_timeout): assert "timed out" in message.error +async def test_sandbox_enforces_server_runtime_ceiling( + sandbox_conn, long_timeout, monkeypatch +): + monkeypatch.setattr(sandbox.settings, "sandbox_max_session_runtime_seconds", 0.1) + request = WorkerExecutionRequest( + code="while True: pass", + question="Q?", + local_vars={}, + global_vars={}, + function_names={}, + ) + (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 "timed out" in message.error + + async def test_sandbox_function_call(sandbox_conn, long_timeout): request = WorkerExecutionRequest( code="show(msg=foo())", @@ -188,3 +208,91 @@ 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_requires_token(monkeypatch, tmp_path): + monkeypatch.setattr(sandbox.settings, "sandbox_token", None) + + async def callback(_task): + return None + + runner = sandbox.SandboxRunner.remote(str(tmp_path / "unused.sock"), callback) + request = WorkerExecutionRequest( + code="", question="Q?", local_vars={}, global_vars={}, function_names={} + ) + + with pytest.raises(RuntimeError, match="SANDBOX_TOKEN is required"): + await runner.run(request) + + +async def test_sandbox_rejects_concurrent_session( + sandbox_conn, long_timeout, monkeypatch +): + request = WorkerExecutionRequest( + code="foo()", + question="Q?", + local_vars={}, + global_vars={}, + function_names={ + "self": { + "foo": FunctionDefinition(name="foo", description="", parameters={}) + } + }, + ) + monkeypatch.setattr(sandbox.settings, "sandbox_max_concurrent_sessions", 1) + first_rx, first_tx = sandbox_conn + 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( + os.environ["SANDBOX_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) + + +def test_sandbox_settings_default_to_private_optional_metrics() -> None: + configured = sandbox.SandboxSettings() + + assert configured.sandbox_metrics_enabled is False + assert configured.sandbox_metrics_host == "127.0.0.1" + assert configured.sandbox_socket_mode == "0600" + assert configured.sandbox_max_session_runtime_seconds == 60 + assert configured.sandbox_max_session_memory_bytes == 512 * 1024 * 1024 + + +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 cb6ff05c..e30f36f3 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -190,6 +190,193 @@ handlers always receive a validated model instance; they never receive `None`. 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` 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 + +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, + 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( + harness: AgentHarness, + 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( + 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, + tools=[code_mode], +) +``` + +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, 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 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. +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. 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. + +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, 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. + +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()`. + +### 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` is absent and never silently falls back to local execution. +Remote clients and the sandbox service both require a non-empty +`SANDBOX_TOKEN`; the Unix 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 settings use secure defaults: + +- `SANDBOX_MAX_CONCURRENT_SESSIONS=4` bounds authenticated worker sessions and + concurrent authentication handshakes. +- `SANDBOX_MAX_SESSION_RUNTIME_SECONDS=60` and + `SANDBOX_MAX_SESSION_MEMORY_BYTES=536870912` impose server-owned ceilings even + when a client omits limits. The remote client also bounds each connection by + the requested runtime (or the session ceiling when no runtime is requested) + plus `SANDBOX_TIMEOUT_SLACK_SECONDS=10`, so a hung sandbox cannot stall a + caller indefinitely. +- `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=false` avoids creating an IP listener. If enabled, + `SANDBOX_METRICS_HOST` defaults to `127.0.0.1` and + `SANDBOX_METRICS_PORT` defaults to `8091`. +- Local process IPC and remote length-prefixed JSON messages are bounded; + oversized run requests, callbacks, responses, and frames are rejected. + +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 All limits are disabled by default. Configure only the limits needed by the @@ -301,8 +488,10 @@ 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 inside a tool call also carry its `parent_call_id`; +parallel calls use task-local call context and do not overwrite one another. 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 @@ -376,10 +565,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/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..0669ac32 100644 --- a/hyperforge/src/hyperforge/codemode/model.py +++ b/hyperforge/src/hyperforge/codemode/model.py @@ -1,12 +1,27 @@ -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 pydantic import BaseModel, Field from hyperforge.definition import FunctionDefinition from hyperforge.memory import Context +MAX_PROTOCOL_BYTES = 16 * 1024 * 1024 +MAX_PROTOCOL_DEPTH = 256 +_JSON_STRING_CHUNK_CHARS = 4096 + + +class _ProtocolValueTooLarge(ValueError): + pass + + +@dataclass(frozen=True) +class _PlainModel: + value: BaseModel + class WorkerExecutionRequest(BaseModel): code: str @@ -14,8 +29,8 @@ 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) class RestrictedPythonTask(BaseModel): @@ -49,25 +64,36 @@ class WorkerError(BaseModel): } -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 == "RestrictedPythonTask": + args = msg.get("args") + keyword_args = msg.get("keyword_args") + 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"]), + args=tuple(deserialize(a, _depth=_depth + 1) for a in args), keyword_args={ - k: deserialize(v) for k, v in msg["keyword_args"].items() + k: deserialize(v, _depth=_depth + 1) + for k, v in keyword_args.items() }, ) model = WorkerModels.get(model_name) if model is not None: return cast(WorkerTypes, model.model_validate(msg)) 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 @@ -91,6 +117,324 @@ def serialize(message: Any) -> Any: 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 +) -> bytes: + try: + encoded = bytearray() + for chunk in _iter_json_bytes( + value, + depth=0, + markers=set(), + max_bytes=max_bytes, + include_model_markers=True, + ): + 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 (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, +) -> Iterator[bytes]: + if depth > MAX_PROTOCOL_DEPTH: + raise ValueError("Protocol value is too deeply nested") + if is_dataclass(value) and not isinstance(value, _PlainModel): + raise TypeError("Dataclass values are not supported by the sandbox protocol") + if isinstance(value, _PlainModel): + model = value.value + yield from _iter_json_object_bytes( + _model_items(model, include_marker=False), + marker=id(model), + depth=depth, + markers=markers, + max_bytes=max_bytes, + include_model_markers=False, + ) + 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, + ) + return + elif isinstance(value, BaseModel): + 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, + ) + 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, + ) + yield b"]" + finally: + markers.remove(marker) + return + if isinstance(value, dict): + 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, + ) + 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, +) -> 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, + ) + 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 _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: @@ -125,7 +469,8 @@ def parse(cls, data: dict[str, Any]) -> AnyMessage: ) 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"])) @@ -155,3 +500,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..ca100f44 100644 --- a/hyperforge/src/hyperforge/codemode/sandbox.py +++ b/hyperforge/src/hyperforge/codemode/sandbox.py @@ -1,9 +1,10 @@ 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 multiprocessing import Pipe, Process @@ -13,31 +14,63 @@ 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 = "127.0.0.1" + sandbox_metrics_enabled: bool = False sandbox_token: str | None = None + sandbox_max_concurrent_sessions: int = Field(default=4, gt=0) + sandbox_max_session_runtime_seconds: float = Field( + default=60, gt=0, allow_inf_nan=False + ) + sandbox_max_session_memory_bytes: int = Field(default=512 * 1024 * 1024, 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() +_remote_active_sessions = 0 class InsecureSandbox(Exception): @@ -94,32 +127,72 @@ def __init__( self.debug = debug self.isolated = False self.token = token + 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: - return await self._run_remotely(request) + 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") + try: + runtime = message.run.max_runtime_seconds + if runtime is None: + runtime = settings.sandbox_max_session_runtime_seconds + async with asyncio.timeout( + runtime + settings.sandbox_timeout_slack_seconds + ): + return await self._run_with_remote_admission(message.run) + 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") + async def _run_with_remote_admission(self, request: WorkerExecutionRequest): + await _acquire_remote_session() + try: + return await self._run_remotely(request) + finally: + _release_remote_session() + 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, - ) - 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, + ) + await controller_task + finally: + await self._cancel_callbacks() + 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) @@ -142,6 +215,7 @@ async def _run_isolated(self, request: WorkerExecutionRequest): if process.is_alive(): process.kill() await asyncio.to_thread(process.join) + await self._cancel_callbacks() if not controller_task.done(): controller_task.cancel() await asyncio.gather(controller_task, return_exceptions=True) @@ -169,16 +243,57 @@ async def _run_remotely(self, request: WorkerExecutionRequest): 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") + await self._run_remote_callback(reader, writer, msg.task) finally: + await self._cancel_callbacks() writer.close() await writer.wait_closed() + async def _run_remote_callback( + self, + reader: "SandboxReader", + writer: "SandboxWriter", + task: RestrictedPythonTask, + ) -> None: + 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)) + finally: + for candidate in (callback_task, incoming): + if not candidate.done(): + candidate.cancel() + done, still_pending = await asyncio.wait( + {callback_task, incoming}, timeout=CALLBACK_CANCEL_TIMEOUT + ) + for completed in done: + await asyncio.gather(completed, return_exceptions=True) + if still_pending: + logger.warning("Sandbox callback ignored cancellation") + self._callback_tasks.discard(callback_task) + for remaining in still_pending: + if remaining is callback_task: + self._track_orphaned_callback(callback_task) + else: + remaining.add_done_callback(_consume_task_result) + def run_in_process(self, request: WorkerExecutionRequest): pipe_worker, pipe_restricted = Pipe() worker = PythonAgentWorker(pipe_worker, debug=self.debug) @@ -212,6 +327,46 @@ async def background_task(self, pipe: Connection) -> bool: finally: pipe.close() + async def _cancel_callbacks(self) -> None: + tasks = list(self._callback_tasks) + for task in tasks: + task.cancel() + future = self._callback_futures.pop(task, None) + if future is not None and not future.done(): + future.set_exception(RuntimeError("Sandbox callback cancelled")) + 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") + 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: + self._orphaned_callback_tasks.add(task) + + def completed(done: asyncio.Task[WorkerTypes]) -> None: + 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: + pending = {task for task in self._orphaned_callback_tasks if not task.done()} + if not pending: + return False + + def completed(done: asyncio.Task[WorkerTypes]) -> None: + pending.discard(done) + if not pending: + callback() + + for task in pending: + task.add_done_callback(completed) + return True + def read_write_queue( self, pipe: Connection, @@ -220,7 +375,13 @@ def read_write_queue( ) -> bool: try: while True: - item = pipe.recv() + item = decode_protocol_value( + pipe.recv_bytes(MAX_PACKET_BYTES), + "Local sandbox request", + max_bytes=MAX_PACKET_BYTES, + ) + if not isinstance(item, RestrictedPythonTask): + raise ValueError("Invalid local sandbox request") if item.function == "_error": raise RuntimeError(item.args[0]) if item.function == "_close": @@ -229,12 +390,18 @@ def read_write_queue( def schedule_callback() -> None: 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: + self._callback_tasks.discard(completed) + controller_future = self._callback_futures.pop(completed, None) + if controller_future is None or controller_future.done(): + return try: - future.set_result(completed.result()) - except Exception as exc: - future.set_exception(exc) + controller_future.set_result(completed.result()) + except BaseException as exc: + controller_future.set_exception(exc) task.add_done_callback(copy_result) @@ -243,7 +410,19 @@ def copy_result(completed: asyncio.Task[WorkerTypes]) -> None: result = future.result() except Exception as exc: result = WorkerError(error=str(exc)) - pipe.send(result) + 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) except EOFError: return False @@ -269,16 +448,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 +476,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: @@ -320,14 +514,18 @@ def __init__( self.timed_out = False async def run(self, request: WorkerExecutionRequest): + memory = settings.sandbox_max_session_memory_bytes + if request.max_memory_bytes is not None: + memory = min(memory, request.max_memory_bytes) + request = request.model_copy(update={"max_memory_bytes": memory}) self.task, self.process = self.runner.run_in_process(request) try: await self._start_timeout(self.callback_wait_seconds) - if request.max_runtime_seconds is None: + runtime = settings.sandbox_max_session_runtime_seconds + if request.max_runtime_seconds is not None: + runtime = min(runtime, request.max_runtime_seconds) + async with asyncio.timeout(runtime): clean_exit = await self.task - else: - async with asyncio.timeout(request.max_runtime_seconds): - clean_exit = await self.task await self._stop_timeout() if self.timed_out: raise RuntimeError("Python agent timeout") @@ -348,6 +546,7 @@ async def run(self, request: WorkerExecutionRequest): if self.process.is_alive(): self.process.kill() await asyncio.to_thread(self.process.join) + await self.runner._cancel_callbacks() if not self.task.done(): self.task.cancel() await asyncio.gather(self.task, return_exceptions=True) @@ -357,6 +556,7 @@ async def _fail(self, exc: Exception): if self.process.is_alive(): self.process.kill() await asyncio.to_thread(self.process.join) + await self.runner._cancel_callbacks() if not self.task.done(): self.task.cancel() await asyncio.gather(self.task, return_exceptions=True) @@ -376,7 +576,8 @@ async def _callback(self, item: RestrictedPythonTask) -> WorkerTypes: self.process.join() 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 @@ -405,18 +606,30 @@ 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() + handshake_limit = asyncio.Semaphore(settings.sandbox_max_concurrent_sessions) + async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter): + nonlocal active_sessions reader = SandboxReader(rx) writer = SandboxWriter(tx) + if handshake_limit.locked(): + writer.close() + await writer.wait_closed() + return + 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,6 +647,18 @@ async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter): writer.close() await writer.wait_closed() return + finally: + handshake_limit.release() + + async with admission_lock: + if active_sessions >= settings.sandbox_max_concurrent_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() @@ -447,12 +672,46 @@ async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter): 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() -> None: + global _remote_active_sessions + while True: + with _remote_admission_lock: + if _remote_active_sessions < settings.sandbox_max_concurrent_sessions: + _remote_active_sessions += 1 + return + await asyncio.sleep(0.01) + + +def _release_remote_session() -> None: + global _remote_active_sessions + with _remote_admission_lock: + if _remote_active_sessions <= 0: + raise RuntimeError("Remote sandbox admission counter underflow") + _remote_active_sessions -= 1 diff --git a/hyperforge/src/hyperforge/codemode/worker.py b/hyperforge/src/hyperforge/codemode/worker.py index 89b134bc..96c6bc06 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", @@ -162,7 +168,7 @@ def _process_question_context_sync( global_vars["pdb"] = __import__("pdb") exec(byte_code, global_vars, local_vars) except BaseException as exc: - self.pipe.send( + self._send( RestrictedPythonTask( function="_error", agent="_controller", @@ -171,7 +177,7 @@ def _process_question_context_sync( ) ) finally: - self.pipe.send( + self._send( RestrictedPythonTask( function="_close", agent="_controller", @@ -188,7 +194,7 @@ def save( question: str = "", final_answer: Optional[str] = None, ): - self.pipe.send( + self._send( RestrictedPythonTask( function="save", agent="_controller", @@ -224,7 +230,7 @@ def valid_functions_call( ) agent_id = agents[0] - self.pipe.send( + self._send( RestrictedPythonTask( function=function_name, agent=agent_id, @@ -234,8 +240,13 @@ def valid_functions_call( ) return self._receive() + def _send(self, task: RestrictedPythonTask) -> None: + self.pipe.send_bytes(encode_protocol_value(task, "Local sandbox request")) + def _receive(self) -> Any: - result = self.pipe.recv() + result = decode_protocol_value( + self.pipe.recv_bytes(MAX_PROTOCOL_BYTES), "Local sandbox response" + ) 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 36c901db..d38cdeec 100644 --- a/hyperforge/src/hyperforge/harness_sdk/__init__.py +++ b/hyperforge/src/hyperforge/harness_sdk/__init__.py @@ -39,8 +39,21 @@ HarnessToolCall, ) from .storage import HarnessStorageProtocol, InMemoryHarnessStorage -from .tools import HarnessTool, ToolCallContext, tool -from .tools.codemode import CodemodeInput, CodemodeOutput, codemode +from .tools import HarnessTool, ToolCallContext, ToolInheritancePolicy, tool +from .tools.codemode import ( + CodeModeCapability, + CodeModeDispatch, + CodeModeExecutionLimiter, + CodemodeInput, + CodeModeLimits, + CodemodeOutput, + CodeModeResultAdapter, + CodeModeRunner, + codemode, + context_codemode_result_adapter, + create_codemode_tool, + raw_codemode_result_adapter, +) from .usage import HarnessUsage, UsageLimitExceeded, UsageLimits __all__ = [ @@ -56,6 +69,12 @@ "ChatCompletionToolCallDelta", "ChatCompletionToolCallFunctionDelta", "ChatCompletionUsage", + "CodeModeCapability", + "CodeModeDispatch", + "CodeModeExecutionLimiter", + "CodeModeLimits", + "CodeModeResultAdapter", + "CodeModeRunner", "CodemodeInput", "CodemodeOutput", "HarnessAgentToTool", @@ -79,14 +98,18 @@ "NucliaChatCompletionsClient", "NucliaChatCompletionsError", "ReasoningEffort", + "ToolInheritancePolicy", "UsageLimitExceeded", "UsageLimits", "context_schema", "codemode", + "context_codemode_result_adapter", "create_agent", + "create_codemode_tool", "format_context", "make_context", "register_context", + "raw_codemode_result_adapter", "tool", "published_agent_to_tools", "preload_published_agent_to_tools", diff --git a/hyperforge/src/hyperforge/harness_sdk/execution.py b/hyperforge/src/hyperforge/harness_sdk/execution.py new file mode 100644 index 00000000..89c82b18 --- /dev/null +++ b/hyperforge/src/hyperforge/harness_sdk/execution.py @@ -0,0 +1,24 @@ +from contextvars import ContextVar, Token + +_current_tool_call_id: ContextVar[str | None] = ContextVar( + "hyperforge_current_tool_call_id", default=None +) + + +def current_tool_call_id() -> str | None: + return _current_tool_call_id.get() + + +def set_current_tool_call_id(call_id: str) -> Token[str | None]: + return _current_tool_call_id.set(call_id) + + +def reset_current_tool_call_id(token: Token[str | None]) -> None: + _current_tool_call_id.reset(token) + + +__all__ = [ + "current_tool_call_id", + "reset_current_tool_call_id", + "set_current_tool_call_id", +] diff --git a/hyperforge/src/hyperforge/harness_sdk/harness.py b/hyperforge/src/hyperforge/harness_sdk/harness.py index ba252b36..d0560cb0 100644 --- a/hyperforge/src/hyperforge/harness_sdk/harness.py +++ b/hyperforge/src/hyperforge/harness_sdk/harness.py @@ -14,6 +14,11 @@ from .agents import published_agent_to_tools from .clients import ModelClient, ReasoningEffort from .context import format_context +from .execution import ( + current_tool_call_id, + reset_current_tool_call_id, + set_current_tool_call_id, +) from .models import ( HarnessConversation, HarnessEvent, @@ -23,7 +28,7 @@ HarnessToolCall, ) from .storage import HarnessStorageProtocol, InMemoryHarnessStorage -from .tools import HarnessTool, ToolCallContext +from .tools import HarnessTool, ToolCallContext, ToolInheritancePolicy from .tools.core import DictOutput, SendMessageInput, SpawnAgentInput, create_core_tools from .usage import HarnessUsage, UsageLimitExceeded, UsageLimits @@ -594,6 +599,7 @@ async def emit( turn_id=self._turn_id, agent_id=self.agent_id, parent_agent_id=self.parent_agent_id, + parent_call_id=current_tool_call_id(), category=self.category, tags=self.tags, metadata=self._persisted_metadata(), @@ -967,7 +973,11 @@ async def _execute_tool_call(self, call: HarnessToolCall) -> HarnessMessage: raise ValueError(f"Tool is not active: {call.name}") else: context = ToolCallContext(harness=self, name=call.name, id=call.id) - output = await tool.execute(context, call.arguments) + token = set_current_tool_call_id(call.id) + try: + output = await tool.execute(context, call.arguments) + finally: + reset_current_tool_call_id(token) result = output.model_dump(mode="json") reference = tool.context(output) content = format_context(reference) @@ -1036,7 +1046,11 @@ def _create_child(self, child_id: str, *, include_history: bool) -> AgentHarness model=self._config.model, model_client=self._config.model_client, reasoning_effort=self._config.reasoning_effort, - tools=self._config.tools, + tools=( + tool + for tool in self._config.tools + if tool.inheritance == ToolInheritancePolicy.INHERIT + ), system_prompt=self._config.system_prompt, title=self._config.title, conversation_id=self.conversation_id, 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/__init__.py b/hyperforge/src/hyperforge/harness_sdk/tools/__init__.py index 9bb28864..7d0c592a 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/__init__.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/__init__.py @@ -3,7 +3,15 @@ HarnessTool, ToolCallContext, ToolHandler, + ToolInheritancePolicy, tool, ) -__all__ = ["ContextFactory", "HarnessTool", "ToolCallContext", "ToolHandler", "tool"] +__all__ = [ + "ContextFactory", + "HarnessTool", + "ToolCallContext", + "ToolHandler", + "ToolInheritancePolicy", + "tool", +] diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py index 18e3c442..a78e4029 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -1,4 +1,13 @@ -from typing import Any +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 pydantic import BaseModel @@ -6,24 +15,167 @@ RestrictedPythonTask, SandboxRunner, WorkerExecutionRequest, + encode_protocol_value, ) from hyperforge.codemode.sandbox import settings from hyperforge.definition import FunctionDefinition -from . import HarnessTool, ToolCallContext, tool +from ..context import format_context +from ..execution import ( + current_tool_call_id, + reset_current_tool_call_id, + set_current_tool_call_id, +) +from ..models import HarnessEventType +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", + "Chunk", + "Context", + "Dict", + "List", + "__builtins__", + "_getitem_", + "_getiter_", + "_inplacevar_", + "_iter_unpack_sequence_", + "_unpack_sequence_", + "agent_id", + "codemode", + "dataclass", + "output", + "pdb", + "print", + "printed", + "question", + "save", + } +) +_SENSITIVE_FIELD_PARTS = ( + "api_key", + "apikey", + "authorization", + "cookie", + "credential", + "passwd", + "password", + "private_key", + "privatekey", + "secret", + "token", +) +type CodeModeResultAdapter = Callable[ + [HarnessTool[Any, Any], BaseModel], Any | Awaitable[Any] +] +type CodeModeDispatch = Callable[[RestrictedPythonTask], Awaitable[Any]] class CodemodeInput(BaseModel): code: str + question: str = "" class CodemodeOutput(BaseModel): value: Any = None +def context_codemode_result_adapter( + capability: HarnessTool[Any, Any], output: BaseModel +) -> Any: + """Project a result to the same formatted context shown to the model.""" + return format_context(capability.context(output)) + + +def raw_codemode_result_adapter( + _capability: HarnessTool[Any, Any], output: BaseModel +) -> Any: + """Explicitly expose the capability's complete JSON-mode output.""" + return output.model_dump(mode="json") + + +class CodeModeRunner(Protocol): + """Executes a Code Mode request, routing worker callbacks through dispatch. + + Injections are intended for deterministic tests; production tools use the + default remote sandbox (or the isolated local process when + ``remote_required=False``). An injected runner owns its callback lifecycle: + pending callbacks it leaves behind are not tracked for slot retention. + """ + + async def run( + self, request: WorkerExecutionRequest, dispatch: CodeModeDispatch + ) -> Any: ... + + +@dataclass(frozen=True) +class CodeModeCapability: + tool: HarnessTool[Any, Any] + result_adapter: CodeModeResultAdapter = 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=( @@ -83,6 +235,7 @@ async def dispatch(task: RestrictedPythonTask) -> Any: await runner.run( WorkerExecutionRequest( code=input_value.code, + question=input_value.question, local_vars={}, global_vars={}, function_names={"harness": definitions}, @@ -93,6 +246,385 @@ async def dispatch(task: RestrictedPythonTask) -> Any: return result +def create_codemode_tool( + *, + capabilities: tuple[CodeModeCapability, ...], + 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 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] = {} + for capability in capabilities: + if not isinstance(capability, CodeModeCapability): + raise TypeError("capabilities must contain CodeModeCapability values") + capability_name = capability.tool.name + if ( + not capability_name.isidentifier() + or keyword.iskeyword(capability_name) + or capability_name.startswith("_") + ): + raise ValueError( + f"Code Mode capability name {capability_name!r} is not a valid " + "Python identifier" + ) + if capability_name in RESERVED_CAPABILITY_NAMES: + raise ValueError( + f"Code Mode capability name {capability_name!r} is reserved" + ) + if capability_name in capability_map: + raise ValueError( + f"Code Mode capability names must be unique: {capability_name}" + ) + if ( + "agent_id" in capability.tool.input_model.model_fields + or "agent_id" in capability.tool.parameters.get("properties", {}) + ): + raise ValueError( + f"Code Mode capability {capability_name!r} uses the reserved " + "'agent_id' input field" + ) + capability_map[capability_name] = capability + scoped_capabilities = tuple(capabilities) + + async def execute( + context: ToolCallContext, input_value: CodemodeInput + ) -> CodemodeOutput: + harness = context.harness + source_size = len(input_value.code.encode("utf-8")) + if source_size > limits.max_source_bytes: + raise ValueError( + "Code Mode source exceeds maximum size: " + f"{source_size} > {limits.max_source_bytes} bytes" + ) + socket = settings.sandbox_socket + 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 settings.sandbox_token: + raise RuntimeError( + "Remote Code Mode execution is required but SANDBOX_TOKEN is absent" + ) + execution_limiter.acquire() + return await _execute_scoped_codemode( + harness, + input_value, + scoped_capabilities, + limits, + socket, + execution_limiter, + runner, + ) + + 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, ...], + limits: CodeModeLimits, + socket: str | None, + execution_limiter: CodeModeExecutionLimiter, + runner: CodeModeRunner | None = None, +) -> CodemodeOutput: + capability_map = {capability.tool.name: capability for capability in capabilities} + result = CodemodeOutput() + output_state = _OutputState() + nested_calls = 0 + cumulative_result_bytes = 0 + parent_call_id = current_tool_call_id() + + 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 + + capability = capability_map.get(task.function) + if capability is None: + raise ValueError(f"Unknown Code Mode capability: {task.function}") + nested_calls += 1 + call_id = uuid.uuid4().hex + event_context = _nested_event_context(harness) + marker = { + "codemode": True, + "nested": True, + "parent_call_id": parent_call_id, + **event_context, + } + await harness.emit( + HarnessEventType.TOOL_REQUESTED, + { + "call": { + "id": call_id, + "name": capability.tool.name, + "arguments": _sanitize_task_arguments(capability.tool, task), + }, + **marker, + }, + ) + token = set_current_tool_call_id(call_id) + started = time.perf_counter() + try: + harness.usage.tool_calls += 1 + harness._check_limit("max_tool_calls", harness.usage.tool_calls) + if nested_calls > limits.max_nested_calls: + raise RuntimeError( + "Code Mode nested call limit exceeded: " + f"{nested_calls} > {limits.max_nested_calls}" + ) + arguments = _tool_arguments(capability.tool, task.args, task.keyword_args) + output = await capability.tool.execute( + ToolCallContext(harness=harness, name=capability.tool.name, id=call_id), + arguments, + ) + projected = capability.result_adapter(capability.tool, output) + if inspect.isawaitable(projected): + projected = await projected + normalized, result_bytes = _normalize_worker_value( + projected, + limits.max_result_bytes, + f"Code Mode result from {capability.tool.name}", + ) + cumulative_result_bytes += result_bytes + if cumulative_result_bytes > limits.max_cumulative_result_bytes: + raise RuntimeError( + "Code Mode cumulative result limit exceeded: " + f"{cumulative_result_bytes} > " + f"{limits.max_cumulative_result_bytes} bytes" + ) + sanitized_result = _sanitize_event_value(normalized) + except BaseException as exc: + reset_current_tool_call_id(token) + await harness.emit( + 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 + raise RuntimeError( + f"Code Mode capability {capability.tool.name!r} failed " + f"({type(exc).__name__})" + ) from None + else: + reset_current_tool_call_id(token) + await harness.emit( + 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 + if runner is None: + sandbox_runner = ( + SandboxRunner.remote(socket, dispatch) + if socket is not None + else SandboxRunner.isolated_process(dispatch) + ) + try: + definitions = { + capability.tool.name: FunctionDefinition( + name=capability.tool.name, + description=capability.tool.description, + parameters=capability.tool.parameters.get("properties", {}), + ) + for capability in capabilities + } + definitions[OUTPUT_FUNCTION_NAME] = FunctionDefinition( + name=OUTPUT_FUNCTION_NAME, + description="Set the value returned by Code Mode.", + parameters={"value": {}}, + ) + worker_request = WorkerExecutionRequest( + code=input_value.code, + question=input_value.question, + local_vars={}, + global_vars={}, + function_names={"harness": definitions}, + max_runtime_seconds=harness.usage_limits.max_codemode_runtime_seconds, + max_memory_bytes=harness.usage_limits.max_codemode_memory_bytes, + ) + if runner is not None: + 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() + if not output_state.called: + raise ValueError("Code Mode must call output(value) exactly once") + return result + + +class _OutputState: + """Validate output(value) calls: exactly one call, exactly one value.""" + + def __init__(self) -> None: + self.called = False + + def record(self, task: RestrictedPythonTask, max_output_bytes: int) -> Any: + if self.called: + raise ValueError("output(value) may only be called once") + value = _output_value(task) + normalized, _ = _normalize_worker_value( + value, max_output_bytes, "Code Mode output" + ) + 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]: + if isinstance(value, BaseModel): + raise TypeError( + f"{label} must be projected to serializable worker values, not a model" + ) + try: + _reject_reserved_model_markers(value, label) + encoded = encode_protocol_value(value, label, max_bytes=max_bytes) + return json.loads(encoded), len(encoded) + except (RecursionError, TypeError, ValueError) as exc: + if "exceeds maximum size" in str(exc): + raise ValueError(str(exc)) from exc + raise TypeError(f"{label} is not a serializable worker value") from exc + + +def _reject_reserved_model_markers(value: Any, label: str) -> None: + if isinstance(value, dict): + if "__model__" in value: + raise ValueError(f"{label} contains the reserved '__model__' key") + for nested in value.values(): + _reject_reserved_model_markers(nested, label) + elif isinstance(value, (list, tuple)): + for nested in value: + _reject_reserved_model_markers(nested, label) + + +def _sanitize_event_value(value: Any) -> Any: + if isinstance(value, dict): + return { + str(key): ( + "[REDACTED]" + if any(part in str(key).casefold() for part in _SENSITIVE_FIELD_PARTS) + else _sanitize_event_value(nested) + ) + for key, nested in value.items() + } + if isinstance(value, (list, tuple)): + return [_sanitize_event_value(nested) for nested in value] + if isinstance(value, str): + try: + parsed = json.loads(value, parse_constant=_reject_json_constant) + except (json.JSONDecodeError, RecursionError, ValueError): + return {"type": "string", "utf8_bytes": len(value.encode("utf-8"))} + if isinstance(parsed, (dict, list)): + return _sanitize_event_value(parsed) + return {"type": "string", "utf8_bytes": len(value.encode("utf-8"))} + if isinstance(value, float) and not math.isfinite(value): + return "" + if isinstance(value, (bool, int, float)) or value is None: + return value + return f"<{type(value).__name__}>" + + +def _reject_json_constant(value: str) -> None: + raise ValueError(f"Invalid JSON constant: {value}") + + +def _sanitize_task_arguments( + capability: HarnessTool[Any, Any], task: RestrictedPythonTask +) -> Any: + try: + arguments = _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 = {} + for key in ("actor", "actor_id", "tenant", "tenant_id", "user_id"): + value = harness.execution_context.get(key) + if isinstance(value, (str, int, float, bool)): + context[key] = value + return {"execution_context": context} if context else {} + + +def _scoped_description( + capabilities: tuple[CodeModeCapability, ...], 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." + ) + if not capabilities: + return introduction + definitions = [] + for capability in capabilities: + schema = json.dumps( + capability.tool.parameters, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ) + definitions.append( + f"{capability.tool.name}: {capability.tool.description or 'No description.'} " + f"Arguments JSON Schema: {schema}" + ) + return f"{introduction}\nCapability functions:\n" + "\n".join(definitions) + + def _tool_arguments( tool: HarnessTool[Any, Any], args: tuple[Any, ...], @@ -112,4 +644,18 @@ def _tool_arguments( return arguments -__all__ = ["CodemodeInput", "CodemodeOutput", "codemode"] +__all__ = [ + "CodeModeCapability", + "CodeModeDispatch", + "CodeModeExecutionLimiter", + "CodeModeLimits", + "CodeModeResultAdapter", + "CodeModeRunner", + "CodemodeInput", + "CodemodeOutput", + "RESERVED_CAPABILITY_NAMES", + "codemode", + "context_codemode_result_adapter", + "create_codemode_tool", + "raw_codemode_result_adapter", +] diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/utils.py b/hyperforge/src/hyperforge/harness_sdk/tools/utils.py index 20c4eec1..7c6740bd 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/utils.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/utils.py @@ -2,6 +2,7 @@ import inspect from dataclasses import dataclass, field +from enum import StrEnum from typing import TYPE_CHECKING, Any, Awaitable, Callable, cast, get_type_hints from jsonschema import ValidationError as JsonSchemaValidationError @@ -21,6 +22,11 @@ type ContextFactory[OutputT: BaseModel] = Callable[[OutputT], HarnessContextReference] +class ToolInheritancePolicy(StrEnum): + INHERIT = "inherit" + DO_NOT_INHERIT = "do_not_inherit" + + @dataclass(frozen=True) class ToolCallContext: """Execution context passed to a tool handler. @@ -43,6 +49,7 @@ class HarnessTool[InputT: BaseModel, OutputT: BaseModel]: parameters_schema: dict[str, Any] | None = None context_factory: ContextFactory[OutputT] | None = None lazy_load: bool = False + inheritance: ToolInheritancePolicy = ToolInheritancePolicy.INHERIT input_model: type[InputT] = field(init=False) output_model: type[OutputT] = field(init=False) _parameters: dict[str, Any] = field(init=False, repr=False) @@ -116,6 +123,7 @@ def tool( parameters_schema: dict[str, Any] | None = None, context_factory: ContextFactory[Any] | None = None, lazy_load: bool = False, + inheritance: ToolInheritancePolicy = ToolInheritancePolicy.INHERIT, ) -> Callable[[ToolHandler[Any, Any]], HarnessTool[Any, Any]]: """Create a harness tool from an annotated async handler.""" @@ -128,6 +136,7 @@ def decorate(handler: ToolHandler[Any, Any]) -> HarnessTool[Any, Any]: parameters_schema=parameters_schema, context_factory=context_factory, lazy_load=lazy_load, + inheritance=inheritance, ) return decorate @@ -137,4 +146,11 @@ def _is_model_type(value: Any) -> bool: return isinstance(value, type) and issubclass(value, BaseModel) -__all__ = ["ContextFactory", "HarnessTool", "ToolCallContext", "ToolHandler", "tool"] +__all__ = [ + "ContextFactory", + "HarnessTool", + "ToolCallContext", + "ToolHandler", + "ToolInheritancePolicy", + "tool", +] 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..c954a8f6 100644 --- a/hyperforge/tests/codemode/test_model.py +++ b/hyperforge/tests/codemode/test_model.py @@ -1,9 +1,21 @@ +import json +import tracemalloc +from dataclasses import dataclass +from typing import Annotated + +import pytest +from pydantic import BaseModel, PlainSerializer, ValidationError, field_serializer + from hyperforge.codemode import ( SandboxMessage, WorkerExecutionRequest, + decode_protocol_value, deserialize, + encode_protocol_value, + encode_sandbox_message, serialize, ) +from hyperforge.codemode import model as model_module def test_nested_json_values_round_trip() -> None: @@ -27,3 +39,173 @@ def test_sandbox_run_token_round_trip() -> None: assert isinstance(parsed, SandboxMessage.Run) assert parsed.token == "secret" + + +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 + + +@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..add0abcf 100644 --- a/hyperforge/tests/codemode/test_protocol.py +++ b/hyperforge/tests/codemode/test_protocol.py @@ -1,15 +1,23 @@ import asyncio import tempfile +import threading from pathlib import Path import pytest -from hyperforge.codemode.model import WorkerExecutionRequest, deserialize +from hyperforge.codemode import sandbox +from hyperforge.codemode.model import ( + RestrictedPythonTask, + SandboxMessage, + WorkerExecutionRequest, + deserialize, +) from hyperforge.codemode.sandbox import ( MAX_PACKET_BYTES, SandboxReader, SandboxRunner, SandboxSettings, + SandboxWriter, run_sandbox_server, ) from hyperforge.definition import FunctionDefinition @@ -51,6 +59,292 @@ 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=1, + ) + ) + + +@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) + 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) + + +@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_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)) diff --git a/hyperforge/tests/harness_sdk/test_codemode.py b/hyperforge/tests/harness_sdk/test_codemode.py index 007b4cda..03fa00fc 100644 --- a/hyperforge/tests/harness_sdk/test_codemode.py +++ b/hyperforge/tests/harness_sdk/test_codemode.py @@ -1,19 +1,37 @@ +import asyncio import sys from collections.abc import AsyncIterator from contextvars import ContextVar +from typing import Any import pytest from pydantic import BaseModel +from hyperforge.codemode import RestrictedPythonTask, WorkerExecutionRequest from hyperforge.harness_sdk import ( AgentHarness, + CodeModeCapability, + CodeModeExecutionLimiter, CodemodeInput, + CodeModeLimits, + HarnessContextReference, + HarnessContextType, + HarnessEventType, + HarnessTool, + HarnessToolCall, ModelDelta, ToolCallContext, UsageLimits, codemode, + create_codemode_tool, + raw_codemode_result_adapter, tool, ) +from hyperforge.harness_sdk.tools import codemode as codemode_module + + +def _ctx(harness: Any) -> ToolCallContext: + return ToolCallContext(harness=harness, name="test") class UpperInput(BaseModel): @@ -24,6 +42,11 @@ class UpperOutput(BaseModel): value: str +class SensitiveOutput(BaseModel): + value: str + secret: str + + @tool(description="Uppercase a value") async def upper(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: return UpperOutput(value=input_value.value.upper()) @@ -35,6 +58,29 @@ async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: yield ModelDelta(text="unused") +def local_codemode( + *capabilities: CodeModeCapability, + limits: CodeModeLimits | None = None, + execution_limiter: CodeModeExecutionLimiter | None = None, +) -> HarnessTool: + 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, output: BaseModel) -> 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( @@ -184,3 +230,692 @@ async def test_codemode_enforces_memory_limit() -> None: ToolCallContext(harness=harness, name=codemode.name), CodemodeInput(code="output('x' * (1024 * 1024 * 1024))").model_dump(), ) + + +@pytest.mark.asyncio +async def test_scoped_codemode_uses_only_explicit_hidden_capabilities( + monkeypatch: pytest.MonkeyPatch, +) -> None: + @tool() + async def lower(_: AgentHarness, input_value: UpperInput) -> UpperOutput: + return UpperOutput(value=input_value.value.lower()) + + scoped = local_codemode(CodeModeCapability(upper, result_adapter=value_adapter)) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + tools=[scoped, lower], + ) + monkeypatch.setattr( + harness, + "iter_tools", + lambda *args, **kwargs: pytest.fail("scoped Code Mode discovered tools"), + ) + + result = await scoped.execute( + _ctx(harness), + {"code": "result = upper(value='hello')\noutput(result['value'])"}, + ) + + assert result.value == "HELLO" + assert upper.name not in harness._tools + for unavailable in ("lower", "remember"): + with pytest.raises(RuntimeError, match=f"{unavailable}.*not defined"): + await scoped.execute(_ctx(harness), {"code": f"{unavailable}(value='x')"}) + + +@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"] == { + "value": {"type": "string", "utf8_bytes": 4}, + "secret": "[REDACTED]", + } + + +@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"] == { + "value": {"type": "string", "utf8_bytes": 35}, + "items": [{"type": "string", "utf8_bytes": 14}], + } + assert "secret" not in str(completed.payload) + + +@pytest.mark.asyncio +async def test_scoped_codemode_does_not_persist_non_finite_json_numbers() -> 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"] == {"score": ""} + + +@pytest.mark.asyncio +async def test_scoped_codemode_raw_projection_must_be_explicit() -> 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, secret="raw-secret") + + scoped = local_codemode( + CodeModeCapability(lookup, result_adapter=raw_codemode_result_adapter) + ) + 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", "secret": "raw-secret"} + + +@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="capability 'upper' 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="capability 'upper' 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="output exceeds maximum size"): + await output_limited.execute(_ctx(harness), {"code": "output('hello')"}) + + +@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 == 3 + 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, + HarnessEventType.TOOL_REQUESTED, + HarnessEventType.TOOL_FAILED, + ] + + +@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"]["value"] == { + "type": "string", + "utf8_bytes": 4, + } + assert nested[0].payload["call"]["arguments"]["token"] == "[REDACTED]" + assert nested[1].payload["result"] == { + "value": {"type": "string", "utf8_bytes": 4}, + "token": "[REDACTED]", + } + assert nested[0].payload["execution_context"] == { + "actor_id": "actor-1", + "tenant_id": "tenant-1", + } + + +@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"] == {"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_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=r"capability 'fail' failed \(ValueError\)"): + 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) + + +@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)"}) + + +class RecordingRunner: + def __init__(self, tasks: list[RestrictedPythonTask]) -> None: + self.tasks = tasks + self.requests: list[WorkerExecutionRequest] = [] + self.results: dict[str, Any] = {} + + async def run(self, request: WorkerExecutionRequest, dispatch: Any) -> None: + self.requests.append(request) + for task in self.tasks: + self.results[task.function] = await dispatch(task) + + +def worker_task(function: str, *args: Any, **kwargs: Any) -> RestrictedPythonTask: + return RestrictedPythonTask( + function=function, agent="harness", args=args, keyword_args=kwargs + ) + + +@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 runner.results["upper"] == {"value": "HELLO"} + 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_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: + runner = RecordingRunner([worker_task("upper", "hi"), worker_task("upper", "ok")]) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(upper, result_adapter=value_adapter),), + runner=runner, + limits=CodeModeLimits(max_result_bytes=20, max_cumulative_result_bytes=20), + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(RuntimeError, match="capability 'upper' failed"): + await scoped.execute(_ctx(harness), {"code": "upper('hi'); upper('ok')"}) + + 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, + 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( + "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) + + with pytest.raises(ValueError, match="reserved"): + local_codemode(CodeModeCapability(reserved)) + + +@pytest.mark.parametrize( + "name", ["not-valid", "two words", "class", "1tool", "_private"] +) +def test_scoped_codemode_rejects_invalid_capability_names(name: str) -> None: + invalid = HarnessTool(name, upper.handler) + + with pytest.raises(ValueError, match="valid Python identifier"): + local_codemode(CodeModeCapability(invalid)) + + +def test_scoped_codemode_rejects_agent_id_input_field() -> None: + class AgentInput(BaseModel): + agent_id: str + + @tool() + async def routed(_: AgentHarness, input_value: AgentInput) -> UpperOutput: + return UpperOutput(value=input_value.agent_id) + + with pytest.raises(ValueError, match="reserved 'agent_id' input field"): + local_codemode(CodeModeCapability(routed)) + + +def test_scoped_codemode_requires_immutable_capabilities() -> None: + with pytest.raises(TypeError, match="immutable tuple"): + create_codemode_tool(capabilities=[]) # type: ignore[arg-type] + + +@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_scoped_codemode_description_includes_capability_contract() -> None: + scoped = local_codemode(CodeModeCapability(upper)) + custom = create_codemode_tool( + capabilities=(CodeModeCapability(upper),), + description="Custom orchestration instructions.", + ) + + assert "Uppercase a value" in scoped.description + assert '"required":["value"]' in scoped.description + assert "Custom orchestration instructions." in custom.description + assert '"required":["value"]' in custom.description + + +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(_: AgentHarness, _input_value: UpperInput) -> UpperOutput: + started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + cancellation_seen.set() + await release.wait() + 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()) + 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)"}) + + 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')"}) + assert result.value == "released" diff --git a/hyperforge/tests/harness_sdk/test_harness.py b/hyperforge/tests/harness_sdk/test_harness.py index beb16545..755f21f0 100644 --- a/hyperforge/tests/harness_sdk/test_harness.py +++ b/hyperforge/tests/harness_sdk/test_harness.py @@ -22,8 +22,10 @@ UsageLimitExceeded, UsageLimits, codemode, + create_codemode_tool, tool, ) +from hyperforge.harness_sdk.execution import current_tool_call_id from hyperforge.harness_sdk.harness import ( EMPTY_RESPONSE_RETRY_PROMPT, AgentResult, @@ -94,6 +96,58 @@ def test_codemode_can_be_registered_explicitly() -> None: assert "codemode" in {tool.name for tool in harness.iter_tools()} +def test_scoped_codemode_is_not_inherited_by_default() -> None: + async def execute(_harness: AgentHarness, value: ToolInput) -> ToolOutput: + return ToolOutput(value=value.value) + + ordinary = HarnessTool("ordinary", execute) + scoped = create_codemode_tool(capabilities=()) + harness = AgentHarness( + model="test-model", + model_client=Model(), + tools=[ordinary, scoped], + ) + + child = harness._create_child("child", include_history=False) + + assert "ordinary" in child._tools + assert "codemode" not in child._tools + + +@pytest.mark.asyncio +async def test_current_tool_call_id_is_isolated_between_parallel_calls() -> None: + ready = asyncio.Event() + seen: dict[str, str | None] = {} + count = 0 + + async def execute(harness: AgentHarness, value: ToolInput) -> ToolOutput: + del harness + nonlocal count + count += 1 + if count == 2: + ready.set() + await ready.wait() + await asyncio.sleep(0) + seen[value.value] = current_tool_call_id() + return ToolOutput(value=value.value) + + first = HarnessTool("first", execute) + second = HarnessTool("second", execute) + harness = AgentHarness( + model="test-model", model_client=Model(), tools=[first, second] + ) + + await harness._execute_tool_calls( + [ + HarnessToolCall(id="first-call", name="first", arguments={"value": "a"}), + HarnessToolCall(id="second-call", name="second", arguments={"value": "b"}), + ] + ) + + assert seen == {"a": "first-call", "b": "second-call"} + assert current_tool_call_id() is None + + def test_turn_loop_clears_pending_tool_result_after_non_empty_response() -> None: state = TurnLoopState() state.record_tool_calls() From 254402d8ad8a6f82d3b36352ee4ae6735cdbefd8 Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Wed, 9 Sep 2026 17:35:31 -0400 Subject: [PATCH 2/7] Fix dropped sandbox frame during callback response write --- hyperforge/src/hyperforge/codemode/sandbox.py | 37 +++++--- hyperforge/tests/codemode/test_protocol.py | 88 +++++++++++++++++++ 2 files changed, 112 insertions(+), 13 deletions(-) diff --git a/hyperforge/src/hyperforge/codemode/sandbox.py b/hyperforge/src/hyperforge/codemode/sandbox.py index ca100f44..fccf1276 100644 --- a/hyperforge/src/hyperforge/codemode/sandbox.py +++ b/hyperforge/src/hyperforge/codemode/sandbox.py @@ -230,24 +230,31 @@ 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}") if not isinstance(msg, SandboxMessage.Request): raise RuntimeError("Unexpected sandbox protocol message") - await self._run_remote_callback(reader, writer, msg.task) + incoming = await self._run_remote_callback(reader, writer, msg.task) finally: await self._cancel_callbacks() + if incoming is not None: + incoming.cancel() + await asyncio.gather(incoming, return_exceptions=True) writer.close() await writer.wait_closed() @@ -256,7 +263,7 @@ async def _run_remote_callback( reader: "SandboxReader", writer: "SandboxWriter", task: RestrictedPythonTask, - ) -> None: + ) -> asyncio.Task[SandboxMessage.AnyMessage]: callback_task = asyncio.create_task(self.callback(task)) self._callback_tasks.add(callback_task) incoming = asyncio.create_task(reader.read_message()) @@ -276,23 +283,27 @@ async def _run_remote_callback( 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: - for candidate in (callback_task, incoming): - if not candidate.done(): - candidate.cancel() - done, still_pending = await asyncio.wait( - {callback_task, incoming}, timeout=CALLBACK_CANCEL_TIMEOUT + if not callback_task.done(): + callback_task.cancel() + callback_done, still_pending = await asyncio.wait( + {callback_task}, timeout=CALLBACK_CANCEL_TIMEOUT ) - for completed in done: + for completed in callback_done: await asyncio.gather(completed, return_exceptions=True) if still_pending: logger.warning("Sandbox callback ignored cancellation") self._callback_tasks.discard(callback_task) for remaining in still_pending: - if remaining is callback_task: - self._track_orphaned_callback(callback_task) - else: - remaining.add_done_callback(_consume_task_result) + self._track_orphaned_callback(remaining) + return incoming def run_in_process(self, request: WorkerExecutionRequest): pipe_worker, pipe_restricted = Pipe() diff --git a/hyperforge/tests/codemode/test_protocol.py b/hyperforge/tests/codemode/test_protocol.py index add0abcf..0928d4f7 100644 --- a/hyperforge/tests/codemode/test_protocol.py +++ b/hyperforge/tests/codemode/test_protocol.py @@ -220,6 +220,94 @@ async def write_message(self, _message): await asyncio.wait_for(callback_done.wait(), timeout=0.2) +@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) + 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) + + +@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, From 329ed7bab72fc5f5f7569eaec45ecd8c9ca43fb3 Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Wed, 9 Sep 2026 17:35:43 -0400 Subject: [PATCH 3/7] Document JSON worker values for local Code Mode execution --- docs/harness-sdk.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index e30f36f3..50e4bf92 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -334,6 +334,12 @@ 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()`. +All Code Mode worker transports, including local pool and isolated-process +execution, now move values over the bounded JSON worker protocol instead of +pickle. Values crossing the worker boundary must be JSON worker values; +previously picklable types such as `bytes`, sets, and dataclasses are rejected +before transport. Convert binary data to base64 strings in tool results. + ### Code Mode Security RestrictedPython reduces the available Python language surface. It is not a From e5ed3a8dfebbd7443304362b9748c706e6753fdc Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Thu, 10 Sep 2026 14:31:09 -0400 Subject: [PATCH 4/7] Reconcile scoped codemode tokens with optional sandbox token --- agents/restricted/tests/test_sandbox.py | 38 ++++++++++++++++--------- docs/harness-sdk.md | 9 +++--- 2 files changed, 30 insertions(+), 17 deletions(-) diff --git a/agents/restricted/tests/test_sandbox.py b/agents/restricted/tests/test_sandbox.py index a3b53cd0..da197d1b 100644 --- a/agents/restricted/tests/test_sandbox.py +++ b/agents/restricted/tests/test_sandbox.py @@ -222,19 +222,31 @@ async def test_sandbox_rejects_wrong_token(sandbox_conn): await rx.read_message() -async def test_remote_runner_requires_token(monkeypatch, tmp_path): - monkeypatch.setattr(sandbox.settings, "sandbox_token", None) - - async def callback(_task): - return None - - runner = sandbox.SandboxRunner.remote(str(tmp_path / "unused.sock"), callback) - request = WorkerExecutionRequest( - code="", question="Q?", local_vars={}, global_vars={}, function_names={} - ) - - with pytest.raises(RuntimeError, match="SANDBOX_TOKEN is required"): - await runner.run(request) +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_sandbox_rejects_concurrent_session( diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index 50e4bf92..abc2da1d 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -351,10 +351,11 @@ 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` is absent and never silently falls back to local execution. -Remote clients and the sandbox service both require a non-empty -`SANDBOX_TOKEN`; the Unix socket ACL is an additional control, not a replacement -for token authentication. `remote_required=False` enables the isolated local +`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. From 16309e9545050e1960a1e5ce8f26b974583fb089 Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Thu, 10 Sep 2026 16:11:18 -0400 Subject: [PATCH 5/7] Replace tool call ContextVar with ToolCallContext threading --- docs/harness-sdk.md | 3 ++- .../src/hyperforge/harness_sdk/execution.py | 24 ------------------- .../src/hyperforge/harness_sdk/harness.py | 14 +++-------- .../hyperforge/harness_sdk/tools/codemode.py | 17 ++++++------- .../src/hyperforge/harness_sdk/tools/core.py | 6 ++--- .../src/hyperforge/harness_sdk/tools/utils.py | 22 ++++++++++++++++- hyperforge/tests/harness_sdk/test_harness.py | 9 +++---- 7 files changed, 39 insertions(+), 56 deletions(-) delete mode 100644 hyperforge/src/hyperforge/harness_sdk/execution.py diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index abc2da1d..567e204f 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -496,7 +496,8 @@ publication never waits for the consumer. If the consumer falls behind by 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. Events emitted inside a tool call also carry its `parent_call_id`; -parallel calls use task-local call context and do not overwrite one another. Use +each call receives its own `ToolCallContext` (use `context.emit()` to attribute +events to the current call). Use `async with agent` when a caller may stop consuming early so the active turn and descendants are cleaned up. diff --git a/hyperforge/src/hyperforge/harness_sdk/execution.py b/hyperforge/src/hyperforge/harness_sdk/execution.py deleted file mode 100644 index 89c82b18..00000000 --- a/hyperforge/src/hyperforge/harness_sdk/execution.py +++ /dev/null @@ -1,24 +0,0 @@ -from contextvars import ContextVar, Token - -_current_tool_call_id: ContextVar[str | None] = ContextVar( - "hyperforge_current_tool_call_id", default=None -) - - -def current_tool_call_id() -> str | None: - return _current_tool_call_id.get() - - -def set_current_tool_call_id(call_id: str) -> Token[str | None]: - return _current_tool_call_id.set(call_id) - - -def reset_current_tool_call_id(token: Token[str | None]) -> None: - _current_tool_call_id.reset(token) - - -__all__ = [ - "current_tool_call_id", - "reset_current_tool_call_id", - "set_current_tool_call_id", -] diff --git a/hyperforge/src/hyperforge/harness_sdk/harness.py b/hyperforge/src/hyperforge/harness_sdk/harness.py index d0560cb0..e4a096e6 100644 --- a/hyperforge/src/hyperforge/harness_sdk/harness.py +++ b/hyperforge/src/hyperforge/harness_sdk/harness.py @@ -14,11 +14,6 @@ from .agents import published_agent_to_tools from .clients import ModelClient, ReasoningEffort from .context import format_context -from .execution import ( - current_tool_call_id, - reset_current_tool_call_id, - set_current_tool_call_id, -) from .models import ( HarnessConversation, HarnessEvent, @@ -592,6 +587,7 @@ async def emit( payload: dict[str, Any], *, persist: bool = True, + parent_call_id: str | None = None, ) -> HarnessEvent: event = HarnessEvent( id=uuid.uuid4().hex, @@ -599,7 +595,7 @@ async def emit( turn_id=self._turn_id, agent_id=self.agent_id, parent_agent_id=self.parent_agent_id, - parent_call_id=current_tool_call_id(), + parent_call_id=parent_call_id, category=self.category, tags=self.tags, metadata=self._persisted_metadata(), @@ -973,11 +969,7 @@ async def _execute_tool_call(self, call: HarnessToolCall) -> HarnessMessage: raise ValueError(f"Tool is not active: {call.name}") else: context = ToolCallContext(harness=self, name=call.name, id=call.id) - token = set_current_tool_call_id(call.id) - try: - output = await tool.execute(context, call.arguments) - finally: - reset_current_tool_call_id(token) + output = await tool.execute(context, call.arguments) result = output.model_dump(mode="json") reference = tool.context(output) content = format_context(reference) diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py index a78e4029..5fe58191 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -21,11 +21,6 @@ from hyperforge.definition import FunctionDefinition from ..context import format_context -from ..execution import ( - current_tool_call_id, - reset_current_tool_call_id, - set_current_tool_call_id, -) from ..models import HarnessEventType from . import HarnessTool, ToolCallContext, ToolInheritancePolicy, tool @@ -210,7 +205,7 @@ async def dispatch(task: RestrictedPythonTask) -> Any: harness._check_limit("max_tool_calls", harness.usage.tool_calls) arguments = _tool_arguments(tool, task.args, task.keyword_args) output = await tool.execute( - ToolCallContext(harness=harness, name=tool.name), arguments + ToolCallContext(harness=harness, name=tool.name, id=context.id), arguments ) return output.model_dump(mode="json") @@ -328,6 +323,7 @@ async def execute( socket, execution_limiter, runner, + parent_call_id=context.id, ) execute.__name__ = name @@ -348,13 +344,14 @@ async def _execute_scoped_codemode( socket: str | None, execution_limiter: CodeModeExecutionLimiter, runner: CodeModeRunner | None = None, + *, + parent_call_id: str | None = None, ) -> CodemodeOutput: capability_map = {capability.tool.name: capability for capability in capabilities} result = CodemodeOutput() output_state = _OutputState() nested_calls = 0 cumulative_result_bytes = 0 - parent_call_id = current_tool_call_id() async def dispatch(task: RestrictedPythonTask) -> Any: nonlocal nested_calls, cumulative_result_bytes @@ -384,8 +381,8 @@ async def dispatch(task: RestrictedPythonTask) -> Any: }, **marker, }, + parent_call_id=parent_call_id, ) - token = set_current_tool_call_id(call_id) started = time.perf_counter() try: harness.usage.tool_calls += 1 @@ -417,7 +414,6 @@ async def dispatch(task: RestrictedPythonTask) -> Any: ) sanitized_result = _sanitize_event_value(normalized) except BaseException as exc: - reset_current_tool_call_id(token) await harness.emit( HarnessEventType.TOOL_FAILED, { @@ -427,6 +423,7 @@ async def dispatch(task: RestrictedPythonTask) -> Any: "duration_ms": round((time.perf_counter() - started) * 1000, 3), **marker, }, + parent_call_id=parent_call_id, ) if isinstance(exc, asyncio.CancelledError): raise @@ -435,7 +432,6 @@ async def dispatch(task: RestrictedPythonTask) -> Any: f"({type(exc).__name__})" ) from None else: - reset_current_tool_call_id(token) await harness.emit( HarnessEventType.TOOL_COMPLETED, { @@ -446,6 +442,7 @@ async def dispatch(task: RestrictedPythonTask) -> Any: "result_bytes": result_bytes, **marker, }, + parent_call_id=parent_call_id, ) return normalized diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/core.py b/hyperforge/src/hyperforge/harness_sdk/tools/core.py index 3b4c995d..6c7bc104 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/core.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/core.py @@ -127,7 +127,7 @@ async def remember(context: ToolCallContext, input_value: RememberInput) -> Dict metadata=harness._persisted_metadata(), ) await harness.storage.remember(memory) - await harness.emit( + await context.emit( HarnessEventType.MEMORY_REMEMBERED, {"memory": memory.model_dump(mode="json")}, ) @@ -148,7 +148,7 @@ async def recall(context: ToolCallContext, input_value: RecallInput) -> ListOutp async def forget(context: ToolCallContext, input_value: ForgetInput) -> DictOutput: harness = context.harness await harness.storage.forget(input_value.id) - await harness.emit(HarnessEventType.MEMORY_FORGOTTEN, {"id": input_value.id}) + await context.emit(HarnessEventType.MEMORY_FORGOTTEN, {"id": input_value.id}) return DictOutput(value={"id": input_value.id}) @@ -180,7 +180,7 @@ async def compact(context: ToolCallContext, input_value: CompactInput) -> DictOu role="user", content=f"Conversation summary:\n{input_value.summary}" ), ] - await harness.emit( + await context.emit( HarnessEventType.COMPACTED, {"messages": [message.model_dump(mode="json") for message in harness.messages]}, ) diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/utils.py b/hyperforge/src/hyperforge/harness_sdk/tools/utils.py index 7c6740bd..3b813327 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: @@ -39,6 +44,21 @@ class ToolCallContext: name: str id: str | None = None + async def emit( + self, + event_type: HarnessEventType, + payload: dict[str, Any], + *, + persist: bool = True, + ) -> "HarnessEvent": + """Emit an event attributed to this tool call.""" + return await self.harness.emit( + event_type, + payload, + persist=persist, + parent_call_id=self.id, + ) + @dataclass(frozen=True) class HarnessTool[InputT: BaseModel, OutputT: BaseModel]: diff --git a/hyperforge/tests/harness_sdk/test_harness.py b/hyperforge/tests/harness_sdk/test_harness.py index 755f21f0..eddbb235 100644 --- a/hyperforge/tests/harness_sdk/test_harness.py +++ b/hyperforge/tests/harness_sdk/test_harness.py @@ -25,7 +25,6 @@ create_codemode_tool, tool, ) -from hyperforge.harness_sdk.execution import current_tool_call_id from hyperforge.harness_sdk.harness import ( EMPTY_RESPONSE_RETRY_PROMPT, AgentResult, @@ -115,20 +114,19 @@ async def execute(_harness: AgentHarness, value: ToolInput) -> ToolOutput: @pytest.mark.asyncio -async def test_current_tool_call_id_is_isolated_between_parallel_calls() -> None: +async def test_parallel_tool_calls_receive_their_own_call_context() -> None: ready = asyncio.Event() seen: dict[str, str | None] = {} count = 0 - async def execute(harness: AgentHarness, value: ToolInput) -> ToolOutput: - del harness + async def execute(context: ToolCallContext, value: ToolInput) -> ToolOutput: nonlocal count count += 1 if count == 2: ready.set() await ready.wait() await asyncio.sleep(0) - seen[value.value] = current_tool_call_id() + seen[value.value] = context.id return ToolOutput(value=value.value) first = HarnessTool("first", execute) @@ -145,7 +143,6 @@ async def execute(harness: AgentHarness, value: ToolInput) -> ToolOutput: ) assert seen == {"a": "first-call", "b": "second-call"} - assert current_tool_call_id() is None def test_turn_loop_clears_pending_tool_result_after_non_empty_response() -> None: From f512fcb960e6a6f2b01b2c5bea8d1d1379a3a20b Mon Sep 17 00:00:00 2001 From: Nathan Van Gheem <524413+vangheem@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:07:33 -0400 Subject: [PATCH 6/7] Preserve restricted agent sandbox compatibility --- agents/restricted/tests/test_sandbox.py | 40 +++++++- docs/harness-sdk.md | 34 ++++--- hyperforge/src/hyperforge/codemode/sandbox.py | 99 ++++++++++++------- hyperforge/src/hyperforge/codemode/worker.py | 19 +++- 4 files changed, 130 insertions(+), 62 deletions(-) diff --git a/agents/restricted/tests/test_sandbox.py b/agents/restricted/tests/test_sandbox.py index da197d1b..f96c00c1 100644 --- a/agents/restricted/tests/test_sandbox.py +++ b/agents/restricted/tests/test_sandbox.py @@ -1,6 +1,7 @@ import asyncio import os import tempfile +from concurrent.futures import ProcessPoolExecutor import pytest from hyperforge.codemode import sandbox @@ -249,6 +250,34 @@ async def callback(_task): 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( sandbox_conn, long_timeout, monkeypatch ): @@ -286,14 +315,15 @@ async def test_sandbox_rejects_concurrent_session( assert isinstance(await first_rx.read_message(), SandboxMessage.Done) -def test_sandbox_settings_default_to_private_optional_metrics() -> None: +def test_sandbox_settings_preserve_legacy_metrics_defaults() -> None: configured = sandbox.SandboxSettings() - assert configured.sandbox_metrics_enabled is False - assert configured.sandbox_metrics_host == "127.0.0.1" + 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_session_runtime_seconds == 60 - assert configured.sandbox_max_session_memory_bytes == 512 * 1024 * 1024 + 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: diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index 567e204f..5dd504e7 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -334,11 +334,12 @@ 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()`. -All Code Mode worker transports, including local pool and isolated-process -execution, now move values over the bounded JSON worker protocol instead of -pickle. Values crossing the worker boundary must be JSON worker values; -previously picklable types such as `bytes`, sets, and dataclasses are rejected -before transport. Convert binary data to base64 strings in tool results. +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. 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. ### Code Mode Security @@ -361,21 +362,22 @@ development environments. Sandbox deployment settings use secure defaults: -- `SANDBOX_MAX_CONCURRENT_SESSIONS=4` bounds authenticated worker sessions and - concurrent authentication handshakes. -- `SANDBOX_MAX_SESSION_RUNTIME_SECONDS=60` and - `SANDBOX_MAX_SESSION_MEMORY_BYTES=536870912` impose server-owned ceilings even - when a client omits limits. The remote client also bounds each connection by - the requested runtime (or the session ceiling when no runtime is requested) - plus `SANDBOX_TIMEOUT_SLACK_SECONDS=10`, so a hung sandbox cannot stall a - caller indefinitely. +- `SANDBOX_MAX_CONCURRENT_SESSIONS` bounds authenticated worker sessions and + concurrent authentication handshakes when configured. +- `SANDBOX_MAX_SESSION_RUNTIME_SECONDS` and + `SANDBOX_MAX_SESSION_MEMORY_BYTES` impose server-owned ceilings when + configured. Existing compatibility deployments leave these unset. 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=false` avoids creating an IP listener. If enabled, - `SANDBOX_METRICS_HOST` defaults to `127.0.0.1` and - `SANDBOX_METRICS_PORT` defaults to `8091`. +- `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. diff --git a/hyperforge/src/hyperforge/codemode/sandbox.py b/hyperforge/src/hyperforge/codemode/sandbox.py index fccf1276..9560d403 100644 --- a/hyperforge/src/hyperforge/codemode/sandbox.py +++ b/hyperforge/src/hyperforge/codemode/sandbox.py @@ -43,14 +43,14 @@ class SandboxSettings(BaseSettings): sandbox_socket_mode: str = "0600" sandbox_socket_group: str | None = None sandbox_metrics_port: int = 8091 - sandbox_metrics_host: str = "127.0.0.1" - sandbox_metrics_enabled: bool = False + sandbox_metrics_host: str = "0.0.0.0" + sandbox_metrics_enabled: bool = True sandbox_token: str | None = None - sandbox_max_concurrent_sessions: int = Field(default=4, gt=0) - sandbox_max_session_runtime_seconds: float = Field( - default=60, gt=0, allow_inf_nan=False + 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 = Field(default=512 * 1024 * 1024, gt=0) + 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 ) @@ -154,9 +154,12 @@ async def run(self, request: WorkerExecutionRequest): runtime = message.run.max_runtime_seconds if runtime is None: runtime = settings.sandbox_max_session_runtime_seconds - async with asyncio.timeout( + timeout = ( runtime + settings.sandbox_timeout_slack_seconds - ): + if runtime is not None + else None + ) + async with asyncio.timeout(timeout): return await self._run_with_remote_admission(message.run) except TimeoutError as exc: raise RuntimeError("Codemode execution timed out") from exc @@ -174,8 +177,10 @@ async def _run_with_remote_admission(self, request: WorkerExecutionRequest): 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)) + worker = PythonAgentWorker(pipe_worker, debug=self.debug, json_protocol=False) + controller_task = asyncio.create_task( + self.background_task(pipe_restricted, json_protocol=False) + ) try: await loop.run_in_executor( self.pool, @@ -324,7 +329,9 @@ def run_in_process(self, request: WorkerExecutionRequest): pipe_worker.close() 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: @@ -334,6 +341,7 @@ async def background_task(self, pipe: Connection) -> bool: pipe, loop, context, + json_protocol, ) finally: pipe.close() @@ -383,14 +391,18 @@ def read_write_queue( pipe: Connection, loop: asyncio.AbstractEventLoop, context: contextvars.Context, + json_protocol: bool = True, ) -> bool: try: while True: - item = decode_protocol_value( - pipe.recv_bytes(MAX_PACKET_BYTES), - "Local sandbox request", - max_bytes=MAX_PACKET_BYTES, - ) + 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": @@ -421,19 +433,22 @@ def copy_result(completed: asyncio.Task[WorkerTypes]) -> None: result = future.result() except Exception as exc: result = WorkerError(error=str(exc)) - 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) + 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 @@ -526,14 +541,18 @@ def __init__( async def run(self, request: WorkerExecutionRequest): memory = settings.sandbox_max_session_memory_bytes - if request.max_memory_bytes is not None: + 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}) self.task, self.process = self.runner.run_in_process(request) try: await self._start_timeout(self.callback_wait_seconds) runtime = settings.sandbox_max_session_runtime_seconds - if request.max_runtime_seconds is not None: + 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 self.task @@ -625,17 +644,19 @@ async def run_sandbox_server( active_sessions = 0 admission_lock = asyncio.Lock() - handshake_limit = asyncio.Semaphore(settings.sandbox_max_concurrent_sessions) + max_sessions = 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.locked(): + if handshake_limit is not None and handshake_limit.locked(): writer.close() await writer.wait_closed() return - await handshake_limit.acquire() + if handshake_limit is not None: + await handshake_limit.acquire() try: async with asyncio.timeout(1): msg = await reader.read_message() @@ -659,10 +680,12 @@ async def handler(rx: asyncio.StreamReader, tx: asyncio.StreamWriter): await writer.wait_closed() return finally: - handshake_limit.release() + if handshake_limit is not None: + handshake_limit.release() async with admission_lock: - if active_sessions >= settings.sandbox_max_concurrent_sessions: + active_limit = settings.sandbox_max_concurrent_sessions + if active_limit is not None and active_sessions >= active_limit: await writer.write_message( SandboxMessage.Error(error="Sandbox concurrency limit reached") ) @@ -712,6 +735,8 @@ def _consume_task_result(task: asyncio.Future[Any]) -> None: async def _acquire_remote_session() -> None: global _remote_active_sessions + if settings.sandbox_max_concurrent_sessions is None: + return while True: with _remote_admission_lock: if _remote_active_sessions < settings.sandbox_max_concurrent_sessions: @@ -722,6 +747,8 @@ async def _acquire_remote_session() -> None: def _release_remote_session() -> None: global _remote_active_sessions + if settings.sandbox_max_concurrent_sessions is None: + return with _remote_admission_lock: if _remote_active_sessions <= 0: raise RuntimeError("Remote sandbox admission counter underflow") diff --git a/hyperforge/src/hyperforge/codemode/worker.py b/hyperforge/src/hyperforge/codemode/worker.py index 96c6bc06..fcc8ce4f 100644 --- a/hyperforge/src/hyperforge/codemode/worker.py +++ b/hyperforge/src/hyperforge/codemode/worker.py @@ -112,10 +112,13 @@ 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 def _process_question_context_sync( self, @@ -241,12 +244,18 @@ def valid_functions_call( return self._receive() def _send(self, task: RestrictedPythonTask) -> None: - self.pipe.send_bytes(encode_protocol_value(task, "Local sandbox request")) + if self.json_protocol: + self.pipe.send_bytes(encode_protocol_value(task, "Local sandbox request")) + else: + self.pipe.send(task) def _receive(self) -> Any: - result = decode_protocol_value( - self.pipe.recv_bytes(MAX_PROTOCOL_BYTES), "Local sandbox response" - ) + 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 From 320de047e2ec59797671f8b37786c0bac6d5b139 Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Mon, 14 Sep 2026 06:37:30 -0400 Subject: [PATCH 7/7] Fix sandbox formatting --- hyperforge/src/hyperforge/codemode/sandbox.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/hyperforge/src/hyperforge/codemode/sandbox.py b/hyperforge/src/hyperforge/codemode/sandbox.py index 9560d403..5024b86e 100644 --- a/hyperforge/src/hyperforge/codemode/sandbox.py +++ b/hyperforge/src/hyperforge/codemode/sandbox.py @@ -645,7 +645,9 @@ async def run_sandbox_server( active_sessions = 0 admission_lock = asyncio.Lock() max_sessions = settings.sandbox_max_concurrent_sessions - handshake_limit = asyncio.Semaphore(max_sessions) if max_sessions is not None else None + 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