diff --git a/tests/test_trajectory_compressor.py b/tests/test_trajectory_compressor.py index 8fcbfc38cfef..b4e4c876a3f4 100644 --- a/tests/test_trajectory_compressor.py +++ b/tests/test_trajectory_compressor.py @@ -628,3 +628,67 @@ def test_snap_boundary_falls_back_to_backward(self): {"from": "tool", "value": "a"}, ] assert tc._snap_boundary(trajectory, 1, 0, 1) == 0 + + +# --------------------------------------------------------------------------- +# TrajectoryCompressor — compression must never increase the token count +# --------------------------------------------------------------------------- + + +class TestCompressionNetSavingsGuard: + """When the compressible middle is no larger than the summary that would + replace it, compression cannot help — it must be skipped rather than grow + the trajectory (and burn a summarization call).""" + + def _tiny_middle_trajectory(self): + # Large protected head (system+human), tiny compressible middle. + big = "w " * 400 # ~200 tokens each (1 token / 4 chars) + small = "ok " * 2 + return [ + {"from": "system", "value": big}, # protected (first_system) + {"from": "human", "value": big}, # protected (first_human) + {"from": "gpt", "value": small}, # protected (first_gpt) + {"from": "tool", "value": small}, # protected (first_tool) + {"from": "gpt", "value": small}, # compressible middle + {"from": "tool", "value": small}, # compressible middle + {"from": "gpt", "value": small}, # protected (last 2) + {"from": "human", "value": small}, # protected (last 2) + ] + + def _config(self): + config = CompressionConfig() + config.protect_last_n_turns = 2 + config.summary_target_tokens = 20 + config.target_max_tokens = 100 # trajectory is far over this + return config + + def test_sync_skips_compression_when_middle_smaller_than_summary(self): + tc = _make_compressor(self._config()) + tc._generate_summary = MagicMock( + return_value="[CONTEXT SUMMARY]: " + "blah " * 30 + ) + trajectory = self._tiny_middle_trajectory() + before = sum(tc.count_turn_tokens(trajectory)) + + compressed, metrics = tc.compress_trajectory(trajectory) + + assert metrics.was_compressed is False + assert compressed == trajectory + assert sum(tc.count_turn_tokens(compressed)) == before + tc._generate_summary.assert_not_called() + + @pytest.mark.asyncio + async def test_async_skips_compression_when_middle_smaller_than_summary(self): + tc = _make_compressor(self._config()) + tc._generate_summary_async = AsyncMock( + return_value="[CONTEXT SUMMARY]: " + "blah " * 30 + ) + trajectory = self._tiny_middle_trajectory() + before = sum(tc.count_turn_tokens(trajectory)) + + compressed, metrics = await tc.compress_trajectory_async(trajectory) + + assert metrics.was_compressed is False + assert compressed == trajectory + assert sum(tc.count_turn_tokens(compressed)) == before + tc._generate_summary_async.assert_not_called() diff --git a/trajectory_compressor.py b/trajectory_compressor.py index 45d2386e933c..1b447811f3e3 100644 --- a/trajectory_compressor.py +++ b/trajectory_compressor.py @@ -831,6 +831,18 @@ def compress_trajectory( metrics.still_over_limit = total_tokens > self.config.target_max_tokens return trajectory, metrics + # If the region we can safely compress is no larger than the summary + # that would replace it, compression cannot reduce the token count -- + # it would grow the trajectory and still spend a summarization call. + if ( + sum(turn_tokens[compress_start:compress_until]) + <= self.config.summary_target_tokens + ): + metrics.compressed_tokens = total_tokens + metrics.compressed_turns = len(trajectory) + metrics.still_over_limit = total_tokens > self.config.target_max_tokens + return trajectory, metrics + # Record compression region metrics.turns_compressed_start_idx = compress_start metrics.turns_compressed_end_idx = compress_until @@ -946,6 +958,18 @@ async def compress_trajectory_async( metrics.still_over_limit = total_tokens > self.config.target_max_tokens return trajectory, metrics + # If the region we can safely compress is no larger than the summary + # that would replace it, compression cannot reduce the token count -- + # it would grow the trajectory and still spend a summarization call. + if ( + sum(turn_tokens[compress_start:compress_until]) + <= self.config.summary_target_tokens + ): + metrics.compressed_tokens = total_tokens + metrics.compressed_turns = len(trajectory) + metrics.still_over_limit = total_tokens > self.config.target_max_tokens + return trajectory, metrics + # Record compression region metrics.turns_compressed_start_idx = compress_start metrics.turns_compressed_end_idx = compress_until