feat(mcp): queue local OpenCode ingest jobs
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.
This commit is contained in:
Codex
@@ -0,0 +1,185 @@
|
||||
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()
|
||||
@@ -138,6 +138,54 @@ class RemoteMcpServerTests(unittest.TestCase):
|
||||
self.assertIn("list_file_assets", FakeMCP.tool_names)
|
||||
self.assertIn("get_file_asset_metadata", FakeMCP.tool_names)
|
||||
self.assertIn("how_to_use_wiki_file_assets", FakeMCP.prompt_names)
|
||||
self.assertNotIn("create_ingest_upload", FakeMCP.tool_names)
|
||||
|
||||
def test_create_server_registers_ingest_tools_when_enabled(self):
|
||||
class FakeMCP:
|
||||
tool_names = []
|
||||
|
||||
def __init__(self, _name, **_kwargs):
|
||||
pass
|
||||
|
||||
def tool(self):
|
||||
def decorator(func):
|
||||
self.tool_names.append(func.__name__)
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def resource(self, *_args, **_kwargs):
|
||||
def decorator(func):
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
def prompt(self, *_args, **_kwargs):
|
||||
def decorator(func):
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
with (
|
||||
tempfile.TemporaryDirectory() as tmp,
|
||||
patch.dict("sys.modules", {"mcp.server.fastmcp": type("Mod", (), {"FastMCP": FakeMCP})()}),
|
||||
):
|
||||
root = Path(tmp)
|
||||
server = mcp_server.create_server(
|
||||
vault_path=root / "vault",
|
||||
enable_ingest=True,
|
||||
project_root=root / "project",
|
||||
ingest_runtime_root=root / "runtime",
|
||||
)
|
||||
|
||||
self.assertIsInstance(server, FakeMCP)
|
||||
self.assertIn("create_ingest_upload", FakeMCP.tool_names)
|
||||
self.assertIn("append_ingest_upload_chunk", FakeMCP.tool_names)
|
||||
self.assertIn("complete_ingest_upload", FakeMCP.tool_names)
|
||||
self.assertIn("start_ingest_job", FakeMCP.tool_names)
|
||||
self.assertIn("get_ingest_job_status", FakeMCP.tool_names)
|
||||
self.assertIn("get_ingest_job_result", FakeMCP.tool_names)
|
||||
self.assertIn("cancel_ingest_job", FakeMCP.tool_names)
|
||||
|
||||
def test_run_http_server_fails_before_server_creation_without_token(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
|
||||
@@ -53,13 +53,21 @@ class McpServiceConfigTests(unittest.TestCase):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
env_path = Path(tmp) / ".runtime" / "obsidian-wiki-mcp.env"
|
||||
|
||||
manage_mcp_service.write_token_env_file(env_path, "test-token", host="0.0.0.0", port=8765, path="/mcp")
|
||||
manage_mcp_service.write_token_env_file(
|
||||
env_path,
|
||||
"test-token",
|
||||
host="0.0.0.0",
|
||||
port=8765,
|
||||
path="/mcp",
|
||||
enable_ingest=True,
|
||||
)
|
||||
|
||||
content = env_path.read_text(encoding="utf-8")
|
||||
self.assertIn("OBSIDIAN_MCP_TOKEN=test-token\n", content)
|
||||
self.assertIn("OBSIDIAN_MCP_HOST=0.0.0.0\n", content)
|
||||
self.assertIn("OBSIDIAN_MCP_PORT=8765\n", content)
|
||||
self.assertIn("OBSIDIAN_MCP_PATH=/mcp\n", content)
|
||||
self.assertIn("OBSIDIAN_MCP_ENABLE_INGEST=1\n", content)
|
||||
self.assertEqual(stat.S_IMODE(os.stat(env_path).st_mode), 0o600)
|
||||
self.assertEqual(stat.S_IMODE(os.stat(env_path.parent).st_mode), 0o700)
|
||||
|
||||
@@ -191,13 +199,20 @@ class McpServiceConfigTests(unittest.TestCase):
|
||||
patch.object(manage_mcp_service, "prepare_runtime") as prepare_runtime,
|
||||
patch.object(manage_mcp_service, "run_systemctl") as run_systemctl,
|
||||
):
|
||||
output = manage_mcp_service.start_service(project_root=project_root, host="0.0.0.0", port=9876, path="/wiki-mcp")
|
||||
output = manage_mcp_service.start_service(
|
||||
project_root=project_root,
|
||||
host="0.0.0.0",
|
||||
port=9876,
|
||||
path="/wiki-mcp",
|
||||
enable_ingest=True,
|
||||
)
|
||||
|
||||
env_path = project_root / ".runtime" / "obsidian-wiki-mcp.env"
|
||||
env_content = env_path.read_text(encoding="utf-8")
|
||||
self.assertIn("OBSIDIAN_MCP_TOKEN=created-token\n", env_content)
|
||||
self.assertIn("OBSIDIAN_MCP_PORT=9876\n", env_content)
|
||||
self.assertIn("OBSIDIAN_MCP_PATH=/wiki-mcp\n", env_content)
|
||||
self.assertIn("OBSIDIAN_MCP_ENABLE_INGEST=1\n", env_content)
|
||||
self.assertIn("http://10.0.0.4:9876/wiki-mcp", output)
|
||||
self.assertIn("Authorization: Bearer created-token", output)
|
||||
self.assertIn('"type": "streamable-http"', output)
|
||||
|
||||
Reference in New Issue
Block a user