Codex
b174df642f
Add optional write-capable MCP ingest tools for chunked uploads, SHA-verified completion, and asynchronous local OpenCode wiki-ingest jobs. The tools remain disabled by default and are gated by OBSIDIAN_MCP_ENABLE_INGEST or mcp_service.sh start --enable-ingest. Remote clients can upload files into the vault inbox, start a restricted fixed-prompt ingest job, poll job status/result, and cancel queued or running work. The implementation keeps arbitrary prompt and shell execution out of the MCP surface. Validation: python3 -m unittest discover -v; python3 -m py_compile obsidian_wiki/mcp_ingest.py obsidian_wiki/mcp_server.py scripts/manage_mcp_service.py tests/test_mcp_ingest.py tests/test_mcp_server_remote.py tests/test_mcp_service.py; git diff --check.
186 lines
6.6 KiB
Python
186 lines
6.6 KiB
Python
import base64
|
|
import json
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from obsidian_wiki.mcp_ingest import (
|
|
IngestError,
|
|
IngestManager,
|
|
parse_bool,
|
|
)
|
|
|
|
|
|
class McpIngestTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tempdir = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.tempdir.name)
|
|
self.vault = self.root / "vault"
|
|
self.project = self.root / "project"
|
|
self.runtime = self.project / ".runtime" / "mcp-ingest"
|
|
self.vault.mkdir()
|
|
self.project.mkdir()
|
|
self.manager = IngestManager(
|
|
vault_path=self.vault,
|
|
project_root=self.project,
|
|
runtime_root=self.runtime,
|
|
opencode_path=Path("/usr/bin/opencode"),
|
|
max_file_bytes=100,
|
|
max_chunk_bytes=8,
|
|
)
|
|
|
|
def tearDown(self):
|
|
self.tempdir.cleanup()
|
|
|
|
def test_parse_bool_accepts_enabled_values_only(self):
|
|
self.assertTrue(parse_bool("1"))
|
|
self.assertTrue(parse_bool("true"))
|
|
self.assertTrue(parse_bool("YES"))
|
|
self.assertFalse(parse_bool(""))
|
|
self.assertFalse(parse_bool("0"))
|
|
self.assertFalse(parse_bool("false"))
|
|
|
|
def test_upload_chunk_lifecycle_verifies_offset_size_and_sha(self):
|
|
payload = b"hello ingest"
|
|
upload = self.manager.create_upload(
|
|
file_name="drawing.pdf",
|
|
size_bytes=len(payload),
|
|
sha256="sha256:" + self._sha256(payload),
|
|
project="bjd",
|
|
doc_type_hint="drawing",
|
|
)
|
|
|
|
first = self.manager.append_upload_chunk(
|
|
upload["upload_id"],
|
|
offset=0,
|
|
data_base64=base64.b64encode(payload[:5]).decode("ascii"),
|
|
)
|
|
self.assertEqual(first["received_bytes"], 5)
|
|
|
|
with self.assertRaises(IngestError):
|
|
self.manager.append_upload_chunk(
|
|
upload["upload_id"],
|
|
offset=3,
|
|
data_base64=base64.b64encode(payload[5:]).decode("ascii"),
|
|
)
|
|
|
|
second = self.manager.append_upload_chunk(
|
|
upload["upload_id"],
|
|
offset=5,
|
|
data_base64=base64.b64encode(payload[5:]).decode("ascii"),
|
|
)
|
|
self.assertEqual(second["received_bytes"], len(payload))
|
|
|
|
complete = self.manager.complete_upload(upload["upload_id"])
|
|
|
|
self.assertEqual(complete["state"], "complete")
|
|
self.assertEqual(Path(complete["path"]).read_bytes(), payload)
|
|
|
|
def test_upload_rejects_unsafe_names_oversized_files_and_large_chunks(self):
|
|
with self.assertRaises(IngestError):
|
|
self.manager.create_upload(file_name="../secret.pdf", size_bytes=10, sha256="a" * 64)
|
|
|
|
with self.assertRaises(IngestError):
|
|
self.manager.create_upload(file_name="large.pdf", size_bytes=101, sha256="a" * 64)
|
|
|
|
upload = self.manager.create_upload(file_name="ok.pdf", size_bytes=20, sha256=self._sha256(b"x" * 20))
|
|
with self.assertRaises(IngestError):
|
|
self.manager.append_upload_chunk(
|
|
upload["upload_id"],
|
|
offset=0,
|
|
data_base64=base64.b64encode(b"x" * 9).decode("ascii"),
|
|
)
|
|
|
|
def test_start_ingest_job_constructs_restricted_opencode_command(self):
|
|
payload = b"source data"
|
|
upload = self._complete_upload("source.pdf", payload)
|
|
popen = FakePopen()
|
|
|
|
with patch("obsidian_wiki.mcp_ingest.subprocess.Popen", return_value=popen) as popen_mock:
|
|
job = self.manager.start_ingest_job([upload["upload_id"]], note="customer asked for drawings")
|
|
|
|
args = popen_mock.call_args.args[0]
|
|
self.assertEqual(args[:5], ["/usr/bin/opencode", "run", "--format", "json", "--auto"])
|
|
self.assertIn("--dir", args)
|
|
self.assertEqual(args[args.index("--dir") + 1], str(self.project))
|
|
self.assertIn("--title", args)
|
|
self.assertIn(f"MCP ingest {job['job_id']}", args)
|
|
prompt = args[-1]
|
|
self.assertIn("/wiki-ingest", prompt)
|
|
self.assertIn(str(Path(upload["path"])), prompt)
|
|
self.assertIn("Do not ingest any file outside this list.", prompt)
|
|
self.assertIn("obsidian-wiki file-index rebuild", prompt)
|
|
self.assertNotIn("../secret", prompt)
|
|
|
|
popen.returncode = 0
|
|
self.manager.refresh_jobs()
|
|
status = self.manager.get_job_status(job["job_id"])
|
|
self.assertEqual(status["state"], "succeeded")
|
|
self.assertEqual(status["exit_code"], 0)
|
|
|
|
def test_jobs_queue_one_at_a_time_and_cancel_queued_job(self):
|
|
first = self._complete_upload("first.pdf", b"first")
|
|
second = self._complete_upload("second.pdf", b"second")
|
|
popen = FakePopen(returncode=None)
|
|
|
|
with patch("obsidian_wiki.mcp_ingest.subprocess.Popen", return_value=popen):
|
|
first_job = self.manager.start_ingest_job([first["upload_id"]])
|
|
second_job = self.manager.start_ingest_job([second["upload_id"]])
|
|
|
|
self.assertEqual(self.manager.get_job_status(first_job["job_id"])["state"], "running")
|
|
self.assertEqual(self.manager.get_job_status(second_job["job_id"])["state"], "queued")
|
|
|
|
cancelled = self.manager.cancel_ingest_job(second_job["job_id"])
|
|
|
|
self.assertEqual(cancelled["state"], "cancelled")
|
|
|
|
def _complete_upload(self, name: str, payload: bytes) -> dict:
|
|
upload = self.manager.create_upload(name, len(payload), self._sha256(payload))
|
|
offset = 0
|
|
while offset < len(payload):
|
|
chunk = payload[offset : offset + self.manager.max_chunk_bytes]
|
|
self.manager.append_upload_chunk(
|
|
upload["upload_id"],
|
|
offset=offset,
|
|
data_base64=base64.b64encode(chunk).decode("ascii"),
|
|
)
|
|
offset += len(chunk)
|
|
return self.manager.complete_upload(upload["upload_id"])
|
|
|
|
def _sha256(self, payload: bytes) -> str:
|
|
import hashlib
|
|
|
|
return hashlib.sha256(payload).hexdigest()
|
|
|
|
|
|
class FakePopen:
|
|
def __init__(self, returncode=None):
|
|
self.returncode = returncode
|
|
self.pid = 12345
|
|
self.terminated = False
|
|
|
|
def poll(self):
|
|
return self.returncode
|
|
|
|
def wait(self, timeout=None):
|
|
deadline = time.monotonic() + (timeout or 0)
|
|
while self.returncode is None and (timeout is None or time.monotonic() < deadline):
|
|
time.sleep(0.01)
|
|
if self.returncode is None:
|
|
raise subprocess.TimeoutExpired("fake", timeout)
|
|
return self.returncode
|
|
|
|
def terminate(self):
|
|
self.terminated = True
|
|
self.returncode = -15
|
|
|
|
def kill(self):
|
|
self.returncode = -9
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|