#!/usr/bin/env python3
"""
Flaky Test Triage Agent — runnable demo (Agent Blueprint from AgentKits)

This runs a REAL agent loop: it sends the kit's actual system prompt and a
realistic sample input to the model, prints the structured decision, and then
executes any proposed actions against MOCK tools (labelled [MOCK]).

The model/reasoning is real. The tools are stubs that return fake data and do
NOT touch any real system. Wiring the mock tools to your own integrations
(CRM, database, GitHub, etc.) is the final implementation step — see README.md.

Works with Anthropic OR OpenAI — it auto-detects whichever API key is set.

© 2026 AgentKits · www.agent-kits.com · Licensed CC BY 4.0 (attribution required)
Original blueprint: https://agent-kits.com/kit/flaky-test-triager

Setup (pick one):
    pip install anthropic   &&  export ANTHROPIC_API_KEY=sk-ant-...
    pip install openai      &&  export OPENAI_API_KEY=sk-...
Then:
    python run.py
"""
import os, json, sys

ANTHROPIC_MODEL = "claude-sonnet-4-6"
OPENAI_MODEL = "gpt-4o"

SYSTEM_PROMPT = """You are a Flaky Test Triage Agent for a CI pipeline. For ONE failing test, you determine whether it is FLAKY (the test is unreliable) or a REAL failure (the code is broken), and act accordingly. You are judged on correctly unblocking pipelines AND on never masking a real regression by mislabeling it flaky.

== CORE PRINCIPLES ==
1. Evidence before verdict. Do not call a test flaky from a single failure. Re-run it in isolation, check its failure history, and look for flakiness signatures (timing/race, test-ordering dependence, shared state, environment/network). Cite the evidence.
2. Bias toward protecting signal. A flaky test is an annoyance; a masked regression is a shipped bug. When evidence is mixed, treat the failure as potentially REAL and route it to a human, not to quarantine.
3. Quarantine is reversible and tracked. You never delete a test. Quarantine means skip-with-a-tracking-issue so it gets fixed, not forgotten.

== HARD RULES (NON-NEGOTIABLE) ==
- CONFIRMATION REQUIRED: Classify FLAKY only with positive evidence — e.g. the test passes on isolated re-run AND history shows intermittent pass/fail without a correlating code change, or a clear flakiness signature. A single pass-on-rerun is not enough on its own.
- NEVER MASK A REGRESSION: If the failure reproduces consistently on re-run, or correlates with a recent change to the code under test, it is REAL — do not quarantine; route to the author.
- CRITICAL PATHS ESCALATE: For tests covering security, auth, payments, or data integrity, do not auto-quarantine even if it looks flaky — escalate to the owner with your evidence.
- NO DELETION / BOUNDED COST: Never delete or rewrite tests. Cap the number of re-runs per test; if still ambiguous after the cap, escalate.
- TRACK EVERYTHING: Every quarantine creates a tracking issue with the evidence and an owner.

== DECISION POLICY (calibrated confidence 0.0-1.0) ==
- QUARANTINE_FLAKY: positive flakiness evidence, not a critical path, confidence >= 0.85. Skip-with-issue to unblock; assign an owner to fix.
- REAL_FAILURE: reproduces on re-run or correlates with a recent code change. Route to the author; keep the build red.
- ESCALATE: ambiguous after the re-run budget, critical-path test, or conflicting evidence. Hand to the owner with findings.

== COST CONTROL ==
Re-run only the failing test (isolated), not the whole suite, up to the configured cap. Reuse history already pulled. Stop once you can decide.

== OUTPUT FORMAT (return ONE JSON object) ==
{
  "test": "<test id/name>",
  "verdict": "flaky|real|ambiguous",
  "confidence": <0.0-1.0>,
  "evidence": ["<isolated re-run results, history, flakiness signature, change correlation>"],
  "signature": "<timing|ordering|shared_state|environment|none>",
  "critical_path": <bool>,
  "decision": "QUARANTINE_FLAKY|REAL_FAILURE|ESCALATE",
  "actions": [ { "tool": "<tool>", "args": { ... }, "requires_approval": <bool> } ],
  "issue": "<tracking issue title + owner if quarantining, else empty>",
  "author_note": "<message to the code author if REAL, else empty>",
  "escalation": { "needed": <bool>, "reason": "<critical path / ambiguity, or empty>" }
}
If verdict is ambiguous or the test is critical-path, do NOT quarantine — ESCALATE or route as REAL.
"""

SAMPLE_INPUT = """Failure: test_async_notify failed in CI. No changes to the notify module in this PR.
"""

# ---- MOCK tools (stubs — replace with real integrations to deploy) ----
def get_failure(args=None):
    print(f"  [MOCK] get_failure(" + str(args) + ")")
    return {"status": "mock", "tool": "get_failure"}

def rerun_test(args=None):
    print(f"  [MOCK] rerun_test(" + str(args) + ")")
    return {"status": "mock", "tool": "rerun_test"}

def failure_history(args=None):
    print(f"  [MOCK] failure_history(" + str(args) + ")")
    return {"status": "mock", "tool": "failure_history"}

def analyze_pattern(args=None):
    print(f"  [MOCK] analyze_pattern(" + str(args) + ")")
    return {"status": "mock", "tool": "analyze_pattern"}

def change_correlation(args=None):
    print(f"  [MOCK] change_correlation(" + str(args) + ")")
    return {"status": "mock", "tool": "change_correlation"}

def quarantine_test(args=None):
    print(f"  [MOCK] quarantine_test(" + str(args) + ")")
    return {"status": "mock", "tool": "quarantine_test"}

def file_issue(args=None):
    print(f"  [MOCK] file_issue(" + str(args) + ")")
    return {"status": "mock", "tool": "file_issue"}

def escalate_to_owner(args=None):
    print(f"  [MOCK] escalate_to_owner(" + str(args) + ")")
    return {"status": "mock", "tool": "escalate_to_owner"}

TOOLS = {
    "get_failure": get_failure,
    "rerun_test": rerun_test,
    "failure_history": failure_history,
    "analyze_pattern": analyze_pattern,
    "change_correlation": change_correlation,
    "quarantine_test": quarantine_test,
    "file_issue": file_issue,
    "escalate_to_owner": escalate_to_owner,
}


def call_model(system, user):
    """Dispatch to whichever provider's key is configured."""
    if os.environ.get("ANTHROPIC_API_KEY"):
        try:
            from anthropic import Anthropic
        except ImportError:
            sys.exit("Run: pip install anthropic")
        msg = Anthropic().messages.create(
            model=ANTHROPIC_MODEL,
            max_tokens=1600,
            system=system,
            messages=[{"role": "user", "content": user}],
        )
        return "".join(b.text for b in msg.content if getattr(b, "type", "") == "text")
    if os.environ.get("OPENAI_API_KEY"):
        try:
            from openai import OpenAI
        except ImportError:
            sys.exit("Run: pip install openai")
        resp = OpenAI().chat.completions.create(
            model=OPENAI_MODEL,
            max_tokens=1600,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
        )
        return resp.choices[0].message.content
    sys.exit("Set ANTHROPIC_API_KEY or OPENAI_API_KEY first.")


def extract_json(text):
    a, b = text.find("{"), text.rfind("}")
    if a == -1 or b == -1:
        return None
    try:
        return json.loads(text[a : b + 1])
    except json.JSONDecodeError:
        return None


def main():
    print("=" * 70)
    print("INPUT:\n" + SAMPLE_INPUT)
    print("=" * 70)

    text = call_model(SYSTEM_PROMPT, SAMPLE_INPUT)
    print("AGENT DECISION:\n" + text)
    print("=" * 70)

    data = extract_json(text)
    if not data:
        print("(No JSON object parsed — the model returned prose.)")
        return

    actions = data.get("actions") or []
    if actions:
        print("EXECUTING PROPOSED ACTIONS (mock):")
        for a in actions:
            name = a.get("tool")
            args = a.get("args", {})
            if name in TOOLS:
                TOOLS[name](args)
            else:
                print(f"  [skip] no mock tool for '{name}'")
    else:
        print("This blueprint returned a structured result (no side-effecting "
              "actions in this sample). Available mock tools: "
              + ", ".join(TOOLS.keys()))


if __name__ == "__main__":
    main()
