fix(streaming): do not add close() to CustomStreamWrapper

`tests/unit/test_router_silent_experiment.py` fails on this branch: the shadow
and fallback streaming tests see zero success callbacks where they expect one.

The cause is a name collision. `Router.stream_with_fallbacks` ends with

    finally:
        if hasattr(model_response, "close"):
            model_response.close()

`CustomStreamWrapper` had no `close`, so that branch never ran. Adding one
turned it on, so the fallback path now invokes the terminal-log flush from its
`finally`, and the shadow stream reports its completion through a path the
success callbacks were never accounting for.

The flush API is renamed to `close_stream()`. Same behaviour, same context
manager, but `hasattr(stream, "close")` stays False exactly as before, so no
existing call path changes. `test_router_close_duck_typing_contract` pins that
so it cannot be reintroduced by rename.
This commit is contained in:
JingHao-Leon 2026-10-11 12:36:17 +08:00
parent d745030539
commit 5643b80340
3 changed files with 207 additions and 501 deletions

View file

@ -363,14 +363,21 @@ class CustomStreamWrapper:
def __aiter__(self) -> AsyncIterator["ModelResponseStream"]:
return self
def close(self) -> None:
def close_stream(self) -> None:
"""Deterministically end a stream the consumer stopped reading early.
A stream abandoned without reaching its final chunk would otherwise
never emit its terminal success log (spans, spend, custom callbacks
all miss the call). ``close()`` flushes that log from the chunks
all miss the call). ``close_stream()`` flushes that log from the chunks
received so far and releases the underlying provider stream; safe to
call after the stream has been fully consumed (it is a no-op then).
Deliberately *not* named ``close``: ``Router.stream_with_fallbacks``
duck-types on ``hasattr(model_response, "close")`` and calls it from a
``finally``, so adding a ``close`` here silently starts firing the
terminal success log from the fallback path, which changes how shadow
and fallback streams report completions. Callers that hold a concrete
``CustomStreamWrapper`` use this name or the context manager below.
"""
self._flush_abandoned_stream_logging()
stream_to_close: Final = self.completion_stream # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the underlying stream object is provider-specific
@ -386,7 +393,7 @@ class CustomStreamWrapper:
return self
def __exit__(self, exc_type: object, exc: object, tb: object) -> bool:
self.close()
self.close_stream()
return False
def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None:

File diff suppressed because it is too large Load diff