Skip to content

Commit 7112d72

Browse files
AbhiPrasadAbhiPrasad
authored
fix(langchain): prevent async runs from inheriting ended spans (#624)
fixes #621
1 parent 075b0ee commit 7112d72

5 files changed

Lines changed: 676 additions & 20 deletions

File tree

‎py/src/braintrust/context.py‎

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,12 @@ class ParentSpanIds:
2424
span_parents: list[str]
2525

2626

27+
@dataclass(frozen=True)
28+
class _SpanContextEntry:
29+
span_object: Any
30+
parent: "_SpanContextEntry | None"
31+
32+
2733
class ContextManager(ABC):
2834
"""Abstract base class for managing span context in Braintrust.
2935
@@ -77,29 +83,43 @@ class BraintrustContextManager(ContextManager):
7783
"""Braintrust-only context manager using contextvars when OTEL is not available."""
7884

7985
def __init__(self):
80-
self._current_span: ContextVar[Any | None] = ContextVar("braintrust_current_span", default=None)
86+
self._current_span: ContextVar[_SpanContextEntry | None] = ContextVar("braintrust_current_span", default=None)
87+
88+
def _get_active_entry(self) -> _SpanContextEntry | None:
89+
entry = self._current_span.get()
90+
# Async cleanup can end a span from a copied context without clearing
91+
# the original context. Skip that stale span while preserving any active
92+
# parent that preceded it.
93+
while entry is not None and getattr(entry.span_object, "_logged_end_time", None) is not None:
94+
entry = entry.parent
95+
return entry
8196

8297
def get_current_span_info(self) -> SpanInfo | None:
8398
"""Get information about the currently active span."""
84-
current_span = self._current_span.get()
85-
if not current_span:
99+
entry = self._get_active_entry()
100+
if entry is None:
86101
return None
87102

103+
current_span = entry.span_object
88104
# Return SpanInfo for BT spans
89105
return SpanInfo(trace_id=current_span.root_span_id, span_id=current_span.span_id, span_object=current_span)
90106

91107
def get_parent_span_ids(self) -> ParentSpanIds | None:
92108
"""Get parent information for creating a new Braintrust span."""
93-
current_span = self._current_span.get()
94-
if not current_span:
109+
entry = self._get_active_entry()
110+
if entry is None:
95111
return None
96112

113+
current_span = entry.span_object
97114
# If current span is a BT span, use it as parent
98115
return ParentSpanIds(root_span_id=current_span.root_span_id, span_parents=[current_span.span_id])
99116

100117
def set_current_span(self, span_object: Any) -> Any:
101118
"""Set the current active span."""
102-
return self._current_span.set(span_object)
119+
entry = None
120+
if span_object is not None:
121+
entry = _SpanContextEntry(span_object=span_object, parent=self._get_active_entry())
122+
return self._current_span.set(entry)
103123

104124
def unset_current_span(self, context_token: Any = None) -> None:
105125
"""Unset the current active span."""

‎py/src/braintrust/integrations/langchain/cassettes/0.3.28/test_consecutive_async_invocations_are_separate_traces.yaml‎

Lines changed: 302 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,302 @@
1+
interactions:
2+
- request:
3+
body: '{"messages":[{"content":"What is 1 + 0?","role":"user"}],"model":"gpt-4o-mini","stream":false}'
4+
headers:
5+
accept:
6+
- application/json
7+
accept-encoding:
8+
- gzip, deflate, zstd
9+
connection:
10+
- keep-alive
11+
content-length:
12+
- '94'
13+
content-type:
14+
- application/json
15+
host:
16+
- api.openai.com
17+
user-agent:
18+
- AsyncOpenAI/Python 2.32.0
19+
x-stainless-arch:
20+
- arm64
21+
x-stainless-async:
22+
- async:asyncio
23+
x-stainless-lang:
24+
- python
25+
x-stainless-os:
26+
- MacOS
27+
x-stainless-package-version:
28+
- 2.32.0
29+
x-stainless-raw-response:
30+
- 'true'
31+
x-stainless-retry-count:
32+
- '0'
33+
x-stainless-runtime:
34+
- CPython
35+
x-stainless-runtime-version:
36+
- 3.14.6
37+
method: POST
38+
uri: https://api.openai.com/v1/chat/completions
39+
response:
40+
body:
41+
string: "{\n \"id\": \"chatcmpl-E6yqVnsqAOvHpagpGLwLLoTLqyawV\",\n \"object\":
42+
\"chat.completion\",\n \"created\": 1785332567,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n
43+
\ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\":
44+
\"assistant\",\n \"content\": \"1 + 0 equals 1.\",\n \"refusal\":
45+
null,\n \"annotations\": []\n },\n \"logprobs\": null,\n
46+
\ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\":
47+
15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\n \"prompt_tokens_details\":
48+
{\n \"cached_tokens\": 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\":
49+
{\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\":
50+
0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\":
51+
\"default\",\n \"system_fingerprint\": \"fp_836208cd2a\"\n}\n"
52+
headers:
53+
access-control-expose-headers:
54+
- X-Request-ID
55+
- CF-Ray
56+
- CF-Ray
57+
alt-svc:
58+
- h3=":443"; ma=86400
59+
cf-cache-status:
60+
- DYNAMIC
61+
cf-ray:
62+
- a22c8c483935ddfa-YYZ
63+
connection:
64+
- keep-alive
65+
content-length:
66+
- '821'
67+
content-type:
68+
- application/json
69+
date:
70+
- Wed, 29 Jul 2026 13:42:48 GMT
71+
openai-processing-ms:
72+
- '1015'
73+
openai-version:
74+
- '2020-10-01'
75+
server:
76+
- cloudflare
77+
strict-transport-security:
78+
- max-age=31536000; includeSubDomains; preload
79+
transfer-encoding:
80+
- chunked
81+
x-content-type-options:
82+
- nosniff
83+
x-openai-proxy-wasm:
84+
- v0.1
85+
x-ratelimit-limit-requests:
86+
- '30000'
87+
x-ratelimit-limit-tokens:
88+
- '150000000'
89+
x-ratelimit-remaining-requests:
90+
- '29999'
91+
x-ratelimit-remaining-tokens:
92+
- '149999995'
93+
x-ratelimit-reset-requests:
94+
- 2ms
95+
x-ratelimit-reset-tokens:
96+
- 0s
97+
x-request-id:
98+
- req_dbe641bded204be29da98e20199f6002
99+
status:
100+
code: 200
101+
message: OK
102+
- request:
103+
body: '{"messages":[{"content":"What is 1 + 1?","role":"user"}],"model":"gpt-4o-mini","stream":false}'
104+
headers:
105+
accept:
106+
- application/json
107+
accept-encoding:
108+
- gzip, deflate, zstd
109+
connection:
110+
- keep-alive
111+
content-length:
112+
- '94'
113+
content-type:
114+
- application/json
115+
host:
116+
- api.openai.com
117+
user-agent:
118+
- AsyncOpenAI/Python 2.32.0
119+
x-stainless-arch:
120+
- arm64
121+
x-stainless-async:
122+
- async:asyncio
123+
x-stainless-lang:
124+
- python
125+
x-stainless-os:
126+
- MacOS
127+
x-stainless-package-version:
128+
- 2.32.0
129+
x-stainless-raw-response:
130+
- 'true'
131+
x-stainless-retry-count:
132+
- '0'
133+
x-stainless-runtime:
134+
- CPython
135+
x-stainless-runtime-version:
136+
- 3.14.6
137+
method: POST
138+
uri: https://api.openai.com/v1/chat/completions
139+
response:
140+
body:
141+
string: "{\n \"id\": \"chatcmpl-E6yqWfKjPjvvbeg7Zl04J5l3vQDpc\",\n \"object\":
142+
\"chat.completion\",\n \"created\": 1785332568,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n
143+
\ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\":
144+
\"assistant\",\n \"content\": \"1 + 1 equals 2.\",\n \"refusal\":
145+
null,\n \"annotations\": []\n },\n \"logprobs\": null,\n
146+
\ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\":
147+
15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\n \"prompt_tokens_details\":
148+
{\n \"cached_tokens\": 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\":
149+
{\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\":
150+
0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\":
151+
\"default\",\n \"system_fingerprint\": \"fp_33a9b387c2\"\n}\n"
152+
headers:
153+
access-control-expose-headers:
154+
- X-Request-ID
155+
- CF-Ray
156+
- CF-Ray
157+
alt-svc:
158+
- h3=":443"; ma=86400
159+
cf-cache-status:
160+
- DYNAMIC
161+
cf-ray:
162+
- a22c8c888b87ddfa-YYZ
163+
connection:
164+
- keep-alive
165+
content-length:
166+
- '821'
167+
content-type:
168+
- application/json
169+
date:
170+
- Wed, 29 Jul 2026 13:42:49 GMT
171+
openai-processing-ms:
172+
- '717'
173+
openai-version:
174+
- '2020-10-01'
175+
server:
176+
- cloudflare
177+
strict-transport-security:
178+
- max-age=31536000; includeSubDomains; preload
179+
transfer-encoding:
180+
- chunked
181+
x-content-type-options:
182+
- nosniff
183+
x-openai-proxy-wasm:
184+
- v0.1
185+
x-ratelimit-limit-requests:
186+
- '30000'
187+
x-ratelimit-limit-tokens:
188+
- '150000000'
189+
x-ratelimit-remaining-requests:
190+
- '29999'
191+
x-ratelimit-remaining-tokens:
192+
- '149999992'
193+
x-ratelimit-reset-requests:
194+
- 2ms
195+
x-ratelimit-reset-tokens:
196+
- 0s
197+
x-request-id:
198+
- req_63ce970e8407408dbfb043eb017efd93
199+
status:
200+
code: 200
201+
message: OK
202+
- request:
203+
body: '{"messages":[{"content":"What is 1 + 2?","role":"user"}],"model":"gpt-4o-mini","stream":false}'
204+
headers:
205+
accept:
206+
- application/json
207+
accept-encoding:
208+
- gzip, deflate, zstd
209+
connection:
210+
- keep-alive
211+
content-length:
212+
- '94'
213+
content-type:
214+
- application/json
215+
host:
216+
- api.openai.com
217+
user-agent:
218+
- AsyncOpenAI/Python 2.32.0
219+
x-stainless-arch:
220+
- arm64
221+
x-stainless-async:
222+
- async:asyncio
223+
x-stainless-lang:
224+
- python
225+
x-stainless-os:
226+
- MacOS
227+
x-stainless-package-version:
228+
- 2.32.0
229+
x-stainless-raw-response:
230+
- 'true'
231+
x-stainless-retry-count:
232+
- '0'
233+
x-stainless-runtime:
234+
- CPython
235+
x-stainless-runtime-version:
236+
- 3.14.6
237+
method: POST
238+
uri: https://api.openai.com/v1/chat/completions
239+
response:
240+
body:
241+
string: "{\n \"id\": \"chatcmpl-E6yqXMMW5UVWCocAJ6F4OlSRWpmuX\",\n \"object\":
242+
\"chat.completion\",\n \"created\": 1785332569,\n \"model\": \"gpt-4o-mini-2024-07-18\",\n
243+
\ \"choices\": [\n {\n \"index\": 0,\n \"message\": {\n \"role\":
244+
\"assistant\",\n \"content\": \"1 + 2 equals 3.\",\n \"refusal\":
245+
null,\n \"annotations\": []\n },\n \"logprobs\": null,\n
246+
\ \"finish_reason\": \"stop\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\":
247+
15,\n \"completion_tokens\": 8,\n \"total_tokens\": 23,\n \"prompt_tokens_details\":
248+
{\n \"cached_tokens\": 0,\n \"audio_tokens\": 0\n },\n \"completion_tokens_details\":
249+
{\n \"reasoning_tokens\": 0,\n \"audio_tokens\": 0,\n \"accepted_prediction_tokens\":
250+
0,\n \"rejected_prediction_tokens\": 0\n }\n },\n \"service_tier\":
251+
\"default\",\n \"system_fingerprint\": \"fp_ab0a2ab924\"\n}\n"
252+
headers:
253+
access-control-expose-headers:
254+
- X-Request-ID
255+
- CF-Ray
256+
- CF-Ray
257+
alt-svc:
258+
- h3=":443"; ma=86400
259+
cf-cache-status:
260+
- DYNAMIC
261+
cf-ray:
262+
- a22c8c8e29baddfa-YYZ
263+
connection:
264+
- keep-alive
265+
content-length:
266+
- '821'
267+
content-type:
268+
- application/json
269+
date:
270+
- Wed, 29 Jul 2026 13:42:50 GMT
271+
openai-processing-ms:
272+
- '646'
273+
openai-version:
274+
- '2020-10-01'
275+
server:
276+
- cloudflare
277+
strict-transport-security:
278+
- max-age=31536000; includeSubDomains; preload
279+
transfer-encoding:
280+
- chunked
281+
x-content-type-options:
282+
- nosniff
283+
x-openai-proxy-wasm:
284+
- v0.1
285+
x-ratelimit-limit-requests:
286+
- '30000'
287+
x-ratelimit-limit-tokens:
288+
- '150000000'
289+
x-ratelimit-remaining-requests:
290+
- '29999'
291+
x-ratelimit-remaining-tokens:
292+
- '149999995'
293+
x-ratelimit-reset-requests:
294+
- 2ms
295+
x-ratelimit-reset-tokens:
296+
- 0s
297+
x-request-id:
298+
- req_1f3ebdbc09074c9a84a3b4a4001cf9ad
299+
status:
300+
code: 200
301+
message: OK
302+
version: 1

0 commit comments

Comments
 (0)