cleveragents-core/robot/helper_context_tier_runtime.py

"""Robot Framework helper for context tier runtime integration tests.

Each command exercises the runtime logic added by issue #821:
auto-promotion, staleness enforcement, and budget enforcement.

Usage::

    python helper_context_tier_runtime.py <command>

Commands:
    auto-promote     -- Verify auto-promotion on repeated access.
    staleness        -- Verify enforce_staleness() demotes stale fragments.
    budget-evict     -- Verify budget enforcement on store.
    event-emission   -- Verify tier transition events are emitted.
"""

from __future__ import annotations

import sys
from collections.abc import Callable
from datetime import UTC, datetime, timedelta
from pathlib import Path

# Ensure source tree is importable
_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from cleveragents.application.services.context_tiers import (  # noqa: E402
    ContextTierService,
)
from cleveragents.domain.models.acms.tiers import (  # noqa: E402
    ContextTier,
    TierBudget,
    TieredFragment,
)
from cleveragents.infrastructure.events.models import DomainEvent  # noqa: E402
from cleveragents.infrastructure.events.types import EventType  # noqa: E402


class _EventCollector:
    """Simple event bus replacement that collects emitted events."""

    def __init__(self) -> None:
        self.events: list[DomainEvent] = []

    def emit(self, event: DomainEvent) -> None:
        self.events.append(event)

    def subscribe(
        self,
        event_type: EventType,
        handler: Callable[[DomainEvent], None],
    ) -> None:
        pass


def _make_frag(
    fid: str,
    tier: ContextTier,
    tokens: int = 50,
    last_accessed: datetime | None = None,
) -> TieredFragment:
    frag = TieredFragment(
        fragment_id=fid,
        content="test",
        tier=tier,
        token_count=tokens,
        project_name="test",
    )
    if last_accessed is not None:
        frag.last_accessed = last_accessed
    return frag


def cmd_auto_promote() -> None:
    svc = ContextTierService(settings=None)
    frag = _make_frag("frag-01", ContextTier.COLD)
    svc.store(frag)

    for _ in range(5):
        svc.get("frag-01")

    result = svc._find_fragment("frag-01")
    assert result is not None, "Fragment disappeared"
    assert result.tier == ContextTier.WARM, (
        f"Expected WARM (not HOT — counter reset prevents chain promotion)"
        f" after 5 accesses from cold, got {result.tier}"
    )
    print("context-tier-runtime-auto-promote-ok")


def cmd_staleness() -> None:
    svc = ContextTierService(settings=None)
    stale = datetime.now(tz=UTC) - timedelta(hours=25)
    frag = _make_frag("stale-01", ContextTier.HOT, last_accessed=stale)
    svc.store(frag)

    demoted = svc.enforce_staleness()
    assert "stale-01" in demoted, f"Expected stale-01 in demoted list: {demoted}"

    result = svc._find_fragment("stale-01")
    assert result is not None, "Fragment disappeared"
    assert result.tier == ContextTier.WARM, (
        f"Expected WARM after staleness demotion, got {result.tier}"
    )
    print("context-tier-runtime-staleness-ok")


def cmd_budget_evict() -> None:
    svc = ContextTierService(settings=None)
    svc._budget = TierBudget(max_tokens_hot=100)

    old = datetime.now(tz=UTC) - timedelta(minutes=10)
    svc.store(_make_frag("old-01", ContextTier.HOT, tokens=50, last_accessed=old))
    svc.store(_make_frag("new-01", ContextTier.HOT, tokens=50))
    svc.store(_make_frag("overflow-01", ContextTier.HOT, tokens=50))

    total = sum(f.token_count for f in svc._hot.values())
    assert total <= 100, f"Hot tier over budget: {total} > 100"

    old_frag = svc._find_fragment("old-01")
    assert old_frag is None, "old-01 should have been evicted"
    print("context-tier-runtime-budget-evict-ok")


def cmd_event_emission() -> None:
    collector = _EventCollector()
    svc = ContextTierService(settings=None, event_bus=collector)

    svc.store(_make_frag("evt-01", ContextTier.COLD))
    svc.promote("evt-01")

    promoted = [e for e in collector.events if e.event_type == EventType.TIER_PROMOTED]
    assert len(promoted) >= 1, f"No TIER_PROMOTED events: {collector.events}"
    assert promoted[0].details["fragment_id"] == "evt-01"
    print("context-tier-runtime-event-emission-ok")


_COMMANDS: dict[str, Callable[[], None]] = {
    "auto-promote": cmd_auto_promote,
    "staleness": cmd_staleness,
    "budget-evict": cmd_budget_evict,
    "event-emission": cmd_event_emission,
}


def main() -> None:
    if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
        valid = ", ".join(sorted(_COMMANDS.keys()))
        print(f"Usage: {sys.argv[0]} <{valid}>", file=sys.stderr)
        sys.exit(1)
    cmd = _COMMANDS[sys.argv[1]]
    cmd()


if __name__ == "__main__":
    main()