PicoBot/src/agent/sub_agent.rs

use std::sync::Arc;
use std::time::Instant;

use crate::agent::AgentError;
use crate::agent::AgentLoop;
use crate::agent::system_prompt::build_sub_agent_system_prompt;
use crate::bus::ChatMessage;
use crate::config::LLMProviderConfig;
use crate::providers::{LLMProvider, create_provider};
use crate::skills::SkillsLoader;
use crate::tools::{ToolExecutionContext, ToolRegistry};

const DEFAULT_MAX_ITERATIONS: usize = 99;

#[derive(Debug, Clone)]
pub struct SubAgentConfig {
    pub target: Option<String>,
    pub prompt: String,
    pub context: Option<String>,
    pub mode: ExecutionMode,
    pub allowed_tools: Option<Vec<String>>,
    pub max_iterations: Option<usize>,
    pub timeout_secs: Option<u64>,
    pub plan_item_id: Option<String>,
    pub session_id: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExecutionMode {
    Foreground,
    Background,
}

#[derive(Debug, Clone)]
pub struct SubAgentResult {
    pub task_id: String,
    /// Bounded projection returned to the model; may carry a truncation note.
    pub content: String,
    pub content_truncated: bool,
    /// Untruncated final text, persisted as the durable run result.
    pub full_content: String,
    pub status: TaskStatus,
    pub tool_calls_count: usize,
    pub iterations: usize,
    pub duration_ms: u64,
}

#[derive(Debug, Clone)]
pub enum TaskStatus {
    Completed,
    Failed(String),
    Cancelled,
    TimedOut,
}

#[derive(Debug)]
pub enum SubAgentError {
    ProviderCreation(String),
    Storage(String),
    Other(String),
}

impl std::fmt::Display for SubAgentError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::ProviderCreation(e) => write!(f, "provider creation failed: {}", e),
            Self::Storage(e) => write!(f, "storage error: {}", e),
            Self::Other(e) => write!(f, "{}", e),
        }
    }
}

impl std::error::Error for SubAgentError {}

pub struct SubAgentManager {
    provider_config: LLMProviderConfig,
    full_tools: Arc<ToolRegistry>,
    storage: Option<Arc<crate::storage::Storage>>,
    skills_loader: Option<Arc<SkillsLoader>>,
    work_manager: Option<Arc<crate::work::WorkManager>>,
    catalog: Arc<crate::agent::AgentCatalog>,
    execution_gate: Arc<crate::agent::gate::ExecutionGate>,
    /// Late-bound durable Coordinator.  Signals are only available to runs
    /// whose definition carries a signal contract AND the runtime has an
    /// active Coordinator; resolution happens at delegate time.
    coordinator: std::sync::RwLock<Option<std::sync::Weak<super::coordinator::AgentCoordinator>>>,
}

#[derive(Clone)]
pub(crate) struct ResolvedAgentRun {
    pub provider_config: Arc<LLMProviderConfig>,
    pub tools: Arc<ToolRegistry>,
    pub timeout_secs: u64,
    pub max_iterations: usize,
    pub max_result_chars: usize,
    pub role_prompt: Option<String>,
    pub skills_prompt: Option<String>,
    pub tool_context: ToolExecutionContext,
    /// Named-definition metadata used by the durable Coordinator; `None` for
    /// the legacy transient general Agent.
    pub agent_id: Option<String>,
    pub definition_hash: Option<String>,
    pub llm_profile: Option<String>,
    /// Durable signal contract of the definition; `None` means the run can
    /// never emit signals and must not persist contract state.
    pub signal_contract: Option<crate::agent::definition::SignalContract>,
}

impl SubAgentManager {
    pub fn new(
        provider_config: LLMProviderConfig,
        full_tools: Arc<ToolRegistry>,
        storage: Option<Arc<crate::storage::Storage>>,
        skills_loader: Option<Arc<SkillsLoader>>,
    ) -> Self {
        Self {
            provider_config,
            full_tools,
            storage,
            skills_loader,
            work_manager: None,
            catalog: Arc::new(crate::agent::AgentCatalog::legacy()),
            execution_gate: crate::agent::gate::ExecutionGate::unbounded(),
            coordinator: std::sync::RwLock::new(None),
        }
    }

    /// Bind the durable Coordinator so named runs can resolve the
    /// contract-bound `emit_signal` tool.  Kept as a weak reference: the
    /// Coordinator owns this manager, so a strong cycle must never exist.
    pub fn bind_coordinator(&self, coordinator: &Arc<super::coordinator::AgentCoordinator>) {
        *self.coordinator.write().unwrap() = Some(Arc::downgrade(coordinator));
    }

    fn coordinator(&self) -> Option<Arc<super::coordinator::AgentCoordinator>> {
        self.coordinator
            .read()
            .unwrap()
            .as_ref()
            .and_then(std::sync::Weak::upgrade)
    }

    pub fn with_catalog(mut self, catalog: Arc<crate::agent::AgentCatalog>) -> Self {
        self.catalog = catalog;
        self
    }

    pub fn with_execution_gate(
        mut self,
        execution_gate: Arc<crate::agent::gate::ExecutionGate>,
    ) -> Self {
        self.execution_gate = execution_gate;
        self
    }

    pub fn catalog(&self) -> Arc<crate::agent::AgentCatalog> {
        self.catalog.clone()
    }

    pub fn with_work_manager(mut self, work_manager: Arc<crate::work::WorkManager>) -> Self {
        self.work_manager = Some(work_manager);
        self
    }

    pub(crate) fn resolve_agent(
        &self,
        config: &SubAgentConfig,
        caller: &ToolExecutionContext,
        task_id: &str,
    ) -> Result<ResolvedAgentRun, SubAgentError> {
        let Some(target) = config.target.as_deref() else {
            return Err(SubAgentError::Other(
                "delegate requires a named target Agent; the legacy anonymous general Agent is no longer available"
                    .to_string(),
            ));
        };

        let definition = self
            .catalog
            .get(target)
            .ok_or_else(|| SubAgentError::Other(format!("unknown Agent target '{target}'")))?;
        let root_session_id = caller
            .agent
            .as_ref()
            .map(|context| context.root_session_id.clone())
            .or_else(|| caller.session_id.clone())
            .ok_or_else(|| {
                SubAgentError::Other(
                    "delegate requires a session-bound ToolExecutionContext".to_string(),
                )
            })?;
        let cancellation = caller.cancellation.child_token();
        let execution = if let Some(parent) = caller.agent.as_ref() {
            if !self.catalog.can_delegate(&parent.current_agent_id, target) {
                return Err(SubAgentError::Other(format!(
                    "Agent '{}' is not allowed to delegate to '{target}'",
                    parent.current_agent_id
                )));
            }
            if parent.ancestry.iter().any(|agent| agent == target) {
                return Err(SubAgentError::Other(format!(
                    "delegation cycle rejected: '{target}' is already in the current ancestry"
                )));
            }
            if parent.budget.remaining_runs == 0 || parent.budget.remaining_depth == 0 {
                return Err(SubAgentError::Other(
                    "delegation budget exhausted".to_string(),
                ));
            }
            let next_depth = parent.depth.saturating_add(1);
            if next_depth > self.catalog.max_tree_depth() {
                return Err(SubAgentError::Other(format!(
                    "delegation depth {next_depth} exceeds global limit {}",
                    self.catalog.max_tree_depth()
                )));
            }
            if parent
                .reserve_tree_run(self.catalog.max_runs_per_tree())
                .is_none()
            {
                return Err(SubAgentError::Other(format!(
                    "delegation tree already uses {} runs; max_runs_per_tree is {}",
                    self.catalog.max_runs_per_tree(),
                    self.catalog.max_runs_per_tree()
                )));
            }
            let mut child = crate::agent::AgentExecutionContext::child(
                parent,
                task_id.to_string(),
                target.to_string(),
                config.plan_item_id.clone(),
                cancellation.clone(),
            );
            child.budget.remaining_depth = child
                .budget
                .remaining_depth
                .min(definition.limits.max_depth);
            child.signal_contract = definition
                .signal_contract
                .as_ref()
                .map(|contract| Arc::new(contract.clone()));
            Arc::new(child)
        } else {
            if !self.catalog.root_can_delegate(target) {
                return Err(SubAgentError::Other(format!(
                    "ROOT is not allowed to delegate to '{target}'"
                )));
            }
            Arc::new(crate::agent::AgentExecutionContext {
                root_session_id: root_session_id.clone(),
                root_turn_id: caller.turn_id.clone(),
                run_id: task_id.to_string(),
                execution_id: task_id.to_string(),
                parent_run_id: None,
                caller_agent_id: "ROOT".to_string(),
                current_agent_id: target.to_string(),
                ancestry: vec![target.to_string()],
                depth: 1,
                plan_item_id: config.plan_item_id.clone(),
                cancellation: cancellation.clone(),
                budget: crate::agent::AgentBudget {
                    remaining_runs: self.catalog.max_runs_per_tree().saturating_sub(1),
                    remaining_depth: self
                        .catalog
                        .max_tree_depth()
                        .saturating_sub(1)
                        .min(definition.limits.max_depth),
                },
                tree_runs: Arc::new(std::sync::atomic::AtomicUsize::new(1)),
                signal_contract: definition
                    .signal_contract
                    .as_ref()
                    .map(|contract| Arc::new(contract.clone())),
                emitted_signals: Arc::new(std::sync::Mutex::new(Vec::new())),
            })
        };

        let mut effective_names = definition.tools.clone();
        if let Some(allowed) = config.allowed_tools.as_ref() {
            effective_names.retain(|name| allowed.iter().any(|allowed| allowed == name));
        }
        let has_get_skill = effective_names.iter().any(|name| name == "get_skill");
        let mut names = effective_names;
        names.retain(|name| name != "get_skill");
        let mut runtime_tools = Vec::new();
        let skills_prompt = if has_get_skill {
            let loader = self
                .skills_loader
                .as_ref()
                .ok_or_else(|| SubAgentError::Other("skills loader is unavailable".to_string()))?;
            runtime_tools.push(Arc::new(crate::tools::GetSkillTool::scoped(
                loader.clone(),
                &definition.skills,
            )) as Arc<dyn crate::tools::Tool>);
            let prompt = loader.build_scoped_skills_prompt(&definition.skills);
            (!prompt.is_empty()).then_some(prompt)
        } else {
            None
        };
        let delegate_targets = self.catalog.delegate_targets(target);
        if !delegate_targets.is_empty() {
            let delegate = self.full_tools.get("delegate").ok_or_else(|| {
                SubAgentError::Other("delegate runtime tool is unavailable".to_string())
            })?;
            runtime_tools.push(Arc::new(crate::tools::delegate::ScopedDelegateTool::new(
                delegate,
                delegate_targets,
            )) as Arc<dyn crate::tools::Tool>);
        }
        // The signal tool is contract-bound: it exists only when the
        // definition declares a signal block and the durable Coordinator is
        // live.  If either is missing the run cannot emit signals.
        if definition.signal_contract.is_some() {
            match self.coordinator() {
                Some(coordinator) => {
                    let contract = definition.signal_contract.clone().unwrap();
                    runtime_tools.push(Arc::new(crate::tools::EmitSignalTool::new(
                        coordinator,
                        Arc::new(contract),
                    )) as Arc<dyn crate::tools::Tool>);
                }
                None => {
                    return Err(SubAgentError::Other(
                        "Agent '{}' declares a signal contract but the durable Coordinator is unavailable"
                            .replace("{}", target),
                    ));
                }
            }
        }
        let tools = self
            .full_tools
            .scoped_for_agent(&names, runtime_tools)
            .map_err(SubAgentError::Other)?;

        Ok(ResolvedAgentRun {
            provider_config: definition.provider_config.clone(),
            tools,
            timeout_secs: definition.limits.timeout_secs,
            max_iterations: definition.limits.max_iterations,
            max_result_chars: definition.limits.max_result_chars,
            role_prompt: Some(definition.role_prompt.clone()),
            skills_prompt,
            agent_id: Some(target.to_string()),
            definition_hash: Some(definition.definition_hash.clone()),
            llm_profile: definition.llm_profile.clone(),
            signal_contract: definition.signal_contract.clone(),
            tool_context: ToolExecutionContext::for_session(format!("agent-run:{task_id}"))
                .with_turn_id(
                    caller
                        .turn_id
                        .clone()
                        .unwrap_or_else(|| task_id.to_string()),
                )
                .with_agent(execution)
                .with_cancellation(cancellation)
                .with_execution_gate(self.execution_gate.clone()),
        })
    }

    pub fn build_sub_agent(
        &self,
        config: &SubAgentConfig,
        tools: Arc<ToolRegistry>,
    ) -> Result<AgentLoop, AgentError> {
        self.build_sub_agent_with_provider(config, tools, &self.provider_config)
    }

    fn build_sub_agent_with_provider(
        &self,
        config: &SubAgentConfig,
        tools: Arc<ToolRegistry>,
        provider_config: &LLMProviderConfig,
    ) -> Result<AgentLoop, AgentError> {
        let mut provider = create_provider(provider_config.clone())
            .map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
        if let Some(ref s) = self.storage {
            provider.set_storage(s.clone());
        }
        let provider: Arc<dyn LLMProvider> = Arc::from(provider);

        let max_iterations = config.max_iterations.unwrap_or(DEFAULT_MAX_ITERATIONS);
        let workspace_dir = provider_config.workspace_dir.clone();
        let model_name = provider_config.model_id.clone();
        let input_types = provider_config.input_types.clone();

        let agent = AgentLoop::with_provider_and_tools(
            provider,
            tools,
            max_iterations,
            model_name,
            workspace_dir,
            input_types,
        )
        .with_context_window(provider_config.token_limit);

        Ok(agent)
    }

    /// Execute an already-resolved Agent.  The durable Coordinator owns run
    /// admission and terminal commits; this is the shared execution core.
    pub(crate) async fn execute_resolved(
        &self,
        config: &SubAgentConfig,
        resolved: ResolvedAgentRun,
        task_id: &str,
    ) -> Result<SubAgentResult, SubAgentError> {
        let tools = resolved.tools;
        let timeout_secs = resolved.timeout_secs;
        let timeout_human = format_duration(timeout_secs);
        let mut system_prompt = build_sub_agent_system_prompt(
            &config.prompt,
            &timeout_human,
            &tools,
            &resolved.provider_config.workspace_dir,
            &resolved.provider_config.model_id,
            resolved.skills_prompt,
        );
        if let Some(role_prompt) = resolved.role_prompt {
            system_prompt.push_str("\n\n## Agent Definition\n\n");
            system_prompt.push_str(&role_prompt);
        }
        if let Some(context) = config
            .context
            .as_deref()
            .filter(|value| !value.trim().is_empty())
        {
            system_prompt.push_str("\n\n## 调用方提供的任务上下文\n\n");
            system_prompt.push_str(context);
        }

        let mut effective_config = config.clone();
        effective_config.max_iterations = Some(resolved.max_iterations);
        let max_result_chars = resolved.max_result_chars;

        let (transcript_tx, transcript_rx) = tokio::sync::mpsc::unbounded_channel();
        let agent = self
            .build_sub_agent_with_provider(&effective_config, tools, &resolved.provider_config)
            .map_err(|e| SubAgentError::ProviderCreation(e.to_string()))?
            .with_transcript_sink(transcript_tx);

        let history = vec![
            ChatMessage::system(system_prompt),
            ChatMessage::user(&config.prompt),
        ];

        let start = Instant::now();
        let tool_context = resolved.tool_context;

        let writer = self.spawn_transcript_writer(task_id, transcript_rx);

        let outcome = tokio::select! {
            result = tokio::time::timeout(
                std::time::Duration::from_secs(timeout_secs),
                agent.process_with_context(history, tool_context.clone()),
            ) => match result {
                Ok(inner) => ExecutionOutcome::Finished(Box::new(inner)),
                Err(_elapsed) => ExecutionOutcome::TimedOut,
            },
            _ = tool_context.cancellation.cancelled() => ExecutionOutcome::Cancelled,
        };

        let duration_ms = start.elapsed().as_millis() as u64;

        // Drop the agent (which owns the transcript sender) so the writer can
        // drain, then await the writer before the caller's terminal commit so
        // the persisted transcript is complete first.
        drop(agent);
        if let Err(error) = writer.await {
            tracing::warn!(run_id = task_id, error = %error, "transcript writer failed");
        }

        Ok(match outcome {
            ExecutionOutcome::Finished(result) => match *result {
                Ok(agent_result) => {
                    let (content, truncated) = truncate_sub_agent_result_at(
                        &agent_result.final_response.content,
                        max_result_chars,
                    );
                    let tool_calls_count = agent_result
                        .emitted_messages
                        .iter()
                        .filter(|m| m.tool_calls.is_some())
                        .count();
                    let iterations = agent_result
                        .emitted_messages
                        .iter()
                        .filter(|m| m.role == "assistant" && m.tool_calls.is_some())
                        .count();
                    SubAgentResult {
                        task_id: task_id.to_string(),
                        content,
                        content_truncated: truncated,
                        full_content: agent_result.final_response.content,
                        status: TaskStatus::Completed,
                        tool_calls_count,
                        iterations,
                        duration_ms,
                    }
                }
                Err(error) => SubAgentResult {
                    task_id: task_id.to_string(),
                    content: String::new(),
                    content_truncated: false,
                    full_content: String::new(),
                    status: terminal_status_from_error(error),
                    tool_calls_count: 0,
                    iterations: 0,
                    duration_ms,
                },
            },
            ExecutionOutcome::TimedOut => SubAgentResult {
                task_id: task_id.to_string(),
                content: String::new(),
                content_truncated: false,
                full_content: String::new(),
                status: TaskStatus::TimedOut,
                tool_calls_count: 0,
                iterations: 0,
                duration_ms,
            },
            ExecutionOutcome::Cancelled => SubAgentResult {
                task_id: task_id.to_string(),
                content: String::new(),
                content_truncated: false,
                full_content: String::new(),
                status: TaskStatus::Cancelled,
                tool_calls_count: 0,
                iterations: 0,
                duration_ms,
            },
        })
    }

    /// Spawn a task that drains the transcript channel into
    /// `agent_run_messages`, assigning a monotonically increasing `seq` and
    /// stripping `provider_state` (which must never be persisted or exposed).
    /// With no storage the writer becomes a drain-and-discard no-op.
    fn spawn_transcript_writer(
        &self,
        run_id: &str,
        receiver: tokio::sync::mpsc::UnboundedReceiver<ChatMessage>,
    ) -> tokio::task::JoinHandle<()> {
        let storage = self.storage.clone();
        let run_id = run_id.to_string();
        tokio::spawn(async move {
            let Some(storage) = storage else {
                let mut receiver = receiver;
                while receiver.recv().await.is_some() {}
                return;
            };
            let mut seq = 0i64;
            let mut receiver = receiver;
            while let Some(mut message) = receiver.recv().await {
                message.provider_state = None;
                let now = chrono::Utc::now().timestamp_millis();
                if let Err(error) = storage
                    .append_agent_run_message(&run_id, seq, &message, now)
                    .await
                {
                    tracing::warn!(
                        run_id = %run_id,
                        seq,
                        error = %error,
                        "failed to append transcript message"
                    );
                }
                seq += 1;
            }
        })
    }
}

/// Intermediate outcome of a resolved run, unified so the transcript writer
/// is awaited on every exit path before `execute_resolved` returns.
enum ExecutionOutcome {
    Finished(Box<Result<crate::agent::AgentProcessResult, AgentError>>),
    TimedOut,
    Cancelled,
}

fn terminal_status_from_error(error: AgentError) -> TaskStatus {
    match error {
        AgentError::Cancelled => TaskStatus::Cancelled,
        AgentError::TimedOut => TaskStatus::TimedOut,
        other => TaskStatus::Failed(other.to_string()),
    }
}

fn format_duration(seconds: u64) -> String {
    if seconds < 60 {
        format!("{}s", seconds)
    } else if seconds < 3600 {
        format!("{}m", seconds / 60)
    } else {
        format!("{}h", seconds / 3600)
    }
}

fn truncate_sub_agent_result_at(content: &str, max_chars: usize) -> (String, bool) {
    if content.len() <= max_chars {
        (content.to_string(), false)
    } else {
        let truncate_at = content.floor_char_boundary(max_chars);
        (
            format!(
                "{}\n\n[... 结果已截断，共 {} 字符，完整结果请使用 check_task 查看 ...]",
                &content[..truncate_at],
                content.len()
            ),
            true,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    fn manager() -> SubAgentManager {
        SubAgentManager::new(
            LLMProviderConfig {
                provider_type: "openai".into(),
                name: "test".into(),
                base_url: "http://localhost".into(),
                api_key: "test".into(),
                extra_headers: HashMap::new(),
                model_id: "test".into(),
                temperature: None,
                max_tokens: None,
                model_extra: HashMap::new(),
                max_tool_iterations: 1,
                token_limit: 4096,
                workspace_dir: std::env::temp_dir(),
                input_types: vec!["text".into()],
                price_input_per_million: None,
                price_output_per_million: None,
            },
            Arc::new(ToolRegistry::new()),
            None,
            None,
        )
    }

    fn config(target: Option<&str>) -> SubAgentConfig {
        SubAgentConfig {
            target: target.map(str::to_string),
            prompt: "test".into(),
            context: None,
            mode: ExecutionMode::Foreground,
            allowed_tools: None,
            max_iterations: None,
            timeout_secs: None,
            plan_item_id: None,
            session_id: Some("cli:test:dialog".to_string()),
        }
    }

    #[test]
    fn runtime_injected_tools_are_marked_but_ordinary_tools_are_not() {
        let reload = crate::tools::ReloadConfigTool::new(
            crate::gateway::reload::ReloadHandle::unavailable(),
        );
        assert!(!crate::tools::Tool::runtime_injected(&reload));
    }

    #[test]
    fn resolve_agent_rejects_missing_target() {
        let manager = manager();
        let error =
            match manager.resolve_agent(&config(None), &ToolExecutionContext::default(), "t-1") {
                Ok(_) => panic!("expected rejection"),
                Err(error) => error,
            };
        assert!(matches!(error, SubAgentError::Other(message) if message.contains("named target")));
    }

    #[test]
    fn resolve_agent_rejects_unknown_target_without_catalog() {
        let manager = manager();
        let error = match manager.resolve_agent(
            &config(Some("ghost")),
            &ToolExecutionContext::default(),
            "t-2",
        ) {
            Ok(_) => panic!("expected rejection"),
            Err(error) => error,
        };
        assert!(
            matches!(error, SubAgentError::Other(message) if message.contains("unknown Agent target"))
        );
    }
}