fix: enforce native stream deadlines

datamind/agent/loop_native.py

@@ -815,26 +815,35 @@ async def stream_turn(
         if contract_prompt:
             system_prompt = f"{system_prompt}\n\n{contract_prompt}".strip()
 
+        started = time.monotonic()
         for iteration in range(self._cfg.max_tool_turns):
             allow_tools = (
                 iteration < self._cfg.max_tool_turns - 1
                 and tool_call_count < self._cfg.max_tool_calls
                 and total_input < self._cfg.max_input_tokens
             )
             final: ModelResponse | None = None
-            async for model_event in self._model_client.stream(
-                model=self._cfg.model,
-                max_tokens=self._cfg.max_tokens,
-                temperature=self._cfg.temperature,
-                system=system_prompt or None,
-                tools=self._tools.as_anthropic_tools() if allow_tools and len(self._tools) else None,
-                tool_choice=None if allow_tools else "none",
-                messages=conv,
-            ):
-                if model_event.type == "text" and model_event.delta:
-                    yield AgentEvent(type="text", data={"delta": model_event.delta})
-                elif model_event.type == "done":
-                    final = model_event.response
+            remaining = max(
+                0.1, self._cfg.wall_clock_timeout_s - (time.monotonic() - started)
+            )
+            try:
+                async with asyncio.timeout(remaining):
+                    async for model_event in self._model_client.stream(
+                        model=self._cfg.model,
+                        max_tokens=self._cfg.max_tokens,
+                        temperature=self._cfg.temperature,
+                        system=system_prompt or None,
+                        tools=self._tools.as_anthropic_tools() if allow_tools and len(self._tools) else None,
+                        tool_choice=None if allow_tools else "none",
+                        messages=conv,
+                    ):
+                        if model_event.type == "text" and model_event.delta:
+                            yield AgentEvent(type="text", data={"delta": model_event.delta})
+                        elif model_event.type == "done":
+                            final = model_event.response
+            except TimeoutError:
+                yield AgentEvent(type="error", data={"message": "model stream timed out"})
+                return
             if final is None:
                 yield AgentEvent(type="error", data={"message": "model stream ended without a final response"})
                 return

datamind/tests/test_agent_loop.py

@@ -15,6 +15,7 @@
 from datamind.agent.loop_native import NativeAgentLoop
 from datamind.agent.base import AgentLoopConfig
 from datamind.core.errors import FinalAnswerContractError
+from datamind.core.protocols import ModelStreamEvent
 from datamind.core.tools import ToolRegistry, ToolSpec
 
 
@@ -62,6 +63,23 @@ def __init__(self, script: list[_Message]) -> None:
         self.messages = _FakeMessages(script)
 
 
+class _HangingStreamClient:
+    protocol = "test"
+
+    def __init__(self) -> None:
+        self.cancelled = asyncio.Event()
+
+    async def complete(self, **kwargs: Any):
+        raise AssertionError("stream test should not call complete")
+
+    async def stream(self, **kwargs: Any):
+        try:
+            await asyncio.Event().wait()
+        finally:
+            self.cancelled.set()
+        yield ModelStreamEvent(type="done")
+
+
 # ------------------------------------------------------- tiny tool ---
 
 
@@ -225,6 +243,28 @@ async def test_slow_tool_is_cut_off_to_preserve_contract_finalization_budget():
     assert out["tool_trace"][0]["error_type"] == "TimeoutError"
 
 
+@pytest.mark.asyncio
+async def test_stream_turn_enforces_model_wall_clock_deadline():
+    client = _HangingStreamClient()
+    loop = NativeAgentLoop(
+        client=client,
+        tools=ToolRegistry(),
+        config=AgentLoopConfig(model="m", wall_clock_timeout_s=0.1),
+    )
+
+    async def collect():
+        return [event async for event in loop.stream_turn(user_message="hello")]
+
+    try:
+        events = await asyncio.wait_for(collect(), timeout=0.5)
+    except asyncio.TimeoutError:
+        pytest.fail("stream_turn ignored the configured wall-clock deadline")
+
+    assert events[-1].type == "error"
+    assert "timed out" in events[-1].data["message"]
+    assert client.cancelled.is_set()
+
+
 @pytest.mark.asyncio
 async def test_tool_error_is_surfaced_as_tool_result():
     script = [