diff --git a/datamind/agent/loop_native.py b/datamind/agent/loop_native.py index 7673261..c9881cd 100644 --- a/datamind/agent/loop_native.py +++ b/datamind/agent/loop_native.py @@ -815,6 +815,7 @@ 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 @@ -822,19 +823,27 @@ async def stream_turn( 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 diff --git a/datamind/tests/test_agent_loop.py b/datamind/tests/test_agent_loop.py index b4c5662..c12aded 100644 --- a/datamind/tests/test_agent_loop.py +++ b/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 = [