Эх сурвалжийг харах

fix: preserve provider stream exceptions

Problem: TimingChatClient cleanup could replace an active provider failure or cancellation with an iterator close error.

Risk: Close failures are now secondary only while a provider exception is propagating; normal completion and consumer-close paths continue to surface close failures.
zhenyu.hu 2 долоо хоног өмнө
parent
commit
81156abdbd

+ 7 - 1
src/agent_lab/application/benchmark.py

@@ -282,6 +282,7 @@ class TimingChatClient:
         usage: TokenUsage | None = None
         provider_elapsed_seconds = 0.0
         iterator: AsyncIterator[StreamItem] | None = None
+        provider_exception_active = False
         try:
             stream_chat = getattr(self.inner, "stream_chat")
             iterator = stream_chat(
@@ -301,6 +302,7 @@ class TimingChatClient:
                     )
                     break
                 except BaseException:
+                    provider_exception_active = True
                     provider_elapsed_seconds += max(
                         0.0,
                         self.clock() - wait_started_at,
@@ -327,7 +329,11 @@ class TimingChatClient:
             try:
                 close = getattr(iterator, "aclose", None)
                 if close is not None:
-                    await close()
+                    try:
+                        await close()
+                    except BaseException:
+                        if not provider_exception_active:
+                            raise
             finally:
                 timing = BenchmarkModelCallTiming(
                     call_index=call_index,

+ 111 - 0
tests/test_benchmark_runner.py

@@ -233,6 +233,45 @@ class IncrementalTimingClient:
         return None
 
 
+class CloseFailingProviderIterator:
+    def __init__(
+        self,
+        *,
+        next_error: BaseException | None = None,
+        items: tuple[StreamItem, ...] = (),
+    ) -> None:
+        self.next_error = next_error
+        self.items = list(items)
+
+    def __aiter__(self) -> "CloseFailingProviderIterator":
+        return self
+
+    async def __anext__(self) -> StreamItem:
+        if self.next_error is not None:
+            raise self.next_error
+        if self.items:
+            return self.items.pop(0)
+        raise StopAsyncIteration
+
+    async def aclose(self) -> None:
+        raise RuntimeError("close failed")
+
+
+class IteratorTimingClient:
+    def __init__(self, iterator: AsyncIterator[StreamItem]) -> None:
+        self.iterator = iterator
+
+    def stream_chat(
+        self,
+        messages: list[ChatMessage],
+        tools: list[dict[str, Any]],
+        params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
+    ) -> AsyncIterator[StreamItem]:
+        del messages, tools, params, tool_choice
+        return self.iterator
+
+
 class SelectiveUsageBenchmarkClient:
     def __init__(
         self,
@@ -523,6 +562,78 @@ async def test_timing_client_finishes_timing_when_stream_is_cancelled():
     assert inner.closed is True
 
 
+@pytest.mark.asyncio
+async def test_timing_client_preserves_provider_cancellation_when_close_fails():
+    iterator = CloseFailingProviderIterator(next_error=asyncio.CancelledError())
+    client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
+
+    with pytest.raises(asyncio.CancelledError):
+        async for _ in client.stream_chat(
+            messages=[],
+            tools=[],
+            params=AgentParams(model="benchmark-model"),
+        ):
+            pass
+
+    assert len(client.timings) == 1
+    assert client.timings[0].first_item_kind is None
+    assert client.timings[0].usage is None
+
+
+@pytest.mark.asyncio
+async def test_timing_client_preserves_provider_error_when_close_fails():
+    iterator = CloseFailingProviderIterator(
+        next_error=RuntimeError("provider failed")
+    )
+    client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
+
+    with pytest.raises(RuntimeError, match="provider failed"):
+        async for _ in client.stream_chat(
+            messages=[],
+            tools=[],
+            params=AgentParams(model="benchmark-model"),
+        ):
+            pass
+
+    assert len(client.timings) == 1
+
+
+@pytest.mark.asyncio
+async def test_timing_client_propagates_close_error_after_normal_completion():
+    iterator = CloseFailingProviderIterator()
+    client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
+
+    with pytest.raises(RuntimeError, match="close failed"):
+        async for _ in client.stream_chat(
+            messages=[],
+            tools=[],
+            params=AgentParams(model="benchmark-model"),
+        ):
+            pass
+
+    assert len(client.timings) == 1
+
+
+@pytest.mark.asyncio
+async def test_timing_client_propagates_close_error_when_consumer_closes_stream():
+    iterator = CloseFailingProviderIterator(
+        items=(StreamItem.raw_response_chunk({"provider": "first"}),)
+    )
+    client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
+    stream = client.stream_chat(
+        messages=[],
+        tools=[],
+        params=AgentParams(model="benchmark-model"),
+    )
+
+    assert (await anext(stream)).kind == "raw_chunk"
+    with pytest.raises(RuntimeError, match="close failed"):
+        await stream.aclose()
+
+    assert len(client.timings) == 1
+    assert client.timings[0].first_item_kind == "raw_chunk"
+
+
 @pytest.mark.asyncio
 async def test_timing_client_indexes_concurrent_calls_by_start_order():
     inner = ConcurrentTimingClient()