diff --git a/crates/app/cli/src/main.rs b/crates/app/cli/src/main.rs index 0061c48..d4ff1af 100644 --- a/crates/app/cli/src/main.rs +++ b/crates/app/cli/src/main.rs @@ -1,11 +1,126 @@ -use clap::Parser; +use clap::{Parser, ValueEnum}; use color_eyre::eyre::{Result, eyre}; use config_agent::load_settings; use futures_util::TryStreamExt; use hooks::{HookEvent, HookManager, HookResult}; use llm_ollama::{OllamaClient, OllamaOptions, types::ChatMessage}; use permissions::{PermissionDecision, Tool}; +use serde::Serialize; use std::io::{self, Write}; +use std::time::{SystemTime, UNIX_EPOCH}; + +#[derive(Debug, Clone, Copy, ValueEnum)] +enum OutputFormat { + Text, + Json, + StreamJson, +} + +#[derive(Serialize)] +struct SessionOutput { + session_id: String, + messages: Vec, + stats: Stats, + #[serde(skip_serializing_if = "Option::is_none")] + result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tool: Option, +} + +#[derive(Serialize)] +struct Stats { + total_tokens: u64, + #[serde(skip_serializing_if = "Option::is_none")] + prompt_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + completion_tokens: Option, + duration_ms: u64, +} + +#[derive(Serialize)] +struct StreamEvent { + #[serde(rename = "type")] + event_type: String, + #[serde(skip_serializing_if = "Option::is_none")] + session_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + stats: Option, +} + +fn generate_session_id() -> String { + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis(); + format!("session-{}", timestamp) +} + +fn output_tool_result( + format: OutputFormat, + tool: &str, + result: serde_json::Value, + session_id: &str, +) -> Result<()> { + match format { + OutputFormat::Text => { + // For text, just print the result as-is + if let Some(s) = result.as_str() { + println!("{}", s); + } else { + println!("{}", serde_json::to_string_pretty(&result)?); + } + } + OutputFormat::Json => { + let output = SessionOutput { + session_id: session_id.to_string(), + messages: vec![], + stats: Stats { + total_tokens: 0, + prompt_tokens: None, + completion_tokens: None, + duration_ms: 0, + }, + result: Some(result), + tool: Some(tool.to_string()), + }; + println!("{}", serde_json::to_string(&output)?); + } + OutputFormat::StreamJson => { + // For stream-json, emit session_start, result, and session_end + let session_start = StreamEvent { + event_type: "session_start".to_string(), + session_id: Some(session_id.to_string()), + content: None, + stats: None, + }; + println!("{}", serde_json::to_string(&session_start)?); + + let result_event = StreamEvent { + event_type: "tool_result".to_string(), + session_id: None, + content: Some(serde_json::to_string(&result)?), + stats: None, + }; + println!("{}", serde_json::to_string(&result_event)?); + + let session_end = StreamEvent { + event_type: "session_end".to_string(), + session_id: None, + content: None, + stats: Some(Stats { + total_tokens: 0, + prompt_tokens: None, + completion_tokens: None, + duration_ms: 0, + }), + }; + println!("{}", serde_json::to_string(&session_end)?); + } + } + Ok(()) +} #[derive(clap::Subcommand, Debug)] enum Cmd { @@ -32,6 +147,9 @@ struct Args { /// Override the permission mode (plan, acceptEdits, code) #[arg(long)] mode: Option, + /// Output format (text, json, stream-json) + #[arg(long, value_enum, default_value = "text")] + output_format: OutputFormat, #[arg()] prompt: Vec, #[command(subcommand)] @@ -55,6 +173,10 @@ async fn main() -> Result<()> { // Create hook manager let hook_mgr = HookManager::new("."); + // Generate session ID + let session_id = generate_session_id(); + let output_format = args.output_format; + if let Some(cmd) = args.cmd { match cmd { Cmd::Read { path } => { @@ -74,7 +196,7 @@ async fn main() -> Result<()> { } let s = tools_fs::read_file(&path)?; - println!("{}", s); + output_tool_result(output_format, "Read", serde_json::json!(s), &session_id)?; return Ok(()); } PermissionDecision::Ask => { @@ -341,21 +463,118 @@ async fn main() -> Result<()> { let msgs = vec![ChatMessage { role: "user".into(), - content: prompt, + content: prompt.clone(), }]; - let mut stream = client.chat_stream(&msgs, &opts).await?; - while let Some(chunk) = stream.try_next().await? { - if let Some(m) = chunk.message { - if let Some(c) = m.content { - print!("{c}"); - io::stdout().flush()?; + let start_time = SystemTime::now(); + + // Handle different output formats + match output_format { + OutputFormat::Text => { + // Text format: stream to stdout as before + let mut stream = client.chat_stream(&msgs, &opts).await?; + while let Some(chunk) = stream.try_next().await? { + if let Some(m) = chunk.message { + if let Some(c) = m.content { + print!("{c}"); + io::stdout().flush()?; + } + } + if matches!(chunk.done, Some(true)) { + break; + } } + println!(); // Newline after response } - if matches!(chunk.done, Some(true)) { - break; + OutputFormat::Json => { + // JSON format: collect all chunks, then output final JSON + let mut stream = client.chat_stream(&msgs, &opts).await?; + let mut response = String::new(); + + while let Some(chunk) = stream.try_next().await? { + if let Some(m) = chunk.message { + if let Some(c) = m.content { + response.push_str(&c); + } + } + if matches!(chunk.done, Some(true)) { + break; + } + } + + let duration_ms = start_time.elapsed().unwrap().as_millis() as u64; + + // Rough token estimate (tokens ~= chars / 4) + let estimated_tokens = ((prompt.len() + response.len()) / 4) as u64; + + let output = SessionOutput { + session_id, + messages: vec![ + serde_json::json!({"role": "user", "content": prompt}), + serde_json::json!({"role": "assistant", "content": response}), + ], + stats: Stats { + total_tokens: estimated_tokens, + prompt_tokens: Some((prompt.len() / 4) as u64), + completion_tokens: Some((response.len() / 4) as u64), + duration_ms, + }, + result: None, + tool: None, + }; + + println!("{}", serde_json::to_string(&output)?); + } + OutputFormat::StreamJson => { + // Stream-JSON format: emit session_start, chunks, and session_end + let session_start = StreamEvent { + event_type: "session_start".to_string(), + session_id: Some(session_id.clone()), + content: None, + stats: None, + }; + println!("{}", serde_json::to_string(&session_start)?); + + let mut stream = client.chat_stream(&msgs, &opts).await?; + let mut response = String::new(); + + while let Some(chunk) = stream.try_next().await? { + if let Some(m) = chunk.message { + if let Some(c) = m.content { + response.push_str(&c); + let chunk_event = StreamEvent { + event_type: "chunk".to_string(), + session_id: None, + content: Some(c), + stats: None, + }; + println!("{}", serde_json::to_string(&chunk_event)?); + } + } + if matches!(chunk.done, Some(true)) { + break; + } + } + + let duration_ms = start_time.elapsed().unwrap().as_millis() as u64; + + // Rough token estimate + let estimated_tokens = ((prompt.len() + response.len()) / 4) as u64; + + let session_end = StreamEvent { + event_type: "session_end".to_string(), + session_id: None, + content: None, + stats: Some(Stats { + total_tokens: estimated_tokens, + prompt_tokens: Some((prompt.len() / 4) as u64), + completion_tokens: Some((response.len() / 4) as u64), + duration_ms, + }), + }; + println!("{}", serde_json::to_string(&session_end)?); } } - println!(); // Newline after response + Ok(()) } diff --git a/crates/app/cli/tests/headless.rs b/crates/app/cli/tests/headless.rs new file mode 100644 index 0000000..a29c7d9 --- /dev/null +++ b/crates/app/cli/tests/headless.rs @@ -0,0 +1,145 @@ +use assert_cmd::Command; +use serde_json::Value; +use std::fs; +use tempfile::tempdir; + +#[test] +fn print_json_has_session_id_and_stats() { + let mut cmd = Command::cargo_bin("owlen").unwrap(); + cmd.arg("--output-format") + .arg("json") + .arg("Say hello"); + + let output = cmd.assert().success(); + let stdout = String::from_utf8_lossy(&output.get_output().stdout); + + // Parse JSON output + let json: Value = serde_json::from_str(&stdout).expect("Output should be valid JSON"); + + // Verify session_id exists + assert!(json.get("session_id").is_some(), "JSON output should have session_id"); + let session_id = json["session_id"].as_str().unwrap(); + assert!(!session_id.is_empty(), "session_id should not be empty"); + + // Verify stats exist + assert!(json.get("stats").is_some(), "JSON output should have stats"); + let stats = &json["stats"]; + + // Check for token counts + assert!(stats.get("total_tokens").is_some(), "stats should have total_tokens"); + + // Check for messages + assert!(json.get("messages").is_some(), "JSON output should have messages"); +} + +#[test] +fn stream_json_sequence_is_well_formed() { + let mut cmd = Command::cargo_bin("owlen").unwrap(); + cmd.arg("--output-format") + .arg("stream-json") + .arg("Say hello"); + + let output = cmd.assert().success(); + let stdout = String::from_utf8_lossy(&output.get_output().stdout); + + // Stream-JSON is NDJSON - each line should be valid JSON + let lines: Vec<&str> = stdout.lines().filter(|l| !l.is_empty()).collect(); + + assert!(!lines.is_empty(), "Stream-JSON should produce at least one event"); + + // Each line should be valid JSON + for (i, line) in lines.iter().enumerate() { + let json: Value = serde_json::from_str(line) + .expect(&format!("Line {} should be valid JSON: {}", i, line)); + + // Each event should have a type + assert!(json.get("type").is_some(), "Event should have a type field"); + } + + // First event should be session_start + let first: Value = serde_json::from_str(lines[0]).unwrap(); + assert_eq!(first["type"].as_str().unwrap(), "session_start"); + assert!(first.get("session_id").is_some()); + + // Last event should be session_end or complete + let last: Value = serde_json::from_str(lines[lines.len() - 1]).unwrap(); + let last_type = last["type"].as_str().unwrap(); + assert!( + last_type == "session_end" || last_type == "complete", + "Last event should be session_end or complete, got: {}", + last_type + ); +} + +#[test] +fn text_format_is_default() { + let mut cmd = Command::cargo_bin("owlen").unwrap(); + cmd.arg("Say hello"); + + let output = cmd.assert().success(); + let stdout = String::from_utf8_lossy(&output.get_output().stdout); + + // Text format should not be JSON + assert!(serde_json::from_str::(&stdout).is_err(), + "Default output should be text, not JSON"); +} + +#[test] +fn json_format_with_tool_execution() { + let dir = tempdir().unwrap(); + let file = dir.path().join("test.txt"); + fs::write(&file, "hello world").unwrap(); + + let mut cmd = Command::cargo_bin("owlen").unwrap(); + cmd.arg("--mode") + .arg("code") + .arg("--output-format") + .arg("json") + .arg("read") + .arg(file.to_str().unwrap()); + + let output = cmd.assert().success(); + let stdout = String::from_utf8_lossy(&output.get_output().stdout); + + let json: Value = serde_json::from_str(&stdout).expect("Output should be valid JSON"); + + // Should have result + assert!(json.get("result").is_some()); + + // Should have tool info + assert!(json.get("tool").is_some()); + assert_eq!(json["tool"].as_str().unwrap(), "Read"); +} + +#[test] +fn stream_json_includes_chunk_events() { + let mut cmd = Command::cargo_bin("owlen").unwrap(); + cmd.arg("--output-format") + .arg("stream-json") + .arg("Say hello"); + + let output = cmd.assert().success(); + let stdout = String::from_utf8_lossy(&output.get_output().stdout); + + let lines: Vec<&str> = stdout.lines().filter(|l| !l.is_empty()).collect(); + + // Should have chunk events between session_start and session_end + let chunk_events: Vec<&str> = lines.iter() + .filter(|line| { + if let Ok(json) = serde_json::from_str::(line) { + json["type"].as_str() == Some("chunk") + } else { + false + } + }) + .copied() + .collect(); + + assert!(!chunk_events.is_empty(), "Should have at least one chunk event"); + + // Each chunk should have content + for chunk_line in chunk_events { + let chunk: Value = serde_json::from_str(chunk_line).unwrap(); + assert!(chunk.get("content").is_some(), "Chunk should have content"); + } +}