1use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
65use std::sync::Arc;
66
67use agentkit_core::{
68 CancellationHandle, DataRef, Delta, FinishReason, Item, ItemKind, MetadataMap, Modality, Part,
69 PartId, PartKind, ReasoningPart, SessionId, StructuredPart, TaskId, TextPart, Timestamp,
70 ToolCallId, ToolCallPart, ToolOutput, ToolResultPart, TurnCancellation, Usage,
71};
72use agentkit_task_manager::{
73 PendingLoopUpdates, SimpleTaskManager, TOOL_RESULT_NOT_STARTED_METADATA_KEY, TaskApproval,
74 TaskLaunchKind, TaskLaunchRequest, TaskManager, TaskResolution, TaskStartContext,
75 TaskStartOutcome, TurnTaskUpdate,
76};
77#[cfg(test)]
78use agentkit_task_manager::{
79 TOOL_RESULT_FAILURE_KIND_METADATA_KEY, TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED,
80};
81#[cfg(test)]
82use agentkit_tools_core::ToolContext;
83use agentkit_tools_core::{
84 AllowAllPermissions, ApprovalDecision, ApprovalRequest, BasicToolExecutor, OwnedToolContext,
85 PermissionChecker, ToolCatalogEvent, ToolError, ToolExecutionScope, ToolExecutor, ToolRequest,
86 ToolResources, ToolSource, ToolSpec,
87};
88use async_trait::async_trait;
89use serde::{Deserialize, Serialize};
90use serde_json::Value;
91use thiserror::Error;
92
93mod retry;
94pub use retry::{
95 ProviderClassification, ProviderFailure, ProviderFailureReason, ProviderRetryEvent,
96 ProviderRoute, RetryAccounting, RetryObserver, RetryProgress, UpstreamErrorKind,
97};
98
99const INTERRUPTED_METADATA_KEY: &str = "agentkit.interrupted";
100const INTERRUPT_REASON_METADATA_KEY: &str = "agentkit.interrupt_reason";
101const INTERRUPT_STAGE_METADATA_KEY: &str = "agentkit.interrupt_stage";
102const USER_CANCELLED_REASON: &str = "user_cancelled";
103const DETACHED_NOTIFICATION_TEXT_MAX_CHARS: usize = 512;
104const DETACHED_TEXT_PREVIEW_MAX_CHARS: usize = 160;
105const DETACHED_CALL_ID_MAX_CHARS: usize = 80;
106const MAX_STREAMED_ASSISTANT_CONTENT_BYTES: usize = 1024 * 1024;
107const MAX_STREAMED_ASSISTANT_CONTENT_PARTS: usize = 256;
108
109pub const PROVIDER_FINISH_REASONS_METADATA_KEY: &str = "agentkit.provider_finish_reasons";
111
112pub fn set_provider_finish_reasons<I, S>(metadata: &mut MetadataMap, reasons: I)
114where
115 I: IntoIterator<Item = S>,
116 S: Into<String>,
117{
118 let mut seen = HashSet::new();
119 let reasons = reasons
120 .into_iter()
121 .map(Into::into)
122 .filter(|reason: &String| !reason.is_empty() && seen.insert(reason.clone()))
123 .map(Value::String)
124 .collect::<Vec<_>>();
125 metadata.remove(PROVIDER_FINISH_REASONS_METADATA_KEY);
126 if !reasons.is_empty() {
127 metadata.insert(
128 PROVIDER_FINISH_REASONS_METADATA_KEY.into(),
129 Value::Array(reasons),
130 );
131 }
132}
133
134fn provider_finish_reasons(metadata: &MetadataMap, fallback: &FinishReason) -> Vec<String> {
135 metadata
136 .get(PROVIDER_FINISH_REASONS_METADATA_KEY)
137 .and_then(Value::as_array)
138 .map(|values| {
139 let mut seen = HashSet::new();
140 values
141 .iter()
142 .filter_map(Value::as_str)
143 .filter(|reason| !reason.is_empty() && seen.insert((*reason).to_owned()))
144 .map(str::to_owned)
145 .collect::<Vec<_>>()
146 })
147 .filter(|reasons| !reasons.is_empty())
148 .unwrap_or_else(|| vec![normalized_finish_reason(fallback).into()])
149}
150
151fn normalized_finish_reason(reason: &FinishReason) -> &str {
152 match reason {
153 FinishReason::Completed => "completed",
154 FinishReason::ToolCall => "tool_call",
155 FinishReason::MaxTokens => "max_tokens",
156 FinishReason::Cancelled => "cancelled",
157 FinishReason::Blocked => "blocked",
158 FinishReason::Error => "error",
159 FinishReason::Other(reason) => reason,
160 }
161}
162
163#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)]
165pub enum MessageCaptureError {
166 #[error("message capture max_messages must be nonzero")]
168 ZeroMessages,
169 #[error("message capture max_bytes must be nonzero")]
171 ZeroBytes,
172}
173
174#[derive(Clone, Copy, Debug, PartialEq, Eq)]
176pub struct MessageCapture {
177 max_messages: usize,
178 max_bytes: usize,
179}
180
181impl MessageCapture {
182 pub fn new(max_messages: usize, max_bytes: usize) -> Result<Self, MessageCaptureError> {
184 if max_messages == 0 {
185 return Err(MessageCaptureError::ZeroMessages);
186 }
187 if max_bytes == 0 {
188 return Err(MessageCaptureError::ZeroBytes);
189 }
190 Ok(Self {
191 max_messages,
192 max_bytes,
193 })
194 }
195
196 pub fn max_messages(self) -> usize {
198 self.max_messages
199 }
200
201 pub fn max_bytes(self) -> usize {
203 self.max_bytes
204 }
205}
206
207#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
213pub struct TelemetryConfig {
214 input_messages: Option<MessageCapture>,
215 output_messages: Option<MessageCapture>,
216}
217
218impl TelemetryConfig {
219 pub fn with_input_messages(mut self, capture: MessageCapture) -> Self {
221 self.input_messages = Some(capture);
222 self
223 }
224
225 pub fn with_output_messages(mut self, capture: MessageCapture) -> Self {
227 self.output_messages = Some(capture);
228 self
229 }
230
231 pub fn without_input_messages(mut self) -> Self {
233 self.input_messages = None;
234 self
235 }
236
237 pub fn without_output_messages(mut self) -> Self {
239 self.output_messages = None;
240 self
241 }
242
243 pub fn input_messages(self) -> Option<MessageCapture> {
245 self.input_messages
246 }
247
248 pub fn output_messages(self) -> Option<MessageCapture> {
250 self.output_messages
251 }
252}
253
254#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
256pub struct SessionConsumerCapabilities {
257 #[serde(default)]
259 pub response_attempt_supersession: bool,
260}
261
262impl SessionConsumerCapabilities {
263 pub fn with_response_attempt_supersession(mut self) -> Self {
265 self.response_attempt_supersession = true;
266 self
267 }
268}
269
270#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
285pub struct SessionConfig {
286 pub session_id: SessionId,
288 pub metadata: MetadataMap,
290 pub cache: Option<PromptCacheRequest>,
292 #[serde(default)]
294 pub consumer_capabilities: SessionConsumerCapabilities,
295}
296
297impl SessionConfig {
298 pub fn new(session_id: impl Into<SessionId>) -> Self {
300 Self {
301 session_id: session_id.into(),
302 metadata: MetadataMap::new(),
303 cache: None,
304 consumer_capabilities: SessionConsumerCapabilities::default(),
305 }
306 }
307
308 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
310 self.metadata = metadata;
311 self
312 }
313
314 pub fn with_cache(mut self, cache: PromptCacheRequest) -> Self {
316 self.cache = Some(cache);
317 self
318 }
319
320 pub fn without_cache(mut self) -> Self {
322 self.cache = None;
323 self
324 }
325
326 pub fn with_response_attempt_supersession(mut self) -> Self {
328 self.consumer_capabilities = self
329 .consumer_capabilities
330 .with_response_attempt_supersession();
331 self
332 }
333}
334
335#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
341pub enum PromptCacheMode {
342 Disabled,
344 #[default]
346 BestEffort,
347 Required,
349}
350
351#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
357pub enum PromptCacheRetention {
358 Default,
360 Short,
362 Extended,
364}
365
366#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
368pub enum PromptCacheStrategy {
369 #[default]
371 Automatic,
372 Explicit {
374 breakpoints: Vec<PromptCacheBreakpoint>,
376 },
377}
378
379impl PromptCacheStrategy {
380 pub fn automatic() -> Self {
383 Self::Automatic
384 }
385
386 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
388 Self::Explicit {
389 breakpoints: breakpoints.into_iter().collect(),
390 }
391 }
392}
393
394#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
396pub enum PromptCacheBreakpoint {
397 ToolsEnd,
399 TranscriptItemEnd { index: usize },
401 TranscriptPartEnd {
407 item_index: usize,
408 part_index: usize,
409 },
410}
411
412impl PromptCacheBreakpoint {
413 pub fn tools_end() -> Self {
415 Self::ToolsEnd
416 }
417
418 pub fn transcript_item_end(index: usize) -> Self {
420 Self::TranscriptItemEnd { index }
421 }
422
423 pub fn transcript_part_end(item_index: usize, part_index: usize) -> Self {
425 Self::TranscriptPartEnd {
426 item_index,
427 part_index,
428 }
429 }
430}
431
432#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
434pub struct PromptCacheRequest {
435 pub mode: PromptCacheMode,
437 pub strategy: PromptCacheStrategy,
439 pub retention: Option<PromptCacheRetention>,
441 pub key: Option<String>,
443}
444
445impl PromptCacheRequest {
446 pub fn automatic() -> Self {
448 Self::best_effort(PromptCacheStrategy::automatic())
449 }
450
451 pub fn automatic_required() -> Self {
453 Self::required(PromptCacheStrategy::automatic())
454 }
455
456 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
458 Self::best_effort(PromptCacheStrategy::explicit(breakpoints))
459 }
460
461 pub fn explicit_required(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
463 Self::required(PromptCacheStrategy::explicit(breakpoints))
464 }
465
466 pub fn disabled() -> Self {
468 Self {
469 mode: PromptCacheMode::Disabled,
470 strategy: PromptCacheStrategy::Automatic,
471 retention: None,
472 key: None,
473 }
474 }
475
476 pub fn best_effort(strategy: PromptCacheStrategy) -> Self {
478 Self {
479 mode: PromptCacheMode::BestEffort,
480 strategy,
481 retention: None,
482 key: None,
483 }
484 }
485
486 pub fn required(strategy: PromptCacheStrategy) -> Self {
488 Self {
489 mode: PromptCacheMode::Required,
490 strategy,
491 retention: None,
492 key: None,
493 }
494 }
495
496 pub fn with_mode(mut self, mode: PromptCacheMode) -> Self {
498 self.mode = mode;
499 self
500 }
501
502 pub fn with_strategy(mut self, strategy: PromptCacheStrategy) -> Self {
504 self.strategy = strategy;
505 self
506 }
507
508 pub fn with_retention(mut self, retention: PromptCacheRetention) -> Self {
510 self.retention = Some(retention);
511 self
512 }
513
514 pub fn with_key(mut self, key: impl Into<String>) -> Self {
516 self.key = Some(key.into());
517 self
518 }
519
520 pub fn without_retention(mut self) -> Self {
522 self.retention = None;
523 self
524 }
525
526 pub fn without_key(mut self) -> Self {
528 self.key = None;
529 self
530 }
531
532 pub fn is_enabled(&self) -> bool {
534 !matches!(self.mode, PromptCacheMode::Disabled)
535 }
536}
537
538#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
544pub struct TurnRequest {
545 pub session_id: SessionId,
547 pub turn_id: agentkit_core::TurnId,
549 pub transcript: Vec<Item>,
551 pub available_tools: Vec<ToolSpec>,
553 pub cache: Option<PromptCacheRequest>,
555 pub metadata: MetadataMap,
557}
558
559#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
564pub struct ModelTurnResult {
565 pub finish_reason: FinishReason,
567 pub output_items: Vec<Item>,
569 pub usage: Option<Usage>,
571 pub metadata: MetadataMap,
573 #[serde(default)]
577 pub model: Option<String>,
578 #[serde(default)]
582 pub response_id: Option<String>,
583}
584
585#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
591pub enum ModelTurnEvent {
592 Delta(Delta),
594 ToolCall(ToolCallPart),
596 Usage(Usage),
598 ResponseAttemptSuperseded,
603 Finished(ModelTurnResult),
605}
606
607#[async_trait]
644pub trait ModelAdapter: Send + Sync {
645 type Session: ModelSession;
647
648 async fn start_session(&self, config: SessionConfig) -> Result<Self::Session, LoopError>;
654
655 fn provider_name(&self) -> Option<&str> {
662 None
663 }
664}
665
666#[async_trait]
673pub trait ModelSession: Send {
674 fn set_retry_observer(&mut self, _observer: Option<Arc<dyn RetryObserver>>) {}
677
678 type Turn: ModelTurn;
680
681 async fn begin_turn(
694 &mut self,
695 request: TurnRequest,
696 cancellation: Option<TurnCancellation>,
697 ) -> Result<Self::Turn, LoopError>;
698
699 fn model_name(&self) -> Option<&str> {
705 None
706 }
707
708 fn provider_name(&self) -> Option<&str> {
713 None
714 }
715}
716
717#[async_trait]
723pub trait ModelTurn: Send {
724 fn on_cancelled(&mut self) {}
728
729 async fn next_event(
738 &mut self,
739 cancellation: Option<TurnCancellation>,
740 ) -> Result<Option<ModelTurnEvent>, LoopError>;
741}
742
743pub trait LoopObserver: Send + Sync {
763 fn handle_event(&self, event: ObservedEvent);
768}
769
770#[derive(Clone, Debug, PartialEq)]
776pub struct ObservedEvent {
777 pub session_id: Arc<SessionId>,
779 pub event: AgentEvent,
781}
782
783pub trait TranscriptObserver: Send + Sync {
818 fn on_transcript_event(&self, event: TranscriptEvent<'_>);
823
824 fn on_transcript_replaced(&self, _session_id: &SessionId, _items: &[Item]) {}
827}
828
829#[derive(Clone, Debug)]
833pub struct TranscriptEvent<'a> {
834 pub session_id: &'a SessionId,
836 pub item: &'a Item,
838}
839
840#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
845#[non_exhaustive]
846pub enum MutationPoint {
847 AfterToolResult,
850 AfterTurnEnded,
853 ContextLengthRejected,
855}
856
857#[derive(Clone, Copy, Debug)]
861pub struct ContextRecoveryPolicy {
862 pub max_attempts: u32,
864 pub budget: std::time::Duration,
866}
867
868#[derive(Clone, Copy, Debug)]
870pub struct ContextRecoveryInfo {
871 pub attempt: u32,
873 pub deadline: std::time::Instant,
875}
876
877pub trait EventEmitter: Send + Sync {
880 fn emit(&self, event: AgentEvent);
882}
883
884#[non_exhaustive]
886pub struct LoopCtx<'a> {
887 pub session_id: &'a SessionId,
889 pub turn_id: Option<&'a agentkit_core::TurnId>,
891 pub point: MutationPoint,
893 pub recovery: Option<ContextRecoveryInfo>,
895 pub cancellation: Option<TurnCancellation>,
897 pub emitter: &'a dyn EventEmitter,
899}
900
901pub struct TranscriptCursor<'a> {
909 items: &'a mut Vec<Item>,
910 pub(crate) dirty: bool,
911}
912
913impl<'a> std::ops::Deref for TranscriptCursor<'a> {
914 type Target = Vec<Item>;
915 fn deref(&self) -> &Vec<Item> {
916 self.items
917 }
918}
919
920impl<'a> std::ops::DerefMut for TranscriptCursor<'a> {
921 fn deref_mut(&mut self) -> &mut Vec<Item> {
922 self.dirty = true;
923 self.items
924 }
925}
926
927#[async_trait]
935pub trait LoopMutator: Send + Sync {
936 async fn mutate(
946 &self,
947 cursor: &mut TranscriptCursor<'_>,
948 ctx: LoopCtx<'_>,
949 ) -> Result<(), LoopError> {
950 let _ = (cursor, ctx);
951 Ok(())
952 }
953}
954
955#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
960#[non_exhaustive]
961pub enum AgentEvent {
962 ProviderRetry(ProviderRetryEvent),
964 RunStarted { session_id: SessionId },
966 TurnStarted {
968 session_id: SessionId,
969 turn_id: agentkit_core::TurnId,
970 },
971 InputAccepted {
973 session_id: SessionId,
974 items: Vec<Item>,
975 },
976 ContentDelta(Delta),
978 ToolCallRequested(ToolCallPart),
980 ToolExecutionStarted(ToolCallPart),
982 ToolExecutionProgress(ToolResultPart),
988 ToolResultReceived(ToolResultPart),
997 ApprovalRequired(ApprovalRequest),
999 ApprovalResolved { approved: bool },
1001 ToolCatalogChanged(ToolCatalogEvent),
1003 MutationStarted {
1006 session_id: SessionId,
1007 turn_id: Option<agentkit_core::TurnId>,
1008 mutator: String,
1009 point: MutationPoint,
1010 },
1011 MutationFinished {
1015 session_id: SessionId,
1016 turn_id: Option<agentkit_core::TurnId>,
1017 mutator: String,
1018 dirty: bool,
1019 metadata: MetadataMap,
1020 },
1021 UsageUpdated(Usage),
1023 ResponseAttemptSuperseded,
1028 Warning { message: String },
1030 RunFailed { message: String },
1032 TurnFinished(TurnResult),
1034}
1035
1036#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1058pub struct PendingApproval {
1059 pub request: ApprovalRequest,
1061}
1062
1063impl std::ops::Deref for PendingApproval {
1064 type Target = ApprovalRequest;
1065 fn deref(&self) -> &ApprovalRequest {
1066 &self.request
1067 }
1068}
1069
1070impl PendingApproval {
1071 pub fn approve<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
1073 let call_id = self
1074 .request
1075 .call_id
1076 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
1077 driver.resolve_approval_for(call_id, ApprovalDecision::Approve)
1078 }
1079
1080 pub fn deny<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
1082 let call_id = self
1083 .request
1084 .call_id
1085 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
1086 driver.resolve_approval_for(call_id, ApprovalDecision::Deny { reason: None })
1087 }
1088
1089 pub fn deny_with_reason<S: ModelSession>(
1091 self,
1092 driver: &mut LoopDriver<S>,
1093 reason: impl Into<String>,
1094 ) -> Result<(), LoopError> {
1095 let call_id = self
1096 .request
1097 .call_id
1098 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
1099 driver.resolve_approval_for(
1100 call_id,
1101 ApprovalDecision::Deny {
1102 reason: Some(reason.into()),
1103 },
1104 )
1105 }
1106
1107 pub fn approve_with_patched_input<S: ModelSession>(
1117 self,
1118 driver: &mut LoopDriver<S>,
1119 input: serde_json::Value,
1120 ) -> Result<(), LoopError> {
1121 let call_id = self
1122 .request
1123 .call_id
1124 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
1125 driver.resolve_approval_for_with_patched_input(call_id, input)
1126 }
1127}
1128
1129#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1154pub struct InputRequest {
1155 pub session_id: SessionId,
1157 pub reason: String,
1159}
1160
1161impl InputRequest {
1162 pub fn submit<S: ModelSession>(
1164 self,
1165 driver: &mut LoopDriver<S>,
1166 items: Vec<Item>,
1167 ) -> Result<(), LoopError> {
1168 driver.submit_input(items)
1169 }
1170}
1171
1172#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1177pub struct TurnResult {
1178 pub turn_id: agentkit_core::TurnId,
1180 pub finish_reason: FinishReason,
1182 pub items: Vec<Item>,
1184 pub usage: Option<Usage>,
1186 pub metadata: MetadataMap,
1188}
1189
1190#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1226pub enum LoopInterrupt {
1227 ApprovalRequest(PendingApproval),
1229 AwaitingInput(InputRequest),
1231 AfterToolResult(ToolRoundInfo),
1241}
1242
1243impl LoopInterrupt {
1244 pub fn is_blocking(&self) -> bool {
1250 matches!(self, LoopInterrupt::ApprovalRequest(_))
1251 }
1252}
1253
1254#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1257pub struct ToolRoundInfo {
1258 pub session_id: SessionId,
1260 pub turn_id: agentkit_core::TurnId,
1262 pub transcript_len: usize,
1264}
1265
1266impl ToolRoundInfo {
1267 pub fn submit<S: ModelSession>(
1270 self,
1271 driver: &mut LoopDriver<S>,
1272 items: Vec<Item>,
1273 ) -> Result<(), LoopError> {
1274 driver.submit_input(items)
1275 }
1276}
1277
1278#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1306pub enum LoopStep {
1307 Interrupt(LoopInterrupt),
1309 Finished(TurnResult),
1311}
1312
1313#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1319pub struct LoopSnapshot {
1320 pub session_id: SessionId,
1322 pub transcript: Vec<Item>,
1324 pub pending_input: Vec<Item>,
1326}
1327
1328#[derive(Clone)]
1329struct PendingApprovalToolCall {
1330 request: ApprovalRequest,
1331 decision: Option<ApprovalDecision>,
1332 surfaced: bool,
1333 presentation_turn_id: agentkit_core::TurnId,
1334 task_id: TaskId,
1335 call: ToolCallPart,
1336 tool_request: ToolRequest,
1337 cancellation: Option<TurnCancellation>,
1338}
1339
1340#[derive(Clone, Default)]
1341struct ActiveToolRound {
1342 presentation_turn_id: agentkit_core::TurnId,
1343 task_turn_id: agentkit_core::TurnId,
1344 pending_calls: VecDeque<(ToolCallPart, ToolRequest)>,
1345 cancellation: Option<TurnCancellation>,
1346 background_pending: bool,
1347 foreground_progressed: bool,
1348}
1349
1350#[derive(Default)]
1351struct DriverLifecycle {
1352 active_turn: Option<agentkit_core::TurnId>,
1353}
1354
1355pub struct Agent<M>
1397where
1398 M: ModelAdapter,
1399{
1400 model: M,
1401 tool_sources: Vec<Arc<dyn ToolSource>>,
1402 tool_executor: Option<Arc<dyn ToolExecutor>>,
1403 task_manager: Arc<dyn TaskManager>,
1404 permissions: Arc<dyn PermissionChecker>,
1405 resources: Arc<dyn ToolResources>,
1406 cancellation: Option<CancellationHandle>,
1407 mutators: Vec<Arc<dyn LoopMutator>>,
1408 context_recovery: Option<ContextRecoveryPolicy>,
1409 observers: Vec<Arc<dyn LoopObserver>>,
1410 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1411 transcript: Vec<Item>,
1412 input: Vec<Item>,
1413 telemetry: TelemetryConfig,
1414}
1415
1416impl<M> Agent<M>
1417where
1418 M: ModelAdapter,
1419{
1420 pub fn builder() -> AgentBuilder<M> {
1422 AgentBuilder::default()
1423 }
1424
1425 pub async fn start(&self, config: SessionConfig) -> Result<LoopDriver<M::Session>, LoopError> {
1441 let session_id = config.session_id.clone();
1442 let allow_supersession = config.consumer_capabilities.response_attempt_supersession;
1443 let default_cache = config.cache.clone();
1444 let mut session = self.model.start_session(config).await?;
1445 if !self.observers.is_empty() {
1446 let observers = self.observers.clone();
1447 let observed_session_id = Arc::new(session_id.clone());
1448 session.set_retry_observer(Some(Arc::new(move |event| {
1449 fan_out_observed_event(
1450 &observers,
1451 &observed_session_id,
1452 AgentEvent::ProviderRetry(event),
1453 );
1454 })));
1455 }
1456 let provider_name = self.model.provider_name().map(str::to_owned);
1457 let tool_executor = self
1458 .tool_executor
1459 .clone()
1460 .unwrap_or_else(|| Arc::new(BasicToolExecutor::new(self.tool_sources.clone())));
1461 let driver = LoopDriver {
1462 session_id: session_id.clone(),
1463 observed_session_id: Arc::new(session_id.clone()),
1464 provider_name,
1465 telemetry: self.telemetry,
1466 default_cache,
1467 next_turn_cache: None,
1468 allow_supersession,
1469 session: Some(session),
1470 tool_executor,
1471 task_manager: self.task_manager.clone(),
1472 permissions: self.permissions.clone(),
1473 resources: self.resources.clone(),
1474 cancellation: self.cancellation.clone(),
1475 mutators: self.mutators.clone(),
1476 context_recovery: self.context_recovery,
1477 observers: self.observers.clone(),
1478 transcript_observers: self.transcript_observers.clone(),
1479 transcript: self.transcript.clone(),
1480 pending_input: self.input.clone(),
1481 pending_approvals: BTreeMap::new(),
1482 pending_approval_order: VecDeque::new(),
1483 active_tool_round: None,
1484 pending_round_resume: None,
1485 pending_loop_updates: VecDeque::new(),
1486 next_turn_index: 1,
1487 lifecycle: DriverLifecycle::default(),
1488 background_call_ids: HashSet::new(),
1489 detached_call_ids: HashSet::new(),
1490 interrupted_background_call_ids: HashSet::new(),
1491 tool_cancellations: HashMap::new(),
1492 };
1493 driver.emit(AgentEvent::RunStarted { session_id });
1494 Ok(driver)
1495 }
1496}
1497
1498pub struct AgentBuilder<M>
1504where
1505 M: ModelAdapter,
1506{
1507 model: Option<M>,
1508 tool_sources: Vec<Arc<dyn ToolSource>>,
1509 tool_executor: Option<Arc<dyn ToolExecutor>>,
1510 task_manager: Option<Arc<dyn TaskManager>>,
1511 permissions: Arc<dyn PermissionChecker>,
1512 resources: Arc<dyn ToolResources>,
1513 cancellation: Option<CancellationHandle>,
1514 mutators: Vec<Arc<dyn LoopMutator>>,
1515 context_recovery: Option<ContextRecoveryPolicy>,
1516 observers: Vec<Arc<dyn LoopObserver>>,
1517 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1518 transcript: Vec<Item>,
1519 input: Vec<Item>,
1520 telemetry: TelemetryConfig,
1521}
1522
1523impl<M> Default for AgentBuilder<M>
1524where
1525 M: ModelAdapter,
1526{
1527 fn default() -> Self {
1528 Self {
1529 model: None,
1530 tool_sources: Vec::new(),
1531 tool_executor: None,
1532 task_manager: None,
1533 permissions: Arc::new(AllowAllPermissions),
1534 resources: Arc::new(()),
1535 cancellation: None,
1536 mutators: Vec::new(),
1537 context_recovery: None,
1538 observers: Vec::new(),
1539 transcript_observers: Vec::new(),
1540 transcript: Vec::new(),
1541 input: Vec::new(),
1542 telemetry: TelemetryConfig::default(),
1543 }
1544 }
1545}
1546
1547impl<M> AgentBuilder<M>
1548where
1549 M: ModelAdapter,
1550{
1551 pub fn model(mut self, model: M) -> Self {
1553 self.model = Some(model);
1554 self
1555 }
1556
1557 pub fn add_tool_source<S: ToolSource + 'static>(mut self, source: S) -> Self {
1570 self.tool_sources.push(Arc::new(source));
1571 self
1572 }
1573
1574 pub fn tool_executor(mut self, executor: impl ToolExecutor + 'static) -> Self {
1580 self.tool_executor = Some(Arc::new(executor));
1581 self
1582 }
1583
1584 pub fn task_manager(mut self, manager: impl TaskManager + 'static) -> Self {
1589 self.task_manager = Some(Arc::new(manager));
1590 self
1591 }
1592
1593 pub fn permissions(mut self, permissions: impl PermissionChecker + 'static) -> Self {
1597 self.permissions = Arc::new(permissions);
1598 self
1599 }
1600
1601 pub fn resources(mut self, resources: impl ToolResources + 'static) -> Self {
1603 self.resources = Arc::new(resources);
1604 self
1605 }
1606
1607 pub fn cancellation(mut self, handle: CancellationHandle) -> Self {
1609 self.cancellation = Some(handle);
1610 self
1611 }
1612
1613 pub fn context_recovery(mut self, policy: ContextRecoveryPolicy) -> Self {
1616 self.context_recovery = Some(policy);
1617 self
1618 }
1619
1620 pub fn mutator<L: LoopMutator + 'static>(mut self, mutator: L) -> Self {
1628 self.mutators.push(Arc::new(mutator));
1629 self
1630 }
1631
1632 pub fn observer<O: LoopObserver + 'static>(mut self, observer: O) -> Self {
1636 self.observers.push(Arc::new(observer));
1637 self
1638 }
1639
1640 pub fn transcript_observer<O: TranscriptObserver + 'static>(mut self, observer: O) -> Self {
1649 self.transcript_observers.push(Arc::new(observer));
1650 self
1651 }
1652
1653 pub fn transcript(mut self, transcript: Vec<Item>) -> Self {
1661 self.transcript = transcript;
1662 self
1663 }
1664
1665 pub fn input(mut self, input: Vec<Item>) -> Self {
1674 self.input = input;
1675 self
1676 }
1677
1678 pub fn telemetry(mut self, telemetry: TelemetryConfig) -> Self {
1681 self.telemetry = telemetry;
1682 self
1683 }
1684
1685 pub fn build(self) -> Result<Agent<M>, LoopError> {
1691 let model = self
1692 .model
1693 .ok_or_else(|| LoopError::InvalidState("model adapter is required".into()))?;
1694 Ok(Agent {
1695 model,
1696 tool_sources: self.tool_sources,
1697 tool_executor: self.tool_executor,
1698 task_manager: self
1699 .task_manager
1700 .unwrap_or_else(|| Arc::new(SimpleTaskManager::new())),
1701 permissions: self.permissions,
1702 resources: self.resources,
1703 cancellation: self.cancellation,
1704 mutators: self.mutators,
1705 context_recovery: self.context_recovery,
1706 observers: self.observers,
1707 transcript_observers: self.transcript_observers,
1708 transcript: self.transcript,
1709 input: self.input,
1710 telemetry: self.telemetry,
1711 })
1712 }
1713}
1714
1715pub struct LoopDriver<S>
1746where
1747 S: ModelSession,
1748{
1749 allow_supersession: bool,
1750 session_id: SessionId,
1751 observed_session_id: Arc<SessionId>,
1752 provider_name: Option<String>,
1753 telemetry: TelemetryConfig,
1754 default_cache: Option<PromptCacheRequest>,
1755 next_turn_cache: Option<PromptCacheRequest>,
1756 session: Option<S>,
1757 tool_executor: Arc<dyn ToolExecutor>,
1758 task_manager: Arc<dyn TaskManager>,
1759 permissions: Arc<dyn PermissionChecker>,
1760 resources: Arc<dyn ToolResources>,
1761 cancellation: Option<CancellationHandle>,
1762 mutators: Vec<Arc<dyn LoopMutator>>,
1763 context_recovery: Option<ContextRecoveryPolicy>,
1764 observers: Vec<Arc<dyn LoopObserver>>,
1765 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1766 transcript: Vec<Item>,
1767 pending_input: Vec<Item>,
1768 pending_approvals: BTreeMap<ToolCallId, PendingApprovalToolCall>,
1769 pending_approval_order: VecDeque<ToolCallId>,
1770 active_tool_round: Option<ActiveToolRound>,
1771 pending_round_resume: Option<agentkit_core::TurnId>,
1772 pending_loop_updates: VecDeque<TaskResolution>,
1773 next_turn_index: u64,
1774 lifecycle: DriverLifecycle,
1775 background_call_ids: HashSet<ToolCallId>,
1777 detached_call_ids: HashSet<ToolCallId>,
1785 interrupted_background_call_ids: HashSet<ToolCallId>,
1789 tool_cancellations: HashMap<ToolCallId, TurnCancellation>,
1790}
1791
1792impl<S> LoopDriver<S>
1793where
1794 S: ModelSession,
1795{
1796 fn execute_tool_span(
1797 &self,
1798 request: &ToolRequest,
1799 turn_id: &agentkit_core::TurnId,
1800 launch_kind: &'static str,
1801 ) -> tracing::Span {
1802 tracing::info_span!(
1803 "agent.execute_tool",
1804 "otel.name" = %format!("execute_tool {}", request.tool_name),
1805 "gen_ai.operation.name" = "execute_tool",
1806 "gen_ai.tool.name" = %request.tool_name,
1807 "gen_ai.tool.call.id" = %request.call_id,
1808 "gen_ai.conversation.id" = %self.session_id,
1809 "error.type" = tracing::field::Empty,
1810 session.id = %self.session_id,
1811 turn.id = %turn_id,
1812 launch_kind = launch_kind,
1813 )
1814 }
1815
1816 fn start_task_via_manager(
1817 &self,
1818 task_id: Option<TaskId>,
1819 tool_request: ToolRequest,
1820 kind: TaskLaunchKind,
1821 cancellation: Option<TurnCancellation>,
1822 ) -> impl std::future::Future<Output = Result<TaskStartOutcome, LoopError>> + Send + 'static
1823 {
1824 let task_manager = self.task_manager.clone();
1825 let tool_executor = self.tool_executor.clone();
1826 let permissions = self.permissions.clone();
1827 let resources = self.resources.clone();
1828 let session_id = self.session_id.clone();
1829 let turn_id = tool_request.turn_id.clone();
1830 let metadata = tool_request.metadata.clone();
1831
1832 async move {
1833 task_manager
1834 .start_task(
1835 TaskLaunchRequest {
1836 task_id,
1837 request: tool_request.clone(),
1838 kind,
1839 },
1840 TaskStartContext {
1841 executor: tool_executor.clone(),
1842 tool_context: {
1843 let execution_scope = ToolExecutionScope {
1844 executor: tool_executor,
1845 session_id: session_id.clone(),
1846 turn_id: turn_id.clone(),
1847 permissions: permissions.clone(),
1848 resources: resources.clone(),
1849 cancellation: cancellation.clone(),
1850 };
1851 OwnedToolContext {
1852 session_id,
1853 turn_id,
1854 metadata,
1855 permissions,
1856 resources,
1857 cancellation,
1858 execution_scope: Some(execution_scope),
1859 approved_request: None,
1860 }
1861 },
1862 },
1863 )
1864 .await
1865 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))
1866 }
1867 }
1868
1869 fn register_tool_cancellation(
1870 &mut self,
1871 call_id: &ToolCallId,
1872 cancellation: Option<TurnCancellation>,
1873 ) {
1874 if let Some(cancellation) = cancellation {
1875 self.tool_cancellations
1876 .insert(call_id.clone(), cancellation);
1877 }
1878 }
1879
1880 fn tool_cancellation_for(
1881 &mut self,
1882 call_id: &ToolCallId,
1883 fallback: Option<TurnCancellation>,
1884 ) -> Option<TurnCancellation> {
1885 self.tool_cancellations.get(call_id).cloned().or(fallback)
1886 }
1887
1888 fn clear_tool_cancellation(&mut self, call_id: &ToolCallId) {
1889 self.tool_cancellations.remove(call_id);
1890 }
1891
1892 fn has_pending_interrupts(&self) -> bool {
1893 !self.pending_approvals.is_empty()
1894 }
1895
1896 fn start_logical_turn(&mut self) -> agentkit_core::TurnId {
1897 if let Some(turn_id) = &self.lifecycle.active_turn {
1898 return turn_id.clone();
1899 }
1900 let turn_id = agentkit_core::TurnId::new(format!("turn-{}", self.next_turn_index));
1901 self.next_turn_index += 1;
1902 self.start_logical_turn_with(turn_id)
1903 }
1904
1905 fn start_logical_turn_with(&mut self, turn_id: agentkit_core::TurnId) -> agentkit_core::TurnId {
1906 if let Some(active_turn) = &self.lifecycle.active_turn {
1907 return active_turn.clone();
1908 }
1909 self.lifecycle.active_turn = Some(turn_id.clone());
1910 self.emit(AgentEvent::TurnStarted {
1911 session_id: self.session_id.clone(),
1912 turn_id: turn_id.clone(),
1913 });
1914 turn_id
1915 }
1916
1917 fn finish_logical_turn(&mut self, result: &TurnResult) {
1918 if self.pending_round_resume.as_ref() == Some(&result.turn_id) {
1919 self.pending_round_resume = None;
1920 }
1921 if self.lifecycle.active_turn.as_ref() == Some(&result.turn_id) {
1922 self.lifecycle.active_turn = None;
1923 self.emit(AgentEvent::TurnFinished(result.clone()));
1924 }
1925 }
1926
1927 fn emit_tool_catalog_events(&mut self, events: Vec<ToolCatalogEvent>) {
1928 for event in events {
1929 self.emit(AgentEvent::ToolCatalogChanged(event));
1930 }
1931 }
1932
1933 fn enqueue_pending_approval(
1934 &mut self,
1935 presentation_turn_id: &agentkit_core::TurnId,
1936 task: TaskApproval,
1937 cancellation: Option<TurnCancellation>,
1938 ) {
1939 let call_id = task.tool_request.call_id.clone();
1940 self.background_call_ids.remove(&call_id);
1941 let cancellation = self.tool_cancellation_for(&call_id, cancellation);
1942 let call = ToolCallPart {
1943 id: call_id.clone(),
1944 name: task.tool_request.tool_name.to_string(),
1945 input: task.tool_request.input.clone(),
1946 metadata: task.tool_request.metadata.clone(),
1947 };
1948 let mut request = task.approval;
1949 request.call_id = Some(call_id.clone());
1950 let pending = PendingApprovalToolCall {
1951 request: request.clone(),
1952 decision: None,
1953 surfaced: false,
1954 presentation_turn_id: presentation_turn_id.clone(),
1955 task_id: task.task_id,
1956 call,
1957 tool_request: task.tool_request,
1958 cancellation,
1959 };
1960 self.pending_approvals.insert(call_id.clone(), pending);
1961 if !self.pending_approval_order.iter().any(|id| id == &call_id) {
1962 self.pending_approval_order.push_back(call_id);
1963 }
1964 self.emit(AgentEvent::ApprovalRequired(request));
1965 }
1966
1967 fn take_next_unsurfaced_approval_interrupt(&mut self) -> Option<LoopStep> {
1968 for call_id in self.pending_approval_order.clone() {
1969 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
1970 continue;
1971 };
1972 if pending.decision.is_none() && !pending.surfaced {
1973 pending.surfaced = true;
1974 return Some(LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(
1975 PendingApproval {
1976 request: pending.request.clone(),
1977 },
1978 )));
1979 }
1980 }
1981 None
1982 }
1983
1984 fn next_unresolved_approval_interrupt(&self) -> Option<LoopStep> {
1985 self.pending_approval_order.iter().find_map(|call_id| {
1986 self.pending_approvals.get(call_id).and_then(|pending| {
1987 pending.decision.is_none().then(|| {
1988 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(PendingApproval {
1989 request: pending.request.clone(),
1990 }))
1991 })
1992 })
1993 })
1994 }
1995
1996 fn take_next_resolved_approval(&mut self) -> Option<PendingApprovalToolCall> {
1997 let call_id = self.pending_approval_order.iter().find_map(|call_id| {
1998 self.pending_approvals
1999 .get(call_id)
2000 .and_then(|pending| pending.decision.as_ref().map(|_| call_id.clone()))
2001 })?;
2002 self.pending_approval_order.retain(|id| id != &call_id);
2003 self.pending_approvals.remove(&call_id)
2004 }
2005
2006 fn queue_resolution_interrupt(
2007 &mut self,
2008 presentation_turn_id: &agentkit_core::TurnId,
2009 resolution: TaskResolution,
2010 cancellation: Option<TurnCancellation>,
2011 ) -> Option<LoopStep> {
2012 match resolution {
2013 TaskResolution::Item(item) => {
2014 self.append_tool_result_item(item);
2015 None
2016 }
2017 TaskResolution::Approval(task) => {
2018 self.enqueue_pending_approval(presentation_turn_id, task, cancellation);
2019 self.take_next_unsurfaced_approval_interrupt()
2020 }
2021 }
2022 }
2023
2024 async fn collect_pending_loop_updates(&mut self) -> Result<(), LoopError> {
2025 let PendingLoopUpdates { resolutions } = self
2026 .task_manager
2027 .take_pending_loop_updates()
2028 .await
2029 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2030 self.pending_loop_updates.extend(resolutions);
2031 Ok(())
2032 }
2033
2034 async fn drain_pending_loop_updates(&mut self) -> Result<(bool, Option<LoopStep>), LoopError> {
2035 self.collect_pending_loop_updates().await?;
2036 let mut resolutions = std::mem::take(&mut self.pending_loop_updates);
2037 if !resolutions.is_empty() {
2038 self.start_logical_turn();
2039 }
2040 let mut saw_items = false;
2041 while let Some(resolution) = resolutions.pop_front() {
2042 match resolution {
2043 TaskResolution::Item(item) => {
2044 self.append_tool_result_item(item);
2045 saw_items = true;
2046 }
2047 TaskResolution::Approval(task) => {
2048 let turn_id = self.start_logical_turn();
2049 self.enqueue_pending_approval(&turn_id, task, None);
2050 }
2051 }
2052 }
2053 if let Some(step) = self.finish_cancelled_pending_approval().await? {
2054 return Ok((saw_items, Some(step)));
2055 }
2056 Ok((saw_items, self.take_next_unsurfaced_approval_interrupt()))
2057 }
2058
2059 async fn finish_cancelled_pending_approval(&mut self) -> Result<Option<LoopStep>, LoopError> {
2060 if self.pending_approvals.is_empty() {
2061 return Ok(None);
2062 }
2063 if !self.pending_approvals.values().any(|pending| {
2064 pending
2065 .cancellation
2066 .as_ref()
2067 .is_some_and(TurnCancellation::is_cancelled)
2068 }) {
2069 return Ok(None);
2070 }
2071 self.cancel_pending_approvals().await
2072 }
2073
2074 async fn run_mutators(
2075 &mut self,
2076 point: MutationPoint,
2077 turn_id: Option<&agentkit_core::TurnId>,
2078 cancellation: Option<TurnCancellation>,
2079 ) -> Result<(), LoopError> {
2080 if self.mutators.is_empty() {
2081 return Ok(());
2082 }
2083 if cancellation
2084 .as_ref()
2085 .is_some_and(TurnCancellation::is_cancelled)
2086 {
2087 return Err(LoopError::Cancelled);
2088 }
2089 let mutators = self.mutators.clone();
2090 let session_id = self.session_id.clone();
2091 let observed_session_id = Arc::clone(&self.observed_session_id);
2092 let observers = self.observers.clone();
2093 let emitter = DriverEmitter {
2094 session_id: &observed_session_id,
2095 observers: &observers,
2096 };
2097 let mut cursor = TranscriptCursor {
2098 items: &mut self.transcript,
2099 dirty: false,
2100 };
2101 for mutator in &mutators {
2102 if cancellation
2103 .as_ref()
2104 .is_some_and(TurnCancellation::is_cancelled)
2105 {
2106 return Err(LoopError::Cancelled);
2107 }
2108 let ctx = LoopCtx {
2109 session_id: &session_id,
2110 turn_id,
2111 point,
2112 recovery: None,
2113 cancellation: cancellation.clone(),
2114 emitter: &emitter,
2115 };
2116 mutator.mutate(&mut cursor, ctx).await?;
2117 }
2118 if cursor.dirty {
2119 validate_transcript_invariants(cursor.items)?;
2120 }
2121 Ok(())
2122 }
2123
2124 async fn recover_context(
2127 &mut self,
2128 error: &LoopError,
2129 turn_id: &agentkit_core::TurnId,
2130 cancellation: Option<TurnCancellation>,
2131 recovery: &mut Option<ContextRecoveryInfo>,
2132 ) -> Result<bool, LoopError> {
2133 let Some(policy) = self.context_recovery else {
2134 return Ok(false);
2135 };
2136 if !error.is_context_length_exceeded()
2137 || policy.max_attempts == 0
2138 || policy.budget.is_zero()
2139 || self.mutators.is_empty()
2140 {
2141 return Ok(false);
2142 }
2143 let previous_attempt = recovery.map_or(0, |info| info.attempt);
2144 if previous_attempt >= policy.max_attempts {
2145 return Ok(false);
2146 }
2147 let deadline = match *recovery {
2148 Some(info) => info.deadline,
2149 None => std::time::Instant::now()
2150 .checked_add(policy.budget)
2151 .ok_or_else(|| {
2152 LoopError::InvalidState("context recovery budget is too large".into())
2153 })?,
2154 };
2155 let info = ContextRecoveryInfo {
2156 attempt: previous_attempt + 1,
2157 deadline,
2158 };
2159 *recovery = Some(info);
2160 let mut candidate = self.transcript.clone();
2161 let emitter = DriverEmitter {
2162 session_id: &self.observed_session_id,
2163 observers: &self.observers,
2164 };
2165 let mutation = async {
2166 let mut cursor = TranscriptCursor {
2167 items: &mut candidate,
2168 dirty: false,
2169 };
2170 for mutator in &self.mutators {
2171 mutator
2172 .mutate(
2173 &mut cursor,
2174 LoopCtx {
2175 session_id: &self.session_id,
2176 turn_id: Some(turn_id),
2177 point: MutationPoint::ContextLengthRejected,
2178 recovery: Some(info),
2179 cancellation: cancellation.clone(),
2180 emitter: &emitter,
2181 },
2182 )
2183 .await?;
2184 }
2185 validate_transcript_invariants(cursor.items)
2186 };
2187 recovery_bound(mutation, Some(info), cancellation.clone()).await?;
2188 if cancellation
2189 .as_ref()
2190 .is_some_and(TurnCancellation::is_cancelled)
2191 {
2192 return Err(LoopError::Cancelled);
2193 }
2194 if std::time::Instant::now() >= deadline {
2195 return Err(LoopError::ContextRecoveryDeadlineExceeded);
2196 }
2197 if candidate == self.transcript || !transcript_has_pending_input(&candidate) {
2198 return Ok(false);
2199 }
2200 self.transcript = candidate;
2201 for observer in &self.transcript_observers {
2202 observer.on_transcript_replaced(&self.session_id, &self.transcript);
2203 }
2204 Ok(true)
2205 }
2206
2207 async fn continue_active_tool_round(&mut self) -> Result<Option<LoopStep>, LoopError> {
2208 let Some((presentation_turn_id, task_turn_id, cancellation)) =
2209 self.active_tool_round.as_ref().map(|active| {
2210 (
2211 active.presentation_turn_id.clone(),
2212 active.task_turn_id.clone(),
2213 active.cancellation.clone(),
2214 )
2215 })
2216 else {
2217 return Ok(None);
2218 };
2219 loop {
2220 if cancellation
2221 .as_ref()
2222 .is_some_and(TurnCancellation::is_cancelled)
2223 {
2224 self.task_manager
2225 .on_turn_interrupted(&task_turn_id)
2226 .await
2227 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2228 self.active_tool_round = None;
2229 return self
2230 .finish_cancelled(presentation_turn_id, Vec::new())
2231 .map(Some);
2232 }
2233
2234 let next_call = self
2235 .active_tool_round
2236 .as_mut()
2237 .and_then(|active| active.pending_calls.pop_front());
2238 if let Some((call, tool_request)) = next_call {
2239 use tracing::Instrument;
2240 self.register_tool_cancellation(&call.id, cancellation.clone());
2241 let dispatch_span =
2242 self.execute_tool_span(&tool_request, &presentation_turn_id, "plain");
2243 match self
2244 .start_task_via_manager(
2245 None,
2246 tool_request.clone(),
2247 TaskLaunchKind::Plain,
2248 cancellation.clone(),
2249 )
2250 .instrument(dispatch_span.clone())
2251 .await?
2252 {
2253 TaskStartOutcome::Ready(resolution) => {
2254 let resolution = *resolution;
2255 match resolution {
2256 TaskResolution::Item(item) => {
2257 if !tool_result_not_started(&item) {
2258 self.emit(AgentEvent::ToolExecutionStarted(call.clone()));
2259 }
2260 if tool_result_is_error(&item) {
2261 dispatch_span.record("error.type", "tool_error");
2262 }
2263 if let Some(active) = self.active_tool_round.as_mut() {
2264 active.foreground_progressed = true;
2265 }
2266 self.append_tool_result_item(item);
2267 }
2268 TaskResolution::Approval(task) => {
2269 self.enqueue_pending_approval(
2270 &presentation_turn_id,
2271 task,
2272 cancellation.clone(),
2273 );
2274 }
2275 }
2276 continue;
2277 }
2278 TaskStartOutcome::Pending { kind, .. } => {
2279 self.emit(AgentEvent::ToolExecutionStarted(call.clone()));
2280 if kind == agentkit_task_manager::TaskKind::Background {
2281 self.append_detach_placeholder(call.id.clone(), &call.name);
2282 if let Some(active) = self.active_tool_round.as_mut() {
2283 active.background_pending = true;
2284 }
2285 }
2286 continue;
2287 }
2288 }
2289 }
2290
2291 match self
2292 .task_manager
2293 .wait_for_turn(&task_turn_id, cancellation.clone())
2294 .await
2295 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?
2296 {
2297 Some(TurnTaskUpdate::Resolution(resolution)) => {
2298 let resolution = *resolution;
2299 match resolution {
2300 TaskResolution::Item(item) => {
2301 if let Some(active) = self.active_tool_round.as_mut() {
2302 active.foreground_progressed = true;
2303 }
2304 self.append_tool_result_item(item);
2305 }
2306 TaskResolution::Approval(task) => {
2307 self.enqueue_pending_approval(
2308 &presentation_turn_id,
2309 task,
2310 cancellation.clone(),
2311 );
2312 }
2313 }
2314 }
2315 Some(TurnTaskUpdate::Detached(snapshot)) => {
2316 self.append_detach_placeholder(snapshot.call_id, &snapshot.tool_name);
2317 if let Some(active) = self.active_tool_round.as_mut() {
2318 active.background_pending = true;
2319 active.foreground_progressed = true;
2320 }
2321 }
2322 None => {
2323 if cancellation
2324 .as_ref()
2325 .is_some_and(TurnCancellation::is_cancelled)
2326 {
2327 self.task_manager
2328 .on_turn_interrupted(&task_turn_id)
2329 .await
2330 .map_err(|error| {
2331 LoopError::Tool(ToolError::Internal(error.to_string()))
2332 })?;
2333 self.active_tool_round = None;
2334 return self
2335 .finish_cancelled(presentation_turn_id, Vec::new())
2336 .map(Some);
2337 }
2338 let active = self.active_tool_round.take().ok_or_else(|| {
2339 LoopError::InvalidState("missing active tool round".into())
2340 })?;
2341 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2342 return Ok(Some(step));
2343 }
2344 if let Some(step) = self.next_unresolved_approval_interrupt() {
2345 return Ok(Some(step));
2346 }
2347 if active.background_pending && !active.foreground_progressed {
2348 return Ok(None);
2349 }
2350 let info = ToolRoundInfo {
2357 session_id: self.session_id.clone(),
2358 turn_id: presentation_turn_id.clone(),
2359 transcript_len: self.transcript.len(),
2360 };
2361 self.pending_round_resume = Some(presentation_turn_id);
2362 return Ok(Some(LoopStep::Interrupt(LoopInterrupt::AfterToolResult(
2363 info,
2364 ))));
2365 }
2366 }
2367 }
2368 }
2369
2370 #[tracing::instrument(
2371 name = "agent.turn",
2372 skip_all,
2373 fields(
2374 otel.name = "invoke_agent",
2375 gen_ai.operation.name = "invoke_agent",
2376 gen_ai.conversation.id = %self.session_id,
2377 gen_ai.provider.name = tracing::field::Empty,
2378 session.id = %self.session_id,
2379 turn.id = %turn_id,
2380 transcript.len = self.transcript.len(),
2381 saw_tool_call = tracing::field::Empty,
2382 finish_reason = tracing::field::Empty,
2383 ),
2384 )]
2385 async fn drive_turn(
2386 &mut self,
2387 turn_id: agentkit_core::TurnId,
2388 mutation_point: MutationPoint,
2389 ) -> Result<LoopStep, LoopError> {
2390 let cancellation = self
2391 .cancellation
2392 .as_ref()
2393 .map(CancellationHandle::checkpoint);
2394 match self
2395 .run_mutators(mutation_point, Some(&turn_id), cancellation.clone())
2396 .await
2397 {
2398 Ok(()) => {}
2399 Err(LoopError::Cancelled) => {
2400 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2401 }
2402 Err(error) => return Err(error),
2403 }
2404
2405 if !transcript_has_pending_input(&self.transcript) {
2411 let turn_result = TurnResult {
2412 turn_id,
2413 finish_reason: FinishReason::Completed,
2414 items: Vec::new(),
2415 usage: None,
2416 metadata: MetadataMap::new(),
2417 };
2418 self.finish_logical_turn(&turn_result);
2419 return Ok(LoopStep::Finished(turn_result));
2420 }
2421
2422 if cancellation
2423 .as_ref()
2424 .is_some_and(TurnCancellation::is_cancelled)
2425 {
2426 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2427 }
2428
2429 let catalog_events = self.tool_executor.drain_catalog_events();
2430 self.emit_tool_catalog_events(catalog_events);
2431
2432 let available_tools = self.tool_executor.specs();
2433 let cache = self
2434 .next_turn_cache
2435 .take()
2436 .or_else(|| self.default_cache.clone());
2437
2438 use tracing::Instrument;
2439 let mut recovery = None;
2440 let (mut result, saw_tool_call, latest_usage, chat_span) = 'generation: loop {
2441 let session = self
2442 .session
2443 .as_mut()
2444 .ok_or_else(|| LoopError::InvalidState("model session is not available".into()))?;
2445
2446 let chat_span = tracing::info_span!(
2455 "chat",
2456 "otel.name" = tracing::field::Empty,
2457 "otel.kind" = "client",
2458 "otel.status_code" = tracing::field::Empty,
2459 "error.type" = tracing::field::Empty,
2460 "gen_ai.operation.name" = "chat",
2461 "gen_ai.provider.name" = tracing::field::Empty,
2462 "gen_ai.conversation.id" = %self.session_id,
2463 "gen_ai.request.model" = tracing::field::Empty,
2464 "gen_ai.response.model" = tracing::field::Empty,
2465 "gen_ai.response.id" = tracing::field::Empty,
2466 "gen_ai.response.finish_reasons" = tracing::field::Empty,
2467 "gen_ai.input.messages" = tracing::field::Empty,
2468 "gen_ai.output.messages" = tracing::field::Empty,
2469 "gen_ai.usage.input_tokens" = tracing::field::Empty,
2470 "gen_ai.usage.output_tokens" = tracing::field::Empty,
2471 "gen_ai.usage.cost" = tracing::field::Empty,
2472 );
2473 let initial_provider_name =
2477 effective_provider_name(session.provider_name(), self.provider_name.as_deref());
2478 if let Some(provider) = &initial_provider_name {
2479 chat_span.record("gen_ai.provider.name", provider.as_str());
2480 tracing::Span::current().record("gen_ai.provider.name", provider.as_str());
2481 }
2482 match session.model_name() {
2483 Some(model) => {
2484 chat_span.record("gen_ai.request.model", model);
2485 chat_span.record("otel.name", format!("chat {model}").as_str());
2486 }
2487 None => {
2488 chat_span.record("otel.name", "chat");
2489 }
2490 }
2491
2492 let request = TurnRequest {
2493 session_id: self.session_id.clone(),
2494 turn_id: turn_id.clone(),
2495 transcript: self.transcript.clone(),
2496 available_tools: available_tools.clone(),
2497 cache: cache.clone(),
2498 metadata: MetadataMap::new(),
2499 };
2500 let mut turn = match recovery_bound(
2501 async {
2502 if let Some(capture) = self.telemetry.input_messages {
2503 record_string_array_attribute(
2504 &chat_span,
2505 "gen_ai.input.messages",
2506 capture_messages(
2507 &request.transcript,
2508 capture,
2509 CaptureOrder::NewestTail,
2510 ),
2511 );
2512 }
2513 self.session
2514 .as_mut()
2515 .expect("session checked above")
2516 .begin_turn(request, cancellation.clone())
2517 .await
2518 }
2519 .instrument(chat_span.clone()),
2520 recovery,
2521 cancellation.clone(),
2522 )
2523 .await
2524 {
2525 Ok(turn) => turn,
2526 Err(LoopError::Cancelled) => {
2527 record_inference_error(&chat_span, &LoopError::Cancelled);
2528 self.task_manager
2529 .on_turn_interrupted(&turn_id)
2530 .await
2531 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2532 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2533 }
2534 Err(error) => {
2535 record_inference_error(&chat_span, &error);
2536 drop(chat_span);
2537 match self
2538 .recover_context(&error, &turn_id, cancellation.clone(), &mut recovery)
2539 .await
2540 {
2541 Ok(true) => continue 'generation,
2542 Ok(false) => return Err(error),
2543 Err(LoopError::Cancelled) => {
2544 self.task_manager
2545 .on_turn_interrupted(&turn_id)
2546 .await
2547 .map_err(|error| {
2548 LoopError::Tool(ToolError::Internal(error.to_string()))
2549 })?;
2550 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2551 }
2552 Err(error) => return Err(error),
2553 }
2554 }
2555 };
2556
2557 let session = self.session.as_ref().expect("session checked above");
2560 let provider_name =
2561 effective_provider_name(session.provider_name(), self.provider_name.as_deref());
2562 if let Some(provider) = &provider_name {
2563 chat_span.record("gen_ai.provider.name", provider.as_str());
2564 tracing::Span::current().record("gen_ai.provider.name", provider.as_str());
2565 }
2566 match session.model_name() {
2567 Some(model) => {
2568 chat_span.record("gen_ai.request.model", model);
2569 chat_span.record("otel.name", format!("chat {model}").as_str());
2570 }
2571 None => {
2572 chat_span.record("otel.name", "chat");
2573 }
2574 }
2575
2576 let mut emitted_output = false;
2577 let mut saw_tool_call = false;
2578 let mut finished_result = None;
2579 let mut latest_usage = None;
2580 let mut streamed_content = cancellation
2581 .is_some()
2582 .then(StreamedAssistantContent::default);
2583
2584 while let Some(event) = match recovery_bound(
2585 turn.next_event(cancellation.clone())
2586 .instrument(chat_span.clone()),
2587 recovery,
2588 cancellation.clone(),
2589 )
2590 .await
2591 {
2592 Ok(event) => event,
2593 Err(LoopError::Cancelled) => {
2594 record_inference_error(&chat_span, &LoopError::Cancelled);
2595 turn.on_cancelled();
2596 self.task_manager
2597 .on_turn_interrupted(&turn_id)
2598 .await
2599 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2600 return self.finish_cancelled(
2601 turn_id,
2602 interrupted_stream_items(streamed_content.as_ref()),
2603 );
2604 }
2605 Err(error) => {
2606 record_inference_error(&chat_span, &error);
2607 drop(chat_span);
2608 drop(turn);
2609 if emitted_output || (latest_usage.is_some() && !self.allow_supersession) {
2610 return Err(error);
2611 }
2612 match self
2613 .recover_context(&error, &turn_id, cancellation.clone(), &mut recovery)
2614 .await
2615 {
2616 Ok(true) => {
2617 if latest_usage.is_some() {
2618 self.emit(AgentEvent::ResponseAttemptSuperseded);
2619 }
2620 continue 'generation;
2621 }
2622 Ok(false) => return Err(error),
2623 Err(LoopError::Cancelled) => {
2624 self.task_manager
2625 .on_turn_interrupted(&turn_id)
2626 .await
2627 .map_err(|error| {
2628 LoopError::Tool(ToolError::Internal(error.to_string()))
2629 })?;
2630 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2631 }
2632 Err(error) => return Err(error),
2633 }
2634 }
2635 } {
2636 emitted_output |= matches!(
2639 event,
2640 ModelTurnEvent::Delta(_)
2641 | ModelTurnEvent::ToolCall(_)
2642 | ModelTurnEvent::Finished(_)
2643 );
2644 let attempt_superseded = matches!(event, ModelTurnEvent::ResponseAttemptSuperseded);
2645 if attempt_superseded {
2646 saw_tool_call = false;
2647 latest_usage = None;
2648 if let Some(content) = &mut streamed_content {
2649 content.reset();
2650 }
2651 self.emit(AgentEvent::ResponseAttemptSuperseded);
2652 }
2653 if cancellation
2654 .as_ref()
2655 .is_some_and(TurnCancellation::is_cancelled)
2656 {
2657 record_inference_error(&chat_span, &LoopError::Cancelled);
2658 turn.on_cancelled();
2659 self.task_manager
2660 .on_turn_interrupted(&turn_id)
2661 .await
2662 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2663 return self.finish_cancelled(
2664 turn_id,
2665 interrupted_stream_items(streamed_content.as_ref()),
2666 );
2667 }
2668 if attempt_superseded {
2669 continue;
2670 }
2671 match event {
2672 ModelTurnEvent::Delta(delta) => {
2673 if let Some(content) = &mut streamed_content {
2674 content.apply_delta(&delta);
2675 }
2676 self.emit(AgentEvent::ContentDelta(delta));
2677 }
2678 ModelTurnEvent::Usage(usage) => {
2679 latest_usage = Some(usage.clone());
2680 self.emit(AgentEvent::UsageUpdated(usage));
2681 }
2682 ModelTurnEvent::ToolCall(call) => {
2683 saw_tool_call = true;
2684 if let Some(content) = &mut streamed_content {
2685 content.commit_tool_call(&call);
2686 }
2687 self.emit(AgentEvent::ToolCallRequested(call));
2688 }
2689 ModelTurnEvent::ResponseAttemptSuperseded => {
2690 unreachable!("response-attempt supersession is handled before cancellation")
2691 }
2692 ModelTurnEvent::Finished(result) => {
2693 finished_result = Some(result);
2694 break;
2695 }
2696 }
2697 }
2698
2699 let result = finished_result.ok_or_else(|| {
2700 let error = LoopError::Provider("model turn ended without a Finished event".into());
2701 record_inference_error(&chat_span, &error);
2702 error
2703 })?;
2704 break 'generation (result, saw_tool_call, latest_usage, chat_span);
2705 };
2706 result.usage = merge_usage(result.usage, latest_usage);
2707 if let Some(model) = &result.model {
2708 chat_span.record("gen_ai.response.model", model.as_str());
2709 }
2710 if let Some(id) = &result.response_id {
2711 chat_span.record("gen_ai.response.id", id.as_str());
2712 }
2713 if let Some(tokens) = result
2714 .usage
2715 .as_ref()
2716 .and_then(|usage| usage.tokens.as_ref())
2717 {
2718 record_token_attribute(&chat_span, "gen_ai.usage.input_tokens", tokens.input_tokens);
2719 record_token_attribute(
2720 &chat_span,
2721 "gen_ai.usage.output_tokens",
2722 tokens.output_tokens,
2723 );
2724 }
2725 if let Some(cost) = result.usage.as_ref().and_then(|usage| usage.cost.as_ref()) {
2726 record_f64_attribute(&chat_span, "gen_ai.usage.cost", cost.amount);
2727 }
2728 record_string_array_attribute(
2729 &chat_span,
2730 "gen_ai.response.finish_reasons",
2731 provider_finish_reasons(&result.metadata, &result.finish_reason),
2732 );
2733 if let Some(capture) = self.telemetry.output_messages {
2734 record_string_array_attribute(
2735 &chat_span,
2736 "gen_ai.output.messages",
2737 capture_messages(&result.output_items, capture, CaptureOrder::OldestHead),
2738 );
2739 }
2740 drop(chat_span);
2741 tracing::Span::current().record("saw_tool_call", saw_tool_call);
2742 tracing::Span::current().record(
2743 "finish_reason",
2744 tracing::field::debug(&result.finish_reason),
2745 );
2746 let now = Timestamp::now();
2747 let usage = result.usage.clone();
2748 let finish_reason = result.finish_reason.clone();
2749 let output_items: Vec<Item> = result
2750 .output_items
2751 .drain(..)
2752 .map(|mut item| {
2753 if matches!(item.kind, ItemKind::Assistant) {
2754 if item.usage.is_none() {
2755 item.usage = usage.clone();
2756 }
2757 if item.finish_reason.is_none() {
2758 item.finish_reason = Some(finish_reason.clone());
2759 }
2760 }
2761 if item.created_at.is_none() {
2762 item.created_at = Some(now);
2763 }
2764 item
2765 })
2766 .collect();
2767 self.extend_transcript(output_items.clone());
2768
2769 if saw_tool_call {
2770 let pending_calls = extract_tool_calls(&output_items)
2771 .into_iter()
2772 .map(|call| {
2773 let tool_request = ToolRequest {
2774 call_id: call.id.clone(),
2775 tool_name: agentkit_tools_core::ToolName::new(call.name.clone()),
2776 input: call.input.clone(),
2777 session_id: self.session_id.clone(),
2778 turn_id: turn_id.clone(),
2779 metadata: call.metadata.clone(),
2780 };
2781 (call, tool_request)
2782 })
2783 .collect();
2784 self.active_tool_round = Some(ActiveToolRound {
2785 presentation_turn_id: turn_id.clone(),
2786 task_turn_id: turn_id.clone(),
2787 pending_calls,
2788 cancellation: cancellation.clone(),
2789 background_pending: false,
2790 foreground_progressed: false,
2791 });
2792 if let Some(step) = self.continue_active_tool_round().await? {
2793 return Ok(step);
2794 }
2795 self.finish_logical_turn(&TurnResult {
2796 turn_id,
2797 finish_reason: result.finish_reason,
2798 items: output_items,
2799 usage: result.usage,
2800 metadata: result.metadata,
2801 });
2802 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
2803 InputRequest {
2804 session_id: self.session_id.clone(),
2805 reason: "driver is waiting for input".into(),
2806 },
2807 )));
2808 }
2809
2810 let turn_result = TurnResult {
2811 turn_id,
2812 finish_reason: result.finish_reason,
2813 items: output_items,
2814 usage: result.usage,
2815 metadata: result.metadata,
2816 };
2817 self.finish_logical_turn(&turn_result);
2818 Ok(LoopStep::Finished(turn_result))
2819 }
2820
2821 async fn resume_after_approval(
2822 &mut self,
2823 pending: PendingApprovalToolCall,
2824 ) -> Result<LoopStep, LoopError> {
2825 let decision = pending
2826 .decision
2827 .clone()
2828 .ok_or_else(|| LoopError::InvalidState("pending approval has no decision".into()))?;
2829
2830 match decision {
2831 ApprovalDecision::Approve => {
2832 use tracing::Instrument;
2833 self.emit(AgentEvent::ToolExecutionStarted(pending.call.clone()));
2834 let dispatch_span = self.execute_tool_span(
2835 &pending.tool_request,
2836 &pending.presentation_turn_id,
2837 "approved",
2838 );
2839 let cancellation = self
2840 .cancellation
2841 .as_ref()
2842 .map(CancellationHandle::checkpoint);
2843 self.register_tool_cancellation(&pending.call.id, cancellation.clone());
2844 let start = self
2845 .start_task_via_manager(
2846 Some(pending.task_id.clone()),
2847 pending.tool_request.clone(),
2848 TaskLaunchKind::Approved(pending.request.clone()),
2849 cancellation.clone(),
2850 )
2851 .instrument(dispatch_span.clone())
2852 .await;
2853 let outcome = match start {
2854 Ok(outcome) => outcome,
2855 Err(error) => {
2856 self.append_tool_result_item(Item {
2857 id: None,
2858 kind: ItemKind::Tool,
2859 parts: vec![Part::ToolResult(ToolResultPart {
2860 call_id: pending.call.id.clone(),
2861 output: ToolOutput::Text(format!(
2862 "approved task failed to start: {error}"
2863 )),
2864 is_error: true,
2865 metadata: pending.call.metadata.clone(),
2866 })],
2867 metadata: MetadataMap::new(),
2868 usage: None,
2869 finish_reason: None,
2870 created_at: None,
2871 });
2872 let turn_id = pending.tool_request.turn_id.clone();
2873 if let Err(cleanup_error) =
2874 self.task_manager.on_turn_interrupted(&turn_id).await
2875 {
2876 tracing::debug!(
2877 %cleanup_error,
2878 %turn_id,
2879 "failed to clean up turn after approved task start error"
2880 );
2881 }
2882 return Err(error);
2883 }
2884 };
2885 match outcome {
2886 TaskStartOutcome::Ready(resolution) => {
2887 let resolution = *resolution;
2888 if let TaskResolution::Item(item) = &resolution
2889 && tool_result_is_error(item)
2890 {
2891 dispatch_span.record("error.type", "tool_error");
2892 }
2893 if let Some(step) = self.queue_resolution_interrupt(
2894 &pending.presentation_turn_id,
2895 resolution,
2896 cancellation,
2897 ) {
2898 return Ok(step);
2899 }
2900 }
2901 TaskStartOutcome::Pending { kind, .. } => {
2902 if kind == agentkit_task_manager::TaskKind::Background {
2903 self.append_detach_placeholder(
2904 pending.call.id.clone(),
2905 &pending.call.name,
2906 );
2907 } else {
2908 self.active_tool_round = Some(ActiveToolRound {
2909 presentation_turn_id: pending.presentation_turn_id.clone(),
2910 task_turn_id: pending.tool_request.turn_id.clone(),
2911 pending_calls: VecDeque::new(),
2912 cancellation: cancellation.clone(),
2913 background_pending: false,
2914 foreground_progressed: false,
2915 });
2916 }
2917 }
2918 }
2919 }
2920 ApprovalDecision::Deny { reason } => {
2921 self.append_tool_result_item(Item {
2922 id: None,
2923 kind: ItemKind::Tool,
2924 parts: vec![Part::ToolResult(ToolResultPart {
2925 call_id: pending.call.id.clone(),
2926 output: ToolOutput::Text(
2927 reason.unwrap_or_else(|| "approval denied".into()),
2928 ),
2929 is_error: true,
2930 metadata: pending.call.metadata.clone(),
2931 })],
2932 metadata: MetadataMap::new(),
2933 usage: None,
2934 finish_reason: None,
2935 created_at: None,
2936 });
2937 }
2938 }
2939
2940 if let Some(step) = self.continue_active_tool_round().await? {
2941 Ok(step)
2942 } else if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2943 Ok(step)
2944 } else if let Some(step) = self.next_unresolved_approval_interrupt() {
2945 Ok(step)
2946 } else {
2947 self.drive_turn(pending.presentation_turn_id, MutationPoint::AfterToolResult)
2948 .await
2949 }
2950 }
2951
2952 fn finish_cancelled(
2953 &mut self,
2954 turn_id: agentkit_core::TurnId,
2955 items: Vec<Item>,
2956 ) -> Result<LoopStep, LoopError> {
2957 let pending = self.drain_pending_approval_items();
2958 self.reject_drained_approvals(pending);
2959 self.extend_transcript(items.clone());
2960 self.close_interrupted_tool_calls();
2961 let turn_result = TurnResult {
2962 turn_id,
2963 finish_reason: FinishReason::Cancelled,
2964 items,
2965 usage: None,
2966 metadata: interrupted_metadata("turn"),
2967 };
2968 self.finish_logical_turn(&turn_result);
2969 Ok(LoopStep::Finished(turn_result))
2970 }
2971
2972 pub fn submit_input(&mut self, input: Vec<Item>) -> Result<(), LoopError> {
2982 if self.has_pending_interrupts() {
2983 return Err(LoopError::InvalidState(
2984 "cannot submit input while an interrupt is pending".into(),
2985 ));
2986 }
2987 self.emit(AgentEvent::InputAccepted {
2988 session_id: self.session_id.clone(),
2989 items: input.clone(),
2990 });
2991 self.pending_input.extend(input);
2992 Ok(())
2993 }
2994
2995 pub fn set_next_turn_cache(&mut self, cache: PromptCacheRequest) -> Result<(), LoopError> {
3000 if self.has_pending_interrupts() {
3001 return Err(LoopError::InvalidState(
3002 "cannot update next-turn cache while an interrupt is pending".into(),
3003 ));
3004 }
3005 self.next_turn_cache = Some(cache);
3006 Ok(())
3007 }
3008
3009 #[cfg(test)]
3010 pub(crate) fn submit_input_with_cache(
3011 &mut self,
3012 input: Vec<Item>,
3013 cache: PromptCacheRequest,
3014 ) -> Result<(), LoopError> {
3015 self.set_next_turn_cache(cache)?;
3016 self.submit_input(input)
3017 }
3018
3019 pub fn resolve_approval_for(
3029 &mut self,
3030 call_id: ToolCallId,
3031 decision: ApprovalDecision,
3032 ) -> Result<(), LoopError> {
3033 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
3034 return Err(LoopError::InvalidState(format!(
3035 "no approval request is pending for call {}",
3036 call_id.0
3037 )));
3038 };
3039 pending.decision = Some(decision.clone());
3040 self.emit(AgentEvent::ApprovalResolved {
3041 approved: matches!(decision, ApprovalDecision::Approve),
3042 });
3043 Ok(())
3044 }
3045
3046 pub fn resolve_approval_for_with_patched_input(
3059 &mut self,
3060 call_id: ToolCallId,
3061 input: serde_json::Value,
3062 ) -> Result<(), LoopError> {
3063 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
3064 return Err(LoopError::InvalidState(format!(
3065 "no approval request is pending for call {}",
3066 call_id.0
3067 )));
3068 };
3069 pending.tool_request.input = input;
3070 self.resolve_approval_for(call_id, ApprovalDecision::Approve)
3071 }
3072
3073 pub fn resolve_approval(&mut self, decision: ApprovalDecision) -> Result<(), LoopError> {
3076 let mut unresolved = self
3077 .pending_approval_order
3078 .iter()
3079 .filter(|call_id| {
3080 self.pending_approvals
3081 .get(*call_id)
3082 .is_some_and(|pending| pending.decision.is_none())
3083 })
3084 .cloned();
3085 let Some(call_id) = unresolved.next() else {
3086 return Err(LoopError::InvalidState(
3087 "no approval request is pending".into(),
3088 ));
3089 };
3090 if unresolved.next().is_some() {
3091 return Err(LoopError::InvalidState(
3092 "multiple approvals are pending; use resolve_approval_for".into(),
3093 ));
3094 }
3095 self.resolve_approval_for(call_id, decision)
3096 }
3097
3098 pub fn cancel_pending_approval_for(&mut self, call_id: ToolCallId) -> Result<(), LoopError> {
3103 let Some(pending) = self.drain_pending_approval_for(&call_id) else {
3104 return Err(LoopError::InvalidState(format!(
3105 "no approval request is pending for call {}",
3106 call_id.0
3107 )));
3108 };
3109 let turn_id = pending.presentation_turn_id.clone();
3110 self.reject_drained_approvals(vec![pending]);
3111 if self.pending_approvals.is_empty() && self.active_tool_round.is_none() {
3112 let _ = self.finish_cancelled(turn_id, Vec::new())?;
3113 }
3114 Ok(())
3115 }
3116
3117 pub async fn cancel_pending_approvals(&mut self) -> Result<Option<LoopStep>, LoopError> {
3123 if self.pending_approvals.is_empty() {
3124 return Ok(None);
3125 }
3126 let Some(turn_id) = self
3127 .pending_approval_order
3128 .iter()
3129 .find_map(|call_id| self.pending_approvals.get(call_id))
3130 .map(|pending| pending.presentation_turn_id.clone())
3131 else {
3132 return Ok(None);
3133 };
3134 let mut seen_turns = HashSet::new();
3135 let mut originating_turns = Vec::new();
3136 for pending in self.pending_approvals.values() {
3137 let originating_turn = pending.tool_request.turn_id.clone();
3138 if seen_turns.insert(originating_turn.clone()) {
3139 originating_turns.push(originating_turn);
3140 }
3141 }
3142
3143 let pending = self.drain_pending_approval_items();
3144 self.active_tool_round = None;
3145 let mut cleanup_error = None;
3146 for originating_turn in originating_turns {
3147 if let Err(error) = self
3148 .task_manager
3149 .on_turn_interrupted(&originating_turn)
3150 .await
3151 && cleanup_error.is_none()
3152 {
3153 cleanup_error = Some(LoopError::Tool(ToolError::Internal(error.to_string())));
3154 }
3155 }
3156 self.reject_drained_approvals(pending);
3157 if let Some(error) = cleanup_error {
3158 self.close_interrupted_tool_calls();
3159 self.finish_logical_turn(&TurnResult {
3160 turn_id,
3161 finish_reason: FinishReason::Error,
3162 items: Vec::new(),
3163 usage: None,
3164 metadata: MetadataMap::new(),
3165 });
3166 return Err(error);
3167 }
3168 self.finish_cancelled(turn_id, Vec::new()).map(Some)
3169 }
3170
3171 pub async fn retire_interrupted_turn(&mut self) -> Result<Option<TurnResult>, LoopError> {
3188 let Some(turn_id) = self.lifecycle.active_turn.clone() else {
3189 return Ok(None);
3190 };
3191 self.pending_round_resume = None;
3194 let cleanup = self.cleanup_interrupted_turn().await;
3195 let result = TurnResult {
3196 turn_id,
3197 finish_reason: FinishReason::Cancelled,
3198 items: Vec::new(),
3199 usage: None,
3200 metadata: interrupted_metadata("turn"),
3201 };
3202 self.finish_logical_turn(&result);
3203 cleanup?;
3204 Ok(Some(result))
3205 }
3206
3207 pub fn snapshot(&self) -> LoopSnapshot {
3209 LoopSnapshot {
3210 session_id: self.session_id.clone(),
3211 transcript: self.transcript.clone(),
3212 pending_input: self.pending_input.clone(),
3213 }
3214 }
3215
3216 pub fn wait_for_loop_update(
3222 &self,
3223 ) -> impl std::future::Future<Output = Result<(), LoopError>> + Send + 'static {
3224 let has_collected_update = !self.pending_loop_updates.is_empty();
3225 let task_manager = self.task_manager.clone();
3226 async move {
3227 if has_collected_update {
3228 return Ok(());
3229 }
3230 task_manager
3231 .wait_for_loop_update()
3232 .await
3233 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))
3234 }
3235 }
3236
3237 pub async fn next(&mut self) -> Result<LoopStep, LoopError> {
3258 if self.lifecycle.active_turn.is_none() {
3259 let continuation_turn = self
3260 .pending_approval_order
3261 .iter()
3262 .find_map(|call_id| self.pending_approvals.get(call_id))
3263 .map(|pending| pending.presentation_turn_id.clone())
3264 .or_else(|| {
3265 self.active_tool_round
3266 .as_ref()
3267 .map(|active| active.presentation_turn_id.clone())
3268 })
3269 .or_else(|| self.pending_round_resume.clone());
3270 if let Some(turn_id) = continuation_turn {
3271 self.start_logical_turn_with(turn_id);
3272 } else if !self.pending_input.is_empty() {
3273 self.start_logical_turn();
3274 }
3275 }
3276
3277 let result = self.next_inner().await;
3278 match &result {
3279 Ok(LoopStep::Finished(turn)) => self.finish_logical_turn(turn),
3280 Err(_) => {
3281 if let Some(turn_id) = self.lifecycle.active_turn.clone() {
3282 if let Err(error) = self.cleanup_interrupted_turn().await {
3283 tracing::debug!(%error, "failed to clean up turn after loop error");
3284 }
3285 self.finish_logical_turn(&TurnResult {
3286 turn_id,
3287 finish_reason: FinishReason::Error,
3288 items: Vec::new(),
3289 usage: None,
3290 metadata: MetadataMap::new(),
3291 });
3292 }
3293 }
3294 _ => {}
3295 }
3296 result
3297 }
3298
3299 async fn cleanup_interrupted_turn(&mut self) -> Result<(), LoopError> {
3300 let mut seen_turns = HashSet::new();
3301 let mut interrupted_turns = Vec::new();
3302 if let Some(active) = self.active_tool_round.take()
3303 && seen_turns.insert(active.task_turn_id.clone())
3304 {
3305 interrupted_turns.push(active.task_turn_id);
3306 }
3307 if let Some(turn_id) = self.pending_round_resume.take()
3308 && seen_turns.insert(turn_id.clone())
3309 {
3310 interrupted_turns.push(turn_id);
3311 }
3312 for pending in self.pending_approvals.values() {
3313 let turn_id = pending.tool_request.turn_id.clone();
3314 if seen_turns.insert(turn_id.clone()) {
3315 interrupted_turns.push(turn_id);
3316 }
3317 }
3318
3319 let pending = self.drain_pending_approval_items();
3320 let mut cleanup_error = None;
3321 for turn_id in interrupted_turns {
3322 if let Err(error) = self.task_manager.on_turn_interrupted(&turn_id).await
3323 && cleanup_error.is_none()
3324 {
3325 cleanup_error = Some(LoopError::Tool(ToolError::Internal(error.to_string())));
3326 }
3327 }
3328 self.reject_drained_approvals(pending);
3329 self.close_interrupted_tool_calls();
3330 match cleanup_error {
3331 Some(error) => Err(error),
3332 None => Ok(()),
3333 }
3334 }
3335
3336 async fn next_inner(&mut self) -> Result<LoopStep, LoopError> {
3337 if let Some(pending) = self.take_next_resolved_approval() {
3338 return self.resume_after_approval(pending).await;
3339 }
3340
3341 if let Some(step) = self.finish_cancelled_pending_approval().await? {
3342 return Ok(step);
3343 }
3344
3345 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
3346 return Ok(step);
3347 }
3348
3349 if let Some(step) = self.next_unresolved_approval_interrupt() {
3350 return Ok(step);
3351 }
3352
3353 if let Some(step) = self.continue_active_tool_round().await? {
3354 return Ok(step);
3355 }
3356
3357 if self.pending_round_resume.is_none() && !self.pending_input.is_empty() {
3362 self.collect_pending_loop_updates().await?;
3365 let turn_id = self.start_logical_turn();
3366 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
3367 self.extend_transcript(drained);
3368 return self
3369 .drive_turn(turn_id, MutationPoint::AfterTurnEnded)
3370 .await;
3371 }
3372
3373 let (had_loop_updates, loop_step) = self.drain_pending_loop_updates().await?;
3374 if let Some(step) = loop_step {
3375 return Ok(step);
3376 }
3377
3378 if let Some(turn_id) = self.pending_round_resume.take() {
3383 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
3384 self.extend_transcript(drained);
3385 return self
3386 .drive_turn(turn_id, MutationPoint::AfterToolResult)
3387 .await;
3388 }
3389
3390 if self.pending_input.is_empty() && !had_loop_updates {
3391 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
3392 InputRequest {
3393 session_id: self.session_id.clone(),
3394 reason: "driver is waiting for input".into(),
3395 },
3396 )));
3397 }
3398
3399 let turn_id = self.start_logical_turn();
3400 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
3401 self.extend_transcript(drained);
3402 self.drive_turn(turn_id, MutationPoint::AfterTurnEnded)
3403 .await
3404 }
3405
3406 fn emit(&self, event: AgentEvent) {
3407 fan_out_observed_event(&self.observers, &self.observed_session_id, event);
3408 }
3409
3410 fn append_item(&mut self, mut item: Item) {
3415 if item.created_at.is_none() {
3416 item.created_at = Some(Timestamp::now());
3417 }
3418 for observer in &self.transcript_observers {
3419 observer.on_transcript_event(TranscriptEvent {
3420 session_id: &self.session_id,
3421 item: &item,
3422 });
3423 }
3424 self.transcript.push(item);
3425 }
3426
3427 fn append_detach_placeholder(&mut self, call_id: ToolCallId, tool_name: &str) {
3428 self.background_call_ids.insert(call_id.clone());
3429 if !self.detached_call_ids.insert(call_id.clone()) {
3430 return;
3431 }
3432 let detached_result = ToolResultPart {
3433 call_id: call_id.clone(),
3434 output: ToolOutput::Text(format!(
3435 "Tool {tool_name} is now running in the background. The result will be delivered when it completes."
3436 )),
3437 is_error: false,
3438 metadata: MetadataMap::new(),
3439 };
3440 self.emit(AgentEvent::ToolExecutionProgress(detached_result.clone()));
3441 self.append_item(Item {
3442 id: None,
3443 kind: ItemKind::Tool,
3444 parts: vec![Part::ToolResult(detached_result)],
3445 metadata: MetadataMap::new(),
3446 usage: None,
3447 finish_reason: None,
3448 created_at: None,
3449 });
3450 }
3451
3452 fn append_tool_result_item(&mut self, item: Item) {
3466 for part in &item.parts {
3467 if let Part::ToolResult(result) = part {
3468 if !self
3469 .interrupted_background_call_ids
3470 .contains(&result.call_id)
3471 {
3472 self.emit(AgentEvent::ToolResultReceived(result.clone()));
3473 }
3474 self.background_call_ids.remove(&result.call_id);
3475 self.clear_tool_cancellation(&result.call_id);
3476 }
3477 }
3478 let item = self.maybe_convert_detached(item);
3479 self.append_item(item);
3480 }
3481
3482 fn drain_pending_approval_for(
3483 &mut self,
3484 call_id: &ToolCallId,
3485 ) -> Option<PendingApprovalToolCall> {
3486 let pending = self.pending_approvals.remove(call_id)?;
3487 self.pending_approval_order.retain(|id| id != call_id);
3488 self.clear_tool_cancellation(call_id);
3489 Some(pending)
3490 }
3491
3492 fn drain_pending_approval_items(&mut self) -> Vec<PendingApprovalToolCall> {
3493 let order = std::mem::take(&mut self.pending_approval_order);
3494 let pending = order
3495 .iter()
3496 .filter_map(|call_id| {
3497 let pending = self.pending_approvals.remove(call_id);
3498 self.clear_tool_cancellation(call_id);
3499 pending
3500 })
3501 .collect();
3502 self.pending_approvals.clear();
3503 pending
3504 }
3505
3506 fn reject_drained_approvals(&mut self, pending: Vec<PendingApprovalToolCall>) {
3507 for pending in pending {
3508 self.emit(AgentEvent::ApprovalResolved { approved: false });
3509 self.append_tool_result_item(cancelled_approval_item(pending));
3510 }
3511 }
3512
3513 fn close_interrupted_tool_calls(&mut self) {
3530 for call in unanswered_tool_calls(&self.transcript) {
3531 let call_id = call.id.clone();
3532 let completes_in_background = self.background_call_ids.contains(&call_id);
3533 self.append_tool_result_item(interrupted_tool_result_item(call));
3534 if completes_in_background {
3535 self.detached_call_ids.insert(call_id.clone());
3536 self.interrupted_background_call_ids.insert(call_id);
3537 }
3538 }
3539 }
3540
3541 fn maybe_convert_detached(&mut self, mut item: Item) -> Item {
3542 if !matches!(item.kind, ItemKind::Tool) {
3543 return item;
3544 }
3545 let results: Vec<&ToolResultPart> = item
3546 .parts
3547 .iter()
3548 .filter_map(|p| match p {
3549 Part::ToolResult(r) => Some(r),
3550 _ => None,
3551 })
3552 .collect();
3553 if results.is_empty()
3554 || !results
3555 .iter()
3556 .all(|r| self.detached_call_ids.contains(&r.call_id))
3557 {
3558 return item;
3559 }
3560 let structured_results = results
3561 .iter()
3562 .map(|result| {
3563 Part::structured(serde_json::to_value(result).unwrap_or_else(
3564 |error| serde_json::json!({ "serialization_error": error.to_string() }),
3565 ))
3566 })
3567 .collect::<Vec<_>>();
3568 let failed = results.iter().filter(|result| result.is_error).count();
3569 let with_metadata = results
3570 .iter()
3571 .filter(|result| !result.metadata.is_empty())
3572 .count();
3573 let mut text = format!(
3574 "Background tool results: {} total, {failed} failed, {with_metadata} with metadata. ",
3575 results.len()
3576 );
3577 for (index, result) in results.iter().enumerate() {
3578 self.detached_call_ids.remove(&result.call_id);
3579 self.interrupted_background_call_ids.remove(&result.call_id);
3580 if text.chars().count() >= DETACHED_NOTIFICATION_TEXT_MAX_CHARS {
3581 continue;
3582 }
3583 if index > 0 {
3584 text.push_str("; ");
3585 }
3586 let label = if result.is_error {
3587 "failed"
3588 } else {
3589 "completed"
3590 };
3591 let call_id = truncate_chars(&result.call_id.0, DETACHED_CALL_ID_MAX_CHARS);
3592 let body = render_tool_output_brief(&result.output);
3593 text.push_str(&format!("{call_id} {label}: {body}"));
3594 }
3595 let text = truncate_chars(&text, DETACHED_NOTIFICATION_TEXT_MAX_CHARS);
3596 let mut notification_parts = Vec::with_capacity(1 + structured_results.len());
3597 notification_parts.push(Part::text(text));
3598 notification_parts.extend(structured_results);
3599 item.kind = ItemKind::Notification;
3600 item.parts = notification_parts;
3601 item
3602 }
3603
3604 fn extend_transcript(&mut self, items: impl IntoIterator<Item = Item>) {
3608 let now = Timestamp::now();
3609 for mut item in items {
3610 if item.created_at.is_none() {
3611 item.created_at = Some(now);
3612 }
3613 self.append_item(item);
3614 }
3615 }
3616}
3617
3618fn render_tool_output_brief(output: &ToolOutput) -> String {
3619 match output {
3620 ToolOutput::Text(text) => format!(
3621 "text preview: {}",
3622 truncate_chars(text, DETACHED_TEXT_PREVIEW_MAX_CHARS)
3623 ),
3624 ToolOutput::Structured(_) => "structured payload".into(),
3625 ToolOutput::Parts(parts) => format!("parts payload ({} parts)", parts.len()),
3626 ToolOutput::Files(files) => format!("files payload ({} files)", files.len()),
3627 }
3628}
3629
3630fn truncate_chars(text: &str, max_chars: usize) -> String {
3631 let mut chars = text.chars();
3632 let mut truncated = chars.by_ref().take(max_chars).collect::<String>();
3633 if chars.next().is_some() && max_chars > 0 {
3634 truncated.pop();
3635 truncated.push('…');
3636 }
3637 truncated
3638}
3639
3640fn interrupted_metadata(stage: &str) -> MetadataMap {
3641 let mut metadata = MetadataMap::new();
3642 metadata.insert(INTERRUPTED_METADATA_KEY.into(), true.into());
3643 metadata.insert(
3644 INTERRUPT_REASON_METADATA_KEY.into(),
3645 USER_CANCELLED_REASON.into(),
3646 );
3647 metadata.insert(INTERRUPT_STAGE_METADATA_KEY.into(), stage.into());
3648 metadata
3649}
3650
3651fn record_token_attribute(span: &tracing::Span, key: &'static str, value: u64) {
3652 match i64::try_from(value) {
3653 Ok(value) => record_i64_attribute(span, key, value),
3654 Err(_) => tracing::warn!(attribute = key, value, "token count exceeds OTEL i64 range"),
3655 }
3656}
3657
3658#[cfg(feature = "otel")]
3659fn record_i64_attribute(span: &tracing::Span, key: &'static str, value: i64) {
3660 use tracing_opentelemetry::OpenTelemetrySpanExt;
3661 span.set_attribute(key, value);
3662}
3663
3664#[cfg(not(feature = "otel"))]
3665fn record_i64_attribute(span: &tracing::Span, key: &'static str, value: i64) {
3666 span.record(key, value);
3667}
3668
3669#[cfg(feature = "otel")]
3670fn record_f64_attribute(span: &tracing::Span, key: &'static str, value: f64) {
3671 use tracing_opentelemetry::OpenTelemetrySpanExt;
3672 span.set_attribute(key, value);
3673}
3674
3675#[cfg(not(feature = "otel"))]
3676fn record_f64_attribute(span: &tracing::Span, key: &'static str, value: f64) {
3677 span.record(key, value);
3678}
3679
3680#[cfg(feature = "otel")]
3681fn otel_string_array(values: Vec<String>) -> opentelemetry::Value {
3682 use opentelemetry::{Array, StringValue, Value as OtelValue};
3683 OtelValue::Array(Array::String(
3684 values.into_iter().map(StringValue::from).collect(),
3685 ))
3686}
3687
3688#[cfg(feature = "otel")]
3689fn record_string_array_attribute(span: &tracing::Span, key: &'static str, values: Vec<String>) {
3690 use tracing_opentelemetry::OpenTelemetrySpanExt;
3691 span.set_attribute(key, otel_string_array(values));
3692}
3693
3694#[cfg(not(feature = "otel"))]
3695fn record_string_array_attribute(span: &tracing::Span, key: &'static str, values: Vec<String>) {
3696 span.record(key, tracing::field::debug(&values));
3697}
3698
3699#[derive(Clone, Copy)]
3700enum CaptureOrder {
3701 NewestTail,
3702 OldestHead,
3703}
3704
3705fn effective_provider_name(
3706 session_provider: Option<&str>,
3707 adapter_provider: Option<&str>,
3708) -> Option<String> {
3709 session_provider.or(adapter_provider).map(str::to_owned)
3710}
3711
3712fn merge_usage(final_usage: Option<Usage>, streamed_usage: Option<Usage>) -> Option<Usage> {
3713 match (final_usage, streamed_usage) {
3714 (None, streamed) => streamed,
3715 (Some(final_usage), None) => Some(final_usage),
3716 (Some(mut final_usage), Some(streamed)) => {
3717 if final_usage.tokens.is_none() {
3718 final_usage.tokens = streamed.tokens;
3719 }
3720 if final_usage.cost.is_none() {
3721 final_usage.cost = streamed.cost;
3722 }
3723 for (key, value) in streamed.metadata {
3724 final_usage.metadata.entry(key).or_insert(value);
3725 }
3726 Some(final_usage)
3727 }
3728 }
3729}
3730
3731fn capture_messages(items: &[Item], capture: MessageCapture, order: CaptureOrder) -> Vec<String> {
3732 let mut captured = Vec::new();
3733 let mut used_bytes = 0usize;
3734 let indices: Box<dyn Iterator<Item = usize>> = match order {
3735 CaptureOrder::NewestTail => Box::new((0..items.len()).rev()),
3736 CaptureOrder::OldestHead => Box::new(0..items.len()),
3737 };
3738
3739 for index in indices.take(capture.max_messages) {
3740 let item = &items[index];
3741 let original_bytes = source_content_bytes(item);
3742 let remaining = capture.max_bytes.saturating_sub(used_bytes);
3743 if original_bytes > remaining {
3744 captured.push(
3745 serde_json::json!({
3746 "type": "truncated",
3747 "original_bytes": original_bytes,
3748 })
3749 .to_string(),
3750 );
3751 break;
3752 }
3753 used_bytes += original_bytes;
3754 captured.push(capture_item_json(item, remaining));
3755 }
3756
3757 if matches!(order, CaptureOrder::NewestTail) {
3758 captured.reverse();
3759 }
3760 captured
3761}
3762
3763fn source_content_bytes(item: &Item) -> usize {
3764 item.parts.iter().fold(0, |total, part| {
3765 total.saturating_add(part_source_content_bytes(part))
3766 })
3767}
3768
3769fn part_source_content_bytes(part: &Part) -> usize {
3770 match part {
3771 Part::Text(text) => text.text.len(),
3772 Part::Media(media) => media.mime_type.len(),
3773 Part::File(file) => file
3774 .name
3775 .as_deref()
3776 .map_or(0, str::len)
3777 .saturating_add(file.mime_type.as_deref().map_or(0, str::len)),
3778 Part::Structured(_) => 0,
3779 Part::Reasoning(reasoning) => reasoning.summary.as_deref().map_or(0, str::len),
3780 Part::ToolCall(call) => call.id.0.len().saturating_add(call.name.len()),
3781 Part::ToolResult(result) => result
3782 .call_id
3783 .0
3784 .len()
3785 .saturating_add(tool_output_source_content_bytes(&result.output)),
3786 Part::Custom(custom) => custom.kind.len(),
3787 }
3788}
3789
3790fn tool_output_source_content_bytes(output: &ToolOutput) -> usize {
3791 match output {
3792 ToolOutput::Text(text) => text.len(),
3793 ToolOutput::Structured(_) | ToolOutput::Parts(_) | ToolOutput::Files(_) => 0,
3794 }
3795}
3796
3797struct CaptureBudget {
3798 remaining: usize,
3799}
3800
3801impl CaptureBudget {
3802 fn text(&mut self, text: &str) -> (String, bool) {
3803 let end = floor_char_boundary(text, self.remaining.min(text.len()));
3804 self.remaining = self.remaining.saturating_sub(end);
3805 (text[..end].to_owned(), end < text.len())
3806 }
3807}
3808
3809fn floor_char_boundary(text: &str, mut index: usize) -> usize {
3810 while index > 0 && !text.is_char_boundary(index) {
3811 index -= 1;
3812 }
3813 index
3814}
3815
3816const MAX_CAPTURED_PARTS_PER_ITEM: usize = 256;
3817
3818fn capture_item_json(item: &Item, max_bytes: usize) -> String {
3819 let mut budget = CaptureBudget {
3820 remaining: max_bytes,
3821 };
3822 let mut parts = item
3823 .parts
3824 .iter()
3825 .take(MAX_CAPTURED_PARTS_PER_ITEM)
3826 .map(|part| sanitized_part(part, &mut budget))
3827 .collect::<Vec<_>>();
3828 if item.parts.len() > parts.len() {
3829 parts.push(serde_json::json!({
3830 "type": "truncated",
3831 "reason": "part_limit",
3832 }));
3833 }
3834 serde_json::json!({
3835 "role": item_kind_name(item.kind),
3836 "parts": parts,
3837 })
3838 .to_string()
3839}
3840
3841fn item_kind_name(kind: ItemKind) -> &'static str {
3842 match kind {
3843 ItemKind::System => "system",
3844 ItemKind::Developer => "developer",
3845 ItemKind::User => "user",
3846 ItemKind::Assistant => "assistant",
3847 ItemKind::Tool => "tool",
3848 ItemKind::Context => "context",
3849 ItemKind::Notification => "notification",
3850 }
3851}
3852
3853fn modality_name(modality: Modality) -> &'static str {
3854 match modality {
3855 Modality::Audio => "audio",
3856 Modality::Image => "image",
3857 Modality::Video => "video",
3858 Modality::Binary => "binary",
3859 }
3860}
3861
3862fn omitted_data_ref(data: &DataRef) -> Value {
3863 let kind = match data {
3864 DataRef::InlineText(_) => "inline_text",
3865 DataRef::InlineBytes(_) => "inline_bytes",
3866 DataRef::Uri(_) => "uri",
3867 DataRef::Handle(_) => "handle",
3868 };
3869 serde_json::json!({ "kind": kind, "omitted": true })
3870}
3871
3872fn bounded_field(text: &str, budget: &mut CaptureBudget) -> Value {
3873 let (text, truncated) = budget.text(text);
3874 serde_json::json!({ "value": text, "truncated": truncated })
3875}
3876
3877fn sanitized_part(part: &Part, budget: &mut CaptureBudget) -> Value {
3878 match part {
3879 Part::Text(text) => serde_json::json!({
3880 "type": "text",
3881 "text": bounded_field(&text.text, budget),
3882 }),
3883 Part::Media(media) => serde_json::json!({
3884 "type": "media",
3885 "modality": modality_name(media.modality),
3886 "mime_type": bounded_field(&media.mime_type, budget),
3887 "data": omitted_data_ref(&media.data),
3888 }),
3889 Part::File(file) => serde_json::json!({
3890 "type": "file",
3891 "name": file.name.as_deref().map(|name| bounded_field(name, budget)),
3892 "mime_type": file.mime_type.as_deref().map(|mime| bounded_field(mime, budget)),
3893 "data": omitted_data_ref(&file.data),
3894 }),
3895 Part::Structured(_) => serde_json::json!({
3896 "type": "structured",
3897 "truncated": true,
3898 }),
3899 Part::Reasoning(reasoning) => serde_json::json!({
3900 "type": "reasoning",
3901 "summary": reasoning.summary.as_deref().map(|summary| bounded_field(summary, budget)),
3902 "redacted": reasoning.redacted,
3903 "data": reasoning.data.as_ref().map(omitted_data_ref),
3904 }),
3905 Part::ToolCall(call) => serde_json::json!({
3906 "type": "tool_call",
3907 "id": bounded_field(&call.id.0, budget),
3908 "name": bounded_field(&call.name, budget),
3909 "input": { "truncated": true },
3910 }),
3911 Part::ToolResult(result) => serde_json::json!({
3912 "type": "tool_result",
3913 "call_id": bounded_field(&result.call_id.0, budget),
3914 "is_error": result.is_error,
3915 "output": sanitized_tool_output(&result.output, budget),
3916 }),
3917 Part::Custom(custom) => serde_json::json!({
3918 "type": "custom",
3919 "kind": bounded_field(&custom.kind, budget),
3920 "data": custom.data.as_ref().map(omitted_data_ref),
3921 "value": custom.value.as_ref().map(|_| serde_json::json!({ "truncated": true })),
3922 }),
3923 }
3924}
3925
3926fn sanitized_tool_output(output: &ToolOutput, budget: &mut CaptureBudget) -> Value {
3927 match output {
3928 ToolOutput::Text(text) => serde_json::json!({
3929 "type": "text",
3930 "text": bounded_field(text, budget),
3931 }),
3932 ToolOutput::Structured(_) => serde_json::json!({
3933 "type": "structured",
3934 "truncated": true,
3935 }),
3936 ToolOutput::Parts(parts) => serde_json::json!({
3937 "type": "parts",
3938 "count": parts.len(),
3939 "truncated": true,
3940 }),
3941 ToolOutput::Files(files) => serde_json::json!({
3942 "type": "files",
3943 "count": files.len(),
3944 "truncated": true,
3945 }),
3946 }
3947}
3948
3949#[cfg(test)]
3950mod telemetry_tests {
3951 use super::*;
3952
3953 #[test]
3954 fn message_capture_is_off_by_default_and_independent() {
3955 let default = TelemetryConfig::default();
3956 assert_eq!(default.input_messages(), None);
3957 assert_eq!(default.output_messages(), None);
3958
3959 let capture = MessageCapture::new(2, 1).unwrap();
3960 let input_only = TelemetryConfig::default().with_input_messages(capture);
3961 assert_eq!(input_only.input_messages().unwrap().max_messages(), 2);
3962 assert_eq!(input_only.input_messages().unwrap().max_bytes(), 1);
3963 assert_eq!(input_only.output_messages(), None);
3964 assert_eq!(
3965 MessageCapture::new(0, 1),
3966 Err(MessageCaptureError::ZeroMessages)
3967 );
3968 assert_eq!(
3969 MessageCapture::new(1, 0),
3970 Err(MessageCaptureError::ZeroBytes)
3971 );
3972 }
3973
3974 #[test]
3975 fn one_source_byte_is_not_rejected_for_json_envelope_overhead() {
3976 let items = vec![Item::text(ItemKind::User, "x")];
3977 let captured = capture_messages(
3978 &items,
3979 MessageCapture::new(1, 1).unwrap(),
3980 CaptureOrder::OldestHead,
3981 );
3982 assert_eq!(captured.len(), 1);
3983 let value: Value = serde_json::from_str(&captured[0]).unwrap();
3984 assert_eq!(value["role"], "user");
3985 assert_eq!(value["parts"][0]["text"]["value"], "x");
3986 assert_eq!(value["parts"][0]["text"]["truncated"], false);
3987 }
3988
3989 #[test]
3990 fn multibyte_source_accounting_preserves_utf8_boundaries() {
3991 let items = vec![Item::text(ItemKind::User, "é")];
3992
3993 let exact = capture_messages(
3994 &items,
3995 MessageCapture::new(1, "é".len()).unwrap(),
3996 CaptureOrder::OldestHead,
3997 );
3998 let value: Value = serde_json::from_str(&exact[0]).unwrap();
3999 assert_eq!(value["parts"][0]["text"]["value"], "é");
4000 assert_eq!(value["parts"][0]["text"]["truncated"], false);
4001
4002 let too_small = capture_messages(
4003 &items,
4004 MessageCapture::new(1, 1).unwrap(),
4005 CaptureOrder::OldestHead,
4006 );
4007 let value: Value = serde_json::from_str(&too_small[0]).unwrap();
4008 assert_eq!(value["type"], "truncated");
4009 assert_eq!(value["original_bytes"], 2);
4010 }
4011
4012 #[test]
4013 fn source_bytes_are_aggregated_without_charging_json_envelopes() {
4014 let items = vec![
4015 Item::text(ItemKind::User, "a"),
4016 Item::text(ItemKind::Assistant, "b"),
4017 Item::text(ItemKind::User, "cd"),
4018 ];
4019 let captured = capture_messages(
4020 &items,
4021 MessageCapture::new(3, 3).unwrap(),
4022 CaptureOrder::OldestHead,
4023 );
4024 assert_eq!(captured.len(), 3);
4025 let values = captured
4026 .iter()
4027 .map(|encoded| serde_json::from_str::<Value>(encoded).unwrap())
4028 .collect::<Vec<_>>();
4029 assert_eq!(values[0]["parts"][0]["text"]["value"], "a");
4030 assert_eq!(values[1]["parts"][0]["text"]["value"], "b");
4031 assert_eq!(values[2]["type"], "truncated");
4032 assert_eq!(values[2]["original_bytes"], 2);
4033 }
4034
4035 #[test]
4036 fn tiny_limits_emit_valid_structured_source_byte_truncation() {
4037 let items = vec![Item::text(ItemKind::User, "x".repeat(1_000))];
4038 let captured = capture_messages(
4039 &items,
4040 MessageCapture::new(1, 1).unwrap(),
4041 CaptureOrder::OldestHead,
4042 );
4043 assert_eq!(captured.len(), 1);
4044 let value: Value = serde_json::from_str(&captured[0]).unwrap();
4045 assert_eq!(value["type"], "truncated");
4046 assert_eq!(value["original_bytes"], 1_000);
4047 }
4048
4049 #[test]
4050 fn provider_finish_reason_metadata_round_trips() {
4051 let mut metadata = MetadataMap::new();
4052 set_provider_finish_reasons(&mut metadata, ["end_turn", "", "tool_use", "end_turn"]);
4053 assert_eq!(
4054 provider_finish_reasons(&metadata, &FinishReason::Completed),
4055 vec!["end_turn", "tool_use"]
4056 );
4057 set_provider_finish_reasons(&mut metadata, std::iter::empty::<String>());
4058 assert!(!metadata.contains_key(PROVIDER_FINISH_REASONS_METADATA_KEY));
4059 let fallbacks = [
4060 (FinishReason::Completed, "completed"),
4061 (FinishReason::ToolCall, "tool_call"),
4062 (FinishReason::MaxTokens, "max_tokens"),
4063 (FinishReason::Cancelled, "cancelled"),
4064 (FinishReason::Blocked, "blocked"),
4065 (FinishReason::Error, "error"),
4066 (FinishReason::Other("native".into()), "native"),
4067 ];
4068 for (reason, expected) in fallbacks {
4069 assert_eq!(
4070 provider_finish_reasons(&MetadataMap::new(), &reason),
4071 [expected]
4072 );
4073 }
4074 }
4075
4076 #[test]
4077 fn input_is_newest_tail_output_is_head_and_data_refs_are_omitted() {
4078 let items = vec![
4079 Item::text(ItemKind::User, "old"),
4080 Item::text(ItemKind::User, "middle"),
4081 Item::text(ItemKind::User, "new"),
4082 ];
4083 let capture = MessageCapture::new(2, 10_000).unwrap();
4084 let input = capture_messages(&items, capture, CaptureOrder::NewestTail);
4085 assert!(input[0].contains("middle"));
4086 assert!(input[1].contains("new"));
4087 let output = capture_messages(&items, capture, CaptureOrder::OldestHead);
4088 assert!(output[0].contains("old"));
4089 assert!(output[1].contains("middle"));
4090
4091 let media = Item::new(
4092 ItemKind::User,
4093 vec![Part::media(
4094 agentkit_core::Modality::Image,
4095 "image/png",
4096 agentkit_core::DataRef::uri("https://secret.invalid/image.png"),
4097 )],
4098 );
4099 let encoded = capture_item_json(&media, 10_000);
4100 assert!(!encoded.contains("secret.invalid"));
4101 assert!(encoded.contains("omitted"));
4102 }
4103
4104 #[test]
4105 fn final_usage_wins_and_streamed_usage_fills_only_missing_fields() {
4106 let mut final_metadata = MetadataMap::new();
4107 final_metadata.insert("shared".into(), serde_json::json!("final"));
4108 let final_usage = Usage {
4109 tokens: Some(agentkit_core::TokenUsage::new(1, 2)),
4110 cost: None,
4111 metadata: final_metadata,
4112 };
4113 let mut streamed_metadata = MetadataMap::new();
4114 streamed_metadata.insert("shared".into(), serde_json::json!("streamed"));
4115 streamed_metadata.insert("stream_only".into(), serde_json::json!(true));
4116 let streamed_usage = Usage {
4117 tokens: Some(agentkit_core::TokenUsage::new(10, 20)),
4118 cost: Some(agentkit_core::CostUsage::new(0.5, "USD")),
4119 metadata: streamed_metadata,
4120 };
4121 let merged = merge_usage(Some(final_usage), Some(streamed_usage)).unwrap();
4122 assert_eq!(merged.tokens.unwrap().input_tokens, 1);
4123 assert_eq!(merged.cost.unwrap().amount, 0.5);
4124 assert_eq!(merged.metadata["shared"], "final");
4125 assert_eq!(merged.metadata["stream_only"], true);
4126 }
4127
4128 #[test]
4129 fn session_provider_precedes_adapter_fallback() {
4130 assert_eq!(
4131 effective_provider_name(Some("session"), Some("adapter")).as_deref(),
4132 Some("session")
4133 );
4134 assert_eq!(
4135 effective_provider_name(None, Some("adapter")).as_deref(),
4136 Some("adapter")
4137 );
4138 }
4139}
4140
4141#[cfg(all(test, feature = "otel"))]
4142mod true_otel_integration_tests {
4143 use std::fs;
4144 use std::process::Command;
4145 use std::sync::atomic::{AtomicU64, Ordering};
4146 use std::time::{SystemTime, UNIX_EPOCH};
4147
4148 static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0);
4149
4150 #[test]
4151 fn actual_layer_exports_driven_loop_spans() {
4152 let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
4153 .parent()
4154 .unwrap()
4155 .parent()
4156 .unwrap();
4157 let nonce = SystemTime::now()
4158 .duration_since(UNIX_EPOCH)
4159 .unwrap()
4160 .as_nanos();
4161 let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed);
4162 let temp = std::env::temp_dir().join(format!(
4163 "agentkit-loop-true-otel-{}-{nonce}-{sequence}",
4164 std::process::id()
4165 ));
4166 fs::create_dir(&temp).unwrap();
4167 fs::create_dir(temp.join("src")).unwrap();
4168 let core_path = toml_path(&root.join("crates/agentkit-core"));
4169 let loop_path = toml_path(&root.join("crates/agentkit-loop"));
4170 let async_trait_version = locked_version(root, "async-trait");
4173 let opentelemetry_version = locked_version(root, "opentelemetry");
4174 let serde_json_version = locked_version(root, "serde_json");
4175 let tokio_version = locked_version(root, "tokio");
4176 let tracing_version = locked_version(root, "tracing");
4177 let tracing_otel_version = locked_version(root, "tracing-opentelemetry");
4178 let tracing_subscriber_version = locked_version(root, "tracing-subscriber");
4179 let manifest = format!(
4180 r#"[package]
4181name = "agentkit-loop-true-otel-test"
4182version = "0.0.0"
4183edition = "2024"
4184
4185[dependencies]
4186agentkit-core = {{ path = {core_path} }}
4187agentkit-loop = {{ path = {loop_path}, features = ["otel"] }}
4188async-trait = "={async_trait_version}"
4189opentelemetry = {{ version = "={opentelemetry_version}", default-features = false }}
4190serde_json = "={serde_json_version}"
4191tokio = {{ version = "={tokio_version}", features = ["rt"] }}
4192tracing = "={tracing_version}"
4193tracing-opentelemetry = {{ version = "={tracing_otel_version}", default-features = false }}
4194tracing-subscriber = "={tracing_subscriber_version}"
4195"#
4196 );
4197 fs::write(temp.join("Cargo.toml"), manifest).unwrap();
4198 fs::write(temp.join("src/main.rs"), TRUE_OTEL_HARNESS).unwrap();
4199
4200 let output = Command::new(env!("CARGO"))
4201 .args(["run", "--quiet", "--offline"])
4202 .current_dir(&temp)
4203 .env("CARGO_TARGET_DIR", temp.join("target"))
4204 .output()
4205 .unwrap();
4206 let _ = fs::remove_dir_all(&temp);
4207 assert!(
4208 output.status.success(),
4209 "true OTEL harness failed:\nstdout:\n{}\nstderr:\n{}",
4210 String::from_utf8_lossy(&output.stdout),
4211 String::from_utf8_lossy(&output.stderr)
4212 );
4213 }
4214
4215 fn locked_version(root: &std::path::Path, name: &str) -> String {
4216 let lock = fs::read_to_string(root.join("Cargo.lock")).unwrap();
4217 let expected_name = format!("name = {}", serde_json::to_string(name).unwrap());
4218 let versions = lock
4219 .split("[[package]]")
4220 .filter(|package| {
4221 package
4222 .lines()
4223 .any(|line| line.trim() == expected_name.as_str())
4224 })
4225 .filter_map(|package| {
4226 package.lines().find_map(|line| {
4227 line.trim()
4228 .strip_prefix("version = \"")
4229 .and_then(|version| version.strip_suffix('\"'))
4230 .map(str::to_owned)
4231 })
4232 })
4233 .collect::<Vec<_>>();
4234 assert_eq!(versions.len(), 1, "expected one locked version for {name}");
4235 versions.into_iter().next().unwrap()
4236 }
4237
4238 fn toml_path(path: &std::path::Path) -> String {
4239 serde_json::to_string(&path.to_string_lossy()).unwrap()
4240 }
4241
4242 #[test]
4243 fn toml_path_escapes_windows_separators_and_quotes() {
4244 let encoded = toml_path(std::path::Path::new(r#"C:\Users\name\quoted\"dir"#));
4245 assert_eq!(encoded, r#""C:\\Users\\name\\quoted\\\"dir""#);
4246 }
4247
4248 const TRUE_OTEL_HARNESS: &str = concat!(
4249 r#"
4250use std::borrow::Cow;
4251use std::collections::VecDeque;
4252use std::sync::{Arc, Mutex};
4253use std::time::SystemTime;
4254use agentkit_core::{CostUsage, DataRef, FinishReason, Item, ItemKind, MetadataMap, Modality, Part, TokenUsage, TurnCancellation, Usage};
4255use agentkit_loop::{Agent, LoopError, LoopStep, MessageCapture, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, TelemetryConfig, TurnRequest, set_provider_finish_reasons};
4256use async_trait::async_trait;
4257use opentelemetry::trace::{Span, SpanBuilder, SpanContext, Status, Tracer};
4258use opentelemetry::{Array, Context, KeyValue, Value};
4259use tracing_subscriber::layer::SubscriberExt;
4260
4261#[derive(Clone, Debug)]
4262struct Exported { name: String, attributes: Vec<KeyValue>, capture_id: usize, status: Status }
4263#[derive(Clone, Default)]
4264struct MemoryTracer { exported: Arc<Mutex<Vec<Exported>>> }
4265struct MemorySpan { name: String, attributes: Vec<KeyValue>, exported: Arc<Mutex<Vec<Exported>>>, ended: bool, capture_id: usize, status: Status }
4266impl Tracer for MemoryTracer {
4267 type Span = MemorySpan;
4268 fn build_with_context(&self, builder: SpanBuilder, _: &Context) -> MemorySpan {
4269 static NEXT_CAPTURE: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
4270 MemorySpan { name: builder.name.into_owned(), attributes: builder.attributes.unwrap_or_default(), exported: self.exported.clone(), ended: false, capture_id: NEXT_CAPTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed), status: Status::Unset }
4271 }
4272}
4273impl Span for MemorySpan {
4274 fn add_event_with_timestamp<T>(&mut self, _: T, _: SystemTime, _: Vec<KeyValue>) where T: Into<Cow<'static, str>> {}
4275 fn span_context(&self) -> &SpanContext { &SpanContext::NONE }
4276 fn is_recording(&self) -> bool { !self.ended }
4277 fn set_attribute(&mut self, attribute: KeyValue) { self.attributes.push(attribute); }
4278 fn set_status(&mut self, status: Status) { self.status = status; }
4279 fn update_name<T>(&mut self, name: T) where T: Into<Cow<'static, str>> { self.name = name.into().into_owned(); }
4280 fn add_link(&mut self, _: SpanContext, _: Vec<KeyValue>) {}
4281 fn end_with_timestamp(&mut self, _: SystemTime) {
4282 if !self.ended {
4283 self.ended = true;
4284 self.exported.lock().unwrap().push(Exported { name: self.name.clone(), attributes: self.attributes.clone(), capture_id: self.capture_id, status: self.status.clone() });
4285 }
4286 }
4287}
4288
4289#[derive(Clone, Copy)]
4290enum BeginMode { Success, Error, Cancelled }
4291#[derive(Clone)]
4292struct ScriptedAdapter { adapter_provider: &'static str, final_usage: bool, overflow: bool, begin_mode: BeginMode }
4293struct ScriptedSession { selected_provider: Option<&'static str>, model: &'static str, final_usage: bool, overflow: bool, begin_mode: BeginMode }
4294struct ScriptedTurn { events: VecDeque<ModelTurnEvent> }
4295#[async_trait]
4296impl ModelAdapter for ScriptedAdapter {
4297 type Session = ScriptedSession;
4298 async fn start_session(&self, _: SessionConfig) -> Result<Self::Session, LoopError> {
4299 Ok(ScriptedSession { selected_provider: Some("before-begin"), model: "before-model", final_usage: self.final_usage, overflow: self.overflow, begin_mode: self.begin_mode })
4300 }
4301 fn provider_name(&self) -> Option<&str> { Some(self.adapter_provider) }
4302}
4303#[async_trait]
4304impl ModelSession for ScriptedSession {
4305 type Turn = ScriptedTurn;
4306 async fn begin_turn(&mut self, _: TurnRequest, _: Option<TurnCancellation>) -> Result<Self::Turn, LoopError> {
4307 match self.begin_mode {
4308 BeginMode::Error => return Err(LoopError::InvalidState("begin failed".into())),
4309 BeginMode::Cancelled => return Err(LoopError::Cancelled),
4310 BeginMode::Success => {}
4311 }
4312 self.selected_provider = if self.final_usage { Some("session-after-begin") } else { None };
4313 self.model = if self.final_usage { "model-after-begin" } else { "fallback-model-after-begin" };
4314 let mut stream_meta = MetadataMap::new();
4315 stream_meta.insert("stream_only".into(), serde_json::json!(true));
4316 stream_meta.insert("shared".into(), serde_json::json!("stream"));
4317 let streamed = Usage {
4318 tokens: Some(TokenUsage::new(if self.overflow { i64::MAX as u64 + 1 } else { 30 }, 40)),
4319 cost: Some(CostUsage::new(0.75, "USD")),
4320 metadata: stream_meta,
4321 };
4322 let final_usage = if self.final_usage {
4323 let mut metadata = MetadataMap::new();
4324 metadata.insert("shared".into(), serde_json::json!("final"));
4325 Some(Usage { tokens: Some(TokenUsage::new(if self.overflow { i64::MAX as u64 + 1 } else { 1 }, 2)), cost: None, metadata })
4326 } else { None };
4327 let mut result_metadata = MetadataMap::new();
4328 set_provider_finish_reasons(&mut result_metadata, ["native", "native", "done"]);
4329 let outputs = vec![
4330 Item::text(ItemKind::Assistant, "first"),
4331 Item::text(ItemKind::Assistant, "second"),
4332 Item::text(ItemKind::Assistant, "third"),
4333 ];
4334 Ok(ScriptedTurn { events: VecDeque::from([
4335 ModelTurnEvent::Usage(streamed),
4336 ModelTurnEvent::Finished(ModelTurnResult { finish_reason: FinishReason::Completed, output_items: outputs, usage: final_usage, metadata: result_metadata, model: Some(self.model.into()), response_id: Some("response-id".into()) }),
4337 ]) })
4338 }
4339 fn model_name(&self) -> Option<&str> { Some(self.model) }
4340 fn provider_name(&self) -> Option<&str> { self.selected_provider }
4341}
4342#[async_trait]
4343impl ModelTurn for ScriptedTurn {
4344 async fn next_event(&mut self, _: Option<TurnCancellation>) -> Result<Option<ModelTurnEvent>, LoopError> { Ok(self.events.pop_front()) }
4345}
4346
4347fn attr<'a>(span: &'a Exported, key: &str) -> Option<&'a Value> {
4348 span.attributes.iter().rev().find(|a| a.key.as_str() == key).map(|a| &a.value)
4349}
4350fn operation(span: &Exported) -> Option<&str> {
4351 match attr(span, "gen_ai.operation.name") { Some(Value::String(value)) => Some(value.as_str()), _ => None }
4352}
4353fn json_array(span: &Exported, key: &str) -> Vec<serde_json::Value> {
4354 match attr(span, key) {
4355 Some(Value::Array(Array::String(values))) => values.iter().map(|v| serde_json::from_str(v.as_str()).unwrap()).collect(),
4356 other => panic!("{key} was not Array<String>: {other:?}"),
4357 }
4358}
4359
4360fn run_attempt(adapter: ScriptedAdapter) -> (Vec<Exported>, Result<LoopStep, LoopError>) {
4361 let tracer = MemoryTracer::default();
4362 let subscriber = tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer.clone()));
4363 let runtime = tokio::runtime::Builder::new_current_thread().build().unwrap();
4364 let result = tracing::subscriber::with_default(subscriber, || runtime.block_on(async {
4365 let media = Item::new(ItemKind::User, vec![
4366 Part::media(Modality::Image, "image/png", DataRef::inline_bytes([1, 2, 3])),
4367 Part::media(Modality::Audio, "audio/wav", DataRef::uri("https://secret.invalid/audio")),
4368 ]);
4369 let agent = Agent::builder().model(adapter).transcript(vec![
4370 Item::text(ItemKind::System, "old"), Item::text(ItemKind::User, "middle"), media,
4371 ]).input(vec![Item::text(ItemKind::User, "newest")]).telemetry(
4372 TelemetryConfig::default()
4373 .with_input_messages(MessageCapture::new(3, 100_000).unwrap())
4374 .with_output_messages(MessageCapture::new(2, 100_000).unwrap())
4375 ).build().unwrap();
4376 let mut driver = agent.start(SessionConfig::new("otel-test")).await.unwrap();
4377 driver.next().await
4378 }));
4379 (tracer.exported.lock().unwrap().clone(), result)
4380}
4381fn run(adapter: ScriptedAdapter) -> (Vec<Exported>, agentkit_loop::TurnResult) {
4382 let (spans, result) = run_attempt(adapter);
4383 let result = match result.unwrap() { LoopStep::Finished(result) => result, other => panic!("unexpected step: {other:?}") };
4384 (spans, result)
4385}
4386
4387fn main() {
4388 recovery_attempts_have_distinct_otel_spans();
4389 let (spans, result) = run(ScriptedAdapter { adapter_provider: "adapter", final_usage: true, overflow: true, begin_mode: BeginMode::Success });
4390 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
4391 let agent = spans.iter().find(|s| operation(s) == Some("invoke_agent")).unwrap();
4392 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("session-after-begin".into())));
4393 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("model-after-begin".into())));
4394 assert!(attr(chat, "gen_ai.usage.input_tokens").is_none());
4395 assert_eq!(attr(chat, "gen_ai.usage.output_tokens"), Some(&Value::I64(2)));
4396 assert_eq!(attr(chat, "gen_ai.usage.cost"), Some(&Value::F64(0.75)));
4397 assert_eq!(attr(agent, "gen_ai.provider.name"), Some(&Value::String("session-after-begin".into())));
4398 assert!(attr(agent, "gen_ai.usage.input_tokens").is_none());
4399 assert!(attr(agent, "gen_ai.usage.cost").is_none());
4400 match attr(chat, "gen_ai.response.finish_reasons") {
4401 Some(Value::Array(Array::String(values))) => assert_eq!(values.iter().map(|v| v.as_str()).collect::<Vec<_>>(), ["native", "done"]),
4402 other => panic!("finish reasons were not Array<String>: {other:?}"),
4403 }
4404 assert_eq!(result.usage.as_ref().unwrap().metadata["shared"], "final");
4405 assert_eq!(result.usage.as_ref().unwrap().metadata["stream_only"], true);
4406 let input = json_array(chat, "gen_ai.input.messages");
4407 assert_eq!(input.len(), 3);
4408 assert!(input[0].to_string().contains("middle"));
4409 assert!(input[1].to_string().contains("omitted"));
4410 assert!(input[2].to_string().contains("newest"));
4411 assert!(!input.iter().any(|message| message.to_string().contains("old")));
4412 let output = json_array(chat, "gen_ai.output.messages");
4413 assert_eq!(output.len(), 2);
4414 assert!(output[0].to_string().contains("first"));
4415 assert!(output[1].to_string().contains("second"));
4416 let encoded = format!("{input:?}{output:?}");
4417 assert!(!encoded.contains("secret.invalid"));
4418 assert!(!encoded.contains("[1,2,3]"));
4419
4420 let (spans, result) = run(ScriptedAdapter { adapter_provider: "adapter-fallback", final_usage: false, overflow: false, begin_mode: BeginMode::Success });
4421 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
4422 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("adapter-fallback".into())));
4423 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("fallback-model-after-begin".into())));
4424 assert_eq!(attr(chat, "gen_ai.usage.input_tokens"), Some(&Value::I64(30)));
4425 assert_eq!(attr(chat, "gen_ai.usage.cost"), Some(&Value::F64(0.75)));
4426 assert_eq!(result.usage.unwrap().metadata["stream_only"], true);
4427
4428 for begin_mode in [BeginMode::Error, BeginMode::Cancelled] {
4429 let (spans, result) = run_attempt(ScriptedAdapter { adapter_provider: "adapter-before-error", final_usage: false, overflow: false, begin_mode });
4430 match begin_mode {
4431 BeginMode::Error => assert!(matches!(result, Err(LoopError::InvalidState(_)))),
4432 BeginMode::Cancelled => assert!(matches!(result, Ok(LoopStep::Finished(_)))),
4433 BeginMode::Success => unreachable!(),
4434 }
4435 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
4436 let agent = spans.iter().find(|s| operation(s) == Some("invoke_agent")).unwrap();
4437 assert_eq!(spans.iter().filter(|s| operation(s) == Some("chat")).count(), 1);
4438 assert!(matches!(chat.status, Status::Error { .. }));
4439 let expected_error = match begin_mode { BeginMode::Cancelled => "cancelled", _ => "model_error" };
4440 assert_eq!(attr(chat, "error.type"), Some(&Value::String(expected_error.into())));
4441 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("before-begin".into())));
4442 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("before-model".into())));
4443 assert_eq!(attr(agent, "gen_ai.provider.name"), Some(&Value::String("before-begin".into())));
4444 }
4445}
4446"#,
4447 include_str!("context_recovery_otel_harness.rs")
4448 );
4449}
4450
4451#[derive(Default)]
4452struct StreamedAssistantContent {
4453 parts: Vec<StreamedPart>,
4454 retained_bytes: usize,
4455 overflowed: bool,
4456}
4457
4458enum StreamedPart {
4459 Open {
4460 id: PartId,
4461 kind: PartKind,
4462 text: String,
4463 bytes: Vec<u8>,
4464 structured: Option<Value>,
4465 structured_bytes: usize,
4466 metadata: MetadataMap,
4467 metadata_bytes: usize,
4468 retained_bytes: usize,
4469 },
4470 Committed {
4471 part: Part,
4472 retained_bytes: usize,
4473 },
4474}
4475
4476impl StreamedAssistantContent {
4477 fn apply_delta(&mut self, delta: &Delta) {
4478 if self.overflowed {
4479 return;
4480 }
4481
4482 match delta {
4483 Delta::BeginPart { part_id, kind } => {
4484 let index = self.open_part_index(part_id);
4485 let retained_bytes = part_id.0.len();
4486 if !self.reserve_slot(index, retained_bytes) {
4487 return;
4488 }
4489 let open = StreamedPart::Open {
4490 id: part_id.clone(),
4491 kind: *kind,
4492 text: String::new(),
4493 bytes: Vec::new(),
4494 structured: None,
4495 structured_bytes: 0,
4496 metadata: MetadataMap::new(),
4497 metadata_bytes: 0,
4498 retained_bytes,
4499 };
4500 if let Some(index) = index {
4501 self.parts[index] = open;
4502 } else {
4503 self.parts.push(open);
4504 }
4505 }
4506 Delta::AppendText { part_id, chunk } => {
4507 let Some(index) = self.open_part_index(part_id) else {
4508 return;
4509 };
4510 if self.grow_slot(index, chunk.len())
4511 && let StreamedPart::Open { text, .. } = &mut self.parts[index]
4512 {
4513 text.push_str(chunk);
4514 }
4515 }
4516 Delta::AppendBytes { part_id, chunk } => {
4517 let Some(index) = self.open_part_index(part_id) else {
4518 return;
4519 };
4520 if matches!(
4521 self.parts[index],
4522 StreamedPart::Open {
4523 kind: PartKind::Media,
4524 ..
4525 }
4526 ) {
4527 return;
4528 }
4529 if self.grow_slot(index, chunk.len())
4530 && let StreamedPart::Open { bytes, .. } = &mut self.parts[index]
4531 {
4532 bytes.extend_from_slice(chunk);
4533 }
4534 }
4535 Delta::ReplaceStructured { part_id, value } => {
4536 let Some(index) = self.open_part_index(part_id) else {
4537 return;
4538 };
4539 let old_bytes = match &self.parts[index] {
4540 StreamedPart::Open {
4541 structured_bytes, ..
4542 } => *structured_bytes,
4543 StreamedPart::Committed { .. } => unreachable!(),
4544 };
4545 let limit = self.available_after_replacing(old_bytes);
4546 let Some(new_bytes) = serialized_size_with_limit(value, limit) else {
4547 self.overflow();
4548 return;
4549 };
4550 self.replace_slot_bytes(index, old_bytes, new_bytes);
4551 if let StreamedPart::Open {
4552 structured,
4553 structured_bytes,
4554 ..
4555 } = &mut self.parts[index]
4556 {
4557 *structured = Some(value.clone());
4558 *structured_bytes = new_bytes;
4559 }
4560 }
4561 Delta::SetMetadata { part_id, metadata } => {
4562 let Some(index) = self.open_part_index(part_id) else {
4563 return;
4564 };
4565 let old_bytes = match &self.parts[index] {
4566 StreamedPart::Open { metadata_bytes, .. } => *metadata_bytes,
4567 StreamedPart::Committed { .. } => unreachable!(),
4568 };
4569 let limit = self.available_after_replacing(old_bytes);
4570 let Some(new_bytes) = serialized_size_with_limit(metadata, limit) else {
4571 self.overflow();
4572 return;
4573 };
4574 self.replace_slot_bytes(index, old_bytes, new_bytes);
4575 if let StreamedPart::Open {
4576 metadata: target,
4577 metadata_bytes,
4578 ..
4579 } = &mut self.parts[index]
4580 {
4581 *target = metadata.clone();
4582 *metadata_bytes = new_bytes;
4583 }
4584 }
4585 Delta::CommitPart { part } => self.commit_part(part),
4586 }
4587 }
4588
4589 fn open_part_index(&self, id: &PartId) -> Option<usize> {
4590 self.parts.iter().position(
4591 |part| matches!(part, StreamedPart::Open { id: open_id, .. } if open_id == id),
4592 )
4593 }
4594
4595 fn commit_part(&mut self, part: &Part) {
4596 if self.overflowed {
4597 return;
4598 }
4599
4600 let duplicate_tool_call = if let Part::ToolCall(call) = part {
4601 self.parts.iter().position(|slot| {
4602 matches!(
4603 slot,
4604 StreamedPart::Committed {
4605 part: Part::ToolCall(existing),
4606 ..
4607 } if existing.id == call.id
4608 )
4609 })
4610 } else {
4611 None
4612 };
4613
4614 let kind = part_kind(part);
4615 let matching_open = self
4616 .parts
4617 .iter()
4618 .position(|slot| slot.open_matches_part(part));
4619 let mut open_with_kind = self.parts.iter().enumerate().filter_map(|(index, slot)| {
4620 matches!(slot, StreamedPart::Open { kind: open_kind, .. } if *open_kind == kind)
4621 .then_some(index)
4622 });
4623 let only_open_with_kind = match (open_with_kind.next(), open_with_kind.next()) {
4624 (Some(index), None) => Some(index),
4625 _ => None,
4626 };
4627 let index = duplicate_tool_call
4628 .or(matching_open)
4629 .or(only_open_with_kind);
4630 let limit = self.available_for_slot(index);
4631 let Some(retained_bytes) = serialized_size_with_limit(part, limit) else {
4632 self.overflow();
4633 return;
4634 };
4635 if !self.reserve_slot(index, retained_bytes) {
4636 return;
4637 }
4638 let committed = StreamedPart::Committed {
4639 part: part.clone(),
4640 retained_bytes,
4641 };
4642 if let Some(index) = index {
4643 self.parts[index] = committed;
4644 } else {
4645 self.parts.push(committed);
4646 }
4647 }
4648
4649 fn commit_tool_call(&mut self, call: &ToolCallPart) {
4650 if self.overflowed {
4651 return;
4652 }
4653 let duplicate_tool_call = self.parts.iter().position(|slot| {
4654 matches!(
4655 slot,
4656 StreamedPart::Committed {
4657 part: Part::ToolCall(existing),
4658 ..
4659 } if existing.id == call.id
4660 )
4661 });
4662 let mut open_tool_calls = self.parts.iter().enumerate().filter_map(|(index, slot)| {
4663 matches!(
4664 slot,
4665 StreamedPart::Open {
4666 kind: PartKind::ToolCall,
4667 ..
4668 }
4669 )
4670 .then_some(index)
4671 });
4672 let only_open_tool_call = match (open_tool_calls.next(), open_tool_calls.next()) {
4673 (Some(index), None) => Some(index),
4674 _ => None,
4675 };
4676 let index = duplicate_tool_call.or(only_open_tool_call);
4677 let limit = self.available_for_slot(index);
4678 let borrowed_part = BorrowedPart::ToolCall(call);
4679 let Some(retained_bytes) = serialized_size_with_limit(&borrowed_part, limit) else {
4680 self.overflow();
4681 return;
4682 };
4683 if !self.reserve_slot(index, retained_bytes) {
4684 return;
4685 }
4686 let committed = StreamedPart::Committed {
4687 part: Part::ToolCall(call.clone()),
4688 retained_bytes,
4689 };
4690 if let Some(index) = index {
4691 self.parts[index] = committed;
4692 } else {
4693 self.parts.push(committed);
4694 }
4695 }
4696
4697 fn reserve_slot(&mut self, index: Option<usize>, retained_bytes: usize) -> bool {
4698 if index.is_none() && self.parts.len() >= MAX_STREAMED_ASSISTANT_CONTENT_PARTS {
4699 self.overflow();
4700 return false;
4701 }
4702 let replaced_bytes = index.map_or(0, |index| self.parts[index].retained_bytes());
4703 let Some(total) = self
4704 .retained_bytes
4705 .checked_sub(replaced_bytes)
4706 .and_then(|bytes| bytes.checked_add(retained_bytes))
4707 .filter(|bytes| *bytes <= MAX_STREAMED_ASSISTANT_CONTENT_BYTES)
4708 else {
4709 self.overflow();
4710 return false;
4711 };
4712 self.retained_bytes = total;
4713 true
4714 }
4715
4716 fn grow_slot(&mut self, index: usize, additional_bytes: usize) -> bool {
4717 let Some(total) = self
4718 .retained_bytes
4719 .checked_add(additional_bytes)
4720 .filter(|bytes| *bytes <= MAX_STREAMED_ASSISTANT_CONTENT_BYTES)
4721 else {
4722 self.overflow();
4723 return false;
4724 };
4725 self.retained_bytes = total;
4726 *self.parts[index].retained_bytes_mut() += additional_bytes;
4727 true
4728 }
4729
4730 fn available_after_replacing(&self, replaced_bytes: usize) -> usize {
4731 MAX_STREAMED_ASSISTANT_CONTENT_BYTES - (self.retained_bytes - replaced_bytes)
4732 }
4733
4734 fn available_for_slot(&self, index: Option<usize>) -> usize {
4735 self.available_after_replacing(index.map_or(0, |index| self.parts[index].retained_bytes()))
4736 }
4737
4738 fn replace_slot_bytes(&mut self, index: usize, old_bytes: usize, new_bytes: usize) {
4739 self.retained_bytes = self.retained_bytes - old_bytes + new_bytes;
4740 let retained_bytes = self.parts[index].retained_bytes_mut();
4741 *retained_bytes = *retained_bytes - old_bytes + new_bytes;
4742 }
4743
4744 fn overflow(&mut self) {
4745 self.parts.clear();
4746 self.retained_bytes = 0;
4747 self.overflowed = true;
4748 }
4749
4750 fn reset(&mut self) {
4751 self.parts.clear();
4752 self.retained_bytes = 0;
4753 self.overflowed = false;
4754 }
4755
4756 fn interrupted_items(&self) -> Vec<Item> {
4757 if self.overflowed {
4758 return interrupted_assistant_items();
4759 }
4760 let parts = self
4761 .parts
4762 .iter()
4763 .filter_map(StreamedPart::preserved_part)
4764 .collect::<Vec<_>>();
4765 if parts.is_empty() {
4766 return interrupted_assistant_items();
4767 }
4768 vec![Item::new(ItemKind::Assistant, parts).with_metadata(interrupted_metadata("assistant"))]
4769 }
4770}
4771
4772impl StreamedPart {
4773 fn retained_bytes(&self) -> usize {
4774 match self {
4775 Self::Open { retained_bytes, .. } | Self::Committed { retained_bytes, .. } => {
4776 *retained_bytes
4777 }
4778 }
4779 }
4780
4781 fn retained_bytes_mut(&mut self) -> &mut usize {
4782 match self {
4783 Self::Open { retained_bytes, .. } | Self::Committed { retained_bytes, .. } => {
4784 retained_bytes
4785 }
4786 }
4787 }
4788
4789 fn open_matches_part(&self, committed: &Part) -> bool {
4790 match (self, committed) {
4791 (
4792 Self::Open {
4793 kind: PartKind::Text,
4794 text,
4795 ..
4796 },
4797 Part::Text(part),
4798 ) => text == &part.text,
4799 (
4800 Self::Open {
4801 kind: PartKind::Reasoning,
4802 text,
4803 bytes,
4804 ..
4805 },
4806 Part::Reasoning(part),
4807 ) => {
4808 part.summary.as_deref() == (!text.is_empty()).then_some(text.as_str())
4809 && match &part.data {
4810 Some(DataRef::InlineBytes(data)) => data == bytes,
4811 None => bytes.is_empty(),
4812 _ => false,
4813 }
4814 }
4815 (
4816 Self::Open {
4817 kind: PartKind::Structured,
4818 structured: Some(value),
4819 ..
4820 },
4821 Part::Structured(part),
4822 ) => value == &part.value,
4823 _ => false,
4824 }
4825 }
4826
4827 fn preserved_part(&self) -> Option<Part> {
4828 match self {
4829 Self::Committed { part, .. } if useful_streamed_part(part) => Some(part.clone()),
4830 Self::Committed { .. } => None,
4831 Self::Open {
4832 kind,
4833 text,
4834 bytes,
4835 structured,
4836 metadata,
4837 ..
4838 } => match kind {
4839 PartKind::Text if !text.is_empty() => Some(Part::Text(
4840 TextPart::new(text).with_metadata(metadata.clone()),
4841 )),
4842 PartKind::Reasoning if !text.is_empty() || !bytes.is_empty() => {
4843 Some(Part::Reasoning(ReasoningPart {
4844 summary: (!text.is_empty()).then(|| text.clone()),
4845 data: (!bytes.is_empty()).then(|| DataRef::inline_bytes(bytes.clone())),
4846 redacted: false,
4847 metadata: metadata.clone(),
4848 }))
4849 }
4850 PartKind::Structured => structured.clone().map(|value| {
4851 Part::Structured(StructuredPart::new(value).with_metadata(metadata.clone()))
4852 }),
4853 PartKind::Media => None,
4856 _ => None,
4857 },
4858 }
4859 }
4860}
4861
4862#[derive(Serialize)]
4863enum BorrowedPart<'a> {
4864 ToolCall(&'a ToolCallPart),
4865}
4866
4867struct PayloadSizeWriter {
4868 remaining: usize,
4869 written: usize,
4870}
4871
4872impl std::io::Write for PayloadSizeWriter {
4873 fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
4874 if bytes.len() > self.remaining {
4875 return Err(std::io::Error::other(
4876 "streamed assistant payload limit exceeded",
4877 ));
4878 }
4879 self.remaining -= bytes.len();
4880 self.written += bytes.len();
4881 Ok(bytes.len())
4882 }
4883
4884 fn flush(&mut self) -> std::io::Result<()> {
4885 Ok(())
4886 }
4887}
4888
4889fn serialized_size_with_limit<T: Serialize + ?Sized>(value: &T, limit: usize) -> Option<usize> {
4890 let mut writer = PayloadSizeWriter {
4891 remaining: limit,
4892 written: 0,
4893 };
4894 serde_json::to_writer(&mut writer, value).ok()?;
4895 Some(writer.written)
4896}
4897
4898fn part_kind(part: &Part) -> PartKind {
4899 match part {
4900 Part::Text(_) => PartKind::Text,
4901 Part::Media(_) => PartKind::Media,
4902 Part::File(_) => PartKind::File,
4903 Part::Structured(_) => PartKind::Structured,
4904 Part::Reasoning(_) => PartKind::Reasoning,
4905 Part::ToolCall(_) => PartKind::ToolCall,
4906 Part::ToolResult(_) => PartKind::ToolResult,
4907 Part::Custom(_) => PartKind::Custom,
4908 }
4909}
4910
4911fn useful_streamed_part(part: &Part) -> bool {
4912 match part {
4913 Part::Text(text) => !text.text.is_empty(),
4914 Part::Reasoning(reasoning) => {
4915 reasoning
4916 .summary
4917 .as_ref()
4918 .is_some_and(|text| !text.is_empty())
4919 || reasoning.data.is_some()
4920 || reasoning.redacted
4921 }
4922 Part::ToolResult(_) => false,
4923 _ => true,
4924 }
4925}
4926
4927fn interrupted_stream_items(content: Option<&StreamedAssistantContent>) -> Vec<Item> {
4928 content.map_or_else(
4929 interrupted_assistant_items,
4930 StreamedAssistantContent::interrupted_items,
4931 )
4932}
4933
4934fn interrupted_assistant_items() -> Vec<Item> {
4935 vec![
4936 Item::new(
4937 ItemKind::Assistant,
4938 vec![Part::Text(TextPart {
4939 text: "Previous assistant response was interrupted by the user before completion."
4940 .into(),
4941 metadata: interrupted_metadata("assistant"),
4942 })],
4943 )
4944 .with_metadata(interrupted_metadata("assistant")),
4945 ]
4946}
4947
4948fn unanswered_tool_calls(transcript: &[Item]) -> Vec<ToolCallPart> {
4950 let mut open: Vec<ToolCallPart> = Vec::new();
4951 for part in transcript.iter().flat_map(|item| &item.parts) {
4952 match part {
4953 Part::ToolCall(call) => open.push(call.clone()),
4954 Part::ToolResult(result) => open.retain(|call| call.id != result.call_id),
4955 _ => {}
4956 }
4957 }
4958 open
4959}
4960
4961fn interrupted_tool_result_item(call: ToolCallPart) -> Item {
4966 Item {
4967 id: None,
4968 kind: ItemKind::Tool,
4969 parts: vec![Part::ToolResult(ToolResultPart {
4970 call_id: call.id,
4971 output: ToolOutput::Text("tool call interrupted before it reported a result".into()),
4972 is_error: true,
4973 metadata: interrupted_metadata("tool"),
4974 })],
4975 metadata: interrupted_metadata("tool"),
4976 usage: None,
4977 finish_reason: None,
4978 created_at: None,
4979 }
4980}
4981
4982fn cancelled_approval_item(pending: PendingApprovalToolCall) -> Item {
4983 Item {
4984 id: None,
4985 kind: ItemKind::Tool,
4986 parts: vec![Part::ToolResult(ToolResultPart {
4987 call_id: pending.call.id,
4988 output: ToolOutput::Text("approval cancelled".into()),
4989 is_error: true,
4990 metadata: pending.call.metadata,
4991 })],
4992 metadata: MetadataMap::new(),
4993 usage: None,
4994 finish_reason: None,
4995 created_at: None,
4996 }
4997}
4998
4999fn transcript_has_pending_input(transcript: &[Item]) -> bool {
5005 matches!(
5006 transcript.last().map(|item| item.kind),
5007 Some(ItemKind::User | ItemKind::Tool | ItemKind::Notification)
5008 )
5009}
5010
5011fn extract_tool_calls(items: &[Item]) -> Vec<ToolCallPart> {
5012 let mut calls = Vec::new();
5013 for item in items {
5014 for part in &item.parts {
5015 if let Part::ToolCall(call) = part {
5016 calls.push(call.clone());
5017 }
5018 }
5019 }
5020 calls
5021}
5022
5023fn tool_result_is_error(item: &Item) -> bool {
5024 item.parts
5025 .iter()
5026 .any(|part| matches!(part, Part::ToolResult(result) if result.is_error))
5027}
5028
5029fn tool_result_not_started(item: &Item) -> bool {
5030 item.parts.iter().any(|part| {
5031 matches!(
5032 part,
5033 Part::ToolResult(result)
5034 if result
5035 .metadata
5036 .get(TOOL_RESULT_NOT_STARTED_METADATA_KEY)
5037 .and_then(Value::as_bool)
5038 == Some(true)
5039 )
5040 })
5041}
5042
5043#[derive(Debug, Error)]
5045pub enum LoopError {
5046 #[error(transparent)]
5048 ProviderFailure(Box<ProviderFailure>),
5049 #[error("invalid driver state: {0}")]
5051 InvalidState(String),
5052 #[error("turn cancelled")]
5054 Cancelled,
5055 #[error("context recovery deadline exceeded")]
5057 ContextRecoveryDeadlineExceeded,
5058 #[error("provider error: {0}")]
5060 Provider(String),
5061 #[error("tool error: {0}")]
5063 Tool(#[from] ToolError),
5064 #[error("mutator error: {0}")]
5066 Mutator(String),
5067 #[error("unsupported operation: {0}")]
5069 Unsupported(String),
5070}
5071
5072fn record_inference_error(span: &tracing::Span, error: &LoopError) {
5074 let kind = if error.is_context_length_exceeded() {
5075 "context_length_exceeded"
5076 } else {
5077 match error {
5078 LoopError::Cancelled => "cancelled",
5079 LoopError::ContextRecoveryDeadlineExceeded => "context_recovery_deadline_exceeded",
5080 _ => "model_error",
5081 }
5082 };
5083 span.record("error.type", kind);
5084 span.record("otel.status_code", "ERROR");
5085}
5086
5087async fn recovery_bound<T>(
5089 future: impl std::future::Future<Output = Result<T, LoopError>>,
5090 recovery: Option<ContextRecoveryInfo>,
5091 cancellation: Option<TurnCancellation>,
5092) -> Result<T, LoopError> {
5093 let Some(info) = recovery else {
5094 return future.await;
5095 };
5096 let cancelled = async {
5097 match cancellation {
5098 Some(cancellation) => cancellation.cancelled().await,
5099 None => std::future::pending::<()>().await,
5100 }
5101 };
5102 use std::future::Future;
5103 use std::task::Poll;
5104
5105 let mut cancelled = std::pin::pin!(cancelled);
5106 let mut deadline = std::pin::pin!(futures_timer::Delay::new(
5107 info.deadline
5108 .saturating_duration_since(std::time::Instant::now()),
5109 ));
5110 let mut future = std::pin::pin!(future);
5111 std::future::poll_fn(|cx| {
5112 if cancelled.as_mut().poll(cx).is_ready() {
5116 return Poll::Ready(Err(LoopError::Cancelled));
5117 }
5118 if std::time::Instant::now() >= info.deadline || deadline.as_mut().poll(cx).is_ready() {
5119 return Poll::Ready(Err(LoopError::ContextRecoveryDeadlineExceeded));
5120 }
5121 future.as_mut().poll(cx)
5122 })
5123 .await
5124}
5125
5126impl LoopError {
5127 pub fn is_context_length_exceeded(&self) -> bool {
5131 self.provider_failure().is_some_and(|failure| {
5132 failure.upstream.code == UpstreamErrorKind::ContextLengthExceeded
5133 && !matches!(
5134 failure.reason,
5135 ProviderFailureReason::Cancelled | ProviderFailureReason::ReplayUnsafe
5136 )
5137 && !matches!(
5138 failure.last_attempt_reason,
5139 Some(ProviderFailureReason::Cancelled | ProviderFailureReason::ReplayUnsafe)
5140 )
5141 })
5142 }
5143
5144 pub fn provider_failure(&self) -> Option<&ProviderFailure> {
5146 match self {
5147 Self::ProviderFailure(failure) => Some(failure),
5148 _ => None,
5149 }
5150 }
5151}
5152
5153struct DriverEmitter<'a> {
5158 session_id: &'a Arc<SessionId>,
5159 observers: &'a [Arc<dyn LoopObserver>],
5160}
5161
5162impl<'a> EventEmitter for DriverEmitter<'a> {
5163 fn emit(&self, event: AgentEvent) {
5164 fan_out_observed_event(self.observers, self.session_id, event);
5165 }
5166}
5167
5168fn fan_out_observed_event(
5169 observers: &[Arc<dyn LoopObserver>],
5170 session_id: &Arc<SessionId>,
5171 event: AgentEvent,
5172) {
5173 if observers.is_empty() {
5174 return;
5175 }
5176 let observed = ObservedEvent {
5177 session_id: Arc::clone(session_id),
5178 event,
5179 };
5180 let last = observers.len() - 1;
5181 for observer in &observers[..last] {
5182 observer.handle_event(observed.clone());
5183 }
5184 observers[last].handle_event(observed);
5185}
5186
5187fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> {
5192 let mut pending: HashSet<ToolCallId> = HashSet::new();
5193 let mut seen_calls: HashSet<ToolCallId> = HashSet::new();
5194 let mut seen_results: HashSet<ToolCallId> = HashSet::new();
5195 for item in transcript {
5196 for part in &item.parts {
5197 match part {
5198 Part::ToolCall(call) => {
5199 if !seen_calls.insert(call.id.clone()) {
5200 return Err(LoopError::Mutator(format!(
5201 "transcript invariant violation: duplicate tool_use: {}",
5202 call.id.0
5203 )));
5204 }
5205 pending.insert(call.id.clone());
5206 }
5207 Part::ToolResult(result) => {
5208 if !pending.remove(&result.call_id) {
5209 let kind = if seen_results.contains(&result.call_id) {
5210 "duplicate"
5211 } else {
5212 "orphaned"
5213 };
5214 return Err(LoopError::Mutator(format!(
5215 "transcript invariant violation: {kind} tool_result: {}",
5216 result.call_id.0
5217 )));
5218 }
5219 seen_results.insert(result.call_id.clone());
5220 }
5221 _ => {}
5222 }
5223 }
5224 }
5225 if !pending.is_empty() {
5226 let missing: Vec<String> = pending.into_iter().map(|id| id.0).collect();
5227 return Err(LoopError::Mutator(format!(
5228 "transcript invariant violation: tool_use(s) without matching tool_result: {}",
5229 missing.join(", ")
5230 )));
5231 }
5232 Ok(())
5233}
5234
5235#[cfg(test)]
5236mod tests {
5237 include!("context_recovery_tests.rs");
5238 use std::collections::VecDeque;
5239 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
5240 use std::sync::{Arc as StdArc, Mutex as StdMutex};
5241
5242 use agentkit_core::{
5243 CancellationController, ItemKind, Part, TextPart, ToolCallId, ToolCallPart, ToolOutput,
5244 ToolResultPart,
5245 };
5246 use agentkit_task_manager::{
5247 AsyncTaskManager, RoutingDecision, TaskEvent, TaskManager, TaskManagerError,
5248 TaskManagerHandle, TaskRoutingPolicy,
5249 };
5250 use agentkit_tools_core::{
5251 FileSystemPermissionRequest, PermissionCode, PermissionDecision, PermissionDenial, Tool,
5252 ToolAnnotations, ToolCatalogEvent, ToolExecutionOutcome, ToolName, ToolRegistry,
5253 ToolResult, ToolSpec,
5254 };
5255 use serde_json::{Value, json};
5256 use tokio::sync::Notify;
5257 use tokio::time::{Duration, timeout};
5258
5259 use super::*;
5260
5261 struct FakeAdapter;
5262 struct SupersedingAdapter;
5263 struct InterruptedStreamAdapter {
5264 controller: CancellationController,
5265 scenario: InterruptedStreamScenario,
5266 }
5267 struct SlowAdapter;
5268 struct RecordingAdapter {
5269 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
5270 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
5271 }
5272 struct MultiToolAdapter;
5273 struct DualApprovalAdapter;
5274
5275 struct FakeSession;
5276 struct SupersedingSession {
5277 supersession_enabled: bool,
5278 }
5279 struct InterruptedStreamSession {
5280 controller: CancellationController,
5281 scenario: InterruptedStreamScenario,
5282 }
5283 struct SlowSession;
5284 struct RecordingSession {
5285 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
5286 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
5287 }
5288 struct MultiToolSession;
5289 struct DualApprovalSession;
5290
5291 struct FakeTurn {
5292 events: VecDeque<ModelTurnEvent>,
5293 }
5294
5295 struct SupersedingTurn {
5296 events: VecDeque<ModelTurnEvent>,
5297 }
5298
5299 #[derive(Clone, Copy)]
5300 enum InterruptedStreamScenario {
5301 PreserveContent,
5302 CancelOnSupersession,
5303 OverflowThenSupersession,
5304 }
5305
5306 struct InterruptedStreamTurn {
5307 controller: CancellationController,
5308 events: VecDeque<ModelTurnEvent>,
5309 cancel_on_supersession: bool,
5310 }
5311
5312 struct SlowTurn {
5313 emitted: bool,
5314 }
5315
5316 struct RecordingTurn {
5317 emitted: bool,
5318 }
5319 struct MultiToolTurn {
5320 events: VecDeque<ModelTurnEvent>,
5321 }
5322 struct DualApprovalTurn {
5323 events: VecDeque<ModelTurnEvent>,
5324 }
5325
5326 struct TestTaskManager<T> {
5327 inner: T,
5328 start_error: Option<&'static str>,
5329 approved_start_error: Option<&'static str>,
5330 pending_update_error: Option<(usize, &'static str)>,
5331 pending_update_calls: AtomicUsize,
5332 interrupted: Option<StdArc<StdMutex<Vec<agentkit_core::TurnId>>>>,
5333 interrupt_error: Option<&'static str>,
5334 }
5335
5336 impl<T> TestTaskManager<T> {
5337 fn new(inner: T) -> Self {
5338 Self {
5339 inner,
5340 start_error: None,
5341 approved_start_error: None,
5342 pending_update_error: None,
5343 pending_update_calls: AtomicUsize::new(0),
5344 interrupted: None,
5345 interrupt_error: None,
5346 }
5347 }
5348
5349 fn fail_start(mut self, message: &'static str) -> Self {
5350 self.start_error = Some(message);
5351 self
5352 }
5353
5354 fn fail_approved_start(mut self, message: &'static str) -> Self {
5355 self.approved_start_error = Some(message);
5356 self
5357 }
5358
5359 fn fail_pending_update_on(mut self, call: usize, message: &'static str) -> Self {
5360 self.pending_update_error = Some((call, message));
5361 self
5362 }
5363
5364 fn record_interrupts(
5365 mut self,
5366 interrupted: StdArc<StdMutex<Vec<agentkit_core::TurnId>>>,
5367 ) -> Self {
5368 self.interrupted = Some(interrupted);
5369 self
5370 }
5371
5372 fn fail_interrupt(mut self, message: &'static str) -> Self {
5373 self.interrupt_error = Some(message);
5374 self
5375 }
5376 }
5377
5378 #[async_trait]
5379 impl<T: TaskManager> TaskManager for TestTaskManager<T> {
5380 async fn start_task(
5381 &self,
5382 request: TaskLaunchRequest,
5383 ctx: TaskStartContext,
5384 ) -> Result<TaskStartOutcome, TaskManagerError> {
5385 if let Some(message) = self.start_error.or_else(|| {
5386 matches!(&request.kind, TaskLaunchKind::Approved(_))
5387 .then_some(self.approved_start_error)
5388 .flatten()
5389 }) {
5390 return Err(TaskManagerError::Internal(message.into()));
5391 }
5392 self.inner.start_task(request, ctx).await
5393 }
5394
5395 async fn wait_for_turn(
5396 &self,
5397 turn_id: &agentkit_core::TurnId,
5398 cancellation: Option<TurnCancellation>,
5399 ) -> Result<Option<TurnTaskUpdate>, TaskManagerError> {
5400 self.inner.wait_for_turn(turn_id, cancellation).await
5401 }
5402
5403 async fn take_pending_loop_updates(&self) -> Result<PendingLoopUpdates, TaskManagerError> {
5404 if let Some((call, message)) = self.pending_update_error
5405 && self.pending_update_calls.fetch_add(1, Ordering::SeqCst) == call
5406 {
5407 return Err(TaskManagerError::Internal(message.into()));
5408 }
5409 self.inner.take_pending_loop_updates().await
5410 }
5411
5412 async fn on_turn_interrupted(
5413 &self,
5414 turn_id: &agentkit_core::TurnId,
5415 ) -> Result<(), TaskManagerError> {
5416 if let Some(interrupted) = &self.interrupted {
5417 interrupted.lock().unwrap().push(turn_id.clone());
5418 }
5419 if let Some(message) = self.interrupt_error {
5420 return Err(TaskManagerError::Internal(message.into()));
5421 }
5422 self.inner.on_turn_interrupted(turn_id).await
5423 }
5424
5425 fn handle(&self) -> TaskManagerHandle {
5426 self.inner.handle()
5427 }
5428 }
5429
5430 struct DelayedApprovalExecutor {
5431 entered: StdArc<AtomicBool>,
5432 release: StdArc<Notify>,
5433 approved_entered: Option<StdArc<AtomicBool>>,
5434 approved_release: Option<StdArc<Notify>>,
5435 cancellation: Option<CancellationController>,
5436 spec: ToolSpec,
5437 }
5438
5439 impl DelayedApprovalExecutor {
5440 fn new(entered: StdArc<AtomicBool>, release: StdArc<Notify>) -> Self {
5441 Self {
5442 entered,
5443 release,
5444 approved_entered: None,
5445 approved_release: None,
5446 cancellation: None,
5447 spec: ToolSpec {
5448 name: ToolName::new("echo"),
5449 description: "delayed approval".into(),
5450 input_schema: json!({
5451 "type": "object",
5452 "properties": {
5453 "value": { "type": "string" }
5454 },
5455 "required": ["value"],
5456 "additionalProperties": false
5457 }),
5458 output_schema: None,
5459 annotations: ToolAnnotations::default(),
5460 metadata: MetadataMap::new(),
5461 },
5462 }
5463 }
5464
5465 fn cancelling_on_approval(mut self, controller: CancellationController) -> Self {
5466 self.cancellation = Some(controller);
5467 self
5468 }
5469
5470 fn blocking_after_approval(
5471 mut self,
5472 entered: StdArc<AtomicBool>,
5473 release: StdArc<Notify>,
5474 ) -> Self {
5475 self.approved_entered = Some(entered);
5476 self.approved_release = Some(release);
5477 self
5478 }
5479 }
5480
5481 #[async_trait]
5482 impl ToolExecutor for DelayedApprovalExecutor {
5483 fn specs(&self) -> Vec<ToolSpec> {
5484 vec![self.spec.clone()]
5485 }
5486
5487 async fn execute(
5488 &self,
5489 request: ToolRequest,
5490 _ctx: &mut ToolContext<'_>,
5491 ) -> ToolExecutionOutcome {
5492 self.entered.store(true, Ordering::SeqCst);
5493 self.release.notified().await;
5494 if let Some(controller) = &self.cancellation {
5495 controller.interrupt();
5496 }
5497 ToolExecutionOutcome::Interrupted(
5498 agentkit_tools_core::ToolInterruption::ApprovalRequired(ApprovalRequest {
5499 task_id: None,
5500 call_id: None,
5501 id: "approval:delayed".into(),
5502 request_kind: "delayed.approval".into(),
5503 reason: agentkit_tools_core::ApprovalReason::PolicyRequiresConfirmation,
5504 summary: "delayed approval".into(),
5505 metadata: request.metadata,
5506 }),
5507 )
5508 }
5509
5510 async fn execute_approved(
5511 &self,
5512 request: ToolRequest,
5513 approved_request: &ApprovalRequest,
5514 ctx: &mut ToolContext<'_>,
5515 ) -> ToolExecutionOutcome {
5516 let (Some(entered), Some(release)) = (&self.approved_entered, &self.approved_release)
5517 else {
5518 return self.execute(request, ctx).await;
5519 };
5520 let _ = approved_request;
5521 entered.store(true, Ordering::SeqCst);
5522 release.notified().await;
5523 ToolExecutionOutcome::Completed(ToolResult {
5524 result: ToolResultPart {
5525 call_id: request.call_id,
5526 output: ToolOutput::Text("approved-ok".into()),
5527 is_error: false,
5528 metadata: MetadataMap::new(),
5529 },
5530 duration: None,
5531 metadata: MetadataMap::new(),
5532 })
5533 }
5534 }
5535
5536 #[async_trait]
5537 impl ModelAdapter for FakeAdapter {
5538 type Session = FakeSession;
5539
5540 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
5541 Ok(FakeSession)
5542 }
5543 }
5544
5545 #[async_trait]
5546 impl ModelAdapter for SupersedingAdapter {
5547 type Session = SupersedingSession;
5548
5549 async fn start_session(&self, config: SessionConfig) -> Result<Self::Session, LoopError> {
5550 Ok(SupersedingSession {
5551 supersession_enabled: config.consumer_capabilities.response_attempt_supersession,
5552 })
5553 }
5554 }
5555
5556 #[async_trait]
5557 impl ModelAdapter for InterruptedStreamAdapter {
5558 type Session = InterruptedStreamSession;
5559
5560 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
5561 Ok(InterruptedStreamSession {
5562 controller: self.controller.clone(),
5563 scenario: self.scenario,
5564 })
5565 }
5566 }
5567
5568 #[async_trait]
5569 impl ModelAdapter for SlowAdapter {
5570 type Session = SlowSession;
5571
5572 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
5573 Ok(SlowSession)
5574 }
5575 }
5576
5577 #[async_trait]
5578 impl ModelAdapter for RecordingAdapter {
5579 type Session = RecordingSession;
5580
5581 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
5582 Ok(RecordingSession {
5583 seen_descriptions: self.seen_descriptions.clone(),
5584 seen_caches: self.seen_caches.clone(),
5585 })
5586 }
5587 }
5588
5589 #[async_trait]
5590 impl ModelAdapter for MultiToolAdapter {
5591 type Session = MultiToolSession;
5592
5593 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
5594 Ok(MultiToolSession)
5595 }
5596 }
5597
5598 #[async_trait]
5599 impl ModelAdapter for DualApprovalAdapter {
5600 type Session = DualApprovalSession;
5601
5602 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
5603 Ok(DualApprovalSession)
5604 }
5605 }
5606
5607 #[async_trait]
5608 impl ModelSession for SupersedingSession {
5609 type Turn = SupersedingTurn;
5610
5611 async fn begin_turn(
5612 &mut self,
5613 _request: TurnRequest,
5614 _cancellation: Option<TurnCancellation>,
5615 ) -> Result<Self::Turn, LoopError> {
5616 assert!(self.supersession_enabled);
5617 Ok(SupersedingTurn {
5618 events: VecDeque::from([
5619 ModelTurnEvent::Usage(Usage::default()),
5620 ModelTurnEvent::ToolCall(ToolCallPart::new(
5621 "discarded-call",
5622 "discarded-tool",
5623 json!({}),
5624 )),
5625 ModelTurnEvent::ResponseAttemptSuperseded,
5626 ModelTurnEvent::Finished(ModelTurnResult {
5627 finish_reason: FinishReason::Completed,
5628 output_items: vec![Item::text(ItemKind::Assistant, "replacement")],
5629 usage: None,
5630 metadata: MetadataMap::new(),
5631 model: None,
5632 response_id: None,
5633 }),
5634 ]),
5635 })
5636 }
5637 }
5638
5639 #[async_trait]
5640 impl ModelSession for InterruptedStreamSession {
5641 type Turn = InterruptedStreamTurn;
5642
5643 async fn begin_turn(
5644 &mut self,
5645 _request: TurnRequest,
5646 _cancellation: Option<TurnCancellation>,
5647 ) -> Result<Self::Turn, LoopError> {
5648 let events = match self.scenario {
5649 InterruptedStreamScenario::PreserveContent => {
5650 let mut reasoning_metadata = MetadataMap::new();
5651 reasoning_metadata.insert("provider.detail".into(), true.into());
5652 VecDeque::from([
5653 ModelTurnEvent::Delta(Delta::BeginPart {
5654 part_id: PartId::new("text"),
5655 kind: PartKind::Text,
5656 }),
5657 ModelTurnEvent::Delta(Delta::AppendText {
5658 part_id: PartId::new("text"),
5659 chunk: "partial answer".into(),
5660 }),
5661 ModelTurnEvent::Delta(Delta::CommitPart {
5662 part: Part::text("partial answer"),
5663 }),
5664 ModelTurnEvent::Delta(Delta::BeginPart {
5665 part_id: PartId::new("reasoning"),
5666 kind: PartKind::Reasoning,
5667 }),
5668 ModelTurnEvent::Delta(Delta::AppendText {
5669 part_id: PartId::new("reasoning"),
5670 chunk: "partial thought".into(),
5671 }),
5672 ModelTurnEvent::Delta(Delta::SetMetadata {
5673 part_id: PartId::new("reasoning"),
5674 metadata: reasoning_metadata,
5675 }),
5676 ModelTurnEvent::Delta(Delta::BeginPart {
5677 part_id: PartId::new("structured"),
5678 kind: PartKind::Structured,
5679 }),
5680 ModelTurnEvent::Delta(Delta::ReplaceStructured {
5681 part_id: PartId::new("structured"),
5682 value: json!({ "complete": false }),
5683 }),
5684 ModelTurnEvent::Delta(Delta::BeginPart {
5685 part_id: PartId::new("media"),
5686 kind: PartKind::Media,
5687 }),
5688 ModelTurnEvent::Delta(Delta::AppendBytes {
5689 part_id: PartId::new("media"),
5690 chunk: vec![1, 2, 3],
5691 }),
5692 ModelTurnEvent::ToolCall(ToolCallPart::new(
5693 "partial-call",
5694 "unfinished-tool",
5695 json!({}),
5696 )),
5697 ])
5698 }
5699 InterruptedStreamScenario::CancelOnSupersession => VecDeque::from([
5700 ModelTurnEvent::Delta(Delta::BeginPart {
5701 part_id: PartId::new("stale"),
5702 kind: PartKind::Text,
5703 }),
5704 ModelTurnEvent::Delta(Delta::AppendText {
5705 part_id: PartId::new("stale"),
5706 chunk: "discard me".into(),
5707 }),
5708 ModelTurnEvent::ResponseAttemptSuperseded,
5709 ]),
5710 InterruptedStreamScenario::OverflowThenSupersession => VecDeque::from([
5711 ModelTurnEvent::Delta(Delta::BeginPart {
5712 part_id: PartId::new("oversized"),
5713 kind: PartKind::Text,
5714 }),
5715 ModelTurnEvent::Delta(Delta::AppendText {
5716 part_id: PartId::new("oversized"),
5717 chunk: "x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES),
5718 }),
5719 ModelTurnEvent::ResponseAttemptSuperseded,
5720 ModelTurnEvent::Delta(Delta::BeginPart {
5721 part_id: PartId::new("fresh"),
5722 kind: PartKind::Text,
5723 }),
5724 ModelTurnEvent::Delta(Delta::AppendText {
5725 part_id: PartId::new("fresh"),
5726 chunk: "preserve me".into(),
5727 }),
5728 ]),
5729 };
5730 Ok(InterruptedStreamTurn {
5731 controller: self.controller.clone(),
5732 events,
5733 cancel_on_supersession: matches!(
5734 self.scenario,
5735 InterruptedStreamScenario::CancelOnSupersession
5736 ),
5737 })
5738 }
5739 }
5740
5741 #[async_trait]
5742 impl ModelSession for FakeSession {
5743 type Turn = FakeTurn;
5744
5745 async fn begin_turn(
5746 &mut self,
5747 request: TurnRequest,
5748 _cancellation: Option<TurnCancellation>,
5749 ) -> Result<Self::Turn, LoopError> {
5750 let has_tool_result = request.transcript.iter().any(|item| {
5751 item.kind == ItemKind::Tool
5752 && item
5753 .parts
5754 .iter()
5755 .any(|part| matches!(part, Part::ToolResult(_)))
5756 });
5757 let tool_name = request
5758 .available_tools
5759 .first()
5760 .map(|tool| tool.name.0.clone())
5761 .unwrap_or_else(|| "echo".into());
5762
5763 let events = if has_tool_result {
5764 let result_text = request
5765 .transcript
5766 .iter()
5767 .rev()
5768 .find_map(|item| {
5769 item.parts.iter().find_map(|part| match (item.kind, part) {
5770 (ItemKind::Notification, Part::Text(text)) => Some(text.text.clone()),
5771 (
5772 _,
5773 Part::ToolResult(ToolResultPart {
5774 output: ToolOutput::Text(text),
5775 ..
5776 }),
5777 ) => Some(text.clone()),
5778 _ => None,
5779 })
5780 })
5781 .unwrap_or_else(|| "missing".into());
5782
5783 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
5784 model: None,
5785 response_id: None,
5786 finish_reason: FinishReason::Completed,
5787 output_items: vec![Item {
5788 id: None,
5789 kind: ItemKind::Assistant,
5790 parts: vec![Part::Text(TextPart {
5791 text: format!("tool said: {result_text}"),
5792 metadata: MetadataMap::new(),
5793 })],
5794 metadata: MetadataMap::new(),
5795 usage: None,
5796 finish_reason: None,
5797 created_at: None,
5798 }],
5799 usage: None,
5800 metadata: MetadataMap::new(),
5801 })])
5802 } else {
5803 VecDeque::from([
5804 ModelTurnEvent::ToolCall(agentkit_core::ToolCallPart {
5805 id: ToolCallId::new("call-1"),
5806 name: tool_name.clone(),
5807 input: json!({ "value": "pong" }),
5808 metadata: MetadataMap::new(),
5809 }),
5810 ModelTurnEvent::Finished(ModelTurnResult {
5811 model: None,
5812 response_id: None,
5813 finish_reason: FinishReason::ToolCall,
5814 output_items: vec![Item {
5815 id: None,
5816 kind: ItemKind::Assistant,
5817 parts: vec![Part::ToolCall(agentkit_core::ToolCallPart {
5818 id: ToolCallId::new("call-1"),
5819 name: tool_name,
5820 input: json!({ "value": "pong" }),
5821 metadata: MetadataMap::new(),
5822 })],
5823 metadata: MetadataMap::new(),
5824 usage: None,
5825 finish_reason: None,
5826 created_at: None,
5827 }],
5828 usage: None,
5829 metadata: MetadataMap::new(),
5830 }),
5831 ])
5832 };
5833
5834 Ok(FakeTurn { events })
5835 }
5836 }
5837
5838 #[async_trait]
5839 impl ModelSession for SlowSession {
5840 type Turn = SlowTurn;
5841
5842 async fn begin_turn(
5843 &mut self,
5844 request: TurnRequest,
5845 cancellation: Option<TurnCancellation>,
5846 ) -> Result<Self::Turn, LoopError> {
5847 let should_block = request
5848 .transcript
5849 .iter()
5850 .rev()
5851 .find(|item| item.kind == ItemKind::User)
5852 .is_some_and(|item| {
5853 item.parts.iter().any(|part| match part {
5854 Part::Text(text) => text.text == "do the long task",
5855 _ => false,
5856 })
5857 });
5858
5859 if should_block && let Some(cancellation) = cancellation {
5860 cancellation.cancelled().await;
5861 return Err(LoopError::Cancelled);
5862 }
5863
5864 Ok(SlowTurn { emitted: false })
5865 }
5866 }
5867
5868 #[async_trait]
5869 impl ModelSession for RecordingSession {
5870 type Turn = RecordingTurn;
5871
5872 async fn begin_turn(
5873 &mut self,
5874 request: TurnRequest,
5875 _cancellation: Option<TurnCancellation>,
5876 ) -> Result<Self::Turn, LoopError> {
5877 let descriptions = request
5878 .available_tools
5879 .iter()
5880 .map(|tool| tool.description.clone())
5881 .collect::<Vec<_>>();
5882 self.seen_descriptions.lock().unwrap().push(descriptions);
5883 self.seen_caches.lock().unwrap().push(request.cache.clone());
5884
5885 Ok(RecordingTurn { emitted: false })
5886 }
5887 }
5888
5889 #[async_trait]
5890 impl ModelSession for MultiToolSession {
5891 type Turn = MultiToolTurn;
5892
5893 async fn begin_turn(
5894 &mut self,
5895 request: TurnRequest,
5896 _cancellation: Option<TurnCancellation>,
5897 ) -> Result<Self::Turn, LoopError> {
5898 let has_tool_result = request.transcript.iter().any(|item| {
5899 item.kind == ItemKind::Tool
5900 && item
5901 .parts
5902 .iter()
5903 .any(|part| matches!(part, Part::ToolResult(_)))
5904 });
5905
5906 let events = if has_tool_result {
5907 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
5908 model: None,
5909 response_id: None,
5910 finish_reason: FinishReason::Completed,
5911 output_items: vec![Item {
5912 id: None,
5913 kind: ItemKind::Assistant,
5914 parts: vec![Part::Text(TextPart {
5915 text: "mixed tools finished".into(),
5916 metadata: MetadataMap::new(),
5917 })],
5918 metadata: MetadataMap::new(),
5919 usage: None,
5920 finish_reason: None,
5921 created_at: None,
5922 }],
5923 usage: None,
5924 metadata: MetadataMap::new(),
5925 })])
5926 } else {
5927 let foreground = agentkit_core::ToolCallPart {
5928 id: ToolCallId::new("call-foreground"),
5929 name: "foreground-wait".into(),
5930 input: json!({}),
5931 metadata: MetadataMap::new(),
5932 };
5933 let background = agentkit_core::ToolCallPart {
5934 id: ToolCallId::new("call-background"),
5935 name: "background-wait".into(),
5936 input: json!({}),
5937 metadata: MetadataMap::new(),
5938 };
5939 VecDeque::from([
5940 ModelTurnEvent::ToolCall(foreground.clone()),
5941 ModelTurnEvent::ToolCall(background.clone()),
5942 ModelTurnEvent::Finished(ModelTurnResult {
5943 model: None,
5944 response_id: None,
5945 finish_reason: FinishReason::ToolCall,
5946 output_items: vec![Item {
5947 id: None,
5948 kind: ItemKind::Assistant,
5949 parts: vec![Part::ToolCall(foreground), Part::ToolCall(background)],
5950 metadata: MetadataMap::new(),
5951 usage: None,
5952 finish_reason: None,
5953 created_at: None,
5954 }],
5955 usage: None,
5956 metadata: MetadataMap::new(),
5957 }),
5958 ])
5959 };
5960
5961 Ok(MultiToolTurn { events })
5962 }
5963 }
5964
5965 #[async_trait]
5966 impl ModelSession for DualApprovalSession {
5967 type Turn = DualApprovalTurn;
5968
5969 async fn begin_turn(
5970 &mut self,
5971 request: TurnRequest,
5972 _cancellation: Option<TurnCancellation>,
5973 ) -> Result<Self::Turn, LoopError> {
5974 let tool_results = request
5975 .transcript
5976 .iter()
5977 .flat_map(|item| item.parts.iter())
5978 .filter(|part| matches!(part, Part::ToolResult(_)))
5979 .count();
5980
5981 let events = if tool_results >= 2 {
5982 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
5983 model: None,
5984 response_id: None,
5985 finish_reason: FinishReason::Completed,
5986 output_items: vec![Item {
5987 id: None,
5988 kind: ItemKind::Assistant,
5989 parts: vec![Part::Text(TextPart {
5990 text: "both approvals finished".into(),
5991 metadata: MetadataMap::new(),
5992 })],
5993 metadata: MetadataMap::new(),
5994 usage: None,
5995 finish_reason: None,
5996 created_at: None,
5997 }],
5998 usage: None,
5999 metadata: MetadataMap::new(),
6000 })])
6001 } else {
6002 let first = agentkit_core::ToolCallPart {
6003 id: ToolCallId::new("call-1"),
6004 name: "echo".into(),
6005 input: json!({ "value": "first" }),
6006 metadata: MetadataMap::new(),
6007 };
6008 let second = agentkit_core::ToolCallPart {
6009 id: ToolCallId::new("call-2"),
6010 name: "echo".into(),
6011 input: json!({ "value": "second" }),
6012 metadata: MetadataMap::new(),
6013 };
6014 VecDeque::from([
6015 ModelTurnEvent::ToolCall(first.clone()),
6016 ModelTurnEvent::ToolCall(second.clone()),
6017 ModelTurnEvent::Finished(ModelTurnResult {
6018 model: None,
6019 response_id: None,
6020 finish_reason: FinishReason::ToolCall,
6021 output_items: vec![Item {
6022 id: None,
6023 kind: ItemKind::Assistant,
6024 parts: vec![Part::ToolCall(first), Part::ToolCall(second)],
6025 metadata: MetadataMap::new(),
6026 usage: None,
6027 finish_reason: None,
6028 created_at: None,
6029 }],
6030 usage: None,
6031 metadata: MetadataMap::new(),
6032 }),
6033 ])
6034 };
6035
6036 Ok(DualApprovalTurn { events })
6037 }
6038 }
6039
6040 #[async_trait]
6041 impl ModelTurn for FakeTurn {
6042 async fn next_event(
6043 &mut self,
6044 _cancellation: Option<TurnCancellation>,
6045 ) -> Result<Option<ModelTurnEvent>, LoopError> {
6046 Ok(self.events.pop_front())
6047 }
6048 }
6049
6050 #[async_trait]
6051 impl ModelTurn for SupersedingTurn {
6052 async fn next_event(
6053 &mut self,
6054 _cancellation: Option<TurnCancellation>,
6055 ) -> Result<Option<ModelTurnEvent>, LoopError> {
6056 Ok(self.events.pop_front())
6057 }
6058 }
6059
6060 #[async_trait]
6061 impl ModelTurn for InterruptedStreamTurn {
6062 async fn next_event(
6063 &mut self,
6064 _cancellation: Option<TurnCancellation>,
6065 ) -> Result<Option<ModelTurnEvent>, LoopError> {
6066 if let Some(event) = self.events.pop_front() {
6067 if self.cancel_on_supersession
6068 && matches!(event, ModelTurnEvent::ResponseAttemptSuperseded)
6069 {
6070 self.controller.interrupt();
6071 }
6072 return Ok(Some(event));
6073 }
6074 self.controller.interrupt();
6075 Err(LoopError::Cancelled)
6076 }
6077 }
6078
6079 #[async_trait]
6080 impl ModelTurn for SlowTurn {
6081 async fn next_event(
6082 &mut self,
6083 cancellation: Option<TurnCancellation>,
6084 ) -> Result<Option<ModelTurnEvent>, LoopError> {
6085 if let Some(cancellation) = cancellation
6086 && cancellation.is_cancelled()
6087 {
6088 return Err(LoopError::Cancelled);
6089 }
6090
6091 if self.emitted {
6092 Ok(None)
6093 } else {
6094 self.emitted = true;
6095 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
6096 model: None,
6097 response_id: None,
6098 finish_reason: FinishReason::Completed,
6099 output_items: vec![Item {
6100 id: None,
6101 kind: ItemKind::Assistant,
6102 parts: vec![Part::Text(TextPart {
6103 text: "done".into(),
6104 metadata: MetadataMap::new(),
6105 })],
6106 metadata: MetadataMap::new(),
6107 usage: None,
6108 finish_reason: None,
6109 created_at: None,
6110 }],
6111 usage: None,
6112 metadata: MetadataMap::new(),
6113 })))
6114 }
6115 }
6116 }
6117
6118 #[async_trait]
6119 impl ModelTurn for RecordingTurn {
6120 async fn next_event(
6121 &mut self,
6122 _cancellation: Option<TurnCancellation>,
6123 ) -> Result<Option<ModelTurnEvent>, LoopError> {
6124 if self.emitted {
6125 Ok(None)
6126 } else {
6127 self.emitted = true;
6128 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
6129 model: None,
6130 response_id: None,
6131 finish_reason: FinishReason::Completed,
6132 output_items: vec![Item {
6133 id: None,
6134 kind: ItemKind::Assistant,
6135 parts: vec![Part::Text(TextPart {
6136 text: "done".into(),
6137 metadata: MetadataMap::new(),
6138 })],
6139 metadata: MetadataMap::new(),
6140 usage: None,
6141 finish_reason: None,
6142 created_at: None,
6143 }],
6144 usage: None,
6145 metadata: MetadataMap::new(),
6146 })))
6147 }
6148 }
6149 }
6150
6151 #[async_trait]
6152 impl ModelTurn for MultiToolTurn {
6153 async fn next_event(
6154 &mut self,
6155 _cancellation: Option<TurnCancellation>,
6156 ) -> Result<Option<ModelTurnEvent>, LoopError> {
6157 Ok(self.events.pop_front())
6158 }
6159 }
6160
6161 #[async_trait]
6162 impl ModelTurn for DualApprovalTurn {
6163 async fn next_event(
6164 &mut self,
6165 _cancellation: Option<TurnCancellation>,
6166 ) -> Result<Option<ModelTurnEvent>, LoopError> {
6167 Ok(self.events.pop_front())
6168 }
6169 }
6170
6171 #[derive(Clone)]
6172 struct EchoTool {
6173 spec: ToolSpec,
6174 }
6175
6176 #[derive(Clone)]
6177 struct FailingTool {
6178 spec: ToolSpec,
6179 }
6180
6181 #[derive(Clone)]
6182 struct RunThenDenyTool {
6183 spec: ToolSpec,
6184 }
6185
6186 impl Default for EchoTool {
6187 fn default() -> Self {
6188 Self {
6189 spec: ToolSpec {
6190 name: ToolName::new("echo"),
6191 description: "Echo back a value".into(),
6192 input_schema: json!({
6193 "type": "object",
6194 "properties": {
6195 "value": { "type": "string" }
6196 },
6197 "required": ["value"],
6198 "additionalProperties": false
6199 }),
6200 output_schema: None,
6201 annotations: ToolAnnotations::default(),
6202 metadata: MetadataMap::new(),
6203 },
6204 }
6205 }
6206 }
6207
6208 impl Default for FailingTool {
6209 fn default() -> Self {
6210 Self {
6211 spec: ToolSpec {
6212 name: ToolName::new("failing"),
6213 description: "Always fails after execution starts".into(),
6214 input_schema: json!({
6215 "type": "object",
6216 "properties": {
6217 "value": { "type": "string" }
6218 },
6219 "additionalProperties": true
6220 }),
6221 output_schema: None,
6222 annotations: ToolAnnotations::default(),
6223 metadata: MetadataMap::new(),
6224 },
6225 }
6226 }
6227 }
6228
6229 impl Default for RunThenDenyTool {
6230 fn default() -> Self {
6231 Self {
6232 spec: ToolSpec {
6233 name: ToolName::new("run_then_deny"),
6234 description: "Runs, then returns a permission-denied error".into(),
6235 input_schema: json!({
6236 "type": "object",
6237 "properties": {
6238 "value": { "type": "string" }
6239 },
6240 "additionalProperties": true
6241 }),
6242 output_schema: None,
6243 annotations: ToolAnnotations::default(),
6244 metadata: MetadataMap::new(),
6245 },
6246 }
6247 }
6248 }
6249
6250 #[derive(Clone)]
6251 struct DynamicSpecTool {
6252 spec: ToolSpec,
6253 version: StdArc<AtomicUsize>,
6254 }
6255
6256 impl DynamicSpecTool {
6257 fn new(version: StdArc<AtomicUsize>) -> Self {
6258 Self {
6259 spec: ToolSpec {
6260 name: ToolName::new("dynamic"),
6261 description: "dynamic version 0".into(),
6262 input_schema: json!({
6263 "type": "object",
6264 "properties": {},
6265 "additionalProperties": false
6266 }),
6267 output_schema: None,
6268 annotations: ToolAnnotations::default(),
6269 metadata: MetadataMap::new(),
6270 },
6271 version,
6272 }
6273 }
6274 }
6275
6276 #[async_trait]
6277 impl Tool for EchoTool {
6278 fn spec(&self) -> &ToolSpec {
6279 &self.spec
6280 }
6281
6282 fn proposed_requests(
6283 &self,
6284 request: &agentkit_tools_core::ToolRequest,
6285 ) -> Result<
6286 Vec<Box<dyn agentkit_tools_core::PermissionRequest>>,
6287 agentkit_tools_core::ToolError,
6288 > {
6289 Ok(vec![Box::new(FileSystemPermissionRequest::Read {
6290 path: "/tmp/echo".into(),
6291 metadata: request.metadata.clone(),
6292 })])
6293 }
6294
6295 async fn invoke(
6296 &self,
6297 request: agentkit_tools_core::ToolRequest,
6298 _ctx: &mut ToolContext<'_>,
6299 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
6300 let value = request
6301 .input
6302 .get("value")
6303 .and_then(Value::as_str)
6304 .ok_or_else(|| {
6305 agentkit_tools_core::ToolError::InvalidInput("missing value".into())
6306 })?;
6307
6308 Ok(ToolResult {
6309 result: ToolResultPart {
6310 call_id: request.call_id,
6311 output: ToolOutput::Text(value.into()),
6312 is_error: false,
6313 metadata: MetadataMap::new(),
6314 },
6315 duration: None,
6316 metadata: MetadataMap::new(),
6317 })
6318 }
6319 }
6320
6321 #[async_trait]
6322 impl Tool for FailingTool {
6323 fn spec(&self) -> &ToolSpec {
6324 &self.spec
6325 }
6326
6327 async fn invoke(
6328 &self,
6329 _request: agentkit_tools_core::ToolRequest,
6330 _ctx: &mut ToolContext<'_>,
6331 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
6332 Err(agentkit_tools_core::ToolError::ExecutionFailed(
6333 "runtime failed".into(),
6334 ))
6335 }
6336 }
6337
6338 #[async_trait]
6339 impl Tool for RunThenDenyTool {
6340 fn spec(&self) -> &ToolSpec {
6341 &self.spec
6342 }
6343
6344 async fn invoke(
6345 &self,
6346 _request: agentkit_tools_core::ToolRequest,
6347 _ctx: &mut ToolContext<'_>,
6348 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
6349 Err(agentkit_tools_core::ToolError::PermissionDenied(
6350 PermissionDenial {
6351 code: PermissionCode::CustomPolicyDenied,
6352 message: "remote 403".into(),
6353 metadata: MetadataMap::new(),
6354 },
6355 ))
6356 }
6357 }
6358
6359 #[async_trait]
6360 impl Tool for DynamicSpecTool {
6361 fn spec(&self) -> &ToolSpec {
6362 &self.spec
6363 }
6364
6365 fn current_spec(&self) -> Option<ToolSpec> {
6366 let mut spec = self.spec.clone();
6367 spec.description = format!("dynamic version {}", self.version.load(Ordering::SeqCst));
6368 Some(spec)
6369 }
6370
6371 async fn invoke(
6372 &self,
6373 request: agentkit_tools_core::ToolRequest,
6374 _ctx: &mut ToolContext<'_>,
6375 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
6376 Ok(ToolResult {
6377 result: ToolResultPart {
6378 call_id: request.call_id,
6379 output: ToolOutput::Text("ok".into()),
6380 is_error: false,
6381 metadata: MetadataMap::new(),
6382 },
6383 duration: None,
6384 metadata: MetadataMap::new(),
6385 })
6386 }
6387 }
6388
6389 struct DenyFsReads;
6390
6391 impl PermissionChecker for DenyFsReads {
6392 fn evaluate(
6393 &self,
6394 request: &dyn agentkit_tools_core::PermissionRequest,
6395 ) -> PermissionDecision {
6396 if request.kind() == "filesystem.read" {
6397 return PermissionDecision::Deny(PermissionDenial {
6398 code: PermissionCode::PathNotAllowed,
6399 message: "reads denied in test".into(),
6400 metadata: MetadataMap::new(),
6401 });
6402 }
6403
6404 PermissionDecision::Allow
6405 }
6406 }
6407
6408 struct ApproveFsReads;
6409
6410 impl PermissionChecker for ApproveFsReads {
6411 fn evaluate(
6412 &self,
6413 request: &dyn agentkit_tools_core::PermissionRequest,
6414 ) -> PermissionDecision {
6415 if request.kind() == "filesystem.read" {
6416 return PermissionDecision::RequireApproval(ApprovalRequest {
6417 task_id: None,
6418 call_id: None,
6419 id: "approval:fs-read".into(),
6420 request_kind: request.kind().into(),
6421 reason: agentkit_tools_core::ApprovalReason::SensitivePath,
6422 summary: request.summary(),
6423 metadata: request.metadata().clone(),
6424 });
6425 }
6426
6427 PermissionDecision::Allow
6428 }
6429 }
6430
6431 struct KeepRecentMutator {
6432 keep: usize,
6433 }
6434
6435 #[async_trait]
6436 impl LoopMutator for KeepRecentMutator {
6437 async fn mutate(
6438 &self,
6439 cursor: &mut TranscriptCursor<'_>,
6440 ctx: LoopCtx<'_>,
6441 ) -> Result<(), LoopError> {
6442 if cursor.len() < 2 {
6443 return Ok(());
6444 }
6445 let drop = cursor.len().saturating_sub(self.keep);
6446 ctx.emitter.emit(AgentEvent::MutationStarted {
6447 session_id: ctx.session_id.clone(),
6448 turn_id: ctx.turn_id.cloned(),
6449 mutator: "keep-recent".into(),
6450 point: ctx.point,
6451 });
6452 cursor.drain(..drop);
6453 ctx.emitter.emit(AgentEvent::MutationFinished {
6454 session_id: ctx.session_id.clone(),
6455 turn_id: ctx.turn_id.cloned(),
6456 mutator: "keep-recent".into(),
6457 dirty: true,
6458 metadata: MetadataMap::new(),
6459 });
6460 Ok(())
6461 }
6462 }
6463
6464 struct PointRecordingMutator {
6467 points: StdArc<StdMutex<Vec<MutationPoint>>>,
6468 }
6469
6470 #[async_trait]
6471 impl LoopMutator for PointRecordingMutator {
6472 async fn mutate(
6473 &self,
6474 _cursor: &mut TranscriptCursor<'_>,
6475 ctx: LoopCtx<'_>,
6476 ) -> Result<(), LoopError> {
6477 self.points.lock().unwrap().push(ctx.point);
6478 Ok(())
6479 }
6480 }
6481
6482 struct RecordingObserver {
6483 events: StdArc<StdMutex<Vec<AgentEvent>>>,
6484 }
6485
6486 impl LoopObserver for RecordingObserver {
6487 fn handle_event(&self, event: ObservedEvent) {
6488 let event = event.event;
6489 self.events.lock().unwrap().push(event);
6490 }
6491 }
6492
6493 #[test]
6494 fn session_consumer_capabilities_are_typed_and_serde_defaulted() {
6495 let config = SessionConfig::new("session").with_response_attempt_supersession();
6496 assert!(config.consumer_capabilities.response_attempt_supersession);
6497
6498 let decoded: SessionConfig = serde_json::from_value(json!({
6499 "session_id": "session",
6500 "metadata": {},
6501 "cache": null
6502 }))
6503 .unwrap();
6504 assert_eq!(
6505 decoded.consumer_capabilities,
6506 SessionConsumerCapabilities::default()
6507 );
6508 }
6509
6510 #[tokio::test]
6511 async fn response_attempt_supersession_is_forwarded_and_resets_attempt_state() {
6512 let events = StdArc::new(StdMutex::new(Vec::new()));
6513 let agent = Agent::builder()
6514 .model(SupersedingAdapter)
6515 .observer(RecordingObserver {
6516 events: events.clone(),
6517 })
6518 .build()
6519 .unwrap();
6520 let mut driver = agent
6521 .start(SessionConfig::new("supersession-session").with_response_attempt_supersession())
6522 .await
6523 .unwrap();
6524 driver
6525 .submit_input(vec![Item::text(ItemKind::User, "hello")])
6526 .unwrap();
6527
6528 let LoopStep::Finished(result) = run_until_finished(&mut driver).await else {
6529 panic!("turn did not finish");
6530 };
6531 assert!(result.usage.is_none());
6532
6533 let events = events.lock().unwrap();
6534 let tool_call = events
6535 .iter()
6536 .position(|event| matches!(event, AgentEvent::ToolCallRequested(_)))
6537 .unwrap();
6538 let superseded = events
6539 .iter()
6540 .position(|event| matches!(event, AgentEvent::ResponseAttemptSuperseded))
6541 .unwrap();
6542 assert!(tool_call < superseded);
6543 }
6544
6545 #[tokio::test]
6546 async fn cancellation_preserves_streamed_content_in_result_and_transcript() {
6547 let controller = CancellationController::new();
6548 let agent = Agent::builder()
6549 .model(InterruptedStreamAdapter {
6550 controller: controller.clone(),
6551 scenario: InterruptedStreamScenario::PreserveContent,
6552 })
6553 .cancellation(controller.handle())
6554 .build()
6555 .unwrap();
6556 let mut driver = agent
6557 .start(SessionConfig::new("preserve-interrupted-stream"))
6558 .await
6559 .unwrap();
6560 driver
6561 .submit_input(vec![Item::text(ItemKind::User, "start")])
6562 .unwrap();
6563
6564 let LoopStep::Finished(result) = run_until_finished(&mut driver).await else {
6565 panic!("turn did not finish");
6566 };
6567 assert_eq!(result.finish_reason, FinishReason::Cancelled);
6568 assert_eq!(result.items.len(), 1);
6569 assert_eq!(
6570 result.items[0].metadata.get(INTERRUPTED_METADATA_KEY),
6571 Some(&Value::Bool(true))
6572 );
6573 assert_eq!(result.items[0].parts.len(), 4);
6574 assert!(matches!(
6575 &result.items[0].parts[0],
6576 Part::Text(text) if text.text == "partial answer"
6577 ));
6578 assert!(matches!(
6579 &result.items[0].parts[1],
6580 Part::Reasoning(reasoning)
6581 if reasoning.summary.as_deref() == Some("partial thought")
6582 && reasoning.metadata.get("provider.detail") == Some(&Value::Bool(true))
6583 ));
6584 assert!(matches!(
6585 &result.items[0].parts[2],
6586 Part::Structured(structured) if structured.value == json!({ "complete": false })
6587 ));
6588 assert!(matches!(
6589 &result.items[0].parts[3],
6590 Part::ToolCall(call) if call.id == ToolCallId::new("partial-call")
6591 ));
6592 assert!(
6593 !result.items[0]
6594 .parts
6595 .iter()
6596 .any(|part| matches!(part, Part::Media(_)))
6597 );
6598 assert!(!result.items[0].parts.iter().any(|part| {
6599 matches!(part, Part::Text(text) if text.text.contains("response was interrupted"))
6600 }));
6601
6602 let transcript = driver.snapshot().transcript;
6603 let preserved = transcript
6604 .iter()
6605 .find(|item| {
6606 item.kind == ItemKind::Assistant
6607 && item.metadata.get(INTERRUPTED_METADATA_KEY) == Some(&Value::Bool(true))
6608 })
6609 .expect("preserved assistant item in transcript");
6610 assert_eq!(preserved.parts, result.items[0].parts);
6611 assert!(transcript.iter().any(|item| {
6612 item.parts.iter().any(|part| {
6613 matches!(
6614 part,
6615 Part::ToolResult(tool_result)
6616 if tool_result.call_id == ToolCallId::new("partial-call")
6617 && tool_result.is_error
6618 )
6619 })
6620 }));
6621 validate_transcript_invariants(&transcript).unwrap();
6622 }
6623
6624 #[test]
6625 fn committed_part_matches_the_correct_interleaved_open_part() {
6626 let mut content = StreamedAssistantContent::default();
6627 for delta in [
6628 Delta::BeginPart {
6629 part_id: PartId::new("first"),
6630 kind: PartKind::Text,
6631 },
6632 Delta::AppendText {
6633 part_id: PartId::new("first"),
6634 chunk: "one".into(),
6635 },
6636 Delta::BeginPart {
6637 part_id: PartId::new("second"),
6638 kind: PartKind::Text,
6639 },
6640 Delta::AppendText {
6641 part_id: PartId::new("second"),
6642 chunk: "two".into(),
6643 },
6644 Delta::CommitPart {
6645 part: Part::text("two"),
6646 },
6647 ] {
6648 content.apply_delta(&delta);
6649 }
6650
6651 let items = content.interrupted_items();
6652 assert!(matches!(
6653 items[0].parts.as_slice(),
6654 [Part::Text(first), Part::Text(second)]
6655 if first.text == "one" && second.text == "two"
6656 ));
6657 }
6658
6659 #[test]
6660 fn ambiguous_committed_part_does_not_duplicate_identical_open_parts() {
6661 let mut content = StreamedAssistantContent::default();
6662 for part_id in [PartId::new("first"), PartId::new("second")] {
6663 content.apply_delta(&Delta::BeginPart {
6664 part_id: part_id.clone(),
6665 kind: PartKind::Text,
6666 });
6667 content.apply_delta(&Delta::AppendText {
6668 part_id,
6669 chunk: "same".into(),
6670 });
6671 }
6672 content.apply_delta(&Delta::CommitPart {
6673 part: Part::text("same"),
6674 });
6675
6676 let items = content.interrupted_items();
6677 assert_eq!(items[0].parts.len(), 2);
6678 assert!(
6679 items[0]
6680 .parts
6681 .iter()
6682 .all(|part| matches!(part, Part::Text(text) if text.text == "same"))
6683 );
6684 }
6685
6686 #[test]
6687 fn deltas_without_begin_part_are_not_guessed() {
6688 let mut content = StreamedAssistantContent::default();
6689 content.apply_delta(&Delta::AppendText {
6690 part_id: PartId::new("missing"),
6691 chunk: "orphan".into(),
6692 });
6693
6694 let items = content.interrupted_items();
6695 assert!(matches!(
6696 items[0].parts.as_slice(),
6697 [Part::Text(text)] if text.text.contains("response was interrupted")
6698 ));
6699 }
6700
6701 #[test]
6702 fn streamed_content_byte_budget_covers_every_retained_payload() {
6703 fn open_content(kind: PartKind) -> StreamedAssistantContent {
6704 let mut content = StreamedAssistantContent::default();
6705 content.apply_delta(&Delta::BeginPart {
6706 part_id: PartId::new("part"),
6707 kind,
6708 });
6709 content
6710 }
6711
6712 let mut content = open_content(PartKind::Text);
6713 content.apply_delta(&Delta::AppendText {
6714 part_id: PartId::new("part"),
6715 chunk: "x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES),
6716 });
6717 assert!(content.overflowed);
6718
6719 let mut content = open_content(PartKind::Reasoning);
6720 content.apply_delta(&Delta::AppendBytes {
6721 part_id: PartId::new("part"),
6722 chunk: vec![0; MAX_STREAMED_ASSISTANT_CONTENT_BYTES],
6723 });
6724 assert!(content.overflowed);
6725
6726 let mut content = open_content(PartKind::Structured);
6727 content.apply_delta(&Delta::ReplaceStructured {
6728 part_id: PartId::new("part"),
6729 value: Value::String("x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES)),
6730 });
6731 assert!(content.overflowed);
6732
6733 let mut metadata = MetadataMap::new();
6734 metadata.insert(
6735 "large".into(),
6736 Value::String("x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES)),
6737 );
6738 let mut content = open_content(PartKind::Text);
6739 content.apply_delta(&Delta::SetMetadata {
6740 part_id: PartId::new("part"),
6741 metadata,
6742 });
6743 assert!(content.overflowed);
6744
6745 let mut content = StreamedAssistantContent::default();
6746 content.apply_delta(&Delta::CommitPart {
6747 part: Part::text("x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES)),
6748 });
6749 assert!(content.overflowed);
6750
6751 let mut content = StreamedAssistantContent::default();
6752 content.commit_tool_call(&ToolCallPart::new(
6753 "call",
6754 "tool",
6755 Value::String("x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES)),
6756 ));
6757 assert!(content.overflowed);
6758 }
6759
6760 #[test]
6761 fn streamed_content_part_budget_releases_all_parts() {
6762 let mut content = StreamedAssistantContent::default();
6763 for index in 0..MAX_STREAMED_ASSISTANT_CONTENT_PARTS {
6764 content.apply_delta(&Delta::BeginPart {
6765 part_id: PartId::new(format!("part-{index}")),
6766 kind: PartKind::Text,
6767 });
6768 }
6769 assert_eq!(content.parts.len(), MAX_STREAMED_ASSISTANT_CONTENT_PARTS);
6770
6771 content.apply_delta(&Delta::BeginPart {
6772 part_id: PartId::new("overflow"),
6773 kind: PartKind::Text,
6774 });
6775
6776 assert!(content.overflowed);
6777 assert!(content.parts.is_empty());
6778 assert_eq!(content.retained_bytes, 0);
6779 }
6780
6781 #[test]
6782 fn streamed_content_overflow_uses_generic_interruption_fallback() {
6783 let mut content = StreamedAssistantContent::default();
6784 content.apply_delta(&Delta::BeginPart {
6785 part_id: PartId::new("text"),
6786 kind: PartKind::Text,
6787 });
6788 content.apply_delta(&Delta::AppendText {
6789 part_id: PartId::new("text"),
6790 chunk: "release me".into(),
6791 });
6792 content.apply_delta(&Delta::AppendText {
6793 part_id: PartId::new("text"),
6794 chunk: "x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES),
6795 });
6796 content.apply_delta(&Delta::BeginPart {
6797 part_id: PartId::new("ignored after overflow"),
6798 kind: PartKind::Text,
6799 });
6800
6801 assert!(content.parts.is_empty());
6802 assert_eq!(content.retained_bytes, 0);
6803 let items = content.interrupted_items();
6804 assert!(matches!(
6805 items[0].parts.as_slice(),
6806 [Part::Text(text)]
6807 if text.text.contains("response was interrupted")
6808 && !text.text.contains("release me")
6809 ));
6810 }
6811
6812 #[tokio::test]
6813 async fn supersession_is_processed_before_concurrent_cancellation() {
6814 let controller = CancellationController::new();
6815 let events = StdArc::new(StdMutex::new(Vec::new()));
6816 let agent = Agent::builder()
6817 .model(InterruptedStreamAdapter {
6818 controller: controller.clone(),
6819 scenario: InterruptedStreamScenario::CancelOnSupersession,
6820 })
6821 .cancellation(controller.handle())
6822 .observer(RecordingObserver {
6823 events: events.clone(),
6824 })
6825 .build()
6826 .unwrap();
6827 let mut driver = agent
6828 .start(SessionConfig::new("supersede-then-cancel").with_response_attempt_supersession())
6829 .await
6830 .unwrap();
6831 driver
6832 .submit_input(vec![Item::text(ItemKind::User, "start")])
6833 .unwrap();
6834
6835 let LoopStep::Finished(result) = run_until_finished(&mut driver).await else {
6836 panic!("turn did not finish");
6837 };
6838 assert_eq!(result.finish_reason, FinishReason::Cancelled);
6839 assert_eq!(result.items.len(), 1);
6840 assert!(matches!(
6841 result.items[0].parts.as_slice(),
6842 [Part::Text(text)]
6843 if text.text.contains("response was interrupted")
6844 && !text.text.contains("discard me")
6845 ));
6846 assert!(!driver.snapshot().transcript.iter().any(|item| {
6847 item.parts
6848 .iter()
6849 .any(|part| matches!(part, Part::Text(text) if text.text.contains("discard me")))
6850 }));
6851
6852 let events = events.lock().unwrap();
6853 let superseded = events
6854 .iter()
6855 .position(|event| matches!(event, AgentEvent::ResponseAttemptSuperseded))
6856 .expect("supersession event");
6857 let finished = events
6858 .iter()
6859 .position(|event| matches!(event, AgentEvent::TurnFinished(_)))
6860 .expect("turn-finished event");
6861 assert!(superseded < finished);
6862 }
6863
6864 #[tokio::test]
6865 async fn supersession_resets_streamed_content_budget() {
6866 let controller = CancellationController::new();
6867 let agent = Agent::builder()
6868 .model(InterruptedStreamAdapter {
6869 controller: controller.clone(),
6870 scenario: InterruptedStreamScenario::OverflowThenSupersession,
6871 })
6872 .cancellation(controller.handle())
6873 .build()
6874 .unwrap();
6875 let mut driver = agent
6876 .start(
6877 SessionConfig::new("overflow-then-supersede").with_response_attempt_supersession(),
6878 )
6879 .await
6880 .unwrap();
6881 driver
6882 .submit_input(vec![Item::text(ItemKind::User, "start")])
6883 .unwrap();
6884
6885 let LoopStep::Finished(result) = run_until_finished(&mut driver).await else {
6886 panic!("turn did not finish");
6887 };
6888 assert_eq!(result.finish_reason, FinishReason::Cancelled);
6889 assert!(matches!(
6890 result.items[0].parts.as_slice(),
6891 [Part::Text(text)] if text.text == "preserve me"
6892 ));
6893 }
6894
6895 fn turn_lifecycle_events(
6896 events: &[AgentEvent],
6897 ) -> Vec<(agentkit_core::TurnId, Option<FinishReason>)> {
6898 events
6899 .iter()
6900 .filter_map(|event| match event {
6901 AgentEvent::TurnStarted { turn_id, .. } => Some((turn_id.clone(), None)),
6902 AgentEvent::TurnFinished(turn) => {
6903 Some((turn.turn_id.clone(), Some(turn.finish_reason.clone())))
6904 }
6905 _ => None,
6906 })
6907 .collect()
6908 }
6909
6910 struct CatalogExecutor {
6911 version: AtomicUsize,
6912 events: StdMutex<Vec<ToolCatalogEvent>>,
6913 }
6914
6915 impl CatalogExecutor {
6916 fn new() -> Self {
6917 Self {
6918 version: AtomicUsize::new(0),
6919 events: StdMutex::new(Vec::new()),
6920 }
6921 }
6922
6923 fn publish_change(&self, version: usize, event: ToolCatalogEvent) {
6924 self.version.store(version, Ordering::SeqCst);
6925 self.events.lock().unwrap().push(event);
6926 }
6927 }
6928
6929 #[async_trait]
6930 impl ToolExecutor for CatalogExecutor {
6931 fn specs(&self) -> Vec<ToolSpec> {
6932 vec![ToolSpec {
6933 name: ToolName::new("dynamic"),
6934 description: format!("dynamic version {}", self.version.load(Ordering::SeqCst)),
6935 input_schema: json!({
6936 "type": "object",
6937 "properties": {},
6938 "additionalProperties": false
6939 }),
6940 output_schema: None,
6941 annotations: ToolAnnotations::default(),
6942 metadata: MetadataMap::new(),
6943 }]
6944 }
6945
6946 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
6947 std::mem::take(&mut *self.events.lock().unwrap())
6948 }
6949
6950 async fn execute(
6951 &self,
6952 request: ToolRequest,
6953 _ctx: &mut ToolContext<'_>,
6954 ) -> ToolExecutionOutcome {
6955 ToolExecutionOutcome::Completed(ToolResult {
6956 result: ToolResultPart {
6957 call_id: request.call_id,
6958 output: ToolOutput::Text("dynamic-ok".into()),
6959 is_error: false,
6960 metadata: MetadataMap::new(),
6961 },
6962 duration: None,
6963 metadata: MetadataMap::new(),
6964 })
6965 }
6966 }
6967
6968 #[derive(Clone)]
6969 struct BlockingTool {
6970 spec: ToolSpec,
6971 entered: StdArc<AtomicBool>,
6972 release: StdArc<Notify>,
6973 output: &'static str,
6974 }
6975
6976 impl BlockingTool {
6977 fn new(
6978 name: &str,
6979 entered: StdArc<AtomicBool>,
6980 release: StdArc<Notify>,
6981 output: &'static str,
6982 ) -> Self {
6983 Self {
6984 spec: ToolSpec {
6985 name: ToolName::new(name),
6986 description: format!("blocking tool {name}"),
6987 input_schema: json!({
6988 "type": "object",
6989 "properties": {},
6990 "additionalProperties": false
6991 }),
6992 output_schema: None,
6993 annotations: ToolAnnotations::default(),
6994 metadata: MetadataMap::new(),
6995 },
6996 entered,
6997 release,
6998 output,
6999 }
7000 }
7001 }
7002
7003 #[async_trait]
7004 impl Tool for BlockingTool {
7005 fn spec(&self) -> &ToolSpec {
7006 &self.spec
7007 }
7008
7009 async fn invoke(
7010 &self,
7011 request: agentkit_tools_core::ToolRequest,
7012 _ctx: &mut ToolContext<'_>,
7013 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
7014 self.entered.store(true, Ordering::SeqCst);
7015 self.release.notified().await;
7016 Ok(ToolResult {
7017 result: ToolResultPart {
7018 call_id: request.call_id,
7019 output: ToolOutput::Text(self.output.into()),
7020 is_error: false,
7021 metadata: MetadataMap::new(),
7022 },
7023 duration: None,
7024 metadata: MetadataMap::new(),
7025 })
7026 }
7027 }
7028
7029 struct NameRoutingPolicy {
7030 routes: Vec<(String, RoutingDecision)>,
7031 }
7032
7033 impl NameRoutingPolicy {
7034 fn new(routes: impl IntoIterator<Item = (impl Into<String>, RoutingDecision)>) -> Self {
7035 Self {
7036 routes: routes
7037 .into_iter()
7038 .map(|(name, decision)| (name.into(), decision))
7039 .collect(),
7040 }
7041 }
7042 }
7043
7044 impl TaskRoutingPolicy for NameRoutingPolicy {
7045 fn route(&self, request: &ToolRequest) -> RoutingDecision {
7046 self.routes
7047 .iter()
7048 .find(|(name, _)| name == &request.tool_name.0)
7049 .map(|(_, decision)| *decision)
7050 .unwrap_or(RoutingDecision::Foreground)
7051 }
7052 }
7053
7054 async fn wait_for_task_event(handle: &TaskManagerHandle) -> TaskEvent {
7055 timeout(Duration::from_secs(1), handle.next_event())
7056 .await
7057 .expect("timed out waiting for task event")
7058 .expect("task event stream ended unexpectedly")
7059 }
7060
7061 async fn wait_until_entered(flag: &AtomicBool) {
7062 timeout(Duration::from_secs(1), async {
7063 while !flag.load(Ordering::SeqCst) {
7064 tokio::task::yield_now().await;
7065 }
7066 })
7067 .await
7068 .expect("task never entered execution");
7069 }
7070
7071 async fn wait_until_completed(handle: &TaskManagerHandle) {
7072 timeout(Duration::from_secs(1), async {
7073 while handle.list_completed().await.is_empty() {
7074 tokio::task::yield_now().await;
7075 }
7076 })
7077 .await
7078 .expect("task never completed");
7079 }
7080
7081 #[tokio::test]
7082 async fn loop_continues_after_completed_tool_call() {
7083 let events = StdArc::new(StdMutex::new(Vec::new()));
7084 let tools = ToolRegistry::new().with(EchoTool::default());
7085 let agent = Agent::builder()
7086 .model(FakeAdapter)
7087 .add_tool_source(tools)
7088 .permissions(AllowAllPermissions)
7089 .observer(RecordingObserver {
7090 events: events.clone(),
7091 })
7092 .build()
7093 .unwrap();
7094
7095 let mut driver = agent
7096 .start(SessionConfig {
7097 session_id: SessionId::new("session-1"),
7098 metadata: MetadataMap::new(),
7099 cache: None,
7100 consumer_capabilities: SessionConsumerCapabilities::default(),
7101 })
7102 .await
7103 .unwrap();
7104
7105 driver
7106 .submit_input(vec![Item {
7107 id: None,
7108 kind: ItemKind::User,
7109 parts: vec![Part::Text(TextPart {
7110 text: "ping".into(),
7111 metadata: MetadataMap::new(),
7112 })],
7113 metadata: MetadataMap::new(),
7114 usage: None,
7115 finish_reason: None,
7116 created_at: None,
7117 }])
7118 .unwrap();
7119
7120 let result = run_until_finished(&mut driver).await;
7121
7122 match result {
7123 LoopStep::Finished(turn) => {
7124 assert_eq!(turn.finish_reason, FinishReason::Completed);
7125 assert_eq!(turn.items.len(), 1);
7126 match &turn.items[0].parts[0] {
7127 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
7128 other => panic!("unexpected part: {other:?}"),
7129 }
7130 }
7131 other => panic!("unexpected loop step: {other:?}"),
7132 }
7133
7134 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
7135 assert_eq!(lifecycle.len(), 2);
7136 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7137 assert_eq!(lifecycle[0].1, None);
7138 assert_eq!(lifecycle[1].1, Some(FinishReason::Completed));
7139 }
7140
7141 async fn run_until_finished<S: ModelSession + Send>(driver: &mut LoopDriver<S>) -> LoopStep {
7145 loop {
7146 match driver.next().await.unwrap() {
7147 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue,
7148 step => return step,
7149 }
7150 }
7151 }
7152
7153 #[tokio::test]
7160 async fn post_tool_continuation_reports_after_tool_result_mutation_point() {
7161 let points = StdArc::new(StdMutex::new(Vec::<MutationPoint>::new()));
7162 let tools = ToolRegistry::new().with(EchoTool::default());
7163 let agent = Agent::builder()
7164 .model(FakeAdapter)
7165 .add_tool_source(tools)
7166 .permissions(AllowAllPermissions)
7167 .mutator(PointRecordingMutator {
7168 points: points.clone(),
7169 })
7170 .build()
7171 .unwrap();
7172
7173 let mut driver = agent
7174 .start(SessionConfig {
7175 session_id: SessionId::new("session-mutation-point"),
7176 metadata: MetadataMap::new(),
7177 cache: None,
7178 consumer_capabilities: SessionConsumerCapabilities::default(),
7179 })
7180 .await
7181 .unwrap();
7182
7183 driver
7184 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7185 .unwrap();
7186
7187 let _ = run_until_finished(&mut driver).await;
7189
7190 let recorded = points.lock().unwrap().clone();
7191 assert_eq!(
7192 recorded.first(),
7193 Some(&MutationPoint::AfterTurnEnded),
7194 "first drive of a fresh turn must report AfterTurnEnded, got {recorded:?}"
7195 );
7196 assert!(
7197 recorded.contains(&MutationPoint::AfterToolResult),
7198 "post-tool continuation must report AfterToolResult, got {recorded:?}"
7199 );
7200 }
7201
7202 #[tokio::test]
7203 async fn no_work_awaiting_input_emits_no_turn_lifecycle() {
7204 let events = StdArc::new(StdMutex::new(Vec::new()));
7205 let agent = Agent::builder()
7206 .model(SlowAdapter)
7207 .observer(RecordingObserver {
7208 events: events.clone(),
7209 })
7210 .build()
7211 .unwrap();
7212 let mut driver = agent
7213 .start(SessionConfig::new("session-no-work"))
7214 .await
7215 .unwrap();
7216
7217 for _ in 0..2 {
7218 assert!(matches!(
7219 driver.next().await.unwrap(),
7220 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
7221 ));
7222 }
7223
7224 assert!(turn_lifecycle_events(&events.lock().unwrap()).is_empty());
7225 }
7226
7227 #[tokio::test]
7228 async fn normal_turn_emits_one_matched_lifecycle_pair() {
7229 let events = StdArc::new(StdMutex::new(Vec::new()));
7230 let agent = Agent::builder()
7231 .model(SlowAdapter)
7232 .observer(RecordingObserver {
7233 events: events.clone(),
7234 })
7235 .build()
7236 .unwrap();
7237 let mut driver = agent
7238 .start(SessionConfig::new("session-normal-lifecycle"))
7239 .await
7240 .unwrap();
7241 driver
7242 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7243 .unwrap();
7244
7245 assert!(matches!(
7246 driver.next().await.unwrap(),
7247 LoopStep::Finished(TurnResult {
7248 finish_reason: FinishReason::Completed,
7249 ..
7250 })
7251 ));
7252
7253 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
7254 assert_eq!(lifecycle.len(), 2);
7255 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7256 assert_eq!(lifecycle[0].1, None);
7257 assert_eq!(lifecycle[1].1, Some(FinishReason::Completed));
7258 }
7259
7260 #[tokio::test]
7261 async fn post_start_error_emits_terminal_error_without_run_failed() {
7262 let events = StdArc::new(StdMutex::new(Vec::new()));
7263 let agent = Agent::builder()
7264 .model(SlowAdapter)
7265 .mutator(ErrorMutator)
7266 .observer(RecordingObserver {
7267 events: events.clone(),
7268 })
7269 .build()
7270 .unwrap();
7271 let mut driver = agent
7272 .start(SessionConfig::new("session-error-lifecycle"))
7273 .await
7274 .unwrap();
7275 driver
7276 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7277 .unwrap();
7278
7279 assert!(matches!(
7280 driver.next().await,
7281 Err(LoopError::Mutator(message)) if message == "boom"
7282 ));
7283
7284 let events = events.lock().unwrap();
7285 let lifecycle = turn_lifecycle_events(&events);
7286 assert_eq!(lifecycle.len(), 2);
7287 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7288 assert_eq!(lifecycle[0].1, None);
7289 assert_eq!(lifecycle[1].1, Some(FinishReason::Error));
7290 assert!(
7291 !events
7292 .iter()
7293 .any(|event| matches!(event, AgentEvent::RunFailed { .. }))
7294 );
7295 }
7296
7297 #[tokio::test]
7298 async fn active_tool_error_repairs_state_and_retry_uses_fresh_lifecycle() {
7299 let events = StdArc::new(StdMutex::new(Vec::new()));
7300 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
7301 let agent = Agent::builder()
7302 .model(FakeAdapter)
7303 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
7304 .task_manager(
7305 TestTaskManager::new(SimpleTaskManager::new())
7306 .fail_start("original start failure")
7307 .record_interrupts(interrupted.clone())
7308 .fail_interrupt("cleanup failure"),
7309 )
7310 .observer(RecordingObserver {
7311 events: events.clone(),
7312 })
7313 .build()
7314 .unwrap();
7315 let mut driver = agent
7316 .start(SessionConfig::new("session-active-tool-error"))
7317 .await
7318 .unwrap();
7319 driver
7320 .submit_input(vec![Item::text(ItemKind::User, "first")])
7321 .unwrap();
7322
7323 let error = driver.next().await.unwrap_err();
7324 assert!(error.to_string().contains("original start failure"));
7325 assert!(!error.to_string().contains("cleanup failure"));
7326 assert!(driver.active_tool_round.is_none());
7327 assert!(driver.pending_round_resume.is_none());
7328 assert!(unanswered_tool_calls(&driver.snapshot().transcript).is_empty());
7329 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
7330 assert_eq!(interrupted.lock().unwrap().len(), 1);
7331
7332 driver
7333 .submit_input(vec![Item::text(ItemKind::User, "retry")])
7334 .unwrap();
7335 assert!(matches!(
7336 driver.next().await.unwrap(),
7337 LoopStep::Finished(TurnResult {
7338 finish_reason: FinishReason::Completed,
7339 ..
7340 })
7341 ));
7342
7343 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
7344 assert_eq!(lifecycle.len(), 4, "{lifecycle:?}");
7345 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7346 assert_eq!(lifecycle[1].1, Some(FinishReason::Error));
7347 assert_eq!(lifecycle[2].0, lifecycle[3].0);
7348 assert_eq!(lifecycle[3].1, Some(FinishReason::Completed));
7349 assert_ne!(lifecycle[0].0, lifecycle[2].0);
7350 }
7351
7352 #[tokio::test]
7353 async fn continuation_error_clears_resume_and_retry_uses_fresh_lifecycle() {
7354 let events = StdArc::new(StdMutex::new(Vec::new()));
7355 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
7356 let agent = Agent::builder()
7357 .model(FakeAdapter)
7358 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
7359 .task_manager(
7360 TestTaskManager::new(SimpleTaskManager::new())
7361 .fail_pending_update_on(1, "original continuation failure")
7362 .record_interrupts(interrupted.clone())
7363 .fail_interrupt("cleanup failure"),
7364 )
7365 .observer(RecordingObserver {
7366 events: events.clone(),
7367 })
7368 .build()
7369 .unwrap();
7370 let mut driver = agent
7371 .start(SessionConfig::new("session-continuation-error"))
7372 .await
7373 .unwrap();
7374 driver
7375 .submit_input(vec![Item::text(ItemKind::User, "first")])
7376 .unwrap();
7377
7378 assert!(matches!(
7379 driver.next().await.unwrap(),
7380 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_))
7381 ));
7382 let error = driver.next().await.unwrap_err();
7383 assert!(error.to_string().contains("original continuation failure"));
7384 assert!(!error.to_string().contains("cleanup failure"));
7385 assert!(driver.pending_round_resume.is_none());
7386 assert!(driver.active_tool_round.is_none());
7387 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
7388 assert_eq!(interrupted.lock().unwrap().len(), 1);
7389
7390 driver
7391 .submit_input(vec![Item::text(ItemKind::User, "retry")])
7392 .unwrap();
7393 assert!(matches!(
7394 driver.next().await.unwrap(),
7395 LoopStep::Finished(TurnResult {
7396 finish_reason: FinishReason::Completed,
7397 ..
7398 })
7399 ));
7400
7401 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
7402 assert_eq!(lifecycle.len(), 4, "{lifecycle:?}");
7403 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7404 assert_eq!(lifecycle[1].1, Some(FinishReason::Error));
7405 assert_eq!(lifecycle[2].0, lifecycle[3].0);
7406 assert_eq!(lifecycle[3].1, Some(FinishReason::Completed));
7407 assert_ne!(lifecycle[0].0, lifecycle[2].0);
7408 }
7409
7410 #[test]
7411 fn pending_input_requires_input_bearing_tail_role() {
7412 assert!(!transcript_has_pending_input(&[]));
7413 assert!(!transcript_has_pending_input(&[Item::text(
7414 ItemKind::System,
7415 "system"
7416 )]));
7417 assert!(!transcript_has_pending_input(&[Item::text(
7418 ItemKind::Developer,
7419 "developer"
7420 )]));
7421 assert!(!transcript_has_pending_input(&[Item::text(
7422 ItemKind::Context,
7423 "context"
7424 )]));
7425 assert!(!transcript_has_pending_input(&[Item::text(
7426 ItemKind::Assistant,
7427 "assistant"
7428 )]));
7429
7430 assert!(transcript_has_pending_input(&[Item::text(
7431 ItemKind::User,
7432 "user"
7433 )]));
7434 assert!(transcript_has_pending_input(&[Item::notification(
7435 "background update"
7436 )]));
7437 assert!(transcript_has_pending_input(&[Item {
7438 id: None,
7439 kind: ItemKind::Tool,
7440 parts: vec![Part::ToolResult(ToolResultPart {
7441 call_id: ToolCallId::new("call-test"),
7442 output: ToolOutput::Text("ok".into()),
7443 is_error: false,
7444 metadata: MetadataMap::new(),
7445 })],
7446 metadata: MetadataMap::new(),
7447 usage: None,
7448 finish_reason: None,
7449 created_at: None,
7450 }]));
7451 }
7452
7453 struct DropTrailingUserMutator;
7459 struct ErrorMutator;
7460
7461 #[async_trait]
7462 impl LoopMutator for ErrorMutator {
7463 async fn mutate(
7464 &self,
7465 _cursor: &mut TranscriptCursor<'_>,
7466 _ctx: LoopCtx<'_>,
7467 ) -> Result<(), LoopError> {
7468 Err(LoopError::Mutator("boom".into()))
7469 }
7470 }
7471
7472 #[async_trait]
7473 impl LoopMutator for DropTrailingUserMutator {
7474 async fn mutate(
7475 &self,
7476 cursor: &mut TranscriptCursor<'_>,
7477 _ctx: LoopCtx<'_>,
7478 ) -> Result<(), LoopError> {
7479 if cursor.last().map(|item| item.kind) == Some(ItemKind::User) {
7480 cursor.pop();
7481 }
7482 Ok(())
7483 }
7484 }
7485
7486 struct RejectAssistantPrefillAdapter {
7491 saw_assistant_tail: StdArc<AtomicBool>,
7492 }
7493
7494 struct RejectAssistantPrefillSession {
7495 saw_assistant_tail: StdArc<AtomicBool>,
7496 }
7497
7498 #[async_trait]
7499 impl ModelAdapter for RejectAssistantPrefillAdapter {
7500 type Session = RejectAssistantPrefillSession;
7501
7502 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
7503 Ok(RejectAssistantPrefillSession {
7504 saw_assistant_tail: self.saw_assistant_tail.clone(),
7505 })
7506 }
7507 }
7508
7509 #[async_trait]
7510 impl ModelSession for RejectAssistantPrefillSession {
7511 type Turn = FakeTurn;
7512
7513 async fn begin_turn(
7514 &mut self,
7515 request: TurnRequest,
7516 _cancellation: Option<TurnCancellation>,
7517 ) -> Result<Self::Turn, LoopError> {
7518 if request.transcript.last().map(|item| item.kind) == Some(ItemKind::Assistant) {
7519 self.saw_assistant_tail.store(true, Ordering::SeqCst);
7520 return Err(LoopError::Provider(
7521 "conversation must end with a user message".into(),
7522 ));
7523 }
7524 Ok(FakeTurn {
7525 events: VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
7526 model: None,
7527 response_id: None,
7528 finish_reason: FinishReason::Completed,
7529 output_items: vec![Item::text(ItemKind::Assistant, "ok")],
7530 usage: None,
7531 metadata: MetadataMap::new(),
7532 })]),
7533 })
7534 }
7535 }
7536
7537 #[tokio::test]
7545 async fn drive_does_not_dispatch_without_valid_trailing_input() {
7546 let saw_assistant_tail = StdArc::new(AtomicBool::new(false));
7547 let events = StdArc::new(StdMutex::new(Vec::new()));
7548 let agent = Agent::builder()
7549 .model(RejectAssistantPrefillAdapter {
7550 saw_assistant_tail: saw_assistant_tail.clone(),
7551 })
7552 .mutator(DropTrailingUserMutator)
7553 .observer(RecordingObserver {
7554 events: events.clone(),
7555 })
7556 .transcript(vec![
7559 Item::text(ItemKind::User, "kickoff"),
7560 Item::text(ItemKind::Assistant, "prior reply"),
7561 ])
7562 .build()
7563 .unwrap();
7564
7565 let mut driver = agent
7566 .start(SessionConfig {
7567 session_id: SessionId::new("session-no-valid-input"),
7568 metadata: MetadataMap::new(),
7569 cache: None,
7570 consumer_capabilities: SessionConsumerCapabilities::default(),
7571 })
7572 .await
7573 .unwrap();
7574
7575 driver
7576 .submit_input(vec![Item::text(ItemKind::User, "follow up")])
7577 .unwrap();
7578
7579 let outcome = driver.next().await;
7581
7582 assert!(
7583 !saw_assistant_tail.load(Ordering::SeqCst),
7584 "loop dispatched a model turn whose transcript ends in an assistant \
7585 message (outcome: {outcome:?}); with no valid trailing input the turn \
7586 must finish instead of driving"
7587 );
7588 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
7589 assert_eq!(lifecycle.len(), 2);
7590 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7591 assert_eq!(lifecycle[0].1, None);
7592 assert_eq!(lifecycle[1].1, Some(FinishReason::Completed));
7593 }
7594
7595 #[tokio::test]
7596 async fn loop_uses_injected_permission_checker() {
7597 let events = StdArc::new(StdMutex::new(Vec::new()));
7598 let tools = ToolRegistry::new().with(EchoTool::default());
7599 let agent = Agent::builder()
7600 .model(FakeAdapter)
7601 .add_tool_source(tools)
7602 .permissions(DenyFsReads)
7603 .observer(RecordingObserver {
7604 events: events.clone(),
7605 })
7606 .build()
7607 .unwrap();
7608
7609 let mut driver = agent
7610 .start(SessionConfig {
7611 session_id: SessionId::new("session-2"),
7612 metadata: MetadataMap::new(),
7613 cache: None,
7614 consumer_capabilities: SessionConsumerCapabilities::default(),
7615 })
7616 .await
7617 .unwrap();
7618
7619 driver
7620 .submit_input(vec![Item {
7621 id: None,
7622 kind: ItemKind::User,
7623 parts: vec![Part::Text(TextPart {
7624 text: "ping".into(),
7625 metadata: MetadataMap::new(),
7626 })],
7627 metadata: MetadataMap::new(),
7628 usage: None,
7629 finish_reason: None,
7630 created_at: None,
7631 }])
7632 .unwrap();
7633
7634 let result = run_until_finished(&mut driver).await;
7635
7636 match result {
7637 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
7638 Part::Text(text) => assert!(text.text.contains("tool permission denied")),
7639 other => panic!("unexpected part: {other:?}"),
7640 },
7641 other => panic!("unexpected loop step: {other:?}"),
7642 }
7643
7644 assert!(
7645 events
7646 .lock()
7647 .unwrap()
7648 .iter()
7649 .all(|event| !matches!(event, AgentEvent::ToolExecutionStarted(_))),
7650 "denied tools must not be reported as started"
7651 );
7652 }
7653
7654 #[tokio::test]
7655 async fn failed_tool_execution_still_reports_started() {
7656 let events = StdArc::new(StdMutex::new(Vec::new()));
7657 let tools = ToolRegistry::new().with(FailingTool::default());
7658 let agent = Agent::builder()
7659 .model(FakeAdapter)
7660 .add_tool_source(tools)
7661 .permissions(AllowAllPermissions)
7662 .observer(RecordingObserver {
7663 events: events.clone(),
7664 })
7665 .build()
7666 .unwrap();
7667
7668 let mut driver = agent
7669 .start(SessionConfig {
7670 session_id: SessionId::new("session-failing-start-event"),
7671 metadata: MetadataMap::new(),
7672 cache: None,
7673 consumer_capabilities: SessionConsumerCapabilities::default(),
7674 })
7675 .await
7676 .unwrap();
7677
7678 driver
7679 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7680 .unwrap();
7681
7682 match run_until_finished(&mut driver).await {
7683 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
7684 other => panic!("unexpected loop step: {other:?}"),
7685 }
7686
7687 let events = events.lock().unwrap();
7688 assert!(events.iter().any(|event| matches!(
7689 event,
7690 AgentEvent::ToolExecutionStarted(call) if call.name == "failing"
7691 )));
7692 assert!(events.iter().any(|event| matches!(
7693 event,
7694 AgentEvent::ToolResultReceived(result) if result.is_error
7695 )));
7696 }
7697
7698 #[tokio::test]
7699 async fn run_then_deny_tool_execution_still_reports_started() {
7700 let events = StdArc::new(StdMutex::new(Vec::new()));
7701 let tools = ToolRegistry::new().with(RunThenDenyTool::default());
7702 let agent = Agent::builder()
7703 .model(FakeAdapter)
7704 .add_tool_source(tools)
7705 .permissions(AllowAllPermissions)
7706 .observer(RecordingObserver {
7707 events: events.clone(),
7708 })
7709 .build()
7710 .unwrap();
7711
7712 let mut driver = agent
7713 .start(SessionConfig {
7714 session_id: SessionId::new("session-run-then-deny-start-event"),
7715 metadata: MetadataMap::new(),
7716 cache: None,
7717 consumer_capabilities: SessionConsumerCapabilities::default(),
7718 })
7719 .await
7720 .unwrap();
7721
7722 driver
7723 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7724 .unwrap();
7725
7726 match run_until_finished(&mut driver).await {
7727 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
7728 other => panic!("unexpected loop step: {other:?}"),
7729 }
7730
7731 let events = events.lock().unwrap();
7732 assert!(events.iter().any(|event| matches!(
7733 event,
7734 AgentEvent::ToolExecutionStarted(call) if call.name == "run_then_deny"
7735 )));
7736 assert!(events.iter().any(|event| matches!(
7739 event,
7740 AgentEvent::ToolResultReceived(result)
7741 if result.is_error
7742 && result
7743 .metadata
7744 .get(TOOL_RESULT_FAILURE_KIND_METADATA_KEY)
7745 .and_then(Value::as_str)
7746 == Some(TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED)
7747 && !result
7748 .metadata
7749 .contains_key(TOOL_RESULT_NOT_STARTED_METADATA_KEY)
7750 )));
7751 }
7752
7753 #[tokio::test]
7754 async fn async_task_manager_background_round_requires_explicit_continue() {
7755 let events = StdArc::new(StdMutex::new(Vec::new()));
7756 let entered = StdArc::new(AtomicBool::new(false));
7757 let release = StdArc::new(Notify::new());
7758 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
7759 "background-wait",
7760 RoutingDecision::Background,
7761 )]));
7762 let handle = task_manager.handle();
7763 let tools = ToolRegistry::new().with(BlockingTool::new(
7764 "background-wait",
7765 entered.clone(),
7766 release.clone(),
7767 "background-done",
7768 ));
7769 let agent = Agent::builder()
7770 .model(FakeAdapter)
7771 .add_tool_source(tools)
7772 .permissions(AllowAllPermissions)
7773 .task_manager(task_manager)
7774 .observer(RecordingObserver {
7775 events: events.clone(),
7776 })
7777 .build()
7778 .unwrap();
7779
7780 let mut driver = agent
7781 .start(SessionConfig {
7782 session_id: SessionId::new("session-background"),
7783 metadata: MetadataMap::new(),
7784 cache: None,
7785 consumer_capabilities: SessionConsumerCapabilities::default(),
7786 })
7787 .await
7788 .unwrap();
7789
7790 driver
7791 .submit_input(vec![Item {
7792 id: None,
7793 kind: ItemKind::User,
7794 parts: vec![Part::Text(TextPart {
7795 text: "ping".into(),
7796 metadata: MetadataMap::new(),
7797 })],
7798 metadata: MetadataMap::new(),
7799 usage: None,
7800 finish_reason: None,
7801 created_at: None,
7802 }])
7803 .unwrap();
7804
7805 let first = driver.next().await.unwrap();
7806 match first {
7807 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
7808 other => panic!("unexpected first loop step: {other:?}"),
7809 }
7810
7811 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
7812 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
7813 assert_eq!(lifecycle.len(), 2);
7814 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7815 assert_eq!(lifecycle[1].1, Some(FinishReason::ToolCall));
7816
7817 match wait_for_task_event(&handle).await {
7818 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "background-wait"),
7819 other => panic!("unexpected task event: {other:?}"),
7820 }
7821 wait_until_entered(entered.as_ref()).await;
7822 release.notify_waiters();
7823
7824 match wait_for_task_event(&handle).await {
7825 TaskEvent::Completed(_, result) => {
7826 assert_eq!(result.output, ToolOutput::Text("background-done".into()))
7827 }
7828 other => panic!("unexpected completion event: {other:?}"),
7829 }
7830
7831 let resumed = driver.next().await.unwrap();
7832 match resumed {
7833 LoopStep::Finished(turn) => {
7834 assert_eq!(turn.finish_reason, FinishReason::Completed);
7835 match &turn.items[0].parts[0] {
7836 Part::Text(text) => assert_eq!(
7837 text.text,
7838 "tool said: Background tool results: 1 total, 0 failed, 0 with metadata. \
7839 call-1 completed: text preview: background-done"
7840 ),
7841 other => panic!("unexpected part after resume: {other:?}"),
7842 }
7843 }
7844 other => panic!("unexpected resumed step: {other:?}"),
7845 }
7846
7847 let events = events.lock().unwrap();
7848 let lifecycle = turn_lifecycle_events(&events);
7849 assert_eq!(lifecycle.len(), 4);
7850 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7851 assert_eq!(lifecycle[2].0, lifecycle[3].0);
7852 assert_ne!(lifecycle[0].0, lifecycle[2].0);
7853 assert_eq!(lifecycle[3].1, Some(FinishReason::Completed));
7854
7855 let terminal_results: Vec<_> = events
7856 .iter()
7857 .filter_map(|event| match event {
7858 AgentEvent::ToolResultReceived(result)
7859 if result.call_id == ToolCallId::new("call-1") =>
7860 {
7861 Some(result)
7862 }
7863 _ => None,
7864 })
7865 .collect();
7866 assert_eq!(
7867 terminal_results.len(),
7868 1,
7869 "background completion must emit one terminal result event per call: {events:?}"
7870 );
7871 }
7872
7873 #[tokio::test]
7874 async fn detached_parts_notification_preserves_full_output_and_metadata() {
7875 let agent = Agent::builder().model(FakeAdapter).build().unwrap();
7876 let mut driver = agent
7877 .start(SessionConfig::new("session-detached-parts"))
7878 .await
7879 .unwrap();
7880 let call_id = ToolCallId::new("parts-call");
7881 driver.detached_call_ids.insert(call_id.clone());
7882 let parts = vec![
7883 Part::text("part text"),
7884 Part::structured(json!({
7885 "nested": [1, 2, 3]
7886 })),
7887 ];
7888 let mut metadata = MetadataMap::new();
7889 metadata.insert("source".into(), json!("background"));
7890 let result = ToolResultPart {
7891 call_id,
7892 output: ToolOutput::Parts(parts.clone()),
7893 is_error: true,
7894 metadata: metadata.clone(),
7895 };
7896 let mut item_metadata = MetadataMap::new();
7897 item_metadata.insert("delivery".into(), json!("deferred"));
7898 let item = Item::new(ItemKind::Tool, vec![Part::ToolResult(result.clone())])
7899 .with_metadata(item_metadata.clone());
7900
7901 let converted = driver.maybe_convert_detached(item);
7902 let (text, structured) = match converted.parts.as_slice() {
7903 [Part::Text(text), Part::Structured(structured)] => (text, structured),
7904 other => panic!("unexpected converted parts: {other:?}"),
7905 };
7906 assert_eq!(converted.kind, ItemKind::Notification);
7907 assert_eq!(converted.metadata, item_metadata);
7908 assert_eq!(structured.value, serde_json::to_value(&result).unwrap());
7909 assert_eq!(
7910 text.text,
7911 "Background tool results: 1 total, 1 failed, 1 with metadata. \
7912 parts-call failed: parts payload (2 parts)"
7913 );
7914 assert!(!text.text.contains("part text"));
7915 assert!(!text.text.contains("background"));
7916 }
7917
7918 #[tokio::test]
7919 async fn detached_files_notification_preserves_full_output() {
7920 let agent = Agent::builder().model(FakeAdapter).build().unwrap();
7921 let mut driver = agent
7922 .start(SessionConfig::new("session-detached-files"))
7923 .await
7924 .unwrap();
7925 let call_id = ToolCallId::new("files-call");
7926 driver.detached_call_ids.insert(call_id.clone());
7927 let files = vec![
7928 agentkit_core::FilePart::named("report.txt", DataRef::inline_text("full file body"))
7929 .with_mime_type("text/plain"),
7930 agentkit_core::FilePart::named(
7931 "remote.json",
7932 DataRef::uri("https://example.test/remote.json"),
7933 ),
7934 ];
7935 let mut result_metadata = MetadataMap::new();
7936 result_metadata.insert("archive".into(), json!(true));
7937 let result = ToolResultPart::success(call_id, ToolOutput::Files(files.clone()))
7938 .with_metadata(result_metadata);
7939 let mut item_metadata = MetadataMap::new();
7940 item_metadata.insert("delivery".into(), json!("deferred"));
7941 let item = Item::new(ItemKind::Tool, vec![Part::ToolResult(result.clone())])
7942 .with_metadata(item_metadata.clone());
7943
7944 let converted = driver.maybe_convert_detached(item);
7945 let (text, structured) = match converted.parts.as_slice() {
7946 [Part::Text(text), Part::Structured(structured)] => (text, structured),
7947 other => panic!("unexpected converted files: {other:?}"),
7948 };
7949 assert_eq!(converted.kind, ItemKind::Notification);
7950 assert_eq!(converted.metadata, item_metadata);
7951 assert_eq!(structured.value, serde_json::to_value(&result).unwrap());
7952 assert_eq!(
7953 text.text,
7954 "Background tool results: 1 total, 0 failed, 1 with metadata. \
7955 files-call completed: files payload (2 files)"
7956 );
7957 assert!(!text.text.contains("full file body"));
7958 assert!(!text.text.contains("remote.json"));
7959 }
7960
7961 #[test]
7962 fn detached_result_summaries_are_bounded_and_do_not_serialize_structured_payloads() {
7963 let long_text = "é".repeat(DETACHED_TEXT_PREVIEW_MAX_CHARS + 20);
7964 let text_summary = render_tool_output_brief(&ToolOutput::Text(long_text.clone()));
7965 assert_eq!(
7966 text_summary.chars().count(),
7967 "text preview: ".chars().count() + DETACHED_TEXT_PREVIEW_MAX_CHARS
7968 );
7969 assert!(text_summary.ends_with('…'));
7970 assert!(!text_summary.contains(&long_text));
7971
7972 let secret = "structured payload must remain out of notification text";
7973 let structured = ToolOutput::Structured(json!({ "secret": secret }));
7974 assert_eq!(render_tool_output_brief(&structured), "structured payload");
7975
7976 let oversized = "x".repeat(DETACHED_NOTIFICATION_TEXT_MAX_CHARS + 20);
7977 let bounded = truncate_chars(&oversized, DETACHED_NOTIFICATION_TEXT_MAX_CHARS);
7978 assert_eq!(
7979 bounded.chars().count(),
7980 DETACHED_NOTIFICATION_TEXT_MAX_CHARS
7981 );
7982 assert!(bounded.ends_with('…'));
7983 }
7984
7985 #[tokio::test]
7986 async fn detached_tool_placeholder_is_progress_not_terminal_result() {
7987 let events = StdArc::new(StdMutex::new(Vec::new()));
7988 let entered = StdArc::new(AtomicBool::new(false));
7989 let release = StdArc::new(Notify::new());
7990 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
7991 "detaching-wait",
7992 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10)),
7993 )]));
7994 let handle = task_manager.handle();
7995 let tools = ToolRegistry::new().with(BlockingTool::new(
7996 "detaching-wait",
7997 entered.clone(),
7998 release.clone(),
7999 "detached-done",
8000 ));
8001 let agent = Agent::builder()
8002 .model(FakeAdapter)
8003 .add_tool_source(tools)
8004 .permissions(AllowAllPermissions)
8005 .task_manager(task_manager)
8006 .observer(RecordingObserver {
8007 events: events.clone(),
8008 })
8009 .build()
8010 .unwrap();
8011
8012 let mut driver = agent
8013 .start(SessionConfig {
8014 session_id: SessionId::new("session-detached-progress"),
8015 metadata: MetadataMap::new(),
8016 cache: None,
8017 consumer_capabilities: SessionConsumerCapabilities::default(),
8018 })
8019 .await
8020 .unwrap();
8021
8022 driver
8023 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8024 .unwrap();
8025
8026 match driver.next().await.unwrap() {
8027 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => {}
8028 other => panic!("unexpected detach step: {other:?}"),
8029 }
8030
8031 match wait_for_task_event(&handle).await {
8032 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "detaching-wait"),
8033 other => panic!("unexpected task event: {other:?}"),
8034 }
8035 match wait_for_task_event(&handle).await {
8036 TaskEvent::Detached(snapshot) => assert_eq!(snapshot.tool_name, "detaching-wait"),
8037 other => panic!("unexpected detach event: {other:?}"),
8038 }
8039 wait_until_entered(entered.as_ref()).await;
8040 release.notify_waiters();
8041
8042 match wait_for_task_event(&handle).await {
8043 TaskEvent::Completed(_, result) => {
8044 assert_eq!(result.output, ToolOutput::Text("detached-done".into()))
8045 }
8046 other => panic!("unexpected completion event: {other:?}"),
8047 }
8048
8049 match driver.next().await.unwrap() {
8050 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
8051 other => panic!("unexpected resumed step: {other:?}"),
8052 }
8053
8054 let events = events.lock().unwrap();
8055 assert!(events.iter().any(|event| matches!(
8056 event,
8057 AgentEvent::ToolExecutionProgress(result)
8058 if result.call_id == ToolCallId::new("call-1") && !result.is_error
8059 )));
8060 let terminal_results: Vec<_> = events
8061 .iter()
8062 .filter_map(|event| match event {
8063 AgentEvent::ToolResultReceived(result)
8064 if result.call_id == ToolCallId::new("call-1") =>
8065 {
8066 Some(result)
8067 }
8068 _ => None,
8069 })
8070 .collect();
8071 assert_eq!(
8072 terminal_results.len(),
8073 1,
8074 "detached call must emit one terminal result event: {events:?}"
8075 );
8076 }
8077
8078 #[tokio::test]
8079 async fn cancelled_background_approval_auto_resolves_when_drained() {
8080 let controller = CancellationController::new();
8081 let events = StdArc::new(StdMutex::new(Vec::new()));
8082 let entered = StdArc::new(AtomicBool::new(false));
8083 let release = StdArc::new(Notify::new());
8084 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
8085 "echo",
8086 RoutingDecision::Background,
8087 )]));
8088 let handle = task_manager.handle();
8089 let agent = Agent::builder()
8090 .model(FakeAdapter)
8091 .tool_executor(DelayedApprovalExecutor::new(
8092 entered.clone(),
8093 release.clone(),
8094 ))
8095 .task_manager(task_manager)
8096 .cancellation(controller.handle())
8097 .observer(RecordingObserver {
8098 events: events.clone(),
8099 })
8100 .build()
8101 .unwrap();
8102
8103 let mut driver = agent
8104 .start(SessionConfig {
8105 session_id: SessionId::new("session-cancel-delayed-background-approval"),
8106 metadata: MetadataMap::new(),
8107 cache: None,
8108 consumer_capabilities: SessionConsumerCapabilities::default(),
8109 })
8110 .await
8111 .unwrap();
8112
8113 driver
8114 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8115 .unwrap();
8116
8117 match driver.next().await.unwrap() {
8118 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
8119 other => panic!("unexpected first step: {other:?}"),
8120 }
8121
8122 match wait_for_task_event(&handle).await {
8123 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "echo"),
8124 other => panic!("unexpected task event: {other:?}"),
8125 }
8126
8127 wait_until_entered(entered.as_ref()).await;
8128 controller.interrupt();
8129 release.notify_waiters();
8130 wait_until_completed(&handle).await;
8131
8132 match driver.next().await.unwrap() {
8133 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
8134 other => panic!("cancelled background approval should finish cancelled, got {other:?}"),
8135 }
8136
8137 let events = events.lock().unwrap();
8138 assert!(
8139 events
8140 .iter()
8141 .any(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false }))
8142 );
8143 assert!(events.iter().any(|event| matches!(
8144 event,
8145 AgentEvent::ToolResultReceived(result)
8146 if result.call_id == ToolCallId::new("call-1") && result.is_error
8147 )));
8148 }
8149
8150 #[tokio::test]
8151 async fn approved_foreground_task_waits_for_result_before_model_continuation() {
8152 let entered = StdArc::new(AtomicBool::new(false));
8153 let release = StdArc::new(Notify::new());
8154 let approved_entered = StdArc::new(AtomicBool::new(false));
8155 let approved_release = StdArc::new(Notify::new());
8156 let route_count = StdArc::new(AtomicUsize::new(0));
8157 let routing_count = route_count.clone();
8158 let task_manager = AsyncTaskManager::new().routing(move |_request: &ToolRequest| {
8159 if routing_count.fetch_add(1, Ordering::SeqCst) == 0 {
8160 RoutingDecision::Background
8161 } else {
8162 RoutingDecision::Foreground
8163 }
8164 });
8165 let handle = task_manager.handle();
8166 let agent = Agent::builder()
8167 .model(FakeAdapter)
8168 .tool_executor(
8169 DelayedApprovalExecutor::new(entered.clone(), release.clone())
8170 .blocking_after_approval(approved_entered.clone(), approved_release.clone()),
8171 )
8172 .task_manager(task_manager)
8173 .build()
8174 .unwrap();
8175 let mut driver = agent
8176 .start(SessionConfig::new("session-approved-foreground"))
8177 .await
8178 .unwrap();
8179 driver
8180 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8181 .unwrap();
8182
8183 assert!(matches!(
8184 driver.next().await.unwrap(),
8185 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
8186 ));
8187 let task_turn = match wait_for_task_event(&handle).await {
8188 TaskEvent::Started(snapshot) => snapshot.turn_id,
8189 other => panic!("unexpected task event: {other:?}"),
8190 };
8191 wait_until_entered(entered.as_ref()).await;
8192 release.notify_one();
8193 wait_until_completed(&handle).await;
8194
8195 let pending = match driver.next().await.unwrap() {
8196 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
8197 other => panic!("unexpected delayed approval step: {other:?}"),
8198 };
8199 let presentation_turn = driver.lifecycle.active_turn.clone().unwrap();
8200 assert_ne!(presentation_turn, task_turn);
8201 pending.approve(&mut driver).unwrap();
8202
8203 let info = {
8204 let next = driver.next();
8205 tokio::pin!(next);
8206 tokio::select! {
8207 () = wait_until_entered(approved_entered.as_ref()) => {}
8208 result = &mut next => {
8209 panic!("model continued before approved foreground result: {result:?}")
8210 }
8211 }
8212 assert!(
8213 timeout(Duration::from_millis(10), &mut next).await.is_err(),
8214 "model continued while approved foreground work was blocked"
8215 );
8216 approved_release.notify_one();
8217 let step = timeout(Duration::from_secs(1), &mut next)
8218 .await
8219 .expect("approved foreground result was not delivered")
8220 .unwrap();
8221 match step {
8222 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info,
8223 other => panic!("unexpected approved foreground step: {other:?}"),
8224 }
8225 };
8226 assert_eq!(info.turn_id, presentation_turn);
8227
8228 let turn = match driver.next().await.unwrap() {
8229 LoopStep::Finished(turn) => turn,
8230 other => panic!("model did not continue after approved result: {other:?}"),
8231 };
8232 assert_eq!(turn.finish_reason, FinishReason::Completed);
8233 assert_eq!(turn.turn_id, presentation_turn);
8234 assert!(driver.snapshot().transcript.iter().any(|item| {
8235 item.kind == ItemKind::Notification
8236 && item.parts.iter().any(
8237 |part| matches!(part, Part::Text(text) if text.text.contains("approved-ok")),
8238 )
8239 }));
8240 }
8241
8242 #[tokio::test]
8243 async fn approved_foreground_then_detach_waits_and_keeps_one_placeholder() {
8244 let entered = StdArc::new(AtomicBool::new(false));
8245 let release = StdArc::new(Notify::new());
8246 let approved_entered = StdArc::new(AtomicBool::new(false));
8247 let approved_release = StdArc::new(Notify::new());
8248 let route_count = StdArc::new(AtomicUsize::new(0));
8249 let routing_count = route_count.clone();
8250 let task_manager = AsyncTaskManager::new().routing(move |_request: &ToolRequest| {
8251 if routing_count.fetch_add(1, Ordering::SeqCst) == 0 {
8252 RoutingDecision::Background
8253 } else {
8254 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10))
8255 }
8256 });
8257 let handle = task_manager.handle();
8258 let agent = Agent::builder()
8259 .model(FakeAdapter)
8260 .tool_executor(
8261 DelayedApprovalExecutor::new(entered.clone(), release.clone())
8262 .blocking_after_approval(approved_entered.clone(), approved_release.clone()),
8263 )
8264 .task_manager(task_manager)
8265 .build()
8266 .unwrap();
8267 let mut driver = agent
8268 .start(SessionConfig::new("session-approved-foreground-detach"))
8269 .await
8270 .unwrap();
8271 driver
8272 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8273 .unwrap();
8274
8275 assert!(matches!(
8276 driver.next().await.unwrap(),
8277 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
8278 ));
8279 let _ = wait_for_task_event(&handle).await;
8280 wait_until_entered(entered.as_ref()).await;
8281 release.notify_one();
8282 wait_until_completed(&handle).await;
8283
8284 let pending = match driver.next().await.unwrap() {
8285 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
8286 other => panic!("unexpected delayed approval step: {other:?}"),
8287 };
8288 let presentation_turn = driver.lifecycle.active_turn.clone().unwrap();
8289 pending.approve(&mut driver).unwrap();
8290
8291 let step = timeout(Duration::from_secs(1), driver.next())
8292 .await
8293 .expect("approved task did not detach")
8294 .unwrap();
8295 assert!(approved_entered.load(Ordering::SeqCst));
8296 match step {
8297 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => {
8298 assert_eq!(info.turn_id, presentation_turn);
8299 }
8300 other => panic!("unexpected approved detach step: {other:?}"),
8301 }
8302 let placeholders = driver
8303 .snapshot()
8304 .transcript
8305 .iter()
8306 .filter(|item| item.kind == ItemKind::Tool)
8307 .flat_map(|item| &item.parts)
8308 .filter(|part| {
8309 matches!(
8310 part,
8311 Part::ToolResult(result) if result.call_id == ToolCallId::new("call-1")
8312 )
8313 })
8314 .count();
8315 assert_eq!(placeholders, 1, "detach appended a second tool result");
8316
8317 approved_release.notify_one();
8318 wait_until_completed(&handle).await;
8319 let turn = match driver.next().await.unwrap() {
8320 LoopStep::Finished(turn) => turn,
8321 other => panic!("model did not continue after detached result: {other:?}"),
8322 };
8323 assert_eq!(turn.finish_reason, FinishReason::Completed);
8324 assert_eq!(turn.turn_id, presentation_turn);
8325 let transcript = driver.snapshot().transcript;
8326 assert_eq!(
8327 transcript
8328 .iter()
8329 .filter(|item| item.kind == ItemKind::Tool)
8330 .flat_map(|item| &item.parts)
8331 .filter(|part| {
8332 matches!(
8333 part,
8334 Part::ToolResult(result)
8335 if result.call_id == ToolCallId::new("call-1")
8336 )
8337 })
8338 .count(),
8339 1
8340 );
8341 assert!(transcript.iter().any(|item| {
8342 item.kind == ItemKind::Notification
8343 && item.parts.iter().any(
8344 |part| matches!(part, Part::Text(text) if text.text.contains("approved-ok")),
8345 )
8346 }));
8347 }
8348
8349 #[tokio::test]
8350 async fn approving_detached_background_call_keeps_one_placeholder() {
8351 let entered = StdArc::new(AtomicBool::new(false));
8352 let release = StdArc::new(Notify::new());
8353 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
8354 "echo",
8355 RoutingDecision::Background,
8356 )]));
8357 let handle = task_manager.handle();
8358 let agent = Agent::builder()
8359 .model(FakeAdapter)
8360 .tool_executor(DelayedApprovalExecutor::new(
8361 entered.clone(),
8362 release.clone(),
8363 ))
8364 .task_manager(task_manager)
8365 .build()
8366 .unwrap();
8367 let mut driver = agent
8368 .start(SessionConfig::new("session-approve-detached-background"))
8369 .await
8370 .unwrap();
8371 driver
8372 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8373 .unwrap();
8374
8375 assert!(matches!(
8376 driver.next().await.unwrap(),
8377 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
8378 ));
8379 let _ = wait_for_task_event(&handle).await;
8380 wait_until_entered(entered.as_ref()).await;
8381 release.notify_waiters();
8382 wait_until_completed(&handle).await;
8383
8384 let pending = match driver.next().await.unwrap() {
8385 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
8386 other => panic!("unexpected delayed approval step: {other:?}"),
8387 };
8388 pending.approve(&mut driver).unwrap();
8389 release.notify_one();
8390 let _ = driver.next().await.unwrap();
8391
8392 let placeholders = driver
8393 .snapshot()
8394 .transcript
8395 .iter()
8396 .flat_map(|item| &item.parts)
8397 .filter(|part| {
8398 matches!(
8399 part,
8400 Part::ToolResult(result) if result.call_id == ToolCallId::new("call-1")
8401 )
8402 })
8403 .count();
8404 assert_eq!(placeholders, 1, "approval appended a second detach result");
8405 }
8406
8407 #[tokio::test]
8408 async fn failed_background_approval_cleanup_clears_queued_resume() {
8409 let events = StdArc::new(StdMutex::new(Vec::new()));
8410 let entered = StdArc::new(AtomicBool::new(false));
8411 let release = StdArc::new(Notify::new());
8412 let inner = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
8413 "echo",
8414 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10)),
8415 )]));
8416 let handle = inner.handle();
8417 let agent = Agent::builder()
8418 .model(FakeAdapter)
8419 .tool_executor(DelayedApprovalExecutor::new(
8420 entered.clone(),
8421 release.clone(),
8422 ))
8423 .task_manager(TestTaskManager::new(inner).fail_interrupt("cleanup failure"))
8424 .observer(RecordingObserver {
8425 events: events.clone(),
8426 })
8427 .build()
8428 .unwrap();
8429 let mut driver = agent
8430 .start(SessionConfig::new(
8431 "session-failed-detached-background-approval-cleanup",
8432 ))
8433 .await
8434 .unwrap();
8435 driver
8436 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8437 .unwrap();
8438
8439 let old_turn_id = match driver.next().await.unwrap() {
8440 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info.turn_id,
8441 other => panic!("unexpected detach step: {other:?}"),
8442 };
8443 assert_eq!(driver.pending_round_resume.as_ref(), Some(&old_turn_id));
8444 driver
8445 .submit_input(vec![Item::text(ItemKind::User, "fresh input")])
8446 .unwrap();
8447 let _ = wait_for_task_event(&handle).await;
8448 wait_until_entered(entered.as_ref()).await;
8449 release.notify_waiters();
8450 wait_until_completed(&handle).await;
8451
8452 assert!(matches!(
8453 driver.next().await.unwrap(),
8454 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))
8455 ));
8456 let error = driver.cancel_pending_approvals().await.unwrap_err();
8457 assert!(error.to_string().contains("cleanup failure"));
8458
8459 assert!(driver.lifecycle.active_turn.is_none());
8460 assert!(driver.pending_round_resume.is_none());
8461 assert_eq!(driver.snapshot().pending_input.len(), 1);
8462 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
8463 let [started, finished] = &lifecycle[lifecycle.len() - 2..] else {
8464 panic!("missing terminal lifecycle events: {lifecycle:?}");
8465 };
8466 assert_eq!(started.0, finished.0);
8467 assert_eq!(finished.1, Some(FinishReason::Error));
8468 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
8469
8470 let fresh_turn = match driver.next().await.unwrap() {
8471 LoopStep::Finished(turn) => turn,
8472 other => panic!("fresh input did not start a new turn: {other:?}"),
8473 };
8474 assert_ne!(fresh_turn.turn_id, old_turn_id);
8475 }
8476
8477 #[tokio::test]
8478 async fn fresh_input_runs_before_delayed_background_approval() {
8479 let entered = StdArc::new(AtomicBool::new(false));
8480 let release = StdArc::new(Notify::new());
8481 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
8482 "echo",
8483 RoutingDecision::Background,
8484 )]));
8485 let handle = task_manager.handle();
8486 let agent = Agent::builder()
8487 .model(FakeAdapter)
8488 .tool_executor(DelayedApprovalExecutor::new(
8489 entered.clone(),
8490 release.clone(),
8491 ))
8492 .task_manager(task_manager)
8493 .build()
8494 .unwrap();
8495 let mut driver = agent
8496 .start(SessionConfig::new(
8497 "session-input-before-background-approval",
8498 ))
8499 .await
8500 .unwrap();
8501 driver
8502 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8503 .unwrap();
8504
8505 assert!(matches!(
8506 driver.next().await.unwrap(),
8507 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
8508 ));
8509 let _ = wait_for_task_event(&handle).await;
8510 wait_until_entered(entered.as_ref()).await;
8511 release.notify_waiters();
8512 wait_until_completed(&handle).await;
8513
8514 driver
8515 .submit_input(vec![Item::text(ItemKind::User, "fresh input")])
8516 .unwrap();
8517 let fresh_turn = match driver.next().await.unwrap() {
8518 LoopStep::Finished(turn) => turn.turn_id,
8519 other => panic!("fresh input was not driven first: {other:?}"),
8520 };
8521 assert!(driver.pending_approvals.is_empty());
8522 assert!(driver.snapshot().pending_input.is_empty());
8523 timeout(Duration::from_millis(100), driver.wait_for_loop_update())
8524 .await
8525 .expect("collected background update did not wake the loop")
8526 .unwrap();
8527
8528 let approval = match driver.next().await.unwrap() {
8529 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(approval)) => approval,
8530 other => panic!("delayed approval was not presented separately: {other:?}"),
8531 };
8532 let approval_turn = driver.lifecycle.active_turn.clone().unwrap();
8533 assert_ne!(fresh_turn, approval_turn);
8534 assert_eq!(
8535 driver
8536 .snapshot()
8537 .transcript
8538 .iter()
8539 .filter(|item| {
8540 item.kind == ItemKind::User
8541 && item.parts.iter().any(
8542 |part| matches!(part, Part::Text(text) if text.text == "fresh input"),
8543 )
8544 })
8545 .count(),
8546 1,
8547 "fresh input must not be replayed while presenting the approval"
8548 );
8549 approval.deny(&mut driver).unwrap();
8550 }
8551
8552 #[tokio::test]
8553 async fn delayed_background_approval_interrupts_originating_task_turn() {
8554 let entered = StdArc::new(AtomicBool::new(false));
8555 let release = StdArc::new(Notify::new());
8556 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
8557 let inner = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
8558 "echo",
8559 RoutingDecision::Background,
8560 )]));
8561 let handle = inner.handle();
8562 let agent = Agent::builder()
8563 .model(FakeAdapter)
8564 .tool_executor(DelayedApprovalExecutor::new(
8565 entered.clone(),
8566 release.clone(),
8567 ))
8568 .task_manager(TestTaskManager::new(inner).record_interrupts(interrupted.clone()))
8569 .build()
8570 .unwrap();
8571 let mut driver = agent
8572 .start(SessionConfig::new(
8573 "session-background-approval-origin-turn",
8574 ))
8575 .await
8576 .unwrap();
8577 driver
8578 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8579 .unwrap();
8580
8581 assert!(matches!(
8582 driver.next().await.unwrap(),
8583 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
8584 ));
8585 let _ = wait_for_task_event(&handle).await;
8586 wait_until_entered(entered.as_ref()).await;
8587 release.notify_waiters();
8588 wait_until_completed(&handle).await;
8589 assert!(matches!(
8590 driver.next().await.unwrap(),
8591 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))
8592 ));
8593
8594 let presentation_turn = driver.lifecycle.active_turn.clone().unwrap();
8595 let task_turn = driver
8596 .pending_approvals
8597 .values()
8598 .next()
8599 .unwrap()
8600 .tool_request
8601 .turn_id
8602 .clone();
8603 assert_ne!(presentation_turn, task_turn);
8604 assert!(matches!(
8605 driver.cancel_pending_approvals().await.unwrap(),
8606 Some(LoopStep::Finished(TurnResult {
8607 finish_reason: FinishReason::Cancelled,
8608 ..
8609 }))
8610 ));
8611 assert_eq!(interrupted.lock().unwrap().as_slice(), &[task_turn]);
8612 assert!(driver.lifecycle.active_turn.is_none());
8613 }
8614
8615 #[tokio::test]
8616 async fn approved_background_start_error_interrupts_originating_turn() {
8617 let events = StdArc::new(StdMutex::new(Vec::new()));
8618 let entered = StdArc::new(AtomicBool::new(false));
8619 let release = StdArc::new(Notify::new());
8620 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
8621 let inner = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
8622 "echo",
8623 RoutingDecision::Background,
8624 )]));
8625 let handle = inner.handle();
8626 let agent = Agent::builder()
8627 .model(FakeAdapter)
8628 .tool_executor(DelayedApprovalExecutor::new(
8629 entered.clone(),
8630 release.clone(),
8631 ))
8632 .task_manager(
8633 TestTaskManager::new(inner)
8634 .fail_approved_start("original approved start failure")
8635 .record_interrupts(interrupted.clone())
8636 .fail_interrupt("cleanup failure"),
8637 )
8638 .observer(RecordingObserver {
8639 events: events.clone(),
8640 })
8641 .build()
8642 .unwrap();
8643 let mut driver = agent
8644 .start(SessionConfig::new("session-approved-start-error"))
8645 .await
8646 .unwrap();
8647 driver
8648 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8649 .unwrap();
8650
8651 assert!(matches!(
8652 driver.next().await.unwrap(),
8653 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
8654 ));
8655 let _ = wait_for_task_event(&handle).await;
8656 wait_until_entered(entered.as_ref()).await;
8657 release.notify_waiters();
8658 wait_until_completed(&handle).await;
8659 let pending = match driver.next().await.unwrap() {
8660 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
8661 other => panic!("unexpected delayed approval step: {other:?}"),
8662 };
8663 let task_turn = driver
8664 .pending_approvals
8665 .values()
8666 .next()
8667 .unwrap()
8668 .tool_request
8669 .turn_id
8670 .clone();
8671 let call_id = pending.request.call_id.clone().expect("approval call id");
8672 assert!(driver.detached_call_ids.contains(&call_id));
8673 pending.approve(&mut driver).unwrap();
8674
8675 let error = driver.next().await.unwrap_err();
8676 assert!(
8677 error
8678 .to_string()
8679 .contains("original approved start failure")
8680 );
8681 assert!(!error.to_string().contains("cleanup failure"));
8682 assert_eq!(interrupted.lock().unwrap().as_slice(), &[task_turn]);
8683 assert!(driver.lifecycle.active_turn.is_none());
8684 assert!(!driver.detached_call_ids.contains(&call_id));
8685 assert!(!driver.background_call_ids.contains(&call_id));
8686 assert!(!driver.tool_cancellations.contains_key(&call_id));
8687 assert!(events.lock().unwrap().iter().any(|event| matches!(
8688 event,
8689 AgentEvent::ToolResultReceived(result)
8690 if result.call_id == call_id && result.is_error
8691 )));
8692 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
8693 }
8694
8695 #[tokio::test]
8696 async fn loop_can_cancel_a_turn_and_continue_after_new_input() {
8697 let controller = CancellationController::new();
8698 let agent = Agent::builder()
8699 .model(SlowAdapter)
8700 .cancellation(controller.handle())
8701 .build()
8702 .unwrap();
8703
8704 let mut driver = agent
8705 .start(SessionConfig {
8706 session_id: SessionId::new("session-cancel"),
8707 metadata: MetadataMap::new(),
8708 cache: None,
8709 consumer_capabilities: SessionConsumerCapabilities::default(),
8710 })
8711 .await
8712 .unwrap();
8713
8714 driver
8715 .submit_input(vec![Item {
8716 id: None,
8717 kind: ItemKind::User,
8718 parts: vec![Part::Text(TextPart {
8719 text: "do the long task".into(),
8720 metadata: MetadataMap::new(),
8721 })],
8722 metadata: MetadataMap::new(),
8723 usage: None,
8724 finish_reason: None,
8725 created_at: None,
8726 }])
8727 .unwrap();
8728
8729 let cancelled = tokio::join!(async { driver.next().await }, async {
8730 tokio::task::yield_now().await;
8731 controller.interrupt();
8732 })
8733 .0
8734 .unwrap();
8735
8736 match cancelled {
8737 LoopStep::Finished(turn) => {
8738 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
8739 assert_eq!(turn.items.len(), 1);
8740 assert_eq!(turn.items[0].kind, ItemKind::Assistant);
8741 assert_eq!(
8742 turn.items[0].metadata.get(INTERRUPTED_METADATA_KEY),
8743 Some(&Value::Bool(true))
8744 );
8745 }
8746 other => panic!("unexpected loop step: {other:?}"),
8747 }
8748
8749 driver
8750 .submit_input(vec![Item {
8751 id: None,
8752 kind: ItemKind::User,
8753 parts: vec![Part::Text(TextPart {
8754 text: "try again".into(),
8755 metadata: MetadataMap::new(),
8756 })],
8757 metadata: MetadataMap::new(),
8758 usage: None,
8759 finish_reason: None,
8760 created_at: None,
8761 }])
8762 .unwrap();
8763
8764 let result = driver.next().await.unwrap();
8765 match result {
8766 LoopStep::Finished(turn) => {
8767 assert_eq!(turn.finish_reason, FinishReason::Completed);
8768 }
8769 other => panic!("unexpected loop step after retry: {other:?}"),
8770 }
8771 }
8772
8773 #[tokio::test]
8774 async fn loop_interrupt_cancels_foreground_tasks_but_keeps_background_tasks_running() {
8775 let controller = CancellationController::new();
8776 let fg_entered = StdArc::new(AtomicBool::new(false));
8777 let fg_release = StdArc::new(Notify::new());
8778 let bg_entered = StdArc::new(AtomicBool::new(false));
8779 let bg_release = StdArc::new(Notify::new());
8780 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([
8781 ("foreground-wait", RoutingDecision::Foreground),
8782 ("background-wait", RoutingDecision::Background),
8783 ]));
8784 let handle = task_manager.handle();
8785 let tools = ToolRegistry::new()
8786 .with(BlockingTool::new(
8787 "foreground-wait",
8788 fg_entered.clone(),
8789 fg_release,
8790 "foreground-done",
8791 ))
8792 .with(BlockingTool::new(
8793 "background-wait",
8794 bg_entered.clone(),
8795 bg_release.clone(),
8796 "background-done",
8797 ));
8798 let agent = Agent::builder()
8799 .model(MultiToolAdapter)
8800 .add_tool_source(tools)
8801 .permissions(AllowAllPermissions)
8802 .cancellation(controller.handle())
8803 .task_manager(task_manager)
8804 .build()
8805 .unwrap();
8806
8807 let mut driver = agent
8808 .start(SessionConfig {
8809 session_id: SessionId::new("session-mixed-cancel"),
8810 metadata: MetadataMap::new(),
8811 cache: None,
8812 consumer_capabilities: SessionConsumerCapabilities::default(),
8813 })
8814 .await
8815 .unwrap();
8816
8817 driver
8818 .submit_input(vec![Item {
8819 id: None,
8820 kind: ItemKind::User,
8821 parts: vec![Part::Text(TextPart {
8822 text: "run both".into(),
8823 metadata: MetadataMap::new(),
8824 })],
8825 metadata: MetadataMap::new(),
8826 usage: None,
8827 finish_reason: None,
8828 created_at: None,
8829 }])
8830 .unwrap();
8831
8832 let cancelled = tokio::join!(async { driver.next().await }, async {
8833 let _ = wait_for_task_event(&handle).await;
8834 let _ = wait_for_task_event(&handle).await;
8835 wait_until_entered(fg_entered.as_ref()).await;
8836 wait_until_entered(bg_entered.as_ref()).await;
8837 controller.interrupt();
8838 })
8839 .0
8840 .unwrap();
8841
8842 match cancelled {
8843 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
8844 other => panic!("unexpected loop step after interrupt: {other:?}"),
8845 }
8846
8847 match wait_for_task_event(&handle).await {
8848 TaskEvent::Cancelled(snapshot) => assert_eq!(snapshot.tool_name, "foreground-wait"),
8849 other => panic!("unexpected post-interrupt event: {other:?}"),
8850 }
8851
8852 let running = handle.list_running().await;
8853 assert_eq!(running.len(), 1);
8854 assert_eq!(running[0].tool_name, "background-wait");
8855
8856 bg_release.notify_waiters();
8857 match wait_for_task_event(&handle).await {
8858 TaskEvent::Completed(snapshot, result) => {
8859 assert_eq!(snapshot.tool_name, "background-wait");
8860 assert_eq!(result.output, ToolOutput::Text("background-done".into()));
8861 }
8862 other => panic!("unexpected background completion event: {other:?}"),
8863 }
8864 }
8865
8866 #[tokio::test]
8867 async fn a_cancelled_turn_answers_the_tool_call_it_abandoned() {
8868 let controller = CancellationController::new();
8873 let entered = StdArc::new(AtomicBool::new(false));
8874 let release = StdArc::new(Notify::new());
8875 let items = StdArc::new(StdMutex::new(Vec::<Item>::new()));
8876 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
8877 "wait",
8878 RoutingDecision::Foreground,
8879 )]));
8880 let agent = Agent::builder()
8881 .model(FakeAdapter)
8882 .add_tool_source(ToolRegistry::new().with(BlockingTool::new(
8883 "wait",
8884 entered.clone(),
8885 release,
8886 "done",
8887 )))
8888 .permissions(AllowAllPermissions)
8889 .cancellation(controller.handle())
8890 .task_manager(task_manager)
8891 .transcript_observer(RecordingTranscriptObserver {
8892 items: items.clone(),
8893 })
8894 .build()
8895 .unwrap();
8896
8897 let mut driver = agent
8898 .start(SessionConfig {
8899 session_id: SessionId::new("session-cancel-mid-call"),
8900 metadata: MetadataMap::new(),
8901 cache: None,
8902 consumer_capabilities: SessionConsumerCapabilities::default(),
8903 })
8904 .await
8905 .unwrap();
8906
8907 driver
8908 .submit_input(vec![Item {
8909 id: None,
8910 kind: ItemKind::User,
8911 parts: vec![Part::Text(TextPart {
8912 text: "run the tool".into(),
8913 metadata: MetadataMap::new(),
8914 })],
8915 metadata: MetadataMap::new(),
8916 usage: None,
8917 finish_reason: None,
8918 created_at: None,
8919 }])
8920 .unwrap();
8921
8922 let cancelled = tokio::join!(async { driver.next().await }, async {
8923 wait_until_entered(entered.as_ref()).await;
8924 controller.interrupt();
8925 })
8926 .0
8927 .unwrap();
8928
8929 match cancelled {
8930 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
8931 other => panic!("unexpected loop step after interrupt: {other:?}"),
8932 }
8933
8934 let transcript = driver.snapshot().transcript;
8935 assert!(
8936 unanswered_tool_calls(&transcript).is_empty(),
8937 "the cancelled turn left a tool call unanswered: {transcript:?}"
8938 );
8939 validate_transcript_invariants(&transcript)
8940 .expect("a cancelled turn must leave a resumable transcript");
8941
8942 let persisted = items.lock().unwrap().clone();
8943 let results: Vec<&ToolResultPart> = persisted
8944 .iter()
8945 .flat_map(|item| &item.parts)
8946 .filter_map(|part| match part {
8947 Part::ToolResult(result) => Some(result),
8948 _ => None,
8949 })
8950 .collect();
8951 assert_eq!(results.len(), 1, "{persisted:?}");
8952 assert_eq!(results[0].call_id, ToolCallId::new("call-1"));
8953 assert!(results[0].is_error);
8954 assert_eq!(
8955 results[0].metadata.get(INTERRUPTED_METADATA_KEY),
8956 Some(&Value::Bool(true))
8957 );
8958 }
8959
8960 #[tokio::test]
8961 async fn regression_cancelled_background_completion_emits_one_terminal_result() {
8962 let events = StdArc::new(StdMutex::new(Vec::new()));
8963 let agent = Agent::builder()
8964 .model(FakeAdapter)
8965 .observer(RecordingObserver {
8966 events: events.clone(),
8967 })
8968 .build()
8969 .unwrap();
8970 let mut driver = agent
8971 .start(SessionConfig::new("session-cancelled-background-event"))
8972 .await
8973 .unwrap();
8974
8975 driver.append_item(Item::new(
8976 ItemKind::Assistant,
8977 vec![Part::ToolCall(ToolCallPart {
8978 id: ToolCallId::new("call-1"),
8979 name: "wait".into(),
8980 input: json!({}),
8981 metadata: MetadataMap::new(),
8982 })],
8983 ));
8984 driver.background_call_ids.insert(ToolCallId::new("call-1"));
8985 driver.close_interrupted_tool_calls();
8986 driver.append_tool_result_item(Item::new(
8987 ItemKind::Tool,
8988 vec![Part::ToolResult(ToolResultPart {
8989 call_id: ToolCallId::new("call-1"),
8990 output: ToolOutput::Text("background-done".into()),
8991 is_error: false,
8992 metadata: MetadataMap::new(),
8993 })],
8994 ));
8995
8996 let events = events.lock().unwrap();
8997 let terminal_results = events
8998 .iter()
8999 .filter(|event| {
9000 matches!(
9001 event,
9002 AgentEvent::ToolResultReceived(result)
9003 if result.call_id == ToolCallId::new("call-1")
9004 )
9005 })
9006 .count();
9007 assert_eq!(
9008 terminal_results, 1,
9009 "a cancelled background call emitted multiple terminal results: {events:?}"
9010 );
9011 }
9012
9013 #[tokio::test]
9014 async fn regression_cancelled_queued_approval_is_answered_once() {
9015 let controller = CancellationController::new();
9016 let entered = StdArc::new(AtomicBool::new(false));
9017 let release = StdArc::new(Notify::new());
9018 release.notify_one();
9019 let events = StdArc::new(StdMutex::new(Vec::new()));
9020 let agent = Agent::builder()
9021 .model(FakeAdapter)
9022 .tool_executor(
9023 DelayedApprovalExecutor::new(entered, release)
9024 .cancelling_on_approval(controller.clone()),
9025 )
9026 .cancellation(controller.handle())
9027 .observer(RecordingObserver {
9028 events: events.clone(),
9029 })
9030 .build()
9031 .unwrap();
9032 let mut driver = agent
9033 .start(SessionConfig::new("session-cancelled-queued-approval"))
9034 .await
9035 .unwrap();
9036 driver
9037 .submit_input(vec![Item::text(ItemKind::User, "ping")])
9038 .unwrap();
9039
9040 match driver.next().await.unwrap() {
9041 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
9042 other => panic!("unexpected first cancellation step: {other:?}"),
9043 }
9044 match driver.next().await.unwrap() {
9045 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
9046 other => panic!("unexpected post-cancellation step: {other:?}"),
9047 }
9048
9049 let events = events.lock().unwrap();
9050 let terminal_results = events
9051 .iter()
9052 .filter(|event| {
9053 matches!(
9054 event,
9055 AgentEvent::ToolResultReceived(result)
9056 if result.call_id == ToolCallId::new("call-1")
9057 )
9058 })
9059 .count();
9060 assert_eq!(
9061 terminal_results, 1,
9062 "a cancelled queued approval was answered more than once: {events:?}"
9063 );
9064 drop(events);
9065
9066 let transcript = driver.snapshot().transcript;
9067 assert!(
9068 !transcript.iter().any(|item| {
9069 item.kind == ItemKind::Notification
9070 && item.parts.iter().any(|part| {
9071 matches!(part, Part::Text(text) if text.text.contains("Background tool call"))
9072 })
9073 }),
9074 "a queued approval was misreported as a background call: {transcript:?}"
9075 );
9076 }
9077
9078 #[tokio::test]
9079 async fn regression_cancelled_unstarted_call_is_not_tracked_as_detached() {
9080 let agent = Agent::builder().model(FakeAdapter).build().unwrap();
9081 let mut driver = agent
9082 .start(SessionConfig::new("session-cancelled-unstarted-call"))
9083 .await
9084 .unwrap();
9085
9086 driver.append_item(Item::new(
9087 ItemKind::Assistant,
9088 vec![Part::ToolCall(ToolCallPart {
9089 id: ToolCallId::new("call-never-started"),
9090 name: "wait".into(),
9091 input: json!({}),
9092 metadata: MetadataMap::new(),
9093 })],
9094 ));
9095 driver
9096 .finish_cancelled(agentkit_core::TurnId::new("turn-cancelled"), Vec::new())
9097 .unwrap();
9098
9099 assert!(
9100 !driver
9101 .detached_call_ids
9102 .contains(&ToolCallId::new("call-never-started")),
9103 "an unstarted call can never deliver a detached result"
9104 );
9105 }
9106
9107 #[tokio::test]
9108 async fn loop_resumes_after_approved_tool_request() {
9109 let tools = ToolRegistry::new().with(EchoTool::default());
9110 let agent = Agent::builder()
9111 .model(FakeAdapter)
9112 .add_tool_source(tools)
9113 .permissions(ApproveFsReads)
9114 .build()
9115 .unwrap();
9116
9117 let mut driver = agent
9118 .start(SessionConfig {
9119 session_id: SessionId::new("session-approval"),
9120 metadata: MetadataMap::new(),
9121 cache: None,
9122 consumer_capabilities: SessionConsumerCapabilities::default(),
9123 })
9124 .await
9125 .unwrap();
9126
9127 driver
9128 .submit_input(vec![Item {
9129 id: None,
9130 kind: ItemKind::User,
9131 parts: vec![Part::Text(TextPart {
9132 text: "ping".into(),
9133 metadata: MetadataMap::new(),
9134 })],
9135 metadata: MetadataMap::new(),
9136 usage: None,
9137 finish_reason: None,
9138 created_at: None,
9139 }])
9140 .unwrap();
9141
9142 let first = driver.next().await.unwrap();
9143 match first {
9144 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
9145 assert!(pending.request.task_id.is_some());
9146 assert_eq!(pending.request.id.0, "approval:fs-read");
9147 pending.approve(&mut driver).unwrap();
9148 }
9149 other => panic!("unexpected loop step: {other:?}"),
9150 }
9151 let second = driver.next().await.unwrap();
9152 match second {
9153 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
9154 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
9155 other => panic!("unexpected part: {other:?}"),
9156 },
9157 other => panic!("unexpected loop step after approval: {other:?}"),
9158 }
9159 }
9160
9161 #[tokio::test]
9162 async fn approval_gated_tool_does_not_start_before_approval() {
9163 let events = StdArc::new(StdMutex::new(Vec::new()));
9164 let tools = ToolRegistry::new().with(EchoTool::default());
9165 let agent = Agent::builder()
9166 .model(FakeAdapter)
9167 .add_tool_source(tools)
9168 .permissions(ApproveFsReads)
9169 .observer(RecordingObserver {
9170 events: events.clone(),
9171 })
9172 .build()
9173 .unwrap();
9174
9175 let mut driver = agent
9176 .start(SessionConfig {
9177 session_id: SessionId::new("session-approval-start-event"),
9178 metadata: MetadataMap::new(),
9179 cache: None,
9180 consumer_capabilities: SessionConsumerCapabilities::default(),
9181 })
9182 .await
9183 .unwrap();
9184
9185 driver
9186 .submit_input(vec![Item::text(ItemKind::User, "ping")])
9187 .unwrap();
9188
9189 let pending = match driver.next().await.unwrap() {
9190 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
9191 other => panic!("unexpected loop step: {other:?}"),
9192 };
9193
9194 assert!(
9195 events
9196 .lock()
9197 .unwrap()
9198 .iter()
9199 .all(|event| !matches!(event, AgentEvent::ToolExecutionStarted(_))),
9200 "tool start must not be reported before approval"
9201 );
9202
9203 pending.approve(&mut driver).unwrap();
9204 match driver.next().await.unwrap() {
9205 LoopStep::Finished(_) => {}
9206 other => panic!("unexpected loop step after approval: {other:?}"),
9207 }
9208
9209 let started = events
9210 .lock()
9211 .unwrap()
9212 .iter()
9213 .filter(|event| matches!(event, AgentEvent::ToolExecutionStarted(_)))
9214 .count();
9215 assert_eq!(started, 1);
9216 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
9217 assert_eq!(lifecycle.len(), 2);
9218 assert_eq!(lifecycle[0].0, lifecycle[1].0);
9219 assert_eq!(lifecycle[1].1, Some(FinishReason::Completed));
9220 }
9221
9222 #[tokio::test]
9223 async fn cancelling_pending_approval_resolves_it_and_pairs_tool_result() {
9224 let controller = CancellationController::new();
9225 let events = StdArc::new(StdMutex::new(Vec::new()));
9226 let tools = ToolRegistry::new().with(EchoTool::default());
9227 let agent = Agent::builder()
9228 .model(FakeAdapter)
9229 .add_tool_source(tools)
9230 .permissions(ApproveFsReads)
9231 .cancellation(controller.handle())
9232 .observer(RecordingObserver {
9233 events: events.clone(),
9234 })
9235 .build()
9236 .unwrap();
9237
9238 let mut driver = agent
9239 .start(SessionConfig {
9240 session_id: SessionId::new("session-cancel-pending-approval"),
9241 metadata: MetadataMap::new(),
9242 cache: None,
9243 consumer_capabilities: SessionConsumerCapabilities::default(),
9244 })
9245 .await
9246 .unwrap();
9247
9248 driver
9249 .submit_input(vec![Item::text(ItemKind::User, "ping")])
9250 .unwrap();
9251
9252 match driver.next().await.unwrap() {
9253 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) => {}
9254 other => panic!("unexpected loop step: {other:?}"),
9255 }
9256
9257 controller.interrupt();
9258
9259 match driver.next().await.unwrap() {
9260 LoopStep::Finished(turn) => {
9261 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
9262 }
9263 other => panic!("unexpected loop step after cancel: {other:?}"),
9264 }
9265
9266 let events = events.lock().unwrap();
9267 assert!(
9268 events
9269 .iter()
9270 .any(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false })),
9271 "pending approval cancellation should close approval UI state"
9272 );
9273 assert!(
9274 events.iter().any(|event| matches!(
9275 event,
9276 AgentEvent::ToolResultReceived(result)
9277 if result.call_id == ToolCallId::new("call-1") && result.is_error
9278 )),
9279 "pending approval cancellation should pair the assistant tool_use"
9280 );
9281 drop(events);
9282
9283 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
9284 }
9285
9286 #[tokio::test]
9287 async fn cancelling_sole_foreground_approval_for_call_finishes_turn() {
9288 let events = StdArc::new(StdMutex::new(Vec::new()));
9289 let agent = Agent::builder()
9290 .model(FakeAdapter)
9291 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
9292 .permissions(ApproveFsReads)
9293 .observer(RecordingObserver {
9294 events: events.clone(),
9295 })
9296 .build()
9297 .unwrap();
9298 let mut driver = agent
9299 .start(SessionConfig::new("session-cancel-foreground-approval-for"))
9300 .await
9301 .unwrap();
9302 driver
9303 .submit_input(vec![Item::text(ItemKind::User, "ping")])
9304 .unwrap();
9305
9306 let call_id = match driver.next().await.unwrap() {
9307 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
9308 pending.request.call_id.expect("approval call id")
9309 }
9310 other => panic!("unexpected loop step: {other:?}"),
9311 };
9312 driver.cancel_pending_approval_for(call_id).unwrap();
9313
9314 assert!(driver.lifecycle.active_turn.is_none());
9315 assert!(driver.pending_approvals.is_empty());
9316 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
9317 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
9318 assert_eq!(lifecycle.len(), 2, "{lifecycle:?}");
9319 assert_eq!(lifecycle[0].0, lifecycle[1].0);
9320 assert_eq!(lifecycle[1].1, Some(FinishReason::Cancelled));
9321 }
9322
9323 #[tokio::test]
9324 async fn resolved_approval_runs_even_if_cancellation_also_fired() {
9325 let controller = CancellationController::new();
9326 let tools = ToolRegistry::new().with(EchoTool::default());
9327 let agent = Agent::builder()
9328 .model(FakeAdapter)
9329 .add_tool_source(tools)
9330 .permissions(ApproveFsReads)
9331 .cancellation(controller.handle())
9332 .build()
9333 .unwrap();
9334
9335 let mut driver = agent
9336 .start(SessionConfig {
9337 session_id: SessionId::new("session-resolved-approval-cancel-race"),
9338 metadata: MetadataMap::new(),
9339 cache: None,
9340 consumer_capabilities: SessionConsumerCapabilities::default(),
9341 })
9342 .await
9343 .unwrap();
9344
9345 driver
9346 .submit_input(vec![Item::text(ItemKind::User, "ping")])
9347 .unwrap();
9348
9349 let pending = match driver.next().await.unwrap() {
9350 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
9351 other => panic!("unexpected loop step: {other:?}"),
9352 };
9353
9354 controller.interrupt();
9355 pending.approve(&mut driver).unwrap();
9356
9357 match driver.next().await.unwrap() {
9358 LoopStep::Finished(turn) => {
9359 assert_eq!(turn.finish_reason, FinishReason::Completed);
9360 match &turn.items[0].parts[0] {
9361 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
9362 other => panic!("unexpected part after approval: {other:?}"),
9363 }
9364 }
9365 other => panic!("unexpected loop step after approved cancel race: {other:?}"),
9366 }
9367 }
9368
9369 #[tokio::test]
9370 async fn loop_resumes_with_patched_input_on_approval() {
9371 let tools = ToolRegistry::new().with(EchoTool::default());
9372 let agent = Agent::builder()
9373 .model(FakeAdapter)
9374 .add_tool_source(tools)
9375 .permissions(ApproveFsReads)
9376 .build()
9377 .unwrap();
9378
9379 let mut driver = agent
9380 .start(SessionConfig {
9381 session_id: SessionId::new("session-approval-patched"),
9382 metadata: MetadataMap::new(),
9383 cache: None,
9384 consumer_capabilities: SessionConsumerCapabilities::default(),
9385 })
9386 .await
9387 .unwrap();
9388
9389 driver
9390 .submit_input(vec![Item {
9391 id: None,
9392 kind: ItemKind::User,
9393 parts: vec![Part::Text(TextPart {
9394 text: "ping".into(),
9395 metadata: MetadataMap::new(),
9396 })],
9397 metadata: MetadataMap::new(),
9398 usage: None,
9399 finish_reason: None,
9400 created_at: None,
9401 }])
9402 .unwrap();
9403
9404 match driver.next().await.unwrap() {
9405 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
9406 pending
9407 .approve_with_patched_input(&mut driver, json!({ "value": "patched" }))
9408 .unwrap();
9409 }
9410 other => panic!("unexpected loop step: {other:?}"),
9411 }
9412 match driver.next().await.unwrap() {
9413 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
9414 Part::Text(text) => assert_eq!(text.text, "tool said: patched"),
9415 other => panic!("unexpected part: {other:?}"),
9416 },
9417 other => panic!("unexpected loop step after approval: {other:?}"),
9418 }
9419 }
9420
9421 #[tokio::test]
9422 async fn loop_tracks_multiple_pending_approvals_by_call_id() {
9423 let tools = ToolRegistry::new().with(EchoTool::default());
9424 let agent = Agent::builder()
9425 .model(DualApprovalAdapter)
9426 .add_tool_source(tools)
9427 .permissions(ApproveFsReads)
9428 .build()
9429 .unwrap();
9430
9431 let mut driver = agent
9432 .start(SessionConfig {
9433 session_id: SessionId::new("session-dual-approval"),
9434 metadata: MetadataMap::new(),
9435 cache: None,
9436 consumer_capabilities: SessionConsumerCapabilities::default(),
9437 })
9438 .await
9439 .unwrap();
9440
9441 driver
9442 .submit_input(vec![Item {
9443 id: None,
9444 kind: ItemKind::User,
9445 parts: vec![Part::Text(TextPart {
9446 text: "run both approvals".into(),
9447 metadata: MetadataMap::new(),
9448 })],
9449 metadata: MetadataMap::new(),
9450 usage: None,
9451 finish_reason: None,
9452 created_at: None,
9453 }])
9454 .unwrap();
9455
9456 let pending_first = match driver.next().await.unwrap() {
9457 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
9458 assert_eq!(
9459 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
9460 Some("call-1")
9461 );
9462 pending
9463 }
9464 other => panic!("unexpected first loop step: {other:?}"),
9465 };
9466
9467 let pending_second = match driver.next().await.unwrap() {
9468 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
9469 assert_eq!(
9470 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
9471 Some("call-2")
9472 );
9473 pending
9474 }
9475 other => panic!("unexpected second loop step: {other:?}"),
9476 };
9477
9478 pending_second.approve(&mut driver).unwrap();
9479 match driver.next().await.unwrap() {
9480 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
9481 assert_eq!(
9482 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
9483 Some("call-1")
9484 );
9485 }
9486 other => panic!("unexpected step after approving second request: {other:?}"),
9487 }
9488
9489 pending_first.approve(&mut driver).unwrap();
9490 match driver.next().await.unwrap() {
9491 LoopStep::Finished(turn) => {
9492 assert_eq!(turn.finish_reason, FinishReason::Completed);
9493 match &turn.items[0].parts[0] {
9494 Part::Text(text) => assert_eq!(text.text, "both approvals finished"),
9495 other => panic!("unexpected final part: {other:?}"),
9496 }
9497 }
9498 other => panic!("unexpected final loop step: {other:?}"),
9499 }
9500 }
9501
9502 #[tokio::test]
9503 async fn failed_pending_approval_cleanup_repairs_and_finishes_error() {
9504 let events = StdArc::new(StdMutex::new(Vec::new()));
9505 let agent = Agent::builder()
9506 .model(FakeAdapter)
9507 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
9508 .permissions(ApproveFsReads)
9509 .task_manager(
9510 TestTaskManager::new(SimpleTaskManager::new())
9511 .fail_interrupt("interrupt cleanup failed"),
9512 )
9513 .observer(RecordingObserver {
9514 events: events.clone(),
9515 })
9516 .build()
9517 .unwrap();
9518 let mut driver = agent
9519 .start(SessionConfig::new("session-failed-approval-cleanup"))
9520 .await
9521 .unwrap();
9522 driver
9523 .submit_input(vec![Item::text(ItemKind::User, "ping")])
9524 .unwrap();
9525
9526 assert!(matches!(
9527 driver.next().await.unwrap(),
9528 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))
9529 ));
9530 let error = driver.cancel_pending_approvals().await.unwrap_err();
9531 assert!(error.to_string().contains("interrupt cleanup failed"));
9532 assert!(driver.lifecycle.active_turn.is_none());
9533 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
9534
9535 let events = events.lock().unwrap();
9536 assert!(events.iter().any(|event| matches!(
9537 event,
9538 AgentEvent::ToolResultReceived(result)
9539 if result.call_id == ToolCallId::new("call-1") && result.is_error
9540 )));
9541 assert!(events.iter().any(|event| matches!(
9542 event,
9543 AgentEvent::TurnFinished(turn) if turn.finish_reason == FinishReason::Error
9544 )));
9545 }
9546
9547 #[tokio::test]
9548 async fn cancelling_all_pending_approvals_interrupts_every_originating_turn() {
9549 let events = StdArc::new(StdMutex::new(Vec::new()));
9550 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
9551 let tools = ToolRegistry::new().with(EchoTool::default());
9552 let agent = Agent::builder()
9553 .model(DualApprovalAdapter)
9554 .add_tool_source(tools)
9555 .permissions(ApproveFsReads)
9556 .task_manager(
9557 TestTaskManager::new(SimpleTaskManager::new())
9558 .record_interrupts(interrupted.clone()),
9559 )
9560 .observer(RecordingObserver {
9561 events: events.clone(),
9562 })
9563 .build()
9564 .unwrap();
9565
9566 let mut driver = agent
9567 .start(SessionConfig {
9568 session_id: SessionId::new("session-dual-approval-cancel"),
9569 metadata: MetadataMap::new(),
9570 cache: None,
9571 consumer_capabilities: SessionConsumerCapabilities::default(),
9572 })
9573 .await
9574 .unwrap();
9575
9576 driver
9577 .submit_input(vec![Item {
9578 id: None,
9579 kind: ItemKind::User,
9580 parts: vec![Part::Text(TextPart {
9581 text: "run both approvals".into(),
9582 metadata: MetadataMap::new(),
9583 })],
9584 metadata: MetadataMap::new(),
9585 usage: None,
9586 finish_reason: None,
9587 created_at: None,
9588 }])
9589 .unwrap();
9590
9591 for expected_call in ["call-1", "call-2"] {
9592 match driver.next().await.unwrap() {
9593 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
9594 assert_eq!(
9595 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
9596 Some(expected_call)
9597 );
9598 }
9599 other => panic!("unexpected approval step: {other:?}"),
9600 }
9601 }
9602
9603 let first_origin = driver
9604 .pending_approvals
9605 .get(&ToolCallId::new("call-1"))
9606 .unwrap()
9607 .tool_request
9608 .turn_id
9609 .clone();
9610 let second_origin = agentkit_core::TurnId::new("second-originating-turn");
9611 driver
9612 .pending_approvals
9613 .get_mut(&ToolCallId::new("call-2"))
9614 .unwrap()
9615 .tool_request
9616 .turn_id = second_origin.clone();
9617
9618 match driver.cancel_pending_approvals().await.unwrap() {
9619 Some(LoopStep::Finished(turn)) => {
9620 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
9621 }
9622 other => panic!("unexpected cancellation result: {other:?}"),
9623 }
9624 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
9625 let interrupted = interrupted
9626 .lock()
9627 .unwrap()
9628 .iter()
9629 .cloned()
9630 .collect::<HashSet<_>>();
9631 assert_eq!(interrupted, HashSet::from([first_origin, second_origin]));
9632
9633 let events = events.lock().unwrap();
9634 let cancelled = events
9635 .iter()
9636 .filter(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false }))
9637 .count();
9638 assert_eq!(cancelled, 2);
9639 assert!(events.iter().any(|event| matches!(
9640 event,
9641 AgentEvent::TurnFinished(turn) if turn.finish_reason == FinishReason::Cancelled
9642 )));
9643 for expected_call in ["call-1", "call-2"] {
9644 assert!(events.iter().any(|event| matches!(
9645 event,
9646 AgentEvent::ToolResultReceived(result)
9647 if result.call_id == ToolCallId::new(expected_call) && result.is_error
9648 )));
9649 }
9650 }
9651
9652 #[tokio::test]
9653 async fn loop_compacts_transcript_before_new_turns() {
9654 let events = StdArc::new(StdMutex::new(Vec::new()));
9655 let agent = Agent::builder()
9656 .model(FakeAdapter)
9657 .mutator(KeepRecentMutator { keep: 1 })
9658 .observer(RecordingObserver {
9659 events: events.clone(),
9660 })
9661 .build()
9662 .unwrap();
9663
9664 let mut driver = agent
9665 .start(SessionConfig {
9666 session_id: SessionId::new("session-4"),
9667 metadata: MetadataMap::new(),
9668 cache: None,
9669 consumer_capabilities: SessionConsumerCapabilities::default(),
9670 })
9671 .await
9672 .unwrap();
9673
9674 for text in ["first", "second"] {
9675 driver
9676 .submit_input(vec![Item {
9677 id: None,
9678 kind: ItemKind::User,
9679 parts: vec![Part::Text(TextPart {
9680 text: text.into(),
9681 metadata: MetadataMap::new(),
9682 })],
9683 metadata: MetadataMap::new(),
9684 usage: None,
9685 finish_reason: None,
9686 created_at: None,
9687 }])
9688 .unwrap();
9689 let _ = driver.next().await.unwrap();
9690 }
9691
9692 let events = events.lock().unwrap();
9693 assert!(
9694 events
9695 .iter()
9696 .any(|event| matches!(event, AgentEvent::MutationFinished { dirty: true, .. }))
9697 );
9698 }
9699
9700 #[test]
9701 fn transcript_validation_rejects_orphaned_tool_result() {
9702 let transcript = vec![Item {
9703 id: None,
9704 kind: ItemKind::Tool,
9705 parts: vec![Part::ToolResult(ToolResultPart {
9706 call_id: "call-1".into(),
9707 output: ToolOutput::Text("result".into()),
9708 is_error: false,
9709 metadata: MetadataMap::new(),
9710 })],
9711 metadata: MetadataMap::new(),
9712 usage: None,
9713 finish_reason: None,
9714 created_at: None,
9715 }];
9716
9717 let error = validate_transcript_invariants(&transcript).unwrap_err();
9718 assert!(error.to_string().contains("orphaned tool_result"));
9719 }
9720
9721 #[test]
9722 fn transcript_validation_rejects_duplicate_tool_result() {
9723 let transcript = vec![
9724 Item {
9725 id: None,
9726 kind: ItemKind::Assistant,
9727 parts: vec![Part::ToolCall(ToolCallPart {
9728 id: "call-1".into(),
9729 name: "lookup".into(),
9730 input: serde_json::json!({}),
9731 metadata: MetadataMap::new(),
9732 })],
9733 metadata: MetadataMap::new(),
9734 usage: None,
9735 finish_reason: None,
9736 created_at: None,
9737 },
9738 Item {
9739 id: None,
9740 kind: ItemKind::Tool,
9741 parts: vec![Part::ToolResult(ToolResultPart {
9742 call_id: "call-1".into(),
9743 output: ToolOutput::Text("result".into()),
9744 is_error: false,
9745 metadata: MetadataMap::new(),
9746 })],
9747 metadata: MetadataMap::new(),
9748 usage: None,
9749 finish_reason: None,
9750 created_at: None,
9751 },
9752 Item {
9753 id: None,
9754 kind: ItemKind::Tool,
9755 parts: vec![Part::ToolResult(ToolResultPart {
9756 call_id: "call-1".into(),
9757 output: ToolOutput::Text("again".into()),
9758 is_error: false,
9759 metadata: MetadataMap::new(),
9760 })],
9761 metadata: MetadataMap::new(),
9762 usage: None,
9763 finish_reason: None,
9764 created_at: None,
9765 },
9766 ];
9767
9768 let error = validate_transcript_invariants(&transcript).unwrap_err();
9769 assert!(error.to_string().contains("duplicate tool_result"));
9770 }
9771
9772 #[tokio::test]
9773 async fn loop_refreshes_tool_specs_each_turn() {
9774 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
9775 let version = StdArc::new(AtomicUsize::new(1));
9776 let tools = ToolRegistry::new().with(DynamicSpecTool::new(version.clone()));
9777 let agent = Agent::builder()
9778 .model(RecordingAdapter {
9779 seen_descriptions: seen_descriptions.clone(),
9780 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
9781 })
9782 .add_tool_source(tools)
9783 .permissions(AllowAllPermissions)
9784 .build()
9785 .unwrap();
9786
9787 let mut driver = agent
9788 .start(SessionConfig {
9789 session_id: SessionId::new("session-dynamic-tools"),
9790 metadata: MetadataMap::new(),
9791 cache: None,
9792 consumer_capabilities: SessionConsumerCapabilities::default(),
9793 })
9794 .await
9795 .unwrap();
9796
9797 for text in ["first", "second"] {
9798 driver
9799 .submit_input(vec![Item {
9800 id: None,
9801 kind: ItemKind::User,
9802 parts: vec![Part::Text(TextPart {
9803 text: text.into(),
9804 metadata: MetadataMap::new(),
9805 })],
9806 metadata: MetadataMap::new(),
9807 usage: None,
9808 finish_reason: None,
9809 created_at: None,
9810 }])
9811 .unwrap();
9812
9813 let _ = driver.next().await.unwrap();
9814 if text == "first" {
9815 version.store(2, Ordering::SeqCst);
9816 }
9817 }
9818
9819 let seen_descriptions = seen_descriptions.lock().unwrap();
9820 assert_eq!(seen_descriptions.len(), 2);
9821 assert_eq!(seen_descriptions[0], vec!["dynamic version 1".to_string()]);
9822 assert_eq!(seen_descriptions[1], vec!["dynamic version 2".to_string()]);
9823 }
9824
9825 #[tokio::test]
9826 async fn loop_emits_catalog_change_and_uses_updated_specs_next_turn() {
9827 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
9828 let events = StdArc::new(StdMutex::new(Vec::new()));
9829 let executor = StdArc::new(CatalogExecutor::new());
9830 let executor_for_agent: Arc<dyn ToolExecutor> = executor.clone();
9831 let agent = Agent::builder()
9832 .model(RecordingAdapter {
9833 seen_descriptions: seen_descriptions.clone(),
9834 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
9835 })
9836 .tool_executor(executor_for_agent)
9837 .permissions(AllowAllPermissions)
9838 .observer(RecordingObserver {
9839 events: events.clone(),
9840 })
9841 .build()
9842 .unwrap();
9843
9844 let mut driver = agent
9845 .start(SessionConfig {
9846 session_id: SessionId::new("session-catalog-events"),
9847 metadata: MetadataMap::new(),
9848 cache: None,
9849 consumer_capabilities: SessionConsumerCapabilities::default(),
9850 })
9851 .await
9852 .unwrap();
9853
9854 driver
9855 .submit_input(vec![Item::text(ItemKind::User, "first")])
9856 .unwrap();
9857 let _ = driver.next().await.unwrap();
9858
9859 executor.publish_change(
9860 1,
9861 ToolCatalogEvent {
9862 source: "mcp:mock".into(),
9863 added: vec!["dynamic".into()],
9864 removed: Vec::new(),
9865 changed: Vec::new(),
9866 },
9867 );
9868
9869 driver
9870 .submit_input(vec![Item::text(ItemKind::User, "second")])
9871 .unwrap();
9872 let _ = driver.next().await.unwrap();
9873
9874 let seen_descriptions = seen_descriptions.lock().unwrap();
9875 assert_eq!(seen_descriptions.len(), 2);
9876 assert_eq!(seen_descriptions[0], vec!["dynamic version 0".to_string()]);
9877 assert_eq!(seen_descriptions[1], vec!["dynamic version 1".to_string()]);
9878
9879 let events = events.lock().unwrap();
9880 assert!(events.iter().any(|event| matches!(
9881 event,
9882 AgentEvent::ToolCatalogChanged(ToolCatalogEvent {
9883 source,
9884 added,
9885 removed,
9886 changed,
9887 }) if source == "mcp:mock"
9888 && added == &vec!["dynamic".to_string()]
9889 && removed.is_empty()
9890 && changed.is_empty()
9891 )));
9892 }
9893
9894 #[tokio::test]
9895 async fn loop_passes_session_default_and_next_turn_cache_requests() {
9896 let seen_caches = StdArc::new(StdMutex::new(Vec::new()));
9897 let agent = Agent::builder()
9898 .model(RecordingAdapter {
9899 seen_descriptions: StdArc::new(StdMutex::new(Vec::new())),
9900 seen_caches: seen_caches.clone(),
9901 })
9902 .permissions(AllowAllPermissions)
9903 .build()
9904 .unwrap();
9905
9906 let default_cache = PromptCacheRequest::best_effort(PromptCacheStrategy::Automatic)
9907 .with_retention(PromptCacheRetention::Short);
9908 let override_cache = PromptCacheRequest::required(PromptCacheStrategy::Explicit {
9909 breakpoints: vec![PromptCacheBreakpoint::TranscriptItemEnd { index: 0 }],
9910 });
9911
9912 let mut driver = agent
9913 .start(SessionConfig {
9914 session_id: SessionId::new("session-cache"),
9915 metadata: MetadataMap::new(),
9916 cache: Some(default_cache.clone()),
9917 consumer_capabilities: SessionConsumerCapabilities::default(),
9918 })
9919 .await
9920 .unwrap();
9921
9922 driver
9923 .submit_input(vec![Item {
9924 id: None,
9925 kind: ItemKind::User,
9926 parts: vec![Part::Text(TextPart {
9927 text: "first".into(),
9928 metadata: MetadataMap::new(),
9929 })],
9930 metadata: MetadataMap::new(),
9931 usage: None,
9932 finish_reason: None,
9933 created_at: None,
9934 }])
9935 .unwrap();
9936 let _ = driver.next().await.unwrap();
9937
9938 driver
9939 .submit_input_with_cache(
9940 vec![Item {
9941 id: None,
9942 kind: ItemKind::User,
9943 parts: vec![Part::Text(TextPart {
9944 text: "second".into(),
9945 metadata: MetadataMap::new(),
9946 })],
9947 metadata: MetadataMap::new(),
9948 usage: None,
9949 finish_reason: None,
9950 created_at: None,
9951 }],
9952 override_cache.clone(),
9953 )
9954 .unwrap();
9955 let _ = driver.next().await.unwrap();
9956
9957 let seen = seen_caches.lock().unwrap();
9958 assert_eq!(seen.len(), 2);
9959 assert_eq!(seen[0], Some(default_cache));
9960 assert_eq!(seen[1], Some(override_cache));
9961 }
9962
9963 struct RetirementCountingAdapter(StdArc<AtomicUsize>);
9965 struct RetirementCountingSession(StdArc<AtomicUsize>);
9966
9967 #[async_trait]
9968 impl ModelAdapter for RetirementCountingAdapter {
9969 type Session = RetirementCountingSession;
9970
9971 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
9972 Ok(RetirementCountingSession(self.0.clone()))
9973 }
9974 }
9975
9976 #[async_trait]
9977 impl ModelSession for RetirementCountingSession {
9978 type Turn = FakeTurn;
9979
9980 async fn begin_turn(
9981 &mut self,
9982 request: TurnRequest,
9983 cancellation: Option<TurnCancellation>,
9984 ) -> Result<Self::Turn, LoopError> {
9985 self.0.fetch_add(1, Ordering::SeqCst);
9986 FakeSession.begin_turn(request, cancellation).await
9987 }
9988 }
9989
9990 #[tokio::test]
9991 async fn retire_interrupted_turn_after_tool_result_starts_fresh_without_continuing() {
9992 let calls = StdArc::new(AtomicUsize::new(0));
9993 let events = StdArc::new(StdMutex::new(Vec::new()));
9994 let agent = Agent::builder()
9995 .model(RetirementCountingAdapter(calls.clone()))
9996 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
9997 .permissions(AllowAllPermissions)
9998 .observer(RecordingObserver {
9999 events: events.clone(),
10000 })
10001 .build()
10002 .unwrap();
10003 let mut driver = agent
10004 .start(SessionConfig::new("retire-after-tool"))
10005 .await
10006 .unwrap();
10007 assert!(driver.retire_interrupted_turn().await.unwrap().is_none());
10008 driver
10009 .submit_input(vec![Item::text(ItemKind::User, "ping")])
10010 .unwrap();
10011 let turn_id = match driver.next().await.unwrap() {
10012 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info.turn_id,
10013 other => panic!("expected AfterToolResult, got {other:?}"),
10014 };
10015 let transcript = driver.snapshot().transcript;
10016 let queued = vec![Item::text(ItemKind::User, "fresh request")];
10017 driver.submit_input(queued.clone()).unwrap();
10018 let retired = driver.retire_interrupted_turn().await.unwrap().unwrap();
10019 assert_eq!(retired.turn_id, turn_id);
10020 assert_eq!(retired.finish_reason, FinishReason::Cancelled);
10021 assert!(retired.items.is_empty());
10022 assert!(driver.retire_interrupted_turn().await.unwrap().is_none());
10023 assert_eq!(calls.load(Ordering::SeqCst), 1);
10024 assert_eq!(driver.snapshot().transcript, transcript);
10025 assert_eq!(driver.snapshot().pending_input, queued);
10026 assert!(driver.pending_round_resume.is_none());
10027 assert!(driver.lifecycle.active_turn.is_none());
10028 let fresh = match driver.next().await.unwrap() {
10029 LoopStep::Finished(turn) => turn,
10030 other => panic!("expected fresh turn to finish, got {other:?}"),
10031 };
10032 assert_ne!(fresh.turn_id, turn_id);
10033 assert_eq!(fresh.finish_reason, FinishReason::Completed);
10034 assert_eq!(calls.load(Ordering::SeqCst), 2);
10035 assert!(driver.snapshot().pending_input.is_empty());
10036 assert!(
10037 driver
10038 .snapshot()
10039 .transcript
10040 .iter()
10041 .any(|item| { item.kind == ItemKind::User && item.parts == queued[0].parts })
10042 );
10043 assert!(matches!(
10044 driver.next().await.unwrap(),
10045 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
10046 ));
10047 assert!(driver.retire_interrupted_turn().await.unwrap().is_none());
10048 assert_eq!(calls.load(Ordering::SeqCst), 2);
10049 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
10050 assert_eq!(
10051 lifecycle,
10052 vec![
10053 (turn_id.clone(), None),
10054 (turn_id, Some(FinishReason::Cancelled)),
10055 (fresh.turn_id.clone(), None),
10056 (fresh.turn_id, Some(FinishReason::Completed)),
10057 ]
10058 );
10059 }
10060
10061 #[tokio::test]
10062 async fn retire_interrupted_turn_denies_pending_approval_even_on_cleanup_error() {
10063 for fail_cleanup in [false, true] {
10064 let calls = StdArc::new(AtomicUsize::new(0));
10065 let events = StdArc::new(StdMutex::new(Vec::new()));
10066 let manager = TestTaskManager::new(SimpleTaskManager::new());
10067 let manager = if fail_cleanup {
10068 manager.fail_interrupt("retirement cleanup failed")
10069 } else {
10070 manager
10071 };
10072 let agent = Agent::builder()
10073 .model(RetirementCountingAdapter(calls.clone()))
10074 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
10075 .permissions(ApproveFsReads)
10076 .task_manager(manager)
10077 .observer(RecordingObserver {
10078 events: events.clone(),
10079 })
10080 .build()
10081 .unwrap();
10082 let mut driver = agent
10083 .start(SessionConfig::new("retire-approval"))
10084 .await
10085 .unwrap();
10086 driver
10087 .submit_input(vec![Item::text(ItemKind::User, "ping")])
10088 .unwrap();
10089 assert!(matches!(
10090 driver.next().await.unwrap(),
10091 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))
10092 ));
10093 let turn_id = driver.lifecycle.active_turn.clone().unwrap();
10094 let result = driver.retire_interrupted_turn().await;
10095 if fail_cleanup {
10096 assert!(
10097 result
10098 .unwrap_err()
10099 .to_string()
10100 .contains("retirement cleanup failed")
10101 );
10102 } else {
10103 assert_eq!(
10104 result.unwrap().unwrap().finish_reason,
10105 FinishReason::Cancelled
10106 );
10107 }
10108 assert!(driver.pending_approvals.is_empty());
10109 assert!(driver.pending_approval_order.is_empty());
10110 assert!(driver.active_tool_round.is_none());
10111 assert!(driver.lifecycle.active_turn.is_none());
10112 assert!(driver.retire_interrupted_turn().await.unwrap().is_none());
10113 assert!(matches!(
10114 driver.next().await.unwrap(),
10115 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
10116 ));
10117 assert_eq!(calls.load(Ordering::SeqCst), 1);
10118 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
10119 let events = events.lock().unwrap();
10120 assert_eq!(
10121 turn_lifecycle_events(&events),
10122 vec![
10123 (turn_id.clone(), None),
10124 (turn_id, Some(FinishReason::Cancelled))
10125 ]
10126 );
10127 assert_eq!(events.iter().filter(|event| matches!(event,
10128 AgentEvent::ToolResultReceived(result) if result.call_id == ToolCallId::new("call-1") && result.is_error
10129 )).count(), 1);
10130 }
10131 }
10132
10133 #[tokio::test]
10134 async fn retire_interrupted_turn_preserves_detached_background_task() {
10135 let calls = StdArc::new(AtomicUsize::new(0));
10136 let entered = StdArc::new(AtomicBool::new(false));
10137 let release = StdArc::new(Notify::new());
10138 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
10139 "detaching-wait",
10140 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10)),
10141 )]));
10142 let handle = task_manager.handle();
10143 let agent = Agent::builder()
10144 .model(RetirementCountingAdapter(calls.clone()))
10145 .add_tool_source(ToolRegistry::new().with(BlockingTool::new(
10146 "detaching-wait",
10147 entered.clone(),
10148 release.clone(),
10149 "background-done",
10150 )))
10151 .permissions(AllowAllPermissions)
10152 .task_manager(task_manager)
10153 .build()
10154 .unwrap();
10155 let mut driver = agent
10156 .start(SessionConfig::new("retire-background"))
10157 .await
10158 .unwrap();
10159 driver
10160 .submit_input(vec![Item::text(ItemKind::User, "ping")])
10161 .unwrap();
10162 assert!(matches!(
10163 driver.next().await.unwrap(),
10164 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_))
10165 ));
10166 assert!(matches!(
10167 wait_for_task_event(&handle).await,
10168 TaskEvent::Started(_)
10169 ));
10170 assert!(matches!(
10171 wait_for_task_event(&handle).await,
10172 TaskEvent::Detached(_)
10173 ));
10174 wait_until_entered(entered.as_ref()).await;
10175 let transcript = driver.snapshot().transcript;
10176 assert_eq!(
10177 driver
10178 .retire_interrupted_turn()
10179 .await
10180 .unwrap()
10181 .unwrap()
10182 .finish_reason,
10183 FinishReason::Cancelled
10184 );
10185 assert_eq!(calls.load(Ordering::SeqCst), 1);
10186 assert_eq!(driver.snapshot().transcript, transcript);
10187 let running = handle.list_running().await;
10188 assert_eq!(running.len(), 1);
10189 assert_eq!(running[0].tool_name, "detaching-wait");
10190 release.notify_one();
10191 match wait_for_task_event(&handle).await {
10192 TaskEvent::Completed(_, result) => {
10193 assert_eq!(result.output, ToolOutput::Text("background-done".into()))
10194 }
10195 other => panic!("retired turn must preserve background completion, got {other:?}"),
10196 }
10197 }
10198
10199 #[tokio::test]
10200 async fn loop_yields_after_tool_result_between_rounds() {
10201 let tools = ToolRegistry::new().with(EchoTool::default());
10202 let agent = Agent::builder()
10203 .model(FakeAdapter)
10204 .add_tool_source(tools)
10205 .permissions(AllowAllPermissions)
10206 .build()
10207 .unwrap();
10208
10209 let mut driver = agent
10210 .start(SessionConfig {
10211 session_id: SessionId::new("yield-session"),
10212 metadata: MetadataMap::new(),
10213 cache: None,
10214 consumer_capabilities: SessionConsumerCapabilities::default(),
10215 })
10216 .await
10217 .unwrap();
10218
10219 driver
10220 .submit_input(vec![Item::text(ItemKind::User, "ping")])
10221 .unwrap();
10222
10223 let step = driver.next().await.unwrap();
10226 let info = match step {
10227 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info,
10228 other => panic!("expected AfterToolResult, got {other:?}"),
10229 };
10230 assert_eq!(info.session_id, SessionId::new("yield-session"));
10231 assert_eq!(info.transcript_len, 3);
10233
10234 let interrupt = LoopInterrupt::AfterToolResult(info.clone());
10236 assert!(!interrupt.is_blocking());
10237
10238 driver
10240 .submit_input(vec![Item::text(ItemKind::User, "also: report back")])
10241 .unwrap();
10242
10243 let step = driver.next().await.unwrap();
10246 match step {
10247 LoopStep::Finished(turn) => {
10248 assert_eq!(turn.finish_reason, FinishReason::Completed);
10249 }
10250 other => panic!("expected Finished, got {other:?}"),
10251 }
10252
10253 let snapshot = driver.snapshot();
10255 let has_injected_message = snapshot.transcript.iter().any(|item| {
10256 item.kind == ItemKind::User
10257 && item.parts.iter().any(|part| match part {
10258 Part::Text(text) => text.text == "also: report back",
10259 _ => false,
10260 })
10261 });
10262 assert!(
10263 has_injected_message,
10264 "injected user message should be in transcript, got: {:?}",
10265 snapshot.transcript
10266 );
10267 }
10268
10269 struct RecordingTranscriptObserver {
10270 items: StdArc<StdMutex<Vec<Item>>>,
10271 }
10272
10273 impl TranscriptObserver for RecordingTranscriptObserver {
10274 fn on_transcript_event(&self, event: TranscriptEvent<'_>) {
10275 self.items.lock().unwrap().push(event.item.clone());
10276 }
10277 }
10278
10279 #[tokio::test]
10280 async fn observers_see_full_tool_round() {
10281 let events = StdArc::new(StdMutex::new(Vec::<AgentEvent>::new()));
10287 let items = StdArc::new(StdMutex::new(Vec::<Item>::new()));
10288 let agent = Agent::builder()
10289 .model(FakeAdapter)
10290 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
10291 .permissions(AllowAllPermissions)
10292 .observer(RecordingObserver {
10293 events: events.clone(),
10294 })
10295 .transcript_observer(RecordingTranscriptObserver {
10296 items: items.clone(),
10297 })
10298 .build()
10299 .unwrap();
10300
10301 let mut driver = agent
10302 .start(SessionConfig {
10303 session_id: SessionId::new("observer-session"),
10304 metadata: MetadataMap::new(),
10305 cache: None,
10306 consumer_capabilities: SessionConsumerCapabilities::default(),
10307 })
10308 .await
10309 .unwrap();
10310
10311 driver
10312 .submit_input(vec![Item {
10313 id: None,
10314 kind: ItemKind::User,
10315 parts: vec![Part::Text(TextPart {
10316 text: "ping".into(),
10317 metadata: MetadataMap::new(),
10318 })],
10319 metadata: MetadataMap::new(),
10320 usage: None,
10321 finish_reason: None,
10322 created_at: None,
10323 }])
10324 .unwrap();
10325
10326 let result = run_until_finished(&mut driver).await;
10327 assert!(matches!(result, LoopStep::Finished(_)), "got {result:?}");
10328
10329 let events = events.lock().unwrap().clone();
10332 let tool_call_id = events.iter().find_map(|e| match e {
10333 AgentEvent::ToolCallRequested(c) => Some(c.id.clone()),
10334 _ => None,
10335 });
10336 let tool_results: Vec<_> = events
10337 .iter()
10338 .filter_map(|e| match e {
10339 AgentEvent::ToolResultReceived(r) => Some(r.clone()),
10340 _ => None,
10341 })
10342 .collect();
10343 assert_eq!(tool_results.len(), 1, "events: {events:?}");
10344 assert_eq!(Some(tool_results[0].call_id.clone()), tool_call_id);
10345 assert!(!tool_results[0].is_error);
10346
10347 let items = items.lock().unwrap().clone();
10351 assert_eq!(items.len(), 4, "items: {items:?}");
10352 assert_eq!(items[0].kind, ItemKind::User);
10353 assert_eq!(items[1].kind, ItemKind::Assistant);
10354 assert!(
10355 items[1]
10356 .parts
10357 .iter()
10358 .any(|p| matches!(p, Part::ToolCall(_)))
10359 );
10360 assert_eq!(items[2].kind, ItemKind::Tool);
10361 assert!(
10362 items[2]
10363 .parts
10364 .iter()
10365 .any(|p| matches!(p, Part::ToolResult(_)))
10366 );
10367 assert_eq!(items[3].kind, ItemKind::Assistant);
10368 }
10369
10370 #[test]
10371 fn convenience_cache_builders_construct_expected_defaults() {
10372 let cache = PromptCacheRequest::automatic()
10373 .with_retention(PromptCacheRetention::Short)
10374 .with_key("workspace:demo");
10375 let session = SessionConfig::new("demo").with_cache(cache.clone());
10376
10377 assert_eq!(session.session_id, SessionId::new("demo"));
10378 assert_eq!(session.cache, Some(cache));
10379
10380 let explicit = PromptCacheRequest::explicit([
10381 PromptCacheBreakpoint::tools_end(),
10382 PromptCacheBreakpoint::transcript_item_end(2),
10383 PromptCacheBreakpoint::transcript_part_end(3, 1),
10384 ]);
10385
10386 assert_eq!(explicit.mode, PromptCacheMode::BestEffort);
10387 assert_eq!(
10388 explicit.strategy,
10389 PromptCacheStrategy::Explicit {
10390 breakpoints: vec![
10391 PromptCacheBreakpoint::ToolsEnd,
10392 PromptCacheBreakpoint::TranscriptItemEnd { index: 2 },
10393 PromptCacheBreakpoint::TranscriptPartEnd {
10394 item_index: 3,
10395 part_index: 1,
10396 },
10397 ],
10398 }
10399 );
10400 }
10401}