mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy/gemini): thread savings-profile kwargs into apply() (#1994)
## Description
The native Gemini / Vertex-google compression handlers build the
pipeline call like this
(`headroom/proxy/handlers/gemini.py`, three sites —
`handle_gemini_generate_content` ~L487,
`handle_google_cloudcode_stream` ~L843, `handle_gemini_count_tokens`
~L1104):
```python
result = await self._run_compression_in_executor(
lambda: self.openai_pipeline.apply(
messages=messages,
model=model,
model_limit=context_limit,
context=extract_user_query(messages),
waste_messages=waste_messages,
), # <-- no **proxy_pipeline_kwargs(self.config)
...
)
```
Every other live compression path threads
`proxy_pipeline_kwargs(self.config)` into `apply()`
— `handlers/openai.py` (chat + responses), `handlers/anthropic.py`, and
the `/v1/compress`
endpoint. The Gemini handler never imports or calls it, so the savings
profile and the
ProxyConfig compression knobs never reach the pipeline for Gemini/Vertex
requests.
The proxy pipeline is built with only `transforms` + `provider`
(`server.py`), and
`ContentRouter` reads the accuracy-sensitive knobs per-call from
`**kwargs`. With the kwargs
missing, Gemini falls back to router defaults instead of the
profile/config values:
- `min_tokens_to_compress` → hardcoded **50** instead of the
coding-profile **25** / `config.min_tokens_to_crush`
- `protect_recent` → router default instead of the profile / configured
value
- `target_ratio` → **None** instead of the CLI default **0.4** used
everywhere else
- `max_items_after_crush`, `smart_crusher_with_compaction`,
`force_kompress` → router defaults
So Gemini/Vertex requests compress with a materially different (and
inconsistent) posture than
Claude/Codex/Cursor, and any user-tuned `HEADROOM_SAVINGS_PROFILE` /
`HEADROOM_TARGET_RATIO` / `HEADROOM_MIN_TOKENS` /
`HEADROOM_PROTECT_RECENT` is ignored on this
path.
This is the exact bug **#1534** fixed for the OpenAI chat path.
Closes: no issue filed — found while auditing profile/config threading
across the provider handlers.
## Fix
Import `proxy_pipeline_kwargs` (as `openai.py`/`anthropic.py` do) and
add
`**proxy_pipeline_kwargs(self.config)` to all three Gemini
`openai_pipeline.apply(...)` calls.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `headroom/proxy/handlers/gemini.py`: import `proxy_pipeline_kwargs`;
thread `**proxy_pipeline_kwargs(self.config)` into the three
`openai_pipeline.apply(...)` call sites (generateContent, Cloud Code
stream, countTokens).
- `tests/test_proxy/test_gemini_savings_profile.py`: drive the native
`/v1beta/models/{model}:generateContent` route with
`savings_profile="agent-90"` and assert the profile knobs
(`compress_user_messages`, `target_ratio`, `min_tokens_to_compress`,
`compress_system_messages`) reach `apply()`.
## Testing
- [x] New regression test added
(`tests/test_proxy/test_gemini_savings_profile.py`), mirroring the #1534
chat-path test
- [x] Linting passes (`ruff check`) and formatting is clean (`ruff
format --check`) — run with the CI-pinned `ruff==0.15.17`
- [ ] Full `pytest` deferred to CI (local-OOM reason below).
```text
$ uvx ruff@0.15.17 check headroom/proxy/handlers/gemini.py tests/test_proxy/test_gemini_savings_profile.py
All checks passed!
$ uvx ruff@0.15.17 format --check headroom/proxy/handlers/gemini.py tests/test_proxy/test_gemini_savings_profile.py
2 files already formatted
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.10, headroom from this branch.
Importing `headroom` pulls in the torch/transformers stack and a full
`pytest` gets OOM-killed on this box, so I verified the kwargs threading
with a dependency-free script and left the full pytest to CI.
- Exact command / steps: mechanically reproduced the call site —
`apply(**base)` (old) vs `apply(**base,
**proxy_pipeline_kwargs(config))` (new) — and captured the kwargs each
produced.
- Observed result: the old call omits the profile knobs entirely; the
new call threads them:
```text
OLD gemini apply() kwargs: ['context', 'messages', 'model', 'model_limit', 'waste_messages']
-> profile knobs DROPPED (savings profile / config ignored)
NEW gemini apply() kwargs: ['compress_system_messages', 'compress_user_messages', 'context',
'max_items_after_crush', 'messages', 'min_tokens_to_compress', 'model', 'model_limit',
'protect_recent', 'target_ratio', 'waste_messages']
-> profile knobs THREADED: target_ratio=0.10, min_tokens_to_compress=120, compress_user/system=True
GEMINI KWARGS THREADING VERIFIED
```
- Not tested: a full request through a live Gemini/Vertex upstream
(needs the heavy stack + a key). The new test drives the native route
with a mocked upstream and asserts the kwargs reach `apply()`. Full
local `pytest` deferred to CI (OOM, per above).
## 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
- [ ] 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
- [ ] New and existing unit tests pass locally with my changes — ran
lint + a standalone logic check; full pytest deferred to CI (local OOM,
disclosed above)
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
- No new dependencies; one import plus three call-site kwargs and a
test.
- @JerrettDavis tagging you — this is the Gemini sibling of the #1534
chat-path fix; the profile/config knobs are currently ignored for the
whole Gemini/Vertex surface, so it may be worth a look when you have a
moment.
Co-authored-by: JD Davis <mxjerrett@gmail.com>
This commit is contained in:
parent
ec55ddcfb3
commit
38306a331c
3 changed files with 92 additions and 0 deletions
|
|
@ -38,6 +38,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
* **cache/semantic:** key entries by the full-context hash, not the trailing query text. `SemanticCache.put` stored each response under `sha256(query)[:16]` where `query` is only the last user message, and the exact-match branch of `get` returned the slot without checking the stored entry's `messages_hash`. Two requests that share a trailing message ("continue", "yes", "run the tests") but differ in earlier context therefore collided on one slot — the second overwrote the first, and the first's hash then resolved to the second's cached response (wrong data served). Entries are now keyed by `messages_hash` when present, and `get` verifies `entry.messages_hash` before returning.
|
||||
* **proxy/openai:** stop PRE_SEND from reintroducing `tools: []` after the direct #728 fix. The OpenAI request handler now mirrors the existing `tools or _original_tools is not None` body-write guard during PRE_SEND write-back, so providers that reject empty tool arrays no longer see a tools field when the client omitted it, while explicit client `tools: []` remains preserved ([#1983](https://github.com/headroomlabs-ai/headroom/issues/1983)).
|
||||
* **proxy/openai:** keep the exact Responses function name `terminal` resident during OpenAI tool-search deferral so cache-mode optimization stops forwarding `terminal.terminal` and triggering the reserved-namespace 400 on Codex Responses ([#1946](https://github.com/headroomlabs-ai/headroom/issues/1946)).
|
||||
* **proxy/gemini:** thread the savings-profile kwargs into the native Gemini/Vertex compression paths. `handle_gemini_generate_content`, `handle_google_cloudcode_stream`, and `handle_gemini_count_tokens` called `openai_pipeline.apply()` without `proxy_pipeline_kwargs(self.config)`, so `HEADROOM_SAVINGS_PROFILE` and the ProxyConfig knobs (`target_ratio`/`min_tokens_to_compress`/`protect_recent`/`max_items_after_crush`/...) were silently dropped on the Gemini path — those requests compressed with router defaults instead of the configured profile, diverging from the Claude/Codex/Cursor paths. This is the same fix #1534 made for the OpenAI chat path; it now covers Gemini too.
|
||||
* **wrap:** `headroom wrap claude` no longer installs RTK or lean-ctx by default. Claude context-tool setup is now explicit via `--context-tool`, `--no-context-tool` remains accepted, and other wrap commands keep their current defaults ([#1915](https://github.com/headroomlabs-ai/headroom/issues/1915)).
|
||||
* **proxy/openai:** thread the savings-profile kwargs into the live `/v1/chat/completions` compression path. The chat handler called `openai_pipeline.apply()` without `proxy_pipeline_kwargs(config)`, so `HEADROOM_SAVINGS_PROFILE=agent-90` (and the individual `compress_user_messages`/`target_ratio`/`min_tokens_to_compress`/... knobs) were silently dropped — OpenAI-compatible clients like OpenCode kept protecting user messages and missed the configured profile. Both the token-mode and non-token chat branches now pass the profile kwargs, matching `handlers/anthropic.py` and the dedicated OpenAI compress endpoint ([#1534](https://github.com/headroomlabs-ai/headroom/issues/1534)).
|
||||
* **proxy:** forward Codex Desktop `/v1/responses` posts byte-faithfully so they stop returning upstream `400 {"detail":"Bad Request"}`. `handle_openai_responses` decoded the inbound body to inspect it but always re-serialized a canonical body on the way out, and it never stripped the inbound `content-encoding` header — so a `content-encoding: zstd` Codex Desktop request was forwarded as already-decoded JSON still advertising `zstd`, and the upstream ChatGPT Codex endpoint rejected it. The handler now keeps the original decoded bytes and forwards them verbatim whenever nothing (compression or memory injection) mutated the request, and drops the stale `content-encoding` header, mirroring the byte-faithful passthrough the chat and Anthropic paths already use ([#1542](https://github.com/headroomlabs-ai/headroom/issues/1542)).
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ if TYPE_CHECKING:
|
|||
from fastapi import Request
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
|
||||
from headroom.agent_savings import proxy_pipeline_kwargs
|
||||
from headroom.copilot_auth import build_copilot_upstream_url
|
||||
from headroom.proxy.auth_mode import classify_client
|
||||
from headroom.proxy.compression_decision import CompressionDecision
|
||||
|
|
@ -490,6 +491,7 @@ class GeminiHandlerMixin:
|
|||
model_limit=context_limit,
|
||||
context=extract_user_query(messages),
|
||||
waste_messages=waste_messages,
|
||||
**proxy_pipeline_kwargs(self.config),
|
||||
),
|
||||
timeout=COMPRESSION_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
|
@ -846,6 +848,7 @@ class GeminiHandlerMixin:
|
|||
model_limit=context_limit,
|
||||
context=extract_user_query(messages),
|
||||
waste_messages=waste_messages,
|
||||
**proxy_pipeline_kwargs(self.config),
|
||||
),
|
||||
timeout=COMPRESSION_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
|
@ -1106,6 +1109,7 @@ class GeminiHandlerMixin:
|
|||
model=model,
|
||||
model_limit=context_limit,
|
||||
context=extract_user_query(messages),
|
||||
**proxy_pipeline_kwargs(self.config),
|
||||
),
|
||||
timeout=COMPRESSION_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
|
|
|||
87
tests/test_proxy/test_gemini_savings_profile.py
Normal file
87
tests/test_proxy/test_gemini_savings_profile.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
"""Regression test: the native Gemini generateContent compression path must
|
||||
thread the proxy savings-profile kwargs (``proxy_pipeline_kwargs(config)``) into
|
||||
``openai_pipeline.apply`` — the same way ``handlers/openai.py`` (#1534) and
|
||||
``handlers/anthropic.py`` already do.
|
||||
|
||||
Before the fix the three Gemini/Vertex ``openai_pipeline.apply(...)`` call sites
|
||||
passed only ``messages``/``model``/``model_limit``/``context``/``waste_messages``,
|
||||
so ``HEADROOM_SAVINGS_PROFILE`` and the ProxyConfig compression knobs
|
||||
(``target_ratio``/``min_tokens_to_compress``/``protect_recent``/...) were
|
||||
silently dropped on the Gemini path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
fastapi = pytest.importorskip("fastapi")
|
||||
pytest.importorskip("httpx")
|
||||
|
||||
from fastapi.testclient import TestClient # noqa: E402
|
||||
|
||||
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
|
||||
|
||||
|
||||
def _make_fake_gemini_response() -> MagicMock:
|
||||
"""A minimal stand-in for the httpx response returned by _retry_request."""
|
||||
resp = MagicMock()
|
||||
resp.status_code = 200
|
||||
resp.headers = {"content-type": "application/json"}
|
||||
resp.content = b'{"candidates":[{"content":{"parts":[{"text":"ok"}]}}],"usageMetadata":{"promptTokenCount":100,"candidatesTokenCount":2}}'
|
||||
resp.json.return_value = {
|
||||
"candidates": [{"content": {"parts": [{"text": "ok"}]}}],
|
||||
"usageMetadata": {"promptTokenCount": 100, "candidatesTokenCount": 2},
|
||||
}
|
||||
return resp
|
||||
|
||||
|
||||
def test_gemini_generate_content_threads_savings_profile_kwargs_into_apply():
|
||||
"""With HEADROOM_SAVINGS_PROFILE=agent-90, the native Gemini path must pass
|
||||
the profile knobs (compress_user_messages, target_ratio, ...) to apply()."""
|
||||
config = ProxyConfig(
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
savings_profile="agent-90",
|
||||
)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def recording_apply(**kwargs):
|
||||
captured.update(kwargs)
|
||||
sent = kwargs["messages"]
|
||||
return SimpleNamespace(
|
||||
messages=sent,
|
||||
transforms_applied=[],
|
||||
timing={},
|
||||
tokens_before=4000,
|
||||
tokens_after=400,
|
||||
waste_signals=None,
|
||||
)
|
||||
|
||||
# A large user message so the compression decision actually fires.
|
||||
big = "word " * 4000
|
||||
|
||||
app = create_app(config)
|
||||
with TestClient(app) as client:
|
||||
proxy = client.app.state.proxy
|
||||
proxy.openai_pipeline.apply = MagicMock(side_effect=recording_apply)
|
||||
proxy._retry_request = AsyncMock(return_value=_make_fake_gemini_response())
|
||||
|
||||
resp = client.post(
|
||||
"/v1beta/models/gemini-2.0-flash:generateContent?key=test-key",
|
||||
json={"contents": [{"parts": [{"text": big}]}]},
|
||||
)
|
||||
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert proxy.openai_pipeline.apply.call_count >= 1, "compression apply() never ran"
|
||||
|
||||
# The agent-90 profile knobs must be present on the apply() call.
|
||||
assert captured.get("compress_user_messages") is True
|
||||
assert captured.get("target_ratio") == 0.10
|
||||
assert captured.get("min_tokens_to_compress") == 120
|
||||
assert captured.get("compress_system_messages") is True
|
||||
Loading…
Add table
Add a link
Reference in a new issue