shkit 1.2.0__py3-none-any.whl

healing_kit/services/dependency_graph.py

@@ -0,0 +1,141 @@
+"""Dependency graph builder and event batching for thundering herd prevention."""
+
+from collections import defaultdict, deque
+from dataclasses import dataclass, field
+
+from healing_kit.models.events import FailureEvent, RootCauseEvent
+from healing_kit.utils.error_hash import compute_error_hash
+
+
+@dataclass
+class DependencyGraph:
+    """DAG representing task and table dependencies."""
+
+    # task_key → list of upstream task_keys
+    upstream: dict[str, list[str]] = field(default_factory=lambda: defaultdict(list))
+    # task_key → list of downstream task_keys
+    downstream: dict[str, list[str]] = field(default_factory=lambda: defaultdict(list))
+    # All task keys in the graph
+    nodes: set[str] = field(default_factory=set)
+
+
+class DependencyGraphBuilder:
+    """
+    Builds a dependency graph from Databricks Jobs API task dependencies
+    and Unity Catalog lineage (table-level upstream/downstream).
+    """
+
+    def build_from_tasks(self, tasks: list[dict]) -> DependencyGraph:
+        """
+        Build DAG from a list of task definitions (from Jobs API).
+        Each task has 'task_key' and optional 'depends_on' list.
+        """
+        graph = DependencyGraph()
+
+        for task in tasks:
+            key = task.get("task_key", "")
+            graph.nodes.add(key)
+            for dep in task.get("depends_on", []):
+                dep_key = dep.get("task_key", "") if isinstance(dep, dict) else dep
+                graph.upstream[key].append(dep_key)
+                graph.downstream[dep_key].append(key)
+
+        return graph
+
+    def find_root_ancestors(self, graph: DependencyGraph, failed_keys: set[str]) -> set[str]:
+        """
+        Traverse upstream to find root ancestors of failed tasks.
+        A root ancestor is a failed task with no failed upstream dependencies.
+        """
+        roots = set()
+
+        for key in failed_keys:
+            # Walk upstream until we find a failed task with no failed parents
+            current = key
+            visited = set()
+
+            while current not in visited:
+                visited.add(current)
+                upstream_failed = [
+                    u for u in graph.upstream.get(current, [])
+                    if u in failed_keys
+                ]
+                if not upstream_failed:
+                    roots.add(current)
+                    break
+                current = upstream_failed[0]  # Follow first failed parent
+
+        return roots
+
+    def get_all_downstream(self, graph: DependencyGraph, task_key: str) -> list[str]:
+        """BFS to find all transitively downstream tasks."""
+        downstream = []
+        visited = set()
+        queue = deque([task_key])
+
+        while queue:
+            current = queue.popleft()
+            if current in visited:
+                continue
+            visited.add(current)
+
+            for child in graph.downstream.get(current, []):
+                if child not in visited:
+                    downstream.append(child)
+                    queue.append(child)
+
+        return downstream
+
+
+class EventBatcher:
+    """
+    Batches failure events by dependency graph, deduplicating cascading
+    downstream failures into single root-cause events.
+    """
+
+    def __init__(self):
+        self.graph_builder = DependencyGraphBuilder()
+
+    def deduplicate_to_root_causes(
+        self, events: list[FailureEvent], tasks: list[dict]
+    ) -> list[RootCauseEvent]:
+        """
+        Given a batch of failure events and the pipeline's task definitions,
+        identify true root causes and group derived failures.
+        """
+        if not events:
+            return []
+
+        graph = self.graph_builder.build_from_tasks(tasks)
+        failed_keys = {e.task_key for e in events}
+        roots = self.graph_builder.find_root_ancestors(graph, failed_keys)
+
+        # Map events by task_key for lookup
+        event_map = {e.task_key: e for e in events}
+
+        root_cause_events = []
+        for root_key in roots:
+            root_event = event_map.get(root_key)
+            if not root_event:
+                continue
+
+            # Find all downstream failures derived from this root
+            all_downstream = self.graph_builder.get_all_downstream(graph, root_key)
+            derived = [event_map[k] for k in all_downstream if k in event_map and k != root_key]
+
+            error_hash = compute_error_hash(
+                error_type=root_event.error_message[:50],
+                notebook_path=root_key,
+                affected_tables=[],
+            )
+
+            root_cause_events.append(RootCauseEvent(
+                root_run_id=root_event.run_id,
+                root_job_id=root_event.job_id,
+                root_task_key=root_key,
+                error_hash=error_hash,
+                derived_failures=derived,
+                total_downstream_impact=len(derived),
+            ))
+
+        return root_cause_events

healing_kit/services/diagnosis_engine.py

@@ -0,0 +1,165 @@
+"""AI Diagnosis Engine — structured LLM invocation with schema validation."""
+
+import json
+from typing import Optional
+
+from healing_kit.clients.databricks_client import DatabricksClient
+from healing_kit.models.diagnosis import ActionId, DiagnosisResponse
+from healing_kit.models.evidence import EvidencePackage
+from healing_kit.services.model_router import ModelRouter
+from healing_kit.services.token_budget import TokenBudgetEnforcer
+
+SYSTEM_PROMPT = """You are a Databricks pipeline failure diagnosis expert.
+You receive structured evidence from a failed pipeline task and must return a JSON diagnosis.
+
+RULES:
+- NEVER suggest DROP, DELETE, TRUNCATE, or ALTER on user tables.
+- NEVER return free text. Only parseable JSON.
+- confidence_score must be between 0.0 and 1.0.
+- action_id must be one of: CLUSTER_RESIZE, RETRY, SCHEMA_FIX, CODE_PATCH, ESCALATE, UNKNOWN.
+
+REQUIRED JSON OUTPUT:
+{
+    "root_cause": "specific human-readable explanation",
+    "confidence_score": 0.0-1.0,
+    "action_id": "CLUSTER_RESIZE | RETRY | SCHEMA_FIX | CODE_PATCH | ESCALATE | UNKNOWN",
+    "action_params": {},
+    "reasoning": "2-3 sentence explanation citing evidence",
+    "evidence_used": ["list of evidence sources you consulted"],
+    "diagnostic_query": "SELECT query to show bad rows (or null)",
+    "code_issue": "specific line/logic issue (or null)"
+}"""
+
+
+class DiagnosisEngine:
+    """
+    Sends enriched context to LLM, validates response schema,
+    integrates with model router and token budget.
+    """
+
+    def __init__(self, client: DatabricksClient, model_router: ModelRouter, token_budget: Optional[TokenBudgetEnforcer] = None):
+        self.client = client
+        self.router = model_router
+        self.budget = token_budget
+
+    def diagnose(self, evidence: EvidencePackage, hint_action_id: str = "UNKNOWN") -> DiagnosisResponse:
+        """
+        Invoke the appropriate LLM with structured evidence.
+        Returns validated DiagnosisResponse.
+        """
+        # Check token budget
+        model_tier = self.router.get_model_tier(ActionId(hint_action_id))
+        if self.budget and not self.budget.can_invoke_model(model_tier):
+            return DiagnosisResponse(
+                root_cause="Token budget exhausted — degraded mode active",
+                confidence_score=0.0,
+                action_id=ActionId.ESCALATE,
+                reasoning="System is in degraded mode due to token budget limits",
+            )
+
+        # Select model
+        model = self.router.get_model(ActionId(hint_action_id))
+        if model is None:
+            return DiagnosisResponse(
+                root_cause="No LLM needed for this action type",
+                confidence_score=1.0,
+                action_id=ActionId(hint_action_id),
+            )
+
+        # Build prompt
+        prompt = self._build_prompt(evidence)
+
+        # Call LLM
+        try:
+            response_text = self.client.invoke_model(
+                endpoint_name=model,
+                messages=[
+                    {"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user", "content": prompt},
+                ],
+                max_tokens=2000,
+                temperature=0.1,
+            )
+        except Exception as e:
+            return DiagnosisResponse(
+                root_cause=f"LLM invocation failed: {str(e)[:200]}",
+                confidence_score=0.0,
+                action_id=ActionId.ESCALATE,
+                reasoning="Model serving endpoint error",
+            )
+
+        # Record token usage (estimate)
+        if self.budget:
+            estimated_tokens = len(prompt.split()) + len(response_text.split())
+            self.budget.record_usage(estimated_tokens, estimated_tokens * 0.00001)
+
+        # Parse and validate
+        return self._parse_response(response_text)
+
+    def _build_prompt(self, evidence: EvidencePackage) -> str:
+        """Build the structured prompt from the evidence package."""
+        sections = [f"FAILED TASK: {evidence.task_key}", f"JOB: {evidence.job_id}", f"RUN: {evidence.run_id}"]
+
+        if evidence.driver_stdout:
+            if evidence.driver_stdout.schema_errors:
+                sections.append("SCHEMA ERRORS:\n" + "\n".join(evidence.driver_stdout.schema_errors[:5]))
+            if evidence.driver_stdout.data_source_exceptions:
+                sections.append("EXCEPTIONS:\n" + "\n".join(evidence.driver_stdout.data_source_exceptions[:10]))
+            if evidence.driver_stdout.missing_table_errors:
+                sections.append("MISSING TABLES:\n" + "\n".join(evidence.driver_stdout.missing_table_errors[:5]))
+
+        if evidence.spark_event_logs and evidence.spark_event_logs.first_errors:
+            sections.append("NOTEBOOK SOURCE / SPARK LOGS:\n" + "\n".join(evidence.spark_event_logs.first_errors[:1])[:2000])
+
+        if evidence.task_metrics:
+            m = evidence.task_metrics
+            sections.append(f"METRICS: spill={m.spill_to_disk_bytes}B, GC={m.gc_overhead_pct}%, shuffle_r={m.shuffle_read_bytes}B")
+
+        if evidence.git_context:
+            sections.append(f"GIT: {evidence.git_context.last_commit_author} — {evidence.git_context.last_commit_message}")
+
+        if evidence.cluster_events and evidence.cluster_events.termination_reason:
+            sections.append(f"CLUSTER: terminated={evidence.cluster_events.termination_reason}")
+
+        if evidence.missing_sources:
+            sections.append(f"UNAVAILABLE SOURCES: {', '.join(evidence.missing_sources)}")
+
+        return "\n\n".join(sections)
+
+    def _parse_response(self, response_text: str) -> DiagnosisResponse:
+        """Parse and validate the LLM JSON response."""
+        try:
+            cleaned = response_text.strip()
+            if cleaned.startswith("```"):
+                cleaned = cleaned.split("\n", 1)[1] if "\n" in cleaned else cleaned[3:]
+            if cleaned.endswith("```"):
+                cleaned = cleaned[:-3]
+
+            data = json.loads(cleaned.strip())
+
+            # Validate action_id
+            action_id_str = data.get("action_id", "UNKNOWN")
+            valid_actions = [a.value for a in ActionId]
+            if action_id_str not in valid_actions:
+                action_id_str = "UNKNOWN"
+
+            # Validate confidence
+            confidence = float(data.get("confidence_score", 0))
+            confidence = max(0.0, min(1.0, confidence))
+
+            return DiagnosisResponse(
+                root_cause=data.get("root_cause", "Unknown"),
+                confidence_score=confidence,
+                action_id=ActionId(action_id_str),
+                action_params=data.get("action_params", {}),
+                reasoning=data.get("reasoning", ""),
+                evidence_used=data.get("evidence_used", []),
+            )
+
+        except (json.JSONDecodeError, KeyError, ValueError):
+            return DiagnosisResponse(
+                root_cause=f"LLM response not parseable: {response_text[:200]}",
+                confidence_score=0.0,
+                action_id=ActionId.ESCALATE,
+                reasoning="Response failed JSON schema validation",
+            )

healing_kit/services/identity.py

@@ -0,0 +1,61 @@
+"""Approver authorization (#3).
+
+When a user clicks Approve/Reject, we must confirm — by governance — that the
+clicker actually has access to this workspace before honoring the action. The
+clicker is identified by their SSO login or Teams email, then verified against
+the workspace via SCIM (user exists + is active, and optionally is in an
+approvers group). If they don't have access we refuse with a clear message.
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ApproverDecision:
+    """Result of authorizing an approver."""
+
+    authorized: bool
+    identity: str
+    user_id: str = ""
+    display_name: str = ""
+    reason: str = ""
+
+
+# Message returned to the relay/handler when the clicker isn't entitled.
+ACCESS_DENIED_MESSAGE = "You don't have enough access to approve this action."
+
+
+def authorize_approver(client, identity: str, approvers_group: str = "") -> ApproverDecision:
+    """Authorize an approver by SSO login / Teams email against workspace governance.
+
+    client: a DatabricksClient (with find_user_by_identity / user_in_group).
+    identity: the clicker's login or email (verified upstream by the relay's SSO).
+    approvers_group: optional group the user must belong to (segregation of duties).
+    """
+    ident = (identity or "").strip()
+    if not ident:
+        return ApproverDecision(False, ident, reason="No approver identity supplied")
+
+    user = client.find_user_by_identity(ident)
+    if not user:
+        return ApproverDecision(False, ident, reason=ACCESS_DENIED_MESSAGE)
+
+    if user.get("active") is False:
+        return ApproverDecision(
+            False, ident, user_id=str(user.get("id", "")),
+            reason="User exists but is deactivated.",
+        )
+
+    if approvers_group and not client.user_in_group(user, approvers_group):
+        return ApproverDecision(
+            False, ident, user_id=str(user.get("id", "")),
+            display_name=user.get("displayName", ""),
+            reason=f"User is not in the '{approvers_group}' approvers group.",
+        )
+
+    return ApproverDecision(
+        True, ident,
+        user_id=str(user.get("id", "")),
+        display_name=user.get("displayName", ""),
+        reason="ok",
+    )

healing_kit/services/model_router.py

@@ -0,0 +1,52 @@
+"""Deterministic model routing by action_id."""
+
+from healing_kit.models.diagnosis import ActionId
+
+
+class ModelRouter:
+    """
+    Routes LLM calls to the appropriate model based on action_id.
+
+    Routing is fully deterministic and auditable:
+      - RETRY, CLUSTER_RESIZE → lightweight model (cost-efficient)
+      - SCHEMA_FIX, CODE_PATCH, UNKNOWN → powerful model (needs reasoning)
+      - CACHE_HIT, ESCALATE → no LLM call (zero tokens)
+    """
+
+    def __init__(self, lightweight_model: str, powerful_model: str):
+        self.lightweight_model = lightweight_model
+        self.powerful_model = powerful_model
+
+        self._routing_table = {
+            ActionId.RETRY: lightweight_model,
+            ActionId.CLUSTER_RESIZE: lightweight_model,
+            ActionId.SCHEMA_FIX: powerful_model,
+            ActionId.CODE_PATCH: powerful_model,
+            ActionId.UNKNOWN: powerful_model,
+            ActionId.CACHE_HIT: None,
+            ActionId.ESCALATE: None,
+        }
+
+    def get_model(self, action_id: ActionId) -> str | None:
+        """
+        Get the model endpoint name for a given action_id.
+        Returns None if no LLM call is needed.
+        """
+        if isinstance(action_id, str):
+            action_id = ActionId(action_id)
+        return self._routing_table.get(action_id)
+
+    def get_model_tier(self, action_id: ActionId) -> str:
+        """
+        Get the tier ('lightweight', 'powerful', or 'none') for budget checking.
+        """
+        model = self.get_model(action_id)
+        if model is None:
+            return "none"
+        if model == self.lightweight_model:
+            return "lightweight"
+        return "powerful"
+
+    def requires_llm(self, action_id: ActionId) -> bool:
+        """Check if this action requires an LLM call."""
+        return self.get_model(action_id) is not None

healing_kit/services/query_guard.py

@@ -0,0 +1,168 @@
+"""Query guard — allowlist + classification-aware masking for diagnostic SELECTs.
+
+Closes the prompt-injection -> data-exfiltration chain (#2). The LLM-emitted
+diagnostic query is:
+  1. checked read-only (single SELECT/CTE, no DDL/DML/comments/stacking),
+  2. parsed with sqlglot so every *referenced table* is matched against an
+     explicit allowlist (deny-by-default) — far stronger than keyword filtering,
+  3. inspected for sensitive columns; result values are masked BY DEFAULT and
+     only sent raw to Teams when the data classification + policy permit it.
+
+Pure logic, no Spark — unit-tested with plain strings/rows.
+"""
+
+from dataclasses import dataclass, field
+
+import sqlglot
+from sqlglot import exp
+
+from healing_kit.utils.sql_safety import is_safe_select
+
+# Evidence policies (what may leave the workspace into Teams):
+POLICY_SCHEMA_ONLY = "schema_only"      # never send values; only column names + row count
+POLICY_MASKED = "masked"                # send masked values (DEFAULT)
+POLICY_RAW_IF_ALLOWED = "raw_if_allowed"  # raw only if no sensitive column referenced
+
+# Conservative default set of sensitive column-name fragments (case-insensitive).
+DEFAULT_SENSITIVE_FRAGMENTS = (
+    "ssn", "social_security", "email", "e_mail", "phone", "mobile", "dob",
+    "birth", "address", "passport", "license", "credit", "card", "cvv",
+    "iban", "account_number", "acct_no", "salary", "password", "secret",
+    "token", "national_id", "tax_id", "mrn", "patient", "diagnosis_code",
+)
+
+
+@dataclass
+class GuardDecision:
+    """Outcome of evaluating a diagnostic query."""
+
+    allowed: bool
+    reason: str
+    tables: list[str] = field(default_factory=list)
+    columns: list[str] = field(default_factory=list)
+    sensitive_columns: list[str] = field(default_factory=list)
+    has_star: bool = False
+    # How rows must be handled if the query runs: schema_only | masked | raw
+    effective_policy: str = POLICY_MASKED
+
+
+class QueryGuard:
+    """Allowlist + masking gate for LLM-generated diagnostic queries."""
+
+    def __init__(
+        self,
+        allowed_tables,
+        sensitive_fragments=None,
+        policy: str = POLICY_MASKED,
+        default_catalog: str | None = None,
+        default_schema: str | None = None,
+        dialect: str = "databricks",
+    ):
+        # allowed_tables: iterable of "catalog.schema.table", "catalog.schema.*",
+        # "catalog.*", or "*" (allow all — discouraged). Stored lowercased.
+        self.allowed = {t.strip().lower() for t in allowed_tables if t and t.strip()}
+        self.sensitive_fragments = tuple(
+            f.lower() for f in (sensitive_fragments or DEFAULT_SENSITIVE_FRAGMENTS)
+        )
+        self.policy = policy if policy in (POLICY_SCHEMA_ONLY, POLICY_MASKED, POLICY_RAW_IF_ALLOWED) else POLICY_MASKED
+        self.default_catalog = (default_catalog or "").lower()
+        self.default_schema = (default_schema or "").lower()
+        self.dialect = dialect
+
+    # ─── table allowlisting ───
+
+    def _qualify(self, table: exp.Table) -> str:
+        catalog = (table.catalog or self.default_catalog).lower()
+        schema = (table.db or self.default_schema).lower()
+        name = (table.name or "").lower()
+        return ".".join(p for p in (catalog, schema, name) if p)
+
+    def _table_allowed(self, qualified: str) -> bool:
+        if "*" in self.allowed:
+            return True
+        if qualified in self.allowed:
+            return True
+        parts = qualified.split(".")
+        # Try progressively broader wildcards: cat.sch.*, cat.*
+        for i in range(len(parts) - 1, 0, -1):
+            if ".".join(parts[:i] + ["*"]) in self.allowed:
+                return True
+        return False
+
+    def _is_sensitive(self, col_name: str) -> bool:
+        c = col_name.lower()
+        return any(frag in c for frag in self.sensitive_fragments)
+
+    def evaluate(self, query: str) -> GuardDecision:
+        """Decide whether a query may run and how its rows must be handled."""
+        if not is_safe_select(query):
+            return GuardDecision(False, "Not a single read-only SELECT/CTE")
+
+        try:
+            parsed = sqlglot.parse_one(query, read=self.dialect)
+        except Exception as ex:  # unparseable -> reject
+            return GuardDecision(False, f"Unparseable SQL: {str(ex)[:80]}")
+
+        tables = sorted({self._qualify(t) for t in parsed.find_all(exp.Table)})
+        if not tables:
+            return GuardDecision(False, "No table referenced")
+
+        not_allowed = [t for t in tables if not self._table_allowed(t)]
+        if not_allowed:
+            return GuardDecision(
+                False,
+                f"Table(s) not on evidence allowlist: {', '.join(not_allowed)}",
+                tables=tables,
+            )
+
+        has_star = bool(list(parsed.find_all(exp.Star)))
+        columns = sorted({c.name.lower() for c in parsed.find_all(exp.Column) if c.name})
+        sensitive = sorted({c for c in columns if self._is_sensitive(c)})
+
+        # Decide the effective row policy.
+        if self.policy == POLICY_SCHEMA_ONLY:
+            effective = POLICY_SCHEMA_ONLY
+        elif self.policy == POLICY_RAW_IF_ALLOWED:
+            # SELECT * could expose unknown sensitive columns -> downgrade to masked.
+            effective = "masked" if (sensitive or has_star) else "raw"
+        else:  # POLICY_MASKED (default)
+            effective = "masked"
+
+        return GuardDecision(
+            allowed=True,
+            reason="ok",
+            tables=tables,
+            columns=columns,
+            sensitive_columns=sensitive,
+            has_star=has_star,
+            effective_policy=effective,
+        )
+
+    # ─── output handling ───
+
+    @staticmethod
+    def _mask_value(value) -> str:
+        s = str(value)
+        if len(s) <= 2:
+            return "**"
+        return s[:2] + "*" * min(len(s) - 2, 6)
+
+    def render_rows(self, columns: list[str], rows: list[list], decision: GuardDecision, max_rows: int = 8) -> str:
+        """Render query rows for Teams according to the decision's effective policy."""
+        if decision.effective_policy == POLICY_SCHEMA_ONLY:
+            return "columns: " + ", ".join(columns) + f"\n({len(rows)} row(s) — values withheld by policy)"
+
+        sensitive_set = set(decision.sensitive_columns)
+        lines = [" | ".join(columns)]
+        for row in rows[:max_rows]:
+            cells = []
+            for col, val in zip(columns, row):
+                if decision.effective_policy == "raw":
+                    cells.append(str(val)[:20])
+                else:  # masked: mask sensitive columns, lightly truncate the rest
+                    if col.lower() in sensitive_set or self._is_sensitive(col):
+                        cells.append(self._mask_value(val))
+                    else:
+                        cells.append(str(val)[:20])
+            lines.append(" | ".join(cells))
+        return "\n".join(lines)

healing_kit/services/resolution_verifier.py

@@ -0,0 +1,100 @@
+"""Resolution verifier (#5) — close the loop after a fix is applied.
+
+"Self-healing" only earns the name if an applied fix is *verified*: re-run the
+failed task, confirm it now succeeds, and roll back / escalate if it doesn't.
+This module owns that loop so the approval handler doesn't blindly mark things
+resolved. RETRY/CLUSTER_RESIZE are verified by re-running; SCHEMA_FIX/CODE_PATCH
+must supply an ``apply_fn`` (and ideally a ``rollback_fn``) and are gated so they
+never silently no-op.
+"""
+
+import time
+from dataclasses import dataclass
+
+# Action types we can verify by simply re-running the task.
+RERUNNABLE = {"RETRY", "CLUSTER_RESIZE"}
+# Action types that change code/schema and must apply a fix before re-running.
+MUTATING = {"SCHEMA_FIX", "CODE_PATCH"}
+
+
+@dataclass
+class VerificationResult:
+    action_id: str
+    applied: bool
+    verified: bool
+    rolled_back: bool
+    final_state: str
+    detail: str = ""
+
+
+class ResolutionVerifier:
+    def __init__(self, client, poll_seconds: int = 15, max_polls: int = 40):
+        self.client = client
+        self.poll_seconds = poll_seconds
+        self.max_polls = max_polls
+
+    def _poll_task_state(self, run_id: int, task_key: str) -> str:
+        """Poll runs/get until the task reaches a terminal state. Returns the
+        result_state ('SUCCESS'/'FAILED'/...) or 'TIMEOUT'."""
+        for _ in range(self.max_polls):
+            try:
+                run = self.client._get("/api/2.1/jobs/runs/get", {"run_id": int(run_id)})
+            except Exception as ex:
+                return f"ERROR: {str(ex)[:80]}"
+            for t in run.get("tasks", []):
+                if t.get("task_key") == task_key:
+                    state = t.get("state", {})
+                    life = state.get("life_cycle_state")
+                    if life in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR"):
+                        return state.get("result_state", "UNKNOWN")
+            time.sleep(self.poll_seconds)
+        return "TIMEOUT"
+
+    def verify(self, action_id: str, run_id, task_key: str, apply_fn=None, rollback_fn=None) -> VerificationResult:
+        """Apply (if needed), re-run, and verify a resolution."""
+        action = str(action_id).upper()
+
+        if action in MUTATING:
+            if apply_fn is None:
+                # Never claim success for a fix we can't actually apply.
+                return VerificationResult(action, applied=False, verified=False, rolled_back=False,
+                                          final_state="STAGED",
+                                          detail=f"{action} requires an apply function; staged for manual execution.")
+            try:
+                apply_fn()
+            except Exception as ex:
+                return VerificationResult(action, applied=False, verified=False, rolled_back=False,
+                                          final_state="APPLY_FAILED", detail=str(ex)[:120])
+            applied = True
+        elif action in RERUNNABLE:
+            applied = True
+        else:
+            return VerificationResult(action, applied=False, verified=False, rolled_back=False,
+                                      final_state="NOT_VERIFIABLE",
+                                      detail=f"{action} is not auto-verifiable (e.g. ESCALATE).")
+
+        # Re-run the failed task and watch for the outcome.
+        try:
+            self.client.repair_run(int(run_id), [task_key])
+        except Exception as ex:
+            return VerificationResult(action, applied=applied, verified=False, rolled_back=False,
+                                      final_state="RERUN_FAILED", detail=str(ex)[:120])
+
+        state = self._poll_task_state(run_id, task_key)
+        if state == "SUCCESS":
+            return VerificationResult(action, applied=applied, verified=True, rolled_back=False,
+                                      final_state="VERIFIED", detail="Task re-ran and succeeded.")
+
+        # Did not succeed — roll back mutating fixes if we can.
+        rolled_back = False
+        if action in MUTATING and rollback_fn is not None:
+            try:
+                rollback_fn()
+                rolled_back = True
+            except Exception as ex:
+                return VerificationResult(action, applied=applied, verified=False, rolled_back=False,
+                                          final_state="ROLLBACK_FAILED", detail=str(ex)[:120])
+
+        return VerificationResult(action, applied=applied, verified=False, rolled_back=rolled_back,
+                                  final_state="ESCALATE",
+                                  detail=f"Re-run ended in {state}; escalating to humans.")