fix(provenance): enforce tool_catalog.tool_count

src/agentrust_trace/provenance.py

@@ -66,6 +66,14 @@ def _as_object(value: Any, field: str) -> dict[str, Any]:
     return value
 
 
+def _tool_count(catalog: dict[str, Any]) -> int:
+    """Return the required catalog count as a JSON integer."""
+    value = catalog.get("tool_count")
+    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
+        raise ProvenanceError("tool_catalog.tool_count must be a non-negative integer")
+    return value
+
+
 class ToolCatalogMismatch(ProvenanceError):
     """The server offered a tool set the record does not describe.
 
@@ -344,6 +352,7 @@ def verify_record(
     catalog = _as_object(record.get("tool_catalog"), "tool_catalog")
     if not _DIGEST_RE.match(str(catalog.get("hash", ""))):
         raise ProvenanceError("tool_catalog.hash is not a sha256: digest")
+    _tool_count(catalog)
 
     # Freshness. `issued_at` has been required and type-checked since the format
     # existed, with an error message explaining that a record with no issue time
@@ -443,3 +452,9 @@ def check_tool_catalog(record: dict[str, Any], tools: list[dict[str, Any]]) -> N
             "The signature may be perfectly valid; this is about the server, not the "
             "document."
         )
+    declared_count = _tool_count(catalog)
+    if declared_count != len(tools):
+        raise ProvenanceError(
+            f"tool_catalog.tool_count declares {declared_count} tools, but the matching "
+            f"catalog contains {len(tools)}"
+        )

tests/test_provenance_tool_count.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import pytest
+
+from agentrust_trace.provenance import (
+    ProvenanceError,
+    ToolCatalogMismatch,
+    build_record,
+    check_tool_catalog,
+    sign_record,
+    verify_record,
+)
+from agentrust_trace.sign import generate_key, key_to_jwk
+
+TOOLS = [
+    {
+        "name": "search",
+        "description": "search the docs",
+        "input_schema": {"type": "object"},
+    },
+    {
+        "name": "fetch",
+        "description": "fetch a page",
+        "input_schema": {"type": "object"},
+    },
+]
+ARTIFACT = {
+    "package": "pkg:npm/%40acme/mcp-search@2.1.0",
+    "digest": "sha256:" + "a" * 64,
+}
+_MISSING = object()
+
+
+def _signed_with_count(count: object = _MISSING):
+    key = generate_key()
+    record = build_record(
+        kind="publisher-asserted",
+        publisher="did:web:acme.example",
+        tools=TOOLS,
+        artifact=ARTIFACT,
+    )
+    if count is _MISSING:
+        del record["tool_catalog"]["tool_count"]
+    else:
+        record["tool_catalog"]["tool_count"] = count
+    return sign_record(record, key), key_to_jwk(key)
+
+
+@pytest.mark.parametrize("bad_count", [_MISSING, -1, True, "2", None])
+def test_verify_record_refuses_malformed_or_missing_tool_count(bad_count) -> None:
+    record, trusted = _signed_with_count(bad_count)
+    with pytest.raises(ProvenanceError, match="tool_catalog.tool_count must be"):
+        verify_record(record, trusted)
+
+
+@pytest.mark.parametrize("bad_count", [2.0, "2", True, -1, None])
+def test_check_tool_catalog_alone_refuses_malformed_count(bad_count) -> None:
+    record, _ = _signed_with_count(bad_count)
+    with pytest.raises(ProvenanceError, match="tool_catalog.tool_count must be"):
+        check_tool_catalog(record, TOOLS)
+
+
+def test_wrong_positive_count_is_detected_when_live_catalog_is_checked() -> None:
+    record, trusted = _signed_with_count(999)
+
+    verify_record(record, trusted)
+    with pytest.raises(ProvenanceError, match="declares 999 tools") as excinfo:
+        check_tool_catalog(record, TOOLS)
+    assert not isinstance(excinfo.value, ToolCatalogMismatch)
+
+
+def test_correct_count_and_hash_still_verify() -> None:
+    record, trusted = _signed_with_count(len(TOOLS))
+    verify_record(record, trusted)
+    check_tool_catalog(record, TOOLS)
+
+
+def test_hash_mismatch_remains_a_tool_catalog_mismatch() -> None:
+    record, trusted = _signed_with_count(len(TOOLS))
+    verify_record(record, trusted)
+
+    offered = [TOOLS[0]]
+    with pytest.raises(ToolCatalogMismatch, match="about the server, not the document"):
+        check_tool_catalog(record, offered)
+
+
+def test_hash_mismatch_outranks_a_malformed_count() -> None:
+    record, _ = _signed_with_count("2")
+    with pytest.raises(ToolCatalogMismatch):
+        check_tool_catalog(record, [TOOLS[0]])