src/inferscale/agentic.py

from __future__ import annotations

import heapq
import math
import random
from dataclasses import asdict, dataclass
from statistics import mean, median
from typing import Any

from .latency import AnalyticalLatencyModel
from .metrics import percentile
from .prediction import OnlineToolGapPredictor, PREDICTOR_SCOPES
from .profiles import get_accelerator, get_model


ROUTING_POLICIES = {"least_load", "session_affinity", "bounded_affinity"}
RETENTION_POLICIES = {"evict", "retain", "ttl", "offload", "gap_aware", "adaptive"}

TOOL_GAP_MULTIPLIERS = {
    "filesystem": 0.45,
    "search": 0.80,
    "database": 1.25,
    "remote_api": 2.20,
}
TOOL_WEIGHTS = (0.28, 0.32, 0.22, 0.18)


@dataclass
class AgentSessionConfig:
    model: str = "Qwen2.5-3B"
    accelerator: str = "L4"
    quantization: str = "int8"
    seed: int = 7
    duration_s: float = 30.0
    session_rate_rps: float = 0.45
    replicas: int = 2
    turns_mean: float = 4.0
    turns_cv: float = 0.25
    initial_prompt_tokens_mean: int = 640
    append_tokens_mean: int = 180
    token_cv: float = 0.35
    output_tokens_mean: int = 72
    output_tokens_cv: float = 0.45
    tool_gap_mean_s: float = 1.5
    tool_gap_cv: float = 0.75
    retention_policy: str = "ttl"
    kv_ttl_s: float = 3.0
    routing_policy: str = "session_affinity"
    kv_memory_fraction: float = 0.85
    kv_capacity_override_gb: float = 0.0
    host_memory_gb: float = 32.0
    host_bandwidth_gbps: float = 32.0
    host_transfer_base_ms: float = 0.15
    affinity_slack_ms: float = 150.0
    gap_aware_threshold_s: float = 1.5
    adaptive_predictor_scope: str = "per_tool_ema"
    adaptive_alpha: float = 0.30
    adaptive_min_observations: int = 2
    tool_regime_shift_fraction: float = 0.0
    tool_regime_shift_multiplier: float = 1.0
    slo_turn_ttft_ms: float = 500.0
    slo_session_e2e_ms: float = 30000.0
    timeline_points: int = 240

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "AgentSessionConfig":
        allowed = cls.__dataclass_fields__.keys()
        cfg = cls(**{k: data[k] for k in allowed if k in data})
        if cfg.routing_policy not in ROUTING_POLICIES:
            raise ValueError(f"Unsupported agent routing policy: {cfg.routing_policy}")
        if cfg.retention_policy not in RETENTION_POLICIES:
            raise ValueError(f"Unsupported KV retention policy: {cfg.retention_policy}")
        cfg.replicas = max(1, min(int(cfg.replicas), 8))
        cfg.session_rate_rps = max(float(cfg.session_rate_rps), 0.01)
        cfg.duration_s = max(float(cfg.duration_s), 1.0)
        cfg.kv_ttl_s = max(float(cfg.kv_ttl_s), 0.0)
        cfg.kv_memory_fraction = min(max(float(cfg.kv_memory_fraction), 0.01), 0.98)
        cfg.kv_capacity_override_gb = max(float(cfg.kv_capacity_override_gb), 0.0)
        cfg.host_memory_gb = max(float(cfg.host_memory_gb), 0.0)
        cfg.host_bandwidth_gbps = max(float(cfg.host_bandwidth_gbps), 0.1)
        cfg.host_transfer_base_ms = max(float(cfg.host_transfer_base_ms), 0.0)
        cfg.affinity_slack_ms = max(float(cfg.affinity_slack_ms), 0.0)
        cfg.gap_aware_threshold_s = max(float(cfg.gap_aware_threshold_s), 0.0)
        if cfg.adaptive_predictor_scope not in PREDICTOR_SCOPES:
            raise ValueError(f"Unsupported adaptive predictor scope: {cfg.adaptive_predictor_scope}")
        cfg.adaptive_alpha = min(max(float(cfg.adaptive_alpha), 0.01), 1.0)
        cfg.adaptive_min_observations = max(1, int(cfg.adaptive_min_observations))
        cfg.tool_regime_shift_fraction = min(max(float(cfg.tool_regime_shift_fraction), 0.0), 0.95)
        cfg.tool_regime_shift_multiplier = max(float(cfg.tool_regime_shift_multiplier), 0.1)
        return cfg

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class TurnSpec:
    session_id: int
    turn_index: int
    append_tokens: int
    output_tokens: int
    tool_gap_after_s: float
    tool_kind: str = "generic"
    shifted_regime: bool = False


@dataclass
class SessionSpec:
    session_id: int
    arrival_time: float
    initial_prompt_tokens: int
    turns: list[TurnSpec]


@dataclass
class CacheEntry:
    session_id: int
    tokens: int
    size_gb: float
    last_access: float
    expiry_time: float | None
    generation: int
    source_replica: int | None = None
    available_time: float = 0.0


@dataclass
class ReplicaState:
    replica_id: int
    busy: bool = False
    busy_until: float = 0.0
    queue: list[tuple[int, TurnSpec, float]] | None = None
    cache: dict[int, CacheEntry] | None = None
    current_session: int | None = None

    def __post_init__(self) -> None:
        if self.queue is None:
            self.queue = []
        if self.cache is None:
            self.cache = {}


@dataclass
class SessionRuntime:
    spec: SessionSpec
    context_tokens: int
    completed_turns: int = 0
    completion_time: float | None = None
    last_replica: int | None = None


def _sample_positive_lognormal(rng: random.Random, mean_value: float, cv: float, minimum: float) -> float:
    mean_value = max(float(mean_value), minimum)
    cv = max(float(cv), 0.0)
    if mean_value <= 0.0:
        return minimum
    if cv <= 1e-12:
        return mean_value
    sigma2 = math.log1p(cv * cv)
    sigma = math.sqrt(sigma2)
    mu = math.log(mean_value) - sigma2 / 2.0
    return max(minimum, rng.lognormvariate(mu, sigma))


def _sample_int(rng: random.Random, mean_value: float, cv: float, minimum: int = 1) -> int:
    return max(minimum, int(round(_sample_positive_lognormal(rng, mean_value, cv, minimum))))


def generate_agent_sessions(cfg: AgentSessionConfig) -> list[SessionSpec]:
    """Generate a deterministic multi-turn program trace.

    All policies consume this same trace when the seed/config are unchanged. Tool
    gaps are sampled up front so policy comparisons use common random numbers.
    """
    rng = random.Random(cfg.seed ^ 0xA63E17)
    sessions: list[SessionSpec] = []
    now = 0.0
    sid = 0
    while True:
        now += rng.expovariate(cfg.session_rate_rps)
        if now > cfg.duration_s:
            break
        turns = max(2, _sample_int(rng, cfg.turns_mean, cfg.turns_cv, 2))
        initial = _sample_int(rng, cfg.initial_prompt_tokens_mean, cfg.token_cv, 32)
        specs = []
        tool_names = tuple(TOOL_GAP_MULTIPLIERS)
        shifted = cfg.tool_regime_shift_fraction > 0.0 and now >= cfg.duration_s * cfg.tool_regime_shift_fraction
        for turn_idx in range(turns):
            append = 0 if turn_idx == 0 else _sample_int(rng, cfg.append_tokens_mean, cfg.token_cv, 8)
            output = _sample_int(rng, cfg.output_tokens_mean, cfg.output_tokens_cv, 1)
            tool_kind = rng.choices(tool_names, weights=TOOL_WEIGHTS, k=1)[0]
            if turn_idx == turns - 1:
                gap = 0.0
            else:
                mean_gap = cfg.tool_gap_mean_s * TOOL_GAP_MULTIPLIERS[tool_kind]
                # The optional shift deliberately changes only the slower external
                # tools. This creates a non-stationary workload where a global
                # duration average is less informative than tool-aware history.
                if shifted and tool_kind in {"database", "remote_api"}:
                    mean_gap *= cfg.tool_regime_shift_multiplier
                gap = _sample_positive_lognormal(rng, mean_gap, cfg.tool_gap_cv, 0.0)
            specs.append(TurnSpec(sid, turn_idx, append, output, gap, tool_kind, shifted))
        sessions.append(SessionSpec(sid, now, initial, specs))
        sid += 1
    return sessions


class AgentSessionSimulator:
    """Discrete-event simulator for stateful multi-turn serving.

    The module deliberately isolates session locality / KV-retention effects from
    dynamic batching. Each replica is a serial analytical service station; the
    existing Serving Lab remains the place to study continuous batching. This
    separation keeps the agent experiment interpretable while still modelling
    session dependencies, tool gaps, routing, memory pressure, and cross-turn KV.
    """

    def __init__(self, cfg: AgentSessionConfig, sessions: list[SessionSpec] | None = None):
        self.cfg = cfg
        self.model = get_model(cfg.model)
        self.accelerator = get_accelerator(cfg.accelerator)
        self.latency = AnalyticalLatencyModel(self.model, self.accelerator, cfg.quantization)
        remaining = max(0.0, self.accelerator.vram_gb - self.latency.model_weight_gb - 1.2)
        automatic_capacity = remaining * cfg.kv_memory_fraction
        self.kv_capacity_gb = (
            min(automatic_capacity, cfg.kv_capacity_override_gb)
            if cfg.kv_capacity_override_gb > 0.0
            else automatic_capacity
        )
        self.replicas = [ReplicaState(i) for i in range(cfg.replicas)]
        self.host_cache: dict[int, CacheEntry] = {}
        specs = sessions if sessions is not None else generate_agent_sessions(cfg)
        self.sessions = {s.session_id: SessionRuntime(s, s.initial_prompt_tokens) for s in specs}
        self.events: list[tuple[float, int, str, tuple[Any, ...]]] = []
        self.event_seq = 0
        self.now = 0.0
        self.turn_rows: list[dict[str, Any]] = []
        self.timeline: list[dict[str, Any]] = []
        self.peak_kv_gb = 0.0
        self.peak_replica_kv_gb = 0.0
        self.hbm_gb_seconds = 0.0
        self.host_gb_seconds = 0.0
        self.peak_host_gb = 0.0
        self.last_memory_time = 0.0
        self.pressure_evictions = 0
        self.ttl_evictions = 0
        self.host_pressure_evictions = 0
        self.host_cache_hits = 0
        self.offload_bytes = 0.0
        self.restore_bytes = 0.0
        self.host_transfer_latencies_ms: list[float] = []
        self.recompute_tokens = 0
        self.cache_hits = 0
        self.cache_eligible_turns = 0
        self.affinity_routes = 0
        self.route_opportunities = 0
        self.failed_turns = 0
        self.tool_gap_total_s = 0.0
        self.tool_gap_predictor = OnlineToolGapPredictor(
            initial_mean_s=cfg.tool_gap_mean_s,
            alpha=cfg.adaptive_alpha,
            min_observations=cfg.adaptive_min_observations,
            scope=cfg.adaptive_predictor_scope,
        )
        self.prediction_rows: list[dict[str, Any]] = []
        for session in specs:
            if session.turns:
                self._push(session.arrival_time, "turn_ready", session.session_id, 0)

    def _push(self, when: float, kind: str, *payload: Any) -> None:
        self.event_seq += 1
        heapq.heappush(self.events, (float(when), self.event_seq, kind, tuple(payload)))

    def _used_gb(self) -> float:
        return sum(entry.size_gb for replica in self.replicas for entry in replica.cache.values())

    def _used_host_gb(self) -> float:
        return sum(entry.size_gb for entry in self.host_cache.values())

    def _integrate_memory(self, new_time: float) -> None:
        new_time = max(float(new_time), self.last_memory_time)
        elapsed = new_time - self.last_memory_time
        used = self._used_gb()
        host_used = self._used_host_gb()
        self.hbm_gb_seconds += used * elapsed
        self.host_gb_seconds += host_used * elapsed
        self.last_memory_time = new_time
        self.peak_kv_gb = max(self.peak_kv_gb, used)
        self.peak_replica_kv_gb = max(
            self.peak_replica_kv_gb,
            max((sum(entry.size_gb for entry in replica.cache.values()) for replica in self.replicas), default=0.0),
        )
        self.peak_host_gb = max(self.peak_host_gb, host_used)

    def _record_timeline(self) -> None:
        max_points = max(int(self.cfg.timeline_points), 40)
        if self.timeline and self.now - self.timeline[-1]["time_s"] < max(self.cfg.duration_s / max_points, 0.05):
            return
        self.timeline.append(
            {
                "time_s": self.now,
                "queued_turns": sum(len(r.queue) for r in self.replicas),
                "busy_replicas": sum(1 for r in self.replicas if r.busy),
                "kv_used_gb": self._used_gb(),
                "host_kv_gb": self._used_host_gb(),
                "resident_sessions": sum(len(r.cache) for r in self.replicas),
                "host_sessions": len(self.host_cache),
            }
        )
        if len(self.timeline) > max_points * 2:
            self.timeline = self.timeline[::2]

    def _entry(self, replica: ReplicaState, session_id: int) -> CacheEntry | None:
        return replica.cache.get(session_id)

    def _cache_replica(self, session_id: int) -> ReplicaState | None:
        for replica in self.replicas:
            if session_id in replica.cache:
                return replica
        return None

    def _replica_load_key(self, replica: ReplicaState) -> tuple[int, float, int]:
        return (len(replica.queue), max(replica.busy_until, self.now), replica.replica_id)

    def _route(self, session_id: int) -> ReplicaState:
        cached = self._cache_replica(session_id)
        if cached is not None:
            self.route_opportunities += 1
        least_loaded = min(self.replicas, key=self._replica_load_key)
        if self.cfg.routing_policy == "session_affinity" and cached is not None:
            self.affinity_routes += 1
            return cached
        if self.cfg.routing_policy == "bounded_affinity" and cached is not None:
            cached_ready = max(cached.busy_until, self.now)
            best_ready = max(least_loaded.busy_until, self.now)
            queue_gap = max(0, len(cached.queue) - len(least_loaded.queue))
            # Queue length is converted to a modest reference penalty because
            # exact queued service time is intentionally not known at routing time.
            estimated_extra_ms = max(0.0, cached_ready - best_ready) * 1000.0 + queue_gap * 75.0
            if estimated_extra_ms <= self.cfg.affinity_slack_ms:
                self.affinity_routes += 1
                return cached
        return least_loaded

    def _remove_cache(self, replica: ReplicaState, session_id: int, reason: str) -> None:
        if session_id not in replica.cache:
            return
        self._integrate_memory(self.now)
        del replica.cache[session_id]
        if reason == "pressure":
            self.pressure_evictions += 1
        elif reason == "ttl":
            self.ttl_evictions += 1
        self._record_timeline()

    def _host_transfer_seconds(self, size_gb: float) -> float:
        return self.cfg.host_transfer_base_ms / 1000.0 + size_gb / self.cfg.host_bandwidth_gbps

    def _remove_host_cache(self, session_id: int) -> None:
        if session_id not in self.host_cache:
            return
        self._integrate_memory(self.now)
        del self.host_cache[session_id]
        self._record_timeline()

    def _ensure_host_capacity(self, session_id: int, target_gb: float) -> bool:
        if target_gb > self.cfg.host_memory_gb + 1e-12:
            return False
        while self._used_host_gb() + target_gb > self.cfg.host_memory_gb + 1e-12:
            victims = [entry for sid, entry in self.host_cache.items() if sid != session_id]
            if not victims:
                return False
            victim = min(victims, key=lambda e: (e.last_access, e.session_id))
            self._integrate_memory(self.now)
            del self.host_cache[victim.session_id]
            self.host_pressure_evictions += 1
        return True

    def _offload_cache(self, replica: ReplicaState, session_id: int, tokens: int) -> bool:
        size_gb = tokens * self.latency.kv_bytes_per_token() / 1e9
        if self.cfg.host_memory_gb <= 0.0 or not self._ensure_host_capacity(session_id, size_gb):
            self._remove_cache(replica, session_id, "complete")
            return False
        transfer_s = self._host_transfer_seconds(size_gb)
        self._integrate_memory(self.now)
        replica.cache.pop(session_id, None)
        self.host_cache[session_id] = CacheEntry(
            session_id=session_id,
            tokens=tokens,
            size_gb=size_gb,
            last_access=self.now,
            expiry_time=None,
            generation=1,
            source_replica=replica.replica_id,
            available_time=self.now + transfer_s,
        )
        self.offload_bytes += size_gb * 1e9
        self.host_transfer_latencies_ms.append(transfer_s * 1000.0)
        self.peak_host_gb = max(self.peak_host_gb, self._used_host_gb())
        self._record_timeline()
        return True

    def _restore_host_entry(self, replica: ReplicaState, session_id: int) -> tuple[bool, float]:
        entry = self.host_cache.get(session_id)
        if entry is None:
            return False, 0.0
        if not self._ensure_capacity(replica, session_id, entry.size_gb):
            return False, 0.0
        wait_s = max(0.0, entry.available_time - self.now)
        copy_s = self._host_transfer_seconds(entry.size_gb)
        # The host entry remains resident while an unfinished offload or restore
        # is exposed. Account for that residency explicitly because the event
        # clock advances only when the turn completes.
        self.host_gb_seconds += entry.size_gb * (wait_s + copy_s)
        self._integrate_memory(self.now)
        del self.host_cache[session_id]
        self.restore_bytes += entry.size_gb * 1e9
        self.host_transfer_latencies_ms.append((wait_s + copy_s) * 1000.0)
        self.host_cache_hits += 1
        self._record_timeline()
        return True, wait_s + copy_s

    def _ensure_capacity(self, replica: ReplicaState, session_id: int, target_gb: float) -> bool:
        current = replica.cache.get(session_id)
        current_gb = current.size_gb if current else 0.0
        additional = max(0.0, target_gb - current_gb)
        if additional <= 1e-12:
            return True
        # Capacity is per replica; evict least-recently-used inactive session KV.
        while sum(e.size_gb for e in replica.cache.values()) + additional > self.kv_capacity_gb + 1e-12:
            victims = [e for sid, e in replica.cache.items() if sid != session_id]
            if not victims:
                return False
            victim = min(victims, key=lambda e: (e.last_access, e.session_id))
            self._remove_cache(replica, victim.session_id, "pressure")
        return True

    def _put_cache(self, replica: ReplicaState, session_id: int, tokens: int, keep: bool) -> bool:
        size_gb = tokens * self.latency.kv_bytes_per_token() / 1e9
        if size_gb > self.kv_capacity_gb + 1e-12:
            return False
        if not self._ensure_capacity(replica, session_id, size_gb):
            return False
        self._integrate_memory(self.now)
        previous = replica.cache.get(session_id)
        generation = (previous.generation + 1) if previous else 1
        expiry: float | None = None
        if keep and self.cfg.retention_policy in {"ttl", "adaptive"}:
            expiry = self.now + self.cfg.kv_ttl_s
        replica.cache[session_id] = CacheEntry(session_id, tokens, size_gb, self.now, expiry, generation)
        self.peak_kv_gb = max(self.peak_kv_gb, self._used_gb())
        self.peak_replica_kv_gb = max(
            self.peak_replica_kv_gb,
            sum(entry.size_gb for entry in replica.cache.values()),
        )
        if expiry is not None:
            self._push(expiry, "cache_expire", replica.replica_id, session_id, generation)
        self._record_timeline()
        return True

    def _start_next(self, replica: ReplicaState) -> None:
        if replica.busy or not replica.queue:
            return
        session_id, turn, ready_time = replica.queue.pop(0)
        runtime = self.sessions[session_id]
        cache = self._entry(replica, session_id)
        hbm_hit = turn.turn_index > 0 and cache is not None
        host_hit = False
        restore_s = 0.0
        if turn.turn_index > 0:
            self.cache_eligible_turns += 1
        if hbm_hit:
            self.cache_hits += 1
            prefill_tokens = max(turn.append_tokens, 1)
            cache_source = "hbm"
        elif turn.turn_index > 0 and session_id in self.host_cache:
            host_hit, restore_s = self._restore_host_entry(replica, session_id)
            if host_hit:
                prefill_tokens = max(turn.append_tokens, 1)
                cache_source = "host"
            else:
                prefill_tokens = max(runtime.context_tokens + turn.append_tokens, 1)
                self.recompute_tokens += runtime.context_tokens
                cache_source = "miss"
        else:
            prefill_tokens = max(runtime.context_tokens + turn.append_tokens, 1)
            if turn.turn_index > 0:
                self.recompute_tokens += runtime.context_tokens
            cache_source = "miss"

        context_before_decode = runtime.context_tokens + turn.append_tokens
        projected_tokens = context_before_decode + turn.output_tokens
        projected_gb = projected_tokens * self.latency.kv_bytes_per_token() / 1e9
        if not self._ensure_capacity(replica, session_id, projected_gb):
            self.failed_turns += 1
            self._push(self.now, "turn_failed", session_id, turn.turn_index, replica.replica_id, ready_time)
            self._start_next(replica)
            return

        # Model the current turn's KV as resident while it executes, even when
        # the selected policy will evict it immediately after the turn.
        if not self._put_cache(replica, session_id, projected_tokens, keep=False):
            self.failed_turns += 1
            self._push(self.now, "turn_failed", session_id, turn.turn_index, replica.replica_id, ready_time)
            self._start_next(replica)
            return

        prefill_s = self.latency.prefill_seconds([prefill_tokens])
        first_decode_s = self.latency.decode_step_seconds([context_before_decode])
        midpoint = context_before_decode + max(turn.output_tokens // 2, 1)
        decode_step_s = self.latency.decode_step_seconds([midpoint])
        service_s = restore_s + prefill_s + turn.output_tokens * decode_step_s
        ttft_s = (self.now - ready_time) + restore_s + prefill_s + first_decode_s

        replica.busy = True
        replica.current_session = session_id
        replica.busy_until = self.now + service_s
        self._push(
            replica.busy_until,
            "turn_complete",
            session_id,
            turn.turn_index,
            replica.replica_id,
            ready_time,
            hbm_hit or host_hit,
            cache_source,
            restore_s,
            prefill_tokens,
            ttft_s,
            service_s,
            projected_tokens,
        )

    def _finish_turn(
        self,
        session_id: int,
        turn_index: int,
        replica_id: int,
        ready_time: float,
        cache_hit: bool,
        cache_source: str,
        restore_s: float,
        prefill_tokens: int,
        ttft_s: float,
        service_s: float,
        projected_tokens: int,
    ) -> None:
        replica = self.replicas[replica_id]
        runtime = self.sessions[session_id]
        turn = runtime.spec.turns[turn_index]
        runtime.context_tokens = projected_tokens
        runtime.completed_turns += 1
        runtime.last_replica = replica_id
        replica.busy = False
        replica.current_session = None
        replica.busy_until = self.now

        is_final = turn_index == len(runtime.spec.turns) - 1
        state_action = "evict"
        if is_final:
            self._remove_cache(replica, session_id, "complete")
            self._remove_host_cache(session_id)
        elif self.cfg.retention_policy == "evict":
            self._remove_cache(replica, session_id, "complete")
            self._remove_host_cache(session_id)
        elif self.cfg.retention_policy in {"retain", "ttl"}:
            self._put_cache(replica, session_id, projected_tokens, keep=True)
            state_action = "retain_hbm"
        elif self.cfg.retention_policy == "offload":
            if self._offload_cache(replica, session_id, projected_tokens):
                state_action = "offload_host"
            else:
                state_action = "evict_host_full"
        elif self.cfg.retention_policy == "gap_aware":
            # This is intentionally an oracle upper-bound policy: the simulated
            # tool gap is already known from the generated program trace.
            if turn.tool_gap_after_s <= self.cfg.gap_aware_threshold_s:
                self._put_cache(replica, session_id, projected_tokens, keep=True)
                state_action = "retain_hbm_short_gap"
            elif self._offload_cache(replica, session_id, projected_tokens):
                state_action = "offload_host_long_gap"
            else:
                state_action = "evict_host_full"
        elif self.cfg.retention_policy == "adaptive":
            predicted_gap_s, prediction_source, history_count = self.tool_gap_predictor.predict(turn.tool_kind)
            predicted_retain = predicted_gap_s <= self.cfg.gap_aware_threshold_s
            oracle_retain = turn.tool_gap_after_s <= self.cfg.gap_aware_threshold_s
            if predicted_retain:
                self._put_cache(replica, session_id, projected_tokens, keep=True)
                state_action = "adaptive_retain_hbm"
            elif self._offload_cache(replica, session_id, projected_tokens):
                state_action = "adaptive_offload_host"
            else:
                state_action = "adaptive_evict_host_full"
            self.prediction_rows.append({
                "session_id": session_id,
                "turn_index": turn_index + 1,
                "tool_kind": turn.tool_kind,
                "shifted_regime": turn.shifted_regime,
                "predicted_gap_s": predicted_gap_s,
                "actual_gap_s": turn.tool_gap_after_s,
                "absolute_error_s": abs(predicted_gap_s - turn.tool_gap_after_s),
                "prediction_source": prediction_source,
                "history_count": history_count,
                "predicted_action": "retain" if predicted_retain else "offload",
                "oracle_action": "retain" if oracle_retain else "offload",
                "action_match": predicted_retain == oracle_retain,
            })

        e2e_ms = (self.now - ready_time) * 1000.0
        self.turn_rows.append(
            {
                "session_id": session_id,
                "turn_index": turn_index + 1,
                "replica": replica_id,
                "ready_time": ready_time,
                "completion_time": self.now,
                "cache_hit": cache_hit,
                "cache_source": cache_source,
                "restore_ms": restore_s * 1000.0,
                "state_action": state_action,
                "tool_kind": turn.tool_kind,
                "tool_gap_after_s": turn.tool_gap_after_s,
                "shifted_regime": turn.shifted_regime,
                "prefill_tokens": prefill_tokens,
                "context_tokens_after": projected_tokens,
                "output_tokens": turn.output_tokens,
                "ttft_ms": ttft_s * 1000.0,
                "e2e_ms": e2e_ms,
                "queue_ms": max(0.0, e2e_ms - service_s * 1000.0),
            }
        )

        if is_final:
            runtime.completion_time = self.now
        else:
            self.tool_gap_total_s += turn.tool_gap_after_s
            observed_at = self.now + turn.tool_gap_after_s
            self._push(observed_at, "tool_observed", turn.tool_kind, turn.tool_gap_after_s)
            self._push(observed_at, "turn_ready", session_id, turn_index + 1)
        self._start_next(replica)

    def run(self) -> dict[str, Any]:
        self._record_timeline()
        while self.events:
            event_time, _, kind, payload = heapq.heappop(self.events)
            self._integrate_memory(event_time)
            self.now = event_time
            if kind == "tool_observed":
                self.tool_gap_predictor.observe(str(payload[0]), float(payload[1]))
            elif kind == "turn_ready":
                session_id, turn_index = int(payload[0]), int(payload[1])
                turn = self.sessions[session_id].spec.turns[turn_index]
                replica = self._route(session_id)
                replica.queue.append((session_id, turn, self.now))
                self._start_next(replica)
            elif kind == "turn_complete":
                self._finish_turn(
                    int(payload[0]), int(payload[1]), int(payload[2]), float(payload[3]), bool(payload[4]),
                    str(payload[5]), float(payload[6]), int(payload[7]), float(payload[8]), float(payload[9]), int(payload[10]),
                )
            elif kind == "cache_expire":
                replica_id, session_id, generation = map(int, payload)
                replica = self.replicas[replica_id]
                entry = replica.cache.get(session_id)
                if entry is not None and entry.generation == generation and entry.expiry_time is not None and entry.expiry_time <= self.now + 1e-12:
                    self._remove_cache(replica, session_id, "ttl")
            elif kind == "turn_failed":
                session_id, turn_index, replica_id, ready_time = int(payload[0]), int(payload[1]), int(payload[2]), float(payload[3])
                self.turn_rows.append({
                    "session_id": session_id, "turn_index": turn_index + 1, "replica": replica_id,
                    "ready_time": ready_time, "completion_time": self.now, "cache_hit": False,
                    "prefill_tokens": 0, "context_tokens_after": self.sessions[session_id].context_tokens,
                    "output_tokens": 0, "ttft_ms": 0.0, "e2e_ms": 0.0, "queue_ms": 0.0, "failed": True,
                })
            self._record_timeline()

        self._integrate_memory(self.now)
        self._record_timeline()
        completed_sessions = [s for s in self.sessions.values() if s.completion_time is not None]
        successful_turns = [row for row in self.turn_rows if not row.get("failed")]
        ttfts = [float(row["ttft_ms"]) for row in successful_turns]
        turn_e2e = [float(row["e2e_ms"]) for row in successful_turns]
        session_e2e = [
            (s.completion_time - s.spec.arrival_time) * 1000.0
            for s in completed_sessions
            if s.completion_time is not None
        ]
        session_slo = sum(1 for v in session_e2e if v <= self.cfg.slo_session_e2e_ms)
        turn_slo = sum(1 for v in ttfts if v <= self.cfg.slo_turn_ttft_ms)
        turns_generated = sum(len(s.spec.turns) for s in self.sessions.values())
        completion_horizon = max([s.completion_time or 0.0 for s in self.sessions.values()] + [self.cfg.duration_s, 1e-9])
        mean_kv_gb = self.hbm_gb_seconds / completion_horizon
        summary = {
            "sessions_generated": len(self.sessions),
            "sessions_completed": len(completed_sessions),
            "turns_generated": turns_generated,
            "turns_completed": len(successful_turns),
            "turns_failed": self.failed_turns,
            "session_completion_rate": len(completed_sessions) / len(self.sessions) if self.sessions else 0.0,
            "turn_completion_rate": len(successful_turns) / turns_generated if turns_generated else 0.0,
            "session_throughput_rps": len(completed_sessions) / completion_horizon,
            "turn_throughput_rps": len(successful_turns) / completion_horizon,
            "turn_ttft_slo_attainment": turn_slo / turns_generated if turns_generated else 0.0,
            "session_slo_attainment": session_slo / len(self.sessions) if self.sessions else 0.0,
            "simulated_makespan_s": completion_horizon,
        }
        latency = {
            "turn_ttft_ms": {"p50": percentile(ttfts, 0.50), "p95": percentile(ttfts, 0.95), "p99": percentile(ttfts, 0.99)},
            "turn_e2e_ms": {"p50": percentile(turn_e2e, 0.50), "p95": percentile(turn_e2e, 0.95), "p99": percentile(turn_e2e, 0.99)},
            "session_e2e_ms": {"p50": percentile(session_e2e, 0.50), "p95": percentile(session_e2e, 0.95), "p99": percentile(session_e2e, 0.99)},
        }
        host_restore_p95_ms = percentile(self.host_transfer_latencies_ms, 0.95)
        total_reuse_hits = self.cache_hits + self.host_cache_hits
        prediction_errors = [float(row["absolute_error_s"]) for row in self.prediction_rows]
        prediction_matches = [bool(row["action_match"]) for row in self.prediction_rows]
        resource = {
            "replicas": self.cfg.replicas,
            "kv_capacity_gb_per_replica": self.kv_capacity_gb,
            "peak_kv_gb": self.peak_kv_gb,
            "peak_replica_kv_gb": self.peak_replica_kv_gb,
            "mean_kv_gb": mean_kv_gb,
            "hbm_gb_seconds": self.hbm_gb_seconds,
            "peak_host_kv_gb": self.peak_host_gb,
            "mean_host_kv_gb": self.host_gb_seconds / completion_horizon,
            "host_gb_seconds": self.host_gb_seconds,
            "cross_turn_cache_hits": total_reuse_hits,
            "hbm_cache_hits": self.cache_hits,
            "host_cache_hits": self.host_cache_hits,
            "cross_turn_cache_eligible": self.cache_eligible_turns,
            "cross_turn_cache_hit_rate": total_reuse_hits / self.cache_eligible_turns if self.cache_eligible_turns else 0.0,
            "hbm_cache_hit_rate": self.cache_hits / self.cache_eligible_turns if self.cache_eligible_turns else 0.0,
            "host_cache_hit_rate": self.host_cache_hits / self.cache_eligible_turns if self.cache_eligible_turns else 0.0,
            "routing_locality_rate": self.affinity_routes / self.route_opportunities if self.route_opportunities else 0.0,
            "recomputed_history_tokens": self.recompute_tokens,
            "pressure_evictions": self.pressure_evictions,
            "ttl_evictions": self.ttl_evictions,
            "host_pressure_evictions": self.host_pressure_evictions,
            "offloaded_gb": self.offload_bytes / 1e9,
            "restored_gb": self.restore_bytes / 1e9,
            "p95_host_transfer_ms": host_restore_p95_ms,
            "tool_gap_total_s": self.tool_gap_total_s,
            "adaptive_prediction_count": len(self.prediction_rows),
            "adaptive_prediction_mae_s": mean(prediction_errors) if prediction_errors else 0.0,
            "adaptive_prediction_p95_abs_error_s": percentile(prediction_errors, 0.95),
            "adaptive_oracle_action_agreement": (
                sum(1 for match in prediction_matches if match) / len(prediction_matches)
                if prediction_matches else 0.0
            ),
        }
        return {
            "config": self.cfg.to_dict(),
            "provenance": {
                "simulator": "InferScale-Sim",
                "mode": "stateful-agent-session-simulation",
                "latency_profile_type": "analytical-reference",
                "agent_service_model": "serial-per-replica-reference",
                "host_tier_model": "serialized-reference-transfer",
                "gap_aware_policy": "oracle-upper-bound" if self.cfg.retention_policy == "gap_aware" else "not-active",
                "adaptive_policy": (
                    "online-tool-gap-ewma-no-lookahead" if self.cfg.retention_policy == "adaptive" else "not-active"
                ),
                "warning": "Agent-session mode isolates routing/KV-retention effects and does not model dynamic batching within each replica.",
            },
            "summary": summary,
            "latency": latency,
            "resource": resource,
            "turns": successful_turns[:3000],
            "prediction": {
                "rows": self.prediction_rows[:3000],
                "predictor": self.tool_gap_predictor.snapshot(),
            },
            "sessions": [
                {
                    "session_id": s.spec.session_id,
                    "arrival_time": s.spec.arrival_time,
                    "turns": len(s.spec.turns),
                    "completion_time": s.completion_time,
                    "e2e_ms": (s.completion_time - s.spec.arrival_time) * 1000.0 if s.completion_time is not None else None,
                }
                for s in list(self.sessions.values())[:1000]
            ],
            "timeline": self.timeline,
        }


def run_agent_session_simulation(config: dict[str, Any], sessions: list[SessionSpec] | None = None) -> dict[str, Any]:
    cfg = AgentSessionConfig.from_dict(config)
    return AgentSessionSimulator(cfg, sessions=sessions).run()


def _same_trace(cfg: AgentSessionConfig) -> list[SessionSpec]:
    return generate_agent_sessions(cfg)


def compare_agent_policies(config: dict[str, Any]) -> dict[str, Any]:
    base = AgentSessionConfig.from_dict(config)
    trace = _same_trace(base)
    policies = [
        ("Stateless / least-load", "evict", "least_load", 0.0),
        ("Retain / least-load", "retain", "least_load", base.kv_ttl_s),
        ("TTL / affinity", "ttl", "session_affinity", base.kv_ttl_s),
        ("Retain / affinity", "retain", "session_affinity", base.kv_ttl_s),
    ]
    rows = []
    results = []
    for label, retention, routing, ttl in policies:
        cfg = AgentSessionConfig.from_dict(base.to_dict())
        cfg.retention_policy = retention
        cfg.routing_policy = routing
        cfg.kv_ttl_s = ttl
        result = run_agent_session_simulation(cfg.to_dict(), trace)
        results.append(result)
        rows.append(
            {
                "label": label,
                "retention": retention,
                "routing": routing,
                "p95_turn_ttft_ms": result["latency"]["turn_ttft_ms"]["p95"],
                "p95_session_e2e_ms": result["latency"]["session_e2e_ms"]["p95"],
                "cache_hit_rate": result["resource"]["cross_turn_cache_hit_rate"],
                "routing_locality_rate": result["resource"]["routing_locality_rate"],
                "recomputed_history_tokens": result["resource"]["recomputed_history_tokens"],
                "peak_kv_gb": result["resource"]["peak_kv_gb"],
                "mean_kv_gb": result["resource"]["mean_kv_gb"],
                "hbm_gb_seconds": result["resource"]["hbm_gb_seconds"],
                "pressure_evictions": result["resource"]["pressure_evictions"],
                "ttl_evictions": result["resource"]["ttl_evictions"],
                "sessions_completed": result["summary"]["sessions_completed"],
            }
        )
    return {"protocol": "common-agent-program-trace", "candidate_count": len(rows), "rows": rows, "results": results}


def _pareto(rows: list[dict[str, Any]], x: str, y: str) -> set[int]:
    # Both x and y are minimized.
    front: set[int] = set()
    for idx, row in enumerate(rows):
        dominated = False
        for jdx, other in enumerate(rows):
            if idx == jdx:
                continue
            if other[x] <= row[x] and other[y] <= row[y] and (other[x] < row[x] or other[y] < row[y]):
                dominated = True
                break
        if not dominated:
            front.add(idx)
    return front


def ttl_retention_sweep(config: dict[str, Any], ttl_values: list[float] | None = None) -> dict[str, Any]:
    base = AgentSessionConfig.from_dict(config)
    base.retention_policy = "ttl"
    base.routing_policy = "session_affinity"
    trace = _same_trace(base)
    values = ttl_values or [0.0, 0.25, 0.5, 1.0, 2.0, 3.0, 5.0, 8.0, 13.0]
    cleaned = sorted({max(0.0, min(float(v), 120.0)) for v in values})
    rows = []
    for ttl in cleaned:
        cfg = AgentSessionConfig.from_dict(base.to_dict())
        cfg.kv_ttl_s = ttl
        result = run_agent_session_simulation(cfg.to_dict(), trace)
        rows.append(
            {
                "ttl_s": ttl,
                "p95_turn_ttft_ms": result["latency"]["turn_ttft_ms"]["p95"],
                "p95_session_e2e_ms": result["latency"]["session_e2e_ms"]["p95"],
                "cache_hit_rate": result["resource"]["cross_turn_cache_hit_rate"],
                "recomputed_history_tokens": result["resource"]["recomputed_history_tokens"],
                "mean_kv_gb": result["resource"]["mean_kv_gb"],
                "hbm_gb_seconds": result["resource"]["hbm_gb_seconds"],
                "pressure_evictions": result["resource"]["pressure_evictions"],
                "ttl_evictions": result["resource"]["ttl_evictions"],
            }
        )
    front = _pareto(rows, "p95_turn_ttft_ms", "mean_kv_gb")
    # Collapse numerically equivalent frontier points to the shortest TTL. A
    # longer retention horizon with indistinguishable latency/residency is not a
    # distinct engineering trade-off.
    unique_front: set[int] = set()
    seen: set[tuple[float, float]] = set()
    for idx in sorted(front, key=lambda i: rows[i]["ttl_s"]):
        key = (round(rows[idx]["p95_turn_ttft_ms"], 6), round(rows[idx]["mean_kv_gb"], 6))
        if key not in seen:
            unique_front.add(idx)
            seen.add(key)
    for idx, row in enumerate(rows):
        row["pareto"] = idx in unique_front
    return {
        "protocol": "common-agent-program-trace",
        "objective": "minimize-p95-turn-ttft-and-mean-kv-residency",
        "rows": rows,
        "pareto_count": len(unique_front),
    }


def compare_agent_memory_policies(config: dict[str, Any]) -> dict[str, Any]:
    """Compare cache-residency and routing strategies on one common agent trace."""
    base = AgentSessionConfig.from_dict(config)
    trace = _same_trace(base)
    policies = [
        ("Stateless / least-load", "evict", "least_load"),
        ("TTL / strict affinity", "ttl", "session_affinity"),
        ("TTL / bounded affinity", "ttl", "bounded_affinity"),
        ("Host offload / bounded affinity", "offload", "bounded_affinity"),
        ("Gap-aware tiering / bounded affinity", "gap_aware", "bounded_affinity"),
    ]
    rows: list[dict[str, Any]] = []
    results: list[dict[str, Any]] = []
    for label, retention, routing in policies:
        cfg = AgentSessionConfig.from_dict(base.to_dict())
        cfg.retention_policy = retention
        cfg.routing_policy = routing
        result = run_agent_session_simulation(cfg.to_dict(), trace)
        results.append(result)
        rows.append(
            {
                "label": label,
                "retention": retention,
                "routing": routing,
                "p95_turn_ttft_ms": result["latency"]["turn_ttft_ms"]["p95"],
                "p95_session_e2e_ms": result["latency"]["session_e2e_ms"]["p95"],
                "turn_slo_attainment": result["summary"]["turn_ttft_slo_attainment"],
                "session_slo_attainment": result["summary"]["session_slo_attainment"],
                "cache_hit_rate": result["resource"]["cross_turn_cache_hit_rate"],
                "hbm_hit_rate": result["resource"]["hbm_cache_hit_rate"],
                "host_hit_rate": result["resource"]["host_cache_hit_rate"],
                "routing_locality_rate": result["resource"]["routing_locality_rate"],
                "recomputed_history_tokens": result["resource"]["recomputed_history_tokens"],
                "mean_hbm_gb": result["resource"]["mean_kv_gb"],
                "mean_host_gb": result["resource"]["mean_host_kv_gb"],
                "hbm_gb_seconds": result["resource"]["hbm_gb_seconds"],
                "host_gb_seconds": result["resource"]["host_gb_seconds"],
                "p95_host_transfer_ms": result["resource"]["p95_host_transfer_ms"],
                "offloaded_gb": result["resource"]["offloaded_gb"],
                "pressure_evictions": result["resource"]["pressure_evictions"],
                "host_pressure_evictions": result["resource"]["host_pressure_evictions"],
                "turns_failed": result["summary"]["turns_failed"],
            }
        )
    return {
        "protocol": "common-agent-program-trace",
        "study": "agent-memory-tiering",
        "candidate_count": len(rows),
        "rows": rows,
        "results": results,
        "note": "Gap-aware tiering uses realized simulated tool gaps and is an oracle upper bound, not a deployable predictor.",
    }


def agent_memory_budget_sweep(
    config: dict[str, Any], budget_multipliers: list[float] | None = None
) -> dict[str, Any]:
    """Stress state policies under finite per-replica HBM KV budgets.

    Budgets are derived from the unconstrained peak working set of the exact same
    program trace, avoiding arbitrary fractions of total GPU VRAM that would be
    too loose for small models.
    """
    base = AgentSessionConfig.from_dict(config)
    trace = _same_trace(base)

    reference_cfg = AgentSessionConfig.from_dict(base.to_dict())
    reference_cfg.retention_policy = "retain"
    reference_cfg.routing_policy = "session_affinity"
    reference_cfg.kv_capacity_override_gb = 0.0
    reference = run_agent_session_simulation(reference_cfg.to_dict(), trace)
    reference_peak = max(float(reference["resource"]["peak_replica_kv_gb"]), 0.002)

    multipliers = budget_multipliers or [0.35, 0.50, 0.75, 1.00, 1.50]
    cleaned = sorted({max(0.10, min(float(v), 3.0)) for v in multipliers})
    policies = [
        ("TTL / bounded affinity", "ttl", "bounded_affinity"),
        ("Host offload / bounded affinity", "offload", "bounded_affinity"),
        ("Gap-aware tiering / bounded affinity", "gap_aware", "bounded_affinity"),
    ]

    rows: list[dict[str, Any]] = []
    for multiplier in cleaned:
        budget = reference_peak * multiplier
        for label, retention, routing in policies:
            cfg = AgentSessionConfig.from_dict(base.to_dict())
            cfg.retention_policy = retention
            cfg.routing_policy = routing
            cfg.kv_capacity_override_gb = budget
            result = run_agent_session_simulation(cfg.to_dict(), trace)
            rows.append(
                {
                    "policy": label,
                    "budget_multiplier": multiplier,
                    "budget_gb_per_replica": budget,
                    "p95_turn_ttft_ms": result["latency"]["turn_ttft_ms"]["p95"],
                    "p95_session_e2e_ms": result["latency"]["session_e2e_ms"]["p95"],
                    "turn_slo_attainment": result["summary"]["turn_ttft_slo_attainment"],
                    "session_slo_attainment": result["summary"]["session_slo_attainment"],
                    "cache_hit_rate": result["resource"]["cross_turn_cache_hit_rate"],
                    "recomputed_history_tokens": result["resource"]["recomputed_history_tokens"],
                    "mean_hbm_gb": result["resource"]["mean_kv_gb"],
                    "mean_host_gb": result["resource"]["mean_host_kv_gb"],
                    "p95_host_transfer_ms": result["resource"]["p95_host_transfer_ms"],
                    "pressure_evictions": result["resource"]["pressure_evictions"],
                    "host_pressure_evictions": result["resource"]["host_pressure_evictions"],
                    "turns_failed": result["summary"]["turns_failed"],
                }
            )
    return {
        "protocol": "common-agent-program-trace",
        "study": "finite-hbm-budget-stress",
        "reference_peak_replica_kv_gb": reference_peak,
        "rows": rows,
        "policy_count": len(policies),
        "budget_count": len(cleaned),
    }


def agent_affinity_sweep(config: dict[str, Any], slack_values_ms: list[float] | None = None) -> dict[str, Any]:
    """Sweep how much queue imbalance the router tolerates for KV locality."""
    base = AgentSessionConfig.from_dict(config)
    base.retention_policy = "ttl"
    base.routing_policy = "bounded_affinity"
    trace = _same_trace(base)
    values = slack_values_ms or [0.0, 25.0, 75.0, 150.0, 300.0, 600.0, 1200.0]
    cleaned = sorted({max(0.0, min(float(v), 5000.0)) for v in values})
    rows: list[dict[str, Any]] = []
    for slack in cleaned:
        cfg = AgentSessionConfig.from_dict(base.to_dict())
        cfg.affinity_slack_ms = slack
        result = run_agent_session_simulation(cfg.to_dict(), trace)
        rows.append(
            {
                "affinity_slack_ms": slack,
                "p95_turn_ttft_ms": result["latency"]["turn_ttft_ms"]["p95"],
                "p95_session_e2e_ms": result["latency"]["session_e2e_ms"]["p95"],
                "cache_hit_rate": result["resource"]["cross_turn_cache_hit_rate"],
                "routing_locality_rate": result["resource"]["routing_locality_rate"],
                "recomputed_history_tokens": result["resource"]["recomputed_history_tokens"],
                "turn_slo_attainment": result["summary"]["turn_ttft_slo_attainment"],
                "session_slo_attainment": result["summary"]["session_slo_attainment"],
                "pressure_evictions": result["resource"]["pressure_evictions"],
            }
        )
    return {
        "protocol": "common-agent-program-trace",
        "study": "bounded-affinity-routing-frontier",
        "rows": rows,
    }


def _prediction_phase_metrics(result: dict[str, Any]) -> dict[str, float]:
    rows = result.get("prediction", {}).get("rows", [])
    before = [row for row in rows if not row.get("shifted_regime")]
    after = [row for row in rows if row.get("shifted_regime")]

    def phase(items: list[dict[str, Any]]) -> tuple[float, float]:
        if not items:
            return 0.0, 0.0
        mae = mean(float(row["absolute_error_s"]) for row in items)
        agreement = sum(1 for row in items if row.get("action_match")) / len(items)
        return mae, agreement

    pre_mae, pre_agreement = phase(before)
    post_mae, post_agreement = phase(after)
    return {
        "pre_shift_mae_s": pre_mae,
        "post_shift_mae_s": post_mae,
        "pre_shift_action_agreement": pre_agreement,
        "post_shift_action_agreement": post_agreement,
    }


def _rolling_prediction_curve(result: dict[str, Any], window: int = 16) -> list[dict[str, Any]]:
    rows = result.get("prediction", {}).get("rows", [])
    if not rows:
        return []
    points: list[dict[str, Any]] = []
    width = max(4, min(int(window), 64))
    for idx in range(len(rows)):
        start = max(0, idx - width + 1)
        chunk = rows[start : idx + 1]
        points.append(
            {
                "observation": idx + 1,
                "rolling_mae_s": mean(float(row["absolute_error_s"]) for row in chunk),
                "rolling_action_agreement": sum(1 for row in chunk if row.get("action_match")) / len(chunk),
                "tool_kind": rows[idx].get("tool_kind", "generic"),
                "shifted_regime": bool(rows[idx].get("shifted_regime")),
            }
        )
    return points


def adaptive_tiering_study(
    config: dict[str, Any],
    *,
    horizon_s: float = 120.0,
    shift_fraction: float = 0.55,
    shift_multiplier: float = 2.5,
    alpha: float = 0.30,
) -> dict[str, Any]:
    """Compare fixed, adaptive, and oracle KV tiering on one non-stationary trace.

    The adaptive candidates never inspect the realized future tool duration at
    decision time. They learn an EWMA online from completed tool calls. The
    oracle candidate intentionally sees the realized gap and serves only as an
    upper bound.
    """
    base = AgentSessionConfig.from_dict(config)
    base.duration_s = max(float(horizon_s), 30.0)
    base.tool_regime_shift_fraction = min(max(float(shift_fraction), 0.05), 0.95)
    base.tool_regime_shift_multiplier = max(float(shift_multiplier), 1.0)
    base.adaptive_alpha = min(max(float(alpha), 0.01), 1.0)
    base.routing_policy = "bounded_affinity"
    trace = generate_agent_sessions(base)

    candidates = [
        ("Fixed TTL", "ttl", "per_tool_ema"),
        ("Always host offload", "offload", "per_tool_ema"),
        ("Adaptive global EWMA", "adaptive", "global_ema"),
        ("Adaptive per-tool EWMA", "adaptive", "per_tool_ema"),
        ("Oracle gap-aware", "gap_aware", "per_tool_ema"),
    ]
    rows: list[dict[str, Any]] = []
    results: dict[str, dict[str, Any]] = {}
    for label, retention, scope in candidates:
        cfg = AgentSessionConfig.from_dict(base.to_dict())
        cfg.retention_policy = retention
        cfg.adaptive_predictor_scope = scope
        result = run_agent_session_simulation(cfg.to_dict(), trace)
        results[label] = result
        phase = _prediction_phase_metrics(result)
        rows.append(
            {
                "label": label,
                "retention": retention,
                "predictor_scope": scope if retention == "adaptive" else "not-active",
                "p95_turn_ttft_ms": result["latency"]["turn_ttft_ms"]["p95"],
                "p95_session_e2e_ms": result["latency"]["session_e2e_ms"]["p95"],
                "session_slo_attainment": result["summary"]["session_slo_attainment"],
                "cache_hit_rate": result["resource"]["cross_turn_cache_hit_rate"],
                "mean_hbm_gb": result["resource"]["mean_kv_gb"],
                "mean_host_gb": result["resource"]["mean_host_kv_gb"],
                "recomputed_history_tokens": result["resource"]["recomputed_history_tokens"],
                "prediction_count": result["resource"].get("adaptive_prediction_count", 0),
                "prediction_mae_s": result["resource"].get("adaptive_prediction_mae_s", 0.0),
                "oracle_action_agreement": result["resource"].get("adaptive_oracle_action_agreement", 0.0),
                **phase,
            }
        )

    global_result = results["Adaptive global EWMA"]
    tool_result = results["Adaptive per-tool EWMA"]
    shift_observation = next(
        (
            idx + 1
            for idx, row in enumerate(tool_result.get("prediction", {}).get("rows", []))
            if row.get("shifted_regime")
        ),
        0,
    )
    return {
        "protocol": "common-nonstationary-agent-program-trace",
        "study": "online-predictive-kv-tiering",
        "horizon_s": base.duration_s,
        "shift_fraction": base.tool_regime_shift_fraction,
        "shift_multiplier": base.tool_regime_shift_multiplier,
        "shift_time_s": base.duration_s * base.tool_regime_shift_fraction,
        "shift_observation": shift_observation,
        "threshold_s": base.gap_aware_threshold_s,
        "alpha": base.adaptive_alpha,
        "rows": rows,
        "learning_curves": {
            "global": _rolling_prediction_curve(global_result),
            "per_tool": _rolling_prediction_curve(tool_result),
        },
        "tool_profiles": TOOL_GAP_MULTIPLIERS,
        "note": (
            "The oracle gap-aware candidate sees realized future tool gaps and is an upper bound. "
            "Adaptive candidates learn only from tool calls that have already completed."
        ),
    }


def adaptive_alpha_sweep(
    config: dict[str, Any],
    alpha_values: list[float] | None = None,
    *,
    horizon_s: float = 120.0,
    shift_fraction: float = 0.55,
    shift_multiplier: float = 2.5,
) -> dict[str, Any]:
    """Sweep EWMA adaptation speed on one shifted trace.

    Lower alpha values are stable but adapt slowly; high values react quickly but
    are noisier. The study reports both pre- and post-shift prediction error and
    the resulting serving metrics.
    """
    base = AgentSessionConfig.from_dict(config)
    base.duration_s = max(float(horizon_s), 30.0)
    base.tool_regime_shift_fraction = min(max(float(shift_fraction), 0.05), 0.95)
    base.tool_regime_shift_multiplier = max(float(shift_multiplier), 1.0)
    base.retention_policy = "adaptive"
    base.routing_policy = "bounded_affinity"
    base.adaptive_predictor_scope = "per_tool_ema"
    trace = generate_agent_sessions(base)
    values = alpha_values or [0.05, 0.10, 0.20, 0.30, 0.50, 0.75, 1.00]
    cleaned = sorted({min(max(float(value), 0.01), 1.0) for value in values})
    rows: list[dict[str, Any]] = []
    for alpha in cleaned:
        cfg = AgentSessionConfig.from_dict(base.to_dict())
        cfg.adaptive_alpha = alpha
        result = run_agent_session_simulation(cfg.to_dict(), trace)
        phase = _prediction_phase_metrics(result)
        rows.append(
            {
                "alpha": alpha,
                "prediction_mae_s": result["resource"]["adaptive_prediction_mae_s"],
                "oracle_action_agreement": result["resource"]["adaptive_oracle_action_agreement"],
                "p95_turn_ttft_ms": result["latency"]["turn_ttft_ms"]["p95"],
                "p95_session_e2e_ms": result["latency"]["session_e2e_ms"]["p95"],
                "session_slo_attainment": result["summary"]["session_slo_attainment"],
                "mean_hbm_gb": result["resource"]["mean_kv_gb"],
                "mean_host_gb": result["resource"]["mean_host_kv_gb"],
                **phase,
            }
        )
    best_post = min(rows, key=lambda row: (row["post_shift_mae_s"], row["p95_turn_ttft_ms"])) if rows else None
    return {
        "protocol": "common-nonstationary-agent-program-trace",
        "study": "adaptive-ewma-rate-sweep",
        "shift_fraction": base.tool_regime_shift_fraction,
        "shift_multiplier": base.tool_regime_shift_multiplier,
        "rows": rows,
        "best_post_shift_alpha": best_post["alpha"] if best_post else None,
        "note": "Alpha controls adaptation speed; this is an online heuristic study, not a learned model benchmark.",
    }