use std::future::Future;
use std::process::ExitStatus;
use anyhow::Result;
use serde::Deserialize;
use tokio::io::AsyncRead;
use tokio::task::JoinHandle;
use crate::bridge::protocol::{self, TurnInput};
pub fn read_turn_or_error() -> TurnInput {
match protocol::read_turn() {
Some(turn) if turn.has_content() => turn,
Some(_) => {
emit_error("turn is empty (no message and no attachments)");
TurnInput::default()
}
None => {
emit_error("no turn object on stdin");
TurnInput::default()
}
}
}
pub fn message_and_session(turn: &TurnInput) -> (String, String) {
(turn.message.clone(), turn.session_id.clone())
}
pub async fn with_session_resume_fallback<T, F, Fut>(
tool: &str,
message: &str,
session_id: &str,
op: F,
) -> Result<T>
where
F: Fn(String, String) -> Fut,
Fut: Future<Output = Result<T>>,
{
if session_id.is_empty() {
return op(message.to_string(), String::new()).await;
}
match op(message.to_string(), session_id.to_string()).await {
Ok(v) => Ok(v),
Err(e) => {
eprintln!("[{tool}] resume failed ({e:#}), retrying as new session");
op(message.to_string(), String::new()).await
}
}
}
pub fn emit_partial_with_session(text: &str, session_id: Option<&str>) -> Result<()> {
let mut obj = serde_json::json!({"type":"partial","text":text});
if let Some(sid) = session_id.filter(|s| !s.is_empty()) {
obj["session_id"] = serde_json::json!(sid);
}
println!("{obj}");
std::io::Write::flush(&mut std::io::stdout()).ok();
Ok(())
}
pub fn emit_result(text: &str, session_id: Option<&str>) -> Result<()> {
let mut obj = serde_json::json!({"type":"result","text":text});
if let Some(sid) = session_id.filter(|s| !s.is_empty()) {
obj["session_id"] = serde_json::json!(sid);
}
println!("{obj}");
std::io::Write::flush(&mut std::io::stdout()).ok();
Ok(())
}
pub fn emit_result_or_session(text: Option<&str>, session_id: Option<&str>) -> Result<()> {
let body = text.unwrap_or("");
let has_sid = session_id.is_some_and(|s| !s.is_empty());
if body.is_empty() && !has_sid {
return Ok(());
}
emit_result(body, session_id)
}
pub fn emit_error(text: &str) {
emit_error_with_session(text, None);
}
pub fn emit_error_with_session(text: &str, session_id: Option<&str>) {
let mut obj = serde_json::json!({"type":"error","message":text});
if let Some(sid) = session_id.filter(|s| !s.is_empty()) {
obj["session_id"] = serde_json::json!(sid);
}
println!("{obj}");
std::io::Write::flush(&mut std::io::stdout()).ok();
}
pub struct SessionEmitter {
inbound: String,
discovered: Option<String>,
buffered_partials: Vec<String>,
}
impl SessionEmitter {
pub fn new(inbound_session_id: &str) -> Self {
Self {
inbound: inbound_session_id.to_string(),
discovered: None,
buffered_partials: Vec::new(),
}
}
pub fn current_id(&self) -> Option<&str> {
if !self.inbound.is_empty() {
Some(self.inbound.as_str())
} else {
self.discovered.as_deref().filter(|s| !s.is_empty())
}
}
pub fn discover(&mut self, id: Option<&str>) -> Result<()> {
let Some(id) = id.filter(|s| !s.is_empty()) else {
return Ok(());
};
if self.inbound.is_empty() && self.discovered.is_none() {
self.discovered = Some(id.to_string());
self.flush_partials()?;
}
Ok(())
}
pub fn emit_partial(&mut self, text: &str) -> Result<()> {
if let Some(sid) = self.current_id() {
emit_partial_with_session(text, Some(sid))
} else {
self.buffered_partials.push(text.to_string());
Ok(())
}
}
pub fn emit_result_opt(&mut self, text: Option<&str>, session_id: Option<&str>) -> Result<()> {
self.discover(session_id)?;
let sid = self
.current_id()
.map(str::to_string)
.or_else(|| session_id.filter(|s| !s.is_empty()).map(str::to_string));
self.flush_partials()?;
emit_result_or_session(text.filter(|t| !t.is_empty()), sid.as_deref())
}
pub fn emit_error(&mut self, text: &str, session_id: Option<&str>) -> Result<()> {
self.discover(session_id)?;
let sid = self
.current_id()
.map(str::to_string)
.or_else(|| session_id.filter(|s| !s.is_empty()).map(str::to_string));
self.flush_partials()?;
emit_error_with_session(text, sid.as_deref());
Ok(())
}
pub fn finish_without_session(&mut self) -> Result<()> {
self.flush_partials()
}
fn flush_partials(&mut self) -> Result<()> {
if self.buffered_partials.is_empty() {
return Ok(());
}
let sid = self.current_id().map(str::to_string);
let pending = std::mem::take(&mut self.buffered_partials);
for text in pending {
emit_partial_with_session(&text, sid.as_deref())?;
}
Ok(())
}
}
pub fn spawn_capped_drain<R>(reader: R) -> JoinHandle<String>
where
R: AsyncRead + Unpin + Send + 'static,
{
tokio::spawn(async move {
use tokio::io::{AsyncReadExt, BufReader};
let mut buf = Vec::new();
BufReader::new(reader)
.take(crate::bridge::MAX_CLI_CAPTURE_BYTES as u64)
.read_to_end(&mut buf)
.await
.ok();
String::from_utf8_lossy(&buf).into_owned()
})
}
pub fn ensure_success(tool: &str, status: ExitStatus, stderr: &str, recovered: bool) -> Result<()> {
if !status.success() && !recovered {
let detail = if stderr.is_empty() {
"(no output)"
} else {
stderr
};
anyhow::bail!(
"{tool} exited with status {:?}\nstderr: {detail}",
status.code()
);
}
Ok(())
}
#[derive(Debug, Deserialize)]
pub struct StreamJsonEvent {
#[serde(rename = "type")]
pub event_type: Option<String>,
pub result: Option<String>,
pub session_id: Option<String>,
pub subtype: Option<String>,
pub is_error: Option<bool>,
pub errors: Option<Vec<String>>,
pub stop_reason: Option<String>,
pub message: Option<StreamMessage>,
}
impl StreamJsonEvent {
pub fn is_result_error(&self) -> bool {
self.is_error == Some(true) || self.subtype.as_deref() == Some("error_during_execution")
}
pub fn result_error_message(&self) -> Option<String> {
if !self.is_result_error() {
return None;
}
if let Some(errors) = &self.errors {
let joined = errors
.iter()
.map(|s| s.trim())
.filter(|s| !s.is_empty())
.collect::<Vec<_>>()
.join("; ");
if !joined.is_empty() {
return Some(joined);
}
}
if let Some(reason) = self
.stop_reason
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty())
{
return Some(reason.to_string());
}
self.result
.as_ref()
.map(|s| s.trim())
.filter(|s| !s.is_empty())
.map(str::to_string)
}
}
#[derive(Debug, Deserialize)]
pub struct StreamMessage {
pub content: Option<Vec<StreamContentBlock>>,
}
impl StreamMessage {
pub fn text(&self) -> String {
self.content
.as_deref()
.unwrap_or_default()
.iter()
.filter(|b| b.block_type.as_deref() == Some("text"))
.filter_map(|b| b.text.as_deref())
.collect::<Vec<_>>()
.join("")
}
}
#[derive(Debug, Deserialize)]
pub struct StreamContentBlock {
#[serde(rename = "type")]
pub block_type: Option<String>,
pub text: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn stream_message_text_concatenates_text_blocks() {
let msg = StreamMessage {
content: Some(vec![
StreamContentBlock {
block_type: Some("text".into()),
text: Some("Hello".into()),
},
StreamContentBlock {
block_type: Some("tool_use".into()),
text: Some("ignored".into()),
},
StreamContentBlock {
block_type: Some("text".into()),
text: Some(" World".into()),
},
]),
};
assert_eq!(msg.text(), "Hello World");
}
#[test]
fn stream_message_text_ignores_non_text_block_types() {
let msg = StreamMessage {
content: Some(vec![
StreamContentBlock {
block_type: Some("image".into()),
text: Some("should be ignored".into()),
},
StreamContentBlock {
block_type: Some("tool_use".into()),
text: Some("also ignored".into()),
},
]),
};
assert_eq!(
msg.text(),
"",
"non-text blocks must not contribute to output"
);
}
#[test]
fn stream_message_text_empty_when_no_content() {
let msg = StreamMessage { content: None };
assert_eq!(msg.text(), "");
}
#[test]
fn stream_message_text_skips_text_block_with_none_text_field() {
let msg = StreamMessage {
content: Some(vec![StreamContentBlock {
block_type: Some("text".into()),
text: None,
}]),
};
assert_eq!(
msg.text(),
"",
"None text field in text block must produce empty output"
);
}
#[test]
fn result_error_message_prefers_errors_array() {
let json = r#"{"type":"result","subtype":"error_during_execution","is_error":true,"errors":["LLM 404"],"stop_reason":"provider_stop:LLM 404","session_id":"sess-1"}"#;
let event: StreamJsonEvent = serde_json::from_str(json).unwrap();
assert!(event.is_result_error());
assert_eq!(event.result_error_message().as_deref(), Some("LLM 404"));
}
#[test]
fn result_error_message_falls_back_to_stop_reason() {
let json = r#"{"type":"result","is_error":true,"stop_reason":"provider_stop:boom","session_id":"sess-1"}"#;
let event: StreamJsonEvent = serde_json::from_str(json).unwrap();
assert_eq!(
event.result_error_message().as_deref(),
Some("provider_stop:boom")
);
}
#[cfg(unix)]
mod unix_tests {
use super::*;
use std::os::unix::process::ExitStatusExt;
fn exit(code: i32) -> std::process::ExitStatus {
std::process::ExitStatus::from_raw(code << 8)
}
#[test]
fn ensure_success_ok_for_status_zero() {
assert!(ensure_success("tool", exit(0), "", false).is_ok());
}
#[test]
fn ensure_success_err_when_nonzero_and_not_recovered() {
let result = ensure_success("tool", exit(1), "stderr output", false);
assert!(
result.is_err(),
"non-zero exit without recovery must be Err"
);
}
#[test]
fn ensure_success_ok_when_nonzero_but_recovered() {
assert!(
ensure_success("tool", exit(1), "stderr", true).is_ok(),
"recovered=true must suppress non-zero exit error"
);
}
#[test]
fn ensure_success_err_message_includes_tool_name() {
let err = ensure_success("my-tool", exit(2), "", false).unwrap_err();
assert!(
err.to_string().contains("my-tool"),
"error must include tool name"
);
}
}
}