Claude Agents Architecture & Orchestration

import asyncio
import json
import os
import shutil
import tempfile
import unittest
from pathlib import Path
from typing import Any, Dict, List
from anthropic import AsyncAnthropic

# Initialize Async Client (Requires ANTHROPIC_API_KEY in environment)
client = AsyncAnthropic(api_key=os.getenv("ANTHROPIC_API_KEY", "sk-ant-mock-key"))

# Explicit 5.5 Model Tier Registry
MODEL_TIERS = {
    "opus_supervisor": "claude-opus-5-5",
    "sonnet_patcher": "claude-sonnet-5-5",
    "haiku_triage": "claude-haiku-5-5",
    "claude_code_subagent": "claude-sonnet-5-5"
}

# ---------------------------------------------------------------------------
# 1. CLAUDE CODE SUBAGENT GENERATOR (.claude/agents/*.md)
# ---------------------------------------------------------------------------
def provision_claude_code_subagents(workspace_dir: Path) -> List[Path]:
    """
    Creates declarative Claude Code subagent markdown files in `.claude/agents/`.
    Claude Code automatically discovers these files for `/agents` and Task delegation.
    """
    agents_dir = workspace_dir / ".claude" / "agents"
    agents_dir.mkdir(parents=True, exist_ok=True)

    subagent_specs = {
        "log-triager.md": (
            "---\n"
            "name: log-triager\n"
            "description: High-speed CI/CD stacktrace analyzer and root-cause classifier. Invoke first on build failures.\n"
            "model: claude-haiku-5-5\n"
            "tools: Read, Grep, Glob\n"
            "---\n"
            "You are a Claude Haiku 5.5 Telemetry Triager. Inspect logs, isolate the exact file and line number, "
            "and return a concise JSON root-cause summary without modifying files.\n"
        ),
        "ast-patcher.md": (
            "---\n"
            "name: ast-patcher\n"
            "description: Autonomous code repair specialist for surgical AST patches and unit test verification.\n"
            "model: claude-sonnet-5-5\n"
            "tools: Read, Edit, Write, Bash, Grep\n"
            "---\n"
            "You are a Claude Sonnet 5.5 Senior Patch Engineer. Apply minimal, non-breaking code diffs to resolve "
            "the root cause identified by `log-triager` and verify with pytest.\n"
        ),
        "security-auditor.md": (
            "---\n"
            "name: security-auditor\n"
            "description: Deep architectural and security gatekeeper for reviewing critical diffs prior to merge.\n"
            "model: claude-opus-5-5\n"
            "tools: Read, Grep, Bash\n"
            "---\n"
            "You are a Claude Opus 5.5 Principal Security Auditor. Deliberate over edge cases, race conditions, "
            "and injection vectors before approving any patch.\n"
        )
    }

    created_files: List[Path] = []
    for filename, content in subagent_specs.items():
        file_path = agents_dir / filename
        file_path.write_text(content, encoding="utf-8")
        created_files.append(file_path)
    return created_files

# ---------------------------------------------------------------------------
# 2. PROGRAMMATIC API & HEADLESS CLAUDE CODE SUBAGENT DISPATCHER
# ---------------------------------------------------------------------------
async def run_subagent(role: str, objective: str) -> Dict[str, Any]:
    """Spawns an isolated subagent mapped to Haiku 5.5, Sonnet 5.5, or Opus 5.5."""
    selected_model = MODEL_TIERS.get(role, "claude-sonnet-5-5")

    # If live API key is configured, invoke Anthropic Messages API with Prompt Caching
    if os.getenv("ANTHROPIC_API_KEY"):
        thinking_config = (
            {"type": "enabled", "budget_tokens": 4096}
            if selected_model in ("claude-opus-5-5", "claude-sonnet-5-5")
            else {"type": "disabled"}
        )
        resp = await client.messages.create(
            model=selected_model,
            max_tokens=8192,
            thinking=thinking_config,
            system=[{
                "type": "text",
                "text": f"You are specialized subagent [{role}] running on {selected_model}. Output strict JSON.",
                "cache_control": {"type": "ephemeral"}
            }],
            messages=[{"role": "user", "content": objective}]
        )
        text_blocks = [b.text for b in resp.content if b.type == "text"]
        return {"role": role, "model": selected_model, "output": "\n".join(text_blocks)}

    # Deterministic local fallback for zero-cost CI unit testing
    if role == "haiku_triage":
        return {
            "role": role,
            "model": selected_model,
            "root_cause": "ZeroDivisionError in calculate_burn_rate when elapsed_hours == 0"
        }
    elif role == "opus_supervisor":
        return {
            "role": role,
            "model": selected_model,
            "dag_plan": ["1. Triage via Haiku 5.5", "2. Patch via Sonnet 5.5", "3. Verify via Claude Code CLI"],
            "security_signoff": "APPROVED"
        }
    return {
        "role": role,
        "model": selected_model,
        "patched_code": (
            "def calculate_burn_rate(tokens: int, elapsed_hours: float) -> float:\n"
            "    return round(tokens / max(elapsed_hours, 0.0001), 2)"
        )
    }

async def invoke_headless_claude_code_cli(prompt: str, agent_name: str = "ast-patcher") -> Dict[str, Any]:
    """
    Demonstrates spawning a headless Claude Code CLI subagent subprocess:
    `claude -p "<prompt>" --agent <agent_name> --output-format json --max-turns 3`
    """
    cli_binary = shutil.which("claude")
    cmd = [
        cli_binary or "claude",
        "-p", prompt,
        "--agent", agent_name,
        "--model", MODEL_TIERS["claude_code_subagent"],
        "--output-format", "json",
        "--max-turns", "3"
    ]
    # In environments without the CLI binary installed, return structured telemetry
    if not cli_binary or not os.getenv("ANTHROPIC_API_KEY"):
        return {
            "execution_mode": "simulated_headless_cli",
            "command": " ".join(cmd),
            "subagent": agent_name,
            "model": MODEL_TIERS["claude_code_subagent"],
            "status": "READY"
        }

    proc = await asyncio.create_subprocess_exec(
        *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await proc.communicate()
    return {
        "execution_mode": "live_headless_cli",
        "returncode": proc.returncode,
        "stdout": stdout.decode("utf-8", errors="replace"),
        "stderr": stderr.decode("utf-8", errors="replace")
    }

def execute_sandbox_verification(patched_code: str) -> Dict[str, Any]:
    """Executes runtime regression checks against the Sonnet 5.5 patched code block."""
    namespace: Dict[str, Any] = {}
    try:
        exec(patched_code, namespace)
        fn = namespace.get("calculate_burn_rate")
        assert fn is not None, "Missing calculate_burn_rate function"
        assert fn(1000, 2.0) == 500.0, "Standard burn-rate math check failed"
        assert fn(500, 0.0) > 0, "Zero-division guardrail failed"
        return {"status": "PASSED", "assertions_verified": 3}
    except Exception as exc:
        return {"status": "FAILED", "error": str(exc)}

# ---------------------------------------------------------------------------
# 3. OPUS 5.5 SUPERVISOR ORCHESTRATION LOOP
# ---------------------------------------------------------------------------
async def orchestrate_healing_swarm(incident_payload: str, workspace_dir: Path) -> Dict[str, Any]:
    """
    Opus 5.5 Supervisor provisions `.claude/agents/*.md`, dispatches Haiku 5.5 and Sonnet 5.5
    subagents in parallel, invokes headless Claude Code CLI, and validates the fix.
    """
    # Step 1: Provision declarative Claude Code subagents on disk
    agent_files = provision_claude_code_subagents(workspace_dir)

    # Step 2: Opus 5.5 Supervisor formulates the remediation DAG
    supervisor_plan = await run_subagent("opus_supervisor", f"Plan remediation for: {incident_payload}")

    # Step 3: Parallel Fan-Out across Haiku 5.5 (Triage), Sonnet 5.5 (Patch), & Claude Code CLI
    triage_task = run_subagent("haiku_triage", f"Analyze stacktrace: {incident_payload}")
    patch_task = run_subagent("sonnet_patcher", f"Generate safe AST patch for: {incident_payload}")
    cli_task = invoke_headless_claude_code_cli("Verify AST patch and run test suite", agent_name="ast-patcher")

    triage_res, patch_res, cli_res = await asyncio.gather(triage_task, patch_task, cli_task)
    patch_code = patch_res.get("patched_code", "")

    # Step 4: Execute sandbox regression gate
    verification = execute_sandbox_verification(patch_code)
    return {
        "supervisor": supervisor_plan,
        "claude_code_agents_provisioned": [f.name for f in agent_files],
        "triage_report": triage_res,
        "patch_worker": patch_res,
        "claude_code_cli": cli_res,
        "verification": verification
    }

# ---------------------------------------------------------------------------
# 4. EMBEDDED AUTOMATED TEST SUITE (Run via `python3 claude_5_5_self_healing_swarm.py`)
# ---------------------------------------------------------------------------
class TestClaude55AgentSwarm(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self.temp_dir = Path(tempfile.mkdtemp(prefix="claude_code_swarm_"))

    def tearDown(self):
        shutil.rmtree(self.temp_dir, ignore_errors=True)

    async def test_full_5_5_triad_orchestration(self):
        """Verifies Opus 5.5 Supervisor, Haiku 5.5 Triager, Sonnet 5.5 Patcher, and Claude Code subagents."""
        res = await orchestrate_healing_swarm(
            "ZeroDivisionError: division by zero in metrics.py",
            self.temp_dir
        )
        self.assertEqual(res["supervisor"]["model"], "claude-opus-5-5")
        self.assertEqual(res["triage_report"]["model"], "claude-haiku-5-5")
        self.assertEqual(res["patch_worker"]["model"], "claude-sonnet-5-5")
        self.assertEqual(res["verification"]["status"], "PASSED")
        self.assertIn("log-triager.md", res["claude_code_agents_provisioned"])
        self.assertIn("ast-patcher.md", res["claude_code_agents_provisioned"])
        self.assertIn("security-auditor.md", res["claude_code_agents_provisioned"])

    def test_claude_code_subagent_markdown_frontmatter(self):
        """Verifies `.claude/agents/*.md` files contain valid YAML frontmatter and 5.5 model bindings."""
        files = provision_claude_code_subagents(self.temp_dir)
        self.assertEqual(len(files), 3)
        content_map = {f.name: f.read_text(encoding="utf-8") for f in files}
        self.assertIn("model: claude-haiku-5-5", content_map["log-triager.md"])
        self.assertIn("model: claude-sonnet-5-5", content_map["ast-patcher.md"])
        self.assertIn("model: claude-opus-5-5", content_map["security-auditor.md"])

    def test_sandbox_regression_guard(self):
        """Verifies sandbox rejects broken patches and passes valid patches."""
        bad_code = "def calculate_burn_rate(tokens, elapsed_hours):\n    return tokens / elapsed_hours"
        self.assertEqual(execute_sandbox_verification(bad_code)["status"], "FAILED")

        good_code = "def calculate_burn_rate(tokens, elapsed_hours):\n    return round(tokens / max(elapsed_hours, 0.0001), 2)"
        self.assertEqual(execute_sandbox_verification(good_code)["status"], "PASSED")

if __name__ == "__main__":
    print("[*] Executing Claude 5.5 Triad Swarm & Claude Code Subagent Verification Suite...")
    unittest.main(verbosity=2)