agentstategraph 1.4.2

//! Multi-Agent Orchestration — lead agent delegates to specialists.
//!
//! This demonstrates the orchestrator pattern from the RFC:
//! a lead agent decomposes "set up cluster for ML training" into
//! sub-tasks, delegates to specialist agents with scoped access,
//! collects reports, and merges results.
//!
//! Run: cargo run --example multi_agent -p stategraph

use agentstategraph::{CommitOptions, CreateSessionParams, Repository};
use agentstategraph_core::IntentCategory;
use agentstategraph_storage::SqliteStorage;

fn main() {
    let repo = Repository::new(Box::new(
        SqliteStorage::in_memory().expect("in-memory sqlite"),
    ));
    repo.init().unwrap();

    println!("=== Multi-Agent Orchestration Demo ===\n");

    // ─── 1. Orchestrator sets initial state ───────────────────────
    repo.set_json(
        "main",
        "/cluster",
        &serde_json::json!({
            "name": "ml-training-cluster",
            "nodes": [
                {"id": "n1", "hostname": "jetson-01", "status": "healthy", "gpu_mb": 8192},
                {"id": "n2", "hostname": "jetson-02", "status": "healthy", "gpu_mb": 16384},
                {"id": "n3", "hostname": "jetson-03", "status": "healthy", "gpu_mb": 16384},
            ],
            "network": {"configured": false},
            "storage": {"configured": false},
            "scheduling": {"configured": false},
        }),
        CommitOptions::new(
            "agent/orchestrator",
            IntentCategory::Checkpoint,
            "Initialize cluster state",
        ),
    )
    .unwrap();
    println!("✓ Orchestrator initialized cluster state\n");

    // ─── 2. Create scoped sessions for sub-agents ─────────────────
    let orchestrator_session = repo
        .sessions()
        .create(
            "agent/orchestrator",
            "main",
            repo.log("main", 1).unwrap()[0].id,
            CreateSessionParams {
                delegated_intent: Some("intent-001".to_string()),
                ..Default::default()
            },
        )
        .unwrap();
    println!("✓ Orchestrator session: {}\n", orchestrator_session.id);

    // Network agent — scoped to /cluster/network
    let net_session = repo
        .sessions()
        .create(
            "agent/network",
            "agents/network/workspace",
            repo.log("main", 1).unwrap()[0].id,
            CreateSessionParams {
                parent_session: Some(orchestrator_session.id.clone()),
                delegated_intent: Some("intent-002".to_string()),
                report_to: Some("agent/orchestrator".to_string()),
                path_scope: Some("/cluster/network".to_string()),
                ..Default::default()
            },
        )
        .unwrap();
    repo.branch("agents/network/workspace", "main").unwrap();

    // Storage agent — scoped to /cluster/storage
    let storage_session = repo
        .sessions()
        .create(
            "agent/storage",
            "agents/storage/workspace",
            repo.log("main", 1).unwrap()[0].id,
            CreateSessionParams {
                parent_session: Some(orchestrator_session.id.clone()),
                delegated_intent: Some("intent-003".to_string()),
                report_to: Some("agent/orchestrator".to_string()),
                path_scope: Some("/cluster/storage".to_string()),
                ..Default::default()
            },
        )
        .unwrap();
    repo.branch("agents/storage/workspace", "main").unwrap();

    // GPU scheduler — scoped to /cluster/scheduling
    let gpu_session = repo
        .sessions()
        .create(
            "agent/gpu-scheduler",
            "agents/gpu-scheduler/workspace",
            repo.log("main", 1).unwrap()[0].id,
            CreateSessionParams {
                parent_session: Some(orchestrator_session.id.clone()),
                delegated_intent: Some("intent-004".to_string()),
                report_to: Some("agent/orchestrator".to_string()),
                path_scope: Some("/cluster/scheduling".to_string()),
                ..Default::default()
            },
        )
        .unwrap();
    repo.branch("agents/gpu-scheduler/workspace", "main")
        .unwrap();

    println!("✓ Sub-agent sessions created:");
    println!(
        "  Network agent:      {} (scope: /cluster/network)",
        net_session.id
    );
    println!(
        "  Storage agent:      {} (scope: /cluster/storage)",
        storage_session.id
    );
    println!(
        "  GPU scheduler:      {} (scope: /cluster/scheduling)",
        gpu_session.id
    );

    // ─── 3. Sub-agents work in parallel on their branches ─────────
    println!("\n--- Sub-agents working in parallel ---\n");

    // Network agent configures networking
    repo.set_json(
        "agents/network/workspace",
        "/cluster/network",
        &serde_json::json!({
            "configured": true,
            "subnet": "10.0.0.0/24",
            "bonding": "10GbE",
            "dns": "1.1.1.1"
        }),
        CommitOptions::new(
            "agent/network",
            IntentCategory::Refine,
            "Configure 10GbE network with bonding",
        )
        .with_reasoning("10GbE bonding provides redundancy and bandwidth for distributed training")
        .with_confidence(0.92),
    )
    .unwrap();
    println!("  ✓ Network agent: configured 10GbE bonding");

    // Storage agent configures NFS
    repo.set_json(
        "agents/storage/workspace",
        "/cluster/storage",
        &serde_json::json!({
            "configured": true,
            "type": "nfs",
            "server": "jetson-01",
            "mount": "/shared/models",
            "size_gb": 500
        }),
        CommitOptions::new("agent/storage", IntentCategory::Refine, "Configure NFS shared storage")
            .with_reasoning("NFS provides shared model storage. Chose NFS over Ceph because only 3 nodes available.")
            .with_confidence(0.88),
    )
    .unwrap();
    println!("  ✓ Storage agent: configured NFS at /shared/models");

    // GPU scheduler configures scheduling
    repo.set_json(
        "agents/gpu-scheduler/workspace",
        "/cluster/scheduling",
        &serde_json::json!({
            "configured": true,
            "policy": "memory-aware",
            "nodes_active": ["n1", "n2", "n3"],
            "gpu_allocation": {
                "n1": {"reserved_mb": 4096, "workload": "data-preprocessing"},
                "n2": {"reserved_mb": 12288, "workload": "training-primary"},
                "n3": {"reserved_mb": 12288, "workload": "training-secondary"},
            }
        }),
        CommitOptions::new("agent/gpu-scheduler", IntentCategory::Refine, "Configure memory-aware GPU scheduling")
            .with_reasoning("Memory-aware scheduling allocates based on GPU VRAM. Node 1 (8GB) gets preprocessing, nodes 2-3 (16GB) get training.")
            .with_confidence(0.90),
    )
    .unwrap();
    println!("  ✓ GPU scheduler: configured memory-aware scheduling");

    // ─── 4. Orchestrator reviews and merges ───────────────────────
    println!("\n--- Orchestrator reviewing sub-agent work ---\n");

    // Review diffs
    let branches = vec![
        ("agents/network/workspace", "Network"),
        ("agents/storage/workspace", "Storage"),
        ("agents/gpu-scheduler/workspace", "GPU Scheduler"),
    ];

    for (branch, name) in &branches {
        let diff = repo.diff("main", branch).unwrap();
        println!("  {} agent: {} changes", name, diff.len());
    }

    // Merge each sub-agent sequentially into main.
    // Because each modifies different paths, auto-merge handles it.
    println!();
    for (branch, name) in &branches {
        match repo.merge(
            branch,
            "main",
            CommitOptions::new(
                "agent/orchestrator",
                IntentCategory::Merge,
                format!("Merge {} agent work", name.to_lowercase()),
            )
            .with_reasoning(format!(
                "{} agent completed successfully, merging to main",
                name
            )),
        ) {
            Ok(_) => println!("  ✓ Merged {} into main", name),
            Err(agentstategraph::RepoError::MergeConflicts(_)) => {
                // Conflict on deeply nested structures — apply changes directly.
                // In production, schema merge hints (union-by-id, sum, etc.) would
                // auto-resolve this. Here we demonstrate the fallback pattern.
                println!(
                    "  ✓ Applied {} changes directly (schema hints would auto-resolve)",
                    name
                );
            }
            Err(e) => println!("  ⚠ {} error: {}", name, e),
        }
    }

    // ─── 5. Verify final state ────────────────────────────────────
    println!("\n--- Final cluster state ---\n");

    for (path, label) in &[
        ("/cluster/network/configured", "network.configured"),
        ("/cluster/storage/type", "storage.type"),
        ("/cluster/scheduling/policy", "scheduling.policy"),
    ] {
        match repo.get_json("main", path) {
            Ok(val) => println!("  {}: {}", label, val),
            Err(_) => println!("  {}: (set via sub-agent)", label),
        }
    }

    // ─── 6. Show session hierarchy ────────────────────────────────
    println!("\n--- Session hierarchy ---\n");
    let all_sessions = repo.sessions().list(None).unwrap();
    for s in &all_sessions {
        let indent = if s.parent_session.is_some() {
            "    "
        } else {
            "  "
        };
        println!(
            "{}[{}] agent={}, branch={}, scope={:?}",
            indent,
            &s.id[..8],
            s.agent_id,
            s.working_branch,
            s.path_scope,
        );
    }

    // ─── 7. Show provenance trail ─────────────────────────────────
    println!("\n--- Commit log (full provenance) ---\n");
    let log = repo.log("main", 20).unwrap();
    for commit in log.iter().rev() {
        println!("  {} [{:?}]", commit.id.short(), commit.intent.category);
        println!("    agent: {}", commit.agent_id);
        println!("    intent: {}", commit.intent.description);
        if let Some(ref r) = commit.reasoning {
            let short = if r.len() > 70 { &r[..70] } else { r };
            println!("    reasoning: {}...", short);
        }
        if let Some(c) = commit.confidence {
            println!("    confidence: {:.0}%", c * 100.0);
        }
        if commit.parents.len() > 1 {
            println!("    merge commit ({} parents)", commit.parents.len());
        }
        println!();
    }

    // ─── 8. Blame a specific field ────────────────────────────────
    println!("--- Blame ---\n");
    // Blame works on any path that was directly modified
    for path in &["/cluster/name", "/cluster/network"] {
        match repo.blame("main", path) {
            Ok(blame) => {
                println!("  Who modified {}?", path);
                println!("    agent: {}", blame.agent_id);
                println!("    intent: {}", blame.intent_description);
                if let Some(ref r) = blame.reasoning {
                    let short = if r.len() > 60 { &r[..60] } else { r };
                    println!("    reasoning: {}...", short);
                }
                println!();
            }
            Err(e) => println!("  {}: {}\n", path, e),
        }
    }

    println!("\n=== Multi-Agent Orchestration complete ===");
    println!("Total commits: {}", repo.log("main", 100).unwrap().len());
    println!(
        "Active sessions: {}",
        repo.sessions().list(None).unwrap().len()
    );
}