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