fix(llmobs): eagerly finalize stream to prevent nested agent spans

[ncybul]

Description

The claude agent sdk closes async generators lazily which can lead to unwanted nested agent spans. This PR attempts to fix this by finishing the agent span when the ResultMessage chunk is received.

Testing

Repro script:

import asyncio
from claude_agent_sdk import (
    ClaudeSDKClient,
    ResultMessage,
)


async def run_query(client, prompt):
    await client.query(prompt)
    async for msg in client.receive_messages():
        if isinstance(msg, ResultMessage):
            break


async def main():
    async with ClaudeSDKClient() as client:
        await run_query(client, "What is 2+2?")
        await run_query(client, "What is 3+3?")


asyncio.run(main())

Before

The first agent span is left open indefinitely, leading to the second agent span being treated as its child (nested underneath it within the same trace).

After

Agent spans are eagerly closed when ResultMessage is received, leading to two separate traces.

Risks

Additional Notes

[datadog-prod-us1-6]

 

✨ Fix all issues with BitsAI

⚠️ Warnings

🚦 8 Pipeline jobs failed

DataDog/apm-reliability/dd-trace-py | build linux: [arm64, cp315-cp315, v113741589-d2b8243-musllinux_1_2_aarch64]     

🔧 Fix in code (Fix with Cursor).

NotImplementedError: This version of CPython is not supported yet.

DataDog/apm-reliability/dd-trace-py | build linux serverless: [arm64, cp315-cp315, v113741357-d2b8243-manylinux2014_aarch64, 1]     

🔄 Retry job. This looks flaky and may succeed on retry.

0/554 nodes are available: insufficient cpu and memory. Several nodes had untolerated taint(s) and did not match Pod's affinity/selector.

DataDog/apm-reliability/dd-trace-py | build linux: [amd64, cp315-cp315, v113741238-d2b8243-manylinux2014_x86_64]     

🔄 Retry job. This looks flaky and may succeed on retry.

Startup probe failed: HTTP probe failed with statuscode: 503

View all 8 failed jobs.

ℹ️ Info

No other issues found (see more)

🧪 All tests passed
❄️ No new flaky tests detected

Useful? React with 👍 / 👎

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 24f78a6 | Docs | Datadog PR Page | Give us feedback!}

[cit-pr-commenter-54b7da]

Codeowners resolved as

ddtrace/contrib/internal/claude_agent_sdk/_streaming.py                 @DataDog/ml-observability
releasenotes/notes/fix-claude-finalize-stream-f9837e6baceccfae.yaml     @DataDog/apm-python
tests/contrib/claude_agent_sdk/test_claude_agent_sdk_llmobs.py          @DataDog/ml-observability

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 24f78a60e1

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[sabrenner]

code lgtm, i can QA in a bit but don't block on that. just one dumb question, and responding to codex's - i think we can ignore it for now, but if i'm reading it wrong, feel free to address it and i can re-review!