fix(moa): carry completed responses through the managed Relay path

Under managed Relay execution the provider factory runs lazily inside
provider_stream() on the Relay session's event loop. The MoA facade's
auxiliary call_llm(stream=True) is invoked from that callback, so the
eager final_response check added for the non-managed path never fires:
the inner ManagedLlmStream is returned to the outer stream, which then
synchronously iterates it on the same loop thread and dies with
RuntimeError: Cannot run the event loop while another loop is running
(the completed response effectively trapped one level deeper).

stream_current() now detects a running event loop and returns the raw
factory result instead of nesting a ManagedLlmStream: the outer managed
stream already provides Relay tracking for the enclosing attempt, and
its own completed_response_predicate traps the completed response as
final_response — the same contract the main streaming worker consumes
(chat_completion_helpers reads stream.final_response after the chunk
loop). Nested managed streams remain supported for genuinely streaming
providers via the outer stream's own iteration.

Adds managed-execution regressions using the retained relay_turn
fixture: direct completed-response trapping, and the nested
facade-shaped stream_current call.
This commit is contained in:
AKAZIK-py 2026-07-31 11:53:17 +08:00 committed by Teknium
parent 52705e496f
commit c23ff21d7c
No known key found for this signature in database
2 changed files with 93 additions and 0 deletions

View File

@ -247,6 +247,14 @@ async def execute_current_async(
)
def _has_running_event_loop() -> bool:
try:
asyncio.get_running_loop()
except RuntimeError:
return False
return True
def stream_current(
request: dict[str, Any],
stream_factory: Callable[[dict[str, Any]], Any],
@ -272,6 +280,17 @@ def stream_current(
turn = relay_runtime.active_turn()
if turn is None:
return stream_factory(request)
if _has_running_event_loop():
# Managed provider callbacks execute on the Relay session's event
# loop. A nested ManagedLlmStream built here would be synchronously
# iterated on that same loop thread, which asyncio forbids
# ("Cannot run the event loop while another loop is running").
# Return the raw factory result instead: the outer managed stream
# already provides Relay tracking for the enclosing attempt, and its
# own completed_response_predicate traps a completed response (e.g.
# the MoA facade's auxiliary ``call_llm(stream=True)`` returning a
# full response when an adapter ignores ``stream=True``).
return stream_factory(request)
managed = stream(
request,
stream_factory,

View File

@ -795,3 +795,77 @@ def test_stream_current_streams_iterators_with_predicate(tmp_path, monkeypatch):
relay_runtime.SESSION_COORDINATOR.release_conversation(lease)
relay_runtime._reset_for_tests()
def _completed_response(content: str = "done") -> SimpleNamespace:
return SimpleNamespace(
model="test-model",
choices=[
SimpleNamespace(
message=SimpleNamespace(
role="assistant",
content=content,
tool_calls=None,
),
finish_reason="stop",
)
],
usage=None,
)
def _choices_predicate(value) -> bool:
return hasattr(value, "choices")
def test_stream_managed_traps_direct_completed_response(relay_turn):
"""Managed path: a factory returning a completed response (adapter
ignoring stream=True) is trapped as final_response instead of iterated."""
relay, turn = relay_turn
del relay, turn
stream = relay_llm.stream(
{"model": "test-model", "messages": []},
lambda request: _completed_response(),
session_id="session-1",
name="test-provider",
model_name="test-model",
finalizer=lambda: {},
completed_response_predicate=_choices_predicate,
)
assert list(stream) == []
assert stream.final_response is not None
assert stream.final_response.choices[0].message.content == "done"
def test_stream_current_inside_managed_callback_returns_raw(relay_turn):
"""Managed path: an auxiliary stream_current() call made from inside a
managed provider callback (the MoA facade's call_llm(stream=True) shape)
must return the raw factory result; the outer stream traps a completed
response as its final_response instead of crashing on a nested event
loop or surfacing an empty stream."""
relay, turn = relay_turn
del relay, turn
def outer_factory(request):
return relay_llm.stream_current(
{"model": "test-model", "messages": []},
lambda inner_request: _completed_response(),
name="moa-aggregator",
model_name="test-model",
finalizer=lambda: {},
completed_response_predicate=_choices_predicate,
)
stream = relay_llm.stream(
{"model": "test-model", "messages": []},
outer_factory,
session_id="session-1",
name="moa",
model_name="test-model",
finalizer=lambda: {},
completed_response_predicate=_choices_predicate,
)
assert list(stream) == []
assert stream.final_response is not None
assert stream.final_response.choices[0].message.content == "done"