nac
76a480a0b7
Add a stdio MCP server for searching and reading safe engineering source files from the wiki file asset index. The implementation uses registry-backed ACL checks, engineering-file whitelisting, asset:// resource URIs, read-only SQLite access, CLI entrypoints, optional MCP packaging, and unittest coverage for search, metadata, resource reads, blocked sources, bot ACLs, and size limits. Validation: python3 -m unittest discover -v; python3 -m py_compile obsidian_wiki/config.py obsidian_wiki/file_assets.py obsidian_wiki/mcp_assets.py obsidian_wiki/mcp_server.py obsidian_wiki/cli.py tests/test_mcp_assets.py; real BJD vault search/read smoke tests.
192 lines
7.5 KiB
Python
192 lines
7.5 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from obsidian_wiki.file_assets import connect, init_db, upsert_asset
|
|
from obsidian_wiki.mcp_assets import (
|
|
AssetAccessError,
|
|
AssetTooLargeError,
|
|
asset_uri_for_file_id,
|
|
file_id_from_asset_uri,
|
|
get_file_asset_metadata,
|
|
read_file_asset_resource,
|
|
search_file_assets,
|
|
)
|
|
|
|
|
|
class McpAssetTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tempdir = tempfile.TemporaryDirectory()
|
|
self.vault = Path(self.tempdir.name)
|
|
(self.vault / "_meta").mkdir()
|
|
self.db_path = self.vault / "_meta" / "file-assets.sqlite"
|
|
self.registry_path = self.vault / "_meta" / "source-registry.jsonl"
|
|
init_db(self.db_path)
|
|
|
|
self.drawing_id = "sha256:" + "a" * 64
|
|
self.order_id = "sha256:" + "b" * 64
|
|
self.complaint_id = "sha256:" + "c" * 64
|
|
self._write_asset_file("drawings/a.pdf", b"%PDF allowed drawing")
|
|
self._write_asset_file("orders/order.xlsx", b"customer price data")
|
|
self._write_asset_file("quality/photo.png", b"complaint evidence")
|
|
|
|
self._insert_asset(
|
|
self.drawing_id,
|
|
archive_path="_sources/archive/drawings/a.pdf",
|
|
original_name="B3501_11101c_X2.pdf",
|
|
doc_type="drawing",
|
|
source_type="pdf",
|
|
product_lines=["BTB"],
|
|
business_categories=["products/btb", "drawings/customer"],
|
|
data_classification="technical-confidential",
|
|
visibility="internal",
|
|
summary="B3501 customer drawing.",
|
|
)
|
|
self._insert_asset(
|
|
self.order_id,
|
|
archive_path="_sources/archive/orders/order.xlsx",
|
|
original_name="订单项.xlsx",
|
|
doc_type="order_data",
|
|
source_type="spreadsheet",
|
|
business_categories=["orders/raw-excel"],
|
|
data_classification="commercial-confidential",
|
|
visibility="pii",
|
|
summary="Raw customer order lines.",
|
|
)
|
|
self._insert_asset(
|
|
self.complaint_id,
|
|
archive_path="_sources/archive/quality/photo.png",
|
|
original_name="客户反馈来料标签错误图片.png",
|
|
doc_type="evidence_media",
|
|
source_type="image",
|
|
business_categories=["quality/customer-complaints"],
|
|
data_classification="quality-confidential",
|
|
visibility="internal",
|
|
summary="Raw complaint evidence image.",
|
|
)
|
|
|
|
def tearDown(self):
|
|
self.tempdir.cleanup()
|
|
|
|
def _write_asset_file(self, rel_path, data):
|
|
path = self.vault / "_sources" / "archive" / rel_path
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_bytes(data)
|
|
|
|
def _insert_asset(self, file_id, **fields):
|
|
record = {
|
|
"record_type": "source_archive",
|
|
"source_id": file_id,
|
|
"project": "bjd",
|
|
"title": Path(fields["original_name"]).stem,
|
|
"size_bytes": len((self.vault / fields["archive_path"]).read_bytes()),
|
|
"related_pages": ["projects/bjd/references/example.md"],
|
|
"allowed_bots": ["bjd-general-bot"],
|
|
**fields,
|
|
}
|
|
with self.registry_path.open("a", encoding="utf-8") as fh:
|
|
fh.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
asset = {
|
|
"file_id": file_id,
|
|
"sha256": file_id.removeprefix("sha256:"),
|
|
"archive_path": fields["archive_path"],
|
|
"original_name": fields["original_name"],
|
|
"file_ext": Path(fields["original_name"]).suffix.lower(),
|
|
"mime_type": "application/pdf",
|
|
"source_type": fields["source_type"],
|
|
"doc_type": fields["doc_type"],
|
|
"project": "bjd",
|
|
"title": record["title"],
|
|
"summary": fields["summary"],
|
|
"visibility": fields["visibility"],
|
|
"business_categories": json.dumps(fields.get("business_categories", []), ensure_ascii=False),
|
|
"product_lines": json.dumps(fields.get("product_lines", []), ensure_ascii=False),
|
|
"customers": "[]",
|
|
"competitors": "[]",
|
|
"standards": "[]",
|
|
"market_tags": "[]",
|
|
"bjd_part_numbers": "[]",
|
|
"competitor_part_numbers": "[]",
|
|
"series": "[]",
|
|
"pitch": "[]",
|
|
"pin_count": "[]",
|
|
"related_pages": json.dumps(record["related_pages"], ensure_ascii=False),
|
|
"catalog_paths": "[]",
|
|
"extracted_text": None,
|
|
"classification_confidence": 0.9,
|
|
"classification_evidence": "[]",
|
|
"needs_review": 0,
|
|
"size_bytes": record["size_bytes"],
|
|
"modified_at": None,
|
|
"archived_at": "2026-07-05T00:00:00Z",
|
|
}
|
|
with connect(self.db_path) as conn:
|
|
upsert_asset(conn, asset, status="active")
|
|
|
|
def test_search_returns_allowed_engineering_assets_only(self):
|
|
rows = search_file_assets(self.vault, "B3501 drawing", db_path=self.db_path)
|
|
|
|
self.assertEqual([row["file_id"] for row in rows], [self.drawing_id])
|
|
self.assertEqual(rows[0]["resource_uri"], asset_uri_for_file_id(self.drawing_id))
|
|
self.assertNotIn("absolute_path", rows[0])
|
|
|
|
def test_search_blocks_sensitive_order_and_complaint_sources(self):
|
|
self.assertEqual(search_file_assets(self.vault, "订单", db_path=self.db_path), [])
|
|
self.assertEqual(search_file_assets(self.vault, "来料标签错误", db_path=self.db_path), [])
|
|
|
|
def test_asset_uri_round_trips_file_id(self):
|
|
uri = asset_uri_for_file_id(self.drawing_id)
|
|
|
|
self.assertEqual(uri, "asset://source/sha256/" + "a" * 64)
|
|
self.assertEqual(file_id_from_asset_uri(uri), self.drawing_id)
|
|
|
|
def test_resource_read_rechecks_acl_and_returns_file_bytes(self):
|
|
data = read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.drawing_id),
|
|
db_path=self.db_path,
|
|
)
|
|
|
|
self.assertEqual(data["name"], "B3501_11101c_X2.pdf")
|
|
self.assertEqual(data["mime_type"], "application/pdf")
|
|
self.assertEqual(data["bytes"], b"%PDF allowed drawing")
|
|
|
|
def test_resource_read_rejects_blocked_asset_even_by_direct_uri(self):
|
|
with self.assertRaises(AssetAccessError):
|
|
read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.order_id),
|
|
db_path=self.db_path,
|
|
)
|
|
|
|
def test_resource_read_enforces_max_size(self):
|
|
with self.assertRaises(AssetTooLargeError):
|
|
read_file_asset_resource(
|
|
self.vault,
|
|
asset_uri_for_file_id(self.drawing_id),
|
|
db_path=self.db_path,
|
|
max_file_bytes=5,
|
|
)
|
|
|
|
def test_metadata_uses_safe_fields(self):
|
|
metadata = get_file_asset_metadata(self.vault, self.drawing_id, db_path=self.db_path)
|
|
|
|
self.assertEqual(metadata["file_id"], self.drawing_id)
|
|
self.assertEqual(metadata["doc_type"], "drawing")
|
|
self.assertEqual(metadata["resource_uri"], asset_uri_for_file_id(self.drawing_id))
|
|
self.assertNotIn("absolute_path", metadata)
|
|
|
|
def test_registry_allowed_bots_are_enforced_without_profile_file(self):
|
|
with self.assertRaises(AssetAccessError):
|
|
get_file_asset_metadata(
|
|
self.vault,
|
|
self.drawing_id,
|
|
db_path=self.db_path,
|
|
profile_name="public-product-bot",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|