shkit 1.2.0__py3-none-any.whl

healing_kit/services/token_budget.py

@@ -0,0 +1,137 @@
+"""Token budget enforcer — cost governance for LLM calls."""
+
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum
+
+
+class BudgetMode(str, Enum):
+    NORMAL = "normal"
+    WARNING = "warning"
+    DEGRADED = "degraded"
+
+
+@dataclass
+class BudgetState:
+    """Current budget window state."""
+
+    window_id: str
+    tokens_used: int
+    cost_usd: float
+    hourly_budget: int
+    current_mode: BudgetMode
+    spend_pct: float
+
+
+class TokenBudgetEnforcer:
+    """
+    Tracks hourly LLM token spend and enforces mode transitions.
+
+    Modes:
+      - Normal (< 70% budget): Full AI healing active
+      - Warning (70-90%): Lightweight model only
+      - Degraded (> 90%): No LLM calls, log + page on-call
+    """
+
+    def __init__(self, spark_session, catalog: str, schema: str = "healing_schema", hourly_budget: int = 50000):
+        self.spark = spark_session
+        self.table = f"{catalog}.{schema}.token_budget"
+        self.hourly_budget = hourly_budget
+
+    def _current_window_id(self) -> str:
+        """Get the current hourly window ID (YYYY-MM-DD-HH)."""
+        return datetime.utcnow().strftime("%Y-%m-%d-%H")
+
+    def get_current_state(self) -> BudgetState:
+        """Get the current budget window state."""
+        window_id = self._current_window_id()
+        df = self.spark.sql(f"""
+            SELECT window_id, tokens_used, cost_usd, hourly_budget, current_mode
+            FROM {self.table}
+            WHERE window_id = '{window_id}'
+        """)
+        rows = df.collect()
+
+        if not rows:
+            return BudgetState(
+                window_id=window_id,
+                tokens_used=0,
+                cost_usd=0.0,
+                hourly_budget=self.hourly_budget,
+                current_mode=BudgetMode.NORMAL,
+                spend_pct=0.0,
+            )
+
+        row = rows[0]
+        tokens = row["tokens_used"]
+        budget = row["hourly_budget"]
+        spend_pct = (tokens / budget * 100) if budget > 0 else 0
+
+        return BudgetState(
+            window_id=row["window_id"],
+            tokens_used=tokens,
+            cost_usd=row["cost_usd"],
+            hourly_budget=budget,
+            current_mode=BudgetMode(row["current_mode"]),
+            spend_pct=spend_pct,
+        )
+
+    def get_current_mode(self) -> BudgetMode:
+        """Return current mode based on spend percentage."""
+        state = self.get_current_state()
+        return self._compute_mode(state.spend_pct)
+
+    def can_invoke_model(self, model_tier: str) -> bool:
+        """
+        Check if a model invocation is allowed under current budget mode.
+
+        model_tier: 'lightweight' or 'powerful'
+        """
+        mode = self.get_current_mode()
+        if mode == BudgetMode.DEGRADED:
+            return False
+        if mode == BudgetMode.WARNING and model_tier == "powerful":
+            return False
+        return True
+
+    def record_usage(self, tokens_used: int, cost_usd: float) -> BudgetMode:
+        """
+        Record token usage and return the resulting mode.
+        Creates the window entry if it doesn't exist.
+        """
+        window_id = self._current_window_id()
+        now = datetime.utcnow().isoformat()
+
+        self.spark.sql(f"""
+            MERGE INTO {self.table} AS target
+            USING (SELECT '{window_id}' AS window_id) AS source
+            ON target.window_id = source.window_id
+            WHEN MATCHED THEN UPDATE SET
+                tokens_used = target.tokens_used + {tokens_used},
+                cost_usd = target.cost_usd + {cost_usd},
+                updated_at = '{now}'
+            WHEN NOT MATCHED THEN INSERT
+                (window_id, window_start, window_end, tokens_used, cost_usd, hourly_budget, current_mode, created_at, updated_at)
+            VALUES ('{window_id}', '{now}', '{now}', {tokens_used}, {cost_usd}, {self.hourly_budget}, 'normal', '{now}', '{now}')
+        """)
+
+        # Evaluate and update mode
+        state = self.get_current_state()
+        new_mode = self._compute_mode(state.spend_pct)
+
+        if new_mode != state.current_mode:
+            self.spark.sql(f"""
+                UPDATE {self.table}
+                SET current_mode = '{new_mode.value}', mode_changed_at = '{now}', updated_at = '{now}'
+                WHERE window_id = '{window_id}'
+            """)
+
+        return new_mode
+
+    def _compute_mode(self, spend_pct: float) -> BudgetMode:
+        """Determine mode from spend percentage."""
+        if spend_pct >= 90:
+            return BudgetMode.DEGRADED
+        elif spend_pct >= 70:
+            return BudgetMode.WARNING
+        return BudgetMode.NORMAL

healing_kit/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility functions."""

healing_kit/utils/error_hash.py

@@ -0,0 +1,15 @@
+"""SHA-256 error fingerprinting for resolution cache keys."""
+
+import hashlib
+
+
+def compute_error_hash(error_type: str, notebook_path: str, affected_tables: list[str]) -> str:
+    """
+    Compute a deterministic SHA-256 hash of the failure fingerprint.
+
+    The hash is built from: error_type + notebook_path + sorted affected table names.
+    This ensures the same failure pattern always produces the same cache key.
+    """
+    sorted_tables = sorted(t.strip() for t in affected_tables if t.strip())
+    payload = f"{error_type}||{notebook_path}||{'|'.join(sorted_tables)}"
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()

healing_kit/utils/hmac_tokens.py

@@ -0,0 +1,86 @@
+"""HMAC-SHA256 signed one-time tokens for HITL approval flow."""
+
+import hashlib
+import hmac
+import json
+import time
+from dataclasses import dataclass
+
+TOKEN_TTL_SECONDS = 900  # 15 minutes
+
+
+@dataclass
+class TokenPayload:
+    """Decoded token payload."""
+
+    run_id: str
+    approver_email: str
+    action_id: str
+    issued_at: float
+    expires_at: float
+
+
+def generate_token(run_id: str, approver_email: str, action_id: str, secret_key: bytes) -> str:
+    """
+    Generate a signed HMAC-SHA256 token with 15-minute TTL.
+
+    The token encodes: run_id + approver_email + action_id + timestamp.
+    """
+    issued_at = time.time()
+    expires_at = issued_at + TOKEN_TTL_SECONDS
+
+    payload = json.dumps({
+        "run_id": run_id,
+        "approver_email": approver_email,
+        "action_id": action_id,
+        "issued_at": issued_at,
+        "expires_at": expires_at,
+    }, separators=(",", ":"))
+
+    signature = hmac.HMAC(secret_key, payload.encode(), hashlib.sha256).hexdigest()
+    # Token = base64-like: payload_hex.signature
+    payload_hex = payload.encode().hex()
+    return f"{payload_hex}.{signature}"
+
+
+def validate_token(token: str, secret_key: bytes) -> TokenPayload:
+    """
+    Validate a signed token. Raises ValueError if invalid or expired.
+
+    Checks:
+    1. Token format is valid
+    2. HMAC signature matches
+    3. Token has not expired (15-min TTL)
+    """
+    parts = token.split(".")
+    if len(parts) != 2:
+        raise ValueError("Invalid token format")
+
+    payload_hex, signature = parts
+
+    # Reconstruct payload
+    try:
+        payload_bytes = bytes.fromhex(payload_hex)
+        payload_str = payload_bytes.decode()
+    except (ValueError, UnicodeDecodeError):
+        raise ValueError("Invalid token encoding")
+
+    # Verify signature
+    expected_sig = hmac.HMAC(secret_key, payload_bytes, hashlib.sha256).hexdigest()
+    if not hmac.compare_digest(signature, expected_sig):
+        raise ValueError("Invalid token signature — possible tampering")
+
+    # Parse payload
+    data = json.loads(payload_str)
+
+    # Check expiry
+    if time.time() > data["expires_at"]:
+        raise ValueError(f"Token expired at {data['expires_at']}")
+
+    return TokenPayload(
+        run_id=data["run_id"],
+        approver_email=data["approver_email"],
+        action_id=data["action_id"],
+        issued_at=data["issued_at"],
+        expires_at=data["expires_at"],
+    )

healing_kit/utils/sql_safety.py

@@ -0,0 +1,84 @@
+"""Safe SQL helpers — literal escaping and a word-boundary SELECT guard.
+
+These exist because the runtime previously (a) interpolated LLM/free-text output
+directly into INSERT statements and (b) used a substring keyword blocklist that
+both let crafted statements through and false-positived on legitimate columns
+like ``created_at`` (contains "CREATE") and ``updated_at`` (contains "UPDATE").
+"""
+
+import re
+
+# DDL/DML keywords that must never appear in a diagnostic query. Matched as whole
+# words (see ``is_safe_select``) so ``created_at`` / ``updated_at`` are allowed.
+_FORBIDDEN_KEYWORDS = (
+    "INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "TRUNCATE",
+    "CREATE", "MERGE", "GRANT", "REVOKE", "REPLACE", "COPY",
+    "CALL", "EXECUTE", "SET", "USE",
+)
+
+_KEYWORD_RE = re.compile(
+    r"\b(" + "|".join(_FORBIDDEN_KEYWORDS) + r")\b", re.IGNORECASE
+)
+# A SQL line/block comment can be used to smuggle a second statement past naive
+# checks; reject anything containing comment markers.
+_COMMENT_RE = re.compile(r"(--|/\*|\*/|#)")
+
+
+def sql_literal(value) -> str:
+    """Render a Python value as a safe SQL literal.
+
+    - ``None`` -> ``NULL``
+    - ``bool`` -> ``TRUE``/``FALSE``
+    - ``int``/``float`` -> bare numeric
+    - everything else -> single-quoted string with quotes doubled
+    """
+    if value is None:
+        return "NULL"
+    if isinstance(value, bool):
+        return "TRUE" if value else "FALSE"
+    if isinstance(value, (int, float)):
+        return repr(value)
+    return "'" + str(value).replace("'", "''") + "'"
+
+
+# Short alias used by the runtime / notebook code paths.
+sql_lit = sql_literal
+
+
+def is_safe_select(query: str) -> bool:
+    """Return True only if ``query`` is a single read-only SELECT statement.
+
+    Rejects: empty/non-SELECT queries, any forbidden DDL/DML keyword (whole-word),
+    multiple statements (a non-trailing semicolon), and SQL comments.
+    """
+    if not query or not isinstance(query, str):
+        return False
+
+    stripped = query.strip().rstrip(";").strip()
+    if not stripped:
+        return False
+
+    # No comments (could hide a second statement or smuggle keywords).
+    if _COMMENT_RE.search(stripped):
+        return False
+
+    # Only one statement allowed: no semicolons left after stripping a trailing one.
+    if ";" in stripped:
+        return False
+
+    upper = stripped.upper()
+    if not (upper.startswith("SELECT") or upper.startswith("WITH")):
+        return False
+
+    if _KEYWORD_RE.search(stripped):
+        return False
+
+    return True
+
+
+def ensure_limit(query: str, limit: int = 20) -> str:
+    """Append a LIMIT clause if the query doesn't already have one."""
+    stripped = query.strip().rstrip(";")
+    if re.search(r"\bLIMIT\b", stripped, re.IGNORECASE):
+        return stripped
+    return f"{stripped} LIMIT {limit}"

iic/__init__.py

@@ -0,0 +1,51 @@
+"""Incident Intelligence Core (IIC).
+
+A deterministic + AI hybrid engine that converts raw data-pipeline failures into
+structured, prioritized, explainable incident knowledge.
+
+Design principles
+-----------------
+1. Deterministic first, AI second — the LLM only runs after the full structured
+   context (the :class:`~iic.models.IncidentDNA`) has been built.
+2. Every failure becomes a structured object — one ``IncidentReport`` per root cause.
+3. No external workflow dependencies — Databricks + logs only, no Jira/ServiceNow.
+4. No "chatty AI" — the diagnosis engine returns structured output only.
+5. Every decision is traceable from evidence.
+
+The 11-stage pipeline lives in :mod:`iic.runtime.incident_engine`.
+"""
+
+from __future__ import annotations
+
+__version__ = "1.2.0"
+
+# NOTE: `import iic` is intentionally lightweight and does NOT arm anything.
+# The self-arming tripwire is installed by `iic_autoload.pth` at interpreter
+# startup (it runs `import iic.runtime.bootstrap`). See docs/SELF_ARMING.md.
+
+# Operator/console entry points — thin lazy wrappers so `import iic` stays light
+# (the heavy bits load only when these are actually called from a notebook).
+
+def console(*args, **kwargs):
+    """Fold pending occurrences and render the antibody ledger."""
+    from iic._console import console as _fn
+    return _fn(*args, **kwargs)
+
+
+def console_record(*args, **kwargs):
+    """Record a human-confirmed resolution for a pattern (append-only)."""
+    from iic._console import console_record as _fn
+    return _fn(*args, **kwargs)
+
+
+def onboard(*args, **kwargs):
+    """One-time first-run setup: secret scope, keys, ACL, dirs, doctor, test card."""
+    from iic._console import onboard as _fn
+    return _fn(*args, **kwargs)
+
+
+def doctor(*args, **kwargs):
+    """Verify the secret scope, keys, volume write, and webhook. Exit 0 iff healthy."""
+    from iic._doctor import doctor as _fn
+    return _fn(*args, **kwargs)
+

iic/__main__.py

@@ -0,0 +1,18 @@
+"""``python -m iic doctor [--check-principal <name>]`` — the support entry point."""
+
+from __future__ import annotations
+
+import sys
+
+
+def main(argv=None) -> int:
+    argv = argv if argv is not None else sys.argv[1:]
+    if argv and argv[0] == "doctor":
+        from iic._doctor import main as doctor_main
+        return doctor_main(argv[1:])
+    print("usage: python -m iic doctor [--check-principal <name>]")
+    return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main())

iic/_console.py

@@ -0,0 +1,235 @@
+"""In-workspace console (``iic.console``) and first-run onboarding (``iic.onboard``).
+
+Runs from a Databricks notebook on serverless. The console is the ONLY writer of
+``antibodies.yaml`` besides a human editing it: it folds ``.iic_pending/*`` into the
+ledger with the canonical merge rule (:mod:`iic.runtime.ledger`), renders the ledger,
+and records resolutions (append-only). ``onboard`` does the one-time setup: create
+the ``iic`` secret scope, write the keys, grant READ, create the volume dirs, run
+doctor, and send a test card. Everything is best-effort / fail-open.
+"""
+
+from __future__ import annotations
+
+import glob
+import json
+import os
+
+from iic.runtime.antibodies import load_ledger
+from iic.runtime.constants import (
+    ANTIBODIES_FILENAME,
+    DEFAULT_SECRET_SCOPE,
+    PENDING_DIRNAME,
+)
+from iic.runtime.ledger import dump_ledger, merge_ledgers
+
+# ── ledger console ──
+
+def _resolve_base(base_dir):
+    if base_dir:
+        return base_dir
+    try:
+        from iic.runtime.scope_config import load_settings
+        return (load_settings() or {}).get("volume_path")
+    except Exception:
+        return None
+
+
+def _read_markers(base_dir):
+    out = []
+    try:
+        for p in glob.glob(os.path.join(base_dir, PENDING_DIRNAME, "*.json")):
+            try:
+                with open(p) as f:
+                    out.append((p, json.load(f)))
+            except Exception:
+                continue
+    except Exception:
+        pass
+    return out
+
+
+def _write_ledger(base_dir, ledger):
+    with open(os.path.join(base_dir, ANTIBODIES_FILENAME), "w") as f:
+        f.write(dump_ledger(ledger))
+
+
+def fold(base_dir=None) -> dict:
+    """Fold ``.iic_pending/*`` into ``antibodies.yaml`` (the only non-human writer)."""
+    base = _resolve_base(base_dir)
+    if not base:
+        raise RuntimeError("no volume_path — configure the 'iic' secret scope or pass base_dir")
+    existing = load_ledger(base)
+    markers = _read_markers(base)
+    merged, stats = merge_ledgers(existing, {}, [m for _, m in markers])
+    _write_ledger(base, merged)
+    for p, _ in markers:
+        try:
+            os.remove(p)
+        except Exception:
+            pass
+    fresh = stats["new_keys"]
+    print(f"[console] folded {len(markers)} occurrence(s); "
+          + (f"new pattern(s): {fresh}" if fresh else "no new patterns"))
+    return merged
+
+
+def render(ledger: dict) -> str:
+    if not ledger:
+        return "No patterns recorded yet. Break something to see your first incident."
+    resolved = [(k, e) for k, e in ledger.items() if str((e or {}).get("resolution", "")).strip()]
+    unresolved = [(k, e) for k, e in ledger.items() if not str((e or {}).get("resolution", "")).strip()]
+    lines = [f"Antibody ledger — {len(ledger)} pattern(s): "
+             f"{len(resolved)} resolved, {len(unresolved)} awaiting a fix", ""]
+    if unresolved:
+        lines.append("⚠️  Awaiting resolution (record a fix for these):")
+        for k, e in sorted(unresolved, key=lambda kv: -_int((kv[1] or {}).get("times_seen"))):
+            lines.append(f"  • {k}  (seen {_int((e or {}).get('times_seen'))}×)  "
+                         f"e.g. {(e or {}).get('example', '')}")
+        lines.append("")
+    if resolved:
+        lines.append("♻️  Resolved (these fixes show on the card):")
+        for k, e in sorted(resolved):
+            lines.append(f"  • {k}  → {(e or {}).get('resolution', '')}")
+    return "\n".join(lines)
+
+
+def record_resolution(pattern_id, resolution, *, base_dir=None, overwrite=False) -> bool:
+    """Append-only: set a resolution for a pattern. Refuses to overwrite a non-empty
+    resolution unless ``overwrite=True``. Returns True if written."""
+    base = _resolve_base(base_dir)
+    if not base:
+        raise RuntimeError("no volume_path — configure the 'iic' secret scope or pass base_dir")
+    ledger = load_ledger(base)
+    entry = dict(ledger.get(pattern_id) or {})
+    if str(entry.get("resolution", "")).strip() and not overwrite:
+        print(f"[console] '{pattern_id}' already has a resolution; pass overwrite=True to replace it")
+        return False
+    entry.setdefault("times_seen", entry.get("times_seen", 0))
+    entry["resolution"] = resolution
+    ledger[pattern_id] = entry
+    _write_ledger(base, ledger)
+    print(f"[console] recorded resolution for '{pattern_id}'")
+    return True
+
+
+def console(base_dir=None) -> dict:
+    """Fold pending occurrences, then render the ledger. Returns the ledger dict."""
+    ledger = fold(base_dir)
+    print(render(ledger))
+    print("\nRecord a fix:  iic.console_record('<pattern_id>', 'the fix that worked')")
+    return ledger
+
+
+def console_record(pattern_id, resolution, *, base_dir=None, overwrite=False) -> bool:
+    """Alias matching the hint printed by console()."""
+    return record_resolution(pattern_id, resolution, base_dir=base_dir, overwrite=overwrite)
+
+
+def _int(x) -> int:
+    try:
+        return int(x or 0)
+    except Exception:
+        return 0
+
+
+# ── onboarding ──
+
+_SCOPE_KEYS = ("teams_webhook", "volume_path", "host", "pat",
+               "github_repo", "github_dispatch_token", "dedup_ttl_seconds")
+
+
+def _workspace_client():
+    try:
+        from databricks.sdk import WorkspaceClient
+        return WorkspaceClient()
+    except Exception:
+        return None
+
+
+def _prompt_answers() -> dict:  # pragma: no cover - interactive
+    print("IIC onboarding — answer a few questions (blank to skip optional ones):")
+    a = {
+        "teams_webhook": input("  Teams webhook URL (required): ").strip(),
+        "volume_path": input("  Volume path for memory, e.g. /Volumes/cat/sch/libs (required): ").strip(),
+        "host": input("  Workspace host for 'View Run' links (optional): ").strip(),
+        "pat": input("  Workspace PAT to enable enrichment (optional): ").strip(),
+        "github_repo": input("  GitHub owner/repo for incident archiving (optional): ").strip(),
+        "github_dispatch_token": input("  GitHub dispatch token (optional): ").strip(),
+    }
+    principals = input("  Group/SP names to grant READ (comma-separated): ").strip()
+    a["read_principals"] = [p.strip() for p in principals.split(",") if p.strip()]
+    return a
+
+
+def _send_test_card(webhook) -> None:
+    if not webhook:
+        return
+    try:
+        import requests
+        card = {"type": "message", "attachments": [{
+            "contentType": "application/vnd.microsoft.card.adaptive",
+            "content": {"type": "AdaptiveCard", "version": "1.4",
+                        "body": [{"type": "TextBlock", "weight": "Bolder",
+                                  "text": "✅ IIC connected — you'll see incident cards here."}]}}]}
+        requests.post(webhook, json=card, headers={"Content-Type": "application/json"}, timeout=5)
+    except Exception:
+        pass
+
+
+def onboard(*, scope=None, answers=None) -> bool:  # pragma: no cover - Databricks/interactive
+    """One-time setup: create the secret scope, write keys, grant READ, create the
+    volume dirs, run doctor, and send a test card. Pass ``answers`` to skip prompts."""
+    scope = scope or os.environ.get("IIC_SECRET_SCOPE", DEFAULT_SECRET_SCOPE)
+    a = answers or _prompt_answers()
+    if not (a.get("teams_webhook") and a.get("volume_path")):
+        print("❌ teams_webhook and volume_path are required.")
+        return False
+    w = _workspace_client()
+    if w is None:
+        print("❌ could not construct a Databricks WorkspaceClient — run this in your workspace.")
+        return False
+
+    try:
+        existing = [s.name for s in w.secrets.list_scopes()]
+        if scope not in existing:
+            w.secrets.create_scope(scope=scope)
+            print(f"✅ created secret scope '{scope}'")
+        else:
+            print(f"✅ secret scope '{scope}' already exists")
+    except Exception as ex:
+        print(f"⚠️ could not list/create scope ({str(ex)[:120]}); assuming it exists")
+
+    for key in _SCOPE_KEYS:
+        val = a.get(key)
+        if val:
+            try:
+                w.secrets.put_secret(scope=scope, key=key, string_value=str(val))
+            except Exception as ex:
+                print(f"⚠️ put_secret {key} failed: {str(ex)[:100]}")
+    print(f"✅ wrote secrets to scope '{scope}'")
+
+    for principal in a.get("read_principals", []):
+        try:
+            from databricks.sdk.service.workspace import AclPermission
+            w.secrets.put_acl(scope=scope, principal=principal, permission=AclPermission.READ)
+            print(f"✅ granted READ on '{scope}' to {principal}")
+        except Exception as ex:
+            print(f"⚠️ put_acl {principal} failed: {str(ex)[:100]}")
+    print("ℹ️  EVERY identity that runs monitored jobs must have READ on this scope, "
+          "or the agent can't load config there.")
+
+    vol = a.get("volume_path")
+    for d in (vol, os.path.join(vol, PENDING_DIRNAME), os.path.join(vol, ".iic_seen")):
+        try:
+            os.makedirs(d, exist_ok=True)
+        except Exception:
+            pass
+
+    try:
+        from iic._doctor import doctor
+        doctor()
+    except Exception:
+        pass
+    _send_test_card(a.get("teams_webhook"))
+    print("✅ onboarding complete — break something to see your first card.")
+    return True