huskies/server/src/agents/runtime/claude_code.rs

//! Claude Code runtime — launches Claude Code CLI sessions as agent backends.
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use portable_pty::ChildKiller;
use tokio::sync::broadcast;

use crate::agent_log::AgentLogWriter;
use crate::io::watcher::WatcherEvent;
use crate::slog_warn;

use super::{AgentEvent, AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus};

/// Agent runtime that spawns the `claude` CLI in a PTY and streams JSON events.
///
/// This is the default runtime (`runtime = "claude-code"` in project.toml).
/// It wraps the existing PTY-based execution logic, preserving all streaming,
/// token tracking, and inactivity timeout behaviour.
pub struct ClaudeCodeRuntime {
    child_killers: Arc<Mutex<HashMap<String, Box<dyn ChildKiller + Send + Sync>>>>,
    watcher_tx: broadcast::Sender<WatcherEvent>,
}

impl ClaudeCodeRuntime {
    pub fn new(
        child_killers: Arc<Mutex<HashMap<String, Box<dyn ChildKiller + Send + Sync>>>>,
        watcher_tx: broadcast::Sender<WatcherEvent>,
    ) -> Self {
        Self {
            child_killers,
            watcher_tx,
        }
    }
}

impl AgentRuntime for ClaudeCodeRuntime {
    async fn start(
        &self,
        ctx: RuntimeContext,
        tx: broadcast::Sender<AgentEvent>,
        event_log: Arc<Mutex<Vec<AgentEvent>>>,
        log_writer: Option<Arc<Mutex<AgentLogWriter>>>,
    ) -> Result<RuntimeResult, String> {
        let pty_result = super::super::pty::run_agent_pty_streaming(
            &ctx.story_id,
            &ctx.agent_name,
            &ctx.command,
            &ctx.args,
            &ctx.prompt,
            &ctx.cwd,
            &tx,
            &event_log,
            log_writer.clone(),
            ctx.inactivity_timeout_secs,
            Arc::clone(&self.child_killers),
            self.watcher_tx.clone(),
            ctx.session_id_to_resume.as_deref(),
        )
        .await;

        match pty_result {
            Ok(result) => Ok(RuntimeResult {
                session_id: result.session_id,
                token_usage: result.token_usage,
            }),
            Err(e) if ctx.session_id_to_resume.is_some() && ctx.fresh_prompt.is_some() => {
                // Resume failed — fall back to a fresh session without --resume.
                slog_warn!(
                    "[agents] Resume failed for {}:{}, retrying without --resume: {}",
                    ctx.story_id,
                    ctx.agent_name,
                    e
                );
                let fresh = ctx.fresh_prompt.unwrap();
                let fallback_result = super::super::pty::run_agent_pty_streaming(
                    &ctx.story_id,
                    &ctx.agent_name,
                    &ctx.command,
                    &ctx.args,
                    &fresh,
                    &ctx.cwd,
                    &tx,
                    &event_log,
                    log_writer,
                    ctx.inactivity_timeout_secs,
                    Arc::clone(&self.child_killers),
                    self.watcher_tx.clone(),
                    None, // no --resume on fallback
                )
                .await?;
                Ok(RuntimeResult {
                    session_id: fallback_result.session_id,
                    token_usage: fallback_result.token_usage,
                })
            }
            Err(e) => Err(e),
        }
    }

    fn stop(&self) {
        // Stopping is handled externally by the pool via kill_child_for_key().
        // The ChildKillerGuard in pty.rs deregisters automatically on process exit.
    }

    fn get_status(&self) -> RuntimeStatus {
        // Lifecycle status is tracked by the pool; the runtime itself is stateless.
        RuntimeStatus::Idle
    }
}