github-featured/chopratejas--headroom
2
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix(proxy): complete stateless Responses and buffered CCR lifecycle (#2997)

Browse Source 
## Description

Consolidates the related OpenAI Responses ZDR/stateless continuation and
buffered CCR response-lifecycle corrections on current main. It
preserves client storage policy, makes Headroom-owned continuations
stateless across HTTP and WebSocket, and prevents buffered streaming
paths from committing a false HTTP 200 before the real upstream outcome
is known.

Closes #2675

## 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

- Preserves explicit and omitted Responses `store` policy instead of
forcing provider storage or disabling memory tools.
- Replays normalized input, replayable outputs, encrypted reasoning
content, and Headroom function outputs without `previous_response_id`.
- Applies the same stateless continuation policy to HTTP and WebSocket.
- Prevents transparent memory execution after client-visible WebSocket
output.
- Delays buffered CCR ASGI status/headers until the operation resolves
for Anthropic Messages and OpenAI Responses.
- Preserves real 429/5xx status and retry headers.
- Converts malformed non-JSON/non-SSE upstream 200 replies to a
sanitized 502 protocol error.
- Preserves valid JSON-to-SSE synthesis and existing SSE adaptation.
- Removes unreachable task cleanup left behind after replacing the old
keepalive polling loop with a direct awaited operation.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
118 passed across the changed HTTP/WS ZDR, lifecycle, and both-provider CCR suites
11526 tests collected with no collection errors
ruff check .: All checks passed
ruff format --check .: 1411 files already formatted
mypy headroom/proxy/handlers/anthropic.py headroom/proxy/handlers/openai.py:
Success: no issues found in 2 source files
```

Exact-head CI is entirely green on
`cbc2739c0c633ea477727ec3eb2f8a3862fa08f3`.

## Real Behavior Proof

- Environment: macOS arm64/Python 3.13 locally; GitHub-hosted Ubuntu
matrix pending.
- Exact command / steps: exercise `store=false` Responses memory calls
over HTTP and WebSocket; exercise buffered Anthropic and Responses
requests returning successful JSON/SSE, delayed 429 responses,
exceptions, and malformed successful bodies; invoke returned ASGI
responses and inspect emitted status, headers, and body order.
- Observed result: stateless continuations omit provider response IDs
and retain `store=false`; no ASGI start event is emitted before the
buffered outcome; real failures preserve status/headers; malformed 200
responses become sanitized 502 errors.
- Not tested: live ZDR tenant and live Anthropic/OpenAI upstream
credentials are unavailable in repository CI; wire contracts are
exercised through deterministic upstream doubles.

## Runtime Rollout Safety

- Rollout-managed feature(s): Responses memory continuation and buffered
CCR handling.
- Minimum rollout channel: normal patch release after full CI
qualification.
- Stable/default behavior changed: memory continuation no longer
requires provider storage; buffered CCR waits before committing response
status.
- Kill switch / disable path: disable memory/CCR using existing proxy
configuration (`--no-ccr` for CCR); ordinary non-buffered paths are
unchanged.
- Unsafe override required: none.
- Qualification impact: full Python matrix plus focused HTTP/WS
lifecycle suites must pass; patch coverage must not rely on unreachable
cleanup.
- Rollback path: human revert of this PR restores prior
continuation/buffering behavior; no persisted data migration is
introduced.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review — exact-head CI is entirely
green

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation — inline
protocol/lifecycle documentation; no separate user guide required
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md` — it is generated by
release-please from my Conventional Commit PR title (a CI guard enforces
this)

## Screenshots (if applicable)

Not applicable; proxy protocol behavior.

## Additional Notes

Human review only. No merge or auto-merge is configured. This supersedes
narrower #2995 and incorporates the complete intent of #2705, #2959, and
#2968 without falsely closing those PRs. It does not claim the broader
event-level streaming-splice guarantees requested by #1877. Refreshed
from main after #2996; the MCP cap `mcp>=1.28.1,<2.0.0` is preserved.
This commit is contained in:
JD Davis
2026-08-13 21:12:38 -05:00
committed by GitHub
parent a708c0571e
commit 8a1d38bc5d
7 changed files with 1274 additions and 454 deletions
Download Patch File Download Diff File Expand all files Collapse all files
headroom/proxy/handlers/anthropic.py
+60 -113
View File
@@ -4195,6 +4195,25 @@ class AnthropicHandlerMixin:
headers=relay_headers,
)
if buffered_stream_ccr and response.status_code == 200 and not resp_json:
logger.warning(
f"[{request_id}] CCR: rejecting malformed buffered 200 reply "
f"(content-type={response.headers.get('content-type')!r}, "
f"body_bytes={len(response.content)})"
)
return Response(
content=json.dumps(
{
"error": {
"type": "upstream_protocol_error",
"message": "Upstream returned an invalid buffered response.",
}
}
),
status_code=502,
media_type="application/json",
)
if buffered_stream_ccr and response.status_code == 200 and resp_json:
sse_headers = {
k: v
@@ -4293,122 +4312,50 @@ class AnthropicHandlerMixin:
class _BufferedCCRResponse(Response):
async def __call__(self, scope, receive, send): # noqa: ANN001
await asyncio.sleep(0)
loop = asyncio.get_running_loop()
keepalive_deadline = loop.time() + 1.0
started = False
# Send nothing until the buffered operation resolves.
# The previous keepalive preamble committed
# `200 text/event-stream` after 1s, i.e. before the
# outcome was known: any upstream reply that then
# failed to become SSE (non-200, unparseable body)
# reached the client as a 200 whose body carried no
# `message_start`, which Claude Code reports as "API
# returned an empty or malformed response (HTTP 200) —
# check for a proxy or gateway intercepting the
# request". The real status was lost with it, so
# client-side 429/5xx backoff never fired. Clients
# budget minutes for a turn (Claude Code sends
# `x-stainless-timeout: 600`), so waiting is free.
try:
while True:
timeout = (
0.25
if started
else max(0.0, keepalive_deadline - loop.time())
)
done, _ = await asyncio.wait({operation}, timeout=timeout)
if done:
try:
result = operation.result()
except Exception as e:
await record_failed(provider=provider_name)
logger.error(
f"[{request_id}] Request failed: {type(e).__name__}: {e}"
)
if not started:
await send(
{
"type": "http.response.start",
"status": 502,
"headers": [
(b"content-type", b"application/json")
],
}
)
await send(
{
"type": "http.response.body",
"body": json.dumps(
{
"type": "error",
"error": {
"type": "api_error",
"message": "An error occurred while processing your request. Please try again.",
},
}
).encode(),
"more_body": False,
}
)
return
await send(
{
"type": "http.response.body",
"body": b'event: error\ndata: {"type":"error","error":{"type":"api_error","message":"An error occurred while processing the request."}}\n\n',
"more_body": False,
}
)
return
if not started:
await result(scope, receive, send)
return
body_iterator = getattr(result, "body_iterator", None)
if body_iterator is not None:
async for chunk in body_iterator:
await send(
{
"type": "http.response.body",
"body": chunk,
"more_body": True,
}
)
await send(
{
"type": "http.response.body",
"body": b"",
"more_body": False,
}
)
return
await send(
result = await operation
except Exception as e:
await record_failed(provider=provider_name)
logger.error(
f"[{request_id}] Request failed: {type(e).__name__}: {e}"
)
await send(
{
"type": "http.response.start",
"status": 502,
"headers": [(b"content-type", b"application/json")],
}
)
await send(
{
"type": "http.response.body",
"body": json.dumps(
{
"type": "http.response.body",
"body": b'event: error\ndata: {"type":"error","error":{"type":"api_error","message":"An error occurred while processing the request."}}\n\n',
"more_body": False,
"type": "error",
"error": {
"type": "api_error",
"message": "An error occurred while processing your request. Please try again.",
},
}
)
return
if not started:
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [
(b"content-type", b"text/event-stream")
],
}
)
started = True
await send(
{
"type": "http.response.body",
"body": b'event: ping\ndata: {"type":"ping"}\n\n',
"more_body": True,
}
)
except asyncio.CancelledError:
raise
finally:
if not operation.done():
operation.cancel()
try:
await operation
except asyncio.CancelledError:
pass
except Exception:
pass
).encode(),
"more_body": False,
}
)
return
await result(scope, receive, send)
return _BufferedCCRResponse(media_type="text/event-stream")
return await _buffered_ccr_operation()
headroom/proxy/handlers/openai.py
+280 -273
View File
@@ -763,36 +763,9 @@ def _compact_openai_responses_tools(
return compact_tools(payload)
def _responses_request_allows_memory_tool_continuation(payload: dict[str, Any]) -> bool:
"""Return whether Responses memory tools may rely on stored continuations.
Headroom memory tools use ``previous_response_id`` continuations after a
tool call. Those continuations require the originating response to be
stored. When a client explicitly sends ``store=false``, preserve that
contract and skip the Responses memory-tool injection path instead of
mutating the request.
"""
return payload.get("store") is not False
def _ensure_responses_store_for_memory_tools(
payload: dict[str, Any],
*,
memory_tools_injected: bool,
) -> bool:
"""Return True when memory-tool injection requires and receives store=true."""
if memory_tools_injected and payload.get("store") is not True:
payload["store"] = True
return True
return False
def _allow_responses_memory_tools(is_chatgpt_auth: bool) -> bool:
# ChatGPT Codex rejects Responses payloads unless store=false. The
# transparent memory-tool continuation flow needs stored responses, so keep
# it on the regular API path only.
# Preserve the ChatGPT Codex route's existing store policy and memory-tool
# exclusion while API Responses memory continuations stay stateless.
return not is_chatgpt_auth
@@ -1047,6 +1020,23 @@ def _responses_input_to_items(input_data: Any) -> list[dict[str, Any]]:
return []
def _responses_stateless_output_items(output_items: Any) -> list[dict[str, Any]]:
"""Return response items that can be replayed without provider state."""
if not isinstance(output_items, list):
return []
return [
item
for item in output_items
if isinstance(item, dict)
and not (item.get("type") == "reasoning" and not item.get("encrypted_content"))
]
def _responses_stateless_input_items(input_data: Any) -> list[dict[str, Any]]:
"""Normalize input and remove reasoning items without encrypted content."""
return _responses_stateless_output_items(_responses_input_to_items(input_data))
def _dedup_responses_output_items(
items: list[dict[str, Any]],
output_types: frozenset[str],
@@ -5345,10 +5335,7 @@ class OpenAIHandlerMixin:
else:
memory_tool_defs_responses.append(t)
if (
responses_memory_tools_allowed
and _responses_request_allows_memory_tool_continuation(body)
):
if responses_memory_tools_allowed:
resp_tools = body.get("tools") or []
resp_tools, mem_tools_injected = _apply_sticky_mem_tools_resp(
provider="openai",
@@ -5360,23 +5347,16 @@ class OpenAIHandlerMixin:
)
if mem_tools_injected:
body["tools"] = resp_tools
include = body.get("include")
if isinstance(include, list):
if "reasoning.encrypted_content" not in include:
body["include"] = [*include, "reasoning.encrypted_content"]
elif include is None:
body["include"] = ["reasoning.encrypted_content"]
body_mutation_tracker.mark_mutated("responses_memory_tools")
logger.info(
f"[{request_id}] Memory: Injected memory tools (openai/responses)"
)
if _ensure_responses_store_for_memory_tools(
body,
memory_tools_injected=True,
):
body_mutation_tracker.mark_mutated("responses_memory_store")
logger.info(
f"[{request_id}] Memory: forced store=true for Responses memory tool continuation"
)
elif self.memory_handler.config.inject_tools:
logger.info(
"[%s] Memory: skipped Responses memory tools because client set store=false",
request_id,
)
except Exception as e:
logger.warning(f"[{request_id}] Memory injection failed (responses): {e}")
elif self.memory_handler and memory_user_id and _bypass:
@@ -5808,6 +5788,7 @@ class OpenAIHandlerMixin:
total_input_tokens = original_tokens # fallback
output_tokens = 0
cache_read_tokens = 0
resp_json = None
try:
resp_json = response.json()
usage = resp_json.get("usage", {})
@@ -5826,7 +5807,13 @@ class OpenAIHandlerMixin:
details = usage.get("input_tokens_details")
if isinstance(details, dict):
cache_read_tokens = _usage_int(details.get("cached_tokens"))
except (KeyError, TypeError, AttributeError) as e:
except (
json.JSONDecodeError,
ValueError,
KeyError,
TypeError,
AttributeError,
) as e:
logger.debug(
f"[{request_id}] Failed to extract cached tokens from OpenAI passthrough response: {e}"
)
@@ -5996,17 +5983,17 @@ class OpenAIHandlerMixin:
)
if tool_outputs:
# Make continuation request with tool results
response_id = resp_json.get("id")
# Make a stateless continuation with the complete
# item history instead of relying on retained state.
continuation_body = {
"model": model,
"input": tool_outputs,
**body,
"input": [
*_responses_stateless_input_items(body.get("input")),
*_responses_stateless_output_items(output_items),
*tool_outputs,
],
}
if response_id:
continuation_body["previous_response_id"] = response_id
existing_tools = body.get("tools")
if existing_tools:
continuation_body["tools"] = existing_tools
continuation_body.pop("previous_response_id", None)
cont_response = await self._retry_request(
"POST", url, headers, continuation_body
@@ -6144,6 +6131,25 @@ class OpenAIHandlerMixin:
headers=sse_headers,
)
if buffered_stream_ccr and response.status_code == 200 and not resp_json:
logger.warning(
f"[{request_id}] CCR: rejecting malformed buffered Responses 200 "
f"reply (content-type={response.headers.get('content-type')!r}, "
f"body_bytes={len(response.content)})"
)
return Response(
content=json.dumps(
{
"error": {
"type": "upstream_protocol_error",
"message": "Upstream returned an invalid buffered response.",
}
}
),
status_code=502,
media_type="application/json",
)
# Inline marker resolution, non-streaming only. Runs
# outside the has_ccr_tool_calls gate above on purpose:
# the #2509 case has no retrieve tool call at all.
@@ -6171,118 +6177,44 @@ class OpenAIHandlerMixin:
class _BufferedCCRResponse(Response):
async def __call__(self, scope, receive, send): # noqa: ANN001
await asyncio.sleep(0)
loop = asyncio.get_running_loop()
keepalive_deadline = loop.time() + 1.0
started = False
# Send nothing until the buffered operation resolves —
# see the AnthropicHandler twin for the full rationale.
# Committing `200 text/event-stream` on a keepalive timer,
# before the outcome is known, turns every non-200 or
# unparseable upstream reply into a 200 with no usable
# body and discards the status the client needs to back
# off on.
try:
while True:
timeout = (
0.25 if started else max(0.0, keepalive_deadline - loop.time())
)
done, _ = await asyncio.wait({operation}, timeout=timeout)
if done:
try:
result = operation.result()
except Exception as e:
await record_failed(provider="openai")
logger.error(
f"[{request_id}] OpenAI responses request failed: {type(e).__name__}: {e}"
)
if not started:
await send(
{
"type": "http.response.start",
"status": 502,
"headers": [
(b"content-type", b"application/json")
],
}
)
await send(
{
"type": "http.response.body",
"body": json.dumps(
{
"error": {
"message": "An error occurred while processing your request. Please try again.",
"type": "server_error",
"code": "proxy_error",
}
}
).encode(),
"more_body": False,
}
)
return
await send(
{
"type": "http.response.body",
"body": b'event: error\ndata: {"type":"error","error":{"message":"An error occurred while processing the request."}}\n\n',
"more_body": False,
}
)
return
if not started:
await result(scope, receive, send)
return
body_iterator = getattr(result, "body_iterator", None)
if body_iterator is not None:
async for chunk in body_iterator:
await send(
{
"type": "http.response.body",
"body": chunk,
"more_body": True,
}
)
await send(
{
"type": "http.response.body",
"body": b"",
"more_body": False,
}
)
return
await send(
result = await operation
except Exception as e:
await record_failed(provider="openai")
logger.error(
f"[{request_id}] OpenAI responses request failed: {type(e).__name__}: {e}"
)
await send(
{
"type": "http.response.start",
"status": 502,
"headers": [(b"content-type", b"application/json")],
}
)
await send(
{
"type": "http.response.body",
"body": json.dumps(
{
"type": "http.response.body",
"body": b'event: error\ndata: {"type":"error","error":{"message":"An error occurred while processing the request."}}\n\n',
"more_body": False,
"error": {
"message": "An error occurred while processing your request. Please try again.",
"type": "server_error",
"code": "proxy_error",
}
}
)
return
if not started:
await send(
{
"type": "http.response.start",
"status": 200,
"headers": [(b"content-type", b"text/event-stream")],
}
)
started = True
await send(
{
"type": "http.response.body",
"body": b'event: ping\ndata: {"type":"ping"}\n\n',
"more_body": True,
}
)
except asyncio.CancelledError:
raise
finally:
if not operation.done():
operation.cancel()
try:
await operation
except asyncio.CancelledError:
pass
except Exception:
pass
).encode(),
"more_body": False,
}
)
return
await result(scope, receive, send)
return _BufferedCCRResponse(media_type="text/event-stream")
return await _buffered_ccr_operation()
@@ -6827,6 +6759,8 @@ class OpenAIHandlerMixin:
return json.dumps(flattened, ensure_ascii=False)
body: dict[str, Any] = {}
current_response_input: list[dict[str, Any]] = []
current_response_template: dict[str, Any] = {}
tokens_saved = 0
# Session-scoped accumulator for tokens we *attempted* to
# compress (extracted units + schema). Drives the active-
@@ -7025,7 +6959,8 @@ class OpenAIHandlerMixin:
try:
# Unwrap response.create envelope to access the response body
ws_response_body = frame_body.get("response", frame_body)
if not isinstance(ws_response_body, dict):
return frame_raw
# Per-project memory routing (GH #462). For WS,
# ``ws_response_body`` carries ``instructions`` —
# that's the system-prompt-equivalent we feed to the
@@ -7171,6 +7106,15 @@ class OpenAIHandlerMixin:
)
if mem_injected:
ws_response_body["tools"] = ws_tools
include = ws_response_body.get("include")
if isinstance(include, list):
if "reasoning.encrypted_content" not in include:
ws_response_body["include"] = [
*include,
"reasoning.encrypted_content",
]
elif include is None:
ws_response_body["include"] = ["reasoning.encrypted_content"]
# Add memory instruction so the model uses
# memory tools as persistent cross-session knowledge.
@@ -7210,6 +7154,12 @@ class OpenAIHandlerMixin:
body.get("type") == "response.create" or ("type" not in body and "input" in body)
):
first_msg_raw = await _prepare_memory_frame(body, first_msg_raw)
first_response_body = body.get("response", body)
current_response_input = _responses_input_to_items(
first_response_body.get("input")
if isinstance(first_response_body, dict)
else None
)
# Hot-fix follow-up to PR #406 — inline Rust compression on the
# WS first frame before forwarding upstream. PR #406 enabled
@@ -7469,6 +7419,19 @@ class OpenAIHandlerMixin:
)
first_msg_raw = _normalize_ws_response_create_for_upstream(first_msg_raw)
try:
final_first_body = json.loads(first_msg_raw)
except json.JSONDecodeError:
final_first_body = None
if isinstance(final_first_body, dict):
final_first_response = final_first_body.get("response", final_first_body)
if isinstance(final_first_response, dict):
current_response_template = dict(final_first_response)
current_response_template.pop("input", None)
current_response_template.pop("previous_response_id", None)
current_response_input = _responses_input_to_items(
final_first_response.get("input")
)
_first_upstream_body: Any = None
try:
_first_upstream_body = json.loads(first_msg_raw)
@@ -7807,6 +7770,8 @@ class OpenAIHandlerMixin:
nonlocal ws_client_frames_total, ws_cancel_frames
nonlocal ws_frames_compressed
nonlocal ws_last_client_frame_type, ws_client_disconnect_seen
nonlocal current_response_input
nonlocal current_response_template
client_frame_index = 1
try:
while True:
@@ -7859,6 +7824,14 @@ class OpenAIHandlerMixin:
and _inbound_frame_body.get("type") == "response.create"
):
ws_response_create_frames += 1
inbound_response = _inbound_frame_body.get(
"response", _inbound_frame_body
)
current_response_input = _responses_input_to_items(
inbound_response.get("input")
if isinstance(inbound_response, dict)
else None
)
msg = await _prepare_memory_frame(_inbound_frame_body, msg)
(
msg,
@@ -7903,6 +7876,27 @@ class OpenAIHandlerMixin:
)
msg = _normalize_ws_response_create_for_upstream(msg)
if ws_last_client_frame_type == "response.create":
try:
outbound_body = json.loads(msg)
outbound_body = (
outbound_body.get("response", outbound_body)
if isinstance(outbound_body, dict)
else {}
)
if isinstance(outbound_body, dict):
current_response_template = dict(outbound_body)
current_response_template.pop("input", None)
current_response_template.pop(
"previous_response_id", None
)
candidate_input = _responses_input_to_items(
outbound_body.get("input")
)
if candidate_input:
current_response_input = candidate_input
except (json.JSONDecodeError, AttributeError, TypeError):
current_response_input = []
_outbound_frame_body: Any = None
try:
_outbound_frame_body = json.loads(msg)
@@ -7959,10 +7953,11 @@ class OpenAIHandlerMixin:
"""Relay upstream→client with transparent memory tool handling.
Uses a buffer-then-decide approach:
1. Buffer events until first output item arrives
2. If first output is a memory tool → suppress entire response,
execute tools silently, send continuation upstream
3. If first output is non-memory → flush buffer, stream normally
1. Buffer events until the first non-reasoning output item arrives
2. If a memory tool arrives, suppress the response and execute it
silently before sending a stateless continuation upstream
3. Flush ordinary output as soon as the response is known not to
begin with a reasoning-only prelude
4. Continuation response events are relayed to Codex seamlessly
This prevents orphaned response.created events from confusing Codex.
@@ -7987,25 +7982,22 @@ class OpenAIHandlerMixin:
nonlocal ws_recorded_overhead_ms_total, ws_recorded_ttfb_ms
nonlocal ws_upstream_frames_total, ws_last_upstream_frame_type
nonlocal ws_ttfb_ms
memory_enabled = bool(
self.memory_handler and memory_user_id and ws_memory_tools_allowed
)
nonlocal current_response_input
# Per-response state (reset after each response.completed)
event_buffer: list[str] = []
decided = False
suppress_response = False
pending_fcs: list[dict[str, Any]] = []
resp_id: str | None = None
response_output_items: list[dict[str, Any]] = []
def _reset() -> None:
nonlocal decided, suppress_response, resp_id
nonlocal decided, suppress_response
event_buffer.clear()
decided = False
suppress_response = False
pending_fcs.clear()
resp_id = None
response_output_items.clear()
response_started_ms: float | None = None
@@ -8233,6 +8225,11 @@ class OpenAIHandlerMixin:
ws_cache_write_tokens_total += usage_cache_write_tokens
ws_uncached_input_tokens_total += usage_uncached_tokens
memory_enabled = bool(
self.memory_handler
and memory_user_id
and ws_memory_tools_allowed
)
if not memory_enabled:
if event_type == "response.completed":
response_completed_seen = True
@@ -8240,17 +8237,25 @@ class OpenAIHandlerMixin:
await websocket.send_text(msg_str)
continue
# --- Phase 1: Buffer until first output item ---
if not decided:
event_buffer.append(msg_str)
if event_type == "response.output_item.added":
item = event.get("item", {})
if event_type == "response.output_item.done":
item = event.get("item", {})
if isinstance(item, dict):
response_output_items.append(item)
if (
item.get("type") == "function_call"
and item.get("name") in MEMORY_TOOL_NAMES
):
# Memory tool first → suppress entire response
pending_fcs.append(item)
if not decided:
event_buffer.append(msg_str)
if event_type == "response.output_item.added":
item = event.get("item", {})
if (
isinstance(item, dict)
and item.get("type") == "function_call"
and item.get("name") in MEMORY_TOOL_NAMES
):
suppress_response = True
decided = True
event_buffer.clear()
@@ -8258,106 +8263,108 @@ class OpenAIHandlerMixin:
f"[{request_id}] WS Memory: Detected "
f"{item.get('name')} — suppressing response"
)
else:
# Non-memory first → flush buffer, pass through
elif not (
isinstance(item, dict)
and item.get("type") == "reasoning"
):
decided = True
for buf in event_buffer:
await websocket.send_text(buf)
event_buffer.clear()
continue
elif event_type == "response.completed":
# No output items at all — flush
decided = True
for buf in event_buffer:
await websocket.send_text(buf)
event_buffer.clear()
await _record_ws_response_metrics()
_reset()
response_completed_seen = True
continue
# --- Phase 2a: Suppress mode (memory response) ---
if suppress_response:
if event_type == "response.output_item.done":
item = event.get("item", {})
if (
item.get("type") == "function_call"
and item.get("name") in MEMORY_TOOL_NAMES
):
pending_fcs.append(item)
elif event_type == "response.completed":
response_completed_seen = True
await _record_ws_response_metrics()
resp = event.get("response", {})
resp_id = resp.get("id")
if pending_fcs:
logger.info(
f"[{request_id}] WS Memory: Executing "
f"{len(pending_fcs)} tool(s) transparently"
)
# Execute memory tool calls
tool_outputs: list[dict[str, Any]] = []
for fc in pending_fcs:
call_id = fc.get("call_id", fc.get("id", ""))
fc_name = fc.get("name", "")
args_str = fc.get("arguments") or "{}"
try:
fc_args = json.loads(args_str)
except (json.JSONDecodeError, TypeError):
fc_args = {}
await self.memory_handler._ensure_initialized()
if self.memory_handler._backend:
result = await self.memory_handler._execute_memory_tool(
fc_name,
fc_args,
memory_user_id,
"openai",
)
else:
result = json.dumps(
{"error": "backend not ready"}
)
tool_outputs.append(
{
"type": "function_call_output",
"call_id": call_id,
"output": result,
}
)
logger.info(
f"[{request_id}] WS Memory: Executed "
f"{fc_name} for user {memory_user_id}"
)
# Send continuation upstream
cont: dict[str, Any] = {
"type": "response.create",
"response": {"input": tool_outputs},
}
if resp_id:
cont["response"]["previous_response_id"] = resp_id
await upstream.send(
_normalize_ws_response_create_for_upstream(
json.dumps(cont)
)
)
logger.info(
f"[{request_id}] WS Memory: Sent continuation "
f"with {len(tool_outputs)} result(s)"
)
for buf in event_buffer:
await websocket.send_text(buf)
_reset()
continue
if not decided:
continue
if not suppress_response:
await websocket.send_text(msg_str)
if event_type == "response.completed":
response_completed_seen = True
await _record_ws_response_metrics()
_reset()
# All events suppressed in this mode
continue
# --- Phase 2b: Pass-through mode ---
await websocket.send_text(msg_str)
if event_type != "response.completed":
continue
response_completed_seen = True
await _record_ws_response_metrics()
logger.info(
f"[{request_id}] WS Memory: Executing "
f"{len(pending_fcs)} tool(s) transparently"
)
# Execute memory tool calls.
tool_outputs: list[dict[str, Any]] = []
for fc in pending_fcs:
call_id = fc.get("call_id", fc.get("id", ""))
fc_name = fc.get("name", "")
args_str = fc.get("arguments") or "{}"
try:
fc_args = json.loads(args_str)
except (json.JSONDecodeError, TypeError):
fc_args = {}
await self.memory_handler._ensure_initialized()
if self.memory_handler._backend:
result = await self.memory_handler._execute_memory_tool(
fc_name,
fc_args,
memory_user_id,
"openai",
)
else:
result = json.dumps({"error": "backend not ready"})
tool_outputs.append(
{
"type": "function_call_output",
"call_id": call_id,
"output": result,
}
)
logger.info(
f"[{request_id}] WS Memory: Executed "
f"{fc_name} for user {memory_user_id}"
)
stateless_input = _responses_stateless_input_items(
current_response_input
)
stateless_output = _responses_stateless_output_items(
response_output_items
)
current_response_input = [
*stateless_input,
*stateless_output,
*tool_outputs,
]
continuation_response = {
**current_response_template,
"input": current_response_input,
}
continuation_response.pop("previous_response_id", None)
cont: dict[str, Any] = {
"type": "response.create",
"response": continuation_response,
}
continuation_raw = _strip_codex_lite_metadata(json.dumps(cont))
await upstream.send(
_normalize_ws_response_create_for_upstream(continuation_raw)
)
logger.info(
f"[{request_id}] WS Memory: Sent continuation "
f"with {len(tool_outputs)} result(s)"
)
_reset()
except asyncio.CancelledError:
raise
tests/test_openai_codex_routing.py
+232 -2
View File
@@ -2,10 +2,12 @@ import asyncio
import base64
import json
import sys
from copy import deepcopy
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import anyio
import httpx
import pytest
from fastapi import Request
@@ -14,6 +16,7 @@ from headroom.proxy.handlers.openai import (
_is_allowed_websocket_origin,
_openai_responses_unit_cache_key,
_resolve_codex_routing_headers,
_responses_stateless_output_items,
)
@@ -135,6 +138,30 @@ def test_openai_responses_unit_cache_key_includes_target_ratio() -> None:
assert aggressive_key != balanced_key
def test_responses_stateless_output_items_drop_unencrypted_reasoning() -> None:
assert _responses_stateless_output_items(None) == []
assert _responses_stateless_output_items(
[
{"type": "reasoning", "id": "rs-unusable", "summary": []},
{
"type": "reasoning",
"id": "rs-reusable",
"summary": [],
"encrypted_content": "encrypted",
},
{"type": "function_call", "call_id": "call-1"},
]
) == [
{
"type": "reasoning",
"id": "rs-reusable",
"summary": [],
"encrypted_content": "encrypted",
},
{"type": "function_call", "call_id": "call-1"},
]
class _DummyMetrics:
async def record_request(self, **kwargs): # noqa: ANN003
return None
@@ -274,6 +301,106 @@ class _MemoryToolsOnlyHandler:
return False
class _MemoryContinuationHandler(_MemoryToolsOnlyHandler):
async def _ensure_initialized(self) -> None:
self._backend = True
async def _execute_memory_tool(
self,
name: str,
args: dict,
user_id: str,
provider: str,
) -> str:
assert (name, args, user_id, provider) == (
"memory_search",
{},
"user-1",
"openai",
)
return '{"memories": []}'
def has_memory_tool_calls(self, response: dict, provider: str) -> bool:
assert provider == "openai"
return any(
item.get("name") == "memory_search"
for item in response.get("output", [])
if isinstance(item, dict)
)
class _ZdrResponsesHandler(_DummyOpenAIHandler):
def __init__(self) -> None:
super().__init__()
self.memory_handler = _MemoryContinuationHandler()
self.requests: list[dict] = []
async def _retry_request(self, method: str, url: str, headers: dict, body: dict, **kwargs):
assert (method, url) == ("POST", "https://api.openai.com/v1/responses")
self.requests.append(deepcopy(body))
request = httpx.Request(method, url)
if len(self.requests) == 1:
return httpx.Response(
200,
json={
"id": "resp-initial",
"output": [
{
"type": "reasoning",
"id": "reasoning-1",
"summary": [],
"encrypted_content": "encrypted-1",
},
{
"type": "function_call",
"id": "fc-1",
"call_id": "call-1",
"name": "memory_search",
"arguments": "{}",
},
],
"usage": {"input_tokens": 3, "output_tokens": 2},
},
request=request,
)
if any(
isinstance(item, dict)
and item.get("type") == "reasoning"
and not item.get("encrypted_content")
for item in body.get("input", [])
if isinstance(body.get("input"), list)
):
return httpx.Response(
400,
json={
"error": {"message": "Reasoning item is not reusable without encrypted_content"}
},
request=request,
)
if "previous_response_id" in body:
return httpx.Response(
400,
json={
"error": {
"message": "Unknown parameter: 'previous_response_id'.",
"type": "invalid_request_error",
"param": "previous_response_id",
"code": "unsupported_parameter",
}
},
request=request,
)
return httpx.Response(
200,
json={
"id": "resp-final",
"output": [{"type": "message", "id": "message-1"}],
"usage": {"input_tokens": 8, "output_tokens": 3},
},
request=request,
)
def _build_request(body: dict, headers: dict[str, str]) -> Request:
payload = json.dumps(body).encode("utf-8")
@@ -423,7 +550,7 @@ def test_handle_openai_responses_chatgpt_codex_timeout_fails_open(monkeypatch):
assert body["store"] is False
def test_handle_openai_responses_api_auth_store_false_skips_memory_tools(monkeypatch):
def test_handle_openai_responses_api_auth_store_false_injects_stateless_memory_tools(monkeypatch):
request = _build_request(
{"model": "gpt-4o-mini", "input": "hello", "store": False},
{"Authorization": "Bearer sk-test", "x-headroom-user-id": "user-1"},
@@ -444,10 +571,113 @@ def test_handle_openai_responses_api_auth_store_false_skips_memory_tools(monkeyp
_, url, _, body = handler.captured_request
assert url == "https://api.openai.com/v1/responses"
assert body["store"] is False
assert "tools" not in body
assert [tool["name"] for tool in body["tools"]] == ["memory_search"]
assert body["include"] == ["reasoning.encrypted_content"]
assert memory_handler.compute_calls == 1
@pytest.mark.parametrize("store", [pytest.param(None, id="omitted"), True, False])
@pytest.mark.parametrize(
"include",
[
pytest.param(None, id="omitted"),
pytest.param(["response.output_text.done"], id="missing-marker"),
pytest.param(
["response.output_text.done", "reasoning.encrypted_content"],
id="existing-marker",
),
pytest.param("not-a-list", id="non-list"),
],
)
def test_openai_responses_memory_continuation_is_zdr_safe(store, include, monkeypatch):
body = {
"model": "gpt-5.4",
"previous_response_id": "resp-inherited",
"input": [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "hello"}],
},
{
"type": "reasoning",
"id": "prior-reasoning",
"summary": [],
"encrypted_content": "prior-encrypted",
},
{
"type": "reasoning",
"id": "prior-unencrypted-reasoning",
"summary": [],
},
],
}
if include is not None:
body["include"] = include
if store is not None:
body["store"] = store
request = _build_request(
body,
{"Authorization": "Bearer sk-test", "x-headroom-user-id": "user-1"},
)
handler = _ZdrResponsesHandler()
monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer())
response = anyio.run(handler.handle_openai_responses, request)
assert response.status_code == 200
assert len(handler.requests) == 2
first_body, continuation_body = handler.requests
expected_include = (
["reasoning.encrypted_content"]
if include is None
else (
include
if not isinstance(include, list)
else (
include
if "reasoning.encrypted_content" in include
else [*include, "reasoning.encrypted_content"]
)
)
)
assert first_body["include"] == expected_include
assert continuation_body["include"] == expected_include
assert first_body["previous_response_id"] == "resp-inherited"
assert ("store" in first_body) is (store is not None)
if store is not None:
assert first_body["store"] is store
else:
assert "store" not in first_body
assert continuation_body["input"] == [
body["input"][0],
body["input"][1],
{
"type": "reasoning",
"id": "reasoning-1",
"summary": [],
"encrypted_content": "encrypted-1",
},
{
"type": "function_call",
"id": "fc-1",
"call_id": "call-1",
"name": "memory_search",
"arguments": "{}",
},
{
"type": "function_call_output",
"call_id": "call-1",
"output": '{"memories": []}',
},
]
assert "previous_response_id" not in continuation_body
assert ("store" in continuation_body) is (store is not None)
if store is not None:
assert continuation_body["store"] is store
def test_handle_openai_responses_routes_api_key_auth_direct_to_openai(monkeypatch):
request = _build_request(
{"model": "gpt-4o-mini", "input": "hello"},
tests/test_openai_codex_ws_lifecycle.py
+504
View File
@@ -70,6 +70,47 @@ class _DummyMetrics:
self.codex_ws_frames.append(dict(kwargs))
class _MemoryWsHandler:
def __init__(self) -> None:
self.config = SimpleNamespace(
inject_context=False,
inject_tools=True,
project_root_override="",
)
self._backend = False
def compute_memory_tool_definitions(self, provider: str) -> list[dict]:
assert provider == "openai"
return [
{
"type": "function",
"function": {
"name": "memory_search",
"description": "Search memory.",
"parameters": {"type": "object", "properties": {}},
},
}
]
async def _ensure_initialized(self) -> None:
self._backend = True
async def _execute_memory_tool(
self,
name: str,
args: dict,
user_id: str,
provider: str,
) -> str:
assert (name, args, user_id, provider) == (
"memory_search",
{},
user_id,
"openai",
)
return '{"memories": []}'
class _DummyOpenAIHandler(OpenAIHandlerMixin):
OPENAI_API_URL = "https://api.openai.com"
@@ -1739,3 +1780,466 @@ async def test_ws_recognized_client_with_real_path_is_not_restamped():
# the caller already self-identifies via its User-Agent.
assert "x-client" not in {k.lower() for k in client_ws.headers}
assert handler.ws_sessions.active_count() == 0
@pytest.mark.asyncio
@pytest.mark.parametrize("store", [True, False])
@pytest.mark.parametrize(
"include",
[
pytest.param(None, id="omitted"),
pytest.param(["response.output_text.done"], id="missing-marker"),
pytest.param(
["response.output_text.done", "reasoning.encrypted_content"],
id="existing-marker",
),
pytest.param("not-a-list", id="non-list"),
],
)
async def test_ws_memory_continuation_replays_history_without_previous_response_id(include, store):
function_call = {
"type": "function_call",
"id": "fc-1",
"call_id": "call-1",
"name": "memory_search",
"arguments": "{}",
}
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r-1"}}),
json.dumps({"type": "response.output_item.added", "item": function_call}),
json.dumps({"type": "response.output_item.done", "item": function_call}),
json.dumps({"type": "response.completed", "response": {"id": "r-1"}}),
]
upstream = _FakeUpstream(upstream_events)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(
frames=[
json.dumps(
{
"type": "response.create",
"response": {
"model": "gpt-5.4",
"input": "remember this",
"store": store,
},
}
)
],
hold_after_initial=True,
)
if include is not None:
client_ws._frames[0] = json.dumps(
{
"type": "response.create",
"response": {
"model": "gpt-5.4",
"input": "remember this",
"store": store,
"include": include,
},
}
)
client_ws.headers["x-headroom-user-id"] = "user-1"
handler = _DummyOpenAIHandler()
handler.memory_handler = _MemoryWsHandler()
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
assert len(upstream.sent) >= 2
expected_include = (
["reasoning.encrypted_content"]
if include is None
else (
include
if not isinstance(include, list)
else (
include
if "reasoning.encrypted_content" in include
else [*include, "reasoning.encrypted_content"]
)
)
)
assert json.loads(upstream.sent[0])["response"]["include"] == expected_include
continuation = json.loads(upstream.sent[1])
assert "previous_response_id" not in continuation["response"]
assert continuation["response"]["model"] == "gpt-5.4"
assert continuation["response"]["store"] is store
assert continuation["response"]["include"] == expected_include
assert continuation["response"]["tools"]
assert continuation["response"]["instructions"]
assert continuation["response"]["input"] == [
{"role": "user", "content": "remember this"},
function_call,
{
"type": "function_call_output",
"call_id": "call-1",
"output": '{"memories": []}',
},
]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"initial_frame",
[
pytest.param("not-json", id="initial-non-json"),
pytest.param(
json.dumps({"type": "response.create", "response": []}),
id="initial-non-mapping-response",
),
],
)
async def test_ws_memory_frame_shape_guards_fail_open(initial_frame):
later_frames = [
json.dumps(
{
"type": "response.create",
"response": {"model": "gpt-5.4", "input": []},
}
),
json.dumps({"type": "response.create", "response": "invalid"}),
]
frames = [initial_frame, *later_frames]
upstream = _FakeUpstream([], hold_after_events=True)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(frames=frames, hold_after_initial=True)
client_ws.headers["x-headroom-user-id"] = "user-1"
handler = _DummyOpenAIHandler()
handler.memory_handler = _MemoryWsHandler()
async def _trigger_disconnect() -> None:
await asyncio.sleep(0.05)
client_ws.trigger_disconnect()
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
trigger_task = asyncio.create_task(_trigger_disconnect())
try:
await asyncio.wait_for(handler.handle_openai_responses_ws(client_ws), timeout=2.0)
finally:
trigger_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await trigger_task
assert upstream.sent[0] == initial_frame
forwarded_valid = json.loads(upstream.sent[1])["response"]
assert forwarded_valid["model"] == "gpt-5.4"
assert forwarded_valid["input"] == []
assert forwarded_valid["tools"]
assert forwarded_valid["include"] == ["reasoning.encrypted_content"]
assert upstream.sent[2] == later_frames[1]
@pytest.mark.asyncio
async def test_ws_memory_enabled_non_memory_response_streams_completion():
message_item = {
"type": "message",
"id": "message-1",
"role": "assistant",
"content": [{"type": "output_text", "text": "hello"}],
}
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r-1"}}),
json.dumps({"type": "response.output_item.added", "item": message_item}),
json.dumps({"type": "response.output_item.done", "item": message_item}),
json.dumps({"type": "response.completed", "response": {"id": "r-1"}}),
]
upstream = _FakeUpstream(upstream_events)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(frames=[_first_frame()])
client_ws.headers["x-headroom-user-id"] = "user-1"
handler = _DummyOpenAIHandler()
handler.memory_handler = _MemoryWsHandler()
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
forwarded_initial = json.loads(upstream.sent[0])["response"]
assert forwarded_initial["model"] == "gpt-5.4"
assert forwarded_initial["input"] == "hi"
assert forwarded_initial["tools"]
assert forwarded_initial["include"] == ["reasoning.encrypted_content"]
assert client_ws.sent_text == upstream_events
assert len(upstream.sent) == 1
@pytest.mark.asyncio
async def test_ws_late_memory_call_after_streamed_message_passes_through():
message_item = {
"type": "message",
"id": "message-1",
"role": "assistant",
"content": [{"type": "output_text", "text": "searching"}],
}
function_call = {
"type": "function_call",
"id": "fc-1",
"call_id": "call-1",
"name": "memory_search",
"arguments": "{}",
}
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r-1"}}),
json.dumps({"type": "response.output_item.added", "item": message_item}),
json.dumps({"type": "response.output_item.done", "item": message_item}),
json.dumps({"type": "response.output_item.added", "item": function_call}),
json.dumps({"type": "response.output_item.done", "item": function_call}),
json.dumps({"type": "response.completed", "response": {"id": "r-1"}}),
]
upstream = _FakeUpstream(upstream_events)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(frames=[_first_frame()])
client_ws.headers["x-headroom-user-id"] = "user-1"
handler = _DummyOpenAIHandler()
handler.memory_handler = _MemoryWsHandler()
executed: list[tuple[str, dict, str, str]] = []
async def _execute_memory_tool(name, args, user_id, provider):
executed.append((name, args, user_id, provider))
return '{"memories": []}'
handler.memory_handler._execute_memory_tool = _execute_memory_tool
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
assert client_ws.sent_text == upstream_events
assert len(upstream.sent) == 1
assert executed == []
@pytest.mark.asyncio
async def test_ws_memory_continuation_handles_invalid_item_arguments_and_unavailable_backend():
function_call = {
"type": "function_call",
"id": "fc-1",
"call_id": "call-1",
"name": "memory_search",
"arguments": "{malformed",
}
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r-1"}}),
json.dumps({"type": "response.output_item.done", "item": "invalid"}),
json.dumps({"type": "response.output_item.added", "item": function_call}),
json.dumps({"type": "response.output_item.done", "item": function_call}),
json.dumps({"type": "response.completed", "response": {"id": "r-1"}}),
]
upstream = _FakeUpstream(upstream_events)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(frames=[_first_frame()])
client_ws.headers["x-headroom-user-id"] = "user-1"
handler = _DummyOpenAIHandler()
handler.memory_handler = _MemoryWsHandler()
async def _leave_backend_unavailable():
return None
handler.memory_handler._ensure_initialized = _leave_backend_unavailable
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
assert len(upstream.sent) == 2
continuation = json.loads(upstream.sent[1])["response"]
assert continuation["input"] == [
{"role": "user", "content": "hi"},
function_call,
{
"type": "function_call_output",
"call_id": "call-1",
"output": '{"error": "backend not ready"}',
},
]
assert continuation["input"][-1] == {
"type": "function_call_output",
"call_id": "call-1",
"output": '{"error": "backend not ready"}',
}
@pytest.mark.asyncio
async def test_ws_memory_continuation_normalizes_malformed_arguments():
function_call = {
"type": "function_call",
"id": "fc-1",
"call_id": "call-1",
"name": "memory_search",
"arguments": "{malformed",
}
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r-1"}}),
json.dumps({"type": "response.output_item.done", "item": "invalid"}),
json.dumps({"type": "response.output_item.added", "item": function_call}),
json.dumps({"type": "response.output_item.done", "item": function_call}),
json.dumps({"type": "response.completed", "response": {"id": "r-1"}}),
]
upstream = _FakeUpstream(upstream_events)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(frames=[_first_frame()])
client_ws.headers["x-headroom-user-id"] = "user-1"
handler = _DummyOpenAIHandler()
handler.memory_handler = _MemoryWsHandler()
executed: list[tuple[str, dict, str, str]] = []
async def _execute_memory_tool(name, args, user_id, provider):
executed.append((name, args, user_id, provider))
return '{"memories": []}'
handler.memory_handler._execute_memory_tool = _execute_memory_tool
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
assert executed == [("memory_search", {}, "user-1", "openai")]
@pytest.mark.asyncio
async def test_ws_memory_tools_preserve_explicit_store_false_while_injecting():
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r-1"}}),
json.dumps({"type": "response.completed", "response": {"id": "r-1"}}),
]
upstream = _FakeUpstream(upstream_events)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(
frames=[
json.dumps(
{
"type": "response.create",
"response": {
"model": "gpt-5.4",
"input": "use stateless memory",
"store": False,
},
}
)
],
hold_after_initial=True,
)
client_ws.headers["x-headroom-user-id"] = "user-1"
handler = _DummyOpenAIHandler()
handler.memory_handler = _MemoryWsHandler()
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
assert len(upstream.sent) == 1
initial = json.loads(upstream.sent[0])["response"]
assert initial["store"] is False
assert [tool["name"] for tool in initial["tools"]] == ["memory_search"]
assert initial["include"] == ["reasoning.encrypted_content"]
@pytest.mark.asyncio
async def test_ws_memory_continuation_continues_pre_stream_and_passes_late_call():
function_call_one = {
"type": "function_call",
"id": "fc-1",
"call_id": "call-1",
"name": "memory_search",
"arguments": "{}",
}
function_call_two = {
"type": "function_call",
"id": "fc-2",
"call_id": "call-2",
"name": "memory_search",
"arguments": "{}",
}
reasoning_without_encryption = {
"type": "reasoning",
"id": "reasoning-1",
"summary": [],
}
reasoning_with_encryption = {
"type": "reasoning",
"id": "reasoning-2",
"summary": [],
"encrypted_content": "encrypted-2",
}
message_item = {
"type": "message",
"id": "message-2",
"role": "assistant",
"content": [{"type": "output_text", "text": "searching"}],
}
upstream_events = [
json.dumps({"type": "response.created", "response": {"id": "r-1"}}),
json.dumps({"type": "response.output_item.added", "item": reasoning_without_encryption}),
json.dumps({"type": "response.output_item.done", "item": reasoning_without_encryption}),
json.dumps({"type": "response.output_item.added", "item": function_call_one}),
json.dumps({"type": "response.output_item.done", "item": function_call_one}),
json.dumps({"type": "response.completed", "response": {"id": "r-1"}}),
json.dumps({"type": "response.created", "response": {"id": "r-2"}}),
json.dumps({"type": "response.output_item.added", "item": reasoning_with_encryption}),
json.dumps({"type": "response.output_item.done", "item": reasoning_with_encryption}),
json.dumps({"type": "response.output_item.added", "item": message_item}),
json.dumps({"type": "response.output_item.done", "item": message_item}),
json.dumps({"type": "response.output_item.added", "item": function_call_two}),
json.dumps({"type": "response.output_item.done", "item": function_call_two}),
json.dumps({"type": "response.completed", "response": {"id": "r-2"}}),
]
upstream = _FakeUpstream(upstream_events)
fake_ws_mod = _make_fake_websockets_module(upstream)
client_ws = _FakeWebSocket(
frames=[
json.dumps(
{
"type": "response.create",
"response": {
"model": "gpt-5.4",
"input": "remember this",
"client_metadata": {
"ws_request_header_x_openai_internal_codex_responses_lite": "true",
"keep": "yes",
},
},
}
)
],
hold_after_initial=True,
)
client_ws.headers["x-headroom-user-id"] = "user-1"
handler = _DummyOpenAIHandler()
handler.memory_handler = _MemoryWsHandler()
executed: list[tuple[str, dict, str, str]] = []
async def _execute_memory_tool(name, args, user_id, provider):
executed.append((name, args, user_id, provider))
return '{"memories": []}'
handler.memory_handler._execute_memory_tool = _execute_memory_tool
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
await handler.handle_openai_responses_ws(client_ws)
assert len(upstream.sent) == 2
first_continuation = json.loads(upstream.sent[1])["response"]["input"]
assert reasoning_without_encryption not in first_continuation
assert function_call_one in first_continuation
assert {
"type": "function_call_output",
"call_id": "call-1",
"output": '{"memories": []}',
} in first_continuation
assert json.loads(upstream.sent[1])["response"]["client_metadata"] == {"keep": "yes"}
second_response = [json.loads(frame) for frame in client_ws.sent_text]
assert [event["type"] for event in second_response] == [
"response.created",
"response.output_item.added",
"response.output_item.done",
"response.output_item.added",
"response.output_item.done",
"response.output_item.added",
"response.output_item.done",
"response.completed",
]
assert second_response[0]["response"]["id"] == "r-2"
assert second_response[2]["item"] == reasoning_with_encryption
assert second_response[3]["item"] == message_item
assert second_response[4]["item"] == message_item
assert second_response[5]["item"] == function_call_two
assert second_response[6]["item"] == function_call_two
assert second_response[7]["response"]["id"] == "r-2"
assert executed == [("memory_search", {}, "user-1", "openai")]
tests/test_openai_responses_context_compaction.py
-25
View File
@@ -7,7 +7,6 @@ from headroom.proxy.handlers.openai import (
OpenAIHandlerMixin,
_compact_openai_responses_tools,
_openai_responses_context_budget,
_responses_request_allows_memory_tool_continuation,
)
from headroom.transforms.content_router import (
CompressionStrategy,
@@ -438,30 +437,6 @@ def test_content_router_retries_kompress_when_structured_strategy_noops(monkeypa
assert strategy_chain == ["smart_crusher", "kompress"]
def test_responses_memory_tools_skip_explicit_store_false() -> None:
"""Regression: explicit store=false must block Responses memory-tool injection."""
payload = {"model": "gpt-5.5", "input": "remember this", "store": False}
assert _responses_request_allows_memory_tool_continuation(payload) is False
assert payload["store"] is False
def test_responses_memory_tools_allow_default_and_stored_requests() -> None:
no_memory_payload = {"model": "gpt-5.5", "input": "plain", "store": False}
already_stored_payload = {"model": "gpt-5.5", "input": "plain", "store": True}
default_store_payload = {"model": "gpt-5.5", "input": "plain"}
assert _responses_request_allows_memory_tool_continuation(no_memory_payload) is False
assert no_memory_payload["store"] is False
assert _responses_request_allows_memory_tool_continuation(already_stored_payload) is True
assert already_stored_payload["store"] is True
assert _responses_request_allows_memory_tool_continuation(default_store_payload) is True
assert "store" not in default_store_payload
def test_responses_turn_hook_message_fold_is_applied_and_counted() -> None:
"""On the Responses path a turn hook may fold the `input` items (in place),
not just tools. The fold must be written back to the outbound payload AND its
tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py
+142 -24
View File
@@ -52,10 +52,6 @@ def _message_response(content: list[dict], *, stop_reason: str = "end_turn") ->
}
def _is_client_visible_sse(body: bytes) -> bool:
return b"event:" in body or b"data:" in body
class _ContinuationClient:
def __init__(self, response_json: dict) -> None:
self.response_json = response_json
@@ -355,7 +351,14 @@ def test_unresolved_ccr_only_streams_through_as_200() -> None:
@pytest.mark.asyncio
async def test_buffered_ccr_emits_keepalive_before_delayed_upstream() -> None:
async def test_buffered_ccr_withholds_output_until_delayed_upstream_resolves() -> None:
"""Nothing is sent — no status, no body — until the buffered result exists.
The response used to commit ``200 text/event-stream`` on a 1s keepalive
timer, which made every later failure unreportable: the client saw a 200
with no ``message_start`` and the real status was gone. See
``test_buffered_ccr_preserves_late_failure_status_and_headers``.
"""
config = _make_config()
final_response = _message_response([{"type": "text", "text": "done"}])
started = asyncio.Event()
@@ -367,8 +370,17 @@ async def test_buffered_ccr_emits_keepalive_before_delayed_upstream() -> None:
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": "wait"}],
}
request_delivered = False
async def receive():
# Mirror a real ASGI server: the body arrives once, then the channel
# stays open because the client is still connected. Returning instantly
# on every call spins `StreamingResponse.listen_for_disconnect`, which
# never yields, so the response body would never be scheduled.
nonlocal request_delivered
if request_delivered:
await asyncio.Event().wait()
request_delivered = True
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
scope = {
@@ -400,24 +412,23 @@ async def test_buffered_ccr_emits_keepalive_before_delayed_upstream() -> None:
await started.wait()
response = await asyncio.wait_for(asyncio.shield(task), 1)
events: list[dict] = []
first_visible_body = asyncio.Event()
async def send(message): # noqa: ANN001
events.append(message)
if message["type"] == "http.response.body" and _is_client_visible_sse(
message["body"]
):
first_visible_body.set()
response_task = asyncio.create_task(response(scope, receive, send))
await asyncio.wait_for(first_visible_body.wait(), 2)
assert not release.is_set()
# Longer than the deleted 1.0s keepalive deadline: an unresolved
# upstream must still have produced no ASGI message at all.
await asyncio.sleep(1.1)
assert events == []
release.set()
await response_task
bodies = [event["body"] for event in events if event["type"] == "http.response.body"]
assert bodies[0] == b'event: ping\ndata: {"type":"ping"}\n\n'
assert b"done" in b"".join(bodies)
start = next(event for event in events if event["type"] == "http.response.start")
assert start["status"] == 200
bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body")
assert b"event: ping" not in bodies
assert b"done" in bodies
@pytest.mark.asyncio
@@ -480,7 +491,115 @@ async def test_buffered_ccr_preserves_early_failure_status_and_headers() -> None
@pytest.mark.asyncio
async def test_buffered_ccr_late_failure_emits_sanitized_error_event() -> None:
async def test_buffered_ccr_preserves_late_failure_status_and_headers() -> None:
"""The reported failure: a non-200 landing after the old keepalive deadline.
Headroom had already committed ``200 text/event-stream`` by then, so the 429
reached Claude Code as a 200 whose body carried no ``message_start`` — shown
as "API returned an empty or malformed response (HTTP 200) — check for a
proxy or gateway intercepting the request" — and ``retry-after`` was dropped,
so the client never backed off.
"""
config = _make_config()
started = asyncio.Event()
release = asyncio.Event()
body = {
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": "fail late"}],
}
async def receive():
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
scope = {
"type": "http",
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"path": "/v1/messages",
"raw_path": b"/v1/messages",
"query_string": b"",
"headers": [(b"x-api-key", b"test-key"), (b"anthropic-version", b"2023-06-01")],
"server": ("testserver", 80),
"client": ("testclient", 123),
"root_path": "",
}
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app):
proxy = app.state.proxy
async def delayed_failure(*args, **kwargs): # noqa: ANN002, ANN003
started.set()
await release.wait()
return httpx.Response(
429,
headers={"retry-after": "7"},
json={"error": {"message": "slow down"}},
)
proxy._retry_request = delayed_failure
task = asyncio.create_task(proxy.handle_anthropic_messages(Request(scope, receive)))
await started.wait()
response = await asyncio.wait_for(asyncio.shield(task), 1)
events: list[dict] = []
async def send(message): # noqa: ANN001
events.append(message)
response_task = asyncio.create_task(response(scope, receive, send))
# Past the deleted 1.0s keepalive deadline before the upstream fails.
await asyncio.sleep(1.1)
assert events == []
release.set()
await response_task
start = next(event for event in events if event["type"] == "http.response.start")
assert start["status"] == 429
assert dict(start["headers"])[b"retry-after"] == b"7"
assert b"slow down" in b"".join(
event["body"] for event in events if event["type"] == "http.response.body"
)
def test_buffered_ccr_rejects_malformed_success_as_502() -> None:
"""A non-SSE, non-JSON 200 is an upstream protocol error, not success."""
config = _make_config()
with patch("headroom.proxy.server.AnyLLMBackend"):
app = create_app(config)
with TestClient(app) as client:
proxy = app.state.proxy
proxy._retry_request = AsyncMock(
return_value=httpx.Response(
200,
content=b"<html>gateway timeout</html>",
headers={"content-type": "text/html"},
)
)
response = client.post(
"/v1/messages",
headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
json={
"model": "claude-sonnet-4-6",
"max_tokens": 64,
"stream": True,
"tools": [create_ccr_tool_definition("anthropic")],
"messages": [{"role": "user", "content": "fail safely"}],
},
)
assert response.status_code == 502
assert response.json()["error"]["type"] == "upstream_protocol_error"
assert b"gateway timeout" not in response.content
@pytest.mark.asyncio
async def test_buffered_ccr_late_failure_returns_sanitized_json_error() -> None:
"""A slow crash gets the same 502 the fast one does, not a downgraded 200."""
config = _make_config()
started = asyncio.Event()
release = asyncio.Event()
@@ -533,24 +652,23 @@ async def test_buffered_ccr_late_failure_emits_sanitized_error_event() -> None:
await started.wait()
response = await asyncio.wait_for(asyncio.shield(task), 1)
events: list[dict] = []
first_body = asyncio.Event()
async def send(message): # noqa: ANN001
events.append(message)
if message["type"] == "http.response.body" and message["body"]:
first_body.set()
response_task = asyncio.create_task(response(scope, receive, send))
await asyncio.wait_for(first_body.wait(), 2)
await asyncio.sleep(0)
assert events == []
release.set()
await response_task
record_failed.assert_awaited_once_with(provider="anthropic")
proxy_logger.removeHandler(log_handler)
bodies = [event["body"] for event in events if event["type"] == "http.response.body"]
assert bodies[0] == b'event: ping\ndata: {"type":"ping"}\n\n'
assert b"An error occurred while processing the request." in bodies[-1]
assert b"boom" not in bodies[-1]
start = next(event for event in events if event["type"] == "http.response.start")
assert start["status"] == 502
bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body")
assert b"An error occurred while processing your request." in bodies
assert b"boom" not in bodies
assert events[-1]["more_body"] is False
assert any(
record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage()
tests/test_proxy/test_openai_responses_ccr.py
+56 -17
View File
@@ -266,7 +266,8 @@ def test_streaming_request_without_retrieve_tool_uses_normal_stream_path():
@pytest.mark.asyncio
async def test_buffered_responses_ccr_emits_keepalive_before_delayed_upstream():
async def test_buffered_responses_ccr_withholds_output_until_upstream_resolves():
"""Nothing is sent — no status, no body — until the buffered result exists."""
app = _make_app()
body = {
"model": "gpt-5-codex",
@@ -276,8 +277,17 @@ async def test_buffered_responses_ccr_emits_keepalive_before_delayed_upstream():
}
started = asyncio.Event()
release = asyncio.Event()
request_delivered = False
async def receive():
# Mirror a real ASGI server: the body arrives once, then the channel
# stays open because the client is still connected. Returning instantly
# on every call spins `StreamingResponse.listen_for_disconnect`, which
# never yields, so the response body would never be scheduled.
nonlocal request_delivered
if request_delivered:
await asyncio.Event().wait()
request_delivered = True
return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False}
scope = {
@@ -307,22 +317,23 @@ async def test_buffered_responses_ccr_emits_keepalive_before_delayed_upstream():
await started.wait()
response = await asyncio.wait_for(asyncio.shield(task), 1)
events: list[dict] = []
first_body = asyncio.Event()
async def send(message): # noqa: ANN001
events.append(message)
if message["type"] == "http.response.body" and message["body"]:
first_body.set()
response_task = asyncio.create_task(response(scope, receive, send))
await asyncio.wait_for(first_body.wait(), 2)
assert not release.is_set()
# Longer than the deleted 1.0s keepalive deadline: an unresolved upstream
# must still have produced no ASGI message at all.
await asyncio.sleep(1.1)
assert events == []
release.set()
await response_task
bodies = [event["body"] for event in events if event["type"] == "http.response.body"]
assert bodies[0] == b'event: ping\ndata: {"type":"ping"}\n\n'
assert b"Resolved!" in b"".join(bodies)
start = next(event for event in events if event["type"] == "http.response.start")
assert start["status"] == 200
bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body")
assert b"event: ping" not in bodies
assert b"Resolved!" in bodies
@pytest.mark.asyncio
@@ -382,6 +393,35 @@ async def test_buffered_responses_ccr_preserves_early_failure_status_and_headers
)
def test_buffered_responses_ccr_rejects_malformed_success_as_502():
"""A malformed upstream 200 must not become a successful streamed turn."""
app = _make_app()
with TestClient(app) as client:
server = app.state.proxy
server._retry_request = AsyncMock(
return_value=httpx.Response(
200,
content=b"<html>gateway timeout</html>",
headers={"content-type": "text/html"},
request=httpx.Request("POST", "https://api.openai.com/v1/responses"),
)
)
response = client.post(
"/v1/responses",
headers={"authorization": "Bearer sk-test"},
json={
"model": "gpt-5-codex",
"input": "fail safely",
"stream": True,
"tools": [_RETRIEVE_TOOL],
},
)
assert response.status_code == 502
assert response.json()["error"]["type"] == "upstream_protocol_error", response.text
assert b"gateway timeout" not in response.content
@pytest.mark.asyncio
async def test_buffered_responses_ccr_late_failure_emits_sanitized_error_event():
app = _make_app()
@@ -431,24 +471,23 @@ async def test_buffered_responses_ccr_late_failure_emits_sanitized_error_event()
await started.wait()
response = await asyncio.wait_for(asyncio.shield(task), 1)
events: list[dict] = []
first_body = asyncio.Event()
async def send(message): # noqa: ANN001
events.append(message)
if message["type"] == "http.response.body" and message["body"]:
first_body.set()
response_task = asyncio.create_task(response(scope, receive, send))
await asyncio.wait_for(first_body.wait(), 2)
await asyncio.sleep(0)
assert events == []
release.set()
await response_task
record_failed.assert_awaited_once_with(provider="openai")
proxy_logger.removeHandler(log_handler)
bodies = [event["body"] for event in events if event["type"] == "http.response.body"]
assert bodies[0] == b'event: ping\ndata: {"type":"ping"}\n\n'
assert b"An error occurred while processing the request." in bodies[-1]
assert b"boom" not in bodies[-1]
start = next(event for event in events if event["type"] == "http.response.start")
assert start["status"] == 502
bodies = b"".join(event["body"] for event in events if event["type"] == "http.response.body")
assert b"An error occurred while processing your request." in bodies
assert b"boom" not in bodies
assert events[-1]["more_body"] is False
assert any(
record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage()