Two-Tier Agent Architecture

async fn claude_call(prompt: &str, step_name: &str, trace_dir: Option<&PathBuf>) -> Result<String> {
    let mut tb = TraceBuilder::new(step_name, prompt);

    let output = tokio::process::Command::new("claude")
        .args(["-p", prompt])
        .env_remove("CLAUDECODE")
        .output()
        .await
        .context("failed to run `claude` CLI — is it installed and on PATH?")?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        tb.record_event(EventKind::Error, &stderr);
        let trace = tb.finish("");
        if let Some(dir) = trace_dir {
            let _ = trace::write_trace(&trace, dir);
        }
        anyhow::bail!("claude CLI failed: {stderr}");
    }

    let response = String::from_utf8_lossy(&output.stdout).trim().to_string();
    let trace = tb.finish(&response);
    if let Some(dir) = trace_dir {
        let _ = trace::write_trace(&trace, dir);
    }

    Ok(response)
}

info!(task, "ambiguous task → asking Claude to classify");
let prompt = format!(
    "Classify this task as either SIMPLE, STANDARD, or BUGFIX.\n\n\
     SIMPLE = documentation changes, typo fixes, renames, formatting, trivial edits.\n\
     STANDARD = new features, refactors, integrations, anything that needs tests.\n\
     BUGFIX = fixing bugs, crashes, errors, regressions, investigating broken behavior.\n\n\
     Task: {task}\n\n\
     Reply with ONLY the word SIMPLE, STANDARD, or BUGFIX, nothing else."
);

match claude_call(&prompt, "classify", trace_dir).await {
    Ok(response) => {
        let upper = response.trim().to_uppercase();
        if upper.contains("SIMPLE") {
            info!(task, "agent classified as Simple");
            TaskComplexity::Simple
        } else if upper.contains("BUGFIX") {
            info!(task, "agent classified as BugFix");
            TaskComplexity::BugFix
        } else {
            info!(task, "agent classified as Standard (default for ambiguous)");
            TaskComplexity::Standard
        }
    }
    Err(e) => {
        warn!(task, error = %e, "classification claude_call failed → defaulting to Standard");
        TaskComplexity::Standard
    }
}

#[derive(Debug, Clone)]
pub struct MagpieConfig {
    pub provider: String,
    pub model: String,
    pub max_turns: u32,
    /// When set, agent calls are traced to JSONL files in this directory.
    pub trace_dir: Option<PathBuf>,
}

impl Default for MagpieConfig {
    fn default() -> Self {
        Self {
            provider: "claude-code".to_string(),
            model: "default".to_string(),
            max_turns: 10,
            trace_dir: None,
        }
    }
}

pub struct MagpieAgent {
    config: MagpieConfig,
    agent: Agent,
    session_manager: Arc<SessionManager>,
}

impl MagpieAgent {
    pub fn new(config: MagpieConfig) -> Result<Self> {
        let agent = Agent::new();
        let session_manager = Arc::new(SessionManager::instance());
        Ok(Self {
            config,
            agent,
            session_manager,
        })
    }
}

pub async fn run(
    &self,
    prompt: &str,
    working_dir: Option<&PathBuf>,
    step_name: Option<&str>,
) -> Result<String> {
    if prompt.trim().is_empty() {
        bail!("prompt must not be empty");
    }

    let step_label = step_name.unwrap_or("agent");
    let mut tb = TraceBuilder::new(step_label, prompt);

    // Create provider
    let provider = create_with_named_model(
        &self.config.provider,
        &self.config.model,
        Vec::<ExtensionConfig>::new(),
    )
    .await?;

    // Create session — use caller-provided dir or fall back to process CWD
    let working_dir = match working_dir {
        Some(dir) => dir.clone(),
        None => std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
    };
    let session = self
        .session_manager
        .create_session(working_dir, "magpie".to_string(), SessionType::Hidden)
        .await?;

    // Wire provider to agent
    self.agent.update_provider(provider, &session.id).await?;

    let session_config = SessionConfig {
        id: session.id.clone(),
        schedule_id: None,
        max_turns: Some(self.config.max_turns),
        retry_config: None,
    };

    let user_message = Message::user().with_text(prompt);

    let mut stream = self.agent.reply(user_message, session_config, None).await?;

    // Collect assistant text from the stream, classifying events for tracing.
    //
    // The claude-code provider streams tokens as individual AgentEvent::Message
    // events, so we concatenate directly — tokens already carry their own spacing
    // (e.g. " the", " add").  Inserting newlines between fragments would corrupt
    // single-line outputs like commit messages and branch slugs.
    let mut response = String::new();
    while let Some(event) = stream.next().await {
        match event {
            Ok(AgentEvent::Message(msg)) => {
                for content in msg.content.iter() {
                    if let Some(text) = content.as_text() {
                        response.push_str(text);
                        tb.record_event(EventKind::Text, text);
                    } else {
                        // Classify non-text content for tracing
                        let display = format!("{content}");
                        if display.starts_with("[ToolRequest") {
                            tb.record_event(EventKind::ToolRequest, &display);
                        } else if display.starts_with("[ToolResponse") {
                            tb.record_event(EventKind::ToolResponse, &display);
                        } else if display.starts_with("[Thinking") {
                            tb.record_event(EventKind::Thinking, &display);
                        }
                    }
                }
            }
            Err(e) => {
                tb.record_event(EventKind::Error, &e.to_string());
            }
            _ => {}
        }
    }

    let call_trace = tb.finish(&response);
    if let Some(ref dir) = self.config.trace_dir {
        let _ = trace::write_trace(&call_trace, dir);
    }

    Ok(response)
}

pub async fn execute(
    &self,
    ctx: &StepContext,
    step_name: &str,
    sandbox: &dyn Sandbox,
) -> Result<StepContext> {
    let mut prompt = self.prompt.clone();

    // Prepend metadata context (e.g. chat history) if configured
    if let Some(ref key) = self.context_metadata_key {
        if let Some(value) = ctx.metadata.get(key) {
            prompt = format!("Context from conversation:\n```\n{value}\n```\n\n{prompt}");
        }
    }

    if self.include_last_output {
        if let Some(ref output) = ctx.last_output {
            prompt = format!("Previous step output:\n```\n{output}\n```\n\n{prompt}");
        }
    }

    let response = match sandbox.name() {
        // Remote sandbox: agent can't use Goose locally, so exec `claude -p` inside the sandbox.
        "daytona" => {
            let max_turns = self.max_turns.unwrap_or(10);
            let output = sandbox
                .exec(
                    "claude",
                    &["-p", &prompt, "--max-turns", &max_turns.to_string()],
                )
                .await?;
            if output.exit_code != 0 {
                anyhow::bail!(
                    "claude -p failed inside sandbox (exit {}): {}",
                    output.exit_code,
                    output.stdout.trim()
                );
            }
            output.stdout.trim().to_string()
        }
        // Local sandbox: use the full Goose agent loop (current behavior).
        _ => {
            let mut config = MagpieConfig::default();
            if let Some(turns) = self.max_turns {
                config.max_turns = turns;
            }

            // Read trace_dir from context metadata (injected by pipeline)
            if let Some(dir) = ctx.metadata.get("trace_dir") {
                config.trace_dir = Some(PathBuf::from(dir));
            }

            let agent = MagpieAgent::new(config)?;
            let working_dir = PathBuf::from(sandbox.working_dir());
            agent
                .run(&prompt, Some(&working_dir), Some(step_name))
                .await?
        }
    };

    let mut new_ctx = ctx.clone();
    new_ctx.last_output = Some(response);
    new_ctx.last_exit_code = Some(0);
    Ok(new_ctx)
}