use std::sync::Arc;
use aion_core::{
ActivityEvent, AgentOutcome, InterventionCapabilities, InterventionCommand, Payload, StopKind,
};
use aion_integrations::contract::AgentSession;
use aion_integrations::error::HarnessError;
use async_trait::async_trait;
use futures::stream::BoxStream;
use tokio::io::{AsyncBufReadExt, AsyncRead, BufReader, Lines};
use tokio::sync::{Mutex, mpsc, oneshot};
use crate::demux::{self, EventIdentity};
enum ChildEnd {
NoChild,
Exited(i32),
Signalled,
StatusUnreadable {
detail: String,
},
}
struct TerminalOutcome {
final_line: Option<String>,
end: ChildEnd,
announced_stop: Option<StopKind>,
}
pub struct CliSession {
capabilities: InterventionCapabilities,
events: Option<mpsc::UnboundedReceiver<ActivityEvent>>,
result: Option<oneshot::Receiver<TerminalOutcome>>,
pump: Option<tokio::task::JoinHandle<()>>,
identity: EventIdentity,
child: Arc<Mutex<Option<tokio::process::Child>>>,
}
impl CliSession {
pub(crate) fn start<R>(
stdout: R,
identity: EventIdentity,
child: Option<tokio::process::Child>,
) -> Self
where
R: AsyncRead + Unpin + Send + 'static,
{
let (event_tx, event_rx) = mpsc::unbounded_channel();
let (result_tx, result_rx) = oneshot::channel();
let child = Arc::new(Mutex::new(child));
let pump = tokio::spawn(pump_loop(
BufReader::new(stdout).lines(),
identity.clone(),
event_tx,
result_tx,
Arc::clone(&child),
));
Self {
capabilities: InterventionCapabilities::none(),
events: Some(event_rx),
result: Some(result_rx),
pump: Some(pump),
identity,
child,
}
}
}
async fn pump_loop<R>(
mut lines: Lines<BufReader<R>>,
identity: EventIdentity,
event_tx: mpsc::UnboundedSender<ActivityEvent>,
result_tx: oneshot::Sender<TerminalOutcome>,
child: Arc<Mutex<Option<tokio::process::Child>>>,
) where
R: AsyncRead + Unpin + Send,
{
let mut worker_seq: u64 = 0;
let mut final_line: Option<String> = None;
let mut announced_stop: Option<StopKind> = None;
loop {
match lines.next_line().await {
Ok(Some(line)) => {
if !line.trim().is_empty() {
final_line = Some(line.clone());
}
let event = demux::line_to_event(&identity, worker_seq, &line);
worker_seq = worker_seq.saturating_add(1);
if let aion_core::ActivityEventKind::Stop { reason } = &event.kind {
announced_stop = Some(reason.clone());
}
let _ = event_tx.send(event);
}
Ok(None) => break,
Err(error) => {
tracing::debug!(%error, "cli adapter: reader pump ended on a stdout read error");
break;
}
}
}
let end = wait_child_end(&child).await;
if matches!(end, ChildEnd::Signalled) {
let stop = demux::synthesized_stop_event(
&identity,
worker_seq,
StopKind::Other {
reason: "terminated by signal".to_owned(),
},
);
let _ = event_tx.send(stop);
}
let _ = result_tx.send(TerminalOutcome {
final_line,
end,
announced_stop,
});
}
async fn wait_child_end(child: &Arc<Mutex<Option<tokio::process::Child>>>) -> ChildEnd {
let mut guard = child.lock().await;
let Some(child) = guard.as_mut() else {
return ChildEnd::NoChild;
};
match child.wait().await {
Ok(status) => match status.code() {
Some(code) => ChildEnd::Exited(code),
None => ChildEnd::Signalled,
},
Err(error) => ChildEnd::StatusUnreadable {
detail: error.to_string(),
},
}
}
#[async_trait]
impl AgentSession for CliSession {
fn capabilities(&self) -> &InterventionCapabilities {
&self.capabilities
}
fn events(&mut self) -> BoxStream<'static, ActivityEvent> {
match self.events.take() {
Some(receiver) => Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(
receiver,
)),
None => Box::pin(futures::stream::empty()),
}
}
async fn intervene(&self, cmd: InterventionCommand) -> Result<(), HarnessError> {
Err(HarnessError::capability_not_supported(format!(
"{:?}",
cmd.kind.primitive()
)))
}
async fn wait_result(mut self) -> Result<Payload, HarnessError> {
let receiver = self
.result
.take()
.ok_or_else(|| HarnessError::protocol("run result already awaited"))?;
let outcome = receiver.await.map_err(|_recv| {
HarnessError::transport("cli agent result channel closed before the run finished")
})?;
terminal_outcome_into_payload(outcome, &self.identity)
}
}
fn terminal_outcome_into_payload(
outcome: TerminalOutcome,
identity: &EventIdentity,
) -> Result<Payload, HarnessError> {
let TerminalOutcome {
final_line,
end,
announced_stop,
} = outcome;
let stop = match end {
ChildEnd::Exited(code) if code != 0 => {
return Err(HarnessError::harness(format!(
"cli agent exited with a non-zero status: {code}"
)));
}
ChildEnd::StatusUnreadable { detail } => {
return Err(HarnessError::transport(format!(
"the cli agent's exit status could not be read, so the run has no honest \
terminal classification: {detail}"
)));
}
ChildEnd::Signalled => StopKind::Other {
reason: "terminated by signal".to_owned(),
},
ChildEnd::Exited(_) | ChildEnd::NoChild => announced_stop.unwrap_or(StopKind::EndTurn),
};
let session_id = format!(
"{}:{}:{}:{}",
identity.workflow_id, identity.run_id, identity.activity_id, identity.attempt
);
let final_line = final_line.unwrap_or_default();
AgentOutcome::new(final_line.clone(), final_line, &stop, session_id)
.into_payload()
.map_err(|source| {
HarnessError::protocol(format!("cli agent result is not encodable: {source}"))
})
}
impl Drop for CliSession {
fn drop(&mut self) {
if let Some(pump) = self.pump.take() {
pump.abort();
}
if let Ok(mut guard) = self.child.try_lock()
&& let Some(child) = guard.as_mut()
{
let _ = child.start_kill();
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests;