codex-app-server/cloud/models.py

"""Model fallback routing and pre-dispatch rule routing (v0).

Zero third-party dependencies; standard library only.

This module decides *whether* a failed model attempt may fall back to
the next model in a priority list, and *which* model to try next. It
never runs a model turn itself and never talks to any provider.

Rule routing (v0) is deliberately narrower than that: it only fills
in the model for tasks that did not specify one, using structural
signals already stored on the task (``plan``, ``goal_mode`` /
``verify_command``). It never changes ``effort``, ``approval``,
``plan``, ``executor`` or any quota, never switches provider/gateway,
and never invents a quality ranking between models. A policy that
cannot be parsed, or a tier with no model, falls back to the existing
worker-default chain instead of failing the task.
"""

import json
import os

FALLBACKABLE = {404, 429, 500, 502, 503, 504, "timeout"}


def _is_fallbackable_code(value):
    """Return True for a single status code / error token."""
    if isinstance(value, bool):
        return False
    if isinstance(value, int):
        if value in FALLBACKABLE:
            return True
        # Any other 5xx is a server-side failure of the same kind.
        return 500 <= value <= 599
    if isinstance(value, str):
        text = value.strip().lower()
        if text in FALLBACKABLE:
            return True
        if text.isdigit():
            return _is_fallbackable_code(int(text))
        return False
    return False


def is_fallbackable(error) -> bool:
    """Return True if ``error`` may trigger a model fallback.

    ``error`` may be:

    * an int HTTP status code (404, 429, 5xx -> True; 401, 400 -> False);
    * the string ``"timeout"``;
    * an object / exception carrying a ``status`` or ``code`` attribute
      holding one of the above.

    Authentication failures (401) and bad-request errors (400) are
    deliberately *not* fallbackable: retrying them on another model
    cannot fix bad credentials or a bad request.
    """
    if error is None:
        return False
    if isinstance(error, (int, str)) and not isinstance(error, bool):
        return _is_fallbackable_code(error)
    for attr in ("status", "code"):
        value = getattr(error, attr, None)
        if value is not None:
            return _is_fallbackable_code(value)
    return False


class ModelRouter:
    """Walk a priority list of models, one fallback at a time."""

    def __init__(self, priority_list):
        if not priority_list:
            raise ValueError("priority_list must not be empty")
        self.priority_list = list(priority_list)

    def first(self):
        """Return the first (highest-priority) model."""
        return self.priority_list[0]

    def next_after(self, current_model, error):
        """Return the model after ``current_model``, or None.

        Returns None when ``error`` is not fallbackable, when
        ``current_model`` is not in the priority list, or when there
        is no model left after ``current_model``.
        """
        if not is_fallbackable(error):
            return None
        try:
            index = self.priority_list.index(current_model)
        except ValueError:
            return None
        if index + 1 < len(self.priority_list):
            return self.priority_list[index + 1]
        return None

    def route(self, attempts):
        """Return the next model to try given ``attempts``, or None.

        ``attempts`` is a list of ``(model, error)`` pairs in the order
        they were tried. With no attempts yet, the next model to try
        is the first in the priority list. A non-fallbackable error
        anywhere in the chain stops routing (returns None), as does
        exhausting the priority list.
        """
        if not attempts:
            return self.first()
        last_model = None
        for model, error in attempts:
            if not is_fallbackable(error):
                return None
            last_model = model
        return self.next_after(last_model, attempts[-1][1])


ROUTING_MODES = ("off", "shadow", "active")
ROUTING_TIERS = ("plan", "goal", "default")


def prompt_len_bucket(prompt) -> str:
    """Bucket a prompt by length for logging only, never for routing.

    Length is recorded so shadow reports can test whether it correlates
    with outcomes; v0 deliberately does not route on it (a long prompt
    may just be a pasted log, not a hard task).
    """
    length = len(prompt) if isinstance(prompt, str) else 0
    if length < 500:
        return "<500"
    if length <= 3000:
        return "500-3000"
    return ">3000"


def classify_task(task) -> str:
    """Classify a task into one routing tier, first match wins.

    Only structural signals are used, in this order: ``plan`` (read-only,
    never mutates files) before ``goal`` (multi-round goal loop, possibly
    unbounded rounds), everything else is ``default``. Keyword, prompt
    length, priority, ``depends_on`` count and executor are intentionally
    excluded: they are weak or unrelated proxies for difficulty.
    """
    if not isinstance(task, dict):
        return "default"
    if task.get("plan"):
        return "plan"
    if task.get("goal_mode") or task.get("verify_command"):
        return "goal"
    return "default"


def _clean_model_name(value):
    if isinstance(value, str) and value.strip():
        return value.strip()
    return None


def _clean_tiers(raw):
    """Return {tier: model} keeping only valid non-empty entries."""
    tiers = {}
    if not isinstance(raw, dict):
        return tiers
    for tier in ROUTING_TIERS:
        model = _clean_model_name(raw.get(tier))
        if isinstance(model, str):
            # A tier entry may itself be a dict {"model": ...}; accept
            # that shape too so policy authors are not punished for it.
            tiers[tier] = model
        elif isinstance(raw.get(tier), dict):
            model = _clean_model_name(raw[tier].get("model"))
            if model:
                tiers[tier] = model
    return tiers


class RoutingPolicy:
    """A parsed, versioned routing policy (or an empty one).

    ``mode`` is one of ``off`` / ``shadow`` / ``active``. ``tiers`` maps
    a tier name to a model name; a tier absent from the map falls back
    to the worker default. ``fallback_models`` is parsed and validated
    but deliberately NOT applied by v0: the worker's explicit
    ``--fallback-models`` remains the only source that can enable a real
    fallback retry, so a policy file alone can never start burning a
    second attempt's quota.
    """

    def __init__(self, version=None, mode="shadow", tiers=None,
                 fallback_models=None):
        self.version = version if isinstance(version, str) \
            and version.strip() else None
        self.mode = mode if mode in ROUTING_MODES else "off"
        self.tiers = dict(tiers or {})
        self.fallback_models = []
        for item in fallback_models or []:
            model = _clean_model_name(item)
            if model and model not in self.fallback_models:
                self.fallback_models.append(model)


def parse_routing_policy(data):
    """Parse one policy document, or return None when it is unusable.

    Unusable means: not a dict, no usable ``tiers`` map, or a ``mode``
    outside ``off``/``shadow``/``active``. Returning None (instead of
    raising) is the fail-closed path: the caller falls back to the
    worker-default chain and the task is never rejected for a bad
    policy file.
    """
    if not isinstance(data, dict):
        return None
    mode = data.get("mode", "shadow")
    if mode not in ROUTING_MODES:
        return None
    tiers = _clean_tiers(data.get("tiers"))
    if not tiers:
        return None
    fallback_models = data.get("fallback_models")
    if fallback_models is not None and not isinstance(fallback_models,
                                                     list):
        return None
    return RoutingPolicy(version=data.get("version"), mode=mode,
                         tiers=tiers, fallback_models=fallback_models)


def load_routing_policy(env=None):
    """Load a policy from env, following the cost-table convention.

    ``CLOUD_ROUTING_POLICY_FILE`` is read first, then
    ``CLOUD_ROUTING_POLICY_JSON`` overrides its fields per tier;
    ``CLOUD_ROUTING_MODE`` overrides the mode when valid. Any source
    that fails to parse is ignored, never fatal. Returns None when no
    usable policy was found, so the caller keeps the old behaviour.
    """
    env = env if env is not None else os.environ
    merged = None
    path = env.get("CLOUD_ROUTING_POLICY_FILE") or ""
    if path:
        try:
            with open(path, encoding="utf-8") as f:
                merged = json.load(f)
        except (OSError, ValueError):
            merged = None
    raw = env.get("CLOUD_ROUTING_POLICY_JSON") or ""
    if raw:
        try:
            override = json.loads(raw)
        except ValueError:
            override = None
        if isinstance(override, dict):
            if not isinstance(merged, dict):
                merged = override
            else:
                merged = dict(merged)
                merged.update(override)
                if isinstance(merged.get("tiers"), dict) and \
                        isinstance(override.get("tiers"), dict):
                    tiers = dict(merged["tiers"])
                    tiers.update(override["tiers"])
                    merged["tiers"] = tiers
    policy = parse_routing_policy(merged)
    if policy is None:
        return None
    mode_override = env.get("CLOUD_ROUTING_MODE") or ""
    if mode_override in ROUTING_MODES:
        policy.mode = mode_override
    return policy


class RoutingDecision:
    """One pre-dispatch model decision, kept auditable by construction."""

    def __init__(self, task_id, tier, chosen_model, suggested_model,
                 source, policy_version, applied, reason, signals):
        self.task_id = task_id
        self.tier = tier
        self.chosen_model = chosen_model
        self.suggested_model = suggested_model
        self.source = source
        self.policy_version = policy_version
        self.applied = applied
        self.reason = reason
        self.signals = dict(signals or {})

    def to_event(self) -> dict:
        return {
            "type": "model_routing",
            "task_id": self.task_id,
            "tier": self.tier,
            "chosen_model": self.chosen_model,
            "suggested_model": self.suggested_model,
            "source": self.source,
            "policy_version": self.policy_version,
            "applied": self.applied,
            "reason": self.reason,
        }


def decide_routing(task, worker_model, policy=None,
                   worker_model_explicit=False):
    """Decide the model for one task, in the documented priority order.

    1. task-level ``model`` (explicit) -- never overridden;
    2. worker CLI ``--model`` (explicit) -- a pinned deployment is never
       silently changed by a policy file;
    3. active policy tier model -- only when the task named no model;
    4. worker default chain (config file / ``DEFAULT_MODEL``).

    Shadow mode returns the worker default as ``chosen_model`` and the
    policy tier model as ``suggested_model`` with ``applied=False``:
    it observes without changing behaviour. This function is pure: it
    never reads env, files, or a provider, so it is unit-testable with
    plain dicts.
    """
    task = task if isinstance(task, dict) else {}
    tier = classify_task(task)
    signals = {
        "plan": bool(task.get("plan")),
        "goal_mode": bool(task.get("goal_mode")),
        "prompt_len_bucket": prompt_len_bucket(task.get("prompt")),
        "effort": task.get("effort"),
    }
    task_model = _clean_model_name(task.get("model"))
    if task_model:
        return RoutingDecision(
            task_id=task.get("id"), tier=tier, chosen_model=task_model,
            suggested_model=None, source="explicit", policy_version=None,
            applied=True, reason="task model explicit", signals=signals)
    worker_model = _clean_model_name(worker_model)
    if worker_model_explicit and worker_model:
        return RoutingDecision(
            task_id=task.get("id"), tier=tier, chosen_model=worker_model,
            suggested_model=None, source="worker_explicit",
            policy_version=None, applied=True,
            reason="worker CLI model explicit", signals=signals)
    if policy is None or policy.mode == "off":
        return RoutingDecision(
            task_id=task.get("id"), tier=tier, chosen_model=worker_model,
            suggested_model=None, source="worker_default",
            policy_version=getattr(policy, "version", None),
            applied=True, reason="routing off or no policy",
            signals=signals)
    tier_model = policy.tiers.get(tier)
    if not tier_model:
        return RoutingDecision(
            task_id=task.get("id"), tier=tier, chosen_model=worker_model,
            suggested_model=None, source="worker_default",
            policy_version=policy.version, applied=True,
            reason="%s tier has no policy model" % tier, signals=signals)
    if policy.mode == "shadow":
        return RoutingDecision(
            task_id=task.get("id"), tier=tier, chosen_model=worker_model,
            suggested_model=tier_model, source="shadow",
            policy_version=policy.version, applied=False,
            reason="%s tier suggests %s (shadow, not applied)"
            % (tier, tier_model), signals=signals)
    # mode == "active"
    return RoutingDecision(
        task_id=task.get("id"), tier=tier, chosen_model=tier_model,
        suggested_model=None, source="policy",
        policy_version=policy.version, applied=True,
        reason="%s tier -> %s" % (tier, tier_model), signals=signals)


def record_routing(task_record, decision) -> dict:
    """Append a ``model_routing`` event to a task record, like fallback.

    Mirrors ``record_fallback``: the event is appended to
    ``task_record["events"]`` (created if absent) so a routed model is
    never silent. Returns the event.
    """
    event = decision.to_event() if isinstance(decision, RoutingDecision) \
        else dict(decision)
    events = task_record.get("events")
    if not isinstance(events, list):
        events = []
        task_record["events"] = events
    events.append(event)
    return event


def record_fallback(task_record, from_model, to_model, reason) -> dict:
    """Record a model fallback on ``task_record`` and return the event.

    Fallback must be recorded as an event, never silently downgraded:
    the returned event is appended to ``task_record["events"]``
    (created if absent) so the downgrade is always visible in the
    task's history.
    """
    event = {
        "type": "model_fallback",
        "task_id": task_record.get("id"),
        "from_model": from_model,
        "to_model": to_model,
        "reason": reason,
    }
    events = task_record.get("events")
    if not isinstance(events, list):
        events = []
        task_record["events"] = events
    events.append(event)
    return event