From 3bb4d7e5ca5f9518a495f2e2f653fe0bf6711ed8 Mon Sep 17 00:00:00 2001
From: SwiftWinds <12981958+SwiftWinds@users.noreply.github.com>
Date: Fri, 3 Apr 2026 12:23:44 -0700
Subject: [PATCH] feat(responses): add universal server-side context compaction

When input tokens exceed a configurable threshold, the Responses API
now automatically summarizes conversation history before sending it to
the LLM. This reduces token usage for long conversations while
preserving context continuity across all providers.

- Add compaction module with summarization logic and threshold checks
- Integrate compaction into the async response handler, prepending a
  compaction output item when summarization occurs
- Support compaction input items in the transformation layer so
  follow-up turns can consume prior summaries
- Forward custom_llm_provider through compaction to fix provider
  resolution for models using provider-prefixed names
- Add mock unit test validating the two-call flow (summarize then
  complete) and response structure

Made-with: Cursor
---
 litellm/responses/compaction.py               | 124 ++++++++++++++++++
 .../handler.py                                |  30 ++++-
 .../transformation.py                         |  10 ++
 .../test_google_ai_studio_responses_api.py    | 115 ++++++++++++++++
 4 files changed, 278 insertions(+), 1 deletion(-)
 create mode 100644 litellm/responses/compaction.py

diff --git a/litellm/responses/compaction.py b/litellm/responses/compaction.py
new file mode 100644
index 00000000000..26c84197467
--- /dev/null
+++ b/litellm/responses/compaction.py
@@ -0,0 +1,124 @@
+"""
+Universal server-side context compaction for the Responses API.
+
+Summarizes conversation history when input tokens exceed a configured threshold,
+using the same model the caller is already using. Works across all providers.
+"""
+
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+import litellm
+from litellm._logging import verbose_logger
+
+SUMMARIZATION_SYSTEM_PROMPT = (
+    "You have written a partial transcript for the initial task above. "
+    "Please write a summary of the transcript. The purpose of this summary is "
+    "to provide continuity so you can continue to make progress towards solving "
+    "the task in a future context, where the raw history above may not be "
+    "accessible and will be replaced with this summary. Write down anything "
+    "that would be helpful, including the state, next steps, learnings etc. "
+    "You must wrap your summary in a <summary></summary> block."
+)
+
+MIN_COMPACT_THRESHOLD = 1000
+
+
+def _get_compact_threshold(context_management: List[Dict[str, Any]]) -> Optional[int]:
+    """Extract the compact_threshold from a context_management list, if present."""
+    for entry in context_management:
+        if not isinstance(entry, dict):
+            continue
+        if entry.get("type") == "compaction":
+            threshold = entry.get("compact_threshold")
+            if threshold is not None:
+                return max(int(threshold), MIN_COMPACT_THRESHOLD)
+            return None
+    return None
+
+
+def _extract_summary(text: str) -> str:
+    """Pull content out of <summary>...</summary> tags, falling back to the full text."""
+    match = re.search(r"<summary>(.*?)</summary>", text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    return text.strip()
+
+
+def _serialize_messages_for_summary(messages: List[Dict[str, Any]]) -> str:
+    """Render a message list into a readable transcript for the summarizer."""
+    parts: List[str] = []
+    for msg in messages:
+        role = msg.get("role", "unknown")
+        content = msg.get("content", "")
+        if isinstance(content, list):
+            text_parts = []
+            for block in content:
+                if isinstance(block, dict):
+                    text_parts.append(block.get("text", str(block)))
+                else:
+                    text_parts.append(str(block))
+            content = "\n".join(text_parts)
+        parts.append(f"[{role}]: {content}")
+    return "\n\n".join(parts)
+
+
+async def maybe_compact_context(
+    messages: List[Dict[str, Any]],
+    model: str,
+    context_management: List[Dict[str, Any]],
+    custom_llm_provider: Optional[str] = None,
+) -> Tuple[List[Dict[str, Any]], Optional[str]]:
+    """
+    Check whether compaction should trigger and, if so, summarize ALL messages.
+
+    Returns:
+        (messages, summary_text):
+        - If compaction triggered: messages is [{role: "user", content: summary}],
+          summary_text is the raw summary string.
+        - If not triggered: messages is unchanged, summary_text is None.
+    """
+    threshold = _get_compact_threshold(context_management)
+    if threshold is None:
+        return messages, None
+
+    token_count = litellm.token_counter(model=model, messages=messages)
+    verbose_logger.debug(
+        "compaction: token_count=%d, threshold=%d", token_count, threshold
+    )
+
+    if token_count <= threshold:
+        return messages, None
+
+    verbose_logger.info(
+        "compaction: triggering summarization (tokens=%d > threshold=%d)",
+        token_count,
+        threshold,
+    )
+
+    transcript = _serialize_messages_for_summary(messages)
+
+    summarization_messages: List[Dict[str, Any]] = [
+        {"role": "user", "content": transcript},
+        {"role": "user", "content": SUMMARIZATION_SYSTEM_PROMPT},
+    ]
+
+    acompletion_kwargs: Dict[str, Any] = {
+        "model": model,
+        "messages": summarization_messages,
+    }
+    if custom_llm_provider is not None:
+        acompletion_kwargs["custom_llm_provider"] = custom_llm_provider
+
+    summary_response = await litellm.acompletion(**acompletion_kwargs)
+
+    summary_text_raw = summary_response.choices[0].message.content or ""
+    summary_text = _extract_summary(summary_text_raw)
+
+    verbose_logger.debug("compaction: summary length=%d chars", len(summary_text))
+
+    compacted_messages: List[Dict[str, Any]] = [
+        {"role": "user", "content": summary_text},
+    ]
+
+    return compacted_messages, summary_text
diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py
index 5faa8b587c9..208ae26b880 100644
--- a/litellm/responses/litellm_completion_transformation/handler.py
+++ b/litellm/responses/litellm_completion_transformation/handler.py
@@ -2,9 +2,10 @@
 Handler for transforming responses api requests to litellm.completion requests
 """
 
-from typing import Any, Coroutine, Dict, Optional, Union
+from typing import Any, Coroutine, Dict, List, Optional, Union
 
 import litellm
+from litellm.responses.compaction import maybe_compact_context
 from litellm.responses.litellm_completion_transformation.streaming_iterator import (
     LiteLLMCompletionStreamingIterator,
 )
@@ -109,6 +110,17 @@ class LiteLLMCompletionTransformationHandler:
         acompletion_args.update(kwargs)
         acompletion_args.update(litellm_completion_request)
 
+        context_management = acompletion_args.pop("context_management", None)
+        summary_text: Optional[str] = None
+        if context_management:
+            compacted_messages, summary_text = await maybe_compact_context(
+                messages=acompletion_args["messages"],
+                model=acompletion_args["model"],
+                context_management=context_management,
+                custom_llm_provider=acompletion_args.get("custom_llm_provider"),
+            )
+            acompletion_args["messages"] = compacted_messages
+
         litellm_completion_response: Union[
             ModelResponse, litellm.CustomStreamWrapper
         ] = await litellm.acompletion(
@@ -122,6 +134,11 @@ class LiteLLMCompletionTransformationHandler:
                 responses_api_request=responses_api_request,
             )
 
+            if summary_text is not None:
+                responses_api_response.output = _prepend_compaction_output(
+                    summary_text, responses_api_response.output
+                )
+
             return responses_api_response
 
         elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper):
@@ -138,3 +155,14 @@ class LiteLLMCompletionTransformationHandler:
         raise ValueError(
             f"Unexpected response type: {type(litellm_completion_response)}"
         )
+
+
+def _prepend_compaction_output(
+    summary_text: str, existing_output: List[Any]
+) -> List[Any]:
+    """Prepend a compaction output item before the existing output items."""
+    compaction_item = {
+        "type": "compaction",
+        "content": summary_text,
+    }
+    return [compaction_item] + list(existing_output)
diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py
index 9075373f1cf..9b85e8eb63a 100644
--- a/litellm/responses/litellm_completion_transformation/transformation.py
+++ b/litellm/responses/litellm_completion_transformation/transformation.py
@@ -376,6 +376,16 @@ class LiteLLMCompletionResponsesConfig:
         elif isinstance(input, list):
             existing_tool_call_ids: Set[str] = set()
             for _input in input:
+                if isinstance(_input, dict) and _input.get("type") == "compaction":
+                    messages.clear()
+                    compaction_content = _input.get("content", "")
+                    messages.append(
+                        ChatCompletionSystemMessage(
+                            role="system", content=compaction_content
+                        )
+                    )
+                    continue
+
                 chat_completion_messages = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
                     input_item=_input
                 )
diff --git a/tests/llm_responses_api_testing/test_google_ai_studio_responses_api.py b/tests/llm_responses_api_testing/test_google_ai_studio_responses_api.py
index 13801fce7d1..bbebd44b8ba 100644
--- a/tests/llm_responses_api_testing/test_google_ai_studio_responses_api.py
+++ b/tests/llm_responses_api_testing/test_google_ai_studio_responses_api.py
@@ -249,6 +249,121 @@ async def test_gemini_3_responses_api_streaming_with_thought_signatures():
     print(f"✅ Collected {len(chunks)} streaming chunks")
 
 
+@pytest.mark.asyncio
+async def test_mock_google_ai_studio_compaction():
+    """
+    Test that universal compaction works for Google AI Studio via the Responses API.
+
+    Sends a very large input ("cats " * 100_000) with a low compact_threshold (50_000).
+    Mocks litellm.acompletion so no real API call is made:
+      - 1st call: summarization (triggered by compaction)
+      - 2nd call: the actual completion using ONLY the summary
+    Validates the response output has a compaction item followed by a text item.
+    """
+    request_model = "gemini/gemini-2.5-flash"
+    large_input = "cats " * 100_000
+
+    summary_response = litellm.ModelResponse(
+        id="summary-id",
+        created=1000000000,
+        model=request_model,
+        object="chat.completion",
+        choices=[
+            litellm.utils.Choices(
+                index=0,
+                message=litellm.utils.Message(
+                    role="assistant",
+                    content="<summary>The user repeated the word cats many times.</summary>",
+                ),
+                finish_reason="stop",
+            )
+        ],
+    )
+
+    final_response = litellm.ModelResponse(
+        id="final-id",
+        created=1000000001,
+        model=request_model,
+        object="chat.completion",
+        choices=[
+            litellm.utils.Choices(
+                index=0,
+                message=litellm.utils.Message(
+                    role="assistant",
+                    content="Based on the summary, you were talking about cats.",
+                ),
+                finish_reason="stop",
+            )
+        ],
+    )
+
+    call_count = 0
+
+    async def mock_acompletion(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            return summary_response
+        return final_response
+
+    with patch("litellm.acompletion", new_callable=AsyncMock) as mock_ac:
+        mock_ac.side_effect = mock_acompletion
+
+        response = await litellm.aresponses(
+            model=request_model,
+            input=large_input,
+            context_management=[
+                {"type": "compaction", "compact_threshold": 50000}
+            ],
+        )
+
+        assert call_count == 2, f"Expected 2 acompletion calls, got {call_count}"
+
+        # 1st call: summarization — messages should contain the large input
+        first_call_kwargs = mock_ac.call_args_list[0][1]
+        first_msgs = first_call_kwargs.get("messages", [])
+        assert any(
+            "cats" in str(m.get("content", "")) for m in first_msgs
+        ), "Summarization call should contain the original input"
+
+        # 2nd call: actual completion — messages should contain ONLY the summary
+        second_call_kwargs = mock_ac.call_args_list[1][1]
+        second_msgs = second_call_kwargs.get("messages", [])
+        assert len(second_msgs) == 1, (
+            f"After compaction, completion should receive exactly 1 message (the summary), "
+            f"got {len(second_msgs)}"
+        )
+        assert "cats many times" in str(second_msgs[0].get("content", "")), (
+            "Completion call should see the extracted summary, not the original input"
+        )
+
+        # Validate response structure
+        from litellm.types.llms.openai import ResponsesAPIResponse
+
+        assert isinstance(response, ResponsesAPIResponse)
+        assert len(response.output) >= 2, (
+            f"Response output should have at least 2 items (compaction + text), "
+            f"got {len(response.output)}"
+        )
+
+        # First output item should be the compaction block
+        compaction_item = response.output[0]
+        if isinstance(compaction_item, dict):
+            assert compaction_item["type"] == "compaction"
+            assert "cats many times" in compaction_item["content"]
+        else:
+            assert getattr(compaction_item, "type", None) == "compaction"
+
+        # Second output item should be the text response
+        text_item = response.output[1]
+        if isinstance(text_item, dict):
+            assert text_item.get("type") == "message"
+        else:
+            assert getattr(text_item, "type", None) == "message"
+
+        print("compaction test passed: response output =", json.dumps(response.output, indent=2, default=str))
+
+
 class TestGoogleAIStudioResponsesAPITest(BaseResponsesAPITest):
     def get_base_completion_call_args(self):
         #litellm._turn_on_debug()