refactor(streaming): Extract deduplication and auto-continue logic into helpers

Improve readability of stream_completion_with_tools (~1000 line function):

- Add deduplicate_tool_calls() helper with closure for previous-message check
- Add should_auto_continue() with AutoContinueReason enum for clearer control flow
- Replace inline deduplication loop with helper call (-19 lines)
- Replace complex auto-continue conditional with match on reason enum (-13 lines)
- Add section comments for major phases (State Init, Pre-loop, Main Loop, Auto-Continue, Post-Loop)
- Add comprehensive tests for new helpers

Net reduction: 82 deletions, behavior unchanged (172+ tests pass)

Agent: carmack
This commit is contained in:
Dhanji R. Prasanna
2026-01-13 11:44:06 +05:30
parent dc45987e8d
commit a09967eb27
2 changed files with 216 additions and 82 deletions

View File

@@ -1,7 +1,8 @@
//! Streaming completion logic for the Agent.
//!
//! This module handles the streaming response from LLM providers,
//! including tool call detection, execution, and auto-continue logic.
//! This module provides state management and helper functions for streaming
//! LLM responses, including tool call detection, deduplication, and
//! auto-continue decision logic.
use crate::context_window::ContextWindow;
use crate::streaming_parser::StreamingToolParser;
@@ -385,6 +386,92 @@ pub fn format_rehydrate_summary(result: &str) -> String {
}
}
// =============================================================================
// Tool Call Deduplication
// =============================================================================
/// Result of deduplicating a batch of tool calls.
/// Each tool call is paired with an optional duplicate marker.
pub type DeduplicatedTools = Vec<(ToolCall, Option<String>)>;
/// Deduplicate tool calls, detecting sequential duplicates within a chunk
/// and duplicates against the previous message.
///
/// Returns each tool call paired with `Some("DUP IN CHUNK")` or `Some("DUP IN MSG")`
/// if it's a duplicate, or `None` if it should be executed.
pub fn deduplicate_tool_calls<F>(
tool_calls: Vec<ToolCall>,
check_previous_message: F,
) -> DeduplicatedTools
where
F: Fn(&ToolCall) -> Option<String>,
{
let mut last_tool_in_chunk: Option<ToolCall> = None;
let mut result = Vec::with_capacity(tool_calls.len());
for tool_call in tool_calls {
let duplicate_type = if let Some(ref last) = last_tool_in_chunk {
// Check for sequential duplicate within this chunk
if are_tool_calls_duplicate(last, &tool_call) {
Some("DUP IN CHUNK".to_string())
} else {
None
}
} else {
// First tool in chunk - check against previous message
check_previous_message(&tool_call)
};
last_tool_in_chunk = Some(tool_call.clone());
result.push((tool_call, duplicate_type));
}
result
}
// =============================================================================
// Auto-Continue Decision Logic
// =============================================================================
/// Reasons why the streaming loop should auto-continue.
#[derive(Debug, Clone, PartialEq)]
pub enum AutoContinueReason {
/// Tools were executed and we're in autonomous mode
ToolsExecuted,
/// LLM emitted an incomplete (truncated) tool call
IncompleteToolCall,
/// LLM emitted a tool call that wasn't executed
UnexecutedToolCall,
/// Response was truncated due to max_tokens limit
MaxTokensTruncation,
}
/// Determine if the streaming loop should auto-continue.
/// Returns `Some(reason)` if it should continue, `None` otherwise.
pub fn should_auto_continue(
is_autonomous: bool,
any_tool_executed: bool,
has_incomplete_tool_call: bool,
has_unexecuted_tool_call: bool,
was_truncated: bool,
) -> Option<AutoContinueReason> {
if !is_autonomous {
return None;
}
if any_tool_executed {
Some(AutoContinueReason::ToolsExecuted)
} else if has_incomplete_tool_call {
Some(AutoContinueReason::IncompleteToolCall)
} else if has_unexecuted_tool_call {
Some(AutoContinueReason::UnexecutedToolCall)
} else if was_truncated {
Some(AutoContinueReason::MaxTokensTruncation)
} else {
None
}
}
/// Determine if a response is essentially empty (whitespace or timing only)
pub fn is_empty_response(response: &str) -> bool {
response.trim().is_empty()
@@ -474,4 +561,72 @@ mod tests {
assert_eq!(lines.len(), 3);
assert_eq!(lines[0], "line1");
}
#[test]
fn test_deduplicate_tool_calls_no_duplicates() {
let tools = vec![
ToolCall { tool: "shell".to_string(), args: serde_json::json!({"command": "ls"}) },
ToolCall { tool: "read_file".to_string(), args: serde_json::json!({"path": "foo.rs"}) },
];
let result = deduplicate_tool_calls(tools, |_| None);
assert_eq!(result.len(), 2);
assert!(result[0].1.is_none());
assert!(result[1].1.is_none());
}
#[test]
fn test_deduplicate_tool_calls_sequential_duplicate() {
let tools = vec![
ToolCall { tool: "shell".to_string(), args: serde_json::json!({"command": "ls"}) },
ToolCall { tool: "shell".to_string(), args: serde_json::json!({"command": "ls"}) },
];
let result = deduplicate_tool_calls(tools, |_| None);
assert_eq!(result.len(), 2);
assert!(result[0].1.is_none(), "First should not be duplicate");
assert_eq!(result[1].1, Some("DUP IN CHUNK".to_string()));
}
#[test]
fn test_deduplicate_tool_calls_previous_message_duplicate() {
let tools = vec![
ToolCall { tool: "shell".to_string(), args: serde_json::json!({"command": "ls"}) },
];
// Simulate finding a duplicate in previous message
let result = deduplicate_tool_calls(tools, |_| Some("DUP IN MSG".to_string()));
assert_eq!(result.len(), 1);
assert_eq!(result[0].1, Some("DUP IN MSG".to_string()));
}
#[test]
fn test_should_auto_continue_not_autonomous() {
// Never auto-continue in interactive mode
assert_eq!(should_auto_continue(false, true, false, false, false), None);
assert_eq!(should_auto_continue(false, false, true, false, false), None);
}
#[test]
fn test_should_auto_continue_autonomous() {
use AutoContinueReason::*;
// Tools executed
assert_eq!(should_auto_continue(true, true, false, false, false), Some(ToolsExecuted));
// Incomplete tool call
assert_eq!(should_auto_continue(true, false, true, false, false), Some(IncompleteToolCall));
// Unexecuted tool call
assert_eq!(should_auto_continue(true, false, false, true, false), Some(UnexecutedToolCall));
// Max tokens truncation
assert_eq!(should_auto_continue(true, false, false, false, true), Some(MaxTokensTruncation));
// Nothing special - no auto-continue
assert_eq!(should_auto_continue(true, false, false, false, false), None);
}
}