diff --git a/src/cookbook_serve_lifecycle.py b/src/cookbook_serve_lifecycle.py index fcdacbe7a..f2700cf7d 100644 --- a/src/cookbook_serve_lifecycle.py +++ b/src/cookbook_serve_lifecycle.py @@ -161,11 +161,13 @@ async def _tick() -> None: # Re-read state once before writing so we capture any updates from # concurrent UI syncs. stopped_any = False + successfully_stopped_sids = set() for sid, host, port in to_stop: ok = await _stop_serve(sid, host, port) logger.info(f"cookbook_serve_lifecycle: stop {sid} (host={host or 'local'}): {'ok' if ok else 'failed'}") if ok: stopped_any = True + successfully_stopped_sids.add(sid) # Drop the auto-registered endpoint so the model picker and # the chat router don't keep pointing at a dead server. for t in tasks: @@ -188,12 +190,11 @@ async def _tick() -> None: except Exception: fresh = state fresh_tasks = tasks - stopped_sids = {sid for sid, _, _ in to_stop} for ft in fresh_tasks: if not isinstance(ft, dict): continue ft_sid = ft.get("sessionId") or ft.get("id") - if ft_sid in stopped_sids: + if ft_sid in successfully_stopped_sids: ft["status"] = "stopped" ft["_scheduledStopAtMs"] = None ft["_lastStatusFlipAt"] = now_ms diff --git a/tests/test_cookbook_serve_lifecycle.py b/tests/test_cookbook_serve_lifecycle.py new file mode 100644 index 000000000..cd32d4314 --- /dev/null +++ b/tests/test_cookbook_serve_lifecycle.py @@ -0,0 +1,50 @@ +import json + +import pytest + +from src import cookbook_serve_lifecycle as lifecycle + + +@pytest.mark.asyncio +async def test_tick_persists_only_successfully_stopped_serves(tmp_path, monkeypatch): + state_path = tmp_path / "cookbook_state.json" + state_path.write_text( + json.dumps({ + "tasks": [ + { + "id": "stop-succeeds", + "type": "serve", + "status": "running", + "_scheduledStopAtMs": 0, + }, + { + "id": "stop-fails", + "type": "serve", + "status": "running", + "_scheduledStopAtMs": 0, + }, + ] + }), + encoding="utf-8", + ) + + async def fake_stop_serve(session_id, remote_host="", ssh_port=""): + return session_id == "stop-succeeds" + + async def fake_delete_endpoint(task): + return None + + monkeypatch.setattr(lifecycle, "COOKBOOK_STATE_FILE", str(state_path)) + monkeypatch.setattr(lifecycle, "_stop_serve", fake_stop_serve) + monkeypatch.setattr(lifecycle, "_delete_endpoint_for_task", fake_delete_endpoint) + + await lifecycle._tick() + + tasks = { + task["id"]: task + for task in json.loads(state_path.read_text(encoding="utf-8"))["tasks"] + } + assert tasks["stop-succeeds"]["status"] == "stopped" + assert tasks["stop-succeeds"]["_scheduledStopAtMs"] is None + assert tasks["stop-fails"]["status"] == "running" + assert tasks["stop-fails"]["_scheduledStopAtMs"] == 0