diff --git a/docs/content/docs/ccr.mdx b/docs/content/docs/ccr.mdx index 18c4ff2c4..4d668b07c 100644 --- a/docs/content/docs/ccr.mdx +++ b/docs/content/docs/ccr.mdx @@ -73,13 +73,15 @@ When the LLM calls `headroom_retrieve`: The client never sees CCR tool calls on the Anthropic and OpenAI proxy paths; Headroom resolves them transparently there. - - Native Gemini requests do not yet run the server-side CCR response handler, so - `headroom_retrieve` is not resolved transparently on that path today. Google's - OpenAI-compatible Gemini endpoint can also return - `finish_reason=MALFORMED_FUNCTION_CALL` on large function-response continuations - after CCR retrieval. If you need fully transparent CCR resolution today, use the - Anthropic or OpenAI proxy paths. See [issue #2041](https://github.com/headroomlabs-ai/headroom/issues/2041). + + Buffered native Gemini requests resolve `headroom_retrieve` server-side and + return the model's final response. Streaming native Gemini requests keep the + existing forwarding behavior. When a response contains `headroom_retrieve` + alongside a client-owned function call, Headroom preserves both calls for the + client instead of resolving the mixed response. Google's OpenAI-compatible Gemini endpoint can + also return `finish_reason=MALFORMED_FUNCTION_CALL` on large function-response + continuations after CCR retrieval; that separate limitation remains tracked in + [issue #2041](https://github.com/headroomlabs-ai/headroom/issues/2041). ## Phase 4: Context Tracker diff --git a/headroom/ccr/response_handler.py b/headroom/ccr/response_handler.py index ff93aa981..7c959e9a6 100644 --- a/headroom/ccr/response_handler.py +++ b/headroom/ccr/response_handler.py @@ -58,6 +58,7 @@ class CCRToolResult: content: str success: bool items_retrieved: int = 0 + tool_name: str | None = None @dataclass @@ -211,6 +212,7 @@ class CCRResponseHandler: tool_call_id=ccr_call.tool_call_id, content=content, success=False, + tool_name=ccr_call.tool_name, ) # Retrieval is by hash: always return the full original content. @@ -229,6 +231,7 @@ class CCRResponseHandler: content=content, success=True, items_retrieved=entry.original_item_count, + tool_name=ccr_call.tool_name, ) miss_status = ( @@ -249,6 +252,7 @@ class CCRResponseHandler: tool_call_id=ccr_call.tool_call_id, content=content, success=False, + tool_name=ccr_call.tool_name, ) except Exception as e: @@ -264,6 +268,7 @@ class CCRResponseHandler: tool_call_id=ccr_call.tool_call_id, content=content, success=False, + tool_name=ccr_call.tool_name, ) def _create_tool_result_message( @@ -337,14 +342,13 @@ class CCRResponseHandler: response_data = json.loads(result.content) except json.JSONDecodeError: response_data = {"content": result.content} - parts.append( - { - "functionResponse": { - "name": result.tool_call_id, # tool_call_id contains the function name for Google - "response": response_data, - } - } - ) + function_response = { + "name": result.tool_name or result.tool_call_id, + "response": response_data, + } + if result.tool_name and result.tool_call_id != result.tool_name: + function_response["id"] = result.tool_call_id + parts.append({"functionResponse": function_response}) return { "role": "user", "parts": parts, diff --git a/headroom/ccr/tool_calls.py b/headroom/ccr/tool_calls.py index 46e5feacb..44d75ff34 100644 --- a/headroom/ccr/tool_calls.py +++ b/headroom/ccr/tool_calls.py @@ -14,6 +14,7 @@ class CCRToolCall: tool_call_id: str hash_key: str + tool_name: str | None = None def extract_tool_calls(response: dict[str, Any], provider: str) -> list[dict[str, Any]]: @@ -88,6 +89,8 @@ def tool_call_id_for_provider(tool_call: dict[str, Any], provider: str) -> str: if provider == "google": function_call = tool_call.get("functionCall", {}) if isinstance(function_call, dict): + if function_call.get("id"): + return str(function_call["id"]) name = function_call.get("name", CCR_TOOL_NAME) return str(name) return CCR_TOOL_NAME @@ -111,11 +114,14 @@ def parse_ccr_tool_calls( other_calls.append(tool_call) continue + tool_name = None + tool_call_id = tool_call_id_for_provider(tool_call, provider) + if provider == "google": + function_call = tool_call.get("functionCall", {}) + if isinstance(function_call, dict) and function_call.get("id"): + tool_name = str(function_call.get("name", CCR_TOOL_NAME)) ccr_calls.append( - CCRToolCall( - tool_call_id=tool_call_id_for_provider(tool_call, provider), - hash_key=hash_key, - ) + CCRToolCall(tool_call_id=tool_call_id, hash_key=hash_key, tool_name=tool_name) ) return ccr_calls, other_calls diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index 941d31bd3..3311fe69c 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -36,6 +36,14 @@ def _usage_int(value: Any, default: int = 0) -> int: return int(value) +class _GeminiContinuationError(Exception): + def __init__(self, status_code: int, content: bytes, headers: dict[str, str]) -> None: + super().__init__(f"Gemini continuation failed with HTTP {status_code}") + self.status_code = status_code + self.content = content + self.headers = headers + + class GeminiHandlerMixin: """Mixin providing Gemini API handler methods for HeadroomProxy.""" @@ -630,6 +638,68 @@ class GeminiHandlerMixin: except Exception as e: logger.warning(f"[{request_id}] Memory injection failed (gemini): {e}") + query_params = dict(request.query_params) + is_streaming = query_params.get("alt") == "sse" or request.url.path.endswith( + ":streamGenerateContent" + ) + native_tools = body.get("tools") + native_function_declarations = None + + def rebuild_tools(function_declarations: list[dict]) -> list[dict]: + rebuilt_tools = [] + replaced = False + declaration_tools = [ + tool for tool in body.get("tools") or [] if "functionDeclarations" in tool + ] + later_names = { + declaration.get("name") + for tool in declaration_tools[1:] + for declaration in tool["functionDeclarations"] + } + first_declarations = [ + declaration + for declaration in function_declarations + if declaration.get("name") not in later_names + ] + for tool in body.get("tools") or []: + if "functionDeclarations" in tool and not replaced: + rebuilt_tools.append({**tool, "functionDeclarations": first_declarations}) + replaced = True + else: + rebuilt_tools.append(tool) + if not replaced: + rebuilt_tools.append({"functionDeclarations": function_declarations}) + return rebuilt_tools + + ccr_inject_tool = getattr(self.config, "ccr_inject_tool", True) + ccr_inject_system_instructions = getattr( + self.config, "ccr_inject_system_instructions", False + ) + if ccr_inject_tool and tokens_saved > 0 and not is_streaming: + from headroom.ccr import CCRToolInjector + + seen_names = set() + native_function_declarations = [] + for tool in native_tools or []: + for declaration in tool.get("functionDeclarations", []): + name = declaration.get("name") + if name not in seen_names: + native_function_declarations.append(declaration) + seen_names.add(name) + injector = CCRToolInjector( + provider="google", + inject_tool=True, + inject_system_instructions=ccr_inject_system_instructions, + ) + optimized_messages, injected_funcs, was_injected = injector.process_request( + optimized_messages, native_function_declarations + ) + if was_injected: + native_function_declarations = injected_funcs + body["tools"] = rebuild_tools(injected_funcs) + elif native_function_declarations is not None: + native_function_declarations = list(native_function_declarations) + # Convert back to Gemini format if optimized if optimized_messages != messages: optimized_contents, optimized_system = self._messages_to_gemini_contents( @@ -644,12 +714,6 @@ class GeminiHandlerMixin: elif "systemInstruction" in body: del body["systemInstruction"] - # Check if streaming requested via query param - query_params = dict(request.query_params) - is_streaming = query_params.get("alt") == "sse" or request.url.path.endswith( - ":streamGenerateContent" - ) - # Build URL - model is extracted from path. Vertex publisher # routes use the request's full path under the Vertex base URL; # native Gemini uses the public Gemini API shape. @@ -703,6 +767,8 @@ class GeminiHandlerMixin: total_input_tokens = optimized_tokens # fallback output_tokens = 0 cache_read_tokens = 0 + resp_json = None + response_content = response.content try: resp_json = response.json() usage = resp_json.get("usageMetadata", {}) @@ -736,6 +802,77 @@ class GeminiHandlerMixin: f"[{request_id}] Failed to extract cached tokens from Gemini response: {e}" ) + if ( + response.status_code == 200 + and isinstance(resp_json, dict) + and self.ccr_response_handler + and getattr(getattr(self.ccr_response_handler, "config", None), "enabled", True) + and self.ccr_response_handler.has_ccr_tool_calls(resp_json, "google") + ): + + async def api_call_fn( + native_contents: list[dict], + function_declarations: list[dict] | None, + ) -> dict[str, Any]: + continuation_body = {**body, "contents": native_contents} + if function_declarations is not None: + continuation_body["tools"] = rebuild_tools(function_declarations) + continuation_headers = { + key: value + for key, value in headers.items() + if key.lower() + not in ("accept-encoding", "content-encoding", "content-length") + } + continuation = await self._retry_request( + "POST", url, continuation_headers, continuation_body + ) + if continuation.status_code >= 400: + return { + "_headroom_continuation_error": { + "status_code": continuation.status_code, + "content": continuation.content, + "headers": dict(continuation.headers), + } + } + try: + return continuation.json() + except (json.JSONDecodeError, ValueError, TypeError): + return { + "_headroom_continuation_error": { + "status_code": continuation.status_code, + "content": continuation.content, + "headers": dict(continuation.headers), + } + } + + final_resp_json = await self.ccr_response_handler.handle_response( + resp_json, + body.get("contents", []), + native_function_declarations, + api_call_fn, + provider="google", + ) + continuation_error = final_resp_json.get("_headroom_continuation_error") + if isinstance(continuation_error, dict): + raise _GeminiContinuationError( + continuation_error["status_code"], + continuation_error["content"], + continuation_error["headers"], + ) + from headroom.ccr.response_handler import RESIDUAL_CCR_ERROR + + if ( + self.ccr_response_handler.residual_ccr_status(final_resp_json, "google") + == RESIDUAL_CCR_ERROR + ): + raise RuntimeError("Gemini CCR continuation left an unresolved retrieval") + resp_json = final_resp_json + response_content = json.dumps(resp_json).encode() + usage = resp_json.get("usageMetadata", {}) + total_input_tokens = usage.get("promptTokenCount", total_input_tokens) + output_tokens = usage.get("candidatesTokenCount", output_tokens) + cache_read_tokens = usage.get("cachedContentTokenCount", cache_read_tokens) + uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens) # optimized_tokens carries Gemini's own promptTokenCount, which is @@ -822,10 +959,16 @@ class GeminiHandlerMixin: response_headers["x-headroom-compression-failed"] = "true" return Response( - content=response.content, + content=response_content, status_code=response.status_code, headers=response_headers, ) + except _GeminiContinuationError as e: + await self.metrics.record_failed(provider=provider_name) + response_headers = dict(e.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + return Response(content=e.content, status_code=e.status_code, headers=response_headers) except Exception as e: await self.metrics.record_failed(provider=provider_name) logger.error(f"[{request_id}] Gemini request failed: {type(e).__name__}: {e}") diff --git a/tests/test_ccr_response_handler_extra.py b/tests/test_ccr_response_handler_extra.py index 20883ee19..36d2849fb 100644 --- a/tests/test_ccr_response_handler_extra.py +++ b/tests/test_ccr_response_handler_extra.py @@ -137,6 +137,27 @@ def test_create_tool_result_message_google_and_generic_formats() -> None: assert invalid_google["parts"][0]["functionResponse"]["response"] == {"content": "not-json"} +def test_create_tool_result_message_google_preserves_call_id() -> None: + handler = CCRResponseHandler() + message = handler._create_tool_result_message( + [ + CCRToolResult( + tool_call_id="call-1", + tool_name="headroom_retrieve", + content='{"count": 1}', + success=True, + ) + ], + "google", + ) + + assert message["parts"][0]["functionResponse"] == { + "name": "headroom_retrieve", + "id": "call-1", + "response": {"count": 1}, + } + + def test_extract_assistant_message_google_and_generic() -> None: handler = CCRResponseHandler() google_message = handler._extract_assistant_message( diff --git a/tests/test_proxy_handlers_batch.py b/tests/test_proxy_handlers_batch.py index 79c07182e..5188b718a 100644 --- a/tests/test_proxy_handlers_batch.py +++ b/tests/test_proxy_handlers_batch.py @@ -6,7 +6,10 @@ from types import SimpleNamespace import pytest +from headroom.cache.compression_store import CompressionEntry +from headroom.ccr import response_handler as response_handler_module from headroom.proxy.handlers import batch as batch_module +from headroom.proxy.handlers import gemini as gemini_module from headroom.proxy.handlers.gemini import GeminiHandlerMixin @@ -148,11 +151,443 @@ class FakeRequest: self.headers = headers or {} self.method = method self.url = SimpleNamespace(path=path, query=query) + self.query_params = {} async def body(self) -> bytes: return self._body +class NativeGeminiHandler(DummyBatchHandler): + def __init__(self, responses: list[FakeResponse]) -> None: + super().__init__() + self.config.optimize = True + self.config.ccr_inject_tool = True + self.config.ccr_inject_system_instructions = False + self.memory_handler = None + self.rate_limiter = None + self.usage_reporter = None + self.responses = iter(responses) + self.sent_bodies: list[dict] = [] + from headroom.ccr.response_handler import CCRResponseHandler + + self.ccr_response_handler = CCRResponseHandler() + self.openai_pipeline = SimpleNamespace( + apply=lambda **kwargs: SimpleNamespace( + messages=[ + { + "role": "user", + "content": "compressed [100 items compressed to 1. Retrieve more: hash=aaaaaaaaaaaaaaaaaaaaaaaa]", + } + ], + timing={}, + tokens_before=10, + tokens_after=5, + transforms_applied=[], + waste_signals=SimpleNamespace(to_dict=lambda: {}), + ) + ) + + def _gemini_contents_to_messages( + self, contents, system_instruction=None, *, include_function_responses=False + ): # noqa: ANN001, ANN201 + return GeminiHandlerMixin._gemini_contents_to_messages( + self, + contents, + system_instruction, + include_function_responses=include_function_responses, + ) + + def _messages_to_gemini_contents(self, messages): # noqa: ANN001, ANN201 + return GeminiHandlerMixin._messages_to_gemini_contents(self, messages) + + async def _retry_request(self, method, url, headers, body, **kwargs): # noqa: ANN001, ANN201 + self.sent_bodies.append(body) + return next(self.responses) + + async def _run_compression_in_executor(self, fn, *, timeout): # noqa: ANN001, ANN201 + return fn() + + +def install_native_gemini_compression(monkeypatch: pytest.MonkeyPatch) -> None: + class Decision: + should_compress = True + passthrough_reason = "" + + def apply_to_tags(self, tags) -> None: # noqa: ANN001 + return None + + monkeypatch.setattr(gemini_module.CompressionDecision, "decide", lambda **kwargs: Decision()) + + +def native_gemini_request(tools=None) -> dict: # noqa: ANN001 + return { + "contents": [{"role": "user", "parts": [{"text": "compressed input"}]}], + "generationConfig": {"temperature": 0.2}, + **({"tools": tools} if tools is not None else {}), + } + + +def native_ccr_response() -> FakeResponse: + return FakeResponse( + json_data={ + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "functionCall": { + "name": "headroom_retrieve", + "id": "call-1", + "args": {"hash": "aaaaaaaaaaaaaaaaaaaaaaaa"}, + } + } + ], + } + } + ], + "usageMetadata": {"promptTokenCount": 5}, + } + ) + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_continuation(monkeypatch: pytest.MonkeyPatch) -> None: + install_native_gemini_compression(monkeypatch) + from headroom.ccr.response_handler import CCRToolResult + + final = FakeResponse( + json_data={ + "candidates": [{"content": {"role": "model", "parts": [{"text": "final answer"}]}}] + } + ) + handler = NativeGeminiHandler([native_ccr_response(), final]) + handler.ccr_response_handler._execute_retrieval = lambda call: CCRToolResult( + call.tool_call_id, + json.dumps({"hash": call.hash_key, "original_content": [{"type": "code"}]}), + True, + 1, + "headroom_retrieve", + ) + + response = await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request()), + headers={"content-type": "application/json", "x-goog-api-key": "secret"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + assert response.status_code == 200 + assert ( + json.loads(response.body)["candidates"][0]["content"]["parts"][0]["text"] == "final answer" + ), response.body + assert len(handler.sent_bodies) == 2 + continuation = handler.sent_bodies[1]["contents"] + assert continuation[-2]["role"] == "model" + assert continuation[-2]["parts"][0]["functionCall"]["name"] == "headroom_retrieve" + assert continuation[-1]["role"] == "user" + assert continuation[-1]["parts"][0]["functionResponse"]["name"] == "headroom_retrieve" + assert continuation[-1]["parts"][0]["functionResponse"]["id"] == "call-1" + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_tools(monkeypatch: pytest.MonkeyPatch) -> None: + install_native_gemini_compression(monkeypatch) + handler = NativeGeminiHandler( + [FakeResponse(json_data={"candidates": [{"content": {"parts": [{"text": "answer"}]}}]})] + ) + tools = [ + {"functionDeclarations": [{"name": "client_tool"}]}, + {"functionDeclarations": [{"name": "second_tool"}]}, + {"googleSearch": {}}, + {"codeExecution": {}}, + ] + + await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request(tools)), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + forwarded_tools = handler.sent_bodies[0]["tools"] + assert forwarded_tools[2:] == tools[2:] + declarations = forwarded_tools[0]["functionDeclarations"] + assert {item["name"] for item in declarations} == {"client_tool", "headroom_retrieve"} + assert forwarded_tools[1]["functionDeclarations"] == [{"name": "second_tool"}] + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_does_not_duplicate_existing_declaration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_native_gemini_compression(monkeypatch) + tools = [ + {"functionDeclarations": [{"name": "client_tool"}]}, + {"functionDeclarations": [{"name": "headroom_retrieve"}]}, + ] + handler = NativeGeminiHandler( + [FakeResponse(json_data={"candidates": [{"content": {"parts": [{"text": "answer"}]}}]})] + ) + + await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request(tools)), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + names = [ + declaration["name"] + for tool in handler.sent_bodies[0]["tools"] + for declaration in tool.get("functionDeclarations", []) + ] + assert names.count("headroom_retrieve") == 1 + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_does_not_inject_into_streaming_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_native_gemini_compression(monkeypatch) + handler = NativeGeminiHandler([FakeResponse()]) + captured: dict[str, object] = {} + + async def fake_stream(*args, **kwargs): # noqa: ANN002, ANN003, ANN202 + captured["body"] = args[2] + return FakeResponse() + + monkeypatch.setattr(handler, "_stream_response", fake_stream, raising=False) + tools = [{"functionDeclarations": [{"name": "client_tool"}]}] + await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request(tools)), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:streamGenerateContent", + ), + "gemini-2.5-flash", + ) + + streamed_tools = captured["body"]["tools"] # type: ignore[index] + names = [ + declaration["name"] + for tool in streamed_tools + for declaration in tool.get("functionDeclarations", []) + ] + assert names == ["client_tool"] + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_mixed(monkeypatch: pytest.MonkeyPatch) -> None: + install_native_gemini_compression(monkeypatch) + response_json = { + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "headroom_retrieve", + "args": {"hash": "aaaaaaaaaaaaaaaaaaaaaaaa"}, + } + }, + {"functionCall": {"name": "client_tool", "args": {}}}, + ] + } + } + ] + } + handler = NativeGeminiHandler([FakeResponse(json_data=response_json)]) + + response = await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request()), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + assert response.status_code == 200 + assert len(handler.sent_bodies) == 1 + assert json.loads(response.body) == response_json + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_non_ccr_function_call_is_not_intercepted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_native_gemini_compression(monkeypatch) + response_json = { + "candidates": [ + {"content": {"parts": [{"functionCall": {"name": "client_tool", "args": {}}}]}} + ] + } + handler = NativeGeminiHandler([FakeResponse(json_data=response_json)]) + + response = await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request()), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + assert response.status_code == 200 + assert len(handler.sent_bodies) == 1 + assert response.body == b"{}" + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_continuation_error_preserves_upstream_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_native_gemini_compression(monkeypatch) + handler = NativeGeminiHandler( + [ + native_ccr_response(), + FakeResponse(status_code=503, content=b"busy", headers={"retry-after": "2"}), + ] + ) + + response = await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request()), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + assert response.status_code == 503 + assert response.body == b"busy" + assert response.headers["retry-after"] == "2" + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_continuation_non_json_preserves_upstream_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_native_gemini_compression(monkeypatch) + handler = NativeGeminiHandler( + [native_ccr_response(), FakeResponse(status_code=200, content=b"upstream")] + ) + + response = await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request()), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + assert response.status_code == 200 + assert response.body == b"upstream" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "original_content", + [[{"type": "code", "text": "print('x')"}], "plain text", {"key": "value"}, 42], + ids=["code-aware-array", "kompress-text", "mcp-object", "mcp-scalar"], +) +async def test_gemini_native_ccr_uses_real_retrieval_result_shape( + monkeypatch: pytest.MonkeyPatch, original_content +) -> None: # noqa: ANN001 + install_native_gemini_compression(monkeypatch) + entry = CompressionEntry( + hash="a" * 24, + original_content=json.dumps(original_content), + compressed_content="compressed", + original_tokens=10, + compressed_tokens=2, + original_item_count=1, + compressed_item_count=1, + tool_name="headroom_retrieve", + tool_call_id="headroom_retrieve", + query_context=None, + created_at=0, + ) + + class Store: + def get_entry_status(self, hash_key, clean_expired=True): # noqa: ANN001, ARG002 + return {"status": "available", "default_ttl_seconds": 1800} + + def retrieve(self, hash_key): # noqa: ANN001, ARG002 + return entry + + monkeypatch.setattr(response_handler_module, "get_compression_store", lambda: Store()) + handler = NativeGeminiHandler( + [ + native_ccr_response(), + FakeResponse(json_data={"candidates": [{"content": {"parts": [{"text": "done"}]}}]}), + ] + ) + + response = await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request()), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + assert response.status_code == 200 + function_response = handler.sent_bodies[1]["contents"][-1]["parts"][0]["functionResponse"] + assert function_response["response"]["original_content"] == json.dumps(original_content) + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_preserves_non_ccr_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + install_native_gemini_compression(monkeypatch) + handler = NativeGeminiHandler([FakeResponse(status_code=503, content=b"busy")]) + + response = await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request()), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + assert response.status_code == 503 + assert response.body == b"busy" + + +@pytest.mark.asyncio +async def test_gemini_native_ccr_residual(monkeypatch: pytest.MonkeyPatch) -> None: + install_native_gemini_compression(monkeypatch) + from headroom.ccr.response_handler import CCRToolResult + + handler = NativeGeminiHandler([native_ccr_response()] * 4) + handler.ccr_response_handler._execute_retrieval = lambda call: CCRToolResult( + "headroom_retrieve", "still unresolved", True, 0 + ) + + response = await handler.handle_gemini_generate_content( + FakeRequest( + json.dumps(native_gemini_request()), + headers={"content-type": "application/json"}, + path="/v1beta/models/gemini-2.5-flash:generateContent", + ), + "gemini-2.5-flash", + ) + + assert response.status_code == 502 + + def install_batch_support_modules( monkeypatch: pytest.MonkeyPatch, *,