feat(stream): add method to detect incomplete tool calls and handle parsing errors

This commit is contained in:
asepharyana
2026-07-15 00:49:50 +07:00
parent 4eba9d0a2f
commit 732d6039dc
3 changed files with 122 additions and 4 deletions
+60
View File
@@ -157,6 +157,28 @@ impl StreamedTurn {
msg
}
/// Find the first named tool call whose accumulated `arguments` do not
/// parse as valid JSON.
///
/// Why: a connection that closes mid-stream (no `[DONE]` event) still
/// leaves partial argument text in the accumulator — e.g. a `write`
/// tool call cut off mid-string. Parsing that fragment always fails,
/// so a parse failure at end-of-stream is a reliable signal that the
/// response was truncated, not that the model legitimately finished
/// without sending `[DONE]`.
///
/// Return: `Some((name, parse_error))` for the first bad tool call, or
/// `None` if every tool call's arguments are complete, parsable JSON.
pub fn incomplete_tool_call(&self) -> Option<(&str, String)> {
self.tool_calls.iter()
.filter(|tc| !tc.name.is_empty())
.find_map(|tc| {
serde_json::from_str::<Value>(&tc.arguments)
.err()
.map(|e| (tc.name.as_str(), e.to_string()))
})
}
/// Reserved accessor for callers that want to branch mid-stream before the turn
/// completes; the current wiring only inspects the final `build_assistant_message()`.
#[allow(dead_code)]
@@ -176,3 +198,41 @@ impl Default for StreamedTurn {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn tool_call(name: &str, arguments: &str) -> ParsedToolCall {
ParsedToolCall {
id: "call_1".to_string(),
name: name.to_string(),
arguments: arguments.to_string(),
is_complete: false,
}
}
#[test]
fn incomplete_tool_call_flags_truncated_json() {
let mut turn = StreamedTurn::new();
turn.tool_calls.push(tool_call("write", "{\"path\": \"a.txt\", \"content\": \"unterm"));
let bad = turn.incomplete_tool_call();
assert_eq!(bad.map(|(name, _)| name), Some("write"));
}
#[test]
fn incomplete_tool_call_accepts_complete_json() {
let mut turn = StreamedTurn::new();
turn.tool_calls.push(tool_call("write", "{\"path\": \"a.txt\", \"content\": \"done\"}"));
assert!(turn.incomplete_tool_call().is_none());
}
#[test]
fn incomplete_tool_call_ignores_calls_without_a_name() {
// A slot reserved by `apply_event` (via index padding) but never
// filled with a name shouldn't be mistaken for a truncated call.
let mut turn = StreamedTurn::new();
turn.tool_calls.push(tool_call("", "not json at all"));
assert!(turn.incomplete_tool_call().is_none());
}
}
+49 -4
View File
@@ -54,8 +54,10 @@ fn build_subagent_tools(allowed_tools: &[String]) -> (Vec<Box<dyn crate::tool::T
/// Why: matches the main agent's credential resolution exactly, so
/// subagents automatically inherit the same provider settings.
///
/// Return: `(api_key, model, optional_base_url)`.
fn resolve_provider_config() -> (String, String, Option<String>) {
/// Return: `(api_key, model, optional_base_url, provider_name)`. `api_key`
/// is empty when every resolution path was exhausted — callers must check
/// for this before issuing requests (see `run_subagent`).
fn resolve_provider_config() -> (String, String, Option<String>, String) {
let settings = crate::model::settings::Settings::load();
let app_config = crate::model::app_config::AppConfig::load();
@@ -79,7 +81,21 @@ fn resolve_provider_config() -> (String, String, Option<String>) {
}
}
(api_key, model, base_url)
(api_key, model, base_url, settings.provider)
}
/// Reject an empty API key with an actionable error instead of letting the
/// caller send a request that is guaranteed to fail once it reaches the network.
///
/// Return: `Ok(())` if `api_key` is non-empty, `Err` with a message naming
/// `provider` and where to fix it otherwise.
fn require_api_key(api_key: &str, provider: &str) -> anyhow::Result<()> {
if api_key.is_empty() {
anyhow::bail!(
"no API key configured for provider '{provider}' — set one in Settings or ~/.claude/settings.json"
);
}
Ok(())
}
// ─── Subagent-level tool gating (mirrors Harness checks) ───
@@ -325,7 +341,20 @@ pub fn run_subagent(ctx: &SubagentContext, tx: &mpsc::Sender<SubagentEvent>) ->
// Cache provider config once before the loop instead of re-resolving
// from disk on every step (Settings::load + AppConfig::load each parse
// JSON files, and the config cannot change between steps).
let (api_key, model, base_url) = resolve_provider_config();
let (api_key, model, base_url, provider) = resolve_provider_config();
// Fail fast on a missing key instead of sending a doomed request: an
// empty api_key still reaches the network (base_url falls back to a
// default endpoint), so without this check every step burns a full
// 10-retry timeout/backoff cycle against a server that was never going
// to authenticate, and the real cause (no key configured) never
// surfaces past a buried WARN log.
if let Err(error) = require_api_key(&api_key, &provider) {
let error = error.to_string();
let _ = tx.blocking_send(SubagentEvent::StepFailed { step: 0, error: error.clone() });
anyhow::bail!(error);
}
let client = crate::service::provider::LlmClient::new(api_key, model, base_url);
for step in 0..ctx.max_steps {
@@ -559,3 +588,19 @@ pub fn run_subagent(ctx: &SubagentContext, tx: &mpsc::Sender<SubagentEvent>) ->
let _ = tx.blocking_send(SubagentEvent::Completed { output: output.clone() });
Ok(output)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn require_api_key_rejects_empty_key_with_provider_named_in_message() {
let err = require_api_key("", "claude").unwrap_err();
assert!(err.to_string().contains("claude"));
}
#[test]
fn require_api_key_accepts_non_empty_key() {
assert!(require_api_key("sk-live-abc123", "claude").is_ok());
}
}
+13
View File
@@ -318,6 +318,19 @@ impl LlmClient {
}
}
// The connection closed without an explicit `[DONE]` event. Some
// providers legitimately omit it, so EOF alone isn't an error —
// but if it leaves a tool call's arguments as unparsable JSON, the
// response was truncated mid-generation, not finished. Report that
// honestly instead of silently double-stringifying the fragment
// into a tool call that will misbehave (e.g. a `write` call with a
// half-written file body).
if let Some((name, err)) = turn.incomplete_tool_call() {
anyhow::bail!(
"stream ended before tool call '{name}' arguments were complete: {err}"
);
}
turn.is_complete = true;
Ok((turn.build_assistant_message(), usage))
}