opencrabs 0.3.33

//! spawn_agent tool — creates a child agent with forked context.
//!
//! Sub-agent progress is streamed to `~/.opencrabs/tmp/subagents/<agent_id>.json`
//! so the main orchestrator can track status without session_search.

use super::manager::{SubAgent, SubAgentManager, SubAgentState};
use super::status::AgentStatus;
use crate::brain::tools::error::{Result, ToolError};
use crate::brain::tools::r#trait::{Tool, ToolCapability, ToolExecutionContext, ToolResult};
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

/// Tool that spawns a child agent to handle a sub-task.
pub struct SpawnAgentTool {
    manager: Arc<SubAgentManager>,
    parent_registry: Arc<crate::brain::tools::ToolRegistry>,
}

impl SpawnAgentTool {
    pub fn new(
        manager: Arc<SubAgentManager>,
        parent_registry: Arc<crate::brain::tools::ToolRegistry>,
    ) -> Self {
        Self {
            manager,
            parent_registry,
        }
    }
}

#[async_trait]
impl Tool for SpawnAgentTool {
    fn name(&self) -> &str {
        "spawn_agent"
    }

    fn description(&self) -> &str {
        "Spawn a child agent to handle a sub-task autonomously. The child gets its own session \
         and runs in the background. Returns an agent_id you can use with wait_agent, send_input, \
         close_agent, or resume_agent. Use this to delegate independent work items."
    }

    fn input_schema(&self) -> Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "prompt": {
                    "type": "string",
                    "description": "The task/instruction for the child agent to execute"
                },
                "label": {
                    "type": "string",
                    "description": "Short human-readable label for this sub-agent (e.g., 'refactor-auth', 'test-runner')"
                },
                "agent_type": {
                    "type": "string",
                    "description": "Agent specialization: 'general' (full tools), 'explore' (read-only), 'plan' (read+bash), 'code' (full write), 'research' (web+read). Default: general",
                    "enum": ["general", "explore", "plan", "code", "research"]
                }
            },
            "required": ["prompt"]
        })
    }

    fn capabilities(&self) -> Vec<ToolCapability> {
        vec![ToolCapability::SystemModification]
    }

    fn requires_approval(&self) -> bool {
        true
    }

    async fn execute(&self, input: Value, context: &ToolExecutionContext) -> Result<ToolResult> {
        let prompt = input
            .get("prompt")
            .and_then(|v| v.as_str())
            .ok_or_else(|| ToolError::InvalidInput("'prompt' is required".into()))?
            .to_string();

        let label = input
            .get("label")
            .and_then(|v| v.as_str())
            .unwrap_or("sub-agent")
            .to_string();

        let agent_type = super::AgentType::parse(
            input
                .get("agent_type")
                .and_then(|v| v.as_str())
                .unwrap_or("general"),
        );

        // We need a ServiceContext to create a session for the child
        let service_context = context
            .service_context
            .as_ref()
            .ok_or_else(|| ToolError::Execution("No service context available".into()))?
            .clone();

        // Create a new session for the child agent
        let session_service = crate::services::SessionService::new(service_context.clone());
        let child_session = session_service
            .create_session(Some(format!("subagent: {}", label)))
            .await
            .map_err(|e| ToolError::Execution(format!("Failed to create child session: {}", e)))?;

        let child_session_id = child_session.id;
        let agent_id = SubAgentManager::generate_id();

        // Create cancel token and input channel for the child
        let cancel_token = CancellationToken::new();
        let (input_tx, input_rx) = mpsc::unbounded_channel::<String>();

        // Load config and extract model override before entering block scope
        let config = crate::config::Config::load()
            .map_err(|e| ToolError::Execution(format!("Config load failed: {}", e)))?;
        let model_override = config.agent.subagent_model.clone();

        // Build a minimal AgentService for the child
        let child_service = {
            // Use subagent-specific provider if configured, otherwise inherit parent's
            let provider = if let Some(ref provider_name) = config.agent.subagent_provider {
                match crate::brain::provider::create_provider_by_name(&config, provider_name).await
                {
                    Ok(p) => {
                        tracing::info!("Sub-agent using configured provider '{}'", provider_name);
                        p
                    }
                    Err(e) => {
                        tracing::warn!(
                            "Sub-agent provider '{}' failed: {e}, falling back to parent",
                            provider_name
                        );
                        crate::brain::provider::create_provider(&config)
                            .await
                            .map_err(|e| {
                                ToolError::Execution(format!("Failed to create provider: {}", e))
                            })?
                    }
                }
            } else {
                crate::brain::provider::create_provider(&config)
                    .await
                    .map_err(|e| {
                        ToolError::Execution(format!("Failed to create provider: {}", e))
                    })?
            };

            // Build filtered tool registry based on agent type
            let child_registry = agent_type.build_registry(&self.parent_registry);

            let agent =
                crate::brain::agent::AgentService::new(provider, service_context.clone(), &config)
                    .await
                    .with_tool_registry(Arc::new(child_registry))
                    .with_auto_approve_tools(true) // children auto-approve (parent already approved spawn)
                    .with_working_directory(context.working_dir());

            Arc::new(agent)
        };

        // Prepend agent type system prompt to the user's task
        let full_prompt = format!("{}\n\n{}", agent_type.system_prompt(), prompt);

        // Create the status file in Pending state before spawning.
        let _agent_status = AgentStatus::new(
            &agent_id,
            &label,
            &child_session_id.to_string(),
            &full_prompt,
        )
        .map_err(|e| ToolError::Execution(format!("Failed to create status file: {e}")))?;

        // Spawn background task with input loop
        let cancel_clone = cancel_token.clone();
        let manager = self.manager.clone();
        let agent_id_clone = agent_id.clone();
        let prompt_clone = full_prompt;
        let label_clone = label.clone();
        let mut input_rx = input_rx;

        let handle = tokio::spawn(async move {
            tracing::info!("Sub-agent {} starting: {}", agent_id_clone, prompt_clone);

            // Transition to Running state.
            let mut status = AgentStatus::read(&agent_id_clone).unwrap_or_else(|| {
                AgentStatus::new(
                    &agent_id_clone,
                    &label_clone,
                    &child_session_id.to_string(),
                    &prompt_clone,
                )
                .expect("status file")
            });
            if !matches!(
                status.state,
                super::status::AgentState::Completed | super::status::AgentState::Failed
            ) && let Err(e) = status.mark_running()
            {
                tracing::warn!("Failed to write running status: {e}");
            }

            // Reload with correct state.

            let mut current_prompt = prompt_clone;
            let mut iteration: usize = 0;

            // Run prompt → wait for input → run again loop
            let final_output = loop {
                iteration += 1;
                let result = child_service
                    .send_message_with_tools_and_mode(
                        child_session_id,
                        current_prompt,
                        model_override.clone(),
                        Some(cancel_clone.clone()),
                    )
                    .await;

                match result {
                    Ok(response) => {
                        // Extract a short summary of what the agent did this turn.
                        let summary = if response.stop_reason
                            == Some(crate::brain::provider::types::StopReason::ToolUse)
                        {
                            "tool call(s) completed".to_string()
                        } else {
                            response.content.chars().take(120).collect::<String>()
                        };

                        status
                            .update_progress(iteration, None, Some(summary))
                            .unwrap_or_else(|e| tracing::warn!("status write failed: {e}"));

                        manager.update_output(&agent_id_clone, response.content.clone());
                        // Flip to AwaitingInput so wait_agent can observe
                        // round-boundary progress instead of blocking on
                        // task-join semantics (the task never terminates
                        // at a round — only on input/cancel — so the old
                        // `handle.await` in wait.rs always hit its
                        // timeout_secs and the LLM gave up the turn).
                        manager.mark_awaiting_input(&agent_id_clone);
                        tracing::info!(
                            "Sub-agent {} round {} complete, waiting for input",
                            agent_id_clone,
                            iteration
                        );

                        // Wait for follow-up input or shutdown
                        let next = tokio::select! {
                            msg = input_rx.recv() => msg,
                            _ = cancel_clone.cancelled() => {
                                tracing::info!("Sub-agent {} cancelled while waiting for input", agent_id_clone);
                                None
                            }
                        };

                        match next {
                            Some(text) => {
                                manager.mark_running_again(&agent_id_clone);
                                tracing::info!(
                                    "Sub-agent {} received follow-up input",
                                    agent_id_clone
                                );
                                current_prompt = text;
                            }
                            None => break response.content,
                        }
                    }
                    Err(e) => {
                        tracing::error!("Sub-agent {} failed: {}", agent_id_clone, e);
                        let _ = status.mark_failed(e.to_string());
                        manager.mark_failed(&agent_id_clone, e.to_string());
                        return;
                    }
                }
            };

            let _ = status.mark_completed(final_output.chars().take(200).collect());
            manager.mark_completed(&agent_id_clone, final_output);
        });

        // Register in manager
        self.manager.insert(SubAgent {
            id: agent_id.clone(),
            label: label.clone(),
            session_id: child_session_id,
            state: SubAgentState::Running,
            cancel_token,
            join_handle: Some(handle),
            input_tx: Some(input_tx),
            output: None,
            spawned_at: chrono::Utc::now(),
        });

        Ok(ToolResult::success(format!(
            "Spawned sub-agent '{}' with id: {}\nSession: {}\nPrompt: {}",
            label, agent_id, child_session_id, prompt
        )))
    }
}