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
This commit is contained in:
@@ -386,13 +386,22 @@ impl LlmClient {
|
|||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
byte_buf.extend_from_slice(&chunk_buf[..n]);
|
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) {
|
let valid_len = match std::str::from_utf8(&byte_buf) {
|
||||||
Ok(s) => s.len(),
|
Ok(s) => s.len(),
|
||||||
Err(e) => e.valid_up_to(),
|
Err(e) => {
|
||||||
};
|
let n = e.valid_up_to();
|
||||||
if valid_len == 0 {
|
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;
|
continue;
|
||||||
}
|
}
|
||||||
|
n
|
||||||
|
}
|
||||||
|
};
|
||||||
let text =
|
let text =
|
||||||
String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned();
|
String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned();
|
||||||
byte_buf.drain(..valid_len);
|
byte_buf.drain(..valid_len);
|
||||||
|
|||||||
@@ -89,30 +89,47 @@ impl Tool for Bash {
|
|||||||
let start = std::time::Instant::now();
|
let start = std::time::Instant::now();
|
||||||
let timeout = Duration::from_millis(timeout_ms);
|
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<Vec<u8>> {
|
||||||
|
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<Vec<u8>> {
|
||||||
|
use std::io::Read;
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
child_stderr.read_to_end(&mut buf)?;
|
||||||
|
Ok(buf)
|
||||||
|
});
|
||||||
|
|
||||||
loop {
|
loop {
|
||||||
match child.try_wait() {
|
match child.try_wait() {
|
||||||
Ok(Some(status)) => {
|
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 elapsed = start.elapsed().as_secs_f64();
|
||||||
let output = child
|
let stdout_str = String::from_utf8_lossy(&stdout).to_string();
|
||||||
.wait_with_output()
|
let stderr_str = String::from_utf8_lossy(&stderr).to_string();
|
||||||
.map_err(|e| anyhow::anyhow!("failed to collect output: {e}"))?;
|
let combined = if stderr_str.is_empty() {
|
||||||
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
|
stdout_str
|
||||||
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
|
|
||||||
let combined = if stderr.is_empty() {
|
|
||||||
stdout
|
|
||||||
} else {
|
} else {
|
||||||
format!("{stdout}\n{stderr}")
|
format!("{stdout_str}\n{stderr_str}")
|
||||||
};
|
};
|
||||||
let trimmed = combined.trim().to_string();
|
let trimmed = combined.trim().to_string();
|
||||||
|
|
||||||
if status.success() {
|
if status.success() {
|
||||||
debug!(elapsed_secs = elapsed, "bash command completed successfully");
|
|
||||||
return Ok(if trimmed.is_empty() {
|
return Ok(if trimmed.is_empty() {
|
||||||
format!("Command completed in {elapsed:.2}s (exit code 0)")
|
format!("Command completed in {elapsed:.2}s (exit code 0)")
|
||||||
} else {
|
} else {
|
||||||
format!("{trimmed}\n\nExit code: 0 ({elapsed:.2}s)")
|
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!(
|
return Ok(format!(
|
||||||
"{}\n\nExit code: {} ({:.2}s)",
|
"{}\n\nExit code: {} ({:.2}s)",
|
||||||
trimmed,
|
trimmed,
|
||||||
@@ -124,7 +141,6 @@ impl Tool for Bash {
|
|||||||
if start.elapsed() > timeout {
|
if start.elapsed() > timeout {
|
||||||
let _ = child.kill();
|
let _ = child.kill();
|
||||||
let _ = child.wait();
|
let _ = child.wait();
|
||||||
warn!(timeout_ms, "bash command timed out");
|
|
||||||
anyhow::bail!("command timed out after {timeout_ms}ms");
|
anyhow::bail!("command timed out after {timeout_ms}ms");
|
||||||
}
|
}
|
||||||
std::thread::sleep(Duration::from_millis(10));
|
std::thread::sleep(Duration::from_millis(10));
|
||||||
|
|||||||
@@ -96,19 +96,40 @@ fn draw_usage_widget(frame: &mut Frame, area: Rect, state: &crate::state::AppSta
|
|||||||
let lines: Vec<Line> = if let Some(ref rt) = state.session_runtime {
|
let lines: Vec<Line> = if let Some(ref rt) = state.session_runtime {
|
||||||
let now_ms = chrono::Utc::now().timestamp_millis();
|
let now_ms = chrono::Utc::now().timestamp_millis();
|
||||||
let summary = compute_usage_summary(&rt.usage, rt.session_start, now_ms);
|
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(
|
Line::from(Span::styled(
|
||||||
format!(" {:>6}: {} tok", "total", summary.total_tokens),
|
format!(" {:>6}: {} tok", "total", summary.total_tokens),
|
||||||
Style::default()
|
Style::default()
|
||||||
.fg(Theme::TEXT)
|
.fg(Theme::TEXT)
|
||||||
.add_modifier(Modifier::BOLD),
|
.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),
|
format!(" {:>6}: {} tok", "main", summary.main_tokens),
|
||||||
Style::default().fg(Theme::TEXT_DIM),
|
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(
|
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),
|
Style::default().fg(Theme::TEXT_DIM),
|
||||||
)),
|
)),
|
||||||
Line::from(Span::styled(
|
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),
|
Style::default().fg(Theme::TEXT_DIM),
|
||||||
)),
|
)),
|
||||||
]
|
]);
|
||||||
|
|
||||||
|
items
|
||||||
} else {
|
} else {
|
||||||
vec![Line::from(Span::styled(
|
vec![Line::from(Span::styled(
|
||||||
" No active session.",
|
" No active session.",
|
||||||
|
|||||||
@@ -48,32 +48,7 @@ pub fn draw_status_bar(frame: &mut Frame, area: Rect, state: &crate::state::AppS
|
|||||||
status_badge,
|
status_badge,
|
||||||
];
|
];
|
||||||
|
|
||||||
let max_tokens = crate::state::resolve_context_window(&state.app_config, &state.settings);
|
let right_str = String::new();
|
||||||
|
|
||||||
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 left_line = Line::from(left_spans);
|
let left_line = Line::from(left_spans);
|
||||||
let right_line = Line::from(Span::styled(
|
let right_line = Line::from(Span::styled(
|
||||||
|
|||||||
Reference in New Issue
Block a user