Skip to content

Commit e4b79a7

Browse files
authored
fix(query): keep tool turns active until final stop
Suppresses TurnComplete for intermediate tool-use turns and recoverable max_tokens continuations so chat mode stays active until the model reaches a terminal stop.\n\nVerification:\n- git diff --check\n- cargo test -p claurst-query --manifest-path src-rust/Cargo.toml
1 parent c5f6c32 commit e4b79a7

1 file changed

Lines changed: 40 additions & 6 deletions

File tree

  • src-rust/crates/query/src

src-rust/crates/query/src/lib.rs

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -689,6 +689,14 @@ const MAX_TOKENS_RECOVERY_MSG: &str =
689689
you were doing. Pick up mid-thought if that is where the cut happened. \
690690
Break remaining work into smaller pieces.";
691691

692+
fn should_emit_turn_complete(stop: &str, max_tokens_recovery_count: u32) -> bool {
693+
match stop {
694+
"tool_use" => false,
695+
"max_tokens" => max_tokens_recovery_count >= MAX_TOKENS_RECOVERY_LIMIT,
696+
_ => true,
697+
}
698+
}
699+
692700
// Spinner verbs are imported from claurst_core::spinner
693701

694702
/// Run the agentic query loop.
@@ -1700,12 +1708,14 @@ pub async fn run_query_loop(
17001708
}
17011709
}
17021710

1703-
if let Some(ref tx) = event_tx {
1704-
let _ = tx.send(QueryEvent::TurnComplete {
1705-
turn,
1706-
stop_reason: stop.to_string(),
1707-
usage: Some(usage.clone()),
1708-
});
1711+
if should_emit_turn_complete(stop, max_tokens_recovery_count) {
1712+
if let Some(ref tx) = event_tx {
1713+
let _ = tx.send(QueryEvent::TurnComplete {
1714+
turn,
1715+
stop_reason: stop.to_string(),
1716+
usage: Some(usage.clone()),
1717+
});
1718+
}
17091719
}
17101720

17111721
// Helper closure for firing the Stop hook.
@@ -2459,6 +2469,30 @@ mod tests {
24592469
serde_json::json!(10_000)
24602470
);
24612471
}
2472+
2473+
#[test]
2474+
fn turn_complete_emission_skips_intermediate_tool_turns() {
2475+
assert!(!should_emit_turn_complete("tool_use", 0));
2476+
}
2477+
2478+
#[test]
2479+
fn turn_complete_emission_skips_recoverable_max_tokens_turns() {
2480+
assert!(!should_emit_turn_complete("max_tokens", 0));
2481+
assert!(!should_emit_turn_complete("max_tokens", 1));
2482+
assert!(!should_emit_turn_complete("max_tokens", 2));
2483+
assert!(should_emit_turn_complete(
2484+
"max_tokens",
2485+
MAX_TOKENS_RECOVERY_LIMIT
2486+
));
2487+
}
2488+
2489+
#[test]
2490+
fn turn_complete_emission_keeps_terminal_stop_reasons() {
2491+
assert!(should_emit_turn_complete("end_turn", 0));
2492+
assert!(should_emit_turn_complete("stop_sequence", 0));
2493+
assert!(should_emit_turn_complete("content_filtered", 0));
2494+
assert!(should_emit_turn_complete("unknown_stop", 0));
2495+
}
24622496
}
24632497

24642498
/// Stream handler that forwards events to an unbounded channel.

0 commit comments

Comments
 (0)