diff --git a/crates/hi-agent/src/agent/compaction_turn.rs b/crates/hi-agent/src/agent/compaction_turn.rs index 681a647b..5525fea4 100644 --- a/crates/hi-agent/src/agent/compaction_turn.rs +++ b/crates/hi-agent/src/agent/compaction_turn.rs @@ -436,7 +436,7 @@ impl crate::Agent { return Err(err); } }; - self.add_usage(completion.usage); + self.add_side_usage(completion.usage); let _ = self.persist(); self.emit_usage(ui); diff --git a/crates/hi-agent/src/agent/curate_turn.rs b/crates/hi-agent/src/agent/curate_turn.rs index 3c3cde50..ed19f5bf 100644 --- a/crates/hi-agent/src/agent/curate_turn.rs +++ b/crates/hi-agent/src/agent/curate_turn.rs @@ -96,7 +96,7 @@ impl crate::Agent { return; } }; - self.add_usage(completion.usage); + self.add_side_usage(completion.usage); let _ = self.persist(); if out.trim().is_empty() { for c in &completion.content { diff --git a/crates/hi-agent/src/agent/explore_turn.rs b/crates/hi-agent/src/agent/explore_turn.rs index a06a9251..9aef9f47 100644 --- a/crates/hi-agent/src/agent/explore_turn.rs +++ b/crates/hi-agent/src/agent/explore_turn.rs @@ -93,7 +93,7 @@ impl crate::Agent { } }; // Fold the child's token usage into the parent's session totals. - self.add_usage(*child.totals()); + self.add_side_usage(*child.totals()); ui.subagent_note(&format!("↳ explore subagent {n} done")); result } diff --git a/crates/hi-agent/src/agent/lifecycle.rs b/crates/hi-agent/src/agent/lifecycle.rs index 771fcfeb..8519201d 100644 --- a/crates/hi-agent/src/agent/lifecycle.rs +++ b/crates/hi-agent/src/agent/lifecycle.rs @@ -455,15 +455,27 @@ impl crate::Agent { out } + /// Usage from a *main-conversation* model call: counts toward all totals + /// and refreshes the context gauge with the request's occupancy. pub(crate) fn add_usage(&mut self, usage: Usage) { - self.totals.add(usage); - self.last_turn_usage.add(usage); + self.add_side_usage(usage); let effective_input = usage.effective_input_tokens(); if effective_input > 0 { self.context_used = effective_input; } } + /// Usage from a *side* model call (finalize, skeptic, curate, memory, + /// goal planning, explore children, compaction summarize): counts toward + /// totals and the turn's spend, but must not touch `context_used` — that + /// gauge tracks the main conversation's occupancy and drives + /// auto-compaction, and a ~3K-token side request at the end of a 150K + /// session would reset it to 2%, silently disabling the next compaction. + pub(crate) fn add_side_usage(&mut self, usage: Usage) { + self.totals.add(usage); + self.last_turn_usage.add(usage); + } + pub(crate) fn reset_last_turn_usage(&mut self, user_prompt_tokens: u64) { self.last_turn_usage = Usage::default(); self.last_user_prompt_tokens = user_prompt_tokens; diff --git a/crates/hi-agent/src/agent/memory_turn.rs b/crates/hi-agent/src/agent/memory_turn.rs index 61cc2291..1ce27064 100644 --- a/crates/hi-agent/src/agent/memory_turn.rs +++ b/crates/hi-agent/src/agent/memory_turn.rs @@ -115,7 +115,7 @@ impl crate::Agent { return; } }; - self.add_usage(completion.usage); + self.add_side_usage(completion.usage); let _ = self.persist(); // Fall back to the final content if the provider didn't stream text. diff --git a/crates/hi-agent/src/agent/plan_goal.rs b/crates/hi-agent/src/agent/plan_goal.rs index e5dec03b..33585381 100644 --- a/crates/hi-agent/src/agent/plan_goal.rs +++ b/crates/hi-agent/src/agent/plan_goal.rs @@ -65,7 +65,7 @@ impl crate::Agent { return Err(err); } }; - self.add_usage(completion.usage); + self.add_side_usage(completion.usage); // Fall back to the completion content if the provider returned text only in // the final object rather than via stream deltas. if text.trim().is_empty() { diff --git a/crates/hi-agent/src/agent/skeptic.rs b/crates/hi-agent/src/agent/skeptic.rs index 51be5033..86593921 100644 --- a/crates/hi-agent/src/agent/skeptic.rs +++ b/crates/hi-agent/src/agent/skeptic.rs @@ -157,7 +157,7 @@ impl crate::Agent { return SkepticVerdict::Approve; // fail-open on a provider error } }; - self.add_usage(completion.usage); + self.add_side_usage(completion.usage); if text.trim().is_empty() { text = content_text(&completion.content); } diff --git a/crates/hi-agent/src/agent/turn.rs b/crates/hi-agent/src/agent/turn.rs index 1372c19b..c26ff75b 100644 --- a/crates/hi-agent/src/agent/turn.rs +++ b/crates/hi-agent/src/agent/turn.rs @@ -2384,11 +2384,14 @@ If the task is already complete, stop and give your final recap." let mut completed = vec![false; calls.len()]; let mut completion_order: Vec = Vec::with_capacity(calls.len()); let mut scheduler_forced_skip = false; - // Pre-pass: handle `record_decision` calls serially. They mutate - // agent state (`self.decisions`) and aren't real tool dispatches, - // so they can't run in the parallel `execute` stream (no `&mut - // self` there). They're instantaneous and have no deps that - // matter, so handling them up front is safe. + // Pre-pass: resolve calls blocked by read-only intent up front. + // They produce instant synthetic error results and mutate + // nothing, so completing them out of dep order is safe. + // (`explore`/`delegate`/`record_decision` used to run here too, + // but they *do* have deps that matter — running a subagent + // before an earlier `write` in the same batch handed it a stale + // tree — so they now dispatch inside the dep-aware scheduler + // loop below.) for (i, (id, name, arguments)) in calls.iter().enumerate() { if read_only_blocks_tool(read_only_intent, name) { ui.tool_call(name, arguments); @@ -2405,68 +2408,7 @@ If the task is already complete, stop and give your final recap." results[i] = Some((id.clone(), content)); completed[i] = true; completion_order.push(i); - continue; - } - if name == "explore" { - // Read-only subagent: runs a bounded child turn (own context, - // read-only tools) sharing this agent's provider, then returns a - // concise answer. Handled here (not in the parallel `execute` - // stream) because it needs `&mut self` and is itself an agent turn. - ui.tool_call(name, arguments); - let content = self.handle_explore(arguments, &mut *ui).await; - emit_tool_output( - &mut *ui, - name, - &hi_tools::ToolOutput { - content: content.clone(), - display: None, - plan: None, - }, - ); - results[i] = Some((id.clone(), content)); - completed[i] = true; - completion_order.push(i); - continue; } - if name == "delegate" { - // Write-capable subagent: runs a child turn that edits + verifies - // in the same tree, checkpoint-guarded so a failed delegation rolls - // back. Like `explore`, handled here because it needs `&mut self` - // and is itself an agent turn. - // - // Capture the turn baseline + checkpoint BEFORE it mutates the - // tree — otherwise the later lazy snapshot (verify gate) would - // record delegate's own output as the baseline, making the - // parent's verify + changed-files see "no changes", and leaving - // no pre-delegate checkpoint for `/undo` to isolate this turn. - self.ensure_turn_snapshot(&mut turn_snapshot).await; - self.ensure_turn_checkpoint(&mut turn_checkpoint_created, ui) - .await; - ui.tool_call(name, arguments); - let content = self.handle_delegate(arguments, &mut *ui).await; - emit_tool_output( - &mut *ui, - name, - &hi_tools::ToolOutput { - content: content.clone(), - display: None, - plan: None, - }, - ); - results[i] = Some((id.clone(), content)); - completed[i] = true; - completion_order.push(i); - continue; - } - if name != "record_decision" { - continue; - } - ui.tool_call(name, arguments); - let content = self.handle_record_decision(arguments); - ui.tool_result(name, &content); - results[i] = Some((id.clone(), content)); - completed[i] = true; - completion_order.push(i); } let mut done = completion_order.len(); // Proactive per-edit checks: kicked off in the background as @@ -2618,6 +2560,62 @@ If the task is already complete, stop and give your final recap." sched_max_concurrent = sched_max_concurrent.max(1); continue; } + // Self-dispatched calls: `explore`/`delegate` run a child + // agent turn and `record_decision` mutates agent state, so + // all three need `&mut self` and can't join the parallel + // `execute` stream. Run one alone when it's ready — the dep + // graph then guarantees earlier mutations in the batch have + // landed before a subagent sees the tree. + let self_idx = ready.iter().copied().find(|&i| { + matches!( + calls[i].1.as_str(), + "explore" | "delegate" | "record_decision" + ) + }); + if let Some(i) = self_idx { + let (id, name, arguments) = &calls[i]; + if name == "delegate" { + // Write-capable subagent: capture the turn baseline + + // checkpoint BEFORE it mutates the tree — otherwise the + // later lazy snapshot (verify gate) would record + // delegate's own output as the baseline, making the + // parent's verify + changed-files see "no changes", and + // leaving no pre-delegate checkpoint for `/undo` to + // isolate this turn. + self.ensure_turn_snapshot(&mut turn_snapshot).await; + self.ensure_turn_checkpoint(&mut turn_checkpoint_created, ui) + .await; + } + ui.tool_call(name, arguments); + let content = match name.as_str() { + "explore" => self.handle_explore(arguments, &mut *ui).await, + "delegate" => self.handle_delegate(arguments, &mut *ui).await, + _ => self.handle_record_decision(arguments), + }; + emit_tool_output( + &mut *ui, + name, + &hi_tools::ToolOutput { + content: content.clone(), + display: None, + plan: None, + }, + ); + results[i] = Some((id.clone(), content)); + if name == "delegate" { + // The child turn edits the same tree — a dependent + // read after it must re-walk. + self.invalidate_snapshot(); + } + completed[i] = true; + completion_order.push(i); + done += 1; + // Runs alone, like bash. + sched_tool_calls += 1; + sched_serial_runs += 1; + sched_max_concurrent = sched_max_concurrent.max(1); + continue; + } // Run all ready non-bash calls concurrently. Record the // completion order as the ready order (within a concurrent // batch, relative order doesn't matter — none depend on @@ -2824,6 +2822,12 @@ If the task is already complete, stop and give your final recap." done += 1; } } + // Consume an interrupt that landed while (or just after) the + // batch's last call finished — the loop above only polls the + // flag between rounds, so a leftover flag would spuriously + // cancel the next round's (or even the next turn's) batch. + self.interrupt + .store(false, std::sync::atomic::Ordering::Relaxed); debug_assert_eq!( done, calls.len(), @@ -3221,7 +3225,9 @@ If the task is already complete, stop and give your final recap." } }; - self.add_usage(completion.usage); + // Side call: spend counts, but its small request must not clobber the + // main conversation's context gauge (see add_side_usage). + self.add_side_usage(completion.usage); self.emit_usage(ui); // Fall back to the final content if the provider didn't stream text. diff --git a/crates/hi-agent/src/goal.rs b/crates/hi-agent/src/goal.rs index 01678bc9..2deff56c 100644 --- a/crates/hi-agent/src/goal.rs +++ b/crates/hi-agent/src/goal.rs @@ -282,7 +282,11 @@ impl Goal { .any(|s| s.status == GoalStatus::Active) { self.status = GoalStatus::Failed; - } else if self.status == GoalStatus::Done { + } else { + // Per the contract above: not all done, not failed-with-no-active + // → Active. This must also revive a previously `Failed` goal whose + // plan was revised to activate new work — leaving it `Failed` + // would permanently disable auto-drive despite live sub-goals. self.status = GoalStatus::Active; } } diff --git a/crates/hi-agent/src/heuristics.rs b/crates/hi-agent/src/heuristics.rs index 5742c221..97a846b0 100644 --- a/crates/hi-agent/src/heuristics.rs +++ b/crates/hi-agent/src/heuristics.rs @@ -271,15 +271,18 @@ fn strip_chatml_tokens(text: &str) -> String { let mut last = 0; let mut i = 0; while i < bytes.len() { + // Search the suffix from after the `<|` prefix so the prefix's own `|` + // can't match `|>` (the text `<|>` produced a reversed slice range and + // panicked). if bytes[i] == b'<' && i + 1 < bytes.len() && bytes[i + 1] == b'|' - && let Some(end) = text[i..].find("|>") + && let Some(end) = text[i + 2..].find("|>") { - let inner = &text[i + 2..i + end]; + let inner = &text[i + 2..i + 2 + end]; if !inner.is_empty() && inner.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { out.push_str(&text[last..i]); - last = i + end + 2; + last = i + 2 + end + 2; i = last; continue; } diff --git a/crates/hi-agent/src/memory.rs b/crates/hi-agent/src/memory.rs index b13aedea..19c0c0ba 100644 --- a/crates/hi-agent/src/memory.rs +++ b/crates/hi-agent/src/memory.rs @@ -339,6 +339,11 @@ fn looks_like_path_or_command(bullet: &str) -> Option { for word in body.split_whitespace() { let word = word.trim_matches(|c: char| !c.is_alphanumeric() && c != '/' && c != '.' && c != '-'); + // A trailing '.' is sentence punctuation, not part of a path: without + // this, "src/main.rs." never exists and any prose word ending a + // sentence ("messages.") reads as having an (empty) extension — either + // way the bullet gets dropped as an implausible path. + let word = word.trim_end_matches('.'); if word.contains('/') || has_extension(word) { return Some(word.to_string()); } @@ -347,9 +352,16 @@ fn looks_like_path_or_command(bullet: &str) -> Option { } fn has_extension(s: &str) -> bool { - Path::new(s) - .extension() - .is_some_and(|e| e.to_string_lossy().len() <= 6) + Path::new(s).extension().is_some_and(|e| { + let e = e.to_string_lossy(); + // Real file extensions (rs, py, json, mp4, …) are short alphanumerics + // with at least one letter. This keeps version numbers ("v1.2") and + // abbreviations from classifying prose as a path to be verified. + !e.is_empty() + && e.len() <= 6 + && e.chars().all(|c| c.is_ascii_alphanumeric()) + && e.chars().any(|c| c.is_ascii_alphabetic()) + }) } /// Verify distilled bullets that look like paths or commands against the current @@ -582,15 +594,35 @@ impl Drop for LockGuard { } } +/// A lock older than this is presumed abandoned (the holder crashed or was +/// killed before its `Drop` ran). Memory writes are sub-second, so minutes of +/// age means no live holder — without this, one SIGKILL would disable memory +/// persistence for every future session until the file is deleted by hand. +const LOCK_STALE_AFTER: std::time::Duration = std::time::Duration::from_secs(10 * 60); + fn take_lock(parent: &Path) -> Result { let lock = parent.join(".memory.lock"); - match OpenOptions::new().write(true).create_new(true).open(&lock) { - Ok(_) => Ok(LockGuard(lock)), - Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => { - Err("another session is updating memory".to_string()) + for _ in 0..2 { + match OpenOptions::new().write(true).create_new(true).open(&lock) { + Ok(_) => return Ok(LockGuard(lock)), + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => { + let stale = fs::metadata(&lock) + .and_then(|m| m.modified()) + .ok() + .and_then(|t| t.elapsed().ok()) + .is_some_and(|age| age > LOCK_STALE_AFTER); + if !stale { + return Err("another session is updating memory".to_string()); + } + // Break the stale lock and retry the exclusive create once (a + // concurrent session may break it first — losing that race + // lands in AlreadyExists again and errors out normally). + let _ = fs::remove_file(&lock); + } + Err(e) => return Err(format!("couldn't take memory lock: {e}")), } - Err(e) => Err(format!("couldn't take memory lock: {e}")), } + Err("another session is updating memory".to_string()) } // Forward the message types so callers don't need a separate import. diff --git a/crates/hi-agent/src/transcript.rs b/crates/hi-agent/src/transcript.rs index f305d8c2..edfde252 100644 --- a/crates/hi-agent/src/transcript.rs +++ b/crates/hi-agent/src/transcript.rs @@ -594,43 +594,32 @@ impl Transcript { /// only the latest verification output stays in context instead of /// accumulating one nudge per failed round. /// - /// If a prior nudge of `kind` exists: pops from the tail — trailing `Tool` - /// and `Assistant` messages that belong to the prior cycle, then the prior - /// nudge itself — and pushes the new one. If no prior nudge of `kind` is - /// found, **nothing is popped** and the new nudge is simply appended (this - /// is the first round: there's no prior cycle to replace, so the model's - /// just-completed assistant turn and its tool results must stay). + /// A prior nudge is replaceable only when it *immediately precedes* the + /// trailing run of `Tool`/`Assistant` messages — i.e. the tail is exactly + /// the prior nudge's response cycle. Then that cycle and the nudge are + /// popped and the fresh nudge pushed. Otherwise (first round of this turn, + /// or any other message in between) the nudge is simply appended: deciding + /// off a marker anywhere in history — as this used to — meant a stale + /// nudge from an *earlier* turn erased the current turn's work and left + /// two consecutive user messages. pub(crate) fn replace_last_nudge(&mut self, kind: NudgeKind, text: impl Into) { let marker = marker_for(kind); - // First, find whether a prior nudge of this kind exists. If not, just - // append — don't strip the model's just-finished turn. - let has_prior = self.messages.iter().any(|m| { + let body = text.into(); + let tagged = format!("{}\n{}", marker, body); + let msgs = self.make_mut(); + let mut idx = msgs.len(); + while idx > 0 && matches!(msgs[idx - 1].role, Role::Tool | Role::Assistant) { + idx -= 1; + } + let prior_nudge = idx.checked_sub(1).filter(|&i| { + let m = &msgs[i]; m.role == Role::User && m.content .iter() .any(|c| matches!(c, Content::Text(t) if t.starts_with(marker))) }); - let body = text.into(); - let tagged = format!("{}\n{}", marker, body); - let msgs = self.make_mut(); - if has_prior { - while let Some(last) = msgs.last() { - if last.role == Role::User - && last - .content - .iter() - .any(|c| matches!(c, Content::Text(t) if t.starts_with(marker))) - { - msgs.pop(); - break; - } - match last.role { - Role::Tool | Role::Assistant => { - msgs.pop(); - } - _ => break, - } - } + if let Some(i) = prior_nudge { + msgs.truncate(i); } msgs.push(Message::user(tagged)); } diff --git a/crates/hi-ai/src/anthropic.rs b/crates/hi-ai/src/anthropic.rs index 25acb9a4..94bbadb5 100644 --- a/crates/hi-ai/src/anthropic.rs +++ b/crates/hi-ai/src/anthropic.rs @@ -19,6 +19,10 @@ use crate::types::{ const API_VERSION: &str = "2023-06-01"; +/// Upper bound on content blocks in one response — the per-event `index` +/// arrives straight from the provider's SSE JSON. +const MAX_CONTENT_BLOCKS: usize = 512; + pub struct AnthropicProvider { http: reqwest::Client, base_url: String, @@ -105,6 +109,11 @@ impl Provider for AnthropicProvider { } "content_block_start" => { let index = data["index"].as_u64().unwrap_or(0) as usize; + // The index comes straight off the wire — bound it so a + // corrupt frame can't force a huge `resize_with` allocation. + if index >= MAX_CONTENT_BLOCKS { + continue; + } if blocks.len() <= index { blocks.resize_with(index + 1, || None); } @@ -155,9 +164,17 @@ impl Provider for AnthropicProvider { .collect(); if completion.usage.input_tokens == 0 { completion.usage.input_tokens = estimate_messages_tokens(&request.messages); + completion.usage.estimated = true; } if completion.usage.output_tokens == 0 { completion.usage.output_tokens = estimate_completion_output_tokens(&completion.content); + completion.usage.estimated = true; + } + // Keep the occupancy gauge alive on the estimate path too (matches the + // OpenAI path's backfill): a proxy that omits `message_start` usage + // would otherwise leave it at 0 all session. + if completion.usage.context_occupancy == 0 { + completion.usage.context_occupancy = completion.usage.input_tokens; } Ok(completion) } diff --git a/crates/hi-ai/src/openai/stream.rs b/crates/hi-ai/src/openai/stream.rs index 2f1ef20a..f2a19a75 100644 --- a/crates/hi-ai/src/openai/stream.rs +++ b/crates/hi-ai/src/openai/stream.rs @@ -17,6 +17,11 @@ use crate::types::{Completion, Content, StreamEvent, Usage, estimate_messages_to const SPECIAL_TOKEN_PREFIX: &str = "<|"; const SPECIAL_TOKEN_SUFFIX: &str = "|>"; +/// Upper bound on parallel tool calls in one response. The per-delta `index` +/// arrives straight from the provider's SSE JSON; without a bound, one corrupt +/// frame (`"index": 9999999999`) forces a multi-GB allocation. +const MAX_STREAM_TOOL_CALLS: usize = 512; + /// How long to keep reading after `finish_reason` for the trailing usage frame. /// With `stream_options.include_usage`, spec-compliant providers send usage in /// a separate chunk AFTER the finish chunk (empty `choices`) — breaking on @@ -188,9 +193,22 @@ impl<'a> StreamingTextFilter<'a> { let mut i = 0; while i < bytes.len() { if bytes[i] == b'<' && i + 1 < bytes.len() && bytes[i + 1] == b'|' { - if let Some(end) = combined[i..].find(SPECIAL_TOKEN_SUFFIX) { - let token_end = i + end + SPECIAL_TOKEN_SUFFIX.len(); - let inner = &combined[i + SPECIAL_TOKEN_PREFIX.len()..i + end]; + let inner_start = i + SPECIAL_TOKEN_PREFIX.len(); + // A token identifier is `[A-Za-z0-9_]+`, so if the char right + // after `<|` can't start one (e.g. the text `<|>`), this can + // never become a special token — pass it through instead of + // buffering. Searching the suffix from `inner_start` also keeps + // the prefix's own `|` from matching `|>` (which produced a + // reversed slice range and panicked on the 3-char input `<|>`). + if inner_start < bytes.len() + && !(bytes[inner_start].is_ascii_alphanumeric() || bytes[inner_start] == b'_') + { + i += 1; + continue; + } + if let Some(end) = combined[inner_start..].find(SPECIAL_TOKEN_SUFFIX) { + let token_end = inner_start + end + SPECIAL_TOKEN_SUFFIX.len(); + let inner = &combined[inner_start..inner_start + end]; if is_token_identifier(inner) { out.push_str(&combined[last..i]); last = token_end; @@ -565,15 +583,19 @@ fn strip_special_tokens(text: &str) -> String { let mut last = 0; let mut i = 0; while i < bytes.len() { + // Search the suffix from after the `<|` prefix so the prefix's own `|` + // can't match `|>` (the text `<|>` produced a reversed slice range and + // panicked). if bytes[i] == b'<' && i + 1 < bytes.len() && bytes[i + 1] == b'|' - && let Some(end) = text[i..].find(SPECIAL_TOKEN_SUFFIX) + && let Some(end) = text[i + SPECIAL_TOKEN_PREFIX.len()..].find(SPECIAL_TOKEN_SUFFIX) { - let inner = &text[i + SPECIAL_TOKEN_PREFIX.len()..i + end]; + let inner_start = i + SPECIAL_TOKEN_PREFIX.len(); + let inner = &text[inner_start..inner_start + end]; if is_token_identifier(inner) { out.push_str(&text[last..i]); - last = i + end + SPECIAL_TOKEN_SUFFIX.len(); + last = inner_start + end + SPECIAL_TOKEN_SUFFIX.len(); i = last; continue; } @@ -671,12 +693,23 @@ where if data == "[DONE]" { break; } - let chunk = serde_json::from_str::(&data).with_context(|| { - format!( - "malformed SSE JSON chunk: {}", - data.chars().take(160).collect::() - ) - })?; + let chunk = match serde_json::from_str::(&data) { + Ok(chunk) => chunk, + // Same trade-off as the read-error arm above: once the answer has + // finished, the only thing left is the trailing usage frame, and a + // provider that pads the tail with a heartbeat / vendor sentinel / + // truncated frame must not cost us (and re-bill) a fully-streamed + // completion. + Err(_) if stream_complete => break, + Err(err) => { + return Err(err).with_context(|| { + format!( + "malformed SSE JSON chunk: {}", + data.chars().take(160).collect::() + ) + }); + } + }; if let Some(error) = chunk.error { let message = stream_error_message(&error); return Err(ProviderError::new(classify_stream_api_error(&message), message).into()); @@ -722,10 +755,45 @@ where if let Some(deltas) = delta.tool_calls { progressed = true; for tcd in deltas { - if tool_calls.len() <= tcd.index { - tool_calls.resize_with(tcd.index + 1, ToolCallBuilder::default); + let index = match tcd.index { + // The index comes straight off the wire — bound it so a + // corrupt frame can't force a huge `resize_with` + // allocation (or a capacity-overflow abort). + Some(i) if i >= MAX_STREAM_TOOL_CALLS => continue, + Some(i) => i, + // Some local OpenAI-compat servers omit `index` (the + // same ones that omit `id` — see ToolCallBuilder:: + // finish). serde's default of 0 would concatenate every + // call into tool_calls[0] ("readread", `{..}{..}`). + // These servers send each call complete in one delta, + // so: a delta announcing a name (or a different id) + // when the current call already has one starts a new + // call; anything else continues the current call. + None => { + let delta_id = tcd.id.as_deref().unwrap_or(""); + let delta_name = tcd + .function + .as_ref() + .and_then(|f| f.name.as_deref()) + .unwrap_or(""); + match tool_calls.last() { + None => 0, + Some(last) + if (!delta_name.is_empty() && !last.name.is_empty()) + || (!delta_id.is_empty() + && !last.id.is_empty() + && last.id != delta_id) => + { + tool_calls.len() + } + Some(_) => tool_calls.len() - 1, + } + } + }; + if tool_calls.len() <= index { + tool_calls.resize_with(index + 1, ToolCallBuilder::default); } - let builder = &mut tool_calls[tcd.index]; + let builder = &mut tool_calls[index]; if let Some(id) = tcd.id && !id.is_empty() { @@ -956,8 +1024,11 @@ struct Delta { #[derive(Deserialize)] struct ToolCallDelta { + /// `None` when the server omits `index` entirely — routing then falls back + /// to the new-call heuristic in the delta loop instead of collapsing every + /// call into slot 0. #[serde(default)] - index: usize, + index: Option, #[serde(default)] id: Option, #[serde(default)] diff --git a/crates/hi-cli/src/config.rs b/crates/hi-cli/src/config.rs index 1be0f51e..426b991d 100644 --- a/crates/hi-cli/src/config.rs +++ b/crates/hi-cli/src/config.rs @@ -338,101 +338,67 @@ impl serde::Serialize for Config { } } -#[derive(Clone, Debug, Default, Deserialize)] +// Serialized with `skip_serializing_if` on every field so a saved config omits +// unset keys instead of filling with `model = ""` lines. Keep the attribute on +// each new field: a field missing it is fine, but a field missing from +// serialization entirely (as with the old hand-written `Serialize` impl) is +// silently deleted from the user's config file on every save. +#[derive(Clone, Debug, Default, Serialize, Deserialize)] pub struct Profile { + #[serde(skip_serializing_if = "Option::is_none")] pub provider: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub model: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub base_url: Option, /// MCP endpoint used for metadata discovery, when supported by the provider. + #[serde(skip_serializing_if = "Option::is_none")] pub mcp_url: Option, /// A literal API key (written by the setup wizard). + #[serde(skip_serializing_if = "Option::is_none")] pub api_key: Option, /// Name of an env var holding the API key for this profile. + #[serde(skip_serializing_if = "Option::is_none")] pub api_key_env: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub max_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub thinking_budget: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub tool_mode: Option, - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub compat: Option, /// Advertise only the essential tool subset instead of the full set. Small /// (~3B) local models can't reliably plan over the full ~20-tool schema; /// this restores usable tool-calling. Defaults to off. - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub minimal_tools: Option, /// Verifier-gated skill auto-curation: after a verified turn, distill a /// reusable technique into a learned skill. Defaults to off. - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub curate_skills: Option, /// Advertise the read-only `explore` subagent tool. On by default; set to /// false to disable (e.g. for a very small local model). - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub explore_subagents: Option, /// Advertise the write-capable `delegate` subagent tool. Off by default (the /// riskier tier); set to true to enable. - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub write_subagents: Option, /// Model id that decomposes a `/goal ` into sub-goals. Defaults to /// `pipe/glm-5.2-fast` on the pipenetwork profile; `None` disables planning. - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub planner_model: Option, /// Model id for the `/goal team` skeptic gate (reviews a turn before it /// advances a sub-goal). `None` (default) disables the gate. - #[serde(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] pub skeptic_model: Option, /// Other profile names to fall back to, in order, when this one returns /// nothing or errors. + #[serde(skip_serializing_if = "Option::is_none")] pub fallback: Option>, } -// Serialize `Profile` with clean output: omit `None` fields so the TOML -// doesn't fill with `model = ""` lines. We can't put `skip_serializing_if` on -// each field above (it'd require repeating it 9 times), so we implement a -// custom serializer that skips None values. -impl serde::Serialize for Profile { - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let mut s = serializer.serialize_struct("Profile", 11)?; - if let Some(v) = &self.provider { - s.serialize_field("provider", v)?; - } - if let Some(v) = &self.model { - s.serialize_field("model", v)?; - } - if let Some(v) = &self.base_url { - s.serialize_field("base_url", v)?; - } - if let Some(v) = &self.mcp_url { - s.serialize_field("mcp_url", v)?; - } - if let Some(v) = &self.api_key { - s.serialize_field("api_key", v)?; - } - if let Some(v) = &self.api_key_env { - s.serialize_field("api_key_env", v)?; - } - if let Some(v) = &self.max_tokens { - s.serialize_field("max_tokens", v)?; - } - if let Some(v) = &self.thinking_budget { - s.serialize_field("thinking_budget", v)?; - } - if let Some(v) = &self.tool_mode { - s.serialize_field("tool_mode", v)?; - } - if let Some(v) = &self.compat { - s.serialize_field("compat", v)?; - } - if let Some(v) = &self.fallback { - s.serialize_field("fallback", v)?; - } - s.end() - } -} - /// Fully-resolved settings used to build a provider and run the agent. #[derive(Debug)] pub struct Settings { @@ -900,6 +866,23 @@ pub fn writable_config_path(explicit: Option<&Path>) -> Option { default_config_path() } +/// Mask an API key (or env var name) for display: first and last four +/// characters with an ellipsis. Char-based, so a key containing multi-byte +/// characters (e.g. pasted with a stray curly quote) can't panic a byte slice. +pub fn mask_key(key: &str) -> String { + if key.is_empty() { + return "(none)".to_string(); + } + let chars: Vec = key.chars().collect(); + if chars.len() > 8 { + let head: String = chars[..4].iter().collect(); + let tail: String = chars[chars.len() - 4..].iter().collect(); + format!("{head}…{tail}") + } else { + "***".to_string() + } +} + /// Serialize `config` to TOML and write it to `path`, creating parent dirs. /// Sets 0600 permissions on Unix so API keys in the file aren't world-readable. pub fn save_config_to(config: &Config, path: &Path) -> Result<()> { @@ -978,6 +961,22 @@ impl ProfileForm { p } + /// Build a `Profile` for an *edit*: start from the existing profile so the + /// fields the form doesn't cover (max_tokens, thinking_budget, tool_mode, + /// compat, fallback, subagent/planner settings, mcp_url, …) survive, and + /// overwrite only what the form actually edits. + pub fn apply_to(&self, existing: &Profile) -> Profile { + let form = self.to_profile(); + Profile { + provider: form.provider, + model: form.model, + base_url: form.base_url, + api_key: form.api_key, + api_key_env: form.api_key_env, + ..existing.clone() + } + } + /// Populate the form from an existing profile (for editing). pub fn from_profile(name: &str, p: &Profile) -> Self { Self { diff --git a/crates/hi-cli/src/main.rs b/crates/hi-cli/src/main.rs index bfd9954b..a178a2fd 100644 --- a/crates/hi-cli/src/main.rs +++ b/crates/hi-cli/src/main.rs @@ -70,16 +70,7 @@ async fn main() -> Result<()> { ); println!("tool_mode: {:?}", settings.tool_mode); println!("compat: {:?}", settings.compat); - let api_key_display = if settings.api_key.len() > 8 { - format!( - "{}...{}", - &settings.api_key[..4], - &settings.api_key[settings.api_key.len() - 4..] - ) - } else { - "***".to_string() - }; - println!("api_key: {}", api_key_display); + println!("api_key: {}", config::mask_key(&settings.api_key)); return Ok(()); } Err(err) => { @@ -456,12 +447,12 @@ async fn main() -> Result<()> { let path = config::writable_config_path(config_path.as_deref()) .context("could not determine config path")?; let mut file = file.lock().unwrap(); - let mut profile = form.to_profile(); - if profile.mcp_url.is_none() - && let Some(existing) = file.profiles.get(&data.name) - { - profile.mcp_url = existing.mcp_url.clone(); - } + // Editing an existing profile must not wipe the fields the form + // doesn't cover (max_tokens, fallback, tool_mode, …). + let profile = match file.profiles.get(&data.name) { + Some(existing) => form.apply_to(existing), + None => form.to_profile(), + }; config::upsert_profile(&mut file, &data.name, profile, &path)?; // Return the updated profile list. Ok(profile_infos(&file)) diff --git a/crates/hi-cli/src/repl.rs b/crates/hi-cli/src/repl.rs index 4fa2fd67..9c901e48 100644 --- a/crates/hi-cli/src/repl.rs +++ b/crates/hi-cli/src/repl.rs @@ -295,7 +295,7 @@ pub(crate) async fn repl( let arg = arg.trim(); // --- Subcommands --- if arg == "add" { - match provider_add_prompt(config, &mut editor) { + match provider_add_prompt(config, config_path.as_deref(), &mut editor) { Ok(name) => { println!( "\x1b[2msaved profile '{name}' — /provider {name} to use\x1b[0m" @@ -309,7 +309,12 @@ pub(crate) async fn repl( } if let Some(edit_name) = arg.strip_prefix("edit") { let edit_name = edit_name.trim(); - match provider_edit_prompt(config, edit_name, &mut editor) { + match provider_edit_prompt( + config, + config_path.as_deref(), + edit_name, + &mut editor, + ) { Ok(name) => { println!("\x1b[2msaved profile '{name}'\x1b[0m"); } @@ -341,7 +346,9 @@ pub(crate) async fn repl( ); continue; } - let path = match config::writable_config_path(None) { + let path = match config::writable_config_path( + config_path.as_deref(), + ) { Some(p) => p, None => { eprintln!("\x1b[33mcould not determine config path\x1b[0m"); @@ -759,6 +766,7 @@ fn rl_prompt(editor: &mut crate::complete::ReplEditor, message: &str) -> Result< /// config file. Returns the profile name. fn provider_add_prompt( config: &mut config::Config, + config_path: Option<&Path>, editor: &mut crate::complete::ReplEditor, ) -> Result { use config::{ProfileForm, ProviderName, upsert_profile, writable_config_path}; @@ -842,7 +850,7 @@ fn provider_add_prompt( }; let profile = form.to_profile(); - let path = writable_config_path(None).context("could not determine config path")?; + let path = writable_config_path(config_path).context("could not determine config path")?; upsert_profile(config, &name, profile, &path)?; Ok(name) } @@ -850,6 +858,7 @@ fn provider_add_prompt( /// Interactively edit an existing profile. `name` may be empty to prompt for it. fn provider_edit_prompt( config: &mut config::Config, + config_path: Option<&Path>, name: &str, editor: &mut crate::complete::ReplEditor, ) -> Result { @@ -903,17 +912,7 @@ fn provider_edit_prompt( // API key. let key_label = if form.store_as_env { "env var" } else { "key" }; - let masked = if form.api_key.len() > 8 { - format!( - "{}…{}", - &form.api_key[..4], - &form.api_key[form.api_key.len() - 4..] - ) - } else if form.api_key.is_empty() { - "(none)".to_string() - } else { - "***".to_string() - }; + let masked = config::mask_key(&form.api_key); let new_key = rl_prompt( editor, &format!("API key/{key_label} (current: {masked}): "), @@ -935,11 +934,8 @@ fn provider_edit_prompt( form.base_url = new_url; } - let mut profile = form.to_profile(); - if profile.mcp_url.is_none() { - profile.mcp_url = existing.mcp_url.clone(); - } - let path = writable_config_path(None).context("could not determine config path")?; + let profile = form.apply_to(existing); + let path = writable_config_path(config_path).context("could not determine config path")?; upsert_profile(config, &name, profile, &path)?; Ok(name) } diff --git a/crates/hi-cli/src/session.rs b/crates/hi-cli/src/session.rs index 9bca5509..a99058b0 100644 --- a/crates/hi-cli/src/session.rs +++ b/crates/hi-cli/src/session.rs @@ -5,7 +5,7 @@ //! sessions (pi-style) are a future extension; this is a linear log. use std::fs::{self, File, OpenOptions}; -use std::io::{BufWriter, Write}; +use std::io::Write; use std::path::{Path, PathBuf}; use std::time::{SystemTime, UNIX_EPOCH}; @@ -73,28 +73,42 @@ impl JsonlSession { Self { path } } - /// Persist checkpoint refs so a resumed session knows where it branched. - #[allow(dead_code)] - pub fn record_checkpoints(&mut self, refs: &[String]) -> Result<()> { - if refs.is_empty() { - return Ok(()); - } + /// Append a fully-formatted payload (one or more `\n`-terminated JSONL + /// lines) with a single `write_all` on the `O_APPEND` fd. A buffered + /// writer would split records larger than its buffer across multiple + /// `write()` calls, letting a concurrent appender (a second `hi -c` in the + /// same project, or a fleet child on `--session-file`) interleave mid-line + /// — and `load_history` silently drops unparseable lines. + fn append(&self, payload: &str) -> Result<()> { if let Some(parent) = self.path.parent() { fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?; } - let file = OpenOptions::new() + let mut file = OpenOptions::new() .create(true) .append(true) .open(&self.path) .with_context(|| format!("opening {}", self.path.display()))?; - let mut writer = BufWriter::new(file); - let line = serde_json::to_string(&SessionMeta::Checkpoints { - refs: refs.to_vec(), - })?; - writeln!(writer, "{line}")?; - writer.flush()?; + file.write_all(payload.as_bytes()) + .with_context(|| format!("appending to {}", self.path.display()))?; Ok(()) } + + fn append_meta(&self, meta: &SessionMeta) -> Result<()> { + let mut line = serde_json::to_string(meta)?; + line.push('\n'); + self.append(&line) + } + + /// Persist checkpoint refs so a resumed session knows where it branched. + #[allow(dead_code)] + pub fn record_checkpoints(&mut self, refs: &[String]) -> Result<()> { + if refs.is_empty() { + return Ok(()); + } + self.append_meta(&SessionMeta::Checkpoints { + refs: refs.to_vec(), + }) + } } impl SessionSink for JsonlSession { @@ -106,97 +120,40 @@ impl SessionSink for JsonlSession { if messages.is_empty() && usage.is_zero() { return Ok(()); } - if let Some(parent) = self.path.parent() { - fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?; - } - let file = OpenOptions::new() - .create(true) - .append(true) - .open(&self.path) - .with_context(|| format!("opening {}", self.path.display()))?; - let mut writer = BufWriter::new(file); + let mut payload = String::new(); for message in messages { - let line = serde_json::to_string(message)?; - writeln!(writer, "{line}")?; + payload.push_str(&serde_json::to_string(message)?); + payload.push('\n'); } - let line = serde_json::to_string(&SessionMeta::Usage { + payload.push_str(&serde_json::to_string(&SessionMeta::Usage { input_tokens: usage.input_tokens, output_tokens: usage.output_tokens, cache_read_tokens: usage.cache_read_tokens, cache_creation_tokens: usage.cache_creation_tokens, estimated: usage.estimated, - })?; - writeln!(writer, "{line}")?; - writer.flush()?; - Ok(()) + })?); + payload.push('\n'); + self.append(&payload) } fn record_compaction(&mut self, messages: &[Message]) -> Result<()> { - if let Some(parent) = self.path.parent() { - fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?; - } - let file = OpenOptions::new() - .create(true) - .append(true) - .open(&self.path) - .with_context(|| format!("opening {}", self.path.display()))?; - let mut writer = BufWriter::new(file); - let line = serde_json::to_string(&SessionMeta::Compaction { + self.append_meta(&SessionMeta::Compaction { messages: messages.to_vec(), - })?; - writeln!(writer, "{line}")?; - writer.flush()?; - Ok(()) + }) } fn record_goal(&mut self, goal: &hi_agent::Goal) -> Result<()> { - if let Some(parent) = self.path.parent() { - fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?; - } - let file = OpenOptions::new() - .create(true) - .append(true) - .open(&self.path) - .with_context(|| format!("opening {}", self.path.display()))?; - let mut writer = BufWriter::new(file); - let line = serde_json::to_string(&SessionMeta::Goal { goal: goal.clone() })?; - writeln!(writer, "{line}")?; - writer.flush()?; - Ok(()) + self.append_meta(&SessionMeta::Goal { goal: goal.clone() }) } fn clear_goal(&mut self) -> Result<()> { - if let Some(parent) = self.path.parent() { - fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?; - } - let file = OpenOptions::new() - .create(true) - .append(true) - .open(&self.path) - .with_context(|| format!("opening {}", self.path.display()))?; - let mut writer = BufWriter::new(file); - let line = serde_json::to_string(&SessionMeta::GoalCleared)?; - writeln!(writer, "{line}")?; - writer.flush()?; - Ok(()) + self.append_meta(&SessionMeta::GoalCleared) } fn record_decisions(&mut self, decisions: &hi_agent::DecisionLog) -> Result<()> { - if let Some(parent) = self.path.parent() { - fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?; - } - let file = OpenOptions::new() - .create(true) - .append(true) - .open(&self.path) - .with_context(|| format!("opening {}", self.path.display()))?; - let mut writer = BufWriter::new(file); - let line = serde_json::to_string(&SessionMeta::Decisions { + self.append_meta(&SessionMeta::Decisions { decisions: decisions.entries().to_vec(), - })?; - writeln!(writer, "{line}")?; - writer.flush()?; - Ok(()) + }) } fn record_state_replacement( @@ -205,23 +162,11 @@ impl SessionSink for JsonlSession { goal: Option<&hi_agent::Goal>, decisions: &hi_agent::DecisionLog, ) -> Result<()> { - if let Some(parent) = self.path.parent() { - fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?; - } - let file = OpenOptions::new() - .create(true) - .append(true) - .open(&self.path) - .with_context(|| format!("opening {}", self.path.display()))?; - let mut writer = BufWriter::new(file); - let line = serde_json::to_string(&SessionMeta::StateReplacement { + self.append_meta(&SessionMeta::StateReplacement { messages: messages.to_vec(), goal: goal.cloned(), decisions: decisions.entries().to_vec(), - })?; - writeln!(writer, "{line}")?; - writer.flush()?; - Ok(()) + }) } } diff --git a/crates/hi-lsp/src/client.rs b/crates/hi-lsp/src/client.rs index d382303e..7a5d1ee5 100644 --- a/crates/hi-lsp/src/client.rs +++ b/crates/hi-lsp/src/client.rs @@ -3,12 +3,12 @@ use std::collections::HashMap; use std::path::Path; use std::sync::Mutex as StdMutex; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::{Duration, Instant}; use anyhow::{Context, Result, bail}; use serde_json::{Value, json}; -use tokio::io::BufReader; +use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::process::{Child, ChildStdout}; use tokio::sync::Mutex; @@ -18,11 +18,38 @@ use crate::protocol::{read_message, request_timeout, write_message}; /// keyed by document URI. Updated as notifications arrive during requests. pub type DiagnosticsMap = StdMutex>>; +/// Once data for a message has started arriving, how long a complete frame may +/// take to finish before we declare the stream unrecoverable. Local pipes +/// deliver even multi-MB `publishDiagnostics` payloads in milliseconds, so +/// this only fires on a genuinely wedged server. +const MESSAGE_GRACE: Duration = Duration::from_secs(10); + +/// Outcome of one bounded read attempt on the server's stdout. +enum ReadOutcome { + /// A complete JSON-RPC message. + Message(Vec), + /// No data arrived within the budget; the stream position is untouched. + Idle, + /// The stream closed (server exited or pipe broke at a frame boundary). + Closed, + /// A read stalled mid-frame — the stream position is unknown, so this + /// client is unusable and must be respawned (see [`LspClient::is_poisoned`]). + Poisoned, +} + /// One running language server. Owns the child process and its stdio. pub struct LspClient { child: Mutex, stdin: Mutex, stdout: Mutex>, + /// Serializes whole request/drain round-trips. Without it, two concurrent + /// requests read from the same stream and each can consume (and drop) the + /// other's response, leaving the loser to spin until its timeout. + io: Mutex<()>, + /// Set when a read stopped mid-frame, leaving the JSON-RPC stream position + /// unknown: every later read would misparse body bytes as headers. The + /// manager checks this and respawns the server. + poisoned: AtomicBool, next_id: AtomicU64, versions: StdMutex>, /// Diagnostics pushed by the server, keyed by document URI. @@ -48,6 +75,8 @@ impl LspClient { child: Mutex::new(child), stdin: Mutex::new(stdin), stdout: Mutex::new(BufReader::new(stdout)), + io: Mutex::new(()), + poisoned: AtomicBool::new(false), next_id: AtomicU64::new(1), versions: StdMutex::new(HashMap::new()), pushed_diagnostics: StdMutex::new(HashMap::new()), @@ -73,17 +102,61 @@ impl LspClient { Ok(()) } + /// Read one message from stdout, spending at most `budget` *waiting for + /// data to appear*. Split into two phases because `read_message` is not + /// cancellation-safe: a timeout that fires mid-frame throws away consumed + /// bytes and permanently desyncs the stream. + /// + /// Phase 1 waits on `fill_buf`, which consumes nothing — cancelling it is + /// harmless. Once bytes are available, phase 2 commits to reading the full + /// frame under its own generous grace; only if *that* stalls is the client + /// marked poisoned (respawned by the manager on the next query). + async fn read_one(&self, budget: Duration) -> ReadOutcome { + let mut stdout = self.stdout.lock().await; + match tokio::time::timeout(budget, stdout.fill_buf()).await { + Err(_) => return ReadOutcome::Idle, + Ok(Err(_)) => return ReadOutcome::Closed, + Ok(Ok([])) => return ReadOutcome::Closed, + Ok(Ok(_)) => {} + } + match tokio::time::timeout(MESSAGE_GRACE, read_message(&mut stdout)).await { + Ok(Some(msg)) => ReadOutcome::Message(msg), + Ok(None) => ReadOutcome::Closed, + Err(_) => { + self.poisoned.store(true, Ordering::SeqCst); + ReadOutcome::Poisoned + } + } + } + + /// Record a `publishDiagnostics` notification into `pushed_diagnostics`. + /// Other notifications are dropped. + fn capture_notification(&self, v: &Value) { + if v.get("method").and_then(|m| m.as_str()) == Some("textDocument/publishDiagnostics") + && let Some(params) = v.get("params") + && let Some(uri) = params.get("uri").and_then(|u| u.as_str()) + { + let diags = params + .get("diagnostics") + .and_then(|d| d.as_array()) + .cloned() + .unwrap_or_default(); + self.pushed_diagnostics + .lock() + .unwrap() + .insert(uri.to_string(), diags); + } + } + /// Send a request and wait for the matching response. /// - /// NOTE: this acquires `self.stdout` on each read, so concurrent requests - /// on the same `LspClient` serialize. `LspManager` now clones an `Arc` - /// handle and drops the servers lock before calling, so different languages - /// run concurrently — but two calls to the *same* client still serialize - /// here. A response whose `id` doesn't match is currently dropped (only - /// `publishDiagnostics` notifications are captured). If pipelining is ever - /// added, a background reader task feeding a per-id channel would be needed - /// to avoid losing out-of-order responses. + /// The whole round-trip holds the `io` lock, so concurrent requests (and + /// drains) on the same client serialize instead of stealing each other's + /// messages off the shared stream. `LspManager` clones an `Arc` handle and + /// drops the servers lock before calling, so different languages still run + /// concurrently. pub async fn request(&self, method: &str, params: Option) -> Result { + let _io = self.io.lock().await; let id = self.next_id.fetch_add(1, Ordering::SeqCst); let body = json!({ "jsonrpc": "2.0", @@ -101,40 +174,25 @@ impl LspClient { if remaining.is_zero() { bail!("LSP request `{method}` timed out"); } - let msg = tokio::time::timeout(remaining, async { - let mut stdout = self.stdout.lock().await; - read_message(&mut stdout).await - }) - .await - .context("LSP read timed out")?; - let msg = match msg { - Some(m) => m, - None => bail!("LSP server closed the stream"), - }; - let v: Value = serde_json::from_slice(&msg)?; - if v.get("id").and_then(|i| i.as_u64()) == Some(id) { - if let Some(err) = v.get("error") { - bail!("LSP error on `{method}`: {err}"); + match self.read_one(remaining).await { + ReadOutcome::Idle => continue, // deadline re-checked above + ReadOutcome::Closed => bail!("LSP server closed the stream"), + ReadOutcome::Poisoned => bail!( + "LSP stream lost sync during `{method}`; the server will be restarted on the next query" + ), + ReadOutcome::Message(msg) => { + let v: Value = serde_json::from_slice(&msg)?; + if v.get("id").and_then(|i| i.as_u64()) == Some(id) { + if let Some(err) = v.get("error") { + bail!("LSP error on `{method}`: {err}"); + } + return Ok(v["result"].clone()); + } + // Capture pushed diagnostics — the server sends these as + // notifications after didOpen/didChange, not as responses. + self.capture_notification(&v); } - return Ok(v["result"].clone()); - } - // Capture pushed diagnostics — the server sends these as - // notifications after didOpen/didChange, not as responses. - if v.get("method").and_then(|m| m.as_str()) == Some("textDocument/publishDiagnostics") - && let Some(params) = v.get("params") - && let Some(uri) = params.get("uri").and_then(|u| u.as_str()) - { - let diags = params - .get("diagnostics") - .and_then(|d| d.as_array()) - .cloned() - .unwrap_or_default(); - self.pushed_diagnostics - .lock() - .unwrap() - .insert(uri.to_string(), diags); } - // Other notifications are dropped. } } @@ -177,40 +235,28 @@ impl LspClient { Ok(()) } - /// Read any pending notifications from stdout without blocking, capturing - /// `publishDiagnostics` into `pushed_diagnostics`. Times out after `wait` - /// if no data is available. + /// Read any pending notifications from stdout, capturing + /// `publishDiagnostics` into `pushed_diagnostics`. Returns after `wait` + /// with no data. Holds the `io` lock so it can't race a concurrent + /// request's read loop and eat its response; the two-phase `read_one` + /// means an expiring budget can't cancel a frame mid-read (which used to + /// desync the stream when a large diagnostics payload straddled the + /// deadline). pub async fn drain_notifications(&self, wait: Duration) { + let _io = self.io.lock().await; let deadline = Instant::now() + wait; loop { let remaining = deadline.saturating_duration_since(Instant::now()); if remaining.is_zero() { return; } - let msg = tokio::time::timeout(remaining, async { - let mut stdout = self.stdout.lock().await; - read_message(&mut stdout).await - }) - .await; - let msg = match msg { - Ok(Some(m)) => m, - _ => return, - }; - if let Ok(v) = serde_json::from_slice::(&msg) - && v.get("method").and_then(|m| m.as_str()) - == Some("textDocument/publishDiagnostics") - && let Some(params) = v.get("params") - && let Some(uri) = params.get("uri").and_then(|u| u.as_str()) - { - let diags = params - .get("diagnostics") - .and_then(|d| d.as_array()) - .cloned() - .unwrap_or_default(); - self.pushed_diagnostics - .lock() - .unwrap() - .insert(uri.to_string(), diags); + match self.read_one(remaining).await { + ReadOutcome::Message(msg) => { + if let Ok(v) = serde_json::from_slice::(&msg) { + self.capture_notification(&v); + } + } + ReadOutcome::Idle | ReadOutcome::Closed | ReadOutcome::Poisoned => return, } } } @@ -222,8 +268,12 @@ impl LspClient { /// strong ref. The `child`/`stdin`/`stdout` are already behind `Mutex`es, /// so no `&mut` is actually required. pub async fn shutdown(&self) -> Result<()> { - let _ = self.request("shutdown", None).await; - let _ = self.notify("exit", Value::Null).await; + // Skip the graceful JSON-RPC goodbye on a desynced stream — the + // `shutdown` request would only misread frames until its timeout. + if !self.is_poisoned() { + let _ = self.request("shutdown", None).await; + let _ = self.notify("exit", Value::Null).await; + } // Give the server a moment to exit gracefully, then force-kill so a // stubborn server can't hang the shutdown indefinitely. `kill_on_drop` // would eventually clean up, but only when the `LspClient` is dropped — @@ -242,6 +292,13 @@ impl LspClient { matches!(self.child.lock().await.try_wait(), Ok(None)) } + /// Whether the JSON-RPC stream has lost sync (a read stalled mid-frame). + /// A poisoned client's child may still be alive, but no further messages + /// can be exchanged reliably — the manager respawns it on the next query. + pub fn is_poisoned(&self) -> bool { + self.poisoned.load(Ordering::SeqCst) + } + /// Get the diagnostics the server has pushed for a URI (via /// `publishDiagnostics`), if any. Returns a clone of the raw JSON values. pub fn get_pushed_diagnostics(&self, uri: &str) -> Vec { diff --git a/crates/hi-lsp/src/manager.rs b/crates/hi-lsp/src/manager.rs index e733ad07..2fc50da6 100644 --- a/crates/hi-lsp/src/manager.rs +++ b/crates/hi-lsp/src/manager.rs @@ -91,6 +91,10 @@ impl LspManager { let _ = client.shutdown().await; } self.running.lock().unwrap().clear(); + // The dedup hashes describe documents open on the servers we just + // shut down. Without clearing, a later `/lsp on` would skip the + // didOpen for "already synced" files the fresh servers never saw. + self.synced.lock().unwrap().clear(); } else { // Warm up the server for the project's primary language. if let Some(lang) = detect_project_language(&self.root) @@ -159,29 +163,42 @@ impl LspManager { /// wins, and later spawners detect the existing live server and explicitly /// shut down their duplicate (rather than relying on `kill_on_drop`). async fn ensure(&self, lang: Language) -> Result<()> { - // Fast path: already running and alive. + // Fast path: already running, alive, and its stream is intact. { let servers = self.servers.lock().await; if let Some(client) = servers.get(&lang) { - if client.is_alive().await { + if client.is_alive().await && !client.is_poisoned() { return Ok(()); } - // Crashed — fall through to respawn. Drop the lock first by - // removing under a fresh acquisition below. + // Crashed or desynced — fall through to respawn. Drop the lock + // first by removing under a fresh acquisition below. } else { // Not present — spawn. } } - // Remove any dead entry, then spawn outside the lock. - { + // Remove any dead/poisoned entry, then spawn outside the lock. + let stale = { let mut servers = self.servers.lock().await; - if let Some(client) = servers.get(&lang) { - if client.is_alive().await { - return Ok(()); // raced with another ensure; it's alive + match servers.get(&lang) { + Some(client) if client.is_alive().await && !client.is_poisoned() => { + return Ok(()); // raced with another ensure; it's healthy + } + Some(_) => { + let stale = servers.remove(&lang); + self.running.lock().unwrap().remove(&lang); + // The old server had documents open that the replacement + // won't know about — drop the dedup hashes so every file + // re-syncs (didOpen) on its next query instead of being + // skipped as "unchanged" against a server that never saw it. + self.synced.lock().unwrap().clear(); + stale } - servers.remove(&lang); - self.running.lock().unwrap().remove(&lang); + None => None, } + }; + // A poisoned child may still be alive — reap it deterministically. + if let Some(stale) = stale { + let _ = stale.shutdown().await; } if !server_available(lang) { bail!("{}", install_hint(lang)); @@ -189,17 +206,18 @@ impl LspManager { let (cmd, args) = server_command(lang); let client = LspClient::spawn(cmd, &args, &self.root).await?; let mut servers = self.servers.lock().await; - // If another task raced and inserted a live server, keep it and + // If another task raced and inserted a healthy server, keep it and // explicitly shut down the duplicate we just spawned (rather than // relying on `kill_on_drop` at drop time, which is correct but // non-deterministic about *when* the orphaned child is reaped). if let Some(existing) = servers.get(&lang) { - if existing.is_alive().await { + if existing.is_alive().await && !existing.is_poisoned() { drop(servers); // release before awaiting shutdown let _ = client.shutdown().await; return Ok(()); } servers.remove(&lang); + self.synced.lock().unwrap().clear(); } servers.insert(lang, Arc::new(client)); self.running.lock().unwrap().insert(lang, true); @@ -258,11 +276,18 @@ impl LspManager { .with_context(|| format!("no LSP server for {lang:?} after ensure"))? .clone() }; - if already_open { + let result = if already_open { client.clear_pushed_diagnostics(&uri); - client.did_change(&uri, text).await?; + client.did_change(&uri, text).await } else { - client.did_open(&uri, lang.language_id(), text).await?; + client.did_open(&uri, lang.language_id(), text).await + }; + if let Err(e) = result { + // The hash was inserted optimistically above, but the server never + // received the notify — leaving it would make every future sync + // skip this content as "unchanged". + self.synced.lock().unwrap().remove(&uri); + return Err(e); } Ok(()) } diff --git a/crates/hi-tools/src/background.rs b/crates/hi-tools/src/background.rs index c7d0175b..4a41ea14 100644 --- a/crates/hi-tools/src/background.rs +++ b/crates/hi-tools/src/background.rs @@ -235,10 +235,21 @@ async fn drive( /// cap by front-trimming on a char boundary (and shifting the read cursor). async fn pump(pipe: Option, proc: &BgProc) { let Some(pipe) = pipe else { return }; - let mut reader = BufReader::new(pipe).lines(); - while let Ok(Some(line)) = reader.next_line().await { + // Read raw bytes and lossy-decode per line: `next_line()` errors on the + // first invalid-UTF-8 byte, which would stop draining the pipe — output + // after that point would be lost, and a child still writing would block on + // a full pipe buffer. + let mut reader = BufReader::new(pipe); + let mut bytes = Vec::new(); + loop { + bytes.clear(); + match reader.read_until(b'\n', &mut bytes).await { + Ok(0) | Err(_) => break, + Ok(_) => {} + } + let line = String::from_utf8_lossy(&bytes); let mut inner = proc.inner.lock().unwrap(); - inner.output.push_str(&line); + inner.output.push_str(line.trim_end_matches(['\r', '\n'])); inner.output.push('\n'); if inner.output.len() > MAX_BG_BUFFER { let overflow = inner.output.len() - MAX_BG_BUFFER; diff --git a/crates/hi-tools/src/read.rs b/crates/hi-tools/src/read.rs index 9ce23bb7..f3ab0d94 100644 --- a/crates/hi-tools/src/read.rs +++ b/crates/hi-tools/src/read.rs @@ -495,8 +495,11 @@ pub(crate) fn is_binary(bytes: &[u8]) -> bool { } /// Read a file as UTF-8 text, bailing with a clear message if it's binary -/// (same heuristic as `read`). Used by `edit`/`multi_edit` so a binary file -/// produces a helpful error instead of an opaque `read_to_string` UTF-8 panic. +/// (same heuristic as `read`) or not valid UTF-8. Used by the preserving-edit +/// paths (`edit`/`multi_edit`/`apply_patch`), which write the decoded string +/// back to disk — a lossy decode here would silently replace every invalid +/// byte in the whole file with U+FFFD on the write-back, corrupting e.g. +/// Latin-1 files even on lines the edit never touched. pub(crate) async fn read_text_file(path: &str) -> Result { let bytes = tokio::fs::read(path) .await @@ -508,7 +511,13 @@ pub(crate) async fn read_text_file(path: &str) -> Result { bytes.len() ); } - Ok(String::from_utf8_lossy(&bytes).into_owned()) + String::from_utf8(bytes).map_err(|e| { + anyhow::anyhow!( + "{path} is not valid UTF-8 (first invalid byte at offset {}) — editing it in place \ + would corrupt its encoding. Use `bash` (e.g. sed/iconv) to modify it.", + e.utf8_error().valid_up_to() + ) + }) } #[derive(Deserialize)] diff --git a/crates/hi-tools/src/tools.rs b/crates/hi-tools/src/tools.rs index 4e3e7e39..a2010e38 100644 --- a/crates/hi-tools/src/tools.rs +++ b/crates/hi-tools/src/tools.rs @@ -1509,9 +1509,20 @@ async fn read_lines( return; }; use tokio::io::{AsyncBufReadExt, BufReader}; - let mut reader = BufReader::new(pipe).lines(); - while let Ok(Some(line)) = reader.next_line().await { - let mut with_nl = line; + // Read raw bytes and lossy-decode per line: `next_line()` errors on the + // first invalid-UTF-8 byte, which would stop draining the pipe — output + // after that point would be lost, and a child still writing would block on + // a full pipe buffer until the timeout kills it. + let mut reader = BufReader::new(pipe); + let mut bytes = Vec::new(); + loop { + bytes.clear(); + match reader.read_until(b'\n', &mut bytes).await { + Ok(0) | Err(_) => break, + Ok(_) => {} + } + let line = String::from_utf8_lossy(&bytes); + let mut with_nl = line.trim_end_matches(['\r', '\n']).to_string(); with_nl.push('\n'); if let Ok(mut cb) = on_line.lock() { (*cb)(&with_nl); diff --git a/crates/hi-tools/src/web.rs b/crates/hi-tools/src/web.rs index bf66ffed..3b070f60 100644 --- a/crates/hi-tools/src/web.rs +++ b/crates/hi-tools/src/web.rs @@ -557,8 +557,14 @@ async fn resolve_download( source: &str, filename: Option<&str>, ) -> Result<(DownloadTarget, String)> { - // Full URL — use directly. + // Full URL — validate like `web_fetch` before handing it to curl/aria2c: + // the URL is model-supplied, and without this check `web_download` reaches + // cloud metadata services / localhost / private ranges and writes the + // response into a workspace file the model can read back. (The downloader + // still follows redirects without re-validation; the initial-hop check + // blocks the direct form of the attack.) if source.starts_with("http://") || source.starts_with("https://") { + validate_url(source)?; let name = filename .map(str::to_string) .or_else(|| { diff --git a/crates/hi-tui/src/app/render.rs b/crates/hi-tui/src/app/render.rs index 2f5c46df..4581f40e 100644 --- a/crates/hi-tui/src/app/render.rs +++ b/crates/hi-tui/src/app/render.rs @@ -764,8 +764,10 @@ impl crate::App { .and_then(|i| self.input.history.get(i)) .map(|s| s.replace('\n', " ")) .unwrap_or_default(); - let preview = if preview.len() > 60 { - format!("{}…", &preview[..60]) + // Char-based truncation: history entries are arbitrary input, + // and a byte slice panics on a multi-byte char at the cut. + let preview = if preview.chars().count() > 60 { + format!("{}…", preview.chars().take(60).collect::()) } else { preview }; diff --git a/crates/hi-tui/src/app/transcript.rs b/crates/hi-tui/src/app/transcript.rs index 2a72b5a0..cfc63884 100644 --- a/crates/hi-tui/src/app/transcript.rs +++ b/crates/hi-tui/src/app/transcript.rs @@ -378,12 +378,15 @@ impl crate::App { // Merge consecutive same-tool explore results into one line, so a // burst of reads renders as `⏺ read 6 files · 743 lines` instead of // six separate lines. A run continues only while the tool name is - // the same and no other event has broken the chain. + // the same AND the run's summary line is still the last transcript + // entry — events that commit lines without resetting the run + // (assistant text, status) would otherwise get overwritten by the + // in-place update below. + let last_pos = (self.trimmed + self.transcript.len() as u64).checked_sub(1); let merge = self .explore_run .as_ref() - .map(|r| r.tool == name) - .unwrap_or(false); + .is_some_and(|r| r.tool == name && Some(r.line_pos) == last_pos); if merge { let run = self.explore_run.as_mut().unwrap(); run.count += 1; @@ -395,12 +398,14 @@ impl crate::App { self.replace_last_line(line); return; } - // Start a new run. + // Start a new run; its summary line is about to be pushed at the + // current end of the transcript. self.explore_run = Some(ExploreRun { tool: name.to_string(), count: 1, lines: n, all_empty: n == 0, + line_pos: self.trimmed + self.transcript.len() as u64, }); let line = self.render_explore_run(&header); self.push(Line::styled(line, Style::default().fg(Color::Cyan))); diff --git a/crates/hi-tui/src/lib.rs b/crates/hi-tui/src/lib.rs index 82aafd3b..e6c4d359 100644 --- a/crates/hi-tui/src/lib.rs +++ b/crates/hi-tui/src/lib.rs @@ -334,6 +334,11 @@ pub(crate) struct ExploreRun { pub lines: u32, /// Whether every result so far was empty (`(no output)`). pub all_empty: bool, + /// Absolute transcript position (`trimmed` + local index) of this run's + /// summary line. Merging is only valid while that line is still the *last* + /// transcript entry — otherwise the in-place update would overwrite + /// whatever landed after it (e.g. committed assistant prose). + pub line_pos: u64, } impl TranscriptEntry { diff --git a/crates/hi-tui/src/loops.rs b/crates/hi-tui/src/loops.rs index 54ebbaf2..e5de39d4 100644 --- a/crates/hi-tui/src/loops.rs +++ b/crates/hi-tui/src/loops.rs @@ -986,10 +986,25 @@ async fn run_firing(launcher: &FleetLauncher, spec: &LoopSpec) -> Result status.map_err(|e| format!("wait failed: {e}"))?, + Err(_) => { + let _ = child.start_kill(); + let _ = child.wait().await; + return Err(format!( + "child closed stdout but didn't exit within {FIRING_TIMEOUT_SECS}s" + )); + } + }; if !status.success() { return Err(format!( "agent run failed ({}): {}",