From 7ea226509c8e11ee1bab21d216e0413715988fb9 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Mon, 20 Jul 2026 16:01:18 +0700 Subject: [PATCH] feat(llm): improve UTF-8 handling in response processing to prevent infinite loops feat(shell): enhance output capturing by using threads for stdout and stderr feat(tui): update usage widget to display token counts and provider/model information refactor(tui): simplify status bar rendering by removing unnecessary token calculations --- apps/infrastructure/src/llm/provider.rs | 17 ++++++++--- apps/infrastructure/src/tools/shell.rs | 38 ++++++++++++++++++------- apps/interfaces/tui/src/view/sidebar.rs | 31 +++++++++++++++++--- apps/interfaces/tui/src/view/status.rs | 27 +----------------- 4 files changed, 68 insertions(+), 45 deletions(-) diff --git a/apps/infrastructure/src/llm/provider.rs b/apps/infrastructure/src/llm/provider.rs index d9686b0..4a9b265 100644 --- a/apps/infrastructure/src/llm/provider.rs +++ b/apps/infrastructure/src/llm/provider.rs @@ -386,13 +386,22 @@ impl LlmClient { break; } byte_buf.extend_from_slice(&chunk_buf[..n]); + + // Drain any bytes that are not valid UTF-8 to prevent + // infinite loop when a non-UTF-8 sequence is received. let valid_len = match std::str::from_utf8(&byte_buf) { Ok(s) => s.len(), - Err(e) => e.valid_up_to(), + Err(e) => { + let n = e.valid_up_to(); + if n == 0 { + // No valid UTF-8 prefix; skip the first byte (likely + // a partial multi-byte sequence or stray byte). + byte_buf.drain(..1); + continue; + } + n + } }; - if valid_len == 0 { - continue; - } let text = String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned(); byte_buf.drain(..valid_len); diff --git a/apps/infrastructure/src/tools/shell.rs b/apps/infrastructure/src/tools/shell.rs index 90efb8f..d11fd5b 100644 --- a/apps/infrastructure/src/tools/shell.rs +++ b/apps/infrastructure/src/tools/shell.rs @@ -89,30 +89,47 @@ impl Tool for Bash { let start = std::time::Instant::now(); let timeout = Duration::from_millis(timeout_ms); + let mut child_stdout = child.stdout.take() + .ok_or_else(|| anyhow::anyhow!("failed to capture stdout"))?; + let mut child_stderr = child.stderr.take() + .ok_or_else(|| anyhow::anyhow!("failed to capture stderr"))?; + + let stdout_handle = std::thread::spawn(move || -> std::io::Result> { + use std::io::Read; + let mut buf = Vec::new(); + child_stdout.read_to_end(&mut buf)?; + Ok(buf) + }); + let stderr_handle = std::thread::spawn(move || -> std::io::Result> { + use std::io::Read; + let mut buf = Vec::new(); + child_stderr.read_to_end(&mut buf)?; + Ok(buf) + }); + loop { match child.try_wait() { Ok(Some(status)) => { + let stdout = stdout_handle.join().unwrap_or(Ok(Vec::new()))?; + let stderr = stderr_handle.join().unwrap_or(Ok(Vec::new()))?; + let elapsed = start.elapsed().as_secs_f64(); - let output = child - .wait_with_output() - .map_err(|e| anyhow::anyhow!("failed to collect output: {e}"))?; - let stdout = String::from_utf8_lossy(&output.stdout).to_string(); - let stderr = String::from_utf8_lossy(&output.stderr).to_string(); - let combined = if stderr.is_empty() { - stdout + let stdout_str = String::from_utf8_lossy(&stdout).to_string(); + let stderr_str = String::from_utf8_lossy(&stderr).to_string(); + let combined = if stderr_str.is_empty() { + stdout_str } else { - format!("{stdout}\n{stderr}") + format!("{stdout_str}\n{stderr_str}") }; let trimmed = combined.trim().to_string(); + if status.success() { - debug!(elapsed_secs = elapsed, "bash command completed successfully"); return Ok(if trimmed.is_empty() { format!("Command completed in {elapsed:.2}s (exit code 0)") } else { format!("{trimmed}\n\nExit code: 0 ({elapsed:.2}s)") }); } - warn!(exit_code = status.code().unwrap_or(-1), elapsed_secs = elapsed, "bash command failed"); return Ok(format!( "{}\n\nExit code: {} ({:.2}s)", trimmed, @@ -124,7 +141,6 @@ impl Tool for Bash { if start.elapsed() > timeout { let _ = child.kill(); let _ = child.wait(); - warn!(timeout_ms, "bash command timed out"); anyhow::bail!("command timed out after {timeout_ms}ms"); } std::thread::sleep(Duration::from_millis(10)); diff --git a/apps/interfaces/tui/src/view/sidebar.rs b/apps/interfaces/tui/src/view/sidebar.rs index 090475a..d861b1f 100644 --- a/apps/interfaces/tui/src/view/sidebar.rs +++ b/apps/interfaces/tui/src/view/sidebar.rs @@ -96,19 +96,40 @@ fn draw_usage_widget(frame: &mut Frame, area: Rect, state: &crate::state::AppSta let lines: Vec = if let Some(ref rt) = state.session_runtime { let now_ms = chrono::Utc::now().timestamp_millis(); let summary = compute_usage_summary(&rt.usage, rt.session_start, now_ms); - vec![ + let max_tokens = crate::state::resolve_context_window(&state.app_config, &state.settings); + let current_tokens = state.cached_token_count; + + let mut items = vec![ Line::from(Span::styled( format!(" {:>6}: {} tok", "total", summary.total_tokens), Style::default() .fg(Theme::TEXT) .add_modifier(Modifier::BOLD), )), - Line::from(Span::styled( + ]; + + if summary.self_learning_tokens > 0 { + items.push(Line::from(Span::styled( format!(" {:>6}: {} tok", "main", summary.main_tokens), Style::default().fg(Theme::TEXT_DIM), + ))); + items.push(Line::from(Span::styled( + format!(" {:>6}: {} tok", "learn", summary.self_learning_tokens), + Style::default().fg(Theme::TEXT_DIM), + ))); + } + + items.extend(vec![ + Line::from(Span::styled( + format!(" {:>6}: {}/{}", "ctx", current_tokens, max_tokens), + Style::default().fg(Theme::TEXT_DIM), )), Line::from(Span::styled( - format!(" {:>6}: {} tok", "learn", summary.self_learning_tokens), + format!(" {:>6}: {}", "prov", state.settings.provider), + Style::default().fg(Theme::TEXT_DIM), + )), + Line::from(Span::styled( + format!(" {:>6}: {}", "model", state.settings.model), Style::default().fg(Theme::TEXT_DIM), )), Line::from(Span::styled( @@ -122,7 +143,9 @@ fn draw_usage_widget(frame: &mut Frame, area: Rect, state: &crate::state::AppSta ), Style::default().fg(Theme::TEXT_DIM), )), - ] + ]); + + items } else { vec![Line::from(Span::styled( " No active session.", diff --git a/apps/interfaces/tui/src/view/status.rs b/apps/interfaces/tui/src/view/status.rs index 59e176b..c73989f 100644 --- a/apps/interfaces/tui/src/view/status.rs +++ b/apps/interfaces/tui/src/view/status.rs @@ -48,32 +48,7 @@ pub fn draw_status_bar(frame: &mut Frame, area: Rect, state: &crate::state::AppS status_badge, ]; - let max_tokens = crate::state::resolve_context_window(&state.app_config, &state.settings); - - let right_str = if let Some(ref rt) = state.session_runtime { - // Use cached token count — recomputed lazily only when new messages - // arrive (token_count_dirty flag), not on every render frame. - // This eliminates the expensive tiktoken_rs call from the hot path. - let current_tokens = state.cached_token_count; - - let mut parts = Vec::new(); - if rt.usage.last_tokens_in > 0 || rt.usage.last_tokens_out > 0 { - parts.push(format!( - "↑{} ↓{}", - rt.usage.last_tokens_in, rt.usage.last_tokens_out - )); - } - parts.push(format!("{current_tokens}/{max_tokens}")); - parts.push(state.settings.provider.clone()); - parts.push(state.settings.model.clone()); - - format!(" {} ", parts.join(" · ")) - } else { - format!( - " 0/{max_tokens} · {} · {} ", - state.settings.provider, state.settings.model - ) - }; + let right_str = String::new(); let left_line = Line::from(left_spans); let right_line = Line::from(Span::styled(