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