use std::{
collections::BTreeMap, ffi::OsString, os::unix::process::ExitStatusExt, path::PathBuf,
process::Stdio,
};
use async_trait::async_trait;
use tokio::process::{Child, Command};
use crate::{
cli::tui::TerminalControl,
config::app::ProviderId,
domain::errors::{AgentError, AgentResult, ErrorCode},
providers::{
descriptor::{LoginAction, descriptor},
probe::{ProviderProbe, ProviderProbeResult, ProviderProbeRuntime, ToolLaunch},
},
};
#[derive(Clone, PartialEq, Eq)]
pub(crate) struct InheritedCommandSpec {
pub(crate) program: PathBuf,
pub(crate) args: Vec<OsString>,
pub(crate) env: BTreeMap<OsString, OsString>,
pub(crate) cwd: PathBuf,
pub(crate) stdio: InheritedStdio,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum InheritedStdio {
Inherit,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum InheritedExit {
Success,
Nonzero,
Signaled,
Interrupted,
}
#[async_trait]
pub(crate) trait InheritedChild: Send {
async fn wait(&mut self) -> AgentResult<InheritedExit>;
fn terminate(&mut self) -> AgentResult<()>;
}
#[async_trait]
pub(crate) trait InheritedSpawner: Send + Sync {
async fn spawn(&self, spec: &InheritedCommandSpec) -> AgentResult<Box<dyn InheritedChild>>;
}
#[async_trait]
pub(crate) trait InterruptSignal: Send + Sync {
async fn wait(&self) -> AgentResult<()>;
}
#[async_trait]
pub(crate) trait ProviderReprobe: Send + Sync {
async fn reprobe(&self, provider: ProviderId) -> AgentResult<ProviderProbeResult>;
}
#[async_trait]
impl ProviderReprobe for ProviderProbe<'_> {
async fn reprobe(&self, provider: ProviderId) -> AgentResult<ProviderProbeResult> {
self.probe(provider).await
}
}
pub(crate) struct LoginHandoff<'a> {
runtime: ProviderProbeRuntime,
terminal: &'a mut dyn TerminalControl,
spawner: &'a dyn InheritedSpawner,
interrupt: &'a dyn InterruptSignal,
reprobe: &'a dyn ProviderReprobe,
}
impl<'a> LoginHandoff<'a> {
pub(crate) fn new(
runtime: ProviderProbeRuntime,
terminal: &'a mut dyn TerminalControl,
spawner: &'a dyn InheritedSpawner,
interrupt: &'a dyn InterruptSignal,
reprobe: &'a dyn ProviderReprobe,
) -> Self {
Self {
runtime,
terminal,
spawner,
interrupt,
reprobe,
}
}
pub(crate) async fn run(self, provider: ProviderId) -> AgentResult<ProviderProbeResult> {
let Self {
runtime,
terminal,
spawner,
interrupt,
reprobe,
} = self;
let cwd = runtime.validated_cwd()?;
let mut terminal = SuspendedTerminal::new(terminal)?;
let action_result =
run_actions(provider, &runtime, &cwd, &mut terminal, spawner, interrupt).await;
terminal.resume_once()?;
action_result?;
reprobe.reprobe(provider).await
}
}
struct SuspendedTerminal<'a> {
terminal: &'a mut dyn TerminalControl,
resume_attempted: bool,
}
impl<'a> SuspendedTerminal<'a> {
fn new(terminal: &'a mut dyn TerminalControl) -> AgentResult<Self> {
terminal.suspend()?;
Ok(Self {
terminal,
resume_attempted: false,
})
}
fn show_instruction(&mut self, instruction: &'static str) -> AgentResult<()> {
self.terminal.show_instruction(instruction)
}
fn resume_once(&mut self) -> AgentResult<()> {
if self.resume_attempted {
return Ok(());
}
self.resume_attempted = true;
self.terminal.resume()
}
}
impl Drop for SuspendedTerminal<'_> {
fn drop(&mut self) {
if !self.resume_attempted {
self.resume_attempted = true;
let _ = self.terminal.resume();
}
}
}
async fn run_actions(
provider: ProviderId,
runtime: &ProviderProbeRuntime,
cwd: &std::path::Path,
terminal: &mut SuspendedTerminal<'_>,
spawner: &dyn InheritedSpawner,
interrupt: &dyn InterruptSignal,
) -> AgentResult<()> {
for action in descriptor(provider).login {
let spec = match action {
LoginAction::InheritedCommand { program, args } => {
if *program != "claude" {
return Err(login_invariant_failed());
}
let Some(claude) = &runtime.claude else {
break;
};
inherited_spec(claude, args, true, cwd)?
}
LoginAction::InteractivePi { instruction } => {
terminal.show_instruction(instruction)?;
inherited_spec(&runtime.pi, &[], false, cwd)?
}
};
let outcome = match run_child(spawner, interrupt, &spec).await {
Ok(outcome) => outcome,
Err(_) => break,
};
if outcome != InheritedExit::Success {
break;
}
}
Ok(())
}
fn inherited_spec(
launch: &ToolLaunch,
args: &[&str],
managed_claude: bool,
cwd: &std::path::Path,
) -> AgentResult<InheritedCommandSpec> {
let command = launch.command(args, managed_claude, cwd)?;
Ok(InheritedCommandSpec {
program: command.program,
args: command.args,
env: command.env,
cwd: command.cwd.ok_or_else(login_invariant_failed)?,
stdio: InheritedStdio::Inherit,
})
}
async fn run_child(
spawner: &dyn InheritedSpawner,
interrupt: &dyn InterruptSignal,
spec: &InheritedCommandSpec,
) -> AgentResult<InheritedExit> {
let mut child = spawner.spawn(spec).await?;
enum Selection {
Exited(AgentResult<InheritedExit>),
Interrupted(AgentResult<()>),
}
let selection = {
let wait = child.wait();
tokio::pin!(wait);
tokio::select! {
biased;
result = &mut wait => Selection::Exited(result),
signal = interrupt.wait() => Selection::Interrupted(signal),
}
};
match selection {
Selection::Exited(result) => result,
Selection::Interrupted(signal) => {
let terminated = child.terminate();
let reaped = child.wait().await;
signal?;
terminated?;
reaped?;
Ok(InheritedExit::Interrupted)
}
}
}
#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct TokioInheritedSpawner;
#[async_trait]
impl InheritedSpawner for TokioInheritedSpawner {
async fn spawn(&self, spec: &InheritedCommandSpec) -> AgentResult<Box<dyn InheritedChild>> {
let InheritedStdio::Inherit = spec.stdio;
let child = Command::new(&spec.program)
.args(&spec.args)
.envs(&spec.env)
.current_dir(&spec.cwd)
.stdin(Stdio::inherit())
.stdout(Stdio::inherit())
.stderr(Stdio::inherit())
.kill_on_drop(true)
.spawn()
.map_err(|_| child_failed("inherited login command failed to start"))?;
Ok(Box::new(TokioInheritedChild { child }))
}
}
struct TokioInheritedChild {
child: Child,
}
#[async_trait]
impl InheritedChild for TokioInheritedChild {
async fn wait(&mut self) -> AgentResult<InheritedExit> {
let status = self
.child
.wait()
.await
.map_err(|_| child_failed("inherited login command wait failed"))?;
if status.success() {
Ok(InheritedExit::Success)
} else if status.signal().is_some() {
Ok(InheritedExit::Signaled)
} else {
Ok(InheritedExit::Nonzero)
}
}
fn terminate(&mut self) -> AgentResult<()> {
self.child
.start_kill()
.map_err(|_| child_failed("inherited login command termination failed"))
}
}
#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct TokioInterruptSignal;
#[async_trait]
impl InterruptSignal for TokioInterruptSignal {
async fn wait(&self) -> AgentResult<()> {
tokio::signal::ctrl_c()
.await
.map_err(|_| child_failed("interrupt listener failed"))
}
}
fn child_failed(message: &'static str) -> AgentError {
AgentError::new(ErrorCode::InvalidMessage, message)
}
fn login_invariant_failed() -> AgentError {
AgentError::new(
ErrorCode::InvalidMessage,
"unsupported provider login action",
)
}