fix: add close and aclose methods to GeneratorWrapper

[yorickvP]

Currently, close and aclose methods are missing on GeneratorWrapper, making it impossible to close generators from client code without accessing the .generator directly, and missing span uploads in the process.

This commit adds the missing methods for this usecase.

Disclaimer: Experimental PR review

Greptile Summary

This PR fixes a gap in _ContextPreservedSyncGeneratorWrapper and _ContextPreservedAsyncGeneratorWrapper where the close() / aclose() methods were absent, forcing callers to reach into .generator directly and bypassing Langfuse span finalisation. The change also refactors the duplicated span-teardown logic into a shared _end_span() helper protected by an _ended idempotency flag, which is a clean improvement.

Key changes:

• Adds _ended: bool flag and _end_span() helper to both wrappers, eliminating the duplicated span-finalisation block that previously lived inline in __next__ / __anext__.
• Adds close() to _ContextPreservedSyncGeneratorWrapper and aclose() to _ContextPreservedAsyncGeneratorWrapper, each calling _end_span() before delegating to the underlying generator.
• The ordering (_end_span() first, then generator teardown) is intentional — the span is finalized with the items collected so far, even if generator cleanup later raises.
• Minor consistency gap: __next__ runs the generator via self.context.run() / asyncio.create_task(..., context=…) to preserve contextvars, but the new close() / aclose() delegate to the underlying generator directly, so any finally blocks in the generator that read context variables will observe the caller's context rather than the preserved one.

Confidence Score: 5/5

• Safe to merge — fixes a real gap with no regressions; two P2 context-preservation suggestions are edge cases that don't affect the primary fix.
• All findings are P2 style/best-practice suggestions. The primary goal (prevent missing span uploads when a generator is closed early) is correctly achieved. The _ended guard prevents double-finalisation. No P0 or P1 issues found.
• No files require special attention beyond the two context-preservation suggestions in langfuse/_client/observe.py.

Important Files Changed

Filename	Overview
langfuse/_client/observe.py	Adds close() / aclose() to sync/async generator wrappers with a shared _end_span() helper and idempotency guard; minor P2 gap where generator teardown doesn't run inside the preserved context, unlike __next__/__anext__.

Sequence Diagram

sequenceDiagram
    participant Client
    participant SyncWrapper as _ContextPreservedSyncGeneratorWrapper
    participant Generator

    Client->>SyncWrapper: close()
    SyncWrapper->>SyncWrapper: _end_span() [idempotent via _ended flag]
    SyncWrapper->>SyncWrapper: span.update(output=items).end()
    SyncWrapper->>Generator: generator.close()
    Generator-->>SyncWrapper: (GeneratorExit injected)
    SyncWrapper-->>Client: returns

    note over SyncWrapper: __next__ exhaustion or error<br/>also calls _end_span() — _ended<br/>flag prevents double-finalisation

Loading

_{Reviews (1): Last reviewed commit: "fix: add _end_span method to prevent dou..." | Re-trigger Greptile}

_{(5/5) You can turn off certain types of comments like style here!}

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[yorickvP]

Addressed the P2 issue from greptile by preserving the generator context during close as well.