fix(sdk): keep the workflow span current across a streamed agent run - #5708
Open
mmabrouk wants to merge 1 commit into
Open
fix(sdk): keep the workflow span current across a streamed agent run#5708mmabrouk wants to merge 1 commit into
mmabrouk wants to merge 1 commit into
Conversation
Since commit 8e99171 (2026-07-07) an agent run on the streaming path has recorded no tokens and no cost on its workflow span. Measured: zero out of roughly 1,300 streaming runs across two independent stacks since 2026-07-08. Batch runs were never affected, which is why the bug stayed invisible. That commit added an SSE keepalive so a proxy would not drop a connection during a silent gap, such as a tool call that runs for minutes. To race a keepalive timer against each chunk, it began driving the response generator with one asyncio task per pull. Every task gets its own copy of the context. The instrumentation attaches the OpenTelemetry context and activates the workflow span from inside the generator body, so that activation landed in the first pull's private copy and every later pull started from a copy that had never seen it. Chunk one recorded. Everything after it, including the `finally` block where usage is written, ran against a non-recording span and was silently discarded. Two changes. The stream now captures one context when the generator is first driven and creates every pull task with it, so all pulls share a single context. That restores the property the old plain iteration had, keeps the existing lock-step backpressure with no read-ahead, and leaves cancellation and cleanup alone. It also removes the "Token created in a different Context" detach errors visible in the pre-fix logs. `record_usage` no longer depends on the ambient context. The handler captures the workflow span at entry and passes that reference, so this class of bug cannot silence usage again. The existing early return on a falsy total and the truthy-cost guard are unchanged, and an override that takes only the usage argument still works. Tests: 1,905 passing through the canonical runner, and 100 passing in the service's own agent tests, which confirms the recorder seam stays backward compatible. The two new test files fail against the unfixed code and pass against the fix. The stream test asserts the same recording span is current on every chunk, across a keepalive gap, and that no frame is dropped or duplicated. Claude-Session: https://claude.ai/code/session_01RkWWQUNNzRbaB5jnCAdjYA
dosubot
Bot
added
the
size:M
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
dosubot
Bot
added
bug
|
Important Review skippedNo new commits to review since the last review. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change binds usage recording to the workflow span and preserves execution context for Vercel SSE pulls. Tests cover streaming, batch execution, keepalive frames, cancellation, span attributes, and legacy recorders. ChangesWorkflow span usage recording
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant AgentHandler
participant WorkflowSpan
participant UsageRecorder
participant SSEAdapter
participant UpstreamIterator
AgentHandler->>WorkflowSpan: capture active workflow span
AgentHandler->>UsageRecorder: create span-bound recorder
SSEAdapter->>WorkflowSpan: preserve execution context
SSEAdapter->>UpstreamIterator: pull stream item
UpstreamIterator-->>SSEAdapter: item or keepalive timeout
AgentHandler->>UsageRecorder: record usage
UsageRecorder->>WorkflowSpan: write usage attributes
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Member
Author
|
This is one of three independent fixes for the same reported problem: an agent run shows a cost in the playground and nothing in the trace. They can be reviewed and merged separately, in any order.
Live end-to-end verification of all three together on a running stack is in progress, and I will post the result here. @coderabbitai review |
This was referenced Aug 3, 2026
|
✅ Action performedReview finished.
|
Contributor
Railway Preview Environment
|
This was referenced Aug 3, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The symptom
An agent run on the streaming path records no tokens and no cost on its workflow span. Batch runs are fine, which is why this stayed invisible.
Measured: zero out of roughly 1,300 streaming runs across two independent stacks since 2026-07-08. The playground still showed a cost, because the playground reads the number the harness streams to the browser, not the trace.
The cause
Commit 8e99171 (2026-07-07, "Add time-based run limits to the runner") added an SSE keepalive so a proxy would not drop the connection during a silent gap, such as a tool call that runs for minutes. To race a keepalive timer against each chunk, it changed
async for chunk in aiterinto a hand-driven loop that wraps every pull inasyncio.ensure_future(iterator.__anext__()).Every asyncio task gets its own copy of the context. The instrumentation attaches the OpenTelemetry context and activates the workflow span from inside the generator body, so that activation landed in the first pull's private copy. Every later pull started from a copy that had never seen it.
Chunk one recorded. Everything after it ran against a
NonRecordingSpan, including thefinallyblock whererecord_usagewrites usage. The writes were silently discarded.Observed directly in a test against the unfixed code:
The fix
Share one context across every pull. The stream captures
contextvars.copy_context()once when the generator is first driven and creates every pull task with it. That restores the property plain iteration had, where the whole upstream ran in the consumer's single context. Backpressure stays lock-step with no read-ahead, and cancellation and cleanup are untouched. It also removes the "Token created in a different Context" detach errors that show up in the pre-fix logs.Stop depending on ambient context for usage. The handler captures the workflow span at entry and passes that reference into
record_usage, so this class of bug cannot silence usage again. The existing early return on a falsy total and the truthy-cost guard are unchanged. An override that takes only the usage argument still works, so the service's own recorders are unaffected.Verification
run-tests.py --layer unit: 1,905 passed, 10 xfailed.Notes for the reviewer
decorators/routing.pydo not use a per-pull task and are unaffected._bind_workflow_spaninspects the recorder's signature to stay compatible with(usage)-only overrides. The cleaner version widens the callback type, which would have meant touchingservices/osscallers and was left out of scope.finallythat awaits may be truncated. Worth a separate look.Related
Part of a set of three independent fixes for the same reported problem, that an agent run shows a cost in the playground and none in the trace. The other two are the API ingest fix and the runner usage fix.