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..a9a3651be 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 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.0.9"
+ 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..669273784 100644
--- a/py/src/braintrust/integrations/agentscope/test_agentscope.py
+++ b/py/src/braintrust/integrations/agentscope/test_agentscope.py
@@ -311,6 +311,135 @@ 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):
+ 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
+async def test_agentscope_team_pipeline_reply_stream_creates_parent_span(memory_logger):
+ assert not memory_logger.pop()
+
+ 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__()
+ 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["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), [
+ (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]
+
+
+@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."""
+ 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 871b571fb..c7d97ab83 100644
--- a/py/src/braintrust/integrations/agentscope/tracing.py
+++ b/py/src/braintrust/integrations/agentscope/tracing.py
@@ -4,7 +4,7 @@
import inspect
import time
from contextlib import aclosing
-from contextvars import ContextVar
+from contextvars import Context, ContextVar, copy_context
from typing import Any
from braintrust.integrations.utils import (
@@ -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,16 +272,96 @@ 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()
+
+
+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(
+ name="TeamPipeline.reply_stream",
+ type=SpanTypeAttribute.TASK,
+ input=_args_kwargs_input(args, kwargs),
+ metadata=_team_pipeline_metadata(instance),
+ )
+ )
+ stream_context.run(span.set_current)
+ stream = None
+ last_event = None
+ try:
+ stream = stream_context.run(wrapped, *args, **kwargs)
+ while True:
+ try:
+ event = await _AwaitInContext(stream.__anext__(), stream_context)
+ except StopAsyncIteration:
+ break
+ last_event = event
+ yield event
+ except Exception as exc:
+ span.log(error=exc)
+ raise
+ finally:
+ try:
+ 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()