fix(chains): release the client and response when a stream is abandoned

[EnesYilmazcode]

🚀 What

predict_async_stream returns the response iterator from inside the async with self._client_async() block, so the client context exits while the caller is still reading, releasing the concurrency slot early and leaving the aiohttp response checked out of the pool with nothing to release it.

This hands ownership of the client context and the response to the generator the caller iterates, so both are released when it is exhausted or closed.

BEFORE                                   AFTER

_rpc()                                   _rpc()
 └─ async with _client_async()            └─ async with _client_async()
      ├─ post() -> response                    ├─ post() -> response
      └─ return iter_any() ───┐                └─ pop_all() ──────┐
    (context exits here)      │                                   │
      ✗ slot released         │                     generator ────┘
      ✗ response orphaned     │                      ├─ yields chunks
                              ▼                      └─ on close/exhaust:
    caller reads stream ──────┘                           ✓ response released
                                                          ✓ slot released

The practical effect is a stream that ends by cancellation, which is what happens whenever an end user closes the tab mid-generation: uvicorn cancels the parent chainlet's response generator, that propagates into async for data in await self.predict_async_stream(...), and the abandoned response never returns its connection. Each one permanently consumes a slot in TCPConnector(limit=concurrency_limit). After concurrency_limit of them the stub cannot reach that dependency at all until the session is cycled an hour later. ongoing_requests reads 0 throughout, so the gauge reports idle while the stub is wedged.

Measured against a local server with the real StubBase, concurrency_limit=2, two consumers cancelled mid-stream:

	next call to the dependency	ongoing_requests mid-stream
before	wedged, still blocked at 6s	0
after	succeeds immediately	1

gc.collect() does not rescue it, so this is not something the garbage collector cleans up later.

💻 How

contextlib.AsyncExitStack acquires the client context and the response, pop_all() transfers ownership once the request has succeeded, and the generator releases them in its own async with. Errors raised during setup still close the stack before propagating, so retries do not accumulate connections.

The request stays eager rather than moving into the generator, because RPCOptions.retries is documented as only retrying before results stream back, and code_gen.py emits async for data in await self.predict_async_stream(...), which cannot await a plain async generator.

One behavior change worth flagging: the slot is now held for the life of the stream, which is what predict_async already does. Concurrency is unchanged in count, since TCPConnector(limit=...) already capped it, but a caller that exceeds the limit now waits on the semaphore instead of failing with TimeoutError after timeout_sec. At the default concurrency_limit=300 that needs 300 concurrent streams to one dependency; at concurrency_limit=1 a chainlet that calls the same stub while consuming a stream from it would deadlock.

🔬 Testing

Three tests in truss-chains/tests/test_remote_chainlet.py, all against a loopback aiohttp server, no credentials or network. Each fails on main and passes here:

• the slot is held while streaming and released on aclose() (assert 0 == 1 on main)
• cancelling the consumer mid-stream releases the slot
• after concurrency_limit cancellations the stub can still reach its dependency (hangs out the timeout on main)

truss-chains unit suite: 177 passed, 1 skipped. ruff check, ruff format --check and mypy clean on both files.

[michaelfeil]

@EnesYilmazcode can't you check request.is_disconnected() - most of the time stoping the generator is a good idea, e.g. if you wanna sent a callback on the last chunk etc cancellation can be disruptive.

[EnesYilmazcode]

is_disconnected() is server side, on the Starlette request. This is the chains stub, the aiohttp client a chainlet uses to call a dependency, so there's no request to check there.

It also doesn't cancel anything. The cancel already happens when uvicorn kills the parent's generator, this just releases the response and the slot afterwards. Full consumption is untouched, last chunk still lands before cleanup.

They're complementary really. Even with disconnect detection something has to give the connection back, and right now nothing does.