Skip to content

πŸ› MethodProxy breaks async generator methods β€” streaming and observability are mutually exclusive#37

Closed
MFA-X-AI wants to merge 1 commit into
mainfrom
fahreza/method-proxy
Closed

πŸ› MethodProxy breaks async generator methods β€” streaming and observability are mutually exclusive#37
MFA-X-AI wants to merge 1 commit into
mainfrom
fahreza/method-proxy

Conversation

@MFA-X-AI

Copy link
Copy Markdown
Member

Every module resolved through the exchange is wrapped in the observability proxy, and MethodProxy.__call__ unconditionally does await self._method(...). Async generator methods β€” LLMProtocol.generate_stream being the flagship β€” return an async generator object, which can't be awaited, so any streaming call through the proxy raises TypeError: object async_generator can't be used in 'await' expression. Since the proxy wrapping is unconditional, streaming currently works only by bypassing the exchange entirely (e.g. unwrapping _obj), which silently exits the observability layer. This PR makes the proxy transparent for all four method shapes: plain async, async generator, and sync generator methods all forward correctly and emit CALL/RESULT/EXCEPTION events.

# before: TypeError through any proxied module
async for chunk in proxy.generate_stream(messages):  # πŸ’₯
    ...

Changes

  • MethodProxy.__call__ is no longer unconditionally async β€” it dispatches on the wrapped method's kind via inspect.isasyncgenfunction / inspect.isgeneratorfunction and returns a matching wrapper (coroutine, async generator, or sync generator)
  • Generator wrappers emit CALL before the first chunk, EXCEPTION (then re-raise) if the stream fails mid-flight, and RESULT with {"stream": True, "chunks": n} after exhaustion β€” chunk contents are not recorded, keeping event volume bounded
  • Early consumer close (GeneratorExit) is treated as normal termination, not an error
  • Added tests/core/test_proxy_generator_methods.py covering async/sync generators, mid-stream exceptions, early close, listener-free operation, unchanged plain-method behavior, and MockLLM.generate_stream end-to-end 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.

MethodProxy.__call__ unconditionally awaited the wrapped method, so any
async generator method (e.g. LLMProtocol.generate_stream) called through
the observability proxy raised TypeError. Dispatch on the method kind at
call time and return a matching wrapper that emits CALL/EXCEPTION/RESULT
around iteration, with RESULT carrying a chunk count.
@MFA-X-AI MFA-X-AI closed this Jul 13, 2026
@MFA-X-AI
MFA-X-AI deleted the fahreza/method-proxy branch July 13, 2026 07:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant