openreward 0.1.161.dev0__tar.gz → 0.1.161.dev1__tar.gz

{openreward-0.1.161.dev0 → openreward-0.1.161.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openreward
-Version: 0.1.161.dev0
+Version: 0.1.161.dev1
 Summary: Python SDK for the OpenReward platform.
 Author-email: GR Inc <hello@gr.inc>
 License-Expression: MIT

{openreward-0.1.161.dev0 → openreward-0.1.161.dev1}/openreward/environments/_sandbox_tools.py

@@ -137,14 +137,29 @@ async def _bash_impl(
     timeout: Optional[float] = None,
     *,
     root: Optional[str],
+    max_bytes: Optional[int] = None,
 ) -> ToolOutput:
     # Omit rather than forward None, so sandbox.run applies its own default.
-    kwargs = {} if timeout is None else {"timeout": timeout}
+    kwargs: Dict[str, Any] = {} if timeout is None else {"timeout": timeout}
+    if max_bytes is not None:
+        kwargs["max_bytes"] = max_bytes
     try:
-        output, code = await sandbox.run(_rooted_command(command.strip(), root), **kwargs)
+        result = await sandbox.run(_rooted_command(command.strip(), root), **kwargs)
+        output, code = result
+        # A killed command otherwise reads as one that printed nothing and exited 124.
+        notes = []
+        metadata: Dict[str, Any] = {"output": output, "exit_code": code}
+        if getattr(result, "truncated", False):
+            notes.append("output exceeded the limit and was cut off")
+            metadata["truncated"] = True
+        if getattr(result, "timed_out", False):
+            limit = f"{timeout:g}s " if timeout is not None else ""
+            notes.append(f"command hit the {limit}time limit and was killed")
+            metadata["timed_out"] = True
+        suffix = f"\n\n[{'; '.join(notes)}]" if notes else ""
         return ToolOutput(
-            blocks=[TextBlock(text=f"{output}\n\n(exit {code})")],
-            metadata={"output": output, "exit_code": code},
+            blocks=[TextBlock(text=f"{output}\n\n(exit {code}){suffix}")],
+            metadata=metadata,
             finished=False,
         )
     except Exception as e:

{openreward-0.1.161.dev0 → openreward-0.1.161.dev1}/openreward/environments/environment.py

@@ -3,7 +3,8 @@ import functools
 import inspect
 import time
 from abc import ABC, abstractmethod
-from typing import Any, Awaitable, Callable, ClassVar, Optional, Sequence, TypeVar, Union, get_type_hints, overload
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Awaitable, Callable, ClassVar, Optional, Sequence, TypeVar, Union, get_type_hints, overload
 
 from pydantic import BaseModel, ValidationError
 
@@ -25,9 +26,15 @@ _SLOW_LOCK_WAIT_S = 30.0
 @overload
 def tool(fn: Callable[..., Any]) -> Callable[..., Any]: ...
 @overload
-def tool(*, shared: bool = True, concurrent: bool = False) -> Callable[[Callable[..., Any]], Callable[..., Any]]: ...
 def tool(
-    fn: Optional[Callable[..., Any]] = None, *, shared: bool = True, concurrent: bool = False
+    *, shared: bool = True, concurrent: bool = False, exclusive: bool = False
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]: ...
+def tool(
+    fn: Optional[Callable[..., Any]] = None,
+    *,
+    shared: bool = True,
+    concurrent: bool = False,
+    exclusive: bool = False,
 ) -> Callable[..., Any]:
     """Declare a method as a tool.
 
@@ -39,11 +46,17 @@ def tool(
     ``concurrent=True`` only for tools that never grant reward (shell, file
     reads, search); those skip the session lock. ``@terminal`` tools cannot
     be concurrent.
+
+    ``exclusive=True`` is for a tool that ends the episode (a model-called
+    ``submit``): it waits for every in-flight call on the session, concurrent
+    ones included, and holds new calls until it returns — so nothing can
+    change the state it grades. ``@terminal`` tools are always exclusive.
     """
     def mark(f: Callable[..., Any]) -> Callable[..., Any]:
         setattr(f, "_env_tool", True)
         setattr(f, "_env_tool_shared", shared)
         setattr(f, "_env_tool_concurrent", concurrent)
+        setattr(f, "_env_tool_exclusive", exclusive)
         return f
     # @tool(...) returns the marker; bare @tool applies it.
     return mark if fn is None else mark(fn)
@@ -78,8 +91,57 @@ def terminal(fn: Callable[..., Any]) -> Callable[..., Any]:
     return fn
 
 
+def _is_exclusive(fn: Callable[..., Any]) -> bool:
+    return bool(getattr(fn, "_env_tool_exclusive", False) or getattr(fn, "_env_tool_terminal", False))
+
+
 def _is_concurrent(fn: Callable[..., Any]) -> bool:
-    return bool(getattr(fn, "_env_tool_concurrent", False)) and not getattr(fn, "_env_tool_terminal", False)
+    return bool(getattr(fn, "_env_tool_concurrent", False)) and not _is_exclusive(fn)
+
+
+class _ExclusiveGate:
+    """Per-session readers-writer gate: tool calls share it, exclusive tools own it.
+
+    Writer-preferring, so a queued exclusive call is not starved by a stream of
+    concurrent ones. A tool that calls another tool from inside its body would
+    deadlock behind a queued exclusive call; tools don't do that.
+    """
+
+    def __init__(self) -> None:
+        self._cond = asyncio.Condition()
+        self._active = 0
+        self._writing = False
+        self._writers_waiting = 0
+
+    @asynccontextmanager
+    async def shared(self) -> AsyncIterator[None]:
+        async with self._cond:
+            await self._cond.wait_for(lambda: not self._writing and not self._writers_waiting)
+            self._active += 1
+        try:
+            yield
+        finally:
+            async with self._cond:
+                self._active -= 1
+                self._cond.notify_all()
+
+    @asynccontextmanager
+    async def exclusive(self) -> AsyncIterator[None]:
+        async with self._cond:
+            self._writers_waiting += 1
+            try:
+                await self._cond.wait_for(lambda: not self._writing and self._active == 0)
+            finally:
+                self._writers_waiting -= 1
+                # A cancelled writer may have been what held readers back.
+                self._cond.notify_all()
+            self._writing = True
+        try:
+            yield
+        finally:
+            async with self._cond:
+                self._writing = False
+                self._cond.notify_all()
 
 
 def _terminal_arg(name: str, fn: Callable[..., Any]) -> Optional[str]:
@@ -168,6 +230,8 @@ class Environment(ABC):
     _episode_finished: bool = False
     # Per-tool locks serialising non-concurrent calls; created lazily for the same reason.
     _tool_locks: Optional[dict[str, asyncio.Lock]] = None
+    # Shared by every call, owned by exclusive ones; created lazily for the same reason.
+    _exclusive_gate: Optional[_ExclusiveGate] = None
 
     def __init__(self, task_spec: JSONObject = {}, secrets: dict[str, str] = {}) -> None:
         self.task_spec = task_spec
@@ -520,26 +584,47 @@ class Environment(ABC):
         """Validate input against the tool's Pydantic model and invoke it."""
         if self._episode_finished:
             return self._episode_finished_error(name)
+        if self._exclusive_gate is None:
+            self._exclusive_gate = _ExclusiveGate()
+        gate = self._exclusive_gate
+        if _is_exclusive(fn):
+            async with gate.exclusive():
+                return await self._run_after_wait(name, fn, input, start)
         if _is_concurrent(fn):
-            return await self._run_tool_fn(name, fn, input, start)
+            async with gate.shared():
+                # Only an exclusive call can have held us, and it may have ended the episode.
+                if self._episode_finished:
+                    return self._episode_finished_error(name)
+                return await self._run_tool_fn(name, fn, input)
         if self._tool_locks is None:
             self._tool_locks = {}
+        # Tool lock outside the gate, so calls queued on a tool don't hold off an exclusive one.
         async with self._tool_locks.setdefault(name, asyncio.Lock()):
-            waited = time.monotonic() - start
-            if waited > _SLOW_LOCK_WAIT_S:
-                logger.warning("tool_call_waited_for_tool_lock", tool=name, waited_s=round(waited, 1))
-            # The call we queued behind may have ended the episode.
-            if self._episode_finished:
-                return self._episode_finished_error(name)
-            return await self._run_tool_fn(name, fn, input, start)
+            async with gate.shared():
+                return await self._run_after_wait(name, fn, input, start)
+
+    async def _run_after_wait(
+        self, name: str, fn: Callable[..., Any], input: JSONObject, start: float
+    ) -> RunToolOutput:
+        waited = time.monotonic() - start
+        if waited > _SLOW_LOCK_WAIT_S:
+            logger.warning("tool_call_waited_for_tool_lock", tool=name, waited_s=round(waited, 1))
+        # The call we queued behind may have ended the episode.
+        if self._episode_finished:
+            return self._episode_finished_error(name, lock_wait_ms=waited * 1000)
+        return await self._run_tool_fn(name, fn, input, lock_wait_ms=waited * 1000)
 
     async def _run_tool_fn(
         self,
         name: str,
         fn: Callable[..., Any],
         input: JSONObject,
-        start: float,
+        lock_wait_ms: Optional[float] = None,
     ) -> RunToolOutput:
+        """Run the tool; ``duration_ms`` excludes time queued on the tool lock (``lock_wait_ms``)."""
+        start = time.monotonic()
+        # Only serialised calls carry lock_wait_ms, so concurrent tools can be told apart.
+        timing = {} if lock_wait_ms is None else {"lock_wait_ms": lock_wait_ms}
         _, hints, params = _introspect_tool(fn)
         try:
             if not params:
@@ -557,7 +642,7 @@ class Environment(ABC):
                 res = await run_user_callable(fn, inp)
         except Exception:
             duration_ms = (time.monotonic() - start) * 1000
-            logger.exception("tool_call_failed", tool=name, duration_ms=duration_ms)
+            logger.exception("tool_call_failed", tool=name, duration_ms=duration_ms, **timing)
             raise
 
         if not isinstance(res, ToolOutput):
@@ -567,17 +652,17 @@ class Environment(ABC):
         # one ran: the first finished=True wins, so this result is discarded
         # rather than handed back as a second grade.
         if self._episode_finished:
-            return self._episode_finished_error(name)
+            return self._episode_finished_error(name, **timing)
         if res.finished:
             self._episode_finished = True
 
         duration_ms = (time.monotonic() - start) * 1000
-        logger.info("tool_call_completed", tool=name, duration_ms=duration_ms)
+        logger.info("tool_call_completed", tool=name, duration_ms=duration_ms, **timing)
         return RunToolOutput(RunToolSuccess(output=res))
 
     @staticmethod
-    def _episode_finished_error(name: str) -> RunToolOutput:
-        logger.warning("tool_call_after_episode_finished", tool=name)
+    def _episode_finished_error(name: str, **timing: float) -> RunToolOutput:
+        logger.warning("tool_call_after_episode_finished", tool=name, **timing)
         return RunToolOutput(RunToolError(
             reason="episode_finished",
             error=f"{name!r} was not run: the episode has finished (a tool "

{openreward-0.1.161.dev0 → openreward-0.1.161.dev1}/openreward/toolsets/cli.py

@@ -21,7 +21,7 @@ Pydantic param shapes and ``@tool``-decorated wrappers.
 """
 from __future__ import annotations
 
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, ClassVar, Dict, List, Literal, Optional
 
 from pydantic import BaseModel
 
@@ -119,12 +119,26 @@ class CLIToolset(Toolset):
     ``root`` resolution, first match wins: the ``root`` constructor argument
     (for direct instantiation and tests), then ``env.workspace_root``, then the
     class attribute (``None`` = unrooted).
+
+    **bash limits.** Subclass and set ``bash_timeout``, ``max_bash_timeout`` and
+    ``bash_max_bytes`` to replace ``sandbox.run``'s defaults (300s, 50KB)::
+
+        class MyCLIToolset(CLIToolset):
+            bash_timeout = max_bash_timeout = 120
+            bash_max_bytes = 4_000_000
     """
 
     #: Default workspace root; see the class docstring. Subclasses may override.
     #: Not a ClassVar — __init__ overwrites it per instance with the resolved root.
     root: Optional[str] = None
 
+    #: bash timeout (s) when the call sets none; None defers to sandbox.run's default.
+    bash_timeout: ClassVar[Optional[float]] = None
+    #: Cap (s) on any bash timeout, including one the agent passes; None = uncapped.
+    max_bash_timeout: ClassVar[Optional[float]] = None
+    #: sandbox.run max_bytes for bash; None defers to sandbox.run's default.
+    bash_max_bytes: ClassVar[Optional[int]] = None
+
     @classmethod
     def name(cls) -> str:
         return "cli"
@@ -144,8 +158,13 @@ class CLIToolset(Toolset):
     @tool(concurrent=True)
     async def bash(self, params: BashParams) -> ToolOutput:
         """Execute a bash command in the sandbox."""
+        timeout = params.timeout if params.timeout is not None else self.bash_timeout
+        if self.max_bash_timeout is not None:
+            timeout = min(timeout if timeout is not None else self.max_bash_timeout,
+                          self.max_bash_timeout)
         return await _bash_impl(
-            self.sandbox, params.command, params.timeout, root=self.root
+            self.sandbox, params.command, timeout, root=self.root,
+            max_bytes=self.bash_max_bytes,
         )
 
     @tool(concurrent=True)

{openreward-0.1.161.dev0 → openreward-0.1.161.dev1}/openreward.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openreward
-Version: 0.1.161.dev0
+Version: 0.1.161.dev1
 Summary: Python SDK for the OpenReward platform.
 Author-email: GR Inc <hello@gr.inc>
 License-Expression: MIT

{openreward-0.1.161.dev0 → openreward-0.1.161.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "openreward"
-version = "0.1.161.dev0"
+version = "0.1.161.dev1"
 description = "Python SDK for the OpenReward platform."
 readme = "README.md"
 license = "MIT"

{openreward-0.1.161.dev0 → openreward-0.1.161.dev1}/tests/test_cli_toolset_root_and_timeout.py

@@ -146,6 +146,56 @@ def test_bash_timeout_default_is_none():
     assert BashParams(command="x").timeout is None
 
 
+class LimitedToolset(CLIToolset):
+    bash_timeout = 120
+    max_bash_timeout = 120
+    bash_max_bytes = 4_000_000
+
+
+@pytest.mark.asyncio
+async def test_class_limits_apply_when_the_call_sets_none():
+    sb = RecordingSandbox()
+    await LimitedToolset(FakeEnv(sb)).bash(BashParams(command="make"))
+    assert sb.calls[0][1] == {"timeout": 120, "max_bytes": 4_000_000}
+
+
+@pytest.mark.asyncio
+async def test_agent_timeout_is_capped():
+    sb = RecordingSandbox()
+    ts = LimitedToolset(FakeEnv(sb))
+    await ts.bash(BashParams(command="sleep 1", timeout=9999))
+    await ts.bash(BashParams(command="sleep 1", timeout=5))
+    assert [kw["timeout"] for _, kw in sb.calls] == [120, 5]
+
+
+@pytest.mark.asyncio
+async def test_cap_alone_bounds_the_sandbox_default():
+    class CapOnly(CLIToolset):
+        max_bash_timeout = 60
+
+    sb = RecordingSandbox()
+    await CapOnly(FakeEnv(sb)).bash(BashParams(command="make"))
+    assert sb.calls[0][1] == {"timeout": 60}
+
+
+class KilledSandbox(RecordingSandbox):
+    async def run(self, cmd: str, **kwargs) -> RunResult:
+        self.calls.append((cmd, kwargs))
+        return RunResult(output="partial", return_code=124, truncated=True, timed_out=True)
+
+
+@pytest.mark.asyncio
+async def test_timeout_and_truncation_are_reported():
+    sb = KilledSandbox()
+    out = await LimitedToolset(FakeEnv(sb)).bash(BashParams(command="yes"))
+    assert out.blocks[0].text == (
+        "partial\n\n(exit 124)\n\n[output exceeded the limit and was cut off; "
+        "command hit the 120s time limit and was killed]"
+    )
+    assert out.metadata == {"output": "partial", "exit_code": 124,
+                            "truncated": True, "timed_out": True}
+
+
 # ── enforcement ──
 
 def test_impls_refuse_to_run_unrooted_by_omission():

{openreward-0.1.161.dev0 → openreward-0.1.161.dev1}/tests/test_tool_concurrency.py

@@ -229,3 +229,129 @@ def test_file_locks_are_per_sandbox_per_file():
     assert file_lock(a, "/w/f.txt") is file_lock(a, "/w/./f.txt")
     assert file_lock(a, "/w/f.txt") is not file_lock(a, "/w/g.txt")
     assert file_lock(a, "/w/f.txt") is not file_lock(b, "/w/f.txt")
+
+
+@pytest.mark.asyncio
+async def test_logged_duration_excludes_lock_wait(monkeypatch):
+    from openreward.environments import environment as environment_module
+
+    done: list[dict] = []
+    monkeypatch.setattr(
+        environment_module.logger, "info",
+        lambda event, **kw: done.append(kw) if event == "tool_call_completed" else None,
+    )
+    env = SubmitEnv()
+    await asyncio.gather(*(env._call_tool("check", {"answer": "5"}) for _ in range(2)))
+    await env._call_tool("probe", {})
+    first, second, probe = done
+    # Each check runs ~10ms; the second queues ~10ms behind the first.
+    assert first["lock_wait_ms"] < 5 and second["lock_wait_ms"] >= 8
+    assert all(8 <= e["duration_ms"] < 50 for e in (first, second))
+    assert "lock_wait_ms" not in probe
+
+
+# ── exclusive tools ──
+
+class GatedEnv(SubmitEnv):
+    """A shell-like concurrent tool and an exclusive submit that records what overlapped it."""
+
+    def __init__(self, task_spec: JSONObject = {}, secrets: dict[str, str] = {}) -> None:
+        super().__init__(task_spec, secrets)
+        self.shells_running = 0
+        self.overlapped_submit = False
+        self.shells_done = 0
+
+    @tool(concurrent=True)
+    async def shell(self) -> ToolOutput:
+        self.shells_running += 1
+        await asyncio.sleep(0.05)
+        self.shells_running -= 1
+        self.shells_done += 1
+        return ToolOutput(blocks=[TextBlock(text="ok")])
+
+    @tool(exclusive=True)
+    async def final(self) -> ToolOutput:
+        self.overlapped_submit = self.shells_running > 0
+        await asyncio.sleep(0.05)
+        return ToolOutput(blocks=[TextBlock(text="graded")], reward=1.0, finished=True)
+
+    @tool(exclusive=True)
+    async def soft_final(self) -> ToolOutput:
+        """Exclusive but leaves the episode open (e.g. a grader fault)."""
+        self.overlapped_submit = self.shells_running > 0
+        await asyncio.sleep(0.05)
+        return ToolOutput(blocks=[TextBlock(text="try again")])
+
+
+@pytest.mark.asyncio
+async def test_exclusive_waits_for_in_flight_concurrent_calls():
+    env = GatedEnv()
+    shells = [asyncio.create_task(env._call_tool("shell", {})) for _ in range(3)]
+    await asyncio.sleep(0.01)
+    res = await env._call_tool("final", {})
+    await asyncio.gather(*shells)
+    assert res.root.ok and not env.overlapped_submit
+    assert env.shells_done == 3
+
+
+@pytest.mark.asyncio
+async def test_calls_arriving_during_exclusive_are_held_then_refused():
+    env = GatedEnv()
+    final = asyncio.create_task(env._call_tool("final", {}))
+    await asyncio.sleep(0.01)
+    late = await asyncio.gather(env._call_tool("shell", {}), env._call_tool("check", {"answer": "4"}))
+    assert (await final).root.ok
+    assert all(_refused(r) for r in late)
+    assert env.shells_done == 0 and env.attempts == 0
+
+
+@pytest.mark.asyncio
+async def test_held_calls_run_when_exclusive_leaves_episode_open():
+    env = GatedEnv()
+    soft = asyncio.create_task(env._call_tool("soft_final", {}))
+    await asyncio.sleep(0.01)
+    shell = await env._call_tool("shell", {})
+    assert (await soft).root.ok and shell.root.ok
+    assert env.shells_done == 1 and not env.overlapped_submit
+
+
+@pytest.mark.asyncio
+async def test_queued_exclusive_is_not_starved_by_new_concurrent_calls():
+    env = GatedEnv()
+    first = asyncio.create_task(env._call_tool("shell", {}))
+    await asyncio.sleep(0.01)
+    final = asyncio.create_task(env._call_tool("final", {}))
+    await asyncio.sleep(0.01)
+    later = await env._call_tool("shell", {})
+    assert (await first).root.ok and (await final).root.ok
+    assert _refused(later)
+
+
+@pytest.mark.asyncio
+async def test_terminal_tool_is_exclusive():
+    class TerminalEnv(GatedEnv):
+        @terminal
+        @tool
+        async def answer(self, args: Answer) -> ToolOutput:
+            self.overlapped_submit = self.shells_running > 0
+            return ToolOutput(blocks=[TextBlock(text="graded")], reward=1.0, finished=True)
+
+    env = TerminalEnv()
+    shell = asyncio.create_task(env._call_tool("shell", {}))
+    await asyncio.sleep(0.01)
+    await env.call_terminal_tool("4")
+    await shell
+    assert not env.overlapped_submit
+
+
+@pytest.mark.asyncio
+async def test_cancelled_exclusive_releases_held_calls():
+    env = GatedEnv()
+    first = asyncio.create_task(env._call_tool("shell", {}))
+    await asyncio.sleep(0.01)
+    final = asyncio.create_task(env._call_tool("final", {}))
+    await asyncio.sleep(0.01)
+    final.cancel()
+    shell = await asyncio.wait_for(env._call_tool("shell", {}), timeout=1)
+    await first
+    assert shell.root.ok