ZeroClaw 에이전트 런타임의 메모리 안전성과 효율적인 리소스 관리

use tokio::sync::{mpsc, RwLock};
use std::sync::Arc;
use std::time::Duration;

// 에이전트가 처리할 명령 타입 정의
#[derive(Debug)]
enum AgentCommand {
    ProcessTask(String),
    UpdateStatus(String),
    Shutdown,
}

// 에이전트의 상태 구조체
struct AgentState {
    id: String,
    status: String,
    processed_tasks: u64,
}

// 에이전트 실행기 구조체
struct AgentExecutor {
    state: Arc<RwLock<AgentState>>,
    receiver: mpsc::Receiver<AgentCommand>,
}

impl AgentExecutor {
    // 새로운 에이전트 생성자
    fn new(id: String, receiver: mpsc::Receiver<AgentCommand>) -> Self {
        Self {
            state: Arc::new(RwLock::new(AgentState {
                id,
                status: "Initialized".to_string(),
                processed_tasks: 0,
            })),
            receiver,
        }
    }

    // 메시지 수신 및 처리 루프 시작
    async fn run(mut self) {
        println!("Agent {} started.", self.state.read().await.id);
        
        while let Some(cmd) = self.receiver.recv().await {
            match cmd {
                AgentCommand::ProcessTask(task_id) => {
                    // 비동기 작업 시뮬레이션 (예: LLM 추론 요청)
                    let task_id_clone = task_id.clone();
                    let state_clone = Arc::clone(&self.state);
                    
                    // 백그라운드 작업으로 처리하여 메시지 루프 차단 방지
                    tokio::spawn(async move {
                        tokio::time::sleep(Duration::from_millis(100)).await;
                        let mut state = state_clone.write().await;
                        state.processed_tasks += 1;
                        state.status = format!("Processing {}", task_id_clone);
                        println!("Task {} processed by Agent {}. Total: {}", 
                            task_id_clone, state.id, state.processed_tasks);
                    });
                }
                AgentCommand::UpdateStatus(new_status) => {
                    let mut state = self.state.write().await;
                    state.status = new_status;
                }
                AgentCommand::Shutdown => {
                    println!("Agent {} shutting down...", self.state.read().await.id);
                    break;
                }
            }
        }
    }
}

#[tokio::main]
async fn main() {
    // 여러 에이전트를 관리하기 위한 송신자(Sender) 목록 저장
    // 에이전트가 종료될 때를 대비해 Vec으로 관리
    let mut agent_senders = Vec::new();

    // 3개의 에이전트 스폰
    for i in 0..3 {
        let (tx, rx) = mpsc::channel(100); // 버퍼 크기 100
        agent_senders.push(tx);
        
        let executor = AgentExecutor::new(format!("Agent-{}", i), rx);
        tokio::spawn(executor.run());
    }

    // 시스템 전체의 종료 신항 (Ctrl+C 등 대응)
    let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);
    
    // 작업 분산 로직 (시뮬레이션)
    let task_distributor = tokio::spawn(async move {
        let mut task_counter = 0;
        loop {
            // 종료 신호 확인
            if shutdown_rx.try_recv().is_ok() {
                println!("Task distributor stopping...");
                break;
            }

            // 라운드 로빈 방식으로 에이전트들에게 작업 전송
            if !agent_senders.is_empty() {
                let target_index = task_counter % agent_senders.len();
                let task_id = format!("Task-{}", task_counter);
                
                if let Err(_) = agent_senders[target_index].send(AgentCommand::ProcessTask(task_id)).await {
                    println!("Failed to send task. Agent might be dead.");
                }
                
                task_counter += 1;
                tokio::time::sleep(Duration::from_millis(50)).await;
            }
        }
    });

    // 5초 후 시스템 종료 시뮬레이션
    tokio::time::sleep(Duration::from_secs(5)).await;
    
    // 1. 작업 분배 종료
    let _ = shutdown_tx.send(()).await;
    task_distributor.await.unwrap();

    // 2. 모든 에이전트에게 종료 명령 전송
    for tx in agent_senders {
        let _ = tx.send(AgentCommand::Shutdown).await;
    }

    // 리소스 정리 대기
    tokio::time::sleep(Duration::from_millis(500)).await;
    println!("System shutdown complete.");
}