From 5b8c5ee86134e7a906f6782048c313435cb3847a Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 29 Jul 2026 10:03:35 -0400 Subject: [PATCH] fix(langchain): prevent async runs from inheriting ended spans Track native span context as parent-linked entries so async cleanup in a copied context skips ended spans while preserving active ancestors. Add VCR-backed regression coverage for consecutive ainvoke calls across the supported LangChain matrix. --- py/src/braintrust/context.py | 32 +- ...async_invocations_are_separate_traces.yaml | 302 ++++++++++++++++++ ...async_invocations_are_separate_traces.yaml | 302 ++++++++++++++++++ .../integrations/langchain/conftest.py | 15 +- .../integrations/langchain/test_callbacks.py | 45 ++- 5 files changed, 676 insertions(+), 20 deletions(-) create mode 100644 py/src/braintrust/integrations/langchain/cassettes/0.3.28/test_consecutive_async_invocations_are_separate_traces.yaml create mode 100644 py/src/braintrust/integrations/langchain/cassettes/latest/test_consecutive_async_invocations_are_separate_traces.yaml diff --git a/py/src/braintrust/context.py b/py/src/braintrust/context.py index ed596148..19aee99c 100644 --- a/py/src/braintrust/context.py +++ b/py/src/braintrust/context.py @@ -24,6 +24,12 @@ class ParentSpanIds: span_parents: list[str] +@dataclass(frozen=True) +class _SpanContextEntry: + span_object: Any + parent: "_SpanContextEntry | None" + + class ContextManager(ABC): """Abstract base class for managing span context in Braintrust. @@ -77,29 +83,43 @@ class BraintrustContextManager(ContextManager): """Braintrust-only context manager using contextvars when OTEL is not available.""" def __init__(self): - self._current_span: ContextVar[Any | None] = ContextVar("braintrust_current_span", default=None) + self._current_span: ContextVar[_SpanContextEntry | None] = ContextVar("braintrust_current_span", default=None) + + def _get_active_entry(self) -> _SpanContextEntry | None: + entry = self._current_span.get() + # Async cleanup can end a span from a copied context without clearing + # the original context. Skip that stale span while preserving any active + # parent that preceded it. + while entry is not None and getattr(entry.span_object, "_logged_end_time", None) is not None: + entry = entry.parent + return entry def get_current_span_info(self) -> SpanInfo | None: """Get information about the currently active span.""" - current_span = self._current_span.get() - if not current_span: + entry = self._get_active_entry() + if entry is None: return None + current_span = entry.span_object # Return SpanInfo for BT spans return SpanInfo(trace_id=current_span.root_span_id, span_id=current_span.span_id, span_object=current_span) def get_parent_span_ids(self) -> ParentSpanIds | None: """Get parent information for creating a new Braintrust span.""" - current_span = self._current_span.get() - if not current_span: + entry = self._get_active_entry() + if entry is None: return None + current_span = entry.span_object # If current span is a BT span, use it as parent return ParentSpanIds(root_span_id=current_span.root_span_id, span_parents=[current_span.span_id]) def set_current_span(self, span_object: Any) -> Any: """Set the current active span.""" - return self._current_span.set(span_object) + entry = None + if span_object is not None: + entry = _SpanContextEntry(span_object=span_object, parent=self._get_active_entry()) + return self._current_span.set(entry) def unset_current_span(self, context_token: Any = None) -> None: """Unset the current active span.""" diff --git a/py/src/braintrust/integrations/langchain/cassettes/0.3.28/test_consecutive_async_invocations_are_separate_traces.yaml b/py/src/braintrust/integrations/langchain/cassettes/0.3.28/test_consecutive_async_invocations_are_separate_traces.yaml new file mode 100644 index 00000000..6132bfb6 --- /dev/null +++ b/py/src/braintrust/integrations/langchain/cassettes/0.3.28/test_consecutive_async_invocations_are_separate_traces.yaml @@ -0,0 +1,302 @@ +interactions: +- request: + body: '{"messages":[{"content":"What is 1 + 0?","role":"user"}],"model":"gpt-4o-mini","stream":false}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '94' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.32.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.32.0 + x-stainless-raw-response: + - 'true' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-E6yqVnsqAOvHpagpGLwLLoTLqyawV\",\n \"object\": + \"chat.completion\",\n \"created\": 1785332567,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"1 + 0 equals 1.\",\n \"refusal\": + null,\n \"annotations\": []\n },\n \"logprobs\": null,\n + \ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": + 15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\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_836208cd2a\"\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: + - a22c8c483935ddfa-YYZ + connection: + - keep-alive + content-length: + - '821' + content-type: + - application/json + date: + - Wed, 29 Jul 2026 13:42:48 GMT + openai-processing-ms: + - '1015' + 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: + - '149999995' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_dbe641bded204be29da98e20199f6002 + status: + code: 200 + message: OK +- request: + body: '{"messages":[{"content":"What is 1 + 1?","role":"user"}],"model":"gpt-4o-mini","stream":false}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '94' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.32.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.32.0 + x-stainless-raw-response: + - 'true' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-E6yqWfKjPjvvbeg7Zl04J5l3vQDpc\",\n \"object\": + \"chat.completion\",\n \"created\": 1785332568,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"1 + 1 equals 2.\",\n \"refusal\": + null,\n \"annotations\": []\n },\n \"logprobs\": null,\n + \ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": + 15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\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_33a9b387c2\"\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: + - a22c8c888b87ddfa-YYZ + connection: + - keep-alive + content-length: + - '821' + content-type: + - application/json + date: + - Wed, 29 Jul 2026 13:42:49 GMT + openai-processing-ms: + - '717' + 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: + - '149999992' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_63ce970e8407408dbfb043eb017efd93 + status: + code: 200 + message: OK +- request: + body: '{"messages":[{"content":"What is 1 + 2?","role":"user"}],"model":"gpt-4o-mini","stream":false}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '94' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.32.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.32.0 + x-stainless-raw-response: + - 'true' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-E6yqXMMW5UVWCocAJ6F4OlSRWpmuX\",\n \"object\": + \"chat.completion\",\n \"created\": 1785332569,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"1 + 2 equals 3.\",\n \"refusal\": + null,\n \"annotations\": []\n },\n \"logprobs\": null,\n + \ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": + 15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\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_ab0a2ab924\"\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: + - a22c8c8e29baddfa-YYZ + connection: + - keep-alive + content-length: + - '821' + content-type: + - application/json + date: + - Wed, 29 Jul 2026 13:42:50 GMT + openai-processing-ms: + - '646' + 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: + - '149999995' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_1f3ebdbc09074c9a84a3b4a4001cf9ad + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/langchain/cassettes/latest/test_consecutive_async_invocations_are_separate_traces.yaml b/py/src/braintrust/integrations/langchain/cassettes/latest/test_consecutive_async_invocations_are_separate_traces.yaml new file mode 100644 index 00000000..6d681693 --- /dev/null +++ b/py/src/braintrust/integrations/langchain/cassettes/latest/test_consecutive_async_invocations_are_separate_traces.yaml @@ -0,0 +1,302 @@ +interactions: +- request: + body: '{"messages":[{"content":"What is 1 + 0?","role":"user"}],"model":"gpt-4o-mini","stream":false}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '94' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.32.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.32.0 + x-stainless-raw-response: + - 'true' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-E6ypeyIRh9qTBJL7sQzKssXvSBtZO\",\n \"object\": + \"chat.completion\",\n \"created\": 1785332514,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"1 + 0 equals 1.\",\n \"refusal\": + null,\n \"annotations\": []\n },\n \"logprobs\": null,\n + \ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": + 15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\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_836208cd2a\"\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: + - a22c8b34bc1a299a-YYZ + connection: + - keep-alive + content-length: + - '821' + content-type: + - application/json + date: + - Wed, 29 Jul 2026 13:41:54 GMT + openai-processing-ms: + - '454' + 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: + - '149999992' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_dbce064e90d740bfa97031384c13e86f + status: + code: 200 + message: OK +- request: + body: '{"messages":[{"content":"What is 1 + 1?","role":"user"}],"model":"gpt-4o-mini","stream":false}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '94' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.32.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.32.0 + x-stainless-raw-response: + - 'true' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-E6ypeaCOEa1xKlzAAO2AYqctNdW6Q\",\n \"object\": + \"chat.completion\",\n \"created\": 1785332514,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"1 + 1 equals 2.\",\n \"refusal\": + null,\n \"annotations\": []\n },\n \"logprobs\": null,\n + \ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": + 15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\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_33a9b387c2\"\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: + - a22c8b38c855299a-YYZ + connection: + - keep-alive + content-length: + - '821' + content-type: + - application/json + date: + - Wed, 29 Jul 2026 13:41:55 GMT + openai-processing-ms: + - '543' + 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: + - '149999992' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_f4617d7c4c294b60a07bf79b86e1520e + status: + code: 200 + message: OK +- request: + body: '{"messages":[{"content":"What is 1 + 2?","role":"user"}],"model":"gpt-4o-mini","stream":false}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate, zstd + connection: + - keep-alive + content-length: + - '94' + content-type: + - application/json + host: + - api.openai.com + user-agent: + - AsyncOpenAI/Python 2.32.0 + x-stainless-arch: + - arm64 + x-stainless-async: + - async:asyncio + x-stainless-lang: + - python + x-stainless-os: + - MacOS + x-stainless-package-version: + - 2.32.0 + x-stainless-raw-response: + - 'true' + x-stainless-retry-count: + - '0' + x-stainless-runtime: + - CPython + x-stainless-runtime-version: + - 3.14.6 + method: POST + uri: https://api.openai.com/v1/chat/completions + response: + body: + string: "{\n \"id\": \"chatcmpl-E6ypfvpIX4D8YxgpVBU325AaV6nMS\",\n \"object\": + \"chat.completion\",\n \"created\": 1785332515,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n + \ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\": + \"assistant\",\n \"content\": \"1 + 2 equals 3.\",\n \"refusal\": + null,\n \"annotations\": []\n },\n \"logprobs\": null,\n + \ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": + 15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\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_ab0a2ab924\"\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: + - a22c8b3cecbd299a-YYZ + connection: + - keep-alive + content-length: + - '821' + content-type: + - application/json + date: + - Wed, 29 Jul 2026 13:41:56 GMT + openai-processing-ms: + - '843' + 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: + - '149999992' + x-ratelimit-reset-requests: + - 2ms + x-ratelimit-reset-tokens: + - 0s + x-request-id: + - req_0a917711336548aaae109e4bb2cfccc5 + status: + code: 200 + message: OK +version: 1 diff --git a/py/src/braintrust/integrations/langchain/conftest.py b/py/src/braintrust/integrations/langchain/conftest.py index ad394a6e..769db572 100644 --- a/py/src/braintrust/integrations/langchain/conftest.py +++ b/py/src/braintrust/integrations/langchain/conftest.py @@ -1,21 +1,10 @@ -import os - import pytest +from braintrust.conftest import get_vcr_config @pytest.fixture(scope="module") def vcr_config(): - record_mode = "none" if (os.environ.get("CI") or os.environ.get("GITHUB_ACTIONS")) else "once" - return { - "filter_headers": [ - "authorization", - "x-goog-api-key", - "x-api-key", - "api-key", - "openai-api-key", - ], - "record_mode": record_mode, + **get_vcr_config(), "match_on": ["uri", "method"], - "decode_compressed_response": True, } diff --git a/py/src/braintrust/integrations/langchain/test_callbacks.py b/py/src/braintrust/integrations/langchain/test_callbacks.py index 1415d0bd..ceb0da76 100644 --- a/py/src/braintrust/integrations/langchain/test_callbacks.py +++ b/py/src/braintrust/integrations/langchain/test_callbacks.py @@ -8,7 +8,7 @@ import pytest from braintrust import logger from braintrust.integrations.langchain import BraintrustCallbackHandler -from braintrust.logger import flush +from braintrust.logger import NOOP_SPAN, current_span, flush from braintrust.test_helpers import init_test_logger from langchain_core.callbacks import BaseCallbackHandler from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage @@ -157,6 +157,49 @@ def test_llm_calls(logger_memory_logger): ) +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_consecutive_async_invocations_are_separate_traces(logger_memory_logger): + test_logger, memory_logger = logger_memory_logger + assert not memory_logger.pop() + + handler = BraintrustCallbackHandler(logger=test_logger) + prompt = ChatPromptTemplate.from_template("What is 1 + {number}?") + model = ChatOpenAI(model="gpt-4o-mini") + chain: RunnableSerializable[dict[str, str], BaseMessage] = prompt.pipe(model) + + for i in range(2): + result = await chain.ainvoke( + {"number": str(i)}, + config={"callbacks": [cast(BaseCallbackHandler, handler)], "run_name": f"run{i}"}, + ) + assert result.content + assert current_span() == NOOP_SPAN + + with test_logger.start_span(name="parent") as parent_span: + result = await chain.ainvoke( + {"number": "2"}, + config={"callbacks": [cast(BaseCallbackHandler, handler)], "run_name": "run2"}, + ) + assert result.content + assert current_span() == parent_span + + spans = memory_logger.pop() + assert len(spans) == 10 + + langchain_spans = [span for span in spans if span["span_attributes"]["name"] != "parent"] + assert all( + span["context"]["span_origin"]["instrumentation"]["name"] == "langchain-auto" for span in langchain_spans + ) + + roots = [span for span in spans if span.get("span_parents") is None] + assert [span["span_attributes"]["name"] for span in roots] == ["run0", "run1", "parent"] + assert len({span["root_span_id"] for span in roots}) == 3 + + nested_run = next(span for span in spans if span["span_attributes"]["name"] == "run2") + assert nested_run["span_parents"] == [parent_span.span_id] + + @pytest.mark.vcr def test_chain_with_memory(logger_memory_logger): test_logger, memory_logger = logger_memory_logger