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

fix(gemini): resolve native CCR retrieval calls (#2253)

Browse Source 
## Description

Buffered native Gemini requests currently return `headroom_retrieve`
function calls to the client because `GeminiHandlerMixin` never invokes
the shared CCR response handler. This wires native Gemini request and
response translation into the provider handler while reusing the
existing Google CCR extraction, retrieval, round-limit, mixed-tool, and
`functionResponse` machinery.

Streaming native Gemini and Gemini's OpenAI-compatible
`MALFORMED_FUNCTION_CALL` behavior remain separate surfaces.

This follows the current support boundary documented in
https://github.com/headroomlabs-ai/headroom/pull/2044.

Closes #2041

## 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)
- [x] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Invoke `CCRResponseHandler` for successful buffered native Gemini
responses containing `headroom_retrieve`.
- Build Gemini-native continuation requests with the model
`functionCall` and matching user `functionResponse`.
- Inject the existing Google CCR function declaration while preserving
sibling Gemini tool configurations.
- Preserve mixed client-tool responses, streaming requests, non-CCR
responses, and upstream error bodies.
- Preserve Google `functionCall.id` as `functionResponse.id` through the
shared CCR identity contract.
- Leave streaming requests outside buffered CCR injection.
- Fail closed when an exclusive CCR call remains unresolved after
continuation.
- Update the CCR documentation to describe buffered native Gemini
support and the mixed-tool boundary.

## Testing

- [x] Unit tests pass (`uv run pytest tests/test_proxy_handlers_batch.py
-k "gemini_native_ccr or gemini_stream" -q`)
- [x] Linting passes (`uv run ruff check
headroom/proxy/handlers/gemini.py tests/test_proxy_handlers_batch.py`)
- [ ] Type checking passes (`uv run mypy headroom`)
- [x] New tests added for new functionality when applicable
- [ ] Manual testing performed

### Test Output

```text
Focused tests: `14 passed, 22 deselected` for `uv run pytest tests/test_proxy_handlers_batch.py -k "gemini_native_ccr or gemini_stream" -q`; `43 passed` for `uv run pytest tests/test_ccr_response_handler.py tests/test_ccr_response_handler_extra.py -q`. Scoped Ruff check and format check passed for the changed Python files. Full repository format remains blocked by pre-existing formatting outside this target.
```

## Real Behavior Proof

- Environment: Windows, Python 3.12, commit `7b3a92a8`; local
native-shape behavioral harness with no Gemini credentials.
- Exact command / steps: run the focused native Gemini handler tests,
then capture a live `generateContent` request and continuation after a
Gemini credential is available.
- Observed result: local tests prove the buffered `functionCall` to
`functionResponse` continuation, mixed-tool preservation, declaration
preservation, error forwarding, and retrieval-result shapes.
- Not tested: owner-reaching live Gemini continuation and native Gemini
streaming CCR continuation.

## 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] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [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

## Additional Notes

The patch keeps Gemini wire translation in `GeminiHandlerMixin` and
extends the provider-neutral CCR identity fields for Google call ids.
Native streaming continuation and the OpenAI-compatible Gemini round-two
failure are outside this PR.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
Rod Boev
2026-08-12 00:40:49 -04:00
committed by GitHub
parent 5568d738af
commit 2483f57002
6 changed files with 637 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/content/docs/ccr.mdx
+9 -7
View File
@@ -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.
<Callout type="warning" title="Current Gemini limitation">
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).
<Callout type="warning" title="Gemini CCR boundary">
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).
</Callout>
## Phase 4: Context Tracker
headroom/ccr/response_handler.py
+12 -8
View File
@@ -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,
headroom/ccr/tool_calls.py
+10 -4
View File
@@ -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
headroom/proxy/handlers/gemini.py
+150 -7
View File
@@ -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}")
tests/test_ccr_response_handler_extra.py
+21
View File
@@ -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(
tests/test_proxy_handlers_batch.py
+435
View File
@@ -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,
*,