Skip to content

feat(guardrails): add headroom guardrail for message compression - #31407

Merged
krrish-berri-2 merged 23 commits into
litellm_internal_stagingfrom
litellm_headroom_guardrail
Jun 27, 2026
Merged

feat(guardrails): add headroom guardrail for message compression#31407
krrish-berri-2 merged 23 commits into
litellm_internal_stagingfrom
litellm_headroom_guardrail

Conversation

@krrish-berri-2

@krrish-berri-2 krrish-berri-2 commented Jun 26, 2026

Copy link
Copy Markdown
Contributor

Relevant issues

Linear ticket

Pre-Submission checklist

  • I have added meaningful tests
  • My PR passes all unit tests on make test-unit
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

  • Branch creation CI run
    Link:

  • CI run for the last commit
    Link:

  • Merge / cherry-pick CI run
    Links:

Screenshots / Proof of Fix

End-to-end trace on a live proxy with real DB confirmed the full chain works:

[headroom] apply_guardrail called: input_type=request has_structured_messages=True
[headroom-propagate] slg_info=True logging_obj_type=Logging
[SLP] slg_in_metadata=True
[PAYLOAD] guardrail_information=True  ← written to DB

To test: run the proxy with a config that includes the headroom guardrail (with default_on: true or via key metadata), send a /v1/messages request, and verify guardrail_information is non-null in the spend log.

Type

🆕 New Feature + 🐛 Bug Fixes

Changes

Core feature: headroom guardrail for message compression

Adds a new headroom guardrail type that compresses request messages via POST /v1/compress on a configurable headroom service before they reach the LLM. Useful for reducing token costs on long-context conversations.

Config:

guardrails:
  - guardrail_name: headroom-compression
    litellm_params:
      guardrail: headroom
      mode: pre_call
      api_base: https://your-headroom-service/
      default_on: true

Key behaviors:

  • Runs on the unified guardrail path via apply_guardrail, works for both /v1/chat/completions and /v1/messages
  • x-headroom-bypass: true request header skips compression
  • Model forwarded to headroom can be set in config or falls back to the request model
  • API key and base URL configurable via api_key/api_base in guardrail config or HEADROOM_API_KEY/HEADROOM_API_BASE env vars
  • Raises 502 if compression service returns an empty message list or is unreachable
  • Decorated with @log_guardrail_information so guardrail_information is populated in spend logs

Bug fixes made during development:

structured_messages write-back regression: the unified guardrail handler for both OpenAI and Anthropic paths was unconditionally writing back structured_messages even when the guardrail returned inputs unchanged. This clobbered data["messages"] with a filtered copy (e.g. skip-system filtered messages). Fixed with an identity check — only write back when the guardrail actually returned a new object.

anthropic_messages_pt crash on system messages with cache_control: translate_anthropic_to_openai preserves system messages with list content (cache_control blocks) as-is. anthropic_messages_pt rejects these. Fixed by stripping system messages before the reverse-translation call — they live in data["system"] in the Anthropic passthrough path, not data["messages"].

thinking.cache_control rejected by Anthropic: the prompt-caching-scope beta adds cache_control to all content blocks including thinking blocks, which Anthropic rejects. Fixed by stripping cache_control from thinking type content blocks after the reverse-translation.

guardrail_information null in spend logs for /v1/messages: multiple root causes tracked down and fixed. The @log_guardrail_information decorator writes to request_data["metadata"] or request_data["litellm_metadata"] depending on which key is present. The spend log is built from litellm_params["metadata"] via merge_litellm_metadata. For /v1/messages, data only has litellm_metadata (no metadata key), and Logging.litellm_params gets reassigned during the request (creating a new dict that diverges from model_call_details["litellm_params"]). Fix: after apply_guardrail returns, explicitly copy the guardrail info from whichever metadata key it was written to, into both logging_obj.litellm_params["metadata"] and logging_obj.model_call_details["litellm_params"]["metadata"].

New files:

  • litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
  • litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py
  • litellm/types/proxy/guardrails/guardrail_hooks/headroom.py
  • tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
  • tests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py

Adds a headroom guardrail that compresses request messages via POST
/v1/compress before they reach the LLM. The guardrail implements
apply_guardrail so it runs on the unified guardrail path; it receives
pre-built structured_messages (OpenAI format) from the translation
layer, calls the headroom compression service, and returns the
compressed messages as structured_messages.

Set x-headroom-bypass: true on the request to skip compression.

Also adds structured_messages write-back support to the OpenAI and
Anthropic translation handlers: when apply_guardrail returns
structured_messages, those are written to data["messages"] directly
(OpenAI) or reverse-translated via anthropic_messages_pt (Anthropic)
instead of falling through to the existing text-patch path. This is a
prerequisite for any guardrail that needs to replace the full message
list rather than patch individual text spans.
@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@codecov

codecov Bot commented Jun 26, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 90.30303% with 16 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
...xy/guardrails/guardrail_hooks/headroom/__init__.py 47.05% 9 Missing ⚠️
...xy/guardrails/guardrail_hooks/headroom/headroom.py 92.94% 6 Missing ⚠️
...ms/anthropic/chat/guardrail_translation/handler.py 94.44% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented Jun 26, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds a Headroom guardrail for compressing request messages before model calls. The main changes are:

  • New headroom guardrail integration using a configured /v1/compress service
  • Structured message write-back for OpenAI and Anthropic guardrail translation paths
  • Guardrail logging metadata propagation into spend-log metadata for passthrough routes
  • Mocked tests for compression, bypass behavior, error handling, logging sync, and structured message write-back

Confidence Score: 5/5

The changes are isolated to the new Headroom guardrail integration, guardrail translation write-back behavior, and metadata propagation for spend logging.

The implementation is covered by focused mocked tests for compression behavior, bypass handling, service error handling, structured message write-back, and logging metadata synchronization.

T-Rex T-Rex Logs

What T-Rex did

  • Ran headroom guardrail tests and confirmed the base cannot import litellm.proxy.guardrails.guardrail_hooks.headroom.headroom before changes, and after changes observed the chat compression service called with original messages and downstream messages replaced by COMPRESSED, with Anthropic compression invoked and standard_logging_guardrail_information present in route data and logging metadata.
  • Validated headroom options behavior across before and after states: before showed an import failure, confirming no Headroom behavior on the base, and after showed bypass/config/model/env scenarios returning 200 OK, no mock request on bypass, expected forwarded models/Authorization headers/API bases, and 502 HTTPException for empty message list and connection refusal.
  • Saved executed guardrail test artifacts and observed both runs completed with STATUS 200 OK; original OpenAI messages were preserved and Anthropic passthrough did not crash with cached system content; thinking_cache_control_present: False; however, the before artifact already showed those outcomes, so it does not prove the requested regression.

View all artifacts

T-Rex Ran code and verified through T-Rex

Reviews (9): Last reviewed commit: "fix(lint): extract _write_back_structure..." | Re-trigger Greptile

Comment thread litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
Comment thread litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
Comment thread litellm/llms/anthropic/chat/guardrail_translation/handler.py

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 45340607f6

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +183 to +187
data["messages"] = anthropic_messages_pt(
messages=guardrailed_structured_messages,
model=model,
llm_provider="anthropic",
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Handle Anthropic system messages before converting back

When an Anthropic request has the normal top-level system field, get_structured_messages() translates it into a leading OpenAI role: system message before Headroom sees it. Passing the returned list directly into anthropic_messages_pt() here stalls such requests because that helper only advances over user/tool/function/assistant roles, so a leading system message leaves its loop index unchanged. With the default Headroom config where system messages are not skipped, any Anthropic request with a system prompt can hang instead of reaching the model; the system block needs to be stripped or restored to data["system"] separately before this conversion.

Useful? React with 👍 / 👎.

Comment on lines +136 to +138
guardrailed_structured_messages = guardrailed_inputs.get("structured_messages")
if guardrailed_structured_messages is not None:
data["messages"] = guardrailed_structured_messages

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve skipped OpenAI messages on structured writeback

If this guardrail is configured with skip_system_message_in_guardrail or skip_tool_message_in_guardrail, the structured_messages sent to Headroom above has those roles filtered out, but this new writeback treats the returned subset as the full conversation. For an OpenAI chat request with a system prompt and Headroom skip_system_message_in_guardrail: true, the system message is deleted from data["messages"], whereas the old text patchback path preserved skipped messages. Merge the compressed messages back into the original conversation or avoid full replacement when skip filters were applied.

Useful? React with 👍 / 👎.

@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

@greptile review

Comment thread litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py Outdated
@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

@greptile review

@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

@greptile review

Comment thread litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

@greptile review

@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

@greptileai review

@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

@greptile review

)

if litellm_logging_obj is not None:
slg_info = (

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is a really complicated if block. What exactly are you trying to do here? is there a simpler way of achieving this?

"standard_logging_guardrail_information"
)
if existing is None:
log_meta["standard_logging_guardrail_information"] = (

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is so nested. can you simplify what within the logging object you're trying to extract ?

isinstance(block, dict)
and block.get("type") == "thinking"
):
block.pop("cache_control", None)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what is the problem you're trying to solve here? please make sure this is a tested flow

@krrish-berri-2 krrish-berri-2 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

Headroom guardrail + structured_messages write-back look good and should ship. The guardrail logging sync block is the main blocker: it fixes a real platform bug at the wrong layer and needs to move before merge.

Required changes (in order)

1. Remove the inline sync block from the anthropic translation handler

Delete the ~35-line block after apply_guardrail that copies standard_logging_guardrail_information into logging_obj.litellm_params / model_call_details. This is a cross-cutting concern, not anthropic-handler logic.

2. Add a shared helper and call it from @log_guardrail_information

Add _sync_guardrail_info_to_logging_obj(request_data, logging_obj) in litellm/integrations/custom_guardrail.py (next to @log_guardrail_information). It should:

  • Read standard_logging_guardrail_information from request_data["metadata"] or request_data["litellm_metadata"] (same precedence as add_standard_logging_guardrail_information_to_request_data)
  • Append/merge into both logging_obj.litellm_params["metadata"] and logging_obj.model_call_details["litellm_params"]["metadata"] when logging_obj is not None
  • Not swallow exceptions silently

Call it at the end of the @log_guardrail_information async/sync wrappers, after _process_response / _process_error, when kwargs.get("logging_obj") is present. That fixes /v1/messages, /v1/responses, and every other LITELLM_METADATA_ROUTES guardrail in one place (headroom, presidio, generic_guardrail_api, etc.), not just anthropic.

3. Add a regression test for the sync helper

New test file e.g. tests/test_litellm/integrations/test_guardrail_logging_sync.py. Model it on tests/test_litellm/proxy/test_proxy_utils.py::test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params. Assert that when guardrail info is written only to request_data["litellm_metadata"] and logging_obj.litellm_params starts empty, the helper copies it into logging_obj.litellm_params["metadata"] so merge_litellm_metadata surfaces it in spend logs.

4. Fix response.json() error handling in headroom.py

Wrap response.json() in try/except and raise 502 with a clear detail. A 200 with HTML/truncated JSON should not become an unhandled 500.

5. Minor cleanups

  • headroom.py: drop redundant headers.get(BYPASS_HEADER.lower())
  • test_headroom.py::test_init_raises_without_api_base: monkeypatch env/secrets so the test is deterministic
  • Revert the whitespace-only diff in litellm_logging.py
  • Squash the debug-print commits before merge

Keep as-is

  • Headroom guardrail implementation and tests
  • structured_messages identity-check write-back (OpenAI + Anthropic handlers)
  • Anthropic system-message strip before anthropic_messages_pt and cache_control strip on thinking blocks

logging_obj=litellm_logging_obj,
)

if litellm_logging_obj is not None:

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove this block; fix at the decorator layer instead.

This copies standard_logging_guardrail_information from data into logging_obj after every guardrail run. Right symptom, wrong location.

Change: Delete lines 166–200 entirely. Add _sync_guardrail_info_to_logging_obj(request_data, logging_obj) in litellm/integrations/custom_guardrail.py and invoke it from @log_guardrail_information when kwargs.get("logging_obj") is present (after _process_response / _process_error). The decorator already has both request_data and logging_obj in scope for apply_guardrail calls.

Why: On LITELLM_METADATA_ROUTES (/v1/messages, etc.) @log_guardrail_information writes to litellm_metadata on data, but spend logs read from logging_obj.litellm_params via merge_litellm_metadata. That gap affects all unified guardrails on those routes, not just headroom.

Also: Replace except (KeyError, AttributeError): pass with no catch or explicit debug logging.

},
)

body: object = response.json()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wrap response.json() and return 502 on decode failure.

If the headroom service returns HTTP 200 with a non-JSON body, response.json() raises and propagates as an unhandled 500.

Change:

try:
    body: object = response.json()
except Exception:
    raise HTTPException(
        status_code=502,
        detail={
            "error": "Headroom compression service returned non-JSON response",
            "body": response.text,
        },
    )

headers = psr.get("headers")
if not _is_str_object_dict(headers):
return False
value = headers.get(BYPASS_HEADER) or headers.get(BYPASS_HEADER.lower())

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Simplify header lookup.

BYPASS_HEADER is already lowercase. The or headers.get(BYPASS_HEADER.lower()) branch is identical.

Change: value = headers.get(BYPASS_HEADER)

assert "empty message list" in str(exc_info.value.detail)


def test_init_raises_without_api_base():

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make this test deterministic in CI.

If HEADROOM_API_BASE is set in the environment, HeadroomGuardrail(api_base=None) resolves from env and the test never raises.

Change:

def test_init_raises_without_api_base(monkeypatch):
    monkeypatch.delenv("HEADROOM_API_BASE", raising=False)
    monkeypatch.setattr(
        "litellm.proxy.guardrails.guardrail_hooks.headroom.headroom.get_secret_str",
        lambda key: None,
    )
    with pytest.raises(ValueError, match="API base URL"):
        HeadroomGuardrail(api_base=None)

…guardrail.py

- Add _sync_guardrail_info_to_logging_obj in custom_guardrail.py; call it from
  both async and sync wrappers in @log_guardrail_information, fixing
  guardrail_information=null in spend logs for all passthrough routes
  (/v1/messages, /v1/responses, etc.) in one place
- Remove the 35-line inline sync block from the anthropic translation handler
- Wrap response.json() in try/except in headroom.py to 502 on HTML/truncated responses
- Drop redundant headers.get(BYPASS_HEADER.lower()) — header key already lowercase
- Add regression tests for _sync_guardrail_info_to_logging_obj
@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

All required changes from the review addressed in caf42dc:

1. Inline sync block removed from anthropic translation handler. The 35-line block that copied standard_logging_guardrail_information into logging_obj is gone from handler.py.

2. Shared helper added and called from @log_guardrail_information. _sync_guardrail_info_to_logging_obj(request_data, logging_obj) now lives in litellm/integrations/custom_guardrail.py next to the decorator. It reads from request_data["metadata"] or request_data["litellm_metadata"] (same precedence as add_standard_logging_guardrail_information_to_request_data) and merges into both logging_obj.litellm_params["metadata"] and model_call_details["litellm_params"]["metadata"]. Both async_wrapper and sync_wrapper call it after _process_response / _process_error (and on the early-return path when a guardrail already recorded its own entry). kwargs.get("logging_obj") is how the obj is obtained. This fixes guardrail_information for headroom, presidio, generic_guardrail_api and every other guardrail on passthrough routes in one place.

3. Regression test added. tests/test_litellm/integrations/test_guardrail_logging_sync.py has 6 tests covering: sync from litellm_metadata, sync from metadata, precedence, no-op when info absent, no-op when logging_obj is None, and the litellm_params reassignment divergence case (the root cause).

4. response.json() wrapped in try/except. A 200 with HTML or truncated JSON now raises HTTPException 502 with a clear error and truncated body.

5. Minor cleanups done. Dropped headers.get(BYPASS_HEADER.lower()) — BYPASS_HEADER is already lowercase. The whitespace-only diff in litellm_logging.py is reverted. Debug-print commits were already cleaned up in 7d99315.

@krrish-berri-2 krrish-berri-2 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two small follow-ups from the earlier review that are still open in caf42dc6 (the refactor itself looks good)

@@ -5959,7 +5959,6 @@ def get_standard_logging_object_payload(
),
standard_built_in_tools_params=standard_built_in_tools_params,
)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Revert this whitespace-only change.

This PR removes a blank line in get_standard_logging_object_payload; it is unrelated to headroom or guardrail logging sync.

Change: restore the blank line so this file has zero diff vs base:

         )
+
         # emit_standard_logging_payload(payload) - Moved to success_handler to prevent double emitting

Keeps the PR scope clean and avoids noisy merge conflicts on an unrelated 6k-line file.

assert "empty message list" in str(exc_info.value.detail)


def test_init_raises_without_api_base():

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make this test deterministic when HEADROOM_API_BASE is set in CI.

HeadroomGuardrail(api_base=None) falls back to get_secret_str("HEADROOM_API_BASE"), so the test can silently pass without asserting anything if that env var is present.

Change:

def test_init_raises_without_api_base(monkeypatch):
    monkeypatch.delenv("HEADROOM_API_BASE", raising=False)
    monkeypatch.setattr(
        "litellm.proxy.guardrails.guardrail_hooks.headroom.headroom.get_secret_str",
        lambda key: None,
    )
    with pytest.raises(ValueError, match="API base URL"):
        HeadroomGuardrail(api_base=None)

@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

@greptile review

1 similar comment
@krrish-berri-2

Copy link
Copy Markdown
Contributor Author

@greptile review

@krrish-berri-2
krrish-berri-2 merged commit 99b1a32 into litellm_internal_staging Jun 27, 2026
122 checks passed
@krrish-berri-2
krrish-berri-2 deleted the litellm_headroom_guardrail branch June 27, 2026 02:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants