From 88e41b65a15e2e44e1872e0a9171f6c5890dc4c7 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Fri, 10 Jul 2026 13:43:04 +0000 Subject: [PATCH] Extract request log redaction policy (#1968) ## Description Extracts the pure image-base64 request-log redaction decision/transform logic from `request_logger.py` into a dedicated policy module. `RequestLogger` remains the owner of the Prometheus-facing redaction counter and existing request_logger constants remain available for compatibility. Closes # ## Type of Change - [ ] 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) - [ ] Documentation update - [ ] Performance improvement - [x] Code refactoring (no functional changes) ## Changes Made - Added `headroom.proxy.request_log_redaction_policy` with a pure `RedactionResult` outcome. - Kept global redaction metrics/counter side effects in `request_logger.py`. - Added direct policy tests for count reporting, nested image paths, and data URL threshold behavior. - Carried forward the LiteLLM callback compatibility shim needed for current mypy on `main`. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text python -m pytest tests\test_request_log_redaction_policy.py tests\test_image_log_redaction.py 20 passed in 0.31s python -m ruff check . All checks passed! python -m ruff format --check . 1095 files already formatted python -m mypy headroom --ignore-missing-imports Success: no issues found in 409 source files gitleaks protect --staged --no-banner --redact no leaks found ``` ## Real Behavior Proof - Environment: Windows, Python 3.13.13, branch `jd/architecture-slice-23`. - Exact command / steps: ran targeted request-log redaction tests, ruff, ruff format check, mypy, and staged gitleaks scan. - Observed result: redaction behavior remains covered through existing logger tests and new pure policy tests; local lint/type/security checks pass. - Not tested: full proxy runtime; this slice only moves pure redaction policy and keeps the logger entry point intact. ## 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 - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A ## Additional Notes Documentation and changelog updates are N/A for this internal architecture-only refactor. The push reported existing default-branch Dependabot alerts; no staged secret leaks were found for this PR. --- headroom/integrations/litellm_callback.py | 15 ++- .../proxy/request_log_redaction_policy.py | 117 ++++++++++++++++ headroom/proxy/request_logger.py | 127 ++---------------- tests/test_request_log_redaction_policy.py | 51 +++++++ 4 files changed, 193 insertions(+), 117 deletions(-) create mode 100644 headroom/proxy/request_log_redaction_policy.py create mode 100644 tests/test_request_log_redaction_policy.py diff --git a/headroom/integrations/litellm_callback.py b/headroom/integrations/litellm_callback.py index 4c2e4e271..1cf152c59 100644 --- a/headroom/integrations/litellm_callback.py +++ b/headroom/integrations/litellm_callback.py @@ -92,11 +92,18 @@ class HeadroomCallback(_CustomLogger): async def async_pre_call_hook( self, - user_api_key: str, - data: dict[str, Any], - call_type: str, - ) -> dict[str, Any]: + user_api_key_dict: Any = None, + cache: Any = None, + data: dict[Any, Any] | None = None, + call_type: str = "", + *_args: Any, + **_kwargs: Any, + ) -> dict[Any, Any] | None: """Called by LiteLLM before each API call. Compresses messages.""" + if isinstance(cache, dict) and isinstance(data, str): + data, call_type = cache, data + if data is None: + return None if call_type not in ("completion", "acompletion"): return data diff --git a/headroom/proxy/request_log_redaction_policy.py b/headroom/proxy/request_log_redaction_policy.py new file mode 100644 index 000000000..90f8661b6 --- /dev/null +++ b/headroom/proxy/request_log_redaction_policy.py @@ -0,0 +1,117 @@ +"""Pure request-log redaction policy for image-bearing payloads.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +# Phase G PR-G3 - base64 redaction threshold (P4-45). +# +# Anthropic image blocks carry base64-encoded JPEGs/PNGs in +# ``source.data``; OpenAI's vision shape carries them in +# ``image_url.url`` as a ``data:image/...;base64,`` URL. +# The threshold gates "real image payload" against short base64 +# strings (which can appear in arguments, signatures, etc.). +IMAGE_BASE64_REDACT_THRESHOLD_BYTES = 1024 + +# Phase G PR-G3 - replacement-marker format. Operators can grep the +# JSONL for `` bool: + """Return True if ``value`` is an over-threshold image data URL. + + Per M2 remediation the prior bare-base64 density heuristic over-fired on + non-image content. This helper only recognizes explicit image data URLs; + image-bearing JSON-path eligibility is handled by the recursive policy. + """ + if not isinstance(value, str): + return False + if len(value) < IMAGE_BASE64_REDACT_THRESHOLD_BYTES: + return False + return value.startswith(_DATA_IMAGE_URL_PREFIX) + + +def redact_image_base64_value(payload: Any) -> RedactionResult: + """Return ``payload`` with over-threshold image strings redacted.""" + return _redact_value(payload, in_image_path=False) + + +def _redact_value(value: Any, *, in_image_path: bool) -> RedactionResult: + if isinstance(value, str): + should_redact = is_base64_image_payload(value) or ( + in_image_path and len(value) >= IMAGE_BASE64_REDACT_THRESHOLD_BYTES + ) + if not should_redact: + return RedactionResult(value=value, redactions=0) + + byte_len = len(value.encode("utf-8")) + return RedactionResult( + value=IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=byte_len), + redactions=1, + ) + + if isinstance(value, Mapping): + redactions = 0 + redacted: dict[Any, Any] = {} + for key, item in value.items(): + item_result = _redact_value( + item, + in_image_path=(key in IMAGE_BEARING_FIELD_NAMES), + ) + redacted[key] = item_result.value + redactions += item_result.redactions + return RedactionResult(value=redacted, redactions=redactions) + + if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): + redactions = 0 + redacted_items: list[Any] = [] + for item in value: + item_result = _redact_value(item, in_image_path=in_image_path) + redacted_items.append(item_result.value) + redactions += item_result.redactions + return RedactionResult(value=redacted_items, redactions=redactions) + + return RedactionResult(value=value, redactions=0) diff --git a/headroom/proxy/request_logger.py b/headroom/proxy/request_logger.py index 21ea155a8..2f5ad7cda 100644 --- a/headroom/proxy/request_logger.py +++ b/headroom/proxy/request_logger.py @@ -25,7 +25,6 @@ import json import logging import sys from collections import deque -from collections.abc import Mapping, Sequence from dataclasses import asdict from pathlib import Path from threading import Lock @@ -34,55 +33,17 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from ..memory.tracker import ComponentStats +from headroom.proxy import request_log_redaction_policy from headroom.proxy.models import RequestLog -logger = logging.getLogger(__name__) - -# Phase G PR-G3 — base64 redaction threshold (P4-45). -# -# Anthropic image blocks carry base64-encoded JPEGs/PNGs in -# ``source.data``; OpenAI's vision shape carries them in -# ``image_url.url`` as a ``data:image/...;base64,`` URL. -# The threshold gates "real image payload" against short base64 -# strings (which can appear in arguments, signatures, etc.). -IMAGE_BASE64_REDACT_THRESHOLD_BYTES = 1024 - -# Phase G PR-G3 — replacement-marker format. Operators can grep the -# JSONL for `` int: return _redactions_total -def _is_base64_image_payload(value: str) -> bool: - """Return True if ``value`` is an over-threshold base64 image. - - Per M2 remediation the prior bare-base64 density heuristic - over-fired on non-image content (encrypted blobs, signed - tokens, minified JSON, tool outputs). We now only consider a - string an image payload when EITHER: - - 1. It starts with ``data:image/`` (an explicit data URL), - OR - 2. The caller has already established the string lives inside - an image-bearing JSON path (see ``IMAGE_BEARING_FIELD_NAMES``) - AND the string itself is over the byte threshold. - - Case (2) is decided by the caller (``_redact_value``) which - threads ``in_image_path`` through the recursion; this helper - handles case (1) on its own. - """ - if not isinstance(value, str): - return False - if len(value) < IMAGE_BASE64_REDACT_THRESHOLD_BYTES: - return False - return value.startswith(_DATA_IMAGE_URL_PREFIX) - - -def _redact_value(value: Any, *, in_image_path: bool = False) -> Any: - """Recursively redact base64-image payloads in a JSON-ish value. - - Returns a new structure with any over-threshold base64 string - replaced by the placeholder. Non-string, non-container values - pass through unchanged. - - ``in_image_path`` is True when the caller reached this value - via one of the ``IMAGE_BEARING_FIELD_NAMES`` keys; once inside - an image-bearing field, any over-threshold string is treated - as an image payload (M2: prevents redaction of unrelated - base64-shaped content outside known image fields). - """ - global _redactions_total - if isinstance(value, str): - # Always-redact: explicit data URL, regardless of path. - # Also redact when the caller signalled image-bearing path - # AND the string is over threshold (no density check — the - # path tells us it's an image). - should_redact = _is_base64_image_payload(value) or ( - in_image_path and len(value) >= IMAGE_BASE64_REDACT_THRESHOLD_BYTES - ) - if should_redact: - with _redactions_lock: - _redactions_total += 1 - byte_len = len(value.encode("utf-8")) - return IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=byte_len) - return value - if isinstance(value, Mapping): - return { - k: _redact_value( - v, - in_image_path=(k in IMAGE_BEARING_FIELD_NAMES), - ) - for k, v in value.items() - } - if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): - return [_redact_value(item, in_image_path=in_image_path) for item in value] - return value - - def redact_image_base64(payload: Any) -> Any: """Public entry point for base64-image redaction. @@ -178,7 +73,13 @@ def redact_image_base64(payload: Any) -> Any: over-threshold base64 string with a size-only placeholder. Idempotent — applying twice yields the same structure. """ - return _redact_value(payload, in_image_path=False) + global _redactions_total + + result = request_log_redaction_policy.redact_image_base64_value(payload) + if result.redactions: + with _redactions_lock: + _redactions_total += result.redactions + return result.value class RequestLogger: diff --git a/tests/test_request_log_redaction_policy.py b/tests/test_request_log_redaction_policy.py new file mode 100644 index 000000000..f22095591 --- /dev/null +++ b/tests/test_request_log_redaction_policy.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from headroom.proxy.request_log_redaction_policy import ( + IMAGE_BASE64_REDACT_THRESHOLD_BYTES, + IMAGE_BASE64_REPLACEMENT_TEMPLATE, + is_base64_image_payload, + redact_image_base64_value, +) + + +def test_policy_reports_redaction_count_without_side_effects() -> None: + image_payload = "x" * IMAGE_BASE64_REDACT_THRESHOLD_BYTES + non_image_payload = "y" * IMAGE_BASE64_REDACT_THRESHOLD_BYTES + + result = redact_image_base64_value( + { + "source": {"data": image_payload}, + "signature": non_image_payload, + } + ) + + assert result.redactions == 1 + assert result.value["source"]["data"] == IMAGE_BASE64_REPLACEMENT_TEMPLATE.format( + n=len(image_payload) + ) + assert result.value["signature"] == non_image_payload + + +def test_policy_counts_nested_list_redactions() -> None: + data_url = "data:image/png;base64," + ("A" * IMAGE_BASE64_REDACT_THRESHOLD_BYTES) + direct_payload = "B" * IMAGE_BASE64_REDACT_THRESHOLD_BYTES + + result = redact_image_base64_value( + [{"content": [{"image_url": {"url": data_url}}, {"image": direct_payload}]}] + ) + + assert result.redactions == 2 + assert result.value[0]["content"][0]["image_url"]["url"] == ( + IMAGE_BASE64_REPLACEMENT_TEMPLATE.format(n=len(data_url)) + ) + assert result.value[0]["content"][1]["image"] == IMAGE_BASE64_REPLACEMENT_TEMPLATE.format( + n=len(direct_payload) + ) + + +def test_explicit_image_data_url_requires_threshold() -> None: + short_data_url = "data:image/png;base64,abc" + long_data_url = "data:image/png;base64," + ("A" * IMAGE_BASE64_REDACT_THRESHOLD_BYTES) + + assert not is_base64_image_payload(short_data_url) + assert is_base64_image_payload(long_data_url)