diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 5f5e4a510..702764c05 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/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] = {} diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 42ac1a8c0..2780a638e 100644 --- a/headroom/proxy/handlers/openai.py +++ b/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] = {} diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 8d2a2f43c..2f931274f 100644 --- a/headroom/proxy/helpers.py +++ b/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 diff --git a/tests/test_keep_last_turns.py b/tests/test_keep_last_turns.py new file mode 100644 index 000000000..cda4bb036 --- /dev/null +++ b/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