✨ Support streaming methods through the observability proxy

[MFA-X-AI]

Background

My research project uses Xaibo for the agent configs / orchestration, and a custom frontend. Since observability is the main sell there, I decided implementing a standard protocol would be best, hence landed on AG-UI. Basically a server adapter that streams tool calls, steps, and reasoning to any AG-UI client. Thought it would be great if we could put xaibo on the same integration list as LangGraph, CrewAI, Pydantic AI, and the Claude Agent SDK. To get there, there are a few PRs needed (richer generate_stream events → orchestrator event vocabulary → AGUIAdapter), and all of it flows through proxied streaming calls. This fix is prerequisite 1: without it, an AG-UI adapter could stream or be observable, not both.

Current situation

Every module resolved through the exchange is unconditionally wrapped in the observability proxy, and MethodProxy.__call__ always does await self._method(...) — so async generator methods like LLMProtocol.generate_stream raise TypeError: object async_generator can't be used in 'await' expression, meaning no module can consume another module's streaming method through dependency injection at all. The only workaround is reaching past the proxy (getattr(self.llm, "_obj", self.llm)), which silently exits the observability layer for exactly the calls you most want visibility into. The proxy just needs to speak all four Python calling conventions (plain, async, generator, async generator) instead of one.

# before: TypeError through any proxied module
async for chunk in proxy.generate_stream(messages):  # 💥
    ...

Examples from Other Frameworks

Observability layers that wrap LLM calls all hit this and converged on the same fix — dispatch on function kind at wrap time:

• LangSmith @traceable: checks inspect.isasyncgenfunction / isgeneratorfunction, yields through, closes the run on exhaustion
• OpenTelemetry GenAI instrumentation: proxies the returned stream, ends the span when exhausted
• General Python rule: a transparent wrapper must return the same shape it wraps

This PR applies the same pattern to MethodProxy.

Changes

• MethodProxy.__call__ dispatches on method kind and returns a matching wrapper: coroutine, async generator, or sync generator
• Generator wrappers emit CALL on start, EXCEPTION on mid-stream failure, RESULT with {"stream": True, "chunks": n} on exhaustion; early close (GeneratorExit) is normal termination
• Added tests/core/test_proxy_generator_methods.py (7 tests, incl. MockLLM.generate_stream through a Proxy)

Technical Details

• The dispatch must happen at call time, not inside the old async def __call__: a coroutine that returns an async generator would force callers into await-then-iterate, breaking duck-typing against unproxied modules. Plain def __call__ returning the un-awaited coroutine is behaviorally identical for existing callers (await proxy.method(...) works unchanged).
• Existing event semantics for plain async methods are untouched — same events, same payloads, same call_id pairing.
• No public API change; no caller anywhere needs to change.

Verification

uv run pytest tests/core/ -v   # 7 new tests + existing proxy/event tests all pass
uv run pytest --continue-on-collection-errors   # 144 passed, 18 skipped

Without the fix, 6 of the 7 new tests fail with the TypeError above; with it, all pass.

Note: tests/llms/test_anthropic.py::test_anthropic_array_schema_generation fails on main before this change (pre-existing, unrelated). The 12 collection errors are missing optional extras (fastapi, torch, …) and also pre-exist.

[MFA-X-AI]

Thanks for the check, seems like we missed it because it was focusing on enabling the generator. The latest commit changed both generator wrappers to accumulate chunks and emit them in the RESULT event:

{"stream": True, "chunks": 3, "content": "Hello World"}

• content is the joined string when every chunk is a str where chunk boundaries are a provider implementation detail, so the joined text is the faithful record. Non-string chunks (e.g. future typed stream events) are kept as a raw list instead; the proxy never guesses how to merge types it doesn't understand.
• Early close (GeneratorExit - consumer stopped iterating, client disconnected mid-stream) now also emits a RESULT with the partial content and "closed_early": True. Previously an early-closed stream left only a CALL in the log, indistinguishable from a hung one.

As raised in yesterday's call on the log volume concern - we mocked the event sequence each option would emit, using xaibo's Event model, for a ~500-token response streamed OpenAI-style (317 chunks of ~6 chars), and counted the bytes a persisting listener writes as JSON lines:

option	events/call	log size	vs non-streaming generate()
per-delta events	319	119,662 B	43.5×
count-only RESULT (as submitted)	2	789 B	0.3×, but content lost
accumulated RESULT (this change)	2	2,704 B	1.0×
baseline: non-streaming generate()	2	2,752 B	1.0×

The per-delta multiplier is envelope-driven (~370 B of Event envelope around a few-character delta), so it stays at tens-of-× under any realistic chunking - and it would run the synchronous listener loop between every token and the consumer. Accumulating stays 1× regardless and runs listeners once, after the stream ends.

Considered and deliberately left out of this PR:

• Per-chunk events / live token view - natural opt-in follow-up if the debug UI wants it; pairs with a delta-aware listener that updates latest-state instead of appending.
• Size cap for unbounded streams - the buffer equals the response the consumer is already holding, so it only matters for genuinely infinite generators, which xaibo doesn't have today. "truncated": True + cap is the standard fix if they appear.

Tests should be updated to pin the new payload, plus new cases for typed (non-string) chunks, empty stream, and the early-close partial result.

---

Heads-up: where this is going (so the follow-ups have context)

This is prerequisite 1 of the AG-UI plan from the standup. Sequence, dependency-ordered:

1. This PR - proxy speaks all four calling conventions; streaming + observability coexist.
2. Extend ResponseProtocol/Response (issue first, per the standup) - add surface for tool calls, steps, reasoning alongside text/attachments, so orchestrators can report more than final text. As mentioned in the standup, we'll extend the existing protocol. Design note to resolve in the issue: respond-calls are themselves proxied, so delta-shaped methods would def flood the event log, so the delta-aware logger idea should be the answer.
3. AGUIAdapter (server/adapters/agui.py, same pattern as openai.py) - translates those response events into AG-UI SSE via the official ag-ui-protocol package. This is the deliverable: any xaibo agent renders in any AG-UI frontend with tool calls and steps visible.
4. Typed generate_stream events (last, optional for v1) - only token-level text deltas need it; tool calls and usage are already available from non-streaming generate(), so the adapter can ship on whole-message granularity first.

[treo]

The latest commit changed both generator wrappers to accumulate chunks and emit them in the RESULT event:

I'm not sure we want that in this place. The observability aspect of Xaibo is supposed to give you everything that happens in the communication between two modules.

Yes, the default debug logger produces a lot of data that way. It is meant to capture everything that is going on. That is how it helps with debugging in complex situations.

If that overwhelms the viewer, then that is an issue with the viewer that needs to be resolved. I know that the current one needs work to be more useful for long traces.

If we want to have a more compact storage format, we can also do that. But that would just be a different event processor, not something that mangles the events before they can even be handled.

[MFA-X-AI]

Gotcha. Sorry I took yesterday's discussion about chunk-by-chunk logging being super heavy as, better not ship it without a proper handler, so having the proxy accumulate temporarily until that handler existed made sense. Since you're OK with the volume, the latest commit now emits every chunk as its own event. Since this only affects generator methods (today that's generate_stream), plain method calls produce the same event log as before. One thing I hope is aligned - since RESULT means "call completed", these commits updates the chunks to get a new EventType.YIELD instead of overloading it, so a streamed call now reads:

CALL → YIELD (chunk 0) → YIELD (chunk 1) → … → RESULT {"stream": true, "chunks": N}

where each YIELD carries the raw chunk as its result and is logged the instant the callee yields it, no merging or reshaping before listeners see it. Chunk boundaries and per-chunk arrival times are on the log (since each event already has time).

One thing I noticed was that the current implementation final RESULT keeps the count-only payload from the first version of this PR. For my self reference, the count is now thecompletion marker. "What was returned" lives on the YIELD events, and since a generator's output is its sequence of yields, there's no return value left to record. chunks doubles as an integrity check so a listener can verify it saw that many YIELDs. I think that makes sense.

The derived data (the joined text) is kept off the capture layer, but that's a judgment call: it means re-buffering in the proxy and duplicating what's already on the log, and the tradeoff is that a log/UI reader has to reassemble the content from YIELDs until a consolidating listener exists. If you'd rather have the joined content on RESULT too (looking at other frameworks, MLflow records both, raw chunks as events plus the concatenated output), let me know.

• Early close (GeneratorExit) still emits RESULT with "closed_early": true, so a consumer-abandoned stream is distinguishable from a hung one.
• A mid-stream exception now leaves the chunks that escaped before the failure on the log - previously they were lost entirely.

I've added tests that pin the new sequence, typed (non-string) chunks, empty stream, early close, and mid-stream failure for both the async and sync generator wrappers, so every branch in each has coverage.

Additional Notes

One small UI change rode along: the debug UI's call-grouping assumed exactly two events per call_id - open on the first, close on whatever arrives second. A streamed call would have broken it:

CALL           → group opens
YIELD (chunk)  → treated as the closing event → group closes early, console warning
RESULT         → no open group left → dropped, console warning

YIELD events are now skipped in grouping, so the trace view stays intact — but nothing renders them yet.

What gets recorded — every chunk hits the log the moment it's yielded:

sequenceDiagram
    participant M as llm.generate_stream()
    participant P as MethodProxy
    participant L as event log

    P->>L: CALL
    M-->>P: yield "Hel"
    P->>L: YIELD "Hel"
    M-->>P: yield "lo"
    P->>L: YIELD "lo"
    M-->>P: stream ends
    P->>L: RESULT {stream: true, chunks: 2}

Loading

What the viewer renders today — just the grouped bookends:

sequenceDiagram
    participant C as caller
    participant M as llm module

    C->>M: generate_stream(messages)
    M-->>C: RESULT {stream: true, chunks: 2}
    Note over C,M: YIELD rows skipped — stream content not shown (yet)

Loading

So stream content stays raw-log-only until either the joined text goes onto RESULT (the MLflow-style option above) or the consolidating listener lands. Rendering yields in the viewer itself (a latest-state view, so a 300-chunk stream isn't 300 rows). As you've mentioned, we can put that into another PR for trace-volume ergonomics. The same for the compact storage format, a compacting/delta-aware event processor as a separate listener, not in the proxy.

On the log-volume numbers, some comparsion from other frameworks for the capture-everything / compact-downstream split:

• LangChain: the callback bus fires per-token on_llm_new_token at the capture layer — every delta reaches every subscribed handler, same role as our proxy.
• LangSmith: its tracer is just another handler on that bus, and it accumulates into a single run at persistence — the role our consolidating listener would play.
• OTel GenAI SIG: had the same volume debate (Capture GenAI prompts and completions as events or attributes open-telemetry/semantic-conventions#2010) and settled it in 2025 the other way — full message content goes on span attributes, i.e. accumulated onto the completion record. Per-chunk streaming capture is the part still open there (#1964).

This will come up again in the ResponseProtocol extension: orchestrators will call the response handler once per delta (text chunk, tool-call step), and those calls go through the same proxy — so the same volume question lands there, with the same answer: the capture layer stays complete, compaction lives in the listener.