src/agent_harness/lm_studio_management.py

"""LM Studio server lifecycle and exclusive model-residency management.

The `lms` CLI is the sole server lifecycle controller. Model inspection, loading,
and unloading use LM Studio's native v1 REST API so every applied configuration is
captured in experiment telemetry.
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
import json
import os
from pathlib import Path
import shutil
import subprocess
import time
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


class LMStudioManagementError(RuntimeError):
    """Raised when server lifecycle or exclusive residency cannot be verified."""


@dataclass(frozen=True, slots=True)
class ResidencyTransition:
    target_model: str | None
    target_context_length: int | None
    before_instances: tuple[str, ...]
    unloaded_instances: tuple[str, ...]
    loaded_instance: str | None
    after_instances: tuple[str, ...]
    load_response: dict[str, Any] | None
    elapsed_seconds: float
    reused: bool

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


class LMStudioServer:
    """Control only the local LM Studio API server via `lms server`."""

    def __init__(
        self,
        port: int = 1234,
        cli_path: Path | None = None,
        command_timeout_seconds: float = 30.0,
    ):
        discovered = shutil.which("lms")
        default = Path.home() / ".lmstudio" / "bin" / "lms"
        self.cli_path = Path(cli_path or discovered or default).resolve()
        self.port = port
        self.command_timeout_seconds = command_timeout_seconds

    def _run(self, *arguments: str) -> subprocess.CompletedProcess[str]:
        if not self.cli_path.exists():
            raise LMStudioManagementError(f"LM Studio CLI is unavailable: {self.cli_path}")
        try:
            return subprocess.run(
                [str(self.cli_path), *arguments],
                text=True,
                capture_output=True,
                check=False,
                timeout=self.command_timeout_seconds,
            )
        except (OSError, subprocess.SubprocessError) as exc:
            raise LMStudioManagementError(f"lms {' '.join(arguments)} failed: {exc}") from exc

    def status(self) -> dict[str, Any]:
        result = self._run("server", "status")
        # Current lms releases render human-facing status on stderr even when
        # the command succeeds. Treat both captured streams as the CLI record.
        rendered = result.stdout + "\n" + result.stderr
        normalized = rendered.lower()
        return {
            "running": (
                result.returncode == 0
                and "server is running" in normalized
                and "server is not running" not in normalized
            ),
            "returncode": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "command": [str(self.cli_path), "server", "status"],
        }

    def _api_ready(self) -> bool:
        request = Request(
            f"http://127.0.0.1:{self.port}/api/v1/models",
            method="GET",
            headers={"Content-Type": "application/json"},
        )
        try:
            with urlopen(request, timeout=min(self.command_timeout_seconds, 2.0)):
                return True
        except HTTPError as exc:
            # Authentication failures still prove that the REST listener is ready.
            return exc.code in {401, 403}
        except (URLError, OSError):
            return False

    def _wait_until_ready(self) -> dict[str, Any]:
        deadline = time.monotonic() + self.command_timeout_seconds
        last_status: dict[str, Any] | None = None
        while time.monotonic() < deadline:
            last_status = self.status()
            if last_status["running"] and self._api_ready():
                return last_status
            time.sleep(0.25)
        raise LMStudioManagementError(
            "lms reported startup activity, but the official REST API never became ready: "
            + repr(last_status)
        )

    def ensure_running(self) -> dict[str, Any]:
        status = self.status()
        if status["running"]:
            return {
                "action": "already_running",
                "status": self._wait_until_ready(),
            }
        started = self._run("server", "start", "--port", str(self.port))
        if started.returncode != 0:
            raise LMStudioManagementError(
                "lms server start failed: " + (started.stderr.strip() or started.stdout.strip())
            )
        return {
            "action": "started",
            "start_stdout": started.stdout,
            "start_stderr": started.stderr,
            "status": self._wait_until_ready(),
        }

    def stop(self) -> dict[str, Any]:
        result = self._run("server", "stop")
        if result.returncode != 0:
            raise LMStudioManagementError(
                "lms server stop failed: " + (result.stderr.strip() or result.stdout.strip())
            )
        return {"stdout": result.stdout, "stderr": result.stderr, "returncode": result.returncode}


class LMStudioResidencyManager:
    """Use native v1 REST endpoints to enforce exactly one or zero loaded models."""

    def __init__(self, base_url: str, api_token_env: str, timeout_seconds: float = 120.0):
        self.base_url = base_url.rstrip("/")
        self.api_token_env = api_token_env
        self.timeout_seconds = timeout_seconds

    def _headers(self) -> dict[str, str]:
        headers = {"Content-Type": "application/json"}
        token = os.environ.get(self.api_token_env, "").strip()
        if token:
            headers["Authorization"] = f"Bearer {token}"
        return headers

    def _request(
        self, method: str, endpoint: str, payload: dict[str, Any] | None = None
    ) -> dict[str, Any]:
        request = Request(
            self.base_url + endpoint,
            data=None if payload is None else json.dumps(payload).encode("utf-8"),
            method=method,
            headers=self._headers(),
        )
        try:
            with urlopen(request, timeout=self.timeout_seconds) as response:
                body = response.read().decode("utf-8")
        except HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise LMStudioManagementError(
                f"LM Studio returned HTTP {exc.code} for {endpoint}: {detail}"
            ) from exc
        except URLError as exc:
            raise LMStudioManagementError(
                f"Cannot reach LM Studio management API at {self.base_url}: {exc.reason}"
            ) from exc
        try:
            value = json.loads(body)
        except json.JSONDecodeError as exc:
            raise LMStudioManagementError(f"LM Studio returned non-JSON data for {endpoint}") from exc
        if not isinstance(value, dict):
            raise LMStudioManagementError(f"Unexpected LM Studio response for {endpoint}")
        if isinstance(value.get("error"), dict):
            raise LMStudioManagementError(f"LM Studio management error for {endpoint}: {value['error']}")
        return value

    def models(self) -> tuple[dict[str, Any], ...]:
        value = self._request("GET", "/api/v1/models")
        models = value.get("models")
        if not isinstance(models, list):
            raise LMStudioManagementError("/api/v1/models response has no models array")
        return tuple(item for item in models if isinstance(item, dict))

    def loaded_instances(self) -> tuple[dict[str, Any], ...]:
        result: list[dict[str, Any]] = []
        for model in self.models():
            instances = model.get("loaded_instances", [])
            if not isinstance(instances, list):
                continue
            for instance in instances:
                if isinstance(instance, dict) and instance.get("id"):
                    result.append(
                        {
                            "model_key": str(model.get("key", "")),
                            "type": str(model.get("type", "")),
                            "instance_id": str(instance["id"]),
                            "config": dict(instance.get("config", {})),
                        }
                    )
        return tuple(result)

    def _unload(self, instance_id: str) -> dict[str, Any]:
        value = self._request(
            "POST", "/api/v1/models/unload", {"instance_id": instance_id}
        )
        if value.get("instance_id") != instance_id:
            raise LMStudioManagementError(
                f"unload acknowledgement mismatch for {instance_id}: {value}"
            )
        return value

    def unload_all(self) -> ResidencyTransition:
        started = time.monotonic()
        before = self.loaded_instances()
        unloaded: list[str] = []
        for instance in before:
            self._unload(instance["instance_id"])
            unloaded.append(instance["instance_id"])
        after = self.loaded_instances()
        if after:
            raise LMStudioManagementError(f"models remained loaded after unload-all: {after}")
        return ResidencyTransition(
            target_model=None,
            target_context_length=None,
            before_instances=tuple(item["instance_id"] for item in before),
            unloaded_instances=tuple(unloaded),
            loaded_instance=None,
            after_instances=(),
            load_response=None,
            elapsed_seconds=time.monotonic() - started,
            reused=not before,
        )

    def ensure_exclusive(self, model_key: str, context_length: int) -> ResidencyTransition:
        started = time.monotonic()
        before = self.loaded_instances()
        if len(before) == 1:
            current = before[0]
            current_context = current.get("config", {}).get("context_length")
            if current["model_key"] == model_key and current_context == context_length:
                return ResidencyTransition(
                    target_model=model_key,
                    target_context_length=context_length,
                    before_instances=(current["instance_id"],),
                    unloaded_instances=(),
                    loaded_instance=current["instance_id"],
                    after_instances=(current["instance_id"],),
                    load_response=None,
                    elapsed_seconds=time.monotonic() - started,
                    reused=True,
                )

        unloaded: list[str] = []
        for instance in before:
            self._unload(instance["instance_id"])
            unloaded.append(instance["instance_id"])
        response = self._request(
            "POST",
            "/api/v1/models/load",
            {
                "model": model_key,
                "context_length": context_length,
                "echo_load_config": True,
            },
        )
        instance_id = response.get("instance_id")
        if response.get("status") != "loaded" or not isinstance(instance_id, str):
            raise LMStudioManagementError(f"model load did not succeed for {model_key}: {response}")
        load_config = response.get("load_config", {})
        if not isinstance(load_config, dict) or load_config.get("context_length") != context_length:
            raise LMStudioManagementError(
                f"LM Studio did not apply context_length={context_length}: {response}"
            )
        after = self.loaded_instances()
        if len(after) != 1:
            raise LMStudioManagementError(f"exclusive residency failed for {model_key}: {after}")
        only = after[0]
        if (
            only["model_key"] != model_key
            or only["instance_id"] != instance_id
            or only.get("config", {}).get("context_length") != context_length
        ):
            raise LMStudioManagementError(
                f"loaded runtime does not match requested model/context: {after}"
            )
        return ResidencyTransition(
            target_model=model_key,
            target_context_length=context_length,
            before_instances=tuple(item["instance_id"] for item in before),
            unloaded_instances=tuple(unloaded),
            loaded_instance=instance_id,
            after_instances=(instance_id,),
            load_response=response,
            elapsed_seconds=time.monotonic() - started,
            reused=False,
        )