Repository navigation
feat(agui): add a shared reader for handler streams - #164
Conversation
|
bugbot run |
jeffdupont
left a comment
There was a problem hiding this comment.
Reviewed at b7a9f3a. pytest packages/client/tests: 1536 passed. Findings 1–4 and 6 were reproduced with probe scripts against this branch; AG-UI behaviour was checked against the published @ag-ui/client@1.0.2 / @ag-ui/core@1.0.2.
#161, #163, #162 and #165 all build on this reader, so fixing these here fixes them for every provider.
Must fix
1. Every run reuses the same message ids, so chat UIs merge replies (agui.py:110-112, agui.py:326-328)
TurnEmitter._mint and RunReader._mint count from 1 per run, so two runs both emit msg-1, toolmsg-2. @ag-ui/client reuses an existing message on TEXT_MESSAGE_START with a known id, and its message list survives across runs, so the second answer is appended to the first.
Fix: f"{prefix}-{uuid.uuid4()}" in both _mints (uuid is already imported).
2. With capture_content off (the default), TOOL_CALL_RESULT fails the AG-UI schema (agui.py:194-210, agui.py:509-511)
Content is omitted entirely when gated, but ToolCallResultEventSchema in @ag-ui/core@1.0.2 makes content required. Probe output: {'type': 'TOOL_CALL_RESULT', 'messageId': 'toolmsg-1', 'toolCallId': 'a'}.
Fix: send "content": "" when gated; in RunReader, fill "" when the handler sent no string. (Unverified for transports that skip the client's schema parse.)
3. RUN_FINISHED can be emitted while a tool call is still open (agui.py:350-379, agui.py:492-500)
finish closes open text but not open tool calls: TOOL_CALL_START then finish gives RUN_STARTED, TOOL_CALL_START, RUN_FINISHED, which @ag-ui/client rejects ("Cannot send 'RUN_FINISHED' while tool calls are still active"). _tool_result also accepts a result before END without adding the END, unlike TurnEmitter.tool_result (agui.py:192).
Fix: in finish, emit TOOL_CALL_END for _started - _ended; in _tool_result, emit END first if missing. ~6 lines plus one test.
4. Partially streamed arguments are published truncated, and the fix currently lives in two downstream PRs (agui.py:164-184)
tool_end(arguments=...) ignores the full arguments once any delta was streamed: '{"q":' then end gives invalid JSON. #161 and #163 each add the same _arg_text suffix fix to this file, so whichever merges second conflicts, and #165 stays on the old behaviour.
Fix: land that change here once, with a test, and drop it from #161/#163. Note the prefix check only works when the provider's JSON is byte-identical to compact json.dumps (see #163 review), so the docstring's "cannot publish truncated JSON" promise needs weakening or a "streamed text wins" fallback.
5. No spec, no JS counterpart, unregistered usage id (agui.py:537)
ai-sdks-monorepo has no AG-UI spec (nothing for agui/ag-ui/RUN_STARTED). client.agui.stream is not in api-surface/usage-ids.txt (TESTING.md §3.27 requires it in the same PR), and no api-surface list names agui.stream. No js-ai-sdk PR or branch exists.
Fix: monorepo spec section for the event contract (ordering, capture gate, error behaviour incl. no RUN_ERROR, toolFailed metadata), register the id/surface, and a JS ticket.
Should fix
6. Closing the stream early doesn't close the handler's generator (agui.py:556; same pattern client.py:234, :256, tracking.py:310)
No aclosing around the inner loops (only bind_conversation_id adds it when a conversation id is bound, conversation.py:197-199). Probe: take two events, aclose(), and the handler's finally never runs and no duration/success/error metric is tracked. Pre-existing for config().stream(), but AG-UI endpoints are where clients disconnect. test_closing_early_does_not_finish_the_run (test_agui.py:458) doesn't catch it.
Fix: async with aclosing(...) around the three loops; extend the test to assert the handler's finally ran.
7. Callers can't pass the client's threadId / runId (agui.py:307-309, :523-530)
An AG-UI server receives these in RunAgentInput and should echo them; here each turn is a fresh thread. Fix: optional thread_id= / run_id= on stream(), uuid fallback.
8. The judge call is copy-pasted and the copies already differ (agui.py:571-585 vs client.py:263-279)
The agui copy checks response is a string; the client copy doesn't. Fix: one shared _judge_stream(...).
Nice to have
- Dead guards:
if not started,if done is None,and sink(agui.py:567-572) can't flip, sinceexecute_and_streamalways yieldsdoneor raises. - "Drop chunk after text" is implemented three times (
tracking.py:311-316,agui.py:240-242,agui.py:382); keep it intracking.py. - A non-string delta (
{"delta": 5}) becomes a non-string chunktext(agui.py:248-251); add anisinstance(delta, str)check. - Every provider repeats "flush pending,
end_text(),fail_open()" at several exits; aTurnEmitter.close()returning both would remove most of that.
Cross-provider: the provider PRs disagree on (a) whether text closes before a tool starts (#161 closes it, #163/#165 keep it open with parentMessageId), (b) the fallback when a call has no id (tool name vs ""), and (c) what tool-result content is (raw value, str(), _tool_result_text). Worth pinning one rule in this PR's docstrings or the spec.
Verdict: fix 1–5 first. Not checked: mypy; whether every @ag-ui/client transport runs the schema parse behind #2.
|
Addressed the review on Must fix:
Should fix: |
|
bugbot run |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit c8d198c. Configure here.
jeffdupont
left a comment
There was a problem hiding this comment.
Re-checked at c8d198c. Findings 1–4, 7, 8 and 11 are fixed. One new issue from this change, plus one note:
- Blocking:
aclosing(stream_gen)(inline) needsaclose(). A handler whosestream()returns a plain async iterator, not an async generator, now raisesAttributeErrorwhen the loop exits. - Not blocking, already on
main: if the caller closes the stream early,GeneratorExitskipsexcept Exception, so neither success nor error is tracked.mainbehaves the same (tracking.py:323), so a follow-up is fine.
| if event.get("output") is not None: | ||
| full_text = event["output"] | ||
| raw_usage = event.get("usage", {}) | ||
| async with aclosing(stream_gen) as agen: |
There was a problem hiding this comment.
contextlib.aclosing calls aclose() on exit, and a plain async iterator doesn't have it, so this raises AttributeError for a handler that returns one. Close it only when it can be closed, e.g. if hasattr(stream_gen, "aclose"), or wrap a plain iterator in an async generator first. A test with a class that implements only __aiter__/__anext__ would lock it in.
|
Fixed in
The early-close / |
|
@cursor review |
Handlers can yield the AG-UI text and tool subset. config().stream() still yields chunk then done. stream() is the opt-in entry. Co-authored-by: alexis <alexis@launchdarkly.com>
A handler that raises mid-call can still report the lifecycle without the tool result. Co-authored-by: alexis <alexis@launchdarkly.com>
Message ids are unique across runs, tool results always carry content, and a run cannot finish while a tool call is open. Argument text is completed only when it extends what was streamed. Disconnects close the handler, and callers can echo threadId and runId. Co-authored-by: alexis <alexis@launchdarkly.com>
contextlib.aclosing always awaits aclose. A handler may return a plain async iterator, and that call raised AttributeError after the loop had already succeeded. The tracker now closes the stream only when aclose is present. Co-authored-by: alexis <alexis@launchdarkly.com>
ae82b2c to
37c7387
Compare
Handlers can yield the AG-UI text and tool subset.
config().stream()still yieldschunkthendone.launchdarkly_ai_server.agui.streamis the opt-in reader: it brackets a run withRUN_STARTEDandRUN_FINISHED, andcapture_contentgates argument deltas and tool results.Verified with
pytest packages/client/testsandmypy packages/client/src.