From 1dc0d28f7a707a9ab839db8e9b6fa84b28fc59a1 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 28 Sep 2026 16:00:47 -0400 Subject: [PATCH 1/6] fix(agentscope): trace TeamPipeline reply streams --- ...line_reply_stream_creates_parent_span.yaml | 377 ++++++++++++++++++ .../integrations/agentscope/integration.py | 2 + .../integrations/agentscope/patchers.py | 11 + .../agentscope/test_agentscope.py | 40 ++ .../integrations/agentscope/tracing.py | 69 +++- 5 files changed, 489 insertions(+), 10 deletions(-) create mode 100644 py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_team_pipeline_reply_stream_creates_parent_span.yaml diff --git a/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_team_pipeline_reply_stream_creates_parent_span.yaml b/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_team_pipeline_reply_stream_creates_parent_span.yaml new file mode 100644 index 000000000..a3a2bd0dd --- /dev/null +++ b/py/src/braintrust/integrations/agentscope/cassettes/latest/test_agentscope_team_pipeline_reply_stream_creates_parent_span.yaml @@ -0,0 +1,377 @@ +interactions: +- request: + body: '{"messages":[{"role":"system","name":"system","content":[{"type":"text","text":"Delegate + the task to the most suitable team member, then summarize their reply."}]},{"role":"user","name":"user","content":[{"type":"text","text":"Ask + the researcher to explain why recorded tests are useful."}]},{"role":"user","content":[{"type":"text","text":"Treat + the following as the ground truth at this point of the conversation. Anything + stated earlier is outdated, and a later reminder, if any, supersedes this one:\n2026-09-28T19:45:06\nUTC\n"}]}],"model":"gpt-4o-mini","stream":false,"temperature":0.0,"tools":[{"type":"function","function":{"name":"TeamAssign","description":"Assign + a task to a team member and get its reply back as the result. The members are:\n- + Researcher: Answers factual questions.","parameters":{"type":"object","properties":{"member":{"type":"string","enum":["Researcher"],"description":"The + name of the member to assign to."},"prompt":{"type":"string","description":"The + complete task for the member. It cannot see your context, so include everything + it needs."}},"required":["member","prompt"]}}}]}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '1200' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.31.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.31.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.7 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-ETBZbS4oyZ3ZJZnYeqsdCmSx0Zdbm\",\n \"object\": + \"chat.completion\",\n \"created\": 1790624707,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": null,\n \"tool_calls\": [\n {\n + \ \"id\": \"call_4xuSRzuMi2cQMnIxg8vkFlo4\",\n \"type\": + \"function\",\n \"function\": {\n \"name\": \"TeamAssign\",\n + \ \"arguments\": \"{\\\"member\\\":\\\"Researcher\\\",\\\"prompt\\\":\\\"Explain + why recorded tests are useful.\\\"}\"\n }\n }\n ],\n + \ \"refusal\": null,\n \"annotations\": []\n },\n \"logprobs\": + null,\n \"finish_reason\": \"tool_calls\"\n }\n ],\n \"usage\": + {\n \"prompt_tokens\": 203,\n \"completion_tokens\": 25,\n \"total_tokens\": + 228,\n \"prompt_tokens_details\": {\n \"cached_tokens\": 0,\n \"audio_tokens\": + 0\n },\n \"completion_tokens_details\": {\n \"reasoning_tokens\": + 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\": 0,\n + \ \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\": + \"default\",\n \"system_fingerprint\": \"fp_d9f006541a\"\n}\n" + headers: + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a4253f24d88660a9-YYZ + connection: + - keep-alive + content-length: + - '1146' + content-type: + - application/json + date: + - Mon, 28 Sep 2026 19:45:08 GMT + openai-processing-ms: + - '732' + openai-version: + - '2020-10-01' + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + x-content-type-options: + - nosniff + x-openai-proxy-wasm: + - v0.1 + x-ratelimit-limit-requests: + - '30000' + x-ratelimit-limit-tokens: + - '150000000' + x-ratelimit-remaining-requests: + - '29999' + x-ratelimit-remaining-tokens: + - '149999892' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_eb32bc8414f3416a994de81e6f4ba735 + status: + code: 200 + message: OK +- request: + body: '{"messages":[{"role":"system","name":"system","content":[{"type":"text","text":"Answer + factual questions concisely."}]},{"role":"user","name":"user","content":[{"type":"text","text":"Explain + why recorded tests are useful."}]},{"role":"user","content":[{"type":"text","text":"Treat + the following as the ground truth at this point of the conversation. Anything + stated earlier is outdated, and a later reminder, if any, supersedes this one:\n2026-09-28T19:45:08\nUTC\n"}]}],"model":"gpt-4o-mini","stream":false,"temperature":0.0}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '611' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.31.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.31.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.7 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-ETBZe2wI192o7cAeqL5gA5CcrwGD8\",\n \"object\": + \"chat.completion\",\n \"created\": 1790624710,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"Recorded tests are useful for several + reasons:\\n\\n1. **Consistency**: They provide a standardized way to assess + performance, ensuring that the same conditions are applied across different + test instances.\\n\\n2. **Reproducibility**: Recorded tests can be replayed + multiple times, allowing for consistent results and easier identification + of issues.\\n\\n3. **Analysis and Debugging**: They enable detailed analysis + of test outcomes, making it easier to identify failures and understand the + behavior of the system under test.\\n\\n4. **Documentation**: Recorded tests + serve as documentation of the testing process, providing a reference for future + tests and helping to maintain knowledge within a team.\\n\\n5. **Automation**: + They can be integrated into automated testing frameworks, allowing for efficient + and rapid testing cycles.\\n\\n6. **Training and Onboarding**: Recorded tests + can be used to train new team members, providing them with practical examples + of how the system behaves and how to test it effectively.\\n\\n7. **Performance + Monitoring**: They can help in monitoring performance over time, allowing + teams to track improvements or regressions in system behavior.\\n\\nOverall, + recorded tests enhance the reliability and efficiency of the testing process, + contributing to higher quality software.\",\n \"refusal\": null,\n + \ \"annotations\": []\n },\n \"logprobs\": null,\n \"finish_reason\": + \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 102,\n \"completion_tokens\": + 235,\n \"total_tokens\": 337,\n \"prompt_tokens_details\": {\n \"cached_tokens\": + 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\": + {\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\": + 0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\": + \"default\",\n \"system_fingerprint\": \"fp_1bcbffccf6\"\n}\n" + headers: + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a4253f2babc9dda9-YYZ + connection: + - keep-alive + content-length: + - '2108' + content-type: + - application/json + date: + - Mon, 28 Sep 2026 19:45:11 GMT + openai-processing-ms: + - '3248' + openai-version: + - '2020-10-01' + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + x-content-type-options: + - nosniff + x-openai-proxy-wasm: + - v0.1 + x-ratelimit-limit-requests: + - '30000' + x-ratelimit-limit-tokens: + - '150000000' + x-ratelimit-remaining-requests: + - '29999' + x-ratelimit-remaining-tokens: + - '149999910' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_379edea85a164833a7283488c359173b + status: + code: 200 + message: OK +- request: + body: '{"messages":[{"role":"system","name":"system","content":[{"type":"text","text":"Delegate + the task to the most suitable team member, then summarize their reply."}]},{"role":"user","name":"user","content":[{"type":"text","text":"Ask + the researcher to explain why recorded tests are useful."}]},{"role":"user","content":[{"type":"text","text":"Treat + the following as the ground truth at this point of the conversation. Anything + stated earlier is outdated, and a later reminder, if any, supersedes this one:\n2026-09-28T19:45:06\nUTC\n"}]},{"role":"assistant","name":"Leader","content":null,"tool_calls":[{"id":"call_4xuSRzuMi2cQMnIxg8vkFlo4","type":"function","function":{"name":"TeamAssign","arguments":"{\"member\":\"Researcher\",\"prompt\":\"Explain + why recorded tests are useful.\"}"}}]},{"role":"tool","tool_call_id":"call_4xuSRzuMi2cQMnIxg8vkFlo4","content":"Recorded + tests are useful for several reasons:\n\n1. **Consistency**: They provide a + standardized way to assess performance, ensuring that the same conditions are + applied across different test instances.\n\n2. **Reproducibility**: Recorded + tests can be replayed multiple times, allowing for consistent results and easier + identification of issues.\n\n3. **Analysis and Debugging**: They enable detailed + analysis of test outcomes, making it easier to identify failures and understand + the behavior of the system under test.\n\n4. **Documentation**: Recorded tests + serve as documentation of the testing process, providing a reference for future + tests and helping to maintain knowledge within a team.\n\n5. **Automation**: + They can be integrated into automated testing frameworks, allowing for efficient + and rapid testing cycles.\n\n6. **Training and Onboarding**: Recorded tests + can be used to train new team members, providing them with practical examples + of how the system behaves and how to test it effectively.\n\n7. **Performance + Monitoring**: They can help in monitoring performance over time, allowing teams + to track improvements or regressions in system behavior.\n\nOverall, recorded + tests enhance the reliability and efficiency of the testing process, contributing + to higher quality software.","name":"TeamAssign"}],"model":"gpt-4o-mini","stream":false,"temperature":0.0,"tools":[{"type":"function","function":{"name":"TeamAssign","description":"Assign + a task to a team member and get its reply back as the result. The members are:\n- + Researcher: Answers factual questions.","parameters":{"type":"object","properties":{"member":{"type":"string","enum":["Researcher"],"description":"The + name of the member to assign to."},"prompt":{"type":"string","description":"The + complete task for the member. It cannot see your context, so include everything + it needs."}},"required":["member","prompt"]}}}]}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '2847' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.31.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.31.0 + x-stainless-read-timeout: + - '600' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.7 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-ETBZgWD7fJWSLHPQNYzYNZLQG163Z\",\n \"object\": + \"chat.completion\",\n \"created\": 1790624712,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"The researcher explained that recorded + tests are useful for several reasons:\\n\\n1. **Consistency**: They ensure + standardized assessment conditions across different test instances.\\n2. **Reproducibility**: + Recorded tests can be replayed, allowing for consistent results and easier + issue identification.\\n3. **Analysis and Debugging**: They facilitate detailed + analysis of outcomes, helping to pinpoint failures and understand system behavior.\\n4. + **Documentation**: They serve as a reference for future tests, maintaining + knowledge within a team.\\n5. **Automation**: Recorded tests can be integrated + into automated testing frameworks for efficient testing cycles.\\n6. **Training + and Onboarding**: They provide practical examples for training new team members + on system behavior and testing.\\n7. **Performance Monitoring**: They help + track performance over time, identifying improvements or regressions.\\n\\nOverall, + recorded tests enhance the reliability and efficiency of the testing process, + contributing to higher quality software.\",\n \"refusal\": null,\n + \ \"annotations\": []\n },\n \"logprobs\": null,\n \"finish_reason\": + \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 473,\n \"completion_tokens\": + 184,\n \"total_tokens\": 657,\n \"prompt_tokens_details\": {\n \"cached_tokens\": + 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\": + {\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\": + 0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\": + \"default\",\n \"system_fingerprint\": \"fp_d9f006541a\"\n}\n" + headers: + access-control-expose-headers: + - X-Request-ID + - CF-Ray + - CF-Ray + alt-svc: + - h3=":443"; ma=86400 + cf-cache-status: + - DYNAMIC + cf-ray: + - a4253f4189e560a9-YYZ + connection: + - keep-alive + content-length: + - '1842' + content-type: + - application/json + date: + - Mon, 28 Sep 2026 19:45:13 GMT + openai-processing-ms: + - '1801' + openai-version: + - '2020-10-01' + server: + - cloudflare + strict-transport-security: + - max-age=31536000; includeSubDomains; preload + transfer-encoding: + - chunked + x-content-type-options: + - nosniff + x-openai-proxy-wasm: + - v0.1 + x-ratelimit-limit-requests: + - '30000' + x-ratelimit-limit-tokens: + - '150000000' + x-ratelimit-remaining-requests: + - '29999' + x-ratelimit-remaining-tokens: + - '149999567' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_ca64aca887924369a561803f27bfb4b5 + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/agentscope/integration.py b/py/src/braintrust/integrations/agentscope/integration.py index 48d07c27e..7eb7ffa19 100644 --- a/py/src/braintrust/integrations/agentscope/integration.py +++ b/py/src/braintrust/integrations/agentscope/integration.py @@ -8,6 +8,7 @@ ChatModelPatcher, FanoutPipelinePatcher, SequentialPipelinePatcher, + TeamPipelineReplyStreamPatcher, ToolkitCallToolFunctionPatcher, ToolkitCallToolPatcher, ) @@ -24,6 +25,7 @@ class AgentScopeIntegration(BaseIntegration): AgentReplyPatcher, SequentialPipelinePatcher, FanoutPipelinePatcher, + TeamPipelineReplyStreamPatcher, ToolkitCallToolFunctionPatcher, ToolkitCallToolPatcher, ChatModelPatcher, diff --git a/py/src/braintrust/integrations/agentscope/patchers.py b/py/src/braintrust/integrations/agentscope/patchers.py index a70bd716b..cc18972e2 100644 --- a/py/src/braintrust/integrations/agentscope/patchers.py +++ b/py/src/braintrust/integrations/agentscope/patchers.py @@ -8,6 +8,7 @@ _fanout_pipeline_wrapper, _model_call_wrapper, _sequential_pipeline_wrapper, + _team_pipeline_reply_stream_wrapper, _toolkit_call_tool_function_wrapper, ) @@ -48,6 +49,16 @@ class FanoutPipelinePatcher(FunctionWrapperPatcher): wrapper = _fanout_pipeline_wrapper +class TeamPipelineReplyStreamPatcher(FunctionWrapperPatcher): + """Patch AgentScope 2.x team pipeline streaming replies.""" + + name = "agentscope.pipeline.team_reply_stream" + target_module = "agentscope.pipeline" + target_path = "TeamPipeline.reply_stream" + version_spec = ">=2" + wrapper = _team_pipeline_reply_stream_wrapper + + class ToolkitCallToolFunctionPatcher(FunctionWrapperPatcher): """Patch AgentScope 1.x toolkit execution.""" diff --git a/py/src/braintrust/integrations/agentscope/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index ed58819b6..e8186e9d1 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -311,6 +311,46 @@ def answer(): assert tool_span["input"]["tool_name"] == "answer" +@pytest.mark.skipif(not IS_AGENTSCOPE_V2, reason="AgentScope 2.x TeamPipeline API") +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_agentscope_team_pipeline_reply_stream_creates_parent_span(memory_logger): + from agentscope.pipeline import TeamMember, TeamPipeline + + assert not memory_logger.pop() + + pipeline = TeamPipeline( + leader=_make_agent( + "Leader", + "Delegate the task to the most suitable team member, then summarize their reply.", + ), + members=[ + TeamMember( + agent=_make_agent("Researcher", "Answer factual questions concisely."), + description="Answers factual questions.", + ), + ], + ) + events = [ + event + async for event in pipeline.reply_stream( + _make_user_msg("Ask the researcher to explain why recorded tests are useful.") + ) + ] + + assert events + + spans = memory_logger.pop() + pipeline_span = next(span for span in spans if span["span_attributes"]["name"] == "TeamPipeline.reply_stream") + llm_spans = [span for span in spans if _span_type(span) == SpanTypeAttribute.LLM] + + assert _span_type(pipeline_span) == "task" + assert pipeline_span["context"]["span_origin"]["instrumentation"]["name"] == "agentscope-auto" + assert pipeline_span["output"] + assert llm_spans + assert all(pipeline_span["span_id"] in span["span_parents"] for span in llm_spans) + + def test_setup_agentscope_is_idempotent(): """Repeat setup calls must not double-wrap patched targets.""" from agentscope.model import OpenAIChatModel diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index 871b571fb..01f25486e 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -104,6 +104,20 @@ def _pipeline_metadata(args: Any, kwargs: dict[str, Any]) -> dict[str, Any]: return clean_nones({"agent_names": agent_names}) +def _team_pipeline_metadata(instance: Any) -> dict[str, Any]: + leader = getattr(instance, "leader", None) + members = getattr(instance, "members", None) + member_names = None + if isinstance(members, dict): + member_names = [member.agent.name for member in members.values() if getattr(member, "agent", None)] + return clean_nones( + { + "leader": getattr(leader, "name", None), + "member_names": member_names, + } + ) + + def _extract_metrics(*candidates: Any) -> dict[str, float] | None: for candidate in candidates: usage = _field_value(candidate, "usage") @@ -244,6 +258,7 @@ def _deferred_stream_trace( stack: contextlib.ExitStack, log_fn: Any, on_first_chunk: Any = None, + on_error: Any = None, ) -> Any: """Wrap an async iterator so the span stays open until the stream is consumed. @@ -257,20 +272,54 @@ async def _trace(): with deferred: last_chunk = None first_seen = False - async with aclosing(result) as agen: - async for chunk in agen: - if not first_seen: - first_seen = True - if on_first_chunk is not None: - on_first_chunk() - last_chunk = chunk - yield chunk - if last_chunk is not None: - log_fn(span, last_chunk) + try: + async with aclosing(result) as agen: + async for chunk in agen: + if not first_seen: + first_seen = True + if on_first_chunk is not None: + on_first_chunk() + last_chunk = chunk + yield chunk + if last_chunk is not None: + log_fn(span, last_chunk) + except Exception as exc: + if on_error is not None: + on_error(span, exc) + raise return _trace() +def _team_pipeline_reply_stream_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: dict[str, Any]) -> Any: + """Trace a TeamPipeline reply for the lifetime of its event stream.""" + with contextlib.ExitStack() as stack: + span = stack.enter_context( + start_span( + name="TeamPipeline.reply_stream", + type=SpanTypeAttribute.TASK, + input=_args_kwargs_input(args, kwargs), + metadata=_team_pipeline_metadata(instance), + ) + ) + try: + result = wrapped(*args, **kwargs) + if _is_async_iterator(result): + return _deferred_stream_trace( + result, + span, + stack, + lambda s, event: s.log(output=event), + on_error=lambda s, exc: s.log(error=exc), + ) + + span.log(output=result) + return result + except Exception as exc: + span.log(error=exc) + raise + + async def _toolkit_call_tool_function_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: dict[str, Any]) -> Any: tool_call = args[0] if args else kwargs.get("tool_call") tool_name = _tool_name(tool_call) From 34cc90904ecdd9a5ae6707ab91fa2d4b039cd062 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 28 Sep 2026 16:20:13 -0400 Subject: [PATCH 2/6] fix(agentscope): defer pipeline spans until iteration --- .../agentscope/test_agentscope.py | 17 +++++++ .../integrations/agentscope/tracing.py | 44 +++++++++---------- 2 files changed, 37 insertions(+), 24 deletions(-) diff --git a/py/src/braintrust/integrations/agentscope/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index e8186e9d1..f4b862c21 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -311,6 +311,23 @@ def answer(): assert tool_span["input"]["tool_name"] == "answer" +@pytest.mark.skipif(not IS_AGENTSCOPE_V2, reason="AgentScope 2.x TeamPipeline API") +@pytest.mark.asyncio +async def test_agentscope_unstarted_team_pipeline_reply_stream_does_not_open_span(memory_logger): + from agentscope.pipeline import TeamPipeline + + pipeline = TeamPipeline(leader=_make_agent("Leader", "Delegate tasks."), members=[]) + + with logger.start_span(name="caller") as parent: + stream = pipeline.reply_stream(_make_user_msg("hello")) + assert logger.current_span() is parent + await stream.aclose() + assert logger.current_span() is parent + + spans = memory_logger.pop() + assert [span["span_attributes"]["name"] for span in spans] == ["caller"] + + @pytest.mark.skipif(not IS_AGENTSCOPE_V2, reason="AgentScope 2.x TeamPipeline API") @pytest.mark.vcr @pytest.mark.asyncio diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index 01f25486e..6c7ae2757 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -293,31 +293,27 @@ async def _trace(): def _team_pipeline_reply_stream_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: dict[str, Any]) -> Any: """Trace a TeamPipeline reply for the lifetime of its event stream.""" - with contextlib.ExitStack() as stack: - span = stack.enter_context( - start_span( - name="TeamPipeline.reply_stream", - type=SpanTypeAttribute.TASK, - input=_args_kwargs_input(args, kwargs), - metadata=_team_pipeline_metadata(instance), - ) - ) - try: - result = wrapped(*args, **kwargs) - if _is_async_iterator(result): - return _deferred_stream_trace( - result, - span, - stack, - lambda s, event: s.log(output=event), - on_error=lambda s, exc: s.log(error=exc), - ) - span.log(output=result) - return result - except Exception as exc: - span.log(error=exc) - raise + async def _trace(): + with start_span( + name="TeamPipeline.reply_stream", + type=SpanTypeAttribute.TASK, + input=_args_kwargs_input(args, kwargs), + metadata=_team_pipeline_metadata(instance), + ) as span: + try: + last_event = None + async with aclosing(wrapped(*args, **kwargs)) as stream: + async for event in stream: + last_event = event + yield event + if last_event is not None: + span.log(output=last_event) + except Exception as exc: + span.log(error=exc) + raise + + return _trace() async def _toolkit_call_tool_function_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: dict[str, Any]) -> Any: From a3831fffc6a9dca858e2cf242685611867ef3910 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Mon, 28 Sep 2026 17:12:30 -0400 Subject: [PATCH 3/6] fix(agentscope): restore caller context between stream events --- .../agentscope/test_agentscope.py | 29 +++++++--- .../integrations/agentscope/tracing.py | 56 +++++++++++++------ 2 files changed, 60 insertions(+), 25 deletions(-) diff --git a/py/src/braintrust/integrations/agentscope/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index f4b862c21..d138a4839 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -348,24 +348,37 @@ async def test_agentscope_team_pipeline_reply_stream_creates_parent_span(memory_ ), ], ) - events = [ - event - async for event in pipeline.reply_stream( - _make_user_msg("Ask the researcher to explain why recorded tests are useful.") - ) - ] + with logger.start_span(name="caller") as caller_span: + stream = pipeline.reply_stream(_make_user_msg("Ask the researcher to explain why recorded tests are useful.")) + first_event = await stream.__anext__() + assert logger.current_span() is caller_span + with logger.start_span(name="after-first-event"): + pass + events = [first_event] + async for event in stream: + events.append(event) + assert logger.current_span() is caller_span assert events spans = memory_logger.pop() pipeline_span = next(span for span in spans if span["span_attributes"]["name"] == "TeamPipeline.reply_stream") - llm_spans = [span for span in spans if _span_type(span) == SpanTypeAttribute.LLM] + llm_spans = [ + span + for span in spans + if span["span_attributes"].get("type") is not None and _span_type(span) == SpanTypeAttribute.LLM + ] assert _span_type(pipeline_span) == "task" + assert pipeline_span["span_parents"] == [caller_span.span_id] assert pipeline_span["context"]["span_origin"]["instrumentation"]["name"] == "agentscope-auto" assert pipeline_span["output"] assert llm_spans - assert all(pipeline_span["span_id"] in span["span_parents"] for span in llm_spans) + assert all(pipeline_span["span_id"] in span["span_parents"] for span in llm_spans), [ + (span["span_attributes"]["name"], span["span_parents"]) for span in llm_spans + ] + after_event_span = next(span for span in spans if span["span_attributes"]["name"] == "after-first-event") + assert after_event_span["span_parents"] == [caller_span.span_id] def test_setup_agentscope_is_idempotent(): diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index 6c7ae2757..4a3a9168f 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -1,10 +1,11 @@ """AgentScope-specific span creation and stream aggregation.""" +import asyncio import contextlib import inspect import time from contextlib import aclosing -from contextvars import ContextVar +from contextvars import ContextVar, copy_context from typing import Any from braintrust.integrations.utils import ( @@ -295,23 +296,44 @@ def _team_pipeline_reply_stream_wrapper(wrapped: Any, instance: Any, args: Any, """Trace a TeamPipeline reply for the lifetime of its event stream.""" async def _trace(): - with start_span( - name="TeamPipeline.reply_stream", - type=SpanTypeAttribute.TASK, - input=_args_kwargs_input(args, kwargs), - metadata=_team_pipeline_metadata(instance), - ) as span: + stream_context = copy_context() + span = stream_context.run( + lambda: start_span( + name="TeamPipeline.reply_stream", + type=SpanTypeAttribute.TASK, + input=_args_kwargs_input(args, kwargs), + metadata=_team_pipeline_metadata(instance), + ) + ) + stream = None + try: + last_event = None + stream = wrapped(*args, **kwargs) + while True: + stream_context.run(span.set_current) + try: + event = await asyncio.create_task(stream.__anext__(), context=stream_context) + except StopAsyncIteration: + break + finally: + stream_context.run(span.unset_current) + last_event = event + yield event + if last_event is not None: + span.log(output=last_event) + except Exception as exc: + span.log(error=exc) + raise + finally: try: - last_event = None - async with aclosing(wrapped(*args, **kwargs)) as stream: - async for event in stream: - last_event = event - yield event - if last_event is not None: - span.log(output=last_event) - except Exception as exc: - span.log(error=exc) - raise + if stream is not None: + stream_context.run(span.set_current) + try: + await asyncio.create_task(stream.aclose(), context=stream_context) + finally: + stream_context.run(span.unset_current) + finally: + span.end() return _trace() From 83eadcfa5c09b946602c7996019a7fbfdf0b588a Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 29 Sep 2026 14:46:33 +0000 Subject: [PATCH 4/6] fix(agentscope): run TeamPipeline stream steps in a shared context `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 --- .../agentscope/test_agentscope.py | 28 ++++++++++ .../integrations/agentscope/tracing.py | 51 ++++++++++++++----- 2 files changed, 67 insertions(+), 12 deletions(-) diff --git a/py/src/braintrust/integrations/agentscope/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index d138a4839..5a1bf5fa3 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -381,6 +381,34 @@ async def test_agentscope_team_pipeline_reply_stream_creates_parent_span(memory_ assert after_event_span["span_parents"] == [caller_span.span_id] +@pytest.mark.asyncio +async def test_team_pipeline_reply_stream_shares_context_across_events(memory_logger): + """Context set by the wrapped stream in one step must be visible in the next.""" + from braintrust.integrations.agentscope.tracing import _team_pipeline_reply_stream_wrapper + + async def reply_stream(): + with logger.start_span(name="inner"): + yield "first" + with logger.start_span(name="inner-child"): + pass + yield "second" + + with logger.start_span(name="caller") as caller_span: + stream = _team_pipeline_reply_stream_wrapper(lambda *args, **kwargs: reply_stream(), None, (), {}) + events = [] + async for event in stream: + assert logger.current_span() is caller_span + events.append(event) + assert logger.current_span() is caller_span + + assert events == ["first", "second"] + + spans = {span["span_attributes"]["name"]: span for span in memory_logger.pop()} + assert spans["TeamPipeline.reply_stream"]["span_parents"] == [caller_span.span_id] + assert spans["inner"]["span_parents"] == [spans["TeamPipeline.reply_stream"]["span_id"]] + assert spans["inner-child"]["span_parents"] == [spans["inner"]["span_id"]] + + def test_setup_agentscope_is_idempotent(): """Repeat setup calls must not double-wrap patched targets.""" from agentscope.model import OpenAIChatModel diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index 4a3a9168f..e87eb0b32 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -1,11 +1,10 @@ """AgentScope-specific span creation and stream aggregation.""" -import asyncio import contextlib import inspect import time from contextlib import aclosing -from contextvars import ContextVar, copy_context +from contextvars import Context, ContextVar, copy_context from typing import Any from braintrust.integrations.utils import ( @@ -292,10 +291,44 @@ async def _trace(): return _trace() +class _AwaitInContext: + """Await ``coro`` with every step run inside ``context``. + + Like ``asyncio.create_task(coro, context=context)`` (Python 3.11+), but runs on + any Python version and without scheduling a separate task. + """ + + __slots__ = ("_context", "_coro") + + def __init__(self, coro: Any, context: Context): + self._coro = coro + self._context = context + + def __await__(self) -> Any: + coro, context = self._coro, self._context + step, value = coro.send, None + while True: + try: + yielded = context.run(step, value) + except StopIteration as exc: + return exc.value + try: + value = yield yielded + step = coro.send + except GeneratorExit: + context.run(coro.close) + raise + except BaseException as exc: + step, value = coro.throw, exc + + def _team_pipeline_reply_stream_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: dict[str, Any]) -> Any: """Trace a TeamPipeline reply for the lifetime of its event stream.""" async def _trace(): + # 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() span = stream_context.run( lambda: start_span( @@ -305,18 +338,16 @@ async def _trace(): metadata=_team_pipeline_metadata(instance), ) ) + stream_context.run(span.set_current) stream = None try: last_event = None - stream = wrapped(*args, **kwargs) + stream = stream_context.run(wrapped, *args, **kwargs) while True: - stream_context.run(span.set_current) try: - event = await asyncio.create_task(stream.__anext__(), context=stream_context) + event = await _AwaitInContext(stream.__anext__(), stream_context) except StopAsyncIteration: break - finally: - stream_context.run(span.unset_current) last_event = event yield event if last_event is not None: @@ -327,11 +358,7 @@ async def _trace(): finally: try: if stream is not None: - stream_context.run(span.set_current) - try: - await asyncio.create_task(stream.aclose(), context=stream_context) - finally: - stream_context.run(span.unset_current) + await _AwaitInContext(stream.aclose(), stream_context) finally: span.end() From 6365433f374959b3eb92c5adf7eebf460a00492f Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 29 Sep 2026 14:50:07 +0000 Subject: [PATCH 5/6] fix(agentscope): gate TeamPipeline patcher to agentscope>=2.0.9 TeamPipeline first shipped in 2.0.9; 2.0.0-2.0.8 have no such target. Co-Authored-By: Claude Opus 5.5 --- py/src/braintrust/integrations/agentscope/patchers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/py/src/braintrust/integrations/agentscope/patchers.py b/py/src/braintrust/integrations/agentscope/patchers.py index cc18972e2..a9a3651be 100644 --- a/py/src/braintrust/integrations/agentscope/patchers.py +++ b/py/src/braintrust/integrations/agentscope/patchers.py @@ -50,12 +50,12 @@ class FanoutPipelinePatcher(FunctionWrapperPatcher): class TeamPipelineReplyStreamPatcher(FunctionWrapperPatcher): - """Patch AgentScope 2.x team pipeline streaming replies.""" + """Patch AgentScope team pipeline streaming replies (``TeamPipeline`` was added in 2.0.9).""" name = "agentscope.pipeline.team_reply_stream" target_module = "agentscope.pipeline" target_path = "TeamPipeline.reply_stream" - version_spec = ">=2" + version_spec = ">=2.0.9" wrapper = _team_pipeline_reply_stream_wrapper From c3b70d3bffa5fd4906f68c04146e4735aeac457b Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 29 Sep 2026 15:00:03 +0000 Subject: [PATCH 6/6] fix(agentscope): log TeamPipeline output when stream closes early 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 --- .../agentscope/test_agentscope.py | 61 ++++++++++++++----- .../integrations/agentscope/tracing.py | 7 ++- 2 files changed, 50 insertions(+), 18 deletions(-) diff --git a/py/src/braintrust/integrations/agentscope/test_agentscope.py b/py/src/braintrust/integrations/agentscope/test_agentscope.py index 5a1bf5fa3..669273784 100644 --- a/py/src/braintrust/integrations/agentscope/test_agentscope.py +++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py @@ -311,6 +311,26 @@ def answer(): assert tool_span["input"]["tool_name"] == "answer" +TEAM_PIPELINE_PROMPT = "Ask the researcher to explain why recorded tests are useful." + + +def _make_team_pipeline(): + from agentscope.pipeline import TeamMember, TeamPipeline + + return TeamPipeline( + leader=_make_agent( + "Leader", + "Delegate the task to the most suitable team member, then summarize their reply.", + ), + members=[ + TeamMember( + agent=_make_agent("Researcher", "Answer factual questions concisely."), + description="Answers factual questions.", + ), + ], + ) + + @pytest.mark.skipif(not IS_AGENTSCOPE_V2, reason="AgentScope 2.x TeamPipeline API") @pytest.mark.asyncio async def test_agentscope_unstarted_team_pipeline_reply_stream_does_not_open_span(memory_logger): @@ -332,24 +352,11 @@ async def test_agentscope_unstarted_team_pipeline_reply_stream_does_not_open_spa @pytest.mark.vcr @pytest.mark.asyncio async def test_agentscope_team_pipeline_reply_stream_creates_parent_span(memory_logger): - from agentscope.pipeline import TeamMember, TeamPipeline - assert not memory_logger.pop() - pipeline = TeamPipeline( - leader=_make_agent( - "Leader", - "Delegate the task to the most suitable team member, then summarize their reply.", - ), - members=[ - TeamMember( - agent=_make_agent("Researcher", "Answer factual questions concisely."), - description="Answers factual questions.", - ), - ], - ) + pipeline = _make_team_pipeline() with logger.start_span(name="caller") as caller_span: - stream = pipeline.reply_stream(_make_user_msg("Ask the researcher to explain why recorded tests are useful.")) + stream = pipeline.reply_stream(_make_user_msg(TEAM_PIPELINE_PROMPT)) first_event = await stream.__anext__() assert logger.current_span() is caller_span with logger.start_span(name="after-first-event"): @@ -381,6 +388,30 @@ async def test_agentscope_team_pipeline_reply_stream_creates_parent_span(memory_ assert after_event_span["span_parents"] == [caller_span.span_id] +@pytest.mark.skipif(not IS_AGENTSCOPE_V2, reason="AgentScope 2.x TeamPipeline API") +@pytest.mark.vcr +@pytest.mark.asyncio +@pytest.mark.parametrize( + "vcr_cassette_name", ["test_agentscope_team_pipeline_reply_stream_creates_parent_span"], ids=["shared"] +) +async def test_agentscope_team_pipeline_reply_stream_closed_early_logs_output(memory_logger, vcr_cassette_name): + pipeline = _make_team_pipeline() + with logger.start_span(name="caller") as caller_span: + stream = pipeline.reply_stream(_make_user_msg(TEAM_PIPELINE_PROMPT)) + first_event = await stream.__anext__() + await stream.aclose() + assert logger.current_span() is caller_span + + assert first_event is not None + + spans = memory_logger.pop() + pipeline_span = next(span for span in spans if span["span_attributes"]["name"] == "TeamPipeline.reply_stream") + assert pipeline_span["span_parents"] == [caller_span.span_id] + assert "end" in pipeline_span["metrics"] + assert pipeline_span["output"] + assert "error" not in pipeline_span + + @pytest.mark.asyncio async def test_team_pipeline_reply_stream_shares_context_across_events(memory_logger): """Context set by the wrapped stream in one step must be visible in the next.""" diff --git a/py/src/braintrust/integrations/agentscope/tracing.py b/py/src/braintrust/integrations/agentscope/tracing.py index e87eb0b32..c7d97ab83 100644 --- a/py/src/braintrust/integrations/agentscope/tracing.py +++ b/py/src/braintrust/integrations/agentscope/tracing.py @@ -340,8 +340,8 @@ async def _trace(): ) stream_context.run(span.set_current) stream = None + last_event = None try: - last_event = None stream = stream_context.run(wrapped, *args, **kwargs) while True: try: @@ -350,8 +350,6 @@ async def _trace(): break last_event = event yield event - if last_event is not None: - span.log(output=last_event) except Exception as exc: span.log(error=exc) raise @@ -360,6 +358,9 @@ async def _trace(): if stream is not None: await _AwaitInContext(stream.aclose(), stream_context) finally: + # Also runs when the caller closes the stream early. + if last_event is not None: + span.log(output=last_event) span.end() return _trace()