use std::io::Write as _; use std::process::exit; use serde_json::{json, Value}; use crate::color; use crate::flags::Flags; use crate::native::stream::chat; const DEFAULT_MODEL: &str = "anthropic/claude-sonnet-4.6"; #[derive(Clone, Copy, PartialEq)] enum Verbosity { Quiet, Normal, Verbose, } pub fn run_chat(flags: &Flags, message: Option) { if !chat::is_chat_enabled() { if flags.json { println!( "{}", json!({"success": false, "error": "AI_GATEWAY_API_KEY not set. Set the AI_GATEWAY_API_KEY environment variable to enable chat."}) ); } else { eprintln!( "{} AI_GATEWAY_API_KEY not set. Set the AI_GATEWAY_API_KEY environment variable to enable chat.", color::error_indicator() ); } exit(1); } let verbosity = if flags.quiet { Verbosity::Quiet } else if flags.verbose { Verbosity::Verbose } else { Verbosity::Normal }; let model = flags .model .clone() .unwrap_or_else(|| DEFAULT_MODEL.to_string()); let rt = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime"); let is_tty = std::io::IsTerminal::is_terminal(&std::io::stdin()); match message { Some(msg) => { rt.block_on(run_single_turn( &flags.session, &model, &msg, verbosity, flags.json, )); } None if !is_tty => { let mut input = String::new(); if let Err(e) = std::io::stdin().read_line(&mut input) { if flags.json { println!( "{}", json!({"success": false, "error": format!("Failed to read stdin: {}", e)}) ); } else { eprintln!("{} Failed to read stdin: {}", color::error_indicator(), e); } exit(1); } let input = input.trim(); if input.is_empty() { if flags.json { println!( "{}", json!({"success": false, "error": "No input provided"}) ); } else { eprintln!("{} No input provided", color::error_indicator()); } exit(1); } rt.block_on(run_single_turn( &flags.session, &model, input, verbosity, flags.json, )); } None => { rt.block_on(run_interactive( &flags.session, &model, verbosity, flags.json, )); } } } async fn run_single_turn( session: &str, model: &str, message: &str, verbosity: Verbosity, json_mode: bool, ) { let mut openai_messages: Vec = vec![json!({"role": "system", "content": chat::get_system_prompt()})]; openai_messages.push(json!({"role": "user", "content": message})); let result = run_chat_turn(session, model, &mut openai_messages, verbosity, json_mode).await; if !result { exit(1); } } async fn run_interactive(session: &str, model: &str, verbosity: Verbosity, json_mode: bool) { let mut openai_messages: Vec = vec![json!({"role": "system", "content": chat::get_system_prompt()})]; let gateway_url = std::env::var("AI_GATEWAY_URL") .unwrap_or_else(|_| chat::DEFAULT_AI_GATEWAY_URL.to_string()) .trim_end_matches('/') .to_string(); let api_key = std::env::var("AI_GATEWAY_API_KEY").unwrap_or_default(); let url = format!("{}/v1/chat/completions", gateway_url); let client = chat::http_client(); loop { if !json_mode { eprint!("{} ", color::cyan(">")); let _ = std::io::stderr().flush(); } let mut input = String::new(); match std::io::stdin().read_line(&mut input) { Ok(0) => break, Err(_) => break, Ok(_) => {} } let input = input.trim(); if input.is_empty() { continue; } if matches!(input, "quit" | "exit" | "q") { break; } openai_messages.push(json!({"role": "user", "content": input})); // Compaction check let total_chars = chat::estimate_chars(&openai_messages); if total_chars > chat::COMPACT_THRESHOLD_CHARS && openai_messages.len() > chat::KEEP_RECENT_MESSAGES + 2 { let split = chat::find_safe_split(&openai_messages, chat::KEEP_RECENT_MESSAGES); let to_summarize = &openai_messages[1..split]; if let Some(summary) = chat::summarize_for_compaction(client, &url, &api_key, model, to_summarize).await { let summary_msg = json!({ "role": "system", "content": format!("[Conversation summary]\n{}", summary) }); let recent = openai_messages[split..].to_vec(); openai_messages = vec![openai_messages[0].clone(), summary_msg]; openai_messages.extend(recent); } } let success = run_chat_turn(session, model, &mut openai_messages, verbosity, json_mode).await; if !success && !json_mode { // Continue the loop on error; don't exit interactive mode } if !json_mode { eprintln!(); } } } /// Runs one chat turn: sends messages to the gateway, streams text/tool calls, /// executes tools in a loop until the model is done. Appends assistant and tool /// messages to `openai_messages`. Returns true on success. async fn run_chat_turn( session: &str, model: &str, openai_messages: &mut Vec, verbosity: Verbosity, json_mode: bool, ) -> bool { let gateway_url = std::env::var("AI_GATEWAY_URL") .unwrap_or_else(|_| chat::DEFAULT_AI_GATEWAY_URL.to_string()) .trim_end_matches('/') .to_string(); let api_key = match std::env::var("AI_GATEWAY_API_KEY") { Ok(k) => k, Err(_) => { if json_mode { println!( "{}", json!({"success": false, "error": "AI_GATEWAY_API_KEY not set"}) ); } else { eprintln!("{} AI_GATEWAY_API_KEY not set", color::error_indicator()); } return false; } }; let tools: Value = serde_json::from_str(chat::CHAT_TOOLS).unwrap(); let url = format!("{}/v1/chat/completions", gateway_url); let client = chat::http_client(); let total_deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(300); let tool_timeout = std::time::Duration::from_secs(60); let mut all_text = String::new(); let mut all_tool_calls: Vec = Vec::new(); let mut had_text = false; for _step in 0..50 { if tokio::time::Instant::now() >= total_deadline { if json_mode { println!( "{}", json!({"success": false, "error": "Chat session timed out (5 minute limit)."}) ); } else { eprintln!( "\n{} Chat session timed out (5 minute limit).", color::error_indicator() ); } return false; } let gateway_body = json!({ "model": model, "messages": openai_messages, "tools": tools, "stream": true, }); let gw_response = match client .post(&url) .header("Authorization", format!("Bearer {}", api_key)) .header("Content-Type", "application/json") .body(gateway_body.to_string()) .send() .await { Ok(r) => r, Err(e) => { if json_mode { println!( "{}", json!({"success": false, "error": format!("Gateway request failed: {}", e)}) ); } else { eprintln!( "\n{} Gateway request failed: {}", color::error_indicator(), e ); } return false; } }; if !gw_response.status().is_success() { let body_text = gw_response.text().await.unwrap_or_default(); if json_mode { println!("{}", json!({"success": false, "error": body_text})); } else { eprintln!("\n{} {}", color::error_indicator(), body_text); } return false; } let (text_chunks, tool_calls) = parse_gateway_stream(gw_response, verbosity, json_mode).await; if !text_chunks.is_empty() { let text = text_chunks.join(""); all_text.push_str(&text); if !json_mode { if !had_text && verbosity != Verbosity::Quiet { // Add blank line before text if we showed tool calls if !all_tool_calls.is_empty() { println!(); } } had_text = true; } let mut content = json!(text); if let Some(last) = openai_messages.last() { if last.get("role").and_then(|r| r.as_str()) == Some("assistant") && last.get("tool_calls").is_some() { content = json!(text); } } openai_messages.push(json!({"role": "assistant", "content": content})); } if tool_calls.is_empty() { break; } let tc_values: Vec = tool_calls .iter() .map(|(id, name, args)| { json!({"id": id, "type": "function", "function": {"name": name, "arguments": args}}) }) .collect(); if text_chunks.is_empty() { openai_messages.push(json!({"role": "assistant", "tool_calls": tc_values})); } else { // If we had both text and tool calls in the same response, merge them if let Some(last) = openai_messages.last_mut() { if last.get("role").and_then(|r| r.as_str()) == Some("assistant") && last.get("tool_calls").is_none() { last["tool_calls"] = json!(tc_values); } else { openai_messages.push(json!({"role": "assistant", "tool_calls": tc_values})); } } } for (tc_id, _tc_name, tc_args) in &tool_calls { let input: Value = serde_json::from_str(tc_args).unwrap_or(json!({})); let command = input.get("command").and_then(|c| c.as_str()).unwrap_or(""); if !json_mode && verbosity != Verbosity::Quiet { eprintln!("{}", color::dim(&format!("> {}", command))); } let result = match tokio::time::timeout(tool_timeout, chat::execute_chat_tool(session, command)) .await { Ok(r) => r, Err(_) => "Tool execution timed out after 60 seconds.".to_string(), }; if !json_mode && verbosity == Verbosity::Verbose { for line in result.lines() { eprintln!(" {}", color::dim(line)); } } all_tool_calls.push(json!({ "command": command, "output": result })); openai_messages.push(json!({ "role": "tool", "tool_call_id": tc_id, "content": result })); } } if json_mode { println!( "{}", json!({ "success": true, "text": all_text, "tool_calls": all_tool_calls }) ); } else if !had_text && !json_mode { // Model returned only tool calls with no final text; print newline for clean output println!(); } true } /// Parses the SSE stream from the AI gateway, printing text deltas to stdout in /// real-time. Returns (collected_text_chunks, tool_calls). async fn parse_gateway_stream( gw_response: reqwest::Response, verbosity: Verbosity, json_mode: bool, ) -> (Vec, Vec<(String, String, String)>) { use futures_util::StreamExt as _; let mut text_chunks: Vec = Vec::new(); let mut tool_call_args: std::collections::HashMap = std::collections::HashMap::new(); let mut byte_stream = gw_response.bytes_stream(); let mut buffer = String::new(); while let Some(chunk_result) = byte_stream.next().await { let chunk = match chunk_result { Ok(c) => c, Err(_) => break, }; buffer.push_str(&String::from_utf8_lossy(&chunk)); while let Some(newline_pos) = buffer.find('\n') { let line = buffer[..newline_pos].trim_end_matches('\r').to_string(); buffer = buffer[newline_pos + 1..].to_string(); if line.is_empty() { continue; } let Some(data) = line.strip_prefix("data: ") else { continue; }; if data == "[DONE]" { let tool_calls = collect_tool_calls(&mut tool_call_args); if !json_mode && !text_chunks.is_empty() { // End the streamed text line let _ = std::io::stdout().flush(); } return (text_chunks, tool_calls); } let Ok(sse_json) = serde_json::from_str::(data) else { continue; }; let delta = sse_json .get("choices") .and_then(|c| c.get(0)) .and_then(|c| c.get("delta")); let Some(delta) = delta else { continue }; if let Some(text) = delta.get("content").and_then(|c| c.as_str()) { if !text.is_empty() { text_chunks.push(text.to_string()); if !json_mode && verbosity != Verbosity::Quiet { print!("{}", text); let _ = std::io::stdout().flush(); } } } if let Some(tcs) = delta.get("tool_calls").and_then(|t| t.as_array()) { for tc in tcs { let idx = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize; if let std::collections::hash_map::Entry::Vacant(e) = tool_call_args.entry(idx) { let id = tc .get("id") .and_then(|i| i.as_str()) .unwrap_or("") .to_string(); let name = tc .get("function") .and_then(|f| f.get("name")) .and_then(|n| n.as_str()) .unwrap_or("") .to_string(); e.insert((id, name, String::new())); } if let Some(arg_delta) = tc .get("function") .and_then(|f| f.get("arguments")) .and_then(|a| a.as_str()) { let entry = tool_call_args.get_mut(&idx).unwrap(); entry.2.push_str(arg_delta); } } } } } if !json_mode && !text_chunks.is_empty() { let _ = std::io::stdout().flush(); } let tool_calls = collect_tool_calls(&mut tool_call_args); (text_chunks, tool_calls) } fn collect_tool_calls( map: &mut std::collections::HashMap, ) -> Vec<(String, String, String)> { let mut indices: Vec = map.keys().copied().collect(); indices.sort(); indices .into_iter() .filter_map(|idx| map.remove(&idx)) .collect() }