mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): release _active_streams session lock on setup-phase errors (#1864)
## Description
`_active_streams.add(session_key)` in `_stream_response` ran before the
request setup was protected by cleanup. If header preparation, Copilot
auth, outbound-body serialization, or an `asyncio.CancelledError` from a
client disconnect failed before the streaming generator was created, the
session key stayed in `_active_streams` permanently. Subsequent requests
for the same session were then queued forever as `202 headroom_queued`
responses until the proxy restarted.
Closes the setup-phase leak by wrapping the whole pre-generator path in
a thin `_stream_response` guard and moving the existing implementation
into `_stream_response_inner`. The guard releases the session key via
`_cleanup_mid_turn_stream` on `Exception` or `asyncio.CancelledError`,
while the existing generator `finally` still owns cleanup once streaming
starts.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- Split `headroom/proxy/handlers/streaming.py` so `_stream_response`
becomes the cleanup wrapper and `_stream_response_inner` contains the
existing streaming implementation.
- Added setup-phase cleanup for failures before the streaming
generator's own `finally` can run.
- Maintainer follow-up: moved the `Response` / `StreamingResponse`
runtime import into `_stream_response_inner` so lint passes and the
inner implementation can construct `StreamingResponse`.
## Testing
- [x] Syntax check passes (`python -m py_compile
headroom/proxy/handlers/streaming.py`)
- [x] Linting passes for the touched file (`uv run ruff check
headroom/proxy/handlers/streaming.py`)
- [ ] Full CI passes
- [ ] New tests added for new functionality
- [x] Manual testing performed
### Test Output
```text
python -m py_compile headroom/proxy/handlers/streaming.py
# passed
uv run ruff check headroom/proxy/handlers/streaming.py
All checks passed!
```
The earlier CI failure was caused by the local import remaining in the
outer wrapper after the implementation split. A maintainer follow-up
commit moved that import into `_stream_response_inner`; the full CI
rerun is pending on the updated branch.
## Real Behavior Proof
- **Environment:** GitHub Copilot Chat in VS Code routed through
Headroom proxy.
- **Exact command / steps:** During normal Copilot Chat usage, a
streaming setup-phase failure/client disconnect occurred before
`_stream_response` reached the generator cleanup path.
- **Observed result:** After the setup failure, every later request for
that session returned `202 {"status":202,"event":"headroom_queued"}` and
Copilot Chat treated the response as a hard server error. Restarting the
proxy cleared the in-memory `_active_streams` set and restored the
session.
- **Not tested:** A deterministic end-to-end reproduction of the
original VS Code disconnect timing. The code path was reviewed directly,
and the branch has a pending full CI rerun after the maintainer import
fix.
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
## Checklist
- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable
## Screenshots (if applicable)
N/A
## Additional Notes
Documentation and changelog updates are not required for this narrow
internal proxy cleanup fix. A focused regression test would still be
valuable if we can isolate the setup-phase cancellation path without
making the streaming tests brittle.
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
3076e32172
commit
2ccd831032
1 changed files with 64 additions and 2 deletions
|
|
@ -945,11 +945,73 @@ class StreamingMixin:
|
|||
3. Makes continuation requests until no memory tools remain
|
||||
4. Streams the final response to the client
|
||||
"""
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
session_key = session_key or self._get_session_key(body)
|
||||
self._active_streams.add(session_key)
|
||||
|
||||
# Guard everything up to the generator's own try/finally (which owns
|
||||
# cleanup once streaming starts): any exception here — including
|
||||
# asyncio.CancelledError from a client disconnect mid-setup — must
|
||||
# still release session_key, or it wedges in _active_streams forever
|
||||
# and every later request on this session gets stuck 202-queued.
|
||||
try:
|
||||
return await self._stream_response_inner(
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=body,
|
||||
provider=provider,
|
||||
model=model,
|
||||
request_id=request_id,
|
||||
original_tokens=original_tokens,
|
||||
optimized_tokens=optimized_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
transforms_applied=transforms_applied,
|
||||
tags=tags,
|
||||
optimization_latency=optimization_latency,
|
||||
memory_user_id=memory_user_id,
|
||||
pipeline_timing=pipeline_timing,
|
||||
prefix_tracker=prefix_tracker,
|
||||
original_messages=original_messages,
|
||||
original_body_bytes=original_body_bytes,
|
||||
body_mutated=body_mutated,
|
||||
mutation_reasons=mutation_reasons,
|
||||
memory_request_ctx=memory_request_ctx,
|
||||
outcome_provider=outcome_provider,
|
||||
waste_signals=waste_signals,
|
||||
session_key=session_key,
|
||||
)
|
||||
except (Exception, asyncio.CancelledError):
|
||||
self._cleanup_mid_turn_stream(session_key)
|
||||
raise
|
||||
|
||||
async def _stream_response_inner(
|
||||
self,
|
||||
url: str,
|
||||
headers: dict,
|
||||
body: dict,
|
||||
provider: str,
|
||||
model: str,
|
||||
request_id: str,
|
||||
original_tokens: int,
|
||||
optimized_tokens: int,
|
||||
tokens_saved: int,
|
||||
transforms_applied: list[str],
|
||||
tags: dict[str, str],
|
||||
optimization_latency: float,
|
||||
memory_user_id: str | None,
|
||||
pipeline_timing: dict[str, float] | None,
|
||||
prefix_tracker: Any | None,
|
||||
original_messages: list[dict] | None,
|
||||
original_body_bytes: bytes | None,
|
||||
body_mutated: bool,
|
||||
mutation_reasons: list[str] | None,
|
||||
memory_request_ctx: Any | None,
|
||||
outcome_provider: str | None,
|
||||
waste_signals: dict[str, int] | None,
|
||||
session_key: str,
|
||||
) -> Response | StreamingResponse:
|
||||
"""Actual streaming implementation, guarded by _stream_response's cleanup wrapper."""
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from headroom.proxy.helpers import MAX_SSE_BUFFER_SIZE
|
||||
|
||||
# Identify the harness (codex / claude-code / aider / cursor /
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue