opi_coding_agent/

harness.rs

//! Interactive CLI harness (S8.4).
//!
//! Wires together config, tools, system prompt, hooks, and Agent into a
//! single entry point for the interactive coding agent.

use std::path::{Path, PathBuf};

use opi_agent::Agent;
use opi_agent::event::AgentEvent;
use opi_agent::hooks::AgentHooks;
use opi_agent::loop_types::{AgentError, AgentLoopConfig};
use opi_agent::message::AgentMessage;
use opi_agent::tool::Tool;
use opi_ai::message::Message;
use opi_ai::provider::Provider;

use crate::config::OpiConfig;
use crate::context_files;
use crate::policy::{RunMode, ToolRuntimeConfig, ToolSelection};
use crate::prompt::SystemPromptBuilder;
use crate::session_coordinator::{SessionCoordinator, to_wire_result};
use crate::tool::{BashTool, EditTool, FindTool, GlobTool, GrepTool, LsTool, ReadTool, WriteTool};

/// Optional pre-existing session the harness can adopt instead of creating
/// a new JSONL file. Produced by `--resume` flows.
pub struct ResumeInfo {
    pub path: PathBuf,
    pub session_id: String,
    pub entries: Vec<opi_agent::session::SessionEntry>,
    /// The workspace cwd recorded in the session header. Used to restore the
    /// correct workspace root when resuming from a different directory.
    pub original_cwd: PathBuf,
}

/// Harness wiring config, tools, system prompt, hooks, and Agent.
pub struct CodingHarness {
    agent: Agent,
    config: OpiConfig,
    system_prompt: String,
    session: Option<SessionCoordinator>,
    /// Message count before the current turn — used to slice only new messages for persistence.
    turn_offset: usize,
    /// Images queued from --image CLI flag, injected into the first prompt.
    pending_images: Vec<opi_ai::message::InputContent>,
}

impl CodingHarness {
    /// Create a new harness with the given provider, model, config, and workspace root.
    pub fn new(
        provider: Box<dyn Provider>,
        model: String,
        config: OpiConfig,
        workspace_root: PathBuf,
    ) -> Self {
        Self::new_with_hooks(
            provider,
            model,
            config,
            workspace_root,
            Box::new(CodingAgentHooks),
            None,
            Vec::new(),
            ToolSelection::Default,
        )
    }

    /// Create a new harness with an explicit tool selection.
    pub fn new_with_selection(
        provider: Box<dyn Provider>,
        model: String,
        config: OpiConfig,
        workspace_root: PathBuf,
        tool_selection: ToolSelection,
    ) -> Self {
        Self::new_with_hooks(
            provider,
            model,
            config,
            workspace_root,
            Box::new(CodingAgentHooks),
            None,
            Vec::new(),
            tool_selection,
        )
    }

    /// Create a new harness with already resolved tool runtime config.
    pub fn new_with_tool_config(
        provider: Box<dyn Provider>,
        model: String,
        config: OpiConfig,
        workspace_root: PathBuf,
        tool_config: ToolRuntimeConfig,
    ) -> Self {
        Self::new_with_hooks_and_resume_tool_config(
            provider,
            model,
            config,
            workspace_root,
            Box::new(CodingAgentHooks),
            None,
            Vec::new(),
            None,
            tool_config,
        )
    }

    /// Create a new harness with custom hooks.
    #[allow(clippy::too_many_arguments)]
    pub fn new_with_hooks(
        provider: Box<dyn Provider>,
        model: String,
        config: OpiConfig,
        workspace_root: PathBuf,
        hooks: Box<dyn AgentHooks>,
        user_system_prompt: Option<String>,
        initial_messages: Vec<AgentMessage>,
        tool_selection: ToolSelection,
    ) -> Self {
        Self::new_with_hooks_and_resume(
            provider,
            model,
            config,
            workspace_root,
            hooks,
            user_system_prompt,
            initial_messages,
            None,
            tool_selection,
        )
    }

    /// Create a new harness, optionally adopting an existing session (resume).
    #[allow(clippy::too_many_arguments)]
    pub fn new_with_hooks_and_resume(
        provider: Box<dyn Provider>,
        model: String,
        config: OpiConfig,
        workspace_root: PathBuf,
        hooks: Box<dyn AgentHooks>,
        user_system_prompt: Option<String>,
        initial_messages: Vec<AgentMessage>,
        resume: Option<ResumeInfo>,
        tool_selection: ToolSelection,
    ) -> Self {
        let tool_config = ToolRuntimeConfig::resolve(RunMode::Interactive, true, tool_selection)
            .expect("interactive tool config should be valid");
        Self::new_with_hooks_and_resume_tool_config(
            provider,
            model,
            config,
            workspace_root,
            hooks,
            user_system_prompt,
            initial_messages,
            resume,
            tool_config,
        )
    }

    /// Create a new harness, optionally adopting an existing session (resume),
    /// with already resolved tool runtime config.
    #[allow(clippy::too_many_arguments)]
    pub fn new_with_hooks_and_resume_tool_config(
        provider: Box<dyn Provider>,
        model: String,
        config: OpiConfig,
        workspace_root: PathBuf,
        hooks: Box<dyn AgentHooks>,
        user_system_prompt: Option<String>,
        initial_messages: Vec<AgentMessage>,
        resume: Option<ResumeInfo>,
        tool_config: ToolRuntimeConfig,
    ) -> Self {
        Self::new_with_global_config_dir_tool_config(
            provider,
            model,
            config,
            workspace_root,
            hooks,
            user_system_prompt,
            initial_messages,
            resume,
            tool_config,
            None,
        )
    }

    /// Create a new harness with an explicit global config directory override.
    ///
    /// When `global_config_dir` is `None`, uses the platform default from
    /// [`crate::config::user_config_dir`]. Pass `Some(path)` in tests to
    /// isolate global context file discovery from the real user config dir.
    #[allow(clippy::too_many_arguments)]
    pub fn new_with_global_config_dir(
        provider: Box<dyn Provider>,
        model: String,
        config: OpiConfig,
        workspace_root: PathBuf,
        hooks: Box<dyn AgentHooks>,
        user_system_prompt: Option<String>,
        initial_messages: Vec<AgentMessage>,
        resume: Option<ResumeInfo>,
        tool_selection: ToolSelection,
        global_config_dir: Option<PathBuf>,
    ) -> Self {
        let tool_config = ToolRuntimeConfig::resolve(RunMode::Interactive, true, tool_selection)
            .expect("interactive tool config should be valid");
        Self::new_with_global_config_dir_tool_config(
            provider,
            model,
            config,
            workspace_root,
            hooks,
            user_system_prompt,
            initial_messages,
            resume,
            tool_config,
            global_config_dir,
        )
    }

    /// Create a new harness with an explicit global config directory override
    /// and already resolved tool runtime config.
    #[allow(clippy::too_many_arguments)]
    pub fn new_with_global_config_dir_tool_config(
        provider: Box<dyn Provider>,
        model: String,
        config: OpiConfig,
        workspace_root: PathBuf,
        hooks: Box<dyn AgentHooks>,
        user_system_prompt: Option<String>,
        initial_messages: Vec<AgentMessage>,
        resume: Option<ResumeInfo>,
        tool_config: ToolRuntimeConfig,
        global_config_dir: Option<PathBuf>,
    ) -> Self {
        let tools = Self::build_tools(&workspace_root, &tool_config);
        let tool_defs: Vec<_> = tools.iter().map(|t| t.definition()).collect();
        let mut builder = SystemPromptBuilder::new().tools(tool_defs);
        if let Some(content) = user_system_prompt {
            builder = builder.user_system(content);
        }
        let resolved_global_dir = global_config_dir.unwrap_or_else(crate::config::user_config_dir);
        let context = context_files::discover_context_files(
            &workspace_root,
            Some(resolved_global_dir.as_path()),
        );
        if !context.content.is_empty() {
            builder = builder.context_files(context.content);
        }
        let system_prompt = builder.build();

        let agent_config = AgentLoopConfig {
            max_turns: config.defaults.max_iterations,
            retry: Some(config.retry.clone()),
            thinking: if config.thinking.enabled {
                Some(opi_ai::provider::ThinkingConfig {
                    enabled: true,
                    budget_tokens: Some(config.thinking.budget_tokens as u64),
                })
            } else {
                None
            },
            ..Default::default()
        };

        let mut agent = Agent::new(
            provider,
            tools,
            model.clone(),
            Some(system_prompt.clone()),
            agent_config,
            hooks,
        );

        let initial_len = initial_messages.len();
        if !initial_messages.is_empty() {
            agent.set_initial_messages(initial_messages);
        }

        let cwd = if let Some(ref info) = resume {
            // When resuming, use the workspace cwd from the session header so
            // tools operate in the correct workspace even if the process was
            // launched from a different directory.
            info.original_cwd.to_string_lossy().into_owned()
        } else {
            std::env::current_dir()
                .unwrap_or_default()
                .to_string_lossy()
                .into_owned()
        };
        let compaction_config = opi_agent::compaction::CompactionConfig {
            enabled: config.compaction.enabled,
            threshold_tokens: config.compaction.threshold_tokens,
        };

        let session = if let Some(info) = resume {
            SessionCoordinator::open_existing(
                info.path,
                info.session_id,
                &info.entries,
                initial_len,
                compaction_config,
                model.clone(),
            )
            .ok()
        } else {
            let session_dir = crate::session_cli::session_dir();
            SessionCoordinator::new(&session_dir, &cwd, compaction_config, model.clone()).ok()
        };

        Self {
            agent,
            config,
            system_prompt,
            session,
            turn_offset: initial_len,
            pending_images: Vec::new(),
        }
    }

    /// Add an extra tool to the harness (for testing with mock tools).
    pub fn add_tool(&mut self, tool: Box<dyn Tool>) {
        self.agent.add_tool(tool);
    }

    /// Queue images to be injected into the next prompt.
    pub fn queue_images(&mut self, images: Vec<opi_ai::message::InputContent>) {
        self.pending_images.extend(images);
    }

    /// Take and clear queued images.
    pub fn take_pending_images(&mut self) -> Vec<opi_ai::message::InputContent> {
        std::mem::take(&mut self.pending_images)
    }

    /// Return model picker items from the active provider.
    pub fn model_picker_items(&self) -> Vec<opi_tui::SelectItem> {
        crate::picker::model_picker_items_from_provider(self.agent.provider())
    }

    /// Change the model used by subsequent prompts.
    pub fn set_model(&mut self, model: String) {
        self.agent.set_model(model);
    }

    /// Resume an existing session by ID into this harness.
    pub fn resume_session_id(&mut self, session_id: &str) -> Result<usize, String> {
        let dir = crate::session_cli::session_dir();
        let session =
            crate::session_cli::resume_session(&dir, session_id).map_err(|e| e.to_string())?;
        let messages = crate::session_cli::reconstruct_context(&session.entries);
        let message_count = messages.len();
        self.agent.replace_messages(messages);

        let compaction_config = opi_agent::compaction::CompactionConfig {
            enabled: self.config.compaction.enabled,
            threshold_tokens: self.config.compaction.threshold_tokens,
        };
        self.session = SessionCoordinator::open_existing(
            session.path,
            session.header.id,
            &session.entries,
            message_count,
            compaction_config,
            self.agent.model().to_string(),
        )
        .ok();
        self.turn_offset = message_count;
        Ok(message_count)
    }

    /// Send a user prompt and run the agent loop.
    pub async fn prompt(&mut self, text: &str) -> Result<Vec<AgentMessage>, AgentError> {
        let offset = self.turn_offset;
        let messages = self.agent.prompt(text).await?;
        let new = &messages[offset..];
        self.persist_turn(new, offset);
        let final_messages = self.current_messages();
        self.turn_offset = final_messages.len();
        Ok(final_messages)
    }

    /// Send a user message with arbitrary content (text + images) and run the
    /// agent loop.
    pub async fn prompt_with_content(
        &mut self,
        content: Vec<opi_ai::message::InputContent>,
    ) -> Result<Vec<AgentMessage>, AgentError> {
        let offset = self.turn_offset;
        let messages = self.agent.prompt_with_content(content).await?;
        let new = &messages[offset..];
        self.persist_turn(new, offset);
        let final_messages = self.current_messages();
        self.turn_offset = final_messages.len();
        Ok(final_messages)
    }

    /// Continue the conversation with an additional message.
    pub async fn continue_(&mut self, text: &str) -> Result<Vec<AgentMessage>, AgentError> {
        let offset = self.turn_offset;
        let messages = self.agent.continue_(text).await?;
        let new = &messages[offset..];
        self.persist_turn(new, offset);
        let final_messages = self.current_messages();
        self.turn_offset = final_messages.len();
        Ok(final_messages)
    }

    /// Sum usage across every assistant message produced during a turn.
    ///
    /// A single user prompt can drive multiple provider calls (e.g.
    /// tool-call response followed by a final response). Each emitted
    /// assistant message carries its own `usage`; the cumulative session
    /// total must include all of them, not just the last one.
    fn aggregate_turn_usage(messages: &[AgentMessage]) -> opi_ai::stream::Usage {
        let mut total = opi_ai::stream::Usage::default();
        for m in messages {
            if let AgentMessage::Llm(Message::Assistant(a)) = m {
                total.input_tokens = total.input_tokens.saturating_add(a.usage.input_tokens);
                total.output_tokens = total.output_tokens.saturating_add(a.usage.output_tokens);
                total.cache_read_tokens = total
                    .cache_read_tokens
                    .saturating_add(a.usage.cache_read_tokens);
                total.cache_write_tokens = total
                    .cache_write_tokens
                    .saturating_add(a.usage.cache_write_tokens);
            }
        }
        total
    }

    /// Aggregate usage across all assistant messages in a turn and persist.
    ///
    /// If compaction was triggered during persistence, this also rewrites
    /// the Agent's message buffer to `[summary, ...kept]` so subsequent
    /// provider calls no longer carry the compacted history. Emits
    /// `CompactionStart`/`CompactionEnd` events for subscribers.
    fn persist_turn(&mut self, messages: &[AgentMessage], turn_start_agent_index: usize) {
        if let Some(session) = &mut self.session {
            let usage = Self::aggregate_turn_usage(messages);
            let compaction_reason =
                match session.on_turn_end(messages, &usage, turn_start_agent_index) {
                    Ok(reason) => reason,
                    Err(e) => {
                        self.agent.emit_event(AgentEvent::SessionPersistError {
                            message: format!("session write failed: {e}"),
                        });
                        return;
                    }
                };

            if let Some(reason) = compaction_reason {
                self.agent
                    .emit_event(AgentEvent::CompactionStart { reason });
                match session.execute_compaction(reason) {
                    Ok(Some(out)) => {
                        let wire = to_wire_result(&out);
                        self.agent.replace_messages(out.new_agent_messages);
                        self.agent.emit_event(AgentEvent::CompactionEnd {
                            reason,
                            result: Some(wire),
                            aborted: false,
                            error_message: None,
                        });
                    }
                    Ok(None) => {
                        self.agent.emit_event(AgentEvent::CompactionEnd {
                            reason,
                            result: None,
                            aborted: true,
                            error_message: Some("compaction produced no output".into()),
                        });
                    }
                    Err(e) => {
                        // Compaction marker failed to persist — leave in-memory
                        // state un-compacted (SessionCoordinator already skipped
                        // the mutation) and surface the error to subscribers.
                        self.agent.emit_event(AgentEvent::CompactionEnd {
                            reason,
                            result: None,
                            aborted: true,
                            error_message: Some(format!("compaction persist failed: {e}")),
                        });
                        self.agent.emit_event(AgentEvent::SessionPersistError {
                            message: format!("compaction write failed: {e}"),
                        });
                    }
                }
            }
        }
    }

    /// Return the current message buffer (after any compaction).
    fn current_messages(&self) -> Vec<AgentMessage> {
        // The Agent's `set_initial_messages` / `replace_messages` API doesn't
        // expose a getter, so we re-derive the buffer from what was returned
        // by the loop plus any post-loop mutation. Simplest correct option:
        // ask the Agent via a new getter.
        self.agent.messages_snapshot()
    }

    /// Register an event subscriber.
    pub fn subscribe(&mut self, callback: Box<dyn Fn(&AgentEvent) + Send + Sync>) {
        self.agent.subscribe(callback);
    }

    /// Return the assembled system prompt (for testing).
    pub fn system_prompt(&self) -> &str {
        &self.system_prompt
    }

    /// Return a reference to the config.
    pub fn config(&self) -> &OpiConfig {
        &self.config
    }

    /// Cancel the running operation.
    pub fn cancel(&self) {
        self.agent.abort();
    }

    /// Return a clonable cancellation token for external cancellation.
    pub fn cancel_token(&self) -> tokio_util::sync::CancellationToken {
        self.agent.cancel_token()
    }

    /// Return the session coordinator, if active.
    pub fn session(&self) -> Option<&SessionCoordinator> {
        self.session.as_ref()
    }

    fn build_tools(workspace_root: &Path, tool_config: &ToolRuntimeConfig) -> Vec<Box<dyn Tool>> {
        let read_policy = match tool_config.run_mode {
            RunMode::Interactive => crate::tool::PathPolicy::AllowOutsideWorkspace,
            RunMode::NonInteractive => crate::tool::PathPolicy::WorkspaceOnly,
        };

        let mut tools: Vec<(&str, Box<dyn Tool>)> = vec![
            (
                "read",
                Box::new(ReadTool::new_with_policy(
                    workspace_root.to_path_buf(),
                    read_policy,
                )),
            ),
            (
                "write",
                Box::new(WriteTool::new(workspace_root.to_path_buf())),
            ),
            (
                "edit",
                Box::new(EditTool::new(workspace_root.to_path_buf())),
            ),
            (
                "bash",
                Box::new(BashTool::new(workspace_root.to_path_buf())),
            ),
            (
                "grep",
                Box::new(GrepTool::new(workspace_root.to_path_buf())),
            ),
            (
                "find",
                Box::new(FindTool::new(workspace_root.to_path_buf())),
            ),
            ("ls", Box::new(LsTool::new(workspace_root.to_path_buf()))),
            (
                "glob",
                Box::new(GlobTool::new(workspace_root.to_path_buf())),
            ),
        ];

        tools
            .drain(..)
            .filter(|(name, _)| {
                tool_config
                    .active_tool_names
                    .iter()
                    .any(|active| active == name)
            })
            .map(|(_, tool)| tool)
            .collect()
    }
}

// ---------------------------------------------------------------------------
// Hooks
// ---------------------------------------------------------------------------

/// Shared conversion of agent-level messages to the provider-facing Message
/// stream. Used by every hook in this crate so resume/compaction semantics
/// stay consistent between interactive and non-interactive paths.
///
/// - `AgentMessage::Llm` is forwarded directly.
/// - `AgentMessage::CompactionSummary` is rendered as a synthetic user
///   message so the provider sees a textual marker for context that was
///   compacted away.
/// - Other variants (`BranchSummary`, `Custom`) are dropped — they have no
///   provider-facing representation yet.
pub(crate) fn agent_messages_to_llm(messages: &[AgentMessage]) -> Vec<Message> {
    let mut result = Vec::with_capacity(messages.len());
    for msg in messages {
        match msg {
            AgentMessage::Llm(m) => result.push(m.clone()),
            AgentMessage::CompactionSummary(summary) => {
                result.push(Message::User(opi_ai::message::UserMessage {
                    content: vec![opi_ai::message::InputContent::Text {
                        text: format!(
                            "[Context was compacted. Summary of earlier conversation: {}]",
                            summary.summary
                        ),
                    }],
                    timestamp_ms: 0,
                }));
            }
            _ => {}
        }
    }
    result
}

/// Default hooks for the coding agent -- pass-through message conversion.
pub struct CodingAgentHooks;

impl AgentHooks for CodingAgentHooks {
    fn convert_to_llm(&self, messages: &[AgentMessage]) -> Result<Vec<Message>, AgentError> {
        Ok(agent_messages_to_llm(messages))
    }
}

/// Interactive hooks for the coding agent.
///
/// Tool safety is controlled by active tool selection and extension hooks, not
/// by a core interactive permission popup.
pub struct InteractiveCodingHooks;

impl InteractiveCodingHooks {
    pub fn new(_allow_mutating: bool) -> Self {
        Self
    }
}

impl AgentHooks for InteractiveCodingHooks {
    fn convert_to_llm(&self, messages: &[AgentMessage]) -> Result<Vec<Message>, AgentError> {
        Ok(agent_messages_to_llm(messages))
    }
}