pi/modes/print/
mod.rs

1//! Print (single-shot) mode: send prompts, render output, exit.
2//!
3//! Ports `.references/pi/packages/coding-agent/src/modes/print-mode.ts` into
4//! three independent pieces so the runtime binding is deferred:
5//!
6//! - [`input`] assembles the initial prompt from piped stdin, `@file`
7//!   arguments, and CLI messages.
8//! - [`text`] drains an [`AgentSessionEvent`] stream into final stdout text.
9//! - [`json`] emits the session header followed by a lossless JSONL event
10//!   stream.
11//!
12//! Both renderers are generic over a [`PrintSink`] (protocol stdout + product
13//! stderr) and a [`Stream`] of events, so tests inject in-memory buffers and
14//! the real process binds [`OutputGuardSink`] + the live session subscription.
15
16pub mod input;
17pub mod json;
18pub mod text;
19
20pub use input::{
21    PrintInputError, ProcessFileOptions, ProcessedFiles, PromptSource, build_initial_message,
22    process_file_arguments, read_piped_stdin,
23};
24pub use json::{render_json, render_json_event, render_json_events, render_json_header};
25pub use text::{TextOutcome, TextRenderer, render_text};
26
27use std::io;
28use std::sync::{Arc, Mutex};
29
30use futures::Stream;
31use pi_ai::ImageContent;
32
33use crate::core::agent_session::AgentSessionEvent;
34use crate::core::output_guard::{ProductOutput, flush_raw_stdout, write_raw_stdout};
35use crate::core::sessions::SessionHeader;
36
37/// Output mode for print (single-shot) runs.
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum PrintOutput {
40    /// Write the final assistant text only; errors go to stderr.
41    Text,
42    /// Write the session header + every event as JSONL.
43    Json,
44}
45
46impl PrintOutput {
47    /// Returns `true` for the JSON output mode.
48    #[must_use]
49    pub const fn is_json(self) -> bool {
50        matches!(self, Self::Json)
51    }
52}
53
54/// Options for [`run_print_mode`], mirroring TypeScript `PrintModeOptions`.
55#[derive(Clone, Debug)]
56pub struct PrintModeOptions {
57    /// Output mode: text or JSON.
58    pub mode: PrintOutput,
59    /// Additional prompts sent after the initial message.
60    pub messages: Vec<String>,
61    /// First prompt (may carry `@file` content). `None` skips the initial call.
62    pub initial_message: Option<String>,
63    /// Images attached to the initial prompt.
64    pub initial_images: Vec<ImageContent>,
65}
66
67impl PrintModeOptions {
68    /// Build options for a given output mode with no prompts.
69    #[must_use]
70    pub fn new(mode: PrintOutput) -> Self {
71        Self {
72            mode,
73            messages: Vec::new(),
74            initial_message: None,
75            initial_images: Vec::new(),
76        }
77    }
78}
79
80/// Sink abstraction for print-mode output.
81///
82/// [`write_stdout`](Self::write_stdout) targets the protocol stdout sink
83/// (the [`OutputGuard`](crate::core::output_guard) raw queue in production);
84/// [`write_stderr`](Self::write_stderr) targets product stderr. Both provide
85/// FIFO ordering and backpressure. Implementations must be cheap to borrow
86/// (`&self`) since renderers hold the borrow across an entire event drain.
87pub trait PrintSink {
88    /// Append `text` to the protocol stdout sink.
89    fn write_stdout(&self, text: &str) -> impl Future<Output = io::Result<()>> + Send;
90    /// Append `text` to product stderr.
91    fn write_stderr(&self, text: &str) -> impl Future<Output = io::Result<()>> + Send;
92    /// Flush all pending stdout writes.
93    fn flush(&self) -> impl Future<Output = io::Result<()>> + Send;
94}
95
96/// [`OutputGuard`](crate::core::output_guard)-backed sink for the real process.
97///
98/// Stdout writes go through the bounded raw-stdout queue (FIFO + backpressure);
99/// stderr writes go through [`ProductOutput`], which routes to stderr while
100/// stdout is taken over for protocol frames.
101#[derive(Clone, Copy, Debug)]
102pub struct OutputGuardSink;
103
104impl PrintSink for OutputGuardSink {
105    async fn write_stdout(&self, text: &str) -> io::Result<()> {
106        write_raw_stdout(text).await.map_err(io::Error::other)
107    }
108
109    async fn write_stderr(&self, text: &str) -> io::Result<()> {
110        ProductOutput::write(text);
111        Ok(())
112    }
113
114    async fn flush(&self) -> io::Result<()> {
115        flush_raw_stdout().await.map_err(io::Error::other)
116    }
117}
118
119/// In-memory sink for deterministic tests.
120///
121/// Captures stdout and stderr as bytes behind a shared
122/// [`Arc`]`<`[`Mutex`]`>`, so clones observe the same accumulated output.
123#[derive(Clone, Default)]
124pub struct BufferSink {
125    stdout: Arc<Mutex<Vec<u8>>>,
126    stderr: Arc<Mutex<Vec<u8>>>,
127}
128
129impl BufferSink {
130    /// Create an empty buffer sink.
131    #[must_use]
132    pub fn new() -> Self {
133        Self::default()
134    }
135
136    /// Return the captured stdout bytes as a UTF-8 string (lossy).
137    #[must_use]
138    pub fn stdout_string(&self) -> String {
139        String::from_utf8_lossy(&lock_buffer(&self.stdout)).into_owned()
140    }
141
142    /// Return the captured stderr bytes as a UTF-8 string (lossy).
143    #[must_use]
144    pub fn stderr_string(&self) -> String {
145        String::from_utf8_lossy(&lock_buffer(&self.stderr)).into_owned()
146    }
147}
148
149/// Lock a buffer, recovering from poison so a panicking writer never deadlocks
150/// the sink (matches `output_guard`'s `into_inner` recovery).
151fn lock_buffer(buf: &Mutex<Vec<u8>>) -> std::sync::MutexGuard<'_, Vec<u8>> {
152    buf.lock()
153        .unwrap_or_else(std::sync::PoisonError::into_inner)
154}
155
156impl PrintSink for BufferSink {
157    async fn write_stdout(&self, text: &str) -> io::Result<()> {
158        lock_buffer(&self.stdout).extend_from_slice(text.as_bytes());
159        Ok(())
160    }
161
162    async fn write_stderr(&self, text: &str) -> io::Result<()> {
163        lock_buffer(&self.stderr).extend_from_slice(text.as_bytes());
164        Ok(())
165    }
166
167    async fn flush(&self) -> io::Result<()> {
168        Ok(())
169    }
170}
171
172/// For JSON mode the session header is written first. The renderer and prompt
173/// driver then run concurrently. As soon as the prompt driver settles,
174/// `finish_events` must close or unsubscribe the event producer; the renderer
175/// is still awaited so it drains every event already delivered before success
176/// and terminates after an early prompt failure.
177///
178/// Signal handling, extension binding, and disposal remain the caller's
179/// responsibility (the live runtime host owns them); this function is the
180/// rendering + prompt-driving core so it stays unit-testable with injected
181/// streams and sinks.
182///
183/// # Errors
184///
185/// Propagates the first sink write, serialization, or prompt-driver failure.
186pub async fn run_print_mode<S, F, Fut, C, K>(
187    options: &PrintModeOptions,
188    header: Option<&SessionHeader>,
189    events: S,
190    drive_prompts: F,
191    finish_events: C,
192    sink: &K,
193) -> io::Result<i32>
194where
195    S: Stream<Item = AgentSessionEvent> + Send + Unpin,
196    F: FnOnce() -> Fut,
197    Fut: Future<Output = io::Result<()>>,
198    C: FnOnce(),
199    K: PrintSink,
200{
201    if options.mode.is_json() {
202        json::render_json_header(header, sink).await?;
203    }
204
205    let render = async {
206        match options.mode {
207            PrintOutput::Text => text::render_text(events, sink).await,
208            PrintOutput::Json => {
209                json::render_json_events(events, sink).await?;
210                Ok(0)
211            }
212        }
213    };
214    tokio::pin!(render);
215    let prompts = drive_prompts();
216    tokio::pin!(prompts);
217    let mut finish_events = Some(finish_events);
218
219    let (render_result, prompt_result) = tokio::select! {
220        prompt_result = &mut prompts => {
221            if let Some(finish) = finish_events.take() {
222                finish();
223            }
224            (render.await, prompt_result)
225        }
226        render_result = &mut render => {
227            let prompt_result = prompts.await;
228            if let Some(finish) = finish_events.take() {
229                finish();
230            }
231            (render_result, prompt_result)
232        }
233    };
234    prompt_result?;
235    let exit_code = render_result?;
236
237    sink.flush().await?;
238    Ok(exit_code)
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use futures::stream;
245    use pi_agent::AgentMessage;
246    use pi_ai::{AssistantContent, AssistantMessage, Message, StopReason, TextContent};
247    use std::sync::atomic::{AtomicBool, Ordering};
248
249    type TestResult = Result<(), Box<dyn std::error::Error>>;
250
251    fn assistant(text: &str, reason: StopReason) -> AgentMessage {
252        let mut msg = AssistantMessage::new("api", "provider", "model", 2);
253        if !text.is_empty() {
254            msg.content
255                .push(AssistantContent::Text(TextContent::new(text)));
256        }
257        msg.stop_reason = reason;
258        AgentMessage::Llm(Box::new(Message::Assistant(msg)))
259    }
260
261    #[tokio::test]
262    async fn run_print_mode_text_drives_prompts_and_renders() -> TestResult {
263        let final_msg = assistant("answer", StopReason::Stop);
264        let events = vec![AgentSessionEvent::AgentEnd {
265            messages: vec![final_msg],
266            will_retry: false,
267        }];
268        let options = PrintModeOptions::new(PrintOutput::Text);
269        let sink = BufferSink::default();
270
271        let prompt_called = Arc::new(AtomicBool::new(false));
272        let flag = Arc::clone(&prompt_called);
273        let code = run_print_mode(
274            &options,
275            None,
276            stream::iter(events),
277            move || {
278                let flag = Arc::clone(&flag);
279                async move {
280                    flag.store(true, Ordering::SeqCst);
281                    Ok(())
282                }
283            },
284            || {},
285            &sink,
286        )
287        .await?;
288
289        assert_eq!(code, 0);
290        assert_eq!(sink.stdout_string(), "answer\n");
291        assert!(prompt_called.load(Ordering::SeqCst));
292        Ok(())
293    }
294
295    #[tokio::test]
296    async fn run_print_mode_json_writes_header_events_exit_zero() -> TestResult {
297        let header = SessionHeader::new("sid", "2024-01-01T00:00:00.000Z", "/cwd", None);
298        let events = vec![
299            AgentSessionEvent::AgentStart,
300            AgentSessionEvent::AgentSettled,
301        ];
302        let options = PrintModeOptions::new(PrintOutput::Json);
303        let sink = BufferSink::default();
304
305        let code = run_print_mode(
306            &options,
307            Some(&header),
308            stream::iter(events),
309            || async { Ok(()) },
310            || {},
311            &sink,
312        )
313        .await?;
314
315        assert_eq!(code, 0);
316        let stdout = sink.stdout_string();
317        let lines: Vec<&str> = stdout.lines().collect();
318        assert_eq!(lines.len(), 3);
319        assert!(lines[0].contains("\"type\":\"session\""));
320        assert!(lines[1].contains("\"agent_start\""));
321        assert!(lines[2].contains("\"agent_settled\""));
322        Ok(())
323    }
324
325    #[tokio::test]
326    async fn run_print_mode_text_error_exit_one() -> TestResult {
327        let mut msg = AssistantMessage::new("api", "provider", "model", 2);
328        msg.stop_reason = StopReason::Error;
329        msg.error_message = Some("boom".into());
330        let events = vec![AgentSessionEvent::AgentEnd {
331            messages: vec![AgentMessage::Llm(Box::new(Message::Assistant(msg)))],
332            will_retry: false,
333        }];
334        let options = PrintModeOptions::new(PrintOutput::Text);
335        let sink = BufferSink::default();
336
337        let code = run_print_mode(
338            &options,
339            None,
340            stream::iter(events),
341            || async { Ok(()) },
342            || {},
343            &sink,
344        )
345        .await?;
346
347        assert_eq!(code, 1);
348        assert_eq!(sink.stderr_string(), "boom\n");
349        assert!(sink.stdout_string().is_empty());
350        Ok(())
351    }
352
353    #[tokio::test]
354    async fn prompt_failure_before_events_closes_renderer_and_returns_error() -> TestResult {
355        let options = PrintModeOptions::new(PrintOutput::Text);
356        let sink = BufferSink::default();
357        let (event_tx, event_rx) = tokio::sync::mpsc::unbounded_channel();
358        let events = Box::pin(stream::unfold(event_rx, |mut receiver| async move {
359            receiver.recv().await.map(|event| (event, receiver))
360        }));
361
362        let result = tokio::time::timeout(
363            std::time::Duration::from_millis(100),
364            run_print_mode(
365                &options,
366                None,
367                events,
368                || async { Err(io::Error::other("preflight auth failed")) },
369                move || drop(event_tx),
370                &sink,
371            ),
372        )
373        .await;
374
375        let error = match result {
376            Ok(Err(error)) => error,
377            Ok(Ok(exit_code)) => {
378                return Err(format!(
379                    "prompt failure must be preserved, but print mode returned exit code {exit_code}"
380                )
381                .into());
382            }
383            Err(error) => {
384                return Err(
385                    format!("print mode must not hang after prompt setup fails: {error}").into(),
386                );
387            }
388        };
389        assert!(error.to_string().contains("preflight auth failed"));
390        Ok(())
391    }
392
393    #[tokio::test]
394    async fn buffer_sink_appends_in_order() -> TestResult {
395        let sink = BufferSink::default();
396        sink.write_stdout("a").await?;
397        sink.write_stdout("b").await?;
398        sink.write_stderr("e").await?;
399        assert_eq!(sink.stdout_string(), "ab");
400        assert_eq!(sink.stderr_string(), "e");
401        Ok(())
402    }
403}
pi/modes/print/mod.rs

pi/modes/print/
mod.rs
