From 94af51f621de55c1f8ebbe0dbc6c2a54ad4fd0ed Mon Sep 17 00:00:00 2001
From: teknium1 <teknium1@gmail.com>
Date: Sat, 14 Mar 2026 11:03:25 -0700
Subject: [PATCH] fix: harden trajectory compressor summary content handling

Normalize summary-model content before stripping so empty or non-string
responses do not trigger retry/fallback paths. Adds sync and async
regression tests for None content.
---
 tests/test_trajectory_compressor.py | 34 ++++++++++++++++++++++++++++-
 trajectory_compressor.py            | 33 ++++++++++++++++------------
 2 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/tests/test_trajectory_compressor.py b/tests/test_trajectory_compressor.py
index 75fbd5a29..c95a3af94 100644
--- a/tests/test_trajectory_compressor.py
+++ b/tests/test_trajectory_compressor.py
@@ -1,7 +1,10 @@
 """Tests for trajectory_compressor.py — config, metrics, and compression logic."""
 
 import json
-from unittest.mock import patch, MagicMock
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch, MagicMock
+
+import pytest
 
 from trajectory_compressor import (
     CompressionConfig,
@@ -384,3 +387,32 @@ class TestTokenCounting:
         tc.tokenizer.encode = MagicMock(side_effect=Exception("fail"))
         # Should fallback to len(text) // 4
         assert tc.count_tokens("12345678") == 2
+
+
+class TestGenerateSummary:
+    def test_generate_summary_handles_none_content(self):
+        tc = _make_compressor()
+        tc.client = MagicMock()
+        tc.client.chat.completions.create.return_value = SimpleNamespace(
+            choices=[SimpleNamespace(message=SimpleNamespace(content=None))]
+        )
+        metrics = TrajectoryMetrics()
+
+        summary = tc._generate_summary("Turn content", metrics)
+
+        assert summary == "[CONTEXT SUMMARY]:"
+
+    @pytest.mark.asyncio
+    async def test_generate_summary_async_handles_none_content(self):
+        tc = _make_compressor()
+        tc.async_client = MagicMock()
+        tc.async_client.chat.completions.create = AsyncMock(
+            return_value=SimpleNamespace(
+                choices=[SimpleNamespace(message=SimpleNamespace(content=None))]
+            )
+        )
+        metrics = TrajectoryMetrics()
+
+        summary = await tc._generate_summary_async("Turn content", metrics)
+
+        assert summary == "[CONTEXT SUMMARY]:"
diff --git a/trajectory_compressor.py b/trajectory_compressor.py
index ef81d6e27..1bfed6bfc 100644
--- a/trajectory_compressor.py
+++ b/trajectory_compressor.py
@@ -495,6 +495,21 @@ class TrajectoryCompressor:
             parts.append(f"[Turn {i} - {role.upper()}]:\n{value}")
         
         return "\n\n".join(parts)
+
+    @staticmethod
+    def _coerce_summary_content(content: Any) -> str:
+        """Normalize summary-model output to a safe string."""
+        if not isinstance(content, str):
+            content = str(content) if content else ""
+        return content.strip()
+
+    @staticmethod
+    def _ensure_summary_prefix(summary: str) -> str:
+        """Normalize summary text to include the expected prefix exactly once."""
+        text = (summary or "").strip()
+        if text.startswith("[CONTEXT SUMMARY]:"):
+            return text
+        return "[CONTEXT SUMMARY]:" if not text else f"[CONTEXT SUMMARY]: {text}"
     
     def _generate_summary(self, content: str, metrics: TrajectoryMetrics) -> str:
         """
@@ -545,13 +560,8 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix."""
                         max_tokens=self.config.summary_target_tokens * 2,
                     )
                 
-                summary = response.choices[0].message.content.strip()
-                
-                # Ensure it starts with the prefix
-                if not summary.startswith("[CONTEXT SUMMARY]:"):
-                    summary = "[CONTEXT SUMMARY]: " + summary
-                
-                return summary
+                summary = self._coerce_summary_content(response.choices[0].message.content)
+                return self._ensure_summary_prefix(summary)
                 
             except Exception as e:
                 metrics.summarization_errors += 1
@@ -612,13 +622,8 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix."""
                         max_tokens=self.config.summary_target_tokens * 2,
                     )
                 
-                summary = response.choices[0].message.content.strip()
-                
-                # Ensure it starts with the prefix
-                if not summary.startswith("[CONTEXT SUMMARY]:"):
-                    summary = "[CONTEXT SUMMARY]: " + summary
-                
-                return summary
+                summary = self._coerce_summary_content(response.choices[0].message.content)
+                return self._ensure_summary_prefix(summary)
                 
             except Exception as e:
                 metrics.summarization_errors += 1