feat(mcp): harden remote assets and job execution
Gate source-file downloads and verify registry paths and content hashes. Stream multipart uploads with bounded limits, supervise OpenCode jobs with explicit result markers and timeouts, and preserve managed endpoint security settings. Document the trusted-client boundary and add focused regression coverage. Validation: PYTHONPATH=<staged-tree> python -m unittest discover -s tests (182 tests)
This commit is contained in:
@@ -480,7 +480,7 @@ def search_assets(
|
||||
query: str,
|
||||
*,
|
||||
db_path: Path | None = None,
|
||||
limit: int = 10,
|
||||
limit: int | None = 10,
|
||||
doc_type: str | None = None,
|
||||
source_type: str | None = None,
|
||||
product_line: str | None = None,
|
||||
@@ -527,7 +527,7 @@ def search_assets(
|
||||
rows.append(record)
|
||||
|
||||
rows.sort(key=lambda item: (item["score"], item.get("updated_at") or ""), reverse=True)
|
||||
return rows[:limit]
|
||||
return rows if limit is None else rows[:limit]
|
||||
|
||||
|
||||
def get_asset_by_file_id(
|
||||
|
||||
+36
-10
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
@@ -137,11 +138,12 @@ def search_file_assets(
|
||||
vault = _vault_path(vault_path)
|
||||
registry = load_registry_by_file_id(vault)
|
||||
profile = _load_profile(vault, profile_name)
|
||||
public_limit = min(max(limit, 1), 100)
|
||||
rows = search_assets(
|
||||
vault,
|
||||
query,
|
||||
db_path=db_path,
|
||||
limit=max(limit * 3, limit),
|
||||
limit=None,
|
||||
doc_type=doc_type,
|
||||
source_type=source_type,
|
||||
product_line=product_line,
|
||||
@@ -160,7 +162,7 @@ def search_file_assets(
|
||||
if not ok:
|
||||
continue
|
||||
allowed.append(public_asset_record({**record, **row}, expose_paths=expose_paths))
|
||||
if len(allowed) >= limit:
|
||||
if len(allowed) >= public_limit:
|
||||
break
|
||||
return allowed
|
||||
|
||||
@@ -197,8 +199,11 @@ def read_file_asset_resource(
|
||||
db_path: Path | None = None,
|
||||
profile_name: str = DEFAULT_PROFILE_NAME,
|
||||
max_file_bytes: int | None = None,
|
||||
allow_source_files: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
file_id = file_id_from_asset_uri(uri)
|
||||
if not allow_source_files:
|
||||
raise AssetAccessError("source file reads are disabled")
|
||||
vault = _vault_path(vault_path)
|
||||
row = get_asset_by_file_id(vault, file_id, db_path=db_path, read_only=True)
|
||||
if row is None:
|
||||
@@ -215,13 +220,27 @@ def read_file_asset_resource(
|
||||
if not ok:
|
||||
raise AssetAccessError(f"file asset blocked: {reason}")
|
||||
|
||||
path = safe_asset_path(vault, str(row.get("archive_path") or ""))
|
||||
registry_archive_path = str(registry_record.get("archive_path") or "")
|
||||
db_archive_path = str(row.get("archive_path") or "")
|
||||
path = safe_asset_path(vault, registry_archive_path)
|
||||
db_asset_path = safe_asset_path(vault, db_archive_path)
|
||||
if path != db_asset_path:
|
||||
raise AssetAccessError(
|
||||
"file asset archive_path mismatch between source registry and index"
|
||||
)
|
||||
if not path.is_file():
|
||||
raise AssetNotFoundError(f"archived file missing: {row.get('archive_path')}")
|
||||
raise AssetNotFoundError(f"archived file missing: {registry_archive_path}")
|
||||
limit = DEFAULT_MAX_FILE_BYTES if max_file_bytes is None else max_file_bytes
|
||||
size = path.stat().st_size
|
||||
if limit > 0 and size > limit:
|
||||
raise AssetTooLargeError(f"file exceeds max size: {size} > {limit}")
|
||||
content = path.read_bytes()
|
||||
size = len(content)
|
||||
if limit > 0 and size > limit:
|
||||
raise AssetTooLargeError(f"file exceeds max size: {size} > {limit}")
|
||||
digest = hashlib.sha256(content).hexdigest()
|
||||
if digest != file_id.removeprefix("sha256:"):
|
||||
raise AssetAccessError(f"file asset content hash mismatch: {file_id}")
|
||||
name = str(row.get("original_name") or path.name)
|
||||
return {
|
||||
"file_id": file_id,
|
||||
@@ -229,7 +248,7 @@ def read_file_asset_resource(
|
||||
"name": name,
|
||||
"mime_type": str(row.get("mime_type") or mimetypes.guess_type(name)[0] or "application/octet-stream"),
|
||||
"size_bytes": size,
|
||||
"bytes": path.read_bytes(),
|
||||
"bytes": content,
|
||||
}
|
||||
|
||||
|
||||
@@ -283,7 +302,7 @@ def record_allowed_for_mcp(
|
||||
return False, f"business category {category} blocked"
|
||||
|
||||
if profile is None:
|
||||
return True, "allowed"
|
||||
return False, "bot profile not found"
|
||||
decision = source_decision(profile_name, _source_file_profile(profile, record), record)
|
||||
if not decision.include:
|
||||
return False, decision.reason
|
||||
@@ -323,6 +342,7 @@ def mcp_runtime_options(
|
||||
profile_name: str | None = None,
|
||||
max_file_bytes: int | None = None,
|
||||
expose_paths: bool | None = None,
|
||||
allow_source_files: bool | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"profile_name": profile_name or config_value("OBSIDIAN_MCP_BOT_PROFILE", DEFAULT_PROFILE_NAME),
|
||||
@@ -330,14 +350,20 @@ def mcp_runtime_options(
|
||||
if max_file_bytes is None
|
||||
else max_file_bytes,
|
||||
"expose_paths": _env_flag("OBSIDIAN_MCP_EXPOSE_PATHS", False) if expose_paths is None else expose_paths,
|
||||
"allow_source_files": _env_flag("OBSIDIAN_MCP_ALLOW_SOURCE_FILES", False)
|
||||
if allow_source_files is None
|
||||
else allow_source_files,
|
||||
}
|
||||
|
||||
|
||||
def _load_profile(vault: Path, profile_name: str) -> dict[str, Any] | None:
|
||||
def _load_profile(vault: Path, profile_name: str) -> dict[str, Any]:
|
||||
try:
|
||||
return load_profiles(vault).get(profile_name)
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
profile = load_profiles(vault).get(profile_name)
|
||||
except FileNotFoundError as exc:
|
||||
raise AssetAccessError(f"bot profile not found: {profile_name}") from exc
|
||||
if not isinstance(profile, dict):
|
||||
raise AssetAccessError(f"bot profile not found: {profile_name}")
|
||||
return profile
|
||||
|
||||
|
||||
def _source_file_profile(profile: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
|
||||
|
||||
+292
-82
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
@@ -20,21 +21,27 @@ from threading import RLock
|
||||
from typing import Any
|
||||
|
||||
from obsidian_wiki.file_assets import clean_config_value, db_path_for_vault, registry_path_for_vault
|
||||
from obsidian_wiki.mcp_jobs import ProcessOutcome, ProcessSupervisor, extract_final_answer, parse_marked_result
|
||||
|
||||
DEFAULT_INGEST_MAX_FILE_BYTES = 100 * 1024 * 1024
|
||||
DEFAULT_INGEST_MAX_CHUNK_BYTES = 4 * 1024 * 1024
|
||||
DEFAULT_INGEST_MAX_FILES_PER_JOB = 10
|
||||
DEFAULT_INGEST_MAX_REQUEST_BYTES = 0
|
||||
DEFAULT_UPLOAD_TICKET_SECONDS = 30 * 60
|
||||
MAX_UPLOAD_TICKET_SECONDS = 24 * 60 * 60
|
||||
DEFAULT_OPENCODE_PATH = "opencode"
|
||||
DEFAULT_INGEST_OPENCODE_MODEL = "openai/gpt-5.5"
|
||||
DEFAULT_INGEST_OPENCODE_VARIANT = "high"
|
||||
DEFAULT_OPENCODE_SMOKE_TIMEOUT_SECONDS = 120
|
||||
DEFAULT_INGEST_JOB_TIMEOUT_SECONDS = 30 * 60
|
||||
STREAM_COPY_CHUNK_BYTES = 1024 * 1024
|
||||
|
||||
SHA256_RE = re.compile(r"^(?:sha256:)?([0-9a-fA-F]{64})$")
|
||||
OPENCODE_REQUIRED_RUN_FLAGS = ("--format", "--auto", "--dir", "--title")
|
||||
OPENCODE_SMOKE_SUCCESS = "WIKI_INGEST_SKILL_LOADED"
|
||||
OPENCODE_SMOKE_FAILURE = "SKILL_LOAD_FAILED"
|
||||
INGEST_SUCCESS_MARKER = "MCP_INGEST_OK"
|
||||
INGEST_FAILURE_MARKER = "MCP_INGEST_FAILED"
|
||||
TERMINAL_JOB_STATES = {"succeeded", "failed", "cancelled"}
|
||||
|
||||
|
||||
@@ -73,6 +80,8 @@ class IngestManager:
|
||||
max_file_bytes: int = DEFAULT_INGEST_MAX_FILE_BYTES,
|
||||
max_chunk_bytes: int = DEFAULT_INGEST_MAX_CHUNK_BYTES,
|
||||
max_files_per_job: int = DEFAULT_INGEST_MAX_FILES_PER_JOB,
|
||||
max_request_bytes: int = DEFAULT_INGEST_MAX_REQUEST_BYTES,
|
||||
job_timeout_seconds: int = DEFAULT_INGEST_JOB_TIMEOUT_SECONDS,
|
||||
) -> None:
|
||||
self.vault = Path(clean_config_value(str(vault_path))).expanduser().resolve()
|
||||
self.project_root = Path(clean_config_value(str(project_root))).expanduser().resolve()
|
||||
@@ -88,14 +97,20 @@ class IngestManager:
|
||||
self.max_file_bytes = max_file_bytes
|
||||
self.max_chunk_bytes = max_chunk_bytes
|
||||
self.max_files_per_job = max_files_per_job
|
||||
self.max_request_bytes = max(0, int(max_request_bytes))
|
||||
self.job_timeout_seconds = max(1, int(job_timeout_seconds))
|
||||
self.upload_root = self.runtime_root / "uploads"
|
||||
self.jobs_root = self.runtime_root / "jobs"
|
||||
self.upload_ticket_root = self.runtime_root / "upload-tickets"
|
||||
self._lock = RLock()
|
||||
self._processes: dict[str, subprocess.Popen[Any]] = {}
|
||||
self._closing = False
|
||||
self.upload_root.mkdir(parents=True, exist_ok=True)
|
||||
self.jobs_root.mkdir(parents=True, exist_ok=True)
|
||||
self.upload_ticket_root.mkdir(parents=True, exist_ok=True)
|
||||
self._supervisor = ProcessSupervisor(
|
||||
self.runtime_root,
|
||||
popen_factory=lambda *args, **kwargs: subprocess.Popen(*args, **kwargs),
|
||||
)
|
||||
self._mark_interrupted_jobs()
|
||||
|
||||
def create_upload(
|
||||
@@ -143,23 +158,80 @@ class IngestManager:
|
||||
doc_type_hint: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
data = bytes(payload)
|
||||
digest = hashlib.sha256(data).hexdigest()
|
||||
return self.create_completed_upload_from_stream(
|
||||
file_name=file_name,
|
||||
stream=io.BytesIO(data),
|
||||
size_bytes=len(data),
|
||||
project=project,
|
||||
doc_type_hint=doc_type_hint,
|
||||
)
|
||||
|
||||
def create_completed_upload_from_stream(
|
||||
self,
|
||||
file_name: str,
|
||||
stream: Any,
|
||||
*,
|
||||
size_bytes: int | None = None,
|
||||
project: str | None = None,
|
||||
doc_type_hint: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
name = self._validate_file_name(file_name)
|
||||
declared_size = None if size_bytes is None else int(size_bytes)
|
||||
if declared_size is not None and declared_size < 0:
|
||||
raise IngestError("size_bytes must be non-negative")
|
||||
if declared_size is not None and self.max_file_bytes > 0 and declared_size > self.max_file_bytes:
|
||||
raise IngestError(f"file exceeds max upload size: {declared_size} > {self.max_file_bytes}")
|
||||
|
||||
with self._lock:
|
||||
upload = self.create_upload(
|
||||
file_name=file_name,
|
||||
size_bytes=len(data),
|
||||
sha256=digest,
|
||||
project=project,
|
||||
doc_type_hint=doc_type_hint,
|
||||
)
|
||||
meta = self._load_upload(upload["upload_id"])
|
||||
path = Path(meta["path"])
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_bytes(data)
|
||||
meta["received_bytes"] = len(data)
|
||||
meta["updated_at"] = utc_now()
|
||||
self._write_json(self._upload_meta_path(path.parent), meta)
|
||||
return self.complete_upload(upload["upload_id"])
|
||||
upload_id = uuid.uuid4().hex
|
||||
month = datetime.now(timezone.utc).strftime("%Y-%m")
|
||||
upload_dir = self.upload_root / month / upload_id
|
||||
upload_dir.mkdir(parents=True, exist_ok=False)
|
||||
path = upload_dir / name
|
||||
received = 0
|
||||
hasher = hashlib.sha256()
|
||||
try:
|
||||
with path.open("xb") as destination:
|
||||
while True:
|
||||
chunk = stream.read(STREAM_COPY_CHUNK_BYTES)
|
||||
if not chunk:
|
||||
break
|
||||
if not isinstance(chunk, (bytes, bytearray, memoryview)):
|
||||
raise IngestError("upload stream must return bytes")
|
||||
data = bytes(chunk)
|
||||
received += len(data)
|
||||
if self.max_file_bytes > 0 and received > self.max_file_bytes:
|
||||
raise IngestError(
|
||||
f"file exceeds max upload size: {received} > {self.max_file_bytes}"
|
||||
)
|
||||
if declared_size is not None and received > declared_size:
|
||||
raise IngestError(
|
||||
f"upload size mismatch: received more than declared {declared_size} bytes"
|
||||
)
|
||||
destination.write(data)
|
||||
hasher.update(data)
|
||||
if declared_size is not None and received != declared_size:
|
||||
raise IngestError(f"upload size mismatch: {received} != {declared_size}")
|
||||
now = utc_now()
|
||||
meta = {
|
||||
"upload_id": upload_id,
|
||||
"state": "complete",
|
||||
"file_name": name,
|
||||
"size_bytes": received,
|
||||
"received_bytes": received,
|
||||
"sha256": hasher.hexdigest(),
|
||||
"project": self._safe_optional_text(project),
|
||||
"doc_type_hint": self._safe_optional_text(doc_type_hint),
|
||||
"path": str(path),
|
||||
"created_at": now,
|
||||
"completed_at": now,
|
||||
"updated_at": now,
|
||||
}
|
||||
self._write_json(self._upload_meta_path(upload_dir), meta)
|
||||
return self._public_upload(meta)
|
||||
except Exception:
|
||||
shutil.rmtree(upload_dir, ignore_errors=True)
|
||||
raise
|
||||
|
||||
def create_upload_ticket(
|
||||
self,
|
||||
@@ -312,16 +384,53 @@ class IngestManager:
|
||||
self._write_json(self._upload_meta_path(path.parent), meta)
|
||||
return self._public_upload(meta)
|
||||
|
||||
def discard_upload(self, upload_id: str) -> bool:
|
||||
with self._lock:
|
||||
if any(upload_id in job.get("upload_ids", []) for job in self._iter_jobs()):
|
||||
return False
|
||||
meta = self._load_upload(upload_id)
|
||||
upload_dir = Path(meta["path"]).parent.resolve()
|
||||
if not path_is_relative_to(upload_dir, self.upload_root.resolve()):
|
||||
raise IngestError(f"upload path escapes runtime root: {upload_id}")
|
||||
shutil.rmtree(upload_dir)
|
||||
return True
|
||||
|
||||
def start_ingest_job(self, upload_ids: list[str], note: str | None = None) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
if self._closing:
|
||||
raise IngestError("ingest manager is closed")
|
||||
if not upload_ids:
|
||||
raise IngestError("upload_ids must not be empty")
|
||||
if len(set(upload_ids)) != len(upload_ids):
|
||||
raise IngestError("upload_ids must be unique")
|
||||
if len(upload_ids) > self.max_files_per_job:
|
||||
raise IngestError(f"too many files for one ingest job: {len(upload_ids)} > {self.max_files_per_job}")
|
||||
uploads = [self._load_upload(upload_id) for upload_id in upload_ids]
|
||||
for upload in uploads:
|
||||
if upload["state"] != "complete":
|
||||
raise IngestError(f"upload is not complete: {upload['upload_id']}")
|
||||
requested_upload_ids = tuple(upload_ids)
|
||||
claimed_jobs = [
|
||||
job
|
||||
for job in self._iter_jobs()
|
||||
if job.get("state") in {"queued", "running", "succeeded"}
|
||||
and set(upload_ids).intersection(job.get("upload_ids", []))
|
||||
]
|
||||
for claimed in claimed_jobs:
|
||||
if tuple(claimed.get("upload_ids", [])) == requested_upload_ids:
|
||||
return self._public_job(claimed)
|
||||
if claimed_jobs:
|
||||
claimed_ids = sorted(
|
||||
{
|
||||
upload_id
|
||||
for job in claimed_jobs
|
||||
for upload_id in job.get("upload_ids", [])
|
||||
if upload_id in upload_ids
|
||||
}
|
||||
)
|
||||
raise IngestError(
|
||||
"uploads are already assigned to another ingest job: " + ", ".join(claimed_ids)
|
||||
)
|
||||
job_id = uuid.uuid4().hex
|
||||
job_dir = self.jobs_root / job_id
|
||||
job_dir.mkdir(parents=True, exist_ok=False)
|
||||
@@ -331,6 +440,7 @@ class IngestManager:
|
||||
"upload_ids": upload_ids,
|
||||
"files": [self._public_upload(upload, expose_path=True) for upload in uploads],
|
||||
"note": self._safe_optional_text(note),
|
||||
"queue_sequence": time.time_ns(),
|
||||
"created_at": utc_now(),
|
||||
"updated_at": utc_now(),
|
||||
"started_at": None,
|
||||
@@ -343,7 +453,7 @@ class IngestManager:
|
||||
}
|
||||
self._write_job(job)
|
||||
self._start_next_job_locked()
|
||||
return self.get_job_status(job_id)
|
||||
return self.get_job_status(job_id)
|
||||
|
||||
def check_opencode_readiness(self, *, run_smoke_test: bool = False) -> dict[str, Any]:
|
||||
result: dict[str, Any] = {
|
||||
@@ -367,8 +477,8 @@ class IngestManager:
|
||||
return result
|
||||
|
||||
def get_job_status(self, job_id: str) -> dict[str, Any]:
|
||||
self.refresh_jobs()
|
||||
with self._lock:
|
||||
self.refresh_jobs()
|
||||
job = self._load_job(job_id)
|
||||
return self._public_job(job)
|
||||
|
||||
@@ -376,17 +486,16 @@ class IngestManager:
|
||||
return self.get_job_status(job_id)
|
||||
|
||||
def get_ingest_job_result(self, job_id: str) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
status = self.get_job_status(job_id)
|
||||
return {
|
||||
"job_id": status["job_id"],
|
||||
"state": status["state"],
|
||||
"summary": self._summary_for_status(status),
|
||||
"files": status["files"],
|
||||
"exit_code": status.get("exit_code"),
|
||||
"opencode_output_tail": status.get("log_tail", ""),
|
||||
"error": status.get("error", ""),
|
||||
}
|
||||
status = self.get_job_status(job_id)
|
||||
return {
|
||||
"job_id": status["job_id"],
|
||||
"state": status["state"],
|
||||
"summary": self._summary_for_status(status),
|
||||
"files": status["files"],
|
||||
"exit_code": status.get("exit_code"),
|
||||
"opencode_output_tail": status.get("log_tail", ""),
|
||||
"error": status.get("error", ""),
|
||||
}
|
||||
|
||||
def cancel_ingest_job(self, job_id: str) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
@@ -398,33 +507,32 @@ class IngestManager:
|
||||
self._write_job(job)
|
||||
return self._public_job(job)
|
||||
if job["state"] == "running":
|
||||
process = self._processes.get(job_id)
|
||||
if process is not None:
|
||||
process.terminate()
|
||||
job["state"] = "cancelled"
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
job["exit_code"] = -15
|
||||
self._write_job(job)
|
||||
self._processes.pop(job_id, None)
|
||||
self._start_next_job_locked()
|
||||
return self._public_job(job)
|
||||
self._supervisor.cancel(job_id)
|
||||
return self._public_job(self._load_job(job_id))
|
||||
return self._public_job(job)
|
||||
|
||||
def refresh_jobs(self) -> None:
|
||||
self._supervisor.refresh()
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
for job_id, process in list(self._processes.items()):
|
||||
code = process.poll()
|
||||
if code is None:
|
||||
continue
|
||||
job = self._load_job(job_id)
|
||||
job["state"] = "succeeded" if code == 0 else "failed"
|
||||
job["exit_code"] = code
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
self._processes.pop(job_id, None)
|
||||
self._start_next_job_locked()
|
||||
if self._closing:
|
||||
cancel_queued = False
|
||||
else:
|
||||
self._closing = True
|
||||
cancel_queued = True
|
||||
try:
|
||||
if cancel_queued:
|
||||
with self._lock:
|
||||
for job in self._iter_jobs():
|
||||
if job.get("state") == "queued":
|
||||
job["state"] = "cancelled"
|
||||
job["error"] = "ingest manager is shutting down"
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
finally:
|
||||
self._supervisor.close()
|
||||
|
||||
def cleanup_runtime(self, *, max_age_days: int = 30) -> dict[str, Any]:
|
||||
if max_age_days <= 0:
|
||||
@@ -437,6 +545,7 @@ class IngestManager:
|
||||
"old_uploads_deleted": 0,
|
||||
"old_jobs_deleted": 0,
|
||||
}
|
||||
self.refresh_jobs()
|
||||
with self._lock:
|
||||
active_upload_ids = {
|
||||
str(upload_id)
|
||||
@@ -486,16 +595,26 @@ class IngestManager:
|
||||
return result
|
||||
|
||||
def _start_next_job_locked(self) -> None:
|
||||
if any(job.get("state") == "running" for job in self._iter_jobs()):
|
||||
if self._closing:
|
||||
return
|
||||
queued = [job for job in self._iter_jobs() if job.get("state") == "queued"]
|
||||
if not queued:
|
||||
return
|
||||
queued.sort(key=lambda job: job["created_at"])
|
||||
job = queued[0]
|
||||
self._launch_job_locked(job)
|
||||
while True:
|
||||
jobs = self._iter_jobs()
|
||||
if any(job.get("state") == "running" for job in jobs):
|
||||
return
|
||||
queued = [job for job in jobs if job.get("state") == "queued"]
|
||||
if not queued:
|
||||
return
|
||||
queued.sort(
|
||||
key=lambda job: (
|
||||
job.get("created_at", ""),
|
||||
int(job.get("queue_sequence") or 0),
|
||||
job.get("job_id", ""),
|
||||
)
|
||||
)
|
||||
if self._launch_job_locked(queued[0]):
|
||||
return
|
||||
|
||||
def _launch_job_locked(self, job: dict[str, Any]) -> None:
|
||||
def _launch_job_locked(self, job: dict[str, Any]) -> bool:
|
||||
readiness = self.check_opencode_readiness()
|
||||
job["preflight"] = readiness
|
||||
if not readiness["ok"]:
|
||||
@@ -504,7 +623,7 @@ class IngestManager:
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
return
|
||||
return False
|
||||
prompt = self._build_prompt(job)
|
||||
args = [
|
||||
self.opencode_path,
|
||||
@@ -519,31 +638,83 @@ class IngestManager:
|
||||
f"MCP ingest {job['job_id']}",
|
||||
prompt,
|
||||
]
|
||||
log_path = Path(job["log_path"])
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
log_file = log_path.open("ab")
|
||||
job["state"] = "running"
|
||||
job["started_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
process = self._supervisor.start(
|
||||
job["job_id"],
|
||||
args,
|
||||
cwd=str(self.project_root),
|
||||
stdout=log_file,
|
||||
stderr=subprocess.STDOUT,
|
||||
log_path=job["log_path"],
|
||||
timeout_seconds=self.job_timeout_seconds,
|
||||
on_complete=self._handle_process_outcome,
|
||||
on_callback_error=self._handle_process_callback_error,
|
||||
)
|
||||
except OSError as exc:
|
||||
log_file.close()
|
||||
except (OSError, RuntimeError, ValueError) as exc:
|
||||
job["state"] = "failed"
|
||||
job["error"] = str(exc)
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
return
|
||||
job["state"] = "running"
|
||||
job["started_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
job["pid"] = process.pid
|
||||
self._processes[job["job_id"]] = process
|
||||
self._write_job(job)
|
||||
log_file.close()
|
||||
return False
|
||||
current = self._load_job(job["job_id"])
|
||||
if current.get("state") == "running":
|
||||
current["pid"] = process.pid
|
||||
self._write_job(current)
|
||||
return True
|
||||
|
||||
def _handle_process_callback_error(self, outcome: ProcessOutcome, error: Exception) -> None:
|
||||
with self._lock:
|
||||
try:
|
||||
job = self._load_job(outcome.job_id)
|
||||
except IngestError:
|
||||
return
|
||||
if job.get("state") == "running":
|
||||
job["state"] = "failed"
|
||||
job["exit_code"] = outcome.returncode
|
||||
job["error"] = f"ingest completion handling failed: {type(error).__name__}: {error}"
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
self._start_next_job_locked()
|
||||
|
||||
def _handle_process_outcome(self, outcome: ProcessOutcome) -> None:
|
||||
with self._lock:
|
||||
try:
|
||||
job = self._load_job(outcome.job_id)
|
||||
except IngestError:
|
||||
return
|
||||
if job.get("state") != "running":
|
||||
return
|
||||
final_text = extract_final_answer(job["log_path"])
|
||||
protocol_ok, result_text, protocol_error = parse_marked_result(
|
||||
final_text,
|
||||
success_marker=INGEST_SUCCESS_MARKER,
|
||||
failure_marker=INGEST_FAILURE_MARKER,
|
||||
)
|
||||
job["exit_code"] = outcome.returncode
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
job["result_text"] = result_text
|
||||
if outcome.cancelled:
|
||||
job["state"] = "cancelled"
|
||||
job["error"] = outcome.error
|
||||
elif outcome.timed_out:
|
||||
job["state"] = "failed"
|
||||
job["error"] = f"OpenCode ingest exceeded {self.job_timeout_seconds} seconds"
|
||||
elif outcome.returncode != 0:
|
||||
job["state"] = "failed"
|
||||
job["error"] = outcome.error or protocol_error or f"OpenCode exited with code {outcome.returncode}"
|
||||
elif not protocol_ok:
|
||||
job["state"] = "failed"
|
||||
job["error"] = protocol_error
|
||||
else:
|
||||
job["state"] = "succeeded"
|
||||
job["error"] = ""
|
||||
self._write_job(job)
|
||||
self._start_next_job_locked()
|
||||
|
||||
def _build_prompt(self, job: dict[str, Any]) -> str:
|
||||
file_lines = []
|
||||
@@ -582,6 +753,8 @@ class IngestManager:
|
||||
"",
|
||||
"After ingesting, archive sources into the vault source archive, update source registry/catalog/wiki metadata as normal, and run `obsidian-wiki file-index rebuild`.",
|
||||
"Return a concise success or failure summary, including whether the file asset index was refreshed.",
|
||||
f"If ingest succeeds, the final response must start with {INGEST_SUCCESS_MARKER} on its own line, followed by the summary.",
|
||||
f"If ingest fails, the final response must start with {INGEST_FAILURE_MARKER} on its own line, followed by the reason.",
|
||||
]
|
||||
)
|
||||
|
||||
@@ -752,6 +925,15 @@ class IngestManager:
|
||||
return duplicates
|
||||
|
||||
def _public_job(self, job: dict[str, Any]) -> dict[str, Any]:
|
||||
output = str(job.get("result_text") or extract_final_answer(job["log_path"]))
|
||||
if output.startswith((INGEST_SUCCESS_MARKER, INGEST_FAILURE_MARKER)):
|
||||
_ok, public_text, error = parse_marked_result(
|
||||
output,
|
||||
success_marker=INGEST_SUCCESS_MARKER,
|
||||
failure_marker=INGEST_FAILURE_MARKER,
|
||||
)
|
||||
output = public_text or error
|
||||
output = self._redact_paths(output)
|
||||
return {
|
||||
"job_id": job["job_id"],
|
||||
"state": job["state"],
|
||||
@@ -760,12 +942,40 @@ class IngestManager:
|
||||
"started_at": job.get("started_at"),
|
||||
"finished_at": job.get("finished_at"),
|
||||
"exit_code": job.get("exit_code"),
|
||||
"pid": job.get("pid"),
|
||||
"error": job.get("error", ""),
|
||||
"preflight": job.get("preflight"),
|
||||
"log_tail": self._read_tail(Path(job["log_path"])),
|
||||
"error": self._redact_paths(str(job.get("error", ""))),
|
||||
"preflight": self._public_preflight(job.get("preflight")),
|
||||
"log_tail": output,
|
||||
}
|
||||
|
||||
def _public_preflight(self, preflight: dict[str, Any] | None) -> dict[str, Any] | None:
|
||||
if preflight is None:
|
||||
return None
|
||||
return {
|
||||
key: self._redact_public_value(value)
|
||||
for key, value in preflight.items()
|
||||
if key != "checked_paths"
|
||||
}
|
||||
|
||||
def _redact_public_value(self, value: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {key: self._redact_public_value(item) for key, item in value.items()}
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [self._redact_public_value(item) for item in value]
|
||||
if isinstance(value, str):
|
||||
return self._redact_paths(value)
|
||||
return value
|
||||
|
||||
def _redact_paths(self, value: str) -> str:
|
||||
text = value
|
||||
for path, label in (
|
||||
(self.runtime_root, "<runtime>"),
|
||||
(self.vault, "<vault>"),
|
||||
(self.project_root, "<project>"),
|
||||
(Path.home(), "<home>"),
|
||||
):
|
||||
text = text.replace(str(path), label)
|
||||
return text
|
||||
|
||||
def _tree_mtime(self, path: Path) -> float:
|
||||
if not path.exists():
|
||||
return time.time()
|
||||
|
||||
@@ -0,0 +1,331 @@
|
||||
"""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"
|
||||
+173
-60
@@ -15,17 +15,21 @@ from threading import RLock
|
||||
from typing import Any
|
||||
|
||||
from obsidian_wiki.file_assets import clean_config_value
|
||||
from obsidian_wiki.mcp_jobs import ProcessOutcome, ProcessSupervisor, extract_final_answer, parse_marked_result
|
||||
|
||||
DEFAULT_QUERY_MAX_CHARS = 8000
|
||||
DEFAULT_OPENCODE_PATH = "opencode"
|
||||
DEFAULT_QUERY_OPENCODE_MODEL = "openai/gpt-5.5"
|
||||
DEFAULT_QUERY_OPENCODE_VARIANT = "minimal"
|
||||
DEFAULT_OPENCODE_SMOKE_TIMEOUT_SECONDS = 120
|
||||
DEFAULT_QUERY_JOB_TIMEOUT_SECONDS = 5 * 60
|
||||
|
||||
QUERY_MODES = {"normal", "index_only"}
|
||||
OPENCODE_REQUIRED_RUN_FLAGS = ("--format", "--auto", "--dir", "--title")
|
||||
OPENCODE_SMOKE_SUCCESS = "WIKI_QUERY_SKILL_LOADED"
|
||||
OPENCODE_SMOKE_FAILURE = "SKILL_LOAD_FAILED"
|
||||
QUERY_SUCCESS_MARKER = "MCP_QUERY_OK"
|
||||
QUERY_FAILURE_MARKER = "MCP_QUERY_FAILED"
|
||||
TERMINAL_JOB_STATES = {"succeeded", "failed", "cancelled"}
|
||||
|
||||
|
||||
@@ -56,6 +60,7 @@ class QueryManager:
|
||||
opencode_model: str | None = DEFAULT_QUERY_OPENCODE_MODEL,
|
||||
opencode_variant: str | None = DEFAULT_QUERY_OPENCODE_VARIANT,
|
||||
max_query_chars: int = DEFAULT_QUERY_MAX_CHARS,
|
||||
job_timeout_seconds: int = DEFAULT_QUERY_JOB_TIMEOUT_SECONDS,
|
||||
) -> None:
|
||||
self.vault = Path(clean_config_value(str(vault_path))).expanduser().resolve()
|
||||
self.project_root = Path(clean_config_value(str(project_root))).expanduser().resolve()
|
||||
@@ -69,10 +74,15 @@ class QueryManager:
|
||||
self.opencode_model = clean_config_value(str(opencode_model or "")).strip()
|
||||
self.opencode_variant = clean_config_value(str(opencode_variant or "")).strip()
|
||||
self.max_query_chars = max_query_chars
|
||||
self.job_timeout_seconds = max(1, int(job_timeout_seconds))
|
||||
self.jobs_root = self.runtime_root / "jobs"
|
||||
self._lock = RLock()
|
||||
self._processes: dict[str, subprocess.Popen[Any]] = {}
|
||||
self._closing = False
|
||||
self.jobs_root.mkdir(parents=True, exist_ok=True)
|
||||
self._supervisor = ProcessSupervisor(
|
||||
self.runtime_root,
|
||||
popen_factory=lambda *args, **kwargs: subprocess.Popen(*args, **kwargs),
|
||||
)
|
||||
self._mark_interrupted_jobs()
|
||||
|
||||
def start_query_job(
|
||||
@@ -86,6 +96,8 @@ class QueryManager:
|
||||
text = self._validate_query(query)
|
||||
query_mode = self._validate_mode(mode)
|
||||
with self._lock:
|
||||
if self._closing:
|
||||
raise QueryError("query manager is closed")
|
||||
job_id = uuid.uuid4().hex
|
||||
job_dir = self.jobs_root / job_id
|
||||
job_dir.mkdir(parents=True, exist_ok=False)
|
||||
@@ -96,6 +108,7 @@ class QueryManager:
|
||||
"mode": query_mode,
|
||||
"public_only": bool(public_only),
|
||||
"note": self._safe_optional_text(note),
|
||||
"queue_sequence": time.time_ns(),
|
||||
"created_at": utc_now(),
|
||||
"updated_at": utc_now(),
|
||||
"started_at": None,
|
||||
@@ -108,7 +121,7 @@ class QueryManager:
|
||||
}
|
||||
self._write_job(job)
|
||||
self._start_next_job_locked()
|
||||
return self.get_query_job_status(job_id)
|
||||
return self.get_query_job_status(job_id)
|
||||
|
||||
def check_opencode_readiness(self, *, run_smoke_test: bool = False) -> dict[str, Any]:
|
||||
result: dict[str, Any] = {
|
||||
@@ -132,15 +145,14 @@ class QueryManager:
|
||||
return result
|
||||
|
||||
def get_query_job_status(self, job_id: str) -> dict[str, Any]:
|
||||
self.refresh_jobs()
|
||||
with self._lock:
|
||||
self.refresh_jobs()
|
||||
job = self._load_job(job_id)
|
||||
return self._public_job(job)
|
||||
|
||||
def get_query_job_result(self, job_id: str) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
status = self.get_query_job_status(job_id)
|
||||
return self._query_result_from_status(status)
|
||||
status = self.get_query_job_status(job_id)
|
||||
return self._query_result_from_status(status)
|
||||
|
||||
def ask_wiki(
|
||||
self,
|
||||
@@ -179,33 +191,32 @@ class QueryManager:
|
||||
self._write_job(job)
|
||||
return self._public_job(job)
|
||||
if job["state"] == "running":
|
||||
process = self._processes.get(job_id)
|
||||
if process is not None:
|
||||
process.terminate()
|
||||
job["state"] = "cancelled"
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
job["exit_code"] = -15
|
||||
self._write_job(job)
|
||||
self._processes.pop(job_id, None)
|
||||
self._start_next_job_locked()
|
||||
return self._public_job(job)
|
||||
self._supervisor.cancel(job_id)
|
||||
return self._public_job(self._load_job(job_id))
|
||||
return self._public_job(job)
|
||||
|
||||
def refresh_jobs(self) -> None:
|
||||
self._supervisor.refresh()
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
for job_id, process in list(self._processes.items()):
|
||||
code = process.poll()
|
||||
if code is None:
|
||||
continue
|
||||
job = self._load_job(job_id)
|
||||
job["state"] = "succeeded" if code == 0 else "failed"
|
||||
job["exit_code"] = code
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
self._processes.pop(job_id, None)
|
||||
self._start_next_job_locked()
|
||||
if self._closing:
|
||||
cancel_queued = False
|
||||
else:
|
||||
self._closing = True
|
||||
cancel_queued = True
|
||||
try:
|
||||
if cancel_queued:
|
||||
with self._lock:
|
||||
for job in self._iter_jobs():
|
||||
if job.get("state") == "queued":
|
||||
job["state"] = "cancelled"
|
||||
job["error"] = "query manager is shutting down"
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
finally:
|
||||
self._supervisor.close()
|
||||
|
||||
def cleanup_runtime(self, *, max_age_days: int = 30) -> dict[str, Any]:
|
||||
if max_age_days <= 0:
|
||||
@@ -216,8 +227,8 @@ class QueryManager:
|
||||
"max_age_days": max_age_days,
|
||||
"old_jobs_deleted": 0,
|
||||
}
|
||||
self.refresh_jobs()
|
||||
with self._lock:
|
||||
self.refresh_jobs()
|
||||
for job in self._iter_jobs():
|
||||
if job.get("state") not in TERMINAL_JOB_STATES:
|
||||
continue
|
||||
@@ -232,15 +243,26 @@ class QueryManager:
|
||||
return result
|
||||
|
||||
def _start_next_job_locked(self) -> None:
|
||||
if any(job.get("state") == "running" for job in self._iter_jobs()):
|
||||
if self._closing:
|
||||
return
|
||||
queued = [job for job in self._iter_jobs() if job.get("state") == "queued"]
|
||||
if not queued:
|
||||
return
|
||||
queued.sort(key=lambda job: job["created_at"])
|
||||
self._launch_job_locked(queued[0])
|
||||
while True:
|
||||
jobs = self._iter_jobs()
|
||||
if any(job.get("state") == "running" for job in jobs):
|
||||
return
|
||||
queued = [job for job in jobs if job.get("state") == "queued"]
|
||||
if not queued:
|
||||
return
|
||||
queued.sort(
|
||||
key=lambda job: (
|
||||
job.get("created_at", ""),
|
||||
int(job.get("queue_sequence") or 0),
|
||||
job.get("job_id", ""),
|
||||
)
|
||||
)
|
||||
if self._launch_job_locked(queued[0]):
|
||||
return
|
||||
|
||||
def _launch_job_locked(self, job: dict[str, Any]) -> None:
|
||||
def _launch_job_locked(self, job: dict[str, Any]) -> bool:
|
||||
readiness = self.check_opencode_readiness()
|
||||
job["preflight"] = readiness
|
||||
if not readiness["ok"]:
|
||||
@@ -249,12 +271,12 @@ class QueryManager:
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
return
|
||||
return False
|
||||
args = [
|
||||
self.opencode_path,
|
||||
"run",
|
||||
"--format",
|
||||
"default",
|
||||
"json",
|
||||
"--auto",
|
||||
*self._opencode_model_args(),
|
||||
"--dir",
|
||||
@@ -263,31 +285,83 @@ class QueryManager:
|
||||
f"MCP query {job['job_id']}",
|
||||
self._build_prompt(job),
|
||||
]
|
||||
log_path = Path(job["log_path"])
|
||||
log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
log_file = log_path.open("ab")
|
||||
job["state"] = "running"
|
||||
job["started_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
try:
|
||||
process = subprocess.Popen(
|
||||
process = self._supervisor.start(
|
||||
job["job_id"],
|
||||
args,
|
||||
cwd=str(self.project_root),
|
||||
stdout=log_file,
|
||||
stderr=subprocess.STDOUT,
|
||||
log_path=job["log_path"],
|
||||
timeout_seconds=self.job_timeout_seconds,
|
||||
on_complete=self._handle_process_outcome,
|
||||
on_callback_error=self._handle_process_callback_error,
|
||||
)
|
||||
except OSError as exc:
|
||||
log_file.close()
|
||||
except (OSError, RuntimeError, ValueError) as exc:
|
||||
job["state"] = "failed"
|
||||
job["error"] = str(exc)
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
return
|
||||
job["state"] = "running"
|
||||
job["started_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
job["pid"] = process.pid
|
||||
self._processes[job["job_id"]] = process
|
||||
self._write_job(job)
|
||||
log_file.close()
|
||||
return False
|
||||
current = self._load_job(job["job_id"])
|
||||
if current.get("state") == "running":
|
||||
current["pid"] = process.pid
|
||||
self._write_job(current)
|
||||
return True
|
||||
|
||||
def _handle_process_callback_error(self, outcome: ProcessOutcome, error: Exception) -> None:
|
||||
with self._lock:
|
||||
try:
|
||||
job = self._load_job(outcome.job_id)
|
||||
except QueryError:
|
||||
return
|
||||
if job.get("state") == "running":
|
||||
job["state"] = "failed"
|
||||
job["exit_code"] = outcome.returncode
|
||||
job["error"] = f"query completion handling failed: {type(error).__name__}: {error}"
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
self._write_job(job)
|
||||
self._start_next_job_locked()
|
||||
|
||||
def _handle_process_outcome(self, outcome: ProcessOutcome) -> None:
|
||||
with self._lock:
|
||||
try:
|
||||
job = self._load_job(outcome.job_id)
|
||||
except QueryError:
|
||||
return
|
||||
if job.get("state") != "running":
|
||||
return
|
||||
final_text = extract_final_answer(job["log_path"])
|
||||
protocol_ok, answer, protocol_error = parse_marked_result(
|
||||
final_text,
|
||||
success_marker=QUERY_SUCCESS_MARKER,
|
||||
failure_marker=QUERY_FAILURE_MARKER,
|
||||
)
|
||||
job["exit_code"] = outcome.returncode
|
||||
job["finished_at"] = utc_now()
|
||||
job["updated_at"] = utc_now()
|
||||
job["result_text"] = answer
|
||||
if outcome.cancelled:
|
||||
job["state"] = "cancelled"
|
||||
job["error"] = outcome.error
|
||||
elif outcome.timed_out:
|
||||
job["state"] = "failed"
|
||||
job["error"] = f"OpenCode query exceeded {self.job_timeout_seconds} seconds"
|
||||
elif outcome.returncode != 0:
|
||||
job["state"] = "failed"
|
||||
job["error"] = outcome.error or protocol_error or f"OpenCode exited with code {outcome.returncode}"
|
||||
elif not protocol_ok:
|
||||
job["state"] = "failed"
|
||||
job["error"] = protocol_error
|
||||
else:
|
||||
job["state"] = "succeeded"
|
||||
job["error"] = ""
|
||||
self._write_job(job)
|
||||
self._start_next_job_locked()
|
||||
|
||||
def _build_prompt(self, job: dict[str, Any]) -> str:
|
||||
mode = "index-only" if job.get("mode") == "index_only" else "normal"
|
||||
@@ -313,6 +387,8 @@ class QueryManager:
|
||||
"Business SQL must follow the schema guardrails: only SELECT/WITH SELECT against documented structured tables, never raw_json, raw_snapshots, writes, DDL, PRAGMA, or sync/config operations.",
|
||||
"If Visibility is public only, do not use business facts; continue with filtered wiki retrieval.",
|
||||
"Return the final synthesized answer with wiki citations when available.",
|
||||
f"If the query succeeds, the final response must start with {QUERY_SUCCESS_MARKER} on its own line, followed by the answer.",
|
||||
f"If it fails, the final response must start with {QUERY_FAILURE_MARKER} on its own line, followed by the reason.",
|
||||
"",
|
||||
f"Optional context note: {note}",
|
||||
]
|
||||
@@ -348,6 +424,15 @@ class QueryManager:
|
||||
self._write_job(job)
|
||||
|
||||
def _public_job(self, job: dict[str, Any]) -> dict[str, Any]:
|
||||
output = str(job.get("result_text") or extract_final_answer(job["log_path"]))
|
||||
if output.startswith((QUERY_SUCCESS_MARKER, QUERY_FAILURE_MARKER)):
|
||||
_ok, public_text, error = parse_marked_result(
|
||||
output,
|
||||
success_marker=QUERY_SUCCESS_MARKER,
|
||||
failure_marker=QUERY_FAILURE_MARKER,
|
||||
)
|
||||
output = public_text or error
|
||||
output = self._redact_paths(output)
|
||||
return {
|
||||
"job_id": job["job_id"],
|
||||
"state": job["state"],
|
||||
@@ -359,12 +444,40 @@ class QueryManager:
|
||||
"started_at": job.get("started_at"),
|
||||
"finished_at": job.get("finished_at"),
|
||||
"exit_code": job.get("exit_code"),
|
||||
"pid": job.get("pid"),
|
||||
"error": job.get("error", ""),
|
||||
"preflight": job.get("preflight"),
|
||||
"log_tail": self._read_tail(Path(job["log_path"])),
|
||||
"error": self._redact_paths(str(job.get("error", ""))),
|
||||
"preflight": self._public_preflight(job.get("preflight")),
|
||||
"log_tail": output,
|
||||
}
|
||||
|
||||
def _public_preflight(self, preflight: dict[str, Any] | None) -> dict[str, Any] | None:
|
||||
if preflight is None:
|
||||
return None
|
||||
return {
|
||||
key: self._redact_public_value(value)
|
||||
for key, value in preflight.items()
|
||||
if key != "checked_paths"
|
||||
}
|
||||
|
||||
def _redact_public_value(self, value: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {key: self._redact_public_value(item) for key, item in value.items()}
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [self._redact_public_value(item) for item in value]
|
||||
if isinstance(value, str):
|
||||
return self._redact_paths(value)
|
||||
return value
|
||||
|
||||
def _redact_paths(self, value: str) -> str:
|
||||
text = value
|
||||
for path, label in (
|
||||
(self.runtime_root, "<runtime>"),
|
||||
(self.vault, "<vault>"),
|
||||
(self.project_root, "<project>"),
|
||||
(Path.home(), "<home>"),
|
||||
):
|
||||
text = text.replace(str(path), label)
|
||||
return text
|
||||
|
||||
def _query_result_from_status(self, status: dict[str, Any]) -> dict[str, Any]:
|
||||
output_tail = status.get("log_tail", "")
|
||||
answer = output_tail.strip()
|
||||
@@ -542,7 +655,7 @@ class QueryManager:
|
||||
self.opencode_path,
|
||||
"run",
|
||||
"--format",
|
||||
"default",
|
||||
"json",
|
||||
"--auto",
|
||||
*self._opencode_model_args(),
|
||||
"--dir",
|
||||
|
||||
+452
-193
@@ -3,14 +3,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import hmac
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
from email import policy
|
||||
from email.parser import BytesParser
|
||||
from fnmatch import fnmatch
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
@@ -19,9 +18,11 @@ from obsidian_wiki import __version__
|
||||
from obsidian_wiki.config import config_value, resolve_vault_path
|
||||
from obsidian_wiki.file_assets import clean_config_value, db_path_for_vault, registry_path_for_vault
|
||||
from obsidian_wiki.mcp_ingest import (
|
||||
DEFAULT_INGEST_JOB_TIMEOUT_SECONDS,
|
||||
DEFAULT_INGEST_MAX_CHUNK_BYTES,
|
||||
DEFAULT_INGEST_MAX_FILE_BYTES,
|
||||
DEFAULT_INGEST_MAX_FILES_PER_JOB,
|
||||
DEFAULT_INGEST_MAX_REQUEST_BYTES,
|
||||
DEFAULT_INGEST_OPENCODE_MODEL,
|
||||
DEFAULT_INGEST_OPENCODE_VARIANT,
|
||||
DEFAULT_UPLOAD_TICKET_SECONDS,
|
||||
@@ -32,6 +33,7 @@ from obsidian_wiki.mcp_ingest import (
|
||||
parse_bool,
|
||||
)
|
||||
from obsidian_wiki.mcp_query import (
|
||||
DEFAULT_QUERY_JOB_TIMEOUT_SECONDS,
|
||||
DEFAULT_QUERY_MAX_CHARS,
|
||||
DEFAULT_QUERY_OPENCODE_MODEL,
|
||||
DEFAULT_QUERY_OPENCODE_VARIANT,
|
||||
@@ -56,9 +58,13 @@ DEFAULT_HTTP_PORT = 8765
|
||||
DEFAULT_HTTP_PATH = "/mcp"
|
||||
DEFAULT_INGEST_UPLOAD_PATH = "/mcp-ingest/uploads"
|
||||
DEFAULT_TOKEN_ENV = "OBSIDIAN_MCP_TOKEN"
|
||||
PUBLIC_BASE_URL_ENV = "OBSIDIAN_MCP_PUBLIC_BASE_URL"
|
||||
PUBLIC_UPLOAD_URL_ENV = "OBSIDIAN_MCP_INGEST_PUBLIC_UPLOAD_URL"
|
||||
ALLOWED_HOSTS_ENV = "OBSIDIAN_MCP_ALLOWED_HOSTS"
|
||||
ALLOWED_ORIGINS_ENV = "OBSIDIAN_MCP_ALLOWED_ORIGINS"
|
||||
MCP_SERVER_NAME = "wiki-file-assets"
|
||||
CAPABILITIES_SCHEMA_VERSION = 1
|
||||
TOOLS_VERSION = 1
|
||||
CAPABILITIES_SCHEMA_VERSION = 2
|
||||
TOOLS_VERSION = 2
|
||||
TOOL_FILE_ASSETS = "file-assets"
|
||||
TOOL_INGEST = "ingest"
|
||||
TOOL_QUERY = "query"
|
||||
@@ -178,7 +184,7 @@ class BearerTokenAuthMiddleware:
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
|
||||
class RequestBodyTooLarge(IngestError):
|
||||
class RequestBodyTooLarge(IngestError, OSError):
|
||||
"""Raised when an HTTP ingest upload body exceeds the configured request limit."""
|
||||
|
||||
|
||||
@@ -189,6 +195,7 @@ def create_server(
|
||||
profile_name: str = DEFAULT_PROFILE_NAME,
|
||||
max_file_bytes: int = DEFAULT_MAX_FILE_BYTES,
|
||||
expose_paths: bool = False,
|
||||
allow_source_files: bool = False,
|
||||
host: str = DEFAULT_HTTP_HOST,
|
||||
port: int = DEFAULT_HTTP_PORT,
|
||||
streamable_http_path: str = DEFAULT_HTTP_PATH,
|
||||
@@ -209,7 +216,10 @@ def create_server(
|
||||
ingest_max_file_bytes: int = DEFAULT_INGEST_MAX_FILE_BYTES,
|
||||
ingest_max_chunk_bytes: int = DEFAULT_INGEST_MAX_CHUNK_BYTES,
|
||||
ingest_max_files_per_job: int = DEFAULT_INGEST_MAX_FILES_PER_JOB,
|
||||
ingest_max_request_bytes: int = DEFAULT_INGEST_MAX_REQUEST_BYTES,
|
||||
ingest_job_timeout_seconds: int = DEFAULT_INGEST_JOB_TIMEOUT_SECONDS,
|
||||
query_max_chars: int = DEFAULT_QUERY_MAX_CHARS,
|
||||
query_job_timeout_seconds: int = DEFAULT_QUERY_JOB_TIMEOUT_SECONDS,
|
||||
ingest_manager: IngestManager | None = None,
|
||||
query_manager: QueryManager | None = None,
|
||||
):
|
||||
@@ -227,8 +237,10 @@ def create_server(
|
||||
vault = Path(clean_config_value(str(vault_path))).expanduser()
|
||||
db = _resolve_db_path(vault, db_path)
|
||||
project = Path(project_root).expanduser().resolve() if project_root is not None else Path.cwd().resolve()
|
||||
public_base = _normalize_public_base_url(public_base_url or runtime_config_value("OBSIDIAN_MCP_PUBLIC_BASE_URL", ""))
|
||||
public_upload_url = _upload_url_from_public_base(public_base, DEFAULT_INGEST_UPLOAD_PATH)
|
||||
public_base = _normalize_public_base_url(public_base_url or runtime_config_value(PUBLIC_BASE_URL_ENV, ""))
|
||||
public_upload_url = _normalize_public_upload_url(runtime_config_value(PUBLIC_UPLOAD_URL_ENV, ""))
|
||||
if not public_upload_url:
|
||||
public_upload_url = _upload_url_from_public_base(public_base, DEFAULT_INGEST_UPLOAD_PATH)
|
||||
ingest_runtime = (
|
||||
Path(ingest_runtime_root).expanduser().resolve()
|
||||
if ingest_runtime_root is not None
|
||||
@@ -264,6 +276,7 @@ def create_server(
|
||||
"profile_name": profile_name,
|
||||
"max_file_bytes": max_file_bytes,
|
||||
"expose_paths": expose_paths,
|
||||
"allow_source_files": allow_source_files,
|
||||
"enabled_tools": sorted(
|
||||
([TOOL_FILE_ASSETS] if enable_file_assets else [])
|
||||
+ ([TOOL_INGEST] if enable_ingest else [])
|
||||
@@ -297,15 +310,15 @@ def create_server(
|
||||
return {
|
||||
"server_name": MCP_SERVER_NAME,
|
||||
"server_version": __version__,
|
||||
"schema_version": CAPABILITIES_SCHEMA_VERSION,
|
||||
"capabilities_schema_version": CAPABILITIES_SCHEMA_VERSION,
|
||||
"tools_version": TOOLS_VERSION,
|
||||
"enabled_tools": enabled,
|
||||
"recommended_first_steps": [
|
||||
"Use ask_wiki for ordinary wiki, business facts, or hybrid questions when query is enabled.",
|
||||
"Use search_file_assets before reading source file resources.",
|
||||
"Use ingest_text_note for short text notes when ingest is enabled.",
|
||||
"Use create_ingest_upload_ticket before remote HTTP file upload.",
|
||||
],
|
||||
"recommended_first_steps": _recommended_first_steps(
|
||||
enable_file_assets=enable_file_assets,
|
||||
enable_ingest=enable_ingest,
|
||||
enable_query=enable_query,
|
||||
),
|
||||
"file_index": _file_index_capabilities(vault, db),
|
||||
"file_assets": {
|
||||
"enabled": enable_file_assets,
|
||||
@@ -322,6 +335,7 @@ def create_server(
|
||||
"best_for": "safe engineering source files such as drawings, CAD/3D models, specifications, and test reports",
|
||||
"blocked_content": "raw orders, prices, email correspondence, complaint evidence, claims, and personal data",
|
||||
"max_file_bytes": max_file_bytes,
|
||||
"source_file_download_enabled": bool(enable_file_assets and allow_source_files),
|
||||
},
|
||||
"query": {
|
||||
"enabled": enable_query,
|
||||
@@ -353,6 +367,7 @@ def create_server(
|
||||
"public_only_disables_business_facts": True,
|
||||
"supports_structured_result": enable_query,
|
||||
"max_query_chars": query_max_chars,
|
||||
"job_timeout_seconds": query_job_timeout_seconds,
|
||||
"known_limitations": (
|
||||
[
|
||||
"answer_markdown is produced by the local wiki-query agent; citations and source flags are best-effort structured metadata.",
|
||||
@@ -370,7 +385,7 @@ def create_server(
|
||||
"http_upload_url_warning": (
|
||||
None
|
||||
if public_upload_url or not enable_ingest
|
||||
else "Set OBSIDIAN_MCP_PUBLIC_BASE_URL so remote agents can receive a complete upload URL."
|
||||
else "Set OBSIDIAN_MCP_PUBLIC_BASE_URL or OBSIDIAN_MCP_INGEST_PUBLIC_UPLOAD_URL so remote agents can receive a complete upload URL."
|
||||
),
|
||||
"preferred_upload_auth": "upload_ticket" if enable_ingest else None,
|
||||
"tools": (
|
||||
@@ -391,6 +406,10 @@ def create_server(
|
||||
),
|
||||
"max_file_bytes": ingest_max_file_bytes,
|
||||
"max_files_per_job": ingest_max_files_per_job,
|
||||
"max_request_bytes": _max_upload_body_bytes(ingest_manager, None)
|
||||
if enable_ingest and ingest_manager is not None
|
||||
else ingest_max_request_bytes or None,
|
||||
"job_timeout_seconds": ingest_job_timeout_seconds,
|
||||
"last_ingest_at": _last_ingest_at(ingest_manager) if enable_ingest and ingest_manager is not None else None,
|
||||
"queue": _ingest_queue_summary(ingest_manager) if enable_ingest and ingest_manager is not None else None,
|
||||
"known_limitations": (
|
||||
@@ -404,13 +423,21 @@ def create_server(
|
||||
),
|
||||
},
|
||||
"safety": {
|
||||
"arbitrary_shell": False,
|
||||
"arbitrary_shell": bool(enable_ingest or enable_query),
|
||||
"direct_arbitrary_shell_tool": False,
|
||||
"arbitrary_prompt": False,
|
||||
"fixed_prompt_entrypoints": True,
|
||||
"prompt_restrictions_enforced": False,
|
||||
"shell_sandbox_enforced": False,
|
||||
"opencode_auto_approve": bool(enable_ingest or enable_query),
|
||||
"trusted_clients_required": True,
|
||||
"local_paths_hidden_by_default": not expose_paths,
|
||||
"model_choice_server_side": True,
|
||||
"source_file_download_enabled": bool(enable_file_assets and allow_source_files),
|
||||
},
|
||||
}
|
||||
|
||||
owns_ingest_manager = False
|
||||
if enable_ingest:
|
||||
if ingest_manager is None:
|
||||
ingest_manager = IngestManager(
|
||||
@@ -423,7 +450,10 @@ def create_server(
|
||||
max_file_bytes=ingest_max_file_bytes,
|
||||
max_chunk_bytes=ingest_max_chunk_bytes,
|
||||
max_files_per_job=ingest_max_files_per_job,
|
||||
max_request_bytes=ingest_max_request_bytes,
|
||||
job_timeout_seconds=ingest_job_timeout_seconds,
|
||||
)
|
||||
owns_ingest_manager = True
|
||||
|
||||
@server.tool()
|
||||
def get_ingest_upload_instructions() -> dict[str, Any]:
|
||||
@@ -431,7 +461,7 @@ def create_server(
|
||||
upload_url_warning = (
|
||||
None
|
||||
if public_upload_url
|
||||
else "Set OBSIDIAN_MCP_PUBLIC_BASE_URL so remote agents can receive a complete upload URL."
|
||||
else "Set OBSIDIAN_MCP_PUBLIC_BASE_URL or OBSIDIAN_MCP_INGEST_PUBLIC_UPLOAD_URL so remote agents can receive a complete upload URL."
|
||||
)
|
||||
return {
|
||||
"preferred_method": "use ingest_text_note for short text; for files, call create_ingest_upload_ticket and upload with X-Obsidian-Wiki-Upload-Token",
|
||||
@@ -472,6 +502,7 @@ def create_server(
|
||||
"max_file_bytes": ingest_manager.max_file_bytes,
|
||||
"max_chunk_bytes": ingest_manager.max_chunk_bytes,
|
||||
"max_files_per_job": ingest_manager.max_files_per_job,
|
||||
"max_request_bytes": _max_upload_body_bytes(ingest_manager, None),
|
||||
},
|
||||
"staging": {
|
||||
"note": "uploads are temporary staging files outside the vault; wiki-ingest decides final archive/index placement",
|
||||
@@ -529,7 +560,7 @@ def create_server(
|
||||
"upload_url_warning": (
|
||||
None
|
||||
if public_upload_url
|
||||
else "Set OBSIDIAN_MCP_PUBLIC_BASE_URL so remote agents can receive a complete upload URL."
|
||||
else "Set OBSIDIAN_MCP_PUBLIC_BASE_URL or OBSIDIAN_MCP_INGEST_PUBLIC_UPLOAD_URL so remote agents can receive a complete upload URL."
|
||||
),
|
||||
"token_header": "X-Obsidian-Wiki-Upload-Token",
|
||||
"expires_seconds_max": MAX_UPLOAD_TICKET_SECONDS,
|
||||
@@ -577,19 +608,25 @@ def create_server(
|
||||
@server.tool()
|
||||
def check_ingest_opencode_readiness(run_smoke_test: bool = False) -> dict[str, Any]:
|
||||
"""Check whether local OpenCode can load the wiki-ingest skill for MCP ingest jobs."""
|
||||
return ingest_manager.check_opencode_readiness(run_smoke_test=run_smoke_test)
|
||||
return _public_readiness(ingest_manager, run_smoke_test=run_smoke_test)
|
||||
|
||||
if enable_query:
|
||||
if query_manager is None:
|
||||
query_manager = QueryManager(
|
||||
vault_path=vault,
|
||||
project_root=project,
|
||||
runtime_root=query_runtime,
|
||||
opencode_path=opencode_path,
|
||||
opencode_model=query_opencode_model,
|
||||
opencode_variant=query_opencode_variant,
|
||||
max_query_chars=query_max_chars,
|
||||
)
|
||||
try:
|
||||
query_manager = QueryManager(
|
||||
vault_path=vault,
|
||||
project_root=project,
|
||||
runtime_root=query_runtime,
|
||||
opencode_path=opencode_path,
|
||||
opencode_model=query_opencode_model,
|
||||
opencode_variant=query_opencode_variant,
|
||||
max_query_chars=query_max_chars,
|
||||
job_timeout_seconds=query_job_timeout_seconds,
|
||||
)
|
||||
except Exception:
|
||||
if owns_ingest_manager and ingest_manager is not None:
|
||||
ingest_manager.close()
|
||||
raise
|
||||
|
||||
@server.tool()
|
||||
def ask_wiki(
|
||||
@@ -643,7 +680,7 @@ def create_server(
|
||||
@server.tool()
|
||||
def check_query_opencode_readiness(run_smoke_test: bool = False) -> dict[str, Any]:
|
||||
"""Check whether local OpenCode can load the wiki-query skill for MCP query jobs."""
|
||||
return query_manager.check_opencode_readiness(run_smoke_test=run_smoke_test)
|
||||
return _public_readiness(query_manager, run_smoke_test=run_smoke_test)
|
||||
|
||||
if enable_file_assets:
|
||||
|
||||
@@ -722,6 +759,7 @@ def create_server(
|
||||
db_path=db,
|
||||
profile_name=profile_name,
|
||||
max_file_bytes=max_file_bytes,
|
||||
allow_source_files=allow_source_files,
|
||||
)
|
||||
return resource["bytes"]
|
||||
|
||||
@@ -730,20 +768,40 @@ def create_server(
|
||||
description="Explain how to search and retrieve original source files from the wiki-file-assets MCP server.",
|
||||
)
|
||||
def how_to_use_wiki_file_assets() -> str:
|
||||
download_status = (
|
||||
"Source file downloads are enabled by the server operator. "
|
||||
if allow_source_files
|
||||
else "Source file downloads are disabled; metadata search remains available. "
|
||||
)
|
||||
return (
|
||||
"wiki-file-assets serves original source files archived from an Obsidian wiki vault. "
|
||||
"Use search_file_assets for keyword searches such as product codes, drawing names, sizes, "
|
||||
"or document types. Use list_file_assets to browse allowed engineering files with optional "
|
||||
"filters. Use get_file_asset_metadata with a file_id from search/list results to get the "
|
||||
"asset://source/sha256/... URI, then read that resource to fetch the original file bytes. "
|
||||
"This MCP is intentionally limited to safe engineering files such as drawings, CAD/3D models, "
|
||||
f"{download_status}This MCP is intentionally limited to safe engineering files such as drawings, CAD/3D models, "
|
||||
"specifications, product parameters, and test reports; it blocks raw orders, prices, email "
|
||||
"correspondence, complaint evidence, claims, and personal data."
|
||||
)
|
||||
|
||||
setattr(
|
||||
server,
|
||||
"_obsidian_job_managers",
|
||||
tuple(manager for manager in (ingest_manager, query_manager) if manager is not None),
|
||||
)
|
||||
return server
|
||||
|
||||
|
||||
def close_server_managers(server: Any) -> list[Exception]:
|
||||
errors: list[Exception] = []
|
||||
for manager in getattr(server, "_obsidian_job_managers", ()):
|
||||
try:
|
||||
manager.close()
|
||||
except Exception as exc:
|
||||
errors.append(exc)
|
||||
return errors
|
||||
|
||||
|
||||
def read_token_from_env(env_name: str = DEFAULT_TOKEN_ENV) -> str:
|
||||
token = os.environ.get(env_name, "").strip()
|
||||
if not token:
|
||||
@@ -767,6 +825,30 @@ def runtime_config_int(key: str, default: int) -> int:
|
||||
return default
|
||||
|
||||
|
||||
def runtime_config_list(key: str) -> list[str] | None:
|
||||
raw = os.environ.get(key)
|
||||
if raw is None:
|
||||
configured = config_value(key, "")
|
||||
if not configured:
|
||||
return None
|
||||
raw = configured
|
||||
return [item.strip() for item in clean_config_value(raw).split(",") if item.strip()]
|
||||
|
||||
|
||||
def _public_readiness(manager: Any, *, run_smoke_test: bool) -> dict[str, Any]:
|
||||
result = manager.check_opencode_readiness(run_smoke_test=run_smoke_test)
|
||||
sanitizer = getattr(manager, "_public_preflight", None)
|
||||
if callable(sanitizer):
|
||||
public = sanitizer(result)
|
||||
if isinstance(public, dict):
|
||||
return public
|
||||
return {
|
||||
key: value
|
||||
for key, value in result.items()
|
||||
if key not in {"checked_paths", "errors"}
|
||||
} | ({"errors": ["OpenCode readiness failed"]} if result.get("errors") else {"errors": []})
|
||||
|
||||
|
||||
class IngestUploadRouter:
|
||||
def __init__(
|
||||
self,
|
||||
@@ -782,8 +864,10 @@ class IngestUploadRouter:
|
||||
self.ingest_manager = ingest_manager
|
||||
self.upload_path = upload_path
|
||||
self.bearer_auth_token = bearer_auth_token
|
||||
self.allowed_hosts = allowed_hosts or []
|
||||
self.allowed_origins = allowed_origins or []
|
||||
self.enforce_allowed_hosts = allowed_hosts is not None
|
||||
self.enforce_allowed_origins = allowed_origins is not None
|
||||
self.allowed_hosts = list(allowed_hosts or [])
|
||||
self.allowed_origins = list(allowed_origins or [])
|
||||
|
||||
async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
|
||||
if scope.get("type") != "http" or scope.get("path") != self.upload_path:
|
||||
@@ -809,19 +893,19 @@ class IngestUploadRouter:
|
||||
try:
|
||||
data = await self._handle_upload(scope, receive, upload_ticket=upload_ticket)
|
||||
except RequestBodyTooLarge as exc:
|
||||
await _send_json(send, 413, {"error": str(exc)})
|
||||
await _send_json(send, 413, _upload_error_payload(str(exc), upload_ticket=upload_ticket))
|
||||
return
|
||||
except IngestError as exc:
|
||||
await _send_json(send, 400, {"error": str(exc)})
|
||||
await _send_json(send, 400, _upload_error_payload(str(exc), upload_ticket=upload_ticket))
|
||||
return
|
||||
await _send_json(send, 200, data)
|
||||
|
||||
def _transport_headers_allowed(self, scope: dict[str, Any]) -> bool:
|
||||
host = _header_value(scope, b"host")
|
||||
if self.allowed_hosts and (not host or not _matches_host(host, self.allowed_hosts)):
|
||||
if self.enforce_allowed_hosts and (not host or not _matches_host(host, self.allowed_hosts)):
|
||||
return False
|
||||
origin = _header_value(scope, b"origin")
|
||||
if origin and self.allowed_origins and not _matches_any(origin, self.allowed_origins):
|
||||
if origin and self.enforce_allowed_origins and not _matches_any(origin, self.allowed_origins):
|
||||
return False
|
||||
return True
|
||||
|
||||
@@ -832,7 +916,7 @@ class IngestUploadRouter:
|
||||
return True
|
||||
if _header_value(scope, b"x-obsidian-wiki-upload-token"):
|
||||
return True
|
||||
return bool(_query_param(scope, "upload_token"))
|
||||
return False
|
||||
|
||||
def _authorize_upload(self, scope: dict[str, Any]) -> dict[str, Any] | None:
|
||||
authorization = _header_value(scope, b"authorization")
|
||||
@@ -842,8 +926,6 @@ class IngestUploadRouter:
|
||||
if self.ingest_manager is None:
|
||||
raise IngestError("unauthorized")
|
||||
upload_token = _header_value(scope, b"x-obsidian-wiki-upload-token")
|
||||
if not upload_token:
|
||||
upload_token = _query_param(scope, "upload_token")
|
||||
if not upload_token:
|
||||
raise IngestError("unauthorized")
|
||||
return self.ingest_manager.consume_upload_ticket(upload_token)
|
||||
@@ -858,30 +940,135 @@ class IngestUploadRouter:
|
||||
manager = self.ingest_manager
|
||||
if manager is None:
|
||||
raise IngestError("ingest is not enabled")
|
||||
body = await _read_http_body(receive, max_bytes=_max_upload_body_bytes(manager, upload_ticket))
|
||||
content_type = _header_value(scope, b"content-type")
|
||||
fields, files = _parse_multipart_form(content_type, body)
|
||||
if not files:
|
||||
raise IngestError("multipart field 'files' must include at least one file")
|
||||
if not content_type.lower().startswith("multipart/form-data"):
|
||||
raise IngestError("content-type must be multipart/form-data")
|
||||
max_body_bytes = _max_upload_body_bytes(manager, upload_ticket)
|
||||
content_length = _header_value(scope, b"content-length")
|
||||
if content_length:
|
||||
try:
|
||||
declared_body_bytes = int(content_length)
|
||||
except ValueError as exc:
|
||||
raise IngestError("content-length must be an integer") from exc
|
||||
if declared_body_bytes < 0:
|
||||
raise IngestError("content-length must be non-negative")
|
||||
if max_body_bytes is not None and declared_body_bytes > max_body_bytes:
|
||||
raise RequestBodyTooLarge(
|
||||
f"request body exceeds max upload size: {declared_body_bytes} > {max_body_bytes}"
|
||||
)
|
||||
max_files = int(upload_ticket.get("max_files", manager.max_files_per_job)) if upload_ticket else manager.max_files_per_job
|
||||
if len(files) > max_files:
|
||||
raise IngestError(f"too many files for one ingest job: {len(files)} > {max_files}")
|
||||
project = _first_field(fields, "project") or (upload_ticket or {}).get("project")
|
||||
doc_type_hint = _first_field(fields, "doc_type_hint") or (upload_ticket or {}).get("doc_type_hint")
|
||||
note = _first_field(fields, "note") or (upload_ticket or {}).get("note")
|
||||
start_default = bool((upload_ticket or {}).get("start_job", True))
|
||||
start_job = parse_bool(_first_field(fields, "start_job"), default=start_default)
|
||||
uploads = [
|
||||
manager.create_completed_upload(
|
||||
file_name=file["file_name"],
|
||||
payload=file["payload"],
|
||||
project=project,
|
||||
doc_type_hint=doc_type_hint,
|
||||
)
|
||||
for file in files
|
||||
]
|
||||
job = manager.start_ingest_job([upload["upload_id"] for upload in uploads], note=note) if start_job else None
|
||||
return {"uploads": uploads, "job": job}
|
||||
received_body_bytes = 0
|
||||
|
||||
async def limited_receive() -> dict[str, Any]:
|
||||
nonlocal received_body_bytes
|
||||
message = await receive()
|
||||
if message.get("type") == "http.request":
|
||||
received_body_bytes += len(message.get("body", b""))
|
||||
if max_body_bytes is not None and received_body_bytes > max_body_bytes:
|
||||
raise RequestBodyTooLarge(
|
||||
f"request body exceeds max upload size: {received_body_bytes} > {max_body_bytes}"
|
||||
)
|
||||
return message
|
||||
|
||||
try:
|
||||
from starlette.datastructures import UploadFile
|
||||
from starlette.formparsers import MultiPartException, MultiPartParser
|
||||
from starlette.requests import Request
|
||||
except ImportError as exc:
|
||||
raise IngestError("streaming multipart support is not installed") from exc
|
||||
|
||||
class FilePartTooLarge(MultiPartException):
|
||||
pass
|
||||
|
||||
class LimitedFileMultiPartParser(MultiPartParser):
|
||||
def __init__(self, *args: Any, max_file_part_size: int, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.max_file_part_size = max_file_part_size
|
||||
self.current_file_size = 0
|
||||
|
||||
def on_part_begin(self) -> None:
|
||||
super().on_part_begin()
|
||||
self.current_file_size = 0
|
||||
|
||||
def on_part_data(self, data: bytes, start: int, end: int) -> None:
|
||||
if self._current_part.file is not None:
|
||||
self.current_file_size += end - start
|
||||
if self.max_file_part_size > 0 and self.current_file_size > self.max_file_part_size:
|
||||
raise FilePartTooLarge(
|
||||
f"file part exceeds max upload size: "
|
||||
f"{self.current_file_size} > {self.max_file_part_size}"
|
||||
)
|
||||
super().on_part_data(data, start, end)
|
||||
|
||||
form = None
|
||||
try:
|
||||
request = Request(scope, limited_receive)
|
||||
try:
|
||||
parser = LimitedFileMultiPartParser(
|
||||
request.headers,
|
||||
request.stream(),
|
||||
max_files=max_files,
|
||||
max_fields=8,
|
||||
max_part_size=64 * 1024,
|
||||
max_file_part_size=manager.max_file_bytes,
|
||||
)
|
||||
form = await parser.parse()
|
||||
except FilePartTooLarge as exc:
|
||||
raise RequestBodyTooLarge(str(exc)) from exc
|
||||
except MultiPartException as exc:
|
||||
raise IngestError(f"multipart request body is invalid: {exc}") from exc
|
||||
files = [
|
||||
value
|
||||
for field_name, value in form.multi_items()
|
||||
if field_name == "files" and isinstance(value, UploadFile)
|
||||
]
|
||||
if not files:
|
||||
raise IngestError("multipart field 'files' must include at least one file")
|
||||
if len(files) > max_files:
|
||||
raise IngestError(f"too many files for one ingest job: {len(files)} > {max_files}")
|
||||
|
||||
def first_text_field(name: str) -> str | None:
|
||||
for value in form.getlist(name):
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
return None
|
||||
|
||||
def ticket_bound_field(name: str) -> str | None:
|
||||
ticket_value = (upload_ticket or {}).get(name)
|
||||
return ticket_value if ticket_value not in (None, "") else first_text_field(name)
|
||||
|
||||
project = ticket_bound_field("project")
|
||||
doc_type_hint = ticket_bound_field("doc_type_hint")
|
||||
note = ticket_bound_field("note")
|
||||
start_default = bool((upload_ticket or {}).get("start_job", True))
|
||||
requested_start = parse_bool(first_text_field("start_job"), default=start_default)
|
||||
start_job = requested_start if upload_ticket is None else start_default and requested_start
|
||||
uploads = []
|
||||
try:
|
||||
for file in files:
|
||||
await file.seek(0)
|
||||
upload = await _run_blocking(
|
||||
lambda file=file: manager.create_completed_upload_from_stream(
|
||||
file_name=file.filename or "",
|
||||
stream=file.file,
|
||||
size_bytes=file.size,
|
||||
project=project,
|
||||
doc_type_hint=doc_type_hint,
|
||||
)
|
||||
)
|
||||
uploads.append(upload)
|
||||
job = manager.start_ingest_job([upload["upload_id"] for upload in uploads], note=note) if start_job else None
|
||||
return {"uploads": uploads, "job": job}
|
||||
except BaseException:
|
||||
for upload in reversed(uploads):
|
||||
try:
|
||||
manager.discard_upload(upload["upload_id"])
|
||||
except (IngestError, OSError):
|
||||
pass
|
||||
raise
|
||||
finally:
|
||||
if form is not None:
|
||||
await form.close()
|
||||
|
||||
|
||||
def create_http_app(
|
||||
@@ -905,31 +1092,16 @@ def create_http_app(
|
||||
return BearerTokenAuthMiddleware(app, token, bypass_paths={upload_path})
|
||||
|
||||
|
||||
async def _read_http_body(receive: Any, *, max_bytes: int | None = None) -> bytes:
|
||||
chunks: list[bytes] = []
|
||||
received = 0
|
||||
while True:
|
||||
message = await receive()
|
||||
if message.get("type") == "http.disconnect":
|
||||
break
|
||||
if message.get("type") != "http.request":
|
||||
continue
|
||||
chunk = message.get("body", b"")
|
||||
received += len(chunk)
|
||||
if max_bytes is not None and received > max_bytes:
|
||||
raise RequestBodyTooLarge(f"request body exceeds max upload size: {received} > {max_bytes}")
|
||||
chunks.append(chunk)
|
||||
if not message.get("more_body", False):
|
||||
break
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
def _max_upload_body_bytes(manager: IngestManager, upload_ticket: dict[str, Any] | None) -> int | None:
|
||||
if manager.max_file_bytes <= 0:
|
||||
return None
|
||||
max_files = int(upload_ticket.get("max_files", manager.max_files_per_job)) if upload_ticket else manager.max_files_per_job
|
||||
overhead = max(16 * 1024, max_files * 4096)
|
||||
return manager.max_file_bytes * max_files + overhead
|
||||
derived_limit = None
|
||||
if manager.max_file_bytes > 0:
|
||||
overhead = max(16 * 1024, max_files * 4096)
|
||||
derived_limit = manager.max_file_bytes * max_files + overhead
|
||||
configured_limit = int(getattr(manager, "max_request_bytes", 0) or 0)
|
||||
if configured_limit <= 0:
|
||||
return derived_limit
|
||||
return min(configured_limit, derived_limit) if derived_limit is not None else configured_limit
|
||||
|
||||
|
||||
def _header_value(scope: dict[str, Any], name: bytes) -> str:
|
||||
@@ -937,68 +1109,29 @@ def _header_value(scope: dict[str, Any], name: bytes) -> str:
|
||||
return headers.get(name.lower(), b"").decode("latin1")
|
||||
|
||||
|
||||
def _query_param(scope: dict[str, Any], name: str) -> str:
|
||||
from urllib.parse import parse_qs
|
||||
|
||||
raw = scope.get("query_string", b"")
|
||||
query = raw.decode("utf-8", errors="replace") if isinstance(raw, bytes) else str(raw)
|
||||
values = parse_qs(query).get(name, [])
|
||||
return values[0] if values else ""
|
||||
|
||||
|
||||
def _matches_any(value: str, patterns: list[str]) -> bool:
|
||||
return any(fnmatch(value, pattern) for pattern in patterns)
|
||||
if value in patterns:
|
||||
return True
|
||||
return any(
|
||||
pattern.endswith(":*") and value.startswith(pattern[:-2] + ":")
|
||||
for pattern in patterns
|
||||
)
|
||||
|
||||
|
||||
def _matches_host(value: str, patterns: list[str]) -> bool:
|
||||
if _matches_any(value, patterns):
|
||||
return True
|
||||
host_only = value
|
||||
if value.startswith("[") and "]" in value:
|
||||
host_only = value[1 : value.index("]")]
|
||||
elif ":" in value:
|
||||
host_only = value.split(":", 1)[0]
|
||||
return _matches_any(host_only, patterns)
|
||||
return _matches_any(value, patterns)
|
||||
|
||||
|
||||
def _parse_multipart_form(content_type: str, body: bytes) -> tuple[dict[str, list[str]], list[dict[str, Any]]]:
|
||||
if not content_type.lower().startswith("multipart/form-data"):
|
||||
raise IngestError("content-type must be multipart/form-data")
|
||||
message = BytesParser(policy=policy.default).parsebytes(
|
||||
b"Content-Type: "
|
||||
+ content_type.encode("latin1")
|
||||
+ b"\r\nMIME-Version: 1.0\r\n\r\n"
|
||||
+ body
|
||||
)
|
||||
if not message.is_multipart():
|
||||
raise IngestError("multipart request body is invalid")
|
||||
fields: dict[str, list[str]] = {}
|
||||
files: list[dict[str, Any]] = []
|
||||
for part in message.iter_parts():
|
||||
if part.get_content_disposition() != "form-data":
|
||||
continue
|
||||
field_name = part.get_param("name", header="content-disposition")
|
||||
if not field_name:
|
||||
continue
|
||||
payload = part.get_payload(decode=True)
|
||||
if payload is None:
|
||||
content = part.get_content()
|
||||
payload = content.encode(part.get_content_charset() or "utf-8") if isinstance(content, str) else bytes(content)
|
||||
file_name = part.get_filename()
|
||||
if file_name is None:
|
||||
fields.setdefault(field_name, []).append(payload.decode(part.get_content_charset() or "utf-8", errors="replace"))
|
||||
continue
|
||||
if field_name != "files":
|
||||
continue
|
||||
files.append({"file_name": file_name, "payload": payload})
|
||||
return fields, files
|
||||
|
||||
|
||||
def _first_field(fields: dict[str, list[str]], name: str) -> str | None:
|
||||
values = fields.get(name) or []
|
||||
if not values:
|
||||
return None
|
||||
return values[0]
|
||||
def _upload_error_payload(error: str, *, upload_ticket: dict[str, Any] | None) -> dict[str, Any]:
|
||||
payload: dict[str, Any] = {"error": error}
|
||||
if upload_ticket is not None:
|
||||
payload.update(
|
||||
{
|
||||
"ticket_consumed": True,
|
||||
"retry_action": "create_new_upload_ticket",
|
||||
}
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
async def _send_json(send: Any, status: int, payload: dict[str, Any]) -> None:
|
||||
@@ -1016,6 +1149,33 @@ async def _send_json(send: Any, status: int, payload: dict[str, Any]) -> None:
|
||||
await send({"type": "http.response.body", "body": body})
|
||||
|
||||
|
||||
async def _run_blocking(operation: Any) -> Any:
|
||||
done = threading.Event()
|
||||
result: list[Any] = []
|
||||
errors: list[BaseException] = []
|
||||
|
||||
def run() -> None:
|
||||
try:
|
||||
result.append(operation())
|
||||
except BaseException as exc:
|
||||
errors.append(exc)
|
||||
finally:
|
||||
done.set()
|
||||
|
||||
worker = threading.Thread(target=run, name="mcp-upload-copy", daemon=True)
|
||||
worker.start()
|
||||
try:
|
||||
while not done.is_set():
|
||||
await asyncio.sleep(0.005)
|
||||
except asyncio.CancelledError:
|
||||
done.wait()
|
||||
raise
|
||||
worker.join()
|
||||
if errors:
|
||||
raise errors[0]
|
||||
return result[0]
|
||||
|
||||
|
||||
def run_http_server(
|
||||
*,
|
||||
vault_path: str | Path,
|
||||
@@ -1023,6 +1183,7 @@ def run_http_server(
|
||||
profile_name: str,
|
||||
max_file_bytes: int,
|
||||
expose_paths: bool,
|
||||
allow_source_files: bool = False,
|
||||
host: str,
|
||||
port: int,
|
||||
token_env: str,
|
||||
@@ -1043,7 +1204,10 @@ def run_http_server(
|
||||
ingest_max_file_bytes: int = DEFAULT_INGEST_MAX_FILE_BYTES,
|
||||
ingest_max_chunk_bytes: int = DEFAULT_INGEST_MAX_CHUNK_BYTES,
|
||||
ingest_max_files_per_job: int = DEFAULT_INGEST_MAX_FILES_PER_JOB,
|
||||
ingest_max_request_bytes: int = DEFAULT_INGEST_MAX_REQUEST_BYTES,
|
||||
ingest_job_timeout_seconds: int = DEFAULT_INGEST_JOB_TIMEOUT_SECONDS,
|
||||
query_max_chars: int = DEFAULT_QUERY_MAX_CHARS,
|
||||
query_job_timeout_seconds: int = DEFAULT_QUERY_JOB_TIMEOUT_SECONDS,
|
||||
) -> None:
|
||||
token = read_token_from_env(token_env)
|
||||
try:
|
||||
@@ -1069,6 +1233,8 @@ def run_http_server(
|
||||
max_file_bytes=ingest_max_file_bytes,
|
||||
max_chunk_bytes=ingest_max_chunk_bytes,
|
||||
max_files_per_job=ingest_max_files_per_job,
|
||||
max_request_bytes=ingest_max_request_bytes,
|
||||
job_timeout_seconds=ingest_job_timeout_seconds,
|
||||
)
|
||||
|
||||
query_manager = None
|
||||
@@ -1079,53 +1245,72 @@ def run_http_server(
|
||||
if query_runtime_root is not None
|
||||
else project / ".runtime" / "mcp-query"
|
||||
)
|
||||
query_manager = QueryManager(
|
||||
vault_path=vault_path,
|
||||
project_root=project,
|
||||
runtime_root=query_runtime,
|
||||
opencode_path=opencode_path,
|
||||
opencode_model=query_opencode_model,
|
||||
opencode_variant=query_opencode_variant,
|
||||
max_query_chars=query_max_chars,
|
||||
)
|
||||
try:
|
||||
query_manager = QueryManager(
|
||||
vault_path=vault_path,
|
||||
project_root=project,
|
||||
runtime_root=query_runtime,
|
||||
opencode_path=opencode_path,
|
||||
opencode_model=query_opencode_model,
|
||||
opencode_variant=query_opencode_variant,
|
||||
max_query_chars=query_max_chars,
|
||||
job_timeout_seconds=query_job_timeout_seconds,
|
||||
)
|
||||
except Exception:
|
||||
if ingest_manager is not None:
|
||||
ingest_manager.close()
|
||||
raise
|
||||
|
||||
server = create_server(
|
||||
vault_path=vault_path,
|
||||
db_path=db_path,
|
||||
profile_name=profile_name,
|
||||
max_file_bytes=max_file_bytes,
|
||||
expose_paths=expose_paths,
|
||||
host=host,
|
||||
port=port,
|
||||
streamable_http_path=streamable_http_path,
|
||||
allowed_hosts=allowed_hosts,
|
||||
allowed_origins=allowed_origins,
|
||||
enable_file_assets=enable_file_assets,
|
||||
enable_ingest=enable_ingest,
|
||||
enable_query=enable_query,
|
||||
project_root=project_root,
|
||||
ingest_runtime_root=ingest_runtime_root,
|
||||
query_runtime_root=query_runtime_root,
|
||||
opencode_path=opencode_path,
|
||||
ingest_opencode_model=ingest_opencode_model,
|
||||
ingest_opencode_variant=ingest_opencode_variant,
|
||||
query_opencode_model=query_opencode_model,
|
||||
query_opencode_variant=query_opencode_variant,
|
||||
ingest_max_file_bytes=ingest_max_file_bytes,
|
||||
ingest_max_chunk_bytes=ingest_max_chunk_bytes,
|
||||
ingest_max_files_per_job=ingest_max_files_per_job,
|
||||
query_max_chars=query_max_chars,
|
||||
ingest_manager=ingest_manager,
|
||||
query_manager=query_manager,
|
||||
)
|
||||
app = create_http_app(
|
||||
server,
|
||||
token,
|
||||
ingest_manager=ingest_manager,
|
||||
allowed_hosts=allowed_hosts,
|
||||
allowed_origins=allowed_origins,
|
||||
)
|
||||
uvicorn.run(app, host=host, port=port, log_level="info")
|
||||
server = None
|
||||
try:
|
||||
server = create_server(
|
||||
vault_path=vault_path,
|
||||
db_path=db_path,
|
||||
profile_name=profile_name,
|
||||
max_file_bytes=max_file_bytes,
|
||||
expose_paths=expose_paths,
|
||||
allow_source_files=allow_source_files,
|
||||
host=host,
|
||||
port=port,
|
||||
streamable_http_path=streamable_http_path,
|
||||
allowed_hosts=allowed_hosts,
|
||||
allowed_origins=allowed_origins,
|
||||
enable_file_assets=enable_file_assets,
|
||||
enable_ingest=enable_ingest,
|
||||
enable_query=enable_query,
|
||||
project_root=project_root,
|
||||
ingest_runtime_root=ingest_runtime_root,
|
||||
query_runtime_root=query_runtime_root,
|
||||
opencode_path=opencode_path,
|
||||
ingest_opencode_model=ingest_opencode_model,
|
||||
ingest_opencode_variant=ingest_opencode_variant,
|
||||
query_opencode_model=query_opencode_model,
|
||||
query_opencode_variant=query_opencode_variant,
|
||||
ingest_max_file_bytes=ingest_max_file_bytes,
|
||||
ingest_max_chunk_bytes=ingest_max_chunk_bytes,
|
||||
ingest_max_files_per_job=ingest_max_files_per_job,
|
||||
ingest_max_request_bytes=ingest_max_request_bytes,
|
||||
ingest_job_timeout_seconds=ingest_job_timeout_seconds,
|
||||
query_max_chars=query_max_chars,
|
||||
query_job_timeout_seconds=query_job_timeout_seconds,
|
||||
ingest_manager=ingest_manager,
|
||||
query_manager=query_manager,
|
||||
)
|
||||
app = create_http_app(
|
||||
server,
|
||||
token,
|
||||
ingest_manager=ingest_manager,
|
||||
allowed_hosts=allowed_hosts,
|
||||
allowed_origins=allowed_origins,
|
||||
)
|
||||
uvicorn.run(app, host=host, port=port, log_level="info")
|
||||
finally:
|
||||
if server is not None:
|
||||
close_server_managers(server)
|
||||
else:
|
||||
for manager in (ingest_manager, query_manager):
|
||||
if manager is not None:
|
||||
manager.close()
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
@@ -1134,7 +1319,26 @@ def main(argv: list[str] | None = None) -> int:
|
||||
parser.add_argument("--db", metavar="PATH", help="SQLite file asset index path")
|
||||
parser.add_argument("--profile", default=None, help=f"bot profile name, default: {DEFAULT_PROFILE_NAME}")
|
||||
parser.add_argument("--max-file-bytes", type=int, default=None, help="maximum resource size, 0 disables the limit")
|
||||
parser.add_argument("--expose-paths", action="store_true", help="include archive/absolute paths in tool output")
|
||||
parser.add_argument(
|
||||
"--expose-paths",
|
||||
action="store_true",
|
||||
default=None,
|
||||
help="include archive/absolute paths in tool output",
|
||||
)
|
||||
source_file_group = parser.add_mutually_exclusive_group()
|
||||
source_file_group.add_argument(
|
||||
"--allow-source-files",
|
||||
dest="allow_source_files",
|
||||
action="store_true",
|
||||
default=None,
|
||||
help="allow file resource bytes to be downloaded",
|
||||
)
|
||||
source_file_group.add_argument(
|
||||
"--deny-source-files",
|
||||
dest="allow_source_files",
|
||||
action="store_false",
|
||||
help="deny file resource bytes while retaining metadata tools",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--transport",
|
||||
choices=["stdio", "streamable-http"],
|
||||
@@ -1167,6 +1371,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
profile_name=args.profile,
|
||||
max_file_bytes=args.max_file_bytes,
|
||||
expose_paths=args.expose_paths,
|
||||
allow_source_files=args.allow_source_files,
|
||||
)
|
||||
transport = args.transport or runtime_config_value("OBSIDIAN_MCP_TRANSPORT", "stdio")
|
||||
host = args.host or runtime_config_value("OBSIDIAN_MCP_HOST", DEFAULT_HTTP_HOST)
|
||||
@@ -1183,6 +1388,10 @@ def main(argv: list[str] | None = None) -> int:
|
||||
"OBSIDIAN_MCP_INGEST_MAX_FILES_PER_JOB",
|
||||
DEFAULT_INGEST_MAX_FILES_PER_JOB,
|
||||
)
|
||||
ingest_max_request_bytes = runtime_config_int(
|
||||
"OBSIDIAN_MCP_INGEST_MAX_REQUEST_BYTES",
|
||||
DEFAULT_INGEST_MAX_REQUEST_BYTES,
|
||||
)
|
||||
opencode_path = runtime_config_value("OBSIDIAN_MCP_OPENCODE_PATH", "")
|
||||
if not opencode_path:
|
||||
opencode_path = runtime_config_value("OBSIDIAN_MCP_INGEST_OPENCODE_PATH", DEFAULT_OPENCODE_PATH)
|
||||
@@ -1199,6 +1408,18 @@ def main(argv: list[str] | None = None) -> int:
|
||||
query_opencode_model = runtime_config_value("OBSIDIAN_MCP_QUERY_OPENCODE_MODEL", DEFAULT_QUERY_OPENCODE_MODEL)
|
||||
query_opencode_variant = runtime_config_value("OBSIDIAN_MCP_QUERY_OPENCODE_VARIANT", DEFAULT_QUERY_OPENCODE_VARIANT)
|
||||
query_max_chars = runtime_config_int("OBSIDIAN_MCP_QUERY_MAX_CHARS", DEFAULT_QUERY_MAX_CHARS)
|
||||
ingest_job_timeout_seconds = runtime_config_int(
|
||||
"OBSIDIAN_MCP_INGEST_JOB_TIMEOUT_SECONDS",
|
||||
DEFAULT_INGEST_JOB_TIMEOUT_SECONDS,
|
||||
)
|
||||
query_job_timeout_seconds = runtime_config_int(
|
||||
"OBSIDIAN_MCP_QUERY_JOB_TIMEOUT_SECONDS",
|
||||
DEFAULT_QUERY_JOB_TIMEOUT_SECONDS,
|
||||
)
|
||||
allowed_hosts = args.allowed_hosts if args.allowed_hosts is not None else runtime_config_list(ALLOWED_HOSTS_ENV)
|
||||
allowed_origins = (
|
||||
args.allowed_origins if args.allowed_origins is not None else runtime_config_list(ALLOWED_ORIGINS_ENV)
|
||||
)
|
||||
try:
|
||||
if transport == "stdio":
|
||||
server = create_server(
|
||||
@@ -1207,6 +1428,7 @@ def main(argv: list[str] | None = None) -> int:
|
||||
profile_name=options["profile_name"],
|
||||
max_file_bytes=options["max_file_bytes"],
|
||||
expose_paths=options["expose_paths"],
|
||||
allow_source_files=options["allow_source_files"],
|
||||
enable_file_assets=enable_file_assets,
|
||||
enable_ingest=enable_ingest,
|
||||
enable_query=enable_query,
|
||||
@@ -1221,12 +1443,18 @@ def main(argv: list[str] | None = None) -> int:
|
||||
ingest_max_file_bytes=ingest_max_file_bytes,
|
||||
ingest_max_chunk_bytes=ingest_max_chunk_bytes,
|
||||
ingest_max_files_per_job=ingest_max_files_per_job,
|
||||
ingest_max_request_bytes=ingest_max_request_bytes,
|
||||
ingest_job_timeout_seconds=ingest_job_timeout_seconds,
|
||||
query_max_chars=query_max_chars,
|
||||
query_job_timeout_seconds=query_job_timeout_seconds,
|
||||
)
|
||||
try:
|
||||
server.run(transport="stdio")
|
||||
except TypeError:
|
||||
server.run()
|
||||
try:
|
||||
server.run(transport="stdio")
|
||||
except TypeError:
|
||||
server.run()
|
||||
finally:
|
||||
close_server_managers(server)
|
||||
return 0
|
||||
if transport != "streamable-http":
|
||||
raise RuntimeError(f"unsupported MCP transport: {transport}")
|
||||
@@ -1236,12 +1464,13 @@ def main(argv: list[str] | None = None) -> int:
|
||||
profile_name=options["profile_name"],
|
||||
max_file_bytes=options["max_file_bytes"],
|
||||
expose_paths=options["expose_paths"],
|
||||
allow_source_files=options["allow_source_files"],
|
||||
host=host,
|
||||
port=port,
|
||||
token_env=token_env,
|
||||
streamable_http_path=http_path,
|
||||
allowed_hosts=args.allowed_hosts,
|
||||
allowed_origins=args.allowed_origins,
|
||||
allowed_hosts=allowed_hosts,
|
||||
allowed_origins=allowed_origins,
|
||||
enable_file_assets=enable_file_assets,
|
||||
enable_ingest=enable_ingest,
|
||||
enable_query=enable_query,
|
||||
@@ -1256,7 +1485,10 @@ def main(argv: list[str] | None = None) -> int:
|
||||
ingest_max_file_bytes=ingest_max_file_bytes,
|
||||
ingest_max_chunk_bytes=ingest_max_chunk_bytes,
|
||||
ingest_max_files_per_job=ingest_max_files_per_job,
|
||||
ingest_max_request_bytes=ingest_max_request_bytes,
|
||||
ingest_job_timeout_seconds=ingest_job_timeout_seconds,
|
||||
query_max_chars=query_max_chars,
|
||||
query_job_timeout_seconds=query_job_timeout_seconds,
|
||||
)
|
||||
except (RuntimeError, FileNotFoundError, AssetAccessError, AssetNotFoundError, AssetTooLargeError, IngestError, QueryError) as exc:
|
||||
print(f"error: {exc}", file=sys.stderr)
|
||||
@@ -1283,6 +1515,17 @@ def _normalize_public_base_url(raw_url: str | None) -> str:
|
||||
return urlunsplit((parsed.scheme, parsed.netloc, parsed.path.rstrip("/"), "", ""))
|
||||
|
||||
|
||||
def _normalize_public_upload_url(raw_url: str | None) -> str:
|
||||
value = clean_config_value(str(raw_url or "")).strip().rstrip("/")
|
||||
if not value:
|
||||
return ""
|
||||
parsed = urlsplit(value)
|
||||
if not parsed.scheme or not parsed.netloc:
|
||||
return ""
|
||||
path = parsed.path.rstrip("/") or DEFAULT_INGEST_UPLOAD_PATH
|
||||
return urlunsplit((parsed.scheme, parsed.netloc, path, "", ""))
|
||||
|
||||
|
||||
def _upload_url_from_public_base(public_base_url: str, upload_path: str) -> str | None:
|
||||
if not public_base_url:
|
||||
return None
|
||||
@@ -1290,9 +1533,24 @@ def _upload_url_from_public_base(public_base_url: str, upload_path: str) -> str
|
||||
return urlunsplit((parsed.scheme, parsed.netloc, upload_path, "", ""))
|
||||
|
||||
|
||||
def _recommended_first_steps(*, enable_file_assets: bool, enable_ingest: bool, enable_query: bool) -> list[str]:
|
||||
steps: list[str] = []
|
||||
if enable_query:
|
||||
steps.append("Use ask_wiki for ordinary wiki, business facts, or hybrid questions.")
|
||||
if enable_file_assets:
|
||||
steps.append("Use search_file_assets before reading source file resources.")
|
||||
if enable_ingest:
|
||||
steps.append("Use ingest_text_note for short text notes.")
|
||||
steps.append("Use create_ingest_upload_ticket before remote HTTP file upload.")
|
||||
if not steps:
|
||||
steps.append("No optional MCP tool groups are enabled; ask the operator to enable file-assets, query, or ingest.")
|
||||
return steps
|
||||
|
||||
|
||||
def _file_index_capabilities(vault: Path, db: Path) -> dict[str, Any]:
|
||||
registry = registry_path_for_vault(vault)
|
||||
record_count: int | None = None
|
||||
last_rebuild_at = _path_mtime_iso(db) if db.is_file() else None
|
||||
if db.is_file():
|
||||
try:
|
||||
conn = sqlite3.connect(str(db))
|
||||
@@ -1308,7 +1566,8 @@ def _file_index_capabilities(vault: Path, db: Path) -> dict[str, Any]:
|
||||
"exists": db.is_file(),
|
||||
"registry_exists": registry.is_file(),
|
||||
"record_count": record_count,
|
||||
"last_rebuild_at": _path_mtime_iso(db) if db.is_file() else None,
|
||||
"last_rebuild_at": last_rebuild_at,
|
||||
"last_index_rebuild_at": last_rebuild_at,
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user