🐛 MethodProxy breaks async generator methods — streaming and observability are mutually exclusive

src/xaibo/core/exchange.py

@@ -1,3 +1,4 @@
+import inspect
 from itertools import chain
 from collections import defaultdict
 from typing import Type, Union, Any
@@ -258,18 +259,29 @@ def _emit_event(self, event_type: EventType, result=None, arguments=None, except
                     print("Exception during event handling")
                     traceback.print_exc()
 
-    async def __call__(self, *args, **kwargs):
+    def __call__(self, *args, **kwargs):
         """Forward calls to the wrapped method.
-
+
+        Dispatches on the kind of the wrapped method so the proxy stays
+        transparent: coroutine methods return a coroutine to await, generator
+        methods return a wrapping generator to iterate.
+
         Args:
             *args: Positional arguments to pass to wrapped method
             **kwargs: Keyword arguments to pass to wrapped method
-            
+
         Returns:
             The result of calling the wrapped method
         """
+        if inspect.isasyncgenfunction(self._method):
+            return self._call_async_generator(*args, **kwargs)
+        if inspect.isgeneratorfunction(self._method):
+            return self._call_generator(*args, **kwargs)
+        return self._call_async(*args, **kwargs)
+
+    async def _call_async(self, *args, **kwargs):
         self._call_id += 1
-        
+
         # Emit call event
         self._emit_event(
             EventType.CALL,
@@ -296,6 +308,68 @@ async def __call__(self, *args, **kwargs):
 
         return result
 
+    async def _call_async_generator(self, *args, **kwargs):
+        self._call_id += 1
+
+        # Emit call event
+        self._emit_event(
+            EventType.CALL,
+            arguments={"args": args, "kwargs": kwargs}
+        )
+
+        chunks = 0
+        try:
+            async for chunk in self._method(*args, **kwargs):
+                chunks += 1
+                yield chunk
+        except GeneratorExit:
+            # Consumer stopped iterating early; not an error
+            raise
+        except:
+            self._emit_event(
+                EventType.EXCEPTION,
+                exception=traceback.format_exc()
+            )
+            traceback.print_exc()
+            raise
+
+        # Emit result event
+        self._emit_event(
+            EventType.RESULT,
+            result={"stream": True, "chunks": chunks}
+        )
+
+    def _call_generator(self, *args, **kwargs):
+        self._call_id += 1
+
+        # Emit call event
+        self._emit_event(
+            EventType.CALL,
+            arguments={"args": args, "kwargs": kwargs}
+        )
+
+        chunks = 0
+        try:
+            for chunk in self._method(*args, **kwargs):
+                chunks += 1
+                yield chunk
+        except GeneratorExit:
+            # Consumer stopped iterating early; not an error
+            raise
+        except:
+            self._emit_event(
+                EventType.EXCEPTION,
+                exception=traceback.format_exc()
+            )
+            traceback.print_exc()
+            raise
+
+        # Emit result event
+        self._emit_event(
+            EventType.RESULT,
+            result={"stream": True, "chunks": chunks}
+        )
+
     def __repr__(self):
         return f"MethodProxy({self._method.__name__})"
 

tests/core/test_proxy_generator_methods.py

@@ -0,0 +1,160 @@
+import pytest
+
+from xaibo.core.exchange import Proxy
+from xaibo.core.models.events import Event, EventType
+from xaibo.primitives.modules.llm.mock import MockLLM
+from xaibo.core.models.llm import LLMMessage
+
+
+class DummyStreamer:
+    async def stream_method(self, count, prefix="chunk"):
+        for i in range(count):
+            yield f"{prefix}-{i}"
+
+    async def failing_stream(self):
+        yield "first"
+        raise ValueError("stream broke")
+
+    def sync_stream(self, count):
+        for i in range(count):
+            yield i
+
+    async def regular_method(self):
+        return "hello"
+
+
+@pytest.mark.asyncio
+async def test_proxy_async_generator_streams_chunks():
+    events = []
+
+    def event_handler(event: Event):
+        events.append(event)
+
+    obj = DummyStreamer()
+    proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module")
+
+    chunks = [chunk async for chunk in proxy.stream_method(3, prefix="foo")]
+    assert chunks == ["foo-0", "foo-1", "foo-2"]
+
+    # Should have generated 2 events (call and result)
+    assert len(events) == 2
+
+    call_event = events[0]
+    assert call_event.event_type == EventType.CALL
+    assert call_event.module_class == "DummyStreamer"
+    assert call_event.method_name == "stream_method"
+    assert call_event.arguments == {"args": (3,), "kwargs": {"prefix": "foo"}}
+    assert call_event.agent_id == "test-agent"
+
+    result_event = events[1]
+    assert result_event.event_type == EventType.RESULT
+    assert result_event.method_name == "stream_method"
+    assert result_event.result == {"stream": True, "chunks": 3}
+    assert result_event.call_id == call_event.call_id
+
+
+@pytest.mark.asyncio
+async def test_proxy_async_generator_without_listeners():
+    obj = DummyStreamer()
+    proxy = Proxy(obj)
+
+    chunks = [chunk async for chunk in proxy.stream_method(2)]
+    assert chunks == ["chunk-0", "chunk-1"]
+
+
+@pytest.mark.asyncio
+async def test_proxy_async_generator_exception_mid_stream():
+    events = []
+
+    def event_handler(event: Event):
+        events.append(event)
+
+    obj = DummyStreamer()
+    proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module")
+
+    received = []
+    with pytest.raises(ValueError, match="stream broke"):
+        async for chunk in proxy.failing_stream():
+            received.append(chunk)
+
+    assert received == ["first"]
+
+    event_types = [e.event_type for e in events]
+    assert event_types == [EventType.CALL, EventType.EXCEPTION]
+    assert "stream broke" in events[1].exception
+
+
+@pytest.mark.asyncio
+async def test_proxy_async_generator_early_close():
+    events = []
+
+    def event_handler(event: Event):
+        events.append(event)
+
+    obj = DummyStreamer()
+    proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module")
+
+    gen = proxy.stream_method(10)
+    assert await anext(gen) == "chunk-0"
+    await gen.aclose()
+
+    # Early close is not an error and the stream never completed
+    event_types = [e.event_type for e in events]
+    assert event_types == [EventType.CALL]
+
+
+@pytest.mark.asyncio
+async def test_proxy_sync_generator():
+    events = []
+
+    def event_handler(event: Event):
+        events.append(event)
+
+    obj = DummyStreamer()
+    proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module")
+
+    chunks = list(proxy.sync_stream(3))
+    assert chunks == [0, 1, 2]
+
+    event_types = [e.event_type for e in events]
+    assert event_types == [EventType.CALL, EventType.RESULT]
+    assert events[1].result == {"stream": True, "chunks": 3}
+
+
+@pytest.mark.asyncio
+async def test_proxy_regular_method_unchanged():
+    events = []
+
+    def event_handler(event: Event):
+        events.append(event)
+
+    obj = DummyStreamer()
+    proxy = Proxy(obj, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module")
+
+    result = await proxy.regular_method()
+    assert result == "hello"
+
+    event_types = [e.event_type for e in events]
+    assert event_types == [EventType.CALL, EventType.RESULT]
+    assert events[1].result == "hello"
+
+
+@pytest.mark.asyncio
+async def test_proxy_llm_generate_stream():
+    """The flagship case: LLMProtocol.generate_stream through the observability proxy."""
+    events = []
+
+    def event_handler(event: Event):
+        events.append(event)
+
+    llm = MockLLM({"responses": [{"content": "Hello World"}], "streaming_chunk_size": 5})
+    proxy = Proxy(llm, event_listeners=[("", event_handler)], agent_id="test-agent", caller_id="test-caller", module_id="test-module")
+
+    messages = [LLMMessage.user("Hi")]
+    chunks = [chunk async for chunk in proxy.generate_stream(messages)]
+
+    assert "".join(chunks) == "Hello World"
+    event_types = [e.event_type for e in events]
+    assert event_types == [EventType.CALL, EventType.RESULT]
+    assert events[1].result["stream"] is True
+    assert events[1].result["chunks"] == len(chunks)