agcore/examples/sub_agent_dispatch_demo.rs

//! sub_agent_dispatch_demo —— SubAgent 并行派发示例。
//!
//! 演示：
//! 1. 创建父 session（"主编" agent）
//! 2. 并行 dispatch_all 3 个子 agent（研究员 / 写手 / 审校）
//! 3. 收集子任务结果
//! 4. 验证树形结构：children(parent_id) 应返回 3 个子 ID
//!
//! 运行：`cargo run --example sub_agent_dispatch_demo`

use std::sync::Arc;

use agcore::agent::{Agent, AgentBuilder};
use agcore::engine::{SessionManager, SubTaskResult};
use agcore::llm::hooks::HookExecutor;
use agcore::llm::mock::MockProvider;
use agcore::llm::types::Usage;
use agcore::llm::types::message::{ContentBlock, Message};
use agcore::llm::types::response_v2::{MessageResponse, StopReason};
use agcore::memory::store::InMemoryStore;
use agcore::tools::ToolRegistry;

struct EditorAgent;
struct ResearcherAgent;
struct WriterAgent;
struct ReviewerAgent;

impl Agent for EditorAgent {
    fn name(&self) -> &str {
        "editor"
    }
    fn system_prompt(&self) -> Option<&str> {
        Some("You are an editor coordinating a team.")
    }
}
impl Agent for ResearcherAgent {
    fn name(&self) -> &str {
        "researcher"
    }
    fn system_prompt(&self) -> Option<&str> {
        Some("You are a researcher. Provide 3 key findings.")
    }
}
impl Agent for WriterAgent {
    fn name(&self) -> &str {
        "writer"
    }
    fn system_prompt(&self) -> Option<&str> {
        Some("You are a writer. Draft a section.")
    }
}
impl Agent for ReviewerAgent {
    fn name(&self) -> &str {
        "reviewer"
    }
    fn system_prompt(&self) -> Option<&str> {
        Some("You are a reviewer. Check for accuracy.")
    }
}

fn assistant_text(text: &str) -> MessageResponse {
    MessageResponse {
        id: String::new(),
        model: String::new(),
        message: Message::Assistant {
            content: vec![ContentBlock::Text { text: text.into() }],
        },
        usage: Usage::from_input_output(8, 4),
        stop_reason: StopReason::Stop,
        extra: Default::default(),
    }
}

fn print_result(name: &str, r: &Result<SubTaskResult, agcore::engine::EngineError>) {
    match r {
        Ok(res) => println!(
            "  ✓ {name} (child_id={}): {} tokens",
            &res.child_id[..20.min(res.child_id.len())],
            res.usage.total().total_tokens,
        ),
        Err(e) => println!("  ✗ {name}: {e}"),
    }
}

#[tokio::main]
async fn main() {
    println!("=== SubAgent Dispatch Demo ===\n");

    let store: Arc<dyn agcore::memory::store::MemoryStore> = Arc::new(InMemoryStore::new());
    let provider = Arc::new(MockProvider::new(vec![
        assistant_text("Researcher: finding 1, 2, 3"),
        assistant_text("Writer: section drafted"),
        assistant_text("Reviewer: looks good"),
    ]));
    let bundle = Arc::new(
        AgentBuilder::new()
            .provider(provider)
            .tool_registry(Arc::new(ToolRegistry::new()))
            .hook_executor(Arc::new(HookExecutor::new()))
            .session_memory_backend(store.clone())
            .build()
            .expect("RuntimeBundle"),
    );

    let editor: Arc<dyn Agent> = Arc::new(EditorAgent);
    let researcher: Arc<dyn Agent> = Arc::new(ResearcherAgent);
    let writer: Arc<dyn Agent> = Arc::new(WriterAgent);
    let reviewer: Arc<dyn Agent> = Arc::new(ReviewerAgent);

    let sm = Arc::new(SessionManager::new(store));
    let parent_id = sm.create(editor, bundle.clone()).await.expect("create");
    println!("[1] parent session created: {parent_id}");

    // dispatch_all 3 个子 agent
    println!("[2] dispatching 3 sub-agents in parallel...\n");
    let results = sm
        .dispatch_all(
            &parent_id,
            vec![
                (researcher, "Research topic X".to_string()),
                (writer, "Draft intro section".to_string()),
                (reviewer, "Review draft".to_string()),
            ],
            agcore::engine::DispatchConfig::default(),
        )
        .await;

    print_result("researcher", &results[0]);
    print_result("writer", &results[1]);
    print_result("reviewer", &results[2]);

    let success_count = results.iter().filter(|r| r.is_ok()).count();
    assert_eq!(success_count, 3, "all 3 should succeed");

    // 验证树形
    let children = sm.children(&parent_id).await.expect("children");
    println!("\n[3] children(parent) = {} session(s)", children.len());
    assert_eq!(children.len(), 3);

    println!("\n✓ dispatch_all completed: 3/3 sub-agents succeeded");
}