Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 60 additions & 0 deletions agent/context_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,39 @@
from typing import Any, Dict, List


def automatic_compaction_status_message(
engine: Any,
*,
phase: str,
default_message: str,
**context: Any,
) -> str | None:
"""Resolve host-visible status for an automatic compaction event.

Engines can suppress routine automatic status with
``emit_automatic_compaction_status = False`` or customize it by defining
``get_automatic_compaction_status_message(...)``. Empty strings and
``None`` mean "do not emit a lifecycle status".
"""
if not getattr(engine, "emit_automatic_compaction_status", True):
return None

formatter = getattr(engine, "get_automatic_compaction_status_message", None)
if callable(formatter):
message = formatter(
phase=phase,
default_message=default_message,
**context,
)
else:
message = default_message

if message is None:
return None
message = str(message).strip()
return message or None


class ContextEngine(ABC):
"""Base class all context engines must implement."""

Expand Down Expand Up @@ -65,6 +98,12 @@ def name(self) -> str:
protect_first_n: int = 3
protect_last_n: int = 6

# User-visible lifecycle status for automatic host-triggered compaction.
# Alternative engines that treat compaction as routine background
# maintenance can set this false to keep successful automatic passes silent;
# warnings, errors, and explicit manual commands should still surface.
emit_automatic_compaction_status: bool = True

# -- Core interface ----------------------------------------------------

@abstractmethod
Expand Down Expand Up @@ -124,6 +163,27 @@ def should_defer_preflight_to_real_usage(self, rough_tokens: int) -> bool:
"""
return False

def get_automatic_compaction_status_message(
self,
*,
phase: str,
default_message: str,
**context: Any,
) -> str | None:
"""Return user-visible status for automatic host-triggered compaction.

Return ``None`` to suppress successful automatic lifecycle status for
this compaction event. ``phase`` identifies the host call site (for
example ``"preflight"`` or ``"compress"``). ``context`` contains
best-effort fields such as ``approx_tokens`` and ``threshold_tokens``.

This hook does not control warning/error messages or explicit manual
commands such as ``/compress``.
"""
if not self.emit_automatic_compaction_status:
return None
return default_message

# -- Optional: manual /compress preflight ------------------------------

def has_content_to_compress(self, messages: List[Dict[str, Any]]) -> bool:
Expand Down
15 changes: 14 additions & 1 deletion agent/conversation_compression.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
from pathlib import Path
from typing import Any, Optional, Tuple

from agent.context_engine import automatic_compaction_status_message
from agent.model_metadata import estimate_request_tokens_rough

logger = logging.getLogger(__name__)
Expand Down Expand Up @@ -334,7 +335,19 @@ def compress_context(
f"{approx_tokens:,}" if approx_tokens else "unknown", agent.model,
focus_topic,
)
agent._emit_status(COMPACTION_STATUS)
_compaction_status = COMPACTION_STATUS
if not force:
_compaction_status = automatic_compaction_status_message(
agent.context_compressor,
phase="compress",
default_message=_compaction_status,
approx_tokens=approx_tokens,
message_count=_pre_msg_count,
model=agent.model,
focus_topic=focus_topic,
)
if _compaction_status:
agent._emit_status(_compaction_status)

# ── Compression lock ────────────────────────────────────────────────
# Atomic, state.db-backed lock per session_id. Without this, two
Expand Down
19 changes: 15 additions & 4 deletions agent/turn_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from agent.context_engine import automatic_compaction_status_message
from agent.iteration_budget import IterationBudget
from agent.model_metadata import estimate_request_tokens_rough

Expand Down Expand Up @@ -306,11 +307,21 @@ def build_turn_context(
agent.model,
f"{_compressor.context_length:,}",
)
agent._emit_status(
f"πŸ“¦ Preflight compression: ~{_preflight_tokens:,} tokens "
f">= {_compressor.threshold_tokens:,} threshold. "
"This may take a moment."
_preflight_status = automatic_compaction_status_message(
_compressor,
phase="preflight",
default_message=(
f"πŸ“¦ Preflight compression: ~{_preflight_tokens:,} tokens "
f">= {_compressor.threshold_tokens:,} threshold. "
"This may take a moment."
),
approx_tokens=_preflight_tokens,
threshold_tokens=_compressor.threshold_tokens,
context_length=_compressor.context_length,
model=agent.model,
)
if _preflight_status:
agent._emit_status(_preflight_status)
for _pass in range(3):
_orig_len = len(messages)
messages, active_system_prompt = agent._compress_context(
Expand Down
151 changes: 151 additions & 0 deletions tests/run_agent/test_413_compression.py
Original file line number Diff line number Diff line change
Expand Up @@ -477,6 +477,63 @@ def _fake_compress(messages, current_tokens=None, focus_topic=None):
assert "Compacting context" in events[0][1]
assert events[1] == ("compress", "started")

def test_compress_context_suppresses_automatic_status_when_engine_opts_out(self, agent):
"""Plugin engines can make successful automatic compaction silent."""
# Keep this isolated from the lazy aux-provider feasibility warning,
# which is unrelated to automatic compaction lifecycle status.
agent.compression_enabled = False
events = []
agent.status_callback = lambda ev, msg: events.append((ev, msg))
agent.context_compressor.emit_automatic_compaction_status = False

def _fake_compress(messages, current_tokens=None, focus_topic=None):
events.append(("compress", "started"))
return [{"role": "user", "content": f"{SUMMARY_PREFIX}\nPrevious conversation"}]

with (
patch.object(agent.context_compressor, "compress", side_effect=_fake_compress),
patch.object(agent, "_build_system_prompt", return_value="new system prompt"),
patch("run_agent.estimate_request_tokens_rough", return_value=42),
):
compressed, new_system_prompt = agent._compress_context(
[{"role": "user", "content": "hello"}],
"system prompt",
approx_tokens=1234,
)

assert compressed == [{"role": "user", "content": f"{SUMMARY_PREFIX}\nPrevious conversation"}]
assert new_system_prompt == "new system prompt"
assert events == [("compress", "started")]

def test_compress_context_force_keeps_manual_status_when_engine_opts_out(self, agent):
"""Manual /compress remains visible even for quiet automatic engines."""
# Keep this isolated from the lazy aux-provider feasibility warning,
# which is unrelated to manual compression lifecycle status.
agent.compression_enabled = False
events = []
agent.status_callback = lambda ev, msg: events.append((ev, msg))
agent.context_compressor.emit_automatic_compaction_status = False

def _fake_compress(messages, current_tokens=None, focus_topic=None, force=False):
events.append(("compress", "started"))
return [{"role": "user", "content": f"{SUMMARY_PREFIX}\nPrevious conversation"}]

with (
patch.object(agent.context_compressor, "compress", side_effect=_fake_compress),
patch.object(agent, "_build_system_prompt", return_value="new system prompt"),
patch("run_agent.estimate_request_tokens_rough", return_value=42),
):
agent._compress_context(
[{"role": "user", "content": "hello"}],
"system prompt",
approx_tokens=1234,
force=True,
)

assert events[0][0] == "lifecycle"
assert "Compacting context" in events[0][1]
assert events[1] == ("compress", "started")

def test_preflight_compresses_oversized_history(self, agent):
"""When loaded history exceeds the model's context threshold, compress before API call."""
agent.compression_enabled = True
Expand Down Expand Up @@ -529,6 +586,100 @@ def test_preflight_compresses_oversized_history(self, agent):
for ev, msg in status_messages
)

def test_preflight_suppresses_status_when_context_engine_opts_out(self, agent):
"""LCM-style engines can keep routine automatic preflight maintenance silent."""
agent.compression_enabled = True
agent.context_compressor.context_length = 200_000
agent.context_compressor.threshold_tokens = 100_000
agent.context_compressor.emit_automatic_compaction_status = False

big_history = []
for i in range(20):
big_history.append({"role": "user", "content": f"Message {i} padded"})
big_history.append({"role": "assistant", "content": f"Response {i} padded"})

ok_resp = _mock_response(content="After quiet preflight", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [ok_resp]
status_messages = []
agent.status_callback = lambda ev, msg: status_messages.append((ev, msg))

_rough_calls = {"n": 0}

def _rough_estimate(*_args, **_kwargs):
_rough_calls["n"] += 1
return 114_000 if _rough_calls["n"] == 1 else 40_000

with (
patch("agent.turn_context.estimate_request_tokens_rough", side_effect=_rough_estimate),
patch("agent.conversation_loop.estimate_request_tokens_rough", side_effect=_rough_estimate),
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": f"{SUMMARY_PREFIX}\nPrevious conversation"}],
"new system prompt",
)
result = agent.run_conversation("hello", conversation_history=big_history)

mock_compress.assert_called_once()
assert result["completed"] is True
assert not any(
ev == "lifecycle" and "Preflight compression" in msg
for ev, msg in status_messages
)

def test_preflight_uses_context_engine_custom_status_message(self, agent):
"""Plugin engines can replace generic built-in-compressor wording."""
agent.compression_enabled = True
agent.context_compressor.context_length = 200_000
agent.context_compressor.threshold_tokens = 100_000

def _custom_status(**kwargs):
assert kwargs["phase"] == "preflight"
assert kwargs["approx_tokens"] == 114_000
assert kwargs["threshold_tokens"] == 100_000
return "πŸ”§ LCM context maintenance: preparing compacted context."

agent.context_compressor.get_automatic_compaction_status_message = _custom_status

big_history = []
for i in range(20):
big_history.append({"role": "user", "content": f"Message {i} padded"})
big_history.append({"role": "assistant", "content": f"Response {i} padded"})

ok_resp = _mock_response(content="After custom preflight", finish_reason="stop")
agent.client.chat.completions.create.side_effect = [ok_resp]
status_messages = []
agent.status_callback = lambda ev, msg: status_messages.append((ev, msg))

_rough_calls = {"n": 0}

def _rough_estimate(*_args, **_kwargs):
_rough_calls["n"] += 1
return 114_000 if _rough_calls["n"] == 1 else 40_000

with (
patch("agent.turn_context.estimate_request_tokens_rough", side_effect=_rough_estimate),
patch("agent.conversation_loop.estimate_request_tokens_rough", side_effect=_rough_estimate),
patch.object(agent, "_compress_context") as mock_compress,
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
):
mock_compress.return_value = (
[{"role": "user", "content": f"{SUMMARY_PREFIX}\nPrevious conversation"}],
"new system prompt",
)
result = agent.run_conversation("hello", conversation_history=big_history)

mock_compress.assert_called_once()
assert result["completed"] is True
lifecycle_messages = [msg for ev, msg in status_messages if ev == "lifecycle"]
assert "πŸ”§ LCM context maintenance: preparing compacted context." in lifecycle_messages
assert not any("Preflight compression" in msg for msg in lifecycle_messages)

def test_preflight_defers_when_recent_real_usage_fit(self, agent):
"""A noisy rough estimate should not re-compact a recently fitting request."""
agent.compression_enabled = True
Expand Down
Loading