diff --git a/src/server/desktop_audio.py b/src/server/desktop_audio.py index 48d0aab2..61efddd5 100644 --- a/src/server/desktop_audio.py +++ b/src/server/desktop_audio.py @@ -15,6 +15,7 @@ import binascii import logging from dataclasses import dataclass +from typing import Any import httpx diff --git a/src/server/desktop_gateway_methods.py b/src/server/desktop_gateway_methods.py index 77331703..8346ffb1 100644 --- a/src/server/desktop_gateway_methods.py +++ b/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 ``" --provider

--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: diff --git a/tests/server/test_desktop_gateway.py b/tests/server/test_desktop_gateway.py index 28e50580..157413b0 100644 --- a/tests/server/test_desktop_gateway.py +++ b/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."""