feat(inspection): stage 2 response schema validation

src/cmcp_gateway/catalog/loader.py

@@ -4,9 +4,9 @@
 
 import hashlib
 import json
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import jsonschema
 
@@ -49,6 +49,7 @@ class CatalogEntry:
     added_at: str
     approved_by: str
     catalog_exception: bool = False
+    schema_validation_mode: Literal["redact", "strict", "log"] = field(default="redact")
 
 
 @dataclass
@@ -157,6 +158,14 @@ def load_catalog(catalog_path: str, expected_hash: str | None = None) -> ToolCat
                 f"Stored: {raw['definition_hash']}, computed: {computed_def_hash}"
             )
 
+        raw_mode = raw.get("schema_validation_mode", "redact")
+        if raw_mode not in ("redact", "strict", "log"):
+            raise ConfigError(
+                f"Catalog entry '{tool_name}': invalid schema_validation_mode '{raw_mode}'; "
+                "must be 'redact', 'strict', or 'log'"
+            )
+        schema_validation_mode: Literal["redact", "strict", "log"] = raw_mode
+
         entries[tool_name] = CatalogEntry(
             tool_name=tool_name,
             server=server,
@@ -168,6 +177,7 @@ def load_catalog(catalog_path: str, expected_hash: str | None = None) -> ToolCat
             added_at=raw.get("added_at", ""),
             approved_by=raw.get("approved_by", ""),
             catalog_exception=raw.get("catalog_exception", False),
+            schema_validation_mode=schema_validation_mode,
         )
 
     computed_hash = _catalog_hash(raw_list)

src/cmcp_gateway/inspection/pipeline.py

@@ -13,10 +13,13 @@
 from __future__ import annotations
 
 import hashlib
+import json
 import re
 from dataclasses import dataclass, field
 from typing import Any
 
+import jsonschema
+
 from cmcp_gateway.catalog.loader import CatalogEntry
 
 # ── AGT components (optional — fall back gracefully) ─────────────────────────
@@ -84,6 +87,89 @@ def _stage1_size_check(response_bytes: bytes, max_bytes: int) -> StageResult:
     return StageResult(stage="size", decision="allow")
 
 
+def _stage2_schema_validation(
+    response_bytes: bytes,
+    catalog_entry: CatalogEntry,
+) -> tuple[StageResult, bytes]:
+    """
+    Stage 2: validate response against the catalog entry's approved output_schema.
+
+    Mode is per catalog entry (catalog_entry.schema_validation_mode, default "redact").
+
+    Returns a (StageResult, response_bytes) tuple where response_bytes may be
+    modified (surplus fields stripped) in redact mode.
+    """
+    output_schema = catalog_entry.approved_definition.output_schema
+    if output_schema is None:
+        return StageResult(stage="schema", decision="skip"), response_bytes
+
+    # Only validate JSON responses; non-JSON passes through
+    try:
+        payload: Any = json.loads(response_bytes)
+    except (json.JSONDecodeError, ValueError):
+        return StageResult(stage="schema", decision="allow", reason="non-JSON response; schema check skipped"), response_bytes
+
+    # Identify surplus fields at the top level only
+    surplus: list[str] = []
+    if isinstance(payload, dict) and isinstance(output_schema.get("properties"), dict):
+        allowed_props: set[str] = set(output_schema["properties"].keys())
+        surplus = [k for k in payload if k not in allowed_props]
+
+    mode = catalog_entry.schema_validation_mode
+
+    if not surplus:
+        # No surplus fields — still run jsonschema for type/required violations
+        try:
+            jsonschema.validate(payload, output_schema)
+        except jsonschema.ValidationError as exc:
+            return (
+                StageResult(
+                    stage="schema",
+                    decision="deny",
+                    reason=f"RESPONSE_SCHEMA_VIOLATION: {exc.message}",
+                ),
+                response_bytes,
+            )
+        return StageResult(stage="schema", decision="allow"), response_bytes
+
+    # Surplus fields present — mode determines action
+    if mode == "strict":
+        return (
+            StageResult(
+                stage="schema",
+                decision="deny",
+                reason="RESPONSE_SCHEMA_VIOLATION_STRICT",
+                stripped_fields=surplus,
+            ),
+            response_bytes,
+        )
+
+    if mode == "log":
+        return (
+            StageResult(
+                stage="schema",
+                decision="allow",
+                reason="surplus fields logged",
+                stripped_fields=surplus,
+            ),
+            response_bytes,
+        )
+
+    # mode == "redact" (default): strip surplus fields and return modified bytes
+    assert isinstance(payload, dict)  # guaranteed by surplus check above
+    redacted = {k: v for k, v in payload.items() if k not in surplus}
+    modified_bytes = json.dumps(redacted, separators=(",", ":"), ensure_ascii=False).encode()
+    return (
+        StageResult(
+            stage="schema",
+            decision="allow",
+            reason="surplus fields redacted",
+            stripped_fields=surplus,
+        ),
+        modified_bytes,
+    )
+
+
 def _stage4_injection_detection(
     response_text: str,
     custom_patterns: list[tuple[re.Pattern[str], str]] | None = None,
@@ -214,15 +300,24 @@ def run(
         stripped_fields: list[str] | None = None
         injection_pattern: str | None = None
         sensitivity_tags: list[str] = []
+        modified_response: bytes | None = None
 
         # Stage 1: size check
         s1 = _stage1_size_check(response_bytes, self._max_bytes)
         stage_results["size"] = s1.decision
         if s1.decision == "deny":
             deny_reasons.append(s1.reason or "size exceeded")
 
-        # Stage 2: schema validation (Phase 1 GA — issue #74; skipped here)
-        stage_results["schema"] = "skip"
+        # Stage 2: schema validation (issue #74)
+        s2, response_bytes = _stage2_schema_validation(response_bytes, catalog_entry)
+        stage_results["schema"] = s2.decision
+        if s2.decision == "deny":
+            deny_reasons.append(s2.reason or "schema violation")
+        if s2.stripped_fields:
+            stripped_fields = s2.stripped_fields
+        if s2.decision == "allow" and s2.stripped_fields and s2.reason == "surplus fields redacted":
+            # Redact mode modified the bytes — expose to caller
+            modified_response = response_bytes
 
         # Stage 3: sensitivity classification (AGT CredentialRedactor + catalog)
         try:
@@ -289,5 +384,5 @@ def run(
             injection_pattern_matched=injection_pattern,
             stage_results=stage_results,
             response_payload_hash=response_payload_hash,
-            modified_response=None,
+            modified_response=modified_response,
         )

tests/unit/test_stage2_schema.py

@@ -0,0 +1,195 @@
+"""Tests for Stage 2 response schema validation (issue #74)."""
+
+from __future__ import annotations
+
+import json
+
+from cmcp_gateway.catalog.loader import ApprovedDefinition, CatalogEntry, ServerIdentity
+from cmcp_gateway.inspection.pipeline import InspectionPipeline, _stage2_schema_validation
+
+
+def _make_entry(
+    output_schema: dict | None = None,
+    schema_validation_mode: str = "redact",
+) -> CatalogEntry:
+    return CatalogEntry(
+        tool_name="test.tool",
+        server=ServerIdentity(
+            display_name="Test",
+            url="https://test.example.com",
+            tls_fingerprint="SHA256:AAAA/BBBB==",
+            spiffe_id=None,
+            transport="http-sse",
+            rotation_mode="key-pinned",
+        ),
+        approved_definition=ApprovedDefinition(
+            description="test tool",
+            input_schema={},
+            output_schema=output_schema,
+        ),
+        definition_hash="sha256:" + "0" * 64,
+        compliance_domain="external",
+        requires_baa=False,
+        sensitivity_level="public",
+        added_at="2026-06-01T00:00:00Z",
+        approved_by="test",
+        schema_validation_mode=schema_validation_mode,  # type: ignore[arg-type]
+    )
+
+
+_SCHEMA_WITH_PROPERTIES = {
+    "type": "object",
+    "properties": {
+        "result": {"type": "string"},
+        "count": {"type": "integer"},
+    },
+}
+
+_VALID_PAYLOAD = {"result": "ok", "count": 3}
+_SURPLUS_PAYLOAD = {"result": "ok", "count": 3, "internal_id": "secret", "debug": True}
+
+
+# ── Skip when output_schema is None ──────────────────────────────────────────
+
+def test_skip_when_no_output_schema():
+    entry = _make_entry(output_schema=None)
+    result, out_bytes = _stage2_schema_validation(b'{"result": "ok"}', entry)
+    assert result.decision == "skip"
+    assert result.stage == "schema"
+
+
+def test_pipeline_records_skip_when_no_schema():
+    pipeline = InspectionPipeline()
+    entry = _make_entry(output_schema=None)
+    ir = pipeline.run("c1", entry, b'{"result": "ok"}')
+    assert ir.stage_results["schema"] == "skip"
+    assert ir.modified_response is None
+
+
+# ── Non-JSON response passes through ─────────────────────────────────────────
+
+def test_non_json_response_allows():
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES)
+    result, out_bytes = _stage2_schema_validation(b"plain text response", entry)
+    assert result.decision == "allow"
+    assert out_bytes == b"plain text response"
+
+
+def test_pipeline_non_json_passes_through():
+    pipeline = InspectionPipeline()
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES)
+    ir = pipeline.run("c1", entry, b"plain text response")
+    assert ir.final_decision == "allow"
+    assert ir.stage_results["schema"] == "allow"
+    assert ir.modified_response is None
+
+
+# ── Valid response matching schema: allow ─────────────────────────────────────
+
+def test_valid_response_allows():
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES)
+    payload_bytes = json.dumps(_VALID_PAYLOAD).encode()
+    result, out_bytes = _stage2_schema_validation(payload_bytes, entry)
+    assert result.decision == "allow"
+    assert result.stripped_fields is None or result.stripped_fields == []
+    assert out_bytes == payload_bytes
+
+
+def test_pipeline_valid_response_no_modification():
+    pipeline = InspectionPipeline()
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES)
+    ir = pipeline.run("c1", entry, json.dumps(_VALID_PAYLOAD).encode())
+    assert ir.final_decision == "allow"
+    assert ir.modified_response is None
+
+
+# ── Redact mode: strips surplus fields ───────────────────────────────────────
+
+def test_redact_mode_strips_surplus_fields():
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="redact")
+    payload_bytes = json.dumps(_SURPLUS_PAYLOAD).encode()
+    result, out_bytes = _stage2_schema_validation(payload_bytes, entry)
+    assert result.decision == "allow"
+    assert result.stripped_fields is not None
+    assert set(result.stripped_fields) == {"internal_id", "debug"}
+    redacted = json.loads(out_bytes)
+    assert "internal_id" not in redacted
+    assert "debug" not in redacted
+    assert redacted["result"] == "ok"
+    assert redacted["count"] == 3
+
+
+def test_pipeline_redact_mode_sets_modified_response():
+    pipeline = InspectionPipeline()
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="redact")
+    ir = pipeline.run("c1", entry, json.dumps(_SURPLUS_PAYLOAD).encode())
+    assert ir.final_decision == "allow"
+    assert ir.modified_response is not None
+    redacted = json.loads(ir.modified_response)
+    assert "internal_id" not in redacted
+    assert "debug" not in redacted
+    assert ir.stripped_fields is not None
+    assert set(ir.stripped_fields) == {"internal_id", "debug"}
+
+
+def test_pipeline_redact_mode_no_modified_response_when_no_surplus():
+    pipeline = InspectionPipeline()
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="redact")
+    ir = pipeline.run("c1", entry, json.dumps(_VALID_PAYLOAD).encode())
+    assert ir.modified_response is None
+
+
+# ── Strict mode: deny on surplus fields ──────────────────────────────────────
+
+def test_strict_mode_denies_on_surplus():
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="strict")
+    payload_bytes = json.dumps(_SURPLUS_PAYLOAD).encode()
+    result, out_bytes = _stage2_schema_validation(payload_bytes, entry)
+    assert result.decision == "deny"
+    assert result.reason == "RESPONSE_SCHEMA_VIOLATION_STRICT"
+    assert result.stripped_fields is not None
+    assert set(result.stripped_fields) == {"internal_id", "debug"}
+    # Response bytes unchanged in strict mode
+    assert out_bytes == payload_bytes
+
+
+def test_pipeline_strict_mode_denies():
+    pipeline = InspectionPipeline()
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="strict")
+    ir = pipeline.run("c1", entry, json.dumps(_SURPLUS_PAYLOAD).encode())
+    assert ir.final_decision == "deny"
+    assert ir.stage_results["schema"] == "deny"
+    assert "RESPONSE_SCHEMA_VIOLATION_STRICT" in (ir.deny_reason or "")
+
+
+def test_strict_mode_allows_valid_response():
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="strict")
+    payload_bytes = json.dumps(_VALID_PAYLOAD).encode()
+    result, out_bytes = _stage2_schema_validation(payload_bytes, entry)
+    assert result.decision == "allow"
+
+
+# ── Log mode: passes through but records surplus_fields ──────────────────────
+
+def test_log_mode_allows_with_surplus_logged():
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="log")
+    payload_bytes = json.dumps(_SURPLUS_PAYLOAD).encode()
+    result, out_bytes = _stage2_schema_validation(payload_bytes, entry)
+    assert result.decision == "allow"
+    assert result.stripped_fields is not None
+    assert set(result.stripped_fields) == {"internal_id", "debug"}
+    # Response bytes unmodified in log mode
+    assert out_bytes == payload_bytes
+
+
+def test_pipeline_log_mode_passes_through_unmodified():
+    pipeline = InspectionPipeline()
+    entry = _make_entry(output_schema=_SCHEMA_WITH_PROPERTIES, schema_validation_mode="log")
+    payload_bytes = json.dumps(_SURPLUS_PAYLOAD).encode()
+    ir = pipeline.run("c1", entry, payload_bytes)
+    assert ir.final_decision == "allow"
+    assert ir.stage_results["schema"] == "allow"
+    # Log mode does not set modified_response (bytes unchanged)
+    assert ir.modified_response is None
+    assert ir.stripped_fields is not None
+    assert set(ir.stripped_fields) == {"internal_id", "debug"}