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 SessionId, TaskId, TextPart, Timestamp, ToolCallId, ToolCallPart, ToolOutput, ToolResultPart,
70 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
93const INTERRUPTED_METADATA_KEY: &str = "agentkit.interrupted";
94const INTERRUPT_REASON_METADATA_KEY: &str = "agentkit.interrupt_reason";
95const INTERRUPT_STAGE_METADATA_KEY: &str = "agentkit.interrupt_stage";
96const USER_CANCELLED_REASON: &str = "user_cancelled";
97const DETACHED_NOTIFICATION_TEXT_MAX_CHARS: usize = 512;
98const DETACHED_TEXT_PREVIEW_MAX_CHARS: usize = 160;
99const DETACHED_CALL_ID_MAX_CHARS: usize = 80;
100
101pub const PROVIDER_FINISH_REASONS_METADATA_KEY: &str = "agentkit.provider_finish_reasons";
103
104pub fn set_provider_finish_reasons<I, S>(metadata: &mut MetadataMap, reasons: I)
106where
107 I: IntoIterator<Item = S>,
108 S: Into<String>,
109{
110 let mut seen = HashSet::new();
111 let reasons = reasons
112 .into_iter()
113 .map(Into::into)
114 .filter(|reason: &String| !reason.is_empty() && seen.insert(reason.clone()))
115 .map(Value::String)
116 .collect::<Vec<_>>();
117 metadata.remove(PROVIDER_FINISH_REASONS_METADATA_KEY);
118 if !reasons.is_empty() {
119 metadata.insert(
120 PROVIDER_FINISH_REASONS_METADATA_KEY.into(),
121 Value::Array(reasons),
122 );
123 }
124}
125
126fn provider_finish_reasons(metadata: &MetadataMap, fallback: &FinishReason) -> Vec<String> {
127 metadata
128 .get(PROVIDER_FINISH_REASONS_METADATA_KEY)
129 .and_then(Value::as_array)
130 .map(|values| {
131 let mut seen = HashSet::new();
132 values
133 .iter()
134 .filter_map(Value::as_str)
135 .filter(|reason| !reason.is_empty() && seen.insert((*reason).to_owned()))
136 .map(str::to_owned)
137 .collect::<Vec<_>>()
138 })
139 .filter(|reasons| !reasons.is_empty())
140 .unwrap_or_else(|| vec![normalized_finish_reason(fallback).into()])
141}
142
143fn normalized_finish_reason(reason: &FinishReason) -> &str {
144 match reason {
145 FinishReason::Completed => "completed",
146 FinishReason::ToolCall => "tool_call",
147 FinishReason::MaxTokens => "max_tokens",
148 FinishReason::Cancelled => "cancelled",
149 FinishReason::Blocked => "blocked",
150 FinishReason::Error => "error",
151 FinishReason::Other(reason) => reason,
152 }
153}
154
155#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)]
157pub enum MessageCaptureError {
158 #[error("message capture max_messages must be nonzero")]
160 ZeroMessages,
161 #[error("message capture max_bytes must be nonzero")]
163 ZeroBytes,
164}
165
166#[derive(Clone, Copy, Debug, PartialEq, Eq)]
168pub struct MessageCapture {
169 max_messages: usize,
170 max_bytes: usize,
171}
172
173impl MessageCapture {
174 pub fn new(max_messages: usize, max_bytes: usize) -> Result<Self, MessageCaptureError> {
176 if max_messages == 0 {
177 return Err(MessageCaptureError::ZeroMessages);
178 }
179 if max_bytes == 0 {
180 return Err(MessageCaptureError::ZeroBytes);
181 }
182 Ok(Self {
183 max_messages,
184 max_bytes,
185 })
186 }
187
188 pub fn max_messages(self) -> usize {
190 self.max_messages
191 }
192
193 pub fn max_bytes(self) -> usize {
195 self.max_bytes
196 }
197}
198
199#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
205pub struct TelemetryConfig {
206 input_messages: Option<MessageCapture>,
207 output_messages: Option<MessageCapture>,
208}
209
210impl TelemetryConfig {
211 pub fn with_input_messages(mut self, capture: MessageCapture) -> Self {
213 self.input_messages = Some(capture);
214 self
215 }
216
217 pub fn with_output_messages(mut self, capture: MessageCapture) -> Self {
219 self.output_messages = Some(capture);
220 self
221 }
222
223 pub fn without_input_messages(mut self) -> Self {
225 self.input_messages = None;
226 self
227 }
228
229 pub fn without_output_messages(mut self) -> Self {
231 self.output_messages = None;
232 self
233 }
234
235 pub fn input_messages(self) -> Option<MessageCapture> {
237 self.input_messages
238 }
239
240 pub fn output_messages(self) -> Option<MessageCapture> {
242 self.output_messages
243 }
244}
245
246#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
261pub struct SessionConfig {
262 pub session_id: SessionId,
264 pub metadata: MetadataMap,
266 pub cache: Option<PromptCacheRequest>,
268}
269
270impl SessionConfig {
271 pub fn new(session_id: impl Into<SessionId>) -> Self {
273 Self {
274 session_id: session_id.into(),
275 metadata: MetadataMap::new(),
276 cache: None,
277 }
278 }
279
280 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
282 self.metadata = metadata;
283 self
284 }
285
286 pub fn with_cache(mut self, cache: PromptCacheRequest) -> Self {
288 self.cache = Some(cache);
289 self
290 }
291
292 pub fn without_cache(mut self) -> Self {
294 self.cache = None;
295 self
296 }
297}
298
299#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
305pub enum PromptCacheMode {
306 Disabled,
308 #[default]
310 BestEffort,
311 Required,
313}
314
315#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
321pub enum PromptCacheRetention {
322 Default,
324 Short,
326 Extended,
328}
329
330#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
332pub enum PromptCacheStrategy {
333 #[default]
335 Automatic,
336 Explicit {
338 breakpoints: Vec<PromptCacheBreakpoint>,
340 },
341}
342
343impl PromptCacheStrategy {
344 pub fn automatic() -> Self {
347 Self::Automatic
348 }
349
350 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
352 Self::Explicit {
353 breakpoints: breakpoints.into_iter().collect(),
354 }
355 }
356}
357
358#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
360pub enum PromptCacheBreakpoint {
361 ToolsEnd,
363 TranscriptItemEnd { index: usize },
365 TranscriptPartEnd {
371 item_index: usize,
372 part_index: usize,
373 },
374}
375
376impl PromptCacheBreakpoint {
377 pub fn tools_end() -> Self {
379 Self::ToolsEnd
380 }
381
382 pub fn transcript_item_end(index: usize) -> Self {
384 Self::TranscriptItemEnd { index }
385 }
386
387 pub fn transcript_part_end(item_index: usize, part_index: usize) -> Self {
389 Self::TranscriptPartEnd {
390 item_index,
391 part_index,
392 }
393 }
394}
395
396#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
398pub struct PromptCacheRequest {
399 pub mode: PromptCacheMode,
401 pub strategy: PromptCacheStrategy,
403 pub retention: Option<PromptCacheRetention>,
405 pub key: Option<String>,
407}
408
409impl PromptCacheRequest {
410 pub fn automatic() -> Self {
412 Self::best_effort(PromptCacheStrategy::automatic())
413 }
414
415 pub fn automatic_required() -> Self {
417 Self::required(PromptCacheStrategy::automatic())
418 }
419
420 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
422 Self::best_effort(PromptCacheStrategy::explicit(breakpoints))
423 }
424
425 pub fn explicit_required(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
427 Self::required(PromptCacheStrategy::explicit(breakpoints))
428 }
429
430 pub fn disabled() -> Self {
432 Self {
433 mode: PromptCacheMode::Disabled,
434 strategy: PromptCacheStrategy::Automatic,
435 retention: None,
436 key: None,
437 }
438 }
439
440 pub fn best_effort(strategy: PromptCacheStrategy) -> Self {
442 Self {
443 mode: PromptCacheMode::BestEffort,
444 strategy,
445 retention: None,
446 key: None,
447 }
448 }
449
450 pub fn required(strategy: PromptCacheStrategy) -> Self {
452 Self {
453 mode: PromptCacheMode::Required,
454 strategy,
455 retention: None,
456 key: None,
457 }
458 }
459
460 pub fn with_mode(mut self, mode: PromptCacheMode) -> Self {
462 self.mode = mode;
463 self
464 }
465
466 pub fn with_strategy(mut self, strategy: PromptCacheStrategy) -> Self {
468 self.strategy = strategy;
469 self
470 }
471
472 pub fn with_retention(mut self, retention: PromptCacheRetention) -> Self {
474 self.retention = Some(retention);
475 self
476 }
477
478 pub fn with_key(mut self, key: impl Into<String>) -> Self {
480 self.key = Some(key.into());
481 self
482 }
483
484 pub fn without_retention(mut self) -> Self {
486 self.retention = None;
487 self
488 }
489
490 pub fn without_key(mut self) -> Self {
492 self.key = None;
493 self
494 }
495
496 pub fn is_enabled(&self) -> bool {
498 !matches!(self.mode, PromptCacheMode::Disabled)
499 }
500}
501
502#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
508pub struct TurnRequest {
509 pub session_id: SessionId,
511 pub turn_id: agentkit_core::TurnId,
513 pub transcript: Vec<Item>,
515 pub available_tools: Vec<ToolSpec>,
517 pub cache: Option<PromptCacheRequest>,
519 pub metadata: MetadataMap,
521}
522
523#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
528pub struct ModelTurnResult {
529 pub finish_reason: FinishReason,
531 pub output_items: Vec<Item>,
533 pub usage: Option<Usage>,
535 pub metadata: MetadataMap,
537 #[serde(default)]
541 pub model: Option<String>,
542 #[serde(default)]
546 pub response_id: Option<String>,
547}
548
549#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
555pub enum ModelTurnEvent {
556 Delta(Delta),
558 ToolCall(ToolCallPart),
560 Usage(Usage),
562 Finished(ModelTurnResult),
564}
565
566#[async_trait]
603pub trait ModelAdapter: Send + Sync {
604 type Session: ModelSession;
606
607 async fn start_session(&self, config: SessionConfig) -> Result<Self::Session, LoopError>;
613
614 fn provider_name(&self) -> Option<&str> {
621 None
622 }
623}
624
625#[async_trait]
632pub trait ModelSession: Send {
633 type Turn: ModelTurn;
635
636 async fn begin_turn(
649 &mut self,
650 request: TurnRequest,
651 cancellation: Option<TurnCancellation>,
652 ) -> Result<Self::Turn, LoopError>;
653
654 fn model_name(&self) -> Option<&str> {
660 None
661 }
662
663 fn provider_name(&self) -> Option<&str> {
668 None
669 }
670}
671
672#[async_trait]
678pub trait ModelTurn: Send {
679 async fn next_event(
688 &mut self,
689 cancellation: Option<TurnCancellation>,
690 ) -> Result<Option<ModelTurnEvent>, LoopError>;
691}
692
693pub trait LoopObserver: Send + Sync {
713 fn handle_event(&self, event: ObservedEvent);
718}
719
720#[derive(Clone, Debug, PartialEq)]
726pub struct ObservedEvent {
727 pub session_id: Arc<SessionId>,
729 pub event: AgentEvent,
731}
732
733pub trait TranscriptObserver: Send + Sync {
768 fn on_transcript_event(&self, event: TranscriptEvent<'_>);
773}
774
775#[derive(Clone, Debug)]
778pub struct TranscriptEvent<'a> {
779 pub session_id: &'a SessionId,
781 pub item: &'a Item,
783}
784
785#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
790#[non_exhaustive]
791pub enum MutationPoint {
792 AfterToolResult,
795 AfterTurnEnded,
798}
799
800pub trait EventEmitter: Send + Sync {
803 fn emit(&self, event: AgentEvent);
805}
806
807#[non_exhaustive]
809pub struct LoopCtx<'a> {
810 pub session_id: &'a SessionId,
812 pub turn_id: Option<&'a agentkit_core::TurnId>,
814 pub point: MutationPoint,
816 pub cancellation: Option<TurnCancellation>,
818 pub emitter: &'a dyn EventEmitter,
820}
821
822pub struct TranscriptCursor<'a> {
830 items: &'a mut Vec<Item>,
831 pub(crate) dirty: bool,
832}
833
834impl<'a> std::ops::Deref for TranscriptCursor<'a> {
835 type Target = Vec<Item>;
836 fn deref(&self) -> &Vec<Item> {
837 self.items
838 }
839}
840
841impl<'a> std::ops::DerefMut for TranscriptCursor<'a> {
842 fn deref_mut(&mut self) -> &mut Vec<Item> {
843 self.dirty = true;
844 self.items
845 }
846}
847
848#[async_trait]
856pub trait LoopMutator: Send + Sync {
857 async fn mutate(
862 &self,
863 cursor: &mut TranscriptCursor<'_>,
864 ctx: LoopCtx<'_>,
865 ) -> Result<(), LoopError> {
866 let _ = (cursor, ctx);
867 Ok(())
868 }
869}
870
871#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
876#[non_exhaustive]
877pub enum AgentEvent {
878 RunStarted { session_id: SessionId },
880 TurnStarted {
882 session_id: SessionId,
883 turn_id: agentkit_core::TurnId,
884 },
885 InputAccepted {
887 session_id: SessionId,
888 items: Vec<Item>,
889 },
890 ContentDelta(Delta),
892 ToolCallRequested(ToolCallPart),
894 ToolExecutionStarted(ToolCallPart),
896 ToolExecutionProgress(ToolResultPart),
902 ToolResultReceived(ToolResultPart),
911 ApprovalRequired(ApprovalRequest),
913 ApprovalResolved { approved: bool },
915 ToolCatalogChanged(ToolCatalogEvent),
917 MutationStarted {
920 session_id: SessionId,
921 turn_id: Option<agentkit_core::TurnId>,
922 mutator: String,
923 point: MutationPoint,
924 },
925 MutationFinished {
929 session_id: SessionId,
930 turn_id: Option<agentkit_core::TurnId>,
931 mutator: String,
932 dirty: bool,
933 metadata: MetadataMap,
934 },
935 UsageUpdated(Usage),
937 Warning { message: String },
939 RunFailed { message: String },
941 TurnFinished(TurnResult),
943}
944
945#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
967pub struct PendingApproval {
968 pub request: ApprovalRequest,
970}
971
972impl std::ops::Deref for PendingApproval {
973 type Target = ApprovalRequest;
974 fn deref(&self) -> &ApprovalRequest {
975 &self.request
976 }
977}
978
979impl PendingApproval {
980 pub fn approve<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
982 let call_id = self
983 .request
984 .call_id
985 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
986 driver.resolve_approval_for(call_id, ApprovalDecision::Approve)
987 }
988
989 pub fn deny<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
991 let call_id = self
992 .request
993 .call_id
994 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
995 driver.resolve_approval_for(call_id, ApprovalDecision::Deny { reason: None })
996 }
997
998 pub fn deny_with_reason<S: ModelSession>(
1000 self,
1001 driver: &mut LoopDriver<S>,
1002 reason: impl Into<String>,
1003 ) -> Result<(), LoopError> {
1004 let call_id = self
1005 .request
1006 .call_id
1007 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
1008 driver.resolve_approval_for(
1009 call_id,
1010 ApprovalDecision::Deny {
1011 reason: Some(reason.into()),
1012 },
1013 )
1014 }
1015
1016 pub fn approve_with_patched_input<S: ModelSession>(
1026 self,
1027 driver: &mut LoopDriver<S>,
1028 input: serde_json::Value,
1029 ) -> Result<(), LoopError> {
1030 let call_id = self
1031 .request
1032 .call_id
1033 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
1034 driver.resolve_approval_for_with_patched_input(call_id, input)
1035 }
1036}
1037
1038#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1063pub struct InputRequest {
1064 pub session_id: SessionId,
1066 pub reason: String,
1068}
1069
1070impl InputRequest {
1071 pub fn submit<S: ModelSession>(
1073 self,
1074 driver: &mut LoopDriver<S>,
1075 items: Vec<Item>,
1076 ) -> Result<(), LoopError> {
1077 driver.submit_input(items)
1078 }
1079}
1080
1081#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1086pub struct TurnResult {
1087 pub turn_id: agentkit_core::TurnId,
1089 pub finish_reason: FinishReason,
1091 pub items: Vec<Item>,
1093 pub usage: Option<Usage>,
1095 pub metadata: MetadataMap,
1097}
1098
1099#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1135pub enum LoopInterrupt {
1136 ApprovalRequest(PendingApproval),
1138 AwaitingInput(InputRequest),
1140 AfterToolResult(ToolRoundInfo),
1150}
1151
1152impl LoopInterrupt {
1153 pub fn is_blocking(&self) -> bool {
1159 matches!(self, LoopInterrupt::ApprovalRequest(_))
1160 }
1161}
1162
1163#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1166pub struct ToolRoundInfo {
1167 pub session_id: SessionId,
1169 pub turn_id: agentkit_core::TurnId,
1171 pub transcript_len: usize,
1173}
1174
1175impl ToolRoundInfo {
1176 pub fn submit<S: ModelSession>(
1179 self,
1180 driver: &mut LoopDriver<S>,
1181 items: Vec<Item>,
1182 ) -> Result<(), LoopError> {
1183 driver.submit_input(items)
1184 }
1185}
1186
1187#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1215pub enum LoopStep {
1216 Interrupt(LoopInterrupt),
1218 Finished(TurnResult),
1220}
1221
1222#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1228pub struct LoopSnapshot {
1229 pub session_id: SessionId,
1231 pub transcript: Vec<Item>,
1233 pub pending_input: Vec<Item>,
1235}
1236
1237#[derive(Clone)]
1238struct PendingApprovalToolCall {
1239 request: ApprovalRequest,
1240 decision: Option<ApprovalDecision>,
1241 surfaced: bool,
1242 presentation_turn_id: agentkit_core::TurnId,
1243 task_id: TaskId,
1244 call: ToolCallPart,
1245 tool_request: ToolRequest,
1246 cancellation: Option<TurnCancellation>,
1247}
1248
1249#[derive(Clone, Default)]
1250struct ActiveToolRound {
1251 presentation_turn_id: agentkit_core::TurnId,
1252 task_turn_id: agentkit_core::TurnId,
1253 pending_calls: VecDeque<(ToolCallPart, ToolRequest)>,
1254 cancellation: Option<TurnCancellation>,
1255 background_pending: bool,
1256 foreground_progressed: bool,
1257}
1258
1259#[derive(Default)]
1260struct DriverLifecycle {
1261 active_turn: Option<agentkit_core::TurnId>,
1262}
1263
1264pub struct Agent<M>
1306where
1307 M: ModelAdapter,
1308{
1309 model: M,
1310 tool_sources: Vec<Arc<dyn ToolSource>>,
1311 tool_executor: Option<Arc<dyn ToolExecutor>>,
1312 task_manager: Arc<dyn TaskManager>,
1313 permissions: Arc<dyn PermissionChecker>,
1314 resources: Arc<dyn ToolResources>,
1315 cancellation: Option<CancellationHandle>,
1316 mutators: Vec<Arc<dyn LoopMutator>>,
1317 observers: Vec<Arc<dyn LoopObserver>>,
1318 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1319 transcript: Vec<Item>,
1320 input: Vec<Item>,
1321 telemetry: TelemetryConfig,
1322}
1323
1324impl<M> Agent<M>
1325where
1326 M: ModelAdapter,
1327{
1328 pub fn builder() -> AgentBuilder<M> {
1330 AgentBuilder::default()
1331 }
1332
1333 pub async fn start(&self, config: SessionConfig) -> Result<LoopDriver<M::Session>, LoopError> {
1349 let session_id = config.session_id.clone();
1350 let default_cache = config.cache.clone();
1351 let session = self.model.start_session(config).await?;
1352 let provider_name = self.model.provider_name().map(str::to_owned);
1353 let tool_executor = self
1354 .tool_executor
1355 .clone()
1356 .unwrap_or_else(|| Arc::new(BasicToolExecutor::new(self.tool_sources.clone())));
1357 let driver = LoopDriver {
1358 session_id: session_id.clone(),
1359 observed_session_id: Arc::new(session_id.clone()),
1360 provider_name,
1361 telemetry: self.telemetry,
1362 default_cache,
1363 next_turn_cache: None,
1364 session: Some(session),
1365 tool_executor,
1366 task_manager: self.task_manager.clone(),
1367 permissions: self.permissions.clone(),
1368 resources: self.resources.clone(),
1369 cancellation: self.cancellation.clone(),
1370 mutators: self.mutators.clone(),
1371 observers: self.observers.clone(),
1372 transcript_observers: self.transcript_observers.clone(),
1373 transcript: self.transcript.clone(),
1374 pending_input: self.input.clone(),
1375 pending_approvals: BTreeMap::new(),
1376 pending_approval_order: VecDeque::new(),
1377 active_tool_round: None,
1378 pending_round_resume: None,
1379 pending_loop_updates: VecDeque::new(),
1380 next_turn_index: 1,
1381 lifecycle: DriverLifecycle::default(),
1382 background_call_ids: HashSet::new(),
1383 detached_call_ids: HashSet::new(),
1384 interrupted_background_call_ids: HashSet::new(),
1385 tool_cancellations: HashMap::new(),
1386 };
1387 driver.emit(AgentEvent::RunStarted { session_id });
1388 Ok(driver)
1389 }
1390}
1391
1392pub struct AgentBuilder<M>
1398where
1399 M: ModelAdapter,
1400{
1401 model: Option<M>,
1402 tool_sources: Vec<Arc<dyn ToolSource>>,
1403 tool_executor: Option<Arc<dyn ToolExecutor>>,
1404 task_manager: Option<Arc<dyn TaskManager>>,
1405 permissions: Arc<dyn PermissionChecker>,
1406 resources: Arc<dyn ToolResources>,
1407 cancellation: Option<CancellationHandle>,
1408 mutators: Vec<Arc<dyn LoopMutator>>,
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> Default for AgentBuilder<M>
1417where
1418 M: ModelAdapter,
1419{
1420 fn default() -> Self {
1421 Self {
1422 model: None,
1423 tool_sources: Vec::new(),
1424 tool_executor: None,
1425 task_manager: None,
1426 permissions: Arc::new(AllowAllPermissions),
1427 resources: Arc::new(()),
1428 cancellation: None,
1429 mutators: Vec::new(),
1430 observers: Vec::new(),
1431 transcript_observers: Vec::new(),
1432 transcript: Vec::new(),
1433 input: Vec::new(),
1434 telemetry: TelemetryConfig::default(),
1435 }
1436 }
1437}
1438
1439impl<M> AgentBuilder<M>
1440where
1441 M: ModelAdapter,
1442{
1443 pub fn model(mut self, model: M) -> Self {
1445 self.model = Some(model);
1446 self
1447 }
1448
1449 pub fn add_tool_source<S: ToolSource + 'static>(mut self, source: S) -> Self {
1462 self.tool_sources.push(Arc::new(source));
1463 self
1464 }
1465
1466 pub fn tool_executor(mut self, executor: impl ToolExecutor + 'static) -> Self {
1472 self.tool_executor = Some(Arc::new(executor));
1473 self
1474 }
1475
1476 pub fn task_manager(mut self, manager: impl TaskManager + 'static) -> Self {
1481 self.task_manager = Some(Arc::new(manager));
1482 self
1483 }
1484
1485 pub fn permissions(mut self, permissions: impl PermissionChecker + 'static) -> Self {
1489 self.permissions = Arc::new(permissions);
1490 self
1491 }
1492
1493 pub fn resources(mut self, resources: impl ToolResources + 'static) -> Self {
1495 self.resources = Arc::new(resources);
1496 self
1497 }
1498
1499 pub fn cancellation(mut self, handle: CancellationHandle) -> Self {
1501 self.cancellation = Some(handle);
1502 self
1503 }
1504
1505 pub fn mutator<L: LoopMutator + 'static>(mut self, mutator: L) -> Self {
1513 self.mutators.push(Arc::new(mutator));
1514 self
1515 }
1516
1517 pub fn observer<O: LoopObserver + 'static>(mut self, observer: O) -> Self {
1521 self.observers.push(Arc::new(observer));
1522 self
1523 }
1524
1525 pub fn transcript_observer<O: TranscriptObserver + 'static>(mut self, observer: O) -> Self {
1534 self.transcript_observers.push(Arc::new(observer));
1535 self
1536 }
1537
1538 pub fn transcript(mut self, transcript: Vec<Item>) -> Self {
1546 self.transcript = transcript;
1547 self
1548 }
1549
1550 pub fn input(mut self, input: Vec<Item>) -> Self {
1559 self.input = input;
1560 self
1561 }
1562
1563 pub fn telemetry(mut self, telemetry: TelemetryConfig) -> Self {
1566 self.telemetry = telemetry;
1567 self
1568 }
1569
1570 pub fn build(self) -> Result<Agent<M>, LoopError> {
1576 let model = self
1577 .model
1578 .ok_or_else(|| LoopError::InvalidState("model adapter is required".into()))?;
1579 Ok(Agent {
1580 model,
1581 tool_sources: self.tool_sources,
1582 tool_executor: self.tool_executor,
1583 task_manager: self
1584 .task_manager
1585 .unwrap_or_else(|| Arc::new(SimpleTaskManager::new())),
1586 permissions: self.permissions,
1587 resources: self.resources,
1588 cancellation: self.cancellation,
1589 mutators: self.mutators,
1590 observers: self.observers,
1591 transcript_observers: self.transcript_observers,
1592 transcript: self.transcript,
1593 input: self.input,
1594 telemetry: self.telemetry,
1595 })
1596 }
1597}
1598
1599pub struct LoopDriver<S>
1630where
1631 S: ModelSession,
1632{
1633 session_id: SessionId,
1634 observed_session_id: Arc<SessionId>,
1635 provider_name: Option<String>,
1636 telemetry: TelemetryConfig,
1637 default_cache: Option<PromptCacheRequest>,
1638 next_turn_cache: Option<PromptCacheRequest>,
1639 session: Option<S>,
1640 tool_executor: Arc<dyn ToolExecutor>,
1641 task_manager: Arc<dyn TaskManager>,
1642 permissions: Arc<dyn PermissionChecker>,
1643 resources: Arc<dyn ToolResources>,
1644 cancellation: Option<CancellationHandle>,
1645 mutators: Vec<Arc<dyn LoopMutator>>,
1646 observers: Vec<Arc<dyn LoopObserver>>,
1647 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1648 transcript: Vec<Item>,
1649 pending_input: Vec<Item>,
1650 pending_approvals: BTreeMap<ToolCallId, PendingApprovalToolCall>,
1651 pending_approval_order: VecDeque<ToolCallId>,
1652 active_tool_round: Option<ActiveToolRound>,
1653 pending_round_resume: Option<agentkit_core::TurnId>,
1654 pending_loop_updates: VecDeque<TaskResolution>,
1655 next_turn_index: u64,
1656 lifecycle: DriverLifecycle,
1657 background_call_ids: HashSet<ToolCallId>,
1659 detached_call_ids: HashSet<ToolCallId>,
1667 interrupted_background_call_ids: HashSet<ToolCallId>,
1671 tool_cancellations: HashMap<ToolCallId, TurnCancellation>,
1672}
1673
1674impl<S> LoopDriver<S>
1675where
1676 S: ModelSession,
1677{
1678 fn execute_tool_span(
1679 &self,
1680 request: &ToolRequest,
1681 turn_id: &agentkit_core::TurnId,
1682 launch_kind: &'static str,
1683 ) -> tracing::Span {
1684 tracing::info_span!(
1685 "agent.execute_tool",
1686 "otel.name" = %format!("execute_tool {}", request.tool_name),
1687 "gen_ai.operation.name" = "execute_tool",
1688 "gen_ai.tool.name" = %request.tool_name,
1689 "gen_ai.tool.call.id" = %request.call_id,
1690 "gen_ai.conversation.id" = %self.session_id,
1691 "error.type" = tracing::field::Empty,
1692 session.id = %self.session_id,
1693 turn.id = %turn_id,
1694 launch_kind = launch_kind,
1695 )
1696 }
1697
1698 fn start_task_via_manager(
1699 &self,
1700 task_id: Option<TaskId>,
1701 tool_request: ToolRequest,
1702 kind: TaskLaunchKind,
1703 cancellation: Option<TurnCancellation>,
1704 ) -> impl std::future::Future<Output = Result<TaskStartOutcome, LoopError>> + Send + 'static
1705 {
1706 let task_manager = self.task_manager.clone();
1707 let tool_executor = self.tool_executor.clone();
1708 let permissions = self.permissions.clone();
1709 let resources = self.resources.clone();
1710 let session_id = self.session_id.clone();
1711 let turn_id = tool_request.turn_id.clone();
1712 let metadata = tool_request.metadata.clone();
1713
1714 async move {
1715 task_manager
1716 .start_task(
1717 TaskLaunchRequest {
1718 task_id,
1719 request: tool_request.clone(),
1720 kind,
1721 },
1722 TaskStartContext {
1723 executor: tool_executor.clone(),
1724 tool_context: {
1725 let execution_scope = ToolExecutionScope {
1726 executor: tool_executor,
1727 session_id: session_id.clone(),
1728 turn_id: turn_id.clone(),
1729 permissions: permissions.clone(),
1730 resources: resources.clone(),
1731 cancellation: cancellation.clone(),
1732 };
1733 OwnedToolContext {
1734 session_id,
1735 turn_id,
1736 metadata,
1737 permissions,
1738 resources,
1739 cancellation,
1740 execution_scope: Some(execution_scope),
1741 approved_request: None,
1742 }
1743 },
1744 },
1745 )
1746 .await
1747 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))
1748 }
1749 }
1750
1751 fn register_tool_cancellation(
1752 &mut self,
1753 call_id: &ToolCallId,
1754 cancellation: Option<TurnCancellation>,
1755 ) {
1756 if let Some(cancellation) = cancellation {
1757 self.tool_cancellations
1758 .insert(call_id.clone(), cancellation);
1759 }
1760 }
1761
1762 fn tool_cancellation_for(
1763 &mut self,
1764 call_id: &ToolCallId,
1765 fallback: Option<TurnCancellation>,
1766 ) -> Option<TurnCancellation> {
1767 self.tool_cancellations.get(call_id).cloned().or(fallback)
1768 }
1769
1770 fn clear_tool_cancellation(&mut self, call_id: &ToolCallId) {
1771 self.tool_cancellations.remove(call_id);
1772 }
1773
1774 fn has_pending_interrupts(&self) -> bool {
1775 !self.pending_approvals.is_empty()
1776 }
1777
1778 fn start_logical_turn(&mut self) -> agentkit_core::TurnId {
1779 if let Some(turn_id) = &self.lifecycle.active_turn {
1780 return turn_id.clone();
1781 }
1782 let turn_id = agentkit_core::TurnId::new(format!("turn-{}", self.next_turn_index));
1783 self.next_turn_index += 1;
1784 self.start_logical_turn_with(turn_id)
1785 }
1786
1787 fn start_logical_turn_with(&mut self, turn_id: agentkit_core::TurnId) -> agentkit_core::TurnId {
1788 if let Some(active_turn) = &self.lifecycle.active_turn {
1789 return active_turn.clone();
1790 }
1791 self.lifecycle.active_turn = Some(turn_id.clone());
1792 self.emit(AgentEvent::TurnStarted {
1793 session_id: self.session_id.clone(),
1794 turn_id: turn_id.clone(),
1795 });
1796 turn_id
1797 }
1798
1799 fn finish_logical_turn(&mut self, result: &TurnResult) {
1800 if self.pending_round_resume.as_ref() == Some(&result.turn_id) {
1801 self.pending_round_resume = None;
1802 }
1803 if self.lifecycle.active_turn.as_ref() == Some(&result.turn_id) {
1804 self.lifecycle.active_turn = None;
1805 self.emit(AgentEvent::TurnFinished(result.clone()));
1806 }
1807 }
1808
1809 fn emit_tool_catalog_events(&mut self, events: Vec<ToolCatalogEvent>) {
1810 for event in events {
1811 self.emit(AgentEvent::ToolCatalogChanged(event));
1812 }
1813 }
1814
1815 fn enqueue_pending_approval(
1816 &mut self,
1817 presentation_turn_id: &agentkit_core::TurnId,
1818 task: TaskApproval,
1819 cancellation: Option<TurnCancellation>,
1820 ) {
1821 let call_id = task.tool_request.call_id.clone();
1822 self.background_call_ids.remove(&call_id);
1823 let cancellation = self.tool_cancellation_for(&call_id, cancellation);
1824 let call = ToolCallPart {
1825 id: call_id.clone(),
1826 name: task.tool_request.tool_name.to_string(),
1827 input: task.tool_request.input.clone(),
1828 metadata: task.tool_request.metadata.clone(),
1829 };
1830 let mut request = task.approval;
1831 request.call_id = Some(call_id.clone());
1832 let pending = PendingApprovalToolCall {
1833 request: request.clone(),
1834 decision: None,
1835 surfaced: false,
1836 presentation_turn_id: presentation_turn_id.clone(),
1837 task_id: task.task_id,
1838 call,
1839 tool_request: task.tool_request,
1840 cancellation,
1841 };
1842 self.pending_approvals.insert(call_id.clone(), pending);
1843 if !self.pending_approval_order.iter().any(|id| id == &call_id) {
1844 self.pending_approval_order.push_back(call_id);
1845 }
1846 self.emit(AgentEvent::ApprovalRequired(request));
1847 }
1848
1849 fn take_next_unsurfaced_approval_interrupt(&mut self) -> Option<LoopStep> {
1850 for call_id in self.pending_approval_order.clone() {
1851 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
1852 continue;
1853 };
1854 if pending.decision.is_none() && !pending.surfaced {
1855 pending.surfaced = true;
1856 return Some(LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(
1857 PendingApproval {
1858 request: pending.request.clone(),
1859 },
1860 )));
1861 }
1862 }
1863 None
1864 }
1865
1866 fn next_unresolved_approval_interrupt(&self) -> Option<LoopStep> {
1867 self.pending_approval_order.iter().find_map(|call_id| {
1868 self.pending_approvals.get(call_id).and_then(|pending| {
1869 pending.decision.is_none().then(|| {
1870 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(PendingApproval {
1871 request: pending.request.clone(),
1872 }))
1873 })
1874 })
1875 })
1876 }
1877
1878 fn take_next_resolved_approval(&mut self) -> Option<PendingApprovalToolCall> {
1879 let call_id = self.pending_approval_order.iter().find_map(|call_id| {
1880 self.pending_approvals
1881 .get(call_id)
1882 .and_then(|pending| pending.decision.as_ref().map(|_| call_id.clone()))
1883 })?;
1884 self.pending_approval_order.retain(|id| id != &call_id);
1885 self.pending_approvals.remove(&call_id)
1886 }
1887
1888 fn queue_resolution_interrupt(
1889 &mut self,
1890 presentation_turn_id: &agentkit_core::TurnId,
1891 resolution: TaskResolution,
1892 cancellation: Option<TurnCancellation>,
1893 ) -> Option<LoopStep> {
1894 match resolution {
1895 TaskResolution::Item(item) => {
1896 self.append_tool_result_item(item);
1897 None
1898 }
1899 TaskResolution::Approval(task) => {
1900 self.enqueue_pending_approval(presentation_turn_id, task, cancellation);
1901 self.take_next_unsurfaced_approval_interrupt()
1902 }
1903 }
1904 }
1905
1906 async fn collect_pending_loop_updates(&mut self) -> Result<(), LoopError> {
1907 let PendingLoopUpdates { resolutions } = self
1908 .task_manager
1909 .take_pending_loop_updates()
1910 .await
1911 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
1912 self.pending_loop_updates.extend(resolutions);
1913 Ok(())
1914 }
1915
1916 async fn drain_pending_loop_updates(&mut self) -> Result<(bool, Option<LoopStep>), LoopError> {
1917 self.collect_pending_loop_updates().await?;
1918 let mut resolutions = std::mem::take(&mut self.pending_loop_updates);
1919 if !resolutions.is_empty() {
1920 self.start_logical_turn();
1921 }
1922 let mut saw_items = false;
1923 while let Some(resolution) = resolutions.pop_front() {
1924 match resolution {
1925 TaskResolution::Item(item) => {
1926 self.append_tool_result_item(item);
1927 saw_items = true;
1928 }
1929 TaskResolution::Approval(task) => {
1930 let turn_id = self.start_logical_turn();
1931 self.enqueue_pending_approval(&turn_id, task, None);
1932 }
1933 }
1934 }
1935 if let Some(step) = self.finish_cancelled_pending_approval().await? {
1936 return Ok((saw_items, Some(step)));
1937 }
1938 Ok((saw_items, self.take_next_unsurfaced_approval_interrupt()))
1939 }
1940
1941 async fn finish_cancelled_pending_approval(&mut self) -> Result<Option<LoopStep>, LoopError> {
1942 if self.pending_approvals.is_empty() {
1943 return Ok(None);
1944 }
1945 if !self.pending_approvals.values().any(|pending| {
1946 pending
1947 .cancellation
1948 .as_ref()
1949 .is_some_and(TurnCancellation::is_cancelled)
1950 }) {
1951 return Ok(None);
1952 }
1953 self.cancel_pending_approvals().await
1954 }
1955
1956 async fn run_mutators(
1957 &mut self,
1958 point: MutationPoint,
1959 turn_id: Option<&agentkit_core::TurnId>,
1960 cancellation: Option<TurnCancellation>,
1961 ) -> Result<(), LoopError> {
1962 if self.mutators.is_empty() {
1963 return Ok(());
1964 }
1965 if cancellation
1966 .as_ref()
1967 .is_some_and(TurnCancellation::is_cancelled)
1968 {
1969 return Err(LoopError::Cancelled);
1970 }
1971 let mutators = self.mutators.clone();
1972 let session_id = self.session_id.clone();
1973 let observed_session_id = Arc::clone(&self.observed_session_id);
1974 let observers = self.observers.clone();
1975 let emitter = DriverEmitter {
1976 session_id: &observed_session_id,
1977 observers: &observers,
1978 };
1979 let mut cursor = TranscriptCursor {
1980 items: &mut self.transcript,
1981 dirty: false,
1982 };
1983 for mutator in &mutators {
1984 if cancellation
1985 .as_ref()
1986 .is_some_and(TurnCancellation::is_cancelled)
1987 {
1988 return Err(LoopError::Cancelled);
1989 }
1990 let ctx = LoopCtx {
1991 session_id: &session_id,
1992 turn_id,
1993 point,
1994 cancellation: cancellation.clone(),
1995 emitter: &emitter,
1996 };
1997 mutator.mutate(&mut cursor, ctx).await?;
1998 }
1999 if cursor.dirty {
2000 validate_transcript_invariants(cursor.items)?;
2001 }
2002 Ok(())
2003 }
2004
2005 async fn continue_active_tool_round(&mut self) -> Result<Option<LoopStep>, LoopError> {
2006 let Some((presentation_turn_id, task_turn_id, cancellation)) =
2007 self.active_tool_round.as_ref().map(|active| {
2008 (
2009 active.presentation_turn_id.clone(),
2010 active.task_turn_id.clone(),
2011 active.cancellation.clone(),
2012 )
2013 })
2014 else {
2015 return Ok(None);
2016 };
2017 loop {
2018 if cancellation
2019 .as_ref()
2020 .is_some_and(TurnCancellation::is_cancelled)
2021 {
2022 self.task_manager
2023 .on_turn_interrupted(&task_turn_id)
2024 .await
2025 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2026 self.active_tool_round = None;
2027 return self
2028 .finish_cancelled(presentation_turn_id, Vec::new())
2029 .map(Some);
2030 }
2031
2032 let next_call = self
2033 .active_tool_round
2034 .as_mut()
2035 .and_then(|active| active.pending_calls.pop_front());
2036 if let Some((call, tool_request)) = next_call {
2037 use tracing::Instrument;
2038 self.register_tool_cancellation(&call.id, cancellation.clone());
2039 let dispatch_span =
2040 self.execute_tool_span(&tool_request, &presentation_turn_id, "plain");
2041 match self
2042 .start_task_via_manager(
2043 None,
2044 tool_request.clone(),
2045 TaskLaunchKind::Plain,
2046 cancellation.clone(),
2047 )
2048 .instrument(dispatch_span.clone())
2049 .await?
2050 {
2051 TaskStartOutcome::Ready(resolution) => {
2052 let resolution = *resolution;
2053 match resolution {
2054 TaskResolution::Item(item) => {
2055 if !tool_result_not_started(&item) {
2056 self.emit(AgentEvent::ToolExecutionStarted(call.clone()));
2057 }
2058 if tool_result_is_error(&item) {
2059 dispatch_span.record("error.type", "tool_error");
2060 }
2061 if let Some(active) = self.active_tool_round.as_mut() {
2062 active.foreground_progressed = true;
2063 }
2064 self.append_tool_result_item(item);
2065 }
2066 TaskResolution::Approval(task) => {
2067 self.enqueue_pending_approval(
2068 &presentation_turn_id,
2069 task,
2070 cancellation.clone(),
2071 );
2072 }
2073 }
2074 continue;
2075 }
2076 TaskStartOutcome::Pending { kind, .. } => {
2077 self.emit(AgentEvent::ToolExecutionStarted(call.clone()));
2078 if kind == agentkit_task_manager::TaskKind::Background {
2079 self.append_detach_placeholder(call.id.clone(), &call.name);
2080 if let Some(active) = self.active_tool_round.as_mut() {
2081 active.background_pending = true;
2082 }
2083 }
2084 continue;
2085 }
2086 }
2087 }
2088
2089 match self
2090 .task_manager
2091 .wait_for_turn(&task_turn_id, cancellation.clone())
2092 .await
2093 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?
2094 {
2095 Some(TurnTaskUpdate::Resolution(resolution)) => {
2096 let resolution = *resolution;
2097 match resolution {
2098 TaskResolution::Item(item) => {
2099 if let Some(active) = self.active_tool_round.as_mut() {
2100 active.foreground_progressed = true;
2101 }
2102 self.append_tool_result_item(item);
2103 }
2104 TaskResolution::Approval(task) => {
2105 self.enqueue_pending_approval(
2106 &presentation_turn_id,
2107 task,
2108 cancellation.clone(),
2109 );
2110 }
2111 }
2112 }
2113 Some(TurnTaskUpdate::Detached(snapshot)) => {
2114 self.append_detach_placeholder(snapshot.call_id, &snapshot.tool_name);
2115 if let Some(active) = self.active_tool_round.as_mut() {
2116 active.background_pending = true;
2117 active.foreground_progressed = true;
2118 }
2119 }
2120 None => {
2121 if cancellation
2122 .as_ref()
2123 .is_some_and(TurnCancellation::is_cancelled)
2124 {
2125 self.task_manager
2126 .on_turn_interrupted(&task_turn_id)
2127 .await
2128 .map_err(|error| {
2129 LoopError::Tool(ToolError::Internal(error.to_string()))
2130 })?;
2131 self.active_tool_round = None;
2132 return self
2133 .finish_cancelled(presentation_turn_id, Vec::new())
2134 .map(Some);
2135 }
2136 let active = self.active_tool_round.take().ok_or_else(|| {
2137 LoopError::InvalidState("missing active tool round".into())
2138 })?;
2139 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2140 return Ok(Some(step));
2141 }
2142 if let Some(step) = self.next_unresolved_approval_interrupt() {
2143 return Ok(Some(step));
2144 }
2145 if active.background_pending && !active.foreground_progressed {
2146 return Ok(None);
2147 }
2148 let info = ToolRoundInfo {
2155 session_id: self.session_id.clone(),
2156 turn_id: presentation_turn_id.clone(),
2157 transcript_len: self.transcript.len(),
2158 };
2159 self.pending_round_resume = Some(presentation_turn_id);
2160 return Ok(Some(LoopStep::Interrupt(LoopInterrupt::AfterToolResult(
2161 info,
2162 ))));
2163 }
2164 }
2165 }
2166 }
2167
2168 #[tracing::instrument(
2169 name = "agent.turn",
2170 skip_all,
2171 fields(
2172 otel.name = "invoke_agent",
2173 gen_ai.operation.name = "invoke_agent",
2174 gen_ai.conversation.id = %self.session_id,
2175 gen_ai.provider.name = tracing::field::Empty,
2176 session.id = %self.session_id,
2177 turn.id = %turn_id,
2178 transcript.len = self.transcript.len(),
2179 saw_tool_call = tracing::field::Empty,
2180 finish_reason = tracing::field::Empty,
2181 ),
2182 )]
2183 async fn drive_turn(
2184 &mut self,
2185 turn_id: agentkit_core::TurnId,
2186 mutation_point: MutationPoint,
2187 ) -> Result<LoopStep, LoopError> {
2188 let cancellation = self
2189 .cancellation
2190 .as_ref()
2191 .map(CancellationHandle::checkpoint);
2192 match self
2193 .run_mutators(mutation_point, Some(&turn_id), cancellation.clone())
2194 .await
2195 {
2196 Ok(()) => {}
2197 Err(LoopError::Cancelled) => {
2198 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2199 }
2200 Err(error) => return Err(error),
2201 }
2202
2203 if !transcript_has_pending_input(&self.transcript) {
2209 let turn_result = TurnResult {
2210 turn_id,
2211 finish_reason: FinishReason::Completed,
2212 items: Vec::new(),
2213 usage: None,
2214 metadata: MetadataMap::new(),
2215 };
2216 self.finish_logical_turn(&turn_result);
2217 return Ok(LoopStep::Finished(turn_result));
2218 }
2219
2220 if cancellation
2221 .as_ref()
2222 .is_some_and(TurnCancellation::is_cancelled)
2223 {
2224 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2225 }
2226
2227 let catalog_events = self.tool_executor.drain_catalog_events();
2228 self.emit_tool_catalog_events(catalog_events);
2229
2230 let request = TurnRequest {
2231 session_id: self.session_id.clone(),
2232 turn_id: turn_id.clone(),
2233 transcript: self.transcript.clone(),
2234 available_tools: self.tool_executor.specs(),
2235 cache: self
2236 .next_turn_cache
2237 .take()
2238 .or_else(|| self.default_cache.clone()),
2239 metadata: MetadataMap::new(),
2240 };
2241
2242 let session = self
2243 .session
2244 .as_mut()
2245 .ok_or_else(|| LoopError::InvalidState("model session is not available".into()))?;
2246
2247 let chat_span = tracing::info_span!(
2254 "chat",
2255 "otel.name" = tracing::field::Empty,
2256 "otel.kind" = "client",
2257 "gen_ai.operation.name" = "chat",
2258 "gen_ai.provider.name" = tracing::field::Empty,
2259 "gen_ai.conversation.id" = %self.session_id,
2260 "gen_ai.request.model" = tracing::field::Empty,
2261 "gen_ai.response.model" = tracing::field::Empty,
2262 "gen_ai.response.id" = tracing::field::Empty,
2263 "gen_ai.response.finish_reasons" = tracing::field::Empty,
2264 "gen_ai.input.messages" = tracing::field::Empty,
2265 "gen_ai.output.messages" = tracing::field::Empty,
2266 "gen_ai.usage.input_tokens" = tracing::field::Empty,
2267 "gen_ai.usage.output_tokens" = tracing::field::Empty,
2268 "gen_ai.usage.cost" = tracing::field::Empty,
2269 );
2270 if let Some(capture) = self.telemetry.input_messages {
2271 record_string_array_attribute(
2272 &chat_span,
2273 "gen_ai.input.messages",
2274 capture_messages(&request.transcript, capture, CaptureOrder::NewestTail),
2275 );
2276 }
2277
2278 let initial_provider_name =
2282 effective_provider_name(session.provider_name(), self.provider_name.as_deref());
2283 if let Some(provider) = &initial_provider_name {
2284 chat_span.record("gen_ai.provider.name", provider.as_str());
2285 tracing::Span::current().record("gen_ai.provider.name", provider.as_str());
2286 }
2287 match session.model_name() {
2288 Some(model) => {
2289 chat_span.record("gen_ai.request.model", model);
2290 chat_span.record("otel.name", format!("chat {model}").as_str());
2291 }
2292 None => {
2293 chat_span.record("otel.name", "chat");
2294 }
2295 }
2296
2297 use tracing::Instrument;
2298 let mut turn = match session
2299 .begin_turn(request, cancellation.clone())
2300 .instrument(chat_span.clone())
2301 .await
2302 {
2303 Ok(turn) => turn,
2304 Err(LoopError::Cancelled) => {
2305 self.task_manager
2306 .on_turn_interrupted(&turn_id)
2307 .await
2308 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2309 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2310 }
2311 Err(error) => return Err(error),
2312 };
2313
2314 let provider_name =
2317 effective_provider_name(session.provider_name(), self.provider_name.as_deref());
2318 if let Some(provider) = &provider_name {
2319 chat_span.record("gen_ai.provider.name", provider.as_str());
2320 tracing::Span::current().record("gen_ai.provider.name", provider.as_str());
2321 }
2322 match session.model_name() {
2323 Some(model) => {
2324 chat_span.record("gen_ai.request.model", model);
2325 chat_span.record("otel.name", format!("chat {model}").as_str());
2326 }
2327 None => {
2328 chat_span.record("otel.name", "chat");
2329 }
2330 }
2331
2332 let mut saw_tool_call = false;
2333 let mut finished_result = None;
2334 let mut latest_usage = None;
2335
2336 while let Some(event) = match turn
2337 .next_event(cancellation.clone())
2338 .instrument(chat_span.clone())
2339 .await
2340 {
2341 Ok(event) => event,
2342 Err(LoopError::Cancelled) => {
2343 self.task_manager
2344 .on_turn_interrupted(&turn_id)
2345 .await
2346 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2347 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2348 }
2349 Err(error) => return Err(error),
2350 } {
2351 if cancellation
2352 .as_ref()
2353 .is_some_and(TurnCancellation::is_cancelled)
2354 {
2355 self.task_manager
2356 .on_turn_interrupted(&turn_id)
2357 .await
2358 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2359 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2360 }
2361 match event {
2362 ModelTurnEvent::Delta(delta) => self.emit(AgentEvent::ContentDelta(delta)),
2363 ModelTurnEvent::Usage(usage) => {
2364 latest_usage = Some(usage.clone());
2365 self.emit(AgentEvent::UsageUpdated(usage));
2366 }
2367 ModelTurnEvent::ToolCall(call) => {
2368 saw_tool_call = true;
2369 self.emit(AgentEvent::ToolCallRequested(call.clone()));
2370 }
2371 ModelTurnEvent::Finished(result) => {
2372 finished_result = Some(result);
2373 break;
2374 }
2375 }
2376 }
2377
2378 let mut result = finished_result.ok_or_else(|| {
2379 LoopError::Provider("model turn ended without a Finished event".into())
2380 })?;
2381 result.usage = merge_usage(result.usage, latest_usage);
2382 if let Some(model) = &result.model {
2383 chat_span.record("gen_ai.response.model", model.as_str());
2384 }
2385 if let Some(id) = &result.response_id {
2386 chat_span.record("gen_ai.response.id", id.as_str());
2387 }
2388 if let Some(tokens) = result
2389 .usage
2390 .as_ref()
2391 .and_then(|usage| usage.tokens.as_ref())
2392 {
2393 record_token_attribute(&chat_span, "gen_ai.usage.input_tokens", tokens.input_tokens);
2394 record_token_attribute(
2395 &chat_span,
2396 "gen_ai.usage.output_tokens",
2397 tokens.output_tokens,
2398 );
2399 }
2400 if let Some(cost) = result.usage.as_ref().and_then(|usage| usage.cost.as_ref()) {
2401 record_f64_attribute(&chat_span, "gen_ai.usage.cost", cost.amount);
2402 }
2403 record_string_array_attribute(
2404 &chat_span,
2405 "gen_ai.response.finish_reasons",
2406 provider_finish_reasons(&result.metadata, &result.finish_reason),
2407 );
2408 if let Some(capture) = self.telemetry.output_messages {
2409 record_string_array_attribute(
2410 &chat_span,
2411 "gen_ai.output.messages",
2412 capture_messages(&result.output_items, capture, CaptureOrder::OldestHead),
2413 );
2414 }
2415 drop(chat_span);
2416 tracing::Span::current().record("saw_tool_call", saw_tool_call);
2417 tracing::Span::current().record(
2418 "finish_reason",
2419 tracing::field::debug(&result.finish_reason),
2420 );
2421 let now = Timestamp::now();
2422 let usage = result.usage.clone();
2423 let finish_reason = result.finish_reason.clone();
2424 let output_items: Vec<Item> = result
2425 .output_items
2426 .drain(..)
2427 .map(|mut item| {
2428 if matches!(item.kind, ItemKind::Assistant) {
2429 if item.usage.is_none() {
2430 item.usage = usage.clone();
2431 }
2432 if item.finish_reason.is_none() {
2433 item.finish_reason = Some(finish_reason.clone());
2434 }
2435 }
2436 if item.created_at.is_none() {
2437 item.created_at = Some(now);
2438 }
2439 item
2440 })
2441 .collect();
2442 self.extend_transcript(output_items.clone());
2443
2444 if saw_tool_call {
2445 let pending_calls = extract_tool_calls(&output_items)
2446 .into_iter()
2447 .map(|call| {
2448 let tool_request = ToolRequest {
2449 call_id: call.id.clone(),
2450 tool_name: agentkit_tools_core::ToolName::new(call.name.clone()),
2451 input: call.input.clone(),
2452 session_id: self.session_id.clone(),
2453 turn_id: turn_id.clone(),
2454 metadata: call.metadata.clone(),
2455 };
2456 (call, tool_request)
2457 })
2458 .collect();
2459 self.active_tool_round = Some(ActiveToolRound {
2460 presentation_turn_id: turn_id.clone(),
2461 task_turn_id: turn_id.clone(),
2462 pending_calls,
2463 cancellation: cancellation.clone(),
2464 background_pending: false,
2465 foreground_progressed: false,
2466 });
2467 if let Some(step) = self.continue_active_tool_round().await? {
2468 return Ok(step);
2469 }
2470 self.finish_logical_turn(&TurnResult {
2471 turn_id,
2472 finish_reason: result.finish_reason,
2473 items: output_items,
2474 usage: result.usage,
2475 metadata: result.metadata,
2476 });
2477 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
2478 InputRequest {
2479 session_id: self.session_id.clone(),
2480 reason: "driver is waiting for input".into(),
2481 },
2482 )));
2483 }
2484
2485 let turn_result = TurnResult {
2486 turn_id,
2487 finish_reason: result.finish_reason,
2488 items: output_items,
2489 usage: result.usage,
2490 metadata: result.metadata,
2491 };
2492 self.finish_logical_turn(&turn_result);
2493 Ok(LoopStep::Finished(turn_result))
2494 }
2495
2496 async fn resume_after_approval(
2497 &mut self,
2498 pending: PendingApprovalToolCall,
2499 ) -> Result<LoopStep, LoopError> {
2500 let decision = pending
2501 .decision
2502 .clone()
2503 .ok_or_else(|| LoopError::InvalidState("pending approval has no decision".into()))?;
2504
2505 match decision {
2506 ApprovalDecision::Approve => {
2507 use tracing::Instrument;
2508 self.emit(AgentEvent::ToolExecutionStarted(pending.call.clone()));
2509 let dispatch_span = self.execute_tool_span(
2510 &pending.tool_request,
2511 &pending.presentation_turn_id,
2512 "approved",
2513 );
2514 let cancellation = self
2515 .cancellation
2516 .as_ref()
2517 .map(CancellationHandle::checkpoint);
2518 self.register_tool_cancellation(&pending.call.id, cancellation.clone());
2519 let start = self
2520 .start_task_via_manager(
2521 Some(pending.task_id.clone()),
2522 pending.tool_request.clone(),
2523 TaskLaunchKind::Approved(pending.request.clone()),
2524 cancellation.clone(),
2525 )
2526 .instrument(dispatch_span.clone())
2527 .await;
2528 let outcome = match start {
2529 Ok(outcome) => outcome,
2530 Err(error) => {
2531 self.append_tool_result_item(Item {
2532 id: None,
2533 kind: ItemKind::Tool,
2534 parts: vec![Part::ToolResult(ToolResultPart {
2535 call_id: pending.call.id.clone(),
2536 output: ToolOutput::Text(format!(
2537 "approved task failed to start: {error}"
2538 )),
2539 is_error: true,
2540 metadata: pending.call.metadata.clone(),
2541 })],
2542 metadata: MetadataMap::new(),
2543 usage: None,
2544 finish_reason: None,
2545 created_at: None,
2546 });
2547 let turn_id = pending.tool_request.turn_id.clone();
2548 if let Err(cleanup_error) =
2549 self.task_manager.on_turn_interrupted(&turn_id).await
2550 {
2551 tracing::debug!(
2552 %cleanup_error,
2553 %turn_id,
2554 "failed to clean up turn after approved task start error"
2555 );
2556 }
2557 return Err(error);
2558 }
2559 };
2560 match outcome {
2561 TaskStartOutcome::Ready(resolution) => {
2562 let resolution = *resolution;
2563 if let TaskResolution::Item(item) = &resolution
2564 && tool_result_is_error(item)
2565 {
2566 dispatch_span.record("error.type", "tool_error");
2567 }
2568 if let Some(step) = self.queue_resolution_interrupt(
2569 &pending.presentation_turn_id,
2570 resolution,
2571 cancellation,
2572 ) {
2573 return Ok(step);
2574 }
2575 }
2576 TaskStartOutcome::Pending { kind, .. } => {
2577 if kind == agentkit_task_manager::TaskKind::Background {
2578 self.append_detach_placeholder(
2579 pending.call.id.clone(),
2580 &pending.call.name,
2581 );
2582 } else {
2583 self.active_tool_round = Some(ActiveToolRound {
2584 presentation_turn_id: pending.presentation_turn_id.clone(),
2585 task_turn_id: pending.tool_request.turn_id.clone(),
2586 pending_calls: VecDeque::new(),
2587 cancellation: cancellation.clone(),
2588 background_pending: false,
2589 foreground_progressed: false,
2590 });
2591 }
2592 }
2593 }
2594 }
2595 ApprovalDecision::Deny { reason } => {
2596 self.append_tool_result_item(Item {
2597 id: None,
2598 kind: ItemKind::Tool,
2599 parts: vec![Part::ToolResult(ToolResultPart {
2600 call_id: pending.call.id.clone(),
2601 output: ToolOutput::Text(
2602 reason.unwrap_or_else(|| "approval denied".into()),
2603 ),
2604 is_error: true,
2605 metadata: pending.call.metadata.clone(),
2606 })],
2607 metadata: MetadataMap::new(),
2608 usage: None,
2609 finish_reason: None,
2610 created_at: None,
2611 });
2612 }
2613 }
2614
2615 if let Some(step) = self.continue_active_tool_round().await? {
2616 Ok(step)
2617 } else if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2618 Ok(step)
2619 } else if let Some(step) = self.next_unresolved_approval_interrupt() {
2620 Ok(step)
2621 } else {
2622 self.drive_turn(pending.presentation_turn_id, MutationPoint::AfterToolResult)
2623 .await
2624 }
2625 }
2626
2627 fn finish_cancelled(
2628 &mut self,
2629 turn_id: agentkit_core::TurnId,
2630 items: Vec<Item>,
2631 ) -> Result<LoopStep, LoopError> {
2632 let pending = self.drain_pending_approval_items();
2633 self.reject_drained_approvals(pending);
2634 self.close_interrupted_tool_calls();
2635 self.extend_transcript(items.clone());
2636 let turn_result = TurnResult {
2637 turn_id,
2638 finish_reason: FinishReason::Cancelled,
2639 items,
2640 usage: None,
2641 metadata: interrupted_metadata("turn"),
2642 };
2643 self.finish_logical_turn(&turn_result);
2644 Ok(LoopStep::Finished(turn_result))
2645 }
2646
2647 pub fn submit_input(&mut self, input: Vec<Item>) -> Result<(), LoopError> {
2657 if self.has_pending_interrupts() {
2658 return Err(LoopError::InvalidState(
2659 "cannot submit input while an interrupt is pending".into(),
2660 ));
2661 }
2662 self.emit(AgentEvent::InputAccepted {
2663 session_id: self.session_id.clone(),
2664 items: input.clone(),
2665 });
2666 self.pending_input.extend(input);
2667 Ok(())
2668 }
2669
2670 pub fn set_next_turn_cache(&mut self, cache: PromptCacheRequest) -> Result<(), LoopError> {
2675 if self.has_pending_interrupts() {
2676 return Err(LoopError::InvalidState(
2677 "cannot update next-turn cache while an interrupt is pending".into(),
2678 ));
2679 }
2680 self.next_turn_cache = Some(cache);
2681 Ok(())
2682 }
2683
2684 #[cfg(test)]
2685 pub(crate) fn submit_input_with_cache(
2686 &mut self,
2687 input: Vec<Item>,
2688 cache: PromptCacheRequest,
2689 ) -> Result<(), LoopError> {
2690 self.set_next_turn_cache(cache)?;
2691 self.submit_input(input)
2692 }
2693
2694 pub fn resolve_approval_for(
2704 &mut self,
2705 call_id: ToolCallId,
2706 decision: ApprovalDecision,
2707 ) -> Result<(), LoopError> {
2708 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
2709 return Err(LoopError::InvalidState(format!(
2710 "no approval request is pending for call {}",
2711 call_id.0
2712 )));
2713 };
2714 pending.decision = Some(decision.clone());
2715 self.emit(AgentEvent::ApprovalResolved {
2716 approved: matches!(decision, ApprovalDecision::Approve),
2717 });
2718 Ok(())
2719 }
2720
2721 pub fn resolve_approval_for_with_patched_input(
2734 &mut self,
2735 call_id: ToolCallId,
2736 input: serde_json::Value,
2737 ) -> Result<(), LoopError> {
2738 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
2739 return Err(LoopError::InvalidState(format!(
2740 "no approval request is pending for call {}",
2741 call_id.0
2742 )));
2743 };
2744 pending.tool_request.input = input;
2745 self.resolve_approval_for(call_id, ApprovalDecision::Approve)
2746 }
2747
2748 pub fn resolve_approval(&mut self, decision: ApprovalDecision) -> Result<(), LoopError> {
2751 let mut unresolved = self
2752 .pending_approval_order
2753 .iter()
2754 .filter(|call_id| {
2755 self.pending_approvals
2756 .get(*call_id)
2757 .is_some_and(|pending| pending.decision.is_none())
2758 })
2759 .cloned();
2760 let Some(call_id) = unresolved.next() else {
2761 return Err(LoopError::InvalidState(
2762 "no approval request is pending".into(),
2763 ));
2764 };
2765 if unresolved.next().is_some() {
2766 return Err(LoopError::InvalidState(
2767 "multiple approvals are pending; use resolve_approval_for".into(),
2768 ));
2769 }
2770 self.resolve_approval_for(call_id, decision)
2771 }
2772
2773 pub fn cancel_pending_approval_for(&mut self, call_id: ToolCallId) -> Result<(), LoopError> {
2778 let Some(pending) = self.drain_pending_approval_for(&call_id) else {
2779 return Err(LoopError::InvalidState(format!(
2780 "no approval request is pending for call {}",
2781 call_id.0
2782 )));
2783 };
2784 let turn_id = pending.presentation_turn_id.clone();
2785 self.reject_drained_approvals(vec![pending]);
2786 if self.pending_approvals.is_empty() && self.active_tool_round.is_none() {
2787 let _ = self.finish_cancelled(turn_id, Vec::new())?;
2788 }
2789 Ok(())
2790 }
2791
2792 pub async fn cancel_pending_approvals(&mut self) -> Result<Option<LoopStep>, LoopError> {
2798 if self.pending_approvals.is_empty() {
2799 return Ok(None);
2800 }
2801 let Some(turn_id) = self
2802 .pending_approval_order
2803 .iter()
2804 .find_map(|call_id| self.pending_approvals.get(call_id))
2805 .map(|pending| pending.presentation_turn_id.clone())
2806 else {
2807 return Ok(None);
2808 };
2809 let mut seen_turns = HashSet::new();
2810 let mut originating_turns = Vec::new();
2811 for pending in self.pending_approvals.values() {
2812 let originating_turn = pending.tool_request.turn_id.clone();
2813 if seen_turns.insert(originating_turn.clone()) {
2814 originating_turns.push(originating_turn);
2815 }
2816 }
2817
2818 let pending = self.drain_pending_approval_items();
2819 self.active_tool_round = None;
2820 let mut cleanup_error = None;
2821 for originating_turn in originating_turns {
2822 if let Err(error) = self
2823 .task_manager
2824 .on_turn_interrupted(&originating_turn)
2825 .await
2826 && cleanup_error.is_none()
2827 {
2828 cleanup_error = Some(LoopError::Tool(ToolError::Internal(error.to_string())));
2829 }
2830 }
2831 self.reject_drained_approvals(pending);
2832 if let Some(error) = cleanup_error {
2833 self.close_interrupted_tool_calls();
2834 self.finish_logical_turn(&TurnResult {
2835 turn_id,
2836 finish_reason: FinishReason::Error,
2837 items: Vec::new(),
2838 usage: None,
2839 metadata: MetadataMap::new(),
2840 });
2841 return Err(error);
2842 }
2843 self.finish_cancelled(turn_id, Vec::new()).map(Some)
2844 }
2845
2846 pub fn snapshot(&self) -> LoopSnapshot {
2848 LoopSnapshot {
2849 session_id: self.session_id.clone(),
2850 transcript: self.transcript.clone(),
2851 pending_input: self.pending_input.clone(),
2852 }
2853 }
2854
2855 pub fn wait_for_loop_update(
2861 &self,
2862 ) -> impl std::future::Future<Output = Result<(), LoopError>> + Send + 'static {
2863 let has_collected_update = !self.pending_loop_updates.is_empty();
2864 let task_manager = self.task_manager.clone();
2865 async move {
2866 if has_collected_update {
2867 return Ok(());
2868 }
2869 task_manager
2870 .wait_for_loop_update()
2871 .await
2872 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))
2873 }
2874 }
2875
2876 pub async fn next(&mut self) -> Result<LoopStep, LoopError> {
2897 if self.lifecycle.active_turn.is_none() {
2898 let continuation_turn = self
2899 .pending_approval_order
2900 .iter()
2901 .find_map(|call_id| self.pending_approvals.get(call_id))
2902 .map(|pending| pending.presentation_turn_id.clone())
2903 .or_else(|| {
2904 self.active_tool_round
2905 .as_ref()
2906 .map(|active| active.presentation_turn_id.clone())
2907 })
2908 .or_else(|| self.pending_round_resume.clone());
2909 if let Some(turn_id) = continuation_turn {
2910 self.start_logical_turn_with(turn_id);
2911 } else if !self.pending_input.is_empty() {
2912 self.start_logical_turn();
2913 }
2914 }
2915
2916 let result = self.next_inner().await;
2917 match &result {
2918 Ok(LoopStep::Finished(turn)) => self.finish_logical_turn(turn),
2919 Err(_) => {
2920 if let Some(turn_id) = self.lifecycle.active_turn.clone() {
2921 self.recover_from_next_error().await;
2922 self.finish_logical_turn(&TurnResult {
2923 turn_id,
2924 finish_reason: FinishReason::Error,
2925 items: Vec::new(),
2926 usage: None,
2927 metadata: MetadataMap::new(),
2928 });
2929 }
2930 }
2931 _ => {}
2932 }
2933 result
2934 }
2935
2936 async fn recover_from_next_error(&mut self) {
2937 let mut seen_turns = HashSet::new();
2938 let mut interrupted_turns = Vec::new();
2939 if let Some(active) = self.active_tool_round.take()
2940 && seen_turns.insert(active.task_turn_id.clone())
2941 {
2942 interrupted_turns.push(active.task_turn_id);
2943 }
2944 if let Some(turn_id) = self.pending_round_resume.take()
2945 && seen_turns.insert(turn_id.clone())
2946 {
2947 interrupted_turns.push(turn_id);
2948 }
2949 for pending in self.pending_approvals.values() {
2950 let turn_id = pending.tool_request.turn_id.clone();
2951 if seen_turns.insert(turn_id.clone()) {
2952 interrupted_turns.push(turn_id);
2953 }
2954 }
2955
2956 let pending = self.drain_pending_approval_items();
2957 for turn_id in interrupted_turns {
2958 if let Err(error) = self.task_manager.on_turn_interrupted(&turn_id).await {
2959 tracing::debug!(%error, %turn_id, "failed to clean up turn after loop error");
2960 }
2961 }
2962 self.reject_drained_approvals(pending);
2963 self.close_interrupted_tool_calls();
2964 }
2965
2966 async fn next_inner(&mut self) -> Result<LoopStep, LoopError> {
2967 if let Some(pending) = self.take_next_resolved_approval() {
2968 return self.resume_after_approval(pending).await;
2969 }
2970
2971 if let Some(step) = self.finish_cancelled_pending_approval().await? {
2972 return Ok(step);
2973 }
2974
2975 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2976 return Ok(step);
2977 }
2978
2979 if let Some(step) = self.next_unresolved_approval_interrupt() {
2980 return Ok(step);
2981 }
2982
2983 if let Some(step) = self.continue_active_tool_round().await? {
2984 return Ok(step);
2985 }
2986
2987 if self.pending_round_resume.is_none() && !self.pending_input.is_empty() {
2992 self.collect_pending_loop_updates().await?;
2995 let turn_id = self.start_logical_turn();
2996 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
2997 self.extend_transcript(drained);
2998 return self
2999 .drive_turn(turn_id, MutationPoint::AfterTurnEnded)
3000 .await;
3001 }
3002
3003 let (had_loop_updates, loop_step) = self.drain_pending_loop_updates().await?;
3004 if let Some(step) = loop_step {
3005 return Ok(step);
3006 }
3007
3008 if let Some(turn_id) = self.pending_round_resume.take() {
3013 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
3014 self.extend_transcript(drained);
3015 return self
3016 .drive_turn(turn_id, MutationPoint::AfterToolResult)
3017 .await;
3018 }
3019
3020 if self.pending_input.is_empty() && !had_loop_updates {
3021 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
3022 InputRequest {
3023 session_id: self.session_id.clone(),
3024 reason: "driver is waiting for input".into(),
3025 },
3026 )));
3027 }
3028
3029 let turn_id = self.start_logical_turn();
3030 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
3031 self.extend_transcript(drained);
3032 self.drive_turn(turn_id, MutationPoint::AfterTurnEnded)
3033 .await
3034 }
3035
3036 fn emit(&self, event: AgentEvent) {
3037 fan_out_observed_event(&self.observers, &self.observed_session_id, event);
3038 }
3039
3040 fn append_item(&mut self, mut item: Item) {
3045 if item.created_at.is_none() {
3046 item.created_at = Some(Timestamp::now());
3047 }
3048 for observer in &self.transcript_observers {
3049 observer.on_transcript_event(TranscriptEvent {
3050 session_id: &self.session_id,
3051 item: &item,
3052 });
3053 }
3054 self.transcript.push(item);
3055 }
3056
3057 fn append_detach_placeholder(&mut self, call_id: ToolCallId, tool_name: &str) {
3058 self.background_call_ids.insert(call_id.clone());
3059 if !self.detached_call_ids.insert(call_id.clone()) {
3060 return;
3061 }
3062 let detached_result = ToolResultPart {
3063 call_id: call_id.clone(),
3064 output: ToolOutput::Text(format!(
3065 "Tool {tool_name} is now running in the background. The result will be delivered when it completes."
3066 )),
3067 is_error: false,
3068 metadata: MetadataMap::new(),
3069 };
3070 self.emit(AgentEvent::ToolExecutionProgress(detached_result.clone()));
3071 self.append_item(Item {
3072 id: None,
3073 kind: ItemKind::Tool,
3074 parts: vec![Part::ToolResult(detached_result)],
3075 metadata: MetadataMap::new(),
3076 usage: None,
3077 finish_reason: None,
3078 created_at: None,
3079 });
3080 }
3081
3082 fn append_tool_result_item(&mut self, item: Item) {
3096 for part in &item.parts {
3097 if let Part::ToolResult(result) = part {
3098 if !self
3099 .interrupted_background_call_ids
3100 .contains(&result.call_id)
3101 {
3102 self.emit(AgentEvent::ToolResultReceived(result.clone()));
3103 }
3104 self.background_call_ids.remove(&result.call_id);
3105 self.clear_tool_cancellation(&result.call_id);
3106 }
3107 }
3108 let item = self.maybe_convert_detached(item);
3109 self.append_item(item);
3110 }
3111
3112 fn drain_pending_approval_for(
3113 &mut self,
3114 call_id: &ToolCallId,
3115 ) -> Option<PendingApprovalToolCall> {
3116 let pending = self.pending_approvals.remove(call_id)?;
3117 self.pending_approval_order.retain(|id| id != call_id);
3118 self.clear_tool_cancellation(call_id);
3119 Some(pending)
3120 }
3121
3122 fn drain_pending_approval_items(&mut self) -> Vec<PendingApprovalToolCall> {
3123 let order = std::mem::take(&mut self.pending_approval_order);
3124 let pending = order
3125 .iter()
3126 .filter_map(|call_id| {
3127 let pending = self.pending_approvals.remove(call_id);
3128 self.clear_tool_cancellation(call_id);
3129 pending
3130 })
3131 .collect();
3132 self.pending_approvals.clear();
3133 pending
3134 }
3135
3136 fn reject_drained_approvals(&mut self, pending: Vec<PendingApprovalToolCall>) {
3137 for pending in pending {
3138 self.emit(AgentEvent::ApprovalResolved { approved: false });
3139 self.append_tool_result_item(cancelled_approval_item(pending));
3140 }
3141 }
3142
3143 fn close_interrupted_tool_calls(&mut self) {
3160 for call in unanswered_tool_calls(&self.transcript) {
3161 let call_id = call.id.clone();
3162 let completes_in_background = self.background_call_ids.contains(&call_id);
3163 self.append_tool_result_item(interrupted_tool_result_item(call));
3164 if completes_in_background {
3165 self.detached_call_ids.insert(call_id.clone());
3166 self.interrupted_background_call_ids.insert(call_id);
3167 }
3168 }
3169 }
3170
3171 fn maybe_convert_detached(&mut self, mut item: Item) -> Item {
3172 if !matches!(item.kind, ItemKind::Tool) {
3173 return item;
3174 }
3175 let results: Vec<&ToolResultPart> = item
3176 .parts
3177 .iter()
3178 .filter_map(|p| match p {
3179 Part::ToolResult(r) => Some(r),
3180 _ => None,
3181 })
3182 .collect();
3183 if results.is_empty()
3184 || !results
3185 .iter()
3186 .all(|r| self.detached_call_ids.contains(&r.call_id))
3187 {
3188 return item;
3189 }
3190 let structured_results = results
3191 .iter()
3192 .map(|result| {
3193 Part::structured(serde_json::to_value(result).unwrap_or_else(
3194 |error| serde_json::json!({ "serialization_error": error.to_string() }),
3195 ))
3196 })
3197 .collect::<Vec<_>>();
3198 let failed = results.iter().filter(|result| result.is_error).count();
3199 let with_metadata = results
3200 .iter()
3201 .filter(|result| !result.metadata.is_empty())
3202 .count();
3203 let mut text = format!(
3204 "Background tool results: {} total, {failed} failed, {with_metadata} with metadata. ",
3205 results.len()
3206 );
3207 for (index, result) in results.iter().enumerate() {
3208 self.detached_call_ids.remove(&result.call_id);
3209 self.interrupted_background_call_ids.remove(&result.call_id);
3210 if text.chars().count() >= DETACHED_NOTIFICATION_TEXT_MAX_CHARS {
3211 continue;
3212 }
3213 if index > 0 {
3214 text.push_str("; ");
3215 }
3216 let label = if result.is_error {
3217 "failed"
3218 } else {
3219 "completed"
3220 };
3221 let call_id = truncate_chars(&result.call_id.0, DETACHED_CALL_ID_MAX_CHARS);
3222 let body = render_tool_output_brief(&result.output);
3223 text.push_str(&format!("{call_id} {label}: {body}"));
3224 }
3225 let text = truncate_chars(&text, DETACHED_NOTIFICATION_TEXT_MAX_CHARS);
3226 let mut notification_parts = Vec::with_capacity(1 + structured_results.len());
3227 notification_parts.push(Part::text(text));
3228 notification_parts.extend(structured_results);
3229 item.kind = ItemKind::Notification;
3230 item.parts = notification_parts;
3231 item
3232 }
3233
3234 fn extend_transcript(&mut self, items: impl IntoIterator<Item = Item>) {
3238 let now = Timestamp::now();
3239 for mut item in items {
3240 if item.created_at.is_none() {
3241 item.created_at = Some(now);
3242 }
3243 self.append_item(item);
3244 }
3245 }
3246}
3247
3248fn render_tool_output_brief(output: &ToolOutput) -> String {
3249 match output {
3250 ToolOutput::Text(text) => format!(
3251 "text preview: {}",
3252 truncate_chars(text, DETACHED_TEXT_PREVIEW_MAX_CHARS)
3253 ),
3254 ToolOutput::Structured(_) => "structured payload".into(),
3255 ToolOutput::Parts(parts) => format!("parts payload ({} parts)", parts.len()),
3256 ToolOutput::Files(files) => format!("files payload ({} files)", files.len()),
3257 }
3258}
3259
3260fn truncate_chars(text: &str, max_chars: usize) -> String {
3261 let mut chars = text.chars();
3262 let mut truncated = chars.by_ref().take(max_chars).collect::<String>();
3263 if chars.next().is_some() && max_chars > 0 {
3264 truncated.pop();
3265 truncated.push('…');
3266 }
3267 truncated
3268}
3269
3270fn interrupted_metadata(stage: &str) -> MetadataMap {
3271 let mut metadata = MetadataMap::new();
3272 metadata.insert(INTERRUPTED_METADATA_KEY.into(), true.into());
3273 metadata.insert(
3274 INTERRUPT_REASON_METADATA_KEY.into(),
3275 USER_CANCELLED_REASON.into(),
3276 );
3277 metadata.insert(INTERRUPT_STAGE_METADATA_KEY.into(), stage.into());
3278 metadata
3279}
3280
3281fn record_token_attribute(span: &tracing::Span, key: &'static str, value: u64) {
3282 match i64::try_from(value) {
3283 Ok(value) => record_i64_attribute(span, key, value),
3284 Err(_) => tracing::warn!(attribute = key, value, "token count exceeds OTEL i64 range"),
3285 }
3286}
3287
3288#[cfg(feature = "otel")]
3289fn record_i64_attribute(span: &tracing::Span, key: &'static str, value: i64) {
3290 use tracing_opentelemetry::OpenTelemetrySpanExt;
3291 span.set_attribute(key, value);
3292}
3293
3294#[cfg(not(feature = "otel"))]
3295fn record_i64_attribute(span: &tracing::Span, key: &'static str, value: i64) {
3296 span.record(key, value);
3297}
3298
3299#[cfg(feature = "otel")]
3300fn record_f64_attribute(span: &tracing::Span, key: &'static str, value: f64) {
3301 use tracing_opentelemetry::OpenTelemetrySpanExt;
3302 span.set_attribute(key, value);
3303}
3304
3305#[cfg(not(feature = "otel"))]
3306fn record_f64_attribute(span: &tracing::Span, key: &'static str, value: f64) {
3307 span.record(key, value);
3308}
3309
3310#[cfg(feature = "otel")]
3311fn otel_string_array(values: Vec<String>) -> opentelemetry::Value {
3312 use opentelemetry::{Array, StringValue, Value as OtelValue};
3313 OtelValue::Array(Array::String(
3314 values.into_iter().map(StringValue::from).collect(),
3315 ))
3316}
3317
3318#[cfg(feature = "otel")]
3319fn record_string_array_attribute(span: &tracing::Span, key: &'static str, values: Vec<String>) {
3320 use tracing_opentelemetry::OpenTelemetrySpanExt;
3321 span.set_attribute(key, otel_string_array(values));
3322}
3323
3324#[cfg(not(feature = "otel"))]
3325fn record_string_array_attribute(span: &tracing::Span, key: &'static str, values: Vec<String>) {
3326 span.record(key, tracing::field::debug(&values));
3327}
3328
3329#[derive(Clone, Copy)]
3330enum CaptureOrder {
3331 NewestTail,
3332 OldestHead,
3333}
3334
3335fn effective_provider_name(
3336 session_provider: Option<&str>,
3337 adapter_provider: Option<&str>,
3338) -> Option<String> {
3339 session_provider.or(adapter_provider).map(str::to_owned)
3340}
3341
3342fn merge_usage(final_usage: Option<Usage>, streamed_usage: Option<Usage>) -> Option<Usage> {
3343 match (final_usage, streamed_usage) {
3344 (None, streamed) => streamed,
3345 (Some(final_usage), None) => Some(final_usage),
3346 (Some(mut final_usage), Some(streamed)) => {
3347 if final_usage.tokens.is_none() {
3348 final_usage.tokens = streamed.tokens;
3349 }
3350 if final_usage.cost.is_none() {
3351 final_usage.cost = streamed.cost;
3352 }
3353 for (key, value) in streamed.metadata {
3354 final_usage.metadata.entry(key).or_insert(value);
3355 }
3356 Some(final_usage)
3357 }
3358 }
3359}
3360
3361fn capture_messages(items: &[Item], capture: MessageCapture, order: CaptureOrder) -> Vec<String> {
3362 let mut captured = Vec::new();
3363 let mut used_bytes = 0usize;
3364 let indices: Box<dyn Iterator<Item = usize>> = match order {
3365 CaptureOrder::NewestTail => Box::new((0..items.len()).rev()),
3366 CaptureOrder::OldestHead => Box::new(0..items.len()),
3367 };
3368
3369 for index in indices.take(capture.max_messages) {
3370 let item = &items[index];
3371 let original_bytes = source_content_bytes(item);
3372 let remaining = capture.max_bytes.saturating_sub(used_bytes);
3373 if original_bytes > remaining {
3374 captured.push(
3375 serde_json::json!({
3376 "type": "truncated",
3377 "original_bytes": original_bytes,
3378 })
3379 .to_string(),
3380 );
3381 break;
3382 }
3383 used_bytes += original_bytes;
3384 captured.push(capture_item_json(item, remaining));
3385 }
3386
3387 if matches!(order, CaptureOrder::NewestTail) {
3388 captured.reverse();
3389 }
3390 captured
3391}
3392
3393fn source_content_bytes(item: &Item) -> usize {
3394 item.parts.iter().fold(0, |total, part| {
3395 total.saturating_add(part_source_content_bytes(part))
3396 })
3397}
3398
3399fn part_source_content_bytes(part: &Part) -> usize {
3400 match part {
3401 Part::Text(text) => text.text.len(),
3402 Part::Media(media) => media.mime_type.len(),
3403 Part::File(file) => file
3404 .name
3405 .as_deref()
3406 .map_or(0, str::len)
3407 .saturating_add(file.mime_type.as_deref().map_or(0, str::len)),
3408 Part::Structured(_) => 0,
3409 Part::Reasoning(reasoning) => reasoning.summary.as_deref().map_or(0, str::len),
3410 Part::ToolCall(call) => call.id.0.len().saturating_add(call.name.len()),
3411 Part::ToolResult(result) => result
3412 .call_id
3413 .0
3414 .len()
3415 .saturating_add(tool_output_source_content_bytes(&result.output)),
3416 Part::Custom(custom) => custom.kind.len(),
3417 }
3418}
3419
3420fn tool_output_source_content_bytes(output: &ToolOutput) -> usize {
3421 match output {
3422 ToolOutput::Text(text) => text.len(),
3423 ToolOutput::Structured(_) | ToolOutput::Parts(_) | ToolOutput::Files(_) => 0,
3424 }
3425}
3426
3427struct CaptureBudget {
3428 remaining: usize,
3429}
3430
3431impl CaptureBudget {
3432 fn text(&mut self, text: &str) -> (String, bool) {
3433 let end = floor_char_boundary(text, self.remaining.min(text.len()));
3434 self.remaining = self.remaining.saturating_sub(end);
3435 (text[..end].to_owned(), end < text.len())
3436 }
3437}
3438
3439fn floor_char_boundary(text: &str, mut index: usize) -> usize {
3440 while index > 0 && !text.is_char_boundary(index) {
3441 index -= 1;
3442 }
3443 index
3444}
3445
3446const MAX_CAPTURED_PARTS_PER_ITEM: usize = 256;
3447
3448fn capture_item_json(item: &Item, max_bytes: usize) -> String {
3449 let mut budget = CaptureBudget {
3450 remaining: max_bytes,
3451 };
3452 let mut parts = item
3453 .parts
3454 .iter()
3455 .take(MAX_CAPTURED_PARTS_PER_ITEM)
3456 .map(|part| sanitized_part(part, &mut budget))
3457 .collect::<Vec<_>>();
3458 if item.parts.len() > parts.len() {
3459 parts.push(serde_json::json!({
3460 "type": "truncated",
3461 "reason": "part_limit",
3462 }));
3463 }
3464 serde_json::json!({
3465 "role": item_kind_name(item.kind),
3466 "parts": parts,
3467 })
3468 .to_string()
3469}
3470
3471fn item_kind_name(kind: ItemKind) -> &'static str {
3472 match kind {
3473 ItemKind::System => "system",
3474 ItemKind::Developer => "developer",
3475 ItemKind::User => "user",
3476 ItemKind::Assistant => "assistant",
3477 ItemKind::Tool => "tool",
3478 ItemKind::Context => "context",
3479 ItemKind::Notification => "notification",
3480 }
3481}
3482
3483fn modality_name(modality: Modality) -> &'static str {
3484 match modality {
3485 Modality::Audio => "audio",
3486 Modality::Image => "image",
3487 Modality::Video => "video",
3488 Modality::Binary => "binary",
3489 }
3490}
3491
3492fn omitted_data_ref(data: &DataRef) -> Value {
3493 let kind = match data {
3494 DataRef::InlineText(_) => "inline_text",
3495 DataRef::InlineBytes(_) => "inline_bytes",
3496 DataRef::Uri(_) => "uri",
3497 DataRef::Handle(_) => "handle",
3498 };
3499 serde_json::json!({ "kind": kind, "omitted": true })
3500}
3501
3502fn bounded_field(text: &str, budget: &mut CaptureBudget) -> Value {
3503 let (text, truncated) = budget.text(text);
3504 serde_json::json!({ "value": text, "truncated": truncated })
3505}
3506
3507fn sanitized_part(part: &Part, budget: &mut CaptureBudget) -> Value {
3508 match part {
3509 Part::Text(text) => serde_json::json!({
3510 "type": "text",
3511 "text": bounded_field(&text.text, budget),
3512 }),
3513 Part::Media(media) => serde_json::json!({
3514 "type": "media",
3515 "modality": modality_name(media.modality),
3516 "mime_type": bounded_field(&media.mime_type, budget),
3517 "data": omitted_data_ref(&media.data),
3518 }),
3519 Part::File(file) => serde_json::json!({
3520 "type": "file",
3521 "name": file.name.as_deref().map(|name| bounded_field(name, budget)),
3522 "mime_type": file.mime_type.as_deref().map(|mime| bounded_field(mime, budget)),
3523 "data": omitted_data_ref(&file.data),
3524 }),
3525 Part::Structured(_) => serde_json::json!({
3526 "type": "structured",
3527 "truncated": true,
3528 }),
3529 Part::Reasoning(reasoning) => serde_json::json!({
3530 "type": "reasoning",
3531 "summary": reasoning.summary.as_deref().map(|summary| bounded_field(summary, budget)),
3532 "redacted": reasoning.redacted,
3533 "data": reasoning.data.as_ref().map(omitted_data_ref),
3534 }),
3535 Part::ToolCall(call) => serde_json::json!({
3536 "type": "tool_call",
3537 "id": bounded_field(&call.id.0, budget),
3538 "name": bounded_field(&call.name, budget),
3539 "input": { "truncated": true },
3540 }),
3541 Part::ToolResult(result) => serde_json::json!({
3542 "type": "tool_result",
3543 "call_id": bounded_field(&result.call_id.0, budget),
3544 "is_error": result.is_error,
3545 "output": sanitized_tool_output(&result.output, budget),
3546 }),
3547 Part::Custom(custom) => serde_json::json!({
3548 "type": "custom",
3549 "kind": bounded_field(&custom.kind, budget),
3550 "data": custom.data.as_ref().map(omitted_data_ref),
3551 "value": custom.value.as_ref().map(|_| serde_json::json!({ "truncated": true })),
3552 }),
3553 }
3554}
3555
3556fn sanitized_tool_output(output: &ToolOutput, budget: &mut CaptureBudget) -> Value {
3557 match output {
3558 ToolOutput::Text(text) => serde_json::json!({
3559 "type": "text",
3560 "text": bounded_field(text, budget),
3561 }),
3562 ToolOutput::Structured(_) => serde_json::json!({
3563 "type": "structured",
3564 "truncated": true,
3565 }),
3566 ToolOutput::Parts(parts) => serde_json::json!({
3567 "type": "parts",
3568 "count": parts.len(),
3569 "truncated": true,
3570 }),
3571 ToolOutput::Files(files) => serde_json::json!({
3572 "type": "files",
3573 "count": files.len(),
3574 "truncated": true,
3575 }),
3576 }
3577}
3578
3579#[cfg(test)]
3580mod telemetry_tests {
3581 use super::*;
3582
3583 #[test]
3584 fn message_capture_is_off_by_default_and_independent() {
3585 let default = TelemetryConfig::default();
3586 assert_eq!(default.input_messages(), None);
3587 assert_eq!(default.output_messages(), None);
3588
3589 let capture = MessageCapture::new(2, 1).unwrap();
3590 let input_only = TelemetryConfig::default().with_input_messages(capture);
3591 assert_eq!(input_only.input_messages().unwrap().max_messages(), 2);
3592 assert_eq!(input_only.input_messages().unwrap().max_bytes(), 1);
3593 assert_eq!(input_only.output_messages(), None);
3594 assert_eq!(
3595 MessageCapture::new(0, 1),
3596 Err(MessageCaptureError::ZeroMessages)
3597 );
3598 assert_eq!(
3599 MessageCapture::new(1, 0),
3600 Err(MessageCaptureError::ZeroBytes)
3601 );
3602 }
3603
3604 #[test]
3605 fn one_source_byte_is_not_rejected_for_json_envelope_overhead() {
3606 let items = vec![Item::text(ItemKind::User, "x")];
3607 let captured = capture_messages(
3608 &items,
3609 MessageCapture::new(1, 1).unwrap(),
3610 CaptureOrder::OldestHead,
3611 );
3612 assert_eq!(captured.len(), 1);
3613 let value: Value = serde_json::from_str(&captured[0]).unwrap();
3614 assert_eq!(value["role"], "user");
3615 assert_eq!(value["parts"][0]["text"]["value"], "x");
3616 assert_eq!(value["parts"][0]["text"]["truncated"], false);
3617 }
3618
3619 #[test]
3620 fn multibyte_source_accounting_preserves_utf8_boundaries() {
3621 let items = vec![Item::text(ItemKind::User, "é")];
3622
3623 let exact = capture_messages(
3624 &items,
3625 MessageCapture::new(1, "é".len()).unwrap(),
3626 CaptureOrder::OldestHead,
3627 );
3628 let value: Value = serde_json::from_str(&exact[0]).unwrap();
3629 assert_eq!(value["parts"][0]["text"]["value"], "é");
3630 assert_eq!(value["parts"][0]["text"]["truncated"], false);
3631
3632 let too_small = capture_messages(
3633 &items,
3634 MessageCapture::new(1, 1).unwrap(),
3635 CaptureOrder::OldestHead,
3636 );
3637 let value: Value = serde_json::from_str(&too_small[0]).unwrap();
3638 assert_eq!(value["type"], "truncated");
3639 assert_eq!(value["original_bytes"], 2);
3640 }
3641
3642 #[test]
3643 fn source_bytes_are_aggregated_without_charging_json_envelopes() {
3644 let items = vec![
3645 Item::text(ItemKind::User, "a"),
3646 Item::text(ItemKind::Assistant, "b"),
3647 Item::text(ItemKind::User, "cd"),
3648 ];
3649 let captured = capture_messages(
3650 &items,
3651 MessageCapture::new(3, 3).unwrap(),
3652 CaptureOrder::OldestHead,
3653 );
3654 assert_eq!(captured.len(), 3);
3655 let values = captured
3656 .iter()
3657 .map(|encoded| serde_json::from_str::<Value>(encoded).unwrap())
3658 .collect::<Vec<_>>();
3659 assert_eq!(values[0]["parts"][0]["text"]["value"], "a");
3660 assert_eq!(values[1]["parts"][0]["text"]["value"], "b");
3661 assert_eq!(values[2]["type"], "truncated");
3662 assert_eq!(values[2]["original_bytes"], 2);
3663 }
3664
3665 #[test]
3666 fn tiny_limits_emit_valid_structured_source_byte_truncation() {
3667 let items = vec![Item::text(ItemKind::User, "x".repeat(1_000))];
3668 let captured = capture_messages(
3669 &items,
3670 MessageCapture::new(1, 1).unwrap(),
3671 CaptureOrder::OldestHead,
3672 );
3673 assert_eq!(captured.len(), 1);
3674 let value: Value = serde_json::from_str(&captured[0]).unwrap();
3675 assert_eq!(value["type"], "truncated");
3676 assert_eq!(value["original_bytes"], 1_000);
3677 }
3678
3679 #[test]
3680 fn provider_finish_reason_metadata_round_trips() {
3681 let mut metadata = MetadataMap::new();
3682 set_provider_finish_reasons(&mut metadata, ["end_turn", "", "tool_use", "end_turn"]);
3683 assert_eq!(
3684 provider_finish_reasons(&metadata, &FinishReason::Completed),
3685 vec!["end_turn", "tool_use"]
3686 );
3687 set_provider_finish_reasons(&mut metadata, std::iter::empty::<String>());
3688 assert!(!metadata.contains_key(PROVIDER_FINISH_REASONS_METADATA_KEY));
3689 let fallbacks = [
3690 (FinishReason::Completed, "completed"),
3691 (FinishReason::ToolCall, "tool_call"),
3692 (FinishReason::MaxTokens, "max_tokens"),
3693 (FinishReason::Cancelled, "cancelled"),
3694 (FinishReason::Blocked, "blocked"),
3695 (FinishReason::Error, "error"),
3696 (FinishReason::Other("native".into()), "native"),
3697 ];
3698 for (reason, expected) in fallbacks {
3699 assert_eq!(
3700 provider_finish_reasons(&MetadataMap::new(), &reason),
3701 [expected]
3702 );
3703 }
3704 }
3705
3706 #[test]
3707 fn input_is_newest_tail_output_is_head_and_data_refs_are_omitted() {
3708 let items = vec![
3709 Item::text(ItemKind::User, "old"),
3710 Item::text(ItemKind::User, "middle"),
3711 Item::text(ItemKind::User, "new"),
3712 ];
3713 let capture = MessageCapture::new(2, 10_000).unwrap();
3714 let input = capture_messages(&items, capture, CaptureOrder::NewestTail);
3715 assert!(input[0].contains("middle"));
3716 assert!(input[1].contains("new"));
3717 let output = capture_messages(&items, capture, CaptureOrder::OldestHead);
3718 assert!(output[0].contains("old"));
3719 assert!(output[1].contains("middle"));
3720
3721 let media = Item::new(
3722 ItemKind::User,
3723 vec![Part::media(
3724 agentkit_core::Modality::Image,
3725 "image/png",
3726 agentkit_core::DataRef::uri("https://secret.invalid/image.png"),
3727 )],
3728 );
3729 let encoded = capture_item_json(&media, 10_000);
3730 assert!(!encoded.contains("secret.invalid"));
3731 assert!(encoded.contains("omitted"));
3732 }
3733
3734 #[test]
3735 fn final_usage_wins_and_streamed_usage_fills_only_missing_fields() {
3736 let mut final_metadata = MetadataMap::new();
3737 final_metadata.insert("shared".into(), serde_json::json!("final"));
3738 let final_usage = Usage {
3739 tokens: Some(agentkit_core::TokenUsage::new(1, 2)),
3740 cost: None,
3741 metadata: final_metadata,
3742 };
3743 let mut streamed_metadata = MetadataMap::new();
3744 streamed_metadata.insert("shared".into(), serde_json::json!("streamed"));
3745 streamed_metadata.insert("stream_only".into(), serde_json::json!(true));
3746 let streamed_usage = Usage {
3747 tokens: Some(agentkit_core::TokenUsage::new(10, 20)),
3748 cost: Some(agentkit_core::CostUsage::new(0.5, "USD")),
3749 metadata: streamed_metadata,
3750 };
3751 let merged = merge_usage(Some(final_usage), Some(streamed_usage)).unwrap();
3752 assert_eq!(merged.tokens.unwrap().input_tokens, 1);
3753 assert_eq!(merged.cost.unwrap().amount, 0.5);
3754 assert_eq!(merged.metadata["shared"], "final");
3755 assert_eq!(merged.metadata["stream_only"], true);
3756 }
3757
3758 #[test]
3759 fn session_provider_precedes_adapter_fallback() {
3760 assert_eq!(
3761 effective_provider_name(Some("session"), Some("adapter")).as_deref(),
3762 Some("session")
3763 );
3764 assert_eq!(
3765 effective_provider_name(None, Some("adapter")).as_deref(),
3766 Some("adapter")
3767 );
3768 }
3769}
3770
3771#[cfg(all(test, feature = "otel"))]
3772mod true_otel_integration_tests {
3773 use std::fs;
3774 use std::process::Command;
3775 use std::sync::atomic::{AtomicU64, Ordering};
3776 use std::time::{SystemTime, UNIX_EPOCH};
3777
3778 static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0);
3779
3780 #[test]
3781 fn actual_layer_exports_driven_loop_spans() {
3782 let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
3783 .parent()
3784 .unwrap()
3785 .parent()
3786 .unwrap();
3787 let nonce = SystemTime::now()
3788 .duration_since(UNIX_EPOCH)
3789 .unwrap()
3790 .as_nanos();
3791 let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed);
3792 let temp = std::env::temp_dir().join(format!(
3793 "agentkit-loop-true-otel-{}-{nonce}-{sequence}",
3794 std::process::id()
3795 ));
3796 fs::create_dir(&temp).unwrap();
3797 fs::create_dir(temp.join("src")).unwrap();
3798 let core_path = toml_path(&root.join("crates/agentkit-core"));
3799 let loop_path = toml_path(&root.join("crates/agentkit-loop"));
3800 let async_trait_version = locked_version(root, "async-trait");
3803 let opentelemetry_version = locked_version(root, "opentelemetry");
3804 let serde_json_version = locked_version(root, "serde_json");
3805 let tokio_version = locked_version(root, "tokio");
3806 let tracing_version = locked_version(root, "tracing");
3807 let tracing_otel_version = locked_version(root, "tracing-opentelemetry");
3808 let tracing_subscriber_version = locked_version(root, "tracing-subscriber");
3809 let manifest = format!(
3810 r#"[package]
3811name = "agentkit-loop-true-otel-test"
3812version = "0.0.0"
3813edition = "2024"
3814
3815[dependencies]
3816agentkit-core = {{ path = {core_path} }}
3817agentkit-loop = {{ path = {loop_path}, features = ["otel"] }}
3818async-trait = "={async_trait_version}"
3819opentelemetry = {{ version = "={opentelemetry_version}", default-features = false }}
3820serde_json = "={serde_json_version}"
3821tokio = {{ version = "={tokio_version}", features = ["rt"] }}
3822tracing = "={tracing_version}"
3823tracing-opentelemetry = {{ version = "={tracing_otel_version}", default-features = false }}
3824tracing-subscriber = "={tracing_subscriber_version}"
3825"#
3826 );
3827 fs::write(temp.join("Cargo.toml"), manifest).unwrap();
3828 fs::write(temp.join("src/main.rs"), TRUE_OTEL_HARNESS).unwrap();
3829
3830 let output = Command::new(env!("CARGO"))
3831 .args(["run", "--quiet", "--offline"])
3832 .current_dir(&temp)
3833 .env("CARGO_TARGET_DIR", temp.join("target"))
3834 .output()
3835 .unwrap();
3836 let _ = fs::remove_dir_all(&temp);
3837 assert!(
3838 output.status.success(),
3839 "true OTEL harness failed:\nstdout:\n{}\nstderr:\n{}",
3840 String::from_utf8_lossy(&output.stdout),
3841 String::from_utf8_lossy(&output.stderr)
3842 );
3843 }
3844
3845 fn locked_version(root: &std::path::Path, name: &str) -> String {
3846 let lock = fs::read_to_string(root.join("Cargo.lock")).unwrap();
3847 let expected_name = format!("name = {}", serde_json::to_string(name).unwrap());
3848 let versions = lock
3849 .split("[[package]]")
3850 .filter(|package| {
3851 package
3852 .lines()
3853 .any(|line| line.trim() == expected_name.as_str())
3854 })
3855 .filter_map(|package| {
3856 package.lines().find_map(|line| {
3857 line.trim()
3858 .strip_prefix("version = \"")
3859 .and_then(|version| version.strip_suffix('\"'))
3860 .map(str::to_owned)
3861 })
3862 })
3863 .collect::<Vec<_>>();
3864 assert_eq!(versions.len(), 1, "expected one locked version for {name}");
3865 versions.into_iter().next().unwrap()
3866 }
3867
3868 fn toml_path(path: &std::path::Path) -> String {
3869 serde_json::to_string(&path.to_string_lossy()).unwrap()
3870 }
3871
3872 #[test]
3873 fn toml_path_escapes_windows_separators_and_quotes() {
3874 let encoded = toml_path(std::path::Path::new(r#"C:\Users\name\quoted\"dir"#));
3875 assert_eq!(encoded, r#""C:\\Users\\name\\quoted\\\"dir""#);
3876 }
3877
3878 const TRUE_OTEL_HARNESS: &str = r#"
3879use std::borrow::Cow;
3880use std::collections::VecDeque;
3881use std::sync::{Arc, Mutex};
3882use std::time::SystemTime;
3883use agentkit_core::{CostUsage, DataRef, FinishReason, Item, ItemKind, MetadataMap, Modality, Part, TokenUsage, TurnCancellation, Usage};
3884use agentkit_loop::{Agent, LoopError, LoopStep, MessageCapture, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, TelemetryConfig, TurnRequest, set_provider_finish_reasons};
3885use async_trait::async_trait;
3886use opentelemetry::trace::{Span, SpanBuilder, SpanContext, Status, Tracer};
3887use opentelemetry::{Array, Context, KeyValue, Value};
3888use tracing_subscriber::layer::SubscriberExt;
3889
3890#[derive(Clone, Debug)]
3891struct Exported { name: String, attributes: Vec<KeyValue> }
3892#[derive(Clone, Default)]
3893struct MemoryTracer { exported: Arc<Mutex<Vec<Exported>>> }
3894struct MemorySpan { name: String, attributes: Vec<KeyValue>, exported: Arc<Mutex<Vec<Exported>>>, ended: bool }
3895impl Tracer for MemoryTracer {
3896 type Span = MemorySpan;
3897 fn build_with_context(&self, builder: SpanBuilder, _: &Context) -> MemorySpan {
3898 MemorySpan { name: builder.name.into_owned(), attributes: builder.attributes.unwrap_or_default(), exported: self.exported.clone(), ended: false }
3899 }
3900}
3901impl Span for MemorySpan {
3902 fn add_event_with_timestamp<T>(&mut self, _: T, _: SystemTime, _: Vec<KeyValue>) where T: Into<Cow<'static, str>> {}
3903 fn span_context(&self) -> &SpanContext { &SpanContext::NONE }
3904 fn is_recording(&self) -> bool { !self.ended }
3905 fn set_attribute(&mut self, attribute: KeyValue) { self.attributes.push(attribute); }
3906 fn set_status(&mut self, _: Status) {}
3907 fn update_name<T>(&mut self, name: T) where T: Into<Cow<'static, str>> { self.name = name.into().into_owned(); }
3908 fn add_link(&mut self, _: SpanContext, _: Vec<KeyValue>) {}
3909 fn end_with_timestamp(&mut self, _: SystemTime) {
3910 if !self.ended {
3911 self.ended = true;
3912 self.exported.lock().unwrap().push(Exported { name: self.name.clone(), attributes: self.attributes.clone() });
3913 }
3914 }
3915}
3916
3917#[derive(Clone, Copy)]
3918enum BeginMode { Success, Error, Cancelled }
3919#[derive(Clone)]
3920struct ScriptedAdapter { adapter_provider: &'static str, final_usage: bool, overflow: bool, begin_mode: BeginMode }
3921struct ScriptedSession { selected_provider: Option<&'static str>, model: &'static str, final_usage: bool, overflow: bool, begin_mode: BeginMode }
3922struct ScriptedTurn { events: VecDeque<ModelTurnEvent> }
3923#[async_trait]
3924impl ModelAdapter for ScriptedAdapter {
3925 type Session = ScriptedSession;
3926 async fn start_session(&self, _: SessionConfig) -> Result<Self::Session, LoopError> {
3927 Ok(ScriptedSession { selected_provider: Some("before-begin"), model: "before-model", final_usage: self.final_usage, overflow: self.overflow, begin_mode: self.begin_mode })
3928 }
3929 fn provider_name(&self) -> Option<&str> { Some(self.adapter_provider) }
3930}
3931#[async_trait]
3932impl ModelSession for ScriptedSession {
3933 type Turn = ScriptedTurn;
3934 async fn begin_turn(&mut self, _: TurnRequest, _: Option<TurnCancellation>) -> Result<Self::Turn, LoopError> {
3935 match self.begin_mode {
3936 BeginMode::Error => return Err(LoopError::InvalidState("begin failed".into())),
3937 BeginMode::Cancelled => return Err(LoopError::Cancelled),
3938 BeginMode::Success => {}
3939 }
3940 self.selected_provider = if self.final_usage { Some("session-after-begin") } else { None };
3941 self.model = if self.final_usage { "model-after-begin" } else { "fallback-model-after-begin" };
3942 let mut stream_meta = MetadataMap::new();
3943 stream_meta.insert("stream_only".into(), serde_json::json!(true));
3944 stream_meta.insert("shared".into(), serde_json::json!("stream"));
3945 let streamed = Usage {
3946 tokens: Some(TokenUsage::new(if self.overflow { i64::MAX as u64 + 1 } else { 30 }, 40)),
3947 cost: Some(CostUsage::new(0.75, "USD")),
3948 metadata: stream_meta,
3949 };
3950 let final_usage = if self.final_usage {
3951 let mut metadata = MetadataMap::new();
3952 metadata.insert("shared".into(), serde_json::json!("final"));
3953 Some(Usage { tokens: Some(TokenUsage::new(if self.overflow { i64::MAX as u64 + 1 } else { 1 }, 2)), cost: None, metadata })
3954 } else { None };
3955 let mut result_metadata = MetadataMap::new();
3956 set_provider_finish_reasons(&mut result_metadata, ["native", "native", "done"]);
3957 let outputs = vec![
3958 Item::text(ItemKind::Assistant, "first"),
3959 Item::text(ItemKind::Assistant, "second"),
3960 Item::text(ItemKind::Assistant, "third"),
3961 ];
3962 Ok(ScriptedTurn { events: VecDeque::from([
3963 ModelTurnEvent::Usage(streamed),
3964 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()) }),
3965 ]) })
3966 }
3967 fn model_name(&self) -> Option<&str> { Some(self.model) }
3968 fn provider_name(&self) -> Option<&str> { self.selected_provider }
3969}
3970#[async_trait]
3971impl ModelTurn for ScriptedTurn {
3972 async fn next_event(&mut self, _: Option<TurnCancellation>) -> Result<Option<ModelTurnEvent>, LoopError> { Ok(self.events.pop_front()) }
3973}
3974
3975fn attr<'a>(span: &'a Exported, key: &str) -> Option<&'a Value> {
3976 span.attributes.iter().rev().find(|a| a.key.as_str() == key).map(|a| &a.value)
3977}
3978fn operation(span: &Exported) -> Option<&str> {
3979 match attr(span, "gen_ai.operation.name") { Some(Value::String(value)) => Some(value.as_str()), _ => None }
3980}
3981fn json_array(span: &Exported, key: &str) -> Vec<serde_json::Value> {
3982 match attr(span, key) {
3983 Some(Value::Array(Array::String(values))) => values.iter().map(|v| serde_json::from_str(v.as_str()).unwrap()).collect(),
3984 other => panic!("{key} was not Array<String>: {other:?}"),
3985 }
3986}
3987
3988fn run_attempt(adapter: ScriptedAdapter) -> (Vec<Exported>, Result<LoopStep, LoopError>) {
3989 let tracer = MemoryTracer::default();
3990 let subscriber = tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer.clone()));
3991 let runtime = tokio::runtime::Builder::new_current_thread().build().unwrap();
3992 let result = tracing::subscriber::with_default(subscriber, || runtime.block_on(async {
3993 let media = Item::new(ItemKind::User, vec![
3994 Part::media(Modality::Image, "image/png", DataRef::inline_bytes([1, 2, 3])),
3995 Part::media(Modality::Audio, "audio/wav", DataRef::uri("https://secret.invalid/audio")),
3996 ]);
3997 let agent = Agent::builder().model(adapter).transcript(vec![
3998 Item::text(ItemKind::System, "old"), Item::text(ItemKind::User, "middle"), media,
3999 ]).input(vec![Item::text(ItemKind::User, "newest")]).telemetry(
4000 TelemetryConfig::default()
4001 .with_input_messages(MessageCapture::new(3, 100_000).unwrap())
4002 .with_output_messages(MessageCapture::new(2, 100_000).unwrap())
4003 ).build().unwrap();
4004 let mut driver = agent.start(SessionConfig::new("otel-test")).await.unwrap();
4005 driver.next().await
4006 }));
4007 (tracer.exported.lock().unwrap().clone(), result)
4008}
4009fn run(adapter: ScriptedAdapter) -> (Vec<Exported>, agentkit_loop::TurnResult) {
4010 let (spans, result) = run_attempt(adapter);
4011 let result = match result.unwrap() { LoopStep::Finished(result) => result, other => panic!("unexpected step: {other:?}") };
4012 (spans, result)
4013}
4014
4015fn main() {
4016 let (spans, result) = run(ScriptedAdapter { adapter_provider: "adapter", final_usage: true, overflow: true, begin_mode: BeginMode::Success });
4017 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
4018 let agent = spans.iter().find(|s| operation(s) == Some("invoke_agent")).unwrap();
4019 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("session-after-begin".into())));
4020 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("model-after-begin".into())));
4021 assert!(attr(chat, "gen_ai.usage.input_tokens").is_none());
4022 assert_eq!(attr(chat, "gen_ai.usage.output_tokens"), Some(&Value::I64(2)));
4023 assert_eq!(attr(chat, "gen_ai.usage.cost"), Some(&Value::F64(0.75)));
4024 assert_eq!(attr(agent, "gen_ai.provider.name"), Some(&Value::String("session-after-begin".into())));
4025 assert!(attr(agent, "gen_ai.usage.input_tokens").is_none());
4026 assert!(attr(agent, "gen_ai.usage.cost").is_none());
4027 match attr(chat, "gen_ai.response.finish_reasons") {
4028 Some(Value::Array(Array::String(values))) => assert_eq!(values.iter().map(|v| v.as_str()).collect::<Vec<_>>(), ["native", "done"]),
4029 other => panic!("finish reasons were not Array<String>: {other:?}"),
4030 }
4031 assert_eq!(result.usage.as_ref().unwrap().metadata["shared"], "final");
4032 assert_eq!(result.usage.as_ref().unwrap().metadata["stream_only"], true);
4033 let input = json_array(chat, "gen_ai.input.messages");
4034 assert_eq!(input.len(), 3);
4035 assert!(input[0].to_string().contains("middle"));
4036 assert!(input[1].to_string().contains("omitted"));
4037 assert!(input[2].to_string().contains("newest"));
4038 assert!(!input.iter().any(|message| message.to_string().contains("old")));
4039 let output = json_array(chat, "gen_ai.output.messages");
4040 assert_eq!(output.len(), 2);
4041 assert!(output[0].to_string().contains("first"));
4042 assert!(output[1].to_string().contains("second"));
4043 let encoded = format!("{input:?}{output:?}");
4044 assert!(!encoded.contains("secret.invalid"));
4045 assert!(!encoded.contains("[1,2,3]"));
4046
4047 let (spans, result) = run(ScriptedAdapter { adapter_provider: "adapter-fallback", final_usage: false, overflow: false, begin_mode: BeginMode::Success });
4048 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
4049 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("adapter-fallback".into())));
4050 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("fallback-model-after-begin".into())));
4051 assert_eq!(attr(chat, "gen_ai.usage.input_tokens"), Some(&Value::I64(30)));
4052 assert_eq!(attr(chat, "gen_ai.usage.cost"), Some(&Value::F64(0.75)));
4053 assert_eq!(result.usage.unwrap().metadata["stream_only"], true);
4054
4055 for begin_mode in [BeginMode::Error, BeginMode::Cancelled] {
4056 let (spans, result) = run_attempt(ScriptedAdapter { adapter_provider: "adapter-before-error", final_usage: false, overflow: false, begin_mode });
4057 match begin_mode {
4058 BeginMode::Error => assert!(matches!(result, Err(LoopError::InvalidState(_)))),
4059 BeginMode::Cancelled => assert!(matches!(result, Ok(LoopStep::Finished(_)))),
4060 BeginMode::Success => unreachable!(),
4061 }
4062 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
4063 let agent = spans.iter().find(|s| operation(s) == Some("invoke_agent")).unwrap();
4064 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("before-begin".into())));
4065 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("before-model".into())));
4066 assert_eq!(attr(agent, "gen_ai.provider.name"), Some(&Value::String("before-begin".into())));
4067 }
4068}
4069"#;
4070}
4071
4072fn interrupted_assistant_items() -> Vec<Item> {
4073 vec![Item {
4074 id: None,
4075 kind: ItemKind::Assistant,
4076 parts: vec![Part::Text(TextPart {
4077 text: "Previous assistant response was interrupted by the user before completion."
4078 .into(),
4079 metadata: interrupted_metadata("assistant"),
4080 })],
4081 metadata: interrupted_metadata("assistant"),
4082 usage: None,
4083 finish_reason: None,
4084 created_at: None,
4085 }]
4086}
4087
4088fn unanswered_tool_calls(transcript: &[Item]) -> Vec<ToolCallPart> {
4090 let mut open: Vec<ToolCallPart> = Vec::new();
4091 for part in transcript.iter().flat_map(|item| &item.parts) {
4092 match part {
4093 Part::ToolCall(call) => open.push(call.clone()),
4094 Part::ToolResult(result) => open.retain(|call| call.id != result.call_id),
4095 _ => {}
4096 }
4097 }
4098 open
4099}
4100
4101fn interrupted_tool_result_item(call: ToolCallPart) -> Item {
4106 Item {
4107 id: None,
4108 kind: ItemKind::Tool,
4109 parts: vec![Part::ToolResult(ToolResultPart {
4110 call_id: call.id,
4111 output: ToolOutput::Text("tool call interrupted before it reported a result".into()),
4112 is_error: true,
4113 metadata: interrupted_metadata("tool"),
4114 })],
4115 metadata: interrupted_metadata("tool"),
4116 usage: None,
4117 finish_reason: None,
4118 created_at: None,
4119 }
4120}
4121
4122fn cancelled_approval_item(pending: PendingApprovalToolCall) -> Item {
4123 Item {
4124 id: None,
4125 kind: ItemKind::Tool,
4126 parts: vec![Part::ToolResult(ToolResultPart {
4127 call_id: pending.call.id,
4128 output: ToolOutput::Text("approval cancelled".into()),
4129 is_error: true,
4130 metadata: pending.call.metadata,
4131 })],
4132 metadata: MetadataMap::new(),
4133 usage: None,
4134 finish_reason: None,
4135 created_at: None,
4136 }
4137}
4138
4139fn transcript_has_pending_input(transcript: &[Item]) -> bool {
4145 matches!(
4146 transcript.last().map(|item| item.kind),
4147 Some(ItemKind::User | ItemKind::Tool | ItemKind::Notification)
4148 )
4149}
4150
4151fn extract_tool_calls(items: &[Item]) -> Vec<ToolCallPart> {
4152 let mut calls = Vec::new();
4153 for item in items {
4154 for part in &item.parts {
4155 if let Part::ToolCall(call) = part {
4156 calls.push(call.clone());
4157 }
4158 }
4159 }
4160 calls
4161}
4162
4163fn tool_result_is_error(item: &Item) -> bool {
4164 item.parts
4165 .iter()
4166 .any(|part| matches!(part, Part::ToolResult(result) if result.is_error))
4167}
4168
4169fn tool_result_not_started(item: &Item) -> bool {
4170 item.parts.iter().any(|part| {
4171 matches!(
4172 part,
4173 Part::ToolResult(result)
4174 if result
4175 .metadata
4176 .get(TOOL_RESULT_NOT_STARTED_METADATA_KEY)
4177 .and_then(Value::as_bool)
4178 == Some(true)
4179 )
4180 })
4181}
4182
4183#[derive(Debug, Error)]
4185pub enum LoopError {
4186 #[error("invalid driver state: {0}")]
4188 InvalidState(String),
4189 #[error("turn cancelled")]
4191 Cancelled,
4192 #[error("provider error: {0}")]
4194 Provider(String),
4195 #[error("tool error: {0}")]
4197 Tool(#[from] ToolError),
4198 #[error("mutator error: {0}")]
4200 Mutator(String),
4201 #[error("unsupported operation: {0}")]
4203 Unsupported(String),
4204}
4205
4206struct DriverEmitter<'a> {
4211 session_id: &'a Arc<SessionId>,
4212 observers: &'a [Arc<dyn LoopObserver>],
4213}
4214
4215impl<'a> EventEmitter for DriverEmitter<'a> {
4216 fn emit(&self, event: AgentEvent) {
4217 fan_out_observed_event(self.observers, self.session_id, event);
4218 }
4219}
4220
4221fn fan_out_observed_event(
4222 observers: &[Arc<dyn LoopObserver>],
4223 session_id: &Arc<SessionId>,
4224 event: AgentEvent,
4225) {
4226 if observers.is_empty() {
4227 return;
4228 }
4229 let observed = ObservedEvent {
4230 session_id: Arc::clone(session_id),
4231 event,
4232 };
4233 let last = observers.len() - 1;
4234 for observer in &observers[..last] {
4235 observer.handle_event(observed.clone());
4236 }
4237 observers[last].handle_event(observed);
4238}
4239
4240fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> {
4245 let mut pending: HashSet<ToolCallId> = HashSet::new();
4246 let mut seen_calls: HashSet<ToolCallId> = HashSet::new();
4247 let mut seen_results: HashSet<ToolCallId> = HashSet::new();
4248 for item in transcript {
4249 for part in &item.parts {
4250 match part {
4251 Part::ToolCall(call) => {
4252 if !seen_calls.insert(call.id.clone()) {
4253 return Err(LoopError::Mutator(format!(
4254 "transcript invariant violation: duplicate tool_use: {}",
4255 call.id.0
4256 )));
4257 }
4258 pending.insert(call.id.clone());
4259 }
4260 Part::ToolResult(result) => {
4261 if !pending.remove(&result.call_id) {
4262 let kind = if seen_results.contains(&result.call_id) {
4263 "duplicate"
4264 } else {
4265 "orphaned"
4266 };
4267 return Err(LoopError::Mutator(format!(
4268 "transcript invariant violation: {kind} tool_result: {}",
4269 result.call_id.0
4270 )));
4271 }
4272 seen_results.insert(result.call_id.clone());
4273 }
4274 _ => {}
4275 }
4276 }
4277 }
4278 if !pending.is_empty() {
4279 let missing: Vec<String> = pending.into_iter().map(|id| id.0).collect();
4280 return Err(LoopError::Mutator(format!(
4281 "transcript invariant violation: tool_use(s) without matching tool_result: {}",
4282 missing.join(", ")
4283 )));
4284 }
4285 Ok(())
4286}
4287
4288#[cfg(test)]
4289mod tests {
4290 use std::collections::VecDeque;
4291 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
4292 use std::sync::{Arc as StdArc, Mutex as StdMutex};
4293
4294 use agentkit_core::{
4295 CancellationController, ItemKind, Part, TextPart, ToolCallId, ToolCallPart, ToolOutput,
4296 ToolResultPart,
4297 };
4298 use agentkit_task_manager::{
4299 AsyncTaskManager, RoutingDecision, TaskEvent, TaskManager, TaskManagerError,
4300 TaskManagerHandle, TaskRoutingPolicy,
4301 };
4302 use agentkit_tools_core::{
4303 FileSystemPermissionRequest, PermissionCode, PermissionDecision, PermissionDenial, Tool,
4304 ToolAnnotations, ToolCatalogEvent, ToolExecutionOutcome, ToolName, ToolRegistry,
4305 ToolResult, ToolSpec,
4306 };
4307 use serde_json::{Value, json};
4308 use tokio::sync::Notify;
4309 use tokio::time::{Duration, timeout};
4310
4311 use super::*;
4312
4313 struct FakeAdapter;
4314 struct SlowAdapter;
4315 struct RecordingAdapter {
4316 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
4317 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
4318 }
4319 struct MultiToolAdapter;
4320 struct DualApprovalAdapter;
4321
4322 struct FakeSession;
4323 struct SlowSession;
4324 struct RecordingSession {
4325 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
4326 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
4327 }
4328 struct MultiToolSession;
4329 struct DualApprovalSession;
4330
4331 struct FakeTurn {
4332 events: VecDeque<ModelTurnEvent>,
4333 }
4334
4335 struct SlowTurn {
4336 emitted: bool,
4337 }
4338
4339 struct RecordingTurn {
4340 emitted: bool,
4341 }
4342 struct MultiToolTurn {
4343 events: VecDeque<ModelTurnEvent>,
4344 }
4345 struct DualApprovalTurn {
4346 events: VecDeque<ModelTurnEvent>,
4347 }
4348
4349 struct TestTaskManager<T> {
4350 inner: T,
4351 start_error: Option<&'static str>,
4352 approved_start_error: Option<&'static str>,
4353 pending_update_error: Option<(usize, &'static str)>,
4354 pending_update_calls: AtomicUsize,
4355 interrupted: Option<StdArc<StdMutex<Vec<agentkit_core::TurnId>>>>,
4356 interrupt_error: Option<&'static str>,
4357 }
4358
4359 impl<T> TestTaskManager<T> {
4360 fn new(inner: T) -> Self {
4361 Self {
4362 inner,
4363 start_error: None,
4364 approved_start_error: None,
4365 pending_update_error: None,
4366 pending_update_calls: AtomicUsize::new(0),
4367 interrupted: None,
4368 interrupt_error: None,
4369 }
4370 }
4371
4372 fn fail_start(mut self, message: &'static str) -> Self {
4373 self.start_error = Some(message);
4374 self
4375 }
4376
4377 fn fail_approved_start(mut self, message: &'static str) -> Self {
4378 self.approved_start_error = Some(message);
4379 self
4380 }
4381
4382 fn fail_pending_update_on(mut self, call: usize, message: &'static str) -> Self {
4383 self.pending_update_error = Some((call, message));
4384 self
4385 }
4386
4387 fn record_interrupts(
4388 mut self,
4389 interrupted: StdArc<StdMutex<Vec<agentkit_core::TurnId>>>,
4390 ) -> Self {
4391 self.interrupted = Some(interrupted);
4392 self
4393 }
4394
4395 fn fail_interrupt(mut self, message: &'static str) -> Self {
4396 self.interrupt_error = Some(message);
4397 self
4398 }
4399 }
4400
4401 #[async_trait]
4402 impl<T: TaskManager> TaskManager for TestTaskManager<T> {
4403 async fn start_task(
4404 &self,
4405 request: TaskLaunchRequest,
4406 ctx: TaskStartContext,
4407 ) -> Result<TaskStartOutcome, TaskManagerError> {
4408 if let Some(message) = self.start_error.or_else(|| {
4409 matches!(&request.kind, TaskLaunchKind::Approved(_))
4410 .then_some(self.approved_start_error)
4411 .flatten()
4412 }) {
4413 return Err(TaskManagerError::Internal(message.into()));
4414 }
4415 self.inner.start_task(request, ctx).await
4416 }
4417
4418 async fn wait_for_turn(
4419 &self,
4420 turn_id: &agentkit_core::TurnId,
4421 cancellation: Option<TurnCancellation>,
4422 ) -> Result<Option<TurnTaskUpdate>, TaskManagerError> {
4423 self.inner.wait_for_turn(turn_id, cancellation).await
4424 }
4425
4426 async fn take_pending_loop_updates(&self) -> Result<PendingLoopUpdates, TaskManagerError> {
4427 if let Some((call, message)) = self.pending_update_error
4428 && self.pending_update_calls.fetch_add(1, Ordering::SeqCst) == call
4429 {
4430 return Err(TaskManagerError::Internal(message.into()));
4431 }
4432 self.inner.take_pending_loop_updates().await
4433 }
4434
4435 async fn on_turn_interrupted(
4436 &self,
4437 turn_id: &agentkit_core::TurnId,
4438 ) -> Result<(), TaskManagerError> {
4439 if let Some(interrupted) = &self.interrupted {
4440 interrupted.lock().unwrap().push(turn_id.clone());
4441 }
4442 if let Some(message) = self.interrupt_error {
4443 return Err(TaskManagerError::Internal(message.into()));
4444 }
4445 self.inner.on_turn_interrupted(turn_id).await
4446 }
4447
4448 fn handle(&self) -> TaskManagerHandle {
4449 self.inner.handle()
4450 }
4451 }
4452
4453 struct DelayedApprovalExecutor {
4454 entered: StdArc<AtomicBool>,
4455 release: StdArc<Notify>,
4456 approved_entered: Option<StdArc<AtomicBool>>,
4457 approved_release: Option<StdArc<Notify>>,
4458 cancellation: Option<CancellationController>,
4459 spec: ToolSpec,
4460 }
4461
4462 impl DelayedApprovalExecutor {
4463 fn new(entered: StdArc<AtomicBool>, release: StdArc<Notify>) -> Self {
4464 Self {
4465 entered,
4466 release,
4467 approved_entered: None,
4468 approved_release: None,
4469 cancellation: None,
4470 spec: ToolSpec {
4471 name: ToolName::new("echo"),
4472 description: "delayed approval".into(),
4473 input_schema: json!({
4474 "type": "object",
4475 "properties": {
4476 "value": { "type": "string" }
4477 },
4478 "required": ["value"],
4479 "additionalProperties": false
4480 }),
4481 output_schema: None,
4482 annotations: ToolAnnotations::default(),
4483 metadata: MetadataMap::new(),
4484 },
4485 }
4486 }
4487
4488 fn cancelling_on_approval(mut self, controller: CancellationController) -> Self {
4489 self.cancellation = Some(controller);
4490 self
4491 }
4492
4493 fn blocking_after_approval(
4494 mut self,
4495 entered: StdArc<AtomicBool>,
4496 release: StdArc<Notify>,
4497 ) -> Self {
4498 self.approved_entered = Some(entered);
4499 self.approved_release = Some(release);
4500 self
4501 }
4502 }
4503
4504 #[async_trait]
4505 impl ToolExecutor for DelayedApprovalExecutor {
4506 fn specs(&self) -> Vec<ToolSpec> {
4507 vec![self.spec.clone()]
4508 }
4509
4510 async fn execute(
4511 &self,
4512 request: ToolRequest,
4513 _ctx: &mut ToolContext<'_>,
4514 ) -> ToolExecutionOutcome {
4515 self.entered.store(true, Ordering::SeqCst);
4516 self.release.notified().await;
4517 if let Some(controller) = &self.cancellation {
4518 controller.interrupt();
4519 }
4520 ToolExecutionOutcome::Interrupted(
4521 agentkit_tools_core::ToolInterruption::ApprovalRequired(ApprovalRequest {
4522 task_id: None,
4523 call_id: None,
4524 id: "approval:delayed".into(),
4525 request_kind: "delayed.approval".into(),
4526 reason: agentkit_tools_core::ApprovalReason::PolicyRequiresConfirmation,
4527 summary: "delayed approval".into(),
4528 metadata: request.metadata,
4529 }),
4530 )
4531 }
4532
4533 async fn execute_approved(
4534 &self,
4535 request: ToolRequest,
4536 approved_request: &ApprovalRequest,
4537 ctx: &mut ToolContext<'_>,
4538 ) -> ToolExecutionOutcome {
4539 let (Some(entered), Some(release)) = (&self.approved_entered, &self.approved_release)
4540 else {
4541 return self.execute(request, ctx).await;
4542 };
4543 let _ = approved_request;
4544 entered.store(true, Ordering::SeqCst);
4545 release.notified().await;
4546 ToolExecutionOutcome::Completed(ToolResult {
4547 result: ToolResultPart {
4548 call_id: request.call_id,
4549 output: ToolOutput::Text("approved-ok".into()),
4550 is_error: false,
4551 metadata: MetadataMap::new(),
4552 },
4553 duration: None,
4554 metadata: MetadataMap::new(),
4555 })
4556 }
4557 }
4558
4559 #[async_trait]
4560 impl ModelAdapter for FakeAdapter {
4561 type Session = FakeSession;
4562
4563 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4564 Ok(FakeSession)
4565 }
4566 }
4567
4568 #[async_trait]
4569 impl ModelAdapter for SlowAdapter {
4570 type Session = SlowSession;
4571
4572 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4573 Ok(SlowSession)
4574 }
4575 }
4576
4577 #[async_trait]
4578 impl ModelAdapter for RecordingAdapter {
4579 type Session = RecordingSession;
4580
4581 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4582 Ok(RecordingSession {
4583 seen_descriptions: self.seen_descriptions.clone(),
4584 seen_caches: self.seen_caches.clone(),
4585 })
4586 }
4587 }
4588
4589 #[async_trait]
4590 impl ModelAdapter for MultiToolAdapter {
4591 type Session = MultiToolSession;
4592
4593 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4594 Ok(MultiToolSession)
4595 }
4596 }
4597
4598 #[async_trait]
4599 impl ModelAdapter for DualApprovalAdapter {
4600 type Session = DualApprovalSession;
4601
4602 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4603 Ok(DualApprovalSession)
4604 }
4605 }
4606
4607 #[async_trait]
4608 impl ModelSession for FakeSession {
4609 type Turn = FakeTurn;
4610
4611 async fn begin_turn(
4612 &mut self,
4613 request: TurnRequest,
4614 _cancellation: Option<TurnCancellation>,
4615 ) -> Result<Self::Turn, LoopError> {
4616 let has_tool_result = request.transcript.iter().any(|item| {
4617 item.kind == ItemKind::Tool
4618 && item
4619 .parts
4620 .iter()
4621 .any(|part| matches!(part, Part::ToolResult(_)))
4622 });
4623 let tool_name = request
4624 .available_tools
4625 .first()
4626 .map(|tool| tool.name.0.clone())
4627 .unwrap_or_else(|| "echo".into());
4628
4629 let events = if has_tool_result {
4630 let result_text = request
4631 .transcript
4632 .iter()
4633 .rev()
4634 .find_map(|item| {
4635 item.parts.iter().find_map(|part| match (item.kind, part) {
4636 (ItemKind::Notification, Part::Text(text)) => Some(text.text.clone()),
4637 (
4638 _,
4639 Part::ToolResult(ToolResultPart {
4640 output: ToolOutput::Text(text),
4641 ..
4642 }),
4643 ) => Some(text.clone()),
4644 _ => None,
4645 })
4646 })
4647 .unwrap_or_else(|| "missing".into());
4648
4649 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
4650 model: None,
4651 response_id: None,
4652 finish_reason: FinishReason::Completed,
4653 output_items: vec![Item {
4654 id: None,
4655 kind: ItemKind::Assistant,
4656 parts: vec![Part::Text(TextPart {
4657 text: format!("tool said: {result_text}"),
4658 metadata: MetadataMap::new(),
4659 })],
4660 metadata: MetadataMap::new(),
4661 usage: None,
4662 finish_reason: None,
4663 created_at: None,
4664 }],
4665 usage: None,
4666 metadata: MetadataMap::new(),
4667 })])
4668 } else {
4669 VecDeque::from([
4670 ModelTurnEvent::ToolCall(agentkit_core::ToolCallPart {
4671 id: ToolCallId::new("call-1"),
4672 name: tool_name.clone(),
4673 input: json!({ "value": "pong" }),
4674 metadata: MetadataMap::new(),
4675 }),
4676 ModelTurnEvent::Finished(ModelTurnResult {
4677 model: None,
4678 response_id: None,
4679 finish_reason: FinishReason::ToolCall,
4680 output_items: vec![Item {
4681 id: None,
4682 kind: ItemKind::Assistant,
4683 parts: vec![Part::ToolCall(agentkit_core::ToolCallPart {
4684 id: ToolCallId::new("call-1"),
4685 name: tool_name,
4686 input: json!({ "value": "pong" }),
4687 metadata: MetadataMap::new(),
4688 })],
4689 metadata: MetadataMap::new(),
4690 usage: None,
4691 finish_reason: None,
4692 created_at: None,
4693 }],
4694 usage: None,
4695 metadata: MetadataMap::new(),
4696 }),
4697 ])
4698 };
4699
4700 Ok(FakeTurn { events })
4701 }
4702 }
4703
4704 #[async_trait]
4705 impl ModelSession for SlowSession {
4706 type Turn = SlowTurn;
4707
4708 async fn begin_turn(
4709 &mut self,
4710 request: TurnRequest,
4711 cancellation: Option<TurnCancellation>,
4712 ) -> Result<Self::Turn, LoopError> {
4713 let should_block = request
4714 .transcript
4715 .iter()
4716 .rev()
4717 .find(|item| item.kind == ItemKind::User)
4718 .is_some_and(|item| {
4719 item.parts.iter().any(|part| match part {
4720 Part::Text(text) => text.text == "do the long task",
4721 _ => false,
4722 })
4723 });
4724
4725 if should_block && let Some(cancellation) = cancellation {
4726 cancellation.cancelled().await;
4727 return Err(LoopError::Cancelled);
4728 }
4729
4730 Ok(SlowTurn { emitted: false })
4731 }
4732 }
4733
4734 #[async_trait]
4735 impl ModelSession for RecordingSession {
4736 type Turn = RecordingTurn;
4737
4738 async fn begin_turn(
4739 &mut self,
4740 request: TurnRequest,
4741 _cancellation: Option<TurnCancellation>,
4742 ) -> Result<Self::Turn, LoopError> {
4743 let descriptions = request
4744 .available_tools
4745 .iter()
4746 .map(|tool| tool.description.clone())
4747 .collect::<Vec<_>>();
4748 self.seen_descriptions.lock().unwrap().push(descriptions);
4749 self.seen_caches.lock().unwrap().push(request.cache.clone());
4750
4751 Ok(RecordingTurn { emitted: false })
4752 }
4753 }
4754
4755 #[async_trait]
4756 impl ModelSession for MultiToolSession {
4757 type Turn = MultiToolTurn;
4758
4759 async fn begin_turn(
4760 &mut self,
4761 request: TurnRequest,
4762 _cancellation: Option<TurnCancellation>,
4763 ) -> Result<Self::Turn, LoopError> {
4764 let has_tool_result = request.transcript.iter().any(|item| {
4765 item.kind == ItemKind::Tool
4766 && item
4767 .parts
4768 .iter()
4769 .any(|part| matches!(part, Part::ToolResult(_)))
4770 });
4771
4772 let events = if has_tool_result {
4773 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
4774 model: None,
4775 response_id: None,
4776 finish_reason: FinishReason::Completed,
4777 output_items: vec![Item {
4778 id: None,
4779 kind: ItemKind::Assistant,
4780 parts: vec![Part::Text(TextPart {
4781 text: "mixed tools finished".into(),
4782 metadata: MetadataMap::new(),
4783 })],
4784 metadata: MetadataMap::new(),
4785 usage: None,
4786 finish_reason: None,
4787 created_at: None,
4788 }],
4789 usage: None,
4790 metadata: MetadataMap::new(),
4791 })])
4792 } else {
4793 let foreground = agentkit_core::ToolCallPart {
4794 id: ToolCallId::new("call-foreground"),
4795 name: "foreground-wait".into(),
4796 input: json!({}),
4797 metadata: MetadataMap::new(),
4798 };
4799 let background = agentkit_core::ToolCallPart {
4800 id: ToolCallId::new("call-background"),
4801 name: "background-wait".into(),
4802 input: json!({}),
4803 metadata: MetadataMap::new(),
4804 };
4805 VecDeque::from([
4806 ModelTurnEvent::ToolCall(foreground.clone()),
4807 ModelTurnEvent::ToolCall(background.clone()),
4808 ModelTurnEvent::Finished(ModelTurnResult {
4809 model: None,
4810 response_id: None,
4811 finish_reason: FinishReason::ToolCall,
4812 output_items: vec![Item {
4813 id: None,
4814 kind: ItemKind::Assistant,
4815 parts: vec![Part::ToolCall(foreground), Part::ToolCall(background)],
4816 metadata: MetadataMap::new(),
4817 usage: None,
4818 finish_reason: None,
4819 created_at: None,
4820 }],
4821 usage: None,
4822 metadata: MetadataMap::new(),
4823 }),
4824 ])
4825 };
4826
4827 Ok(MultiToolTurn { events })
4828 }
4829 }
4830
4831 #[async_trait]
4832 impl ModelSession for DualApprovalSession {
4833 type Turn = DualApprovalTurn;
4834
4835 async fn begin_turn(
4836 &mut self,
4837 request: TurnRequest,
4838 _cancellation: Option<TurnCancellation>,
4839 ) -> Result<Self::Turn, LoopError> {
4840 let tool_results = request
4841 .transcript
4842 .iter()
4843 .flat_map(|item| item.parts.iter())
4844 .filter(|part| matches!(part, Part::ToolResult(_)))
4845 .count();
4846
4847 let events = if tool_results >= 2 {
4848 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
4849 model: None,
4850 response_id: None,
4851 finish_reason: FinishReason::Completed,
4852 output_items: vec![Item {
4853 id: None,
4854 kind: ItemKind::Assistant,
4855 parts: vec![Part::Text(TextPart {
4856 text: "both approvals finished".into(),
4857 metadata: MetadataMap::new(),
4858 })],
4859 metadata: MetadataMap::new(),
4860 usage: None,
4861 finish_reason: None,
4862 created_at: None,
4863 }],
4864 usage: None,
4865 metadata: MetadataMap::new(),
4866 })])
4867 } else {
4868 let first = agentkit_core::ToolCallPart {
4869 id: ToolCallId::new("call-1"),
4870 name: "echo".into(),
4871 input: json!({ "value": "first" }),
4872 metadata: MetadataMap::new(),
4873 };
4874 let second = agentkit_core::ToolCallPart {
4875 id: ToolCallId::new("call-2"),
4876 name: "echo".into(),
4877 input: json!({ "value": "second" }),
4878 metadata: MetadataMap::new(),
4879 };
4880 VecDeque::from([
4881 ModelTurnEvent::ToolCall(first.clone()),
4882 ModelTurnEvent::ToolCall(second.clone()),
4883 ModelTurnEvent::Finished(ModelTurnResult {
4884 model: None,
4885 response_id: None,
4886 finish_reason: FinishReason::ToolCall,
4887 output_items: vec![Item {
4888 id: None,
4889 kind: ItemKind::Assistant,
4890 parts: vec![Part::ToolCall(first), Part::ToolCall(second)],
4891 metadata: MetadataMap::new(),
4892 usage: None,
4893 finish_reason: None,
4894 created_at: None,
4895 }],
4896 usage: None,
4897 metadata: MetadataMap::new(),
4898 }),
4899 ])
4900 };
4901
4902 Ok(DualApprovalTurn { events })
4903 }
4904 }
4905
4906 #[async_trait]
4907 impl ModelTurn for FakeTurn {
4908 async fn next_event(
4909 &mut self,
4910 _cancellation: Option<TurnCancellation>,
4911 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4912 Ok(self.events.pop_front())
4913 }
4914 }
4915
4916 #[async_trait]
4917 impl ModelTurn for SlowTurn {
4918 async fn next_event(
4919 &mut self,
4920 cancellation: Option<TurnCancellation>,
4921 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4922 if let Some(cancellation) = cancellation
4923 && cancellation.is_cancelled()
4924 {
4925 return Err(LoopError::Cancelled);
4926 }
4927
4928 if self.emitted {
4929 Ok(None)
4930 } else {
4931 self.emitted = true;
4932 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
4933 model: None,
4934 response_id: None,
4935 finish_reason: FinishReason::Completed,
4936 output_items: vec![Item {
4937 id: None,
4938 kind: ItemKind::Assistant,
4939 parts: vec![Part::Text(TextPart {
4940 text: "done".into(),
4941 metadata: MetadataMap::new(),
4942 })],
4943 metadata: MetadataMap::new(),
4944 usage: None,
4945 finish_reason: None,
4946 created_at: None,
4947 }],
4948 usage: None,
4949 metadata: MetadataMap::new(),
4950 })))
4951 }
4952 }
4953 }
4954
4955 #[async_trait]
4956 impl ModelTurn for RecordingTurn {
4957 async fn next_event(
4958 &mut self,
4959 _cancellation: Option<TurnCancellation>,
4960 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4961 if self.emitted {
4962 Ok(None)
4963 } else {
4964 self.emitted = true;
4965 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
4966 model: None,
4967 response_id: None,
4968 finish_reason: FinishReason::Completed,
4969 output_items: vec![Item {
4970 id: None,
4971 kind: ItemKind::Assistant,
4972 parts: vec![Part::Text(TextPart {
4973 text: "done".into(),
4974 metadata: MetadataMap::new(),
4975 })],
4976 metadata: MetadataMap::new(),
4977 usage: None,
4978 finish_reason: None,
4979 created_at: None,
4980 }],
4981 usage: None,
4982 metadata: MetadataMap::new(),
4983 })))
4984 }
4985 }
4986 }
4987
4988 #[async_trait]
4989 impl ModelTurn for MultiToolTurn {
4990 async fn next_event(
4991 &mut self,
4992 _cancellation: Option<TurnCancellation>,
4993 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4994 Ok(self.events.pop_front())
4995 }
4996 }
4997
4998 #[async_trait]
4999 impl ModelTurn for DualApprovalTurn {
5000 async fn next_event(
5001 &mut self,
5002 _cancellation: Option<TurnCancellation>,
5003 ) -> Result<Option<ModelTurnEvent>, LoopError> {
5004 Ok(self.events.pop_front())
5005 }
5006 }
5007
5008 #[derive(Clone)]
5009 struct EchoTool {
5010 spec: ToolSpec,
5011 }
5012
5013 #[derive(Clone)]
5014 struct FailingTool {
5015 spec: ToolSpec,
5016 }
5017
5018 #[derive(Clone)]
5019 struct RunThenDenyTool {
5020 spec: ToolSpec,
5021 }
5022
5023 impl Default for EchoTool {
5024 fn default() -> Self {
5025 Self {
5026 spec: ToolSpec {
5027 name: ToolName::new("echo"),
5028 description: "Echo back a value".into(),
5029 input_schema: json!({
5030 "type": "object",
5031 "properties": {
5032 "value": { "type": "string" }
5033 },
5034 "required": ["value"],
5035 "additionalProperties": false
5036 }),
5037 output_schema: None,
5038 annotations: ToolAnnotations::default(),
5039 metadata: MetadataMap::new(),
5040 },
5041 }
5042 }
5043 }
5044
5045 impl Default for FailingTool {
5046 fn default() -> Self {
5047 Self {
5048 spec: ToolSpec {
5049 name: ToolName::new("failing"),
5050 description: "Always fails after execution starts".into(),
5051 input_schema: json!({
5052 "type": "object",
5053 "properties": {
5054 "value": { "type": "string" }
5055 },
5056 "additionalProperties": true
5057 }),
5058 output_schema: None,
5059 annotations: ToolAnnotations::default(),
5060 metadata: MetadataMap::new(),
5061 },
5062 }
5063 }
5064 }
5065
5066 impl Default for RunThenDenyTool {
5067 fn default() -> Self {
5068 Self {
5069 spec: ToolSpec {
5070 name: ToolName::new("run_then_deny"),
5071 description: "Runs, then returns a permission-denied error".into(),
5072 input_schema: json!({
5073 "type": "object",
5074 "properties": {
5075 "value": { "type": "string" }
5076 },
5077 "additionalProperties": true
5078 }),
5079 output_schema: None,
5080 annotations: ToolAnnotations::default(),
5081 metadata: MetadataMap::new(),
5082 },
5083 }
5084 }
5085 }
5086
5087 #[derive(Clone)]
5088 struct DynamicSpecTool {
5089 spec: ToolSpec,
5090 version: StdArc<AtomicUsize>,
5091 }
5092
5093 impl DynamicSpecTool {
5094 fn new(version: StdArc<AtomicUsize>) -> Self {
5095 Self {
5096 spec: ToolSpec {
5097 name: ToolName::new("dynamic"),
5098 description: "dynamic version 0".into(),
5099 input_schema: json!({
5100 "type": "object",
5101 "properties": {},
5102 "additionalProperties": false
5103 }),
5104 output_schema: None,
5105 annotations: ToolAnnotations::default(),
5106 metadata: MetadataMap::new(),
5107 },
5108 version,
5109 }
5110 }
5111 }
5112
5113 #[async_trait]
5114 impl Tool for EchoTool {
5115 fn spec(&self) -> &ToolSpec {
5116 &self.spec
5117 }
5118
5119 fn proposed_requests(
5120 &self,
5121 request: &agentkit_tools_core::ToolRequest,
5122 ) -> Result<
5123 Vec<Box<dyn agentkit_tools_core::PermissionRequest>>,
5124 agentkit_tools_core::ToolError,
5125 > {
5126 Ok(vec![Box::new(FileSystemPermissionRequest::Read {
5127 path: "/tmp/echo".into(),
5128 metadata: request.metadata.clone(),
5129 })])
5130 }
5131
5132 async fn invoke(
5133 &self,
5134 request: agentkit_tools_core::ToolRequest,
5135 _ctx: &mut ToolContext<'_>,
5136 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
5137 let value = request
5138 .input
5139 .get("value")
5140 .and_then(Value::as_str)
5141 .ok_or_else(|| {
5142 agentkit_tools_core::ToolError::InvalidInput("missing value".into())
5143 })?;
5144
5145 Ok(ToolResult {
5146 result: ToolResultPart {
5147 call_id: request.call_id,
5148 output: ToolOutput::Text(value.into()),
5149 is_error: false,
5150 metadata: MetadataMap::new(),
5151 },
5152 duration: None,
5153 metadata: MetadataMap::new(),
5154 })
5155 }
5156 }
5157
5158 #[async_trait]
5159 impl Tool for FailingTool {
5160 fn spec(&self) -> &ToolSpec {
5161 &self.spec
5162 }
5163
5164 async fn invoke(
5165 &self,
5166 _request: agentkit_tools_core::ToolRequest,
5167 _ctx: &mut ToolContext<'_>,
5168 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
5169 Err(agentkit_tools_core::ToolError::ExecutionFailed(
5170 "runtime failed".into(),
5171 ))
5172 }
5173 }
5174
5175 #[async_trait]
5176 impl Tool for RunThenDenyTool {
5177 fn spec(&self) -> &ToolSpec {
5178 &self.spec
5179 }
5180
5181 async fn invoke(
5182 &self,
5183 _request: agentkit_tools_core::ToolRequest,
5184 _ctx: &mut ToolContext<'_>,
5185 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
5186 Err(agentkit_tools_core::ToolError::PermissionDenied(
5187 PermissionDenial {
5188 code: PermissionCode::CustomPolicyDenied,
5189 message: "remote 403".into(),
5190 metadata: MetadataMap::new(),
5191 },
5192 ))
5193 }
5194 }
5195
5196 #[async_trait]
5197 impl Tool for DynamicSpecTool {
5198 fn spec(&self) -> &ToolSpec {
5199 &self.spec
5200 }
5201
5202 fn current_spec(&self) -> Option<ToolSpec> {
5203 let mut spec = self.spec.clone();
5204 spec.description = format!("dynamic version {}", self.version.load(Ordering::SeqCst));
5205 Some(spec)
5206 }
5207
5208 async fn invoke(
5209 &self,
5210 request: agentkit_tools_core::ToolRequest,
5211 _ctx: &mut ToolContext<'_>,
5212 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
5213 Ok(ToolResult {
5214 result: ToolResultPart {
5215 call_id: request.call_id,
5216 output: ToolOutput::Text("ok".into()),
5217 is_error: false,
5218 metadata: MetadataMap::new(),
5219 },
5220 duration: None,
5221 metadata: MetadataMap::new(),
5222 })
5223 }
5224 }
5225
5226 struct DenyFsReads;
5227
5228 impl PermissionChecker for DenyFsReads {
5229 fn evaluate(
5230 &self,
5231 request: &dyn agentkit_tools_core::PermissionRequest,
5232 ) -> PermissionDecision {
5233 if request.kind() == "filesystem.read" {
5234 return PermissionDecision::Deny(PermissionDenial {
5235 code: PermissionCode::PathNotAllowed,
5236 message: "reads denied in test".into(),
5237 metadata: MetadataMap::new(),
5238 });
5239 }
5240
5241 PermissionDecision::Allow
5242 }
5243 }
5244
5245 struct ApproveFsReads;
5246
5247 impl PermissionChecker for ApproveFsReads {
5248 fn evaluate(
5249 &self,
5250 request: &dyn agentkit_tools_core::PermissionRequest,
5251 ) -> PermissionDecision {
5252 if request.kind() == "filesystem.read" {
5253 return PermissionDecision::RequireApproval(ApprovalRequest {
5254 task_id: None,
5255 call_id: None,
5256 id: "approval:fs-read".into(),
5257 request_kind: request.kind().into(),
5258 reason: agentkit_tools_core::ApprovalReason::SensitivePath,
5259 summary: request.summary(),
5260 metadata: request.metadata().clone(),
5261 });
5262 }
5263
5264 PermissionDecision::Allow
5265 }
5266 }
5267
5268 struct KeepRecentMutator {
5269 keep: usize,
5270 }
5271
5272 #[async_trait]
5273 impl LoopMutator for KeepRecentMutator {
5274 async fn mutate(
5275 &self,
5276 cursor: &mut TranscriptCursor<'_>,
5277 ctx: LoopCtx<'_>,
5278 ) -> Result<(), LoopError> {
5279 if cursor.len() < 2 {
5280 return Ok(());
5281 }
5282 let drop = cursor.len().saturating_sub(self.keep);
5283 ctx.emitter.emit(AgentEvent::MutationStarted {
5284 session_id: ctx.session_id.clone(),
5285 turn_id: ctx.turn_id.cloned(),
5286 mutator: "keep-recent".into(),
5287 point: ctx.point,
5288 });
5289 cursor.drain(..drop);
5290 ctx.emitter.emit(AgentEvent::MutationFinished {
5291 session_id: ctx.session_id.clone(),
5292 turn_id: ctx.turn_id.cloned(),
5293 mutator: "keep-recent".into(),
5294 dirty: true,
5295 metadata: MetadataMap::new(),
5296 });
5297 Ok(())
5298 }
5299 }
5300
5301 struct PointRecordingMutator {
5304 points: StdArc<StdMutex<Vec<MutationPoint>>>,
5305 }
5306
5307 #[async_trait]
5308 impl LoopMutator for PointRecordingMutator {
5309 async fn mutate(
5310 &self,
5311 _cursor: &mut TranscriptCursor<'_>,
5312 ctx: LoopCtx<'_>,
5313 ) -> Result<(), LoopError> {
5314 self.points.lock().unwrap().push(ctx.point);
5315 Ok(())
5316 }
5317 }
5318
5319 struct RecordingObserver {
5320 events: StdArc<StdMutex<Vec<AgentEvent>>>,
5321 }
5322
5323 impl LoopObserver for RecordingObserver {
5324 fn handle_event(&self, event: ObservedEvent) {
5325 let event = event.event;
5326 self.events.lock().unwrap().push(event);
5327 }
5328 }
5329
5330 fn turn_lifecycle_events(
5331 events: &[AgentEvent],
5332 ) -> Vec<(agentkit_core::TurnId, Option<FinishReason>)> {
5333 events
5334 .iter()
5335 .filter_map(|event| match event {
5336 AgentEvent::TurnStarted { turn_id, .. } => Some((turn_id.clone(), None)),
5337 AgentEvent::TurnFinished(turn) => {
5338 Some((turn.turn_id.clone(), Some(turn.finish_reason.clone())))
5339 }
5340 _ => None,
5341 })
5342 .collect()
5343 }
5344
5345 struct CatalogExecutor {
5346 version: AtomicUsize,
5347 events: StdMutex<Vec<ToolCatalogEvent>>,
5348 }
5349
5350 impl CatalogExecutor {
5351 fn new() -> Self {
5352 Self {
5353 version: AtomicUsize::new(0),
5354 events: StdMutex::new(Vec::new()),
5355 }
5356 }
5357
5358 fn publish_change(&self, version: usize, event: ToolCatalogEvent) {
5359 self.version.store(version, Ordering::SeqCst);
5360 self.events.lock().unwrap().push(event);
5361 }
5362 }
5363
5364 #[async_trait]
5365 impl ToolExecutor for CatalogExecutor {
5366 fn specs(&self) -> Vec<ToolSpec> {
5367 vec![ToolSpec {
5368 name: ToolName::new("dynamic"),
5369 description: format!("dynamic version {}", self.version.load(Ordering::SeqCst)),
5370 input_schema: json!({
5371 "type": "object",
5372 "properties": {},
5373 "additionalProperties": false
5374 }),
5375 output_schema: None,
5376 annotations: ToolAnnotations::default(),
5377 metadata: MetadataMap::new(),
5378 }]
5379 }
5380
5381 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
5382 std::mem::take(&mut *self.events.lock().unwrap())
5383 }
5384
5385 async fn execute(
5386 &self,
5387 request: ToolRequest,
5388 _ctx: &mut ToolContext<'_>,
5389 ) -> ToolExecutionOutcome {
5390 ToolExecutionOutcome::Completed(ToolResult {
5391 result: ToolResultPart {
5392 call_id: request.call_id,
5393 output: ToolOutput::Text("dynamic-ok".into()),
5394 is_error: false,
5395 metadata: MetadataMap::new(),
5396 },
5397 duration: None,
5398 metadata: MetadataMap::new(),
5399 })
5400 }
5401 }
5402
5403 #[derive(Clone)]
5404 struct BlockingTool {
5405 spec: ToolSpec,
5406 entered: StdArc<AtomicBool>,
5407 release: StdArc<Notify>,
5408 output: &'static str,
5409 }
5410
5411 impl BlockingTool {
5412 fn new(
5413 name: &str,
5414 entered: StdArc<AtomicBool>,
5415 release: StdArc<Notify>,
5416 output: &'static str,
5417 ) -> Self {
5418 Self {
5419 spec: ToolSpec {
5420 name: ToolName::new(name),
5421 description: format!("blocking tool {name}"),
5422 input_schema: json!({
5423 "type": "object",
5424 "properties": {},
5425 "additionalProperties": false
5426 }),
5427 output_schema: None,
5428 annotations: ToolAnnotations::default(),
5429 metadata: MetadataMap::new(),
5430 },
5431 entered,
5432 release,
5433 output,
5434 }
5435 }
5436 }
5437
5438 #[async_trait]
5439 impl Tool for BlockingTool {
5440 fn spec(&self) -> &ToolSpec {
5441 &self.spec
5442 }
5443
5444 async fn invoke(
5445 &self,
5446 request: agentkit_tools_core::ToolRequest,
5447 _ctx: &mut ToolContext<'_>,
5448 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
5449 self.entered.store(true, Ordering::SeqCst);
5450 self.release.notified().await;
5451 Ok(ToolResult {
5452 result: ToolResultPart {
5453 call_id: request.call_id,
5454 output: ToolOutput::Text(self.output.into()),
5455 is_error: false,
5456 metadata: MetadataMap::new(),
5457 },
5458 duration: None,
5459 metadata: MetadataMap::new(),
5460 })
5461 }
5462 }
5463
5464 struct NameRoutingPolicy {
5465 routes: Vec<(String, RoutingDecision)>,
5466 }
5467
5468 impl NameRoutingPolicy {
5469 fn new(routes: impl IntoIterator<Item = (impl Into<String>, RoutingDecision)>) -> Self {
5470 Self {
5471 routes: routes
5472 .into_iter()
5473 .map(|(name, decision)| (name.into(), decision))
5474 .collect(),
5475 }
5476 }
5477 }
5478
5479 impl TaskRoutingPolicy for NameRoutingPolicy {
5480 fn route(&self, request: &ToolRequest) -> RoutingDecision {
5481 self.routes
5482 .iter()
5483 .find(|(name, _)| name == &request.tool_name.0)
5484 .map(|(_, decision)| *decision)
5485 .unwrap_or(RoutingDecision::Foreground)
5486 }
5487 }
5488
5489 async fn wait_for_task_event(handle: &TaskManagerHandle) -> TaskEvent {
5490 timeout(Duration::from_secs(1), handle.next_event())
5491 .await
5492 .expect("timed out waiting for task event")
5493 .expect("task event stream ended unexpectedly")
5494 }
5495
5496 async fn wait_until_entered(flag: &AtomicBool) {
5497 timeout(Duration::from_secs(1), async {
5498 while !flag.load(Ordering::SeqCst) {
5499 tokio::task::yield_now().await;
5500 }
5501 })
5502 .await
5503 .expect("task never entered execution");
5504 }
5505
5506 async fn wait_until_completed(handle: &TaskManagerHandle) {
5507 timeout(Duration::from_secs(1), async {
5508 while handle.list_completed().await.is_empty() {
5509 tokio::task::yield_now().await;
5510 }
5511 })
5512 .await
5513 .expect("task never completed");
5514 }
5515
5516 #[tokio::test]
5517 async fn loop_continues_after_completed_tool_call() {
5518 let events = StdArc::new(StdMutex::new(Vec::new()));
5519 let tools = ToolRegistry::new().with(EchoTool::default());
5520 let agent = Agent::builder()
5521 .model(FakeAdapter)
5522 .add_tool_source(tools)
5523 .permissions(AllowAllPermissions)
5524 .observer(RecordingObserver {
5525 events: events.clone(),
5526 })
5527 .build()
5528 .unwrap();
5529
5530 let mut driver = agent
5531 .start(SessionConfig {
5532 session_id: SessionId::new("session-1"),
5533 metadata: MetadataMap::new(),
5534 cache: None,
5535 })
5536 .await
5537 .unwrap();
5538
5539 driver
5540 .submit_input(vec![Item {
5541 id: None,
5542 kind: ItemKind::User,
5543 parts: vec![Part::Text(TextPart {
5544 text: "ping".into(),
5545 metadata: MetadataMap::new(),
5546 })],
5547 metadata: MetadataMap::new(),
5548 usage: None,
5549 finish_reason: None,
5550 created_at: None,
5551 }])
5552 .unwrap();
5553
5554 let result = run_until_finished(&mut driver).await;
5555
5556 match result {
5557 LoopStep::Finished(turn) => {
5558 assert_eq!(turn.finish_reason, FinishReason::Completed);
5559 assert_eq!(turn.items.len(), 1);
5560 match &turn.items[0].parts[0] {
5561 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
5562 other => panic!("unexpected part: {other:?}"),
5563 }
5564 }
5565 other => panic!("unexpected loop step: {other:?}"),
5566 }
5567
5568 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
5569 assert_eq!(lifecycle.len(), 2);
5570 assert_eq!(lifecycle[0].0, lifecycle[1].0);
5571 assert_eq!(lifecycle[0].1, None);
5572 assert_eq!(lifecycle[1].1, Some(FinishReason::Completed));
5573 }
5574
5575 async fn run_until_finished<S: ModelSession + Send>(driver: &mut LoopDriver<S>) -> LoopStep {
5579 loop {
5580 match driver.next().await.unwrap() {
5581 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue,
5582 step => return step,
5583 }
5584 }
5585 }
5586
5587 #[tokio::test]
5594 async fn post_tool_continuation_reports_after_tool_result_mutation_point() {
5595 let points = StdArc::new(StdMutex::new(Vec::<MutationPoint>::new()));
5596 let tools = ToolRegistry::new().with(EchoTool::default());
5597 let agent = Agent::builder()
5598 .model(FakeAdapter)
5599 .add_tool_source(tools)
5600 .permissions(AllowAllPermissions)
5601 .mutator(PointRecordingMutator {
5602 points: points.clone(),
5603 })
5604 .build()
5605 .unwrap();
5606
5607 let mut driver = agent
5608 .start(SessionConfig {
5609 session_id: SessionId::new("session-mutation-point"),
5610 metadata: MetadataMap::new(),
5611 cache: None,
5612 })
5613 .await
5614 .unwrap();
5615
5616 driver
5617 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5618 .unwrap();
5619
5620 let _ = run_until_finished(&mut driver).await;
5622
5623 let recorded = points.lock().unwrap().clone();
5624 assert_eq!(
5625 recorded.first(),
5626 Some(&MutationPoint::AfterTurnEnded),
5627 "first drive of a fresh turn must report AfterTurnEnded, got {recorded:?}"
5628 );
5629 assert!(
5630 recorded.contains(&MutationPoint::AfterToolResult),
5631 "post-tool continuation must report AfterToolResult, got {recorded:?}"
5632 );
5633 }
5634
5635 #[tokio::test]
5636 async fn no_work_awaiting_input_emits_no_turn_lifecycle() {
5637 let events = StdArc::new(StdMutex::new(Vec::new()));
5638 let agent = Agent::builder()
5639 .model(SlowAdapter)
5640 .observer(RecordingObserver {
5641 events: events.clone(),
5642 })
5643 .build()
5644 .unwrap();
5645 let mut driver = agent
5646 .start(SessionConfig::new("session-no-work"))
5647 .await
5648 .unwrap();
5649
5650 for _ in 0..2 {
5651 assert!(matches!(
5652 driver.next().await.unwrap(),
5653 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
5654 ));
5655 }
5656
5657 assert!(turn_lifecycle_events(&events.lock().unwrap()).is_empty());
5658 }
5659
5660 #[tokio::test]
5661 async fn normal_turn_emits_one_matched_lifecycle_pair() {
5662 let events = StdArc::new(StdMutex::new(Vec::new()));
5663 let agent = Agent::builder()
5664 .model(SlowAdapter)
5665 .observer(RecordingObserver {
5666 events: events.clone(),
5667 })
5668 .build()
5669 .unwrap();
5670 let mut driver = agent
5671 .start(SessionConfig::new("session-normal-lifecycle"))
5672 .await
5673 .unwrap();
5674 driver
5675 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5676 .unwrap();
5677
5678 assert!(matches!(
5679 driver.next().await.unwrap(),
5680 LoopStep::Finished(TurnResult {
5681 finish_reason: FinishReason::Completed,
5682 ..
5683 })
5684 ));
5685
5686 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
5687 assert_eq!(lifecycle.len(), 2);
5688 assert_eq!(lifecycle[0].0, lifecycle[1].0);
5689 assert_eq!(lifecycle[0].1, None);
5690 assert_eq!(lifecycle[1].1, Some(FinishReason::Completed));
5691 }
5692
5693 #[tokio::test]
5694 async fn post_start_error_emits_terminal_error_without_run_failed() {
5695 let events = StdArc::new(StdMutex::new(Vec::new()));
5696 let agent = Agent::builder()
5697 .model(SlowAdapter)
5698 .mutator(ErrorMutator)
5699 .observer(RecordingObserver {
5700 events: events.clone(),
5701 })
5702 .build()
5703 .unwrap();
5704 let mut driver = agent
5705 .start(SessionConfig::new("session-error-lifecycle"))
5706 .await
5707 .unwrap();
5708 driver
5709 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5710 .unwrap();
5711
5712 assert!(matches!(
5713 driver.next().await,
5714 Err(LoopError::Mutator(message)) if message == "boom"
5715 ));
5716
5717 let events = events.lock().unwrap();
5718 let lifecycle = turn_lifecycle_events(&events);
5719 assert_eq!(lifecycle.len(), 2);
5720 assert_eq!(lifecycle[0].0, lifecycle[1].0);
5721 assert_eq!(lifecycle[0].1, None);
5722 assert_eq!(lifecycle[1].1, Some(FinishReason::Error));
5723 assert!(
5724 !events
5725 .iter()
5726 .any(|event| matches!(event, AgentEvent::RunFailed { .. }))
5727 );
5728 }
5729
5730 #[tokio::test]
5731 async fn active_tool_error_repairs_state_and_retry_uses_fresh_lifecycle() {
5732 let events = StdArc::new(StdMutex::new(Vec::new()));
5733 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
5734 let agent = Agent::builder()
5735 .model(FakeAdapter)
5736 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
5737 .task_manager(
5738 TestTaskManager::new(SimpleTaskManager::new())
5739 .fail_start("original start failure")
5740 .record_interrupts(interrupted.clone())
5741 .fail_interrupt("cleanup failure"),
5742 )
5743 .observer(RecordingObserver {
5744 events: events.clone(),
5745 })
5746 .build()
5747 .unwrap();
5748 let mut driver = agent
5749 .start(SessionConfig::new("session-active-tool-error"))
5750 .await
5751 .unwrap();
5752 driver
5753 .submit_input(vec![Item::text(ItemKind::User, "first")])
5754 .unwrap();
5755
5756 let error = driver.next().await.unwrap_err();
5757 assert!(error.to_string().contains("original start failure"));
5758 assert!(!error.to_string().contains("cleanup failure"));
5759 assert!(driver.active_tool_round.is_none());
5760 assert!(driver.pending_round_resume.is_none());
5761 assert!(unanswered_tool_calls(&driver.snapshot().transcript).is_empty());
5762 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
5763 assert_eq!(interrupted.lock().unwrap().len(), 1);
5764
5765 driver
5766 .submit_input(vec![Item::text(ItemKind::User, "retry")])
5767 .unwrap();
5768 assert!(matches!(
5769 driver.next().await.unwrap(),
5770 LoopStep::Finished(TurnResult {
5771 finish_reason: FinishReason::Completed,
5772 ..
5773 })
5774 ));
5775
5776 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
5777 assert_eq!(lifecycle.len(), 4, "{lifecycle:?}");
5778 assert_eq!(lifecycle[0].0, lifecycle[1].0);
5779 assert_eq!(lifecycle[1].1, Some(FinishReason::Error));
5780 assert_eq!(lifecycle[2].0, lifecycle[3].0);
5781 assert_eq!(lifecycle[3].1, Some(FinishReason::Completed));
5782 assert_ne!(lifecycle[0].0, lifecycle[2].0);
5783 }
5784
5785 #[tokio::test]
5786 async fn continuation_error_clears_resume_and_retry_uses_fresh_lifecycle() {
5787 let events = StdArc::new(StdMutex::new(Vec::new()));
5788 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
5789 let agent = Agent::builder()
5790 .model(FakeAdapter)
5791 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
5792 .task_manager(
5793 TestTaskManager::new(SimpleTaskManager::new())
5794 .fail_pending_update_on(1, "original continuation failure")
5795 .record_interrupts(interrupted.clone())
5796 .fail_interrupt("cleanup failure"),
5797 )
5798 .observer(RecordingObserver {
5799 events: events.clone(),
5800 })
5801 .build()
5802 .unwrap();
5803 let mut driver = agent
5804 .start(SessionConfig::new("session-continuation-error"))
5805 .await
5806 .unwrap();
5807 driver
5808 .submit_input(vec![Item::text(ItemKind::User, "first")])
5809 .unwrap();
5810
5811 assert!(matches!(
5812 driver.next().await.unwrap(),
5813 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_))
5814 ));
5815 let error = driver.next().await.unwrap_err();
5816 assert!(error.to_string().contains("original continuation failure"));
5817 assert!(!error.to_string().contains("cleanup failure"));
5818 assert!(driver.pending_round_resume.is_none());
5819 assert!(driver.active_tool_round.is_none());
5820 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
5821 assert_eq!(interrupted.lock().unwrap().len(), 1);
5822
5823 driver
5824 .submit_input(vec![Item::text(ItemKind::User, "retry")])
5825 .unwrap();
5826 assert!(matches!(
5827 driver.next().await.unwrap(),
5828 LoopStep::Finished(TurnResult {
5829 finish_reason: FinishReason::Completed,
5830 ..
5831 })
5832 ));
5833
5834 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
5835 assert_eq!(lifecycle.len(), 4, "{lifecycle:?}");
5836 assert_eq!(lifecycle[0].0, lifecycle[1].0);
5837 assert_eq!(lifecycle[1].1, Some(FinishReason::Error));
5838 assert_eq!(lifecycle[2].0, lifecycle[3].0);
5839 assert_eq!(lifecycle[3].1, Some(FinishReason::Completed));
5840 assert_ne!(lifecycle[0].0, lifecycle[2].0);
5841 }
5842
5843 #[test]
5844 fn pending_input_requires_input_bearing_tail_role() {
5845 assert!(!transcript_has_pending_input(&[]));
5846 assert!(!transcript_has_pending_input(&[Item::text(
5847 ItemKind::System,
5848 "system"
5849 )]));
5850 assert!(!transcript_has_pending_input(&[Item::text(
5851 ItemKind::Developer,
5852 "developer"
5853 )]));
5854 assert!(!transcript_has_pending_input(&[Item::text(
5855 ItemKind::Context,
5856 "context"
5857 )]));
5858 assert!(!transcript_has_pending_input(&[Item::text(
5859 ItemKind::Assistant,
5860 "assistant"
5861 )]));
5862
5863 assert!(transcript_has_pending_input(&[Item::text(
5864 ItemKind::User,
5865 "user"
5866 )]));
5867 assert!(transcript_has_pending_input(&[Item::notification(
5868 "background update"
5869 )]));
5870 assert!(transcript_has_pending_input(&[Item {
5871 id: None,
5872 kind: ItemKind::Tool,
5873 parts: vec![Part::ToolResult(ToolResultPart {
5874 call_id: ToolCallId::new("call-test"),
5875 output: ToolOutput::Text("ok".into()),
5876 is_error: false,
5877 metadata: MetadataMap::new(),
5878 })],
5879 metadata: MetadataMap::new(),
5880 usage: None,
5881 finish_reason: None,
5882 created_at: None,
5883 }]));
5884 }
5885
5886 struct DropTrailingUserMutator;
5892 struct ErrorMutator;
5893
5894 #[async_trait]
5895 impl LoopMutator for ErrorMutator {
5896 async fn mutate(
5897 &self,
5898 _cursor: &mut TranscriptCursor<'_>,
5899 _ctx: LoopCtx<'_>,
5900 ) -> Result<(), LoopError> {
5901 Err(LoopError::Mutator("boom".into()))
5902 }
5903 }
5904
5905 #[async_trait]
5906 impl LoopMutator for DropTrailingUserMutator {
5907 async fn mutate(
5908 &self,
5909 cursor: &mut TranscriptCursor<'_>,
5910 _ctx: LoopCtx<'_>,
5911 ) -> Result<(), LoopError> {
5912 if cursor.last().map(|item| item.kind) == Some(ItemKind::User) {
5913 cursor.pop();
5914 }
5915 Ok(())
5916 }
5917 }
5918
5919 struct RejectAssistantPrefillAdapter {
5924 saw_assistant_tail: StdArc<AtomicBool>,
5925 }
5926
5927 struct RejectAssistantPrefillSession {
5928 saw_assistant_tail: StdArc<AtomicBool>,
5929 }
5930
5931 #[async_trait]
5932 impl ModelAdapter for RejectAssistantPrefillAdapter {
5933 type Session = RejectAssistantPrefillSession;
5934
5935 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
5936 Ok(RejectAssistantPrefillSession {
5937 saw_assistant_tail: self.saw_assistant_tail.clone(),
5938 })
5939 }
5940 }
5941
5942 #[async_trait]
5943 impl ModelSession for RejectAssistantPrefillSession {
5944 type Turn = FakeTurn;
5945
5946 async fn begin_turn(
5947 &mut self,
5948 request: TurnRequest,
5949 _cancellation: Option<TurnCancellation>,
5950 ) -> Result<Self::Turn, LoopError> {
5951 if request.transcript.last().map(|item| item.kind) == Some(ItemKind::Assistant) {
5952 self.saw_assistant_tail.store(true, Ordering::SeqCst);
5953 return Err(LoopError::Provider(
5954 "conversation must end with a user message".into(),
5955 ));
5956 }
5957 Ok(FakeTurn {
5958 events: VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
5959 model: None,
5960 response_id: None,
5961 finish_reason: FinishReason::Completed,
5962 output_items: vec![Item::text(ItemKind::Assistant, "ok")],
5963 usage: None,
5964 metadata: MetadataMap::new(),
5965 })]),
5966 })
5967 }
5968 }
5969
5970 #[tokio::test]
5978 async fn drive_does_not_dispatch_without_valid_trailing_input() {
5979 let saw_assistant_tail = StdArc::new(AtomicBool::new(false));
5980 let events = StdArc::new(StdMutex::new(Vec::new()));
5981 let agent = Agent::builder()
5982 .model(RejectAssistantPrefillAdapter {
5983 saw_assistant_tail: saw_assistant_tail.clone(),
5984 })
5985 .mutator(DropTrailingUserMutator)
5986 .observer(RecordingObserver {
5987 events: events.clone(),
5988 })
5989 .transcript(vec![
5992 Item::text(ItemKind::User, "kickoff"),
5993 Item::text(ItemKind::Assistant, "prior reply"),
5994 ])
5995 .build()
5996 .unwrap();
5997
5998 let mut driver = agent
5999 .start(SessionConfig {
6000 session_id: SessionId::new("session-no-valid-input"),
6001 metadata: MetadataMap::new(),
6002 cache: None,
6003 })
6004 .await
6005 .unwrap();
6006
6007 driver
6008 .submit_input(vec![Item::text(ItemKind::User, "follow up")])
6009 .unwrap();
6010
6011 let outcome = driver.next().await;
6013
6014 assert!(
6015 !saw_assistant_tail.load(Ordering::SeqCst),
6016 "loop dispatched a model turn whose transcript ends in an assistant \
6017 message (outcome: {outcome:?}); with no valid trailing input the turn \
6018 must finish instead of driving"
6019 );
6020 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
6021 assert_eq!(lifecycle.len(), 2);
6022 assert_eq!(lifecycle[0].0, lifecycle[1].0);
6023 assert_eq!(lifecycle[0].1, None);
6024 assert_eq!(lifecycle[1].1, Some(FinishReason::Completed));
6025 }
6026
6027 #[tokio::test]
6028 async fn loop_uses_injected_permission_checker() {
6029 let events = StdArc::new(StdMutex::new(Vec::new()));
6030 let tools = ToolRegistry::new().with(EchoTool::default());
6031 let agent = Agent::builder()
6032 .model(FakeAdapter)
6033 .add_tool_source(tools)
6034 .permissions(DenyFsReads)
6035 .observer(RecordingObserver {
6036 events: events.clone(),
6037 })
6038 .build()
6039 .unwrap();
6040
6041 let mut driver = agent
6042 .start(SessionConfig {
6043 session_id: SessionId::new("session-2"),
6044 metadata: MetadataMap::new(),
6045 cache: None,
6046 })
6047 .await
6048 .unwrap();
6049
6050 driver
6051 .submit_input(vec![Item {
6052 id: None,
6053 kind: ItemKind::User,
6054 parts: vec![Part::Text(TextPart {
6055 text: "ping".into(),
6056 metadata: MetadataMap::new(),
6057 })],
6058 metadata: MetadataMap::new(),
6059 usage: None,
6060 finish_reason: None,
6061 created_at: None,
6062 }])
6063 .unwrap();
6064
6065 let result = run_until_finished(&mut driver).await;
6066
6067 match result {
6068 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
6069 Part::Text(text) => assert!(text.text.contains("tool permission denied")),
6070 other => panic!("unexpected part: {other:?}"),
6071 },
6072 other => panic!("unexpected loop step: {other:?}"),
6073 }
6074
6075 assert!(
6076 events
6077 .lock()
6078 .unwrap()
6079 .iter()
6080 .all(|event| !matches!(event, AgentEvent::ToolExecutionStarted(_))),
6081 "denied tools must not be reported as started"
6082 );
6083 }
6084
6085 #[tokio::test]
6086 async fn failed_tool_execution_still_reports_started() {
6087 let events = StdArc::new(StdMutex::new(Vec::new()));
6088 let tools = ToolRegistry::new().with(FailingTool::default());
6089 let agent = Agent::builder()
6090 .model(FakeAdapter)
6091 .add_tool_source(tools)
6092 .permissions(AllowAllPermissions)
6093 .observer(RecordingObserver {
6094 events: events.clone(),
6095 })
6096 .build()
6097 .unwrap();
6098
6099 let mut driver = agent
6100 .start(SessionConfig {
6101 session_id: SessionId::new("session-failing-start-event"),
6102 metadata: MetadataMap::new(),
6103 cache: None,
6104 })
6105 .await
6106 .unwrap();
6107
6108 driver
6109 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6110 .unwrap();
6111
6112 match run_until_finished(&mut driver).await {
6113 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
6114 other => panic!("unexpected loop step: {other:?}"),
6115 }
6116
6117 let events = events.lock().unwrap();
6118 assert!(events.iter().any(|event| matches!(
6119 event,
6120 AgentEvent::ToolExecutionStarted(call) if call.name == "failing"
6121 )));
6122 assert!(events.iter().any(|event| matches!(
6123 event,
6124 AgentEvent::ToolResultReceived(result) if result.is_error
6125 )));
6126 }
6127
6128 #[tokio::test]
6129 async fn run_then_deny_tool_execution_still_reports_started() {
6130 let events = StdArc::new(StdMutex::new(Vec::new()));
6131 let tools = ToolRegistry::new().with(RunThenDenyTool::default());
6132 let agent = Agent::builder()
6133 .model(FakeAdapter)
6134 .add_tool_source(tools)
6135 .permissions(AllowAllPermissions)
6136 .observer(RecordingObserver {
6137 events: events.clone(),
6138 })
6139 .build()
6140 .unwrap();
6141
6142 let mut driver = agent
6143 .start(SessionConfig {
6144 session_id: SessionId::new("session-run-then-deny-start-event"),
6145 metadata: MetadataMap::new(),
6146 cache: None,
6147 })
6148 .await
6149 .unwrap();
6150
6151 driver
6152 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6153 .unwrap();
6154
6155 match run_until_finished(&mut driver).await {
6156 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
6157 other => panic!("unexpected loop step: {other:?}"),
6158 }
6159
6160 let events = events.lock().unwrap();
6161 assert!(events.iter().any(|event| matches!(
6162 event,
6163 AgentEvent::ToolExecutionStarted(call) if call.name == "run_then_deny"
6164 )));
6165 assert!(events.iter().any(|event| matches!(
6168 event,
6169 AgentEvent::ToolResultReceived(result)
6170 if result.is_error
6171 && result
6172 .metadata
6173 .get(TOOL_RESULT_FAILURE_KIND_METADATA_KEY)
6174 .and_then(Value::as_str)
6175 == Some(TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED)
6176 && result
6177 .metadata
6178 .get(TOOL_RESULT_NOT_STARTED_METADATA_KEY)
6179 .is_none()
6180 )));
6181 }
6182
6183 #[tokio::test]
6184 async fn async_task_manager_background_round_requires_explicit_continue() {
6185 let events = StdArc::new(StdMutex::new(Vec::new()));
6186 let entered = StdArc::new(AtomicBool::new(false));
6187 let release = StdArc::new(Notify::new());
6188 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
6189 "background-wait",
6190 RoutingDecision::Background,
6191 )]));
6192 let handle = task_manager.handle();
6193 let tools = ToolRegistry::new().with(BlockingTool::new(
6194 "background-wait",
6195 entered.clone(),
6196 release.clone(),
6197 "background-done",
6198 ));
6199 let agent = Agent::builder()
6200 .model(FakeAdapter)
6201 .add_tool_source(tools)
6202 .permissions(AllowAllPermissions)
6203 .task_manager(task_manager)
6204 .observer(RecordingObserver {
6205 events: events.clone(),
6206 })
6207 .build()
6208 .unwrap();
6209
6210 let mut driver = agent
6211 .start(SessionConfig {
6212 session_id: SessionId::new("session-background"),
6213 metadata: MetadataMap::new(),
6214 cache: None,
6215 })
6216 .await
6217 .unwrap();
6218
6219 driver
6220 .submit_input(vec![Item {
6221 id: None,
6222 kind: ItemKind::User,
6223 parts: vec![Part::Text(TextPart {
6224 text: "ping".into(),
6225 metadata: MetadataMap::new(),
6226 })],
6227 metadata: MetadataMap::new(),
6228 usage: None,
6229 finish_reason: None,
6230 created_at: None,
6231 }])
6232 .unwrap();
6233
6234 let first = driver.next().await.unwrap();
6235 match first {
6236 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
6237 other => panic!("unexpected first loop step: {other:?}"),
6238 }
6239
6240 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
6241 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
6242 assert_eq!(lifecycle.len(), 2);
6243 assert_eq!(lifecycle[0].0, lifecycle[1].0);
6244 assert_eq!(lifecycle[1].1, Some(FinishReason::ToolCall));
6245
6246 match wait_for_task_event(&handle).await {
6247 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "background-wait"),
6248 other => panic!("unexpected task event: {other:?}"),
6249 }
6250 wait_until_entered(entered.as_ref()).await;
6251 release.notify_waiters();
6252
6253 match wait_for_task_event(&handle).await {
6254 TaskEvent::Completed(_, result) => {
6255 assert_eq!(result.output, ToolOutput::Text("background-done".into()))
6256 }
6257 other => panic!("unexpected completion event: {other:?}"),
6258 }
6259
6260 let resumed = driver.next().await.unwrap();
6261 match resumed {
6262 LoopStep::Finished(turn) => {
6263 assert_eq!(turn.finish_reason, FinishReason::Completed);
6264 match &turn.items[0].parts[0] {
6265 Part::Text(text) => assert_eq!(
6266 text.text,
6267 "tool said: Background tool results: 1 total, 0 failed, 0 with metadata. \
6268 call-1 completed: text preview: background-done"
6269 ),
6270 other => panic!("unexpected part after resume: {other:?}"),
6271 }
6272 }
6273 other => panic!("unexpected resumed step: {other:?}"),
6274 }
6275
6276 let events = events.lock().unwrap();
6277 let lifecycle = turn_lifecycle_events(&events);
6278 assert_eq!(lifecycle.len(), 4);
6279 assert_eq!(lifecycle[0].0, lifecycle[1].0);
6280 assert_eq!(lifecycle[2].0, lifecycle[3].0);
6281 assert_ne!(lifecycle[0].0, lifecycle[2].0);
6282 assert_eq!(lifecycle[3].1, Some(FinishReason::Completed));
6283
6284 let terminal_results: Vec<_> = events
6285 .iter()
6286 .filter_map(|event| match event {
6287 AgentEvent::ToolResultReceived(result)
6288 if result.call_id == ToolCallId::new("call-1") =>
6289 {
6290 Some(result)
6291 }
6292 _ => None,
6293 })
6294 .collect();
6295 assert_eq!(
6296 terminal_results.len(),
6297 1,
6298 "background completion must emit one terminal result event per call: {events:?}"
6299 );
6300 }
6301
6302 #[tokio::test]
6303 async fn detached_parts_notification_preserves_full_output_and_metadata() {
6304 let agent = Agent::builder().model(FakeAdapter).build().unwrap();
6305 let mut driver = agent
6306 .start(SessionConfig::new("session-detached-parts"))
6307 .await
6308 .unwrap();
6309 let call_id = ToolCallId::new("parts-call");
6310 driver.detached_call_ids.insert(call_id.clone());
6311 let parts = vec![
6312 Part::text("part text"),
6313 Part::structured(json!({
6314 "nested": [1, 2, 3]
6315 })),
6316 ];
6317 let mut metadata = MetadataMap::new();
6318 metadata.insert("source".into(), json!("background"));
6319 let result = ToolResultPart {
6320 call_id,
6321 output: ToolOutput::Parts(parts.clone()),
6322 is_error: true,
6323 metadata: metadata.clone(),
6324 };
6325 let mut item_metadata = MetadataMap::new();
6326 item_metadata.insert("delivery".into(), json!("deferred"));
6327 let item = Item::new(ItemKind::Tool, vec![Part::ToolResult(result.clone())])
6328 .with_metadata(item_metadata.clone());
6329
6330 let converted = driver.maybe_convert_detached(item);
6331 let (text, structured) = match converted.parts.as_slice() {
6332 [Part::Text(text), Part::Structured(structured)] => (text, structured),
6333 other => panic!("unexpected converted parts: {other:?}"),
6334 };
6335 assert_eq!(converted.kind, ItemKind::Notification);
6336 assert_eq!(converted.metadata, item_metadata);
6337 assert_eq!(structured.value, serde_json::to_value(&result).unwrap());
6338 assert_eq!(
6339 text.text,
6340 "Background tool results: 1 total, 1 failed, 1 with metadata. \
6341 parts-call failed: parts payload (2 parts)"
6342 );
6343 assert!(!text.text.contains("part text"));
6344 assert!(!text.text.contains("background"));
6345 }
6346
6347 #[tokio::test]
6348 async fn detached_files_notification_preserves_full_output() {
6349 let agent = Agent::builder().model(FakeAdapter).build().unwrap();
6350 let mut driver = agent
6351 .start(SessionConfig::new("session-detached-files"))
6352 .await
6353 .unwrap();
6354 let call_id = ToolCallId::new("files-call");
6355 driver.detached_call_ids.insert(call_id.clone());
6356 let files = vec![
6357 agentkit_core::FilePart::named("report.txt", DataRef::inline_text("full file body"))
6358 .with_mime_type("text/plain"),
6359 agentkit_core::FilePart::named(
6360 "remote.json",
6361 DataRef::uri("https://example.test/remote.json"),
6362 ),
6363 ];
6364 let mut result_metadata = MetadataMap::new();
6365 result_metadata.insert("archive".into(), json!(true));
6366 let result = ToolResultPart::success(call_id, ToolOutput::Files(files.clone()))
6367 .with_metadata(result_metadata);
6368 let mut item_metadata = MetadataMap::new();
6369 item_metadata.insert("delivery".into(), json!("deferred"));
6370 let item = Item::new(ItemKind::Tool, vec![Part::ToolResult(result.clone())])
6371 .with_metadata(item_metadata.clone());
6372
6373 let converted = driver.maybe_convert_detached(item);
6374 let (text, structured) = match converted.parts.as_slice() {
6375 [Part::Text(text), Part::Structured(structured)] => (text, structured),
6376 other => panic!("unexpected converted files: {other:?}"),
6377 };
6378 assert_eq!(converted.kind, ItemKind::Notification);
6379 assert_eq!(converted.metadata, item_metadata);
6380 assert_eq!(structured.value, serde_json::to_value(&result).unwrap());
6381 assert_eq!(
6382 text.text,
6383 "Background tool results: 1 total, 0 failed, 1 with metadata. \
6384 files-call completed: files payload (2 files)"
6385 );
6386 assert!(!text.text.contains("full file body"));
6387 assert!(!text.text.contains("remote.json"));
6388 }
6389
6390 #[test]
6391 fn detached_result_summaries_are_bounded_and_do_not_serialize_structured_payloads() {
6392 let long_text = "é".repeat(DETACHED_TEXT_PREVIEW_MAX_CHARS + 20);
6393 let text_summary = render_tool_output_brief(&ToolOutput::Text(long_text.clone()));
6394 assert_eq!(
6395 text_summary.chars().count(),
6396 "text preview: ".chars().count() + DETACHED_TEXT_PREVIEW_MAX_CHARS
6397 );
6398 assert!(text_summary.ends_with('…'));
6399 assert!(!text_summary.contains(&long_text));
6400
6401 let secret = "structured payload must remain out of notification text";
6402 let structured = ToolOutput::Structured(json!({ "secret": secret }));
6403 assert_eq!(render_tool_output_brief(&structured), "structured payload");
6404
6405 let oversized = "x".repeat(DETACHED_NOTIFICATION_TEXT_MAX_CHARS + 20);
6406 let bounded = truncate_chars(&oversized, DETACHED_NOTIFICATION_TEXT_MAX_CHARS);
6407 assert_eq!(
6408 bounded.chars().count(),
6409 DETACHED_NOTIFICATION_TEXT_MAX_CHARS
6410 );
6411 assert!(bounded.ends_with('…'));
6412 }
6413
6414 #[tokio::test]
6415 async fn detached_tool_placeholder_is_progress_not_terminal_result() {
6416 let events = StdArc::new(StdMutex::new(Vec::new()));
6417 let entered = StdArc::new(AtomicBool::new(false));
6418 let release = StdArc::new(Notify::new());
6419 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
6420 "detaching-wait",
6421 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10)),
6422 )]));
6423 let handle = task_manager.handle();
6424 let tools = ToolRegistry::new().with(BlockingTool::new(
6425 "detaching-wait",
6426 entered.clone(),
6427 release.clone(),
6428 "detached-done",
6429 ));
6430 let agent = Agent::builder()
6431 .model(FakeAdapter)
6432 .add_tool_source(tools)
6433 .permissions(AllowAllPermissions)
6434 .task_manager(task_manager)
6435 .observer(RecordingObserver {
6436 events: events.clone(),
6437 })
6438 .build()
6439 .unwrap();
6440
6441 let mut driver = agent
6442 .start(SessionConfig {
6443 session_id: SessionId::new("session-detached-progress"),
6444 metadata: MetadataMap::new(),
6445 cache: None,
6446 })
6447 .await
6448 .unwrap();
6449
6450 driver
6451 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6452 .unwrap();
6453
6454 match driver.next().await.unwrap() {
6455 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => {}
6456 other => panic!("unexpected detach step: {other:?}"),
6457 }
6458
6459 match wait_for_task_event(&handle).await {
6460 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "detaching-wait"),
6461 other => panic!("unexpected task event: {other:?}"),
6462 }
6463 match wait_for_task_event(&handle).await {
6464 TaskEvent::Detached(snapshot) => assert_eq!(snapshot.tool_name, "detaching-wait"),
6465 other => panic!("unexpected detach event: {other:?}"),
6466 }
6467 wait_until_entered(entered.as_ref()).await;
6468 release.notify_waiters();
6469
6470 match wait_for_task_event(&handle).await {
6471 TaskEvent::Completed(_, result) => {
6472 assert_eq!(result.output, ToolOutput::Text("detached-done".into()))
6473 }
6474 other => panic!("unexpected completion event: {other:?}"),
6475 }
6476
6477 match driver.next().await.unwrap() {
6478 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
6479 other => panic!("unexpected resumed step: {other:?}"),
6480 }
6481
6482 let events = events.lock().unwrap();
6483 assert!(events.iter().any(|event| matches!(
6484 event,
6485 AgentEvent::ToolExecutionProgress(result)
6486 if result.call_id == ToolCallId::new("call-1") && !result.is_error
6487 )));
6488 let terminal_results: Vec<_> = events
6489 .iter()
6490 .filter_map(|event| match event {
6491 AgentEvent::ToolResultReceived(result)
6492 if result.call_id == ToolCallId::new("call-1") =>
6493 {
6494 Some(result)
6495 }
6496 _ => None,
6497 })
6498 .collect();
6499 assert_eq!(
6500 terminal_results.len(),
6501 1,
6502 "detached call must emit one terminal result event: {events:?}"
6503 );
6504 }
6505
6506 #[tokio::test]
6507 async fn cancelled_background_approval_auto_resolves_when_drained() {
6508 let controller = CancellationController::new();
6509 let events = StdArc::new(StdMutex::new(Vec::new()));
6510 let entered = StdArc::new(AtomicBool::new(false));
6511 let release = StdArc::new(Notify::new());
6512 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
6513 "echo",
6514 RoutingDecision::Background,
6515 )]));
6516 let handle = task_manager.handle();
6517 let agent = Agent::builder()
6518 .model(FakeAdapter)
6519 .tool_executor(DelayedApprovalExecutor::new(
6520 entered.clone(),
6521 release.clone(),
6522 ))
6523 .task_manager(task_manager)
6524 .cancellation(controller.handle())
6525 .observer(RecordingObserver {
6526 events: events.clone(),
6527 })
6528 .build()
6529 .unwrap();
6530
6531 let mut driver = agent
6532 .start(SessionConfig {
6533 session_id: SessionId::new("session-cancel-delayed-background-approval"),
6534 metadata: MetadataMap::new(),
6535 cache: None,
6536 })
6537 .await
6538 .unwrap();
6539
6540 driver
6541 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6542 .unwrap();
6543
6544 match driver.next().await.unwrap() {
6545 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
6546 other => panic!("unexpected first step: {other:?}"),
6547 }
6548
6549 match wait_for_task_event(&handle).await {
6550 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "echo"),
6551 other => panic!("unexpected task event: {other:?}"),
6552 }
6553
6554 wait_until_entered(entered.as_ref()).await;
6555 controller.interrupt();
6556 release.notify_waiters();
6557 wait_until_completed(&handle).await;
6558
6559 match driver.next().await.unwrap() {
6560 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
6561 other => panic!("cancelled background approval should finish cancelled, got {other:?}"),
6562 }
6563
6564 let events = events.lock().unwrap();
6565 assert!(
6566 events
6567 .iter()
6568 .any(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false }))
6569 );
6570 assert!(events.iter().any(|event| matches!(
6571 event,
6572 AgentEvent::ToolResultReceived(result)
6573 if result.call_id == ToolCallId::new("call-1") && result.is_error
6574 )));
6575 }
6576
6577 #[tokio::test]
6578 async fn approved_foreground_task_waits_for_result_before_model_continuation() {
6579 let entered = StdArc::new(AtomicBool::new(false));
6580 let release = StdArc::new(Notify::new());
6581 let approved_entered = StdArc::new(AtomicBool::new(false));
6582 let approved_release = StdArc::new(Notify::new());
6583 let route_count = StdArc::new(AtomicUsize::new(0));
6584 let routing_count = route_count.clone();
6585 let task_manager = AsyncTaskManager::new().routing(move |_request: &ToolRequest| {
6586 if routing_count.fetch_add(1, Ordering::SeqCst) == 0 {
6587 RoutingDecision::Background
6588 } else {
6589 RoutingDecision::Foreground
6590 }
6591 });
6592 let handle = task_manager.handle();
6593 let agent = Agent::builder()
6594 .model(FakeAdapter)
6595 .tool_executor(
6596 DelayedApprovalExecutor::new(entered.clone(), release.clone())
6597 .blocking_after_approval(approved_entered.clone(), approved_release.clone()),
6598 )
6599 .task_manager(task_manager)
6600 .build()
6601 .unwrap();
6602 let mut driver = agent
6603 .start(SessionConfig::new("session-approved-foreground"))
6604 .await
6605 .unwrap();
6606 driver
6607 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6608 .unwrap();
6609
6610 assert!(matches!(
6611 driver.next().await.unwrap(),
6612 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
6613 ));
6614 let task_turn = match wait_for_task_event(&handle).await {
6615 TaskEvent::Started(snapshot) => snapshot.turn_id,
6616 other => panic!("unexpected task event: {other:?}"),
6617 };
6618 wait_until_entered(entered.as_ref()).await;
6619 release.notify_one();
6620 wait_until_completed(&handle).await;
6621
6622 let pending = match driver.next().await.unwrap() {
6623 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
6624 other => panic!("unexpected delayed approval step: {other:?}"),
6625 };
6626 let presentation_turn = driver.lifecycle.active_turn.clone().unwrap();
6627 assert_ne!(presentation_turn, task_turn);
6628 pending.approve(&mut driver).unwrap();
6629
6630 let info = {
6631 let next = driver.next();
6632 tokio::pin!(next);
6633 tokio::select! {
6634 () = wait_until_entered(approved_entered.as_ref()) => {}
6635 result = &mut next => {
6636 panic!("model continued before approved foreground result: {result:?}")
6637 }
6638 }
6639 assert!(
6640 timeout(Duration::from_millis(10), &mut next).await.is_err(),
6641 "model continued while approved foreground work was blocked"
6642 );
6643 approved_release.notify_one();
6644 let step = timeout(Duration::from_secs(1), &mut next)
6645 .await
6646 .expect("approved foreground result was not delivered")
6647 .unwrap();
6648 match step {
6649 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info,
6650 other => panic!("unexpected approved foreground step: {other:?}"),
6651 }
6652 };
6653 assert_eq!(info.turn_id, presentation_turn);
6654
6655 let turn = match driver.next().await.unwrap() {
6656 LoopStep::Finished(turn) => turn,
6657 other => panic!("model did not continue after approved result: {other:?}"),
6658 };
6659 assert_eq!(turn.finish_reason, FinishReason::Completed);
6660 assert_eq!(turn.turn_id, presentation_turn);
6661 assert!(driver.snapshot().transcript.iter().any(|item| {
6662 item.kind == ItemKind::Notification
6663 && item.parts.iter().any(
6664 |part| matches!(part, Part::Text(text) if text.text.contains("approved-ok")),
6665 )
6666 }));
6667 }
6668
6669 #[tokio::test]
6670 async fn approved_foreground_then_detach_waits_and_keeps_one_placeholder() {
6671 let entered = StdArc::new(AtomicBool::new(false));
6672 let release = StdArc::new(Notify::new());
6673 let approved_entered = StdArc::new(AtomicBool::new(false));
6674 let approved_release = StdArc::new(Notify::new());
6675 let route_count = StdArc::new(AtomicUsize::new(0));
6676 let routing_count = route_count.clone();
6677 let task_manager = AsyncTaskManager::new().routing(move |_request: &ToolRequest| {
6678 if routing_count.fetch_add(1, Ordering::SeqCst) == 0 {
6679 RoutingDecision::Background
6680 } else {
6681 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10))
6682 }
6683 });
6684 let handle = task_manager.handle();
6685 let agent = Agent::builder()
6686 .model(FakeAdapter)
6687 .tool_executor(
6688 DelayedApprovalExecutor::new(entered.clone(), release.clone())
6689 .blocking_after_approval(approved_entered.clone(), approved_release.clone()),
6690 )
6691 .task_manager(task_manager)
6692 .build()
6693 .unwrap();
6694 let mut driver = agent
6695 .start(SessionConfig::new("session-approved-foreground-detach"))
6696 .await
6697 .unwrap();
6698 driver
6699 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6700 .unwrap();
6701
6702 assert!(matches!(
6703 driver.next().await.unwrap(),
6704 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
6705 ));
6706 let _ = wait_for_task_event(&handle).await;
6707 wait_until_entered(entered.as_ref()).await;
6708 release.notify_one();
6709 wait_until_completed(&handle).await;
6710
6711 let pending = match driver.next().await.unwrap() {
6712 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
6713 other => panic!("unexpected delayed approval step: {other:?}"),
6714 };
6715 let presentation_turn = driver.lifecycle.active_turn.clone().unwrap();
6716 pending.approve(&mut driver).unwrap();
6717
6718 let step = timeout(Duration::from_secs(1), driver.next())
6719 .await
6720 .expect("approved task did not detach")
6721 .unwrap();
6722 assert!(approved_entered.load(Ordering::SeqCst));
6723 match step {
6724 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => {
6725 assert_eq!(info.turn_id, presentation_turn);
6726 }
6727 other => panic!("unexpected approved detach step: {other:?}"),
6728 }
6729 let placeholders = driver
6730 .snapshot()
6731 .transcript
6732 .iter()
6733 .filter(|item| item.kind == ItemKind::Tool)
6734 .flat_map(|item| &item.parts)
6735 .filter(|part| {
6736 matches!(
6737 part,
6738 Part::ToolResult(result) if result.call_id == ToolCallId::new("call-1")
6739 )
6740 })
6741 .count();
6742 assert_eq!(placeholders, 1, "detach appended a second tool result");
6743
6744 approved_release.notify_one();
6745 wait_until_completed(&handle).await;
6746 let turn = match driver.next().await.unwrap() {
6747 LoopStep::Finished(turn) => turn,
6748 other => panic!("model did not continue after detached result: {other:?}"),
6749 };
6750 assert_eq!(turn.finish_reason, FinishReason::Completed);
6751 assert_eq!(turn.turn_id, presentation_turn);
6752 let transcript = driver.snapshot().transcript;
6753 assert_eq!(
6754 transcript
6755 .iter()
6756 .filter(|item| item.kind == ItemKind::Tool)
6757 .flat_map(|item| &item.parts)
6758 .filter(|part| {
6759 matches!(
6760 part,
6761 Part::ToolResult(result)
6762 if result.call_id == ToolCallId::new("call-1")
6763 )
6764 })
6765 .count(),
6766 1
6767 );
6768 assert!(transcript.iter().any(|item| {
6769 item.kind == ItemKind::Notification
6770 && item.parts.iter().any(
6771 |part| matches!(part, Part::Text(text) if text.text.contains("approved-ok")),
6772 )
6773 }));
6774 }
6775
6776 #[tokio::test]
6777 async fn approving_detached_background_call_keeps_one_placeholder() {
6778 let entered = StdArc::new(AtomicBool::new(false));
6779 let release = StdArc::new(Notify::new());
6780 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
6781 "echo",
6782 RoutingDecision::Background,
6783 )]));
6784 let handle = task_manager.handle();
6785 let agent = Agent::builder()
6786 .model(FakeAdapter)
6787 .tool_executor(DelayedApprovalExecutor::new(
6788 entered.clone(),
6789 release.clone(),
6790 ))
6791 .task_manager(task_manager)
6792 .build()
6793 .unwrap();
6794 let mut driver = agent
6795 .start(SessionConfig::new("session-approve-detached-background"))
6796 .await
6797 .unwrap();
6798 driver
6799 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6800 .unwrap();
6801
6802 assert!(matches!(
6803 driver.next().await.unwrap(),
6804 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
6805 ));
6806 let _ = wait_for_task_event(&handle).await;
6807 wait_until_entered(entered.as_ref()).await;
6808 release.notify_waiters();
6809 wait_until_completed(&handle).await;
6810
6811 let pending = match driver.next().await.unwrap() {
6812 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
6813 other => panic!("unexpected delayed approval step: {other:?}"),
6814 };
6815 pending.approve(&mut driver).unwrap();
6816 release.notify_one();
6817 let _ = driver.next().await.unwrap();
6818
6819 let placeholders = driver
6820 .snapshot()
6821 .transcript
6822 .iter()
6823 .flat_map(|item| &item.parts)
6824 .filter(|part| {
6825 matches!(
6826 part,
6827 Part::ToolResult(result) if result.call_id == ToolCallId::new("call-1")
6828 )
6829 })
6830 .count();
6831 assert_eq!(placeholders, 1, "approval appended a second detach result");
6832 }
6833
6834 #[tokio::test]
6835 async fn failed_background_approval_cleanup_clears_queued_resume() {
6836 let events = StdArc::new(StdMutex::new(Vec::new()));
6837 let entered = StdArc::new(AtomicBool::new(false));
6838 let release = StdArc::new(Notify::new());
6839 let inner = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
6840 "echo",
6841 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10)),
6842 )]));
6843 let handle = inner.handle();
6844 let agent = Agent::builder()
6845 .model(FakeAdapter)
6846 .tool_executor(DelayedApprovalExecutor::new(
6847 entered.clone(),
6848 release.clone(),
6849 ))
6850 .task_manager(TestTaskManager::new(inner).fail_interrupt("cleanup failure"))
6851 .observer(RecordingObserver {
6852 events: events.clone(),
6853 })
6854 .build()
6855 .unwrap();
6856 let mut driver = agent
6857 .start(SessionConfig::new(
6858 "session-failed-detached-background-approval-cleanup",
6859 ))
6860 .await
6861 .unwrap();
6862 driver
6863 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6864 .unwrap();
6865
6866 let old_turn_id = match driver.next().await.unwrap() {
6867 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info.turn_id,
6868 other => panic!("unexpected detach step: {other:?}"),
6869 };
6870 assert_eq!(driver.pending_round_resume.as_ref(), Some(&old_turn_id));
6871 driver
6872 .submit_input(vec![Item::text(ItemKind::User, "fresh input")])
6873 .unwrap();
6874 let _ = wait_for_task_event(&handle).await;
6875 wait_until_entered(entered.as_ref()).await;
6876 release.notify_waiters();
6877 wait_until_completed(&handle).await;
6878
6879 assert!(matches!(
6880 driver.next().await.unwrap(),
6881 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))
6882 ));
6883 let error = driver.cancel_pending_approvals().await.unwrap_err();
6884 assert!(error.to_string().contains("cleanup failure"));
6885
6886 assert!(driver.lifecycle.active_turn.is_none());
6887 assert!(driver.pending_round_resume.is_none());
6888 assert_eq!(driver.snapshot().pending_input.len(), 1);
6889 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
6890 let [started, finished] = &lifecycle[lifecycle.len() - 2..] else {
6891 panic!("missing terminal lifecycle events: {lifecycle:?}");
6892 };
6893 assert_eq!(started.0, finished.0);
6894 assert_eq!(finished.1, Some(FinishReason::Error));
6895 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
6896
6897 let fresh_turn = match driver.next().await.unwrap() {
6898 LoopStep::Finished(turn) => turn,
6899 other => panic!("fresh input did not start a new turn: {other:?}"),
6900 };
6901 assert_ne!(fresh_turn.turn_id, old_turn_id);
6902 }
6903
6904 #[tokio::test]
6905 async fn fresh_input_runs_before_delayed_background_approval() {
6906 let entered = StdArc::new(AtomicBool::new(false));
6907 let release = StdArc::new(Notify::new());
6908 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
6909 "echo",
6910 RoutingDecision::Background,
6911 )]));
6912 let handle = task_manager.handle();
6913 let agent = Agent::builder()
6914 .model(FakeAdapter)
6915 .tool_executor(DelayedApprovalExecutor::new(
6916 entered.clone(),
6917 release.clone(),
6918 ))
6919 .task_manager(task_manager)
6920 .build()
6921 .unwrap();
6922 let mut driver = agent
6923 .start(SessionConfig::new(
6924 "session-input-before-background-approval",
6925 ))
6926 .await
6927 .unwrap();
6928 driver
6929 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6930 .unwrap();
6931
6932 assert!(matches!(
6933 driver.next().await.unwrap(),
6934 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
6935 ));
6936 let _ = wait_for_task_event(&handle).await;
6937 wait_until_entered(entered.as_ref()).await;
6938 release.notify_waiters();
6939 wait_until_completed(&handle).await;
6940
6941 driver
6942 .submit_input(vec![Item::text(ItemKind::User, "fresh input")])
6943 .unwrap();
6944 let fresh_turn = match driver.next().await.unwrap() {
6945 LoopStep::Finished(turn) => turn.turn_id,
6946 other => panic!("fresh input was not driven first: {other:?}"),
6947 };
6948 assert!(driver.pending_approvals.is_empty());
6949 assert!(driver.snapshot().pending_input.is_empty());
6950 timeout(Duration::from_millis(100), driver.wait_for_loop_update())
6951 .await
6952 .expect("collected background update did not wake the loop")
6953 .unwrap();
6954
6955 let approval = match driver.next().await.unwrap() {
6956 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(approval)) => approval,
6957 other => panic!("delayed approval was not presented separately: {other:?}"),
6958 };
6959 let approval_turn = driver.lifecycle.active_turn.clone().unwrap();
6960 assert_ne!(fresh_turn, approval_turn);
6961 assert_eq!(
6962 driver
6963 .snapshot()
6964 .transcript
6965 .iter()
6966 .filter(|item| {
6967 item.kind == ItemKind::User
6968 && item.parts.iter().any(
6969 |part| matches!(part, Part::Text(text) if text.text == "fresh input"),
6970 )
6971 })
6972 .count(),
6973 1,
6974 "fresh input must not be replayed while presenting the approval"
6975 );
6976 approval.deny(&mut driver).unwrap();
6977 }
6978
6979 #[tokio::test]
6980 async fn delayed_background_approval_interrupts_originating_task_turn() {
6981 let entered = StdArc::new(AtomicBool::new(false));
6982 let release = StdArc::new(Notify::new());
6983 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
6984 let inner = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
6985 "echo",
6986 RoutingDecision::Background,
6987 )]));
6988 let handle = inner.handle();
6989 let agent = Agent::builder()
6990 .model(FakeAdapter)
6991 .tool_executor(DelayedApprovalExecutor::new(
6992 entered.clone(),
6993 release.clone(),
6994 ))
6995 .task_manager(TestTaskManager::new(inner).record_interrupts(interrupted.clone()))
6996 .build()
6997 .unwrap();
6998 let mut driver = agent
6999 .start(SessionConfig::new(
7000 "session-background-approval-origin-turn",
7001 ))
7002 .await
7003 .unwrap();
7004 driver
7005 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7006 .unwrap();
7007
7008 assert!(matches!(
7009 driver.next().await.unwrap(),
7010 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
7011 ));
7012 let _ = wait_for_task_event(&handle).await;
7013 wait_until_entered(entered.as_ref()).await;
7014 release.notify_waiters();
7015 wait_until_completed(&handle).await;
7016 assert!(matches!(
7017 driver.next().await.unwrap(),
7018 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))
7019 ));
7020
7021 let presentation_turn = driver.lifecycle.active_turn.clone().unwrap();
7022 let task_turn = driver
7023 .pending_approvals
7024 .values()
7025 .next()
7026 .unwrap()
7027 .tool_request
7028 .turn_id
7029 .clone();
7030 assert_ne!(presentation_turn, task_turn);
7031 assert!(matches!(
7032 driver.cancel_pending_approvals().await.unwrap(),
7033 Some(LoopStep::Finished(TurnResult {
7034 finish_reason: FinishReason::Cancelled,
7035 ..
7036 }))
7037 ));
7038 assert_eq!(interrupted.lock().unwrap().as_slice(), &[task_turn]);
7039 assert!(driver.lifecycle.active_turn.is_none());
7040 }
7041
7042 #[tokio::test]
7043 async fn approved_background_start_error_interrupts_originating_turn() {
7044 let events = StdArc::new(StdMutex::new(Vec::new()));
7045 let entered = StdArc::new(AtomicBool::new(false));
7046 let release = StdArc::new(Notify::new());
7047 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
7048 let inner = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
7049 "echo",
7050 RoutingDecision::Background,
7051 )]));
7052 let handle = inner.handle();
7053 let agent = Agent::builder()
7054 .model(FakeAdapter)
7055 .tool_executor(DelayedApprovalExecutor::new(
7056 entered.clone(),
7057 release.clone(),
7058 ))
7059 .task_manager(
7060 TestTaskManager::new(inner)
7061 .fail_approved_start("original approved start failure")
7062 .record_interrupts(interrupted.clone())
7063 .fail_interrupt("cleanup failure"),
7064 )
7065 .observer(RecordingObserver {
7066 events: events.clone(),
7067 })
7068 .build()
7069 .unwrap();
7070 let mut driver = agent
7071 .start(SessionConfig::new("session-approved-start-error"))
7072 .await
7073 .unwrap();
7074 driver
7075 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7076 .unwrap();
7077
7078 assert!(matches!(
7079 driver.next().await.unwrap(),
7080 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))
7081 ));
7082 let _ = wait_for_task_event(&handle).await;
7083 wait_until_entered(entered.as_ref()).await;
7084 release.notify_waiters();
7085 wait_until_completed(&handle).await;
7086 let pending = match driver.next().await.unwrap() {
7087 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
7088 other => panic!("unexpected delayed approval step: {other:?}"),
7089 };
7090 let task_turn = driver
7091 .pending_approvals
7092 .values()
7093 .next()
7094 .unwrap()
7095 .tool_request
7096 .turn_id
7097 .clone();
7098 let call_id = pending.request.call_id.clone().expect("approval call id");
7099 assert!(driver.detached_call_ids.contains(&call_id));
7100 pending.approve(&mut driver).unwrap();
7101
7102 let error = driver.next().await.unwrap_err();
7103 assert!(
7104 error
7105 .to_string()
7106 .contains("original approved start failure")
7107 );
7108 assert!(!error.to_string().contains("cleanup failure"));
7109 assert_eq!(interrupted.lock().unwrap().as_slice(), &[task_turn]);
7110 assert!(driver.lifecycle.active_turn.is_none());
7111 assert!(!driver.detached_call_ids.contains(&call_id));
7112 assert!(!driver.background_call_ids.contains(&call_id));
7113 assert!(!driver.tool_cancellations.contains_key(&call_id));
7114 assert!(events.lock().unwrap().iter().any(|event| matches!(
7115 event,
7116 AgentEvent::ToolResultReceived(result)
7117 if result.call_id == call_id && result.is_error
7118 )));
7119 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
7120 }
7121
7122 #[tokio::test]
7123 async fn loop_can_cancel_a_turn_and_continue_after_new_input() {
7124 let controller = CancellationController::new();
7125 let agent = Agent::builder()
7126 .model(SlowAdapter)
7127 .cancellation(controller.handle())
7128 .build()
7129 .unwrap();
7130
7131 let mut driver = agent
7132 .start(SessionConfig {
7133 session_id: SessionId::new("session-cancel"),
7134 metadata: MetadataMap::new(),
7135 cache: None,
7136 })
7137 .await
7138 .unwrap();
7139
7140 driver
7141 .submit_input(vec![Item {
7142 id: None,
7143 kind: ItemKind::User,
7144 parts: vec![Part::Text(TextPart {
7145 text: "do the long task".into(),
7146 metadata: MetadataMap::new(),
7147 })],
7148 metadata: MetadataMap::new(),
7149 usage: None,
7150 finish_reason: None,
7151 created_at: None,
7152 }])
7153 .unwrap();
7154
7155 let cancelled = tokio::join!(async { driver.next().await }, async {
7156 tokio::task::yield_now().await;
7157 controller.interrupt();
7158 })
7159 .0
7160 .unwrap();
7161
7162 match cancelled {
7163 LoopStep::Finished(turn) => {
7164 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
7165 assert_eq!(turn.items.len(), 1);
7166 assert_eq!(turn.items[0].kind, ItemKind::Assistant);
7167 assert_eq!(
7168 turn.items[0].metadata.get(INTERRUPTED_METADATA_KEY),
7169 Some(&Value::Bool(true))
7170 );
7171 }
7172 other => panic!("unexpected loop step: {other:?}"),
7173 }
7174
7175 driver
7176 .submit_input(vec![Item {
7177 id: None,
7178 kind: ItemKind::User,
7179 parts: vec![Part::Text(TextPart {
7180 text: "try again".into(),
7181 metadata: MetadataMap::new(),
7182 })],
7183 metadata: MetadataMap::new(),
7184 usage: None,
7185 finish_reason: None,
7186 created_at: None,
7187 }])
7188 .unwrap();
7189
7190 let result = driver.next().await.unwrap();
7191 match result {
7192 LoopStep::Finished(turn) => {
7193 assert_eq!(turn.finish_reason, FinishReason::Completed);
7194 }
7195 other => panic!("unexpected loop step after retry: {other:?}"),
7196 }
7197 }
7198
7199 #[tokio::test]
7200 async fn loop_interrupt_cancels_foreground_tasks_but_keeps_background_tasks_running() {
7201 let controller = CancellationController::new();
7202 let fg_entered = StdArc::new(AtomicBool::new(false));
7203 let fg_release = StdArc::new(Notify::new());
7204 let bg_entered = StdArc::new(AtomicBool::new(false));
7205 let bg_release = StdArc::new(Notify::new());
7206 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([
7207 ("foreground-wait", RoutingDecision::Foreground),
7208 ("background-wait", RoutingDecision::Background),
7209 ]));
7210 let handle = task_manager.handle();
7211 let tools = ToolRegistry::new()
7212 .with(BlockingTool::new(
7213 "foreground-wait",
7214 fg_entered.clone(),
7215 fg_release,
7216 "foreground-done",
7217 ))
7218 .with(BlockingTool::new(
7219 "background-wait",
7220 bg_entered.clone(),
7221 bg_release.clone(),
7222 "background-done",
7223 ));
7224 let agent = Agent::builder()
7225 .model(MultiToolAdapter)
7226 .add_tool_source(tools)
7227 .permissions(AllowAllPermissions)
7228 .cancellation(controller.handle())
7229 .task_manager(task_manager)
7230 .build()
7231 .unwrap();
7232
7233 let mut driver = agent
7234 .start(SessionConfig {
7235 session_id: SessionId::new("session-mixed-cancel"),
7236 metadata: MetadataMap::new(),
7237 cache: None,
7238 })
7239 .await
7240 .unwrap();
7241
7242 driver
7243 .submit_input(vec![Item {
7244 id: None,
7245 kind: ItemKind::User,
7246 parts: vec![Part::Text(TextPart {
7247 text: "run both".into(),
7248 metadata: MetadataMap::new(),
7249 })],
7250 metadata: MetadataMap::new(),
7251 usage: None,
7252 finish_reason: None,
7253 created_at: None,
7254 }])
7255 .unwrap();
7256
7257 let cancelled = tokio::join!(async { driver.next().await }, async {
7258 let _ = wait_for_task_event(&handle).await;
7259 let _ = wait_for_task_event(&handle).await;
7260 wait_until_entered(fg_entered.as_ref()).await;
7261 wait_until_entered(bg_entered.as_ref()).await;
7262 controller.interrupt();
7263 })
7264 .0
7265 .unwrap();
7266
7267 match cancelled {
7268 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
7269 other => panic!("unexpected loop step after interrupt: {other:?}"),
7270 }
7271
7272 match wait_for_task_event(&handle).await {
7273 TaskEvent::Cancelled(snapshot) => assert_eq!(snapshot.tool_name, "foreground-wait"),
7274 other => panic!("unexpected post-interrupt event: {other:?}"),
7275 }
7276
7277 let running = handle.list_running().await;
7278 assert_eq!(running.len(), 1);
7279 assert_eq!(running[0].tool_name, "background-wait");
7280
7281 bg_release.notify_waiters();
7282 match wait_for_task_event(&handle).await {
7283 TaskEvent::Completed(snapshot, result) => {
7284 assert_eq!(snapshot.tool_name, "background-wait");
7285 assert_eq!(result.output, ToolOutput::Text("background-done".into()));
7286 }
7287 other => panic!("unexpected background completion event: {other:?}"),
7288 }
7289 }
7290
7291 #[tokio::test]
7292 async fn a_cancelled_turn_answers_the_tool_call_it_abandoned() {
7293 let controller = CancellationController::new();
7298 let entered = StdArc::new(AtomicBool::new(false));
7299 let release = StdArc::new(Notify::new());
7300 let items = StdArc::new(StdMutex::new(Vec::<Item>::new()));
7301 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
7302 "wait",
7303 RoutingDecision::Foreground,
7304 )]));
7305 let agent = Agent::builder()
7306 .model(FakeAdapter)
7307 .add_tool_source(ToolRegistry::new().with(BlockingTool::new(
7308 "wait",
7309 entered.clone(),
7310 release,
7311 "done",
7312 )))
7313 .permissions(AllowAllPermissions)
7314 .cancellation(controller.handle())
7315 .task_manager(task_manager)
7316 .transcript_observer(RecordingTranscriptObserver {
7317 items: items.clone(),
7318 })
7319 .build()
7320 .unwrap();
7321
7322 let mut driver = agent
7323 .start(SessionConfig {
7324 session_id: SessionId::new("session-cancel-mid-call"),
7325 metadata: MetadataMap::new(),
7326 cache: None,
7327 })
7328 .await
7329 .unwrap();
7330
7331 driver
7332 .submit_input(vec![Item {
7333 id: None,
7334 kind: ItemKind::User,
7335 parts: vec![Part::Text(TextPart {
7336 text: "run the tool".into(),
7337 metadata: MetadataMap::new(),
7338 })],
7339 metadata: MetadataMap::new(),
7340 usage: None,
7341 finish_reason: None,
7342 created_at: None,
7343 }])
7344 .unwrap();
7345
7346 let cancelled = tokio::join!(async { driver.next().await }, async {
7347 wait_until_entered(entered.as_ref()).await;
7348 controller.interrupt();
7349 })
7350 .0
7351 .unwrap();
7352
7353 match cancelled {
7354 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
7355 other => panic!("unexpected loop step after interrupt: {other:?}"),
7356 }
7357
7358 let transcript = driver.snapshot().transcript;
7359 assert!(
7360 unanswered_tool_calls(&transcript).is_empty(),
7361 "the cancelled turn left a tool call unanswered: {transcript:?}"
7362 );
7363 validate_transcript_invariants(&transcript)
7364 .expect("a cancelled turn must leave a resumable transcript");
7365
7366 let persisted = items.lock().unwrap().clone();
7367 let results: Vec<&ToolResultPart> = persisted
7368 .iter()
7369 .flat_map(|item| &item.parts)
7370 .filter_map(|part| match part {
7371 Part::ToolResult(result) => Some(result),
7372 _ => None,
7373 })
7374 .collect();
7375 assert_eq!(results.len(), 1, "{persisted:?}");
7376 assert_eq!(results[0].call_id, ToolCallId::new("call-1"));
7377 assert!(results[0].is_error);
7378 assert_eq!(
7379 results[0].metadata.get(INTERRUPTED_METADATA_KEY),
7380 Some(&Value::Bool(true))
7381 );
7382 }
7383
7384 #[tokio::test]
7385 async fn regression_cancelled_background_completion_emits_one_terminal_result() {
7386 let events = StdArc::new(StdMutex::new(Vec::new()));
7387 let agent = Agent::builder()
7388 .model(FakeAdapter)
7389 .observer(RecordingObserver {
7390 events: events.clone(),
7391 })
7392 .build()
7393 .unwrap();
7394 let mut driver = agent
7395 .start(SessionConfig::new("session-cancelled-background-event"))
7396 .await
7397 .unwrap();
7398
7399 driver.append_item(Item::new(
7400 ItemKind::Assistant,
7401 vec![Part::ToolCall(ToolCallPart {
7402 id: ToolCallId::new("call-1"),
7403 name: "wait".into(),
7404 input: json!({}),
7405 metadata: MetadataMap::new(),
7406 })],
7407 ));
7408 driver.background_call_ids.insert(ToolCallId::new("call-1"));
7409 driver.close_interrupted_tool_calls();
7410 driver.append_tool_result_item(Item::new(
7411 ItemKind::Tool,
7412 vec![Part::ToolResult(ToolResultPart {
7413 call_id: ToolCallId::new("call-1"),
7414 output: ToolOutput::Text("background-done".into()),
7415 is_error: false,
7416 metadata: MetadataMap::new(),
7417 })],
7418 ));
7419
7420 let events = events.lock().unwrap();
7421 let terminal_results = events
7422 .iter()
7423 .filter(|event| {
7424 matches!(
7425 event,
7426 AgentEvent::ToolResultReceived(result)
7427 if result.call_id == ToolCallId::new("call-1")
7428 )
7429 })
7430 .count();
7431 assert_eq!(
7432 terminal_results, 1,
7433 "a cancelled background call emitted multiple terminal results: {events:?}"
7434 );
7435 }
7436
7437 #[tokio::test]
7438 async fn regression_cancelled_queued_approval_is_answered_once() {
7439 let controller = CancellationController::new();
7440 let entered = StdArc::new(AtomicBool::new(false));
7441 let release = StdArc::new(Notify::new());
7442 release.notify_one();
7443 let events = StdArc::new(StdMutex::new(Vec::new()));
7444 let agent = Agent::builder()
7445 .model(FakeAdapter)
7446 .tool_executor(
7447 DelayedApprovalExecutor::new(entered, release)
7448 .cancelling_on_approval(controller.clone()),
7449 )
7450 .cancellation(controller.handle())
7451 .observer(RecordingObserver {
7452 events: events.clone(),
7453 })
7454 .build()
7455 .unwrap();
7456 let mut driver = agent
7457 .start(SessionConfig::new("session-cancelled-queued-approval"))
7458 .await
7459 .unwrap();
7460 driver
7461 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7462 .unwrap();
7463
7464 match driver.next().await.unwrap() {
7465 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
7466 other => panic!("unexpected first cancellation step: {other:?}"),
7467 }
7468 match driver.next().await.unwrap() {
7469 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
7470 other => panic!("unexpected post-cancellation step: {other:?}"),
7471 }
7472
7473 let events = events.lock().unwrap();
7474 let terminal_results = events
7475 .iter()
7476 .filter(|event| {
7477 matches!(
7478 event,
7479 AgentEvent::ToolResultReceived(result)
7480 if result.call_id == ToolCallId::new("call-1")
7481 )
7482 })
7483 .count();
7484 assert_eq!(
7485 terminal_results, 1,
7486 "a cancelled queued approval was answered more than once: {events:?}"
7487 );
7488 drop(events);
7489
7490 let transcript = driver.snapshot().transcript;
7491 assert!(
7492 !transcript.iter().any(|item| {
7493 item.kind == ItemKind::Notification
7494 && item.parts.iter().any(|part| {
7495 matches!(part, Part::Text(text) if text.text.contains("Background tool call"))
7496 })
7497 }),
7498 "a queued approval was misreported as a background call: {transcript:?}"
7499 );
7500 }
7501
7502 #[tokio::test]
7503 async fn regression_cancelled_unstarted_call_is_not_tracked_as_detached() {
7504 let agent = Agent::builder().model(FakeAdapter).build().unwrap();
7505 let mut driver = agent
7506 .start(SessionConfig::new("session-cancelled-unstarted-call"))
7507 .await
7508 .unwrap();
7509
7510 driver.append_item(Item::new(
7511 ItemKind::Assistant,
7512 vec![Part::ToolCall(ToolCallPart {
7513 id: ToolCallId::new("call-never-started"),
7514 name: "wait".into(),
7515 input: json!({}),
7516 metadata: MetadataMap::new(),
7517 })],
7518 ));
7519 driver
7520 .finish_cancelled(agentkit_core::TurnId::new("turn-cancelled"), Vec::new())
7521 .unwrap();
7522
7523 assert!(
7524 !driver
7525 .detached_call_ids
7526 .contains(&ToolCallId::new("call-never-started")),
7527 "an unstarted call can never deliver a detached result"
7528 );
7529 }
7530
7531 #[tokio::test]
7532 async fn loop_resumes_after_approved_tool_request() {
7533 let tools = ToolRegistry::new().with(EchoTool::default());
7534 let agent = Agent::builder()
7535 .model(FakeAdapter)
7536 .add_tool_source(tools)
7537 .permissions(ApproveFsReads)
7538 .build()
7539 .unwrap();
7540
7541 let mut driver = agent
7542 .start(SessionConfig {
7543 session_id: SessionId::new("session-approval"),
7544 metadata: MetadataMap::new(),
7545 cache: None,
7546 })
7547 .await
7548 .unwrap();
7549
7550 driver
7551 .submit_input(vec![Item {
7552 id: None,
7553 kind: ItemKind::User,
7554 parts: vec![Part::Text(TextPart {
7555 text: "ping".into(),
7556 metadata: MetadataMap::new(),
7557 })],
7558 metadata: MetadataMap::new(),
7559 usage: None,
7560 finish_reason: None,
7561 created_at: None,
7562 }])
7563 .unwrap();
7564
7565 let first = driver.next().await.unwrap();
7566 match first {
7567 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
7568 assert!(pending.request.task_id.is_some());
7569 assert_eq!(pending.request.id.0, "approval:fs-read");
7570 pending.approve(&mut driver).unwrap();
7571 }
7572 other => panic!("unexpected loop step: {other:?}"),
7573 }
7574 let second = driver.next().await.unwrap();
7575 match second {
7576 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
7577 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
7578 other => panic!("unexpected part: {other:?}"),
7579 },
7580 other => panic!("unexpected loop step after approval: {other:?}"),
7581 }
7582 }
7583
7584 #[tokio::test]
7585 async fn approval_gated_tool_does_not_start_before_approval() {
7586 let events = StdArc::new(StdMutex::new(Vec::new()));
7587 let tools = ToolRegistry::new().with(EchoTool::default());
7588 let agent = Agent::builder()
7589 .model(FakeAdapter)
7590 .add_tool_source(tools)
7591 .permissions(ApproveFsReads)
7592 .observer(RecordingObserver {
7593 events: events.clone(),
7594 })
7595 .build()
7596 .unwrap();
7597
7598 let mut driver = agent
7599 .start(SessionConfig {
7600 session_id: SessionId::new("session-approval-start-event"),
7601 metadata: MetadataMap::new(),
7602 cache: None,
7603 })
7604 .await
7605 .unwrap();
7606
7607 driver
7608 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7609 .unwrap();
7610
7611 let pending = match driver.next().await.unwrap() {
7612 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
7613 other => panic!("unexpected loop step: {other:?}"),
7614 };
7615
7616 assert!(
7617 events
7618 .lock()
7619 .unwrap()
7620 .iter()
7621 .all(|event| !matches!(event, AgentEvent::ToolExecutionStarted(_))),
7622 "tool start must not be reported before approval"
7623 );
7624
7625 pending.approve(&mut driver).unwrap();
7626 match driver.next().await.unwrap() {
7627 LoopStep::Finished(_) => {}
7628 other => panic!("unexpected loop step after approval: {other:?}"),
7629 }
7630
7631 let started = events
7632 .lock()
7633 .unwrap()
7634 .iter()
7635 .filter(|event| matches!(event, AgentEvent::ToolExecutionStarted(_)))
7636 .count();
7637 assert_eq!(started, 1);
7638 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
7639 assert_eq!(lifecycle.len(), 2);
7640 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7641 assert_eq!(lifecycle[1].1, Some(FinishReason::Completed));
7642 }
7643
7644 #[tokio::test]
7645 async fn cancelling_pending_approval_resolves_it_and_pairs_tool_result() {
7646 let controller = CancellationController::new();
7647 let events = StdArc::new(StdMutex::new(Vec::new()));
7648 let tools = ToolRegistry::new().with(EchoTool::default());
7649 let agent = Agent::builder()
7650 .model(FakeAdapter)
7651 .add_tool_source(tools)
7652 .permissions(ApproveFsReads)
7653 .cancellation(controller.handle())
7654 .observer(RecordingObserver {
7655 events: events.clone(),
7656 })
7657 .build()
7658 .unwrap();
7659
7660 let mut driver = agent
7661 .start(SessionConfig {
7662 session_id: SessionId::new("session-cancel-pending-approval"),
7663 metadata: MetadataMap::new(),
7664 cache: None,
7665 })
7666 .await
7667 .unwrap();
7668
7669 driver
7670 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7671 .unwrap();
7672
7673 match driver.next().await.unwrap() {
7674 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) => {}
7675 other => panic!("unexpected loop step: {other:?}"),
7676 }
7677
7678 controller.interrupt();
7679
7680 match driver.next().await.unwrap() {
7681 LoopStep::Finished(turn) => {
7682 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
7683 }
7684 other => panic!("unexpected loop step after cancel: {other:?}"),
7685 }
7686
7687 let events = events.lock().unwrap();
7688 assert!(
7689 events
7690 .iter()
7691 .any(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false })),
7692 "pending approval cancellation should close approval UI state"
7693 );
7694 assert!(
7695 events.iter().any(|event| matches!(
7696 event,
7697 AgentEvent::ToolResultReceived(result)
7698 if result.call_id == ToolCallId::new("call-1") && result.is_error
7699 )),
7700 "pending approval cancellation should pair the assistant tool_use"
7701 );
7702 drop(events);
7703
7704 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
7705 }
7706
7707 #[tokio::test]
7708 async fn cancelling_sole_foreground_approval_for_call_finishes_turn() {
7709 let events = StdArc::new(StdMutex::new(Vec::new()));
7710 let agent = Agent::builder()
7711 .model(FakeAdapter)
7712 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
7713 .permissions(ApproveFsReads)
7714 .observer(RecordingObserver {
7715 events: events.clone(),
7716 })
7717 .build()
7718 .unwrap();
7719 let mut driver = agent
7720 .start(SessionConfig::new("session-cancel-foreground-approval-for"))
7721 .await
7722 .unwrap();
7723 driver
7724 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7725 .unwrap();
7726
7727 let call_id = match driver.next().await.unwrap() {
7728 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
7729 pending.request.call_id.expect("approval call id")
7730 }
7731 other => panic!("unexpected loop step: {other:?}"),
7732 };
7733 driver.cancel_pending_approval_for(call_id).unwrap();
7734
7735 assert!(driver.lifecycle.active_turn.is_none());
7736 assert!(driver.pending_approvals.is_empty());
7737 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
7738 let lifecycle = turn_lifecycle_events(&events.lock().unwrap());
7739 assert_eq!(lifecycle.len(), 2, "{lifecycle:?}");
7740 assert_eq!(lifecycle[0].0, lifecycle[1].0);
7741 assert_eq!(lifecycle[1].1, Some(FinishReason::Cancelled));
7742 }
7743
7744 #[tokio::test]
7745 async fn resolved_approval_runs_even_if_cancellation_also_fired() {
7746 let controller = CancellationController::new();
7747 let tools = ToolRegistry::new().with(EchoTool::default());
7748 let agent = Agent::builder()
7749 .model(FakeAdapter)
7750 .add_tool_source(tools)
7751 .permissions(ApproveFsReads)
7752 .cancellation(controller.handle())
7753 .build()
7754 .unwrap();
7755
7756 let mut driver = agent
7757 .start(SessionConfig {
7758 session_id: SessionId::new("session-resolved-approval-cancel-race"),
7759 metadata: MetadataMap::new(),
7760 cache: None,
7761 })
7762 .await
7763 .unwrap();
7764
7765 driver
7766 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7767 .unwrap();
7768
7769 let pending = match driver.next().await.unwrap() {
7770 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
7771 other => panic!("unexpected loop step: {other:?}"),
7772 };
7773
7774 controller.interrupt();
7775 pending.approve(&mut driver).unwrap();
7776
7777 match driver.next().await.unwrap() {
7778 LoopStep::Finished(turn) => {
7779 assert_eq!(turn.finish_reason, FinishReason::Completed);
7780 match &turn.items[0].parts[0] {
7781 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
7782 other => panic!("unexpected part after approval: {other:?}"),
7783 }
7784 }
7785 other => panic!("unexpected loop step after approved cancel race: {other:?}"),
7786 }
7787 }
7788
7789 #[tokio::test]
7790 async fn loop_resumes_with_patched_input_on_approval() {
7791 let tools = ToolRegistry::new().with(EchoTool::default());
7792 let agent = Agent::builder()
7793 .model(FakeAdapter)
7794 .add_tool_source(tools)
7795 .permissions(ApproveFsReads)
7796 .build()
7797 .unwrap();
7798
7799 let mut driver = agent
7800 .start(SessionConfig {
7801 session_id: SessionId::new("session-approval-patched"),
7802 metadata: MetadataMap::new(),
7803 cache: None,
7804 })
7805 .await
7806 .unwrap();
7807
7808 driver
7809 .submit_input(vec![Item {
7810 id: None,
7811 kind: ItemKind::User,
7812 parts: vec![Part::Text(TextPart {
7813 text: "ping".into(),
7814 metadata: MetadataMap::new(),
7815 })],
7816 metadata: MetadataMap::new(),
7817 usage: None,
7818 finish_reason: None,
7819 created_at: None,
7820 }])
7821 .unwrap();
7822
7823 match driver.next().await.unwrap() {
7824 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
7825 pending
7826 .approve_with_patched_input(&mut driver, json!({ "value": "patched" }))
7827 .unwrap();
7828 }
7829 other => panic!("unexpected loop step: {other:?}"),
7830 }
7831 match driver.next().await.unwrap() {
7832 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
7833 Part::Text(text) => assert_eq!(text.text, "tool said: patched"),
7834 other => panic!("unexpected part: {other:?}"),
7835 },
7836 other => panic!("unexpected loop step after approval: {other:?}"),
7837 }
7838 }
7839
7840 #[tokio::test]
7841 async fn loop_tracks_multiple_pending_approvals_by_call_id() {
7842 let tools = ToolRegistry::new().with(EchoTool::default());
7843 let agent = Agent::builder()
7844 .model(DualApprovalAdapter)
7845 .add_tool_source(tools)
7846 .permissions(ApproveFsReads)
7847 .build()
7848 .unwrap();
7849
7850 let mut driver = agent
7851 .start(SessionConfig {
7852 session_id: SessionId::new("session-dual-approval"),
7853 metadata: MetadataMap::new(),
7854 cache: None,
7855 })
7856 .await
7857 .unwrap();
7858
7859 driver
7860 .submit_input(vec![Item {
7861 id: None,
7862 kind: ItemKind::User,
7863 parts: vec![Part::Text(TextPart {
7864 text: "run both approvals".into(),
7865 metadata: MetadataMap::new(),
7866 })],
7867 metadata: MetadataMap::new(),
7868 usage: None,
7869 finish_reason: None,
7870 created_at: None,
7871 }])
7872 .unwrap();
7873
7874 let pending_first = match driver.next().await.unwrap() {
7875 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
7876 assert_eq!(
7877 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
7878 Some("call-1")
7879 );
7880 pending
7881 }
7882 other => panic!("unexpected first loop step: {other:?}"),
7883 };
7884
7885 let pending_second = match driver.next().await.unwrap() {
7886 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
7887 assert_eq!(
7888 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
7889 Some("call-2")
7890 );
7891 pending
7892 }
7893 other => panic!("unexpected second loop step: {other:?}"),
7894 };
7895
7896 pending_second.approve(&mut driver).unwrap();
7897 match driver.next().await.unwrap() {
7898 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
7899 assert_eq!(
7900 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
7901 Some("call-1")
7902 );
7903 }
7904 other => panic!("unexpected step after approving second request: {other:?}"),
7905 }
7906
7907 pending_first.approve(&mut driver).unwrap();
7908 match driver.next().await.unwrap() {
7909 LoopStep::Finished(turn) => {
7910 assert_eq!(turn.finish_reason, FinishReason::Completed);
7911 match &turn.items[0].parts[0] {
7912 Part::Text(text) => assert_eq!(text.text, "both approvals finished"),
7913 other => panic!("unexpected final part: {other:?}"),
7914 }
7915 }
7916 other => panic!("unexpected final loop step: {other:?}"),
7917 }
7918 }
7919
7920 #[tokio::test]
7921 async fn failed_pending_approval_cleanup_repairs_and_finishes_error() {
7922 let events = StdArc::new(StdMutex::new(Vec::new()));
7923 let agent = Agent::builder()
7924 .model(FakeAdapter)
7925 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
7926 .permissions(ApproveFsReads)
7927 .task_manager(
7928 TestTaskManager::new(SimpleTaskManager::new())
7929 .fail_interrupt("interrupt cleanup failed"),
7930 )
7931 .observer(RecordingObserver {
7932 events: events.clone(),
7933 })
7934 .build()
7935 .unwrap();
7936 let mut driver = agent
7937 .start(SessionConfig::new("session-failed-approval-cleanup"))
7938 .await
7939 .unwrap();
7940 driver
7941 .submit_input(vec![Item::text(ItemKind::User, "ping")])
7942 .unwrap();
7943
7944 assert!(matches!(
7945 driver.next().await.unwrap(),
7946 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))
7947 ));
7948 let error = driver.cancel_pending_approvals().await.unwrap_err();
7949 assert!(error.to_string().contains("interrupt cleanup failed"));
7950 assert!(driver.lifecycle.active_turn.is_none());
7951 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
7952
7953 let events = events.lock().unwrap();
7954 assert!(events.iter().any(|event| matches!(
7955 event,
7956 AgentEvent::ToolResultReceived(result)
7957 if result.call_id == ToolCallId::new("call-1") && result.is_error
7958 )));
7959 assert!(events.iter().any(|event| matches!(
7960 event,
7961 AgentEvent::TurnFinished(turn) if turn.finish_reason == FinishReason::Error
7962 )));
7963 }
7964
7965 #[tokio::test]
7966 async fn cancelling_all_pending_approvals_interrupts_every_originating_turn() {
7967 let events = StdArc::new(StdMutex::new(Vec::new()));
7968 let interrupted = StdArc::new(StdMutex::new(Vec::new()));
7969 let tools = ToolRegistry::new().with(EchoTool::default());
7970 let agent = Agent::builder()
7971 .model(DualApprovalAdapter)
7972 .add_tool_source(tools)
7973 .permissions(ApproveFsReads)
7974 .task_manager(
7975 TestTaskManager::new(SimpleTaskManager::new())
7976 .record_interrupts(interrupted.clone()),
7977 )
7978 .observer(RecordingObserver {
7979 events: events.clone(),
7980 })
7981 .build()
7982 .unwrap();
7983
7984 let mut driver = agent
7985 .start(SessionConfig {
7986 session_id: SessionId::new("session-dual-approval-cancel"),
7987 metadata: MetadataMap::new(),
7988 cache: None,
7989 })
7990 .await
7991 .unwrap();
7992
7993 driver
7994 .submit_input(vec![Item {
7995 id: None,
7996 kind: ItemKind::User,
7997 parts: vec![Part::Text(TextPart {
7998 text: "run both approvals".into(),
7999 metadata: MetadataMap::new(),
8000 })],
8001 metadata: MetadataMap::new(),
8002 usage: None,
8003 finish_reason: None,
8004 created_at: None,
8005 }])
8006 .unwrap();
8007
8008 for expected_call in ["call-1", "call-2"] {
8009 match driver.next().await.unwrap() {
8010 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
8011 assert_eq!(
8012 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
8013 Some(expected_call)
8014 );
8015 }
8016 other => panic!("unexpected approval step: {other:?}"),
8017 }
8018 }
8019
8020 let first_origin = driver
8021 .pending_approvals
8022 .get(&ToolCallId::new("call-1"))
8023 .unwrap()
8024 .tool_request
8025 .turn_id
8026 .clone();
8027 let second_origin = agentkit_core::TurnId::new("second-originating-turn");
8028 driver
8029 .pending_approvals
8030 .get_mut(&ToolCallId::new("call-2"))
8031 .unwrap()
8032 .tool_request
8033 .turn_id = second_origin.clone();
8034
8035 match driver.cancel_pending_approvals().await.unwrap() {
8036 Some(LoopStep::Finished(turn)) => {
8037 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
8038 }
8039 other => panic!("unexpected cancellation result: {other:?}"),
8040 }
8041 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
8042 let interrupted = interrupted
8043 .lock()
8044 .unwrap()
8045 .iter()
8046 .cloned()
8047 .collect::<HashSet<_>>();
8048 assert_eq!(interrupted, HashSet::from([first_origin, second_origin]));
8049
8050 let events = events.lock().unwrap();
8051 let cancelled = events
8052 .iter()
8053 .filter(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false }))
8054 .count();
8055 assert_eq!(cancelled, 2);
8056 assert!(events.iter().any(|event| matches!(
8057 event,
8058 AgentEvent::TurnFinished(turn) if turn.finish_reason == FinishReason::Cancelled
8059 )));
8060 for expected_call in ["call-1", "call-2"] {
8061 assert!(events.iter().any(|event| matches!(
8062 event,
8063 AgentEvent::ToolResultReceived(result)
8064 if result.call_id == ToolCallId::new(expected_call) && result.is_error
8065 )));
8066 }
8067 }
8068
8069 #[tokio::test]
8070 async fn loop_compacts_transcript_before_new_turns() {
8071 let events = StdArc::new(StdMutex::new(Vec::new()));
8072 let agent = Agent::builder()
8073 .model(FakeAdapter)
8074 .mutator(KeepRecentMutator { keep: 1 })
8075 .observer(RecordingObserver {
8076 events: events.clone(),
8077 })
8078 .build()
8079 .unwrap();
8080
8081 let mut driver = agent
8082 .start(SessionConfig {
8083 session_id: SessionId::new("session-4"),
8084 metadata: MetadataMap::new(),
8085 cache: None,
8086 })
8087 .await
8088 .unwrap();
8089
8090 for text in ["first", "second"] {
8091 driver
8092 .submit_input(vec![Item {
8093 id: None,
8094 kind: ItemKind::User,
8095 parts: vec![Part::Text(TextPart {
8096 text: text.into(),
8097 metadata: MetadataMap::new(),
8098 })],
8099 metadata: MetadataMap::new(),
8100 usage: None,
8101 finish_reason: None,
8102 created_at: None,
8103 }])
8104 .unwrap();
8105 let _ = driver.next().await.unwrap();
8106 }
8107
8108 let events = events.lock().unwrap();
8109 assert!(
8110 events
8111 .iter()
8112 .any(|event| matches!(event, AgentEvent::MutationFinished { dirty: true, .. }))
8113 );
8114 }
8115
8116 #[test]
8117 fn transcript_validation_rejects_orphaned_tool_result() {
8118 let transcript = vec![Item {
8119 id: None,
8120 kind: ItemKind::Tool,
8121 parts: vec![Part::ToolResult(ToolResultPart {
8122 call_id: "call-1".into(),
8123 output: ToolOutput::Text("result".into()),
8124 is_error: false,
8125 metadata: MetadataMap::new(),
8126 })],
8127 metadata: MetadataMap::new(),
8128 usage: None,
8129 finish_reason: None,
8130 created_at: None,
8131 }];
8132
8133 let error = validate_transcript_invariants(&transcript).unwrap_err();
8134 assert!(error.to_string().contains("orphaned tool_result"));
8135 }
8136
8137 #[test]
8138 fn transcript_validation_rejects_duplicate_tool_result() {
8139 let transcript = vec![
8140 Item {
8141 id: None,
8142 kind: ItemKind::Assistant,
8143 parts: vec![Part::ToolCall(ToolCallPart {
8144 id: "call-1".into(),
8145 name: "lookup".into(),
8146 input: serde_json::json!({}),
8147 metadata: MetadataMap::new(),
8148 })],
8149 metadata: MetadataMap::new(),
8150 usage: None,
8151 finish_reason: None,
8152 created_at: None,
8153 },
8154 Item {
8155 id: None,
8156 kind: ItemKind::Tool,
8157 parts: vec![Part::ToolResult(ToolResultPart {
8158 call_id: "call-1".into(),
8159 output: ToolOutput::Text("result".into()),
8160 is_error: false,
8161 metadata: MetadataMap::new(),
8162 })],
8163 metadata: MetadataMap::new(),
8164 usage: None,
8165 finish_reason: None,
8166 created_at: None,
8167 },
8168 Item {
8169 id: None,
8170 kind: ItemKind::Tool,
8171 parts: vec![Part::ToolResult(ToolResultPart {
8172 call_id: "call-1".into(),
8173 output: ToolOutput::Text("again".into()),
8174 is_error: false,
8175 metadata: MetadataMap::new(),
8176 })],
8177 metadata: MetadataMap::new(),
8178 usage: None,
8179 finish_reason: None,
8180 created_at: None,
8181 },
8182 ];
8183
8184 let error = validate_transcript_invariants(&transcript).unwrap_err();
8185 assert!(error.to_string().contains("duplicate tool_result"));
8186 }
8187
8188 #[tokio::test]
8189 async fn loop_refreshes_tool_specs_each_turn() {
8190 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
8191 let version = StdArc::new(AtomicUsize::new(1));
8192 let tools = ToolRegistry::new().with(DynamicSpecTool::new(version.clone()));
8193 let agent = Agent::builder()
8194 .model(RecordingAdapter {
8195 seen_descriptions: seen_descriptions.clone(),
8196 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
8197 })
8198 .add_tool_source(tools)
8199 .permissions(AllowAllPermissions)
8200 .build()
8201 .unwrap();
8202
8203 let mut driver = agent
8204 .start(SessionConfig {
8205 session_id: SessionId::new("session-dynamic-tools"),
8206 metadata: MetadataMap::new(),
8207 cache: None,
8208 })
8209 .await
8210 .unwrap();
8211
8212 for text in ["first", "second"] {
8213 driver
8214 .submit_input(vec![Item {
8215 id: None,
8216 kind: ItemKind::User,
8217 parts: vec![Part::Text(TextPart {
8218 text: text.into(),
8219 metadata: MetadataMap::new(),
8220 })],
8221 metadata: MetadataMap::new(),
8222 usage: None,
8223 finish_reason: None,
8224 created_at: None,
8225 }])
8226 .unwrap();
8227
8228 let _ = driver.next().await.unwrap();
8229 if text == "first" {
8230 version.store(2, Ordering::SeqCst);
8231 }
8232 }
8233
8234 let seen_descriptions = seen_descriptions.lock().unwrap();
8235 assert_eq!(seen_descriptions.len(), 2);
8236 assert_eq!(seen_descriptions[0], vec!["dynamic version 1".to_string()]);
8237 assert_eq!(seen_descriptions[1], vec!["dynamic version 2".to_string()]);
8238 }
8239
8240 #[tokio::test]
8241 async fn loop_emits_catalog_change_and_uses_updated_specs_next_turn() {
8242 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
8243 let events = StdArc::new(StdMutex::new(Vec::new()));
8244 let executor = StdArc::new(CatalogExecutor::new());
8245 let executor_for_agent: Arc<dyn ToolExecutor> = executor.clone();
8246 let agent = Agent::builder()
8247 .model(RecordingAdapter {
8248 seen_descriptions: seen_descriptions.clone(),
8249 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
8250 })
8251 .tool_executor(executor_for_agent)
8252 .permissions(AllowAllPermissions)
8253 .observer(RecordingObserver {
8254 events: events.clone(),
8255 })
8256 .build()
8257 .unwrap();
8258
8259 let mut driver = agent
8260 .start(SessionConfig {
8261 session_id: SessionId::new("session-catalog-events"),
8262 metadata: MetadataMap::new(),
8263 cache: None,
8264 })
8265 .await
8266 .unwrap();
8267
8268 driver
8269 .submit_input(vec![Item::text(ItemKind::User, "first")])
8270 .unwrap();
8271 let _ = driver.next().await.unwrap();
8272
8273 executor.publish_change(
8274 1,
8275 ToolCatalogEvent {
8276 source: "mcp:mock".into(),
8277 added: vec!["dynamic".into()],
8278 removed: Vec::new(),
8279 changed: Vec::new(),
8280 },
8281 );
8282
8283 driver
8284 .submit_input(vec![Item::text(ItemKind::User, "second")])
8285 .unwrap();
8286 let _ = driver.next().await.unwrap();
8287
8288 let seen_descriptions = seen_descriptions.lock().unwrap();
8289 assert_eq!(seen_descriptions.len(), 2);
8290 assert_eq!(seen_descriptions[0], vec!["dynamic version 0".to_string()]);
8291 assert_eq!(seen_descriptions[1], vec!["dynamic version 1".to_string()]);
8292
8293 let events = events.lock().unwrap();
8294 assert!(events.iter().any(|event| matches!(
8295 event,
8296 AgentEvent::ToolCatalogChanged(ToolCatalogEvent {
8297 source,
8298 added,
8299 removed,
8300 changed,
8301 }) if source == "mcp:mock"
8302 && added == &vec!["dynamic".to_string()]
8303 && removed.is_empty()
8304 && changed.is_empty()
8305 )));
8306 }
8307
8308 #[tokio::test]
8309 async fn loop_passes_session_default_and_next_turn_cache_requests() {
8310 let seen_caches = StdArc::new(StdMutex::new(Vec::new()));
8311 let agent = Agent::builder()
8312 .model(RecordingAdapter {
8313 seen_descriptions: StdArc::new(StdMutex::new(Vec::new())),
8314 seen_caches: seen_caches.clone(),
8315 })
8316 .permissions(AllowAllPermissions)
8317 .build()
8318 .unwrap();
8319
8320 let default_cache = PromptCacheRequest::best_effort(PromptCacheStrategy::Automatic)
8321 .with_retention(PromptCacheRetention::Short);
8322 let override_cache = PromptCacheRequest::required(PromptCacheStrategy::Explicit {
8323 breakpoints: vec![PromptCacheBreakpoint::TranscriptItemEnd { index: 0 }],
8324 });
8325
8326 let mut driver = agent
8327 .start(SessionConfig {
8328 session_id: SessionId::new("session-cache"),
8329 metadata: MetadataMap::new(),
8330 cache: Some(default_cache.clone()),
8331 })
8332 .await
8333 .unwrap();
8334
8335 driver
8336 .submit_input(vec![Item {
8337 id: None,
8338 kind: ItemKind::User,
8339 parts: vec![Part::Text(TextPart {
8340 text: "first".into(),
8341 metadata: MetadataMap::new(),
8342 })],
8343 metadata: MetadataMap::new(),
8344 usage: None,
8345 finish_reason: None,
8346 created_at: None,
8347 }])
8348 .unwrap();
8349 let _ = driver.next().await.unwrap();
8350
8351 driver
8352 .submit_input_with_cache(
8353 vec![Item {
8354 id: None,
8355 kind: ItemKind::User,
8356 parts: vec![Part::Text(TextPart {
8357 text: "second".into(),
8358 metadata: MetadataMap::new(),
8359 })],
8360 metadata: MetadataMap::new(),
8361 usage: None,
8362 finish_reason: None,
8363 created_at: None,
8364 }],
8365 override_cache.clone(),
8366 )
8367 .unwrap();
8368 let _ = driver.next().await.unwrap();
8369
8370 let seen = seen_caches.lock().unwrap();
8371 assert_eq!(seen.len(), 2);
8372 assert_eq!(seen[0], Some(default_cache));
8373 assert_eq!(seen[1], Some(override_cache));
8374 }
8375
8376 #[tokio::test]
8377 async fn loop_yields_after_tool_result_between_rounds() {
8378 let tools = ToolRegistry::new().with(EchoTool::default());
8379 let agent = Agent::builder()
8380 .model(FakeAdapter)
8381 .add_tool_source(tools)
8382 .permissions(AllowAllPermissions)
8383 .build()
8384 .unwrap();
8385
8386 let mut driver = agent
8387 .start(SessionConfig {
8388 session_id: SessionId::new("yield-session"),
8389 metadata: MetadataMap::new(),
8390 cache: None,
8391 })
8392 .await
8393 .unwrap();
8394
8395 driver
8396 .submit_input(vec![Item::text(ItemKind::User, "ping")])
8397 .unwrap();
8398
8399 let step = driver.next().await.unwrap();
8402 let info = match step {
8403 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info,
8404 other => panic!("expected AfterToolResult, got {other:?}"),
8405 };
8406 assert_eq!(info.session_id, SessionId::new("yield-session"));
8407 assert_eq!(info.transcript_len, 3);
8409
8410 let interrupt = LoopInterrupt::AfterToolResult(info.clone());
8412 assert!(!interrupt.is_blocking());
8413
8414 driver
8416 .submit_input(vec![Item::text(ItemKind::User, "also: report back")])
8417 .unwrap();
8418
8419 let step = driver.next().await.unwrap();
8422 match step {
8423 LoopStep::Finished(turn) => {
8424 assert_eq!(turn.finish_reason, FinishReason::Completed);
8425 }
8426 other => panic!("expected Finished, got {other:?}"),
8427 }
8428
8429 let snapshot = driver.snapshot();
8431 let has_injected_message = snapshot.transcript.iter().any(|item| {
8432 item.kind == ItemKind::User
8433 && item.parts.iter().any(|part| match part {
8434 Part::Text(text) => text.text == "also: report back",
8435 _ => false,
8436 })
8437 });
8438 assert!(
8439 has_injected_message,
8440 "injected user message should be in transcript, got: {:?}",
8441 snapshot.transcript
8442 );
8443 }
8444
8445 struct RecordingTranscriptObserver {
8446 items: StdArc<StdMutex<Vec<Item>>>,
8447 }
8448
8449 impl TranscriptObserver for RecordingTranscriptObserver {
8450 fn on_transcript_event(&self, event: TranscriptEvent<'_>) {
8451 self.items.lock().unwrap().push(event.item.clone());
8452 }
8453 }
8454
8455 #[tokio::test]
8456 async fn observers_see_full_tool_round() {
8457 let events = StdArc::new(StdMutex::new(Vec::<AgentEvent>::new()));
8463 let items = StdArc::new(StdMutex::new(Vec::<Item>::new()));
8464 let agent = Agent::builder()
8465 .model(FakeAdapter)
8466 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
8467 .permissions(AllowAllPermissions)
8468 .observer(RecordingObserver {
8469 events: events.clone(),
8470 })
8471 .transcript_observer(RecordingTranscriptObserver {
8472 items: items.clone(),
8473 })
8474 .build()
8475 .unwrap();
8476
8477 let mut driver = agent
8478 .start(SessionConfig {
8479 session_id: SessionId::new("observer-session"),
8480 metadata: MetadataMap::new(),
8481 cache: None,
8482 })
8483 .await
8484 .unwrap();
8485
8486 driver
8487 .submit_input(vec![Item {
8488 id: None,
8489 kind: ItemKind::User,
8490 parts: vec![Part::Text(TextPart {
8491 text: "ping".into(),
8492 metadata: MetadataMap::new(),
8493 })],
8494 metadata: MetadataMap::new(),
8495 usage: None,
8496 finish_reason: None,
8497 created_at: None,
8498 }])
8499 .unwrap();
8500
8501 let result = run_until_finished(&mut driver).await;
8502 assert!(matches!(result, LoopStep::Finished(_)), "got {result:?}");
8503
8504 let events = events.lock().unwrap().clone();
8507 let tool_call_id = events.iter().find_map(|e| match e {
8508 AgentEvent::ToolCallRequested(c) => Some(c.id.clone()),
8509 _ => None,
8510 });
8511 let tool_results: Vec<_> = events
8512 .iter()
8513 .filter_map(|e| match e {
8514 AgentEvent::ToolResultReceived(r) => Some(r.clone()),
8515 _ => None,
8516 })
8517 .collect();
8518 assert_eq!(tool_results.len(), 1, "events: {events:?}");
8519 assert_eq!(Some(tool_results[0].call_id.clone()), tool_call_id);
8520 assert!(!tool_results[0].is_error);
8521
8522 let items = items.lock().unwrap().clone();
8526 assert_eq!(items.len(), 4, "items: {items:?}");
8527 assert_eq!(items[0].kind, ItemKind::User);
8528 assert_eq!(items[1].kind, ItemKind::Assistant);
8529 assert!(
8530 items[1]
8531 .parts
8532 .iter()
8533 .any(|p| matches!(p, Part::ToolCall(_)))
8534 );
8535 assert_eq!(items[2].kind, ItemKind::Tool);
8536 assert!(
8537 items[2]
8538 .parts
8539 .iter()
8540 .any(|p| matches!(p, Part::ToolResult(_)))
8541 );
8542 assert_eq!(items[3].kind, ItemKind::Assistant);
8543 }
8544
8545 #[test]
8546 fn convenience_cache_builders_construct_expected_defaults() {
8547 let cache = PromptCacheRequest::automatic()
8548 .with_retention(PromptCacheRetention::Short)
8549 .with_key("workspace:demo");
8550 let session = SessionConfig::new("demo").with_cache(cache.clone());
8551
8552 assert_eq!(session.session_id, SessionId::new("demo"));
8553 assert_eq!(session.cache, Some(cache));
8554
8555 let explicit = PromptCacheRequest::explicit([
8556 PromptCacheBreakpoint::tools_end(),
8557 PromptCacheBreakpoint::transcript_item_end(2),
8558 PromptCacheBreakpoint::transcript_part_end(3, 1),
8559 ]);
8560
8561 assert_eq!(explicit.mode, PromptCacheMode::BestEffort);
8562 assert_eq!(
8563 explicit.strategy,
8564 PromptCacheStrategy::Explicit {
8565 breakpoints: vec![
8566 PromptCacheBreakpoint::ToolsEnd,
8567 PromptCacheBreakpoint::TranscriptItemEnd { index: 2 },
8568 PromptCacheBreakpoint::TranscriptPartEnd {
8569 item_index: 3,
8570 part_index: 1,
8571 },
8572 ],
8573 }
8574 );
8575 }
8576}