%pip install -q "anthropic>=0.91.0" python-dotenv

import json
import os
import time
from pathlib import Path
 
from anthropic import Anthropic
from dotenv import load_dotenv
from utilities import wait_for_idle_status
 
load_dotenv()
client = Anthropic()
MODEL = os.getenv("COOKBOOK_MODEL", "claude-opus-4-6")
FIXTURE = Path("example_data/sre")

# A real skill is usually a folder on disk (SKILL.md plus any helper
# scripts or reference docs) that you zip and upload. For this tutorial
# the SKILL.md is small enough to keep inline.
RUNBOOK_SKILL = """\
---
name: incident-runbooks
description: How to triage production incidents using the team runbooks.
---
 
# Incident runbooks
 
When an alert references a service, locate that service's recent logs
and identify the failure signature (the repeating error class, exit
code, or status pattern).
 
Consult the team runbooks before proposing any fix. Runbooks are
organised by failure signature — for example `oom.md`, `5xx.md`,
`latency.md`. Each one lists the triage steps for that class of
failure and the configuration that usually needs to change.
 
Any fix to infrastructure code must be opened as a pull request that
cites the runbook you followed. Do not patch live resources directly.
"""
 
skill = client.beta.skills.create(
    display_title="incident-runbooks",
    files=[("incident-runbooks/SKILL.md", RUNBOOK_SKILL.encode(), "text/markdown")],
)
print(f"skill: {skill.id} (version {skill.latest_version})")

SRE_SYSTEM_PROMPT = """\
You are an on-call SRE agent. Each user message is a PagerDuty alert
payload. Triage it to root cause and ship the minimal safe fix.
 
The session workspace contains the recent logs, the infrastructure
repo, and the team runbooks for the alerting service. Explore it to
find what you need.
 
Workflow for every alert:
1. Read the logs and identify the failure signature.
2. Find the root cause in the infrastructure repo, save a copy of the
   original file, edit it in place, then produce a unified diff with
   `diff -u`.
3. open_pull_request(title, body, diff) with the fix.
4. request_approval(summary) and wait for the human's decision.
5. Only if the result is "approved", merge_pull_request(pr_number).
   Otherwise stop and report.
 
Never call merge_pull_request unless request_approval returned
"approved". Keep the fix minimal — do not refactor unrelated config.
"""
 
agent = client.beta.agents.create(
    name="cookbook-sre-responder",
    model=MODEL,
    system=SRE_SYSTEM_PROMPT,
    skills=[{"type": "custom", "skill_id": skill.id, "version": skill.latest_version}],
    tools=[
        {
            "type": "agent_toolset_20260401",
            "default_config": {
                "enabled": True,
                "permission_policy": {"type": "always_allow"},
            },
            "configs": [
                {"name": "web_search", "enabled": False},
                {"name": "web_fetch", "enabled": False},
            ],
        },
        {
            "type": "custom",
            "name": "open_pull_request",
            "description": "Open a pull request against the infra repo with the proposed fix.",
            "input_schema": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "body": {"type": "string"},
                    "diff": {"type": "string", "description": "Unified diff of the change."},
                },
                "required": ["title", "body", "diff"],
            },
        },
        {
            "type": "custom",
            "name": "request_approval",
            "description": "Ask the on-call human to approve the proposed PR before merging.",
            "input_schema": {
                "type": "object",
                "properties": {
                    "summary": {"type": "string"},
                },
                "required": ["summary"],
            },
        },
        {
            "type": "custom",
            "name": "merge_pull_request",
            "description": "Merge an approved pull request.",
            "input_schema": {
                "type": "object",
                "properties": {
                    "pr_number": {"type": "integer"},
                },
                "required": ["pr_number"],
            },
        },
    ],
)
print(f"agent: {agent.id} v{agent.version}")

env = client.beta.environments.create(
    name="cookbook-sre-env",
    config={"type": "cloud", "networking": {"type": "limited"}},
)
 
 
def upload(path: Path, mime: str) -> str:
    with path.open("rb") as f:
        return client.beta.files.upload(file=(path.name, f, mime)).id
 
 
log_id = upload(FIXTURE / "logs/checkout-svc.log", "text/plain")
manifest_id = upload(FIXTURE / "infra/k8s/checkout-deploy.yaml", "text/yaml")
runbook_id = upload(FIXTURE / "runbooks/oom.md", "text/markdown")
 
RESOURCES = [
    {"type": "file", "file_id": log_id, "mount_path": "logs/checkout-svc.log"},
    {"type": "file", "file_id": manifest_id, "mount_path": "infra/k8s/checkout-deploy.yaml"},
    {"type": "file", "file_id": runbook_id, "mount_path": "runbooks/oom.md"},
]
print(f"environment: {env.id}")

def handle_pagerduty_webhook(payload: dict) -> str:
    incident = payload["event"]["data"]
    session = client.beta.sessions.create(
        environment_id=env.id,
        agent={"type": "agent", "id": agent.id, "version": agent.version},
        resources=RESOURCES,
        title=f"[{incident['service']['summary']}] {incident['title']}",
    )
    client.beta.sessions.events.send(
        session.id,
        events=[
            {
                "type": "user.message",
                "content": [{"type": "text", "text": json.dumps(payload, indent=2)}],
            }
        ],
    )
    return session.id
 
 
with (FIXTURE / "alert.json").open() as f:
    alert = json.load(f)
 
session_id = handle_pagerduty_webhook(alert)
print(f"session: {session_id}")

prs: list[dict] = []
pending_approvals: list[dict] = []
seen_events: set[str] = set()
 
 
def handle_custom_tool(name: str, args: dict) -> dict:
    if name == "open_pull_request":
        n = len(prs) + 1
        prs.append({"number": n, "merged": False, **args})
        print(f"\n── PR #{n}: {args['title']} ──")
        return {"pr_number": n, "url": f"mock://infra/pull/{n}"}
    if name == "merge_pull_request":
        prs[args["pr_number"] - 1]["merged"] = True
        return {"merged": True}
    raise ValueError(f"unhandled tool {name}")
 
 
def run_until_approval_or_end(session_id: str) -> str | None:
    """Poll the session's event log, servicing custom tools, until either
    a request_approval call arrives (return its event_id so the caller
    can respond) or the agent ends its turn (return None)."""
    custom_calls: dict[str, object] = {}
    responded: set[str] = set()
    while True:
        idle_stop = None
        for ev in client.beta.sessions.events.list(session_id):
            if ev.id in seen_events:
                continue
            seen_events.add(ev.id)
            if ev.type == "agent.message":
                for block in ev.content:
                    if block.type == "text":
                        print(block.text, end="")
            elif ev.type == "agent.tool_use":
                print(f"\n  [{ev.name}]")
            elif ev.type == "agent.custom_tool_use":
                custom_calls[ev.id] = ev
                print(f"\n→ {ev.name}")
            elif ev.type == "session.status_idle":
                idle_stop = ev.stop_reason
            elif ev.type == "session.status_terminated":
                return None
        if idle_stop is None:
            time.sleep(1.0)
            continue
        if idle_stop.type == "end_turn":
            return None
        if idle_stop.type == "requires_action":
            for event_id in idle_stop.event_ids:
                if event_id in responded:
                    continue
                call = custom_calls[event_id]
                if call.name == "request_approval":
                    pending_approvals.append({"event_id": event_id, **call.input})
                    return event_id
                result = handle_custom_tool(call.name, call.input)
                client.beta.sessions.events.send(
                    session_id,
                    events=[
                        {
                            "type": "user.custom_tool_result",
                            "custom_tool_use_id": event_id,
                            "content": [{"type": "text", "text": json.dumps(result)}],
                        }
                    ],
                )
                responded.add(event_id)
        time.sleep(1.0)
 
 
approval_event_id = run_until_approval_or_end(session_id)

pr = prs[0]
print(pr["body"], "\n")
print(pr["diff"])
print("\n" + "─" * 60)
print("APPROVAL REQUESTED:", pending_approvals[0]["summary"])

client.beta.sessions.events.send(
    session_id,
    events=[
        {
            "type": "user.custom_tool_result",
            "custom_tool_use_id": approval_event_id,
            "content": [{"type": "text", "text": json.dumps({"decision": "approved"})}],
        }
    ],
)
 
run_until_approval_or_end(session_id)
print(f"\n\nPR #{pr['number']} merged: {prs[0]['merged']}")

wait_for_idle_status(client, session_id)
client.beta.sessions.archive(session_id)
client.beta.environments.archive(env.id)
client.beta.agents.archive(agent.id)
client.beta.skills.versions.delete(skill.latest_version, skill_id=skill.id)
client.beta.skills.delete(skill.id)
print("archived")