Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 58 additions & 8 deletions mcp-server/src/streamops_mcp/agent/schemas/diagnosis.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from enum import StrEnum

from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, model_validator


class Confidence(StrEnum):
Expand All @@ -29,7 +29,9 @@ class SourceRecord(BaseModel):
"""

source_id: str = Field(description="Unique identifier for this source (e.g., 'src-001')")
tool_name: str = Field(description="MCP tool that produced this data (e.g., 'query_flink_jobs')")
tool_name: str = Field(
description="MCP tool that produced this data (e.g., 'query_flink_jobs')"
)
retrieved_at: str = Field(description="ISO-8601 timestamp when the data was retrieved")
raw_output: str = Field(description="Verbatim or summarized tool output")

Expand All @@ -44,7 +46,9 @@ class ClaimRecord(BaseModel):
"""

claim_id: str = Field(description="Unique identifier for this claim (e.g., 'C01')")
text: str = Field(description="The factual claim (e.g., 'Consumer lag is 45,000 on partition 2')")
text: str = Field(
description="The factual claim (e.g., 'Consumer lag is 45,000 on partition 2')"
)
source_id: str = Field(description="References SourceRecord.source_id that produced this claim")
confidence: Confidence = Field(
description=(
Expand Down Expand Up @@ -86,7 +90,9 @@ class AffectedComponent(BaseModel):
name: str = Field(description="Component identifier (e.g., 'kafka-consumer', 'flink-operator')")
role: str = Field(description="What this component does in the pipeline")
status: str = Field(description="Current state: healthy, degraded, failing, unknown")
evidence: str = Field(description="Specific metric or log entry supporting the status assessment")
evidence: str = Field(
description="Specific metric or log entry supporting the status assessment"
)


class RootCause(BaseModel):
Expand All @@ -111,7 +117,9 @@ class DiagnosisReport(BaseModel):

"""

anomaly_type: str = Field(description="Category: latency_spike, throughput_drop, backpressure, checkpoint_failure, memory_pressure, error_burst")
anomaly_type: str = Field(
description="Category: latency_spike, throughput_drop, backpressure, checkpoint_failure, memory_pressure, error_burst"
)
detected_at: str = Field(description="ISO-8601 timestamp when the anomaly was first detected")
sources: list[SourceRecord] = Field(
default_factory=list,
Expand All @@ -129,10 +137,52 @@ class DiagnosisReport(BaseModel):
description="Components involved in or affected by the anomaly"
)
root_cause: RootCause = Field(description="Identified or suspected root cause")
tools_used: list[str] = Field(
description="MCP tools called during investigation (audit trail)"
)
tools_used: list[str] = Field(description="MCP tools called during investigation (audit trail)")
raw_evidence: list[str] = Field(
default_factory=list,
description="Key data points collected during investigation",
)

@model_validator(mode="after")
def _check_attribution_integrity(self):
"""Enforce claim-source attribution integrity so nothing is untraceable.

Every sourced claim must reference a real source, and every conflict must
reference real claims. This prevents the "attribution lost" failure the
claim-source pattern exists to stop: a dangling source_id would leave a
finding with no traceable origin. UNSOURCED claims are exempt, since they
explicitly have no backing data source.
"""
source_ids = [s.source_id for s in self.sources]
dup_sources = sorted({sid for sid in source_ids if source_ids.count(sid) > 1})
if dup_sources:
raise ValueError(
f"Duplicate source_id(s), references would be ambiguous: {dup_sources}"
)
source_id_set = set(source_ids)

claim_ids = [c.claim_id for c in self.claims]
dup_claims = sorted({cid for cid in claim_ids if claim_ids.count(cid) > 1})
if dup_claims:
raise ValueError(
f"Duplicate claim_id(s), conflict references would be ambiguous: {dup_claims}"
)
claim_id_set = set(claim_ids)

for claim in self.claims:
if claim.confidence == Confidence.UNSOURCED:
continue
if claim.source_id not in source_id_set:
raise ValueError(
f"Claim '{claim.claim_id}' references unknown source_id "
f"'{claim.source_id}'; attribution would be untraceable"
)

for conflict in self.conflicts:
for cid in (conflict.claim_a_id, conflict.claim_b_id):
if cid not in claim_id_set:
raise ValueError(
f"Conflict '{conflict.conflict_id}' references unknown claim_id '{cid}'"
)

return self
128 changes: 115 additions & 13 deletions mcp-server/tests/test_schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import json

import pytest
from pydantic import ValidationError

from streamops_mcp.agent.schemas import (
ClaimRecord,
Expand All @@ -20,7 +21,6 @@


class TestDiagnosisReport:

def test_valid_diagnosis(self):
# Arrange
data = {
Expand Down Expand Up @@ -113,7 +113,6 @@ def test_round_trip_serialization(self):


class TestIncidentReport:

def test_valid_incident(self):
# Arrange
data = {
Expand Down Expand Up @@ -163,7 +162,6 @@ def test_json_schema_includes_severity_enum(self):


class TestRecommendedAction:

def test_action_with_downtime(self):
# Arrange
action = RecommendedAction(
Expand Down Expand Up @@ -191,7 +189,6 @@ def test_action_no_downtime(self):


class TestSourceRecord:

def test_valid_source(self):
# Arrange + Act
source = SourceRecord(
Expand All @@ -207,7 +204,6 @@ def test_valid_source(self):


class TestClaimRecord:

def test_valid_claim(self):
# Arrange + Act
claim = ClaimRecord(
Expand Down Expand Up @@ -242,8 +238,116 @@ def test_claim_references_source(self):
assert claim.confidence == Confidence.MEDIUM


class TestConflictRecord:
class TestAttributionIntegrity:
"""DiagnosisReport enforces claim-source and conflict referential integrity (#81)."""

@staticmethod
def _report(sources, claims, conflicts):
return dict(
anomaly_type="latency_spike",
detected_at="2026-06-18T15:00:00Z",
affected_components=[],
root_cause=RootCause(summary="s", confidence="low", reasoning="r"),
tools_used=["query_metrics"],
sources=sources,
claims=claims,
conflicts=conflicts,
)

def test_valid_attribution_passes(self):
# Arrange
sources = [
SourceRecord(
source_id="src-001",
tool_name="query_metrics",
retrieved_at="2026-06-18T15:00:00Z",
raw_output="lag=45000",
)
]
claims = [
ClaimRecord(
claim_id="C01",
text="Consumer lag is 45000",
source_id="src-001",
confidence=Confidence.HIGH,
),
ClaimRecord(
claim_id="C02",
text="Backpressure rising",
source_id="src-001",
confidence=Confidence.MEDIUM,
),
]
conflicts = [
ConflictRecord(
conflict_id="conf-001", topic="lag trend", claim_a_id="C01", claim_b_id="C02"
)
]

# Act
report = DiagnosisReport(**self._report(sources, claims, conflicts))

# Assert
assert len(report.claims) == 2

def test_claim_referencing_unknown_source_rejected(self):
# Arrange
sources = [
SourceRecord(
source_id="src-001", tool_name="query_metrics", retrieved_at="t", raw_output="o"
)
]
claims = [
ClaimRecord(claim_id="C01", text="x", source_id="src-999", confidence=Confidence.HIGH)
]

# Act / Assert
with pytest.raises(ValidationError, match="unknown source_id"):
DiagnosisReport(**self._report(sources, claims, []))

def test_unsourced_claim_needs_no_source(self):
# Arrange: an UNSOURCED claim legitimately has no backing source
claims = [
ClaimRecord(
claim_id="C01", text="speculative", source_id="", confidence=Confidence.UNSOURCED
)
]

# Act
report = DiagnosisReport(**self._report([], claims, []))

# Assert
assert report.claims[0].confidence == Confidence.UNSOURCED

def test_conflict_referencing_unknown_claim_rejected(self):
# Arrange
sources = [
SourceRecord(source_id="src-001", tool_name="t", retrieved_at="t", raw_output="o")
]
claims = [
ClaimRecord(claim_id="C01", text="x", source_id="src-001", confidence=Confidence.HIGH)
]
conflicts = [
ConflictRecord(conflict_id="conf-001", topic="t", claim_a_id="C01", claim_b_id="C99")
]

# Act / Assert
with pytest.raises(ValidationError, match="unknown claim_id"):
DiagnosisReport(**self._report(sources, claims, conflicts))

def test_duplicate_source_id_rejected(self):
# Arrange
sources = [
SourceRecord(source_id="src-001", tool_name="t", retrieved_at="t", raw_output="a"),
SourceRecord(source_id="src-001", tool_name="t", retrieved_at="t", raw_output="b"),
]

# Act / Assert
with pytest.raises(ValidationError, match="Duplicate source_id"):
DiagnosisReport(**self._report(sources, [], []))


class TestConflictRecord:
def test_valid_conflict(self):
# Arrange + Act
conflict = ConflictRecord(
Expand Down Expand Up @@ -286,7 +390,6 @@ def test_conflict_with_notes(self):


class TestDiagnosisWithAttribution:

def test_full_diagnosis_with_claims_and_sources(self):
# Arrange
data = {
Expand Down Expand Up @@ -449,7 +552,6 @@ def test_round_trip_with_attribution(self):


class TestConfidence:

def test_confidence_enum_values(self):
# Assert
assert Confidence.HIGH.value == "HIGH"
Expand Down Expand Up @@ -565,15 +667,16 @@ def test_mixed_confidence_claims_in_diagnosis(self):

# Act
high_claims = [c for c in report.claims if c.confidence == Confidence.HIGH]
low_claims = [c for c in report.claims if c.confidence in (Confidence.LOW, Confidence.UNSOURCED)]
low_claims = [
c for c in report.claims if c.confidence in (Confidence.LOW, Confidence.UNSOURCED)
]

# Assert
assert len(high_claims) == 1
assert len(low_claims) == 2


class TestIncidentReportLowConfidence:

def test_incident_with_low_confidence_claims(self):
# Arrange
data = {
Expand Down Expand Up @@ -630,7 +733,6 @@ def test_incident_defaults_to_empty_low_confidence(self):


class TestDraftOnlyContract:

def test_requires_human_approval_defaults_true(self):
# Arrange
data = {
Expand Down Expand Up @@ -683,7 +785,6 @@ def test_requires_human_approval_in_json_schema(self):


class TestMonitorToDiagnosticHandoff:

def test_valid_handoff(self):
# Arrange + Act
handoff = MonitorToDiagnosticHandoff(
Expand All @@ -699,6 +800,7 @@ def test_truncates_oversized_context(self, monkeypatch):
# Arrange
monkeypatch.setenv("STREAMOPS_AGENT_HANDOFF_MAX_CONTEXT_CHARS", "100")
from streamops_mcp.config import StreamOpsConfig

test_config = StreamOpsConfig()
monkeypatch.setattr("streamops_mcp.agent.schemas.handoff.config", test_config)

Expand Down Expand Up @@ -729,7 +831,6 @@ def test_round_trip_serialization(self):


class TestDiagnosticToReportHandoff:

def test_valid_handoff(self):
# Arrange
diagnosis = DiagnosisReport(
Expand Down Expand Up @@ -761,6 +862,7 @@ def test_rejects_oversized_diagnosis(self, monkeypatch):
# Arrange
monkeypatch.setenv("STREAMOPS_AGENT_HANDOFF_MAX_CONTEXT_CHARS", "50")
from streamops_mcp.config import StreamOpsConfig

test_config = StreamOpsConfig()
monkeypatch.setattr("streamops_mcp.agent.schemas.handoff.config", test_config)

Expand Down
Loading