"""``LeanCtxEngine`` — lean-ctx as a Hermes context engine. Replaces the built-in ``ContextCompressor``: deterministic, prompt-cache friendly compaction of the message window plus native lean-ctx recall tools. All engine logic that *can* live in the daemon does (Single Source of Truth); this class is a thin, fault-tolerant adapter over the ``/v1`` API. """ from __future__ import annotations import json import logging from typing import Any, Dict, List, Optional try: # Real host contract wins whenever a Hermes checkout is importable. from agent.context_engine import ContextEngine # type: ignore except Exception: # pragma: no cover - exercised outside Hermes from ._hermes_compat import ContextEngine # type: ignore from . import compaction, presets from . import tokens as _tokens from . import tools as _tools from .config import LeanCtxConfig from .schemas import recall_hint from .transport import ToolGateway logger = logging.getLogger(__name__) _ENGINE_NAME = "lean-ctx" _OFFLOAD_MAX_CHARS = 8_000 def _first_int(d: Dict[str, Any], *keys: str) -> Optional[int]: for key in keys: val = d.get(key) if isinstance(val, bool): continue if isinstance(val, int): return val if isinstance(val, float): return int(val) return None class LeanCtxEngine(ContextEngine): """Context engine backed by the lean-ctx daemon.""" def __init__( self, context_length: Optional[int] = None, *, config: Optional[LeanCtxConfig] = None, hermes_home: Optional[str] = None, **kwargs: Any, ) -> None: cfg = config or LeanCtxConfig.from_env() resolved_ctx = int(context_length or cfg.context_length) if resolved_ctx != cfg.context_length: cfg = cfg.with_context_length(resolved_ctx) self._config = cfg self._hermes_home = hermes_home self._session_id: Optional[str] = None self._gateway = ToolGateway(cfg) # Initialise the host base in whatever shape it expects, then assert our # own attribute invariants so they exist regardless of the base. try: super().__init__(context_length=resolved_ctx) # type: ignore[misc] except TypeError: try: super().__init__() # type: ignore[misc] except Exception: # pragma: no cover - exotic host base pass self.last_prompt_tokens = 0 self.last_completion_tokens = 0 self.last_total_tokens = 0 self.context_length = resolved_ctx self.threshold_tokens = cfg.threshold_tokens() self.compression_count = 0 # --- identity ----------------------------------------------------------- @property def name(self) -> str: return _ENGINE_NAME @property def config(self) -> LeanCtxConfig: return self._config # --- token accounting --------------------------------------------------- def update_from_response(self, usage: Dict[str, Any]) -> None: if not isinstance(usage, dict): return pt = _first_int(usage, "prompt_tokens", "input_tokens") ct = _first_int(usage, "completion_tokens", "output_tokens") tt = _first_int(usage, "total_tokens") if pt is not None: self.last_prompt_tokens = pt if ct is not None: self.last_completion_tokens = ct if tt is not None: self.last_total_tokens = tt elif pt is not None or ct is not None: self.last_total_tokens = self.last_prompt_tokens + self.last_completion_tokens def should_compress(self, prompt_tokens: Optional[int] = None) -> bool: if self.threshold_tokens <= 0: return False tokens = prompt_tokens if tokens is None: tokens = self.last_prompt_tokens or self.last_total_tokens return int(tokens or 0) >= self.threshold_tokens def should_compress_preflight(self, messages: List[Dict[str, Any]]) -> bool: if self.threshold_tokens <= 0: return False return _tokens.count_messages_tokens(messages) >= self.threshold_tokens # --- compaction --------------------------------------------------------- def compress( self, messages: List[Dict[str, Any]], current_tokens: Optional[int] = None, focus_topic: Optional[str] = None, ) -> List[Dict[str, Any]]: if not isinstance(messages, list) or not messages: return messages try: input_tokens = _tokens.count_messages_tokens(messages) result: Optional[List[Dict[str, Any]]] = None if self._config.use_core_compaction: # Preferred: the daemon's deterministic core tool owns compaction # (Single Source of Truth) and offloads raw turns server-side. result = self._compress_via_daemon(messages, focus_topic) if result is None: # Fallback: daemon unreachable / tool missing — compact locally. result = self._compress_local(messages, focus_topic) if result is None: return messages # nothing to compact if _tokens.count_messages_tokens(result) < input_tokens: self.compression_count += 1 return result except Exception as exc: # never break the agent loop logger.warning("lean-ctx engine: compress() failed, returning input unchanged: %s", exc) return messages def _compress_via_daemon( self, messages: List[Dict[str, Any]], focus_topic: Optional[str], ) -> Optional[List[Dict[str, Any]]]: """Compact through the daemon's ``ctx_transcript_compact`` core tool. Returns the validated message list when the daemon handled the request, or ``None`` (→ local fallback) when it is unreachable, the tool is missing, or the response fails our hard OpenAI-sequence invariants. """ if not self._gateway.is_available(): return None args: Dict[str, Any] = { "messages": messages, "fresh_tail_tokens": self._config.protect_tokens(), "protect_min_messages": self._config.protect_min_messages, } if focus_topic: args["focus_topic"] = focus_topic raw = self._gateway.call_text("ctx_transcript_compact", args) if not raw: return None try: payload = json.loads(raw) except (ValueError, TypeError): return None if not isinstance(payload, dict): return None new_messages = payload.get("messages") if not isinstance(new_messages, list) or not new_messages: return None if not all(isinstance(m, dict) for m in new_messages): return None # Hard invariant: a tool_call/tool_result pair must never be split. if compaction.tool_pairing_errors(new_messages): logger.warning( "lean-ctx engine: daemon compaction broke tool pairing; using local fallback" ) return None # Safety: compaction must never grow the window. if _tokens.count_messages_tokens(new_messages) > _tokens.count_messages_tokens(messages): return None return new_messages def _compress_local( self, messages: List[Dict[str, Any]], focus_topic: Optional[str], ) -> Optional[List[Dict[str, Any]]]: """Pure-Python compaction used when the daemon path is unavailable.""" plan = compaction.plan_compaction( messages, protect_tokens=self._config.protect_tokens(), protect_min_messages=self._config.protect_min_messages, token_counter=_tokens.count_messages_tokens, ) if plan.nothing_to_do: return None self._offload(plan.to_summarize) summary = compaction.build_summary_message( plan.to_summarize, focus_topic=focus_topic, recall_hint=recall_hint() if self._config.enable_tools else "", ) return compaction.assemble(plan, summary) def _offload(self, to_summarize: List[Dict[str, Any]]) -> None: """Persist offloaded turns to lean-ctx so they remain recoverable. Only used by the local fallback path; the daemon core tool offloads server-side, so this avoids double-writing the same turns. """ if not to_summarize or not self._gateway.is_available(): return digest = compaction.serialize_transcript(to_summarize, max_chars=_OFFLOAD_MAX_CHARS) if not digest: return self._gateway.call_text("ctx_session", {"action": "finding", "value": digest}) # --- model / lifecycle -------------------------------------------------- def update_model( self, model: str, context_length: Optional[int] = None, **kwargs: Any, ) -> None: new_ctx = context_length or presets.context_length_for(model) if new_ctx: self._config = self._config.with_context_length(int(new_ctx)) self.context_length = self._config.context_length self.threshold_tokens = self._config.threshold_tokens() def on_session_start(self, session_id: str, **kwargs: Any) -> None: self._session_id = session_id # Restore prior cross-session state (task / findings / decisions) so that # recall and subsequent compaction summaries reflect earlier sessions. # Best-effort: a fresh project with no history simply no-ops. if self._gateway.is_available(): self._gateway.call_text("ctx_session", {"action": "resume"}) def on_session_end(self, session_id: str, messages: List[Dict[str, Any]]) -> None: # Durable cross-session persistence: record a session summary and write a # deterministic handoff ledger the next session can pull (ctx_handoff). if not self._gateway.is_available(): return self._gateway.call_text("ctx_summary", {"action": "record"}) self._gateway.call_text("ctx_handoff", {"action": "create"}) def on_session_reset(self) -> None: self.last_prompt_tokens = 0 self.last_completion_tokens = 0 self.last_total_tokens = 0 self._session_id = None # --- native tools ------------------------------------------------------- def get_tool_schemas(self) -> List[Dict[str, Any]]: return _tools.get_tool_schemas(self._config) def handle_tool_call(self, name: str, args: Dict[str, Any], **kwargs: Any) -> str: return _tools.handle_tool_call(self._gateway, name, args, **kwargs) # --- status ------------------------------------------------------------- def get_status(self) -> Dict[str, Any]: status: Dict[str, Any] = { "name": self.name, "engine": "lean-ctx", "base_url": self._config.base_url, "daemon_available": self._gateway.is_available(), "session_id": self._session_id, "context_length": self.context_length, "threshold_tokens": self.threshold_tokens, "compression_count": self.compression_count, "core_compaction": self._config.use_core_compaction, "tools_enabled": self._config.enable_tools, "last_prompt_tokens": self.last_prompt_tokens, "last_completion_tokens": self.last_completion_tokens, "last_total_tokens": self.last_total_tokens, } metrics = self._gateway.get_metrics() if metrics: for key in ( "total_tokens_saved", "tokens_saved", "saved_tokens", "net_saved_tokens", "savings", "saved_usd", "compression_ratio", ): if key in metrics: status[f"leanctx_{key}"] = metrics[key] return status