acp_cli/bridge/

mod.rs

pub mod commands;
pub mod events;

pub use commands::BridgeCommand;
pub use events::{
    BridgeEvent, PermissionKind, PermissionOption, PermissionOutcome, PromptResult, ToolCallInfo,
};

use std::path::PathBuf;

use agent_client_protocol::{self as acp, Agent as _};
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};

use crate::client::BridgedAcpClient;
use crate::error::{AcpCliError, Result};

/// Core bridge that manages the ACP agent process lifecycle.
///
/// Spawns the agent in a dedicated blocking thread with a `LocalSet` (required
/// because ACP futures are `!Send`), and exposes an async command/event API
/// for the main thread to drive prompts and receive streamed output.
pub struct AcpBridge {
    cmd_tx: mpsc::Sender<BridgeCommand>,
    pub evt_rx: mpsc::UnboundedReceiver<BridgeEvent>,
    handle: JoinHandle<std::result::Result<(), AcpCliError>>,
}

/// Handle used to send cancel/shutdown commands independently of `evt_rx`.
///
/// Obtained via `AcpBridge::cancel_handle()` so the caller can borrow `evt_rx`
/// mutably while still being able to send cancel commands.
#[derive(Clone)]
pub struct BridgeCancelHandle {
    cmd_tx: mpsc::Sender<BridgeCommand>,
}

impl BridgeCancelHandle {
    /// Request cancellation of the current prompt (best-effort).
    pub async fn cancel(&self) -> Result<()> {
        let _ = self.cmd_tx.send(BridgeCommand::Cancel).await;
        Ok(())
    }
}

impl AcpBridge {
    /// Start the ACP bridge by spawning the agent process in a background thread.
    ///
    /// The agent is launched via `command` with the given `args`, and the working
    /// directory for the ACP session is set to `cwd`.
    pub async fn start(command: String, args: Vec<String>, cwd: PathBuf) -> Result<Self> {
        let (cmd_tx, cmd_rx) = mpsc::channel::<BridgeCommand>(16);
        let (evt_tx, evt_rx) = mpsc::unbounded_channel::<BridgeEvent>();

        let handle = tokio::task::spawn_blocking(move || {
            let rt = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .map_err(|e| AcpCliError::Connection(format!("runtime: {e}")))?;
            let local = tokio::task::LocalSet::new();
            local.block_on(&rt, acp_thread_main(cmd_rx, evt_tx, command, args, cwd))
        });

        Ok(Self {
            cmd_tx,
            evt_rx,
            handle,
        })
    }

    /// Obtain a lightweight cancel handle that can be used while `evt_rx` is
    /// borrowed mutably.
    pub fn cancel_handle(&self) -> BridgeCancelHandle {
        BridgeCancelHandle {
            cmd_tx: self.cmd_tx.clone(),
        }
    }

    /// Send a prompt to the agent and wait for the result.
    ///
    /// Text content is streamed in real-time via `BridgeEvent::TextChunk` events
    /// on `evt_rx`. The returned `PromptResult.content` will be empty because
    /// the main thread is expected to collect content from those events.
    pub async fn prompt(&self, messages: Vec<String>) -> Result<PromptResult> {
        let reply_rx = self.send_prompt(messages).await?;
        reply_rx
            .await
            .map_err(|_| AcpCliError::Connection("bridge reply dropped".into()))?
    }

    /// Send a prompt command without awaiting the reply.
    ///
    /// Returns a oneshot receiver that will resolve when the prompt completes.
    /// This allows the caller to drive `evt_rx` concurrently with waiting for
    /// the prompt result, avoiding borrow conflicts on `self`.
    pub async fn send_prompt(
        &self,
        messages: Vec<String>,
    ) -> Result<tokio::sync::oneshot::Receiver<Result<PromptResult>>> {
        let (reply_tx, reply_rx) = tokio::sync::oneshot::channel();
        self.cmd_tx
            .send(BridgeCommand::Prompt {
                messages,
                reply: reply_tx,
            })
            .await
            .map_err(|_| AcpCliError::Connection("bridge channel closed".into()))?;
        Ok(reply_rx)
    }

    /// Request cancellation of the current prompt (best-effort).
    pub async fn cancel(&self) -> Result<()> {
        let _ = self.cmd_tx.send(BridgeCommand::Cancel).await;
        Ok(())
    }

    /// Set the session mode on the ACP connection.
    pub async fn set_mode(&self, mode: String) -> Result<()> {
        let (reply_tx, reply_rx) = tokio::sync::oneshot::channel();
        self.cmd_tx
            .send(BridgeCommand::SetMode {
                mode,
                reply: reply_tx,
            })
            .await
            .map_err(|_| AcpCliError::Connection("bridge channel closed".into()))?;
        reply_rx
            .await
            .map_err(|_| AcpCliError::Connection("bridge reply dropped".into()))?
    }

    /// Set a session config option on the ACP connection.
    pub async fn set_config(&self, key: String, value: String) -> Result<()> {
        let (reply_tx, reply_rx) = tokio::sync::oneshot::channel();
        self.cmd_tx
            .send(BridgeCommand::SetConfig {
                key,
                value,
                reply: reply_tx,
            })
            .await
            .map_err(|_| AcpCliError::Connection("bridge channel closed".into()))?;
        reply_rx
            .await
            .map_err(|_| AcpCliError::Connection("bridge reply dropped".into()))?
    }

    /// Gracefully shut down the bridge, killing the agent process and joining
    /// the background thread.
    pub async fn shutdown(self) -> Result<()> {
        let _ = self.cmd_tx.send(BridgeCommand::Shutdown).await;
        match self.handle.await {
            Ok(Ok(())) => Ok(()),
            Ok(Err(e)) => Err(e),
            Err(e) => Err(AcpCliError::Connection(format!("join: {e}"))),
        }
    }
}

/// Main loop running inside `spawn_blocking` + `LocalSet`.
///
/// Spawns the agent child process, establishes the ACP connection, initializes
/// the protocol, creates a session, then enters a command loop that processes
/// `BridgeCommand` messages from the main thread.
async fn acp_thread_main(
    mut cmd_rx: mpsc::Receiver<BridgeCommand>,
    evt_tx: mpsc::UnboundedSender<BridgeEvent>,
    command: String,
    args: Vec<String>,
    cwd: PathBuf,
) -> Result<()> {
    // 1. Spawn agent process
    let args_refs: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
    let mut cmd = tokio::process::Command::new(&command);
    cmd.args(&args_refs)
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::inherit())
        .kill_on_drop(true);

    // Remove ANTHROPIC_API_KEY to prevent OAuth setup tokens (sk-ant-oat01-*)
    // from being misidentified as API keys by the Claude Agent SDK.
    cmd.env_remove("ANTHROPIC_API_KEY");

    // Inject auth token for non-OAuth tokens only.
    // OAuth tokens (sk-ant-oat01-*) must NOT be injected via env var —
    // the Claude Agent SDK's auth flow doesn't support the oauth-2025-04-20
    // beta header, causing 401. Let the SDK resolve OAuth tokens itself
    // from Keychain / ~/.claude.json.
    if let Some(token) = resolve_claude_auth_token()
        && !token.starts_with("sk-ant-oat01-")
    {
        cmd.env("ANTHROPIC_AUTH_TOKEN", &token);
    }

    let mut child = cmd
        .spawn()
        .map_err(|e| AcpCliError::Connection(format!("spawn {command}: {e}")))?;

    let stdin = child
        .stdin
        .take()
        .ok_or_else(|| AcpCliError::Connection("agent has no stdin".into()))?;
    let stdout = child
        .stdout
        .take()
        .ok_or_else(|| AcpCliError::Connection("agent has no stdout".into()))?;

    let client = BridgedAcpClient::new(evt_tx.clone());

    let (conn, handle_io) =
        acp::ClientSideConnection::new(client, stdin.compat_write(), stdout.compat(), |fut| {
            tokio::task::spawn_local(fut);
        });

    // Drive the I/O loop in the background on the local task set.
    tokio::task::spawn_local(async move {
        if let Err(e) = handle_io.await {
            eprintln!("[acp-cli] I/O error: {e}");
        }
    });

    let result = async {
        // 2. Initialize
        conn.initialize(
            acp::InitializeRequest::new(acp::ProtocolVersion::V1).client_info(
                acp::Implementation::new("acp-cli", env!("CARGO_PKG_VERSION")),
            ),
        )
        .await
        .map_err(|e| AcpCliError::Connection(format!("initialize: {e}")))?;

        // 3. Create session
        let session = conn
            .new_session(acp::NewSessionRequest::new(cwd))
            .await
            .map_err(|e| AcpCliError::Connection(format!("new_session: {e}")))?;

        let session_id = session.session_id;
        let _ = evt_tx.send(BridgeEvent::SessionCreated {
            session_id: session_id.0.to_string(),
        });

        // 4. Command loop
        while let Some(cmd) = cmd_rx.recv().await {
            match cmd {
                BridgeCommand::Prompt { messages, reply } => {
                    let content_blocks: Vec<acp::ContentBlock> =
                        messages.into_iter().map(|m| m.into()).collect();
                    let result = conn
                        .prompt(acp::PromptRequest::new(session_id.clone(), content_blocks))
                        .await;
                    match result {
                        Ok(response) => {
                            let stop_reason = serde_json::to_value(response.stop_reason)
                                .ok()
                                .and_then(|v| v.as_str().map(String::from))
                                .unwrap_or_else(|| "unknown".to_string());
                            let _ = evt_tx.send(BridgeEvent::PromptDone {
                                stop_reason: stop_reason.clone(),
                            });
                            // Content was already streamed via session_notification -> TextChunk.
                            // The main thread collects content from BridgeEvent::TextChunk events.
                            let _ = reply.send(Ok(PromptResult {
                                content: String::new(),
                                stop_reason,
                            }));
                        }
                        Err(e) => {
                            let _ = reply.send(Err(AcpCliError::Agent(format!("{e}"))));
                        }
                    }
                }
                BridgeCommand::Cancel => {
                    // ACP cancel not yet implemented in SDK
                }
                BridgeCommand::SetMode { mode, reply } => {
                    let mode_id = acp::SessionModeId::new(mode);
                    let request = acp::SetSessionModeRequest::new(session_id.clone(), mode_id);
                    match conn.set_session_mode(request).await {
                        Ok(_) => {
                            let _ = reply.send(Ok(()));
                        }
                        Err(e) => {
                            let _ = reply
                                .send(Err(AcpCliError::Agent(format!("set_session_mode: {e}"))));
                        }
                    }
                }
                BridgeCommand::SetConfig { key, value, reply } => {
                    let config_id = acp::SessionConfigId::new(key);
                    let value_id = acp::SessionConfigValueId::new(value);
                    let request = acp::SetSessionConfigOptionRequest::new(
                        session_id.clone(),
                        config_id,
                        value_id,
                    );
                    match conn.set_session_config_option(request).await {
                        Ok(_) => {
                            let _ = reply.send(Ok(()));
                        }
                        Err(e) => {
                            let _ = reply.send(Err(AcpCliError::Agent(format!(
                                "set_session_config_option: {e}"
                            ))));
                        }
                    }
                }
                BridgeCommand::Shutdown => break,
            }
        }
        Ok(())
    }
    .await;

    // Cleanup: always reap child to avoid zombie accumulation.
    reap_child_process(&mut child).await;
    result
}

async fn reap_child_process(child: &mut tokio::process::Child) {
    if !matches!(child.try_wait(), Ok(Some(_))) {
        let _ = child.start_kill();
    }
    let _ = child.wait().await;
}

// ---------------------------------------------------------------------------
// Claude auth token resolution
// ---------------------------------------------------------------------------

/// Resolve an OAuth token for claude-agent-acp, checking (in order):
/// 1. `ANTHROPIC_AUTH_TOKEN` env var (already set externally)
/// 2. `~/.acp-cli/config.json` → `auth_token` (set via `acp-cli init`)
/// 3. `~/.claude.json` → `oauthAccount.accessToken`
/// 4. macOS Keychain (`security find-generic-password`)
fn resolve_claude_auth_token() -> Option<String> {
    // 1. Already set externally
    if let Some(t) = std::env::var("ANTHROPIC_AUTH_TOKEN")
        .ok()
        .filter(|t| !t.is_empty())
    {
        return Some(t);
    }

    // 2. acp-cli config
    let config = crate::config::AcpCliConfig::load();
    if let Some(token) = config.auth_token.filter(|t| !t.is_empty()) {
        return Some(token);
    }

    // 3. ~/.claude.json
    if let Some(token) = read_claude_json_token() {
        return Some(token);
    }

    // 4. macOS Keychain
    #[cfg(target_os = "macos")]
    if let Some(token) = read_keychain_token() {
        return Some(token);
    }

    None
}

/// Read the OAuth access token from `~/.claude.json`.
fn read_claude_json_token() -> Option<String> {
    let path = dirs::home_dir()?.join(".claude.json");
    let content = std::fs::read_to_string(path).ok()?;
    let json: serde_json::Value = serde_json::from_str(&content).ok()?;

    json.pointer("/oauthAccount/accessToken")
        .or_else(|| json.get("accessToken"))
        .and_then(|v| v.as_str())
        .filter(|s| !s.is_empty())
        .map(|s| s.to_string())
}

/// Read the Claude Code OAuth token from the macOS Keychain.
#[cfg(target_os = "macos")]
fn read_keychain_token() -> Option<String> {
    // Try known service names used by Claude Code
    for service in &["Claude Code", "claude.ai", "anthropic.claude"] {
        let output = std::process::Command::new("security")
            .args(["find-generic-password", "-s", service, "-w"])
            .stderr(std::process::Stdio::null())
            .output()
            .ok()?;
        if output.status.success() {
            let token = String::from_utf8(output.stdout).ok()?.trim().to_string();
            if !token.is_empty() {
                return Some(token);
            }
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::{AcpCliError, BridgeCommand, BridgeEvent, acp_thread_main, reap_child_process};
    use tokio::sync::mpsc;

    fn exited_child_command() -> tokio::process::Command {
        if cfg!(windows) {
            let mut c = tokio::process::Command::new("cmd");
            c.arg("/C").arg("exit 0");
            c
        } else {
            let mut c = tokio::process::Command::new("sh");
            c.arg("-c").arg("exit 0");
            c
        }
    }

    fn running_child_command() -> tokio::process::Command {
        if cfg!(windows) {
            let mut c = tokio::process::Command::new("cmd");
            c.arg("/C").arg("ping -n 30 127.0.0.1 >NUL");
            c
        } else {
            let mut c = tokio::process::Command::new("sh");
            c.arg("-c").arg("sleep 10");
            c
        }
    }

    #[tokio::test]
    async fn reap_child_process_handles_already_exited_child() {
        let mut child = exited_child_command().spawn().expect("spawn child");

        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        reap_child_process(&mut child).await;

        let status = child.wait().await.expect("wait after reap");
        assert!(status.success());
    }

    #[tokio::test]
    async fn reap_child_process_kills_running_child() {
        let mut child = running_child_command().spawn().expect("spawn child");

        reap_child_process(&mut child).await;
        let status = child.wait().await.expect("wait after kill");
        assert!(!status.success());
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn initialize_error_still_reaps_child_process() {
        use std::os::unix::fs::PermissionsExt;
        use tempfile::tempdir;

        let temp = tempdir().expect("create tempdir");
        let pid_file = temp.path().join("agent.pid");
        let script = temp.path().join("fake-acp-agent.sh");
        std::fs::write(
            &script,
            format!(
                "#!/bin/sh\necho $$ > \"{}\"\nexec 1>&-\nsleep 30\n",
                pid_file.display()
            ),
        )
        .expect("write script");
        let mut perms = std::fs::metadata(&script)
            .expect("script metadata")
            .permissions();
        perms.set_mode(0o755);
        std::fs::set_permissions(&script, perms).expect("chmod script");

        let (_cmd_tx, cmd_rx) = mpsc::channel::<BridgeCommand>(2);
        let (evt_tx, _evt_rx) = mpsc::unbounded_channel::<BridgeEvent>();

        let local = tokio::task::LocalSet::new();
        let result = local
            .run_until(tokio::time::timeout(
                std::time::Duration::from_secs(5),
                acp_thread_main(
                    cmd_rx,
                    evt_tx,
                    script.to_string_lossy().to_string(),
                    vec![],
                    temp.path().to_path_buf(),
                ),
            ))
            .await
            .expect("acp_thread_main should not hang")
            .expect_err("initialize should fail for non-ACP output");

        assert!(
            matches!(result, AcpCliError::Connection(_)),
            "expected connection error, got: {result:?}"
        );

        let pid_raw = std::fs::read_to_string(&pid_file).expect("pid file");
        let pid = pid_raw.trim().parse::<i32>().expect("parse pid");

        // Verify child is gone; if still alive, kill it to avoid test leak.
        let alive = unsafe { libc::kill(pid, 0) == 0 };
        if alive {
            let _ = unsafe { libc::kill(pid, libc::SIGKILL) };
        }
        assert!(!alive, "child process {pid} should have been reaped");
    }
}