cowork-local/application/model_routing/routing_application_service.py

"""RoutingApplicationService - one routing flow for every surface (R03-T03).

Before this service, the same routing algorithm existed three times:

* ``ui/chat_panel.py::_apply_routing``       (Cowork chat)
* ``ui/co4e_tab.py::_apply_co4e_routing``    (Co4E studio)
* ``ui/folder_tab.py::_ai_apply_routing``    (AI-Edit)

The three copies had already drifted - each one resolves the "current model"
differently and each one has its own private notion of what to do when the user
declines - and every one of them lives inside a Qt widget, so none of the logic
could be tested without building a window.

This module is the single implementation. It is pure Python: no Qt import, no
config access, no network. The presentation layer supplies a confirm callback
and renders the notice; everything else happens here.

Modes (:class:`RoutingMode`)
----------------------------
* ``OFF``      - never switch. The user's pinned model always wins.
* ``AUTO``     - switch silently when the best candidate clears the gain threshold.
* ``MANUAL``   - propose the switch and switch only if the confirm callback approves.
* ``FALLBACK`` - never switch pre-emptively; switch only AFTER the current model
  fails, to the next-best candidate. This is the mode a user wants when they
  trust their own model choice but still want the turn to survive an outage.

Migration note (ADR-001 section 4): the scoring/ranking engine is NOT rewritten.
This service depends on the small :class:`RoutingPort` interface, and production
wires the existing, already-tested ``core.routing.service.RoutingService`` into
it. Tests wire a fake.
"""
from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, List, Optional, Protocol, Sequence, Tuple


class RoutingMode(str, Enum):
    """Per-surface routing behaviour.

    The first three values match ``core.routing.models.SwitchMode`` string for
    string, so a mode read from the existing config round-trips unchanged.
    """

    OFF = "off"
    AUTO = "auto"
    MANUAL = "manual"
    FALLBACK = "fallback"

    @classmethod
    def parse(cls, raw: Any) -> "RoutingMode":
        """Best-effort parse of a config value.

        Unknown or empty values become ``OFF``: routing is an optimisation, and
        the safe reading of a corrupt setting is "leave the user's model alone"
        rather than "silently move their work to another model".
        """
        try:
            return cls(str(raw or "off").strip().lower())
        except ValueError:
            return cls.OFF


@dataclass(frozen=True)
class RoutingDecision:
    """The outcome of routing one turn - an immutable instruction for the caller.

    ``provider``/``model`` are ALWAYS filled with what the turn should actually
    run on, switched or not, so a call site never has to re-derive the fallback
    itself (the bug that made the three UI copies diverge).
    """

    mode: RoutingMode
    provider: str
    model: str
    switched: bool = False
    task_type: str = ""
    score_gain: float = 0.0
    reason: str = ""
    declined: bool = False  # Manual mode: a switch was offered and refused
    # What the turn would have run on without routing. Carried so the Manual
    # confirm dialog can show "from X to Y" without re-deriving the current
    # model itself - re-deriving it differently per screen is exactly how the
    # three legacy copies drifted apart.
    previous_provider: str = ""
    previous_model: str = ""

    @property
    def should_notify(self) -> bool:
        """True when the UI should show the "switched model" notice - i.e. only
        when a switch really happened."""
        return self.switched

    def target(self) -> Tuple[str, str]:
        """``(provider, model)`` to run this turn on."""
        return self.provider, self.model

    @property
    def from_model(self) -> str:
        """Candidate key (``provider/model``) of the model being switched away
        from, or "" when nothing was selected yet.

        Named to match ``core.routing.models.SwitchDecision`` so the existing
        Manual-mode dialog (``ui/routing_toggle.py::confirm_switch``) accepts
        this object unchanged - the dialog moves to the new shape in EPIC R08.
        """
        if not self.previous_model:
            return ""
        return f"{self.previous_provider}/{self.previous_model}"

    @property
    def to_model(self) -> str:
        """Candidate key (``provider/model``) of the model to run on. See
        :attr:`from_model` for why the name matches the legacy decision."""
        return f"{self.provider}/{self.model}" if self.model else ""


def is_valid_mode(raw: Any) -> bool:
    """True when ``raw`` names a mode the routing service understands.

    Distinct from :func:`normalize_mode` because callers need to tell "the user
    chose off" apart from "this stored value is unrecognised" - the per-workspace
    lookup falls back to the global setting only in the second case.
    """
    try:
        RoutingMode(str(raw or "").strip().lower())
    except ValueError:
        return False
    return True


def normalize_mode(raw: Any) -> str:
    """Canonical mode string for persistence, or ``"off"`` when unrecognised.

    Exists so the mode vocabulary is defined exactly once. It used to be
    hard-coded as a ``("off", "auto", "manual")`` tuple in four separate places
    (config.py twice, state.py twice); adding FALLBACK meant finding all four,
    and missing one silently downgraded the user's choice back to "off".
    """
    return RoutingMode.parse(raw).value


class RoutingPort(Protocol):
    """The slice of the routing engine this service needs.

    Declared as a Protocol so the application layer states its requirement
    without importing the implementation - which is what lets the whole service
    be tested against a 20-line fake, and lets ``core.routing`` be replaced later
    without touching this file.
    """

    def route(self, surface: str, prompt: str, current_provider: str, current_model: str,
              *, mode_override: Optional[str] = None,
              required_capabilities: Optional[List[str]] = None,
              task_type: Optional[Any] = None) -> Any:
        """Return a route result exposing ``should_switch``, ``target()``,
        ``task_type`` and ``decision``."""


# Presentation supplies this to ask the human. Receives the proposal so the
# dialog can explain it; returns True to approve. Manual mode only.
ConfirmFn = Callable[[RoutingDecision], bool]


class RoutingApplicationService:
    """Decides which provider/model one turn runs on.

    Args:
        router: the scoring engine (see :class:`RoutingPort`).
        mode_reader: ``surface -> mode string``; production passes the per-workspace
            lookup ``AppContext.project_routing_mode``. Injected rather than read
            from config here so this layer stays free of config plumbing that
            EPIC R02 is rewriting in parallel.
    """

    def __init__(self, router: RoutingPort,
                 mode_reader: Optional[Callable[[str], str]] = None) -> None:
        self._router = router
        self._mode_reader = mode_reader

    # -- main entry point -------------------------------------------------- #
    def route_turn(
        self,
        surface: str,
        prompt: str,
        current_provider: str,
        current_model: str,
        *,
        mode: Optional[str] = None,
        confirm: Optional[ConfirmFn] = None,
        required_capabilities: Optional[Sequence[str]] = None,
        task_type: Optional[Any] = None,
    ) -> RoutingDecision:
        """Decide what to run this turn on. Never raises.

        A routing failure must never block a message: any unexpected error
        degrades to "keep the current model", which is exactly what all three
        legacy copies did with a bare ``except`` - made explicit and testable here.
        """
        resolved_mode = RoutingMode.parse(mode if mode is not None else self._read_mode(surface))
        keep = self._keep(resolved_mode, current_provider, current_model,
                          reason="routing off - keeping current model")

        # An empty prompt carries no signal to classify, so routing cannot make a
        # meaningful choice; the same guard exists in all three legacy copies.
        if resolved_mode is RoutingMode.OFF or not (prompt or "").strip():
            return keep

        # FALLBACK never switches up front - it only reacts to a failure, which
        # the caller reports through fallback_after_failure().
        if resolved_mode is RoutingMode.FALLBACK:
            return self._keep(resolved_mode, current_provider, current_model,
                              reason="fallback mode - switching only after a failure")

        try:
            result = self._router.route(
                surface, prompt, current_provider, current_model,
                mode_override=resolved_mode.value,
                required_capabilities=list(required_capabilities) if required_capabilities else None,
                task_type=task_type,
            )
        except Exception:  # noqa: BLE001 - routing must never break a turn
            return self._keep(resolved_mode, current_provider, current_model,
                              reason="routing engine failed - keeping current model")

        proposal = self._to_decision(result, resolved_mode, current_provider, current_model)
        if not proposal.switched:
            return proposal

        # Manual mode: the proposal only becomes a switch once a human approves.
        if resolved_mode is RoutingMode.MANUAL:
            if confirm is None or not self._ask(confirm, proposal):
                return self._keep(resolved_mode, current_provider, current_model,
                                  reason="switch declined - keeping current model",
                                  task_type=proposal.task_type, declined=True)
        return proposal

    # -- failure recovery -------------------------------------------------- #
    def fallback_after_failure(
        self,
        surface: str,
        prompt: str,
        failed_provider: str,
        failed_model: str,
        *,
        mode: Optional[str] = None,
        required_capabilities: Optional[Sequence[str]] = None,
        task_type: Optional[Any] = None,
    ) -> Optional[RoutingDecision]:
        """Pick a replacement after ``failed_provider/failed_model`` failed.

        Returns None when there is nothing to fall back to, so the caller can
        surface the original error instead of retrying forever. Available in
        AUTO and FALLBACK; OFF and MANUAL keep the user's model on failure too,
        because silently moving work to another model is exactly what those two
        modes exist to prevent.
        """
        resolved_mode = RoutingMode.parse(mode if mode is not None else self._read_mode(surface))
        if resolved_mode not in (RoutingMode.AUTO, RoutingMode.FALLBACK):
            return None

        try:
            # Asked in AUTO so the engine ranks candidates rather than short-
            # circuiting on FALLBACK's "never switch up front" rule; the failed
            # model is passed as current so any positive gain beats it.
            result = self._router.route(
                surface, prompt, failed_provider, failed_model,
                mode_override=RoutingMode.AUTO.value,
                required_capabilities=list(required_capabilities) if required_capabilities else None,
                task_type=task_type,
            )
        except Exception:  # noqa: BLE001 - a broken router must not mask the real error
            return None

        decision = self._to_decision(result, resolved_mode, failed_provider, failed_model)
        # A "switch" back to the model that just failed would retry the outage.
        if not decision.switched or (decision.provider, decision.model) == (failed_provider, failed_model):
            return None
        return RoutingDecision(
            mode=resolved_mode, provider=decision.provider, model=decision.model,
            switched=True, task_type=decision.task_type, score_gain=decision.score_gain,
            reason=f"{failed_provider}/{failed_model} failed - falling back to "
                   f"{decision.provider}/{decision.model}",
            previous_provider=failed_provider, previous_model=failed_model,
        )

    # -- internals --------------------------------------------------------- #
    def _read_mode(self, surface: str) -> str:
        """Per-surface mode from the injected reader ('off' when none supplied)."""
        if self._mode_reader is None:
            return RoutingMode.OFF.value
        try:
            return self._mode_reader(surface) or RoutingMode.OFF.value
        except Exception:  # noqa: BLE001 - a config read must not break a turn
            return RoutingMode.OFF.value

    @staticmethod
    def _keep(mode: RoutingMode, provider: str, model: str, *, reason: str,
              task_type: str = "", declined: bool = False) -> RoutingDecision:
        """A no-switch decision that still names the model to run on."""
        return RoutingDecision(mode=mode, provider=provider, model=model, switched=False,
                               task_type=task_type, reason=reason, declined=declined,
                               previous_provider=provider, previous_model=model)

    @staticmethod
    def _ask(confirm: ConfirmFn, proposal: RoutingDecision) -> bool:
        """Run the confirm callback, treating any failure as "declined".

        The callback opens a modal dialog in production; if that raises (window
        already closing, for instance) the safe answer is to keep the user's own
        model rather than to switch without consent.
        """
        try:
            return bool(confirm(proposal))
        except Exception:  # noqa: BLE001
            return False

    @staticmethod
    def _to_decision(result: Any, mode: RoutingMode,
                     current_provider: str, current_model: str) -> RoutingDecision:
        """Translate the engine's route result into a :class:`RoutingDecision`.

        Defensive about the result shape on purpose: this is the seam between the
        new layer and a legacy module still under refactor, and a missing
        attribute must degrade to "keep current model" instead of raising into
        the middle of a chat turn.
        """
        inner = getattr(result, "decision", None)
        task_type = getattr(getattr(result, "task_type", None), "value", "") or ""
        gain = float(getattr(inner, "score_gain", 0.0) or 0.0)
        reason = str(getattr(inner, "reason", "") or "")

        target = None
        if getattr(result, "should_switch", False):
            getter = getattr(result, "target", None)
            target = getter() if callable(getter) else None

        if not target:
            return RoutingDecision(mode=mode, provider=current_provider, model=current_model,
                                   switched=False, task_type=task_type, score_gain=gain,
                                   reason=reason or "no better model - keeping current",
                                   previous_provider=current_provider,
                                   previous_model=current_model)

        provider, model = target
        return RoutingDecision(mode=mode, provider=provider or current_provider, model=model,
                               switched=True, task_type=task_type, score_gain=gain, reason=reason,
                               previous_provider=current_provider, previous_model=current_model)


__all__ = ["RoutingApplicationService", "RoutingDecision", "RoutingMode",
           "RoutingPort", "normalize_mode", "is_valid_mode"]