obsidian-wiki/obsidian_wiki/mcp_jobs.py

"""Shared process supervision for OpenCode-backed MCP jobs."""

from __future__ import annotations

import json
import os
import signal
import subprocess
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable


class RuntimeLeaseError(RuntimeError):
    """Raised when another manager already owns a runtime directory."""


@dataclass(frozen=True)
class ProcessOutcome:
    job_id: str
    returncode: int
    cancelled: bool = False
    timed_out: bool = False
    error: str = ""


@dataclass
class _ProcessEntry:
    job_id: str
    process: subprocess.Popen[Any]
    deadline: float | None
    on_complete: Callable[[ProcessOutcome], None]
    on_callback_error: Callable[[ProcessOutcome, Exception], None] | None = None
    completing: bool = False


class _RuntimeLease:
    def __init__(self, runtime_root: Path) -> None:
        runtime_root.mkdir(parents=True, exist_ok=True)
        self.path = runtime_root / ".manager.lock"
        self.handle = self.path.open("a+b")
        try:
            self._lock()
        except OSError as exc:
            self.handle.close()
            raise RuntimeLeaseError(f"MCP runtime is already in use: {runtime_root}") from exc

    def _lock(self) -> None:
        if os.name == "nt":  # pragma: no cover - exercised on Windows
            import msvcrt

            self.handle.seek(0)
            if self.handle.read(1) == b"":
                self.handle.write(b"0")
                self.handle.flush()
            self.handle.seek(0)
            msvcrt.locking(self.handle.fileno(), msvcrt.LK_NBLCK, 1)
            return
        import fcntl

        fcntl.flock(self.handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)

    def close(self) -> None:
        if self.handle.closed:
            return
        try:
            if os.name == "nt":  # pragma: no cover - exercised on Windows
                import msvcrt

                self.handle.seek(0)
                msvcrt.locking(self.handle.fileno(), msvcrt.LK_UNLCK, 1)
            else:
                import fcntl

                fcntl.flock(self.handle.fileno(), fcntl.LOCK_UN)
        finally:
            self.handle.close()


class ProcessSupervisor:
    """Own subprocesses, enforce deadlines, and deliver one completion callback."""

    def __init__(
        self,
        runtime_root: str | Path,
        *,
        popen_factory: Callable[..., subprocess.Popen[Any]] = subprocess.Popen,
        poll_interval_seconds: float = 0.1,
    ) -> None:
        self.runtime_root = Path(runtime_root).expanduser().resolve()
        self._lease = _RuntimeLease(self.runtime_root)
        self._popen_factory = popen_factory
        self._poll_interval = max(0.01, float(poll_interval_seconds))
        self._entries: dict[str, _ProcessEntry] = {}
        self._lock = threading.RLock()
        self._condition = threading.Condition(self._lock)
        self._callbacks_in_progress = 0
        self._wake = threading.Event()
        self._closed = False
        self._thread = threading.Thread(target=self._monitor, name="mcp-job-supervisor", daemon=True)
        self._thread.start()

    def start(
        self,
        job_id: str,
        args: list[str],
        *,
        cwd: str | Path,
        log_path: str | Path,
        timeout_seconds: float,
        on_complete: Callable[[ProcessOutcome], None],
        on_callback_error: Callable[[ProcessOutcome, Exception], None] | None = None,
    ) -> subprocess.Popen[Any]:
        with self._lock:
            if self._closed:
                raise RuntimeError("process supervisor is closed")
            if job_id in self._entries:
                raise ValueError(f"job is already supervised: {job_id}")
            log = Path(log_path)
            log.parent.mkdir(parents=True, exist_ok=True)
            log_file = log.open("ab")
            kwargs: dict[str, Any] = {
                "cwd": str(cwd),
                "stdout": log_file,
                "stderr": subprocess.STDOUT,
            }
            if os.name == "nt":  # pragma: no cover - exercised on Windows
                kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
            else:
                kwargs["start_new_session"] = True
            try:
                process = self._popen_factory(args, **kwargs)
            finally:
                log_file.close()
            deadline = time.monotonic() + float(timeout_seconds) if timeout_seconds > 0 else None
            self._entries[job_id] = _ProcessEntry(
                job_id,
                process,
                deadline,
                on_complete,
                on_callback_error,
            )
        self._wake.set()
        return process

    def is_running(self, job_id: str | None = None) -> bool:
        with self._lock:
            if job_id is None:
                return bool(self._entries)
            return job_id in self._entries

    def refresh(self) -> None:
        for entry in self._snapshot():
            code = entry.process.poll()
            if code is not None:
                self._complete_naturally(entry.job_id, code)
            elif entry.deadline is not None and time.monotonic() >= entry.deadline:
                self._complete_stopped(entry.job_id, timed_out=True, grace_seconds=1.0)
        deadline = time.monotonic() + 1.0
        with self._condition:
            while self._callbacks_in_progress and time.monotonic() < deadline:
                self._condition.wait(deadline - time.monotonic())

    def cancel(self, job_id: str, *, grace_seconds: float = 5.0) -> bool:
        return self._complete_stopped(job_id, cancelled=True, grace_seconds=grace_seconds)

    def close(self) -> None:
        with self._lock:
            if self._closed:
                return
            self._closed = True
            job_ids = list(self._entries)
        for job_id in job_ids:
            self._complete_stopped(job_id, cancelled=True, grace_seconds=0.2)
        self._wake.set()
        if threading.current_thread() is not self._thread:
            self._thread.join(timeout=1)
        self._lease.close()

    def _snapshot(self) -> list[_ProcessEntry]:
        with self._lock:
            return [entry for entry in self._entries.values() if not entry.completing]

    def _begin_completion(self, job_id: str) -> _ProcessEntry | None:
        with self._lock:
            entry = self._entries.get(job_id)
            if entry is None or entry.completing:
                return None
            entry.completing = True
            self._callbacks_in_progress += 1
            return entry

    def _complete_naturally(self, job_id: str, returncode: int) -> bool:
        entry = self._begin_completion(job_id)
        if entry is None:
            return False
        self._finish(entry, ProcessOutcome(job_id=job_id, returncode=int(returncode)))
        return True

    def _complete_stopped(
        self,
        job_id: str,
        *,
        cancelled: bool = False,
        timed_out: bool = False,
        grace_seconds: float,
    ) -> bool:
        entry = self._begin_completion(job_id)
        if entry is None:
            return False
        returncode, error = self._terminate_process(entry.process, grace_seconds=grace_seconds)
        self._finish(
            entry,
            ProcessOutcome(
                job_id=job_id,
                returncode=returncode,
                cancelled=cancelled,
                timed_out=timed_out,
                error=error,
            ),
        )
        return True

    def _terminate_process(self, process: subprocess.Popen[Any], *, grace_seconds: float) -> tuple[int, str]:
        error = ""
        if process.poll() is None:
            try:
                if os.name != "nt":
                    os.killpg(process.pid, signal.SIGTERM)
                else:  # pragma: no cover - exercised on Windows
                    process.terminate()
            except (OSError, ProcessLookupError):
                try:
                    process.terminate()
                except OSError as exc:
                    error = str(exc)
            try:
                process.wait(timeout=max(0, grace_seconds))
            except subprocess.TimeoutExpired:
                try:
                    if os.name != "nt":
                        os.killpg(process.pid, signal.SIGKILL)
                    else:  # pragma: no cover - exercised on Windows
                        process.kill()
                except (OSError, ProcessLookupError):
                    try:
                        process.kill()
                    except OSError as exc:
                        error = error or str(exc)
                try:
                    process.wait(timeout=1)
                except subprocess.TimeoutExpired:
                    error = error or "process did not exit after SIGKILL"
        code = process.poll()
        return (int(code) if code is not None else -9), error

    def _finish(self, entry: _ProcessEntry, outcome: ProcessOutcome) -> None:
        with self._lock:
            current = self._entries.get(entry.job_id)
            if current is not entry:
                return
            self._entries.pop(entry.job_id, None)
        try:
            entry.on_complete(outcome)
        except Exception as exc:
            if entry.on_callback_error is not None:
                try:
                    entry.on_callback_error(outcome, exc)
                except Exception:
                    pass
        finally:
            with self._condition:
                self._callbacks_in_progress -= 1
                self._condition.notify_all()
        self._wake.set()

    def _monitor(self) -> None:
        while True:
            with self._lock:
                if self._closed and not self._entries:
                    return
            self.refresh()
            self._wake.wait(self._poll_interval)
            self._wake.clear()


def extract_final_answer(path: str | Path, *, limit: int = 12000) -> str:
    """Extract only assistant text events, preferring the final-answer phase."""
    log_path = Path(path)
    if not log_path.is_file():
        return ""
    last_text = ""
    final_text = ""
    with log_path.open("r", encoding="utf-8", errors="replace") as handle:
        for raw in handle:
            try:
                event = json.loads(raw)
            except json.JSONDecodeError:
                continue
            if event.get("type") != "text":
                continue
            part = event.get("part") or {}
            text = str(part.get("text") or "")
            if not text:
                continue
            last_text = text
            phase = ((part.get("metadata") or {}).get("openai") or {}).get("phase")
            if phase == "final_answer":
                final_text = text
    text = final_text or last_text
    if limit <= 0 or len(text) <= limit:
        return text if limit > 0 else ""
    first_line, separator, remainder = text.partition("\n")
    if not separator or len(first_line) + 1 >= limit:
        return text[:limit]
    tail_size = limit - len(first_line) - 1
    return first_line + "\n" + remainder[-tail_size:]


def parse_marked_result(text: str, *, success_marker: str, failure_marker: str) -> tuple[bool, str, str]:
    """Return protocol success, public text, and an error string."""
    value = str(text or "").strip()
    first, _, remainder = value.partition("\n")
    if first.strip() == success_marker:
        return True, remainder.strip(), ""
    if first.strip() == failure_marker:
        error = remainder.strip() or f"OpenCode returned {failure_marker}"
        return False, "", error
    return False, "", "OpenCode final answer did not include the required result marker"