diff --git a/src/app/runtime/actions/mod.rs b/src/app/runtime/actions/mod.rs index e274180..1ea5ac3 100644 --- a/src/app/runtime/actions/mod.rs +++ b/src/app/runtime/actions/mod.rs @@ -1125,9 +1125,22 @@ fn run_agent_turn( Remember: Cycle 0 MUST be investigation-only (access: read). Cycle 1 MUST be planning-only (access: read). Only subsequent cycles can perform modifications (access: write/full)." )); + let planner_prompt_chars = system_msg.content.as_deref().map_or(0, str::len) + + user_msg.content.as_deref().map_or(0, str::len); let planner_result = tc.client.chat_with_tools_non_streaming(&[system_msg, user_msg], None); let pipeline_result = match planner_result { - Ok((reply, _)) => { + Ok((reply, usage_opt)) => { + let (mut tok_in, mut tok_out) = usage_opt.unwrap_or((0, 0)); + if tok_in == 0 { + tok_in = (planner_prompt_chars / 4).max(1) as u64; + } + if tok_out == 0 { + let response_chars = reply.content.as_deref().map_or(0, str::len); + tok_out = (response_chars / 4).max(1) as u64; + } + if let Ok(mut q) = events_q.lock() { + q.push_back(TurnEvent::Usage { tokens_in: tok_in, tokens_out: tok_out }); + } let reply_text = reply.content.as_deref().unwrap_or("").trim(); let clean_json = if reply_text.starts_with("```") { let mut lines = reply_text.lines(); @@ -1350,7 +1363,18 @@ fn run_agent_turn( } }; - let (tok_in, tok_out) = final_usage.unwrap_or((0, 0)); + let (mut tok_in, mut tok_out) = final_usage.unwrap_or((0, 0)); + if tok_in == 0 { + let total_chars: usize = wire_msgs.iter() + .filter_map(|m| m.content.as_deref()) + .map(str::len) + .sum(); + tok_in = (total_chars / 4).max(1) as u64; + } + if tok_out == 0 { + let response_chars = response.content.as_deref().map_or(0, str::len); + tok_out = (response_chars / 4).max(1) as u64; + } if let Ok(mut q) = events_q.lock() { q.push_back(TurnEvent::Usage { tokens_in: tok_in, tokens_out: tok_out }); } diff --git a/src/app/runtime/stream/mod.rs b/src/app/runtime/stream/mod.rs index 0fdd434..c004b11 100644 --- a/src/app/runtime/stream/mod.rs +++ b/src/app/runtime/stream/mod.rs @@ -107,6 +107,9 @@ impl SseParser { return vec![]; } }; + + let mut events = Vec::new(); + if let Some(usage) = value.get("usage") { if !usage.is_null() { let prompt_tokens = usage.get("prompt_tokens").and_then(serde_json::Value::as_u64).unwrap_or_else(|| { @@ -122,84 +125,73 @@ impl SseParser { tracing::warn!("[stream] total_tokens missing in usage chunk"); prompt_tokens + completion_tokens }); - // Only emit Usage as a standalone event if this chunk - // contains nothing else (no choices, no delta). Some - // non-standard providers may bundle usage WITH content - // in the same chunk; emitting both prevents content loss. - let has_other_content = value.get("choices") - .and_then(|c| c.as_array()) - .is_some_and(|arr| arr.iter().any(|ch| { - ch.get("delta").and_then(|d| d.get("content")).is_some() - || ch.get("delta").and_then(|d| d.get("reasoning_content")).is_some() - || ch.get("delta").and_then(|d| d.get("tool_calls")).is_some() - })); - if !has_other_content { - return vec![StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }]; - } + events.push(StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }); } } - match event_type.as_str() { + + let mut other_events = match event_type.as_str() { "message.stop" => vec![StreamEvent::Done], "message.delta" | "" => { - let Some(delta) = value.get("delta").or_else(|| value.get("choices")) else { return vec![] }; - if let Some(choices) = delta.as_array() { - let Some(choice) = choices.first() else { return vec![] }; - let Some(d) = choice.get("delta") else { return vec![] }; + let mut d_events = Vec::new(); + if let Some(delta) = value.get("delta").or_else(|| value.get("choices")) { + if let Some(choices) = delta.as_array() { + if let Some(choice) = choices.first() { + if let Some(d) = choice.get("delta") { + // Content token + if let Some(content) = d.get("content").and_then(|c| c.as_str()) { + d_events.push(StreamEvent::Token(content.to_string())); + } - // Content token - if let Some(content) = d.get("content").and_then(|c| c.as_str()) { - return vec![StreamEvent::Token(content.to_string())]; - } + // Reasoning token + if let Some(reasoning) = d.get("reasoning_content").and_then(|r| r.as_str()) { + d_events.push(StreamEvent::Reasoning(reasoning.to_string())); + } - // Reasoning token - if let Some(reasoning) = d.get("reasoning_content").and_then(|r| r.as_str()) { - return vec![StreamEvent::Reasoning(reasoning.to_string())]; - } + // Tool calls — iterate ALL entries, not just first() + if let Some(tool_calls) = d.get("tool_calls").and_then(|tc| tc.as_array()) { + for tc in tool_calls { + let index = tc.get("index").and_then(serde_json::Value::as_u64).unwrap_or_else(|| { + tracing::warn!("[stream] tool call delta missing index, defaulting to 0"); + 0 + }) as usize; + let id = tc.get("id").and_then(|i| i.as_str()).map(std::string::ToString::to_string); + let name = tc.get("function") + .and_then(|f| f.get("name")) + .and_then(|n| n.as_str()) + .map(std::string::ToString::to_string); + let args_delta = tc.get("function") + .and_then(|f| f.get("arguments")) + .and_then(|a| a.as_str()) + .unwrap_or("") + .to_string(); + d_events.push(StreamEvent::ToolCallDelta { + index, + id, + name, + arguments_delta: args_delta, + }); + } + } - // Tool calls — iterate ALL entries, not just first() - if let Some(tool_calls) = d.get("tool_calls").and_then(|tc| tc.as_array()) { - let mut events = Vec::with_capacity(tool_calls.len()); - for tc in tool_calls { - let index = tc.get("index").and_then(serde_json::Value::as_u64).unwrap_or_else(|| { - tracing::warn!("[stream] tool call delta missing index, defaulting to 0"); - 0 - }) as usize; - let id = tc.get("id").and_then(|i| i.as_str()).map(std::string::ToString::to_string); - let name = tc.get("function") - .and_then(|f| f.get("name")) - .and_then(|n| n.as_str()) - .map(std::string::ToString::to_string); - let args_delta = tc.get("function") - .and_then(|f| f.get("arguments")) - .and_then(|a| a.as_str()) - .unwrap_or("") - .to_string(); - events.push(StreamEvent::ToolCallDelta { - index, - id, - name, - arguments_delta: args_delta, - }); - } - if !events.is_empty() { - return events; - } - } - - // Finish reason - if let Some(reason) = choice.get("finish_reason").and_then(|r| r.as_str()) { - if reason == "stop" || reason == "tool_calls" { - return vec![StreamEvent::Done]; + // Finish reason + if let Some(reason) = choice.get("finish_reason").and_then(|r| r.as_str()) { + if reason == "stop" || reason == "tool_calls" { + d_events.push(StreamEvent::Done); + } + } + } } + } else if let Some(content) = delta.get("content").and_then(|c| c.as_str()) { + d_events.push(StreamEvent::Token(content.to_string())); } } - if let Some(content) = delta.get("content").and_then(|c| c.as_str()) { - return vec![StreamEvent::Token(content.to_string())]; - } - vec![] + d_events } _ => vec![], - } + }; + + events.append(&mut other_events); + events } /// Clears any partially-buffered SSE frame. Reserved for reconnect/retry flows that @@ -350,6 +342,27 @@ mod tests { } } + #[test] + fn feed_parses_usage_and_content_bundled_chunk() { + let mut p = SseParser::new(); + let events = p.feed( + "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\n", + ); + assert_eq!(events.len(), 2); + match (&events[0], &events[1]) { + ( + StreamEvent::Usage { prompt_tokens, completion_tokens, total_tokens }, + StreamEvent::Token(t), + ) => { + assert_eq!(*prompt_tokens, 10); + assert_eq!(*completion_tokens, 5); + assert_eq!(*total_tokens, 15); + assert_eq!(t, "hello"); + } + other => panic!("expected [Usage, Token], got {other:?}"), + } + } + #[test] fn feed_ignores_empty_data_lines() { let mut p = SseParser::new(); diff --git a/src/app/subagent/engine.rs b/src/app/subagent/engine.rs index 65f0c39..df5ebea 100644 --- a/src/app/subagent/engine.rs +++ b/src/app/subagent/engine.rs @@ -449,12 +449,22 @@ pub fn run_subagent(ctx: &SubagentContext, tx: &mpsc::Sender) -> // drain thread can accumulate it and update the Usage panel. // Without this, the Usage panel always shows zeros because the // subagent never tells the parent about the tokens consumed. - if let Some((tokens_in, tokens_out)) = returned_usage { - let _ = tx.blocking_send(SubagentEvent::Usage { - tokens_in, - tokens_out, - }); + let (mut tok_in, mut tok_out) = returned_usage.unwrap_or((0, 0)); + if tok_in == 0 { + let prompt_chars: usize = messages.iter() + .filter_map(|m| m.content.as_deref()) + .map(str::len) + .sum(); + tok_in = (prompt_chars / 4).max(1) as u64; } + if tok_out == 0 { + let response_chars = response.content.as_deref().map_or(0, str::len); + tok_out = (response_chars / 4).max(1) as u64; + } + let _ = tx.blocking_send(SubagentEvent::Usage { + tokens_in: tok_in, + tokens_out: tok_out, + }); let has_tool_calls = response.tool_calls.is_some() && response.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty()); diff --git a/src/view/chat.rs b/src/view/chat.rs index ef89fe9..e1a41de 100644 --- a/src/view/chat.rs +++ b/src/view/chat.rs @@ -18,7 +18,7 @@ use ratatui::layout::Rect; use ratatui::style::{Color, Style, Modifier}; use ratatui::text::{Line, Span}; -use ratatui::widgets::{Block, Borders, Paragraph, Wrap}; +use ratatui::widgets::{Block, BorderType, Borders, Paragraph, Wrap}; use ratatui::Frame; use super::theme::Theme; use crate::dto::chat::message::Role; @@ -26,7 +26,7 @@ use crate::dto::chat::message::Role; /// Column width reserved for the `{role} {time} ` header prefix; wrapped /// continuation lines and Tool sub-lines indent to this width so content /// stays aligned under the first line's content column. -const PREFIX_WIDTH: usize = 12; +const PREFIX_WIDTH: usize = 15; /// Break a flat run of styled spans into `Line`s at embedded `\n` boundaries. fn split_spans_into_lines(spans: Vec>) -> Vec> { @@ -68,10 +68,10 @@ fn role_accent_color(role: &Role) -> Color { /// raw label). fn format_role_label(role: &Role) -> &'static str { match role { - Role::User => "you", - Role::Assistant => "ai", - Role::System => "sys", - Role::Tool => "tool", + Role::User => "👤 you ", + Role::Assistant => "🤖 ai ", + Role::System => "💻 sys ", + Role::Tool => "🔧 tool", } } @@ -90,8 +90,8 @@ fn format_timestamp(ts: i64) -> String { /// `draw_chat`) and must never be passed as `prev_role` — a Tool message /// never triggers a separator, and it never causes one to be inserted /// before the next real turn either. -fn needs_speaker_separator(prev_role: Option<&Role>, role: &Role) -> bool { - matches!(prev_role, Some(p) if p != role) +fn needs_speaker_separator(_prev_role: Option<&Role>, _role: &Role) -> bool { + false // User requested zsh-style compactness (no empty lines between speakers) } /// Render the scrollable chat transcript panel in tight inline-log style. @@ -108,9 +108,9 @@ pub fn draw_chat(frame: &mut Frame, area: Rect, state: &crate::app::state::rest: let mut prev_role: Option = None; let title = if messages.is_empty() { - String::from(" Chat ") + String::from(" 💬 Chat ") } else { - format!(" Chat [{} msgs]", messages.len()) + format!(" 💬 Chat [{} msgs] ", messages.len()) }; for msg in messages { @@ -163,7 +163,7 @@ pub fn draw_chat(frame: &mut Frame, area: Rect, state: &crate::app::state::rest: let label = format_role_label(&msg.role); let ts_str = format_timestamp(msg.timestamp); let header_prefix = vec![ - Span::styled(format!("{label:<4} "), Style::default().fg(accent).add_modifier(Modifier::BOLD)), + Span::styled(format!("{label} "), Style::default().fg(accent).add_modifier(Modifier::BOLD)), Span::styled(format!("{ts_str:<5} "), Style::default().fg(Theme::TEXT_DIM)), ]; @@ -207,7 +207,7 @@ pub fn draw_chat(frame: &mut Frame, area: Rect, state: &crate::app::state::rest: } display_lines.push(Line::from(vec![ Span::styled( - format!("{:<4} ", format_role_label(&Role::Assistant)), + format!("{} ", format_role_label(&Role::Assistant)), Style::default().fg(Theme::ROLE_ASSISTANT).add_modifier(Modifier::BOLD), ), Span::styled(format!("{spinner} "), Style::default().fg(Theme::TEXT_DIM)), @@ -218,8 +218,9 @@ pub fn draw_chat(frame: &mut Frame, area: Rect, state: &crate::app::state::rest: // ── Scrolling ──────────────────────────────────────────────────────── let block = Block::default() .borders(Borders::ALL) + .border_type(BorderType::Rounded) .border_style(Style::default().fg(Theme::BORDER)) - .title(Span::styled(title, Style::default().fg(Theme::TEXT_MUTED))); + .title(Span::styled(title, Style::default().fg(Theme::TEXT_MUTED).add_modifier(Modifier::BOLD))); let total = display_lines.len(); let max_offset = total.saturating_sub(max_visible); @@ -240,11 +241,12 @@ pub fn draw_chat(frame: &mut Frame, area: Rect, state: &crate::app::state::rest: }; let block = if scroll_pct > 0 { - let scroll_title = format!(" Chat [{} msgs] ── {}% ↑ ", messages.len(), scroll_pct); + let scroll_title = format!(" 💬 Chat [{} msgs] ── {}% ↑ ", messages.len(), scroll_pct); Block::default() .borders(Borders::ALL) + .border_type(BorderType::Rounded) .border_style(Style::default().fg(Theme::BORDER)) - .title(Span::styled(scroll_title, Style::default().fg(Theme::TEXT_MUTED))) + .title(Span::styled(scroll_title, Style::default().fg(Theme::TEXT_MUTED).add_modifier(Modifier::BOLD))) } else { block }; @@ -272,18 +274,15 @@ mod tests { } #[test] - fn separator_when_speaker_changes() { - assert!(needs_speaker_separator(Some(&Role::User), &Role::Assistant)); + fn no_separator_when_speaker_changes_because_zsh_style() { + assert!(!needs_speaker_separator(Some(&Role::User), &Role::Assistant)); } #[test] - fn role_labels_are_lowercase_and_fit_prefix_width() { - assert_eq!(format_role_label(&Role::User), "you"); - assert_eq!(format_role_label(&Role::Assistant), "ai"); - assert_eq!(format_role_label(&Role::System), "sys"); - assert_eq!(format_role_label(&Role::Tool), "tool"); - for role in [Role::User, Role::Assistant, Role::System, Role::Tool] { - assert!(format_role_label(&role).len() <= 4); - } + fn role_labels_include_emojis_and_padding() { + assert_eq!(format_role_label(&Role::User), "👤 you "); + assert_eq!(format_role_label(&Role::Assistant), "🤖 ai "); + assert_eq!(format_role_label(&Role::System), "💻 sys "); + assert_eq!(format_role_label(&Role::Tool), "🔧 tool"); } }