aion_integration_cli/
harness.rs1use std::ffi::OsString;
14use std::path::PathBuf;
15use std::process::Stdio;
16
17use aion_integrations::contract::AgentHarness;
18use aion_integrations::{AgentRunSpec, HarnessError};
19use async_trait::async_trait;
20use tokio::process::{ChildStdout, Command};
21
22use crate::demux::EventIdentity;
23use crate::session::CliSession;
24
25#[derive(Clone, Debug)]
31pub struct CliHarness {
32 program: PathBuf,
33 args: Vec<OsString>,
34 env: Vec<(String, String)>,
35 pass_input_as_arg: bool,
37}
38
39impl CliHarness {
40 #[must_use]
42 pub fn new(program: impl Into<PathBuf>) -> Self {
43 Self {
44 program: program.into(),
45 args: Vec::new(),
46 env: Vec::new(),
47 pass_input_as_arg: false,
48 }
49 }
50
51 #[must_use]
53 pub fn with_arg(mut self, arg: impl Into<OsString>) -> Self {
54 self.args.push(arg.into());
55 self
56 }
57
58 #[must_use]
60 pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
61 self.env.push((key.into(), value.into()));
62 self
63 }
64
65 #[must_use]
71 pub fn passing_input_as_arg(mut self) -> Self {
72 self.pass_input_as_arg = true;
73 self
74 }
75
76 fn spawn(&self, spec: &AgentRunSpec) -> Result<tokio::process::Child, HarnessError> {
78 let mut command = Command::new(&self.program);
79 command
80 .args(&self.args)
81 .envs(self.env.iter().map(|(key, value)| (key, value)))
82 .stdin(Stdio::null())
83 .stdout(Stdio::piped())
84 .stderr(Stdio::inherit())
85 .kill_on_drop(true);
86 if self.pass_input_as_arg {
87 command.arg(input_arg(spec)?);
88 }
89 command.spawn().map_err(|source| {
90 HarnessError::transport(format!(
91 "failed to spawn `{}`: {source}",
92 self.program.display()
93 ))
94 })
95 }
96}
97
98fn input_arg(spec: &AgentRunSpec) -> Result<String, HarnessError> {
100 std::str::from_utf8(spec.input.bytes())
101 .map(str::to_owned)
102 .map_err(|source| HarnessError::protocol(format!("run input is not valid UTF-8: {source}")))
103}
104
105fn take_stdout(child: &mut tokio::process::Child) -> Result<ChildStdout, HarnessError> {
107 child
108 .stdout
109 .take()
110 .ok_or_else(|| HarnessError::transport("spawned cli child has no piped stdout"))
111}
112
113#[async_trait]
114impl AgentHarness for CliHarness {
115 type Session = CliSession;
116
117 async fn start(&self, spec: AgentRunSpec) -> Result<Self::Session, HarnessError> {
118 let mut child = self.spawn(&spec)?;
119 let stdout = take_stdout(&mut child)?;
120 let identity = EventIdentity {
121 workflow_id: spec.workflow_id,
122 run_id: spec.run_id,
123 activity_id: spec.activity_id,
124 attempt: spec.attempt,
125 };
126 Ok(CliSession::start(stdout, identity, Some(child)))
127 }
128}