Repository navigation
fix(agentscope): trace TeamPipeline reply streams - #817
Abhijeet Prasad (AbhiPrasad) merged 6 commits into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1dc0d28f7a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "Codex (@codex) review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "Codex (@codex) address that feedback".
| with contextlib.ExitStack() as stack: | ||
| span = stack.enter_context( | ||
| start_span( | ||
| name="TeamPipeline.reply_stream", | ||
| type=SpanTypeAttribute.TASK, |
There was a problem hiding this comment.
Because reply_stream returns a lazy async generator, this opens and activates the span before any stream code runs, while _deferred_stream_trace transfers its exit callback into another lazy generator. If a caller creates the stream but never iterates it—or calls aclose() before the first iteration—the deferred generator body never runs, so the span is never ended and remains the current parent for subsequent work in that context. Open the span inside the tracing generator so an abandoned, unstarted stream cannot leak an active span.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 34cc90904e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "Codex (@codex) review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "Codex (@codex) address that feedback".
| async with aclosing(wrapped(*args, **kwargs)) as stream: | ||
| async for event in stream: | ||
| last_event = event | ||
| yield event |
There was a problem hiding this comment.
When a caller breaks out of async for after receiving an event and retains the stream without explicitly calling aclose(), Python leaves the async generator suspended at this yield, so the start_span context remains active and the pipeline span remains unfinished. Subsequent work in that task is then incorrectly parented to TeamPipeline.reply_stream until the stream is closed or finalized. Restore the caller's span context before yielding each event, while activating the pipeline span only while advancing the underlying iterator.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a3831fffc6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "Codex (@codex) review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "Codex (@codex) address that feedback".
| except StopAsyncIteration: | ||
| break | ||
| finally: | ||
| stream_context.run(span.unset_current) |
There was a problem hiding this comment.
When an instrumented operation remains open across a TeamPipeline event yield—for example, a streaming model iterator—its current-span token is suspended in stream_context. This unconditional unset_current() resets the context to the state before the pipeline token, discarding that child span; the next iteration then installs the pipeline directly, so subsequent work is incorrectly parented and the child eventually cleans up against stale context. Preserve the copied context's post-yield state while keeping it isolated from the caller instead of resetting and reinstalling the pipeline span for every event.
Useful? React with 👍 / 👎.
| name = "agentscope.pipeline.team_reply_stream" | ||
| target_module = "agentscope.pipeline" | ||
| target_path = "TeamPipeline.reply_stream" | ||
| version_spec = ">=2" |
There was a problem hiding this comment.
Is this maybe a bit wide of a version range?
`asyncio.create_task(..., context=)` is Python 3.11+ and failed pylint on 3.10. Replace it with a small awaitable that runs each coroutine step inside the stream's context, which works on every Python version and avoids scheduling a task per event. Also stop setting/unsetting the pipeline span around every step: the token reset discarded context set by the wrapped stream, so spans opened across events were parented to the pipeline span instead. The span is now made current once inside the private stream context. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
TeamPipeline first shipped in 2.0.9; 2.0.0-2.0.8 have no such target. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6365433f37
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "Codex (@codex) review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "Codex (@codex) address that feedback".
| if last_event is not None: | ||
| span.log(output=last_event) |
There was a problem hiding this comment.
When a caller receives one or more events and then calls aclose() (or exits an aclosing context), GeneratorExit jumps directly to finally, so this normal-exhaustion-only log is skipped and the ended pipeline span has no output despite data having been delivered. Log last_event during finalization as well, and cover the explicit-close provider path with cassette-backed coverage.
AGENTS.md reference: AGENTS.md:L30-L32
Useful? React with 👍 / 👎.
Closing the stream after receiving events skipped the output log, so the pipeline span ended with no output. Log the last event during finalization, and cover the early-close path by replaying the existing TeamPipeline cassette. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c3b70d3bff
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "Codex (@codex) review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "Codex (@codex) address that feedback".
| # The wrapped stream runs entirely in its own context, with the pipeline span | ||
| # current. Context changes it makes persist across events without leaking into | ||
| # the caller's context between them. | ||
| stream_context = copy_context() |
There was a problem hiding this comment.
Because copy_context() runs only when the lazy _trace generator is first advanced, a stream created under one span and consumed later under another is parented to the consumption-time span rather than the span active at the reply_stream() call. This commonly occurs when a producer creates the stream and hands it to another task or returns it from a span scope, and it produces an incorrect trace hierarchy. Capture the context in _team_pipeline_reply_stream_wrapper before returning _trace(), while still deferring span creation until iteration.
Useful? React with 👍 / 👎.
2 participants
Summary
TeamPipeline.reply_streamas a task span for the lifetime of its event stream.>=2.Linear: SDK-437 (sub-issue of SDK-426)
Validation
mise exec -- nox -s 'test_agentscope(latest)' -- --vcr-record=all -k test_agentscope_team_pipeline_reply_stream_creates_parent_span(failed because the pipeline span was missing).mise exec -- nox -s 'test_agentscope(latest)' -- -k test_agentscope_team_pipeline_reply_stream_creates_parent_span(passed).mise exec -- nox -s 'test_agentscope(latest)' 'test_agentscope(1.0.0)'(both sessions passed).