shkit 1.2.0__py3-none-any.whl

healing_kit/runtime/approval.py

@@ -0,0 +1,141 @@
+"""Approval orchestrator (#3 + #5) — called by the approval_handler notebook /
+relay when a user clicks Approve or Reject.
+
+Flow:
+  1. Validate the HMAC token (signature + 15-min TTL).
+  2. Enforce one-time use (token signature recorded in approval_tokens).
+  3. Authorize the approver by governance — resolve their SSO login / Teams email
+     and confirm they have workspace access via SCIM; refuse otherwise (#3).
+  4. On approve: run the ResolutionVerifier (re-run + verify + rollback) so we
+     only mark RESOLVED when the fix is actually verified (#5).
+"""
+
+import hashlib
+from datetime import datetime
+
+from healing_kit.auth import build_client, resolve_auth_from_dbutils
+from healing_kit.services.identity import ACCESS_DENIED_MESSAGE, authorize_approver
+from healing_kit.services.resolution_verifier import ResolutionVerifier
+from healing_kit.utils.hmac_tokens import validate_token
+from healing_kit.utils.sql_safety import sql_lit
+
+
+def _w(dbutils, name, default=""):
+    dbutils.widgets.text(name, default)
+    return dbutils.widgets.get(name)
+
+
+def _resolve_secret(dbutils, scope, host):
+    try:
+        return dbutils.secrets.get(scope=scope, key="hmac_key").encode()
+    except Exception:
+        return hashlib.sha256(("healing_kit::" + host).encode()).digest()
+
+
+def _audit(spark, catalog, schema, cols):
+    try:
+        spark.sql(f"INSERT INTO {catalog}.{schema}.healing_audit_log VALUES ("
+                  + ", ".join(sql_lit(c) for c in cols) + ")")
+    except Exception as ex:
+        print(f"  audit write failed: {str(ex)[:80]}")
+
+
+def _claim_token_once(spark, catalog, schema, jti, run_id) -> bool:
+    """Record the token signature; return False if it was already used."""
+    table = f"{catalog}.{schema}.approval_tokens"
+    now = datetime.utcnow().isoformat()
+    try:
+        used = spark.sql(f"SELECT count(*) AS n FROM {table} WHERE jti = {sql_lit(jti)}").collect()[0]["n"]
+        if used and used > 0:
+            return False
+        spark.sql(f"INSERT INTO {table} VALUES ({sql_lit(jti)}, {sql_lit(run_id)}, {sql_lit(now)})")
+        return True
+    except Exception as ex:
+        # If the table is missing we fail closed (treat as unusable) rather than
+        # silently allowing replay.
+        print(f"  one-time-use check failed: {str(ex)[:80]}")
+        return False
+
+
+def process_approval(spark, dbutils):
+    """Entry point called by the approval_handler notebook driver."""
+    token = _w(dbutils, "token", "")
+    action = _w(dbutils, "action", "")  # approve | reject
+    task_key = _w(dbutils, "task", "")
+    catalog = _w(dbutils, "catalog", "dev_catalog")
+    schema = _w(dbutils, "schema", "healing_schema")
+    secret_scope = _w(dbutils, "secret_scope", "healing_kit")
+    # Identity is supplied by the authenticated relay (SSO) or the clicker's Teams email.
+    approver_identity = _w(dbutils, "approver_identity", "")
+    approvers_group = _w(dbutils, "approvers_group", "")
+
+    auth = resolve_auth_from_dbutils(dbutils, secret_scope)
+    client = build_client(auth)
+    host = auth.config.host
+    secret = _resolve_secret(dbutils, secret_scope, host)
+    now = datetime.utcnow().isoformat()
+
+    # 1. Validate token.
+    try:
+        payload = validate_token(token, secret)
+    except Exception as e:
+        _audit(spark, catalog, schema,
+               [hashlib.sha256(token.encode()).hexdigest()[:16], "unknown", "unknown", "",
+                f"Token validation failed: {str(e)[:100]}", "REJECTED", 0, "none", "",
+                False, False, None, 0, 0, 0, now])
+        return f"REJECTED: {e}"
+
+    run_id = payload.run_id
+    action_id = payload.action_id
+
+    # 2. One-time use.
+    jti = hashlib.sha256(token.encode()).hexdigest()
+    if not _claim_token_once(spark, catalog, schema, jti, run_id):
+        return "REJECTED: token already used or unusable"
+
+    # 3. Authorize the approver by governance (#3).
+    identity = approver_identity or payload.approver_email
+    decision = authorize_approver(client, identity, approvers_group=approvers_group)
+    if not decision.authorized:
+        _audit(spark, catalog, schema,
+               [jti[:16], run_id, "", "", f"Approver denied ({identity}): {decision.reason}",
+                "ACCESS_DENIED", 0, "approval_handler", "", False, False, None, 0, 0, 0, now])
+        print(f"ACCESS DENIED for '{identity}': {decision.reason}")
+        return ACCESS_DENIED_MESSAGE if not decision.reason else decision.reason
+
+    approver = decision.identity
+    print(f"Approver authorized: {approver} ({decision.display_name})")
+
+    # 4a. Reject path.
+    if action == "reject":
+        spark.sql(f"UPDATE {catalog}.{schema}.healing_state SET current_status = 'ESCALATED', "
+                  f"resolved_by = {sql_lit(approver)}, updated_at = {sql_lit(now)} "
+                  f"WHERE run_id = {sql_lit(run_id)}")
+        _audit(spark, catalog, schema,
+               [jti[:16], run_id, "", "", f"Rejected by {approver}", action_id, 0,
+                "approval_handler", "", False, True, "user_rejected", 0, 0, 0, now])
+        return f"REJECTED by {approver}: {task_key}"
+
+    # 4b. Approve path → verify the fix.
+    verifier = ResolutionVerifier(client)
+    # apply_fn/rollback_fn for CODE_PATCH/SCHEMA_FIX are wired once fix-generation lands;
+    # until then those actions return STAGED (never falsely marked resolved).
+    result = verifier.verify(action_id, run_id, task_key)
+    if result.verified:
+        status = "RESOLVED"
+    elif result.final_state in ("ESCALATE", "STAGED"):
+        status = "ESCALATED"
+    else:
+        status = "EXECUTING"
+
+    spark.sql(f"UPDATE {catalog}.{schema}.healing_state SET current_status = {sql_lit(status)}, "
+              f"resolved_by = {sql_lit(approver)}, updated_at = {sql_lit(now)} "
+              f"WHERE run_id = {sql_lit(run_id)}")
+    _audit(spark, catalog, schema,
+           [jti[:16], run_id, "", "", f"Approved by {approver}: {result.final_state} ({result.detail})",
+            action_id, float(getattr(payload, "confidence", 0) or 0), "approval_handler", "",
+            bool(result.verified), False, None, 0, 0, 0, now])
+
+    summary = f"APPROVED by {approver} | {action_id} | {result.final_state} | verified={result.verified}"
+    print(summary)
+    return summary

healing_kit/runtime/maintenance.py

@@ -0,0 +1,52 @@
+"""Audit retention / PII maintenance (#4) — called by the audit_retention notebook.
+
+The audit log is append-only and may contain sensitive text copied from error
+logs. This job anonymizes the free-text fields on rows older than the retention
+window (kept for forensics but scrubbed of payload), and expires stale one-time
+approval tokens. It does not DELETE audit rows (append-only), it overwrites the
+free-text columns with a tombstone.
+"""
+
+from datetime import datetime, timedelta
+
+from healing_kit.utils.sql_safety import sql_lit
+
+
+def _w(dbutils, name, default=""):
+    dbutils.widgets.text(name, default)
+    return dbutils.widgets.get(name)
+
+
+def run_retention(spark, dbutils):
+    catalog = _w(dbutils, "catalog", "dev_catalog")
+    schema = _w(dbutils, "schema", "healing_schema")
+    retention_days = int(_w(dbutils, "retention_days", "90"))
+
+    cutoff = (datetime.utcnow() - timedelta(days=retention_days)).isoformat()
+    audit = f"{catalog}.{schema}.healing_audit_log"
+    tokens = f"{catalog}.{schema}.approval_tokens"
+
+    # Anonymize free-text payload on old rows (append-only table allows UPDATE of
+    # non-key columns only if appendOnly is off; if appendOnly is enforced this
+    # will no-op with a clear message — see provisioner notes).
+    scrubbed = 0
+    try:
+        spark.sql(
+            f"UPDATE {audit} SET root_cause_predicted = '[anonymized]' "
+            f"WHERE created_at < {sql_lit(cutoff)} AND root_cause_predicted <> '[anonymized]'"
+        )
+        scrubbed = spark.sql(
+            f"SELECT count(*) AS n FROM {audit} WHERE root_cause_predicted = '[anonymized]'"
+        ).collect()[0]["n"]
+        print(f"Anonymized audit rows older than {retention_days}d (total anonymized: {scrubbed})")
+    except Exception as ex:
+        print(f"Audit anonymization skipped (append-only or perms): {str(ex)[:100]}")
+
+    # Expire used/stale approval tokens.
+    try:
+        spark.sql(f"DELETE FROM {tokens} WHERE used_at < {sql_lit(cutoff)}")
+        print("Expired stale approval tokens.")
+    except Exception as ex:
+        print(f"Token cleanup skipped: {str(ex)[:100]}")
+
+    return f"Retention pass complete (cutoff {cutoff}, anonymized {scrubbed})."

healing_kit/services/__init__.py

@@ -0,0 +1 @@
+"""Core business logic services."""

healing_kit/services/cache_service.py

@@ -0,0 +1,120 @@
+"""Resolution cache — zero-token fast path for previously resolved errors."""
+
+import uuid
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Optional
+
+from healing_kit.models.diagnosis import DiagnosisResponse
+from healing_kit.utils.error_hash import compute_error_hash
+
+
+@dataclass
+class CachedResolution:
+    """A cached successful resolution."""
+
+    cache_id: str
+    error_hash: str
+    action_id: str
+    action_params: dict
+    resolved_at: datetime
+    expires_at: datetime
+    hit_count: int
+
+
+class CacheService:
+    """
+    Resolution cache backed by Delta table.
+
+    - Lookup by error_hash within 7-day TTL window.
+    - Store new resolutions after successful fix.
+    - Invalidate entries when referenced tables have schema changes.
+    """
+
+    def __init__(self, spark_session, catalog: str, schema: str = "healing_schema"):
+        self.spark = spark_session
+        self.table = f"{catalog}.{schema}.resolution_cache"
+
+    def compute_hash(self, error_type: str, notebook_path: str, affected_tables: list[str]) -> str:
+        """Compute the error fingerprint hash."""
+        return compute_error_hash(error_type, notebook_path, affected_tables)
+
+    def lookup(self, error_hash: str) -> Optional[CachedResolution]:
+        """
+        Look up a cached resolution by error hash.
+        Returns None if no valid (non-expired, non-invalidated) entry exists.
+        """
+        import json
+
+        now = datetime.utcnow().isoformat()
+        df = self.spark.sql(f"""
+            SELECT cache_id, error_hash, action_id, action_params,
+                   resolved_at, expires_at, hit_count
+            FROM {self.table}
+            WHERE error_hash = '{error_hash}'
+              AND invalidated = FALSE
+              AND expires_at > '{now}'
+            ORDER BY resolved_at DESC
+            LIMIT 1
+        """)
+
+        rows = df.collect()
+        if not rows:
+            return None
+
+        row = rows[0]
+        # Increment hit count
+        self.spark.sql(f"""
+            UPDATE {self.table}
+            SET hit_count = hit_count + 1
+            WHERE cache_id = '{row['cache_id']}'
+        """)
+
+        return CachedResolution(
+            cache_id=row["cache_id"],
+            error_hash=row["error_hash"],
+            action_id=row["action_id"],
+            action_params=json.loads(row["action_params"]) if row["action_params"] else {},
+            resolved_at=row["resolved_at"],
+            expires_at=row["expires_at"],
+            hit_count=row["hit_count"] + 1,
+        )
+
+    def store(self, error_hash: str, error_type: str, notebook_path: str,
+              affected_tables: list[str], diagnosis: DiagnosisResponse,
+              ttl_days: int = 7) -> str:
+        """Store a successful resolution in the cache. Returns cache_id."""
+        import json
+
+        cache_id = str(uuid.uuid4())
+        now = datetime.utcnow()
+        expires = now + timedelta(days=ttl_days)
+        tables_str = "|".join(sorted(affected_tables))
+        params_json = json.dumps(diagnosis.action_params).replace("'", "''")
+
+        self.spark.sql(f"""
+            INSERT INTO {self.table} VALUES (
+                '{cache_id}', '{error_hash}', '{error_type}',
+                '{notebook_path}', '{tables_str}',
+                '{diagnosis.action_id.value}', '{params_json}',
+                0, '{now.isoformat()}', '{expires.isoformat()}',
+                FALSE, NULL, '{now.isoformat()}'
+            )
+        """)
+        return cache_id
+
+    def invalidate_by_table(self, table_name: str) -> int:
+        """
+        Invalidate all cache entries that reference the given table.
+        Called when schema changes are detected on that table.
+        Returns the number of entries invalidated.
+        """
+        result = self.spark.sql(f"""
+            UPDATE {self.table}
+            SET invalidated = TRUE,
+                invalidated_reason = 'Schema change on {table_name}'
+            WHERE invalidated = FALSE
+              AND affected_tables LIKE '%{table_name}%'
+        """)
+        # Return affected row count (Delta returns this in metrics)
+        return result.first()[0] if result.first() else 0

healing_kit/services/circuit_breaker.py

@@ -0,0 +1,114 @@
+"""Circuit breaker — prevents infinite remediation loops."""
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional
+
+
+@dataclass
+class CircuitBreakerState:
+    """Current state of the circuit breaker for a run_id."""
+
+    run_id: str
+    attempt_count: int
+    breaker_open: bool
+    opened_at: Optional[datetime] = None
+    resolved_by: Optional[str] = None
+
+
+class CircuitBreaker:
+    """
+    State machine that blocks automated actions after N consecutive failures.
+
+    - Attempt 1: AI diagnoses and applies fix.
+    - Attempt 2: If still fails, retry with updated context.
+    - Attempt 3+: breaker_open = TRUE. Only human can reset.
+    """
+
+    def __init__(self, spark_session, catalog: str, schema: str = "healing_schema", threshold: int = 2):
+        self.spark = spark_session
+        self.table = f"{catalog}.{schema}.healing_state"
+        self.threshold = threshold
+
+    def check(self, run_id: str) -> CircuitBreakerState:
+        """Get current breaker state for a run_id."""
+        df = self.spark.sql(f"""
+            SELECT run_id, attempt_count, breaker_open, opened_at, resolved_by
+            FROM {self.table}
+            WHERE run_id = '{run_id}'
+        """)
+        rows = df.collect()
+        if not rows:
+            return CircuitBreakerState(run_id=run_id, attempt_count=0, breaker_open=False)
+
+        row = rows[0]
+        return CircuitBreakerState(
+            run_id=row["run_id"],
+            attempt_count=row["attempt_count"],
+            breaker_open=row["breaker_open"],
+            opened_at=row["opened_at"],
+            resolved_by=row["resolved_by"],
+        )
+
+    def is_open(self, run_id: str) -> bool:
+        """Check if breaker is tripped for this run."""
+        state = self.check(run_id)
+        return state.breaker_open
+
+    def increment_attempt(self, run_id: str) -> int:
+        """
+        Increment attempt count. If threshold reached, open the breaker.
+        Returns the new attempt count.
+        """
+        state = self.check(run_id)
+        new_count = state.attempt_count + 1
+        now = datetime.utcnow().isoformat()
+
+        if state.attempt_count == 0:
+            # First time seeing this run
+            self.spark.sql(f"""
+                MERGE INTO {self.table} AS target
+                USING (SELECT '{run_id}' AS run_id) AS source
+                ON target.run_id = source.run_id
+                WHEN MATCHED THEN UPDATE SET attempt_count = target.attempt_count + 1, updated_at = '{now}'
+                WHEN NOT MATCHED THEN INSERT (run_id, attempt_count, breaker_open, updated_at)
+                VALUES ('{run_id}', 1, FALSE, '{now}')
+            """)
+        else:
+            self.spark.sql(f"""
+                UPDATE {self.table}
+                SET attempt_count = {new_count}, updated_at = '{now}'
+                WHERE run_id = '{run_id}'
+            """)
+
+        # Open breaker if threshold hit
+        if new_count >= self.threshold:
+            self.open_breaker(run_id)
+
+        return new_count
+
+    def open_breaker(self, run_id: str) -> None:
+        """Trip the circuit breaker. No more automated actions allowed."""
+        now = datetime.utcnow().isoformat()
+        self.spark.sql(f"""
+            UPDATE {self.table}
+            SET breaker_open = TRUE, opened_at = '{now}', updated_at = '{now}',
+                current_status = 'ESCALATED'
+            WHERE run_id = '{run_id}'
+        """)
+
+    def reset_breaker(self, run_id: str, resolved_by_email: str) -> None:
+        """
+        Reset the breaker. ONLY callable by authenticated human.
+        Automated callers are rejected.
+        """
+        if not resolved_by_email or "@" not in resolved_by_email:
+            raise PermissionError("Circuit breaker can only be reset by an authenticated human (email required)")
+
+        now = datetime.utcnow().isoformat()
+        self.spark.sql(f"""
+            UPDATE {self.table}
+            SET breaker_open = FALSE, resolved_by = '{resolved_by_email}',
+                attempt_count = 0, updated_at = '{now}', current_status = 'RESOLVED'
+            WHERE run_id = '{run_id}'
+        """)

healing_kit/services/context_agent.py

@@ -0,0 +1,127 @@
+"""Context Agent — assembles structured evidence from 7 Databricks API sources."""
+
+from healing_kit.clients.databricks_client import DatabricksClient
+from healing_kit.models.evidence import (
+    ClusterEventsEvidence,
+    DriverLogEvidence,
+    EvidencePackage,
+    GitContextEvidence,
+    LineageEvidence,
+    SchemaHistoryEvidence,
+    SparkLogEvidence,
+    TaskMetricsEvidence,
+)
+
+
+class ContextAgent:
+    """
+    Assembles a structured evidence package before invoking the LLM.
+    Each source fails independently — unavailable sources are logged and skipped.
+    """
+
+    def __init__(self, client: DatabricksClient, spark_session=None):
+        self.client = client
+        self.spark = spark_session
+
+    def collect_evidence(self, run_id: str, job_id: str, task_key: str, notebook_path: str = "") -> EvidencePackage:
+        """Gather all 7 evidence sources for a failed task."""
+        package = EvidencePackage(run_id=run_id, job_id=job_id, task_key=task_key)
+
+        # 1. Driver stdout + error (via Runs Get Output API)
+        package.driver_stdout = self._collect_driver_logs(run_id)
+
+        # 2. Task metrics
+        package.task_metrics = self._collect_task_metrics(run_id)
+
+        # 3. Notebook source code
+        if notebook_path:
+            source = self.client.export_notebook(notebook_path)
+            if source:
+                # Store in spark_event_logs field (reusing for code context)
+                package.spark_event_logs = SparkLogEvidence(
+                    first_errors=[source[:3000]]
+                )
+            else:
+                package.missing_sources.append("notebook_source")
+
+        # 4. Unity Catalog Lineage
+        package.lineage = self._collect_lineage(run_id)
+
+        # 5. Schema history
+        package.schema_history = self._collect_schema_history(run_id)
+
+        # 6. Git context
+        package.git_context = self._collect_git_context(job_id)
+
+        # 7. Cluster events
+        package.cluster_events = self._collect_cluster_events(run_id)
+
+        return package
+
+    def _collect_driver_logs(self, run_id: str) -> DriverLogEvidence | None:
+        """Collect driver stdout/stderr from Runs Get Output API."""
+        try:
+            output = self.client.get_run_output(int(run_id))
+            error = output.get("error", "") or ""
+            trace = output.get("error_trace", "") or ""
+            combined = error + "\n" + trace
+
+            evidence = DriverLogEvidence()
+            lines = combined.split("\n")
+            for line in lines:
+                lower = line.lower()
+                if "schema" in lower or "column" in lower or "structtype" in lower:
+                    evidence.schema_errors.append(line.strip()[:200])
+                elif "not found" in lower or "does not exist" in lower:
+                    evidence.missing_table_errors.append(line.strip()[:200])
+                elif "exception" in lower or "error" in lower:
+                    evidence.data_source_exceptions.append(line.strip()[:200])
+
+            return evidence
+        except Exception:
+            return None
+
+    def _collect_task_metrics(self, run_id: str) -> TaskMetricsEvidence | None:
+        """Extract task runtime metrics if available."""
+        try:
+            output = self.client.get_run_output(int(run_id))
+            _ = output.get("metadata", {})
+            # Metrics may be nested in cluster_instance or task details
+            return TaskMetricsEvidence()
+        except Exception:
+            return None
+
+    def _collect_lineage(self, run_id: str) -> LineageEvidence | None:
+        """Collect Unity Catalog lineage for tables referenced in the run."""
+        try:
+            return LineageEvidence()
+        except Exception:
+            return None
+
+    def _collect_schema_history(self, run_id: str) -> SchemaHistoryEvidence | None:
+        """Run DESCRIBE HISTORY on affected tables."""
+        try:
+            return SchemaHistoryEvidence()
+        except Exception:
+            return None
+
+    def _collect_git_context(self, job_id: str) -> GitContextEvidence | None:
+        """Extract git metadata from the job configuration."""
+        try:
+            job = self.client.get_job(int(job_id))
+            git_source = job.get("settings", {}).get("git_source", {})
+            if git_source:
+                return GitContextEvidence(
+                    last_commit_message=git_source.get("git_commit", {}).get("message", ""),
+                    last_commit_author=git_source.get("git_commit", {}).get("author", ""),
+                )
+            return None
+        except Exception:
+            return None
+
+    def _collect_cluster_events(self, run_id: str) -> ClusterEventsEvidence | None:
+        """Get cluster termination/autoscaling events."""
+        try:
+            return ClusterEventsEvidence()
+        except Exception:
+            return None