src/inferscale/execution.py

from __future__ import annotations

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

from .latency import AnalyticalLatencyModel
from .metrics import percentile
from .profiles import get_accelerator, get_model


AGENT_ROLES = ("planner", "retriever", "reasoner", "verifier", "writer")
END = "END"
ROLE_PREFIX_TOKENS = {
    "planner": 512,
    "retriever": 768,
    "reasoner": 896,
    "verifier": 640,
    "writer": 576,
}

# Structured workflow transitions used to make online execution learning
# meaningful. Rows sum to one. The shifted matrix changes the mix of external
# retrieval/verification work without changing the model or hardware profile.
BASE_TRANSITIONS: dict[str, dict[str, float]] = {
    "planner": {"retriever": 0.52, "reasoner": 0.28, "writer": 0.12, END: 0.08},
    "retriever": {"reasoner": 0.50, "verifier": 0.20, "retriever": 0.10, "writer": 0.12, END: 0.08},
    "reasoner": {"verifier": 0.32, "retriever": 0.18, "writer": 0.30, "reasoner": 0.10, END: 0.10},
    "verifier": {"reasoner": 0.24, "retriever": 0.12, "writer": 0.34, END: 0.30},
    "writer": {"verifier": 0.22, "reasoner": 0.10, "retriever": 0.05, END: 0.63},
}
SHIFTED_TRANSITIONS: dict[str, dict[str, float]] = {
    "planner": {"retriever": 0.72, "reasoner": 0.14, "writer": 0.06, END: 0.08},
    "retriever": {"reasoner": 0.30, "verifier": 0.30, "retriever": 0.24, "writer": 0.08, END: 0.08},
    "reasoner": {"verifier": 0.24, "retriever": 0.36, "writer": 0.18, "reasoner": 0.12, END: 0.10},
    "verifier": {"reasoner": 0.18, "retriever": 0.28, "writer": 0.24, END: 0.30},
    "writer": {"verifier": 0.34, "reasoner": 0.08, "retriever": 0.10, END: 0.48},
}

ROLE_GAP_MULTIPLIERS = {
    "planner": 0.35,
    "retriever": 1.45,
    "reasoner": 0.55,
    "verifier": 0.75,
    "writer": 0.30,
}

PREFETCH_POLICIES = {"none", "cumulative", "decayed", "multistep", "utility", "oracle", "oracle_horizon"}


@dataclass
class ExecutionLearningConfig:
    model: str = "Qwen2.5-3B"
    accelerator: str = "L4"
    quantization: str = "int8"
    seed: int = 7
    duration_s: float = 120.0
    workflow_rate_rps: float = 0.18
    dynamic_prompt_tokens_mean: int = 320
    dynamic_prompt_cv: float = 0.40
    output_tokens_mean: int = 64
    output_tokens_cv: float = 0.45
    tool_gap_mean_s: float = 1.2
    tool_gap_cv: float = 0.65
    max_steps: int = 8
    shift_fraction: float = 0.55
    prefix_cache_budget_fraction: float = 0.62
    host_bandwidth_gbps: float = 32.0
    transfer_base_ms: float = 0.15
    prefetch_policy: str = "decayed"
    transition_decay: float = 0.85
    confidence_threshold: float = 0.55
    prior_strength: float = 0.35
    forecast_horizon: int = 3
    prefetch_top_k: int = 2
    forecast_discount: float = 0.75
    forecast_min_score: float = 0.10
    utility_threshold_ms: float = 0.0
    timeline_points: int = 240

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ExecutionLearningConfig":
        allowed = cls.__dataclass_fields__.keys()
        cfg = cls(**{key: data[key] for key in allowed if key in data})
        cfg.duration_s = max(float(cfg.duration_s), 20.0)
        cfg.workflow_rate_rps = max(float(cfg.workflow_rate_rps), 0.01)
        cfg.max_steps = max(2, min(int(cfg.max_steps), 16))
        cfg.shift_fraction = min(max(float(cfg.shift_fraction), 0.05), 0.95)
        cfg.prefix_cache_budget_fraction = min(max(float(cfg.prefix_cache_budget_fraction), 0.10), 1.50)
        cfg.host_bandwidth_gbps = max(float(cfg.host_bandwidth_gbps), 0.1)
        cfg.transfer_base_ms = max(float(cfg.transfer_base_ms), 0.0)
        cfg.transition_decay = min(max(float(cfg.transition_decay), 0.20), 1.0)
        cfg.confidence_threshold = min(max(float(cfg.confidence_threshold), 0.0), 1.0)
        cfg.prior_strength = max(float(cfg.prior_strength), 0.01)
        cfg.forecast_horizon = max(1, min(int(cfg.forecast_horizon), 6))
        cfg.prefetch_top_k = max(1, min(int(cfg.prefetch_top_k), len(AGENT_ROLES)))
        cfg.forecast_discount = min(max(float(cfg.forecast_discount), 0.05), 1.0)
        cfg.forecast_min_score = min(max(float(cfg.forecast_min_score), 0.0), 1.0)
        cfg.utility_threshold_ms = float(cfg.utility_threshold_ms)
        cfg.timeline_points = max(40, min(int(cfg.timeline_points), 1000))
        if cfg.prefetch_policy not in PREFETCH_POLICIES:
            raise ValueError(f"Unsupported execution prefetch policy: {cfg.prefetch_policy}")
        return cfg

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


@dataclass
class WorkflowStep:
    role: str
    dynamic_prompt_tokens: int
    output_tokens: int
    gap_after_s: float
    shifted_regime: bool


@dataclass
class WorkflowSpec:
    workflow_id: int
    arrival_time: float
    steps: list[WorkflowStep]


@dataclass
class PrefixEntry:
    role: str
    size_gb: float
    last_access: float
    available_time: float
    source: str
    generation: int
    used: bool = False


class OnlineTransitionPredictor:
    """First-order online execution model with optional exponential forgetting.

    The predictor never observes a transition before the next role is revealed.
    ``decay=1`` is cumulative counting. Smaller values forget older transitions
    and can adapt after a workflow-regime shift.
    """

    def __init__(self, *, decay: float = 0.85, prior_strength: float = 0.35) -> None:
        self.decay = min(max(float(decay), 0.20), 1.0)
        self.prior_strength = max(float(prior_strength), 0.01)
        self.targets = (*AGENT_ROLES, END)
        self.rows: dict[str, dict[str, float]] = {
            role: {target: self.prior_strength for target in self.targets}
            for role in AGENT_ROLES
        }
        self.observations = 0

    def predict(self, current_role: str) -> tuple[str, float, dict[str, float]]:
        row = self.rows[current_role]
        total = sum(row.values()) or 1.0
        probabilities = {target: value / total for target, value in row.items()}
        best = max(self.targets, key=lambda target: (probabilities[target], -self.targets.index(target)))
        return best, probabilities[best], probabilities

    def forecast(
        self,
        current_role: str,
        *,
        horizon: int = 3,
        discount: float = 0.75,
    ) -> dict[str, Any]:
        """Roll the learned first-order matrix forward without future trace access.

        ``role_scores`` are discounted expected future visits, not calibrated
        probabilities of at-least-one reuse. They are therefore used as a
        ranking signal and are normalized separately for UI/decision thresholds.
        """

        horizon = max(1, min(int(horizon), 8))
        discount = min(max(float(discount), 0.0), 1.0)
        distributions: list[dict[str, float]] = []
        current: dict[str, float] = {target: 0.0 for target in self.targets}
        current[current_role] = 1.0
        role_scores = {role: 0.0 for role in AGENT_ROLES}

        for depth in range(1, horizon + 1):
            nxt = {target: 0.0 for target in self.targets}
            for source, mass in current.items():
                if mass <= 0.0:
                    continue
                if source == END:
                    nxt[END] += mass
                    continue
                _, _, probs = self.predict(source)
                for target, probability in probs.items():
                    nxt[target] += mass * probability
            distributions.append(nxt)
            weight = discount ** (depth - 1)
            for role in AGENT_ROLES:
                role_scores[role] += weight * nxt.get(role, 0.0)
            current = nxt

        total_score = sum(role_scores.values())
        normalized = {
            role: (score / total_score if total_score > 1e-12 else 0.0)
            for role, score in role_scores.items()
        }
        ranked = sorted(AGENT_ROLES, key=lambda role: (-role_scores[role], role))
        return {
            "horizon": horizon,
            "discount": discount,
            "distributions": distributions,
            "role_scores": role_scores,
            "normalized_scores": normalized,
            "ranked_roles": ranked,
        }

    def observe(self, current_role: str, next_role: str) -> None:
        row = self.rows[current_role]
        if self.decay < 1.0:
            for target in self.targets:
                row[target] = max(self.prior_strength * 0.05, row[target] * self.decay)
        row[next_role] = row.get(next_role, 0.0) + 1.0
        self.observations += 1

    def snapshot(self) -> dict[str, Any]:
        matrix = {}
        for role in AGENT_ROLES:
            _, _, probs = self.predict(role)
            matrix[role] = probs
        return {
            "decay": self.decay,
            "prior_strength": self.prior_strength,
            "observations": self.observations,
            "probabilities": matrix,
        }


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 cv <= 1e-12:
        return mean_value
    sigma2 = math.log1p(cv * cv)
    sigma = math.sqrt(sigma2)
    mu = math.log(max(mean_value, 1e-9)) - 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, float(minimum)))))


def _sample_transition(rng: random.Random, role: str, shifted: bool) -> str:
    matrix = SHIFTED_TRANSITIONS if shifted else BASE_TRANSITIONS
    row = matrix[role]
    targets = tuple(row)
    return rng.choices(targets, weights=[row[target] for target in targets], k=1)[0]


def generate_workflows(cfg: ExecutionLearningConfig) -> list[WorkflowSpec]:
    rng = random.Random(cfg.seed ^ 0xE91A7)
    workflows: list[WorkflowSpec] = []
    now = 0.0
    workflow_id = 0
    while True:
        now += rng.expovariate(cfg.workflow_rate_rps)
        if now > cfg.duration_s:
            break
        shifted = now >= cfg.duration_s * cfg.shift_fraction
        role = "planner"
        steps: list[WorkflowStep] = []
        for step_idx in range(cfg.max_steps):
            dynamic_prompt = _sample_int(rng, cfg.dynamic_prompt_tokens_mean, cfg.dynamic_prompt_cv, 32)
            output = _sample_int(rng, cfg.output_tokens_mean, cfg.output_tokens_cv, 1)
            next_role = _sample_transition(rng, role, shifted)
            if next_role == END or step_idx == cfg.max_steps - 1:
                gap = 0.0
            else:
                mean_gap = cfg.tool_gap_mean_s * ROLE_GAP_MULTIPLIERS.get(next_role, 1.0)
                if shifted and next_role in {"retriever", "verifier"}:
                    mean_gap *= 1.35
                gap = _sample_positive_lognormal(rng, mean_gap, cfg.tool_gap_cv, 0.0)
            steps.append(WorkflowStep(role, dynamic_prompt, output, gap, shifted))
            if next_role == END or step_idx == cfg.max_steps - 1:
                break
            role = next_role
        workflows.append(WorkflowSpec(workflow_id, now, steps))
        workflow_id += 1
    return workflows


def _prediction_calibration(rows: list[dict[str, Any]], bins: int = 10) -> dict[str, Any]:
    if not rows:
        return {"count": 0, "brier": 0.0, "log_loss": 0.0, "ece": 0.0, "bins": []}
    targets = (*AGENT_ROLES, END)
    brier_total = 0.0
    log_total = 0.0
    buckets: list[list[tuple[float, float]]] = [[] for _ in range(max(2, bins))]
    for row in rows:
        probs = row.get("probabilities", {})
        actual = str(row.get("actual_next_role", END))
        brier_total += sum((float(probs.get(target, 0.0)) - (1.0 if target == actual else 0.0)) ** 2 for target in targets)
        actual_p = max(float(probs.get(actual, 0.0)), 1e-12)
        log_total += -math.log(actual_p)
        confidence = min(max(float(row.get("confidence", 0.0)), 0.0), 1.0)
        correct = 1.0 if row.get("prediction_correct") else 0.0
        index = min(int(confidence * len(buckets)), len(buckets) - 1)
        buckets[index].append((confidence, correct))
    ece = 0.0
    rendered = []
    for index, bucket in enumerate(buckets):
        if not bucket:
            continue
        avg_conf = mean(value[0] for value in bucket)
        accuracy = mean(value[1] for value in bucket)
        weight = len(bucket) / len(rows)
        ece += weight * abs(avg_conf - accuracy)
        rendered.append(
            {
                "lower": index / len(buckets),
                "upper": (index + 1) / len(buckets),
                "count": len(bucket),
                "mean_confidence": avg_conf,
                "accuracy": accuracy,
            }
        )
    return {
        "count": len(rows),
        "brier": brier_total / len(rows),
        "log_loss": log_total / len(rows),
        "ece": ece,
        "bins": rendered,
    }


class ExecutionLearningSimulator:
    """Reference simulator for online workflow learning and prefix prefetch.

    This lab intentionally isolates cross-workflow *agent-prefix* caching from the
    session-KV simulator. It uses one serial compute station plus one serialized
    prefix-transfer link and a bounded shared HBM prefix cache. Static agent
    prefixes are always available in the modeled host tier; policy decisions
    determine whether likely next prefixes are prefetched into HBM during gaps.
    """

    def __init__(self, cfg: ExecutionLearningConfig, workflows: list[WorkflowSpec] | 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)
        self.workflows = workflows if workflows is not None else generate_workflows(cfg)
        self.predictor = OnlineTransitionPredictor(decay=cfg.transition_decay, prior_strength=cfg.prior_strength)
        self.events: list[tuple[float, int, str, tuple[Any, ...]]] = []
        self.event_seq = 0
        self.now = 0.0
        self.server_busy_until = 0.0
        self.transfer_busy_until = 0.0
        self.queue: list[tuple[int, int, float]] = []
        self.cache: dict[str, PrefixEntry] = {}
        self.prefix_source: dict[str, str] = {}
        self.pending_predictions: dict[tuple[int, int], dict[str, Any]] = {}
        self.step_rows: list[dict[str, Any]] = []
        self.prediction_rows: list[dict[str, Any]] = []
        self.timeline: list[dict[str, Any]] = []
        self.last_memory_time = 0.0
        self.hbm_gb_seconds = 0.0
        self.peak_hbm_gb = 0.0
        self.pressure_evictions = 0
        self.prefetch_attempts = 0
        self.prefetch_correct = 0
        self.prefetch_useful = 0
        self.prefetch_bytes = 0.0
        self.prefetch_useful_bytes = 0.0
        self.unused_prefetch_bytes = 0.0
        self.wrong_step_prefetch_bytes = 0.0
        self.transfer_latencies_ms: list[float] = []
        self.prefix_hits = 0
        self.prefetch_hits = 0
        self.prefill_tokens_saved = 0
        self.workflow_completion: dict[int, float] = {}

        kv_bytes = self.latency.kv_bytes_per_token()
        self.role_size_gb = {role: ROLE_PREFIX_TOKENS[role] * kv_bytes / 1e9 for role in AGENT_ROLES}
        reference_working_set = sum(self.role_size_gb.values())
        self.cache_capacity_gb = max(reference_working_set * cfg.prefix_cache_budget_fraction, min(self.role_size_gb.values()))

        for workflow in self.workflows:
            if workflow.steps:
                self._push(workflow.arrival_time, "step_ready", workflow.workflow_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 entry in self.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
        self.hbm_gb_seconds += self._used_gb() * elapsed
        self.last_memory_time = new_time
        self.peak_hbm_gb = max(self.peak_hbm_gb, self._used_gb())

    def _record_timeline(self) -> None:
        min_gap = max(self.cfg.duration_s / self.cfg.timeline_points, 0.08)
        if self.timeline and self.now - self.timeline[-1]["time_s"] < min_gap:
            return
        self.timeline.append(
            {
                "time_s": self.now,
                "queued_steps": len(self.queue),
                "server_busy": 1 if self.server_busy_until > self.now + 1e-12 else 0,
                "prefix_hbm_gb": self._used_gb(),
                "cached_roles": len(self.cache),
                "transfer_busy": 1 if self.transfer_busy_until > self.now + 1e-12 else 0,
            }
        )
        if len(self.timeline) > self.cfg.timeline_points * 2:
            self.timeline = self.timeline[::2]

    def _mark_unused_prefetch(self, entry: PrefixEntry) -> None:
        if entry.source == "prefetch" and not entry.used:
            self.unused_prefetch_bytes += entry.size_gb * 1e9

    def _candidate_eviction_cost_ms(self, role: str, forecast_scores: dict[str, float]) -> float:
        """Approximate downstream cost of evicting forecast-relevant cached prefixes."""
        size_gb = self.role_size_gb[role]
        existing = self.cache.get(role)
        current_without = self._used_gb() - (existing.size_gb if existing else 0.0)
        overflow = max(0.0, current_without + size_gb - self.cache_capacity_gb)
        if overflow <= 1e-12:
            return 0.0
        cost_ms = 0.0
        reclaimed = 0.0
        candidates = sorted(
            [entry for key, entry in self.cache.items() if key != role],
            key=lambda entry: (entry.last_access, entry.role),
        )
        for victim in candidates:
            if reclaimed >= overflow - 1e-12:
                break
            reuse_score = max(float(forecast_scores.get(victim.role, 0.0)), 0.0)
            recompute_ms = self.latency.prefill_seconds([ROLE_PREFIX_TOKENS[victim.role]]) * 1000.0
            cost_ms += reuse_score * recompute_ms
            reclaimed += victim.size_gb
        return cost_ms

    def _prefetch_utility_ms(self, role: str, forecast_scores: dict[str, float]) -> float:
        score = max(float(forecast_scores.get(role, 0.0)), 0.0)
        saved_ms = self.latency.prefill_seconds([ROLE_PREFIX_TOKENS[role]]) * 1000.0
        transfer_ms = self.cfg.transfer_base_ms + self.role_size_gb[role] / self.cfg.host_bandwidth_gbps * 1000.0
        eviction_ms = self._candidate_eviction_cost_ms(role, forecast_scores)
        return score * saved_ms - transfer_ms - eviction_ms

    def _evict_for(self, role: str, size_gb: float) -> bool:
        existing = self.cache.get(role)
        current_without_role = self._used_gb() - (existing.size_gb if existing else 0.0)
        if size_gb > self.cache_capacity_gb + 1e-12:
            return False
        while current_without_role + size_gb > self.cache_capacity_gb + 1e-12:
            candidates = [entry for key, entry in self.cache.items() if key != role]
            if not candidates:
                return False
            victim = min(candidates, key=lambda entry: (entry.last_access, entry.role))
            self._integrate_memory(self.now)
            self._mark_unused_prefetch(victim)
            self.cache.pop(victim.role, None)
            self.prefix_source.pop(victim.role, None)
            self.pressure_evictions += 1
            current_without_role = self._used_gb() - (existing.size_gb if existing and role in self.cache else 0.0)
        return True

    def _install_prefix(self, role: str, *, available_time: float, source: str) -> int | None:
        size_gb = self.role_size_gb[role]
        if not self._evict_for(role, size_gb):
            return None
        self._integrate_memory(self.now)
        previous = self.cache.get(role)
        if previous is not None:
            self._mark_unused_prefetch(previous)
        generation = (previous.generation + 1) if previous else 1
        self.cache[role] = PrefixEntry(role, size_gb, self.now, available_time, source, generation)
        self.prefix_source[role] = source
        self.peak_hbm_gb = max(self.peak_hbm_gb, self._used_gb())
        self._record_timeline()
        return generation

    def _prefetch(self, role: str) -> dict[str, Any]:
        entry = self.cache.get(role)
        if entry is not None:
            return {"attempted": False, "reason": "already_cached", "role": role, "available_time": entry.available_time}
        size_gb = self.role_size_gb[role]
        transfer_s = self.cfg.transfer_base_ms / 1000.0 + size_gb / self.cfg.host_bandwidth_gbps
        start = max(self.now, self.transfer_busy_until)
        available = start + transfer_s
        generation = self._install_prefix(role, available_time=available, source="prefetch")
        if generation is None:
            return {"attempted": False, "reason": "capacity_reject", "role": role, "available_time": available}
        self.transfer_busy_until = available
        self.prefetch_attempts += 1
        self.prefetch_bytes += size_gb * 1e9
        self.transfer_latencies_ms.append((available - self.now) * 1000.0)
        self._push(available, "prefetch_complete", role, generation)
        return {"attempted": True, "reason": "scheduled", "role": role, "available_time": available, "size_gb": size_gb}

    def _actual_next_role(self, workflow_id: int, step_index: int) -> str:
        workflow = self.workflows[workflow_id]
        if step_index + 1 >= len(workflow.steps):
            return END
        return workflow.steps[step_index + 1].role

    def _actual_future_roles(self, workflow_id: int, step_index: int, horizon: int) -> list[str]:
        workflow = self.workflows[workflow_id]
        roles = [step.role for step in workflow.steps[step_index + 1 : step_index + 1 + max(1, horizon)]]
        return roles

    def _planned_prefetch_roles(
        self,
        workflow_id: int,
        step_index: int,
        current_role: str,
    ) -> tuple[str, float, dict[str, float], dict[str, Any], list[dict[str, Any]]]:
        actual_next = self._actual_next_role(workflow_id, step_index)
        predicted_role, confidence, probabilities = self.predictor.predict(current_role)
        forecast = self.predictor.forecast(
            current_role,
            horizon=self.cfg.forecast_horizon,
            discount=self.cfg.forecast_discount,
        )
        plans: list[dict[str, Any]] = []

        if self.cfg.prefetch_policy == "oracle":
            predicted_role = actual_next
            confidence = 1.0
            probabilities = {target: 0.0 for target in (*AGENT_ROLES, END)}
            probabilities[predicted_role] = 1.0
            if actual_next != END:
                plans = [{"role": actual_next, "score": 1.0, "utility_ms": None, "oracle": True}]
            return predicted_role, confidence, probabilities, forecast, plans

        if self.cfg.prefetch_policy == "oracle_horizon":
            actual_future = self._actual_future_roles(workflow_id, step_index, self.cfg.forecast_horizon)
            predicted_role = actual_next
            confidence = 1.0
            probabilities = {target: 0.0 for target in (*AGENT_ROLES, END)}
            probabilities[predicted_role] = 1.0
            seen: set[str] = set()
            for role in actual_future:
                if role not in seen:
                    plans.append({"role": role, "score": 1.0, "utility_ms": None, "oracle": True})
                    seen.add(role)
                if len(plans) >= self.cfg.prefetch_top_k:
                    break
            return predicted_role, confidence, probabilities, forecast, plans

        if self.cfg.prefetch_policy in {"none", "cumulative", "decayed"}:
            if self.cfg.prefetch_policy != "none" and predicted_role != END and confidence >= self.cfg.confidence_threshold:
                plans = [{"role": predicted_role, "score": confidence, "utility_ms": None, "oracle": False}]
            return predicted_role, confidence, probabilities, forecast, plans

        ranked = forecast["ranked_roles"]
        normalized = forecast["normalized_scores"]
        raw_scores = forecast["role_scores"]
        for role in ranked:
            if len(plans) >= self.cfg.prefetch_top_k:
                break
            normalized_score = float(normalized.get(role, 0.0))
            if normalized_score < self.cfg.forecast_min_score:
                continue
            utility_ms = self._prefetch_utility_ms(role, raw_scores)
            if self.cfg.prefetch_policy == "utility" and utility_ms < self.cfg.utility_threshold_ms:
                continue
            plans.append(
                {
                    "role": role,
                    "score": normalized_score,
                    "raw_score": float(raw_scores.get(role, 0.0)),
                    "utility_ms": utility_ms,
                    "oracle": False,
                }
            )
        return predicted_role, confidence, probabilities, forecast, plans

    def _schedule_prediction(self, workflow_id: int, step_index: int, current_role: str) -> None:
        actual_next = self._actual_next_role(workflow_id, step_index)
        predicted_role, confidence, probabilities, forecast, plans = self._planned_prefetch_roles(
            workflow_id, step_index, current_role
        )
        attempted_roles: list[str] = []
        attempt_sizes: dict[str, float] = {}
        attempt_reasons: dict[str, str] = {}

        for plan in plans:
            role = str(plan["role"])
            result = self._prefetch(role)
            attempt_reasons[role] = str(result.get("reason", "unknown"))
            if result.get("attempted"):
                attempted_roles.append(role)
                attempt_sizes[role] = float(result.get("size_gb", 0.0))

        if self.cfg.prefetch_policy == "none":
            decision_reason = "learn_only"
        elif not plans:
            decision_reason = "no_candidate"
        elif attempted_roles:
            decision_reason = "scheduled"
        else:
            decision_reason = ",".join(sorted(set(attempt_reasons.values()))) or "not_scheduled"

        self.pending_predictions[(workflow_id, step_index)] = {
            "workflow_id": workflow_id,
            "step_index": step_index + 1,
            "current_role": current_role,
            "predicted_role": predicted_role,
            "confidence": confidence,
            "probabilities": probabilities,
            "forecast_horizon": self.cfg.forecast_horizon,
            "forecast_discount": self.cfg.forecast_discount,
            "forecast_role_scores": forecast["role_scores"],
            "forecast_normalized_scores": forecast["normalized_scores"],
            "forecast_roles": [str(plan["role"]) for plan in plans],
            "forecast_plans": plans,
            "prefetch_attempted": bool(attempted_roles),
            "prefetch_roles": attempted_roles,
            "prefetch_reasons": attempt_reasons,
            "prefetch_reason": decision_reason,
            "prefetch_sizes_gb": attempt_sizes,
            "prefetch_size_gb": sum(attempt_sizes.values()),
            "shifted_regime": self.workflows[workflow_id].steps[step_index].shifted_regime,
        }

        if actual_next == END:
            self._observe_transition(workflow_id, step_index, actual_next)

    def _observe_transition(self, workflow_id: int, step_index: int, actual_next: str) -> None:
        key = (workflow_id, step_index)
        pending = self.pending_predictions.pop(key, None)
        workflow = self.workflows[workflow_id]
        current_role = workflow.steps[step_index].role
        if pending is None:
            return
        predicted_role = str(pending["predicted_role"])
        attempted_roles = [str(role) for role in pending.get("prefetch_roles", [])]
        correct = predicted_role == actual_next
        if actual_next != END and actual_next in attempted_roles:
            self.prefetch_correct += 1
        immediate_wrong_bytes = 0.0
        for role in attempted_roles:
            if role != actual_next:
                immediate_wrong_bytes += float(pending.get("prefetch_sizes_gb", {}).get(role, 0.0)) * 1e9
        self.wrong_step_prefetch_bytes += immediate_wrong_bytes

        future_roles = self._actual_future_roles(workflow_id, step_index, int(pending.get("forecast_horizon", 1)))
        actual_future_set = set(future_roles)
        forecast_roles = [str(role) for role in pending.get("forecast_roles", [])]
        forecast_set = set(forecast_roles)
        forecast_recall = (
            len(actual_future_set & forecast_set) / len(actual_future_set)
            if actual_future_set
            else None
        )
        self.prediction_rows.append(
            pending
            | {
                "actual_next_role": actual_next,
                "actual_future_roles": future_roles,
                "prediction_correct": correct,
                "forecast_next_hit": actual_next in forecast_set if actual_next != END else predicted_role == END,
                "forecast_recall": forecast_recall,
                "immediate_wrong_prefetch_gb": immediate_wrong_bytes / 1e9,
                "post_shift": bool(pending["shifted_regime"]),
            }
        )
        if self.cfg.prefetch_policy not in {"oracle", "oracle_horizon"}:
            self.predictor.observe(current_role, actual_next)

    def _start_next(self) -> None:
        if self.server_busy_until > self.now + 1e-12 or not self.queue:
            return
        workflow_id, step_index, ready_time = self.queue.pop(0)
        workflow = self.workflows[workflow_id]
        step = workflow.steps[step_index]
        role = step.role
        entry = self.cache.get(role)
        prefix_tokens = ROLE_PREFIX_TOKENS[role]
        prefix_wait_s = 0.0
        cache_hit = False
        cache_source = "miss"

        if entry is not None and entry.available_time <= self.now + 1e-12:
            cache_hit = True
            cache_source = self.prefix_source.get(role, entry.source)
        elif entry is not None:
            # If a prefetch is almost complete, waiting can be cheaper than
            # recomputing the whole static prefix. Otherwise the step proceeds
            # immediately with a reactive prefill and the pending prefetch is
            # treated as not useful for this invocation.
            remaining = max(0.0, entry.available_time - self.now)
            full_prefix_s = self.latency.prefill_seconds([prefix_tokens])
            if remaining < full_prefix_s:
                prefix_wait_s = remaining
                cache_hit = True
                cache_source = "prefetch_wait"

        if cache_hit:
            self.prefix_hits += 1
            if cache_source.startswith("prefetch"):
                self.prefetch_hits += 1
            self.prefill_tokens_saved += prefix_tokens
            if entry is not None:
                if entry.source == "prefetch" and not entry.used:
                    entry.used = True
                    self.prefetch_useful += 1
                    self.prefetch_useful_bytes += entry.size_gb * 1e9
                entry.last_access = self.now + prefix_wait_s
            prefill_tokens = max(step.dynamic_prompt_tokens, 1)
        else:
            prefill_tokens = max(prefix_tokens + step.dynamic_prompt_tokens, 1)

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

        self.server_busy_until = self.now + service_s
        self._push(
            self.server_busy_until,
            "step_complete",
            workflow_id,
            step_index,
            ready_time,
            ttft_s,
            service_s,
            cache_hit,
            cache_source,
            prefill_tokens,
        )

    def _finish_step(
        self,
        workflow_id: int,
        step_index: int,
        ready_time: float,
        ttft_s: float,
        service_s: float,
        cache_hit: bool,
        cache_source: str,
        prefill_tokens: int,
    ) -> None:
        workflow = self.workflows[workflow_id]
        step = workflow.steps[step_index]
        role = step.role
        # Reactive execution materializes the role prefix after a miss.
        if not cache_hit:
            self._install_prefix(role, available_time=self.now, source="reactive")
        elif role in self.cache:
            self.cache[role].last_access = self.now

        self.step_rows.append(
            {
                "workflow_id": workflow_id,
                "step_index": step_index + 1,
                "role": role,
                "shifted_regime": step.shifted_regime,
                "ready_time": ready_time,
                "completion_time": self.now,
                "cache_hit": cache_hit,
                "cache_source": cache_source,
                "prefix_tokens": ROLE_PREFIX_TOKENS[role],
                "dynamic_prompt_tokens": step.dynamic_prompt_tokens,
                "prefill_tokens": prefill_tokens,
                "output_tokens": step.output_tokens,
                "ttft_ms": ttft_s * 1000.0,
                "e2e_ms": (self.now - ready_time) * 1000.0,
                "queue_ms": max(0.0, (self.now - ready_time - service_s) * 1000.0),
            }
        )

        self._schedule_prediction(workflow_id, step_index, role)
        actual_next = self._actual_next_role(workflow_id, step_index)
        if actual_next == END:
            self.workflow_completion[workflow_id] = self.now
        else:
            gap = step.gap_after_s
            ready_at = self.now + gap
            self._push(ready_at, "transition_observed", workflow_id, step_index, actual_next)
            self._push(ready_at, "step_ready", workflow_id, step_index + 1)
        self.server_busy_until = self.now
        self._start_next()

    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 == "step_ready":
                workflow_id, step_index = int(payload[0]), int(payload[1])
                self.queue.append((workflow_id, step_index, self.now))
                self._start_next()
            elif kind == "step_complete":
                self._finish_step(
                    int(payload[0]), int(payload[1]), float(payload[2]), float(payload[3]), float(payload[4]),
                    bool(payload[5]), str(payload[6]), int(payload[7]),
                )
            elif kind == "transition_observed":
                self._observe_transition(int(payload[0]), int(payload[1]), str(payload[2]))
            elif kind == "prefetch_complete":
                role, generation = str(payload[0]), int(payload[1])
                entry = self.cache.get(role)
                if entry is not None and entry.generation == generation and entry.available_time <= self.now + 1e-12:
                    entry.last_access = self.now
            self._record_timeline()

        self._integrate_memory(self.now)
        self._record_timeline()
        step_ttft = [float(row["ttft_ms"]) for row in self.step_rows]
        step_e2e = [float(row["e2e_ms"]) for row in self.step_rows]
        workflow_e2e = [
            (self.workflow_completion[wf.workflow_id] - wf.arrival_time) * 1000.0
            for wf in self.workflows
            if wf.workflow_id in self.workflow_completion
        ]
        horizon = max([*self.workflow_completion.values(), self.cfg.duration_s, 1e-9])
        predictions = self.prediction_rows
        pre = [row for row in predictions if not row["post_shift"]]
        post = [row for row in predictions if row["post_shift"]]

        def accuracy(rows: list[dict[str, Any]]) -> float:
            return sum(1 for row in rows if row["prediction_correct"]) / len(rows) if rows else 0.0

        # Account for prefetched entries that remained unused through the end of the run.
        for entry in list(self.cache.values()):
            self._mark_unused_prefetch(entry)
        prefetch_precision = self.prefetch_correct / self.prefetch_attempts if self.prefetch_attempts else 0.0
        prefetch_coverage = (
            sum(1 for row in predictions if row.get("prefetch_attempted")) / len(predictions)
            if predictions else 0.0
        )
        prefetch_utilization = self.prefetch_useful / self.prefetch_attempts if self.prefetch_attempts else 0.0
        forecast_recall_values = [
            float(row["forecast_recall"]) for row in predictions if row.get("forecast_recall") is not None
        ]
        forecast_recall = mean(forecast_recall_values) if forecast_recall_values else 0.0
        workflow_count = len(self.workflows)
        completed_count = len(self.workflow_completion)
        summary = {
            "workflows_generated": workflow_count,
            "workflows_completed": completed_count,
            "steps_completed": len(self.step_rows),
            "workflow_completion_rate": completed_count / workflow_count if workflow_count else 0.0,
            "workflow_throughput_rps": completed_count / horizon,
            "simulated_makespan_s": horizon,
        }
        resource = {
            "prefix_cache_capacity_gb": self.cache_capacity_gb,
            "mean_prefix_hbm_gb": self.hbm_gb_seconds / horizon,
            "peak_prefix_hbm_gb": self.peak_hbm_gb,
            "hbm_gb_seconds": self.hbm_gb_seconds,
            "prefix_hit_rate": self.prefix_hits / len(self.step_rows) if self.step_rows else 0.0,
            "prefetch_hit_rate": self.prefetch_hits / len(self.step_rows) if self.step_rows else 0.0,
            "prefill_tokens_saved": self.prefill_tokens_saved,
            "pressure_evictions": self.pressure_evictions,
            "prefetch_attempts": self.prefetch_attempts,
            "prefetch_correct": self.prefetch_correct,
            "prefetch_precision": prefetch_precision,
            "prefetch_coverage": prefetch_coverage,
            "prefetch_utilization": prefetch_utilization,
            "prefetch_attempts_per_prediction": self.prefetch_attempts / len(predictions) if predictions else 0.0,
            "forecast_recall": forecast_recall,
            "prefetch_gb": self.prefetch_bytes / 1e9,
            "prefetch_useful_gb": self.prefetch_useful_bytes / 1e9,
            "unused_prefetch_gb": self.unused_prefetch_bytes / 1e9,
            "wrong_step_prefetch_gb": self.wrong_step_prefetch_bytes / 1e9,
            "p95_prefetch_transfer_ms": percentile(self.transfer_latencies_ms, 0.95),
        }
        prediction = {
            "count": len(predictions),
            "top1_accuracy": accuracy(predictions),
            "pre_shift_accuracy": accuracy(pre),
            "post_shift_accuracy": accuracy(post),
            "calibration": _prediction_calibration(predictions),
            "pre_shift_calibration": _prediction_calibration(pre),
            "post_shift_calibration": _prediction_calibration(post),
            "rows": predictions[:4000],
            "predictor": self.predictor.snapshot(),
        }
        latency = {
            "step_ttft_ms": {"p50": percentile(step_ttft, 0.50), "p95": percentile(step_ttft, 0.95), "p99": percentile(step_ttft, 0.99)},
            "step_e2e_ms": {"p50": percentile(step_e2e, 0.50), "p95": percentile(step_e2e, 0.95), "p99": percentile(step_e2e, 0.99)},
            "workflow_e2e_ms": {"p50": percentile(workflow_e2e, 0.50), "p95": percentile(workflow_e2e, 0.95), "p99": percentile(workflow_e2e, 0.99)},
        }
        return {
            "config": self.cfg.to_dict(),
            "provenance": {
                "simulator": "InferScale-Sim",
                "mode": "online-agent-execution-learning",
                "latency_profile_type": "analytical-reference",
                "execution_model": "online-transition-multistep-prefix-prefetch-reference",
                "prefetch_policy": self.cfg.prefetch_policy,
                "lookahead": "oracle-upper-bound" if self.cfg.prefetch_policy in {"oracle", "oracle_horizon"} else "no-future-transition-lookahead",
                "warning": "Execution Learning isolates cross-workflow static-prefix prediction/prefetch from session-KV retention and dynamic batching.",
            },
            "summary": summary,
            "latency": latency,
            "resource": resource,
            "prediction": prediction,
            "steps": self.step_rows[:4000],
            "timeline": self.timeline,
            "role_prefix_tokens": ROLE_PREFIX_TOKENS,
        }


def run_execution_learning(config: dict[str, Any], workflows: list[WorkflowSpec] | None = None) -> dict[str, Any]:
    cfg = ExecutionLearningConfig.from_dict(config)
    return ExecutionLearningSimulator(cfg, workflows=workflows).run()


def _common_workflows(cfg: ExecutionLearningConfig) -> list[WorkflowSpec]:
    return generate_workflows(cfg)


def _row(label: str, result: dict[str, Any]) -> dict[str, Any]:
    return {
        "label": label,
        "prefetch_policy": result["config"]["prefetch_policy"],
        "p95_step_ttft_ms": result["latency"]["step_ttft_ms"]["p95"],
        "p95_workflow_e2e_ms": result["latency"]["workflow_e2e_ms"]["p95"],
        "workflow_throughput_rps": result["summary"]["workflow_throughput_rps"],
        "prefix_hit_rate": result["resource"]["prefix_hit_rate"],
        "prefetch_precision": result["resource"]["prefetch_precision"],
        "prefetch_coverage": result["resource"]["prefetch_coverage"],
        "prefetch_utilization": result["resource"].get("prefetch_utilization", 0.0),
        "forecast_recall": result["resource"].get("forecast_recall", 0.0),
        "mean_hbm_gb": result["resource"]["mean_prefix_hbm_gb"],
        "pressure_evictions": result["resource"].get("pressure_evictions", 0),
        "wrong_step_prefetch_gb": result["resource"]["wrong_step_prefetch_gb"],
        "unused_prefetch_gb": result["resource"].get("unused_prefetch_gb", 0.0),
        "prefill_tokens_saved": result["resource"]["prefill_tokens_saved"],
        "top1_accuracy": result["prediction"]["top1_accuracy"],
        "brier": result["prediction"].get("calibration", {}).get("brier", 0.0),
        "ece": result["prediction"].get("calibration", {}).get("ece", 0.0),
        "post_shift_ece": result["prediction"].get("post_shift_calibration", {}).get("ece", 0.0),
        "pre_shift_accuracy": result["prediction"]["pre_shift_accuracy"],
        "post_shift_accuracy": result["prediction"]["post_shift_accuracy"],
    }


def execution_prefetch_study(config: dict[str, Any]) -> dict[str, Any]:
    base = ExecutionLearningConfig.from_dict(config)
    workflows = _common_workflows(base)
    candidates = [
        ("Learn only / no prefetch", "none", base.transition_decay),
        ("Cumulative transitions", "cumulative", 1.0),
        ("Decayed transitions", "decayed", base.transition_decay),
        ("Oracle next-role", "oracle", 1.0),
    ]
    rows: list[dict[str, Any]] = []
    results: dict[str, dict[str, Any]] = {}
    for label, policy, decay in candidates:
        cfg = ExecutionLearningConfig.from_dict(base.to_dict())
        cfg.prefetch_policy = policy
        cfg.transition_decay = decay
        result = run_execution_learning(cfg.to_dict(), workflows)
        results[label] = result
        rows.append(_row(label, result))
    return {
        "protocol": "common-shifted-agent-workflow-trace",
        "study": "online-transition-prefetch-policy-comparison",
        "rows": rows,
        "results": results,
        "shift_fraction": base.shift_fraction,
        "note": "Oracle next-role sees the realized next agent and is an upper bound; online policies update only after a transition is observed.",
    }


def _pearson(xs: list[float], ys: list[float]) -> float:
    if len(xs) < 2 or len(xs) != len(ys):
        return 0.0
    mx = mean(xs)
    my = mean(ys)
    numerator = sum((x - mx) * (y - my) for x, y in zip(xs, ys, strict=True))
    sx = math.sqrt(sum((x - mx) ** 2 for x in xs))
    sy = math.sqrt(sum((y - my) ** 2 for y in ys))
    if sx <= 1e-12 or sy <= 1e-12:
        return 0.0
    return numerator / (sx * sy)


def execution_threshold_sweep(
    config: dict[str, Any], thresholds: list[float] | None = None
) -> dict[str, Any]:
    base = ExecutionLearningConfig.from_dict(config)
    base.prefetch_policy = "decayed"
    workflows = _common_workflows(base)
    values = thresholds or [0.0, 0.35, 0.50, 0.60, 0.70, 0.80, 0.90]
    cleaned = sorted({min(max(float(value), 0.0), 1.0) for value in values})
    rows = []
    for threshold in cleaned:
        cfg = ExecutionLearningConfig.from_dict(base.to_dict())
        cfg.confidence_threshold = threshold
        result = run_execution_learning(cfg.to_dict(), workflows)
        rows.append(_row(f"threshold {threshold:.2f}", result) | {"threshold": threshold})
    association = {
        "coverage_vs_ttft_r": _pearson(
            [float(row["prefetch_coverage"]) for row in rows],
            [float(row["p95_step_ttft_ms"]) for row in rows],
        ),
        "precision_vs_wrong_prefetch_r": _pearson(
            [float(row["prefetch_precision"]) for row in rows],
            [float(row["wrong_step_prefetch_gb"]) for row in rows],
        ),
    }
    return {
        "protocol": "common-shifted-agent-workflow-trace",
        "study": "prefetch-confidence-threshold-frontier",
        "rows": rows,
        "association": association,
        "note": "Correlations summarize this simulated sweep only; they are descriptive, not causal estimates.",
    }


def execution_decay_sweep(
    config: dict[str, Any], decay_values: list[float] | None = None
) -> dict[str, Any]:
    base = ExecutionLearningConfig.from_dict(config)
    base.prefetch_policy = "decayed"
    workflows = _common_workflows(base)
    values = decay_values or [0.50, 0.65, 0.75, 0.85, 0.92, 0.97, 1.00]
    cleaned = sorted({min(max(float(value), 0.20), 1.0) for value in values})
    rows = []
    for decay in cleaned:
        cfg = ExecutionLearningConfig.from_dict(base.to_dict())
        cfg.transition_decay = decay
        result = run_execution_learning(cfg.to_dict(), workflows)
        rows.append(_row(f"decay {decay:.2f}", result) | {"decay": decay})
    best_accuracy = max(rows, key=lambda row: (row["post_shift_accuracy"], -row["p95_step_ttft_ms"])) if rows else None
    best_ttft = min(rows, key=lambda row: (row["p95_step_ttft_ms"], -row["post_shift_accuracy"])) if rows else None
    association = {
        "post_shift_accuracy_vs_ttft_r": _pearson(
            [float(row["post_shift_accuracy"]) for row in rows],
            [float(row["p95_step_ttft_ms"]) for row in rows],
        ),
        "post_shift_accuracy_vs_prefix_hit_r": _pearson(
            [float(row["post_shift_accuracy"]) for row in rows],
            [float(row["prefix_hit_rate"]) for row in rows],
        ),
    }
    return {
        "protocol": "common-shifted-agent-workflow-trace",
        "study": "transition-forgetting-rate-sweep",
        "rows": rows,
        "best_post_shift_accuracy_decay": best_accuracy["decay"] if best_accuracy else None,
        "best_ttft_decay": best_ttft["decay"] if best_ttft else None,
        "association": association,
        "note": "Better next-role prediction need not minimize serving latency because cache occupancy and transfer timing mediate the effect.",
    }


def execution_planning_study(config: dict[str, Any]) -> dict[str, Any]:
    """Compare one-step, multi-step, utility-aware, and clairvoyant planning."""
    base = ExecutionLearningConfig.from_dict(config)
    workflows = _common_workflows(base)
    candidates = [
        ("Top-1 decayed", "decayed"),
        ("Multi-step top-k", "multistep"),
        ("Utility-aware multi-step", "utility"),
        ("Oracle future-set", "oracle_horizon"),
    ]
    rows: list[dict[str, Any]] = []
    results: dict[str, dict[str, Any]] = {}
    for label, policy in candidates:
        cfg = ExecutionLearningConfig.from_dict(base.to_dict())
        cfg.prefetch_policy = policy
        result = run_execution_learning(cfg.to_dict(), workflows)
        results[label] = result
        rows.append(_row(label, result))
    best_ttft = min(rows, key=lambda row: row["p95_step_ttft_ms"]) if rows else None
    best_efficiency = max(
        rows,
        key=lambda row: (
            row["prefill_tokens_saved"] / max(row["mean_hbm_gb"], 1e-9),
            -row["p95_step_ttft_ms"],
        ),
    ) if rows else None
    return {
        "protocol": "common-shifted-agent-workflow-trace",
        "study": "multi-step-prefetch-planning",
        "rows": rows,
        "results": results,
        "best_ttft_policy": best_ttft["label"] if best_ttft else None,
        "best_prefill_per_hbm_policy": best_efficiency["label"] if best_efficiency else None,
        "note": "Oracle future-set sees realized future roles only as an information upper bound; online multi-step policies roll the learned transition matrix forward without future trace access.",
    }


def execution_horizon_sweep(
    config: dict[str, Any], horizon_values: list[int] | None = None
) -> dict[str, Any]:
    base = ExecutionLearningConfig.from_dict(config)
    base.prefetch_policy = "utility"
    workflows = _common_workflows(base)
    values = horizon_values or [1, 2, 3, 4, 5]
    cleaned = sorted({max(1, min(int(value), 6)) for value in values})
    rows: list[dict[str, Any]] = []
    for horizon in cleaned:
        cfg = ExecutionLearningConfig.from_dict(base.to_dict())
        cfg.forecast_horizon = horizon
        result = run_execution_learning(cfg.to_dict(), workflows)
        rows.append(_row(f"horizon {horizon}", result) | {"horizon": horizon})
    best_ttft = min(rows, key=lambda row: (row["p95_step_ttft_ms"], row["unused_prefetch_gb"])) if rows else None
    best_utilization = max(rows, key=lambda row: (row["prefetch_utilization"], -row["unused_prefetch_gb"])) if rows else None
    return {
        "protocol": "common-shifted-agent-workflow-trace",
        "study": "forecast-horizon-sweep",
        "rows": rows,
        "best_ttft_horizon": best_ttft["horizon"] if best_ttft else None,
        "best_utilization_horizon": best_utilization["horizon"] if best_utilization else None,
        "note": "Longer forecasts can expose future reuse but may spend bandwidth and cache capacity on prefixes that are not consumed before eviction.",
    }


def execution_budget_sweep(
    config: dict[str, Any], budget_values: list[float] | None = None
) -> dict[str, Any]:
    base = ExecutionLearningConfig.from_dict(config)
    workflows = _common_workflows(base)
    values = budget_values or [0.30, 0.50, 0.75, 1.00]
    cleaned = sorted({min(max(float(value), 0.10), 1.50) for value in values})
    policies = [
        ("Top-1", "decayed"),
        ("Multi-step", "multistep"),
        ("Utility-aware", "utility"),
    ]
    rows: list[dict[str, Any]] = []
    for budget in cleaned:
        for label, policy in policies:
            cfg = ExecutionLearningConfig.from_dict(base.to_dict())
            cfg.prefix_cache_budget_fraction = budget
            cfg.prefetch_policy = policy
            result = run_execution_learning(cfg.to_dict(), workflows)
            rows.append(_row(label, result) | {"budget": budget, "policy_label": label})
    winners = []
    for budget in cleaned:
        group = [row for row in rows if abs(float(row["budget"]) - budget) < 1e-12]
        if group:
            winner = min(group, key=lambda row: (row["p95_step_ttft_ms"], row["unused_prefetch_gb"]))
            winners.append({"budget": budget, "policy": winner["policy_label"], "p95_step_ttft_ms": winner["p95_step_ttft_ms"]})
    return {
        "protocol": "common-shifted-agent-workflow-trace",
        "study": "prefix-cache-budget-policy-sweep",
        "rows": rows,
        "winners": winners,
        "note": "All policies see the same workflows at each cache budget; the sweep exposes when broader forecasting helps versus when it increases cache pressure.",
    }