obsidian-wiki/tests/test_mcp_jobs.py

import json
import subprocess
import tempfile
import threading
import time
import unittest
from pathlib import Path

from obsidian_wiki.mcp_jobs import (
    ProcessSupervisor,
    RuntimeLeaseError,
    extract_final_answer,
    parse_marked_result,
)


class ProcessSupervisorTests(unittest.TestCase):
    def test_long_final_answer_keeps_protocol_marker_when_truncated(self):
        with tempfile.TemporaryDirectory() as tmp:
            log_path = Path(tmp) / "opencode.log"
            answer = "MCP_QUERY_OK\n" + ("x" * 20_000)
            log_path.write_text(
                '{"type":"text","part":{"text":' + json.dumps(answer) + "}}\n",
                encoding="utf-8",
            )

            extracted = extract_final_answer(log_path, limit=12_000)
            ok, public_text, error = parse_marked_result(
                extracted,
                success_marker="MCP_QUERY_OK",
                failure_marker="MCP_QUERY_FAILED",
            )

        self.assertTrue(ok, error)
        self.assertTrue(extracted.startswith("MCP_QUERY_OK\n"))
        self.assertLessEqual(len(extracted), 12_000)
        self.assertEqual(len(public_text), 12_000 - len("MCP_QUERY_OK\n"))
    def test_completion_callback_runs_without_client_polling(self):
        process = FakeProcess(returncode=0)
        completed = threading.Event()
        outcomes = []

        with tempfile.TemporaryDirectory() as tmp:
            supervisor = ProcessSupervisor(
                Path(tmp),
                popen_factory=lambda *_args, **_kwargs: process,
                poll_interval_seconds=0.01,
            )
            supervisor.start(
                "job-1",
                ["opencode", "run"],
                cwd=Path(tmp),
                log_path=Path(tmp) / "job.log",
                timeout_seconds=30,
                on_complete=lambda outcome: (outcomes.append(outcome), completed.set()),
            )

            self.assertTrue(completed.wait(1))
            self.assertEqual(outcomes[0].returncode, 0)
            self.assertFalse(outcomes[0].cancelled)
            self.assertFalse(outcomes[0].timed_out)
            supervisor.close()

    def test_completion_callback_failure_invokes_error_fallback(self):
        process = FakeProcess(returncode=0)
        recovered = threading.Event()
        failures = []

        with tempfile.TemporaryDirectory() as tmp:
            supervisor = ProcessSupervisor(
                Path(tmp),
                popen_factory=lambda *_args, **_kwargs: process,
                poll_interval_seconds=0.01,
            )

            def fail_callback(_outcome):
                raise RuntimeError("callback failed")

            supervisor.start(
                "job-1",
                ["opencode", "run"],
                cwd=Path(tmp),
                log_path=Path(tmp) / "job.log",
                timeout_seconds=30,
                on_complete=fail_callback,
                on_callback_error=lambda outcome, error: (
                    failures.append((outcome, error)),
                    recovered.set(),
                ),
            )

            self.assertTrue(recovered.wait(1))
            self.assertEqual(failures[0][0].job_id, "job-1")
            self.assertEqual(str(failures[0][1]), "callback failed")
            supervisor.close()

    def test_cancel_waits_for_exit_before_reporting_completion(self):
        process = FakeProcess(returncode=None, terminate_returncode=-15)
        outcomes = []

        with tempfile.TemporaryDirectory() as tmp:
            supervisor = ProcessSupervisor(
                Path(tmp),
                popen_factory=lambda *_args, **_kwargs: process,
                poll_interval_seconds=0.01,
            )
            supervisor.start(
                "job-1",
                ["opencode", "run"],
                cwd=Path(tmp),
                log_path=Path(tmp) / "job.log",
                timeout_seconds=30,
                on_complete=outcomes.append,
            )

            self.assertTrue(supervisor.cancel("job-1", grace_seconds=0.1))
            self.assertTrue(process.terminated)
            self.assertEqual(process.wait_calls, 1)
            self.assertEqual(len(outcomes), 1)
            self.assertTrue(outcomes[0].cancelled)
            self.assertEqual(outcomes[0].returncode, -15)
            supervisor.close()

    def test_timeout_terminates_process_and_marks_outcome(self):
        process = FakeProcess(returncode=None, terminate_returncode=-15)
        completed = threading.Event()
        outcomes = []

        with tempfile.TemporaryDirectory() as tmp:
            supervisor = ProcessSupervisor(
                Path(tmp),
                popen_factory=lambda *_args, **_kwargs: process,
                poll_interval_seconds=0.01,
            )
            supervisor.start(
                "job-1",
                ["opencode", "run"],
                cwd=Path(tmp),
                log_path=Path(tmp) / "job.log",
                timeout_seconds=0.01,
                on_complete=lambda outcome: (outcomes.append(outcome), completed.set()),
            )

            self.assertTrue(completed.wait(1))
            self.assertTrue(process.terminated)
            self.assertTrue(outcomes[0].timed_out)
            supervisor.close()

    def test_runtime_root_allows_only_one_supervisor(self):
        with tempfile.TemporaryDirectory() as tmp:
            first = ProcessSupervisor(Path(tmp), poll_interval_seconds=0.01)
            with self.assertRaises(RuntimeLeaseError):
                ProcessSupervisor(Path(tmp), poll_interval_seconds=0.01)
            first.close()

    def test_closed_supervisor_rejects_start_without_spawning(self):
        calls = []
        with tempfile.TemporaryDirectory() as tmp:
            supervisor = ProcessSupervisor(
                Path(tmp),
                popen_factory=lambda *_args, **_kwargs: calls.append(True),
                poll_interval_seconds=0.01,
            )
            supervisor.close()

            with self.assertRaises(RuntimeError):
                supervisor.start(
                    "job-1",
                    ["opencode", "run"],
                    cwd=Path(tmp),
                    log_path=Path(tmp) / "job.log",
                    timeout_seconds=30,
                    on_complete=lambda _outcome: None,
                )

        self.assertEqual(calls, [])


class FakeProcess:
    def __init__(self, *, returncode, terminate_returncode=-15):
        self.returncode = returncode
        self.terminate_returncode = terminate_returncode
        self.pid = 4242
        self.terminated = False
        self.killed = False
        self.wait_calls = 0

    def poll(self):
        return self.returncode

    def wait(self, timeout=None):
        self.wait_calls += 1
        if self.returncode is None:
            raise subprocess.TimeoutExpired("fake", timeout)
        return self.returncode

    def terminate(self):
        self.terminated = True
        self.returncode = self.terminate_returncode

    def kill(self):
        self.killed = True
        self.returncode = -9


if __name__ == "__main__":
    unittest.main()