Skip to content

fix: delegate traced generator send and throw operations #1077

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Chen Yufeiyang (feiiiiii5) wants to merge 1 commit into from
Closed

fix: delegate traced generator send and throw operations #1077

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 51 additions & 44 deletions py/src/braintrust/logger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3205,14 +3205,14 @@ async def wrapper_async_gen(*f_args, **f_kwargs):
# Get max items from environment or default
max_items = BraintrustEnv.MAX_GENERATOR_ITEMS.get(1000)

if trace_io and max_items != 0:
# Collect output up to limit
collected = []
truncated = False

async_gen = f(*f_args, **f_kwargs)
try:
async for value in async_gen:
collect_output = trace_io and max_items != 0
collected = []
truncated = False
async_gen = f(*f_args, **f_kwargs)
try:
value = await async_gen.__anext__()
while True:
if collect_output:
if max_items == -1 or (not truncated and len(collected) < max_items):
collected.append(value)
else:
Expand All @@ -3222,20 +3222,23 @@ async def wrapper_async_gen(*f_args, **f_kwargs):
f"Generator output exceeded limit of {max_items} items, output not logged. "
"Increase BRAINTRUST_MAX_GENERATOR_ITEMS or set to -1 to disable limit."
)
yield value

if not truncated:
_try_log_output(span, collected)
except Exception as e:
# Log partial output on error
if collected and not truncated:
_try_log_output(span, collected)
raise
else:
# Original behavior - no collection
async_gen = f(*f_args, **f_kwargs)
async for value in async_gen:
yield value
try:
sent = yield value
except GeneratorExit:
await async_gen.aclose()

@chatgpt-codex-connector chatgpt-codex-connector Bot Nov 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Avoid double-closing async generators during loop shutdown

If a partially consumed traced async generator survives until asyncio.run() shuts down and its finally block awaits, the loop has registered both this wrapper and async_gen, so shutdown_asyncgens() schedules aclose() on both concurrently. This delegated close races with the loop's close, producing RuntimeError: aclose(): asynchronous generator is already running through the loop's exception handler. This is reproducible with the changed wrapper on Python 3.14; the parent revision does not produce the shutdown error. Coordinate finalization ownership so normal shutdown does not close the same source generator twice.

Useful? React with 👍 / 👎.

raise
except BaseException as exc:
value = await async_gen.athrow(exc)
else:
value = await async_gen.asend(sent)
except StopAsyncIteration:
if collect_output and not truncated:
_try_log_output(span, collected)
except Exception:
# Log partial output on error
if collect_output and collected and not truncated:
_try_log_output(span, collected)
raise

@wraps(f)
def wrapper_sync_gen(*f_args, **f_kwargs):
Expand All @@ -3246,14 +3249,14 @@ def wrapper_sync_gen(*f_args, **f_kwargs):
# Get max items from environment or default
max_items = BraintrustEnv.MAX_GENERATOR_ITEMS.get(1000)

if trace_io and max_items != 0:
# Collect output up to limit
collected = []
truncated = False

sync_gen = f(*f_args, **f_kwargs)
try:
for value in sync_gen:
collect_output = trace_io and max_items != 0
collected = []
truncated = False
sync_gen = f(*f_args, **f_kwargs)
try:
value = next(sync_gen)
while True:
if collect_output:
if max_items == -1 or (not truncated and len(collected) < max_items):
collected.append(value)
else:
Expand All @@ -3263,20 +3266,24 @@ def wrapper_sync_gen(*f_args, **f_kwargs):
f"Generator output exceeded limit of {max_items} items, output not logged. "
"Increase BRAINTRUST_MAX_GENERATOR_ITEMS or set to -1 to disable limit."
)
yield value

if not truncated:
_try_log_output(span, collected)
except Exception as e:
# Log partial output on error
if collected and not truncated:
_try_log_output(span, collected)
raise
else:
# Original behavior - no collection
sync_gen = f(*f_args, **f_kwargs)
for value in sync_gen:
yield value
try:
sent = yield value
except GeneratorExit:
sync_gen.close()
raise
except BaseException as exc:
value = sync_gen.throw(exc)
else:
value = sync_gen.send(sent)
except StopIteration as done:
if collect_output and not truncated:
_try_log_output(span, collected)
return done.value
except Exception:
# Log partial output on error
if collect_output and collected and not truncated:
_try_log_output(span, collected)
raise

if inspect.isasyncgenfunction(f):
return cast(F, wrapper_async_gen)
Expand Down
165 changes: 165 additions & 0 deletions py/src/braintrust/test_logger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2213,6 +2213,171 @@ def sync_number_generator(n: int):
)


@pytest.mark.parametrize("notrace_io", [False, True])
@pytest.mark.parametrize("operation", ["send", "throw"])
def test_traced_sync_generator_delegates_protocol(with_memory_logger, notrace_io, operation):
init_test_logger(__name__)
closed = []

@logger.traced(notrace_io=notrace_io)
def exchange():
try:
try:
value = yield "ready"
except ValueError as exc:
value = str(exc)
yield value
return "complete"
finally:
closed.append(True)

generator = exchange()
try:
assert next(generator) == "ready"
if operation == "send":
assert generator.send("response") == "response"
else:
assert generator.throw(ValueError("response")) == "response"
with pytest.raises(StopIteration) as done:
next(generator)
assert done.value.value == "complete"
finally:
generator.close()
assert closed == [True]

[log] = with_memory_logger.pop()
assert log.get("output") == (None if notrace_io else ["ready", "response"])
assert not log.get("error")


@pytest.mark.parametrize("notrace_io", [False, True])
def test_traced_sync_generator_delegates_close(with_memory_logger, notrace_io):
init_test_logger(__name__)
closed = []

@logger.traced(notrace_io=notrace_io)
def source():
try:
yield "ready"
finally:
closed.append(True)

generator = source()
assert next(generator) == "ready"
generator.close()
assert closed == [True]
[log] = with_memory_logger.pop()
assert "end" in log["metrics"]


@pytest.mark.parametrize("notrace_io", [False, True])
def test_traced_sync_generator_preserves_generator_exit(with_memory_logger, notrace_io):
init_test_logger(__name__)
closed = []

@logger.traced(notrace_io=notrace_io)
def source():
try:
yield "ready"
finally:
closed.append(True)

generator = source()
assert next(generator) == "ready"
error = GeneratorExit("cancel")
with pytest.raises(GeneratorExit, match="^cancel$") as raised:
generator.throw(error)
assert raised.value is error
assert closed == [True]

[log] = with_memory_logger.pop()
assert "GeneratorExit: cancel" in log["error"]


@pytest.mark.parametrize("notrace_io", [False, True])
@pytest.mark.parametrize("operation", ["send", "throw"])
@pytest.mark.asyncio
async def test_traced_async_generator_delegates_protocol(with_memory_logger, notrace_io, operation):
init_test_logger(__name__)
closed = []

@logger.traced(notrace_io=notrace_io)
async def exchange():
try:
try:
value = yield "ready"
except ValueError as exc:
value = str(exc)
yield value
finally:
closed.append(True)

generator = exchange()
try:
assert await generator.__anext__() == "ready"
if operation == "send":
assert await generator.asend("response") == "response"
else:
assert await generator.athrow(ValueError("response")) == "response"
with pytest.raises(StopAsyncIteration):
await generator.__anext__()
finally:
await generator.aclose()
assert closed == [True]

[log] = with_memory_logger.pop()
assert log.get("output") == (None if notrace_io else ["ready", "response"])
assert not log.get("error")


@pytest.mark.parametrize("notrace_io", [False, True])
@pytest.mark.asyncio
async def test_traced_async_generator_delegates_close(with_memory_logger, notrace_io):
init_test_logger(__name__)
closed = []

@logger.traced(notrace_io=notrace_io)
async def source():
try:
yield "ready"
finally:
await asyncio.sleep(0)
closed.append(True)

generator = source()
assert await generator.__anext__() == "ready"
await generator.aclose()
assert closed == [True]
[log] = with_memory_logger.pop()
assert "end" in log["metrics"]


@pytest.mark.parametrize("notrace_io", [False, True])
@pytest.mark.asyncio
async def test_traced_async_generator_preserves_generator_exit(with_memory_logger, notrace_io):
init_test_logger(__name__)
closed = []

@logger.traced(notrace_io=notrace_io)
async def source():
try:
yield "ready"
finally:
await asyncio.sleep(0)
closed.append(True)

generator = source()
assert await generator.__anext__() == "ready"
error = GeneratorExit("cancel")
with pytest.raises(GeneratorExit, match="^cancel$") as raised:
await generator.athrow(error)
assert raised.value is error
assert closed == [True]

[log] = with_memory_logger.pop()
assert "GeneratorExit: cancel" in log["error"]


def test_traced_sync_generator_with_exception(with_memory_logger):
"""Test sync generator that raises an exception."""
init_test_logger(__name__)
Expand Down