mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): keep buffered CCR streams alive (#2479)
## Description Buffered CCR streaming currently waits for the full upstream response before sending any bytes back to the client. On the Anthropic path this shows up as `API Error: Stream idle timeout - no chunks received`, and the same buffer-then-synthesize mechanism still exists on the `/v1/responses` CCR path. This adds a narrow buffered-stream heartbeat layer so the client sees early stream activity while Headroom preserves the existing server-side retrieval round trip and final synthesized provider events. Closes #2465 ## 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 - open buffered CCR streams early and emit client-visible `event: ping` heartbeats while the buffered upstream call is still in flight - preserve the existing terminal Anthropic and Responses synthesis helpers instead of replacing their event-building logic - preserve early non-streaming failure semantics before the first heartbeat, including normal 429 passthrough and normal JSON 502 failures - log late buffered-task exceptions server-side and record one failed provider metric on that post-keepalive branch, while keeping the client-facing SSE error sanitized - add focused delayed-upstream regression coverage for both buffered provider paths, their early-failure branches, and their late-failure branches ## Testing - [x] Unit tests pass (`uv run pytest tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py tests/test_proxy/test_openai_responses_ccr.py -q`) - [x] Linting passes (`uv run ruff check headroom/proxy/handlers/anthropic.py headroom/proxy/handlers/openai.py tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py tests/test_proxy/test_openai_responses_ccr.py`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality when applicable - [ ] Manual testing performed ### Test Output ```text uv run pytest tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py tests/test_proxy/test_openai_responses_ccr.py -q ======================= 17 passed, 1 warning in 42.47s ======================== uv run ruff check headroom/proxy/handlers/anthropic.py headroom/proxy/handlers/openai.py tests/test_proxy/test_anthropic_streaming_ccr_retrieve.py tests/test_proxy/test_openai_responses_ccr.py All checks passed! ``` ## Real Behavior Proof - Environment: Windows, Python via `uv`, proxy handler tests with gated buffered upstream fixtures - Exact command / steps: run the focused Anthropic and Responses CCR suites above; delayed-upstream tests consume the first client-visible SSE event before releasing the upstream, then consume the synthesized final events - Observed result: both buffered paths emitted `event: ping` before upstream release; pre-keepalive 429 responses preserved their real status and headers, pre-keepalive exceptions returned the normal JSON 502 shape, late transport failures recorded one failed provider metric and one server error log before emitting one sanitized SSE error, Anthropic preserved `done`, and Responses preserved `Resolved!` - Not tested: live slow upstream run with Claude Code or a real Responses client ## 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 - [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 - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes `CHANGELOG.md` stays untouched because Headroom generates release notes from conventional commits. The PR should only claim the local buffered-stream contract and focused regression coverage; live client proof remains an owner check on a slow real upstream.
This commit is contained in:
Rod Boev
GitHub
parent
43a7b578a1
commit
a2e42fb877
4 changed files with 1808 additions and 1035 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -4911,379 +4911,506 @@ class OpenAIHandlerMixin:
|
|||
waste_signals=waste_signals_dict,
|
||||
)
|
||||
else:
|
||||
headers = await apply_copilot_api_auth(headers, url=url)
|
||||
response = await self._retry_request(
|
||||
"POST",
|
||||
url,
|
||||
headers,
|
||||
body,
|
||||
original_body_bytes=original_body_bytes,
|
||||
body_mutated=body_mutation_tracker.mutated,
|
||||
mutation_reasons=body_mutation_tracker.reasons,
|
||||
request_id=request_id,
|
||||
forwarder_name="openai_responses",
|
||||
path_for_log=url,
|
||||
)
|
||||
_response_body_for_debug: Any = None
|
||||
_response_raw_for_debug: str | None = None
|
||||
try:
|
||||
_response_body_for_debug = response.json()
|
||||
except Exception:
|
||||
try:
|
||||
_response_raw_for_debug = response.text[:200_000]
|
||||
except Exception:
|
||||
_response_raw_for_debug = None
|
||||
capture_codex_wire_debug(
|
||||
"http_upstream_response",
|
||||
request_id=request_id,
|
||||
transport="http",
|
||||
direction="upstream_to_headroom",
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=dict(response.headers),
|
||||
body=_response_body_for_debug,
|
||||
raw_text=_response_raw_for_debug,
|
||||
status_code=response.status_code,
|
||||
metadata={"stream": stream, "auth_mode": auth_mode.value},
|
||||
)
|
||||
total_latency = (time.time() - start_time) * 1000
|
||||
|
||||
total_input_tokens = original_tokens # fallback
|
||||
output_tokens = 0
|
||||
cache_read_tokens = 0
|
||||
try:
|
||||
resp_json = response.json()
|
||||
usage = resp_json.get("usage", {})
|
||||
|
||||
def _usage_int(value: Any, default: int = 0) -> int:
|
||||
try:
|
||||
return max(int(value), 0)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
total_input_tokens = _usage_int(
|
||||
usage.get("input_tokens"),
|
||||
original_tokens,
|
||||
)
|
||||
output_tokens = _usage_int(usage.get("output_tokens"))
|
||||
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:
|
||||
logger.debug(
|
||||
f"[{request_id}] Failed to extract cached tokens from OpenAI passthrough response: {e}"
|
||||
)
|
||||
|
||||
# CCR Response Handling: intercept headroom_retrieve tool
|
||||
# calls server-side so a Responses API function_call the
|
||||
# downstream caller can't resolve (e.g. Strands, or a
|
||||
# buffered-stream request) never reaches the client. Mirrors
|
||||
# the chat-completions backend-path block (handle_openai_chat
|
||||
# ~2775-2848), adapted for the Responses API's flat
|
||||
# function_call / output[] shape instead of Messages API
|
||||
# tool_calls. Runs before memory tool handling below so a
|
||||
# retrieve call never gets treated as an unresolved tool_call
|
||||
# by the memory-tool branch.
|
||||
if (
|
||||
_ccr_response_handler
|
||||
and resp_json
|
||||
and response.status_code == 200
|
||||
and _ccr_response_handler.has_ccr_tool_calls(resp_json, "openai_responses")
|
||||
):
|
||||
logger.info(
|
||||
f"[{request_id}] CCR: Detected retrieval tool call (responses), handling..."
|
||||
)
|
||||
|
||||
async def api_call_fn(
|
||||
items: list[dict[str, Any]],
|
||||
tls: list[dict[str, Any]] | None,
|
||||
) -> dict[str, Any]:
|
||||
continuation_body = {**body, "input": items}
|
||||
if tls is not None:
|
||||
continuation_body["tools"] = tls
|
||||
# Fresh stateless continuation: resend the full
|
||||
# item history rather than chaining through
|
||||
# previous_response_id, matching how
|
||||
# CCRResponseHandler accumulates `current_messages`
|
||||
# for every other provider. `body["stream"]` is
|
||||
# left as-is: for a buffered_stream_ccr request it
|
||||
# was already forced False above, and continuations
|
||||
# must stay non-streaming so this handler (not
|
||||
# `_stream_response`) can parse the JSON reply.
|
||||
continuation_body.pop("previous_response_id", None)
|
||||
continuation_body["stream"] = False
|
||||
|
||||
continuation_headers = {
|
||||
k: v
|
||||
for k, v in headers.items()
|
||||
if k.lower()
|
||||
not in (
|
||||
"content-encoding",
|
||||
"transfer-encoding",
|
||||
"accept-encoding",
|
||||
"content-length",
|
||||
)
|
||||
}
|
||||
logger.info(
|
||||
f"[{request_id}] CCR: Issuing Responses continuation "
|
||||
f"({len(items)} input items)"
|
||||
)
|
||||
cont_response = await self._retry_request(
|
||||
"POST",
|
||||
url,
|
||||
continuation_headers,
|
||||
continuation_body,
|
||||
request_id=request_id,
|
||||
forwarder_name="openai_responses_ccr_continuation",
|
||||
path_for_log=url,
|
||||
)
|
||||
return cont_response.json()
|
||||
|
||||
try:
|
||||
final_resp_json = await _ccr_response_handler.handle_response(
|
||||
resp_json,
|
||||
_responses_input_to_items(body.get("input")),
|
||||
body.get("tools"),
|
||||
api_call_fn,
|
||||
provider="openai_responses",
|
||||
)
|
||||
resp_json = final_resp_json
|
||||
# Remove encoding headers since content is now
|
||||
# uncompressed JSON we synthesized.
|
||||
ccr_response_headers = {
|
||||
k: v
|
||||
for k, v in response.headers.items()
|
||||
if k.lower() not in ("content-encoding", "content-length")
|
||||
}
|
||||
response = httpx.Response(
|
||||
status_code=200,
|
||||
content=json.dumps(final_resp_json).encode(),
|
||||
headers=ccr_response_headers,
|
||||
)
|
||||
logger.info(
|
||||
f"[{request_id}] CCR: Retrieval handled successfully (responses)"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[{request_id}] CCR: Response handling failed (responses): {e}"
|
||||
)
|
||||
# NO SILENT FALLBACK: re-raise so the client sees a
|
||||
# clear failure instead of an unresolved tool_call
|
||||
# it can't act on. Matches the OpenAI backend-path
|
||||
# block in handle_openai_chat; see
|
||||
# feedback_no_silent_fallbacks.
|
||||
raise
|
||||
|
||||
# Memory: handle memory tool calls in Responses API response
|
||||
if (
|
||||
self.memory_handler
|
||||
and memory_user_id
|
||||
and responses_memory_tools_allowed
|
||||
and resp_json
|
||||
and response.status_code == 200
|
||||
and self.memory_handler.has_memory_tool_calls(resp_json, "openai")
|
||||
):
|
||||
try:
|
||||
# Extract function_call items from output
|
||||
from headroom.proxy.memory_handler import MEMORY_TOOL_NAMES
|
||||
|
||||
output_items = resp_json.get("output", [])
|
||||
memory_fc_items = [
|
||||
item
|
||||
for item in output_items
|
||||
if isinstance(item, dict)
|
||||
and item.get("type") == "function_call"
|
||||
and item.get("name") in MEMORY_TOOL_NAMES
|
||||
]
|
||||
|
||||
# Execute memory tool calls
|
||||
tool_outputs: list[dict[str, Any]] = []
|
||||
for fc in memory_fc_items:
|
||||
call_id = fc.get("call_id", fc.get("id", ""))
|
||||
name = fc.get("name", "")
|
||||
args_str = fc.get("arguments", "{}")
|
||||
try:
|
||||
args = json.loads(args_str)
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
|
||||
await self.memory_handler._ensure_initialized()
|
||||
if self.memory_handler._backend:
|
||||
result = await self.memory_handler._execute_memory_tool(
|
||||
name, args, memory_user_id, "openai"
|
||||
)
|
||||
else:
|
||||
result = json.dumps({"error": "Memory backend not initialized"})
|
||||
|
||||
tool_outputs.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": call_id,
|
||||
"output": result,
|
||||
}
|
||||
)
|
||||
|
||||
if tool_outputs:
|
||||
# Make continuation request with tool results
|
||||
response_id = resp_json.get("id")
|
||||
continuation_body = {
|
||||
"model": model,
|
||||
"input": 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
|
||||
|
||||
cont_response = await self._retry_request(
|
||||
"POST", url, headers, continuation_body
|
||||
)
|
||||
resp_json = cont_response.json()
|
||||
response = cont_response
|
||||
logger.info(
|
||||
f"[{request_id}] Memory: Handled {len(tool_outputs)} "
|
||||
f"tool call(s) with continuation for user {memory_user_id} (responses)"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"[{request_id}] Memory tool handling failed (responses): {e}"
|
||||
)
|
||||
|
||||
if self.cost_tracker:
|
||||
cache_write_tokens = _infer_openai_cache_write_tokens(
|
||||
total_input_tokens,
|
||||
cache_read_tokens,
|
||||
)
|
||||
uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens)
|
||||
# (record_tokens clamps negative savings to 0 universally.)
|
||||
self.cost_tracker.record_tokens(
|
||||
model,
|
||||
tokens_saved,
|
||||
total_input_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
uncached_tokens=uncached_input_tokens,
|
||||
)
|
||||
else:
|
||||
cache_write_tokens = _infer_openai_cache_write_tokens(
|
||||
total_input_tokens,
|
||||
cache_read_tokens,
|
||||
)
|
||||
uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens)
|
||||
|
||||
effective_optimized_tokens = (
|
||||
total_input_tokens if total_input_tokens > 0 else optimized_tokens
|
||||
)
|
||||
effective_original_tokens = max(
|
||||
original_tokens,
|
||||
effective_optimized_tokens + tokens_saved,
|
||||
)
|
||||
|
||||
_resp_log_tags = {
|
||||
**(tags or {}),
|
||||
"auth_mode": auth_mode.value if auth_mode else "payg",
|
||||
"endpoint": "responses_http",
|
||||
}
|
||||
|
||||
# OpenAI Responses HTTP (non-WS, non-streaming). Codex
|
||||
# uses this path when configured for HTTP transport.
|
||||
# Pre-refactor `cache_hit` was hardcoded False on
|
||||
# RequestLog even when cache_read>0 — funnel derives
|
||||
# it correctly.
|
||||
from headroom.proxy.helpers import compute_turn_id
|
||||
|
||||
await self._record_request_outcome(
|
||||
RequestOutcome(
|
||||
async def _buffered_ccr_operation():
|
||||
nonlocal headers
|
||||
headers = await apply_copilot_api_auth(headers, url=url)
|
||||
response = await self._retry_request(
|
||||
"POST",
|
||||
url,
|
||||
headers,
|
||||
body,
|
||||
original_body_bytes=original_body_bytes,
|
||||
body_mutated=body_mutation_tracker.mutated,
|
||||
mutation_reasons=body_mutation_tracker.reasons,
|
||||
request_id=request_id,
|
||||
provider="openai",
|
||||
model=model,
|
||||
status_code=response.status_code,
|
||||
original_tokens=effective_original_tokens,
|
||||
optimized_tokens=effective_optimized_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
attempted_input_tokens=attempted_input_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
uncached_input_tokens=uncached_input_tokens,
|
||||
total_latency_ms=total_latency,
|
||||
overhead_ms=optimization_latency,
|
||||
transforms_applied=tuple(transforms_applied),
|
||||
waste_signals=waste_signals_dict,
|
||||
num_messages=len(messages) if isinstance(messages, list) else 0,
|
||||
tags=_resp_log_tags,
|
||||
turn_id=compute_turn_id(model, body.get("instructions"), messages),
|
||||
request_messages=messages
|
||||
if getattr(self.config, "log_full_messages", False)
|
||||
else None,
|
||||
client=client,
|
||||
forwarder_name="openai_responses",
|
||||
path_for_log=url,
|
||||
)
|
||||
)
|
||||
_response_body_for_debug: Any = None
|
||||
_response_raw_for_debug: str | None = None
|
||||
try:
|
||||
_response_body_for_debug = response.json()
|
||||
except Exception:
|
||||
try:
|
||||
_response_raw_for_debug = response.text[:200_000]
|
||||
except Exception:
|
||||
_response_raw_for_debug = None
|
||||
capture_codex_wire_debug(
|
||||
"http_upstream_response",
|
||||
request_id=request_id,
|
||||
transport="http",
|
||||
direction="upstream_to_headroom",
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=dict(response.headers),
|
||||
body=_response_body_for_debug,
|
||||
raw_text=_response_raw_for_debug,
|
||||
status_code=response.status_code,
|
||||
metadata={"stream": stream, "auth_mode": auth_mode.value},
|
||||
)
|
||||
total_latency = (time.time() - start_time) * 1000
|
||||
|
||||
logger.info(f"[{request_id}] /v1/responses {model}: {total_input_tokens:,} tokens")
|
||||
total_input_tokens = original_tokens # fallback
|
||||
output_tokens = 0
|
||||
cache_read_tokens = 0
|
||||
try:
|
||||
resp_json = response.json()
|
||||
usage = resp_json.get("usage", {})
|
||||
|
||||
# Capture Codex rate-limit window data from response headers
|
||||
from headroom.subscription.codex_rate_limits import (
|
||||
get_codex_rate_limit_state,
|
||||
)
|
||||
def _usage_int(value: Any, default: int = 0) -> int:
|
||||
try:
|
||||
return max(int(value), 0)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
get_codex_rate_limit_state().update_from_headers(dict(response.headers))
|
||||
|
||||
# Remove compression headers
|
||||
response_headers = _sanitize_forwarded_response_headers(response.headers)
|
||||
|
||||
if buffered_stream_ccr and response.status_code == 200 and resp_json:
|
||||
sse_headers = {
|
||||
k: v
|
||||
for k, v in response_headers.items()
|
||||
if k.lower() not in ("content-length", "content-type")
|
||||
}
|
||||
if _ccr_response_handler and _ccr_response_handler.has_ccr_tool_calls(
|
||||
resp_json, "openai_responses"
|
||||
):
|
||||
# Handling above didn't fully resolve the retrieve
|
||||
# call (e.g. max rounds hit, or it was mixed with a
|
||||
# non-CCR tool call). Fail closed rather than stream
|
||||
# a response the client can't act on — matches the
|
||||
# Anthropic buffered path's residual-CCR guard.
|
||||
logger.warning(
|
||||
f"[{request_id}] CCR: Buffered streaming Responses "
|
||||
"reply still contains headroom_retrieve after "
|
||||
"handling; failing closed"
|
||||
total_input_tokens = _usage_int(
|
||||
usage.get("input_tokens"),
|
||||
original_tokens,
|
||||
)
|
||||
output_tokens = _usage_int(usage.get("output_tokens"))
|
||||
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:
|
||||
logger.debug(
|
||||
f"[{request_id}] Failed to extract cached tokens from OpenAI passthrough response: {e}"
|
||||
)
|
||||
|
||||
async def _residual_ccr_error_sse():
|
||||
error_event = {
|
||||
"type": "error",
|
||||
"error": {
|
||||
"message": "Unable to safely complete streamed CCR retrieval.",
|
||||
},
|
||||
# CCR Response Handling: intercept headroom_retrieve tool
|
||||
# calls server-side so a Responses API function_call the
|
||||
# downstream caller can't resolve (e.g. Strands, or a
|
||||
# buffered-stream request) never reaches the client. Mirrors
|
||||
# the chat-completions backend-path block (handle_openai_chat
|
||||
# ~2775-2848), adapted for the Responses API's flat
|
||||
# function_call / output[] shape instead of Messages API
|
||||
# tool_calls. Runs before memory tool handling below so a
|
||||
# retrieve call never gets treated as an unresolved tool_call
|
||||
# by the memory-tool branch.
|
||||
if (
|
||||
_ccr_response_handler
|
||||
and resp_json
|
||||
and response.status_code == 200
|
||||
and _ccr_response_handler.has_ccr_tool_calls(resp_json, "openai_responses")
|
||||
):
|
||||
logger.info(
|
||||
f"[{request_id}] CCR: Detected retrieval tool call (responses), handling..."
|
||||
)
|
||||
|
||||
async def api_call_fn(
|
||||
items: list[dict[str, Any]],
|
||||
tls: list[dict[str, Any]] | None,
|
||||
) -> dict[str, Any]:
|
||||
continuation_body = {**body, "input": items}
|
||||
if tls is not None:
|
||||
continuation_body["tools"] = tls
|
||||
# Fresh stateless continuation: resend the full
|
||||
# item history rather than chaining through
|
||||
# previous_response_id, matching how
|
||||
# CCRResponseHandler accumulates `current_messages`
|
||||
# for every other provider. `body["stream"]` is
|
||||
# left as-is: for a buffered_stream_ccr request it
|
||||
# was already forced False above, and continuations
|
||||
# must stay non-streaming so this handler (not
|
||||
# `_stream_response`) can parse the JSON reply.
|
||||
continuation_body.pop("previous_response_id", None)
|
||||
continuation_body["stream"] = False
|
||||
|
||||
continuation_headers = {
|
||||
k: v
|
||||
for k, v in headers.items()
|
||||
if k.lower()
|
||||
not in (
|
||||
"content-encoding",
|
||||
"transfer-encoding",
|
||||
"accept-encoding",
|
||||
"content-length",
|
||||
)
|
||||
}
|
||||
yield f"event: error\ndata: {json.dumps(error_event)}\n\n".encode()
|
||||
logger.info(
|
||||
f"[{request_id}] CCR: Issuing Responses continuation "
|
||||
f"({len(items)} input items)"
|
||||
)
|
||||
cont_response = await self._retry_request(
|
||||
"POST",
|
||||
url,
|
||||
continuation_headers,
|
||||
continuation_body,
|
||||
request_id=request_id,
|
||||
forwarder_name="openai_responses_ccr_continuation",
|
||||
path_for_log=url,
|
||||
)
|
||||
return cont_response.json()
|
||||
|
||||
try:
|
||||
final_resp_json = await _ccr_response_handler.handle_response(
|
||||
resp_json,
|
||||
_responses_input_to_items(body.get("input")),
|
||||
body.get("tools"),
|
||||
api_call_fn,
|
||||
provider="openai_responses",
|
||||
)
|
||||
resp_json = final_resp_json
|
||||
# Remove encoding headers since content is now
|
||||
# uncompressed JSON we synthesized.
|
||||
ccr_response_headers = {
|
||||
k: v
|
||||
for k, v in response.headers.items()
|
||||
if k.lower() not in ("content-encoding", "content-length")
|
||||
}
|
||||
response = httpx.Response(
|
||||
status_code=200,
|
||||
content=json.dumps(final_resp_json).encode(),
|
||||
headers=ccr_response_headers,
|
||||
)
|
||||
logger.info(
|
||||
f"[{request_id}] CCR: Retrieval handled successfully (responses)"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[{request_id}] CCR: Response handling failed (responses): {e}"
|
||||
)
|
||||
# NO SILENT FALLBACK: re-raise so the client sees a
|
||||
# clear failure instead of an unresolved tool_call
|
||||
# it can't act on. Matches the OpenAI backend-path
|
||||
# block in handle_openai_chat; see
|
||||
# feedback_no_silent_fallbacks.
|
||||
raise
|
||||
|
||||
# Memory: handle memory tool calls in Responses API response
|
||||
if (
|
||||
self.memory_handler
|
||||
and memory_user_id
|
||||
and responses_memory_tools_allowed
|
||||
and resp_json
|
||||
and response.status_code == 200
|
||||
and self.memory_handler.has_memory_tool_calls(resp_json, "openai")
|
||||
):
|
||||
try:
|
||||
# Extract function_call items from output
|
||||
from headroom.proxy.memory_handler import MEMORY_TOOL_NAMES
|
||||
|
||||
output_items = resp_json.get("output", [])
|
||||
memory_fc_items = [
|
||||
item
|
||||
for item in output_items
|
||||
if isinstance(item, dict)
|
||||
and item.get("type") == "function_call"
|
||||
and item.get("name") in MEMORY_TOOL_NAMES
|
||||
]
|
||||
|
||||
# Execute memory tool calls
|
||||
tool_outputs: list[dict[str, Any]] = []
|
||||
for fc in memory_fc_items:
|
||||
call_id = fc.get("call_id", fc.get("id", ""))
|
||||
name = fc.get("name", "")
|
||||
args_str = fc.get("arguments", "{}")
|
||||
try:
|
||||
args = json.loads(args_str)
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
|
||||
await self.memory_handler._ensure_initialized()
|
||||
if self.memory_handler._backend:
|
||||
result = await self.memory_handler._execute_memory_tool(
|
||||
name, args, memory_user_id, "openai"
|
||||
)
|
||||
else:
|
||||
result = json.dumps({"error": "Memory backend not initialized"})
|
||||
|
||||
tool_outputs.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": call_id,
|
||||
"output": result,
|
||||
}
|
||||
)
|
||||
|
||||
if tool_outputs:
|
||||
# Make continuation request with tool results
|
||||
response_id = resp_json.get("id")
|
||||
continuation_body = {
|
||||
"model": model,
|
||||
"input": 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
|
||||
|
||||
cont_response = await self._retry_request(
|
||||
"POST", url, headers, continuation_body
|
||||
)
|
||||
resp_json = cont_response.json()
|
||||
response = cont_response
|
||||
logger.info(
|
||||
f"[{request_id}] Memory: Handled {len(tool_outputs)} "
|
||||
f"tool call(s) with continuation for user {memory_user_id} (responses)"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"[{request_id}] Memory tool handling failed (responses): {e}"
|
||||
)
|
||||
|
||||
if self.cost_tracker:
|
||||
cache_write_tokens = _infer_openai_cache_write_tokens(
|
||||
total_input_tokens,
|
||||
cache_read_tokens,
|
||||
)
|
||||
uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens)
|
||||
# (record_tokens clamps negative savings to 0 universally.)
|
||||
self.cost_tracker.record_tokens(
|
||||
model,
|
||||
tokens_saved,
|
||||
total_input_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
uncached_tokens=uncached_input_tokens,
|
||||
)
|
||||
else:
|
||||
cache_write_tokens = _infer_openai_cache_write_tokens(
|
||||
total_input_tokens,
|
||||
cache_read_tokens,
|
||||
)
|
||||
uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens)
|
||||
|
||||
effective_optimized_tokens = (
|
||||
total_input_tokens if total_input_tokens > 0 else optimized_tokens
|
||||
)
|
||||
effective_original_tokens = max(
|
||||
original_tokens,
|
||||
effective_optimized_tokens + tokens_saved,
|
||||
)
|
||||
|
||||
_resp_log_tags = {
|
||||
**(tags or {}),
|
||||
"auth_mode": auth_mode.value if auth_mode else "payg",
|
||||
"endpoint": "responses_http",
|
||||
}
|
||||
|
||||
# OpenAI Responses HTTP (non-WS, non-streaming). Codex
|
||||
# uses this path when configured for HTTP transport.
|
||||
# Pre-refactor `cache_hit` was hardcoded False on
|
||||
# RequestLog even when cache_read>0 — funnel derives
|
||||
# it correctly.
|
||||
from headroom.proxy.helpers import compute_turn_id
|
||||
|
||||
await self._record_request_outcome(
|
||||
RequestOutcome(
|
||||
request_id=request_id,
|
||||
provider="openai",
|
||||
model=model,
|
||||
status_code=response.status_code,
|
||||
original_tokens=effective_original_tokens,
|
||||
optimized_tokens=effective_optimized_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
attempted_input_tokens=attempted_input_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
uncached_input_tokens=uncached_input_tokens,
|
||||
total_latency_ms=total_latency,
|
||||
overhead_ms=optimization_latency,
|
||||
transforms_applied=tuple(transforms_applied),
|
||||
waste_signals=waste_signals_dict,
|
||||
num_messages=len(messages) if isinstance(messages, list) else 0,
|
||||
tags=_resp_log_tags,
|
||||
turn_id=compute_turn_id(model, body.get("instructions"), messages),
|
||||
request_messages=messages
|
||||
if getattr(self.config, "log_full_messages", False)
|
||||
else None,
|
||||
client=client,
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[{request_id}] /v1/responses {model}: {total_input_tokens:,} tokens"
|
||||
)
|
||||
|
||||
# Capture Codex rate-limit window data from response headers
|
||||
from headroom.subscription.codex_rate_limits import (
|
||||
get_codex_rate_limit_state,
|
||||
)
|
||||
|
||||
get_codex_rate_limit_state().update_from_headers(dict(response.headers))
|
||||
|
||||
# Remove compression headers
|
||||
response_headers = _sanitize_forwarded_response_headers(response.headers)
|
||||
|
||||
if buffered_stream_ccr and response.status_code == 200 and resp_json:
|
||||
sse_headers = {
|
||||
k: v
|
||||
for k, v in response_headers.items()
|
||||
if k.lower() not in ("content-length", "content-type")
|
||||
}
|
||||
if _ccr_response_handler and _ccr_response_handler.has_ccr_tool_calls(
|
||||
resp_json, "openai_responses"
|
||||
):
|
||||
# Handling above didn't fully resolve the retrieve
|
||||
# call (e.g. max rounds hit, or it was mixed with a
|
||||
# non-CCR tool call). Fail closed rather than stream
|
||||
# a response the client can't act on — matches the
|
||||
# Anthropic buffered path's residual-CCR guard.
|
||||
logger.warning(
|
||||
f"[{request_id}] CCR: Buffered streaming Responses "
|
||||
"reply still contains headroom_retrieve after "
|
||||
"handling; failing closed"
|
||||
)
|
||||
|
||||
async def _residual_ccr_error_sse():
|
||||
error_event = {
|
||||
"type": "error",
|
||||
"error": {
|
||||
"message": "Unable to safely complete streamed CCR retrieval.",
|
||||
},
|
||||
}
|
||||
yield f"event: error\ndata: {json.dumps(error_event)}\n\n".encode()
|
||||
|
||||
return StreamingResponse(
|
||||
_residual_ccr_error_sse(),
|
||||
media_type="text/event-stream",
|
||||
headers=sse_headers,
|
||||
status_code=502,
|
||||
)
|
||||
|
||||
async def _buffered_ccr_sse():
|
||||
for event in _openai_responses_to_sse(resp_json):
|
||||
yield event
|
||||
|
||||
return StreamingResponse(
|
||||
_residual_ccr_error_sse(),
|
||||
_buffered_ccr_sse(),
|
||||
media_type="text/event-stream",
|
||||
headers=sse_headers,
|
||||
status_code=502,
|
||||
)
|
||||
|
||||
async def _buffered_ccr_sse():
|
||||
for event in _openai_responses_to_sse(resp_json):
|
||||
yield event
|
||||
|
||||
return StreamingResponse(
|
||||
_buffered_ccr_sse(),
|
||||
media_type="text/event-stream",
|
||||
headers=sse_headers,
|
||||
return Response(
|
||||
content=response.content,
|
||||
status_code=response.status_code,
|
||||
headers=response_headers,
|
||||
)
|
||||
|
||||
return Response(
|
||||
content=response.content,
|
||||
status_code=response.status_code,
|
||||
headers=response_headers,
|
||||
)
|
||||
if buffered_stream_ccr:
|
||||
operation = asyncio.create_task(_buffered_ccr_operation())
|
||||
record_failed = self.metrics.record_failed
|
||||
|
||||
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
|
||||
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(
|
||||
{
|
||||
"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 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
|
||||
|
||||
return _BufferedCCRResponse(media_type="text/event-stream")
|
||||
return await _buffered_ccr_operation()
|
||||
except Exception as e:
|
||||
await self.metrics.record_failed(provider="openai")
|
||||
logger.error(f"[{request_id}] OpenAI responses request failed: {type(e).__name__}: {e}")
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -12,6 +14,7 @@ httpx = pytest.importorskip("httpx")
|
|||
|
||||
from fastapi.responses import StreamingResponse # noqa: E402
|
||||
from fastapi.testclient import TestClient # noqa: E402
|
||||
from starlette.requests import Request # noqa: E402
|
||||
|
||||
from headroom.cache.compression_store import get_compression_store # noqa: E402
|
||||
from headroom.ccr.tool_injection import create_ccr_tool_definition # noqa: E402
|
||||
|
|
@ -49,6 +52,10 @@ 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
|
||||
|
|
@ -299,9 +306,8 @@ def test_unresolved_ccr_only_streams_through_as_200() -> None:
|
|||
"""CCR-only turn that never resolves: the model keeps re-emitting
|
||||
headroom_retrieve so the continuation exhausts its retrieval rounds with a
|
||||
residual marker and no accompanying client tool. Per #2089 the streaming
|
||||
path no longer hard-502s here — it streams the residual headroom_retrieve
|
||||
back as a 200 SSE so the client (which owns the tool) can resolve or retry
|
||||
it, matching the non-streaming path. It must NOT 502."""
|
||||
path streams the residual headroom_retrieve back as a 200 SSE so the client
|
||||
can resolve or retry it, matching the non-streaming path."""
|
||||
config = _make_config()
|
||||
persistent_ccr = _message_response(
|
||||
[
|
||||
|
|
@ -343,8 +349,266 @@ def test_unresolved_ccr_only_streams_through_as_200() -> None:
|
|||
},
|
||||
)
|
||||
|
||||
# Fails closed no longer: residual CCR is handed back to the client as 200 SSE.
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert "text/event-stream" in resp.headers["content-type"]
|
||||
assert "headroom_retrieve" in resp.text
|
||||
assert "Unable to safely complete streamed CCR retrieval" not in resp.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_ccr_emits_keepalive_before_delayed_upstream() -> None:
|
||||
config = _make_config()
|
||||
final_response = _message_response([{"type": "text", "text": "done"}])
|
||||
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": "wait"}],
|
||||
}
|
||||
|
||||
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_retry(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
started.set()
|
||||
await release.wait()
|
||||
return httpx.Response(200, json=final_response)
|
||||
|
||||
proxy._retry_request = delayed_retry
|
||||
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] = []
|
||||
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()
|
||||
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)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_ccr_preserves_early_failure_status_and_headers() -> None:
|
||||
config = _make_config()
|
||||
body = {
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 64,
|
||||
"stream": True,
|
||||
"tools": [create_ccr_tool_definition("anthropic")],
|
||||
"messages": [{"role": "user", "content": "fail early"}],
|
||||
}
|
||||
|
||||
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 early_failure(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
await asyncio.sleep(0.05)
|
||||
return httpx.Response(
|
||||
429,
|
||||
headers={"retry-after": "7"},
|
||||
json={"error": {"message": "slow down"}},
|
||||
)
|
||||
|
||||
proxy._retry_request = early_failure
|
||||
response = await proxy.handle_anthropic_messages(Request(scope, receive))
|
||||
events: list[dict] = []
|
||||
|
||||
async def send(message): # noqa: ANN001
|
||||
events.append(message)
|
||||
|
||||
await response(scope, receive, send)
|
||||
|
||||
start = next(event for event in events if event["type"] == "http.response.start")
|
||||
headers = dict(start["headers"])
|
||||
assert start["status"] == 429
|
||||
assert headers[b"retry-after"] == b"7"
|
||||
assert b": headroom-keepalive\n\n" not in b"".join(
|
||||
event["body"] for event in events if event["type"] == "http.response.body"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_ccr_late_failure_emits_sanitized_error_event() -> None:
|
||||
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": "wait"}],
|
||||
}
|
||||
|
||||
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
|
||||
proxy_logger = logging.getLogger("headroom.proxy")
|
||||
error_records: list[logging.LogRecord] = []
|
||||
log_handler = logging.Handler()
|
||||
log_handler.setLevel(logging.ERROR)
|
||||
log_handler.emit = error_records.append
|
||||
proxy_logger.addHandler(log_handler)
|
||||
|
||||
async def delayed_failure(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
started.set()
|
||||
await release.wait()
|
||||
raise RuntimeError("boom")
|
||||
|
||||
with patch.object(
|
||||
proxy.metrics, "record_failed", new_callable=AsyncMock
|
||||
) as record_failed:
|
||||
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] = []
|
||||
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)
|
||||
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]
|
||||
assert events[-1]["more_body"] is False
|
||||
assert any(
|
||||
record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage()
|
||||
for record in error_records
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_ccr_pre_keepalive_exception_returns_json_error() -> None:
|
||||
config = _make_config()
|
||||
body = {
|
||||
"model": "claude-sonnet-4-6",
|
||||
"max_tokens": 64,
|
||||
"stream": True,
|
||||
"tools": [create_ccr_tool_definition("anthropic")],
|
||||
"messages": [{"role": "user", "content": "fail before keepalive"}],
|
||||
}
|
||||
|
||||
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 early_exception(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
raise RuntimeError("boom")
|
||||
|
||||
proxy._retry_request = early_exception
|
||||
response = await proxy.handle_anthropic_messages(Request(scope, receive))
|
||||
events: list[dict] = []
|
||||
|
||||
async def send(message): # noqa: ANN001
|
||||
events.append(message)
|
||||
|
||||
await response(scope, receive, send)
|
||||
|
||||
start = next(event for event in events if event["type"] == "http.response.start")
|
||||
bodies = [event["body"] for event in events if event["type"] == "http.response.body"]
|
||||
assert start["status"] == 502
|
||||
assert dict(start["headers"])[b"content-type"] == b"application/json"
|
||||
payload = json.loads(bodies[-1].decode())
|
||||
assert (
|
||||
payload["error"]["message"]
|
||||
== "An error occurred while processing your request. Please try again."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10,16 +10,19 @@ tests/test_proxy/test_openai_backend_path.py for that precedent.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
import pytest
|
||||
|
||||
fastapi = pytest.importorskip("fastapi")
|
||||
httpx = pytest.importorskip("httpx")
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock # noqa: E402
|
||||
from unittest.mock import AsyncMock, MagicMock, patch # noqa: E402
|
||||
|
||||
from fastapi.testclient import TestClient # noqa: E402
|
||||
from starlette.requests import Request # noqa: E402
|
||||
|
||||
from headroom.cache.compression_store import reset_compression_store # noqa: E402
|
||||
from headroom.ccr.tool_injection import CCR_TOOL_NAME # noqa: E402
|
||||
|
|
@ -260,3 +263,247 @@ def test_streaming_request_without_retrieve_tool_uses_normal_stream_path():
|
|||
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert stream_called["value"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_responses_ccr_emits_keepalive_before_delayed_upstream():
|
||||
app = _make_app()
|
||||
body = {
|
||||
"model": "gpt-5-codex",
|
||||
"input": "please wait",
|
||||
"tools": [_RETRIEVE_TOOL],
|
||||
"stream": True,
|
||||
}
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
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/responses",
|
||||
"raw_path": b"/v1/responses",
|
||||
"query_string": b"",
|
||||
"headers": [(b"authorization", b"Bearer sk-test")],
|
||||
"server": ("testserver", 80),
|
||||
"client": ("testclient", 123),
|
||||
"root_path": "",
|
||||
}
|
||||
|
||||
with TestClient(app):
|
||||
server = app.state.proxy
|
||||
|
||||
async def delayed_retry(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
started.set()
|
||||
await release.wait()
|
||||
return _final_response("https://api.openai.com/v1/responses")
|
||||
|
||||
server._retry_request = delayed_retry
|
||||
task = asyncio.create_task(server.handle_openai_responses(Request(scope, receive)))
|
||||
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()
|
||||
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)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_responses_ccr_preserves_early_failure_status_and_headers():
|
||||
app = _make_app()
|
||||
body = {
|
||||
"model": "gpt-5-codex",
|
||||
"input": "fail early",
|
||||
"tools": [_RETRIEVE_TOOL],
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
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/responses",
|
||||
"raw_path": b"/v1/responses",
|
||||
"query_string": b"",
|
||||
"headers": [(b"authorization", b"Bearer sk-test")],
|
||||
"server": ("testserver", 80),
|
||||
"client": ("testclient", 123),
|
||||
"root_path": "",
|
||||
}
|
||||
|
||||
with TestClient(app):
|
||||
server = app.state.proxy
|
||||
|
||||
async def early_failure(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
await asyncio.sleep(0.05)
|
||||
return httpx.Response(
|
||||
429,
|
||||
headers={"retry-after": "7"},
|
||||
json={"error": {"message": "slow down"}},
|
||||
request=httpx.Request("POST", "https://api.openai.com/v1/responses"),
|
||||
)
|
||||
|
||||
server._retry_request = early_failure
|
||||
response = await server.handle_openai_responses(Request(scope, receive))
|
||||
events: list[dict] = []
|
||||
|
||||
async def send(message): # noqa: ANN001
|
||||
events.append(message)
|
||||
|
||||
await response(scope, receive, send)
|
||||
|
||||
start = next(event for event in events if event["type"] == "http.response.start")
|
||||
headers = dict(start["headers"])
|
||||
assert start["status"] == 429
|
||||
assert headers[b"retry-after"] == b"7"
|
||||
assert b": headroom-keepalive\n\n" not in b"".join(
|
||||
event["body"] for event in events if event["type"] == "http.response.body"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_responses_ccr_late_failure_emits_sanitized_error_event():
|
||||
app = _make_app()
|
||||
body = {
|
||||
"model": "gpt-5-codex",
|
||||
"input": "please wait",
|
||||
"tools": [_RETRIEVE_TOOL],
|
||||
"stream": True,
|
||||
}
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
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/responses",
|
||||
"raw_path": b"/v1/responses",
|
||||
"query_string": b"",
|
||||
"headers": [(b"authorization", b"Bearer sk-test")],
|
||||
"server": ("testserver", 80),
|
||||
"client": ("testclient", 123),
|
||||
"root_path": "",
|
||||
}
|
||||
|
||||
with TestClient(app):
|
||||
server = app.state.proxy
|
||||
proxy_logger = logging.getLogger("headroom.proxy")
|
||||
error_records: list[logging.LogRecord] = []
|
||||
log_handler = logging.Handler()
|
||||
log_handler.setLevel(logging.ERROR)
|
||||
log_handler.emit = error_records.append
|
||||
proxy_logger.addHandler(log_handler)
|
||||
|
||||
async def delayed_failure(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
started.set()
|
||||
await release.wait()
|
||||
raise RuntimeError("boom")
|
||||
|
||||
with patch.object(server.metrics, "record_failed", new_callable=AsyncMock) as record_failed:
|
||||
server._retry_request = delayed_failure
|
||||
task = asyncio.create_task(server.handle_openai_responses(Request(scope, receive)))
|
||||
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)
|
||||
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]
|
||||
assert events[-1]["more_body"] is False
|
||||
assert any(
|
||||
record.levelno == logging.ERROR and "RuntimeError: boom" in record.getMessage()
|
||||
for record in error_records
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_buffered_responses_ccr_pre_keepalive_exception_returns_json_error():
|
||||
app = _make_app()
|
||||
body = {
|
||||
"model": "gpt-5-codex",
|
||||
"input": "fail before keepalive",
|
||||
"tools": [_RETRIEVE_TOOL],
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
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/responses",
|
||||
"raw_path": b"/v1/responses",
|
||||
"query_string": b"",
|
||||
"headers": [(b"authorization", b"Bearer sk-test")],
|
||||
"server": ("testserver", 80),
|
||||
"client": ("testclient", 123),
|
||||
"root_path": "",
|
||||
}
|
||||
|
||||
with TestClient(app):
|
||||
server = app.state.proxy
|
||||
|
||||
async def early_exception(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
raise RuntimeError("boom")
|
||||
|
||||
server._retry_request = early_exception
|
||||
response = await server.handle_openai_responses(Request(scope, receive))
|
||||
events: list[dict] = []
|
||||
|
||||
async def send(message): # noqa: ANN001
|
||||
events.append(message)
|
||||
|
||||
await response(scope, receive, send)
|
||||
|
||||
start = next(event for event in events if event["type"] == "http.response.start")
|
||||
bodies = [event["body"] for event in events if event["type"] == "http.response.body"]
|
||||
assert start["status"] == 502
|
||||
assert dict(start["headers"])[b"content-type"] == b"application/json"
|
||||
payload = json.loads(bodies[-1].decode())
|
||||
assert (
|
||||
payload["error"]["message"]
|
||||
== "An error occurred while processing your request. Please try again."
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue