"""Context isolation and token budget management. Implements the core context engineering principles for subagent isolation: - ContextMode: ISOLATED (fresh context per subagent), FORK (shared with fork), SHARED - TokenBudget: hard and soft limits enforcement - Compaction: rolling summaries, head/tail truncation at turn boundaries """ from __future__ import annotations import logging from dataclasses import dataclass from typing import Any from cyberbrain26.orchestration.subagent.interface import BudgetExceeded from cyberbrain26.orchestration.subagent.types import ( ContextMode, SubagentSpec, ) logger = logging.getLogger(__name__) # Token estimation for common content types TOKEN_ESTIMATES = { "text": 4, # ~4 tokens per character on average "code": 1.5, # code is roughly 1.5 tokens/char "json": 3, # JSON is generally compact, ~3 tokens/char "html": 5, # HTML is verbose, ~5 tokens/char } DEFAULT_CONTEXT_BUDGET = 80000 # Default per-subagent context window COMPACTION_THRESHOLD = 0.90 # Compact when context > 90% of budget MAX_SUMMARY_TOKENS = 2000 # Structured output max (Claude Guide rule) COMPACTING_LOOKBEAT_TOKENS = 60000 # Keep last N turns verbatim (rolling summary) @dataclass class Reservation: """Token reservation for a subagent invocation.""" granted: int used: int = 0 is_soft_limit: bool = False # Soft limit: log but don't fail class TokenBudgetManager: """Manages token budgets for subagents. Implements: - Hard limit: raise BudgetExceeded if estimated usage exceeds token_budget - Soft limit: log warning if actual exceeds granted (soft limit) - Accounting: track per-invocation reservations """ def __init__(self, spec: SubagentSpec) -> None: self._spec = spec self._reservations: list[Reservation] = [] def estimate_tokens(self, content: str, content_type: str | None = "text") -> int: """Estimate token count from content.""" est = TOKEN_ESTIMATES.get(content_type or "text", 4) return int(len(content) * est) def check_and_reserve(self, estimated: int) -> Reservation: """Check if invocation is within token budget and reserve tokens. Args: estimated: Estimated token usage for this invocation Returns: Reservation with granted/reserved tokens Raises: BudgetExceeded: if estimated > hard token_budget """ hard_limit = self._spec.token_budget if estimated > hard_limit: raise BudgetExceeded( f"Estimated {estimated} tokens > hard limit {hard_limit}: " f"role={self._spec.role}" ) granted = max(0, hard_limit - estimated) reservation = Reservation(granted=granted, used=estimated, is_soft_limit=False) self._reservations.append(reservation) return reservation def record_usage(self, actual: int) -> None: """Record actual token usage after execution. If actual exceeds soft-granted but under hard limit, just logs. If actual exceeds hard limit, BudgetExceeded already raised. """ for reservation in self._reservations: actual_exceeds = actual > reservation.granted if actual_exceeds and not reservation.is_soft_limit: logger.warning( "Subagent %s token overrun: %d > %d (granted)", self._spec.role, actual, reservation.granted, ) reservation.is_soft_limit = True @property def total_reserved(self) -> int: return sum(r.granted for r in self._reservations) @property def used(self) -> int: return sum(r.used for r in self._reservations) class ContextCompactor: """Manages context compaction for long-running subagents. Implements 3-level compaction: 1. Rolling summary (keep last N turns verbatim, summarize older turns) 2. LLM condensation (if model available, condense to structured output) 3. Head/tail truncate (60/40 split with omission marker as fallback) """ def __init__(self, compactor_function: Any = None) -> None: """Initialize with optional external LLM condensation function. Args: compactor_function: Callable that condenses messages via LLM. If None, only levels 1 and 3 are available. """ self._compactor_function = compactor_function def should_compact(self, current_tokens: int, max_context: int = DEFAULT_CONTEXT_BUDGET) -> bool: """Return True if compaction should be triggered. Compaction triggers at threshold (default 90%). """ return current_tokens >= COMPACTION_THRESHOLD * max_context def compact(self, messages: list[dict[str, Any]], max_context: int = DEFAULT_CONTEXT_BUDGET) -> list[dict[str, Any]]: """Compact a conversation before next turn. Args: messages: List of message dicts {type, content, role} max_context: Maximum context window size Returns: Compacted message list """ current_tokens = sum(self._estimate_tokens_message(m) for m in messages) if not self.should_compact(current_tokens, max_context): return messages if self._compactor_function is not None: try: condensed = self._compactor_function(messages) # Ensure structured output ≤ MAX_SUMMARY_TOKENS if len(condensed) > 0: pass # Condensed by LLM, trust it else: condensed = [m for m in messages if self._estimate_tokens_message(m) < 50] condensed.append({"type": "summary", "content": "[Compacted: LLM condensation unavailable]"}) return condensed[:max_context // TOKEN_ESTIMATES["text"]] except Exception as e: logger.warning("LLM condensation failed, using fallback: %s", e, exc_info=True) # Level 3: Head/tail truncation fallback threshold = int(max_context * COMPACTION_THRESHOLD / TOKEN_ESTIMATES["text"]) head_cutoff = max(0, len(messages) - threshold) truncated = messages[:head_cutoff] + [{"type": "marker", "content": "[Compacted: threshold cutoff]", "level": 3}] return truncated @staticmethod def _estimate_tokens_message(message: dict[str, Any]) -> int: """Estimate token count for a single message.""" content = message.get("content", "") type_key = message.get("type", "") est = TOKEN_ESTIMATES.get(type_key, 4) return int(len(str(content)) * est) def compact_at_turn_boundary(self, messages: list[dict[str, Any]]) -> None: """Hook called at turn boundary to compact accumulated context. This is the integration point with LangGraph node callbacks. """ compacted = self.compact(messages) self._pending_compacted = compacted class ContextIsolator: """Manages context isolation for subagents. Provides three isolation modes: ISOLATED: fresh context window per subagent (default, language model evaluates) FORK: forks of existing context for parallel exploration SHARED: shares context with main agent """ def __init__(self, spec: SubagentSpec) -> None: self._spec = spec self._isolated_state: dict[str, Any] = {} self._fork_state: dict[str, Any] = {} def get_initial_context(self, previous_output: Any = None) -> dict[str, Any]: """Get initial context for a fresh isolated subagent. Per Claude Code Guide: "Nothing loads into a context window unless that specific task needs it." """ if self._spec.context_mode == ContextMode.ISOLATED: return { "mode": "isolated", "priority": [], # Sharding priority list "previous_output": str(previous_output)[:200] if previous_output else "", "evidence_index": {}, # Evidence reference index } elif self._spec.context_mode == ContextMode.FORK: return {"mode": "fork", "parent_seed": id(previous_output)} elif self._spec.context_mode == ContextMode.SHARED: return {"mode": "shared"} return {"mode": "unknown"} def get_fork_context(self, parent_messages: list[Any]) -> dict[str, Any]: """Create a fork of parent context for parallel exploration.""" return { "mode": "fork", "source": "shared", "initial": parent_messages, } def record_evidence(self, evidence_ref: Any) -> None: """Record evidence for later retrieval by main agent.""" if evidence_ref not in self._isolated_state.get("evidence_index", []): self._isolated_state.setdefault("evidence_index", []).append(evidence_ref) def get_evidence_refs(self) -> list[Any]: """Return recorded evidence references.""" return list(self._isolated_state.get("evidence_index", []))