Codex
b89d141f03
Preserve curated source-registry identifiers when rebuilding the file asset index so part-number queries such as 101049 can match drawings even when the filename omits the part number. Raise ranking for structured identifier matches above incidental substrings like timestamps, expose safe part-number fields through the MCP result payload, and allow customer-shareable engineering deliverables in the source-file MCP profile path while keeping blocked doc types enforced. Validation: python3 -m unittest discover -v; python3 -m py_compile obsidian_wiki/file_assets.py obsidian_wiki/mcp_assets.py tests/test_file_assets.py tests/test_mcp_assets.py; git diff --check; rebuilt the configured bjd vault file index; verified HTTP MCP search_file_assets and asset resource read for 101049 图纸.
247 lines
10 KiB
Python
247 lines
10 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.",
|
|
bjd_part_numbers=["101049-X0XXX0"],
|
|
series=["0.5-18"],
|
|
pitch=["0.5"],
|
|
pin_count=["18"],
|
|
)
|
|
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": json.dumps(fields.get("bjd_part_numbers", []), ensure_ascii=False),
|
|
"competitor_part_numbers": json.dumps(fields.get("competitor_part_numbers", []), ensure_ascii=False),
|
|
"series": json.dumps(fields.get("series", []), ensure_ascii=False),
|
|
"pitch": json.dumps(fields.get("pitch", []), ensure_ascii=False),
|
|
"pin_count": json.dumps(fields.get("pin_count", []), ensure_ascii=False),
|
|
"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 _write_profile_file(self):
|
|
data = {
|
|
"version": 2,
|
|
"profiles": {
|
|
"bjd-general-bot": {
|
|
"allowed_visibility": ["public", "internal"],
|
|
"allowed_data_classification": [
|
|
"public",
|
|
"internal",
|
|
"technical-confidential",
|
|
],
|
|
"allowed_business_domains": ["products", "engineering"],
|
|
"blocked_doc_types": ["order_data", "evidence_media"],
|
|
"include_source_files": False,
|
|
}
|
|
},
|
|
}
|
|
(self.vault / "_meta" / "bot-profiles.json").write_text(json.dumps(data), encoding="utf-8")
|
|
|
|
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_exposes_safe_identifier_fields(self):
|
|
rows = search_file_assets(self.vault, "101049 图纸", db_path=self.db_path)
|
|
|
|
self.assertEqual([row["file_id"] for row in rows], [self.drawing_id])
|
|
self.assertEqual(rows[0]["bjd_part_numbers"], ["101049-X0XXX0"])
|
|
self.assertEqual(rows[0]["series"], ["0.5-18"])
|
|
self.assertEqual(rows[0]["pitch"], ["0.5"])
|
|
self.assertEqual(rows[0]["pin_count"], ["18"])
|
|
|
|
def test_customer_shareable_engineering_files_work_with_source_file_profile(self):
|
|
self._write_profile_file()
|
|
shareable_id = "sha256:" + "d" * 64
|
|
self._write_asset_file("drawings/customer-shareable.pdf", b"%PDF customer drawing")
|
|
self._insert_asset(
|
|
shareable_id,
|
|
archive_path="_sources/archive/drawings/customer-shareable.pdf",
|
|
original_name="北极大-0.5-18客户图.pdf",
|
|
doc_type="drawing",
|
|
source_type="pdf",
|
|
product_lines=["FPC"],
|
|
business_categories=["products/fpc", "drawings/customer-drawings"],
|
|
data_classification="customer-shareable",
|
|
visibility="public",
|
|
summary="FPC 0.5-18 customer-shareable drawing.",
|
|
bjd_part_numbers=["101049-X0XXX0"],
|
|
series=["0.5-18"],
|
|
)
|
|
|
|
rows = search_file_assets(self.vault, "101049 图纸", db_path=self.db_path, profile_name="bjd-general-bot")
|
|
|
|
self.assertEqual({row["file_id"] for row in rows}, {shareable_id, self.drawing_id})
|
|
|
|
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()
|