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