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