[M5] Formalize root-cause label schema

agentfailbench/models.py

@@ -6,6 +6,8 @@
 
 from pydantic import BaseModel, Field
 
+from runtime.schemas.root_cause import RootCauseCode, RootCauseLabel, root_cause_label
+
 
 class TaskBlock(BaseModel):
     objective: str
@@ -21,11 +23,16 @@ class FailureBlock(BaseModel):
 
 
 class GroundTruthBlock(BaseModel):
-    root_cause: str
+    root_cause: RootCauseCode
     first_detectable_step: int
     final_failure_step: int
     expected_recovery: list[str] = Field(default_factory=list)
 
+    @property
+    def root_cause_label(self) -> RootCauseLabel:
+        """The versioned, categorized label for :attr:`root_cause`."""
+        return root_cause_label(self.root_cause)
+
 
 class RiskBlock(BaseModel):
     severity: str = "medium"

agentfailbench/registry.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 
 from agentfailbench.models import BenchmarkCaseModel
+from runtime.schemas.root_cause import RootCauseLabel
 from runtime.schemas.taxonomy import FailureCategory
 
 
@@ -20,6 +21,11 @@ class BenchmarkCase:
     def failure_category(self) -> str:
         return self.model.failure_category
 
+    @property
+    def root_cause_label(self) -> RootCauseLabel:
+        """Versioned, categorized ground-truth root-cause label for this case."""
+        return self.model.ground_truth.root_cause_label
+
     @property
     def raw(self) -> dict[str, object]:
         return self.model.to_raw()

runtime/schemas/__init__.py

@@ -1,6 +1,7 @@
 """Schema package exports."""
 
 from runtime.schemas.episode import Action, EnvObservation, EpisodeResult, TaskSpec
+from runtime.schemas.root_cause import RootCauseCode, RootCauseLabel, root_cause_label
 from runtime.schemas.taxonomy import FailureCategory, ToolFailureType
 from runtime.schemas.trace import Expectation, Observation, TraceEvent
 
@@ -11,7 +12,10 @@
     "Expectation",
     "FailureCategory",
     "Observation",
+    "RootCauseCode",
+    "RootCauseLabel",
     "TaskSpec",
     "ToolFailureType",
     "TraceEvent",
+    "root_cause_label",
 ]

runtime/schemas/root_cause.py

@@ -0,0 +1,61 @@
+"""Versioned root-cause label schema for AgentFailBench diagnosis output.
+
+``runtime.diagnosis.rules`` and case-file ``ground_truth.root_cause`` values
+previously used ad hoc free-text strings. This module gives those strings a
+closed, versioned vocabulary (:class:`RootCauseCode`) plus a small
+categorized wrapper (:class:`RootCauseLabel`) so new root causes are added
+deliberately alongside their :class:`~runtime.schemas.taxonomy.FailureCategory`.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+from pydantic import BaseModel, model_validator
+
+from runtime.schemas.taxonomy import FailureCategory
+
+ROOT_CAUSE_LABEL_SCHEMA_VERSION = "1.0"
+
+
+class RootCauseCode(StrEnum):
+    """Known root-cause identifiers. Extend as new failure suites land."""
+
+    PLAN_IDENTIFIER_SEMANTICS_CHANGED = "plan_identifier_semantics_changed"
+    TOOL_TRANSPORT_ERROR = "tool_transport_error"
+
+
+ROOT_CAUSE_CATEGORY: dict[RootCauseCode, FailureCategory] = {
+    RootCauseCode.PLAN_IDENTIFIER_SEMANTICS_CHANGED: FailureCategory.TOOL,
+    RootCauseCode.TOOL_TRANSPORT_ERROR: FailureCategory.TOOL,
+}
+
+
+class RootCauseLabel(BaseModel):
+    """A root-cause code paired with its failure category and schema version."""
+
+    code: RootCauseCode
+    category: FailureCategory
+    schema_version: str = ROOT_CAUSE_LABEL_SCHEMA_VERSION
+
+    @model_validator(mode="after")
+    def _category_matches_code(self) -> RootCauseLabel:
+        expected = ROOT_CAUSE_CATEGORY.get(self.code)
+        if expected is not None and self.category != expected:
+            raise ValueError(
+                f"root cause {self.code!r} belongs to category {expected!r}, got {self.category!r}"
+            )
+        return self
+
+
+def root_cause_label(code: RootCauseCode | str) -> RootCauseLabel:
+    """Build the :class:`RootCauseLabel` for ``code``, inferring its category.
+
+    Raises ``ValueError`` (via enum/model validation) if ``code`` is not a
+    recognized :class:`RootCauseCode`.
+    """
+    resolved = RootCauseCode(code)
+    category = ROOT_CAUSE_CATEGORY.get(resolved)
+    if category is None:
+        raise ValueError(f"no category registered for root cause {resolved!r}")
+    return RootCauseLabel(code=resolved, category=category)

tests/unit/test_root_cause.py

@@ -0,0 +1,82 @@
+"""Unit tests for the versioned root-cause label schema."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+from pydantic import ValidationError
+
+from agentfailbench.registry import CaseRegistry
+from runtime.schemas.root_cause import (
+    ROOT_CAUSE_LABEL_SCHEMA_VERSION,
+    RootCauseCode,
+    RootCauseLabel,
+    root_cause_label,
+)
+from runtime.schemas.taxonomy import FailureCategory
+
+CASE_PATH = (
+    Path(__file__).resolve().parents[2]
+    / "agentfailbench"
+    / "failures"
+    / "tool_drift"
+    / "tool-semantic-drift-001.yaml"
+)
+
+
+def test_root_cause_label_infers_category() -> None:
+    label = root_cause_label(RootCauseCode.PLAN_IDENTIFIER_SEMANTICS_CHANGED)
+    assert label.category == FailureCategory.TOOL
+    assert label.schema_version == ROOT_CAUSE_LABEL_SCHEMA_VERSION
+
+
+def test_root_cause_label_accepts_raw_string() -> None:
+    label = root_cause_label("tool_transport_error")
+    assert label.code is RootCauseCode.TOOL_TRANSPORT_ERROR
+    assert label.category == FailureCategory.TOOL
+
+
+def test_root_cause_label_rejects_mismatched_category() -> None:
+    with pytest.raises(ValidationError):
+        RootCauseLabel(
+            code=RootCauseCode.PLAN_IDENTIFIER_SEMANTICS_CHANGED,
+            category=FailureCategory.MEMORY,
+        )
+
+
+def test_root_cause_label_rejects_unknown_code() -> None:
+    with pytest.raises(ValueError):
+        root_cause_label("not_a_real_root_cause")
+
+
+def test_case_registry_validates_root_cause_on_load() -> None:
+    registry = CaseRegistry.from_yaml_file(CASE_PATH)
+    case = registry.get("tool-semantic-drift-001")
+    assert case.model.ground_truth.root_cause == RootCauseCode.PLAN_IDENTIFIER_SEMANTICS_CHANGED
+    assert case.root_cause_label == root_cause_label(
+        RootCauseCode.PLAN_IDENTIFIER_SEMANTICS_CHANGED
+    )
+
+
+def test_case_registry_rejects_unknown_root_cause(tmp_path: Path) -> None:
+    bad_case = tmp_path / "bad-case.yaml"
+    bad_case.write_text(
+        """
+case_id: bad-case-001
+task:
+  objective: update_customer_subscription
+  environment: customer_service_api
+  expected_steps: 8
+failure:
+  category: tool_semantic_drift
+  trigger_step: 4
+ground_truth:
+  root_cause: not_a_real_root_cause
+  first_detectable_step: 5
+  final_failure_step: 8
+""",
+        encoding="utf-8",
+    )
+    with pytest.raises(ValidationError):
+        CaseRegistry.from_yaml_file(bad_case)