serdes_ai_agent/stream.rs
1//! Streaming agent execution.
2//!
3//! This module provides streaming support for agent runs with real
4//! character-by-character streaming from the model.
5
6use crate::agent::{Agent, RegisteredTool};
7use crate::context::{generate_run_id, RunContext, RunUsage};
8use crate::errors::AgentRunError;
9use crate::run::{CompressionStrategy, RunOptions};
10use chrono::Utc;
11use futures::{Stream, StreamExt};
12use serdes_ai_core::messages::{
13 ModelResponseStreamEvent, ToolCallArgs, ToolReturnPart, UserContent,
14};
15use serdes_ai_core::{
16 FinishReason, ModelRequest, ModelRequestPart, ModelResponse, ModelResponsePart,
17};
18use serdes_ai_models::ModelRequestParameters;
19use std::pin::Pin;
20use std::sync::Arc;
21use std::task::{Context, Poll};
22use tokio::sync::mpsc;
23use tokio_util::sync::CancellationToken;
24
25// Conditional tracing - use no-op macros when tracing feature is disabled
26#[cfg(feature = "tracing-integration")]
27use tracing::{debug, error, info, warn};
28
29#[cfg(not(feature = "tracing-integration"))]
30macro_rules! debug {
31 ($($arg:tt)*) => {};
32}
33#[cfg(not(feature = "tracing-integration"))]
34macro_rules! info {
35 ($($arg:tt)*) => {};
36}
37#[cfg(not(feature = "tracing-integration"))]
38macro_rules! error {
39 ($($arg:tt)*) => {};
40}
41#[cfg(not(feature = "tracing-integration"))]
42macro_rules! warn {
43 ($($arg:tt)*) => {};
44}
45
46/// Events emitted during streaming.
47#[derive(Debug, Clone)]
48pub enum AgentStreamEvent {
49 /// Run started.
50 RunStart { run_id: String },
51 /// Context size information (emitted before each model request).
52 ContextInfo {
53 /// Estimated token count (~request_bytes / 4).
54 estimated_tokens: usize,
55 /// Raw request size in bytes (serialized messages + tools).
56 request_bytes: usize,
57 /// Model's context window limit (if known).
58 context_limit: Option<u64>,
59 },
60 /// Context was compressed to fit within limits.
61 ContextCompressed {
62 /// Token count before compression.
63 original_tokens: usize,
64 /// Token count after compression.
65 compressed_tokens: usize,
66 /// Strategy used: "truncate" or "summarize".
67 strategy: String,
68 /// Number of messages before compression.
69 messages_before: usize,
70 /// Number of messages after compression.
71 messages_after: usize,
72 },
73 /// Model request started.
74 RequestStart { step: u32 },
75 /// Text delta.
76 TextDelta { text: String },
77 /// Tool call started.
78 ToolCallStart {
79 tool_name: String,
80 tool_call_id: Option<String>,
81 },
82 /// Tool call arguments delta.
83 ToolCallDelta {
84 delta: String,
85 tool_call_id: Option<String>,
86 },
87 /// Tool call completed (arguments fully received).
88 ToolCallComplete {
89 tool_name: String,
90 tool_call_id: Option<String>,
91 },
92 /// Tool executed.
93 ToolExecuted {
94 tool_name: String,
95 tool_call_id: Option<String>,
96 success: bool,
97 error: Option<String>,
98 },
99 /// Thinking delta (for reasoning models).
100 ThinkingDelta { text: String },
101 /// Model response completed.
102 ResponseComplete { step: u32 },
103 /// Output ready.
104 OutputReady,
105 /// Run completed.
106 RunComplete {
107 run_id: String,
108 /// Complete message history from this run (system prompt, user prompts,
109 /// assistant responses, tool calls and returns).
110 messages: Vec<ModelRequest>,
111 },
112 /// Error occurred.
113 Error { message: String },
114 /// Run was cancelled.
115 Cancelled {
116 /// Partial text accumulated before cancellation.
117 partial_text: Option<String>,
118 /// Partial thinking content accumulated before cancellation.
119 partial_thinking: Option<String>,
120 /// Tool calls that were in progress when cancelled.
121 pending_tools: Vec<String>,
122 },
123}
124
125/// Streaming agent execution.
126///
127/// This provides real streaming by spawning a task that streams from the model
128/// and sends events through a channel.
129///
130/// # Cancellation
131///
132/// Use [`AgentStream::new_with_cancel`] to create a stream with cancellation support.
133/// When the cancellation token is triggered, the stream will:
134/// 1. Stop the model stream
135/// 2. Cancel any pending tool calls
136/// 3. Emit a [`AgentStreamEvent::Cancelled`] event with partial results
137pub struct AgentStream {
138 rx: mpsc::Receiver<Result<AgentStreamEvent, AgentRunError>>,
139 /// Cancellation token for this stream (if cancellation is enabled).
140 cancel_token: Option<CancellationToken>,
141}
142
143/// Canonicalize tool-call arguments in a model response before persisting it.
144fn canonicalize_tool_call_args_in_response(response: &mut ModelResponse) {
145 for part in &mut response.parts {
146 if let ModelResponsePart::ToolCall(tc) = part {
147 let repaired = tc.args.to_json();
148 tc.args = ToolCallArgs::Json(repaired);
149 }
150 }
151}
152
153impl AgentStream {
154 /// Create a new streaming agent run.
155 ///
156 /// This spawns a background task that handles the actual streaming
157 /// and tool execution.
158 pub async fn new<Deps, Output>(
159 agent: &Agent<Deps, Output>,
160 prompt: UserContent,
161 deps: Deps,
162 options: RunOptions,
163 ) -> Result<Self, AgentRunError>
164 where
165 Deps: Send + Sync + 'static,
166 Output: Send + Sync + 'static,
167 {
168 let run_id = generate_run_id();
169 let (tx, rx) = mpsc::channel(64);
170
171 // Clone what we need for the spawned task
172 let model = agent.model_arc();
173 let model_name = model.name().to_string();
174 let model_settings = options
175 .model_settings
176 .clone()
177 .unwrap_or_else(|| agent.model_settings.clone());
178
179 // Get the static system prompt - for streaming we use just the static part
180 // Dynamic prompts are not supported in streaming mode for simplicity
181 let static_system_prompt = agent.static_system_prompt().to_string();
182
183 let tool_definitions = agent.tool_definitions();
184 let _end_strategy = agent.end_strategy;
185 let usage_limits = agent.usage_limits.clone();
186 let run_usage_limits = options.usage_limits.clone();
187
188 // Clone tool executors - now possible because RegisteredTool implements Clone!
189 let tools: Vec<RegisteredTool<Deps>> = agent.tools.to_vec();
190
191 // Wrap deps in Arc for shared access in tool execution
192 let deps = Arc::new(deps);
193
194 let initial_history = options.message_history.clone();
195 let _metadata = options.metadata.clone();
196 let compression_config = options.compression.clone();
197 let run_id_clone = run_id.clone();
198
199 debug!(run_id = %run_id, "AgentStream: spawning streaming task");
200
201 // Spawn the streaming task
202 tokio::spawn(async move {
203 info!(run_id = %run_id_clone, "AgentStream: task started");
204
205 // Emit RunStart
206 debug!("AgentStream: emitting RunStart");
207 if tx
208 .send(Ok(AgentStreamEvent::RunStart {
209 run_id: run_id_clone.clone(),
210 }))
211 .await
212 .is_err()
213 {
214 warn!("AgentStream: receiver dropped before RunStart");
215 return;
216 }
217
218 // Build initial messages
219 let mut messages = initial_history.unwrap_or_default();
220 debug!(
221 initial_messages = messages.len(),
222 "AgentStream: building messages"
223 );
224
225 // Add system prompt if non-empty
226 if !static_system_prompt.is_empty() {
227 let mut req = ModelRequest::new();
228 req.add_system_prompt(static_system_prompt.clone());
229 messages.push(req);
230 }
231
232 // Add user prompt
233 let mut user_req = ModelRequest::new();
234 user_req.add_user_prompt(prompt);
235 messages.push(user_req);
236
237 let mut responses: Vec<ModelResponse> = Vec::new();
238 let mut usage = RunUsage::new();
239 let mut step = 0u32;
240 let mut finished = false;
241 let mut finish_reason: Option<FinishReason>;
242
243 // Main agent loop
244 while !finished {
245 step += 1;
246
247 // Check usage limits
248 if let Some(ref limits) = usage_limits {
249 if let Err(e) = limits.check(&usage) {
250 let _ = tx.send(Err(e.into())).await;
251 return;
252 }
253 }
254
255 if let Some(ref limits) = run_usage_limits {
256 if let Err(e) = limits.check(&usage) {
257 let _ = tx.send(Err(e.into())).await;
258 return;
259 }
260 }
261
262 // Emit RequestStart
263 if tx
264 .send(Ok(AgentStreamEvent::RequestStart { step }))
265 .await
266 .is_err()
267 {
268 return;
269 }
270
271 // Build request parameters
272 let params = ModelRequestParameters::new()
273 .with_tools_arc(tool_definitions.clone())
274 .with_allow_text(true);
275
276 // === Context Size Calculation & Compression ===
277
278 // Calculate context size by serializing (this is the actual request size)
279 let (request_bytes, estimated_tokens) = {
280 let messages_json = serde_json::to_string(&messages).unwrap_or_default();
281 let tools_json = serde_json::to_string(&*tool_definitions).unwrap_or_default();
282 let bytes = messages_json.len() + tools_json.len();
283 (bytes, bytes / 4)
284 };
285
286 // Get context limit from model profile
287 let context_limit = model.profile().context_window;
288
289 // Emit ContextInfo event
290 let _ = tx
291 .send(Ok(AgentStreamEvent::ContextInfo {
292 estimated_tokens,
293 request_bytes,
294 context_limit,
295 }))
296 .await;
297
298 // Check if compression is needed
299 if let Some(ref compression) = compression_config {
300 if let Some(limit) = context_limit {
301 let threshold_tokens = (limit as f64 * compression.threshold) as usize;
302
303 if estimated_tokens > threshold_tokens {
304 let messages_before = messages.len();
305 let original_tokens = estimated_tokens;
306
307 // Apply compression based on strategy
308 let strategy_name = match compression.strategy {
309 CompressionStrategy::Truncate => {
310 // Use TruncateByTokens with keep_first_n=2 (system + first user)
311 use crate::history::{HistoryProcessor, TruncateByTokens};
312 let truncator =
313 TruncateByTokens::new(compression.target_tokens as u64)
314 .keep_first_n(2);
315
316 // Create a minimal context for the processor
317 let temp_ctx = RunContext::new((), &model_name);
318 messages = truncator.process(&temp_ctx, messages).await;
319 "truncate"
320 }
321 CompressionStrategy::Summarize => {
322 // Use the same model to summarize the conversation history
323 // Keep first 2 messages (system + first user) and last few messages
324 // Summarize everything in between
325
326 if messages.len() <= 4 {
327 // Too few messages to summarize, just truncate
328 use crate::history::{HistoryProcessor, TruncateByTokens};
329 let truncator =
330 TruncateByTokens::new(compression.target_tokens as u64)
331 .keep_first_n(2);
332 let temp_ctx = RunContext::new((), &model_name);
333 messages = truncator.process(&temp_ctx, messages).await;
334 "truncate (too few messages)"
335 } else {
336 // Split messages: first 2 (keep), middle (summarize), last 2 (keep)
337 let first_two: Vec<_> =
338 messages.iter().take(2).cloned().collect();
339 let last_two: Vec<_> = messages
340 .iter()
341 .rev()
342 .take(2)
343 .cloned()
344 .collect::<Vec<_>>()
345 .into_iter()
346 .rev()
347 .collect();
348 let middle: Vec<_> = messages
349 .iter()
350 .skip(2)
351 .take(messages.len().saturating_sub(4))
352 .cloned()
353 .collect();
354
355 if middle.is_empty() {
356 // Nothing to summarize
357 "summarize (nothing to compress)"
358 } else {
359 // Build summarization prompt
360 let middle_json = serde_json::to_string_pretty(&middle)
361 .unwrap_or_default();
362 let summary_prompt = format!(
363 "Condense this conversation history into a brief summary while preserving:\n\
364 - Key decisions and conclusions\n\
365 - Important information discovered\n\
366 - Tool calls made and their essential results\n\
367 - Any errors or issues encountered\n\n\
368 Keep the summary concise but complete enough to continue the conversation.\n\n\
369 Conversation to summarize:\n{}\n\n\
370 Respond with ONLY the summary, no preamble.",
371 middle_json
372 );
373
374 // Create a minimal request for summarization
375 let mut summary_req = ModelRequest::new();
376 summary_req.add_user_prompt(summary_prompt);
377
378 // Call the model (non-streaming for simplicity)
379 let summary_params = ModelRequestParameters::new();
380 match model
381 .request(
382 &[summary_req],
383 &model_settings,
384 &summary_params,
385 )
386 .await
387 {
388 Ok(response) => {
389 // Extract text from response
390 let summary_text = response
391 .parts
392 .iter()
393 .filter_map(|p| match p {
394 ModelResponsePart::Text(t) => {
395 Some(t.content.clone())
396 }
397 _ => None,
398 })
399 .collect::<Vec<_>>()
400 .join("\n");
401
402 if !summary_text.is_empty() {
403 // Build new message list: first 2 + summary + last 2
404 let mut new_messages = first_two;
405
406 // Add summary as a "previous context" message
407 let mut summary_msg = ModelRequest::new();
408 summary_msg.add_user_prompt(format!(
409 "[Previous conversation summary]\n{}\n[End of summary - continuing conversation]",
410 summary_text
411 ));
412 new_messages.push(summary_msg);
413
414 new_messages.extend(last_two);
415 messages = new_messages;
416 "summarize"
417 } else {
418 // Fallback to truncate if summary failed
419 use crate::history::{
420 HistoryProcessor, TruncateByTokens,
421 };
422 let truncator = TruncateByTokens::new(
423 compression.target_tokens as u64,
424 )
425 .keep_first_n(2);
426 let temp_ctx =
427 RunContext::new((), &model_name);
428 messages = truncator
429 .process(&temp_ctx, messages)
430 .await;
431 "truncate (summary empty)"
432 }
433 }
434 Err(_e) => {
435 warn!(
436 "Summarization failed, falling back to truncate: {}",
437 _e
438 );
439 use crate::history::{
440 HistoryProcessor, TruncateByTokens,
441 };
442 let truncator = TruncateByTokens::new(
443 compression.target_tokens as u64,
444 )
445 .keep_first_n(2);
446 let temp_ctx = RunContext::new((), &model_name);
447 messages = truncator
448 .process(&temp_ctx, messages)
449 .await;
450 "truncate (summary failed)"
451 }
452 }
453 }
454 }
455 }
456 };
457
458 // Calculate new size
459 let new_bytes = serde_json::to_string(&messages)
460 .map(|s| s.len())
461 .unwrap_or(0);
462 let compressed_tokens = new_bytes / 4;
463
464 // Emit compression event
465 let _ = tx
466 .send(Ok(AgentStreamEvent::ContextCompressed {
467 original_tokens,
468 compressed_tokens,
469 strategy: strategy_name.to_string(),
470 messages_before,
471 messages_after: messages.len(),
472 }))
473 .await;
474 }
475 }
476 }
477 // === End Context Compression ===
478
479 // Make streaming request
480 info!(
481 step = step,
482 message_count = messages.len(),
483 "AgentStream: calling model.request_stream"
484 );
485 let stream_result = model
486 .request_stream(&messages, &model_settings, ¶ms)
487 .await;
488
489 let mut model_stream = match stream_result {
490 Ok(s) => {
491 debug!("AgentStream: model.request_stream succeeded, got stream");
492 s
493 }
494 Err(e) => {
495 error!(error = %e, "AgentStream: model.request_stream failed");
496 let _ = tx
497 .send(Ok(AgentStreamEvent::Error {
498 message: e.to_string(),
499 }))
500 .await;
501 let _ = tx.send(Err(AgentRunError::Model(e))).await;
502 return;
503 }
504 };
505
506 // Collect response parts while streaming
507 let mut response_parts: Vec<ModelResponsePart> = Vec::new();
508 // Track stream events (used by tracing when enabled)
509 let mut stream_event_count = 0u32;
510
511 // Process stream events
512 debug!("AgentStream: starting to process model stream events");
513 while let Some(event_result) = model_stream.next().await {
514 {
515 stream_event_count += 1;
516 let _ = stream_event_count;
517 }
518 match event_result {
519 Ok(event) => {
520 match event {
521 ModelResponseStreamEvent::PartStart(start) => {
522 match &start.part {
523 ModelResponsePart::Text(t) => {
524 if !t.content.is_empty() {
525 let _ = tx
526 .send(Ok(AgentStreamEvent::TextDelta {
527 text: t.content.clone(),
528 }))
529 .await;
530 }
531 }
532 ModelResponsePart::ToolCall(tc) => {
533 let _ = tx
534 .send(Ok(AgentStreamEvent::ToolCallStart {
535 tool_name: tc.tool_name.clone(),
536 tool_call_id: tc.tool_call_id.clone(),
537 }))
538 .await;
539 // If args are already present (non-streaming models),
540 // send them as a delta immediately
541 if let Ok(args_str) = tc.args.to_json_string() {
542 if !args_str.is_empty() && args_str != "{}" {
543 let _ = tx
544 .send(Ok(AgentStreamEvent::ToolCallDelta {
545 delta: args_str,
546 tool_call_id: tc.tool_call_id.clone(),
547 }))
548 .await;
549 }
550 }
551 }
552 ModelResponsePart::Thinking(t) => {
553 if !t.content.is_empty() {
554 let _ = tx
555 .send(Ok(AgentStreamEvent::ThinkingDelta {
556 text: t.content.clone(),
557 }))
558 .await;
559 }
560 }
561 _ => {}
562 }
563 response_parts.push(start.part.clone());
564 }
565 ModelResponseStreamEvent::PartDelta(delta) => {
566 use serdes_ai_core::messages::ModelResponsePartDelta;
567 match &delta.delta {
568 ModelResponsePartDelta::Text(t) => {
569 let _ = tx
570 .send(Ok(AgentStreamEvent::TextDelta {
571 text: t.content_delta.clone(),
572 }))
573 .await;
574 // Update the part
575 if let Some(ModelResponsePart::Text(ref mut text)) =
576 response_parts.get_mut(delta.index)
577 {
578 text.content.push_str(&t.content_delta);
579 }
580 }
581 ModelResponsePartDelta::ToolCall(tc) => {
582 // Get tool_call_id from the existing response part
583 let tool_call_id =
584 response_parts.get(delta.index).and_then(|p| {
585 if let ModelResponsePart::ToolCall(tc) = p {
586 tc.tool_call_id.clone()
587 } else {
588 None
589 }
590 });
591 let _ = tx
592 .send(Ok(AgentStreamEvent::ToolCallDelta {
593 delta: tc.args_delta.clone(),
594 tool_call_id,
595 }))
596 .await;
597 // Update args - accumulate the delta into the tool call
598 if let Some(ModelResponsePart::ToolCall(
599 ref mut tool_call,
600 )) = response_parts.get_mut(delta.index)
601 {
602 tc.apply(tool_call);
603 }
604 }
605 ModelResponsePartDelta::Thinking(t) => {
606 let _ = tx
607 .send(Ok(AgentStreamEvent::ThinkingDelta {
608 text: t.content_delta.clone(),
609 }))
610 .await;
611 if let Some(ModelResponsePart::Thinking(
612 ref mut think,
613 )) = response_parts.get_mut(delta.index)
614 {
615 t.apply(think);
616 }
617 }
618 _ => {}
619 }
620 }
621 ModelResponseStreamEvent::PartEnd(_) => {
622 // Part finished
623 }
624 }
625 }
626 Err(e) => {
627 let _ = tx
628 .send(Ok(AgentStreamEvent::Error {
629 message: e.to_string(),
630 }))
631 .await;
632 let _ = tx.send(Err(AgentRunError::Model(e))).await;
633 return;
634 }
635 }
636 }
637
638 info!(
639 stream_events = stream_event_count,
640 parts = response_parts.len(),
641 "AgentStream: finished processing model stream"
642 );
643
644 // Build the complete response
645 let mut response = ModelResponse {
646 parts: response_parts.clone(),
647 model_name: Some(model.name().to_string()),
648 timestamp: Utc::now(),
649 finish_reason: Some(FinishReason::Stop),
650 usage: None,
651 vendor_id: None,
652 vendor_details: None,
653 kind: "response".to_string(),
654 };
655 canonicalize_tool_call_args_in_response(&mut response);
656
657 finish_reason = response.finish_reason;
658 responses.push(response.clone());
659
660 // Emit ResponseComplete
661 let _ = tx
662 .send(Ok(AgentStreamEvent::ResponseComplete { step }))
663 .await;
664
665 // Check for tool calls that need execution
666 let tool_calls: Vec<_> = response
667 .parts
668 .iter()
669 .filter_map(|p| {
670 if let ModelResponsePart::ToolCall(tc) = p {
671 Some(tc.clone())
672 } else {
673 None
674 }
675 })
676 .collect();
677
678 if !tool_calls.is_empty() {
679 // Add response to messages for proper alternation
680 let mut response_req = ModelRequest::new();
681 response_req
682 .parts
683 .push(ModelRequestPart::ModelResponse(Box::new(response.clone())));
684 messages.push(response_req);
685
686 let mut tool_req = ModelRequest::new();
687
688 for tc in tool_calls {
689 let _ = tx
690 .send(Ok(AgentStreamEvent::ToolCallComplete {
691 tool_name: tc.tool_name.clone(),
692 tool_call_id: tc.tool_call_id.clone(),
693 }))
694 .await;
695
696 usage.record_tool_call();
697
698 // Find the tool by name
699 let tool = tools.iter().find(|t| t.definition.name == tc.tool_name);
700
701 match tool {
702 Some(tool) => {
703 // Create a RunContext for tool execution
704 let tool_ctx =
705 RunContext::with_shared_deps(deps.clone(), model_name.clone())
706 .for_tool(&tc.tool_name, tc.tool_call_id.clone());
707
708 // Execute the tool
709 let result =
710 tool.executor.execute(tc.args.to_json(), &tool_ctx).await;
711
712 match result {
713 Ok(ret) => {
714 let _ = tx
715 .send(Ok(AgentStreamEvent::ToolExecuted {
716 tool_name: tc.tool_name.clone(),
717 tool_call_id: tc.tool_call_id.clone(),
718 success: true,
719 error: None,
720 }))
721 .await;
722
723 // Use ToolReturnPart for successful execution
724 let mut part =
725 ToolReturnPart::new(&tc.tool_name, ret.content);
726 if let Some(id) = tc.tool_call_id.clone() {
727 part = part.with_tool_call_id(id);
728 }
729 tool_req.parts.push(ModelRequestPart::ToolReturn(part));
730 }
731 Err(e) => {
732 let error_msg = e.to_string();
733 let _ = tx
734 .send(Ok(AgentStreamEvent::ToolExecuted {
735 tool_name: tc.tool_name.clone(),
736 tool_call_id: tc.tool_call_id.clone(),
737 success: false,
738 error: Some(error_msg.clone()),
739 }))
740 .await;
741
742 // Use ToolReturnPart with error content for tool errors
743 let mut part = ToolReturnPart::error(
744 &tc.tool_name,
745 format!("Tool error: {}", e),
746 );
747 if let Some(id) = tc.tool_call_id.clone() {
748 part = part.with_tool_call_id(id);
749 }
750 tool_req.parts.push(ModelRequestPart::ToolReturn(part));
751 }
752 }
753 }
754 None => {
755 let error_msg = format!("Unknown tool: {}", tc.tool_name);
756 let _ = tx
757 .send(Ok(AgentStreamEvent::ToolExecuted {
758 tool_name: tc.tool_name.clone(),
759 tool_call_id: tc.tool_call_id.clone(),
760 success: false,
761 error: Some(error_msg.clone()),
762 }))
763 .await;
764
765 // Unknown tool - use ToolReturnPart with error
766 let mut part = ToolReturnPart::error(
767 &tc.tool_name,
768 format!("Unknown tool: {}", tc.tool_name),
769 );
770 if let Some(id) = tc.tool_call_id.clone() {
771 part = part.with_tool_call_id(id);
772 }
773 tool_req.parts.push(ModelRequestPart::ToolReturn(part));
774 }
775 }
776 }
777
778 if !tool_req.parts.is_empty() {
779 messages.push(tool_req);
780 }
781
782 // Continue to let model respond to tool "error"
783 continue;
784 }
785
786 // No tool calls - check finish condition
787 if finish_reason == Some(FinishReason::Stop) {
788 // Add final response to messages for complete history
789 let mut response_req = ModelRequest::new();
790 response_req
791 .parts
792 .push(ModelRequestPart::ModelResponse(Box::new(response.clone())));
793 messages.push(response_req);
794
795 finished = true;
796 let _ = tx.send(Ok(AgentStreamEvent::OutputReady)).await;
797 }
798 }
799
800 // Emit RunComplete
801 let _ = tx
802 .send(Ok(AgentStreamEvent::RunComplete {
803 run_id: run_id_clone,
804 messages,
805 }))
806 .await;
807 });
808
809 Ok(AgentStream {
810 rx,
811 cancel_token: None,
812 })
813 }
814
815 /// Create a new streaming agent run with cancellation support.
816 ///
817 /// The provided `CancellationToken` can be used to cancel the agent run
818 /// mid-execution. When cancelled:
819 /// - The model stream is stopped
820 /// - In-flight tool calls are aborted
821 /// - A `Cancelled` event is emitted with partial results
822 ///
823 /// # Example
824 ///
825 /// ```ignore
826 /// use tokio_util::sync::CancellationToken;
827 ///
828 /// let cancel_token = CancellationToken::new();
829 /// let stream = AgentStream::new_with_cancel(
830 /// &agent,
831 /// "Hello!".into(),
832 /// deps,
833 /// RunOptions::default(),
834 /// cancel_token.clone(),
835 /// ).await?;
836 ///
837 /// // Cancel from another task
838 /// cancel_token.cancel();
839 /// ```
840 pub async fn new_with_cancel<Deps, Output>(
841 agent: &Agent<Deps, Output>,
842 prompt: UserContent,
843 deps: Deps,
844 options: RunOptions,
845 cancel_token: CancellationToken,
846 ) -> Result<Self, AgentRunError>
847 where
848 Deps: Send + Sync + 'static,
849 Output: Send + Sync + 'static,
850 {
851 let run_id = generate_run_id();
852 let (tx, rx) = mpsc::channel(64);
853
854 // Clone what we need for the spawned task
855 let model = agent.model_arc();
856 let model_name = model.name().to_string();
857 let model_settings = options
858 .model_settings
859 .clone()
860 .unwrap_or_else(|| agent.model_settings.clone());
861
862 let static_system_prompt = agent.static_system_prompt().to_string();
863 let tool_definitions = agent.tool_definitions();
864 let _end_strategy = agent.end_strategy;
865 let usage_limits = agent.usage_limits.clone();
866 let run_usage_limits = options.usage_limits.clone();
867 let tools: Vec<RegisteredTool<Deps>> = agent.tools.to_vec();
868 let deps = Arc::new(deps);
869
870 let initial_history = options.message_history.clone();
871 let _metadata = options.metadata.clone();
872 let compression_config = options.compression.clone();
873 let run_id_clone = run_id.clone();
874 let cancel_token_clone = cancel_token.clone();
875
876 debug!(run_id = %run_id, "AgentStream: spawning streaming task with cancellation support");
877
878 tokio::spawn(async move {
879 info!(run_id = %run_id_clone, "AgentStream: task started with cancellation support");
880
881 // Track partial content for cancellation reporting
882 let mut accumulated_text = String::new();
883 let mut accumulated_thinking = String::new();
884 let mut pending_tool_names: Vec<String> = Vec::new();
885
886 // Emit RunStart
887 if tx
888 .send(Ok(AgentStreamEvent::RunStart {
889 run_id: run_id_clone.clone(),
890 }))
891 .await
892 .is_err()
893 {
894 return;
895 }
896
897 // Build initial messages
898 let mut messages = initial_history.unwrap_or_default();
899
900 if !static_system_prompt.is_empty() {
901 let mut req = ModelRequest::new();
902 req.add_system_prompt(static_system_prompt.clone());
903 messages.push(req);
904 }
905
906 let mut user_req = ModelRequest::new();
907 user_req.add_user_prompt(prompt);
908 messages.push(user_req);
909
910 let mut responses: Vec<ModelResponse> = Vec::new();
911 let mut usage = RunUsage::new();
912 let mut step = 0u32;
913 let mut finished = false;
914 let mut finish_reason: Option<FinishReason>;
915
916 // Main agent loop with cancellation support
917 while !finished {
918 // Check for cancellation at the start of each iteration
919 if cancel_token_clone.is_cancelled() {
920 info!(run_id = %run_id_clone, "AgentStream: cancelled at loop start");
921 let _ = tx
922 .send(Ok(AgentStreamEvent::Cancelled {
923 partial_text: if accumulated_text.is_empty() {
924 None
925 } else {
926 Some(accumulated_text)
927 },
928 partial_thinking: if accumulated_thinking.is_empty() {
929 None
930 } else {
931 Some(accumulated_thinking)
932 },
933 pending_tools: pending_tool_names,
934 }))
935 .await;
936 let _ = tx.send(Err(AgentRunError::Cancelled)).await;
937 return;
938 }
939
940 step += 1;
941
942 // Check usage limits
943 if let Some(ref limits) = usage_limits {
944 if let Err(e) = limits.check(&usage) {
945 let _ = tx.send(Err(e.into())).await;
946 return;
947 }
948 }
949
950 if let Some(ref limits) = run_usage_limits {
951 if let Err(e) = limits.check(&usage) {
952 let _ = tx.send(Err(e.into())).await;
953 return;
954 }
955 }
956
957 if tx
958 .send(Ok(AgentStreamEvent::RequestStart { step }))
959 .await
960 .is_err()
961 {
962 return;
963 }
964
965 let params = ModelRequestParameters::new()
966 .with_tools_arc(tool_definitions.clone())
967 .with_allow_text(true);
968
969 // Context size calculation (simplified - full version in main new())
970 let (request_bytes, estimated_tokens) = {
971 let messages_json = serde_json::to_string(&messages).unwrap_or_default();
972 let tools_json = serde_json::to_string(&*tool_definitions).unwrap_or_default();
973 let bytes = messages_json.len() + tools_json.len();
974 (bytes, bytes / 4)
975 };
976
977 let context_limit = model.profile().context_window;
978
979 let _ = tx
980 .send(Ok(AgentStreamEvent::ContextInfo {
981 estimated_tokens,
982 request_bytes,
983 context_limit,
984 }))
985 .await;
986
987 // Context compression (simplified version)
988 if let Some(ref compression) = compression_config {
989 if let Some(limit) = context_limit {
990 let threshold_tokens = (limit as f64 * compression.threshold) as usize;
991 if estimated_tokens > threshold_tokens {
992 use crate::history::{HistoryProcessor, TruncateByTokens};
993 let truncator = TruncateByTokens::new(compression.target_tokens as u64)
994 .keep_first_n(2);
995 let temp_ctx = RunContext::new((), &model_name);
996 messages = truncator.process(&temp_ctx, messages).await;
997 }
998 }
999 }
1000
1001 // Make streaming request with cancellation support
1002 let stream_result = model
1003 .request_stream(&messages, &model_settings, ¶ms)
1004 .await;
1005
1006 let mut model_stream = match stream_result {
1007 Ok(s) => s,
1008 Err(e) => {
1009 let _ = tx
1010 .send(Ok(AgentStreamEvent::Error {
1011 message: e.to_string(),
1012 }))
1013 .await;
1014 let _ = tx.send(Err(AgentRunError::Model(e))).await;
1015 return;
1016 }
1017 };
1018
1019 let mut response_parts: Vec<ModelResponsePart> = Vec::new();
1020
1021 // Process stream events with cancellation check
1022 loop {
1023 tokio::select! {
1024 biased;
1025
1026 _ = cancel_token_clone.cancelled() => {
1027 info!(run_id = %run_id_clone, "AgentStream: cancelled during model stream");
1028 let _ = tx
1029 .send(Ok(AgentStreamEvent::Cancelled {
1030 partial_text: if accumulated_text.is_empty() {
1031 None
1032 } else {
1033 Some(accumulated_text)
1034 },
1035 partial_thinking: if accumulated_thinking.is_empty() {
1036 None
1037 } else {
1038 Some(accumulated_thinking)
1039 },
1040 pending_tools: pending_tool_names,
1041 }))
1042 .await;
1043 let _ = tx.send(Err(AgentRunError::Cancelled)).await;
1044 return;
1045 }
1046
1047 event_result = model_stream.next() => {
1048 match event_result {
1049 Some(Ok(event)) => {
1050 match event {
1051 ModelResponseStreamEvent::PartStart(start) => {
1052 match &start.part {
1053 ModelResponsePart::Text(t) => {
1054 if !t.content.is_empty() {
1055 accumulated_text.push_str(&t.content);
1056 let _ = tx
1057 .send(Ok(AgentStreamEvent::TextDelta {
1058 text: t.content.clone(),
1059 }))
1060 .await;
1061 }
1062 }
1063 ModelResponsePart::ToolCall(tc) => {
1064 pending_tool_names.push(tc.tool_name.clone());
1065 let _ = tx
1066 .send(Ok(AgentStreamEvent::ToolCallStart {
1067 tool_name: tc.tool_name.clone(),
1068 tool_call_id: tc.tool_call_id.clone(),
1069 }))
1070 .await;
1071 if let Ok(args_str) = tc.args.to_json_string() {
1072 if !args_str.is_empty() && args_str != "{}" {
1073 let _ = tx
1074 .send(Ok(AgentStreamEvent::ToolCallDelta {
1075 delta: args_str,
1076 tool_call_id: tc.tool_call_id.clone(),
1077 }))
1078 .await;
1079 }
1080 }
1081 }
1082 ModelResponsePart::Thinking(t) => {
1083 if !t.content.is_empty() {
1084 accumulated_thinking.push_str(&t.content);
1085 let _ = tx
1086 .send(Ok(AgentStreamEvent::ThinkingDelta {
1087 text: t.content.clone(),
1088 }))
1089 .await;
1090 }
1091 }
1092 _ => {}
1093 }
1094 response_parts.push(start.part.clone());
1095 }
1096 ModelResponseStreamEvent::PartDelta(delta) => {
1097 use serdes_ai_core::messages::ModelResponsePartDelta;
1098 match &delta.delta {
1099 ModelResponsePartDelta::Text(t) => {
1100 accumulated_text.push_str(&t.content_delta);
1101 let _ = tx
1102 .send(Ok(AgentStreamEvent::TextDelta {
1103 text: t.content_delta.clone(),
1104 }))
1105 .await;
1106 if let Some(ModelResponsePart::Text(ref mut text)) =
1107 response_parts.get_mut(delta.index)
1108 {
1109 text.content.push_str(&t.content_delta);
1110 }
1111 }
1112 ModelResponsePartDelta::ToolCall(tc) => {
1113 let tool_call_id =
1114 response_parts.get(delta.index).and_then(|p| {
1115 if let ModelResponsePart::ToolCall(tc) = p {
1116 tc.tool_call_id.clone()
1117 } else {
1118 None
1119 }
1120 });
1121 let _ = tx
1122 .send(Ok(AgentStreamEvent::ToolCallDelta {
1123 delta: tc.args_delta.clone(),
1124 tool_call_id,
1125 }))
1126 .await;
1127 if let Some(ModelResponsePart::ToolCall(
1128 ref mut tool_call,
1129 )) = response_parts.get_mut(delta.index)
1130 {
1131 tc.apply(tool_call);
1132 }
1133 }
1134 ModelResponsePartDelta::Thinking(t) => {
1135 accumulated_thinking.push_str(&t.content_delta);
1136 let _ = tx
1137 .send(Ok(AgentStreamEvent::ThinkingDelta {
1138 text: t.content_delta.clone(),
1139 }))
1140 .await;
1141 if let Some(ModelResponsePart::Thinking(
1142 ref mut think,
1143 )) = response_parts.get_mut(delta.index)
1144 {
1145 t.apply(think);
1146 }
1147 }
1148 _ => {}
1149 }
1150 }
1151 ModelResponseStreamEvent::PartEnd(_) => {}
1152 }
1153 }
1154 Some(Err(e)) => {
1155 let _ = tx
1156 .send(Ok(AgentStreamEvent::Error {
1157 message: e.to_string(),
1158 }))
1159 .await;
1160 let _ = tx.send(Err(AgentRunError::Model(e))).await;
1161 return;
1162 }
1163 None => {
1164 // Stream ended normally
1165 break;
1166 }
1167 }
1168 }
1169 }
1170 }
1171
1172 // Build the complete response
1173 let mut response = ModelResponse {
1174 parts: response_parts.clone(),
1175 model_name: Some(model.name().to_string()),
1176 timestamp: Utc::now(),
1177 finish_reason: Some(FinishReason::Stop),
1178 usage: None,
1179 vendor_id: None,
1180 vendor_details: None,
1181 kind: "response".to_string(),
1182 };
1183 canonicalize_tool_call_args_in_response(&mut response);
1184
1185 finish_reason = response.finish_reason;
1186 responses.push(response.clone());
1187
1188 let _ = tx
1189 .send(Ok(AgentStreamEvent::ResponseComplete { step }))
1190 .await;
1191
1192 // Check for tool calls
1193 let tool_calls: Vec<_> = response
1194 .parts
1195 .iter()
1196 .filter_map(|p| {
1197 if let ModelResponsePart::ToolCall(tc) = p {
1198 Some(tc.clone())
1199 } else {
1200 None
1201 }
1202 })
1203 .collect();
1204
1205 if !tool_calls.is_empty() {
1206 let mut response_req = ModelRequest::new();
1207 response_req
1208 .parts
1209 .push(ModelRequestPart::ModelResponse(Box::new(response.clone())));
1210 messages.push(response_req);
1211
1212 let mut tool_req = ModelRequest::new();
1213
1214 for tc in tool_calls {
1215 // Check for cancellation before each tool execution
1216 if cancel_token_clone.is_cancelled() {
1217 info!(run_id = %run_id_clone, "AgentStream: cancelled before tool execution");
1218 let _ = tx
1219 .send(Ok(AgentStreamEvent::Cancelled {
1220 partial_text: if accumulated_text.is_empty() {
1221 None
1222 } else {
1223 Some(accumulated_text)
1224 },
1225 partial_thinking: if accumulated_thinking.is_empty() {
1226 None
1227 } else {
1228 Some(accumulated_thinking)
1229 },
1230 pending_tools: pending_tool_names,
1231 }))
1232 .await;
1233 let _ = tx.send(Err(AgentRunError::Cancelled)).await;
1234 return;
1235 }
1236
1237 let _ = tx
1238 .send(Ok(AgentStreamEvent::ToolCallComplete {
1239 tool_name: tc.tool_name.clone(),
1240 tool_call_id: tc.tool_call_id.clone(),
1241 }))
1242 .await;
1243
1244 usage.record_tool_call();
1245 // Remove from pending after completion
1246 pending_tool_names.retain(|n| n != &tc.tool_name);
1247
1248 let tool = tools.iter().find(|t| t.definition.name == tc.tool_name);
1249
1250 match tool {
1251 Some(tool) => {
1252 let tool_ctx =
1253 RunContext::with_shared_deps(deps.clone(), model_name.clone())
1254 .for_tool(&tc.tool_name, tc.tool_call_id.clone());
1255
1256 let result =
1257 tool.executor.execute(tc.args.to_json(), &tool_ctx).await;
1258
1259 match result {
1260 Ok(ret) => {
1261 let _ = tx
1262 .send(Ok(AgentStreamEvent::ToolExecuted {
1263 tool_name: tc.tool_name.clone(),
1264 tool_call_id: tc.tool_call_id.clone(),
1265 success: true,
1266 error: None,
1267 }))
1268 .await;
1269
1270 let mut part =
1271 ToolReturnPart::new(&tc.tool_name, ret.content);
1272 if let Some(id) = tc.tool_call_id.clone() {
1273 part = part.with_tool_call_id(id);
1274 }
1275 tool_req.parts.push(ModelRequestPart::ToolReturn(part));
1276 }
1277 Err(e) => {
1278 let error_msg = e.to_string();
1279 let _ = tx
1280 .send(Ok(AgentStreamEvent::ToolExecuted {
1281 tool_name: tc.tool_name.clone(),
1282 tool_call_id: tc.tool_call_id.clone(),
1283 success: false,
1284 error: Some(error_msg.clone()),
1285 }))
1286 .await;
1287
1288 let mut part = ToolReturnPart::error(
1289 &tc.tool_name,
1290 format!("Tool error: {}", e),
1291 );
1292 if let Some(id) = tc.tool_call_id.clone() {
1293 part = part.with_tool_call_id(id);
1294 }
1295 tool_req.parts.push(ModelRequestPart::ToolReturn(part));
1296 }
1297 }
1298 }
1299 None => {
1300 let error_msg = format!("Unknown tool: {}", tc.tool_name);
1301 let _ = tx
1302 .send(Ok(AgentStreamEvent::ToolExecuted {
1303 tool_name: tc.tool_name.clone(),
1304 tool_call_id: tc.tool_call_id.clone(),
1305 success: false,
1306 error: Some(error_msg.clone()),
1307 }))
1308 .await;
1309
1310 let mut part = ToolReturnPart::error(
1311 &tc.tool_name,
1312 format!("Unknown tool: {}", tc.tool_name),
1313 );
1314 if let Some(id) = tc.tool_call_id.clone() {
1315 part = part.with_tool_call_id(id);
1316 }
1317 tool_req.parts.push(ModelRequestPart::ToolReturn(part));
1318 }
1319 }
1320 }
1321
1322 if !tool_req.parts.is_empty() {
1323 messages.push(tool_req);
1324 }
1325
1326 continue;
1327 }
1328
1329 if finish_reason == Some(FinishReason::Stop) {
1330 // Add final response to messages for complete history
1331 let mut response_req = ModelRequest::new();
1332 response_req
1333 .parts
1334 .push(ModelRequestPart::ModelResponse(Box::new(response.clone())));
1335 messages.push(response_req);
1336
1337 finished = true;
1338 let _ = tx.send(Ok(AgentStreamEvent::OutputReady)).await;
1339 }
1340 }
1341
1342 let _ = tx
1343 .send(Ok(AgentStreamEvent::RunComplete {
1344 run_id: run_id_clone,
1345 messages,
1346 }))
1347 .await;
1348 });
1349
1350 Ok(AgentStream {
1351 rx,
1352 cancel_token: Some(cancel_token),
1353 })
1354 }
1355
1356 /// Cancel the running agent stream.
1357 ///
1358 /// If this stream was created with cancellation support via
1359 /// [`AgentStream::new_with_cancel`], this will trigger cancellation.
1360 /// The stream will emit a `Cancelled` event with any partial results.
1361 ///
1362 /// If this stream was created without cancellation support (via `new`),
1363 /// this method does nothing.
1364 pub fn cancel(&self) {
1365 if let Some(ref token) = self.cancel_token {
1366 token.cancel();
1367 }
1368 }
1369
1370 /// Check if this stream was cancelled.
1371 ///
1372 /// Returns `true` if a cancellation token was provided and it has been
1373 /// triggered, `false` otherwise.
1374 pub fn is_cancelled(&self) -> bool {
1375 self.cancel_token
1376 .as_ref()
1377 .map(|t| t.is_cancelled())
1378 .unwrap_or(false)
1379 }
1380
1381 /// Get the cancellation token if one was provided.
1382 ///
1383 /// This can be used to share the token with other tasks that need
1384 /// to coordinate cancellation.
1385 pub fn cancellation_token(&self) -> Option<&CancellationToken> {
1386 self.cancel_token.as_ref()
1387 }
1388}
1389
1390impl Stream for AgentStream {
1391 type Item = Result<AgentStreamEvent, AgentRunError>;
1392
1393 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1394 Pin::new(&mut self.rx).poll_recv(cx)
1395 }
1396}
1397
1398#[cfg(test)]
1399mod tests {
1400 use super::*;
1401 use crate::builder::agent;
1402 use futures::{stream, StreamExt};
1403 use serdes_ai_core::messages::{ModelRequestPart, TextPart, ToolCallPart};
1404 use serdes_ai_models::FunctionModel;
1405 use std::sync::{
1406 atomic::{AtomicUsize, Ordering},
1407 Arc,
1408 };
1409
1410 #[test]
1411 fn test_stream_event_debug() {
1412 let event = AgentStreamEvent::TextDelta {
1413 text: "hello".to_string(),
1414 };
1415 let debug = format!("{:?}", event);
1416 assert!(debug.contains("TextDelta"));
1417 }
1418
1419 #[test]
1420 fn test_stream_event_variants() {
1421 let events = [
1422 AgentStreamEvent::RunStart {
1423 run_id: "123".to_string(),
1424 },
1425 AgentStreamEvent::RequestStart { step: 1 },
1426 AgentStreamEvent::TextDelta {
1427 text: "hi".to_string(),
1428 },
1429 AgentStreamEvent::ToolCallStart {
1430 tool_name: "search".to_string(),
1431 tool_call_id: Some("call-1".to_string()),
1432 },
1433 AgentStreamEvent::OutputReady,
1434 AgentStreamEvent::RunComplete {
1435 run_id: "123".to_string(),
1436 messages: vec![],
1437 },
1438 AgentStreamEvent::Cancelled {
1439 partial_text: Some("partial".to_string()),
1440 partial_thinking: None,
1441 pending_tools: vec!["tool1".to_string()],
1442 },
1443 ];
1444
1445 assert_eq!(events.len(), 7);
1446 }
1447
1448 #[test]
1449 fn test_cancelled_event() {
1450 let event = AgentStreamEvent::Cancelled {
1451 partial_text: Some("Hello, I was saying...".to_string()),
1452 partial_thinking: Some("Let me think about this...".to_string()),
1453 pending_tools: vec!["search".to_string(), "fetch".to_string()],
1454 };
1455
1456 let debug = format!("{:?}", event);
1457 assert!(debug.contains("Cancelled"));
1458 assert!(debug.contains("partial_text"));
1459 assert!(debug.contains("pending_tools"));
1460 }
1461
1462 #[test]
1463 fn test_cancelled_event_empty() {
1464 let event = AgentStreamEvent::Cancelled {
1465 partial_text: None,
1466 partial_thinking: None,
1467 pending_tools: vec![],
1468 };
1469
1470 if let AgentStreamEvent::Cancelled {
1471 partial_text,
1472 partial_thinking,
1473 pending_tools,
1474 } = event
1475 {
1476 assert!(partial_text.is_none());
1477 assert!(partial_thinking.is_none());
1478 assert!(pending_tools.is_empty());
1479 } else {
1480 panic!("Expected Cancelled event");
1481 }
1482 }
1483
1484 #[test]
1485 fn test_canonicalize_tool_call_args_in_response_converts_string_args_to_json() {
1486 let mut response = ModelResponse::new();
1487 response.add_part(ModelResponsePart::ToolCall(
1488 serdes_ai_core::messages::ToolCallPart::new(
1489 "demo_tool",
1490 ToolCallArgs::string("{foo: bar,}"),
1491 )
1492 .with_tool_call_id("call_1"),
1493 ));
1494
1495 canonicalize_tool_call_args_in_response(&mut response);
1496
1497 match &response.parts[0] {
1498 ModelResponsePart::ToolCall(tc) => {
1499 assert!(matches!(tc.args, ToolCallArgs::Json(_)));
1500 }
1501 _ => panic!("expected tool call part"),
1502 }
1503 }
1504
1505 #[tokio::test]
1506 async fn test_run_complete_messages_persist_canonical_tool_call_args() {
1507 let call_count = Arc::new(AtomicUsize::new(0));
1508 let model = {
1509 let call_count = Arc::clone(&call_count);
1510 FunctionModel::with_stream(move |_messages, _settings| {
1511 let step = call_count.fetch_add(1, Ordering::SeqCst);
1512 let events = if step == 0 {
1513 vec![
1514 Ok(ModelResponseStreamEvent::part_start(
1515 0,
1516 ModelResponsePart::ToolCall(
1517 ToolCallPart::new("demo_tool", ToolCallArgs::string("{foo: bar,}"))
1518 .with_tool_call_id("call_1"),
1519 ),
1520 )),
1521 Ok(ModelResponseStreamEvent::part_end(0)),
1522 ]
1523 } else {
1524 vec![
1525 Ok(ModelResponseStreamEvent::part_start(
1526 0,
1527 ModelResponsePart::Text(TextPart::new("done")),
1528 )),
1529 Ok(ModelResponseStreamEvent::part_end(0)),
1530 ]
1531 };
1532
1533 Box::pin(stream::iter(events))
1534 })
1535 };
1536
1537 let agent = agent(model)
1538 .tool_fn("demo_tool", "Demo tool", |_ctx, args: serde_json::Value| {
1539 assert!(args.is_object());
1540 Ok(serdes_ai_tools::ToolReturn::text("ok"))
1541 })
1542 .build();
1543
1544 let mut stream = agent
1545 .run_stream("trigger tool then finish", ())
1546 .await
1547 .expect("stream should start");
1548
1549 let mut run_complete_messages = None;
1550 while let Some(event) = stream.next().await {
1551 let event = event.expect("stream event should be ok");
1552 if let AgentStreamEvent::RunComplete { messages, .. } = event {
1553 run_complete_messages = Some(messages);
1554 break;
1555 }
1556 }
1557
1558 let messages = run_complete_messages.expect("expected RunComplete event");
1559
1560 let mut saw_tool_call = false;
1561 for request in &messages {
1562 for request_part in &request.parts {
1563 if let ModelRequestPart::ModelResponse(response) = request_part {
1564 for response_part in &response.parts {
1565 if let ModelResponsePart::ToolCall(tc) = response_part {
1566 saw_tool_call = true;
1567 assert!(
1568 matches!(tc.args, ToolCallArgs::Json(_)),
1569 "tool call args should be canonical JSON in persisted RunComplete messages"
1570 );
1571 }
1572 }
1573 }
1574 }
1575 }
1576
1577 assert!(
1578 saw_tool_call,
1579 "expected at least one tool call in persisted RunComplete messages"
1580 );
1581 }
1582}