odysseus/src/cookbook_serve_lifecycle.py

"""Cookbook serve lifecycle: kills scheduler-owned serves whose end-of-
window has passed.

Pairs with action_cookbook_serve in builtin_actions.py — that action
stamps the task it launches with `_scheduledStopAtMs`, this loop ticks
every 60s and kills any serve whose stamp is in the past.

Single small module. Delete this file + the registration line in app.py
and the feature stops doing anything; scheduler-launched serves just
stay up until the user kills them manually.
"""

from __future__ import annotations

import asyncio
import json
import logging
import time
from pathlib import Path

import httpx
from core.constants import internal_api_base
from src.constants import COOKBOOK_STATE_FILE

logger = logging.getLogger(__name__)


def _internal_headers() -> dict:
    from core.middleware import INTERNAL_TOOL_HEADER, INTERNAL_TOOL_TOKEN
    return {INTERNAL_TOOL_HEADER: INTERNAL_TOOL_TOKEN}


async def _delete_endpoint_for_task(task: dict) -> None:
    """Remove a stopped serve's endpoint using the normal deletion contract."""
    import re
    from urllib.parse import quote

    endpoint_id = str(task.get("_endpointId") or task.get("endpointId") or "").strip()
    serve_session_id = str(task.get("sessionId") or task.get("id") or "").strip()
    if (
        not re.fullmatch(r"(?:local|img)-[0-9a-f]{8}", endpoint_id)
        or not serve_session_id
    ):
        logger.info("cookbook_serve_lifecycle: missing endpoint/session identity")
        return

    try:
        async with httpx.AsyncClient(timeout=8) as client:
            response = await client.delete(
                f"{internal_api_base()}/api/model-endpoints/{quote(endpoint_id, safe='')}",
                params={"expected_cookbook_serve_session_id": serve_session_id},
                headers=_internal_headers(),
            )
        if response.status_code == 200:
            logger.info(
                "cookbook_serve_lifecycle: deleted endpoint %s for serve %s",
                endpoint_id, serve_session_id,
            )
        elif response.status_code in (404, 409):
            logger.info(
                "cookbook_serve_lifecycle: endpoint %s absent or reassigned; keeping it",
                endpoint_id,
            )
        else:
            logger.warning(
                "cookbook_serve_lifecycle: endpoint %s deletion returned HTTP %s",
                endpoint_id, response.status_code,
            )
    except Exception as exc:
        logger.warning("cookbook_serve_lifecycle: cleanup failed: %s", exc)


async def _stop_serve(session_id: str, remote_host: str = "", ssh_port: str = "") -> bool:
    """Kill the tmux session that hosts the serve.

    There's no `/api/model/stop` route — the cookbook UI and the chat
    agent both kill via `/api/shell/exec` running a `tmux kill-session`
    (wrapped in ssh for remote hosts). Mirror that here so the
    lifecycle loop can actually stop scheduler-launched serves at
    window-end. Without this, the action stamped `_scheduledStopAtMs`
    correctly but every kill attempt failed silently (the route
    returned 404 and the result was logged as "failed").
    """
    import shlex
    if remote_host:
        port_flag = f"-p {shlex.quote(str(ssh_port))} " if ssh_port and str(ssh_port) != "22" else ""
        cmd = (
            f"ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no "
            f"{port_flag}{shlex.quote(remote_host)} "
            f"'tmux kill-session -t {shlex.quote(session_id)}'"
        )
    else:
        cmd = f"tmux kill-session -t {shlex.quote(session_id)}"
    try:
        async with httpx.AsyncClient(timeout=15) as client:
            r = await client.post(
                f"{internal_api_base()}/api/shell/exec",
                json={"command": cmd},
                headers=_internal_headers(),
            )
            if r.status_code >= 400:
                return False
            data = r.json() if r.content else {}
            ec = data.get("exit_code")
            # tmux returns non-zero when the session is already gone
            # ("can't find session: ..."). That's still "stop succeeded"
            # from our POV — the goal is no live session at the end.
            if ec in (None, 0):
                return True
            stderr = (data.get("stderr") or "").lower()
            return "no server" in stderr or "can't find session" in stderr or "session not found" in stderr
    except Exception as e:
        logger.warning(f"cookbook_serve_lifecycle: stop {session_id} failed: {e}")
        return False


async def _tick() -> None:
    state_path = Path(COOKBOOK_STATE_FILE)
    if not state_path.exists():
        return
    try:
        state = json.loads(state_path.read_text(encoding="utf-8"))
    except Exception as e:
        logger.warning("cookbook_serve_lifecycle: state file unreadable (%s), skipping tick", e)
        return
    tasks = state.get("tasks") or []
    now_ms = int(time.time() * 1000)
    to_stop = []
    for t in tasks:
        if not isinstance(t, dict):
            continue
        stop_at = t.get("_scheduledStopAtMs")
        if not isinstance(stop_at, (int, float)):
            continue
        if stop_at > now_ms:
            continue
        if (t.get("status") or "").lower() in {"stopped", "ended", "killed", "crashed"}:
            continue
        sid = t.get("sessionId") or t.get("id")
        if not sid:
            continue
        to_stop.append((sid, t.get("remoteHost") or "", t.get("sshPort") or ""))
    if not to_stop:
        return
    # 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:
                if isinstance(t, dict) and (t.get("sessionId") == sid or t.get("id") == sid):
                    if t.get("type") == "serve":
                        await _delete_endpoint_for_task(t)
                    t["status"] = "stopped"
                    t["_scheduledStopAtMs"] = None
                    t["_lastStatusFlipAt"] = now_ms
                    break
    if stopped_any:
        try:
            from core.atomic_io import atomic_write_json
            # Re-read the state file so concurrent UI writes (task adds,
            # status flips, config edits) are not silently overwritten.
            # Apply only our stop mutations to the fresh snapshot.
            try:
                fresh = json.loads(state_path.read_text(encoding="utf-8"))
                fresh_tasks = fresh.get("tasks") or []
            except Exception:
                fresh = state
                fresh_tasks = tasks
            for ft in fresh_tasks:
                if not isinstance(ft, dict):
                    continue
                ft_sid = ft.get("sessionId") or ft.get("id")
                if ft_sid in successfully_stopped_sids:
                    ft["status"] = "stopped"
                    ft["_scheduledStopAtMs"] = None
                    ft["_lastStatusFlipAt"] = now_ms
            fresh["tasks"] = fresh_tasks
            atomic_write_json(state_path, fresh)
        except Exception as e:
            logger.warning(f"cookbook_serve_lifecycle: state write failed: {e}")


async def cookbook_serve_lifecycle_loop() -> None:
    """Forever-loop. Registered as a startup task in app.py."""
    await asyncio.sleep(20)  # let the rest of startup settle
    while True:
        try:
            await _tick()
        except Exception as e:
            logger.warning(f"cookbook_serve_lifecycle tick failed: {e}")
        await asyncio.sleep(60)