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()