1use async_trait::async_trait;
2use futures::{stream::StreamExt, FutureExt};
3use serde_json::{json, Value};
4use std::sync::Arc;
5use std::time::Duration;
6use tokio::sync::mpsc;
7use tokio_util::sync::CancellationToken;
8
9use crate::compaction::{
10 estimate_messages_tokens, CompactionContext, CompactionStrategy, SummarizeCompactionStrategy,
11};
12use crate::event::{HarnessInternalEvent, HarnessUsage, NativeHarnessError, NativeTurnInput};
13use crate::model::{
14 AssistantThinking, CapabilitySupport, ChatMessage, HostedCapability, HostedTool, ModelChunk,
15 ModelClient, ModelClientError, ModelTurnInput,
16};
17use crate::runner::NativeHarness;
18use crate::tools::{
19 bounded::BoundedToolRuntime, ToolFailure, ToolFailureKind, ToolInvocation, ToolOutcome,
20 ToolRuntime, ToolRuntimeError, ToolSpec,
21};
22
23#[derive(Clone)]
29pub struct CompactionPolicy {
30 pub strategy: Arc<dyn CompactionStrategy>,
31 pub model_client: Arc<dyn ModelClient>,
35 pub context_window_tokens: u64,
36}
37
38impl CompactionPolicy {
39 pub fn new(
46 strategy: Arc<dyn CompactionStrategy>,
47 model_client: Arc<dyn ModelClient>,
48 context_window_tokens: u64,
49 ) -> Self {
50 Self {
51 strategy,
52 model_client,
53 context_window_tokens,
54 }
55 }
56
57 pub fn summarizing(model_client: Arc<dyn ModelClient>, context_window_tokens: u64) -> Self {
63 Self::new(
64 Arc::new(SummarizeCompactionStrategy::default()),
65 model_client,
66 context_window_tokens,
67 )
68 }
69}
70
71const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(90);
77
78const DEFAULT_STREAM_MAX_ATTEMPTS: u32 = 6;
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
88#[serde(rename_all = "snake_case")]
89pub enum WebSearchMode {
90 #[default]
92 Off,
93 Auto,
96 Native,
99 Managed,
101}
102
103#[derive(Clone)]
104pub struct AgentLoopHarness<M, R> {
105 model: M,
106 tools: BoundedToolRuntime<R>,
110 max_steps: usize,
111 compaction: Option<CompactionPolicy>,
112 tool_choice: crate::model::ToolChoice,
113 web_search: WebSearchMode,
114 parallel_tool_calls: Option<bool>,
115 stream_idle_timeout: Duration,
116 stream_max_attempts: u32,
117 turn_end_grace: Duration,
118}
119
120impl<M, R: ToolRuntime> AgentLoopHarness<M, R> {
121 pub fn new(model: M, tools: R) -> Self {
122 Self {
123 model,
124 tools: BoundedToolRuntime::new(tools),
125 max_steps: 8,
126 compaction: None,
127 tool_choice: crate::model::ToolChoice::Auto,
128 web_search: WebSearchMode::Off,
129 parallel_tool_calls: None,
130 stream_idle_timeout: DEFAULT_STREAM_IDLE_TIMEOUT,
131 stream_max_attempts: DEFAULT_STREAM_MAX_ATTEMPTS,
132 turn_end_grace: DEFAULT_TURN_END_GRACE,
133 }
134 }
135
136 pub fn with_max_steps(mut self, max_steps: usize) -> Self {
141 self.max_steps = max_steps;
142 self
143 }
144
145 pub fn with_compaction(mut self, policy: CompactionPolicy) -> Self {
151 self.compaction = Some(policy);
152 self
153 }
154
155 pub fn with_tool_choice(mut self, choice: crate::model::ToolChoice) -> Self {
158 self.tool_choice = choice;
159 self
160 }
161
162 pub fn with_web_search(mut self, mode: WebSearchMode) -> Self {
165 self.web_search = mode;
166 self
167 }
168
169 pub fn with_parallel_tool_calls(mut self, parallel: Option<bool>) -> Self {
173 self.parallel_tool_calls = parallel;
174 self
175 }
176
177 pub fn with_stream_resilience(mut self, idle_timeout: Duration, max_attempts: u32) -> Self {
184 self.stream_idle_timeout = idle_timeout;
185 self.stream_max_attempts = max_attempts.max(1);
186 self
187 }
188
189 pub fn with_turn_end_grace(mut self, grace: Duration) -> Self {
204 self.turn_end_grace = grace;
205 self
206 }
207}
208
209#[async_trait]
210impl<M, R> NativeHarness for AgentLoopHarness<M, R>
211where
212 M: ModelClient + Clone + Send + Sync + 'static,
213 R: ToolRuntime + Clone + Send + Sync + 'static,
214{
215 async fn run_turn(
216 &self,
217 input: NativeTurnInput,
218 ) -> Result<mpsc::Receiver<Result<HarnessInternalEvent, NativeHarnessError>>, NativeHarnessError>
219 {
220 let (tx, rx) = mpsc::channel(16);
221 let model = self.model.clone();
222 let tools = self.tools.clone();
223 let max_steps = self.max_steps;
224 let compaction = self.compaction.clone();
225 let tool_choice = self.tool_choice.clone();
226 let web_search = self.web_search;
227 let parallel_tool_calls = self.parallel_tool_calls;
228 let stream_idle_timeout = self.stream_idle_timeout;
229 let stream_max_attempts = self.stream_max_attempts;
230 let turn_end_grace = self.turn_end_grace;
231
232 tokio::spawn(async move {
233 let tx_for_panic = tx.clone();
234 let result = std::panic::AssertUnwindSafe(run_loop(
235 model,
236 tools,
237 RunLoopConfig {
238 max_steps,
239 compaction,
240 tool_choice,
241 web_search,
242 parallel_tool_calls,
243 stream_idle_timeout,
244 stream_max_attempts,
245 turn_end_grace,
246 },
247 input,
248 tx,
249 ))
250 .catch_unwind()
251 .await;
252 if let Err(payload) = result {
253 let detail = panic_payload_to_string(payload.as_ref());
254 tracing::error!(
255 target: "harness::agent_loop",
256 panic = %detail,
257 "native agent loop panicked"
258 );
259 let _ = tx_for_panic
260 .send(Err(NativeHarnessError::Failed(format!(
261 "agent loop panicked: {detail}"
262 ))))
263 .await;
264 }
265 });
266
267 Ok(rx)
268 }
269}
270
271fn cancel_fired(token: Option<&CancellationToken>) -> bool {
273 token.is_some_and(|t| t.is_cancelled())
274}
275
276fn is_silent_stop(text: &str, stop_reason: &str) -> bool {
277 text.trim().is_empty() && matches!(stop_reason, "end_turn" | "max_tokens")
278}
279
280struct RunLoopConfig {
281 max_steps: usize,
282 compaction: Option<CompactionPolicy>,
283 tool_choice: crate::model::ToolChoice,
284 web_search: WebSearchMode,
285 parallel_tool_calls: Option<bool>,
286 stream_idle_timeout: Duration,
287 stream_max_attempts: u32,
288 turn_end_grace: Duration,
292}
293
294pub const DEFAULT_TURN_END_GRACE: Duration = Duration::from_secs(1);
302
303fn resolve_web_search_tools(
304 mode: WebSearchMode,
305 native_support: CapabilitySupport,
306 mut tools: Vec<ToolSpec>,
307) -> Result<(Vec<ToolSpec>, Vec<HostedTool>), String> {
308 let has_managed = tools.iter().any(|tool| tool.name == "web_search");
309 let remove_managed = |tools: &mut Vec<ToolSpec>| {
310 tools.retain(|tool| tool.name != "web_search");
311 };
312
313 match mode {
314 WebSearchMode::Off => {
315 remove_managed(&mut tools);
316 Ok((tools, vec![]))
317 }
318 WebSearchMode::Auto if native_support == CapabilitySupport::Supported => {
319 remove_managed(&mut tools);
320 Ok((tools, vec![HostedTool::WebSearch]))
321 }
322 WebSearchMode::Auto if has_managed => Ok((tools, vec![])),
323 WebSearchMode::Auto => Ok((tools, vec![])),
324 WebSearchMode::Native if native_support == CapabilitySupport::Unsupported => Err(
325 "web search mode is Native, but the model client does not support hosted web search"
326 .into(),
327 ),
328 WebSearchMode::Native => {
329 remove_managed(&mut tools);
330 Ok((tools, vec![HostedTool::WebSearch]))
331 }
332 WebSearchMode::Managed if !has_managed => Err(
333 "web search mode is Managed, but the tool runtime does not provide `web_search`".into(),
334 ),
335 WebSearchMode::Managed => Ok((tools, vec![])),
336 }
337}
338
339async fn run_loop<M, R>(
340 model: M,
341 tools: R,
342 config: RunLoopConfig,
343 input: NativeTurnInput,
344 tx: mpsc::Sender<Result<HarnessInternalEvent, NativeHarnessError>>,
345) where
346 M: ModelClient + Send + Sync,
347 R: ToolRuntime + Clone + Send + Sync + 'static,
348{
349 let system_prompt = input.system_prompt.clone();
350 let cancel_token = input.cancel_token.clone();
351 let context_path = input.context_path.clone();
352 let (tools_snapshot, hosted_tools) = match resolve_web_search_tools(
356 config.web_search,
357 model.hosted_capability(HostedCapability::WebSearch),
358 tools.specs(),
359 ) {
360 Ok(selection) => selection,
361 Err(message) => {
362 let _ = tx
363 .send(Err(NativeHarnessError::ModelBadRequest(message)))
364 .await;
365 return;
366 }
367 };
368 let mut messages: Vec<ChatMessage> = if let Some(ref path) = context_path {
371 crate::context::jsonl::load_context(path).await
372 } else {
373 input.prior_messages
374 };
375 messages.push(ChatMessage::User {
376 content: input.prompt_text,
377 attachments: input.attachments,
378 });
379 let mut ctx_written: usize = match context_path.as_deref() {
383 None => 0,
384 Some(path) => {
385 let start = messages.len() - 1;
386 crate::context::jsonl::append_context(path, &messages[start..]).await;
387 messages.len()
388 }
389 };
390 let mut total_usage = HarnessUsage::default();
394 let mut saw_any_usage = false;
395
396 macro_rules! check_cancel {
402 () => {
403 if cancel_fired(cancel_token.as_ref()) {
404 let _ = tx
405 .send(Ok(HarnessInternalEvent::TurnEnd {
406 stop_reason: "interrupt".into(),
407 usage: saw_any_usage.then(|| total_usage.clone()),
408 pending_tools: vec![],
409 final_messages: if context_path.is_none() {
410 messages.clone()
411 } else {
412 vec![]
413 },
414 }))
415 .await;
416 return;
417 }
418 };
419 }
420
421 for step in 0.. {
422 if config.max_steps != 0 && step >= config.max_steps {
426 break;
427 }
428 check_cancel!();
429 if let Some(policy) = &config.compaction {
435 if policy
436 .strategy
437 .should_compact(&messages, policy.context_window_tokens)
438 {
439 let original_count = messages.len();
440 let original_tokens = estimate_messages_tokens(&messages);
441 let cctx = CompactionContext {
442 system_prompt: system_prompt.clone(),
443 model_client: policy.model_client.clone(),
444 context_window_tokens: policy.context_window_tokens,
445 tools: tools_snapshot.clone(),
446 };
447 match policy.strategy.compact(messages.clone(), &cctx).await {
448 Ok(outcome) => {
449 let compacted_count = outcome.messages.len();
450 let compacted_tokens = estimate_messages_tokens(&outcome.messages);
451 messages = outcome.messages;
452 if let Some(u) = outcome.usage.as_ref() {
458 saw_any_usage = true;
459 total_usage.input_tokens += u.input_tokens;
460 total_usage.output_tokens += u.output_tokens;
461 total_usage.cache_read_input_tokens += u.cache_read_input_tokens;
462 total_usage.cache_creation_input_tokens +=
463 u.cache_creation_input_tokens;
464 total_usage.compaction_input_tokens += u.input_tokens;
465 total_usage.compaction_output_tokens += u.output_tokens;
466 }
467 tracing::info!(
471 target: "harness::compaction",
472 step,
473 original_message_count = original_count,
474 compacted_message_count = compacted_count,
475 original_estimated_tokens = original_tokens,
476 compacted_estimated_tokens = compacted_tokens,
477 context_window_tokens = policy.context_window_tokens,
478 "compaction applied"
479 );
480 if let Some(ref path) = context_path {
482 crate::context::jsonl::rewrite_context(path, &messages).await;
483 ctx_written = messages.len();
484 }
485 if tx
486 .send(Ok(HarnessInternalEvent::CompactionApplied {
487 original_message_count: original_count,
488 compacted_message_count: compacted_count,
489 original_tokens,
490 compacted_tokens,
491 }))
492 .await
493 .is_err()
494 {
495 return;
496 }
497 }
498 Err(e) => {
499 tracing::warn!(
500 target: "harness::compaction",
501 step,
502 error = %e,
503 "compaction skipped; history retained as-is, model call may now fail with context overflow"
504 );
505 }
506 }
507 }
508 }
509
510 const MAX_RETRIES: u32 = 3;
515 const BASE_BACKOFF_MS: u64 = 1_000;
516 const MAX_BACKOFF_MS: u64 = 16_000;
517
518 let model_input = ModelTurnInput {
519 system_prompt: system_prompt.clone(),
520 messages: messages.clone(),
521 tools: tools_snapshot.clone(),
522 hosted_tools: hosted_tools.clone(),
523 tool_choice: config.tool_choice.clone(),
524 parallel_tool_calls: config.parallel_tool_calls,
525 };
526
527 let mut stream_attempt = 0u32;
537 let outcome = 'stream: loop {
538 let stream = {
539 let mut attempt = 0u32;
540 loop {
541 match model.stream(model_input.clone()).await {
542 Ok(s) => break s,
543 Err(e) => {
544 if e.retryable() && attempt < MAX_RETRIES {
545 let delay_ms =
546 (BASE_BACKOFF_MS * (1 << attempt)).min(MAX_BACKOFF_MS);
547 tracing::warn!(
548 attempt,
549 delay_ms,
550 error = %e,
551 "model call failed (retryable) — backing off"
552 );
553 if !backoff_sleep(delay_ms, cancel_token.as_ref()).await {
554 let _ = tx
555 .send(Err(NativeHarnessError::ModelOther(
556 "interrupted during retry backoff".into(),
557 )))
558 .await;
559 return;
560 }
561 attempt += 1;
562 } else {
563 tracing::error!(
566 attempt,
567 error = %e,
568 retryable = e.retryable(),
569 "model call failed — terminating turn"
570 );
571 let _ = tx.send(Err(model_error_to_native(e))).await;
572 return;
573 }
574 }
575 }
576 }
577 };
578
579 match consume_step_stream(
584 stream,
585 &tx,
586 step,
587 cancel_token.as_ref(),
588 config.stream_idle_timeout,
589 )
590 .await
591 {
592 Ok(StepDrain::Complete(o)) => break 'stream o,
593 Ok(StepDrain::Cancelled) => {
594 let _ = tx
595 .send(Ok(HarnessInternalEvent::TurnEnd {
596 stop_reason: "interrupt".into(),
597 usage: saw_any_usage.then(|| total_usage.clone()),
598 pending_tools: vec![],
599 final_messages: if context_path.is_none() {
600 messages.clone()
601 } else {
602 vec![]
603 },
604 }))
605 .await;
606 return;
607 }
608 Err(StepFailure::Model { err, had_progress }) => {
609 if !had_progress
612 && err.retryable()
613 && stream_attempt + 1 < config.stream_max_attempts
614 {
615 let delay_ms =
616 (BASE_BACKOFF_MS * (1 << stream_attempt)).min(MAX_BACKOFF_MS);
617 tracing::warn!(
618 step,
619 stream_attempt,
620 delay_ms,
621 error = %err,
622 "model stream failed before any output — reconnecting"
623 );
624 if !backoff_sleep(delay_ms, cancel_token.as_ref()).await {
625 let _ = tx
626 .send(Err(NativeHarnessError::ModelOther(
627 "interrupted during stream reconnect backoff".into(),
628 )))
629 .await;
630 return;
631 }
632 stream_attempt += 1;
633 continue 'stream;
634 }
635 tracing::error!(
638 step,
639 stream_attempt,
640 error = %err,
641 had_progress,
642 retryable = err.retryable(),
643 "model stream failed — terminating turn"
644 );
645 let _ = tx.send(Err(model_error_to_native(err))).await;
646 return;
647 }
648 Err(StepFailure::ChannelClosed) => return,
649 Err(StepFailure::Fatal(e)) => {
650 let _ = tx.send(Err(e)).await;
651 return;
652 }
653 }
654 };
655
656 if let Some(u) = outcome.usage.as_ref() {
657 saw_any_usage = true;
658 total_usage.input_tokens += u.input_tokens;
659 total_usage.output_tokens += u.output_tokens;
660 total_usage.cache_read_input_tokens += u.cache_read_input_tokens;
661 total_usage.cache_creation_input_tokens += u.cache_creation_input_tokens;
662 }
663
664 match outcome.next {
665 StepNext::Message { text, stop_reason } => {
666 if is_silent_stop(&text, &stop_reason) {
667 let _ = tx
668 .send(Err(NativeHarnessError::ModelOther(format!(
669 "silent_stop: model returned stop_reason={stop_reason} \
670 with empty text and no tool calls"
671 ))))
672 .await;
673 return;
674 }
675 let assistant_text = (!text.trim().is_empty()).then_some(text);
676 messages.push(ChatMessage::Assistant {
677 text: assistant_text,
678 tool_calls: vec![],
679 thinking: outcome.thinking.clone(),
680 usage: outcome.usage.clone(),
681 });
682 if let Some(ref path) = context_path {
684 crate::context::jsonl::append_context(path, &messages[ctx_written..]).await;
685 }
686 let final_msgs = if context_path.is_none() {
689 messages.clone()
690 } else {
691 vec![]
692 };
693 let _ = tx
694 .send(Ok(HarnessInternalEvent::TurnEnd {
695 stop_reason,
696 usage: saw_any_usage.then(|| total_usage.clone()),
697 pending_tools: vec![],
698 final_messages: final_msgs,
699 }))
700 .await;
701 return;
702 }
703 StepNext::ToolCalls {
704 preface,
705 mut invocations,
706 } => {
707 check_cancel!();
708 for inv in &mut invocations {
716 if let Some(repairs) = tools.repair_invocation(inv) {
721 inv.raw_emitted_args = None;
722 tracing::warn!(
723 target: "harness::tool_repair",
724 tool = %inv.name,
725 id = %inv.id,
726 repairs = ?repairs,
727 "schema-guided tool input repair applied"
728 );
729 }
730 }
731 let preface_text = preface.filter(|s| !s.is_empty());
732 messages.push(ChatMessage::Assistant {
740 text: preface_text,
741 tool_calls: invocations.clone(),
742 thinking: outcome.thinking.clone(),
743 usage: outcome.usage.clone(),
744 });
745
746 for inv in &invocations {
752 if tx
753 .send(Ok(HarnessInternalEvent::ToolCall {
754 id: inv.id.clone(),
755 name: inv.name.clone(),
756 input: inv.input.clone(),
757 }))
758 .await
759 .is_err()
760 {
761 return;
762 }
763 }
764
765 let calls = invocations.iter().cloned().map(|inv| {
771 let tools = tools.clone();
772 let cancel_for_task = cancel_token.clone();
773 let invocation_for_task = inv.clone();
774 async move {
775 let call =
776 tools.invoke_cancellable(invocation_for_task, cancel_for_task.as_ref());
777 let outcome = match std::panic::AssertUnwindSafe(call).catch_unwind().await
778 {
779 Ok(outcome) => outcome,
780 Err(payload) => Err(ToolRuntimeError::Runtime(format!(
781 "tool task panicked: {}",
782 panic_payload_to_string(payload.as_ref())
783 ))),
784 };
785 (inv, outcome)
786 }
787 });
788 let join = futures::future::join_all(calls);
789 tokio::pin!(join);
790
791 let (pairs_opt, pending_ids) = if let Some(token) = cancel_token.as_ref() {
796 let grace = input.turn_end_grace.unwrap_or(config.turn_end_grace);
797 tokio::select! {
798 biased;
799 _ = token.cancelled() => {
800 match tokio::time::timeout(grace, &mut join).await {
809 Ok(_) => (None, Vec::new()),
814 Err(_elapsed) => {
815 let pending = invocations.iter().map(|inv| inv.id.clone()).collect();
816 (None, pending)
817 }
818 }
819 },
820 results = &mut join => (Some(results), Vec::new()),
821 }
822 } else {
823 (Some((&mut join).await), Vec::new())
824 };
825 let pairs = match pairs_opt {
826 Some(o) => o,
827 None => {
828 if !pending_ids.is_empty() {
832 tracing::warn!(
833 target: "harness::agent_loop",
834 pending_tools = ?pending_ids,
835 grace_secs = input
836 .turn_end_grace
837 .unwrap_or(config.turn_end_grace)
838 .as_secs_f32(),
839 "turn cancelled with tool futures still unresolved after grace window"
840 );
841 }
842 let _ = tx
843 .send(Ok(HarnessInternalEvent::TurnEnd {
844 stop_reason: "interrupt".into(),
845 usage: saw_any_usage.then(|| total_usage.clone()),
846 pending_tools: pending_ids,
847 final_messages: if context_path.is_none() {
848 messages.clone()
849 } else {
850 vec![]
851 },
852 }))
853 .await;
854 return;
855 }
856 };
857
858 for (inv, outcome) in pairs {
863 let id = inv.id.clone();
864 let outcome = match outcome {
865 Ok(o) => {
866 tracing::info!(
867 target: "harness::tool",
868 step,
869 tool = %inv.name,
870 id = %id,
871 success = o.output.is_ok(),
872 attachments = o.attachments.len(),
873 "tool invocation completed"
874 );
875 o
876 }
877 Err(e) => {
878 let failure = tool_runtime_error_to_failure(
879 &inv,
880 e,
881 tools_snapshot
882 .iter()
883 .find(|s| s.name == inv.name)
884 .map(|s| &s.input_schema),
885 );
886 tracing::warn!(
887 target: "harness::tool",
888 step,
889 tool = %inv.name,
890 id = %id,
891 failure_kind = ?failure.kind,
892 failure_message = %failure.message,
893 "tool invocation failed; returning model-visible ToolResult"
894 );
895 ToolOutcome {
896 output: Err(failure),
897 attachments: vec![],
898 }
899 }
900 };
901 let tool_attachments = outcome.attachments;
902 let output = outcome.output.map_err(|failure| failure.to_string());
903
904 let (tool_content, is_error) = match &output {
910 Ok(value) => (value.to_string(), false),
911 Err(err) => (json!({ "error": err }).to_string(), true),
912 };
913 messages.push(ChatMessage::Tool {
914 tool_call_id: id.clone(),
915 content: tool_content,
916 is_error,
917 attachments: tool_attachments,
918 });
919
920 if tx
921 .send(Ok(HarnessInternalEvent::ToolResult { id, output }))
922 .await
923 .is_err()
924 {
925 return;
926 }
927 }
928 if let Some(ref path) = context_path {
930 crate::context::jsonl::append_context(path, &messages[ctx_written..]).await;
931 ctx_written = messages.len();
932 }
933 }
936 }
937 }
938
939 if let Some(ref path) = context_path {
941 crate::context::jsonl::append_context(path, &messages[ctx_written..]).await;
942 }
943 let final_msgs = if context_path.is_none() {
944 messages
945 } else {
946 vec![]
947 };
948 let _ = tx
949 .send(Ok(HarnessInternalEvent::TurnEnd {
950 stop_reason: "max_turns".into(),
951 usage: saw_any_usage.then(|| total_usage.clone()),
952 pending_tools: vec![],
953 final_messages: final_msgs,
954 }))
955 .await;
956}
957
958async fn backoff_sleep(delay_ms: u64, cancel_token: Option<&CancellationToken>) -> bool {
969 let sleep = tokio::time::sleep(Duration::from_millis(delay_ms));
970 tokio::pin!(sleep);
971 let cancelled = async {
972 if let Some(t) = cancel_token {
973 t.cancelled().await
974 } else {
975 std::future::pending().await
976 }
977 };
978 tokio::select! {
979 _ = &mut sleep => true,
980 _ = cancelled => false,
981 }
982}
983
984fn model_error_to_native(err: ModelClientError) -> NativeHarnessError {
985 match err {
986 ModelClientError::RateLimit(s) => NativeHarnessError::ModelRateLimit(s),
987 ModelClientError::Auth(s) => NativeHarnessError::ModelAuth(s),
988 ModelClientError::ContextOverflow(s) => NativeHarnessError::ModelContextOverflow(s),
989 ModelClientError::BadRequest(s) => NativeHarnessError::ModelBadRequest(s),
990 ModelClientError::ServerError(s) => NativeHarnessError::ModelServerError(s),
991 ModelClientError::Network(s) => NativeHarnessError::ModelNetwork(s),
992 ModelClientError::Other(s) => NativeHarnessError::ModelOther(s),
993 }
994}
995
996fn tool_runtime_error_to_failure(
997 inv: &ToolInvocation,
998 err: ToolRuntimeError,
999 schema: Option<&serde_json::Value>,
1000) -> ToolFailure {
1001 match err {
1002 ToolRuntimeError::Timeout(message) => ToolFailure::new(ToolFailureKind::Timeout, message),
1003 ToolRuntimeError::InvalidInput { tool, message } => {
1004 crate::tools::invalid_input_failure(&tool, message, &inv.input, schema)
1005 }
1006 ToolRuntimeError::UnknownTool(tool) => ToolFailure::new(
1007 ToolFailureKind::InvalidInput,
1008 format!("unknown tool {tool}; choose one of the advertised tools"),
1009 ),
1010 ToolRuntimeError::Runtime(message) => ToolFailure::new(ToolFailureKind::Runtime, message),
1011 }
1012}
1013
1014fn panic_payload_to_string(payload: &(dyn std::any::Any + Send)) -> String {
1015 if let Some(s) = payload.downcast_ref::<&str>() {
1016 (*s).to_string()
1017 } else if let Some(s) = payload.downcast_ref::<String>() {
1018 s.clone()
1019 } else {
1020 "non-string panic payload".into()
1021 }
1022}
1023
1024struct StepOutcome {
1032 next: StepNext,
1033 usage: Option<HarnessUsage>,
1034 thinking: Option<AssistantThinking>,
1035}
1036
1037enum StepDrain {
1042 Complete(StepOutcome),
1043 Cancelled,
1044}
1045
1046enum StepFailure {
1050 Model {
1057 err: ModelClientError,
1058 had_progress: bool,
1059 },
1060 ChannelClosed,
1063 Fatal(NativeHarnessError),
1066}
1067
1068enum StepNext {
1069 Message {
1070 text: String,
1071 stop_reason: String,
1072 },
1073 ToolCalls {
1079 preface: Option<String>,
1080 invocations: Vec<ToolInvocation>,
1081 },
1082}
1083
1084fn stall_failure(idle_timeout: Duration, had_progress: bool) -> StepFailure {
1095 StepFailure::Model {
1096 err: ModelClientError::Network(format!(
1097 "model stream stalled: no output for {}s (connection open but idle)",
1098 idle_timeout.as_secs()
1099 )),
1100 had_progress,
1101 }
1102}
1103
1104async fn consume_step_stream(
1105 mut stream: futures::stream::BoxStream<'static, Result<ModelChunk, ModelClientError>>,
1106 tx: &mpsc::Sender<Result<HarnessInternalEvent, NativeHarnessError>>,
1107 step: usize,
1108 cancel_token: Option<&CancellationToken>,
1109 idle_timeout: Duration,
1110) -> Result<StepDrain, StepFailure> {
1111 let emit_msg_id = format!("msg_native_{step}");
1112 let emit_thinking_id = format!("thinking_native_{step}");
1113 let mut text_buf = String::new();
1114 let mut text_stream_started = false;
1115 let mut thinking_buf = String::new();
1116 let mut thinking_signature: Option<String> = None;
1117 let mut saw_thinking = false;
1118 let mut tool_states: Vec<ToolBuf> = Vec::new();
1119 let mut stop_reason = "end_turn".to_string();
1120 let mut usage: Option<HarnessUsage> = None;
1121 let mut had_progress = false;
1124
1125 loop {
1126 let idle = tokio::time::sleep(idle_timeout);
1133 tokio::pin!(idle);
1134
1135 let item = if let Some(token) = cancel_token {
1141 tokio::select! {
1142 biased;
1143 _ = token.cancelled() => {
1144 return Ok(StepDrain::Cancelled);
1145 }
1146 _ = &mut idle => return Err(stall_failure(idle_timeout, had_progress)),
1147 next = stream.next() => next,
1148 }
1149 } else {
1150 tokio::select! {
1151 _ = &mut idle => return Err(stall_failure(idle_timeout, had_progress)),
1152 next = stream.next() => next,
1153 }
1154 };
1155 let Some(item) = item else { break };
1156 let chunk = match item {
1157 Ok(c) => c,
1158 Err(e) => {
1159 return Err(StepFailure::Model {
1160 err: e,
1161 had_progress,
1162 })
1163 }
1164 };
1165 match chunk {
1166 ModelChunk::TextDelta { msg_id: _, delta } => {
1167 if delta.is_empty() {
1168 continue;
1169 }
1170 text_buf.push_str(&delta);
1171 let mut flush_delta = delta;
1172 if !text_stream_started {
1173 if text_buf.trim().is_empty() {
1174 continue;
1175 }
1176 text_stream_started = true;
1177 flush_delta = text_buf.clone();
1180 }
1181 had_progress = true;
1184 if tx
1188 .send(Ok(HarnessInternalEvent::AssistantTextChunk {
1189 msg_id: emit_msg_id.clone(),
1190 delta: flush_delta,
1191 }))
1192 .await
1193 .is_err()
1194 {
1195 return Err(StepFailure::ChannelClosed);
1196 }
1197 }
1198 ModelChunk::ThinkingDelta {
1199 thinking_id: _,
1200 delta,
1201 signature,
1202 } => {
1203 if let Some(sig) = signature {
1210 if !sig.is_empty() {
1211 thinking_signature = Some(sig);
1212 }
1213 }
1214 if !delta.is_empty() {
1215 saw_thinking = true;
1216 had_progress = true;
1217 thinking_buf.push_str(&delta);
1218 if tx
1219 .send(Ok(HarnessInternalEvent::AssistantThinkingChunk {
1220 msg_id: emit_thinking_id.clone(),
1221 delta,
1222 }))
1223 .await
1224 .is_err()
1225 {
1226 return Err(StepFailure::ChannelClosed);
1227 }
1228 }
1229 }
1230 ModelChunk::ToolCallStart { id, name } => {
1231 had_progress = true;
1237 tool_states.push(ToolBuf {
1238 id,
1239 name,
1240 args_buf: String::new(),
1241 early_input: None,
1242 });
1243 }
1244 ModelChunk::ToolCallInputDelta { id, delta } => {
1245 if let Some(s) = tool_states.iter_mut().find(|s| s.id == id) {
1246 s.args_buf.push_str(&delta);
1247 }
1248 }
1249 ModelChunk::ToolCallEnd { id, input } => {
1250 if let Some(s) = tool_states.iter_mut().find(|s| s.id == id) {
1251 s.early_input = input;
1252 }
1253 }
1254 ModelChunk::Done {
1255 stop_reason: sr,
1256 usage: u,
1257 } => {
1258 stop_reason = sr;
1259 usage = u;
1260 }
1261 }
1262 }
1263
1264 let thinking = if saw_thinking || thinking_signature.is_some() {
1268 Some(AssistantThinking {
1269 text: thinking_buf,
1270 signature: thinking_signature,
1271 })
1272 } else {
1273 None
1274 };
1275
1276 if !tool_states.is_empty() {
1282 let mut invocations = Vec::with_capacity(tool_states.len());
1283 for state in tool_states {
1284 let parsed_input = match state.early_input {
1285 Some(v) => v,
1286 None => {
1287 let trimmed = state.args_buf.trim();
1288 if trimmed.is_empty() {
1289 Value::Object(serde_json::Map::new())
1293 } else {
1294 match serde_json::from_str(trimmed) {
1295 Ok(v) => v,
1296 Err(e) => {
1297 let res = crate::tool_repair::repair_truncated_json(trimmed);
1303 match serde_json::from_str(&res.repaired) {
1304 Ok(v) if res.changed => {
1305 tracing::warn!(
1306 target: "harness::tool_repair",
1307 tool = %state.name,
1308 id = %state.id,
1309 notes = ?res.notes,
1310 "repaired malformed tool arguments"
1311 );
1312 v
1313 }
1314 _ => {
1315 return Err(StepFailure::Fatal(
1316 NativeHarnessError::ModelOther(format!(
1317 "decode tool arguments for {id}: {e}",
1318 id = state.id
1319 )),
1320 ))
1321 }
1322 }
1323 }
1324 }
1325 }
1326 }
1327 };
1328 let raw_emitted_args = raw_args_for_input(&state.args_buf, &parsed_input);
1329 invocations.push(ToolInvocation {
1330 id: state.id,
1331 name: state.name,
1332 input: parsed_input,
1333 raw_emitted_args,
1334 });
1335 }
1336 return Ok(StepDrain::Complete(StepOutcome {
1337 next: StepNext::ToolCalls {
1338 preface: (!text_buf.is_empty()).then_some(text_buf),
1339 invocations,
1340 },
1341 usage,
1342 thinking,
1343 }));
1344 }
1345
1346 Ok(StepDrain::Complete(StepOutcome {
1347 next: StepNext::Message {
1348 text: text_buf,
1349 stop_reason,
1350 },
1351 usage,
1352 thinking,
1353 }))
1354}
1355
1356struct ToolBuf {
1357 id: String,
1358 name: String,
1359 args_buf: String,
1360 early_input: Option<Value>,
1361}
1362
1363fn raw_args_for_input(raw: &str, input: &Value) -> Option<String> {
1364 let trimmed = raw.trim();
1365 if trimmed.is_empty() {
1366 return None;
1367 }
1368 match serde_json::from_str::<Value>(trimmed) {
1369 Ok(parsed) if parsed == *input => Some(trimmed.to_string()),
1370 _ => None,
1371 }
1372}
1373
1374#[cfg(test)]
1375mod tests {
1376 use super::*;
1377 use crate::compaction::{CompactionContext, CompactionError, CompactionStrategy};
1378 use crate::model::{ModelChunk, ModelClient, ModelClientError, ModelResponse};
1379 use crate::tools::{ToolInvocation, ToolOutcome};
1380 use crate::{HarnessInternalEvent, MockToolRuntime, ScriptedModelClient};
1381 use async_trait::async_trait;
1382 use futures::stream::{BoxStream, StreamExt};
1383
1384 fn search_spec() -> crate::tools::ToolSpec {
1385 crate::tools::web_search::web_search_spec()
1386 }
1387
1388 fn ordinary_spec() -> crate::tools::ToolSpec {
1389 crate::tools::ToolSpec {
1390 name: "read".into(),
1391 description: "read".into(),
1392 input_schema: serde_json::json!({"type": "object"}),
1393 }
1394 }
1395
1396 #[test]
1397 fn web_search_mode_routes_exactly_one_search_surface() {
1398 let both = || vec![ordinary_spec(), search_spec()];
1399
1400 let (tools, hosted) =
1401 resolve_web_search_tools(WebSearchMode::Off, CapabilitySupport::Supported, both())
1402 .unwrap();
1403 assert!(tools.iter().all(|tool| tool.name != "web_search"));
1404 assert!(hosted.is_empty());
1405
1406 let (tools, hosted) =
1407 resolve_web_search_tools(WebSearchMode::Auto, CapabilitySupport::Supported, both())
1408 .unwrap();
1409 assert!(tools.iter().all(|tool| tool.name != "web_search"));
1410 assert_eq!(hosted, vec![HostedTool::WebSearch]);
1411
1412 for support in [CapabilitySupport::Unsupported, CapabilitySupport::Unknown] {
1413 let (tools, hosted) =
1414 resolve_web_search_tools(WebSearchMode::Auto, support, both()).unwrap();
1415 assert!(tools.iter().any(|tool| tool.name == "web_search"));
1416 assert!(hosted.is_empty());
1417 }
1418
1419 let (tools, hosted) =
1420 resolve_web_search_tools(WebSearchMode::Native, CapabilitySupport::Unknown, both())
1421 .unwrap();
1422 assert!(tools.iter().all(|tool| tool.name != "web_search"));
1423 assert_eq!(hosted, vec![HostedTool::WebSearch]);
1424
1425 assert!(resolve_web_search_tools(
1426 WebSearchMode::Native,
1427 CapabilitySupport::Unsupported,
1428 both(),
1429 )
1430 .is_err());
1431
1432 let (tools, hosted) =
1433 resolve_web_search_tools(WebSearchMode::Managed, CapabilitySupport::Supported, both())
1434 .unwrap();
1435 assert!(tools.iter().any(|tool| tool.name == "web_search"));
1436 assert!(hosted.is_empty());
1437
1438 assert!(resolve_web_search_tools(
1439 WebSearchMode::Managed,
1440 CapabilitySupport::Supported,
1441 vec![ordinary_spec()],
1442 )
1443 .is_err());
1444 }
1445
1446 #[test]
1447 fn auto_without_any_search_surface_degrades_to_no_search() {
1448 let (tools, hosted) = resolve_web_search_tools(
1449 WebSearchMode::Auto,
1450 CapabilitySupport::Unknown,
1451 vec![ordinary_spec()],
1452 )
1453 .unwrap();
1454 assert_eq!(tools.len(), 1);
1455 assert!(hosted.is_empty());
1456 }
1457 use std::sync::atomic::{AtomicUsize, Ordering};
1458 use std::sync::{Arc, Mutex};
1459
1460 #[derive(Clone)]
1464 struct QueueModelClient {
1465 queue: Arc<Mutex<Vec<ModelResponse>>>,
1466 }
1467
1468 impl QueueModelClient {
1469 fn new(responses: Vec<ModelResponse>) -> Self {
1470 Self {
1471 queue: Arc::new(Mutex::new(responses)),
1472 }
1473 }
1474 }
1475
1476 #[async_trait]
1477 impl ModelClient for QueueModelClient {
1478 fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
1479 CapabilitySupport::Unsupported
1480 }
1481
1482 async fn stream(
1483 &self,
1484 _input: ModelTurnInput,
1485 ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
1486 {
1487 let mut q = self.queue.lock().unwrap();
1488 if q.is_empty() {
1489 return Err(ModelClientError::Other("queue exhausted".into()));
1490 }
1491 let response = q.remove(0);
1492 let chunks = response_to_chunks(response);
1493 Ok(futures::stream::iter(chunks.into_iter().map(Ok)).boxed())
1494 }
1495 }
1496
1497 fn response_to_chunks(response: ModelResponse) -> Vec<ModelChunk> {
1501 match response {
1502 ModelResponse::Message {
1503 text,
1504 stop_reason,
1505 usage,
1506 } => {
1507 let mut out = Vec::new();
1508 if !text.is_empty() {
1509 out.push(ModelChunk::TextDelta {
1510 msg_id: "queue_msg".into(),
1511 delta: text,
1512 });
1513 }
1514 out.push(ModelChunk::Done { stop_reason, usage });
1515 out
1516 }
1517 ModelResponse::ToolCall {
1518 preface,
1519 invocation,
1520 usage,
1521 } => {
1522 let mut out = Vec::new();
1523 if let Some(p) = preface {
1524 if !p.is_empty() {
1525 out.push(ModelChunk::TextDelta {
1526 msg_id: "queue_msg".into(),
1527 delta: p,
1528 });
1529 }
1530 }
1531 out.push(ModelChunk::ToolCallStart {
1532 id: invocation.id.clone(),
1533 name: invocation.name.clone(),
1534 });
1535 out.push(ModelChunk::ToolCallEnd {
1536 id: invocation.id.clone(),
1537 input: Some(invocation.input.clone()),
1538 });
1539 out.push(ModelChunk::Done {
1540 stop_reason: "end_turn".into(),
1541 usage,
1542 });
1543 out
1544 }
1545 }
1546 }
1547
1548 fn usage(input: u64, output: u64, cache_read: u64) -> HarnessUsage {
1549 HarnessUsage {
1550 input_tokens: input,
1551 output_tokens: output,
1552 cache_read_input_tokens: cache_read,
1553 cache_creation_input_tokens: 0,
1554 compaction_input_tokens: 0,
1555 compaction_output_tokens: 0,
1556 }
1557 }
1558
1559 #[tokio::test]
1560 async fn agent_loop_accumulates_usage_across_steps() {
1561 let model = QueueModelClient::new(vec![
1563 ModelResponse::ToolCall {
1564 preface: None,
1565 invocation: ToolInvocation {
1566 id: "tc_1".into(),
1567 name: "bash".into(),
1568 input: serde_json::json!({"command": "pwd"}),
1569 raw_emitted_args: None,
1570 },
1571 usage: Some(usage(10, 5, 0)),
1572 },
1573 ModelResponse::Message {
1574 text: "done".into(),
1575 stop_reason: "end_turn".into(),
1576 usage: Some(usage(20, 15, 4)),
1577 },
1578 ]);
1579 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1580 let mut rx = harness
1581 .run_turn(NativeTurnInput {
1582 prompt_text: "pwd".into(),
1583 system_prompt: None,
1584 attachments: vec![],
1585 cancel_token: None,
1586 prior_messages: vec![],
1587 context_path: None,
1588 turn_end_grace: None,
1589 })
1590 .await
1591 .unwrap();
1592
1593 let mut final_usage = None;
1595 while let Some(item) = rx.recv().await {
1596 if let HarnessInternalEvent::TurnEnd { usage: u, .. } = item.unwrap() {
1597 final_usage = u;
1598 break;
1599 }
1600 }
1601 let u = final_usage.expect("TurnEnd carried usage");
1602 assert_eq!(u.input_tokens, 30);
1603 assert_eq!(u.output_tokens, 20);
1604 assert_eq!(u.cache_read_input_tokens, 4);
1605 }
1606
1607 #[tokio::test]
1608 async fn agent_loop_errors_on_silent_stop() {
1609 let model = QueueModelClient::new(vec![ModelResponse::Message {
1610 text: " \n".into(),
1611 stop_reason: "end_turn".into(),
1612 usage: None,
1613 }]);
1614 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1615 let mut rx = harness
1616 .run_turn(NativeTurnInput {
1617 prompt_text: "say something".into(),
1618 system_prompt: None,
1619 attachments: vec![],
1620 cancel_token: None,
1621 prior_messages: vec![],
1622 context_path: None,
1623 turn_end_grace: None,
1624 })
1625 .await
1626 .unwrap();
1627
1628 match rx.recv().await.unwrap() {
1629 Err(NativeHarnessError::ModelOther(msg)) => {
1630 assert!(msg.contains("silent_stop"));
1631 assert!(msg.contains("stop_reason=end_turn"));
1632 }
1633 other => panic!("expected silent_stop model error, got {other:?}"),
1634 }
1635 }
1636
1637 #[tokio::test]
1638 async fn agent_loop_errors_on_empty_max_tokens_stop() {
1639 let model = QueueModelClient::new(vec![ModelResponse::Message {
1640 text: "".into(),
1641 stop_reason: "max_tokens".into(),
1642 usage: None,
1643 }]);
1644 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1645 let mut rx = harness
1646 .run_turn(NativeTurnInput {
1647 prompt_text: "think".into(),
1648 system_prompt: None,
1649 attachments: vec![],
1650 cancel_token: None,
1651 prior_messages: vec![],
1652 context_path: None,
1653 turn_end_grace: None,
1654 })
1655 .await
1656 .unwrap();
1657
1658 match rx.recv().await.unwrap() {
1659 Err(NativeHarnessError::ModelOther(msg)) => {
1660 assert!(msg.contains("silent_stop"));
1661 assert!(msg.contains("stop_reason=max_tokens"));
1662 }
1663 other => panic!("expected silent_stop model error, got {other:?}"),
1664 }
1665 }
1666
1667 #[tokio::test]
1668 async fn agent_loop_turn_end_usage_is_none_when_no_step_reported() {
1669 let model = QueueModelClient::new(vec![ModelResponse::Message {
1671 text: "ok".into(),
1672 stop_reason: "end_turn".into(),
1673 usage: None,
1674 }]);
1675 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1676 let mut rx = harness
1677 .run_turn(NativeTurnInput {
1678 prompt_text: "noop".into(),
1679 system_prompt: None,
1680 attachments: vec![],
1681 cancel_token: None,
1682 prior_messages: vec![],
1683 context_path: None,
1684 turn_end_grace: None,
1685 })
1686 .await
1687 .unwrap();
1688 let mut saw_usage = None;
1689 while let Some(item) = rx.recv().await {
1690 if let HarnessInternalEvent::TurnEnd { usage, .. } = item.unwrap() {
1691 saw_usage = Some(usage);
1692 break;
1693 }
1694 }
1695 assert_eq!(saw_usage.unwrap(), None);
1696 }
1697
1698 #[derive(Clone)]
1703 struct StreamingFakeClient {
1704 chunks_per_call: Arc<Mutex<Vec<Vec<ModelChunk>>>>,
1705 }
1706
1707 impl StreamingFakeClient {
1708 fn new(per_call: Vec<Vec<ModelChunk>>) -> Self {
1709 Self {
1710 chunks_per_call: Arc::new(Mutex::new(per_call)),
1711 }
1712 }
1713 }
1714
1715 #[async_trait]
1716 impl ModelClient for StreamingFakeClient {
1717 fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
1718 CapabilitySupport::Unsupported
1719 }
1720
1721 async fn stream(
1722 &self,
1723 _input: ModelTurnInput,
1724 ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
1725 {
1726 let mut bucket = self.chunks_per_call.lock().unwrap();
1727 if bucket.is_empty() {
1728 return Err(ModelClientError::Other("queue exhausted".into()));
1729 }
1730 let chunks = bucket.remove(0);
1731 Ok(futures::stream::iter(chunks.into_iter().map(Ok)).boxed())
1732 }
1733 }
1734
1735 #[tokio::test]
1736 async fn agent_loop_forwards_token_chunks_to_harness_output() {
1737 let model = StreamingFakeClient::new(vec![vec![
1738 ModelChunk::TextDelta {
1739 msg_id: "remote_msg".into(),
1740 delta: "Hel".into(),
1741 },
1742 ModelChunk::TextDelta {
1743 msg_id: "remote_msg".into(),
1744 delta: "lo ".into(),
1745 },
1746 ModelChunk::TextDelta {
1747 msg_id: "remote_msg".into(),
1748 delta: "world".into(),
1749 },
1750 ModelChunk::Done {
1751 stop_reason: "end_turn".into(),
1752 usage: None,
1753 },
1754 ]]);
1755 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1756 let mut rx = harness
1757 .run_turn(NativeTurnInput {
1758 prompt_text: "hi".into(),
1759 system_prompt: None,
1760 attachments: vec![],
1761 cancel_token: None,
1762 prior_messages: vec![],
1763 context_path: None,
1764 turn_end_grace: None,
1765 })
1766 .await
1767 .unwrap();
1768
1769 let mut deltas: Vec<String> = Vec::new();
1770 let mut saw_end = false;
1771 while let Some(item) = rx.recv().await {
1772 match item.unwrap() {
1773 HarnessInternalEvent::AssistantTextChunk { msg_id, delta } => {
1774 assert_eq!(msg_id, "msg_native_0");
1778 deltas.push(delta);
1779 }
1780 HarnessInternalEvent::TurnEnd { stop_reason, .. } => {
1781 assert_eq!(stop_reason, "end_turn");
1782 saw_end = true;
1783 break;
1784 }
1785 other => panic!("unexpected event: {other:?}"),
1786 }
1787 }
1788 assert_eq!(deltas, vec!["Hel", "lo ", "world"]);
1789 assert!(saw_end);
1790 }
1791
1792 #[tokio::test]
1793 async fn agent_loop_streaming_tool_call_then_summary() {
1794 let model = StreamingFakeClient::new(vec![
1801 vec![
1802 ModelChunk::TextDelta {
1803 msg_id: "r1".into(),
1804 delta: "running ".into(),
1805 },
1806 ModelChunk::ToolCallStart {
1807 id: "tc_1".into(),
1808 name: "bash".into(),
1809 },
1810 ModelChunk::ToolCallInputDelta {
1811 id: "tc_1".into(),
1812 delta: "{\"command\":".into(),
1813 },
1814 ModelChunk::ToolCallInputDelta {
1815 id: "tc_1".into(),
1816 delta: "\"pwd\"}".into(),
1817 },
1818 ModelChunk::ToolCallEnd {
1819 id: "tc_1".into(),
1820 input: None,
1821 },
1822 ModelChunk::Done {
1823 stop_reason: "end_turn".into(),
1824 usage: None,
1825 },
1826 ],
1827 vec![
1828 ModelChunk::TextDelta {
1829 msg_id: "r2".into(),
1830 delta: "done".into(),
1831 },
1832 ModelChunk::Done {
1833 stop_reason: "end_turn".into(),
1834 usage: None,
1835 },
1836 ],
1837 ]);
1838 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1839 let mut rx = harness
1840 .run_turn(NativeTurnInput {
1841 prompt_text: "pwd".into(),
1842 system_prompt: None,
1843 attachments: vec![],
1844 cancel_token: None,
1845 prior_messages: vec![],
1846 context_path: None,
1847 turn_end_grace: None,
1848 })
1849 .await
1850 .unwrap();
1851
1852 let ev = rx.recv().await.unwrap().unwrap();
1859 assert!(matches!(
1860 ev,
1861 HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "running "
1862 ));
1863 let ev = rx.recv().await.unwrap().unwrap();
1864 let HarnessInternalEvent::ToolCall { name, input, .. } = ev else {
1865 panic!("expected ToolCall");
1866 };
1867 assert_eq!(name, "bash");
1868 assert_eq!(input["command"], "pwd");
1869 let ev = rx.recv().await.unwrap().unwrap();
1870 assert!(matches!(ev, HarnessInternalEvent::ToolResult { .. }));
1871 let ev = rx.recv().await.unwrap().unwrap();
1872 assert!(matches!(
1873 ev,
1874 HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "done"
1875 ));
1876 let ev = rx.recv().await.unwrap().unwrap();
1877 assert!(matches!(ev, HarnessInternalEvent::TurnEnd { .. }));
1878 }
1879
1880 #[tokio::test]
1881 async fn agent_loop_repairs_truncated_tool_arguments() {
1882 let model = StreamingFakeClient::new(vec![
1886 vec![
1887 ModelChunk::ToolCallStart {
1888 id: "tc_trunc".into(),
1889 name: "bash".into(),
1890 },
1891 ModelChunk::ToolCallInputDelta {
1892 id: "tc_trunc".into(),
1893 delta: r#"{"command":"pwd""#.into(), },
1895 ModelChunk::ToolCallEnd {
1896 id: "tc_trunc".into(),
1897 input: None,
1898 },
1899 ModelChunk::Done {
1900 stop_reason: "tool_use".into(),
1901 usage: None,
1902 },
1903 ],
1904 vec![
1905 ModelChunk::TextDelta {
1906 msg_id: "r2".into(),
1907 delta: "done".into(),
1908 },
1909 ModelChunk::Done {
1910 stop_reason: "end_turn".into(),
1911 usage: None,
1912 },
1913 ],
1914 ]);
1915 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
1916 let mut rx = harness
1917 .run_turn(NativeTurnInput {
1918 prompt_text: "pwd".into(),
1919 system_prompt: None,
1920 attachments: vec![],
1921 cancel_token: None,
1922 prior_messages: vec![],
1923 context_path: None,
1924 turn_end_grace: None,
1925 })
1926 .await
1927 .unwrap();
1928
1929 let mut saw_tool_call = false;
1930 let mut saw_turn_end = false;
1931 while let Some(item) = rx.recv().await {
1932 match item.expect("turn must not fail on truncated args") {
1933 HarnessInternalEvent::ToolCall { name, input, .. } => {
1934 assert_eq!(name, "bash");
1935 assert_eq!(input["command"], "pwd", "repaired args reach the wire");
1936 saw_tool_call = true;
1937 }
1938 HarnessInternalEvent::TurnEnd { .. } => {
1939 saw_turn_end = true;
1940 break;
1941 }
1942 _ => {}
1943 }
1944 }
1945 assert!(saw_tool_call, "expected ToolCall with repaired input");
1946 assert!(saw_turn_end);
1947 }
1948
1949 #[derive(Clone)]
1952 struct ProbeToolRuntime {
1953 seen_input: Arc<Mutex<Option<Value>>>,
1954 }
1955
1956 #[async_trait]
1957 impl ToolRuntime for ProbeToolRuntime {
1958 fn specs(&self) -> Vec<crate::tools::ToolSpec> {
1959 vec![crate::tools::ToolSpec {
1960 name: "probe".into(),
1961 description: "records its input".into(),
1962 input_schema: serde_json::json!({
1963 "type": "object",
1964 "properties": {
1965 "pattern": {"type": "string"},
1966 "literal": {"type": "boolean"},
1967 "limit": {"type": "integer"}
1968 },
1969 "required": ["pattern"]
1970 }),
1971 }]
1972 }
1973
1974 async fn invoke(
1975 &self,
1976 invocation: ToolInvocation,
1977 ) -> Result<ToolOutcome, ToolRuntimeError> {
1978 *self.seen_input.lock().unwrap() = Some(invocation.input);
1979 Ok(ToolOutcome {
1980 output: Ok(r#"{"ok":true}"#.into()),
1981 attachments: vec![],
1982 })
1983 }
1984 }
1985
1986 #[tokio::test]
1987 async fn agent_loop_applies_schema_repair_before_dispatch() {
1988 let model = StreamingFakeClient::new(vec![
1992 vec![
1993 ModelChunk::ToolCallStart {
1994 id: "tc_shape".into(),
1995 name: "probe".into(),
1996 },
1997 ModelChunk::ToolCallEnd {
1998 id: "tc_shape".into(),
1999 input: Some(json!({"pattern": "x", "literal": "true", "limit": "30"})),
2000 },
2001 ModelChunk::Done {
2002 stop_reason: "tool_use".into(),
2003 usage: None,
2004 },
2005 ],
2006 vec![
2007 ModelChunk::TextDelta {
2008 msg_id: "r2".into(),
2009 delta: "done".into(),
2010 },
2011 ModelChunk::Done {
2012 stop_reason: "end_turn".into(),
2013 usage: None,
2014 },
2015 ],
2016 ]);
2017 let seen_input = Arc::new(Mutex::new(None));
2018 let tools = ProbeToolRuntime {
2019 seen_input: seen_input.clone(),
2020 };
2021 let harness = AgentLoopHarness::new(model, tools);
2022 let mut rx = harness
2023 .run_turn(NativeTurnInput {
2024 prompt_text: "go".into(),
2025 system_prompt: None,
2026 attachments: vec![],
2027 cancel_token: None,
2028 prior_messages: vec![],
2029 context_path: None,
2030 turn_end_grace: None,
2031 })
2032 .await
2033 .unwrap();
2034
2035 let mut wire_input: Option<Value> = None;
2036 let mut history: Option<Vec<ChatMessage>> = None;
2037 while let Some(item) = rx.recv().await {
2038 match item.unwrap() {
2039 HarnessInternalEvent::ToolCall { input, .. } => wire_input = Some(input),
2040 HarnessInternalEvent::TurnEnd { final_messages, .. } => {
2041 history = Some(final_messages);
2042 break;
2043 }
2044 _ => {}
2045 }
2046 }
2047 let repaired = json!({"pattern": "x", "literal": true, "limit": 30});
2048 assert_eq!(seen_input.lock().unwrap().clone().unwrap(), repaired);
2050 assert_eq!(wire_input.unwrap(), repaired);
2051 let history = history.unwrap();
2052 let assistant_tool_calls = history
2053 .iter()
2054 .find_map(|m| match m {
2055 ChatMessage::Assistant { tool_calls, .. } if !tool_calls.is_empty() => {
2056 Some(tool_calls.clone())
2057 }
2058 _ => None,
2059 })
2060 .expect("assistant message with tool_calls in history");
2061 assert_eq!(assistant_tool_calls[0].input, repaired);
2062 }
2063
2064 #[derive(Clone)]
2065 struct TimeoutToolRuntime;
2066
2067 #[async_trait]
2068 impl ToolRuntime for TimeoutToolRuntime {
2069 fn specs(&self) -> Vec<crate::tools::ToolSpec> {
2070 vec![crate::tools::ToolSpec {
2071 name: "slow".into(),
2072 description: "always times out".into(),
2073 input_schema: serde_json::json!({"type": "object"}),
2074 }]
2075 }
2076
2077 async fn invoke(
2078 &self,
2079 _invocation: ToolInvocation,
2080 ) -> Result<ToolOutcome, ToolRuntimeError> {
2081 Err(ToolRuntimeError::Timeout("tool timed out after 1s".into()))
2082 }
2083 }
2084
2085 #[tokio::test]
2086 async fn agent_loop_tool_timeout_is_model_observable_result() {
2087 let model = StreamingFakeClient::new(vec![
2088 vec![
2089 ModelChunk::ToolCallStart {
2090 id: "tc_timeout".into(),
2091 name: "slow".into(),
2092 },
2093 ModelChunk::ToolCallEnd {
2094 id: "tc_timeout".into(),
2095 input: Some(json!({})),
2096 },
2097 ModelChunk::Done {
2098 stop_reason: "tool_use".into(),
2099 usage: None,
2100 },
2101 ],
2102 vec![
2103 ModelChunk::TextDelta {
2104 msg_id: "r2".into(),
2105 delta: "recovered".into(),
2106 },
2107 ModelChunk::Done {
2108 stop_reason: "end_turn".into(),
2109 usage: None,
2110 },
2111 ],
2112 ]);
2113 let harness = AgentLoopHarness::new(model, TimeoutToolRuntime);
2114 let mut rx = harness
2115 .run_turn(NativeTurnInput {
2116 prompt_text: "run slow".into(),
2117 system_prompt: None,
2118 attachments: vec![],
2119 cancel_token: None,
2120 prior_messages: vec![],
2121 context_path: None,
2122 turn_end_grace: None,
2123 })
2124 .await
2125 .unwrap();
2126
2127 assert!(matches!(
2128 rx.recv().await.unwrap().unwrap(),
2129 HarnessInternalEvent::ToolCall { .. }
2130 ));
2131 match rx.recv().await.unwrap().unwrap() {
2132 HarnessInternalEvent::ToolResult { output, .. } => {
2133 let err = output.unwrap_err();
2134 assert!(err.contains("Timeout"));
2135 assert!(err.contains("tool timed out"));
2136 }
2137 other => panic!("expected timeout ToolResult, got {other:?}"),
2138 }
2139 assert!(matches!(
2140 rx.recv().await.unwrap().unwrap(),
2141 HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "recovered"
2142 ));
2143 assert!(matches!(
2144 rx.recv().await.unwrap().unwrap(),
2145 HarnessInternalEvent::TurnEnd { ref stop_reason, .. } if stop_reason == "end_turn"
2146 ));
2147 }
2148
2149 #[derive(Clone)]
2150 struct RuntimeErrorToolRuntime;
2151
2152 #[async_trait]
2153 impl ToolRuntime for RuntimeErrorToolRuntime {
2154 fn specs(&self) -> Vec<crate::tools::ToolSpec> {
2155 vec![crate::tools::ToolSpec {
2156 name: "flaky".into(),
2157 description: "always returns a runtime failure".into(),
2158 input_schema: serde_json::json!({"type": "object"}),
2159 }]
2160 }
2161
2162 async fn invoke(
2163 &self,
2164 _invocation: ToolInvocation,
2165 ) -> Result<ToolOutcome, ToolRuntimeError> {
2166 Err(ToolRuntimeError::Runtime(
2167 "sandbox exec stream closed".into(),
2168 ))
2169 }
2170 }
2171
2172 #[tokio::test]
2173 async fn agent_loop_tool_runtime_error_is_model_observable_result() {
2174 let model = StreamingFakeClient::new(vec![
2175 vec![
2176 ModelChunk::ToolCallStart {
2177 id: "tc_runtime".into(),
2178 name: "flaky".into(),
2179 },
2180 ModelChunk::ToolCallEnd {
2181 id: "tc_runtime".into(),
2182 input: Some(json!({})),
2183 },
2184 ModelChunk::Done {
2185 stop_reason: "tool_use".into(),
2186 usage: None,
2187 },
2188 ],
2189 vec![
2190 ModelChunk::TextDelta {
2191 msg_id: "r2".into(),
2192 delta: "recovered".into(),
2193 },
2194 ModelChunk::Done {
2195 stop_reason: "end_turn".into(),
2196 usage: None,
2197 },
2198 ],
2199 ]);
2200 let harness = AgentLoopHarness::new(model, RuntimeErrorToolRuntime);
2201 let mut rx = harness
2202 .run_turn(NativeTurnInput {
2203 prompt_text: "run flaky".into(),
2204 system_prompt: None,
2205 attachments: vec![],
2206 cancel_token: None,
2207 prior_messages: vec![],
2208 context_path: None,
2209 turn_end_grace: None,
2210 })
2211 .await
2212 .unwrap();
2213
2214 assert!(matches!(
2215 rx.recv().await.unwrap().unwrap(),
2216 HarnessInternalEvent::ToolCall { .. }
2217 ));
2218 match rx.recv().await.unwrap().unwrap() {
2219 HarnessInternalEvent::ToolResult { id, output } => {
2220 assert_eq!(id, "tc_runtime");
2221 let err = output.unwrap_err();
2222 assert!(err.contains("Runtime"));
2223 assert!(err.contains("sandbox exec stream closed"));
2224 }
2225 other => panic!("expected runtime-error ToolResult, got {other:?}"),
2226 }
2227 assert!(matches!(
2228 rx.recv().await.unwrap().unwrap(),
2229 HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "recovered"
2230 ));
2231 assert!(matches!(
2232 rx.recv().await.unwrap().unwrap(),
2233 HarnessInternalEvent::TurnEnd { ref stop_reason, .. } if stop_reason == "end_turn"
2234 ));
2235 }
2236
2237 #[tokio::test]
2238 async fn agent_loop_invalid_tool_input_is_model_observable_and_bounded() {
2239 let huge_content = "x".repeat(20_000);
2240 let model = StreamingFakeClient::new(vec![
2241 vec![
2242 ModelChunk::ToolCallStart {
2243 id: "tc_bad_write".into(),
2244 name: "write".into(),
2245 },
2246 ModelChunk::ToolCallEnd {
2247 id: "tc_bad_write".into(),
2248 input: Some(json!({"content": huge_content})),
2249 },
2250 ModelChunk::Done {
2251 stop_reason: "tool_use".into(),
2252 usage: None,
2253 },
2254 ],
2255 vec![
2256 ModelChunk::TextDelta {
2257 msg_id: "r2".into(),
2258 delta: "recovered".into(),
2259 },
2260 ModelChunk::Done {
2261 stop_reason: "end_turn".into(),
2262 usage: None,
2263 },
2264 ],
2265 ]);
2266 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
2267 let mut rx = harness
2268 .run_turn(NativeTurnInput {
2269 prompt_text: "write file".into(),
2270 system_prompt: None,
2271 attachments: vec![],
2272 cancel_token: None,
2273 prior_messages: vec![],
2274 context_path: None,
2275 turn_end_grace: None,
2276 })
2277 .await
2278 .unwrap();
2279
2280 assert!(matches!(
2281 rx.recv().await.unwrap().unwrap(),
2282 HarnessInternalEvent::ToolCall { .. }
2283 ));
2284 match rx.recv().await.unwrap().unwrap() {
2285 HarnessInternalEvent::ToolResult { output, .. } => {
2286 let err = output.unwrap_err();
2287 assert!(err.contains("The write tool was called with invalid arguments"));
2290 assert!(err.contains("missing required field `path`"), "{err}");
2291 assert!(err.contains("Received fields: content"));
2292 assert!(err.contains("string(20000 chars"));
2293 assert!(
2294 err.contains("Expected shape"),
2295 "teaching example missing: {err}"
2296 );
2297 assert!(
2298 !err.contains(&"x".repeat(2000)),
2299 "error should not echo full content"
2300 );
2301 }
2302 other => panic!("expected invalid-input ToolResult, got {other:?}"),
2303 }
2304 assert!(matches!(
2305 rx.recv().await.unwrap().unwrap(),
2306 HarnessInternalEvent::AssistantTextChunk { ref delta, .. } if delta == "recovered"
2307 ));
2308 }
2309
2310 struct CountingCompactionStrategy {
2314 calls: Arc<AtomicUsize>,
2315 }
2316
2317 #[async_trait]
2318 impl CompactionStrategy for CountingCompactionStrategy {
2319 fn should_compact(&self, _messages: &[ChatMessage], _context_window_tokens: u64) -> bool {
2320 true
2321 }
2322
2323 async fn compact(
2324 &self,
2325 _messages: Vec<ChatMessage>,
2326 _ctx: &CompactionContext,
2327 ) -> Result<crate::compaction::CompactionOutcome, CompactionError> {
2328 self.calls.fetch_add(1, Ordering::SeqCst);
2329 Ok(crate::compaction::CompactionOutcome {
2332 messages: vec![ChatMessage::User {
2333 content: "<conversation-summary>FOLDED</conversation-summary>".into(),
2334 attachments: vec![],
2335 }],
2336 usage: None,
2337 })
2338 }
2339 }
2340
2341 struct UsageReportingCompactionStrategy {
2345 invoked: Arc<AtomicUsize>,
2346 per_call_usage: HarnessUsage,
2347 }
2348
2349 #[async_trait]
2350 impl CompactionStrategy for UsageReportingCompactionStrategy {
2351 fn should_compact(&self, _: &[ChatMessage], _: u64) -> bool {
2352 self.invoked.load(Ordering::SeqCst) == 0
2356 }
2357 async fn compact(
2358 &self,
2359 messages: Vec<ChatMessage>,
2360 _ctx: &CompactionContext,
2361 ) -> Result<crate::compaction::CompactionOutcome, CompactionError> {
2362 self.invoked.fetch_add(1, Ordering::SeqCst);
2363 Ok(crate::compaction::CompactionOutcome {
2364 messages,
2365 usage: Some(self.per_call_usage.clone()),
2366 })
2367 }
2368 }
2369
2370 #[tokio::test]
2371 async fn agent_loop_attributes_compaction_usage_to_subbucket_and_total() {
2372 let model = StreamingFakeClient::new(vec![vec![
2376 ModelChunk::TextDelta {
2377 msg_id: "m".into(),
2378 delta: "done".into(),
2379 },
2380 ModelChunk::Done {
2381 stop_reason: "end_turn".into(),
2382 usage: Some(usage(100, 30, 0)),
2383 },
2384 ]]);
2385 let invoked = Arc::new(AtomicUsize::new(0));
2386 let strategy = UsageReportingCompactionStrategy {
2387 invoked: invoked.clone(),
2388 per_call_usage: usage(50, 20, 0),
2389 };
2390 let policy = CompactionPolicy::new(
2391 Arc::new(strategy),
2392 Arc::new(ScriptedModelClient),
2393 1, );
2395 let harness = AgentLoopHarness::new(model, MockToolRuntime::new()).with_compaction(policy);
2396 let mut rx = harness
2397 .run_turn(NativeTurnInput {
2398 prompt_text: "hi".into(),
2399 system_prompt: None,
2400 attachments: vec![],
2401 cancel_token: None,
2402 prior_messages: vec![],
2403 context_path: None,
2404 turn_end_grace: None,
2405 })
2406 .await
2407 .unwrap();
2408 let mut final_usage = None;
2409 while let Some(item) = rx.recv().await {
2410 if let HarnessInternalEvent::TurnEnd { usage, .. } = item.unwrap() {
2411 final_usage = usage;
2412 break;
2413 }
2414 }
2415 assert_eq!(invoked.load(Ordering::SeqCst), 1);
2416 let u = final_usage.expect("TurnEnd carried usage");
2417 assert_eq!(u.input_tokens, 150);
2419 assert_eq!(u.output_tokens, 50);
2420 assert_eq!(u.compaction_input_tokens, 50);
2422 assert_eq!(u.compaction_output_tokens, 20);
2423 }
2424
2425 #[derive(Clone)]
2429 struct RecordingFakeClient {
2430 last_messages: Arc<Mutex<Option<Vec<ChatMessage>>>>,
2431 chunks: Vec<ModelChunk>,
2432 }
2433
2434 #[async_trait]
2435 impl ModelClient for RecordingFakeClient {
2436 fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
2437 CapabilitySupport::Unsupported
2438 }
2439
2440 async fn stream(
2441 &self,
2442 input: ModelTurnInput,
2443 ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
2444 {
2445 *self.last_messages.lock().unwrap() = Some(input.messages);
2446 Ok(futures::stream::iter(self.chunks.clone().into_iter().map(Ok)).boxed())
2447 }
2448 }
2449
2450 #[tokio::test]
2451 async fn agent_loop_invokes_compaction_between_steps() {
2452 let calls = Arc::new(AtomicUsize::new(0));
2453 let last_messages = Arc::new(Mutex::new(None::<Vec<ChatMessage>>));
2454 let model = RecordingFakeClient {
2455 last_messages: last_messages.clone(),
2456 chunks: vec![
2457 ModelChunk::TextDelta {
2458 msg_id: "m".into(),
2459 delta: "done".into(),
2460 },
2461 ModelChunk::Done {
2462 stop_reason: "end_turn".into(),
2463 usage: None,
2464 },
2465 ],
2466 };
2467 let summary_client: Arc<dyn ModelClient> = Arc::new(ScriptedModelClient);
2471 let policy = CompactionPolicy::new(
2472 Arc::new(CountingCompactionStrategy {
2473 calls: calls.clone(),
2474 }),
2475 summary_client,
2476 1, );
2478
2479 let harness = AgentLoopHarness::new(model, MockToolRuntime::new()).with_compaction(policy);
2480 let mut rx = harness
2481 .run_turn(NativeTurnInput {
2482 prompt_text: "hello".into(),
2483 system_prompt: None,
2484 attachments: vec![],
2485 cancel_token: None,
2486 prior_messages: vec![],
2487 context_path: None,
2488 turn_end_grace: None,
2489 })
2490 .await
2491 .unwrap();
2492 let mut compaction_event: Option<(usize, usize)> = None;
2493 while let Some(item) = rx.recv().await {
2494 match item.unwrap() {
2495 HarnessInternalEvent::CompactionApplied {
2496 original_message_count,
2497 compacted_message_count,
2498 ..
2499 } => {
2500 compaction_event = Some((original_message_count, compacted_message_count));
2501 }
2502 HarnessInternalEvent::TurnEnd { .. } => break,
2503 _ => {}
2504 }
2505 }
2506 assert_eq!(calls.load(Ordering::SeqCst), 1);
2508 let (orig, comp) = compaction_event.expect("CompactionApplied event emitted");
2510 assert_eq!(orig, 1, "started with 1 message ([User \"hello\"])");
2511 assert_eq!(comp, 1, "spy strategy folded to single User message");
2512 let observed = last_messages.lock().unwrap().clone().expect("model called");
2515 assert_eq!(observed.len(), 1);
2516 match &observed[0] {
2517 ChatMessage::User { content, .. } => {
2518 assert!(content.contains("FOLDED"), "got {content:?}");
2519 }
2520 other => panic!("expected User, got {other:?}"),
2521 }
2522 }
2523
2524 #[test]
2525 fn compaction_policy_summarizing_uses_default_strategy() {
2526 let policy = CompactionPolicy::summarizing(Arc::new(ScriptedModelClient), 100_000);
2527
2528 assert_eq!(policy.context_window_tokens, 100_000);
2529 assert!(!policy.strategy.should_compact(&[], 100_000));
2530 }
2531
2532 #[derive(Clone)]
2536 struct HangingModelClient {
2537 started: Arc<tokio::sync::Notify>,
2538 }
2539
2540 #[async_trait]
2541 impl ModelClient for HangingModelClient {
2542 fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
2543 CapabilitySupport::Unsupported
2544 }
2545
2546 async fn stream(
2547 &self,
2548 _input: ModelTurnInput,
2549 ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
2550 {
2551 let (tx, rx) = mpsc::channel::<Result<ModelChunk, ModelClientError>>(1);
2556 let started = self.started.clone();
2557 tokio::spawn(async move {
2558 started.notify_one();
2563 let _retain = tx; let () = std::future::pending().await;
2565 });
2566 Ok(tokio_stream::wrappers::ReceiverStream::new(rx).boxed())
2567 }
2568 }
2569
2570 #[tokio::test]
2571 async fn agent_loop_cancellation_interrupts_in_flight_stream() {
2572 let started = Arc::new(tokio::sync::Notify::new());
2577 let model = HangingModelClient {
2578 started: started.clone(),
2579 };
2580 let cancel = CancellationToken::new();
2581 let cancel_for_outside = cancel.clone();
2582
2583 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
2584 let mut rx = harness
2585 .run_turn(NativeTurnInput {
2586 prompt_text: "hi".into(),
2587 system_prompt: None,
2588 attachments: vec![],
2589 cancel_token: Some(cancel),
2590 prior_messages: vec![],
2591 context_path: None,
2592 turn_end_grace: None,
2593 })
2594 .await
2595 .unwrap();
2596
2597 started.notified().await;
2602 cancel_for_outside.cancel();
2603
2604 let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
2606 let mut saw_interrupt = false;
2607 while tokio::time::Instant::now() < deadline {
2608 tokio::select! {
2609 item = rx.recv() => {
2610 match item {
2611 Some(Ok(HarnessInternalEvent::TurnEnd { stop_reason, .. })) => {
2612 assert_eq!(stop_reason, "interrupt");
2613 saw_interrupt = true;
2614 break;
2615 }
2616 Some(_) => continue,
2617 None => break,
2618 }
2619 }
2620 _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {}
2621 }
2622 }
2623 assert!(saw_interrupt, "expected TurnEnd{{interrupt}} after cancel");
2624 }
2625
2626 enum StallBehavior {
2633 Complete(Vec<ModelChunk>),
2635 EmitThenHang(Vec<ModelChunk>),
2637 }
2638
2639 #[derive(Clone)]
2640 struct StallingModelClient {
2641 behaviors: Arc<Mutex<Vec<StallBehavior>>>,
2642 calls: Arc<AtomicUsize>,
2643 }
2644
2645 impl StallingModelClient {
2646 fn new(behaviors: Vec<StallBehavior>) -> Self {
2647 Self {
2648 behaviors: Arc::new(Mutex::new(behaviors)),
2649 calls: Arc::new(AtomicUsize::new(0)),
2650 }
2651 }
2652 }
2653
2654 #[async_trait]
2655 impl ModelClient for StallingModelClient {
2656 fn hosted_capability(&self, _capability: HostedCapability) -> CapabilitySupport {
2657 CapabilitySupport::Unsupported
2658 }
2659
2660 async fn stream(
2661 &self,
2662 _input: ModelTurnInput,
2663 ) -> Result<BoxStream<'static, Result<ModelChunk, ModelClientError>>, ModelClientError>
2664 {
2665 self.calls.fetch_add(1, Ordering::SeqCst);
2666 let behavior = {
2669 let mut b = self.behaviors.lock().unwrap();
2670 if b.is_empty() {
2671 StallBehavior::EmitThenHang(vec![])
2672 } else {
2673 b.remove(0)
2674 }
2675 };
2676 let (tx, rx) = mpsc::channel::<Result<ModelChunk, ModelClientError>>(8);
2677 tokio::spawn(async move {
2678 match behavior {
2679 StallBehavior::Complete(chunks) => {
2680 for c in chunks {
2681 if tx.send(Ok(c)).await.is_err() {
2682 return;
2683 }
2684 }
2685 }
2687 StallBehavior::EmitThenHang(chunks) => {
2688 for c in chunks {
2689 if tx.send(Ok(c)).await.is_err() {
2690 return;
2691 }
2692 }
2693 let _retain = tx; let () = std::future::pending().await;
2695 }
2696 }
2697 });
2698 Ok(tokio_stream::wrappers::ReceiverStream::new(rx).boxed())
2699 }
2700 }
2701
2702 #[tokio::test(start_paused = true)]
2703 async fn agent_loop_reconnects_after_stall_before_any_output() {
2704 let model = StallingModelClient::new(vec![
2709 StallBehavior::EmitThenHang(vec![]),
2710 StallBehavior::Complete(vec![
2711 ModelChunk::TextDelta {
2712 msg_id: "m".into(),
2713 delta: "ok".into(),
2714 },
2715 ModelChunk::Done {
2716 stop_reason: "end_turn".into(),
2717 usage: None,
2718 },
2719 ]),
2720 ]);
2721 let calls = model.calls.clone();
2722 let harness = AgentLoopHarness::new(model, MockToolRuntime::new())
2723 .with_stream_resilience(Duration::from_millis(50), 3);
2724 let mut rx = harness
2725 .run_turn(NativeTurnInput {
2726 prompt_text: "hi".into(),
2727 system_prompt: None,
2728 attachments: vec![],
2729 cancel_token: None,
2730 prior_messages: vec![],
2731 context_path: None,
2732 turn_end_grace: None,
2733 })
2734 .await
2735 .unwrap();
2736
2737 let mut text = String::new();
2738 let mut stop = None;
2739 while let Some(item) = rx.recv().await {
2740 match item.expect("no error expected") {
2741 HarnessInternalEvent::AssistantTextChunk { delta, .. } => text.push_str(&delta),
2742 HarnessInternalEvent::TurnEnd { stop_reason, .. } => {
2743 stop = Some(stop_reason);
2744 break;
2745 }
2746 _ => {}
2747 }
2748 }
2749 assert_eq!(stop.as_deref(), Some("end_turn"));
2750 assert_eq!(text, "ok", "text delivered exactly once, no duplication");
2751 assert_eq!(
2752 calls.load(Ordering::SeqCst),
2753 2,
2754 "stream established twice (one reconnect)"
2755 );
2756 }
2757
2758 #[tokio::test(start_paused = true)]
2759 async fn agent_loop_surfaces_error_when_reconnect_budget_exhausted() {
2760 let model = StallingModelClient::new(vec![]); let calls = model.calls.clone();
2764 let harness = AgentLoopHarness::new(model, MockToolRuntime::new())
2765 .with_stream_resilience(Duration::from_millis(50), 2);
2766 let mut rx = harness
2767 .run_turn(NativeTurnInput {
2768 prompt_text: "hi".into(),
2769 system_prompt: None,
2770 attachments: vec![],
2771 cancel_token: None,
2772 prior_messages: vec![],
2773 context_path: None,
2774 turn_end_grace: None,
2775 })
2776 .await
2777 .unwrap();
2778
2779 let mut saw_error = false;
2780 while let Some(item) = rx.recv().await {
2781 match item {
2782 Err(NativeHarnessError::ModelNetwork(msg)) => {
2783 assert!(msg.contains("stalled"), "got {msg:?}");
2784 saw_error = true;
2785 break;
2786 }
2787 Err(other) => panic!("unexpected error variant: {other:?}"),
2788 Ok(_) => {}
2789 }
2790 }
2791 assert!(
2792 saw_error,
2793 "expected ModelNetwork stall error after budget exhausted"
2794 );
2795 assert_eq!(
2796 calls.load(Ordering::SeqCst),
2797 2,
2798 "two establishments (initial + one reconnect)"
2799 );
2800 }
2801
2802 #[tokio::test(start_paused = true)]
2803 async fn agent_loop_does_not_reconnect_after_stall_with_partial_output() {
2804 let model = StallingModelClient::new(vec![StallBehavior::EmitThenHang(vec![
2810 ModelChunk::TextDelta {
2811 msg_id: "m".into(),
2812 delta: "partial".into(),
2813 },
2814 ])]);
2815 let calls = model.calls.clone();
2816 let harness = AgentLoopHarness::new(model, MockToolRuntime::new())
2817 .with_stream_resilience(Duration::from_millis(50), 5);
2818 let mut rx = harness
2819 .run_turn(NativeTurnInput {
2820 prompt_text: "hi".into(),
2821 system_prompt: None,
2822 attachments: vec![],
2823 cancel_token: None,
2824 prior_messages: vec![],
2825 context_path: None,
2826 turn_end_grace: None,
2827 })
2828 .await
2829 .unwrap();
2830
2831 let mut text = String::new();
2832 let mut saw_error = false;
2833 while let Some(item) = rx.recv().await {
2834 match item {
2835 Ok(HarnessInternalEvent::AssistantTextChunk { delta, .. }) => text.push_str(&delta),
2836 Err(NativeHarnessError::ModelNetwork(_)) => {
2837 saw_error = true;
2838 break;
2839 }
2840 Err(other) => panic!("unexpected error variant: {other:?}"),
2841 Ok(_) => {}
2842 }
2843 }
2844 assert!(saw_error, "expected terminal ModelNetwork error");
2845 assert_eq!(
2846 text, "partial",
2847 "partial output delivered once, not replayed"
2848 );
2849 assert_eq!(
2850 calls.load(Ordering::SeqCst),
2851 1,
2852 "no reconnect once output has reached the user"
2853 );
2854 }
2855
2856 #[tokio::test]
2857 async fn agent_loop_accumulates_thinking_chunks_and_signature() {
2858 let model = StreamingFakeClient::new(vec![vec![
2867 ModelChunk::ThinkingDelta {
2868 thinking_id: "th_1".into(),
2869 delta: "let me think...".into(),
2870 signature: None,
2871 },
2872 ModelChunk::ThinkingDelta {
2873 thinking_id: "th_1".into(),
2874 delta: "".into(),
2875 signature: Some("sig_abc".into()),
2876 },
2877 ModelChunk::TextDelta {
2878 msg_id: "m1".into(),
2879 delta: "ok".into(),
2880 },
2881 ModelChunk::Done {
2882 stop_reason: "end_turn".into(),
2883 usage: None,
2884 },
2885 ]]);
2886 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
2887 let mut rx = harness
2888 .run_turn(NativeTurnInput {
2889 prompt_text: "hi".into(),
2890 system_prompt: None,
2891 attachments: vec![],
2892 cancel_token: None,
2893 prior_messages: vec![],
2894 context_path: None,
2895 turn_end_grace: None,
2896 })
2897 .await
2898 .unwrap();
2899
2900 let mut thinking_chunks: Vec<String> = Vec::new();
2901 let mut text_chunks: Vec<String> = Vec::new();
2902 let mut saw_end = false;
2903 while let Some(item) = rx.recv().await {
2904 match item.unwrap() {
2905 HarnessInternalEvent::AssistantThinkingChunk { msg_id, delta } => {
2906 assert_eq!(msg_id, "thinking_native_0");
2907 thinking_chunks.push(delta);
2908 }
2909 HarnessInternalEvent::AssistantTextChunk { msg_id, delta } => {
2910 assert_eq!(msg_id, "msg_native_0");
2911 text_chunks.push(delta);
2912 }
2913 HarnessInternalEvent::TurnEnd { .. } => {
2914 saw_end = true;
2915 break;
2916 }
2917 other => panic!("unexpected event: {other:?}"),
2918 }
2919 }
2920 assert_eq!(thinking_chunks, vec!["let me think..."]);
2923 assert_eq!(text_chunks, vec!["ok"]);
2924 assert!(saw_end);
2925 }
2926
2927 #[tokio::test]
2928 async fn agent_loop_runs_tool_then_final_message() {
2929 let harness = AgentLoopHarness::new(
2930 ScriptedModelClient,
2931 MockToolRuntime::new().with_file("README.md", "hello"),
2932 );
2933 let mut rx = harness
2934 .run_turn(NativeTurnInput {
2935 prompt_text: "read README.md".into(),
2936 system_prompt: None,
2937 attachments: vec![],
2938 cancel_token: None,
2939 prior_messages: vec![],
2940 context_path: None,
2941 turn_end_grace: None,
2942 })
2943 .await
2944 .unwrap();
2945
2946 assert!(matches!(
2947 rx.recv().await.unwrap().unwrap(),
2948 HarnessInternalEvent::AssistantTextChunk { .. }
2949 ));
2950 assert!(matches!(
2951 rx.recv().await.unwrap().unwrap(),
2952 HarnessInternalEvent::ToolCall { ref name, .. } if name == "read"
2953 ));
2954 assert!(matches!(
2955 rx.recv().await.unwrap().unwrap(),
2956 HarnessInternalEvent::ToolResult { .. }
2957 ));
2958 assert!(matches!(
2959 rx.recv().await.unwrap().unwrap(),
2960 HarnessInternalEvent::AssistantTextChunk { .. }
2961 ));
2962 assert!(matches!(
2963 rx.recv().await.unwrap().unwrap(),
2964 HarnessInternalEvent::TurnEnd { .. }
2965 ));
2966 assert!(rx.recv().await.is_none());
2967 }
2968
2969 #[tokio::test]
2976 async fn agent_loop_turn_end_carries_full_message_history() {
2977 let model = QueueModelClient::new(vec![ModelResponse::Message {
2978 text: "second reply".into(),
2979 stop_reason: "end_turn".into(),
2980 usage: None,
2981 }]);
2982 let harness = AgentLoopHarness::new(model, MockToolRuntime::new());
2983 let prior = vec![
2985 ChatMessage::User {
2986 content: "first prompt".into(),
2987 attachments: vec![],
2988 },
2989 ChatMessage::Assistant {
2990 text: Some("first reply".into()),
2991 tool_calls: vec![],
2992 thinking: None,
2993 usage: None,
2994 },
2995 ];
2996 let mut rx = harness
2997 .run_turn(NativeTurnInput {
2998 prompt_text: "second prompt".into(),
2999 system_prompt: None,
3000 attachments: vec![],
3001 cancel_token: None,
3002 prior_messages: prior,
3003 context_path: None,
3004 turn_end_grace: None,
3005 })
3006 .await
3007 .unwrap();
3008 let mut final_messages: Option<Vec<ChatMessage>> = None;
3009 while let Some(item) = rx.recv().await {
3010 if let HarnessInternalEvent::TurnEnd {
3011 final_messages: m, ..
3012 } = item.unwrap()
3013 {
3014 final_messages = Some(m);
3015 break;
3016 }
3017 }
3018 let msgs = final_messages.expect("TurnEnd carried final_messages");
3019 assert_eq!(msgs.len(), 4, "got {msgs:?}");
3021 match &msgs[0] {
3022 ChatMessage::User { content, .. } => assert_eq!(content, "first prompt"),
3023 other => panic!("msgs[0] not user-1: {other:?}"),
3024 }
3025 match &msgs[1] {
3026 ChatMessage::Assistant { text, .. } => {
3027 assert_eq!(text.as_deref(), Some("first reply"));
3028 }
3029 other => panic!("msgs[1] not assistant-1: {other:?}"),
3030 }
3031 match &msgs[2] {
3032 ChatMessage::User { content, .. } => assert_eq!(content, "second prompt"),
3033 other => panic!("msgs[2] not user-2: {other:?}"),
3034 }
3035 match &msgs[3] {
3036 ChatMessage::Assistant { text, .. } => {
3037 assert_eq!(text.as_deref(), Some("second reply"));
3038 }
3039 other => panic!("msgs[3] not assistant-2: {other:?}"),
3040 }
3041 }
3042
3043 #[derive(Clone)]
3048 struct ConcurrencyProbeRuntime {
3049 sleep_for: std::time::Duration,
3050 in_flight: Arc<AtomicUsize>,
3051 max_concurrency: Arc<AtomicUsize>,
3052 call_order: Arc<Mutex<Vec<String>>>,
3053 cancelled: Arc<AtomicUsize>,
3054 }
3055
3056 impl ConcurrencyProbeRuntime {
3057 fn new(sleep_for: std::time::Duration) -> Self {
3058 Self {
3059 sleep_for,
3060 in_flight: Arc::new(AtomicUsize::new(0)),
3061 max_concurrency: Arc::new(AtomicUsize::new(0)),
3062 call_order: Arc::new(Mutex::new(Vec::new())),
3063 cancelled: Arc::new(AtomicUsize::new(0)),
3064 }
3065 }
3066 }
3067
3068 #[async_trait]
3069 impl ToolRuntime for ConcurrencyProbeRuntime {
3070 fn specs(&self) -> Vec<crate::tools::ToolSpec> {
3071 vec![crate::tools::ToolSpec {
3072 name: "slow".into(),
3073 description: "sleeps".into(),
3074 input_schema: serde_json::json!({"type": "object"}),
3075 }]
3076 }
3077
3078 async fn invoke(
3079 &self,
3080 invocation: ToolInvocation,
3081 ) -> Result<ToolOutcome, ToolRuntimeError> {
3082 self.call_order.lock().unwrap().push(invocation.id.clone());
3083 let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
3084 let mut prev = self.max_concurrency.load(Ordering::SeqCst);
3085 while now > prev {
3086 match self.max_concurrency.compare_exchange(
3087 prev,
3088 now,
3089 Ordering::SeqCst,
3090 Ordering::SeqCst,
3091 ) {
3092 Ok(_) => break,
3093 Err(actual) => prev = actual,
3094 }
3095 }
3096 tokio::time::sleep(self.sleep_for).await;
3097 self.in_flight.fetch_sub(1, Ordering::SeqCst);
3098 Ok(ToolOutcome {
3099 output: Ok(serde_json::json!({"slept": true, "id": invocation.id})),
3100 attachments: vec![],
3101 })
3102 }
3103
3104 async fn invoke_cancellable(
3105 &self,
3106 invocation: ToolInvocation,
3107 cancel: Option<&CancellationToken>,
3108 ) -> Result<ToolOutcome, ToolRuntimeError> {
3109 self.call_order.lock().unwrap().push(invocation.id.clone());
3110 let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
3111 let mut prev = self.max_concurrency.load(Ordering::SeqCst);
3112 while now > prev {
3113 match self.max_concurrency.compare_exchange(
3114 prev,
3115 now,
3116 Ordering::SeqCst,
3117 Ordering::SeqCst,
3118 ) {
3119 Ok(_) => break,
3120 Err(actual) => prev = actual,
3121 }
3122 }
3123 if let Some(token) = cancel {
3124 tokio::select! {
3125 _ = token.cancelled() => {
3126 self.cancelled.fetch_add(1, Ordering::SeqCst);
3127 self.in_flight.fetch_sub(1, Ordering::SeqCst);
3128 Err(ToolRuntimeError::Runtime("cancelled".into()))
3129 }
3130 _ = tokio::time::sleep(self.sleep_for) => {
3131 self.in_flight.fetch_sub(1, Ordering::SeqCst);
3132 Ok(ToolOutcome {
3133 output: Ok(serde_json::json!({"slept": true, "id": invocation.id})),
3134 attachments: vec![],
3135 })
3136 }
3137 }
3138 } else {
3139 tokio::time::sleep(self.sleep_for).await;
3140 self.in_flight.fetch_sub(1, Ordering::SeqCst);
3141 Ok(ToolOutcome {
3142 output: Ok(serde_json::json!({"slept": true, "id": invocation.id})),
3143 attachments: vec![],
3144 })
3145 }
3146 }
3147 }
3148
3149 #[tokio::test]
3155 async fn agent_loop_runs_multi_tool_calls_concurrently() {
3156 let model = StreamingFakeClient::new(vec![
3159 vec![
3160 ModelChunk::ToolCallStart {
3161 id: "tc_a".into(),
3162 name: "slow".into(),
3163 },
3164 ModelChunk::ToolCallEnd {
3165 id: "tc_a".into(),
3166 input: Some(json!({})),
3167 },
3168 ModelChunk::ToolCallStart {
3169 id: "tc_b".into(),
3170 name: "slow".into(),
3171 },
3172 ModelChunk::ToolCallEnd {
3173 id: "tc_b".into(),
3174 input: Some(json!({})),
3175 },
3176 ModelChunk::ToolCallStart {
3177 id: "tc_c".into(),
3178 name: "slow".into(),
3179 },
3180 ModelChunk::ToolCallEnd {
3181 id: "tc_c".into(),
3182 input: Some(json!({})),
3183 },
3184 ModelChunk::Done {
3185 stop_reason: "tool_use".into(),
3186 usage: None,
3187 },
3188 ],
3189 vec![
3190 ModelChunk::TextDelta {
3191 msg_id: "remote".into(),
3192 delta: "done".into(),
3193 },
3194 ModelChunk::Done {
3195 stop_reason: "end_turn".into(),
3196 usage: None,
3197 },
3198 ],
3199 ]);
3200 let probe = ConcurrencyProbeRuntime::new(std::time::Duration::from_millis(80));
3201 let max_concurrency = probe.max_concurrency.clone();
3202 let harness = AgentLoopHarness::new(model, probe);
3203
3204 let start = std::time::Instant::now();
3205 let mut rx = harness
3206 .run_turn(NativeTurnInput {
3207 prompt_text: "go".into(),
3208 system_prompt: None,
3209 attachments: vec![],
3210 cancel_token: None,
3211 prior_messages: vec![],
3212 context_path: None,
3213 turn_end_grace: None,
3214 })
3215 .await
3216 .unwrap();
3217 let mut tool_results = 0;
3218 while let Some(item) = rx.recv().await {
3219 match item.unwrap() {
3220 HarnessInternalEvent::ToolResult { .. } => tool_results += 1,
3221 HarnessInternalEvent::TurnEnd { .. } => break,
3222 _ => {}
3223 }
3224 }
3225 let elapsed = start.elapsed();
3226 assert_eq!(
3228 tool_results, 3,
3229 "expected 3 tool results, got {tool_results}"
3230 );
3231 assert_eq!(
3233 max_concurrency.load(Ordering::SeqCst),
3234 3,
3235 "expected max concurrency 3 (parallel dispatch), got {}",
3236 max_concurrency.load(Ordering::SeqCst)
3237 );
3238 assert!(
3242 elapsed < std::time::Duration::from_millis(200),
3243 "elapsed {elapsed:?} suggests sequential execution"
3244 );
3245 }
3246
3247 #[tokio::test]
3251 async fn agent_loop_cancels_in_flight_tool_calls() {
3252 let model = StreamingFakeClient::new(vec![vec![
3253 ModelChunk::ToolCallStart {
3254 id: "tc_slow".into(),
3255 name: "slow".into(),
3256 },
3257 ModelChunk::ToolCallEnd {
3258 id: "tc_slow".into(),
3259 input: Some(json!({})),
3260 },
3261 ModelChunk::Done {
3262 stop_reason: "tool_use".into(),
3263 usage: None,
3264 },
3265 ]]);
3266 let probe = ConcurrencyProbeRuntime::new(std::time::Duration::from_secs(5));
3269 let cancelled_count = probe.cancelled.clone();
3270 let harness = AgentLoopHarness::new(model, probe);
3271
3272 let cancel = CancellationToken::new();
3273 let cancel_for_input = cancel.clone();
3274 let mut rx = harness
3275 .run_turn(NativeTurnInput {
3276 prompt_text: "go".into(),
3277 system_prompt: None,
3278 attachments: vec![],
3279 cancel_token: Some(cancel_for_input),
3280 prior_messages: vec![],
3281 context_path: None,
3282 turn_end_grace: None,
3283 })
3284 .await
3285 .unwrap();
3286
3287 tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3289 cancel.cancel();
3290
3291 let start = std::time::Instant::now();
3292 let mut saw_interrupt = false;
3293 while let Some(item) = rx.recv().await {
3294 if let HarnessInternalEvent::TurnEnd { stop_reason, .. } = item.unwrap() {
3295 assert_eq!(stop_reason, "interrupt");
3296 saw_interrupt = true;
3297 break;
3298 }
3299 }
3300 let elapsed = start.elapsed();
3301 assert!(saw_interrupt, "must see interrupt TurnEnd");
3302 assert!(
3303 elapsed < std::time::Duration::from_millis(200),
3304 "cancel propagation took too long: {elapsed:?}"
3305 );
3306 let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(1);
3307 while cancelled_count.load(Ordering::SeqCst) == 0 && tokio::time::Instant::now() < deadline
3308 {
3309 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
3310 }
3311 assert_eq!(
3312 cancelled_count.load(Ordering::SeqCst),
3313 1,
3314 "tool runtime must observe the cancellation token"
3315 );
3316 }
3317
3318 #[derive(Clone)]
3319 struct DefaultCancellationProbe {
3320 started: Arc<tokio::sync::Notify>,
3321 completed: Arc<std::sync::atomic::AtomicBool>,
3322 }
3323
3324 #[async_trait]
3325 impl ToolRuntime for DefaultCancellationProbe {
3326 fn specs(&self) -> Vec<crate::tools::ToolSpec> {
3327 vec![crate::tools::ToolSpec {
3328 name: "slow".into(),
3329 description: "records a delayed side effect".into(),
3330 input_schema: serde_json::json!({"type": "object"}),
3331 }]
3332 }
3333
3334 async fn invoke(
3335 &self,
3336 _invocation: ToolInvocation,
3337 ) -> Result<ToolOutcome, ToolRuntimeError> {
3338 self.started.notify_one();
3339 tokio::time::sleep(std::time::Duration::from_millis(200)).await;
3340 self.completed.store(true, Ordering::SeqCst);
3341 Ok(ToolOutcome {
3342 output: Ok(serde_json::json!({"completed": true})),
3343 attachments: vec![],
3344 })
3345 }
3346 }
3347
3348 #[tokio::test]
3353 async fn agent_loop_cancel_drops_default_tool_future() {
3354 let model = StreamingFakeClient::new(vec![vec![
3355 ModelChunk::ToolCallStart {
3356 id: "tc_slow".into(),
3357 name: "slow".into(),
3358 },
3359 ModelChunk::ToolCallEnd {
3360 id: "tc_slow".into(),
3361 input: Some(json!({})),
3362 },
3363 ModelChunk::Done {
3364 stop_reason: "tool_use".into(),
3365 usage: None,
3366 },
3367 ]]);
3368 let started = Arc::new(tokio::sync::Notify::new());
3369 let completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
3370 let runtime = DefaultCancellationProbe {
3371 started: started.clone(),
3372 completed: completed.clone(),
3373 };
3374 let cancel = CancellationToken::new();
3375 let harness = AgentLoopHarness::new(model, runtime);
3376 let mut rx = harness
3377 .run_turn(NativeTurnInput {
3378 prompt_text: "go".into(),
3379 system_prompt: None,
3380 attachments: vec![],
3381 cancel_token: Some(cancel.clone()),
3382 prior_messages: vec![],
3383 context_path: None,
3384 turn_end_grace: None,
3385 })
3386 .await
3387 .unwrap();
3388
3389 started.notified().await;
3390 cancel.cancel();
3391
3392 let mut saw_interrupt = false;
3393 while let Some(item) = rx.recv().await {
3394 if let HarnessInternalEvent::TurnEnd { stop_reason, .. } = item.unwrap() {
3395 assert_eq!(stop_reason, "interrupt");
3396 saw_interrupt = true;
3397 break;
3398 }
3399 }
3400 assert!(saw_interrupt);
3401 tokio::time::sleep(std::time::Duration::from_millis(300)).await;
3402 assert!(
3403 !completed.load(Ordering::SeqCst),
3404 "cancelled tool produced a side effect after TurnEnd"
3405 );
3406 }
3407
3408 #[tokio::test]
3413 async fn agent_loop_cancel_during_tools_yields_clean_history() {
3414 let model = StreamingFakeClient::new(vec![vec![
3415 ModelChunk::ToolCallStart {
3416 id: "tc_a".into(),
3417 name: "slow".into(),
3418 },
3419 ModelChunk::ToolCallEnd {
3420 id: "tc_a".into(),
3421 input: Some(json!({})),
3422 },
3423 ModelChunk::ToolCallStart {
3424 id: "tc_b".into(),
3425 name: "slow".into(),
3426 },
3427 ModelChunk::ToolCallEnd {
3428 id: "tc_b".into(),
3429 input: Some(json!({})),
3430 },
3431 ModelChunk::Done {
3432 stop_reason: "tool_use".into(),
3433 usage: None,
3434 },
3435 ]]);
3436 let probe = ConcurrencyProbeRuntime::new(std::time::Duration::from_secs(3));
3437 let harness = AgentLoopHarness::new(model, probe);
3438
3439 let cancel = CancellationToken::new();
3440 let cancel_for_input = cancel.clone();
3441 let mut rx = harness
3442 .run_turn(NativeTurnInput {
3443 prompt_text: "go".into(),
3444 system_prompt: None,
3445 attachments: vec![],
3446 cancel_token: Some(cancel_for_input),
3447 prior_messages: vec![],
3448 context_path: None,
3449 turn_end_grace: None,
3450 })
3451 .await
3452 .unwrap();
3453
3454 tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3455 cancel.cancel();
3456
3457 let mut final_msgs = None;
3458 while let Some(item) = rx.recv().await {
3459 if let HarnessInternalEvent::TurnEnd { final_messages, .. } = item.unwrap() {
3460 final_msgs = Some(final_messages);
3461 break;
3462 }
3463 }
3464 let msgs = final_msgs.expect("interrupt TurnEnd");
3465 assert_eq!(msgs.len(), 2, "expected 2 messages, got {msgs:?}");
3468 match &msgs[1] {
3469 ChatMessage::Assistant { tool_calls, .. } => {
3470 assert_eq!(tool_calls.len(), 2);
3471 assert_eq!(tool_calls[0].id, "tc_a");
3472 assert_eq!(tool_calls[1].id, "tc_b");
3473 }
3474 other => panic!("msgs[1] not assistant: {other:?}"),
3475 }
3476 }
3477
3478 #[derive(Clone)]
3484 struct WriterProbeRuntime {
3485 sleep_for: std::time::Duration,
3486 completed: Arc<std::sync::atomic::AtomicBool>,
3487 }
3488
3489 #[async_trait]
3490 impl ToolRuntime for WriterProbeRuntime {
3491 fn specs(&self) -> Vec<crate::tools::ToolSpec> {
3492 vec![crate::tools::ToolSpec {
3493 name: "slow".into(),
3494 description: "writer that ignores cancel".into(),
3495 input_schema: serde_json::json!({"type": "object"}),
3496 }]
3497 }
3498
3499 async fn invoke(
3500 &self,
3501 _invocation: ToolInvocation,
3502 ) -> Result<ToolOutcome, ToolRuntimeError> {
3503 tokio::time::sleep(self.sleep_for).await;
3504 self.completed.store(true, Ordering::SeqCst);
3505 Ok(ToolOutcome {
3506 output: Ok(json!({"done": true})),
3507 attachments: vec![],
3508 })
3509 }
3510
3511 async fn invoke_cancellable(
3512 &self,
3513 invocation: ToolInvocation,
3514 _cancel: Option<&CancellationToken>,
3515 ) -> Result<ToolOutcome, ToolRuntimeError> {
3516 self.invoke(invocation).await
3517 }
3518 }
3519
3520 fn writer_tool_turn() -> StreamingFakeClient {
3521 StreamingFakeClient::new(vec![vec![
3522 ModelChunk::ToolCallStart {
3523 id: "tc_slow".into(),
3524 name: "slow".into(),
3525 },
3526 ModelChunk::ToolCallEnd {
3527 id: "tc_slow".into(),
3528 input: Some(json!({})),
3529 },
3530 ModelChunk::Done {
3531 stop_reason: "tool_use".into(),
3532 usage: None,
3533 },
3534 ]])
3535 }
3536
3537 #[tokio::test]
3542 async fn turn_end_grace_drains_unresolved_tool_futures() {
3543 let model = writer_tool_turn();
3544 let completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
3545 let runtime = WriterProbeRuntime {
3546 sleep_for: std::time::Duration::from_millis(400),
3547 completed: completed.clone(),
3548 };
3549 let harness = AgentLoopHarness::new(model, runtime);
3550 let cancel = CancellationToken::new();
3551 let mut rx = harness
3552 .run_turn(NativeTurnInput {
3553 prompt_text: "go".into(),
3554 system_prompt: None,
3555 attachments: vec![],
3556 cancel_token: Some(cancel.clone()),
3557 prior_messages: vec![],
3558 context_path: None,
3559 turn_end_grace: Some(std::time::Duration::from_secs(3)),
3560 })
3561 .await
3562 .unwrap();
3563
3564 tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3565 cancel.cancel();
3566
3567 let start = std::time::Instant::now();
3568 let (saw_interrupt, pending) = loop {
3569 match rx.recv().await {
3570 Some(Ok(HarnessInternalEvent::TurnEnd {
3571 stop_reason,
3572 pending_tools,
3573 ..
3574 })) => break (stop_reason == "interrupt", pending_tools),
3575 Some(_) => continue,
3576 None => panic!("channel closed without TurnEnd"),
3577 }
3578 };
3579 let elapsed = start.elapsed();
3580 assert!(saw_interrupt);
3581 assert!(
3582 elapsed >= std::time::Duration::from_millis(350),
3583 "TurnEnd landed at {elapsed:?} — writer future was not drained"
3584 );
3585 assert!(
3586 elapsed < std::time::Duration::from_secs(2),
3587 "TurnEnd landed at {elapsed:?} — grace window not honoured"
3588 );
3589 assert!(
3590 completed.load(Ordering::SeqCst),
3591 "writer must have finished"
3592 );
3593 assert!(
3594 pending.is_empty(),
3595 "drained turn must report no pending tools, got {pending:?}"
3596 );
3597 }
3598
3599 #[tokio::test]
3603 async fn turn_end_grace_expiry_reports_pending_tools() {
3604 let model = writer_tool_turn();
3605 let completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
3606 let runtime = WriterProbeRuntime {
3607 sleep_for: std::time::Duration::from_secs(5),
3608 completed: completed.clone(),
3609 };
3610 let harness = AgentLoopHarness::new(model, runtime);
3611 let cancel = CancellationToken::new();
3612 let mut rx = harness
3613 .run_turn(NativeTurnInput {
3614 prompt_text: "go".into(),
3615 system_prompt: None,
3616 attachments: vec![],
3617 cancel_token: Some(cancel.clone()),
3618 prior_messages: vec![],
3619 context_path: None,
3620 turn_end_grace: Some(std::time::Duration::from_millis(150)),
3621 })
3622 .await
3623 .unwrap();
3624
3625 tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3626 cancel.cancel();
3627
3628 let start = std::time::Instant::now();
3629 let (saw_interrupt, pending) = loop {
3630 match rx.recv().await {
3631 Some(Ok(HarnessInternalEvent::TurnEnd {
3632 stop_reason,
3633 pending_tools,
3634 ..
3635 })) => break (stop_reason == "interrupt", pending_tools),
3636 Some(_) => continue,
3637 None => panic!("channel closed without TurnEnd"),
3638 }
3639 };
3640 let elapsed = start.elapsed();
3641 assert!(saw_interrupt);
3642 assert!(
3643 elapsed >= std::time::Duration::from_millis(120),
3644 "TurnEnd landed at {elapsed:?} — grace window was skipped"
3645 );
3646 assert!(
3647 elapsed < std::time::Duration::from_secs(1),
3648 "TurnEnd landed at {elapsed:?} — should not wait for the writer"
3649 );
3650 assert_eq!(pending, vec!["tc_slow".to_string()]);
3651 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
3653 assert!(!completed.load(Ordering::SeqCst));
3654 }
3655
3656 #[tokio::test]
3659 async fn turn_end_grace_builder_default_overridden_by_per_turn() {
3660 let completed = Arc::new(std::sync::atomic::AtomicBool::new(false));
3661 let harness = AgentLoopHarness::new(
3663 writer_tool_turn(),
3664 WriterProbeRuntime {
3665 sleep_for: std::time::Duration::from_millis(400),
3666 completed: completed.clone(),
3667 },
3668 )
3669 .with_turn_end_grace(std::time::Duration::from_millis(150));
3670 let cancel = CancellationToken::new();
3671 let mut rx = harness
3672 .run_turn(NativeTurnInput {
3673 prompt_text: "go".into(),
3674 system_prompt: None,
3675 attachments: vec![],
3676 cancel_token: Some(cancel.clone()),
3677 prior_messages: vec![],
3678 context_path: None,
3679 turn_end_grace: None,
3680 })
3681 .await
3682 .unwrap();
3683 tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3684 cancel.cancel();
3685 let (_, pending) = loop {
3686 match rx.recv().await {
3687 Some(Ok(HarnessInternalEvent::TurnEnd {
3688 stop_reason,
3689 pending_tools,
3690 ..
3691 })) => {
3692 assert_eq!(stop_reason, "interrupt");
3693 break (true, pending_tools);
3694 }
3695 Some(_) => continue,
3696 None => panic!("channel closed without TurnEnd"),
3697 }
3698 };
3699 assert_eq!(pending, vec!["tc_slow".to_string()]);
3700
3701 let completed2 = Arc::new(std::sync::atomic::AtomicBool::new(false));
3703 let harness = AgentLoopHarness::new(
3704 writer_tool_turn(),
3705 WriterProbeRuntime {
3706 sleep_for: std::time::Duration::from_millis(400),
3707 completed: completed2.clone(),
3708 },
3709 )
3710 .with_turn_end_grace(std::time::Duration::from_millis(150));
3711 let cancel = CancellationToken::new();
3712 let mut rx = harness
3713 .run_turn(NativeTurnInput {
3714 prompt_text: "go".into(),
3715 system_prompt: None,
3716 attachments: vec![],
3717 cancel_token: Some(cancel.clone()),
3718 prior_messages: vec![],
3719 context_path: None,
3720 turn_end_grace: Some(std::time::Duration::from_secs(2)),
3721 })
3722 .await
3723 .unwrap();
3724 tokio::time::sleep(std::time::Duration::from_millis(30)).await;
3725 cancel.cancel();
3726 let (_, pending) = loop {
3727 match rx.recv().await {
3728 Some(Ok(HarnessInternalEvent::TurnEnd { pending_tools, .. })) => {
3729 break (true, pending_tools)
3730 }
3731 Some(_) => continue,
3732 None => panic!("channel closed without TurnEnd"),
3733 }
3734 };
3735 assert!(pending.is_empty());
3736 assert!(completed2.load(Ordering::SeqCst));
3737 }
3738
3739 #[test]
3740 fn default_turn_end_grace_is_unchanged() {
3741 assert_eq!(DEFAULT_TURN_END_GRACE, Duration::from_secs(1));
3744 }
3745}