UN-4237 [FIX] Close streamed LLM responses so executor workers no longer deadlock on httpcore's pool lock

[johnyrahul]

What

• New close_stream() helper in unstract/sdk1/utils/retry_utils.py that closes a streamed completion and the HTTP response behind it. It walks the stream's handle attributes (completion_stream, streaming_response, response, http_response, stream_iterator) breadth-first, to a bounded depth with a visited set, and closes everything it reaches that has a close(). That covers:

  • litellm's ModelResponseIterator (Anthropic, Vertex/Gemini, the generic HTTP handler)
  • the OpenAI SDK Stream (OpenAI, Azure, OpenAI-compatible)
  • plain generators (Bedrock)
  • litellm's Responses-API bridge (responses/<model>, plus models litellm routes there such as gpt-5-pro, *-codex and o3-pro), where the response sits at wrapper -> completion_stream -> streaming_response -> response

  It is duck-typed (no litellm import) and never raises. A failed close is logged at WARNING on the caller's logger, with the adapter info and only the close failure's own traceback (not the provider error that is propagating at the time).

• collect_with_retry (used by LLM.complete() when streaming) and iter_with_retry (used by LLM.stream_complete()) now call it in a finally: on success, on error, and when the caller closes the generator early. They pass their logger and adapter description through.

• The retry backoff now runs after the failed attempt's stream is closed, so the connection isn't held through the sleep.

• LLM.stream_complete() and iter_with_retry document that a caller who stops early must close the generator, e.g. with contextlib.closing(...).

Why

• Executor workers hung "in the LLM call" under load: 28 of 28 hung executions in the 7 Oct 600-user run, and 1.1% of executions in the 6 Oct run. py-spy showed they were deadlocked on their own lock inside httpcore's connection pool, not waiting on the network (UN-4237).
• LiteLLM's sync CustomStreamWrapper has no close(). On some providers (Anthropic among them) it also stops iterating before the HTTP body ends. So a streamed completion that failed, was abandoned, or on Anthropic simply finished left its response checked out of the pool until the garbage collector finalised it. When that happened on a thread inside ConnectionPool.handle_request, the finaliser's PoolByteStream.close() waited on the same non-reentrant threading.Lock the thread already held. The worker hung forever, with its lease still renewing.
• sdk1's existing getattr(gen, "close", None) on the error path was a no-op for these streams.
• This is not fixed upstream. litellm 1.104.0, httpcore 1.0.9 and httpx 0.28.1 are the latest releases, and I reproduced the deadlock against them.

How

• close_stream() releases the pool slot on the calling thread, outside any pool lock. Closing the line generator propagates GeneratorExit down to PoolByteStream.close(); closing an httpx.Response closes its stream directly.
• In both retry loops, the explicit close is called from a finally around each attempt. The time.sleep moved after the try so it runs only once the stream is closed.

Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)

• Low risk. Chunks are fully collected before the close, and stream_chunk_builder only reads the collected chunks, so responses are unchanged. test_anthropic_stream_complete.py still passes.
• Connections aren't reused after a stream. Closing ends the connection, because litellm leaves unread bytes on it. Before this change the connection leaked until GC and wasn't reused either, so this adds no new connection churn. On paths that already close themselves (OpenAI SDK Stream, a drained Responses stream), the extra close is a no-op.
• The walk is bounded (depth 3, visited set) and tolerates attributes that raise, so it can't loop or fail on unexpected objects.
• Not covered: a stream_complete() generator the caller abandons without closing is still finalised by the GC, possibly inside the pool lock. This is now documented on stream_complete. The one known consumer that can abandon it, unstract-cloud's table_extractor plugin, needs contextlib.closing() around its loops; that is a separate change in that repo.

Database Migrations

• None

Env Config

• None

Relevant Docs

• None

Related Issues or PRs

• UN-4237: this bug, with the py-spy evidence
• UN-4223 / UN-4223 [FIX] Kill and restart only a hung PG consumer child instead of losing the whole worker pod #2311: the per-child watchdog stays as the backstop. A deadline checked inside the worker thread can't fire while that thread is blocked on a lock.
• UN-4179: lease renewal cap, still needed as a backstop
• Same deadlock seen in another project: LLM call can deadlock forever: unclosed stream closed by GC under httpcore pool lock ericfitz/tmi-tf-wh#90

Dependencies Versions

• No changes (litellm 1.104.0, httpcore 1.0.9, httpx 0.28.1)

Notes on Testing

• tests/test_stream_connection_release.py runs real litellm against a local server speaking the provider's SSE format, with GC disabled, and checks pool._requests is empty after:
  • Anthropic through LLM.complete(): a completed stream, and a stream failing after content (overloaded_error event)
  • the OpenAI Responses-API bridge (openai/responses/gpt-4o): error after content, early close, and drain
  • forcing gc.collect() inside the pool lock on the next call (the production trigger), which does not deadlock. It runs in a daemon thread with a 10 s bound, so a regression fails instead of hanging the suite.
• tests/utils/test_collect_with_retry.py: unit tests for close_stream() and both retry loops. They cover:
  • a Bedrock-style generator completion_stream, an OpenAI-SDK-shaped Stream, and the Responses-API chain
  • self-referencing objects, and a close attribute that raises
  • the failed-close warning: its level, the adapter info on the caller's logger, and a traceback that leaves out the propagating provider error
• Mutation checks:
  • dropping completion_stream from the walk fails 14 tests
  • capping the walk at depth 1 fails 13
  • logging a failed close at DEBUG fails the warning test
  • against main's retry_utils.py, every new or changed test fails, and the deadlock test fails with the production py-spy stack
• Full sdk1 suite, after merging main at d3a5995: 844 passed, 1 failed. The failure is test_storage_compat.py::test_uploads_are_plain_payload_signed_not_aws_chunked, which fails identically on main and is unrelated.
• Snapshot snapshot.1016 was built from this branch before the review fixes. Still to do: re-run the UN-4136 600-user load test on a snapshot built from the current head and confirm there are no hung executions.

Screenshots

Checklist

I have read and understood the Contribution Guidelines.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VQoJFFiRZyDKyu14pg5hDb

[CLAassistant]

All committers have signed the CLA.

[johnyrahul]

Standardized review (unstract:review-deep, INITIAL) — reviewed at 7632d7ad9, base 02b979753

Verdict: APPROVE WITH NITS

Summary: Critical 0 · High 0 · Medium 4 · Low 3 · Lenses run: 19/19 (4 N/A)

The fix works for Anthropic, plain OpenAI/Azure chat, Bedrock and Vertex. All 13 new or changed tests fail against main's retry_utils.py and pass at HEAD, run against the pinned litellm 1.104.0. The pinned ruff 0.3.4 check and format --check both pass. Finding #1 is worth fixing before merge: one OpenAI path is still not closed. Details are inline.

#	Severity	Finding
1	Medium	close_stream doesn't reach the HTTP response on litellm's Responses-API path (gpt-5-pro, *-codex, o3-pro, responses/<model>)
2	Medium	"When the caller abandons this generator" is not covered when cyclic GC finalises the generator; there is a real stream_complete() consumer in unstract-cloud's worker plugin
3	Medium	Closing the wrapper's own completion_stream (the whole fix for Bedrock) has no test
4	Medium	A failed close, which recreates the leak, is logged only at DEBUG
5	Low	"A drained stream is still open" holds for some providers, not all
6	Low	response is attributed to Vertex only; Gemini sets it too
7	Low	The test leaves a stub magic module in sys.modules

Lens checklist

1. Spec & intent — See findings #1, #2. The repo states no PR-title convention, so no title check was made.
2. Architectural fit — Clean
3. Correctness — See findings #1, #4
4. Security — Clean. No endpoint, parser, auth or untrusted-input surface. Boundary sweep: 1 sink (the debug log line) traced, clean.
5. Data integrity — N/A: no persisted data or migrations
6. Concurrency — See finding #2
7. API compatibility — Clean. No wire change. The retry helpers' only consumer is sdk1's llm.py; I also searched unstract-cloud, llmwhisperer-portal, the verticals repos and omnibench-platform.
8. Reliability — See finding #1
9. Performance & cost — Clean
10. Observability — See finding #4
11. Operational safety — Clean
12. LLM/agent — Clean. stream_chunk_builder reads only the collected chunks.
13. Testing — See findings #3, #7
14. Dependencies — N/A: no dependency changes
15. Code quality — Clean. ruff 0.3.4 check and format --check pass.
16. Doc & comment accuracy — See findings #2, #5, #6
17. Cross-document consistency — Clean. The branch is current with main, and the removed "leak sockets until GC" comments are not restated elsewhere.
18. On-prem — N/A: no chart or on-prem paths
19. Scaling & lifecycle — N/A: no worker, pool or scaling config

Open questions for the author

• Has the UN-4136 600-user load test been re-run on a snapshot build? The PR lists it as still to do.
• Does any Unstract adapter offer Responses-API-routed models (gpt-5-pro, *-codex, o3-pro)? The answer decides how much #1 matters in practice.

Assumptions

• The table_extractor plugin in unstract-cloud runs in executor workers.
• Verification used litellm 1.104.0, the pinned version. On 1.90.3, which some local venvs carry, the deadlock test also passes against main.

[greptile-apps]

via Greptile

Confidence Score: 5/5

[High impact] The PR appears safe to merge; no blocking issue was found.

Summary

The PR closes streamed LLM responses after success, failure, and early generator closure. Retry waits now happen after cleanup.

• Adds a bounded walk through provider stream handles.
• Adds unit and local-server tests for connection release.
• Documents the caller's duty to close an unfinished generator.
• The previous cleanup finding is fixed: reading close now happens inside the guarded block.
• No new actionable issues were found. Tests were inspected, not run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Open streamed attempt] --> B[Consume chunks]
  B --> C[Finish, fail, or close early]
  C --> D[Finally: close reachable stream handles]
  D --> E{Retry allowed?}
  E -->|Yes| F[Wait before retry]
  F --> A
  E -->|No| G[Return or raise]

Loading

_{Reviews (5) · Last reviewed commit: "Merge branch 'main' into ccr-b889a1d6-n3..." · Reviewed by Greptile}

[chandrasekharan-zipstack]

Review notes. Overall the fix looks correct: the walk can't reach a shared httpx client or pool, the original error is never masked, and the tests catch the regression.

Two inline comments are on the close-failure logging in close_stream. The rest are comment cleanup: ticket numbers, litellm version pins and "before the fix" wording that will go stale. The goal is to keep the explanation in one place (close_stream) and keep the other comments generic.

[sonarqubecloud]

Quality Gate passed

Issues
2 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[github-actions]

Unstract test results

Per-group results

Status	Group	Tier	Passed	Failed	Errors	Skipped	Duration (s)
✅	e2e-api-deployment	e2e	3	0	0	0	12.5
✅	e2e-coowners	e2e	1	0	0	0	1.6
✅	e2e-etl	e2e	1	0	0	0	8.2
✅	e2e-login	e2e	2	0	0	0	1.1
✅	e2e-prompt-studio	e2e	1	0	0	0	9.6
✅	e2e-smoke	e2e	2	0	0	0	1.2
✅	e2e-workflow	e2e	1	0	0	0	14.3
✅	frontend	unit	620	0	0	0	12.4
✅	integration-backend	integration	603	0	0	26	59.9
✅	integration-connectors	integration	1	0	0	7	9.7
✅	integration-workers	integration	164	0	0	1	39.6
❌	ui	e2e	0	1	0	0	0.0
✅	unit-backend	unit	1478	0	0	1	34.6
✅	unit-connectors	unit	82	0	0	0	8.7
✅	unit-core	unit	281	0	0	0	2.1
✅	unit-platform-service	unit	18	0	0	0	2.0
✅	unit-rig	unit	120	0	0	0	3.0
✅	unit-runner	unit	12	0	0	0	2.1
✅	unit-sdk1	unit	843	0	0	0	31.5
✅	unit-tool-sidecar	unit	8	0	0	0	1.3
✅	unit-workers	unit	1412	0	0	1	111.2
	TOTAL		5653	1	0	36	366.5

Critical paths

⚠️ Critical paths not yet covered

• workflow-execution-fan-out — Multi-file workflow execution fans out to file-processing workers and rejoins. (declared coverage: no groups declared)

✅ Covered critical paths

• auth-login — covered by e2e-login
• adapter-register-llm — covered by integration-backend
• workflow-author — covered by integration-backend
• co-owner-manage — covered by integration-backend, e2e-coowners
• workflow-create-execute — covered by e2e-workflow
• api-deployment-provision — covered by integration-backend
• api-deployment-auth — covered by integration-backend
• api-deployment-run — covered by e2e-api-deployment
• mcp-server-auth — covered by integration-backend
• mcp-platform-auth — covered by integration-backend
• platform-key-whoami — covered by integration-backend
• prompt-studio-author — covered by integration-backend
• prompt-studio-fetch-response — covered by e2e-prompt-studio
• connector-register-test — covered by integration-backend
• pipeline-etl-execute — covered by e2e-etl
• usage-aggregate-read — covered by integration-backend
• usage-token-tracking — covered by e2e-api-deployment
• callback-result-delivery — covered by e2e-api-deployment