From 4fb07614198e8fcd3210dcd2030d3b7df7348138 Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Mon, 14 Sep 2026 13:17:57 -0400 Subject: [PATCH 1/2] Add scoped Code Mode foundation --- docs/harness-sdk.md | 62 ++ hyperforge/CHANGELOG.md | 4 + .../src/hyperforge/harness_sdk/__init__.py | 23 +- .../src/hyperforge/harness_sdk/harness.py | 8 +- .../hyperforge/harness_sdk/tools/__init__.py | 10 +- .../hyperforge/harness_sdk/tools/codemode.py | 290 +++++++++- .../src/hyperforge/harness_sdk/tools/utils.py | 18 +- hyperforge/tests/harness_sdk/test_codemode.py | 528 +++++++++++++++++- hyperforge/tests/harness_sdk/test_harness.py | 22 + 9 files changed, 953 insertions(+), 12 deletions(-) diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index cb6ff05c..cb80d59a 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -190,6 +190,68 @@ 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` when a tool must remain on +the current agent. + +## Scoped Code Mode + +`create_codemode_tool()` exposes generated Python as one model tool while making +only an explicit, immutable tuple of capabilities callable from that Python. +The factory never discovers tools from the harness, so capabilities can remain +hidden from the model's top-level tool list. + +```python +from hyperforge.harness_sdk import ( + AgentHarness, + CodeModeCapability, + create_codemode_tool, +) + +code_mode = create_codemode_tool( + capabilities=(CodeModeCapability(weather),), +) +agent = AgentHarness( + model="your-model", + model_client=model_client, + tools=[code_mode], +) +``` + +The generated tool description includes each capability's description and +argument JSON Schema. Capability calls use normal `HarnessTool` input and output +validation and count toward ordinary harness tool-call usage. The optional +`question` Code Mode input is available to generated code as `question`. + +By default, capability results are projected through the tool's formatted +model-facing context. Supply a custom `result_adapter` for another safe +projection, or opt into the complete JSON-mode Pydantic output with +`raw_codemode_result_adapter`. Every adapter result is normalized through a +strict JSON round trip before generated code receives it, so tuples become +lists and non-JSON or non-finite values are rejected consistently across +runners. Adapters must not return Pydantic models, including models nested in +containers, or dictionaries containing the reserved `__model__` key. +Values passed to scoped `output(...)` use the same normalization and rejection +rules. + +Capability names must be unique ASCII public Python identifiers, cannot use +reserved worker names, and cannot expose an `agent_id` input field. Scoped Code +Mode tools default to `ToolInheritancePolicy.DO_NOT_INHERIT`. If the outer Code +Mode tool is explicitly configured with `INHERIT`, all of its capabilities must +also be inheritable; the factory rejects combinations that would bypass a +capability's `DO_NOT_INHERIT` policy. + +For tests, `runner=` accepts a `CodeModeRunner`; its `run()` method receives the +worker request and capability dispatch callback. Without an injected runner, +the factory uses the existing `SandboxRunner.remote` when `SANDBOX_SOCKET` is +configured and the existing isolated process otherwise. Runtime and memory +behavior continues to use the existing sandbox implementation and +`UsageLimits`; this foundation adds no separate sandbox limits. + +The existing exported `codemode` tool remains available for compatibility and +continues to discover registered harness tools. Prefer the scoped factory when +the generated program should have an explicit capability boundary. + ## Usage Limits All limits are disabled by default. Configure only the limits needed by the diff --git a/hyperforge/CHANGELOG.md b/hyperforge/CHANGELOG.md index dc15d431..b2496f6d 100644 --- a/hyperforge/CHANGELOG.md +++ b/hyperforge/CHANGELOG.md @@ -1,5 +1,9 @@ # CHANGELOG +## Unreleased + +- Added explicit scoped Code Mode capabilities and child-tool inheritance policy. + ## 1.0.0 - Initial version diff --git a/hyperforge/src/hyperforge/harness_sdk/__init__.py b/hyperforge/src/hyperforge/harness_sdk/__init__.py index 36c901db..4d7f2539 100644 --- a/hyperforge/src/hyperforge/harness_sdk/__init__.py +++ b/hyperforge/src/hyperforge/harness_sdk/__init__.py @@ -39,8 +39,19 @@ 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, + CodemodeInput, + CodemodeOutput, + CodeModeResultAdapter, + CodeModeRunner, + codemode, + context_codemode_result_adapter, + create_codemode_tool, + raw_codemode_result_adapter, +) from .usage import HarnessUsage, UsageLimitExceeded, UsageLimits __all__ = [ @@ -56,6 +67,10 @@ "ChatCompletionToolCallDelta", "ChatCompletionToolCallFunctionDelta", "ChatCompletionUsage", + "CodeModeCapability", + "CodeModeDispatch", + "CodeModeResultAdapter", + "CodeModeRunner", "CodemodeInput", "CodemodeOutput", "HarnessAgentToTool", @@ -79,14 +94,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/harness.py b/hyperforge/src/hyperforge/harness_sdk/harness.py index ba252b36..9dd06a9a 100644 --- a/hyperforge/src/hyperforge/harness_sdk/harness.py +++ b/hyperforge/src/hyperforge/harness_sdk/harness.py @@ -23,7 +23,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 @@ -1036,7 +1036,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/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..8a3a3bc4 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -1,4 +1,9 @@ -from typing import Any +import inspect +import json +import keyword +import uuid +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Protocol from pydantic import BaseModel @@ -10,20 +15,76 @@ from hyperforge.codemode.sandbox import settings from hyperforge.definition import FunctionDefinition -from . import HarnessTool, ToolCallContext, tool +from ..context import format_context +from . import HarnessTool, ToolCallContext, ToolInheritancePolicy, tool CODEMODE_TOOL_NAME = "codemode" OUTPUT_FUNCTION_NAME = "output" +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", + } +) +type CodeModeResultAdapter[OutputT: BaseModel] = Callable[ + [HarnessTool[Any, OutputT], OutputT], 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[OutputT: BaseModel]( + capability: HarnessTool[Any, OutputT], output: OutputT +) -> Any: + """Project a result to the same formatted context shown to the model.""" + return format_context(capability.context(output)) + + +def raw_codemode_result_adapter[OutputT: BaseModel]( + _capability: HarnessTool[Any, OutputT], output: OutputT +) -> Any: + """Explicitly expose the capability's complete JSON-mode output.""" + return output.model_dump(mode="json") + + +class CodeModeRunner(Protocol): + async def run( + self, request: WorkerExecutionRequest, dispatch: CodeModeDispatch + ) -> Any: ... + + +@dataclass(frozen=True) +class CodeModeCapability[OutputT: BaseModel]: + tool: HarnessTool[Any, OutputT] + result_adapter: CodeModeResultAdapter[OutputT] = context_codemode_result_adapter + + @tool( name=CODEMODE_TOOL_NAME, description=( @@ -83,6 +144,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 +155,197 @@ async def dispatch(task: RestrictedPythonTask) -> Any: return result +def create_codemode_tool( + *, + capabilities: tuple[CodeModeCapability[Any], ...], + runner: CodeModeRunner | None = None, + inheritance: ToolInheritancePolicy = ToolInheritancePolicy.DO_NOT_INHERIT, + name: str = CODEMODE_TOOL_NAME, + description: str | None = None, +) -> HarnessTool[CodemodeInput, CodemodeOutput]: + """Create a Code Mode tool from an explicit immutable capability set.""" + if not isinstance(capabilities, tuple): + raise TypeError("capabilities must be an immutable tuple") + capability_map: dict[str, CodeModeCapability[Any]] = {} + 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.isascii() + or not capability_name.isidentifier() + or keyword.iskeyword(capability_name) + or capability_name.startswith("_") + ): + raise ValueError( + f"Code Mode capability name {capability_name!r} must be an ASCII " + "public 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 + if inheritance == ToolInheritancePolicy.INHERIT: + non_inherited = sorted( + capability.tool.name + for capability in capabilities + if capability.tool.inheritance == ToolInheritancePolicy.DO_NOT_INHERIT + ) + if non_inherited: + raise ValueError( + "An inherited Code Mode tool cannot include capabilities marked " + f"DO_NOT_INHERIT: {', '.join(non_inherited)}" + ) + + async def execute( + context: ToolCallContext, input_value: CodemodeInput + ) -> CodemodeOutput: + harness = context.harness + result = CodemodeOutput() + + async def dispatch(task: RestrictedPythonTask) -> Any: + if task.function == OUTPUT_FUNCTION_NAME: + 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") + value = task.args[0] if task.args else task.keyword_args.get("value") + result.value = _normalize_json_value( + value, + label="Code Mode output", + ) + return None + + capability = capability_map.get(task.function) + if capability is None: + raise ValueError(f"Unknown Code Mode capability: {task.function}") + harness.usage.tool_calls += 1 + harness._check_limit("max_tool_calls", harness.usage.tool_calls) + arguments = _scoped_tool_arguments( + capability.tool, task.args, task.keyword_args + ) + output = await capability.tool.execute( + ToolCallContext( + harness=harness, + name=capability.tool.name, + id=uuid.uuid4().hex, + ), + arguments, + ) + projected = capability.result_adapter(capability.tool, output) + if inspect.isawaitable(projected): + projected = await projected + return _normalize_json_value( + projected, + label="Code Mode result adapter", + ) + + 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": {}}, + ) + 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(request, dispatch) + else: + sandbox_runner = ( + SandboxRunner.remote(settings.sandbox_socket, dispatch) + if settings.sandbox_socket is not None + else SandboxRunner.isolated_process(dispatch) + ) + await sandbox_runner.run(request) + return result + + execute.__name__ = name + return HarnessTool( + name=name, + handler=execute, + description=_scoped_description(capabilities, description), + inheritance=inheritance, + ) + + +def _normalize_json_value(value: Any, *, label: str) -> Any: + _reject_reserved_json_values(value, label=label) + try: + return json.loads(json.dumps(value, allow_nan=False)) + except TypeError as exc: + raise TypeError(f"{label} must be JSON-serializable") from exc + except ValueError as exc: + raise ValueError(f"{label} must be a strict JSON value: {exc}") from exc + + +def _reject_reserved_json_values(value: Any, *, label: str) -> None: + if isinstance(value, BaseModel): + raise TypeError(f"{label} must not contain Pydantic model values") + if isinstance(value, dict): + if "__model__" in value: + raise ValueError(f"{label} cannot contain the reserved '__model__' key") + for key, nested in value.items(): + _reject_reserved_json_values(key, label=label) + _reject_reserved_json_values(nested, label=label) + elif isinstance(value, (list, tuple)): + for nested in value: + _reject_reserved_json_values(nested, label=label) + + +def _scoped_description( + capabilities: tuple[CodeModeCapability[Any], ...], description: str | None +) -> str: + introduction = description or ( + "Execute restricted Python code using only the scoped capabilities below; " + "call output(value) 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}: " + f"{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 +365,35 @@ def _tool_arguments( return arguments -__all__ = ["CodemodeInput", "CodemodeOutput", "codemode"] +def _scoped_tool_arguments( + tool: HarnessTool[Any, Any], + args: tuple[Any, ...], + keyword_args: dict[str, Any], +) -> dict[str, Any]: + if not args: + return keyword_args + names = list(tool.parameters.get("properties", {})) + if len(args) > len(names): + raise ValueError(f"Too many positional arguments for {tool.name}") + arguments = dict(zip(names, args, strict=False)) + duplicates = arguments.keys() & keyword_args.keys() + if duplicates: + duplicate = next(iter(duplicates)) + raise ValueError(f"Multiple values for argument {duplicate!r}") + arguments.update(keyword_args) + return arguments + + +__all__ = [ + "CodeModeCapability", + "CodeModeDispatch", + "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/tests/harness_sdk/test_codemode.py b/hyperforge/tests/harness_sdk/test_codemode.py index 007b4cda..dc3898f6 100644 --- a/hyperforge/tests/harness_sdk/test_codemode.py +++ b/hyperforge/tests/harness_sdk/test_codemode.py @@ -1,19 +1,31 @@ import sys from collections.abc import AsyncIterator from contextvars import ContextVar +from typing import Any import pytest -from pydantic import BaseModel +from pydantic import BaseModel, Field, ValidationError +from hyperforge.codemode import RestrictedPythonTask, WorkerExecutionRequest from hyperforge.harness_sdk import ( AgentHarness, + CodeModeCapability, + CodeModeDispatch, CodemodeInput, + HarnessContextReference, + HarnessContextType, + HarnessTool, + HarnessToolCall, ModelDelta, ToolCallContext, + ToolInheritancePolicy, UsageLimits, codemode, + create_codemode_tool, + raw_codemode_result_adapter, tool, ) +from hyperforge.harness_sdk.tools import codemode as codemode_module class UpperInput(BaseModel): @@ -24,17 +36,56 @@ 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()) +@tool( + description="Look up a public value", + context_factory=lambda output: HarnessContextReference( + type=HarnessContextType.STRUCTURED, + content={"value": output.value}, + ), +) +async def lookup(_: ToolCallContext, input_value: UpperInput) -> SensitiveOutput: + return SensitiveOutput(value=input_value.value.upper(), secret="internal") + + class UnusedModel: async def stream(self, **kwargs) -> AsyncIterator[ModelDelta]: del kwargs yield ModelDelta(text="unused") +class RecordingRunner: + def __init__(self, tasks: list[RestrictedPythonTask]) -> None: + self.tasks = tasks + self.requests: list[WorkerExecutionRequest] = [] + self.results: list[tuple[str, Any]] = [] + + async def run( + self, request: WorkerExecutionRequest, dispatch: CodeModeDispatch + ) -> None: + self.requests.append(request) + for task in self.tasks: + self.results.append((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_codemode_calls_registered_tools_and_returns_output() -> None: harness = AgentHarness( @@ -122,9 +173,13 @@ async def test_codemode_supports_common_aggregation_patterns( @pytest.mark.asyncio async def test_codemode_preserves_context_when_calling_tools() -> None: request_context = ContextVar("request_context", default="missing") + call_ids: list[str | None] = [] @tool() - async def read_context(_: ToolCallContext, _input_value: UpperInput) -> UpperOutput: + async def read_context( + context: ToolCallContext, _input_value: UpperInput + ) -> UpperOutput: + call_ids.append(context.id) return UpperOutput(value=request_context.get()) harness = AgentHarness( @@ -133,7 +188,7 @@ async def read_context(_: ToolCallContext, _input_value: UpperInput) -> UpperOut token = request_context.set("available") try: result = await codemode.execute( - ToolCallContext(harness=harness, name=codemode.name), + ToolCallContext(harness=harness, name=codemode.name, id="outer-call"), CodemodeInput( code="result = read_context(value='unused')\noutput(result['value'])" ).model_dump(), @@ -142,6 +197,7 @@ async def read_context(_: ToolCallContext, _input_value: UpperInput) -> UpperOut request_context.reset(token) assert result.value == "available" + assert call_ids == [None] @pytest.mark.asyncio @@ -184,3 +240,469 @@ 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: + seen_call_ids: list[str | None] = [] + + @tool(name="hidden_upper") + async def hidden_upper( + context: ToolCallContext, input_value: UpperInput + ) -> UpperOutput: + seen_call_ids.append(context.id) + return UpperOutput(value=input_value.value.upper()) + + @tool() + async def lower(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + return UpperOutput(value=input_value.value.lower()) + + runner = RecordingRunner( + [worker_task("hidden_upper", "hello"), worker_task("output", "done")] + ) + scoped = create_codemode_tool( + capabilities=( + CodeModeCapability( + hidden_upper, + result_adapter=raw_codemode_result_adapter, + ), + ), + runner=runner, + ) + 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( + ToolCallContext(harness=harness, name=scoped.name, id="outer-call"), + {"code": "hidden_upper('hello'); output('done')", "question": "Why?"}, + ) + + assert result.value == "done" + assert len(seen_call_ids) == 1 + assert seen_call_ids[0] is not None + assert seen_call_ids[0] != "outer-call" + assert hidden_upper.name not in harness._tools + assert set(runner.requests[0].function_names["harness"]) == { + "hidden_upper", + "output", + } + assert runner.requests[0].code == "hidden_upper('hello'); output('done')" + assert runner.requests[0].question == "Why?" + assert runner.results[0] == ("hidden_upper", {"value": "HELLO"}) + + +@pytest.mark.asyncio +async def test_scoped_codemode_runs_hidden_capability_in_isolated_process( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(codemode_module.settings, "sandbox_socket", None) + + @tool() + async def lower(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: + return UpperOutput(value=input_value.value.lower()) + + scoped = create_codemode_tool( + capabilities=( + CodeModeCapability( + upper, + result_adapter=raw_codemode_result_adapter, + ), + ) + ) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + tools=[scoped, lower], + ) + + result = await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + { + "code": "result = upper(value='hello')\noutput(result['value'])", + }, + ) + + assert result.value == "HELLO" + assert upper.name not in harness._tools + with pytest.raises(RuntimeError, match="lower.*not defined"): + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "lower(value='HELLO')"}, + ) + + +@pytest.mark.asyncio +async def test_scoped_codemode_default_projection_uses_formatted_context() -> None: + runner = RecordingRunner([worker_task("lookup", value="hello")]) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(lookup),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "lookup(value='hello')"}, + ) + + assert runner.results == [("lookup", '{"value":"HELLO"}')] + + +@pytest.mark.asyncio +async def test_scoped_codemode_raw_projection_must_be_explicit() -> None: + runner = RecordingRunner([worker_task("lookup", value="hello")]) + scoped = create_codemode_tool( + capabilities=( + CodeModeCapability( + lookup, + result_adapter=raw_codemode_result_adapter, + ), + ), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "lookup(value='hello')"}, + ) + + assert runner.results == [("lookup", {"value": "HELLO", "secret": "internal"})] + + +@pytest.mark.asyncio +async def test_scoped_codemode_rejects_nested_unprojected_models() -> None: + def wrap_model( + _tool: HarnessTool[Any, UpperOutput], output: UpperOutput + ) -> dict[str, BaseModel]: + return {"wrapped": output} + + runner = RecordingRunner([worker_task("upper", value="hello")]) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(upper, result_adapter=wrap_model),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(TypeError, match="must not contain Pydantic model values"): + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "upper(value='hello')"}, + ) + + +@pytest.mark.asyncio +async def test_scoped_codemode_normalizes_adapter_results_as_strict_json() -> None: + def project( + _tool: HarnessTool[Any, UpperOutput], output: UpperOutput + ) -> dict[str, Any]: + return {"value": output.value, "positions": (1, 2)} + + runner = RecordingRunner([worker_task("upper", value="hello")]) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(upper, result_adapter=project),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "upper(value='hello')"}, + ) + + assert runner.results == [("upper", {"value": "HELLO", "positions": [1, 2]})] + + +@pytest.mark.asyncio +async def test_scoped_codemode_rejects_reserved_model_marker() -> None: + def project( + _tool: HarnessTool[Any, UpperOutput], _output: UpperOutput + ) -> dict[str, Any]: + return {"wrapped": {"__model__": "Context"}} + + runner = RecordingRunner([worker_task("upper", value="hello")]) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(upper, result_adapter=project),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match="reserved '__model__' key"): + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "upper(value='hello')"}, + ) + + +@pytest.mark.asyncio +async def test_scoped_codemode_rejects_non_finite_adapter_result() -> None: + def project( + _tool: HarnessTool[Any, UpperOutput], _output: UpperOutput + ) -> dict[str, float]: + return {"score": float("nan")} + + runner = RecordingRunner([worker_task("upper", value="hello")]) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(upper, result_adapter=project),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match="strict JSON value"): + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "upper(value='hello')"}, + ) + + +@pytest.mark.asyncio +async def test_scoped_codemode_normalizes_output_as_strict_json() -> None: + runner = RecordingRunner( + [worker_task("output", {"positions": (1, 2), "complete": True})] + ) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + result = await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "output({'positions': (1, 2), 'complete': True})"}, + ) + + assert result.value == {"positions": [1, 2], "complete": True} + + +@pytest.mark.asyncio +async def test_scoped_codemode_rejects_reserved_model_marker_in_output() -> None: + runner = RecordingRunner( + [worker_task("output", {"wrapped": {"__model__": "Context"}})] + ) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises( + ValueError, + match="Code Mode output cannot contain the reserved '__model__' key", + ): + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "output({'wrapped': {'__model__': 'Context'}})"}, + ) + + +@pytest.mark.asyncio +async def test_scoped_codemode_uses_normal_tool_validation() -> None: + invalid_input_runner = RecordingRunner([worker_task("upper", missing="value")]) + input_scoped = create_codemode_tool( + capabilities=(CodeModeCapability(upper),), + runner=invalid_input_runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match="Invalid upper arguments"): + await input_scoped.execute( + ToolCallContext(harness=harness, name=input_scoped.name), + {"code": "upper(missing='value')"}, + ) + + @tool() + async def invalid_output( + _context: ToolCallContext, input_value: UpperInput + ) -> UpperOutput: + return UpperInput(value=input_value.value) # type: ignore[return-value] + + invalid_output_runner = RecordingRunner( + [worker_task("invalid_output", value="hello")] + ) + output_scoped = create_codemode_tool( + capabilities=(CodeModeCapability(invalid_output),), + runner=invalid_output_runner, + ) + with pytest.raises(ValidationError): + await output_scoped.execute( + ToolCallContext(harness=harness, name=output_scoped.name), + {"code": "invalid_output(value='hello')"}, + ) + + +@pytest.mark.asyncio +async def test_scoped_positional_arguments_follow_advertised_alias_order() -> None: + class AliasedInput(BaseModel): + value: str = Field(alias="public") + + @tool(description="Use an aliased public argument") + async def aliased( + _context: ToolCallContext, input_value: AliasedInput + ) -> UpperOutput: + return UpperOutput(value=input_value.value.upper()) + + runner = RecordingRunner([worker_task("aliased", "hello")]) + scoped = create_codemode_tool( + capabilities=( + CodeModeCapability( + aliased, + result_adapter=raw_codemode_result_adapter, + ), + ), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "aliased('hello')"}, + ) + + definition = runner.requests[0].function_names["harness"]["aliased"] + assert list(definition.parameters) == ["public"] + assert '"public"' in scoped.description + assert runner.results == [("aliased", {"value": "HELLO"})] + + +@pytest.mark.asyncio +async def test_scoped_codemode_counts_nested_calls_as_ordinary_usage() -> None: + runner = RecordingRunner( + [worker_task("upper", value="one"), worker_task("upper", value="two")] + ) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(upper),), + runner=runner, + ) + harness = AgentHarness( + model="test", + model_client=UnusedModel(), + tools=[scoped], + usage_limits=UsageLimits(max_tool_calls=3), + ) + + await harness._execute_tool_calls( + [ + HarnessToolCall( + id="outer-call", + name=scoped.name, + arguments={"code": "upper(value='one'); upper(value='two')"}, + ) + ] + ) + + assert harness.usage.tool_calls == 3 + + +@pytest.mark.asyncio +async def test_scoped_codemode_assigns_unique_nested_call_ids() -> None: + call_ids: list[str | None] = [] + + @tool() + async def record_id( + context: ToolCallContext, input_value: UpperInput + ) -> UpperOutput: + call_ids.append(context.id) + return UpperOutput(value=input_value.value) + + runner = RecordingRunner( + [worker_task("record_id", value="one"), worker_task("record_id", value="two")] + ) + scoped = create_codemode_tool( + capabilities=(CodeModeCapability(record_id),), + runner=runner, + ) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name, id="outer-call"), + {"code": "record_id(value='one'); record_id(value='two')"}, + ) + + assert len(call_ids) == 2 + assert all(call_id is not None for call_id in call_ids) + assert len(set(call_ids)) == 2 + assert "outer-call" not in call_ids + + +@pytest.mark.parametrize("name", ["codemode", "output", "question", "agent_id"]) +def test_scoped_codemode_rejects_reserved_capability_names(name: str) -> None: + reserved = HarnessTool(name, upper.handler) + + with pytest.raises(ValueError, match="reserved"): + create_codemode_tool(capabilities=(CodeModeCapability(reserved),)) + + +@pytest.mark.parametrize( + "name", + [ + "not-valid", + "two words", + "class", + "1tool", + "_private", + "café", + "ｏｕｔｐｕｔ", + ], +) +def test_scoped_codemode_rejects_invalid_capability_names(name: str) -> None: + invalid = HarnessTool(name, upper.handler) + + with pytest.raises(ValueError, match="ASCII public Python identifier"): + create_codemode_tool(capabilities=(CodeModeCapability(invalid),)) + + +def test_scoped_codemode_rejects_duplicate_capability_names() -> None: + duplicate = HarnessTool(upper.name, upper.handler) + + with pytest.raises(ValueError, match="must be unique"): + create_codemode_tool( + capabilities=(CodeModeCapability(upper), CodeModeCapability(duplicate)) + ) + + +def test_scoped_codemode_rejects_agent_id_input_field() -> None: + class AgentInput(BaseModel): + agent_id: str + + @tool() + async def routed(_context: ToolCallContext, input_value: AgentInput) -> UpperOutput: + return UpperOutput(value=input_value.agent_id) + + with pytest.raises(ValueError, match="reserved 'agent_id' input field"): + create_codemode_tool(capabilities=(CodeModeCapability(routed),)) + + +def test_scoped_codemode_requires_immutable_capability_tuple() -> None: + with pytest.raises(TypeError, match="immutable tuple"): + create_codemode_tool(capabilities=[]) # type: ignore[arg-type] + + +def test_inherited_codemode_rejects_non_inherited_capability() -> None: + @tool(inheritance=ToolInheritancePolicy.DO_NOT_INHERIT) + async def parent_only( + _context: ToolCallContext, input_value: UpperInput + ) -> UpperOutput: + return UpperOutput(value=input_value.value) + + with pytest.raises(ValueError, match="cannot include.*DO_NOT_INHERIT"): + create_codemode_tool( + capabilities=(CodeModeCapability(parent_only),), + inheritance=ToolInheritancePolicy.INHERIT, + ) + + +def test_scoped_codemode_description_includes_capability_schema() -> None: + scoped = create_codemode_tool(capabilities=(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 diff --git a/hyperforge/tests/harness_sdk/test_harness.py b/hyperforge/tests/harness_sdk/test_harness.py index beb16545..80c0f621 100644 --- a/hyperforge/tests/harness_sdk/test_harness.py +++ b/hyperforge/tests/harness_sdk/test_harness.py @@ -19,9 +19,11 @@ ModelDelta, NucliaModelClient, ToolCallContext, + ToolInheritancePolicy, UsageLimitExceeded, UsageLimits, codemode, + create_codemode_tool, tool, ) from hyperforge.harness_sdk.harness import ( @@ -94,6 +96,26 @@ 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(_context: ToolCallContext, 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.inheritance == ToolInheritancePolicy.INHERIT + assert scoped.inheritance == ToolInheritancePolicy.DO_NOT_INHERIT + assert "ordinary" in child._tools + assert "codemode" not in child._tools + + def test_turn_loop_clears_pending_tool_result_after_non_empty_response() -> None: state = TurnLoopState() state.record_tool_calls() From d75f099787ea3d9da2c604fd1ab301fd0985f63d Mon Sep 17 00:00:00 2001 From: Sekou Kone Date: Tue, 15 Sep 2026 09:33:46 -0400 Subject: [PATCH 2/2] Enforce scoped Code Mode output contract --- docs/harness-sdk.md | 3 +- .../hyperforge/harness_sdk/tools/codemode.py | 15 +++- hyperforge/tests/harness_sdk/test_codemode.py | 86 +++++++++++++++---- 3 files changed, 85 insertions(+), 19 deletions(-) diff --git a/docs/harness-sdk.md b/docs/harness-sdk.md index cb80d59a..a1b06f18 100644 --- a/docs/harness-sdk.md +++ b/docs/harness-sdk.md @@ -232,7 +232,8 @@ lists and non-JSON or non-finite values are rejected consistently across runners. Adapters must not return Pydantic models, including models nested in containers, or dictionaries containing the reserved `__model__` key. Values passed to scoped `output(...)` use the same normalization and rejection -rules. +rules. Generated code must call `output(value)` exactly once; a missing or +repeated call fails the invocation. Capability names must be unique ASCII public Python identifiers, cannot use reserved worker names, and cannot expose an `agent_id` input field. Scoped Code diff --git a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py index 8a3a3bc4..1c8352e5 100644 --- a/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py +++ b/hyperforge/src/hyperforge/harness_sdk/tools/codemode.py @@ -215,20 +215,29 @@ async def execute( ) -> CodemodeOutput: harness = context.harness result = CodemodeOutput() + output_calls = 0 + output_succeeded = False async def dispatch(task: RestrictedPythonTask) -> Any: + nonlocal output_calls, output_succeeded if task.function == OUTPUT_FUNCTION_NAME: + output_calls += 1 + if output_calls > 1: + raise ValueError("output(value) may only be called once") if task.args and task.keyword_args: raise ValueError( "output accepts either a positional or keyword value" ) if len(task.args) > 1: raise ValueError("output accepts one value") - value = task.args[0] if task.args else task.keyword_args.get("value") + if not task.args and set(task.keyword_args) != {"value"}: + raise ValueError("output requires exactly one 'value' argument") + value = task.args[0] if task.args else task.keyword_args["value"] result.value = _normalize_json_value( value, label="Code Mode output", ) + output_succeeded = True return None capability = capability_map.get(task.function) @@ -286,6 +295,8 @@ async def dispatch(task: RestrictedPythonTask) -> Any: else SandboxRunner.isolated_process(dispatch) ) await sandbox_runner.run(request) + if output_calls != 1 or not output_succeeded: + raise ValueError("Code Mode must call output(value) exactly once") return result execute.__name__ = name @@ -326,7 +337,7 @@ def _scoped_description( ) -> str: introduction = description or ( "Execute restricted Python code using only the scoped capabilities below; " - "call output(value) to return a result." + "call output(value) exactly once to return a result." ) if not capabilities: return introduction diff --git a/hyperforge/tests/harness_sdk/test_codemode.py b/hyperforge/tests/harness_sdk/test_codemode.py index dc3898f6..f44699ed 100644 --- a/hyperforge/tests/harness_sdk/test_codemode.py +++ b/hyperforge/tests/harness_sdk/test_codemode.py @@ -343,7 +343,9 @@ async def lower(_: ToolCallContext, input_value: UpperInput) -> UpperOutput: @pytest.mark.asyncio async def test_scoped_codemode_default_projection_uses_formatted_context() -> None: - runner = RecordingRunner([worker_task("lookup", value="hello")]) + runner = RecordingRunner( + [worker_task("lookup", value="hello"), worker_task("output", "done")] + ) scoped = create_codemode_tool( capabilities=(CodeModeCapability(lookup),), runner=runner, @@ -352,15 +354,17 @@ async def test_scoped_codemode_default_projection_uses_formatted_context() -> No await scoped.execute( ToolCallContext(harness=harness, name=scoped.name), - {"code": "lookup(value='hello')"}, + {"code": "lookup(value='hello'); output('done')"}, ) - assert runner.results == [("lookup", '{"value":"HELLO"}')] + assert runner.results == [("lookup", '{"value":"HELLO"}'), ("output", None)] @pytest.mark.asyncio async def test_scoped_codemode_raw_projection_must_be_explicit() -> None: - runner = RecordingRunner([worker_task("lookup", value="hello")]) + runner = RecordingRunner( + [worker_task("lookup", value="hello"), worker_task("output", "done")] + ) scoped = create_codemode_tool( capabilities=( CodeModeCapability( @@ -374,10 +378,13 @@ async def test_scoped_codemode_raw_projection_must_be_explicit() -> None: await scoped.execute( ToolCallContext(harness=harness, name=scoped.name), - {"code": "lookup(value='hello')"}, + {"code": "lookup(value='hello'); output('done')"}, ) - assert runner.results == [("lookup", {"value": "HELLO", "secret": "internal"})] + assert runner.results == [ + ("lookup", {"value": "HELLO", "secret": "internal"}), + ("output", None), + ] @pytest.mark.asyncio @@ -408,7 +415,9 @@ def project( ) -> dict[str, Any]: return {"value": output.value, "positions": (1, 2)} - runner = RecordingRunner([worker_task("upper", value="hello")]) + runner = RecordingRunner( + [worker_task("upper", value="hello"), worker_task("output", "done")] + ) scoped = create_codemode_tool( capabilities=(CodeModeCapability(upper, result_adapter=project),), runner=runner, @@ -417,10 +426,13 @@ def project( await scoped.execute( ToolCallContext(harness=harness, name=scoped.name), - {"code": "upper(value='hello')"}, + {"code": "upper(value='hello'); output('done')"}, ) - assert runner.results == [("upper", {"value": "HELLO", "positions": [1, 2]})] + assert runner.results == [ + ("upper", {"value": "HELLO", "positions": [1, 2]}), + ("output", None), + ] @pytest.mark.asyncio @@ -481,6 +493,33 @@ async def test_scoped_codemode_normalizes_output_as_strict_json() -> None: assert result.value == {"positions": [1, 2], "complete": True} +@pytest.mark.asyncio +async def test_scoped_codemode_requires_output() -> None: + scoped = create_codemode_tool(capabilities=(), runner=RecordingRunner([])) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match=r"must call output\(value\) exactly once"): + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "pass"}, + ) + + +@pytest.mark.asyncio +async def test_scoped_codemode_rejects_repeated_output() -> None: + runner = RecordingRunner( + [worker_task("output", "first"), worker_task("output", "second")] + ) + scoped = create_codemode_tool(capabilities=(), runner=runner) + harness = AgentHarness(model="test", model_client=UnusedModel()) + + with pytest.raises(ValueError, match=r"output\(value\) may only be called once"): + await scoped.execute( + ToolCallContext(harness=harness, name=scoped.name), + {"code": "output('first'); output('second')"}, + ) + + @pytest.mark.asyncio async def test_scoped_codemode_rejects_reserved_model_marker_in_output() -> None: runner = RecordingRunner( @@ -545,7 +584,9 @@ async def aliased( ) -> UpperOutput: return UpperOutput(value=input_value.value.upper()) - runner = RecordingRunner([worker_task("aliased", "hello")]) + runner = RecordingRunner( + [worker_task("aliased", "hello"), worker_task("output", "done")] + ) scoped = create_codemode_tool( capabilities=( CodeModeCapability( @@ -559,19 +600,26 @@ async def aliased( await scoped.execute( ToolCallContext(harness=harness, name=scoped.name), - {"code": "aliased('hello')"}, + {"code": "aliased('hello'); output('done')"}, ) definition = runner.requests[0].function_names["harness"]["aliased"] assert list(definition.parameters) == ["public"] assert '"public"' in scoped.description - assert runner.results == [("aliased", {"value": "HELLO"})] + assert runner.results == [ + ("aliased", {"value": "HELLO"}), + ("output", None), + ] @pytest.mark.asyncio async def test_scoped_codemode_counts_nested_calls_as_ordinary_usage() -> None: runner = RecordingRunner( - [worker_task("upper", value="one"), worker_task("upper", value="two")] + [ + worker_task("upper", value="one"), + worker_task("upper", value="two"), + worker_task("output", "done"), + ] ) scoped = create_codemode_tool( capabilities=(CodeModeCapability(upper),), @@ -589,7 +637,9 @@ async def test_scoped_codemode_counts_nested_calls_as_ordinary_usage() -> None: HarnessToolCall( id="outer-call", name=scoped.name, - arguments={"code": "upper(value='one'); upper(value='two')"}, + arguments={ + "code": "upper(value='one'); upper(value='two'); output('done')" + }, ) ] ) @@ -609,7 +659,11 @@ async def record_id( return UpperOutput(value=input_value.value) runner = RecordingRunner( - [worker_task("record_id", value="one"), worker_task("record_id", value="two")] + [ + worker_task("record_id", value="one"), + worker_task("record_id", value="two"), + worker_task("output", "done"), + ] ) scoped = create_codemode_tool( capabilities=(CodeModeCapability(record_id),), @@ -619,7 +673,7 @@ async def record_id( await scoped.execute( ToolCallContext(harness=harness, name=scoped.name, id="outer-call"), - {"code": "record_id(value='one'); record_id(value='two')"}, + {"code": "record_id(value='one'); record_id(value='two'); output('done')"}, ) assert len(call_ids) == 2