diff --git a/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_subagent_transfer_does_not_log_generator_exit.yaml b/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_subagent_transfer_does_not_log_generator_exit.yaml new file mode 100644 index 00000000..be37db94 --- /dev/null +++ b/py/src/braintrust/integrations/adk/cassettes/1.14.1/test_adk_subagent_transfer_does_not_log_generator_exit.yaml @@ -0,0 +1,150 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "What is the capital of France? Delegate + this to the specialist."}], "role": "user"}], "systemInstruction": {"parts": + [{"text": "You cannot answer questions yourself. For every request, immediately + call transfer_to_agent with agent_name=''capital_specialist''.\n\nYou are an + agent. Your internal name is \"coordinator\".\n\n The description about you + is \"Routes geography questions to the capital specialist without answering + them.\"\n\n\nYou have a list of other agents to transfer to:\n\n\nAgent name: + capital_specialist\nAgent description: The only agent allowed to answer geography + questions.\n\n\nIf you are the best to answer the question according to your + description, you\ncan answer it.\n\nIf another agent is better for answering + the question according to its\ndescription, call `transfer_to_agent` function + to transfer the\nquestion to that agent. When transferring, do not generate + any text other than\nthe function call.\n\n**NOTE**: the only available agents + for `transfer_to_agent` function are `capital_specialist`.\n"}], "role": "user"}, + "tools": [{"functionDeclarations": [{"description": "Transfer the question to + another agent.\n\nThis tool hands off control to another agent when it''s more + suitable to\nanswer the user''s question according to the agent''s description.\n\nArgs:\n agent_name: + the agent name to transfer to.\n", "name": "transfer_to_agent", "parameters": + {"properties": {"agent_name": {"type": "STRING"}}, "required": ["agent_name"], + "type": "OBJECT"}}]}], "generationConfig": {}}' + headers: + Content-Type: + - application/json + user-agent: + - google-genai-sdk/1.75.0 gl-python/3.14.6 google-adk/1.14.1 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/1.75.0 gl-python/3.14.6 google-adk/1.14.1 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent + response: + body: + string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": + [\n {\n \"functionCall\": {\n \"name\": \"transfer_to_agent\",\n + \ \"args\": {\n \"agent_name\": \"capital_specialist\"\n + \ }\n },\n \"thoughtSignature\": \"CsgDARFNMg9EuJxGSb6Rw2/4urHga+Zn9WrSEdnTxJhJfb0a/9BQdsCh4EzVx5DHs4AYOpfVkx8GKuSmv3g3DQuQhQtaPqVURt+5PHAQ53BQO04DSUnyVOcyp4PYJVJc+eP6zdMPCyNfJiE1hSDPLdma0+0L1H5dirwmjNLVmehQe8OuE39nVqtm3BYQaZ+eyTXRIno9P1YCi4//rtrciIYqzq5GDhpwBoRB/A+wwC1pu5BSJK3vzV/OvS4vFapAwptwkd3iKpAvOh4jlwTElrN+wAsASVGuZ70Fg1nmFrrdCBAq0yW4wOJv7mY4OOrPFYAJFkEJ24on8eo8qR4bLVn5MoYTnBMWQqR4QtbK4dv9bdrKEfmV926a+lB9kOyxCahYATKEklEiq24OKGcH+zaB4hY5fXArGKOpfxYvj7fS1CIXxvFl7I5Ur3E4V2PVIhstsDyGIaue9mBHP2oLfg+CGOmLNz666gPYkrnwZna6gfkxYLDxvMVFpLnIauLWdITWcQbwOQcaTkRkfH/wMir8XKK3loisYHarOVMOqjerNk9hA+MMWki/ZIhEpX3uIqDz0ymJR4WoDHh0BP7PLWvVpHmg/zLA+rvf\"\n + \ }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": + \"STOP\",\n \"index\": 0,\n \"finishMessage\": \"Model generated + function call(s).\"\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": + 296,\n \"candidatesTokenCount\": 22,\n \"totalTokenCount\": 409,\n \"promptTokensDetails\": + [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 296\n + \ }\n ],\n \"thoughtsTokenCount\": 91,\n \"serviceTier\": \"standard\"\n + \ },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"lwFqaoiON-WrjrEPrbb88Ak\"\n}\n" + headers: + Alt-Svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + Content-Type: + - application/json; charset=UTF-8 + Date: + - Wed, 29 Jul 2026 13:35:20 GMT + Server: + - scaffolding on HTTPServer2 + Server-Timing: + - gfet4t7; dur=765 + Transfer-Encoding: + - chunked + Vary: + - Origin + - X-Origin + - Referer + X-Content-Type-Options: + - nosniff + X-Frame-Options: + - SAMEORIGIN + X-Gemini-Service-Tier: + - standard + X-XSS-Protection: + - '0' + content-length: + - '1439' + status: + code: 200 + message: OK +- request: + body: '{"contents": [{"parts": [{"text": "What is the capital of France? Delegate + this to the specialist."}], "role": "user"}, {"parts": [{"text": "For context:"}, + {"text": "[coordinator] called tool `transfer_to_agent` with parameters: {''agent_name'': + ''capital_specialist''}"}], "role": "user"}, {"parts": [{"text": "For context:"}, + {"text": "[coordinator] `transfer_to_agent` tool returned result: {''result'': + None}"}], "role": "user"}], "systemInstruction": {"parts": [{"text": "Answer + geography questions accurately and in one short sentence.\n\nYou are an agent. + Your internal name is \"capital_specialist\".\n\n The description about you + is \"The only agent allowed to answer geography questions.\"\n\n\nYou have a + list of other agents to transfer to:\n\n\nAgent name: coordinator\nAgent description: + Routes geography questions to the capital specialist without answering them.\n\n\nIf + you are the best to answer the question according to your description, you\ncan + answer it.\n\nIf another agent is better for answering the question according + to its\ndescription, call `transfer_to_agent` function to transfer the\nquestion + to that agent. When transferring, do not generate any text other than\nthe function + call.\n\n**NOTE**: the only available agents for `transfer_to_agent` function + are `coordinator`.\n\nIf neither you nor the other agents are best for the question, + transfer to your parent agent coordinator.\n"}], "role": "user"}, "tools": [{"functionDeclarations": + [{"description": "Transfer the question to another agent.\n\nThis tool hands + off control to another agent when it''s more suitable to\nanswer the user''s + question according to the agent''s description.\n\nArgs:\n agent_name: the + agent name to transfer to.\n", "name": "transfer_to_agent", "parameters": {"properties": + {"agent_name": {"type": "STRING"}}, "required": ["agent_name"], "type": "OBJECT"}}]}], + "generationConfig": {}}' + headers: + Content-Type: + - application/json + user-agent: + - google-genai-sdk/1.75.0 gl-python/3.14.6 google-adk/1.14.1 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/1.75.0 gl-python/3.14.6 google-adk/1.14.1 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent + response: + body: + string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": + [\n {\n \"text\": \"The capital of France is Paris.\",\n + \ \"thoughtSignature\": \"CrcDARFNMg+U1J/4jHXZRfFf7mqUm/nZzRdTIDuD9yyxroAVdvYxYiIlZc4++vCuNqT5XQJ+N5SfjvN26U1xReX9a7uMicxcpLZRhs+KZr4L1pSVFas3u6W8ylpFFl5F5FauQqY1tpMLGhF0bLp/MqzOSVtSfHp7ZXfWFNh7KIhodg6UDYlN9sGEE/jzLDJaYGjJBFYGo0MuXuv3nm8Fi2om9MVQ3XOVEUVnPtRKg4vJRrL6YuL1i9A2YO1tO11iVepTOQoGJ5v7I0rJhii9BhhGC/uV8kIvBo2XbxWxdNKTnEQfmhMPeZp1MQWm0kdwJvBAJnEBoPFqeafM1q555MnaMlo2HVj9qjgqUtLXwnrNLp1b6esIBuUqrEyAD7U8VQVJq51i0Ppo6fD0VJRkkbQkTvKWehgHkwipz3lVkpD0NAf3++Ylz+u39hBF0MS7mj8P+Jo2RDY8TnWdY+mK2TpvU0bJgcPyIsk61gNGw74F/HmmH0DGfkhMuwB0jxy/EZgUE50OIZeCJEX9AD62g/CncsPQK+/WvsW4ejxNdtfcQnqPo4WWM7KMacF4y1Ha18qcI1KM+FnGMw==\"\n + \ }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": + \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": + 350,\n \"candidatesTokenCount\": 7,\n \"totalTokenCount\": 436,\n \"promptTokensDetails\": + [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 350\n + \ }\n ],\n \"thoughtsTokenCount\": 79,\n \"serviceTier\": \"standard\"\n + \ },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"mAFqasSWMqOLjrEP1dG6-Qg\"\n}\n" + headers: + Alt-Svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + Content-Type: + - application/json; charset=UTF-8 + Date: + - Wed, 29 Jul 2026 13:35:21 GMT + Server: + - scaffolding on HTTPServer2 + Server-Timing: + - gfet4t7; dur=728 + Transfer-Encoding: + - chunked + Vary: + - Origin + - X-Origin + - Referer + X-Content-Type-Options: + - nosniff + X-Frame-Options: + - SAMEORIGIN + X-Gemini-Service-Tier: + - standard + X-XSS-Protection: + - '0' + content-length: + - '1234' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_subagent_transfer_does_not_log_generator_exit.yaml b/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_subagent_transfer_does_not_log_generator_exit.yaml new file mode 100644 index 00000000..2584a61b --- /dev/null +++ b/py/src/braintrust/integrations/adk/cassettes/latest/test_adk_subagent_transfer_does_not_log_generator_exit.yaml @@ -0,0 +1,172 @@ +interactions: +- request: + body: '{"contents": [{"parts": [{"text": "What is the capital of France? Delegate + this to the specialist."}], "role": "user"}], "systemInstruction": {"parts": + [{"text": "You cannot answer questions yourself. For every request, immediately + call transfer_to_agent with agent_name=''capital_specialist''.\n\nYou are an + agent. Your internal name is \"coordinator\". The description about you is \"Routes + geography questions to the capital specialist without answering them.\".\n\n\nYou + have a list of other agents to transfer to:\n\n\nAgent name: capital_specialist\nAgent + description: The only agent allowed to answer geography questions.\n\n\nIf you + are the best to answer the question according to your description,\nyou can + answer it.\n\nIf another agent is better for answering the question according + to its\ndescription, call `transfer_to_agent` function to transfer the question + to that\nagent. When transferring, do not generate any text other than the function\ncall.\n\n**NOTE**: + the only available agents for `transfer_to_agent` function are\n`capital_specialist`.\n"}], + "role": "user"}, "tools": [{"functionDeclarations": [{"description": "Transfer + the query to another agent.\n\nUse this tool to hand off control to another + agent that is more suitable to\nanswer the user''s query according to the agent''s + description.\n\nArgs:\n agent_name: the agent name to transfer to.", "name": + "transfer_to_agent", "parameters_json_schema": {"properties": {"agent_name": + {"title": "Agent Name", "type": "string", "enum": ["capital_specialist"]}}, + "required": ["agent_name"], "title": "transfer_to_agentParams", "type": "object"}}]}], + "generationConfig": {}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '1646' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.14.0 gl-python/3.14.6 google-adk/2.5.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.14.0 gl-python/3.14.6 google-adk/2.5.0 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent + response: + body: + string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": + [\n {\n \"functionCall\": {\n \"name\": \"transfer_to_agent\",\n + \ \"args\": {\n \"agent_name\": \"capital_specialist\"\n + \ }\n },\n \"thoughtSignature\": \"CqwDARFNMg/j+9eaPm9CUpDwYKmqmHGZqfm7bYvgkuBJs+rzg1obPejEnFWaQf/v0ATSzFnaSan470RJspYICGrXp7jyuC/SAVL2v+vHaRvjXHkGq4Z4XxyohC0CJwuYNoxQZ0RAgNK9rkh0OXiv4FbyMOMed6ATAxYgDwV6ugRoDzdorj7x/lqQG1V9mWwK8viOhS0LQ76K48dIu7BMph5J2Jlr4EAO4JC8KhuAdWYE36PlLBQhz3KuS6AvWPwt8lmW4s0ZM+AwZZ6H1nU6hNlRTNUtiBVj5aU1qrbA/BoX/GUcqFfqHq+ymkK2mgW3HkPEgmv8Qry3UL1C+jzU1gfdL2ou2T8t8jZOzt9DO/wrINAhWx0Xx7w6NlfESCo1qQctLpXzGBgzkuzwVHAgJPP0eF38joQdN2VDH2cXZ2ao82Aq6uc/DU6CEASfoCq5aI5z4WELa4o9mzH3q+S6mc8dc81lTsFTz/4WN+G/3lp/KD9lUWE6eUTiettCwT+5x0kM5ymA3CT6iPf7Rz9NkmUqFUENjDrCngeZnN/muHqSUFIo7ZYCMdQ/Yn7i6ZQ=\"\n + \ }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": + \"STOP\",\n \"index\": 0,\n \"finishMessage\": \"Model generated + function call(s).\"\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": + 301,\n \"candidatesTokenCount\": 22,\n \"totalTokenCount\": 405,\n \"promptTokensDetails\": + [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 301\n + \ }\n ],\n \"thoughtsTokenCount\": 82,\n \"serviceTier\": \"standard\"\n + \ },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"fgFqapXpHLSrsOIPsOff0Ag\"\n}\n" + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '1403' + content-type: + - application/json; charset=UTF-8 + date: + - Wed, 29 Jul 2026 13:34:55 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=832 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + x-content-type-options: + - nosniff + x-frame-options: + - SAMEORIGIN + x-gemini-service-tier: + - standard + x-xss-protection: + - '0' + status: + code: 200 + message: OK +- request: + body: '{"contents": [{"parts": [{"text": "What is the capital of France? Delegate + this to the specialist."}], "role": "user"}, {"parts": [{"text": "For context:"}, + {"text": "[coordinator] called tool `transfer_to_agent` with parameters: {''agent_name'': + ''capital_specialist''}"}], "role": "user"}, {"parts": [{"text": "For context:"}, + {"text": "[coordinator] `transfer_to_agent` tool returned result: {''result'': + None}"}], "role": "user"}], "systemInstruction": {"parts": [{"text": "Answer + geography questions accurately and in one short sentence.\n\nYou are an agent. + Your internal name is \"capital_specialist\". The description about you is \"The + only agent allowed to answer geography questions.\".\n\n\nYou have a list of + other agents to transfer to:\n\n\nAgent name: coordinator\nAgent description: + Routes geography questions to the capital specialist without answering them.\n\n\nIf + you are the best to answer the question according to your description,\nyou + can answer it.\n\nIf another agent is better for answering the question according + to its\ndescription, call `transfer_to_agent` function to transfer the question + to that\nagent. When transferring, do not generate any text other than the function\ncall.\n\n**NOTE**: + the only available agents for `transfer_to_agent` function are\n`coordinator`.\n\nIf + neither you nor the other agents are best for the question, transfer to your + parent agent coordinator.\n"}], "role": "user"}, "tools": [{"functionDeclarations": + [{"description": "Transfer the query to another agent.\n\nUse this tool to hand + off control to another agent that is more suitable to\nanswer the user''s query + according to the agent''s description.\n\nArgs:\n agent_name: the agent name + to transfer to.", "name": "transfer_to_agent", "parameters_json_schema": {"properties": + {"agent_name": {"title": "Agent Name", "type": "string", "enum": ["coordinator"]}}, + "required": ["agent_name"], "title": "transfer_to_agentParams", "type": "object"}}]}], + "generationConfig": {}}' + headers: + accept: + - '*/*' + accept-encoding: + - gzip, deflate + connection: + - keep-alive + content-length: + - '1985' + content-type: + - application/json + host: + - generativelanguage.googleapis.com + user-agent: + - google-genai-sdk/2.14.0 gl-python/3.14.6 google-adk/2.5.0 gl-python/3.14.6 + x-goog-api-client: + - google-genai-sdk/2.14.0 gl-python/3.14.6 google-adk/2.5.0 gl-python/3.14.6 + method: POST + uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent + response: + body: + string: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": + [\n {\n \"text\": \"The capital of France is Paris.\",\n + \ \"thoughtSignature\": \"Cp0DARFNMg+zy0SfhDvFiWb+X7RwVYg2gND5/R8NYbIXVz2uXsKTzOu1zjMeMYCpJBlwGHeJAuWPG1X6+GKB35F2H9d1aJgdvCgjKC8iwXYjNDpkr3wPPmYZfzqB5xX7YvErtnJC3bLcDKaxtgfcrY+VgWhs5IW/ysrUXpiq+DJvX0RiwXgvmRg25sgGm/JRD8Q1hFQCFLwWrUAs8cg4fhSzHTfdSiP2t1WsmVOJrt7a7K6EiKjt88TAbrb8z7HTyxQtQJbNesMPIslrdIZGvGFSEy+no6KGNuN524eCH2/VYMvsQJJLxYYlYNoh6B7G6IkmfP64nqYbsiJcNxG4MJ3OvVsdLhMJsjBdtIfBsSF/lq14c1IZy2dj4rpa4JBmIityz3y1OSbFWcIeaGlaAmQjzpTu2XtdgeTfqopfPY8KLNN8ugzZgkpdOoq/RzOCebpgx9dUX6jr8xhlwbYBpbkfBapRWu5PKt2GkiKVUBUqWgT/RGh+rm4yLMyJ18nYwO0bDQVOdQNXeyRxuSIunvKi6xFCI1Jdheoa8iYfCTY=\"\n + \ }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": + \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": + 352,\n \"candidatesTokenCount\": 7,\n \"totalTokenCount\": 439,\n \"promptTokensDetails\": + [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 352\n + \ }\n ],\n \"thoughtsTokenCount\": 80,\n \"serviceTier\": \"standard\"\n + \ },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"fwFqapjNF4mkjrEP0rKCUQ\"\n}\n" + headers: + alt-svc: + - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000 + content-length: + - '1197' + content-type: + - application/json; charset=UTF-8 + date: + - Wed, 29 Jul 2026 13:34:56 GMT + server: + - scaffolding on HTTPServer2 + server-timing: + - gfet4t7; dur=896 + transfer-encoding: + - chunked + vary: + - Origin + - X-Origin + - Referer + x-content-type-options: + - nosniff + x-frame-options: + - SAMEORIGIN + x-gemini-service-tier: + - standard + x-xss-protection: + - '0' + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/adk/test_adk.py b/py/src/braintrust/integrations/adk/test_adk.py index bf4f7709..263d1c33 100644 --- a/py/src/braintrust/integrations/adk/test_adk.py +++ b/py/src/braintrust/integrations/adk/test_adk.py @@ -336,6 +336,65 @@ async def test_adk_braintrust_integration(memory_logger): assert "72" in response_output, "Response doesn't mention temperature" +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_adk_subagent_transfer_does_not_log_generator_exit(memory_logger): + """Successful LlmAgent delegation must not mark spans as failed during generator cleanup.""" + assert not memory_logger.pop() + + delegation_model = "gemini-2.5-flash" + specialist = Agent( + name="capital_specialist", + model=delegation_model, + description="The only agent allowed to answer geography questions.", + instruction="Answer geography questions accurately and in one short sentence.", + ) + coordinator = Agent( + name="coordinator", + model=delegation_model, + description="Routes geography questions to the capital specialist without answering them.", + instruction=( + "You cannot answer questions yourself. For every request, immediately call " + "transfer_to_agent with agent_name='capital_specialist'." + ), + sub_agents=[specialist], + ) + + app_name = "delegation_app" + user_id = "test-user" + session_id = "test-session-delegation" + runner = await _create_runner( + coordinator, + app_name=app_name, + user_id=user_id, + session_id=session_id, + ) + user_msg = types.Content( + role="user", + parts=[types.Part(text="What is the capital of France? Delegate this to the specialist.")], + ) + + events = [ + event + async for event in runner.run_async( + user_id=user_id, + session_id=session_id, + new_message=user_msg, + ) + ] + + assert any(event.actions.transfer_to_agent == specialist.name for event in events) + final_responses = [event for event in events if event.is_final_response()] + assert final_responses + assert "paris" in final_responses[-1].content.parts[0].text.lower() + + spans = memory_logger.pop() + assert spans + assert all("error" not in span for span in spans), [ + (span["span_attributes"]["name"], span.get("error")) for span in spans if "error" in span + ] + + @pytest.mark.vcr @pytest.mark.asyncio async def test_adk_nested_subagent_tool_calls_are_traced(memory_logger): diff --git a/py/src/braintrust/integrations/adk/tracing.py b/py/src/braintrust/integrations/adk/tracing.py index fa47fc97..9460ff79 100644 --- a/py/src/braintrust/integrations/adk/tracing.py +++ b/py/src/braintrust/integrations/adk/tracing.py @@ -1,10 +1,11 @@ """ADK-specific span creation, metadata extraction, stream handling, and output normalization.""" +import asyncio import contextvars import inspect import logging import time -from contextlib import aclosing +from contextlib import aclosing, contextmanager from functools import lru_cache from itertools import chain from typing import Any @@ -24,6 +25,20 @@ def start_span(*args, **kwargs): return _bt_start_span(*args, **kwargs) +@contextmanager +def _start_stream_span(*args, **kwargs): + """Keep ADK stream cleanup signals from being logged as span failures.""" + cleanup_signal = None + with start_span(*args, **kwargs) as span: + try: + yield span + except (GeneratorExit, asyncio.CancelledError) as exc: + cleanup_signal = exc + + if cleanup_signal is not None: + raise cleanup_signal + + from braintrust.span_types import SpanTypeAttribute @@ -353,7 +368,7 @@ def _run_in_context(*target_args: Any, **target_kwargs: Any) -> Any: async def _agent_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): async def _trace(): - with start_span( + with _start_stream_span( name=f"agent_run [{instance.name}]", type=SpanTypeAttribute.TASK, metadata={"agent_name": instance.name}, @@ -374,7 +389,7 @@ async def _trace(): async def _flow_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any): async def _trace(): - with start_span( + with _start_stream_span( name="call_llm", type=SpanTypeAttribute.TASK, metadata={"flow_class": instance.__class__.__name__}, @@ -413,7 +428,7 @@ async def _trace(): # Create span BEFORE execution so child spans (like mcp_tool) have proper parent # Start with generic name - we'll update it after we see the response - with start_span( + with _start_stream_span( name="llm_call", type=SpanTypeAttribute.LLM, input=captured_request, @@ -477,7 +492,7 @@ async def _runner_run_async_wrapper(wrapped: Any, instance: Any, args: Any, kwar serialized_message = _serialize_content(new_message) if new_message else None async def _trace(): - with start_span( + with _start_stream_span( name=f"invocation [{instance.app_name}]", type=SpanTypeAttribute.TASK, input={"new_message": serialized_message}, diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index f7b47a19..94de65c9 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -4126,6 +4126,22 @@ def test_span_exit_logs_exception_group_sub_exceptions(with_memory_logger): _assert_test_exception_group_contents(logs[0].get("error", "")) +@pytest.mark.parametrize( + "exception_type", + [GeneratorExit, asyncio.CancelledError, KeyboardInterrupt, SystemExit], +) +def test_span_exit_logs_base_exceptions(with_memory_logger, exception_type): + init_test_logger(__name__) + + with pytest.raises(exception_type): + with braintrust.current_logger().start_span(name="base-exception-span"): + raise exception_type + + logs = with_memory_logger.pop() + assert len(logs) == 1 + assert exception_type.__name__ in logs[0]["error"] + + def test_traced_logs_exception_group_sub_exceptions(with_memory_logger): """Verify sub-exceptions are captured when an ExceptionGroup propagates through @traced.""" init_test_logger(__name__)