cowork-local/core/routing/switch_controller.py

"""Decide whether to switch models, and orchestrate Auto vs Manual execution.

Two independent pieces:

* :func:`decide` — a **pure** function turning ``(current model, ranking, mode,
  threshold)`` into a :class:`SwitchDecision`. No I/O, no state; trivially
  testable.
* :class:`PendingSwitchRegistry` — an in-memory, TTL'd, thread-safe store of
  Manual-mode switch proposals awaiting user confirmation, with an
  **idempotent** ``resolve`` (confirming the same ``request_id`` twice never
  runs the task twice).

Flow (from the task spec)::

    task arrives → classify task_type → selector.best_model()
                 → decide() compares best vs current
                 → gain < min_score_gain          → keep current model
                 → gain ok, mode == AUTO           → switch now, run task
                 → gain ok, mode == MANUAL         → create PendingSwitch, ask user
                 → gain ok, mode == OFF            → never switch (keep current)
"""
from __future__ import annotations

import threading
import time
import uuid
from typing import Callable, Dict, List, Optional, Set

from .models import (
    PendingSwitch,
    SwitchDecision,
    SwitchMode,
    SwitchStatus,
    TaskType,
)
from .selector import Ranking

# Executor signature used by the registry: given the resolved model key and
# whether that represents a switch away from the original, run the task and
# return a JSON-serializable result dict.
Executor = Callable[[str, bool], Dict]


def decide(
    current_key: Optional[str],
    ranking: Ranking,
    mode: SwitchMode,
    min_score_gain: float,
    *,
    task_type: Optional[TaskType] = None,
) -> SwitchDecision:
    """Compare the current model against the ranking's best under ``mode``.

    Returns a :class:`SwitchDecision` whose ``should_switch`` is True only when
    routing is enabled, a better candidate exists, and it beats the current
    model by at least ``min_score_gain``. ``reason`` always explains the call
    in words (e.g. *"coding fit 0.82 > current 0.71, gain 0.11"*).
    """
    tt = task_type or ranking.task_type
    best = ranking.best
    tt_name = tt.value if tt else "?"

    base = dict(
        from_model=current_key,
        to_model=best.key if best else None,
        mode=mode,
        task_type=tt.value if tt else None,
    )

    # Routing disabled → never switch.
    if mode == SwitchMode.OFF:
        return SwitchDecision(
            should_switch=False, score_gain=0.0,
            reason="routing off — keeping current model", **base,
        )

    # Nothing assessed / nothing usable → cannot switch.
    if best is None:
        return SwitchDecision(
            should_switch=False, score_gain=0.0,
            reason="no assessed candidate available for this task", **base,
        )

    to_score = best.score
    from_score = ranking.score_of(current_key) if current_key else 0.0
    best_name = best.assessment.metadata.model_id

    # No current model yet (fresh surface) → adopt the best outright.
    if not current_key:
        return SwitchDecision(
            should_switch=to_score > 0.0,
            from_score=0.0, to_score=to_score, score_gain=to_score,
            reason=f"no current model — selecting best-fit {best_name} ({tt_name} fit {to_score:.2f})",
            **base,
        )

    # Current model is already the best-fit → stay put.
    if best.key == current_key:
        return SwitchDecision(
            should_switch=False,
            from_score=from_score, to_score=to_score, score_gain=0.0,
            reason=f"current model is already best-fit for {tt_name} (fit {to_score:.2f})",
            **base,
        )

    gain = round(to_score - from_score, 6)
    if gain < min_score_gain:
        return SwitchDecision(
            should_switch=False,
            from_score=from_score, to_score=to_score, score_gain=gain,
            reason=(
                f"best {best_name} fit {to_score:.2f} vs current {from_score:.2f}, "
                f"gain {gain:.2f} < threshold {min_score_gain:.2f} — keeping current"
            ),
            **base,
        )

    return SwitchDecision(
        should_switch=True,
        from_score=from_score, to_score=to_score, score_gain=gain,
        reason=(
            f"{tt_name} fit {to_score:.2f} > current {from_score:.2f}, "
            f"gain {gain:.2f} — switch to {best_name}"
        ),
        **base,
    )


class PendingSwitchRegistry:
    """Thread-safe, TTL'd store of Manual-mode switch proposals.

    A proposal is created when Manual mode wants to switch; the UI shows it and
    later calls :meth:`resolve` with the user's approve/reject. Idempotency:
    resolving the same ``request_id`` more than once runs the task exactly once
    and returns the cached result to every caller.

    ``clock`` is injectable so tests can drive expiry deterministically.
    """

    # A short cap so a wedged executor can't hang a waiting confirm forever.
    _RESOLVE_WAIT_SEC = 600.0

    def __init__(self, clock: Callable[[], float] = time.time) -> None:
        """``clock`` tiêm được để test kiểm hết hạn mà không phải chờ thật."""
        self._items: Dict[str, PendingSwitch] = {}
        self._events: Dict[str, threading.Event] = {}
        self._running: Set[str] = set()
        self._lock = threading.Lock()
        self._clock = clock

    # -- creation ------------------------------------------------------- #
    def create(
        self,
        decision: SwitchDecision,
        task_payload: Dict,
        timeout_sec: float,
    ) -> PendingSwitch:
        """Register a new pending switch and return it (with a fresh id)."""
        rid = uuid.uuid4().hex
        now = self._clock()
        ps = PendingSwitch(
            request_id=rid,
            task_payload=task_payload,
            decision=decision,
            created_at=now,
            expires_at=now + max(0.0, float(timeout_sec)),
            status=SwitchStatus.PENDING,
        )
        with self._lock:
            self._items[rid] = ps
            self._events[rid] = threading.Event()
        return ps

    # -- lookup --------------------------------------------------------- #
    def get(self, request_id: str) -> Optional[PendingSwitch]:
        """Fetch a pending switch, lazily marking it EXPIRED if its TTL passed."""
        with self._lock:
            ps = self._items.get(request_id)
            if ps is not None:
                self._maybe_expire_locked(ps)
            return ps

    def _maybe_expire_locked(self, ps: PendingSwitch) -> None:
        """Đánh dấu hết hạn nếu đã quá hạn chờ. Gọi trong lúc đang giữ khoá."""
        if ps.status == SwitchStatus.PENDING and self._clock() >= ps.expires_at:
            ps.status = SwitchStatus.EXPIRED

    # -- resolution ----------------------------------------------------- #
    def resolve(self, request_id: str, approve: bool, run: Executor) -> Optional[Dict]:
        """Confirm (``approve=True``) or reject (``approve=False``) a proposal.

        On the FIRST resolution: runs ``run(model_key, switched)`` where
        ``model_key`` is the proposed model when approved, else the current
        model; caches and returns its result. Subsequent resolutions of the
        same id return the cached result **without** re-running (idempotent).

        An already-EXPIRED proposal is forced down the reject path (run with the
        current model) — matching "timeout → keep current model".

        Returns ``None`` if ``request_id`` is unknown.
        """
        with self._lock:
            ps = self._items.get(request_id)
            if ps is None:
                return None
            self._maybe_expire_locked(ps)
            event = self._events[request_id]

            # Already executed → idempotent replay, no matter who asks.
            if ps.result is not None:
                return ps.result

            expired = ps.status == SwitchStatus.EXPIRED
            effective_approve = bool(approve) and not expired

            # First caller to arrive wins the right to execute exactly once.
            i_run = request_id not in self._running
            if i_run:
                self._running.add(request_id)
                ps.status = (
                    SwitchStatus.CONFIRMED if effective_approve else SwitchStatus.REJECTED
                )

        if not i_run:
            # Another thread is executing — wait for it, then replay its result.
            event.wait(timeout=self._RESOLVE_WAIT_SEC)
            with self._lock:
                return self._items[request_id].result

        # Execute outside the lock (the network/LLM call may be slow).
        decision = ps.decision
        model_key = decision.to_model if effective_approve else decision.from_model
        try:
            result = run(model_key or "", bool(effective_approve))
        finally:
            with self._lock:
                self._running.discard(request_id)
        with self._lock:
            ps.result = result
            event.set()
        return result

    # -- maintenance ---------------------------------------------------- #
    def sweep_expired(self) -> List[str]:
        """Mark all overdue PENDING proposals EXPIRED; return their ids."""
        expired: List[str] = []
        with self._lock:
            for rid, ps in self._items.items():
                if ps.status == SwitchStatus.PENDING and self._clock() >= ps.expires_at:
                    ps.status = SwitchStatus.EXPIRED
                    expired.append(rid)
        return expired

    def purge(self, keep_resolved: bool = False) -> int:
        """Drop resolved/expired entries to free memory. Returns count removed.

        With ``keep_resolved=True``, entries that carry a cached ``result`` are
        retained so their idempotent replay still works.
        """
        removed = 0
        with self._lock:
            for rid in list(self._items):
                ps = self._items[rid]
                terminal = ps.status in (
                    SwitchStatus.CONFIRMED, SwitchStatus.REJECTED, SwitchStatus.EXPIRED
                )
                if terminal and not (keep_resolved and ps.result is not None):
                    self._items.pop(rid, None)
                    self._events.pop(rid, None)
                    self._running.discard(rid)
                    removed += 1
        return removed

    def pending_ids(self) -> List[str]:
        """Id các đề nghị còn đang chờ (đã loại những cái vừa hết hạn)."""
        with self._lock:
            return [
                rid for rid, ps in self._items.items()
                if ps.status == SwitchStatus.PENDING
            ]


__all__ = ["decide", "PendingSwitchRegistry", "Executor"]