Worlds fixes (realized after Seahaven was built)

libs/core/kiln_ai/adapters/eval/test_eval_runner_worlds.py

@@ -11,6 +11,7 @@
 """
 
 import asyncio
+import json
 import re
 from typing import ClassVar
 from unittest.mock import patch
@@ -884,7 +885,12 @@ async def test_tool_error_reaches_the_model_and_ends_the_episode(
     generator = ToolCallingGenerator(task, tool_id, allow_error=True)
     with patch.object(BaseV2EvalBridge, "run_task", new=generator):
         await _drain(_runner([cfg], run_config, session_manager))
-    assert generator.outputs["ei_a"] == "boom"
+    # The environment's `error_type` is rendered for the model alongside the message;
+    # `error_message` stays the message alone so an eval matching on error text is not
+    # perturbed by the type.
+    assert json.loads(generator.outputs["ei_a"]) == {
+        "error": {"code": "execution_error", "message": "boom", "details": None}
+    }
     assert generator.errors["ei_a"] == "boom"
     (trace,) = _traces(task)
     assert trace.world_episode.final_state["notes"] == []

libs/core/kiln_ai/tools/test_world_tool.py

@@ -1,6 +1,7 @@
 """Tests for resolving world tool ids through the registry, and the proxy
 they resolve to. The session manager is a fake that records calls."""
 
+import json
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -29,7 +30,11 @@
 from kiln_ai.tools.built_in_tools.math_tools import AddTool
 from kiln_ai.tools.code_tool import PythonCodeTool
 from kiln_ai.tools.tool_registry import tool_from_id_and_project
-from kiln_ai.tools.world_tool import OpenEnvToolProxy, render_tool_result
+from kiln_ai.tools.world_tool import (
+    OpenEnvToolProxy,
+    render_tool_error,
+    render_tool_result,
+)
 from kiln_ai.worlds.session_manager import ToolCallOutcome
 
 SCHEMA = {"type": "object", "properties": {"x": {"type": "string"}}}
@@ -168,6 +173,36 @@ async def test_error_outcome_reaches_model_as_error(
         assert result.is_error and result.output == "tool exploded"
         assert result.error_message == "tool exploded"
 
+    async def test_coded_error_reaches_the_model_as_a_structured_error(
+        self, world, active, session_manager
+    ):
+        """A code is rendered for the model to read, and the call is still a failure.
+
+        Every error on the observation is a failed call: OpenEnv's `error_type`
+        vocabulary, `execution_error` included, means the call did not work. An
+        environment that wants an error read as an ordinary answer returns it as a
+        result instead."""
+        session_manager.outcome = ToolCallOutcome(
+            result=None,
+            error="issue ENG-99 not found",
+            reward=None,
+            done=False,
+            error_code="execution_error",
+            error_details={"id": "ENG-99"},
+        )
+        tool = tool_from_id_and_project(build_world_tool_id(world.id, "lookup"))
+        result = await tool.run(ToolCallContext(), x="q")
+        assert result.is_error
+        assert json.loads(result.output) == {
+            "error": {
+                "code": "execution_error",
+                "message": "issue ENG-99 not found",
+                "details": {"id": "ENG-99"},
+            }
+        }
+        # The message alone, so an eval matching on error text is not perturbed.
+        assert result.error_message == "issue ENG-99 not found"
+
     async def test_fastmcp_is_error_result_is_an_error(
         self, world, active, session_manager
     ):
@@ -209,9 +244,22 @@ async def test_unknown_tool(self, world, active):
             tool_from_id_and_project(build_world_tool_id(world.id, "nope"))
 
 
+class TestRenderError:
+    def test_no_code_renders_the_message_alone(self):
+        # The common case against a conformant environment that reports no type.
+        assert render_tool_error(None, "boom", None) == "boom"
+        assert render_tool_error("", "boom", {"a": 1}) == "boom"
+
+    def test_a_code_renders_all_three_keys(self):
+        # Always all three, so the shape is stable whether or not there are details.
+        assert json.loads(render_tool_error("timeout", "slow", None)) == {
+            "error": {"code": "timeout", "message": "slow", "details": None}
+        }
+
+
 class TestRenderResult:
     def test_shapes(self):
-        assert render_tool_result(None) == ""
+        assert render_tool_result(None) == "null"
         assert render_tool_result("text") == "text"
         assert render_tool_result({"a": 1}) == '{"a": 1}'
         assert render_tool_result([1, 2]) == "[1, 2]"

libs/core/kiln_ai/tools/world_tool.py

@@ -58,24 +58,47 @@ async def run(
             self._ctx.episode, self._tool.name, dict(kwargs)
         )
         if outcome.error is not None:
+            # Every error on the observation is a failed call. OpenEnv's `error_type`
+            # is the only structure the field carries, and its whole vocabulary --
+            # `execution_error` included -- means the call did not work. An
+            # environment that wants an error of its own read as an ordinary answer
+            # returns it as a result, which is what `result` is for.
+            output = render_tool_error(
+                outcome.error_code, outcome.error, outcome.error_details
+            )
             return ToolCallResult(
-                output=outcome.error, is_error=True, error_message=outcome.error
+                output=output, is_error=True, error_message=outcome.error
             )
         output = render_tool_result(outcome.result)
         if isinstance(outcome.result, dict) and outcome.result.get("is_error") is True:
             return ToolCallResult(output=output, is_error=True, error_message=output)
         return ToolCallResult(output=output)
 
 
+def render_tool_error(code: str | None, message: str, details: Any) -> str:
+    """The text the model sees for an environment's tool error.
+
+    A coded error renders as `{"error": {"code", "message", "details"}}`, always all
+    three keys so the shape is stable whether or not there are details; an environment
+    that reports no code renders the message alone."""
+    if not code:
+        return message
+    return json.dumps(
+        {"error": {"code": code, "message": message, "details": details}},
+        ensure_ascii=False,
+    )
+
+
 def render_tool_result(result: Any) -> str:
     """The environment's tool result as the text the model sees.
 
     A FastMCP call result (`{"content": [...], "structured_content", "data", "is_error"}`)
     renders its native `data` when present, else its text content blocks; bare
     MCP-style content blocks are flattened to their text; anything else that is not
-    already a string is serialized as JSON."""
+    already a string is serialized as JSON. `None` renders as `null`, the same text a
+    tool that serializes its own result would show for an empty body."""
     if result is None:
-        return ""
+        return "null"
     if isinstance(result, str):
         return result
     if isinstance(result, dict) and isinstance(result.get("content"), list):

libs/core/kiln_ai/worlds/session_manager.py

@@ -49,21 +49,68 @@
 
 STEP_TIMEOUT_S = 600.0
 """How long one reset, step or state call may take. Tool calls block the LLM loop, so
-this is deliberately generous; an environment enforces its own per-tool timeouts."""
+this is deliberately generous; an environment enforces its own per-tool timeouts.
+A call that stalls the environment's event loop never reaches this ceiling: the
+keepalive below closes the session after PING_TIMEOUT_S of unanswered pings, so a
+stalled loop surfaces as a closed session at 120s, not a step timeout at 600s."""
+
+PING_INTERVAL_S = 20.0
+"""How often the websockets keepalive pings the environment. Stated rather than left to
+the library's default so a default that moves cannot silently change what a run
+measured."""
+
+PING_TIMEOUT_S = 120.0
+"""How long an unanswered ping may go before the keepalive closes the session. A world
+server doing synchronous work stalls its event loop and cannot answer a ping: the
+library's 20s timeout would drop every session on the box at once, so this matches the
+120s a world framework's own client allows for it. Stated for the same reason as the
+interval above."""
 
 
 class OpenEnvError(RuntimeError):
     """The environment answered a request with an error, or could not be reached."""
 
 
+def read_observation_error(error: Any) -> tuple[str | None, str | None, Any]:
+    """An observation's `error` as `(message, code, details)`.
+
+    Kiln never rejects a shape here. OpenEnv declares `{error_type, message}` and
+    forbids extra keys, so `error_type` is read first; a `code` is read after it, for
+    an environment that reports one instead. `details` is whatever a dict carried,
+    which a conformant environment has nowhere to put. An error that is not a dict,
+    or one with no message, degrades to its own text rather than failing the call."""
+    if error is None:
+        return None, None, None
+    if not isinstance(error, dict):
+        return str(error), None, None
+    code = error.get("error_type") or error.get("code")
+    return (
+        str(error.get("message") or error),
+        str(code) if code else None,
+        error.get("details"),
+    )
+
+
 @dataclass(frozen=True)
 class ToolCallOutcome:
-    """What one `step(CallToolAction)` came back with."""
+    """What one `step(CallToolAction)` came back with.
+
+    An environment that reports an error as a dict has whatever structure it carries
+    read out alongside the human-readable message. `error_code` is OpenEnv's own
+    `error_type` -- one of `execution_error`, `invalid_args`, `transport_error`,
+    `tool_not_found`, `timeout` -- and falls back to a `code` for an environment that
+    predates that shape or does not use it. `error_details` is whatever a `details`
+    carried, which a conformant environment does not send: `ToolError` forbids extra
+    keys, so a world with a code and details of its own puts them somewhere this field
+    does not reach. Neither is a shape Kiln insists on; an error that is not a dict at
+    all leaves both None."""
 
     result: Any
     error: str | None
     reward: float | None
     done: bool
+    error_code: str | None = None
+    error_details: Any = None
 
 
 class WorldSessionManager(Protocol):
@@ -216,16 +263,16 @@ async def call_tool(
             done = bool(data.get("done", False))
             session.done = session.done or done
         observation = data.get("observation") or {}
-        error = observation.get("error")
-        if isinstance(error, dict):
-            error = str(error.get("message") or error)
-        elif error is not None:
-            error = str(error)
+        error, error_code, error_details = read_observation_error(
+            observation.get("error")
+        )
         return ToolCallOutcome(
             result=observation.get("result"),
             error=error,
             reward=float(reward) if isinstance(reward, (int, float)) else None,
             done=done,
+            error_code=error_code,
+            error_details=error_details,
         )
 
     async def end_episode(self, episode: WorldEpisode) -> WorldEpisode:
@@ -307,7 +354,13 @@ async def _connect_remote(self, world: World, base_url: str) -> _EnvServer:
 
     async def _open(self, server: _EnvServer) -> ClientConnection:
         try:
-            return await connect(server.ws_url, max_size=None, open_timeout=30)
+            return await connect(
+                server.ws_url,
+                max_size=None,
+                open_timeout=30,
+                ping_interval=PING_INTERVAL_S,
+                ping_timeout=PING_TIMEOUT_S,
+            )
         except (OSError, ConnectionClosed) as e:
             raise OpenEnvError(
                 f"Could not open a session on {server.base_url}: {e}"

libs/core/kiln_ai/worlds/test_session_manager.py

@@ -5,9 +5,11 @@
 
 from kiln_ai.datamodel.project import Project
 from kiln_ai.datamodel.world import World
+from kiln_ai.worlds import session_manager as session_manager_module
 from kiln_ai.worlds.session_manager import (
     OpenEnvError,
     OpenEnvSessionManager,
+    read_observation_error,
 )
 from kiln_ai.worlds.testing import (
     ENV_NAME,
@@ -84,9 +86,11 @@ async def test_call_tool_and_end_episode(self, session_manager, remote_world):
         assert listed.result == ["hi"]
         boom = await session_manager.call_tool(episode, "explode", {})
         assert boom.result is None and boom.error == "boom"
+        assert boom.error_code == "execution_error" and boom.error_details is None
         assert boom.reward == -1.0 and boom.done is True
         missing = await session_manager.call_tool(episode, "nope", {})
         assert missing.error == "no tool 'nope'"
+        assert missing.error_code == "tool_not_found"
         # Rewards and done are tracked on the live session only.
         session = session_manager._sessions[episode.episode_id]
         assert session.rewards == [1.0, 0.0, -1.0] and session.done is True
@@ -177,3 +181,66 @@ async def test_url_change_reconnects(self, session_manager, remote_world):
         remote_world.env_url = f"http://127.0.0.1:{free_port()}"
         with pytest.raises(OpenEnvError, match="did not answer /metadata"):
             await session_manager.world_version(remote_world, {})
+
+
+class TestErrorShapeTolerance:
+    """Kiln reads what an environment sends and never rejects a shape.
+
+    OpenEnv declares `{error_type, message}` and forbids extra keys, so a conformant
+    environment cannot send a code or details of its own. Kiln still reads both, for
+    an environment that reports an error some other way; an unreadable shape degrades
+    to its text rather than failing the call."""
+
+    @pytest.mark.parametrize(
+        "error,expected_code,expected_message",
+        [
+            ({"error_type": "timeout", "message": "slow"}, "timeout", "slow"),
+            ({"code": "not_found", "message": "gone"}, "not_found", "gone"),
+            # `error_type` wins: it is the field the protocol declares.
+            (
+                {"error_type": "invalid_args", "code": "bad", "message": "no"},
+                "invalid_args",
+                "no",
+            ),
+            # No message: the whole dict is the best text there is.
+            ({"error_type": "timeout"}, "timeout", "{'error_type': 'timeout'}"),
+            # Not a dict at all.
+            ("plain", None, "plain"),
+            (42, None, "42"),
+        ],
+    )
+    def test_shapes(self, error, expected_code, expected_message):
+        message, code, _ = read_observation_error(error)
+        assert message == expected_message
+        assert code == expected_code
+
+    def test_no_error_is_no_error(self):
+        assert read_observation_error(None) == (None, None, None)
+
+    def test_details_ride_along(self):
+        _, _, details = read_observation_error(
+            {"code": "not_found", "message": "gone", "details": {"id": 7}}
+        )
+        assert details == {"id": 7}
+
+
+class TestKeepalive:
+    async def test_session_is_opened_with_a_stated_keepalive(
+        self, session_manager, remote_world, monkeypatch
+    ):
+        """A world server doing synchronous work cannot answer a ping, and the
+        library's 20s default would drop the session long before a long tool call
+        finished. Both values are stated so a moving default cannot change a run."""
+        seen: dict[str, object] = {}
+        real_connect = session_manager_module.connect
+
+        async def spy(url, **kwargs):
+            seen.update(kwargs)
+            return await real_connect(url, **kwargs)
+
+        monkeypatch.setattr(session_manager_module, "connect", spy)
+        episode = await session_manager.start_episode(remote_world, {})
+        assert seen["ping_interval"] == session_manager_module.PING_INTERVAL_S
+        assert seen["ping_timeout"] == session_manager_module.PING_TIMEOUT_S
+        assert seen["ping_timeout"] > 20, "the library default is what we are avoiding"
+        await session_manager.release(episode)

libs/core/kiln_ai/worlds/testing.py

@@ -114,7 +114,7 @@ def step(self, action: dict[str, Any]) -> dict[str, Any]:
                 "observation": {
                     "tool_name": name,
                     "result": None,
-                    "error": {"type": "execution_error", "message": "boom"},
+                    "error": {"error_type": "execution_error", "message": "boom"},
                 },
                 "reward": -1.0,
                 "done": True,
@@ -123,7 +123,10 @@ def step(self, action: dict[str, Any]) -> dict[str, Any]:
             "observation": {
                 "tool_name": name,
                 "result": None,
-                "error": {"type": "tool_not_found", "message": f"no tool {name!r}"},
+                "error": {
+                    "error_type": "tool_not_found",
+                    "message": f"no tool {name!r}",
+                },
             },
             "reward": None,
             "done": False,