fix(llm): subprocess hard-timeout deadlock + caller message-dict mutation

[guangyu-reflexio]

Summary

Two latent robustness fixes in the litellm client (both confirmed with a repro test that fails pre-fix / passes post-fix). A third CodeRabbit-flagged item (_build_completion_params/_apply_prompt_caching mutation) was investigated and found to be a false positive — those already copy-on-write — so it's not touched.

1. Subprocess hard-timeout deadlock → false timeout (availability)

_completion_with_hard_timeout joined the child process before draining the mp.Queue result. A large completion payload (>~64KB) fills the OS pipe buffer, so the child blocks on queue.put(...) while the parent blocks on join(...) — the parent then hits its timeout and raises LLMHardTimeoutError for a call that actually succeeded. Fix: drain the queue (bounded by the same hard_timeout budget) before join; kill/terminate + ok/error snapshot semantics preserved. Regression test uses a 2 MB payload.

2. Caller message-dict mutation on system_message merge (correctness)

generate_chat_response did final_messages = list(messages) (a shallow copy sharing the caller's dict objects) then merged the system message into final_messages[0] in place, corrupting the caller's list and re-prepending on reuse/retry. Fix: replace the list slot with a new dict instead of mutating the shared one.

Test plan

• Both regression tests proven fail-pre (via git stash of the fix) / pass-post.
• tests/server/llm/ 490 passed / 60 skipped, no mock snapshot churn; ruff clean; pyright 0 errors on changed files.

Note

• Out of scope (deferred, distinct): the litellm cleartext image-URL "SSRF" is a security-policy question (inherent to the vision feature), left for a separate decision.

Summary by CodeRabbit

• Bug Fixes
  • Prevented chat request system messages from being modified in place, keeping caller-provided message data unchanged.
  • Improved handling of long-running AI responses so large outputs no longer trigger false hard timeouts or deadlocks.
  • Made result collection more reliable when processing large payloads, helping requests complete successfully and on time.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 6373057f-7a43-445a-acab-21ede3c6263e

📥 Commits

Reviewing files that changed from the base of the PR and between a71a5c1 and c5483e4.

📒 Files selected for processing (2)

• reflexio/server/llm/_litellm_text_generation.py
• tests/server/llm/test_litellm_client_unit.py

---

📝 Walkthrough

Walkthrough

Modified generate_chat_response to build a new merged system-message dict instead of mutating caller-provided messages, and reworked _completion_with_hard_timeout to drain the result queue before joining the subprocess, preventing deadlocks with large payloads. Added corresponding regression tests and reformatted two unrelated patch calls.

Changes

Text generation fixes

Layer / File(s)	Summary
Non-mutating system message merge reflexio/server/llm/_litellm_text_generation.py, tests/server/llm/test_litellm_client_unit.py	generate_chat_response replaces the first system message with a newly built dict rather than mutating the caller's dict; a new test verifies the original message and object identity are preserved across repeated calls.
Queue-drain-before-join for hard timeout reflexio/server/llm/_litellm_text_generation.py, tests/server/llm/test_litellm_client_unit.py	_completion_with_hard_timeout uses a monotonic deadline loop to drain the result queue before joining/terminating the child process, distinguishing true timeouts from child exits without results; a new test with a 2MB payload confirms no false deadlock/timeout.
Unrelated test formatting cleanup tests/server/llm/test_litellm_client_unit.py	Two existing structured-output tests reformat patch(...) calls for assert_provider_safe_schema into multiline form without behavior changes.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• ReflexioAI/reflexio#273: Introduced the TextGenerationMixin/_completion_with_hard_timeout subprocess machinery that this PR fixes and adds regression tests for.

Poem

A rabbit hopped through queues at night,
Draining results before the fight,
No more mutated dicts to fear,
Each caller's message stays sincere,
Two megabytes, no deadlock in sight. 🐇

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the two main fixes: hard-timeout deadlock prevention and avoiding caller message dict mutation.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/litellm-robustness

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}