fix(desktop): model chip updates after a switch (session.info republish)

src/server/desktop_audio.py

@@ -15,6 +15,7 @@
 import binascii
 import logging
 from dataclasses import dataclass
+from typing import Any
 
 import httpx
 

src/server/desktop_gateway_methods.py

@@ -58,6 +58,13 @@ def _init_session_info(init: dict[str, Any]) -> dict[str, Any]:
     model = init.get("model")
     if model:
         payload["model"] = model
+    # The renderer's picker only prefers the session's selection when BOTH
+    # model and provider are set (currentPickerSelection); omitting provider
+    # made the picker fall back to the catalog while the composer chip kept
+    # showing the session's model — the two disagreed.
+    provider = init.get("provider")
+    if provider:
+        payload["provider"] = provider
     session_id = init.get("session_id")
     if session_id:
         payload["stored_session_id"] = session_id
@@ -82,6 +89,8 @@ def __init__(self, session_id: str, state: DesktopServeState) -> None:
         self._pending_asks: dict[str, dict[str, Any]] = {}
         self._last_ask_id: str | None = None
         self.sockets: set[WebSocket] = set()
+        # Scheduled session.info refreshes; held so they aren't GC'd mid-flight.
+        self._background: set[asyncio.Task] = set()
 
     # ── lifecycle ────────────────────────────────────────────────────────────
 
@@ -95,6 +104,9 @@ async def start(self, cwd: str, spawn: Any = None) -> None:
         )
 
     async def shutdown(self) -> None:
+        for task in list(self._background):
+            task.cancel()
+        self._background.clear()
         if self.pump_task is not None:
             self.pump_task.cancel()
         if self.agent is not None:
@@ -112,6 +124,47 @@ async def _broadcast(self, type_: str, payload: Any = None) -> None:
         for ws in list(self.sockets):
             await send_event(ws, type_, payload, session_id=self.session_id)
 
+    async def publish_session_info(self, **extra: Any) -> None:
+        """Re-read live settings and broadcast a full ``session.info``.
+
+        The renderer reconciles model/provider/effort from session.info and
+        expects them stamped on every one — a switch that doesn't publish
+        leaves the composer chip showing the session's spawn-time model
+        forever (it reads the session state, not the composer draft).
+
+        NEVER await this from inside ``_pump``/``_route``: it issues a control
+        query whose response is routed BY the pump, so awaiting there would
+        deadlock until the timeout. Schedule it with ``refresh_session_info``.
+        """
+        settings = await self.control_query("get_settings", {})
+        payload: dict[str, Any] = {"running": False, "desktop_contract": DESKTOP_CONTRACT}
+        if isinstance(settings, dict):
+            model = settings.get("fusion") or settings.get("model")
+            if model:
+                payload["model"] = str(model)
+            if settings.get("provider"):
+                payload["provider"] = str(settings["provider"])
+            if settings.get("permission_mode"):
+                payload["approval_mode"] = settings["permission_mode"]
+            effort = settings.get("reasoning_effort") or settings.get("effort")
+            if effort:
+                payload["reasoning_effort"] = str(effort)
+        payload.update(extra)
+        await self._broadcast("session.info", payload)
+
+    def refresh_session_info(self) -> None:
+        """Schedule a session.info republish (safe to call from the pump)."""
+        task = asyncio.create_task(self._safe_publish_session_info())
+        self._background.add(task)
+        task.add_done_callback(self._background.discard)
+
+    async def _safe_publish_session_info(self) -> None:
+        try:
+            await self.publish_session_info()
+        except Exception:  # noqa: BLE001 — a refresh must never kill a session
+            logger.debug("desktop session %s: session.info refresh failed",
+                         self.session_id, exc_info=True)
+
     # ── the pump: agent frames → gateway events ─────────────────────────────
 
     async def _pump(self) -> None:
@@ -160,6 +213,11 @@ async def _route(self, frame: dict[str, Any]) -> None:
             mode = frame.get("permission_mode")
             if mode:
                 await self._broadcast("session.info", {"approval_mode": mode, "running": False})
+            # …and republish the full line (model/provider/effort) — a turn can
+            # change them server-side (fallback model, plan-mode flip).
+            # Scheduled, never awaited: this control response routes through
+            # THIS pump.
+            self.refresh_session_info()
             # Turn end persisted the transcript — nudge sidebars to refresh.
             await self._broadcast("sessions.changed", {})
 
@@ -245,6 +303,29 @@ async def apply_model(self, model: str, provider: str | None,
         return {"ok": True, "value": result.get("model") or model,
                 "warning": result.get("warning")}
 
+    async def _config_set_model(self, value: Any) -> dict[str, Any]:
+        """Parse the renderer's model string and apply it.
+
+        The composer sends ``"<model> --provider <p> --session"`` (the TUI's
+        ``/model`` grammar); scope flags are informational here — this
+        transport applies to the live session either way.
+        """
+        tokens = str(value or "").split()
+        provider: str | None = None
+        parts: list[str] = []
+        i = 0
+        while i < len(tokens):
+            tok = tokens[i]
+            if tok == "--provider":
+                i += 1
+                provider = tokens[i] if i < len(tokens) else None
+            elif tok in ("--global", "--session", "--tui-session"):
+                pass
+            else:
+                parts.append(tok)
+            i += 1
+        return await self.apply_model(" ".join(parts), provider)
+
     async def config_set(self, key: str, value: Any, persist: bool = False) -> dict[str, Any]:
         """Route a settings write to the matching agent control.
 
@@ -263,21 +344,12 @@ async def config_set(self, key: str, value: Any, persist: bool = False) -> dict[
                 "persisted": res.get("persisted"),
             }
         if key == "model":
-            tokens = str(value or "").split()
-            provider: str | None = None
-            parts: list[str] = []
-            i = 0
-            while i < len(tokens):
-                tok = tokens[i]
-                if tok == "--provider":
-                    i += 1
-                    provider = tokens[i] if i < len(tokens) else None
-                elif tok in ("--global", "--session", "--tui-session"):
-                    pass
-                else:
-                    parts.append(tok)
-                i += 1
-            return await self.apply_model(" ".join(parts), provider)
+            result = await self._config_set_model(value)
+            # Publish the REAL post-switch state (a cross-provider switch can
+            # land on a different model than requested), so the chip, picker
+            # and settings all reconcile to the truth.
+            await self.publish_session_info()
+            return result
         if key == "logoColor":
             reply = await self.control_query("set_logo_color", {"name": value})
             ok = isinstance(reply, dict) and reply.get("ok") is True
@@ -288,6 +360,11 @@ async def config_set(self, key: str, value: Any, persist: bool = False) -> dict[
             await self.send_control("set_provider", {"provider": value})
         elif key == "thinking":
             await self.send_control("set_thinking", {"action": value})
+        else:
+            return {"ok": True}
+        # A fire-and-forget control still changed the session — republish so
+        # the composer/picker reconcile instead of showing the old value.
+        self.refresh_session_info()
         return {"ok": True}
 
     async def send_control(self, subtype: str, params: dict[str, Any]) -> None:

tests/server/test_desktop_gateway.py

@@ -75,19 +75,36 @@ def __init__(self) -> None:
         self.inbound: list[dict] = []
         self.queue: asyncio.Queue = asyncio.Queue()
         self.shutdown_called = False
+        self.model = "fake"
+        self.provider = "fakeprov"
 
     async def send_to_agent(self, frame: dict) -> None:
         self.inbound.append(frame)
         if frame.get("type") == "control_request":
             request = frame.get("request") or {}
-            if request.get("subtype") == "resume":
+            subtype = request.get("subtype")
+            reply: dict | None = None
+            if subtype == "resume":
+                reply = {"ok": True}
+            elif subtype == "set_model":
+                # Record the switch so get_settings reports the new state.
+                self.model = request.get("model") or self.model
+                self.provider = request.get("provider") or self.provider
+                reply = {"ok": True, "model": self.model}
+            elif subtype == "get_settings":
+                reply = {
+                    "model": self.model,
+                    "provider": self.provider,
+                    "permission_mode": "default",
+                }
+            if reply is not None:
                 await self.queue.put(
                     {
                         "type": "control_response",
                         "response": {
                             "subtype": "success",
                             "request_id": frame.get("request_id"),
-                            "response": {"ok": True},
+                            "response": reply,
                         },
                     }
                 )
@@ -195,10 +212,12 @@ def test_create_submit_stream_complete(tmp_path: Path) -> None:
         _drain_for_response(ws, 2, events)
         complete = _drain_for_event(ws, "message.complete", events)
 
-        assert agents[0].inbound[-1] == {
+        # The prompt reached the agent. (Not necessarily the LAST frame: turn
+        # end schedules a get_settings refresh for the session.info republish.)
+        assert {
             "type": "user",
             "message": {"role": "user", "content": "hi"},
-        }
+        } in agents[0].inbound
         types = [e["type"] for e in events] + ["message.complete"]
         assert "message.start" in types
         deltas = [e for e in events if e["type"] == "message.delta"]
@@ -280,6 +299,69 @@ def reply_frames():
         assert reply["response"]["updatedInput"] == {"command": "rm -rf /tmp/x"}
 
 
+def test_init_session_info_carries_provider() -> None:
+    """The picker only prefers the session's selection when BOTH model and
+    provider are set; without provider it fell back to the catalog while the
+    composer chip kept the session's model — the two disagreed."""
+    from src.server.desktop_gateway_methods import _init_session_info
+
+    info = _init_session_info({
+        "cwd": "/w", "model": "m1", "provider": "p1",
+        "permissionMode": "default", "session_id": "s1",
+    })
+    assert info["model"] == "m1"
+    assert info["provider"] == "p1"
+
+
+def test_model_switch_publishes_session_info(tmp_path: Path) -> None:
+    """The composer chip reads the SESSION's model (not the draft), so a switch
+    that doesn't republish session.info leaves it showing the spawn-time model
+    forever — the reported bug."""
+    state, agents = _fake_state(tmp_path)
+    with TestClient(build_app(state)) as client, _connect(client) as ws:
+        ws.receive_json()
+        events: list[dict] = []
+        _rpc(ws, 1, "session.create", {"cwd": "/tmp"})
+        sid = _drain_for_response(ws, 1, events)["result"]["session_id"]
+
+        events.clear()
+        _rpc(ws, 2, "config.set", {
+            "session_id": sid, "key": "model",
+            "value": "new-model --provider newprov --session",
+        })
+        result = _drain_for_response(ws, 2, events)["result"]
+        assert result["ok"] is True
+
+        # A session.info carrying the NEW model+provider must have been pushed.
+        infos = [e for e in events if e["type"] == "session.info"]
+        assert infos, "no session.info published after the model switch"
+        latest = infos[-1]["payload"]
+        assert latest["model"] == "new-model"
+        assert latest["provider"] == "newprov"
+
+
+def test_turn_end_republishes_session_info_without_deadlock(tmp_path: Path) -> None:
+    """Turn end refreshes the info line. The refresh issues a control query
+    whose response is routed by the pump, so it must be SCHEDULED — awaiting it
+    inside the pump would deadlock until the control timeout."""
+    state, agents = _fake_state(tmp_path)
+    with TestClient(build_app(state)) as client, _connect(client) as ws:
+        ws.receive_json()
+        events: list[dict] = []
+        _rpc(ws, 1, "session.create", {"cwd": "/tmp"})
+        sid = _drain_for_response(ws, 1, events)["result"]["session_id"]
+
+        events.clear()
+        _rpc(ws, 2, "prompt.submit", {"session_id": sid, "text": "hi"})
+        _drain_for_response(ws, 2, events)
+        # message.complete proves the pump kept draining (no deadlock)…
+        _drain_for_event(ws, "message.complete", events)
+        # …and the scheduled refresh lands with model/provider stamped.
+        info = _drain_for_event(ws, "session.info", events)
+        assert info["payload"]["model"] == "fake"
+        assert info["payload"]["provider"] == "fakeprov"
+
+
 def test_session_create_honors_provider_override(tmp_path: Path) -> None:
     """A composer selection (provider/model) must reach the spawn, so a session
     can use a working provider even when the config default is broken."""