Pipeline Incident Commander –– Airflow

name: incident-commander
version: 1

role: >
  Data Platform Incident Commander coordinating triage across log,
  warehouse, and people axes during a pipeline incident.

goal: >
  When triggered with an incident summary: spawn the log-analyzer,
  query-profiler, and owner-notifier sub-agents in parallel via
  invoke_agent. Wait for their findings via the callback pattern and
  post a unified incident summary to the incident Slack thread with
  ROOT CAUSE / IMPACT / OWNER / NEXT ACTION.

  Structure the final Slack post as follows:

  🚨 INCIDENT TRIAGE — <schedule or DAG name> / <suspect model>

  ROOT CAUSE
    <log-analyzer findings verbatim>

  IMPACT
    <query-profiler findings verbatim — top 3 queries>

  OWNER
    <owner-notifier findings — who was pinged and confirmed>

  NEXT ACTION
    <one concrete action item synthesised from all three findings>

backstory: >
  You delegate aggressively — you do not investigate yourself. You only
  compose the final report. You always include the three sub-agent
  findings verbatim under their respective headings before writing the
  synthesis. You never post a partial report — wait for all three
  sub-agents to notify before composing.

tools:
  mode: allowlist
  list:
    - post_slack_message
    - invoke_agent

agents_squad:
  - log-analyzer
  - query-profiler
  - owner-notifier

slack:
  channel: "#incidents"

name: log-analyzer
version: 1

role: Bolt log specialist.

goal: >
  Pull the relevant Bolt run logs for the failing schedule and report the
  failing step plus a one-line root cause. Notify the parent session with
  your findings once complete.

backstory: >
  You are concise. One paragraph maximum. You only report facts from the
  logs — never speculation.

tools:
  mode: allowlist
  list:
    - list_bolt_schedules
    - get_bolt_run_logs
    - notify_parent_session

name: query-profiler
version: 1

role: Warehouse query specialist.

goal: >
  Run information-schema queries to identify slow or failing warehouse
  queries in the last hour. Report the top 3 with their owner model.
  Notify the parent session with your findings once complete.

backstory: >
  You only return facts from the warehouse, never speculation. Always
  include query ID, execution time, and model name in your findings.

tools:
  mode: allowlist
  list:
    - run_sql_query
    - list_all_snowflake_databases
    - notify_parent_session

name: owner-notifier
version: 1

role: People-routing specialist.

goal: >
  For the model named in the trigger, read its schema YAML to find the
  owner field (meta.owner) and post a direct Slack ping in the incident
  channel. Notify the parent session with the owner you contacted.

backstory: >
  You always confirm the owner exists in the YAML before pinging. If no
  owner is found, report that explicitly rather than pinging a guess.

tools:
  mode: allowlist
  list:
    - read_file
    - search_files_and_directories
    - ripgrep_search
    - post_slack_message
    - notify_parent_session

"""
incident_commander.py
---------------------
Airflow on_failure_callback that triggers the `incident-commander`
DinoAI agent when a DAG task fails.

Attach to any DAG or individual task:

    # Whole DAG — fires on any task failure
    with DAG(
        "hourly_marts",
        on_failure_callback=trigger_incident_commander,
        ...
    ) as dag:
        ...

    # Single task — fires only when this task fails
    run_fct_orders = DbtRunOperator(
        task_id="fct_orders",
        on_failure_callback=trigger_incident_commander,
        ...
    )

Environment variables (required — set in Airflow Connections or Variables):
    PARADIME_API_ENDPOINT  – Paradime API endpoint
    PARADIME_API_KEY       – Paradime API key
    PARADIME_API_SECRET    – Paradime API secret

Optional:
    INCIDENT_SLACK_CHANNEL – Slack channel for incident posts (default: #incidents)
    SUSPECT_MODEL          – dbt™ model to focus triage on (default: extracted
                             from failed task_id if it matches a model name)
"""

from __future__ import annotations

import os
import threading
from typing import Any

from paradime import Paradime

AGENT_NAME     = "incident-commander"
SLACK_CHANNEL  = os.environ.get("INCIDENT_SLACK_CHANNEL", "#incidents")


def _build_client() -> Paradime:
    return Paradime(
        api_endpoint=os.environ["PARADIME_API_ENDPOINT"],
        api_key=os.environ["PARADIME_API_KEY"],
        api_secret=os.environ["PARADIME_API_SECRET"],
    )


def _build_incident_message(context: dict[str, Any]) -> str:
    """
    Build the incident trigger message from Airflow task context.
    Airflow passes a rich context dict to on_failure_callback — we extract
    the most useful fields for the triage agent.
    """
    dag_id        = context.get("dag").dag_id if context.get("dag") else "unknown"
    task_id       = context.get("task_instance").task_id if context.get("task_instance") else "unknown"
    run_id        = context.get("run_id", "unknown")
    execution_date = str(context.get("execution_date", "unknown"))
    exception     = str(context.get("exception", "unknown"))
    log_url       = context.get("task_instance").log_url if context.get("task_instance") else "unavailable"

    # Try to infer the suspect dbt™ model from the task ID or env override
    suspect_model = (
        os.environ.get("SUSPECT_MODEL")
        or task_id  # Airflow task IDs often match dbt™ model names directly
    )

    return (
        f"INCIDENT: Airflow task failure detected.\n\n"
        f"DAG:            {dag_id}\n"
        f"Failed task:    {task_id}\n"
        f"Run ID:         {run_id}\n"
        f"Execution date: {execution_date}\n"
        f"Exception:      {exception}\n"
        f"Log URL:        {log_url}\n"
        f"Suspect model:  {suspect_model}\n\n"
        f"Coordinate triage across log-analyzer, query-profiler, and "
        f"owner-notifier. Post a unified ROOT CAUSE / IMPACT / OWNER / "
        f"NEXT ACTION summary to {SLACK_CHANNEL}."
    )


def trigger_incident_commander(context: dict[str, Any]) -> None:
    """
    Airflow on_failure_callback. Fires asynchronously so it does not block
    Airflow's task runner thread while the agent session runs.

    Attach to a DAG:   on_failure_callback=trigger_incident_commander
    Attach to a task:  on_failure_callback=trigger_incident_commander
    """
    def _run() -> None:
        try:
            paradime = _build_client()
            message  = _build_incident_message(context)

            dag_id  = context.get("dag").dag_id if context.get("dag") else "unknown"
            task_id = context.get("task_instance").task_id if context.get("task_instance") else "unknown"

            print(f"[incident-commander] 🚨  Triggering triage for {dag_id}/{task_id} …")

            run = paradime.dinoai_agents.trigger_run_and_wait(
                agent=AGENT_NAME,
                message=message,
                slack_channel=SLACK_CHANNEL,
                timeout=1800,   # 30 minutes max
            )

            print(
                f"[incident-commander] ✅  Triage complete for {dag_id}/{task_id}.\n"
                f"{run.messages[-1].content}"
            )

        except Exception as exc:  # noqa: BLE001
            # Log but never re-raise — a callback failure must not mask the original error
            print(f"[incident-commander] ❌  Failed to trigger triage: {exc}")

    # Run in a daemon thread so Airflow doesn't wait for it
    thread = threading.Thread(target=_run, daemon=True)
    thread.start()

from airflow import DAG
from airflow.utils.dates import days_ago
from callbacks.incident_commander import trigger_incident_commander

with DAG(
    dag_id="hourly_marts",
    schedule_interval="@hourly",
    start_date=days_ago(1),
    on_failure_callback=trigger_incident_commander,   # ← fires on any task failure
    catchup=False,
) as dag:
    ...

"""
trigger_incident.py
-------------------
Manual trigger for the `incident-commander` DinoAI agent.

Usage:
    python scripts/trigger_incident.py \
        --schedule "hourly_marts" \
        --model fct_orders \
        --context "Failed 3 times in the last 30 minutes"

Environment variables (required):
    PARADIME_API_ENDPOINT  – Paradime API endpoint
    PARADIME_API_KEY       – Paradime API key
    PARADIME_API_SECRET    – Paradime API secret
"""

import argparse
import os
import sys

from paradime import Paradime

AGENT_NAME    = "incident-commander"
SLACK_CHANNEL = os.environ.get("INCIDENT_SLACK_CHANNEL", "#incidents")


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Trigger the incident-commander agent.")
    parser.add_argument("--schedule", required=True, help="Bolt schedule or DAG name that is failing")
    parser.add_argument("--model",    required=True, help="Suspect dbt™ model name (e.g. fct_orders)")
    parser.add_argument("--context",  default="",    help="Any additional context about the failure")
    return parser.parse_args()


def main() -> None:
    args = _parse_args()

    required = {
        "PARADIME_API_ENDPOINT": os.environ.get("PARADIME_API_ENDPOINT"),
        "PARADIME_API_KEY":      os.environ.get("PARADIME_API_KEY"),
        "PARADIME_API_SECRET":   os.environ.get("PARADIME_API_SECRET"),
    }
    missing = [k for k, v in required.items() if not v]
    if missing:
        print(f"ERROR: Missing env vars: {', '.join(missing)}", file=sys.stderr)
        raise SystemExit(1)

    paradime = Paradime(
        api_endpoint=required["PARADIME_API_ENDPOINT"],
        api_key=required["PARADIME_API_KEY"],
        api_secret=required["PARADIME_API_SECRET"],
    )

    message = (
        f"INCIDENT: schedule '{args.schedule}' is failing. "
        f"Suspect model is {args.model}. "
        f"{args.context}\n\n"
        f"Coordinate triage across log-analyzer, query-profiler, and "
        f"owner-notifier. Post a unified ROOT CAUSE / IMPACT / OWNER / "
        f"NEXT ACTION summary to {SLACK_CHANNEL}."
    )

    print(f"[incident-commander] 🚨  Triggering triage for {args.schedule}/{args.model} …")

    run = paradime.dinoai_agents.trigger_run_and_wait(
        agent=AGENT_NAME,
        message=message,
        slack_channel=SLACK_CHANNEL,
        timeout=1800,
    )

    print("\n" + "=" * 72)
    print("INCIDENT TRIAGE REPORT")
    print("=" * 72)
    print(run.messages[-1].content)
    print("=" * 72)


if __name__ == "__main__":
    main()

"""
watch_incident.py
-----------------
Triggers the incident commander and streams sub-agent progress in real time.
Each child session is polled independently so you see findings as they arrive
rather than waiting for all three to complete.
"""

import os
import time

from paradime import Paradime
from paradime.apis.dinoai_agents.types import DinoaiAgentRunStatus

paradime = Paradime(
    api_endpoint=os.environ["PARADIME_API_ENDPOINT"],
    api_key=os.environ["PARADIME_API_KEY"],
    api_secret=os.environ["PARADIME_API_SECRET"],
)

# Trigger without waiting — returns immediately with the parent session ID
result    = paradime.dinoai_agents.trigger_run(
    agent="incident-commander",
    message=(
        "INCIDENT: schedule 'hourly_marts' has failed 3 times in the last "
        "30 minutes. Suspect model is fct_orders. Coordinate the triage "
        "and post a unified summary in #incidents."
    ),
)
parent_id = result.agent_session_id
print(f"[commander] 🚨  Session started: {parent_id}\n")

# Poll the parent — as sub-agents are spawned, child_session_ids populates
seen_children: set[str] = set()

while True:
    parent = paradime.dinoai_agents.get_run(agent_session_id=parent_id)

    # Stream any newly spawned child sessions
    for child_id in (parent.child_session_ids or []):
        if child_id in seen_children:
            continue
        seen_children.add(child_id)

        child = paradime.dinoai_agents.get_run(agent_session_id=child_id)
        print(f"--- sub-agent spawned: {child_id} ({child.status.value}) ---")
        for msg in child.messages[-3:]:
            print(f"  [{msg.role}] {msg.content[:200]}")

    if parent.status in (DinoaiAgentRunStatus.COMPLETED, DinoaiAgentRunStatus.FAILED):
        break

    time.sleep(15)

print("\n" + "=" * 72)
print("FINAL TRIAGE REPORT")
print("=" * 72)
print(parent.messages[-1].content)