headroom icon indicating copy to clipboard operation
headroom copied to clipboard

fix(rust): proxy token-mode — freeze prefix, compress new-turn tool outputs only

Open chopratejas opened this issue 3 months ago • 2 comments

Summary

The Rust proxy now follows exactly one compression rule:

Find the last user message. Freeze every byte before it. In the new turn, walk for tool outputs and route them through the CompressionPipeline. Everything else passes through.

This replaces the ICM-based interceptor (which risked busting prefix caches) with a simpler, safer design that matches what the Python proxy does in token mode.

Why this design

  • Provider prefix caches (Anthropic `cache_control`, OpenAI auto-prefix, OpenAI `prompt_cache_key`) are positional. We never modify prefix bytes → cache hits 100% preserved.
  • User text, system prompts, assistant text are never touched → accuracy preserved.
  • Tool outputs (file reads, search results, build logs, diffs) are where the tokens are. The `CompressionPipeline` already in `headroom-core` compresses them with CCR backup.

What this PR removes

  • ICM is no longer wired into the proxy. It stays in `headroom-core` for future use, but the proxy crate doesn't import it. ICM-style message dropping risks cache busts; token mode rules that out by construction.

What this PR adds

  • `compression/walker.rs` — `compress_blob()` shared helper
  • `compression/pipeline.rs` — build `CompressionPipeline` + `CcrStore` once at startup
  • `compression/anthropic.rs` — token-mode walker (`tool_result` blocks)
  • `compression/openai.rs` — token-mode walker (`role:tool` messages)
  • `compression/responses.rs` — token-mode walker (`function_call_output` items; `previous_response_id` triggers full passthrough)
  • `compression/mod.rs` — endpoint classifier + dispatch

Per-endpoint rules (sacred = never touched)

Anthropic /v1/messages

  • Sacred: `system`, all `cache_control` markers, `tools`, anything in prefix, all assistant content (incl. `thinking`/`tool_use`), images, documents, plain user text
  • Compressed: `tool_result` block content in user messages of the new turn (string OR list-of-blocks; text-blocks compressed, images preserved)

OpenAI /v1/chat/completions

  • Sacred: `system`/`developer`/`user`/`assistant` messages, all non-`messages` fields, anything in prefix
  • Compressed: `role:tool` messages with string content in the new turn

OpenAI /v1/responses

  • Sacred: `instructions`, `prompt_cache_key`, all non-`input` fields, message items of any role, `function_call`, `reasoning` (encrypted_content), `*_call` items, image items, anything in prefix
  • Compressed: `function_call_output` items in new turn
  • Special: `previous_response_id` set → full passthrough (server has the conversation)

Test plan

  • [x] Unit: 35 tests passing (compression module + walker)
  • [x] Mock integration: 8 tests including prefix-byte-identical assertions across all three endpoints
  • [x] Real OpenAI e2e (`HEADROOM_E2E=1`): 3/3 pass
    • `/v1/chat/completions`: 70 prompt tokens received (raw was ~30K)
    • `/v1/responses`: 68 input tokens received (same payload)
    • Prefix cache: 3,456 of 3,518 tokens cached on second turn — proves byte-stable prefix
  • [x] `cargo test --workspace` → 904 passed, 0 failed
  • [x] `cargo clippy --workspace --all-targets -- -D warnings` → clean
  • [x] `make ci-precheck` → green

chopratejas avatar May 02 '26 04:05 chopratejas

This pull request has had no recent activity and is being marked stale. Please rebase, resolve conflicts, or comment if it is still actively being worked.

github-actions[bot] avatar Jun 10 '26 16:06 github-actions[bot]

This is still valuable work, but the branch is conflicted and the last run had failing Docker e2e jobs. Please rebase on latest main, resolve the conflicts, and rerun the Rust plus Docker e2e checks before review resumes.

JerrettDavis avatar Jun 10 '26 17:06 JerrettDavis