Merge 723c3936ef into 8884d87378

headroom/proxy/handlers/anthropic.py

@@ -1144,6 +1144,7 @@ class AnthropicHandlerMixin:
             # point on, `headers` is the upstream-bound copy.
             from headroom.proxy.helpers import (
                 _strip_internal_headers,
+                apply_keep_last_turns,
                 log_outbound_headers,
                 merge_extra_headers,
             )
@@ -1426,6 +1427,25 @@ class AnthropicHandlerMixin:
             else:
                 _hook_ctx = None
 
+            # x-headroom-keep-last-turns: N — trim history before optimization.
+            # Consumed here (after bypass check, after _strip_internal_headers)
+            # so it never leaks upstream.  Fail-open: any malformed value is
+            # silently ignored and the full message list is used instead.
+            _klt_raw = request.headers.get("x-headroom-keep-last-turns", "").strip()
+            if _klt_raw and not _bypass:
+                try:
+                    _klt = int(_klt_raw)
+                    messages, _klt_dropped = apply_keep_last_turns(messages, _klt)
+                    if _klt_dropped:
+                        logger.info(
+                            "[%s] keep-last-turns=%d: dropped %d leading messages",
+                            request_id,
+                            _klt,
+                            _klt_dropped,
+                        )
+                except ValueError:
+                    pass  # malformed value — never break the request
+
             # Apply optimization
             transforms_applied = []
             pipeline_timing: dict[str, float] = {}

headroom/proxy/handlers/openai.py

@@ -3420,6 +3420,7 @@ class OpenAIHandlerMixin:
         # upstream-bound copy.
         from headroom.proxy.helpers import (
             _strip_internal_headers,
+            apply_keep_last_turns,
             log_outbound_headers,
             merge_extra_headers,
         )
@@ -3628,6 +3629,25 @@ class OpenAIHandlerMixin:
             except Exception as e:
                 logger.debug(f"[{request_id}] Hook error: {e}")
 
+        # x-headroom-keep-last-turns: N — trim history before optimization.
+        # Consumed here (after bypass check, after _strip_internal_headers)
+        # so it never leaks upstream.  Fail-open: any malformed value is
+        # silently ignored and the full message list is used instead.
+        _klt_raw = request.headers.get("x-headroom-keep-last-turns", "").strip()
+        if _klt_raw and not _bypass:
+            try:
+                _klt = int(_klt_raw)
+                messages, _klt_dropped = apply_keep_last_turns(messages, _klt)
+                if _klt_dropped:
+                    logger.info(
+                        "[%s] keep-last-turns=%d: dropped %d leading messages",
+                        request_id,
+                        _klt,
+                        _klt_dropped,
+                    )
+            except ValueError:
+                pass  # malformed value — never break the request
+
         # Optimization
         transforms_applied = []
         pipeline_timing: dict[str, float] = {}

headroom/proxy/helpers.py

@@ -3446,3 +3446,30 @@ def inject_tool_search_deferral_openai(
     if deferred == 0:
         return tools  # nothing to defer → don't perturb the request / cache prefix
     return out
+
+
+def apply_keep_last_turns(
+    messages: list[dict[str, Any]],
+    n: int,
+) -> tuple[list[dict[str, Any]], int]:
+    """Trim ``messages`` to the last *n* conversation turns.
+
+    A "turn" is one user+assistant pair.  The trailing user message is
+    always kept regardless of *n* (it is the current request).
+
+    Returns ``(trimmed_messages, n_dropped)`` — the caller can log
+    *n_dropped* and append ``keep_last_turns:{n}:{n_dropped}_dropped``
+    to ``transforms_applied``.  When nothing is dropped (n_dropped == 0)
+    the original list is returned unchanged so callers can detect a no-op
+    with an identity check.
+
+    Invariants:
+    - n < 0 is treated as no-op (invalid, never trim).
+    - The result always contains at least one message (the final user msg).
+    """
+    if n < 0 or not messages:
+        return messages, 0
+    tail = max(0, len(messages) - 1 - n * 2)
+    if tail == 0:
+        return messages, 0
+    return messages[tail:], tail

tests/test_keep_last_turns.py

@@ -0,0 +1,147 @@
+"""Tests for the x-headroom-keep-last-turns feature (issue #2858).
+
+Covers the apply_keep_last_turns() helper that both anthropic.py and
+openai.py call before the optimized_messages assignment.
+"""
+
+from __future__ import annotations
+
+from headroom.proxy.helpers import apply_keep_last_turns
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _turns(n: int) -> list[dict]:
+    """Build a plausible conversation: n user+assistant pairs + trailing user."""
+    msgs: list[dict] = []
+    for i in range(n):
+        msgs.append({"role": "user", "content": f"q{i}"})
+        msgs.append({"role": "assistant", "content": f"a{i}"})
+    msgs.append({"role": "user", "content": "final question"})
+    return msgs
+
+
+# ---------------------------------------------------------------------------
+# No-op cases
+# ---------------------------------------------------------------------------
+
+
+def test_empty_messages_returns_unchanged():
+    result, dropped = apply_keep_last_turns([], n=0)
+    assert result == []
+    assert dropped == 0
+
+
+def test_n_negative_is_noop():
+    msgs = _turns(3)
+    result, dropped = apply_keep_last_turns(msgs, n=-1)
+    assert result is msgs
+    assert dropped == 0
+
+
+def test_n_larger_than_history_is_noop():
+    """When n >= existing turns no messages are dropped."""
+    msgs = _turns(2)  # 5 messages total
+    result, dropped = apply_keep_last_turns(msgs, n=10)
+    assert result is msgs
+    assert dropped == 0
+
+
+def test_exact_turns_match_is_noop():
+    """Asking for exactly the number of turns present keeps everything."""
+    msgs = _turns(3)  # 3 turns → 7 messages
+    result, dropped = apply_keep_last_turns(msgs, n=3)
+    assert result is msgs
+    assert dropped == 0
+
+
+# ---------------------------------------------------------------------------
+# Trimming cases
+# ---------------------------------------------------------------------------
+
+
+def test_n0_keeps_only_final_user_message():
+    """n=0 means 'no prior turns' → only the trailing user message survives."""
+    msgs = _turns(3)  # 7 messages
+    result, dropped = apply_keep_last_turns(msgs, n=0)
+    assert len(result) == 1
+    assert result[0]["content"] == "final question"
+    assert dropped == 6
+
+
+def test_n1_keeps_one_prior_turn_plus_current():
+    """n=1 keeps the most recent user+assistant pair plus the current user message."""
+    msgs = _turns(3)  # q0 a0 q1 a1 q2 a2 q_final (7 msgs)
+    result, dropped = apply_keep_last_turns(msgs, n=1)
+    # Keeps: a2, q_final = indices 5, 6 → but formula gives tail = 7-1-1*2=4
+    # messages[4:] = q2 a2 q_final  (3 messages)
+    assert len(result) == 3
+    assert result[-1]["content"] == "final question"
+    assert dropped == 4
+
+
+def test_n2_keeps_two_prior_turns_plus_current():
+    msgs = _turns(4)  # 9 messages
+    result, dropped = apply_keep_last_turns(msgs, n=2)
+    # tail = 9-1-2*2 = 4; messages[4:] = 5 messages
+    assert len(result) == 5
+    assert result[-1]["content"] == "final question"
+    assert dropped == 4
+
+
+def test_trailing_user_message_always_present():
+    """The current (final) user message is never dropped regardless of n."""
+    for n in range(5):
+        msgs = _turns(5)
+        result, _ = apply_keep_last_turns(msgs, n=n)
+        assert result[-1] == {"role": "user", "content": "final question"}, f"n={n}"
+
+
+def test_total_dropped_plus_kept_equals_original():
+    msgs = _turns(5)  # 11 messages
+    for n in range(6):
+        result, dropped = apply_keep_last_turns(msgs, n=n)
+        assert len(result) + dropped == len(msgs), f"n={n}"
+
+
+# ---------------------------------------------------------------------------
+# Single-message conversation
+# ---------------------------------------------------------------------------
+
+
+def test_single_message_conversation_n0():
+    """A single-message conversation (no history) is never trimmed."""
+    msgs = [{"role": "user", "content": "hello"}]
+    result, dropped = apply_keep_last_turns(msgs, n=0)
+    assert result is msgs
+    assert dropped == 0
+
+
+def test_single_message_conversation_large_n():
+    msgs = [{"role": "user", "content": "hello"}]
+    result, dropped = apply_keep_last_turns(msgs, n=100)
+    assert result is msgs
+    assert dropped == 0
+
+
+# ---------------------------------------------------------------------------
+# Return-value contract
+# ---------------------------------------------------------------------------
+
+
+def test_no_trim_returns_original_object():
+    """When nothing is dropped the original list is returned (identity, not copy)."""
+    msgs = _turns(2)
+    result, dropped = apply_keep_last_turns(msgs, n=100)
+    assert result is msgs
+    assert dropped == 0
+
+
+def test_trim_returns_new_slice():
+    """When trimming, a new slice (not the original object) is returned."""
+    msgs = _turns(3)
+    result, dropped = apply_keep_last_turns(msgs, n=0)
+    assert result is not msgs
+    assert dropped > 0