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

mod claude_code;
mod gemini;
mod openai;

pub use claude_code::ClaudeCodeRuntime;
pub use gemini::GeminiRuntime;
pub use openai::OpenAiRuntime;

use std::sync::{Arc, Mutex};
use tokio::sync::broadcast;

use crate::agent_log::AgentLogWriter;

use super::{AgentEvent, TokenUsage};

/// Context passed to a runtime when launching an agent session.
pub struct RuntimeContext {
    pub story_id: String,
    pub agent_name: String,
    pub command: String,
    pub args: Vec<String>,
    pub prompt: String,
    pub cwd: String,
    pub inactivity_timeout_secs: u64,
    /// Port of the huskies MCP server, used by API-based runtimes (Gemini, OpenAI)
    /// to call back for tool execution.
    pub mcp_port: u16,
    /// When set, resume a previous Claude Code session instead of starting fresh.
    ///
    /// The CLI is invoked as `claude --resume <session_id> [-p <prompt>]` rather
    /// than `claude -p <full_prompt>`.  The agent re-enters the previous
    /// conversation and receives the `prompt` (if non-empty) as a new message.
    pub session_id_to_resume: Option<String>,
}

/// Result returned by a runtime after the agent session completes.
pub struct RuntimeResult {
    pub session_id: Option<String>,
    pub token_usage: Option<TokenUsage>,
}

/// Runtime status reported by the backend.
#[derive(Debug, Clone, PartialEq)]
#[allow(dead_code)]
pub enum RuntimeStatus {
    Idle,
    Running,
    Completed,
    Failed,
}

/// Abstraction over different agent execution backends.
///
/// Implementations:
/// - [`ClaudeCodeRuntime`]: spawns the `claude` CLI via a PTY (default, `runtime = "claude-code"`)
///
/// Future implementations could include OpenAI and Gemini API runtimes.
#[allow(dead_code)]
pub trait AgentRuntime: Send + Sync {
    /// Start the agent and drive it to completion, streaming events through
    /// the provided broadcast sender and event log.
    ///
    /// Returns when the agent session finishes (success or error).
    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>;

    /// Stop the running agent.
    fn stop(&self);

    /// Get the current runtime status.
    fn get_status(&self) -> RuntimeStatus;

    /// Return any events buffered outside the broadcast channel.
    ///
    /// PTY-based runtimes stream directly to the broadcast channel; this
    /// returns empty by default. API-based runtimes may buffer events here.
    fn stream_events(&self) -> Vec<AgentEvent> {
        vec![]
    }
}

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

    #[test]
    fn runtime_context_fields() {
        let ctx = RuntimeContext {
            story_id: "42_story_foo".to_string(),
            agent_name: "coder-1".to_string(),
            command: "claude".to_string(),
            args: vec!["--model".to_string(), "sonnet".to_string()],
            prompt: "Do the thing".to_string(),
            cwd: "/tmp/wt".to_string(),
            inactivity_timeout_secs: 300,
            mcp_port: 3001,
            session_id_to_resume: None,
        };
        assert_eq!(ctx.story_id, "42_story_foo");
        assert_eq!(ctx.agent_name, "coder-1");
        assert_eq!(ctx.command, "claude");
        assert_eq!(ctx.args.len(), 2);
        assert_eq!(ctx.prompt, "Do the thing");
        assert_eq!(ctx.cwd, "/tmp/wt");
        assert_eq!(ctx.inactivity_timeout_secs, 300);
        assert_eq!(ctx.mcp_port, 3001);
    }

    #[test]
    fn runtime_result_fields() {
        let result = RuntimeResult {
            session_id: Some("sess-123".to_string()),
            token_usage: Some(TokenUsage {
                input_tokens: 100,
                output_tokens: 50,
                cache_creation_input_tokens: 0,
                cache_read_input_tokens: 0,
                total_cost_usd: 0.01,
            }),
        };
        assert_eq!(result.session_id, Some("sess-123".to_string()));
        assert!(result.token_usage.is_some());
        let usage = result.token_usage.unwrap();
        assert_eq!(usage.input_tokens, 100);
        assert_eq!(usage.output_tokens, 50);
        assert_eq!(usage.total_cost_usd, 0.01);
    }

    #[test]
    fn runtime_result_no_usage() {
        let result = RuntimeResult {
            session_id: None,
            token_usage: None,
        };
        assert!(result.session_id.is_none());
        assert!(result.token_usage.is_none());
    }

    #[test]
    fn runtime_status_variants() {
        assert_eq!(RuntimeStatus::Idle, RuntimeStatus::Idle);
        assert_ne!(RuntimeStatus::Running, RuntimeStatus::Completed);
        assert_ne!(RuntimeStatus::Failed, RuntimeStatus::Idle);
    }

    #[test]
    fn claude_code_runtime_get_status_returns_idle() {
        use std::collections::HashMap;
        use crate::io::watcher::WatcherEvent;
        let killers = Arc::new(Mutex::new(HashMap::new()));
        let (watcher_tx, _) = broadcast::channel::<WatcherEvent>(16);
        let runtime = ClaudeCodeRuntime::new(killers, watcher_tx);
        assert_eq!(runtime.get_status(), RuntimeStatus::Idle);
    }

    #[test]
    fn claude_code_runtime_stream_events_empty() {
        use std::collections::HashMap;
        use crate::io::watcher::WatcherEvent;
        let killers = Arc::new(Mutex::new(HashMap::new()));
        let (watcher_tx, _) = broadcast::channel::<WatcherEvent>(16);
        let runtime = ClaudeCodeRuntime::new(killers, watcher_tx);
        assert!(runtime.stream_events().is_empty());
    }
}