oino/modes/
print.rs

1//! Print mode: `-p` text output and `--mode json` event streams.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use futures::StreamExt;
7use oino_ai::{AgentMessage, AssistantContent, StopReason};
8use oino_core::events::{AgentSessionEvent, ClientEvent};
9use oino_core::runtime::AgentSessionRuntime;
10use oino_core::session::{PromptDisposition, PromptOptions};
11use tokio::io::AsyncWriteExt;
12use tokio::sync::watch;
13
14use crate::cli::files::InitialMessage;
15use crate::modes::signals;
16
17/// Output flavour.
18#[derive(Clone, Copy, Debug, PartialEq, Eq)]
19pub enum PrintOutput {
20    /// Final assistant text.
21    Text,
22    /// Header + every session event as JSON lines.
23    Json,
24}
25
26/// What to send.
27#[derive(Clone, Debug, Default, PartialEq)]
28pub struct PrintInput {
29    /// First prompt (stdin + `@file`s + first message).
30    pub initial: InitialMessage,
31    /// Remaining positional messages, one prompt each.
32    pub messages: Vec<String>,
33}
34
35fn stop_reason_name(reason: StopReason) -> String {
36    serde_json::to_value(reason)
37        .ok()
38        .and_then(|v| v.as_str().map(str::to_owned))
39        .unwrap_or_default()
40}
41
42/// Forwards session events (JSON mode) and extension errors (both modes) until the hub
43/// closes.
44async fn forward(
45    mut events: oino_core::hub::EventStream,
46    output: PrintOutput,
47    settled: watch::Sender<u64>,
48) {
49    let mut stdout = tokio::io::stdout();
50    while let Some(event) = events.next().await {
51        if matches!(
52            event,
53            ClientEvent::Session(AgentSessionEvent::AgentSettled { .. })
54        ) {
55            settled.send_modify(|n| *n += 1);
56        }
57        match event {
58            ClientEvent::Session(event) if output == PrintOutput::Json => {
59                let line =
60                    oino_rpc::jsonl::serialize_line(&oino_rpc::wire::session_event_json(&event));
61                if stdout.write_all(line.as_bytes()).await.is_err() {
62                    return;
63                }
64                let _ = stdout.flush().await;
65            }
66            ClientEvent::ExtensionError(error) => {
67                crate::eprintln_line(&format!(
68                    "Extension error ({}): {}",
69                    error.error.extension_path, error.error.error
70                ));
71            }
72            _ => {}
73        }
74    }
75}
76
77/// How long to wait for a started run's `agent_settled` to reach the output after the
78/// session went idle (the runtime forwards session events on its own task).
79const SETTLE_GRACE: Duration = Duration::from_secs(5);
80
81async fn prompts(
82    runtime: &AgentSessionRuntime,
83    input: PrintInput,
84    mut settled: watch::Receiver<u64>,
85) -> Result<(), String> {
86    let mut queue: Vec<(String, Vec<oino_ai::ImageContent>)> = Vec::new();
87    if let Some(text) = input.initial.text {
88        queue.push((text, input.initial.images));
89    }
90    queue.extend(input.messages.into_iter().map(|m| (m, Vec::new())));
91    for (text, images) in queue {
92        let session = runtime.session();
93        let before = *settled.borrow_and_update();
94        let disposition = session
95            .prompt(
96                text,
97                PromptOptions {
98                    images,
99                    ..PromptOptions::default()
100                },
101            )
102            .await
103            .map_err(|e| e.to_string())?;
104        session.wait_for_idle().await;
105        runtime.session().wait_for_idle().await;
106        if disposition == PromptDisposition::Started {
107            let _ = tokio::time::timeout(SETTLE_GRACE, settled.wait_for(|n| *n > before)).await;
108        }
109    }
110    Ok(())
111}
112
113/// The text-mode result: stdout text and exit code.
114fn text_result(messages: &[AgentMessage]) -> (String, Option<String>, i32) {
115    let Some(AgentMessage::Assistant(last)) = messages.last() else {
116        return (String::new(), None, 0);
117    };
118    if matches!(last.stop_reason, StopReason::Error | StopReason::Aborted) {
119        let message = last
120            .error_message
121            .clone()
122            .filter(|m| !m.is_empty())
123            .unwrap_or_else(|| format!("Request {}", stop_reason_name(last.stop_reason)));
124        return (String::new(), Some(message), 1);
125    }
126    let mut out = String::new();
127    for block in &last.content {
128        if let AssistantContent::Text(text) = block {
129            out.push_str(&text.text);
130            out.push('\n');
131        }
132    }
133    if last.stop_reason == StopReason::Length {
134        return (
135            out,
136            Some("Response was truncated: the model reached its output token limit.".to_owned()),
137            1,
138        );
139    }
140    (out, None, 0)
141}
142
143/// Resolves once an extension called `ctx.shutdown()` and the session is idle.
144async fn shutdown_requested(runtime: &AgentSessionRuntime) {
145    let mut requested = runtime.shutdown_requested();
146    if requested.wait_for(|on| *on).await.is_err() {
147        return std::future::pending().await;
148    }
149    runtime.session().wait_for_idle().await;
150}
151
152/// Runs print/JSON mode; returns the exit code.
153pub async fn run(runtime: Arc<AgentSessionRuntime>, output: PrintOutput, input: PrintInput) -> i32 {
154    let mut stdout = tokio::io::stdout();
155    if output == PrintOutput::Json {
156        let header = runtime
157            .session()
158            .with_manager(|m| serde_json::to_value(m.header()).ok());
159        if let Some(header) = header {
160            let line = oino_rpc::jsonl::serialize_line(&header);
161            let _ = stdout.write_all(line.as_bytes()).await;
162            let _ = stdout.flush().await;
163        }
164    }
165    let (settled, settled_rx) = watch::channel(0);
166    let forwarder = tokio::spawn(forward(runtime.subscribe(), output, settled));
167    let outcome = tokio::select! {
168        result = prompts(&runtime, input, settled_rx) => Ok(result),
169        () = shutdown_requested(&runtime) => Ok(Ok(())),
170        signal = signals::next() => Err(signal),
171    };
172    let code = match outcome {
173        Err(signal) => signal.exit_code(),
174        Ok(Err(message)) => {
175            crate::eprintln_line(&message);
176            1
177        }
178        Ok(Ok(())) if output == PrintOutput::Text => {
179            let (text, error, code) = text_result(&runtime.session().messages());
180            let _ = stdout.write_all(text.as_bytes()).await;
181            let _ = stdout.flush().await;
182            if let Some(error) = error {
183                crate::eprintln_line(&error);
184            }
185            code
186        }
187        Ok(Ok(())) => 0,
188    };
189    runtime.dispose().await;
190    let _ = forwarder.await;
191    let _ = stdout.flush().await;
192    code
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use oino_ai::AssistantMessage;
199
200    fn assistant(stop: StopReason, text: &str, error: Option<&str>) -> Vec<AgentMessage> {
201        let mut message = AssistantMessage::empty("faux", "faux", "m", 1);
202        message.stop_reason = stop;
203        message.content = vec![AssistantContent::text(text)];
204        message.error_message = error.map(str::to_owned);
205        vec![AgentMessage::Assistant(message)]
206    }
207
208    #[test]
209    fn exit_rules() {
210        assert_eq!(
211            text_result(&assistant(StopReason::Stop, "hi", None)),
212            ("hi\n".into(), None, 0)
213        );
214        assert_eq!(
215            text_result(&assistant(StopReason::Error, "x", Some("boom"))),
216            (String::new(), Some("boom".into()), 1)
217        );
218        assert_eq!(
219            text_result(&assistant(StopReason::Aborted, "x", None)).1,
220            Some("Request aborted".into())
221        );
222        let (out, error, code) = text_result(&assistant(StopReason::Length, "part", None));
223        assert_eq!((out.as_str(), code), ("part\n", 1));
224        assert_eq!(
225            error.as_deref(),
226            Some("Response was truncated: the model reached its output token limit.")
227        );
228        assert_eq!(text_result(&[]), (String::new(), None, 0));
229    }
230}
oino/modes/print.rs

oino/modes/
print.rs
