fix(compression): include system prompt + tool schemas in token estimates (#18265)

The user-visible /compress banner and the post-compression last_prompt_tokens
writeback both counted only the raw message transcript (chars/4). With a 15KB
system prompt and 30 tool schemas (~26KB), a 4-message transcript that looks
like ~45 tokens to the transcript-only estimator is really ~10.5K tokens of
request pressure — a 234x gap.

Two user-facing consequences:
- Banner shows 'Compressing … (~45 tokens)…' while compression is actually
  firing on 10K+ tokens of real pressure, confusing users about why
  compression triggered (reported by @codecovenant on X; #6217).
- Post-compression last_prompt_tokens writeback omits tool schemas, so the
  next should_compress() check compares real usage against a stale
  underestimate — compression triggers late, potentially past the model's
  context limit on small-context models (#14695).

Swap estimate_messages_tokens_rough() for estimate_request_tokens_rough()
at every user-visible banner and at the post-compression writeback.
estimate_request_tokens_rough() already existed for exactly this purpose
and includes system prompt + tool schemas.

Touched call sites:
- run_agent.py: post-compression last_prompt_tokens writeback, post-tool
  call should_compress() fallback when provider usage is missing
- cli.py: /compress banner + summary
- gateway/run.py: gateway /compress banner + summary
- tui_gateway/server.py: TUI /compress status + summary
- acp_adapter/server.py: ACP /compact before/after

Left intentionally alone:
- Session-hygiene fallback and the 'no agent' /status path in gateway/run.py
  — no agent instance is in scope to query for system prompt/tools, and the
  existing 30-50% overestimate wobble on hygiene is safety-accepted.
- Verbose-mode 'Request size' logging — informational only, already counts
  system prompt via api_messages[0].

Also relabels the feedback line from 'Rough transcript estimate' to
'Approx request size' so the metric label matches what it actually measures.

Credits: diagnoses from @devilardis (#14695) and @Jackten (#6217);
user report @codecovenant on X (2026-04-30).

Closes #14695
Closes #6217
This commit is contained in:
Teknium
2026-04-30 23:03:54 -07:00
committed by GitHub
parent 41fa1f1b5c
commit f0dc919f92
9 changed files with 126 additions and 42 deletions

View File

@@ -730,6 +730,7 @@ class TestSlashCommands:
]
state.agent.compression_enabled = True
state.agent._cached_system_prompt = "system"
state.agent.tools = None
original_session_db = object()
state.agent._session_db = original_session_db
@@ -746,7 +747,7 @@ class TestSlashCommands:
with (
patch.object(agent.session_manager, "save_session") as mock_save,
patch(
"agent.model_metadata.estimate_messages_tokens_rough",
"agent.model_metadata.estimate_request_tokens_rough",
side_effect=[40, 12],
),
):

View File

@@ -21,20 +21,21 @@ def test_manual_compress_reports_noop_without_success_banner(capsys):
shell.agent = MagicMock()
shell.agent.compression_enabled = True
shell.agent._cached_system_prompt = ""
shell.agent.tools = None
shell.agent.session_id = shell.session_id # no-op compression: no split
shell.agent._compress_context.return_value = (list(history), "")
def _estimate(messages):
def _estimate(messages, **_kwargs):
assert messages == history
return 100
with patch("agent.model_metadata.estimate_messages_tokens_rough", side_effect=_estimate):
with patch("agent.model_metadata.estimate_request_tokens_rough", side_effect=_estimate):
shell._manual_compress()
output = capsys.readouterr().out
assert "No changes from compression" in output
assert "✅ Compressed" not in output
assert "Rough transcript estimate: ~100 tokens (unchanged)" in output
assert "Approx request size: ~100 tokens (unchanged)" in output
def test_manual_compress_explains_when_token_estimate_rises(capsys):
@@ -49,22 +50,23 @@ def test_manual_compress_explains_when_token_estimate_rises(capsys):
shell.agent = MagicMock()
shell.agent.compression_enabled = True
shell.agent._cached_system_prompt = ""
shell.agent.tools = None
shell.agent.session_id = shell.session_id # no-op: no split
shell.agent._compress_context.return_value = (compressed, "")
def _estimate(messages):
def _estimate(messages, **_kwargs):
if messages == history:
return 100
if messages == compressed:
return 120
raise AssertionError(f"unexpected transcript: {messages!r}")
with patch("agent.model_metadata.estimate_messages_tokens_rough", side_effect=_estimate):
with patch("agent.model_metadata.estimate_request_tokens_rough", side_effect=_estimate):
shell._manual_compress()
output = capsys.readouterr().out
assert "✅ Compressed: 4 → 3 messages" in output
assert "Rough transcript estimate: ~100 → ~120 tokens" in output
assert "Approx request size: ~100 → ~120 tokens" in output
assert "denser summaries" in output
@@ -89,6 +91,7 @@ def test_manual_compress_syncs_session_id_after_split():
shell.agent = MagicMock()
shell.agent.compression_enabled = True
shell.agent._cached_system_prompt = ""
shell.agent.tools = None
# Simulate _compress_context mutating agent.session_id as a side effect.
def _fake_compress(*args, **kwargs):
shell.agent.session_id = new_child_id
@@ -97,7 +100,7 @@ def test_manual_compress_syncs_session_id_after_split():
shell.agent.session_id = old_id # starts in sync
shell._pending_title = "stale title"
with patch("agent.model_metadata.estimate_messages_tokens_rough", return_value=100):
with patch("agent.model_metadata.estimate_request_tokens_rough", return_value=100):
shell._manual_compress()
# CLI session_id must now point at the continuation child, not the parent.
@@ -118,11 +121,12 @@ def test_manual_compress_no_sync_when_session_id_unchanged():
shell.agent = MagicMock()
shell.agent.compression_enabled = True
shell.agent._cached_system_prompt = ""
shell.agent.tools = None
shell.agent.session_id = shell.session_id
shell.agent._compress_context.return_value = (list(history), "")
shell._pending_title = "keep me"
with patch("agent.model_metadata.estimate_messages_tokens_rough", return_value=100):
with patch("agent.model_metadata.estimate_request_tokens_rough", return_value=100):
shell._manual_compress()
# No split → pending title untouched.

View File

@@ -64,11 +64,13 @@ async def test_compress_command_reports_noop_without_success_banner():
agent_instance = MagicMock()
agent_instance.shutdown_memory_provider = MagicMock()
agent_instance.close = MagicMock()
agent_instance._cached_system_prompt = ""
agent_instance.tools = None
agent_instance.context_compressor.has_content_to_compress.return_value = True
agent_instance.session_id = "sess-1"
agent_instance._compress_context.return_value = (list(history), "")
def _estimate(messages):
def _estimate(messages, **_kwargs):
assert messages == history
return 100
@@ -76,13 +78,13 @@ async def test_compress_command_reports_noop_without_success_banner():
patch("gateway.run._resolve_runtime_agent_kwargs", return_value={"api_key": "test-key"}),
patch("gateway.run._resolve_gateway_model", return_value="test-model"),
patch("run_agent.AIAgent", return_value=agent_instance),
patch("agent.model_metadata.estimate_messages_tokens_rough", side_effect=_estimate),
patch("agent.model_metadata.estimate_request_tokens_rough", side_effect=_estimate),
):
result = await runner._handle_compress_command(_make_event())
assert "No changes from compression" in result
assert "Compressed:" not in result
assert "Rough transcript estimate: ~100 tokens (unchanged)" in result
assert "Approx request size: ~100 tokens (unchanged)" in result
agent_instance.shutdown_memory_provider.assert_called_once()
agent_instance.close.assert_called_once()
@@ -99,11 +101,13 @@ async def test_compress_command_explains_when_token_estimate_rises():
agent_instance = MagicMock()
agent_instance.shutdown_memory_provider = MagicMock()
agent_instance.close = MagicMock()
agent_instance._cached_system_prompt = ""
agent_instance.tools = None
agent_instance.context_compressor.has_content_to_compress.return_value = True
agent_instance.session_id = "sess-1"
agent_instance._compress_context.return_value = (compressed, "")
def _estimate(messages):
def _estimate(messages, **_kwargs):
if messages == history:
return 100
if messages == compressed:
@@ -114,12 +118,12 @@ async def test_compress_command_explains_when_token_estimate_rises():
patch("gateway.run._resolve_runtime_agent_kwargs", return_value={"api_key": "test-key"}),
patch("gateway.run._resolve_gateway_model", return_value="test-model"),
patch("run_agent.AIAgent", return_value=agent_instance),
patch("agent.model_metadata.estimate_messages_tokens_rough", side_effect=_estimate),
patch("agent.model_metadata.estimate_request_tokens_rough", side_effect=_estimate),
):
result = await runner._handle_compress_command(_make_event())
assert "Compressed: 4 → 3 messages" in result
assert "Rough transcript estimate: ~100 → ~120 tokens" in result
assert "Approx request size: ~100 → ~120 tokens" in result
assert "denser summaries" in result
agent_instance.shutdown_memory_provider.assert_called_once()
agent_instance.close.assert_called_once()
@@ -143,6 +147,8 @@ async def test_compress_command_appends_warning_when_summary_generation_fails():
agent_instance = MagicMock()
agent_instance.shutdown_memory_provider = MagicMock()
agent_instance.close = MagicMock()
agent_instance._cached_system_prompt = ""
agent_instance.tools = None
agent_instance.context_compressor.has_content_to_compress.return_value = True
# Simulate summary-generation failure: fallback flag set, dropped count
# populated, error string captured.
@@ -154,7 +160,7 @@ async def test_compress_command_appends_warning_when_summary_generation_fails():
agent_instance.session_id = "sess-1"
agent_instance._compress_context.return_value = (compressed, "")
def _estimate(messages):
def _estimate(messages, **_kwargs):
if messages == history:
return 100
if messages == compressed:
@@ -165,7 +171,7 @@ async def test_compress_command_appends_warning_when_summary_generation_fails():
patch("gateway.run._resolve_runtime_agent_kwargs", return_value={"api_key": "***"}),
patch("gateway.run._resolve_gateway_model", return_value="test-model"),
patch("run_agent.AIAgent", return_value=agent_instance),
patch("agent.model_metadata.estimate_messages_tokens_rough", side_effect=_estimate),
patch("agent.model_metadata.estimate_request_tokens_rough", side_effect=_estimate),
):
result = await runner._handle_compress_command(_make_event())
@@ -200,6 +206,8 @@ async def test_compress_command_surfaces_aux_model_failure_even_when_recovered()
agent_instance = MagicMock()
agent_instance.shutdown_memory_provider = MagicMock()
agent_instance.close = MagicMock()
agent_instance._cached_system_prompt = ""
agent_instance.tools = None
agent_instance.context_compressor.has_content_to_compress.return_value = True
# Fallback placeholder was NOT used — recovery succeeded.
agent_instance.context_compressor._last_summary_fallback_used = False
@@ -215,7 +223,7 @@ async def test_compress_command_surfaces_aux_model_failure_even_when_recovered()
agent_instance.session_id = "sess-1"
agent_instance._compress_context.return_value = (compressed, "")
def _estimate(messages):
def _estimate(messages, **_kwargs):
if messages == history:
return 100
if messages == compressed:
@@ -226,7 +234,7 @@ async def test_compress_command_surfaces_aux_model_failure_even_when_recovered()
patch("gateway.run._resolve_runtime_agent_kwargs", return_value={"api_key": "***"}),
patch("gateway.run._resolve_gateway_model", return_value="test-model"),
patch("run_agent.AIAgent", return_value=agent_instance),
patch("agent.model_metadata.estimate_messages_tokens_rough", side_effect=_estimate),
patch("agent.model_metadata.estimate_request_tokens_rough", side_effect=_estimate),
):
result = await runner._handle_compress_command(_make_event())