Skip to main content

rho_coding_agent/app/
automation.rs

1use std::{
2    fmt,
3    io::{self, Read, Write},
4    path::PathBuf,
5    sync::Arc,
6};
7
8use rho_sdk::{
9    CapabilityRequest, PolicyDecision, SessionOptions, SystemPrompt, UserInput, Workspace,
10    WorkspacePolicy,
11};
12
13use crate::{
14    cli::Command,
15    config::Config,
16    credentials::OsCredentialStore,
17    diagnostics::RuntimeDiagnostics,
18    herdr::{HerdrReporter, HerdrState},
19    prompt,
20    providers::build_automation_provider,
21    subagent::{self, Preset, RunState, RunStatus},
22    tools::sdk_registry::{AppToolSet, ToolSetOptions},
23    tui::AttachmentWriter,
24};
25
26use super::{
27    runtime_builder::{build_runtime, configured_context_window, RuntimeBuildOptions},
28    sdk_config::SdkBootstrapOptions,
29};
30
31/// Error returned after an automation run handles an interrupt and completes cleanup.
32#[derive(Debug)]
33pub struct AutomationInterrupted {
34    signal: ShutdownSignal,
35}
36
37impl AutomationInterrupted {
38    fn new(signal: ShutdownSignal) -> Self {
39        Self { signal }
40    }
41
42    /// Returns the conventional process exit code for the received signal.
43    pub fn exit_code(&self) -> u8 {
44        match self.signal {
45            ShutdownSignal::Interrupt => 130,
46            ShutdownSignal::Terminate => 143,
47        }
48    }
49}
50
51impl fmt::Display for AutomationInterrupted {
52    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
53        write!(formatter, "rho run interrupted by {}", self.signal)
54    }
55}
56
57impl std::error::Error for AutomationInterrupted {}
58
59#[derive(Clone, Copy, Debug)]
60enum ShutdownSignal {
61    Interrupt,
62    Terminate,
63}
64
65impl fmt::Display for ShutdownSignal {
66    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
67        match self {
68            Self::Interrupt => formatter.write_str("SIGINT"),
69            Self::Terminate => formatter.write_str("SIGTERM"),
70        }
71    }
72}
73
74#[derive(Debug)]
75struct SubagentCancelled;
76
77impl fmt::Display for SubagentCancelled {
78    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
79        formatter.write_str("subagent cancellation requested")
80    }
81}
82
83impl std::error::Error for SubagentCancelled {}
84
85pub(super) struct Startup<'a> {
86    pub config: &'a Config,
87    pub config_path: PathBuf,
88    pub cwd: PathBuf,
89    pub no_system_prompt: bool,
90    pub no_tools: bool,
91    pub no_subagents: bool,
92    pub preset: Option<Preset>,
93    pub output_file: Option<PathBuf>,
94    pub diagnostics: RuntimeDiagnostics,
95    pub herdr: HerdrReporter,
96}
97
98pub(super) fn prompt_for_command(command: &Option<Command>) -> anyhow::Result<Option<String>> {
99    match command {
100        Some(Command::Run { prompt, stdin, .. }) => {
101            prompt_from_stdin(prompt.clone(), *stdin).map(Some)
102        }
103        Some(Command::Attach { .. } | Command::Login { .. } | Command::Update) | None => Ok(None),
104    }
105}
106
107pub(super) async fn run(prompt_text: String, startup: Startup<'_>) -> anyhow::Result<()> {
108    // The reporter exists before anything that can fail, so a parent process
109    // watching the output file always sees a terminal state — even when the
110    // run dies during startup (bad auth, broken workspace, ...).
111    let mut reporter = startup
112        .output_file
113        .as_ref()
114        .map(|path| {
115            RunReporter::new(
116                path.clone(),
117                startup.preset.as_ref().map(|preset| preset.name.clone()),
118                startup.cwd.clone(),
119                &prompt_text,
120            )
121        })
122        .transpose()?;
123    let result = run_session(prompt_text, &startup, reporter.as_mut()).await;
124    if let Some(reporter) = reporter.as_mut() {
125        reporter.finish(&result);
126    }
127    let answer = result?;
128    let mut stdout = io::stdout().lock();
129    if reporter.is_some() {
130        // The answer already streamed above and is in the result file.
131        writeln!(stdout, "\n[subagent run complete]")?;
132    } else {
133        writeln!(stdout, "{}", answer.text())?;
134    }
135    stdout.flush()?;
136    Ok(())
137}
138
139async fn run_session(
140    prompt_text: String,
141    startup: &Startup<'_>,
142    reporter: Option<&mut RunReporter>,
143) -> anyhow::Result<rho_sdk::RunOutcome> {
144    let sdk_options = SdkBootstrapOptions::from_config(startup.config, &startup.cwd)?;
145    let credentials = crate::auth::provider_credentials::ApplicationCredentialSource::new(
146        Arc::new(OsCredentialStore),
147    );
148    let provider = build_automation_provider(sdk_options.provider, &credentials)?;
149    let subagents_enabled = startup.config.enable_subagents && !startup.no_subagents;
150    let mut tool_set = if startup.no_tools {
151        AppToolSet::disabled()
152    } else {
153        let subagents = subagents_enabled.then(|| startup.cwd.clone());
154        AppToolSet::new(
155            startup.config,
156            startup.diagnostics.clone(),
157            ToolSetOptions::default()
158                .subagents(subagents)
159                .subagent_config_path(startup.config_path.clone()),
160        )
161    };
162    if let Some(allowed) = startup
163        .preset
164        .as_ref()
165        .and_then(|preset| preset.tools.as_ref())
166    {
167        // The preset's tool list is the subagent's permission boundary:
168        // anything not listed is never registered, so it cannot run.
169        tool_set.retain_named(allowed);
170    }
171    let tool_specs = tool_set.specs();
172    let system_prompt = if startup.no_system_prompt {
173        startup.diagnostics.update_prompt_sources(Vec::new());
174        SystemPrompt::None
175    } else {
176        let system_prompt = prompt::system_prompt(&tool_specs, &startup.cwd);
177        startup
178            .diagnostics
179            .update_prompt_sources(system_prompt.sources);
180        let mut text = system_prompt.text;
181        if !subagents_enabled {
182            prompt::append_subagents_disabled_instruction(&mut text);
183        }
184        if let Some(preset) = &startup.preset {
185            if !preset.prompt.is_empty() {
186                text.push_str("\n\n# Subagent instructions\n\n");
187                text.push_str(&preset.prompt);
188            }
189        }
190        SystemPrompt::Custom(text)
191    };
192    startup.diagnostics.update_tools(&tool_specs);
193
194    let workspace = Workspace::new(&sdk_options.workspace.root)?;
195    let context_window = configured_context_window(startup.config);
196    let compaction = sdk_options.runtime.compaction.clone();
197    startup.diagnostics.update_compaction_config(&compaction);
198    let runtime = build_runtime(RuntimeBuildOptions {
199        provider,
200        tools: tool_set.tools(),
201        workspace,
202        workspace_policy: AutomationWorkspacePolicy,
203        system_prompt,
204        reasoning: sdk_options.runtime.reasoning,
205        compaction,
206        context_window,
207    })?;
208    let session = runtime.session(SessionOptions::default()).await?;
209
210    startup
211        .herdr
212        .report_state(HerdrState::Working, None, None)
213        .await;
214    let result = complete_run(&session, prompt_text, reporter).await;
215
216    runtime.shutdown();
217    tool_set.shutdown().await;
218    startup
219        .herdr
220        .report_state(HerdrState::Idle, None, None)
221        .await;
222    startup.herdr.release().await;
223
224    result
225}
226
227async fn complete_run(
228    session: &rho_sdk::Session,
229    prompt_text: String,
230    reporter: Option<&mut RunReporter>,
231) -> anyhow::Result<rho_sdk::RunOutcome> {
232    let mut run = session.start(UserInput::text(prompt_text)).await?;
233    let cancellation = run.cancellation_handle();
234    let cancel_file = reporter
235        .as_ref()
236        .map(|reporter| reporter.cancel_file.clone());
237    tokio::select! {
238        outcome = drive_headless_run(&mut run, reporter) => outcome,
239        signal = shutdown_signal() => {
240            let signal = signal?;
241            cancellation.cancel();
242            let _ = run.outcome().await;
243            Err(AutomationInterrupted::new(signal).into())
244        }
245        cancelled = wait_for_cancel_request(cancel_file) => {
246            cancelled?;
247            cancellation.cancel();
248            let _ = run.outcome().await;
249            Err(SubagentCancelled.into())
250        }
251    }
252}
253
254/// Drains run events with no interactive host attached.
255///
256/// Host input requests cannot be answered headlessly; cancel instead of
257/// leaving the requesting tool suspended until a signal arrives.
258async fn drive_headless_run(
259    run: &mut rho_sdk::Run,
260    mut reporter: Option<&mut RunReporter>,
261) -> anyhow::Result<rho_sdk::RunOutcome> {
262    let mut heartbeat = tokio::time::interval(REPORT_HEARTBEAT);
263    heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
264    loop {
265        let event = tokio::select! {
266            event = run.next_event() => event,
267            _ = heartbeat.tick(), if reporter.is_some() => {
268                if let Some(reporter) = reporter.as_deref_mut() {
269                    reporter.write();
270                }
271                continue;
272            }
273        };
274        let Some(event) = event else {
275            break;
276        };
277        if let Some(reporter) = reporter.as_deref_mut() {
278            reporter.on_event(&event);
279        }
280        if let rho_sdk::RunEvent::HostInputRequested { request } = event {
281            run.cancel();
282            let _ = run.outcome().await;
283            anyhow::bail!(
284                "rho run cannot answer host input request '{}' ({}); run without tools that require interactive input",
285                request.id(),
286                request.title(),
287            );
288        }
289    }
290    Ok(run.outcome().await?)
291}
292
293/// Maintains the `--output-file` status contract for subagent runs and
294/// streams progress to stdout so a watching pane shows live activity.
295struct RunReporter {
296    path: PathBuf,
297    cancel_file: PathBuf,
298    status: RunStatus,
299    attachment: Option<AttachmentWriter>,
300    last_write: std::time::Instant,
301}
302
303/// Longest a status-file write is deferred while text streams.
304const REPORT_THROTTLE: std::time::Duration = std::time::Duration::from_secs(2);
305/// Keeps the status file fresh while a provider or tool call emits no events.
306const REPORT_HEARTBEAT: std::time::Duration = std::time::Duration::from_secs(10);
307const LAST_TEXT_BYTES: usize = 400;
308
309impl RunReporter {
310    fn new(
311        path: PathBuf,
312        preset: Option<String>,
313        cwd: PathBuf,
314        prompt: &str,
315    ) -> anyhow::Result<Self> {
316        let cancel_file = subagent::cancel_file_for(&path);
317        match std::fs::remove_file(&cancel_file) {
318            Ok(()) => {}
319            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
320            Err(error) => return Err(error.into()),
321        }
322        let status = RunStatus {
323            state: RunState::Starting,
324            preset,
325            ..RunStatus::default()
326        };
327        subagent::write_status(&path, &status)?;
328        let attachment = match AttachmentWriter::new(&path, cwd, prompt) {
329            Ok(attachment) => Some(attachment),
330            Err(error) => {
331                let mut status = status;
332                status.attachment_error = Some(format!("could not record attach output: {error}"));
333                subagent::write_status(&path, &status)?;
334                return Ok(Self {
335                    path,
336                    cancel_file,
337                    status,
338                    attachment: None,
339                    last_write: std::time::Instant::now(),
340                });
341            }
342        };
343        Ok(Self {
344            path,
345            cancel_file,
346            status,
347            attachment,
348            last_write: std::time::Instant::now(),
349        })
350    }
351
352    fn on_event(&mut self, event: &rho_sdk::RunEvent) {
353        use rho_sdk::RunEvent;
354
355        if let Some(attachment) = self.attachment.as_mut() {
356            if let Err(error) = attachment.on_event(event) {
357                self.status.attachment_error =
358                    Some(format!("could not record attach output: {error}"));
359                self.attachment = None;
360                self.write();
361            }
362        }
363        match event {
364            RunEvent::StepStarted { step } => {
365                self.status.state = RunState::Running;
366                self.status.turns = *step as u64;
367                self.write();
368            }
369            RunEvent::ToolStarted { name, .. } => {
370                self.status.last_activity = Some(format!("tool: {name}"));
371                self.stream(&format!("\n[tool] {name}\n"));
372                self.write();
373            }
374            RunEvent::AssistantTextDelta { text } => {
375                self.status.last_activity = Some("assistant text".into());
376                append_tail(
377                    self.status.last_text.get_or_insert_with(String::new),
378                    text,
379                    LAST_TEXT_BYTES,
380                );
381                self.stream(text);
382                self.write_throttled();
383            }
384            RunEvent::UsageUpdated { usage } => {
385                self.status.input_tokens = usage.total_input_tokens().unwrap_or(0);
386                self.status.output_tokens = usage.output_tokens.unwrap_or(0);
387            }
388            _ => {}
389        }
390    }
391
392    fn finish(&mut self, result: &anyhow::Result<rho_sdk::RunOutcome>) {
393        match result {
394            Ok(outcome) => {
395                self.status.state = RunState::Ok;
396                self.status.result = Some(outcome.text().to_string());
397                let usage = outcome.usage();
398                self.status.input_tokens = usage.total_input_tokens().unwrap_or(0);
399                self.status.output_tokens = usage.output_tokens.unwrap_or(0);
400            }
401            Err(error)
402                if error.is::<AutomationInterrupted>() || error.is::<SubagentCancelled>() =>
403            {
404                self.status.state = RunState::Stopped;
405                self.status.result = self
406                    .status
407                    .last_text
408                    .as_ref()
409                    .map(|text| format!("(partial, stopped before finishing)\n{text}"));
410            }
411            Err(error) => {
412                self.status.state = RunState::Error;
413                self.status.error = Some(format!("{error:#}"));
414            }
415        }
416        self.write();
417    }
418
419    fn stream(&self, text: &str) {
420        let mut stdout = io::stdout().lock();
421        let _ = stdout.write_all(text.as_bytes());
422        let _ = stdout.flush();
423    }
424
425    fn write_throttled(&mut self) {
426        if self.last_write.elapsed() >= REPORT_THROTTLE {
427            self.write();
428        }
429    }
430
431    fn write(&mut self) {
432        self.last_write = std::time::Instant::now();
433        let _ = subagent::write_status(&self.path, &self.status);
434    }
435}
436
437const CANCEL_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(100);
438
439async fn wait_for_cancel_request(cancel_file: Option<PathBuf>) -> io::Result<()> {
440    let Some(cancel_file) = cancel_file else {
441        return std::future::pending().await;
442    };
443    loop {
444        match tokio::fs::metadata(&cancel_file).await {
445            Ok(_) => return Ok(()),
446            Err(error) if error.kind() == io::ErrorKind::NotFound => {}
447            Err(error) => return Err(error),
448        }
449        tokio::time::sleep(CANCEL_POLL_INTERVAL).await;
450    }
451}
452
453/// Appends to a rolling tail buffer capped at `max` bytes.
454fn append_tail(buffer: &mut String, text: &str, max: usize) {
455    buffer.push_str(text);
456    if buffer.len() > max {
457        let cut = buffer.len() - max;
458        let boundary = (cut..buffer.len())
459            .find(|index| buffer.is_char_boundary(*index))
460            .unwrap_or(buffer.len());
461        buffer.drain(..boundary);
462    }
463}
464
465#[cfg(unix)]
466async fn shutdown_signal() -> io::Result<ShutdownSignal> {
467    use tokio::signal::unix::{signal, SignalKind};
468
469    let mut interrupt = signal(SignalKind::interrupt())?;
470    let mut terminate = signal(SignalKind::terminate())?;
471    tokio::select! {
472        _ = interrupt.recv() => Ok(ShutdownSignal::Interrupt),
473        _ = terminate.recv() => Ok(ShutdownSignal::Terminate),
474    }
475}
476
477#[cfg(not(unix))]
478async fn shutdown_signal() -> io::Result<ShutdownSignal> {
479    tokio::signal::ctrl_c().await?;
480    Ok(ShutdownSignal::Interrupt)
481}
482
483#[derive(Clone, Copy, Debug)]
484struct AutomationWorkspacePolicy;
485
486impl WorkspacePolicy for AutomationWorkspacePolicy {
487    fn evaluate(&self, _request: &CapabilityRequest) -> PolicyDecision {
488        PolicyDecision::Allow
489    }
490}
491
492fn prompt_from_stdin(parts: Vec<String>, read_stdin: bool) -> anyhow::Result<String> {
493    prompt_from_reader(parts, read_stdin, &mut io::stdin())
494}
495
496fn prompt_from_reader(
497    parts: Vec<String>,
498    read_stdin: bool,
499    stdin: &mut impl Read,
500) -> anyhow::Result<String> {
501    let mut chunks = Vec::new();
502    let inline = parts.join(" ").trim().to_string();
503    if !inline.is_empty() {
504        chunks.push(inline);
505    }
506    if read_stdin {
507        let mut buffer = String::new();
508        stdin.read_to_string(&mut buffer)?;
509        let buffer = buffer.trim().to_string();
510        if !buffer.is_empty() {
511            chunks.push(buffer);
512        }
513    }
514
515    let prompt = chunks.join("\n\n");
516    if prompt.is_empty() {
517        anyhow::bail!("rho run requires a prompt argument or --stdin");
518    }
519    Ok(prompt)
520}
521
522#[cfg(test)]
523#[path = "automation_tests.rs"]
524mod tests;