1use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
65use std::sync::Arc;
66
67use agentkit_core::{
68 CancellationHandle, DataRef, Delta, FinishReason, Item, ItemKind, MetadataMap, Modality, Part,
69 SessionId, TaskId, TextPart, Timestamp, ToolCallId, ToolCallPart, ToolOutput, ToolResultPart,
70 TurnCancellation, Usage,
71};
72use agentkit_task_manager::{
73 PendingLoopUpdates, SimpleTaskManager, TOOL_RESULT_NOT_STARTED_METADATA_KEY, TaskApproval,
74 TaskLaunchKind, TaskLaunchRequest, TaskManager, TaskResolution, TaskStartContext,
75 TaskStartOutcome, TurnTaskUpdate,
76};
77#[cfg(test)]
78use agentkit_task_manager::{
79 TOOL_RESULT_FAILURE_KIND_METADATA_KEY, TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED,
80};
81#[cfg(test)]
82use agentkit_tools_core::ToolContext;
83use agentkit_tools_core::{
84 AllowAllPermissions, ApprovalDecision, ApprovalRequest, BasicToolExecutor, OwnedToolContext,
85 PermissionChecker, ToolCatalogEvent, ToolError, ToolExecutionScope, ToolExecutor, ToolRequest,
86 ToolResources, ToolSource, ToolSpec,
87};
88use async_trait::async_trait;
89use serde::{Deserialize, Serialize};
90use serde_json::Value;
91use thiserror::Error;
92
93const INTERRUPTED_METADATA_KEY: &str = "agentkit.interrupted";
94const INTERRUPT_REASON_METADATA_KEY: &str = "agentkit.interrupt_reason";
95const INTERRUPT_STAGE_METADATA_KEY: &str = "agentkit.interrupt_stage";
96const USER_CANCELLED_REASON: &str = "user_cancelled";
97
98pub const PROVIDER_FINISH_REASONS_METADATA_KEY: &str = "agentkit.provider_finish_reasons";
100
101pub fn set_provider_finish_reasons<I, S>(metadata: &mut MetadataMap, reasons: I)
103where
104 I: IntoIterator<Item = S>,
105 S: Into<String>,
106{
107 let mut seen = HashSet::new();
108 let reasons = reasons
109 .into_iter()
110 .map(Into::into)
111 .filter(|reason: &String| !reason.is_empty() && seen.insert(reason.clone()))
112 .map(Value::String)
113 .collect::<Vec<_>>();
114 metadata.remove(PROVIDER_FINISH_REASONS_METADATA_KEY);
115 if !reasons.is_empty() {
116 metadata.insert(
117 PROVIDER_FINISH_REASONS_METADATA_KEY.into(),
118 Value::Array(reasons),
119 );
120 }
121}
122
123fn provider_finish_reasons(metadata: &MetadataMap, fallback: &FinishReason) -> Vec<String> {
124 metadata
125 .get(PROVIDER_FINISH_REASONS_METADATA_KEY)
126 .and_then(Value::as_array)
127 .map(|values| {
128 let mut seen = HashSet::new();
129 values
130 .iter()
131 .filter_map(Value::as_str)
132 .filter(|reason| !reason.is_empty() && seen.insert((*reason).to_owned()))
133 .map(str::to_owned)
134 .collect::<Vec<_>>()
135 })
136 .filter(|reasons| !reasons.is_empty())
137 .unwrap_or_else(|| vec![normalized_finish_reason(fallback).into()])
138}
139
140fn normalized_finish_reason(reason: &FinishReason) -> &str {
141 match reason {
142 FinishReason::Completed => "completed",
143 FinishReason::ToolCall => "tool_call",
144 FinishReason::MaxTokens => "max_tokens",
145 FinishReason::Cancelled => "cancelled",
146 FinishReason::Blocked => "blocked",
147 FinishReason::Error => "error",
148 FinishReason::Other(reason) => reason,
149 }
150}
151
152#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)]
154pub enum MessageCaptureError {
155 #[error("message capture max_messages must be nonzero")]
157 ZeroMessages,
158 #[error("message capture max_bytes must be nonzero")]
160 ZeroBytes,
161}
162
163#[derive(Clone, Copy, Debug, PartialEq, Eq)]
165pub struct MessageCapture {
166 max_messages: usize,
167 max_bytes: usize,
168}
169
170impl MessageCapture {
171 pub fn new(max_messages: usize, max_bytes: usize) -> Result<Self, MessageCaptureError> {
173 if max_messages == 0 {
174 return Err(MessageCaptureError::ZeroMessages);
175 }
176 if max_bytes == 0 {
177 return Err(MessageCaptureError::ZeroBytes);
178 }
179 Ok(Self {
180 max_messages,
181 max_bytes,
182 })
183 }
184
185 pub fn max_messages(self) -> usize {
187 self.max_messages
188 }
189
190 pub fn max_bytes(self) -> usize {
192 self.max_bytes
193 }
194}
195
196#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
202pub struct TelemetryConfig {
203 input_messages: Option<MessageCapture>,
204 output_messages: Option<MessageCapture>,
205}
206
207impl TelemetryConfig {
208 pub fn with_input_messages(mut self, capture: MessageCapture) -> Self {
210 self.input_messages = Some(capture);
211 self
212 }
213
214 pub fn with_output_messages(mut self, capture: MessageCapture) -> Self {
216 self.output_messages = Some(capture);
217 self
218 }
219
220 pub fn without_input_messages(mut self) -> Self {
222 self.input_messages = None;
223 self
224 }
225
226 pub fn without_output_messages(mut self) -> Self {
228 self.output_messages = None;
229 self
230 }
231
232 pub fn input_messages(self) -> Option<MessageCapture> {
234 self.input_messages
235 }
236
237 pub fn output_messages(self) -> Option<MessageCapture> {
239 self.output_messages
240 }
241}
242
243#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
258pub struct SessionConfig {
259 pub session_id: SessionId,
261 pub metadata: MetadataMap,
263 pub cache: Option<PromptCacheRequest>,
265}
266
267impl SessionConfig {
268 pub fn new(session_id: impl Into<SessionId>) -> Self {
270 Self {
271 session_id: session_id.into(),
272 metadata: MetadataMap::new(),
273 cache: None,
274 }
275 }
276
277 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
279 self.metadata = metadata;
280 self
281 }
282
283 pub fn with_cache(mut self, cache: PromptCacheRequest) -> Self {
285 self.cache = Some(cache);
286 self
287 }
288
289 pub fn without_cache(mut self) -> Self {
291 self.cache = None;
292 self
293 }
294}
295
296#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
302pub enum PromptCacheMode {
303 Disabled,
305 #[default]
307 BestEffort,
308 Required,
310}
311
312#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
318pub enum PromptCacheRetention {
319 Default,
321 Short,
323 Extended,
325}
326
327#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
329pub enum PromptCacheStrategy {
330 #[default]
332 Automatic,
333 Explicit {
335 breakpoints: Vec<PromptCacheBreakpoint>,
337 },
338}
339
340impl PromptCacheStrategy {
341 pub fn automatic() -> Self {
344 Self::Automatic
345 }
346
347 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
349 Self::Explicit {
350 breakpoints: breakpoints.into_iter().collect(),
351 }
352 }
353}
354
355#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
357pub enum PromptCacheBreakpoint {
358 ToolsEnd,
360 TranscriptItemEnd { index: usize },
362 TranscriptPartEnd {
368 item_index: usize,
369 part_index: usize,
370 },
371}
372
373impl PromptCacheBreakpoint {
374 pub fn tools_end() -> Self {
376 Self::ToolsEnd
377 }
378
379 pub fn transcript_item_end(index: usize) -> Self {
381 Self::TranscriptItemEnd { index }
382 }
383
384 pub fn transcript_part_end(item_index: usize, part_index: usize) -> Self {
386 Self::TranscriptPartEnd {
387 item_index,
388 part_index,
389 }
390 }
391}
392
393#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
395pub struct PromptCacheRequest {
396 pub mode: PromptCacheMode,
398 pub strategy: PromptCacheStrategy,
400 pub retention: Option<PromptCacheRetention>,
402 pub key: Option<String>,
404}
405
406impl PromptCacheRequest {
407 pub fn automatic() -> Self {
409 Self::best_effort(PromptCacheStrategy::automatic())
410 }
411
412 pub fn automatic_required() -> Self {
414 Self::required(PromptCacheStrategy::automatic())
415 }
416
417 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
419 Self::best_effort(PromptCacheStrategy::explicit(breakpoints))
420 }
421
422 pub fn explicit_required(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
424 Self::required(PromptCacheStrategy::explicit(breakpoints))
425 }
426
427 pub fn disabled() -> Self {
429 Self {
430 mode: PromptCacheMode::Disabled,
431 strategy: PromptCacheStrategy::Automatic,
432 retention: None,
433 key: None,
434 }
435 }
436
437 pub fn best_effort(strategy: PromptCacheStrategy) -> Self {
439 Self {
440 mode: PromptCacheMode::BestEffort,
441 strategy,
442 retention: None,
443 key: None,
444 }
445 }
446
447 pub fn required(strategy: PromptCacheStrategy) -> Self {
449 Self {
450 mode: PromptCacheMode::Required,
451 strategy,
452 retention: None,
453 key: None,
454 }
455 }
456
457 pub fn with_mode(mut self, mode: PromptCacheMode) -> Self {
459 self.mode = mode;
460 self
461 }
462
463 pub fn with_strategy(mut self, strategy: PromptCacheStrategy) -> Self {
465 self.strategy = strategy;
466 self
467 }
468
469 pub fn with_retention(mut self, retention: PromptCacheRetention) -> Self {
471 self.retention = Some(retention);
472 self
473 }
474
475 pub fn with_key(mut self, key: impl Into<String>) -> Self {
477 self.key = Some(key.into());
478 self
479 }
480
481 pub fn without_retention(mut self) -> Self {
483 self.retention = None;
484 self
485 }
486
487 pub fn without_key(mut self) -> Self {
489 self.key = None;
490 self
491 }
492
493 pub fn is_enabled(&self) -> bool {
495 !matches!(self.mode, PromptCacheMode::Disabled)
496 }
497}
498
499#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
505pub struct TurnRequest {
506 pub session_id: SessionId,
508 pub turn_id: agentkit_core::TurnId,
510 pub transcript: Vec<Item>,
512 pub available_tools: Vec<ToolSpec>,
514 pub cache: Option<PromptCacheRequest>,
516 pub metadata: MetadataMap,
518}
519
520#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
525pub struct ModelTurnResult {
526 pub finish_reason: FinishReason,
528 pub output_items: Vec<Item>,
530 pub usage: Option<Usage>,
532 pub metadata: MetadataMap,
534 #[serde(default)]
538 pub model: Option<String>,
539 #[serde(default)]
543 pub response_id: Option<String>,
544}
545
546#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
552pub enum ModelTurnEvent {
553 Delta(Delta),
555 ToolCall(ToolCallPart),
557 Usage(Usage),
559 Finished(ModelTurnResult),
561}
562
563#[async_trait]
600pub trait ModelAdapter: Send + Sync {
601 type Session: ModelSession;
603
604 async fn start_session(&self, config: SessionConfig) -> Result<Self::Session, LoopError>;
610
611 fn provider_name(&self) -> Option<&str> {
618 None
619 }
620}
621
622#[async_trait]
629pub trait ModelSession: Send {
630 type Turn: ModelTurn;
632
633 async fn begin_turn(
646 &mut self,
647 request: TurnRequest,
648 cancellation: Option<TurnCancellation>,
649 ) -> Result<Self::Turn, LoopError>;
650
651 fn model_name(&self) -> Option<&str> {
657 None
658 }
659
660 fn provider_name(&self) -> Option<&str> {
665 None
666 }
667}
668
669#[async_trait]
675pub trait ModelTurn: Send {
676 async fn next_event(
685 &mut self,
686 cancellation: Option<TurnCancellation>,
687 ) -> Result<Option<ModelTurnEvent>, LoopError>;
688}
689
690pub trait LoopObserver: Send + Sync {
710 fn handle_event(&self, event: ObservedEvent);
715}
716
717#[derive(Clone, Debug, PartialEq)]
723pub struct ObservedEvent {
724 pub session_id: Arc<SessionId>,
726 pub event: AgentEvent,
728}
729
730pub trait TranscriptObserver: Send + Sync {
765 fn on_transcript_event(&self, event: TranscriptEvent<'_>);
770}
771
772#[derive(Clone, Debug)]
775pub struct TranscriptEvent<'a> {
776 pub session_id: &'a SessionId,
778 pub item: &'a Item,
780}
781
782#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
787#[non_exhaustive]
788pub enum MutationPoint {
789 AfterToolResult,
792 AfterTurnEnded,
795}
796
797pub trait EventEmitter: Send + Sync {
800 fn emit(&self, event: AgentEvent);
802}
803
804#[non_exhaustive]
806pub struct LoopCtx<'a> {
807 pub session_id: &'a SessionId,
809 pub turn_id: Option<&'a agentkit_core::TurnId>,
811 pub point: MutationPoint,
813 pub cancellation: Option<TurnCancellation>,
815 pub emitter: &'a dyn EventEmitter,
817}
818
819pub struct TranscriptCursor<'a> {
827 items: &'a mut Vec<Item>,
828 pub(crate) dirty: bool,
829}
830
831impl<'a> std::ops::Deref for TranscriptCursor<'a> {
832 type Target = Vec<Item>;
833 fn deref(&self) -> &Vec<Item> {
834 self.items
835 }
836}
837
838impl<'a> std::ops::DerefMut for TranscriptCursor<'a> {
839 fn deref_mut(&mut self) -> &mut Vec<Item> {
840 self.dirty = true;
841 self.items
842 }
843}
844
845#[async_trait]
853pub trait LoopMutator: Send + Sync {
854 async fn mutate(
859 &self,
860 cursor: &mut TranscriptCursor<'_>,
861 ctx: LoopCtx<'_>,
862 ) -> Result<(), LoopError> {
863 let _ = (cursor, ctx);
864 Ok(())
865 }
866}
867
868#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
873#[non_exhaustive]
874pub enum AgentEvent {
875 RunStarted { session_id: SessionId },
877 TurnStarted {
879 session_id: SessionId,
880 turn_id: agentkit_core::TurnId,
881 },
882 InputAccepted {
884 session_id: SessionId,
885 items: Vec<Item>,
886 },
887 ContentDelta(Delta),
889 ToolCallRequested(ToolCallPart),
891 ToolExecutionStarted(ToolCallPart),
893 ToolExecutionProgress(ToolResultPart),
899 ToolResultReceived(ToolResultPart),
908 ApprovalRequired(ApprovalRequest),
910 ApprovalResolved { approved: bool },
912 ToolCatalogChanged(ToolCatalogEvent),
914 MutationStarted {
917 session_id: SessionId,
918 turn_id: Option<agentkit_core::TurnId>,
919 mutator: String,
920 point: MutationPoint,
921 },
922 MutationFinished {
926 session_id: SessionId,
927 turn_id: Option<agentkit_core::TurnId>,
928 mutator: String,
929 dirty: bool,
930 metadata: MetadataMap,
931 },
932 UsageUpdated(Usage),
934 Warning { message: String },
936 RunFailed { message: String },
938 TurnFinished(TurnResult),
940}
941
942#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
964pub struct PendingApproval {
965 pub request: ApprovalRequest,
967}
968
969impl std::ops::Deref for PendingApproval {
970 type Target = ApprovalRequest;
971 fn deref(&self) -> &ApprovalRequest {
972 &self.request
973 }
974}
975
976impl PendingApproval {
977 pub fn approve<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
979 let call_id = self
980 .request
981 .call_id
982 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
983 driver.resolve_approval_for(call_id, ApprovalDecision::Approve)
984 }
985
986 pub fn deny<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
988 let call_id = self
989 .request
990 .call_id
991 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
992 driver.resolve_approval_for(call_id, ApprovalDecision::Deny { reason: None })
993 }
994
995 pub fn deny_with_reason<S: ModelSession>(
997 self,
998 driver: &mut LoopDriver<S>,
999 reason: impl Into<String>,
1000 ) -> Result<(), LoopError> {
1001 let call_id = self
1002 .request
1003 .call_id
1004 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
1005 driver.resolve_approval_for(
1006 call_id,
1007 ApprovalDecision::Deny {
1008 reason: Some(reason.into()),
1009 },
1010 )
1011 }
1012
1013 pub fn approve_with_patched_input<S: ModelSession>(
1023 self,
1024 driver: &mut LoopDriver<S>,
1025 input: serde_json::Value,
1026 ) -> Result<(), LoopError> {
1027 let call_id = self
1028 .request
1029 .call_id
1030 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
1031 driver.resolve_approval_for_with_patched_input(call_id, input)
1032 }
1033}
1034
1035#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1060pub struct InputRequest {
1061 pub session_id: SessionId,
1063 pub reason: String,
1065}
1066
1067impl InputRequest {
1068 pub fn submit<S: ModelSession>(
1070 self,
1071 driver: &mut LoopDriver<S>,
1072 items: Vec<Item>,
1073 ) -> Result<(), LoopError> {
1074 driver.submit_input(items)
1075 }
1076}
1077
1078#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1083pub struct TurnResult {
1084 pub turn_id: agentkit_core::TurnId,
1086 pub finish_reason: FinishReason,
1088 pub items: Vec<Item>,
1090 pub usage: Option<Usage>,
1092 pub metadata: MetadataMap,
1094}
1095
1096#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1132pub enum LoopInterrupt {
1133 ApprovalRequest(PendingApproval),
1135 AwaitingInput(InputRequest),
1137 AfterToolResult(ToolRoundInfo),
1147}
1148
1149impl LoopInterrupt {
1150 pub fn is_blocking(&self) -> bool {
1156 matches!(self, LoopInterrupt::ApprovalRequest(_))
1157 }
1158}
1159
1160#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1163pub struct ToolRoundInfo {
1164 pub session_id: SessionId,
1166 pub turn_id: agentkit_core::TurnId,
1168 pub transcript_len: usize,
1170}
1171
1172impl ToolRoundInfo {
1173 pub fn submit<S: ModelSession>(
1176 self,
1177 driver: &mut LoopDriver<S>,
1178 items: Vec<Item>,
1179 ) -> Result<(), LoopError> {
1180 driver.submit_input(items)
1181 }
1182}
1183
1184#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1212pub enum LoopStep {
1213 Interrupt(LoopInterrupt),
1215 Finished(TurnResult),
1217}
1218
1219#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1225pub struct LoopSnapshot {
1226 pub session_id: SessionId,
1228 pub transcript: Vec<Item>,
1230 pub pending_input: Vec<Item>,
1232}
1233
1234#[derive(Clone)]
1235struct PendingApprovalToolCall {
1236 request: ApprovalRequest,
1237 decision: Option<ApprovalDecision>,
1238 surfaced: bool,
1239 turn_id: agentkit_core::TurnId,
1240 task_id: TaskId,
1241 call: ToolCallPart,
1242 tool_request: ToolRequest,
1243 cancellation: Option<TurnCancellation>,
1244}
1245
1246#[derive(Clone, Default)]
1247struct ActiveToolRound {
1248 turn_id: agentkit_core::TurnId,
1249 pending_calls: VecDeque<(ToolCallPart, ToolRequest)>,
1250 cancellation: Option<TurnCancellation>,
1251 background_pending: bool,
1252 foreground_progressed: bool,
1253}
1254
1255pub struct Agent<M>
1297where
1298 M: ModelAdapter,
1299{
1300 model: M,
1301 tool_sources: Vec<Arc<dyn ToolSource>>,
1302 tool_executor: Option<Arc<dyn ToolExecutor>>,
1303 task_manager: Arc<dyn TaskManager>,
1304 permissions: Arc<dyn PermissionChecker>,
1305 resources: Arc<dyn ToolResources>,
1306 cancellation: Option<CancellationHandle>,
1307 mutators: Vec<Arc<dyn LoopMutator>>,
1308 observers: Vec<Arc<dyn LoopObserver>>,
1309 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1310 transcript: Vec<Item>,
1311 input: Vec<Item>,
1312 telemetry: TelemetryConfig,
1313}
1314
1315impl<M> Agent<M>
1316where
1317 M: ModelAdapter,
1318{
1319 pub fn builder() -> AgentBuilder<M> {
1321 AgentBuilder::default()
1322 }
1323
1324 pub async fn start(&self, config: SessionConfig) -> Result<LoopDriver<M::Session>, LoopError> {
1340 let session_id = config.session_id.clone();
1341 let default_cache = config.cache.clone();
1342 let session = self.model.start_session(config).await?;
1343 let provider_name = self.model.provider_name().map(str::to_owned);
1344 let tool_executor = self
1345 .tool_executor
1346 .clone()
1347 .unwrap_or_else(|| Arc::new(BasicToolExecutor::new(self.tool_sources.clone())));
1348 let driver = LoopDriver {
1349 session_id: session_id.clone(),
1350 observed_session_id: Arc::new(session_id.clone()),
1351 provider_name,
1352 telemetry: self.telemetry,
1353 default_cache,
1354 next_turn_cache: None,
1355 session: Some(session),
1356 tool_executor,
1357 task_manager: self.task_manager.clone(),
1358 permissions: self.permissions.clone(),
1359 resources: self.resources.clone(),
1360 cancellation: self.cancellation.clone(),
1361 mutators: self.mutators.clone(),
1362 observers: self.observers.clone(),
1363 transcript_observers: self.transcript_observers.clone(),
1364 transcript: self.transcript.clone(),
1365 pending_input: self.input.clone(),
1366 pending_approvals: BTreeMap::new(),
1367 pending_approval_order: VecDeque::new(),
1368 active_tool_round: None,
1369 pending_round_resume: None,
1370 next_turn_index: 1,
1371 background_call_ids: HashSet::new(),
1372 detached_call_ids: HashSet::new(),
1373 interrupted_background_call_ids: HashSet::new(),
1374 tool_cancellations: HashMap::new(),
1375 };
1376 driver.emit(AgentEvent::RunStarted { session_id });
1377 Ok(driver)
1378 }
1379}
1380
1381pub struct AgentBuilder<M>
1387where
1388 M: ModelAdapter,
1389{
1390 model: Option<M>,
1391 tool_sources: Vec<Arc<dyn ToolSource>>,
1392 tool_executor: Option<Arc<dyn ToolExecutor>>,
1393 task_manager: Option<Arc<dyn TaskManager>>,
1394 permissions: Arc<dyn PermissionChecker>,
1395 resources: Arc<dyn ToolResources>,
1396 cancellation: Option<CancellationHandle>,
1397 mutators: Vec<Arc<dyn LoopMutator>>,
1398 observers: Vec<Arc<dyn LoopObserver>>,
1399 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1400 transcript: Vec<Item>,
1401 input: Vec<Item>,
1402 telemetry: TelemetryConfig,
1403}
1404
1405impl<M> Default for AgentBuilder<M>
1406where
1407 M: ModelAdapter,
1408{
1409 fn default() -> Self {
1410 Self {
1411 model: None,
1412 tool_sources: Vec::new(),
1413 tool_executor: None,
1414 task_manager: None,
1415 permissions: Arc::new(AllowAllPermissions),
1416 resources: Arc::new(()),
1417 cancellation: None,
1418 mutators: Vec::new(),
1419 observers: Vec::new(),
1420 transcript_observers: Vec::new(),
1421 transcript: Vec::new(),
1422 input: Vec::new(),
1423 telemetry: TelemetryConfig::default(),
1424 }
1425 }
1426}
1427
1428impl<M> AgentBuilder<M>
1429where
1430 M: ModelAdapter,
1431{
1432 pub fn model(mut self, model: M) -> Self {
1434 self.model = Some(model);
1435 self
1436 }
1437
1438 pub fn add_tool_source<S: ToolSource + 'static>(mut self, source: S) -> Self {
1451 self.tool_sources.push(Arc::new(source));
1452 self
1453 }
1454
1455 pub fn tool_executor(mut self, executor: impl ToolExecutor + 'static) -> Self {
1461 self.tool_executor = Some(Arc::new(executor));
1462 self
1463 }
1464
1465 pub fn task_manager(mut self, manager: impl TaskManager + 'static) -> Self {
1470 self.task_manager = Some(Arc::new(manager));
1471 self
1472 }
1473
1474 pub fn permissions(mut self, permissions: impl PermissionChecker + 'static) -> Self {
1478 self.permissions = Arc::new(permissions);
1479 self
1480 }
1481
1482 pub fn resources(mut self, resources: impl ToolResources + 'static) -> Self {
1484 self.resources = Arc::new(resources);
1485 self
1486 }
1487
1488 pub fn cancellation(mut self, handle: CancellationHandle) -> Self {
1490 self.cancellation = Some(handle);
1491 self
1492 }
1493
1494 pub fn mutator<L: LoopMutator + 'static>(mut self, mutator: L) -> Self {
1502 self.mutators.push(Arc::new(mutator));
1503 self
1504 }
1505
1506 pub fn observer<O: LoopObserver + 'static>(mut self, observer: O) -> Self {
1510 self.observers.push(Arc::new(observer));
1511 self
1512 }
1513
1514 pub fn transcript_observer<O: TranscriptObserver + 'static>(mut self, observer: O) -> Self {
1523 self.transcript_observers.push(Arc::new(observer));
1524 self
1525 }
1526
1527 pub fn transcript(mut self, transcript: Vec<Item>) -> Self {
1535 self.transcript = transcript;
1536 self
1537 }
1538
1539 pub fn input(mut self, input: Vec<Item>) -> Self {
1548 self.input = input;
1549 self
1550 }
1551
1552 pub fn telemetry(mut self, telemetry: TelemetryConfig) -> Self {
1555 self.telemetry = telemetry;
1556 self
1557 }
1558
1559 pub fn build(self) -> Result<Agent<M>, LoopError> {
1565 let model = self
1566 .model
1567 .ok_or_else(|| LoopError::InvalidState("model adapter is required".into()))?;
1568 Ok(Agent {
1569 model,
1570 tool_sources: self.tool_sources,
1571 tool_executor: self.tool_executor,
1572 task_manager: self
1573 .task_manager
1574 .unwrap_or_else(|| Arc::new(SimpleTaskManager::new())),
1575 permissions: self.permissions,
1576 resources: self.resources,
1577 cancellation: self.cancellation,
1578 mutators: self.mutators,
1579 observers: self.observers,
1580 transcript_observers: self.transcript_observers,
1581 transcript: self.transcript,
1582 input: self.input,
1583 telemetry: self.telemetry,
1584 })
1585 }
1586}
1587
1588pub struct LoopDriver<S>
1619where
1620 S: ModelSession,
1621{
1622 session_id: SessionId,
1623 observed_session_id: Arc<SessionId>,
1624 provider_name: Option<String>,
1625 telemetry: TelemetryConfig,
1626 default_cache: Option<PromptCacheRequest>,
1627 next_turn_cache: Option<PromptCacheRequest>,
1628 session: Option<S>,
1629 tool_executor: Arc<dyn ToolExecutor>,
1630 task_manager: Arc<dyn TaskManager>,
1631 permissions: Arc<dyn PermissionChecker>,
1632 resources: Arc<dyn ToolResources>,
1633 cancellation: Option<CancellationHandle>,
1634 mutators: Vec<Arc<dyn LoopMutator>>,
1635 observers: Vec<Arc<dyn LoopObserver>>,
1636 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1637 transcript: Vec<Item>,
1638 pending_input: Vec<Item>,
1639 pending_approvals: BTreeMap<ToolCallId, PendingApprovalToolCall>,
1640 pending_approval_order: VecDeque<ToolCallId>,
1641 active_tool_round: Option<ActiveToolRound>,
1642 pending_round_resume: Option<agentkit_core::TurnId>,
1643 next_turn_index: u64,
1644 background_call_ids: HashSet<ToolCallId>,
1646 detached_call_ids: HashSet<ToolCallId>,
1654 interrupted_background_call_ids: HashSet<ToolCallId>,
1658 tool_cancellations: HashMap<ToolCallId, TurnCancellation>,
1659}
1660
1661impl<S> LoopDriver<S>
1662where
1663 S: ModelSession,
1664{
1665 fn execute_tool_span(
1666 &self,
1667 request: &ToolRequest,
1668 turn_id: &agentkit_core::TurnId,
1669 launch_kind: &'static str,
1670 ) -> tracing::Span {
1671 tracing::info_span!(
1672 "agent.execute_tool",
1673 "otel.name" = %format!("execute_tool {}", request.tool_name),
1674 "gen_ai.operation.name" = "execute_tool",
1675 "gen_ai.tool.name" = %request.tool_name,
1676 "gen_ai.tool.call.id" = %request.call_id,
1677 "gen_ai.conversation.id" = %self.session_id,
1678 "error.type" = tracing::field::Empty,
1679 session.id = %self.session_id,
1680 turn.id = %turn_id,
1681 launch_kind = launch_kind,
1682 )
1683 }
1684
1685 fn start_task_via_manager(
1686 &self,
1687 task_id: Option<TaskId>,
1688 tool_request: ToolRequest,
1689 kind: TaskLaunchKind,
1690 cancellation: Option<TurnCancellation>,
1691 ) -> impl std::future::Future<Output = Result<TaskStartOutcome, LoopError>> + Send + 'static
1692 {
1693 let task_manager = self.task_manager.clone();
1694 let tool_executor = self.tool_executor.clone();
1695 let permissions = self.permissions.clone();
1696 let resources = self.resources.clone();
1697 let session_id = self.session_id.clone();
1698 let turn_id = tool_request.turn_id.clone();
1699 let metadata = tool_request.metadata.clone();
1700
1701 async move {
1702 task_manager
1703 .start_task(
1704 TaskLaunchRequest {
1705 task_id,
1706 request: tool_request.clone(),
1707 kind,
1708 },
1709 TaskStartContext {
1710 executor: tool_executor.clone(),
1711 tool_context: {
1712 let execution_scope = ToolExecutionScope {
1713 executor: tool_executor,
1714 session_id: session_id.clone(),
1715 turn_id: turn_id.clone(),
1716 permissions: permissions.clone(),
1717 resources: resources.clone(),
1718 cancellation: cancellation.clone(),
1719 };
1720 OwnedToolContext {
1721 session_id,
1722 turn_id,
1723 metadata,
1724 permissions,
1725 resources,
1726 cancellation,
1727 execution_scope: Some(execution_scope),
1728 approved_request: None,
1729 }
1730 },
1731 },
1732 )
1733 .await
1734 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))
1735 }
1736 }
1737
1738 fn register_tool_cancellation(
1739 &mut self,
1740 call_id: &ToolCallId,
1741 cancellation: Option<TurnCancellation>,
1742 ) {
1743 if let Some(cancellation) = cancellation {
1744 self.tool_cancellations
1745 .insert(call_id.clone(), cancellation);
1746 }
1747 }
1748
1749 fn tool_cancellation_for(
1750 &mut self,
1751 call_id: &ToolCallId,
1752 fallback: Option<TurnCancellation>,
1753 ) -> Option<TurnCancellation> {
1754 self.tool_cancellations.get(call_id).cloned().or(fallback)
1755 }
1756
1757 fn clear_tool_cancellation(&mut self, call_id: &ToolCallId) {
1758 self.tool_cancellations.remove(call_id);
1759 }
1760
1761 fn has_pending_interrupts(&self) -> bool {
1762 !self.pending_approvals.is_empty()
1763 }
1764
1765 fn emit_tool_catalog_events(&mut self, events: Vec<ToolCatalogEvent>) {
1766 for event in events {
1767 self.emit(AgentEvent::ToolCatalogChanged(event));
1768 }
1769 }
1770
1771 fn enqueue_pending_approval(
1772 &mut self,
1773 turn_id: &agentkit_core::TurnId,
1774 task: TaskApproval,
1775 cancellation: Option<TurnCancellation>,
1776 ) {
1777 let call_id = task.tool_request.call_id.clone();
1778 self.background_call_ids.remove(&call_id);
1779 let cancellation = self.tool_cancellation_for(&call_id, cancellation);
1780 let call = ToolCallPart {
1781 id: call_id.clone(),
1782 name: task.tool_request.tool_name.to_string(),
1783 input: task.tool_request.input.clone(),
1784 metadata: task.tool_request.metadata.clone(),
1785 };
1786 let mut request = task.approval;
1787 request.call_id = Some(call_id.clone());
1788 let pending = PendingApprovalToolCall {
1789 request: request.clone(),
1790 decision: None,
1791 surfaced: false,
1792 turn_id: turn_id.clone(),
1793 task_id: task.task_id,
1794 call,
1795 tool_request: task.tool_request,
1796 cancellation,
1797 };
1798 self.pending_approvals.insert(call_id.clone(), pending);
1799 if !self.pending_approval_order.iter().any(|id| id == &call_id) {
1800 self.pending_approval_order.push_back(call_id);
1801 }
1802 self.emit(AgentEvent::ApprovalRequired(request));
1803 }
1804
1805 fn take_next_unsurfaced_approval_interrupt(&mut self) -> Option<LoopStep> {
1806 for call_id in self.pending_approval_order.clone() {
1807 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
1808 continue;
1809 };
1810 if pending.decision.is_none() && !pending.surfaced {
1811 pending.surfaced = true;
1812 return Some(LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(
1813 PendingApproval {
1814 request: pending.request.clone(),
1815 },
1816 )));
1817 }
1818 }
1819 None
1820 }
1821
1822 fn next_unresolved_approval_interrupt(&self) -> Option<LoopStep> {
1823 self.pending_approval_order.iter().find_map(|call_id| {
1824 self.pending_approvals.get(call_id).and_then(|pending| {
1825 pending.decision.is_none().then(|| {
1826 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(PendingApproval {
1827 request: pending.request.clone(),
1828 }))
1829 })
1830 })
1831 })
1832 }
1833
1834 fn take_next_resolved_approval(&mut self) -> Option<PendingApprovalToolCall> {
1835 let call_id = self.pending_approval_order.iter().find_map(|call_id| {
1836 self.pending_approvals
1837 .get(call_id)
1838 .and_then(|pending| pending.decision.as_ref().map(|_| call_id.clone()))
1839 })?;
1840 self.pending_approval_order.retain(|id| id != &call_id);
1841 self.pending_approvals.remove(&call_id)
1842 }
1843
1844 fn queue_resolution_interrupt(
1845 &mut self,
1846 turn_id: &agentkit_core::TurnId,
1847 resolution: TaskResolution,
1848 cancellation: Option<TurnCancellation>,
1849 ) -> Option<LoopStep> {
1850 match resolution {
1851 TaskResolution::Item(item) => {
1852 self.append_tool_result_item(item);
1853 None
1854 }
1855 TaskResolution::Approval(task) => {
1856 self.enqueue_pending_approval(turn_id, task, cancellation);
1857 self.take_next_unsurfaced_approval_interrupt()
1858 }
1859 }
1860 }
1861
1862 async fn drain_pending_loop_updates(&mut self) -> Result<(bool, Option<LoopStep>), LoopError> {
1863 let PendingLoopUpdates { mut resolutions } = self
1864 .task_manager
1865 .take_pending_loop_updates()
1866 .await
1867 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
1868 let mut saw_items = false;
1869 while let Some(resolution) = resolutions.pop_front() {
1870 match resolution {
1871 TaskResolution::Item(item) => {
1872 self.append_tool_result_item(item);
1873 saw_items = true;
1874 }
1875 TaskResolution::Approval(task) => {
1876 self.enqueue_pending_approval(&task.tool_request.turn_id.clone(), task, None);
1877 }
1878 }
1879 }
1880 if let Some(step) = self.finish_cancelled_pending_approval().await? {
1881 return Ok((saw_items, Some(step)));
1882 }
1883 Ok((saw_items, self.take_next_unsurfaced_approval_interrupt()))
1884 }
1885
1886 async fn finish_cancelled_pending_approval(&mut self) -> Result<Option<LoopStep>, LoopError> {
1887 if self.pending_approvals.is_empty() {
1888 return Ok(None);
1889 }
1890 if !self.pending_approvals.values().any(|pending| {
1891 pending
1892 .cancellation
1893 .as_ref()
1894 .is_some_and(TurnCancellation::is_cancelled)
1895 }) {
1896 return Ok(None);
1897 }
1898 self.cancel_pending_approvals().await
1899 }
1900
1901 async fn run_mutators(
1902 &mut self,
1903 point: MutationPoint,
1904 turn_id: Option<&agentkit_core::TurnId>,
1905 cancellation: Option<TurnCancellation>,
1906 ) -> Result<(), LoopError> {
1907 if self.mutators.is_empty() {
1908 return Ok(());
1909 }
1910 if cancellation
1911 .as_ref()
1912 .is_some_and(TurnCancellation::is_cancelled)
1913 {
1914 return Err(LoopError::Cancelled);
1915 }
1916 let mutators = self.mutators.clone();
1917 let session_id = self.session_id.clone();
1918 let observed_session_id = Arc::clone(&self.observed_session_id);
1919 let observers = self.observers.clone();
1920 let emitter = DriverEmitter {
1921 session_id: &observed_session_id,
1922 observers: &observers,
1923 };
1924 let mut cursor = TranscriptCursor {
1925 items: &mut self.transcript,
1926 dirty: false,
1927 };
1928 for mutator in &mutators {
1929 if cancellation
1930 .as_ref()
1931 .is_some_and(TurnCancellation::is_cancelled)
1932 {
1933 return Err(LoopError::Cancelled);
1934 }
1935 let ctx = LoopCtx {
1936 session_id: &session_id,
1937 turn_id,
1938 point,
1939 cancellation: cancellation.clone(),
1940 emitter: &emitter,
1941 };
1942 mutator.mutate(&mut cursor, ctx).await?;
1943 }
1944 if cursor.dirty {
1945 validate_transcript_invariants(cursor.items)?;
1946 }
1947 Ok(())
1948 }
1949
1950 async fn continue_active_tool_round(&mut self) -> Result<Option<LoopStep>, LoopError> {
1951 let Some(_) = self.active_tool_round.as_ref() else {
1952 return Ok(None);
1953 };
1954 loop {
1955 let turn_id = self
1956 .active_tool_round
1957 .as_ref()
1958 .map(|active| active.turn_id.clone())
1959 .ok_or_else(|| LoopError::InvalidState("missing active tool round".into()))?;
1960 let cancellation = self
1961 .active_tool_round
1962 .as_ref()
1963 .and_then(|active| active.cancellation.clone());
1964
1965 if cancellation
1966 .as_ref()
1967 .is_some_and(TurnCancellation::is_cancelled)
1968 {
1969 self.task_manager
1970 .on_turn_interrupted(&turn_id)
1971 .await
1972 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
1973 self.active_tool_round = None;
1974 return self.finish_cancelled(turn_id, Vec::new()).map(Some);
1975 }
1976
1977 let next_call = self
1978 .active_tool_round
1979 .as_mut()
1980 .and_then(|active| active.pending_calls.pop_front());
1981 if let Some((call, tool_request)) = next_call {
1982 use tracing::Instrument;
1983 self.register_tool_cancellation(&call.id, cancellation.clone());
1984 let dispatch_span = self.execute_tool_span(&tool_request, &turn_id, "plain");
1985 match self
1986 .start_task_via_manager(
1987 None,
1988 tool_request.clone(),
1989 TaskLaunchKind::Plain,
1990 cancellation.clone(),
1991 )
1992 .instrument(dispatch_span.clone())
1993 .await?
1994 {
1995 TaskStartOutcome::Ready(resolution) => {
1996 let resolution = *resolution;
1997 match resolution {
1998 TaskResolution::Item(item) => {
1999 if !tool_result_not_started(&item) {
2000 self.emit(AgentEvent::ToolExecutionStarted(call.clone()));
2001 }
2002 if tool_result_is_error(&item) {
2003 dispatch_span.record("error.type", "tool_error");
2004 }
2005 if let Some(active) = self.active_tool_round.as_mut() {
2006 active.foreground_progressed = true;
2007 }
2008 self.append_tool_result_item(item);
2009 }
2010 TaskResolution::Approval(task) => {
2011 self.enqueue_pending_approval(&turn_id, task, cancellation.clone());
2012 }
2013 }
2014 continue;
2015 }
2016 TaskStartOutcome::Pending { kind, .. } => {
2017 self.emit(AgentEvent::ToolExecutionStarted(call.clone()));
2018 if kind == agentkit_task_manager::TaskKind::Background {
2019 self.background_call_ids.insert(call.id.clone());
2020 if let Some(active) = self.active_tool_round.as_mut() {
2021 active.background_pending = true;
2022 }
2023 }
2024 continue;
2025 }
2026 }
2027 }
2028
2029 match self
2030 .task_manager
2031 .wait_for_turn(&turn_id, cancellation.clone())
2032 .await
2033 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?
2034 {
2035 Some(TurnTaskUpdate::Resolution(resolution)) => {
2036 let resolution = *resolution;
2037 match resolution {
2038 TaskResolution::Item(item) => {
2039 if let Some(active) = self.active_tool_round.as_mut() {
2040 active.foreground_progressed = true;
2041 }
2042 self.append_tool_result_item(item);
2043 }
2044 TaskResolution::Approval(task) => {
2045 self.enqueue_pending_approval(&turn_id, task, cancellation.clone());
2046 }
2047 }
2048 }
2049 Some(TurnTaskUpdate::Detached(snapshot)) => {
2050 let detached_call_id = snapshot.call_id.clone();
2068 self.background_call_ids.insert(detached_call_id.clone());
2069 let detached_result = ToolResultPart {
2070 call_id: detached_call_id.clone(),
2071 output: ToolOutput::Text(format!(
2072 "Tool {} is now running in the background. \
2073 The result will be delivered when it completes.",
2074 snapshot.tool_name,
2075 )),
2076 is_error: false,
2077 metadata: MetadataMap::new(),
2078 };
2079 self.emit(AgentEvent::ToolExecutionProgress(detached_result.clone()));
2080 self.append_item(Item {
2081 id: None,
2082 kind: ItemKind::Tool,
2083 parts: vec![Part::ToolResult(detached_result)],
2084 metadata: MetadataMap::new(),
2085 usage: None,
2086 finish_reason: None,
2087 created_at: None,
2088 });
2089 self.detached_call_ids.insert(detached_call_id);
2090 if let Some(active) = self.active_tool_round.as_mut() {
2091 active.background_pending = true;
2092 active.foreground_progressed = true;
2093 }
2094 }
2095 None => {
2096 if cancellation
2097 .as_ref()
2098 .is_some_and(TurnCancellation::is_cancelled)
2099 {
2100 self.task_manager
2101 .on_turn_interrupted(&turn_id)
2102 .await
2103 .map_err(|error| {
2104 LoopError::Tool(ToolError::Internal(error.to_string()))
2105 })?;
2106 self.active_tool_round = None;
2107 return self.finish_cancelled(turn_id, Vec::new()).map(Some);
2108 }
2109 let active = self.active_tool_round.take().ok_or_else(|| {
2110 LoopError::InvalidState("missing active tool round".into())
2111 })?;
2112 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2113 return Ok(Some(step));
2114 }
2115 if let Some(step) = self.next_unresolved_approval_interrupt() {
2116 return Ok(Some(step));
2117 }
2118 if active.background_pending && !active.foreground_progressed {
2119 return Ok(None);
2120 }
2121 let info = ToolRoundInfo {
2128 session_id: self.session_id.clone(),
2129 turn_id: turn_id.clone(),
2130 transcript_len: self.transcript.len(),
2131 };
2132 self.pending_round_resume = Some(turn_id);
2133 return Ok(Some(LoopStep::Interrupt(LoopInterrupt::AfterToolResult(
2134 info,
2135 ))));
2136 }
2137 }
2138 }
2139 }
2140
2141 #[tracing::instrument(
2142 name = "agent.turn",
2143 skip_all,
2144 fields(
2145 otel.name = "invoke_agent",
2146 gen_ai.operation.name = "invoke_agent",
2147 gen_ai.conversation.id = %self.session_id,
2148 gen_ai.provider.name = tracing::field::Empty,
2149 session.id = %self.session_id,
2150 turn.id = %turn_id,
2151 transcript.len = self.transcript.len(),
2152 saw_tool_call = tracing::field::Empty,
2153 finish_reason = tracing::field::Empty,
2154 ),
2155 )]
2156 async fn drive_turn(
2157 &mut self,
2158 turn_id: agentkit_core::TurnId,
2159 emit_started: bool,
2160 mutation_point: MutationPoint,
2161 ) -> Result<LoopStep, LoopError> {
2162 let cancellation = self
2163 .cancellation
2164 .as_ref()
2165 .map(CancellationHandle::checkpoint);
2166 match self
2167 .run_mutators(mutation_point, Some(&turn_id), cancellation.clone())
2168 .await
2169 {
2170 Ok(()) => {}
2171 Err(LoopError::Cancelled) => {
2172 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2173 }
2174 Err(error) => return Err(error),
2175 }
2176
2177 if !transcript_has_pending_input(&self.transcript) {
2183 let turn_result = TurnResult {
2184 turn_id,
2185 finish_reason: FinishReason::Completed,
2186 items: Vec::new(),
2187 usage: None,
2188 metadata: MetadataMap::new(),
2189 };
2190 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
2191 return Ok(LoopStep::Finished(turn_result));
2192 }
2193
2194 if emit_started {
2195 self.emit(AgentEvent::TurnStarted {
2196 session_id: self.session_id.clone(),
2197 turn_id: turn_id.clone(),
2198 });
2199 }
2200 if cancellation
2201 .as_ref()
2202 .is_some_and(TurnCancellation::is_cancelled)
2203 {
2204 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2205 }
2206
2207 let catalog_events = self.tool_executor.drain_catalog_events();
2208 self.emit_tool_catalog_events(catalog_events);
2209
2210 let request = TurnRequest {
2211 session_id: self.session_id.clone(),
2212 turn_id: turn_id.clone(),
2213 transcript: self.transcript.clone(),
2214 available_tools: self.tool_executor.specs(),
2215 cache: self
2216 .next_turn_cache
2217 .take()
2218 .or_else(|| self.default_cache.clone()),
2219 metadata: MetadataMap::new(),
2220 };
2221
2222 let session = self
2223 .session
2224 .as_mut()
2225 .ok_or_else(|| LoopError::InvalidState("model session is not available".into()))?;
2226
2227 let chat_span = tracing::info_span!(
2234 "chat",
2235 "otel.name" = tracing::field::Empty,
2236 "otel.kind" = "client",
2237 "gen_ai.operation.name" = "chat",
2238 "gen_ai.provider.name" = tracing::field::Empty,
2239 "gen_ai.conversation.id" = %self.session_id,
2240 "gen_ai.request.model" = tracing::field::Empty,
2241 "gen_ai.response.model" = tracing::field::Empty,
2242 "gen_ai.response.id" = tracing::field::Empty,
2243 "gen_ai.response.finish_reasons" = tracing::field::Empty,
2244 "gen_ai.input.messages" = tracing::field::Empty,
2245 "gen_ai.output.messages" = tracing::field::Empty,
2246 "gen_ai.usage.input_tokens" = tracing::field::Empty,
2247 "gen_ai.usage.output_tokens" = tracing::field::Empty,
2248 "gen_ai.usage.cost" = tracing::field::Empty,
2249 );
2250 if let Some(capture) = self.telemetry.input_messages {
2251 record_string_array_attribute(
2252 &chat_span,
2253 "gen_ai.input.messages",
2254 capture_messages(&request.transcript, capture, CaptureOrder::NewestTail),
2255 );
2256 }
2257
2258 let initial_provider_name =
2262 effective_provider_name(session.provider_name(), self.provider_name.as_deref());
2263 if let Some(provider) = &initial_provider_name {
2264 chat_span.record("gen_ai.provider.name", provider.as_str());
2265 tracing::Span::current().record("gen_ai.provider.name", provider.as_str());
2266 }
2267 match session.model_name() {
2268 Some(model) => {
2269 chat_span.record("gen_ai.request.model", model);
2270 chat_span.record("otel.name", format!("chat {model}").as_str());
2271 }
2272 None => {
2273 chat_span.record("otel.name", "chat");
2274 }
2275 }
2276
2277 use tracing::Instrument;
2278 let mut turn = match session
2279 .begin_turn(request, cancellation.clone())
2280 .instrument(chat_span.clone())
2281 .await
2282 {
2283 Ok(turn) => turn,
2284 Err(LoopError::Cancelled) => {
2285 self.task_manager
2286 .on_turn_interrupted(&turn_id)
2287 .await
2288 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2289 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2290 }
2291 Err(error) => return Err(error),
2292 };
2293
2294 let provider_name =
2297 effective_provider_name(session.provider_name(), self.provider_name.as_deref());
2298 if let Some(provider) = &provider_name {
2299 chat_span.record("gen_ai.provider.name", provider.as_str());
2300 tracing::Span::current().record("gen_ai.provider.name", provider.as_str());
2301 }
2302 match session.model_name() {
2303 Some(model) => {
2304 chat_span.record("gen_ai.request.model", model);
2305 chat_span.record("otel.name", format!("chat {model}").as_str());
2306 }
2307 None => {
2308 chat_span.record("otel.name", "chat");
2309 }
2310 }
2311
2312 let mut saw_tool_call = false;
2313 let mut finished_result = None;
2314 let mut latest_usage = None;
2315
2316 while let Some(event) = match turn
2317 .next_event(cancellation.clone())
2318 .instrument(chat_span.clone())
2319 .await
2320 {
2321 Ok(event) => event,
2322 Err(LoopError::Cancelled) => {
2323 self.task_manager
2324 .on_turn_interrupted(&turn_id)
2325 .await
2326 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2327 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2328 }
2329 Err(error) => return Err(error),
2330 } {
2331 if cancellation
2332 .as_ref()
2333 .is_some_and(TurnCancellation::is_cancelled)
2334 {
2335 self.task_manager
2336 .on_turn_interrupted(&turn_id)
2337 .await
2338 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2339 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2340 }
2341 match event {
2342 ModelTurnEvent::Delta(delta) => self.emit(AgentEvent::ContentDelta(delta)),
2343 ModelTurnEvent::Usage(usage) => {
2344 latest_usage = Some(usage.clone());
2345 self.emit(AgentEvent::UsageUpdated(usage));
2346 }
2347 ModelTurnEvent::ToolCall(call) => {
2348 saw_tool_call = true;
2349 self.emit(AgentEvent::ToolCallRequested(call.clone()));
2350 }
2351 ModelTurnEvent::Finished(result) => {
2352 finished_result = Some(result);
2353 break;
2354 }
2355 }
2356 }
2357
2358 let mut result = finished_result.ok_or_else(|| {
2359 LoopError::Provider("model turn ended without a Finished event".into())
2360 })?;
2361 result.usage = merge_usage(result.usage, latest_usage);
2362 if let Some(model) = &result.model {
2363 chat_span.record("gen_ai.response.model", model.as_str());
2364 }
2365 if let Some(id) = &result.response_id {
2366 chat_span.record("gen_ai.response.id", id.as_str());
2367 }
2368 if let Some(tokens) = result
2369 .usage
2370 .as_ref()
2371 .and_then(|usage| usage.tokens.as_ref())
2372 {
2373 record_token_attribute(&chat_span, "gen_ai.usage.input_tokens", tokens.input_tokens);
2374 record_token_attribute(
2375 &chat_span,
2376 "gen_ai.usage.output_tokens",
2377 tokens.output_tokens,
2378 );
2379 }
2380 if let Some(cost) = result.usage.as_ref().and_then(|usage| usage.cost.as_ref()) {
2381 record_f64_attribute(&chat_span, "gen_ai.usage.cost", cost.amount);
2382 }
2383 record_string_array_attribute(
2384 &chat_span,
2385 "gen_ai.response.finish_reasons",
2386 provider_finish_reasons(&result.metadata, &result.finish_reason),
2387 );
2388 if let Some(capture) = self.telemetry.output_messages {
2389 record_string_array_attribute(
2390 &chat_span,
2391 "gen_ai.output.messages",
2392 capture_messages(&result.output_items, capture, CaptureOrder::OldestHead),
2393 );
2394 }
2395 drop(chat_span);
2396 tracing::Span::current().record("saw_tool_call", saw_tool_call);
2397 tracing::Span::current().record(
2398 "finish_reason",
2399 tracing::field::debug(&result.finish_reason),
2400 );
2401 let now = Timestamp::now();
2402 let usage = result.usage.clone();
2403 let finish_reason = result.finish_reason.clone();
2404 let output_items: Vec<Item> = result
2405 .output_items
2406 .drain(..)
2407 .map(|mut item| {
2408 if matches!(item.kind, ItemKind::Assistant) {
2409 if item.usage.is_none() {
2410 item.usage = usage.clone();
2411 }
2412 if item.finish_reason.is_none() {
2413 item.finish_reason = Some(finish_reason.clone());
2414 }
2415 }
2416 if item.created_at.is_none() {
2417 item.created_at = Some(now);
2418 }
2419 item
2420 })
2421 .collect();
2422 self.extend_transcript(output_items.clone());
2423
2424 if saw_tool_call {
2425 let pending_calls = extract_tool_calls(&output_items)
2426 .into_iter()
2427 .map(|call| {
2428 let tool_request = ToolRequest {
2429 call_id: call.id.clone(),
2430 tool_name: agentkit_tools_core::ToolName::new(call.name.clone()),
2431 input: call.input.clone(),
2432 session_id: self.session_id.clone(),
2433 turn_id: turn_id.clone(),
2434 metadata: call.metadata.clone(),
2435 };
2436 (call, tool_request)
2437 })
2438 .collect();
2439 self.active_tool_round = Some(ActiveToolRound {
2440 turn_id: turn_id.clone(),
2441 pending_calls,
2442 cancellation: cancellation.clone(),
2443 background_pending: false,
2444 foreground_progressed: false,
2445 });
2446 if let Some(step) = self.continue_active_tool_round().await? {
2447 return Ok(step);
2448 }
2449 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
2450 InputRequest {
2451 session_id: self.session_id.clone(),
2452 reason: "driver is waiting for input".into(),
2453 },
2454 )));
2455 }
2456
2457 let turn_result = TurnResult {
2458 turn_id,
2459 finish_reason: result.finish_reason,
2460 items: output_items,
2461 usage: result.usage,
2462 metadata: result.metadata,
2463 };
2464 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
2465 Ok(LoopStep::Finished(turn_result))
2466 }
2467
2468 async fn resume_after_approval(
2469 &mut self,
2470 pending: PendingApprovalToolCall,
2471 ) -> Result<LoopStep, LoopError> {
2472 let decision = pending
2473 .decision
2474 .clone()
2475 .ok_or_else(|| LoopError::InvalidState("pending approval has no decision".into()))?;
2476
2477 match decision {
2478 ApprovalDecision::Approve => {
2479 use tracing::Instrument;
2480 self.emit(AgentEvent::ToolExecutionStarted(pending.call.clone()));
2481 let dispatch_span =
2482 self.execute_tool_span(&pending.tool_request, &pending.turn_id, "approved");
2483 let cancellation = self
2484 .cancellation
2485 .as_ref()
2486 .map(CancellationHandle::checkpoint);
2487 self.register_tool_cancellation(&pending.call.id, cancellation.clone());
2488 match self
2489 .start_task_via_manager(
2490 Some(pending.task_id.clone()),
2491 pending.tool_request.clone(),
2492 TaskLaunchKind::Approved(pending.request.clone()),
2493 cancellation.clone(),
2494 )
2495 .instrument(dispatch_span.clone())
2496 .await?
2497 {
2498 TaskStartOutcome::Ready(resolution) => {
2499 let resolution = *resolution;
2500 if let TaskResolution::Item(item) = &resolution
2501 && tool_result_is_error(item)
2502 {
2503 dispatch_span.record("error.type", "tool_error");
2504 }
2505 if let Some(step) = self.queue_resolution_interrupt(
2506 &pending.turn_id,
2507 resolution,
2508 cancellation,
2509 ) {
2510 return Ok(step);
2511 }
2512 }
2513 TaskStartOutcome::Pending { kind, .. } => {
2514 if kind == agentkit_task_manager::TaskKind::Background {
2515 self.background_call_ids.insert(pending.call.id.clone());
2516 }
2517 }
2518 }
2519 }
2520 ApprovalDecision::Deny { reason } => {
2521 self.append_tool_result_item(Item {
2522 id: None,
2523 kind: ItemKind::Tool,
2524 parts: vec![Part::ToolResult(ToolResultPart {
2525 call_id: pending.call.id.clone(),
2526 output: ToolOutput::Text(
2527 reason.unwrap_or_else(|| "approval denied".into()),
2528 ),
2529 is_error: true,
2530 metadata: pending.call.metadata.clone(),
2531 })],
2532 metadata: MetadataMap::new(),
2533 usage: None,
2534 finish_reason: None,
2535 created_at: None,
2536 });
2537 }
2538 }
2539
2540 if let Some(step) = self.continue_active_tool_round().await? {
2541 Ok(step)
2542 } else if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2543 Ok(step)
2544 } else if let Some(step) = self.next_unresolved_approval_interrupt() {
2545 Ok(step)
2546 } else {
2547 self.drive_turn(pending.turn_id, false, MutationPoint::AfterToolResult)
2548 .await
2549 }
2550 }
2551
2552 fn finish_cancelled(
2553 &mut self,
2554 turn_id: agentkit_core::TurnId,
2555 items: Vec<Item>,
2556 ) -> Result<LoopStep, LoopError> {
2557 let pending = self.drain_pending_approval_items();
2558 self.reject_drained_approvals(pending);
2559 self.close_interrupted_tool_calls();
2560 self.extend_transcript(items.clone());
2561 let turn_result = TurnResult {
2562 turn_id,
2563 finish_reason: FinishReason::Cancelled,
2564 items,
2565 usage: None,
2566 metadata: interrupted_metadata("turn"),
2567 };
2568 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
2569 Ok(LoopStep::Finished(turn_result))
2570 }
2571
2572 pub fn submit_input(&mut self, input: Vec<Item>) -> Result<(), LoopError> {
2582 if self.has_pending_interrupts() {
2583 return Err(LoopError::InvalidState(
2584 "cannot submit input while an interrupt is pending".into(),
2585 ));
2586 }
2587 self.emit(AgentEvent::InputAccepted {
2588 session_id: self.session_id.clone(),
2589 items: input.clone(),
2590 });
2591 self.pending_input.extend(input);
2592 Ok(())
2593 }
2594
2595 pub fn set_next_turn_cache(&mut self, cache: PromptCacheRequest) -> Result<(), LoopError> {
2600 if self.has_pending_interrupts() {
2601 return Err(LoopError::InvalidState(
2602 "cannot update next-turn cache while an interrupt is pending".into(),
2603 ));
2604 }
2605 self.next_turn_cache = Some(cache);
2606 Ok(())
2607 }
2608
2609 #[cfg(test)]
2610 pub(crate) fn submit_input_with_cache(
2611 &mut self,
2612 input: Vec<Item>,
2613 cache: PromptCacheRequest,
2614 ) -> Result<(), LoopError> {
2615 self.set_next_turn_cache(cache)?;
2616 self.submit_input(input)
2617 }
2618
2619 pub fn resolve_approval_for(
2629 &mut self,
2630 call_id: ToolCallId,
2631 decision: ApprovalDecision,
2632 ) -> Result<(), LoopError> {
2633 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
2634 return Err(LoopError::InvalidState(format!(
2635 "no approval request is pending for call {}",
2636 call_id.0
2637 )));
2638 };
2639 pending.decision = Some(decision.clone());
2640 self.emit(AgentEvent::ApprovalResolved {
2641 approved: matches!(decision, ApprovalDecision::Approve),
2642 });
2643 Ok(())
2644 }
2645
2646 pub fn resolve_approval_for_with_patched_input(
2659 &mut self,
2660 call_id: ToolCallId,
2661 input: serde_json::Value,
2662 ) -> Result<(), LoopError> {
2663 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
2664 return Err(LoopError::InvalidState(format!(
2665 "no approval request is pending for call {}",
2666 call_id.0
2667 )));
2668 };
2669 pending.tool_request.input = input;
2670 self.resolve_approval_for(call_id, ApprovalDecision::Approve)
2671 }
2672
2673 pub fn resolve_approval(&mut self, decision: ApprovalDecision) -> Result<(), LoopError> {
2676 let mut unresolved = self
2677 .pending_approval_order
2678 .iter()
2679 .filter(|call_id| {
2680 self.pending_approvals
2681 .get(*call_id)
2682 .is_some_and(|pending| pending.decision.is_none())
2683 })
2684 .cloned();
2685 let Some(call_id) = unresolved.next() else {
2686 return Err(LoopError::InvalidState(
2687 "no approval request is pending".into(),
2688 ));
2689 };
2690 if unresolved.next().is_some() {
2691 return Err(LoopError::InvalidState(
2692 "multiple approvals are pending; use resolve_approval_for".into(),
2693 ));
2694 }
2695 self.resolve_approval_for(call_id, decision)
2696 }
2697
2698 pub fn cancel_pending_approval_for(&mut self, call_id: ToolCallId) -> Result<(), LoopError> {
2703 let Some(pending) = self.drain_pending_approval_for(&call_id) else {
2704 return Err(LoopError::InvalidState(format!(
2705 "no approval request is pending for call {}",
2706 call_id.0
2707 )));
2708 };
2709 self.reject_drained_approvals(vec![pending]);
2710 Ok(())
2711 }
2712
2713 pub async fn cancel_pending_approvals(&mut self) -> Result<Option<LoopStep>, LoopError> {
2719 if self.pending_approvals.is_empty() {
2720 return Ok(None);
2721 }
2722 let Some(turn_id) = self
2723 .pending_approval_order
2724 .iter()
2725 .find_map(|call_id| self.pending_approvals.get(call_id))
2726 .map(|pending| pending.turn_id.clone())
2727 else {
2728 return Ok(None);
2729 };
2730 let pending = self.drain_pending_approval_items();
2731 self.active_tool_round = None;
2732 self.task_manager
2733 .on_turn_interrupted(&turn_id)
2734 .await
2735 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2736 self.reject_drained_approvals(pending);
2737 self.finish_cancelled(turn_id, Vec::new()).map(Some)
2738 }
2739
2740 pub fn snapshot(&self) -> LoopSnapshot {
2742 LoopSnapshot {
2743 session_id: self.session_id.clone(),
2744 transcript: self.transcript.clone(),
2745 pending_input: self.pending_input.clone(),
2746 }
2747 }
2748
2749 pub async fn next(&mut self) -> Result<LoopStep, LoopError> {
2770 if let Some(pending) = self.take_next_resolved_approval() {
2771 return self.resume_after_approval(pending).await;
2772 }
2773
2774 if let Some(step) = self.finish_cancelled_pending_approval().await? {
2775 return Ok(step);
2776 }
2777
2778 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2779 return Ok(step);
2780 }
2781
2782 if let Some(step) = self.next_unresolved_approval_interrupt() {
2783 return Ok(step);
2784 }
2785
2786 if let Some(step) = self.continue_active_tool_round().await? {
2787 return Ok(step);
2788 }
2789
2790 let (had_loop_updates, loop_step) = self.drain_pending_loop_updates().await?;
2791 if let Some(step) = loop_step {
2792 return Ok(step);
2793 }
2794
2795 if let Some(turn_id) = self.pending_round_resume.take() {
2800 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
2801 self.extend_transcript(drained);
2802 return self
2803 .drive_turn(turn_id, false, MutationPoint::AfterToolResult)
2804 .await;
2805 }
2806
2807 if self.pending_input.is_empty() && !had_loop_updates {
2808 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
2809 InputRequest {
2810 session_id: self.session_id.clone(),
2811 reason: "driver is waiting for input".into(),
2812 },
2813 )));
2814 }
2815
2816 let turn_id = agentkit_core::TurnId::new(format!("turn-{}", self.next_turn_index));
2817 self.next_turn_index += 1;
2818 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
2819 self.extend_transcript(drained);
2820 self.drive_turn(turn_id, true, MutationPoint::AfterTurnEnded)
2821 .await
2822 }
2823
2824 fn emit(&self, event: AgentEvent) {
2825 fan_out_observed_event(&self.observers, &self.observed_session_id, event);
2826 }
2827
2828 fn append_item(&mut self, mut item: Item) {
2833 if item.created_at.is_none() {
2834 item.created_at = Some(Timestamp::now());
2835 }
2836 for observer in &self.transcript_observers {
2837 observer.on_transcript_event(TranscriptEvent {
2838 session_id: &self.session_id,
2839 item: &item,
2840 });
2841 }
2842 self.transcript.push(item);
2843 }
2844
2845 fn append_tool_result_item(&mut self, item: Item) {
2859 for part in &item.parts {
2860 if let Part::ToolResult(result) = part {
2861 if !self
2862 .interrupted_background_call_ids
2863 .contains(&result.call_id)
2864 {
2865 self.emit(AgentEvent::ToolResultReceived(result.clone()));
2866 }
2867 self.background_call_ids.remove(&result.call_id);
2868 self.clear_tool_cancellation(&result.call_id);
2869 }
2870 }
2871 let item = self.maybe_convert_detached(item);
2872 self.append_item(item);
2873 }
2874
2875 fn drain_pending_approval_for(
2876 &mut self,
2877 call_id: &ToolCallId,
2878 ) -> Option<PendingApprovalToolCall> {
2879 let pending = self.pending_approvals.remove(call_id)?;
2880 self.pending_approval_order.retain(|id| id != call_id);
2881 self.clear_tool_cancellation(call_id);
2882 Some(pending)
2883 }
2884
2885 fn drain_pending_approval_items(&mut self) -> Vec<PendingApprovalToolCall> {
2886 let order = std::mem::take(&mut self.pending_approval_order);
2887 let pending = order
2888 .iter()
2889 .filter_map(|call_id| {
2890 let pending = self.pending_approvals.remove(call_id);
2891 self.clear_tool_cancellation(call_id);
2892 pending
2893 })
2894 .collect();
2895 self.pending_approvals.clear();
2896 pending
2897 }
2898
2899 fn reject_drained_approvals(&mut self, pending: Vec<PendingApprovalToolCall>) {
2900 for pending in pending {
2901 self.emit(AgentEvent::ApprovalResolved { approved: false });
2902 self.append_tool_result_item(cancelled_approval_item(pending));
2903 }
2904 }
2905
2906 fn close_interrupted_tool_calls(&mut self) {
2923 for call in unanswered_tool_calls(&self.transcript) {
2924 let call_id = call.id.clone();
2925 let completes_in_background = self.background_call_ids.contains(&call_id);
2926 self.append_tool_result_item(interrupted_tool_result_item(call));
2927 if completes_in_background {
2928 self.detached_call_ids.insert(call_id.clone());
2929 self.interrupted_background_call_ids.insert(call_id);
2930 }
2931 }
2932 }
2933
2934 fn maybe_convert_detached(&mut self, item: Item) -> Item {
2935 if !matches!(item.kind, ItemKind::Tool) {
2936 return item;
2937 }
2938 let results: Vec<&ToolResultPart> = item
2939 .parts
2940 .iter()
2941 .filter_map(|p| match p {
2942 Part::ToolResult(r) => Some(r),
2943 _ => None,
2944 })
2945 .collect();
2946 if results.is_empty()
2947 || !results
2948 .iter()
2949 .all(|r| self.detached_call_ids.contains(&r.call_id))
2950 {
2951 return item;
2952 }
2953 let mut text = String::new();
2954 for result in &results {
2955 self.detached_call_ids.remove(&result.call_id);
2956 self.interrupted_background_call_ids.remove(&result.call_id);
2957 if !text.is_empty() {
2958 text.push_str("\n\n");
2959 }
2960 let label = if result.is_error {
2961 "failed"
2962 } else {
2963 "completed"
2964 };
2965 let body = render_tool_output_brief(&result.output);
2966 text.push_str(&format!(
2967 "Background tool call {} {}: {body}",
2968 result.call_id.0, label
2969 ));
2970 }
2971 Item::notification(text)
2972 }
2973
2974 fn extend_transcript(&mut self, items: impl IntoIterator<Item = Item>) {
2978 let now = Timestamp::now();
2979 for mut item in items {
2980 if item.created_at.is_none() {
2981 item.created_at = Some(now);
2982 }
2983 self.append_item(item);
2984 }
2985 }
2986}
2987
2988fn render_tool_output_brief(output: &ToolOutput) -> String {
2989 match output {
2990 ToolOutput::Text(t) => t.clone(),
2991 ToolOutput::Structured(value) => value.to_string(),
2992 ToolOutput::Parts(parts) => format!("[{} parts]", parts.len()),
2993 ToolOutput::Files(files) => format!("[{} files]", files.len()),
2994 }
2995}
2996
2997fn interrupted_metadata(stage: &str) -> MetadataMap {
2998 let mut metadata = MetadataMap::new();
2999 metadata.insert(INTERRUPTED_METADATA_KEY.into(), true.into());
3000 metadata.insert(
3001 INTERRUPT_REASON_METADATA_KEY.into(),
3002 USER_CANCELLED_REASON.into(),
3003 );
3004 metadata.insert(INTERRUPT_STAGE_METADATA_KEY.into(), stage.into());
3005 metadata
3006}
3007
3008fn record_token_attribute(span: &tracing::Span, key: &'static str, value: u64) {
3009 match i64::try_from(value) {
3010 Ok(value) => record_i64_attribute(span, key, value),
3011 Err(_) => tracing::warn!(attribute = key, value, "token count exceeds OTEL i64 range"),
3012 }
3013}
3014
3015#[cfg(feature = "otel")]
3016fn record_i64_attribute(span: &tracing::Span, key: &'static str, value: i64) {
3017 use tracing_opentelemetry::OpenTelemetrySpanExt;
3018 span.set_attribute(key, value);
3019}
3020
3021#[cfg(not(feature = "otel"))]
3022fn record_i64_attribute(span: &tracing::Span, key: &'static str, value: i64) {
3023 span.record(key, value);
3024}
3025
3026#[cfg(feature = "otel")]
3027fn record_f64_attribute(span: &tracing::Span, key: &'static str, value: f64) {
3028 use tracing_opentelemetry::OpenTelemetrySpanExt;
3029 span.set_attribute(key, value);
3030}
3031
3032#[cfg(not(feature = "otel"))]
3033fn record_f64_attribute(span: &tracing::Span, key: &'static str, value: f64) {
3034 span.record(key, value);
3035}
3036
3037#[cfg(feature = "otel")]
3038fn otel_string_array(values: Vec<String>) -> opentelemetry::Value {
3039 use opentelemetry::{Array, StringValue, Value as OtelValue};
3040 OtelValue::Array(Array::String(
3041 values.into_iter().map(StringValue::from).collect(),
3042 ))
3043}
3044
3045#[cfg(feature = "otel")]
3046fn record_string_array_attribute(span: &tracing::Span, key: &'static str, values: Vec<String>) {
3047 use tracing_opentelemetry::OpenTelemetrySpanExt;
3048 span.set_attribute(key, otel_string_array(values));
3049}
3050
3051#[cfg(not(feature = "otel"))]
3052fn record_string_array_attribute(span: &tracing::Span, key: &'static str, values: Vec<String>) {
3053 span.record(key, tracing::field::debug(&values));
3054}
3055
3056#[derive(Clone, Copy)]
3057enum CaptureOrder {
3058 NewestTail,
3059 OldestHead,
3060}
3061
3062fn effective_provider_name(
3063 session_provider: Option<&str>,
3064 adapter_provider: Option<&str>,
3065) -> Option<String> {
3066 session_provider.or(adapter_provider).map(str::to_owned)
3067}
3068
3069fn merge_usage(final_usage: Option<Usage>, streamed_usage: Option<Usage>) -> Option<Usage> {
3070 match (final_usage, streamed_usage) {
3071 (None, streamed) => streamed,
3072 (Some(final_usage), None) => Some(final_usage),
3073 (Some(mut final_usage), Some(streamed)) => {
3074 if final_usage.tokens.is_none() {
3075 final_usage.tokens = streamed.tokens;
3076 }
3077 if final_usage.cost.is_none() {
3078 final_usage.cost = streamed.cost;
3079 }
3080 for (key, value) in streamed.metadata {
3081 final_usage.metadata.entry(key).or_insert(value);
3082 }
3083 Some(final_usage)
3084 }
3085 }
3086}
3087
3088fn capture_messages(items: &[Item], capture: MessageCapture, order: CaptureOrder) -> Vec<String> {
3089 let mut captured = Vec::new();
3090 let mut used_bytes = 0usize;
3091 let indices: Box<dyn Iterator<Item = usize>> = match order {
3092 CaptureOrder::NewestTail => Box::new((0..items.len()).rev()),
3093 CaptureOrder::OldestHead => Box::new(0..items.len()),
3094 };
3095
3096 for index in indices.take(capture.max_messages) {
3097 let item = &items[index];
3098 let original_bytes = source_content_bytes(item);
3099 let remaining = capture.max_bytes.saturating_sub(used_bytes);
3100 if original_bytes > remaining {
3101 captured.push(
3102 serde_json::json!({
3103 "type": "truncated",
3104 "original_bytes": original_bytes,
3105 })
3106 .to_string(),
3107 );
3108 break;
3109 }
3110 used_bytes += original_bytes;
3111 captured.push(capture_item_json(item, remaining));
3112 }
3113
3114 if matches!(order, CaptureOrder::NewestTail) {
3115 captured.reverse();
3116 }
3117 captured
3118}
3119
3120fn source_content_bytes(item: &Item) -> usize {
3121 item.parts.iter().fold(0, |total, part| {
3122 total.saturating_add(part_source_content_bytes(part))
3123 })
3124}
3125
3126fn part_source_content_bytes(part: &Part) -> usize {
3127 match part {
3128 Part::Text(text) => text.text.len(),
3129 Part::Media(media) => media.mime_type.len(),
3130 Part::File(file) => file
3131 .name
3132 .as_deref()
3133 .map_or(0, str::len)
3134 .saturating_add(file.mime_type.as_deref().map_or(0, str::len)),
3135 Part::Structured(_) => 0,
3136 Part::Reasoning(reasoning) => reasoning.summary.as_deref().map_or(0, str::len),
3137 Part::ToolCall(call) => call.id.0.len().saturating_add(call.name.len()),
3138 Part::ToolResult(result) => result
3139 .call_id
3140 .0
3141 .len()
3142 .saturating_add(tool_output_source_content_bytes(&result.output)),
3143 Part::Custom(custom) => custom.kind.len(),
3144 }
3145}
3146
3147fn tool_output_source_content_bytes(output: &ToolOutput) -> usize {
3148 match output {
3149 ToolOutput::Text(text) => text.len(),
3150 ToolOutput::Structured(_) | ToolOutput::Parts(_) | ToolOutput::Files(_) => 0,
3151 }
3152}
3153
3154struct CaptureBudget {
3155 remaining: usize,
3156}
3157
3158impl CaptureBudget {
3159 fn text(&mut self, text: &str) -> (String, bool) {
3160 let end = floor_char_boundary(text, self.remaining.min(text.len()));
3161 self.remaining = self.remaining.saturating_sub(end);
3162 (text[..end].to_owned(), end < text.len())
3163 }
3164}
3165
3166fn floor_char_boundary(text: &str, mut index: usize) -> usize {
3167 while index > 0 && !text.is_char_boundary(index) {
3168 index -= 1;
3169 }
3170 index
3171}
3172
3173const MAX_CAPTURED_PARTS_PER_ITEM: usize = 256;
3174
3175fn capture_item_json(item: &Item, max_bytes: usize) -> String {
3176 let mut budget = CaptureBudget {
3177 remaining: max_bytes,
3178 };
3179 let mut parts = item
3180 .parts
3181 .iter()
3182 .take(MAX_CAPTURED_PARTS_PER_ITEM)
3183 .map(|part| sanitized_part(part, &mut budget))
3184 .collect::<Vec<_>>();
3185 if item.parts.len() > parts.len() {
3186 parts.push(serde_json::json!({
3187 "type": "truncated",
3188 "reason": "part_limit",
3189 }));
3190 }
3191 serde_json::json!({
3192 "role": item_kind_name(item.kind),
3193 "parts": parts,
3194 })
3195 .to_string()
3196}
3197
3198fn item_kind_name(kind: ItemKind) -> &'static str {
3199 match kind {
3200 ItemKind::System => "system",
3201 ItemKind::Developer => "developer",
3202 ItemKind::User => "user",
3203 ItemKind::Assistant => "assistant",
3204 ItemKind::Tool => "tool",
3205 ItemKind::Context => "context",
3206 ItemKind::Notification => "notification",
3207 }
3208}
3209
3210fn modality_name(modality: Modality) -> &'static str {
3211 match modality {
3212 Modality::Audio => "audio",
3213 Modality::Image => "image",
3214 Modality::Video => "video",
3215 Modality::Binary => "binary",
3216 }
3217}
3218
3219fn omitted_data_ref(data: &DataRef) -> Value {
3220 let kind = match data {
3221 DataRef::InlineText(_) => "inline_text",
3222 DataRef::InlineBytes(_) => "inline_bytes",
3223 DataRef::Uri(_) => "uri",
3224 DataRef::Handle(_) => "handle",
3225 };
3226 serde_json::json!({ "kind": kind, "omitted": true })
3227}
3228
3229fn bounded_field(text: &str, budget: &mut CaptureBudget) -> Value {
3230 let (text, truncated) = budget.text(text);
3231 serde_json::json!({ "value": text, "truncated": truncated })
3232}
3233
3234fn sanitized_part(part: &Part, budget: &mut CaptureBudget) -> Value {
3235 match part {
3236 Part::Text(text) => serde_json::json!({
3237 "type": "text",
3238 "text": bounded_field(&text.text, budget),
3239 }),
3240 Part::Media(media) => serde_json::json!({
3241 "type": "media",
3242 "modality": modality_name(media.modality),
3243 "mime_type": bounded_field(&media.mime_type, budget),
3244 "data": omitted_data_ref(&media.data),
3245 }),
3246 Part::File(file) => serde_json::json!({
3247 "type": "file",
3248 "name": file.name.as_deref().map(|name| bounded_field(name, budget)),
3249 "mime_type": file.mime_type.as_deref().map(|mime| bounded_field(mime, budget)),
3250 "data": omitted_data_ref(&file.data),
3251 }),
3252 Part::Structured(_) => serde_json::json!({
3253 "type": "structured",
3254 "truncated": true,
3255 }),
3256 Part::Reasoning(reasoning) => serde_json::json!({
3257 "type": "reasoning",
3258 "summary": reasoning.summary.as_deref().map(|summary| bounded_field(summary, budget)),
3259 "redacted": reasoning.redacted,
3260 "data": reasoning.data.as_ref().map(omitted_data_ref),
3261 }),
3262 Part::ToolCall(call) => serde_json::json!({
3263 "type": "tool_call",
3264 "id": bounded_field(&call.id.0, budget),
3265 "name": bounded_field(&call.name, budget),
3266 "input": { "truncated": true },
3267 }),
3268 Part::ToolResult(result) => serde_json::json!({
3269 "type": "tool_result",
3270 "call_id": bounded_field(&result.call_id.0, budget),
3271 "is_error": result.is_error,
3272 "output": sanitized_tool_output(&result.output, budget),
3273 }),
3274 Part::Custom(custom) => serde_json::json!({
3275 "type": "custom",
3276 "kind": bounded_field(&custom.kind, budget),
3277 "data": custom.data.as_ref().map(omitted_data_ref),
3278 "value": custom.value.as_ref().map(|_| serde_json::json!({ "truncated": true })),
3279 }),
3280 }
3281}
3282
3283fn sanitized_tool_output(output: &ToolOutput, budget: &mut CaptureBudget) -> Value {
3284 match output {
3285 ToolOutput::Text(text) => serde_json::json!({
3286 "type": "text",
3287 "text": bounded_field(text, budget),
3288 }),
3289 ToolOutput::Structured(_) => serde_json::json!({
3290 "type": "structured",
3291 "truncated": true,
3292 }),
3293 ToolOutput::Parts(parts) => serde_json::json!({
3294 "type": "parts",
3295 "count": parts.len(),
3296 "truncated": true,
3297 }),
3298 ToolOutput::Files(files) => serde_json::json!({
3299 "type": "files",
3300 "count": files.len(),
3301 "truncated": true,
3302 }),
3303 }
3304}
3305
3306#[cfg(test)]
3307mod telemetry_tests {
3308 use super::*;
3309
3310 #[test]
3311 fn message_capture_is_off_by_default_and_independent() {
3312 let default = TelemetryConfig::default();
3313 assert_eq!(default.input_messages(), None);
3314 assert_eq!(default.output_messages(), None);
3315
3316 let capture = MessageCapture::new(2, 1).unwrap();
3317 let input_only = TelemetryConfig::default().with_input_messages(capture);
3318 assert_eq!(input_only.input_messages().unwrap().max_messages(), 2);
3319 assert_eq!(input_only.input_messages().unwrap().max_bytes(), 1);
3320 assert_eq!(input_only.output_messages(), None);
3321 assert_eq!(
3322 MessageCapture::new(0, 1),
3323 Err(MessageCaptureError::ZeroMessages)
3324 );
3325 assert_eq!(
3326 MessageCapture::new(1, 0),
3327 Err(MessageCaptureError::ZeroBytes)
3328 );
3329 }
3330
3331 #[test]
3332 fn one_source_byte_is_not_rejected_for_json_envelope_overhead() {
3333 let items = vec![Item::text(ItemKind::User, "x")];
3334 let captured = capture_messages(
3335 &items,
3336 MessageCapture::new(1, 1).unwrap(),
3337 CaptureOrder::OldestHead,
3338 );
3339 assert_eq!(captured.len(), 1);
3340 let value: Value = serde_json::from_str(&captured[0]).unwrap();
3341 assert_eq!(value["role"], "user");
3342 assert_eq!(value["parts"][0]["text"]["value"], "x");
3343 assert_eq!(value["parts"][0]["text"]["truncated"], false);
3344 }
3345
3346 #[test]
3347 fn multibyte_source_accounting_preserves_utf8_boundaries() {
3348 let items = vec![Item::text(ItemKind::User, "é")];
3349
3350 let exact = capture_messages(
3351 &items,
3352 MessageCapture::new(1, "é".len()).unwrap(),
3353 CaptureOrder::OldestHead,
3354 );
3355 let value: Value = serde_json::from_str(&exact[0]).unwrap();
3356 assert_eq!(value["parts"][0]["text"]["value"], "é");
3357 assert_eq!(value["parts"][0]["text"]["truncated"], false);
3358
3359 let too_small = capture_messages(
3360 &items,
3361 MessageCapture::new(1, 1).unwrap(),
3362 CaptureOrder::OldestHead,
3363 );
3364 let value: Value = serde_json::from_str(&too_small[0]).unwrap();
3365 assert_eq!(value["type"], "truncated");
3366 assert_eq!(value["original_bytes"], 2);
3367 }
3368
3369 #[test]
3370 fn source_bytes_are_aggregated_without_charging_json_envelopes() {
3371 let items = vec![
3372 Item::text(ItemKind::User, "a"),
3373 Item::text(ItemKind::Assistant, "b"),
3374 Item::text(ItemKind::User, "cd"),
3375 ];
3376 let captured = capture_messages(
3377 &items,
3378 MessageCapture::new(3, 3).unwrap(),
3379 CaptureOrder::OldestHead,
3380 );
3381 assert_eq!(captured.len(), 3);
3382 let values = captured
3383 .iter()
3384 .map(|encoded| serde_json::from_str::<Value>(encoded).unwrap())
3385 .collect::<Vec<_>>();
3386 assert_eq!(values[0]["parts"][0]["text"]["value"], "a");
3387 assert_eq!(values[1]["parts"][0]["text"]["value"], "b");
3388 assert_eq!(values[2]["type"], "truncated");
3389 assert_eq!(values[2]["original_bytes"], 2);
3390 }
3391
3392 #[test]
3393 fn tiny_limits_emit_valid_structured_source_byte_truncation() {
3394 let items = vec![Item::text(ItemKind::User, "x".repeat(1_000))];
3395 let captured = capture_messages(
3396 &items,
3397 MessageCapture::new(1, 1).unwrap(),
3398 CaptureOrder::OldestHead,
3399 );
3400 assert_eq!(captured.len(), 1);
3401 let value: Value = serde_json::from_str(&captured[0]).unwrap();
3402 assert_eq!(value["type"], "truncated");
3403 assert_eq!(value["original_bytes"], 1_000);
3404 }
3405
3406 #[test]
3407 fn provider_finish_reason_metadata_round_trips() {
3408 let mut metadata = MetadataMap::new();
3409 set_provider_finish_reasons(&mut metadata, ["end_turn", "", "tool_use", "end_turn"]);
3410 assert_eq!(
3411 provider_finish_reasons(&metadata, &FinishReason::Completed),
3412 vec!["end_turn", "tool_use"]
3413 );
3414 set_provider_finish_reasons(&mut metadata, std::iter::empty::<String>());
3415 assert!(!metadata.contains_key(PROVIDER_FINISH_REASONS_METADATA_KEY));
3416 let fallbacks = [
3417 (FinishReason::Completed, "completed"),
3418 (FinishReason::ToolCall, "tool_call"),
3419 (FinishReason::MaxTokens, "max_tokens"),
3420 (FinishReason::Cancelled, "cancelled"),
3421 (FinishReason::Blocked, "blocked"),
3422 (FinishReason::Error, "error"),
3423 (FinishReason::Other("native".into()), "native"),
3424 ];
3425 for (reason, expected) in fallbacks {
3426 assert_eq!(
3427 provider_finish_reasons(&MetadataMap::new(), &reason),
3428 [expected]
3429 );
3430 }
3431 }
3432
3433 #[test]
3434 fn input_is_newest_tail_output_is_head_and_data_refs_are_omitted() {
3435 let items = vec![
3436 Item::text(ItemKind::User, "old"),
3437 Item::text(ItemKind::User, "middle"),
3438 Item::text(ItemKind::User, "new"),
3439 ];
3440 let capture = MessageCapture::new(2, 10_000).unwrap();
3441 let input = capture_messages(&items, capture, CaptureOrder::NewestTail);
3442 assert!(input[0].contains("middle"));
3443 assert!(input[1].contains("new"));
3444 let output = capture_messages(&items, capture, CaptureOrder::OldestHead);
3445 assert!(output[0].contains("old"));
3446 assert!(output[1].contains("middle"));
3447
3448 let media = Item::new(
3449 ItemKind::User,
3450 vec![Part::media(
3451 agentkit_core::Modality::Image,
3452 "image/png",
3453 agentkit_core::DataRef::uri("https://secret.invalid/image.png"),
3454 )],
3455 );
3456 let encoded = capture_item_json(&media, 10_000);
3457 assert!(!encoded.contains("secret.invalid"));
3458 assert!(encoded.contains("omitted"));
3459 }
3460
3461 #[test]
3462 fn final_usage_wins_and_streamed_usage_fills_only_missing_fields() {
3463 let mut final_metadata = MetadataMap::new();
3464 final_metadata.insert("shared".into(), serde_json::json!("final"));
3465 let final_usage = Usage {
3466 tokens: Some(agentkit_core::TokenUsage::new(1, 2)),
3467 cost: None,
3468 metadata: final_metadata,
3469 };
3470 let mut streamed_metadata = MetadataMap::new();
3471 streamed_metadata.insert("shared".into(), serde_json::json!("streamed"));
3472 streamed_metadata.insert("stream_only".into(), serde_json::json!(true));
3473 let streamed_usage = Usage {
3474 tokens: Some(agentkit_core::TokenUsage::new(10, 20)),
3475 cost: Some(agentkit_core::CostUsage::new(0.5, "USD")),
3476 metadata: streamed_metadata,
3477 };
3478 let merged = merge_usage(Some(final_usage), Some(streamed_usage)).unwrap();
3479 assert_eq!(merged.tokens.unwrap().input_tokens, 1);
3480 assert_eq!(merged.cost.unwrap().amount, 0.5);
3481 assert_eq!(merged.metadata["shared"], "final");
3482 assert_eq!(merged.metadata["stream_only"], true);
3483 }
3484
3485 #[test]
3486 fn session_provider_precedes_adapter_fallback() {
3487 assert_eq!(
3488 effective_provider_name(Some("session"), Some("adapter")).as_deref(),
3489 Some("session")
3490 );
3491 assert_eq!(
3492 effective_provider_name(None, Some("adapter")).as_deref(),
3493 Some("adapter")
3494 );
3495 }
3496}
3497
3498#[cfg(all(test, feature = "otel"))]
3499mod true_otel_integration_tests {
3500 use std::fs;
3501 use std::process::Command;
3502 use std::sync::atomic::{AtomicU64, Ordering};
3503 use std::time::{SystemTime, UNIX_EPOCH};
3504
3505 static TEMP_SEQUENCE: AtomicU64 = AtomicU64::new(0);
3506
3507 #[test]
3508 fn actual_layer_exports_driven_loop_spans() {
3509 let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
3510 .parent()
3511 .unwrap()
3512 .parent()
3513 .unwrap();
3514 let nonce = SystemTime::now()
3515 .duration_since(UNIX_EPOCH)
3516 .unwrap()
3517 .as_nanos();
3518 let sequence = TEMP_SEQUENCE.fetch_add(1, Ordering::Relaxed);
3519 let temp = std::env::temp_dir().join(format!(
3520 "agentkit-loop-true-otel-{}-{nonce}-{sequence}",
3521 std::process::id()
3522 ));
3523 fs::create_dir(&temp).unwrap();
3524 fs::create_dir(temp.join("src")).unwrap();
3525 let core_path = toml_path(&root.join("crates/agentkit-core"));
3526 let loop_path = toml_path(&root.join("crates/agentkit-loop"));
3527 let async_trait_version = locked_version(root, "async-trait");
3530 let opentelemetry_version = locked_version(root, "opentelemetry");
3531 let serde_json_version = locked_version(root, "serde_json");
3532 let tokio_version = locked_version(root, "tokio");
3533 let tracing_version = locked_version(root, "tracing");
3534 let tracing_otel_version = locked_version(root, "tracing-opentelemetry");
3535 let tracing_subscriber_version = locked_version(root, "tracing-subscriber");
3536 let manifest = format!(
3537 r#"[package]
3538name = "agentkit-loop-true-otel-test"
3539version = "0.0.0"
3540edition = "2024"
3541
3542[dependencies]
3543agentkit-core = {{ path = {core_path} }}
3544agentkit-loop = {{ path = {loop_path}, features = ["otel"] }}
3545async-trait = "={async_trait_version}"
3546opentelemetry = {{ version = "={opentelemetry_version}", default-features = false }}
3547serde_json = "={serde_json_version}"
3548tokio = {{ version = "={tokio_version}", features = ["rt"] }}
3549tracing = "={tracing_version}"
3550tracing-opentelemetry = {{ version = "={tracing_otel_version}", default-features = false }}
3551tracing-subscriber = "={tracing_subscriber_version}"
3552"#
3553 );
3554 fs::write(temp.join("Cargo.toml"), manifest).unwrap();
3555 fs::write(temp.join("src/main.rs"), TRUE_OTEL_HARNESS).unwrap();
3556
3557 let output = Command::new(env!("CARGO"))
3558 .args(["run", "--quiet", "--offline"])
3559 .current_dir(&temp)
3560 .env("CARGO_TARGET_DIR", temp.join("target"))
3561 .output()
3562 .unwrap();
3563 let _ = fs::remove_dir_all(&temp);
3564 assert!(
3565 output.status.success(),
3566 "true OTEL harness failed:\nstdout:\n{}\nstderr:\n{}",
3567 String::from_utf8_lossy(&output.stdout),
3568 String::from_utf8_lossy(&output.stderr)
3569 );
3570 }
3571
3572 fn locked_version(root: &std::path::Path, name: &str) -> String {
3573 let lock = fs::read_to_string(root.join("Cargo.lock")).unwrap();
3574 let expected_name = format!("name = {}", serde_json::to_string(name).unwrap());
3575 let versions = lock
3576 .split("[[package]]")
3577 .filter(|package| {
3578 package
3579 .lines()
3580 .any(|line| line.trim() == expected_name.as_str())
3581 })
3582 .filter_map(|package| {
3583 package.lines().find_map(|line| {
3584 line.trim()
3585 .strip_prefix("version = \"")
3586 .and_then(|version| version.strip_suffix('\"'))
3587 .map(str::to_owned)
3588 })
3589 })
3590 .collect::<Vec<_>>();
3591 assert_eq!(versions.len(), 1, "expected one locked version for {name}");
3592 versions.into_iter().next().unwrap()
3593 }
3594
3595 fn toml_path(path: &std::path::Path) -> String {
3596 serde_json::to_string(&path.to_string_lossy()).unwrap()
3597 }
3598
3599 #[test]
3600 fn toml_path_escapes_windows_separators_and_quotes() {
3601 let encoded = toml_path(std::path::Path::new(r#"C:\Users\name\quoted\"dir"#));
3602 assert_eq!(encoded, r#""C:\\Users\\name\\quoted\\\"dir""#);
3603 }
3604
3605 const TRUE_OTEL_HARNESS: &str = r#"
3606use std::borrow::Cow;
3607use std::collections::VecDeque;
3608use std::sync::{Arc, Mutex};
3609use std::time::SystemTime;
3610use agentkit_core::{CostUsage, DataRef, FinishReason, Item, ItemKind, MetadataMap, Modality, Part, TokenUsage, TurnCancellation, Usage};
3611use agentkit_loop::{Agent, LoopError, LoopStep, MessageCapture, ModelAdapter, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, TelemetryConfig, TurnRequest, set_provider_finish_reasons};
3612use async_trait::async_trait;
3613use opentelemetry::trace::{Span, SpanBuilder, SpanContext, Status, Tracer};
3614use opentelemetry::{Array, Context, KeyValue, Value};
3615use tracing_subscriber::layer::SubscriberExt;
3616
3617#[derive(Clone, Debug)]
3618struct Exported { name: String, attributes: Vec<KeyValue> }
3619#[derive(Clone, Default)]
3620struct MemoryTracer { exported: Arc<Mutex<Vec<Exported>>> }
3621struct MemorySpan { name: String, attributes: Vec<KeyValue>, exported: Arc<Mutex<Vec<Exported>>>, ended: bool }
3622impl Tracer for MemoryTracer {
3623 type Span = MemorySpan;
3624 fn build_with_context(&self, builder: SpanBuilder, _: &Context) -> MemorySpan {
3625 MemorySpan { name: builder.name.into_owned(), attributes: builder.attributes.unwrap_or_default(), exported: self.exported.clone(), ended: false }
3626 }
3627}
3628impl Span for MemorySpan {
3629 fn add_event_with_timestamp<T>(&mut self, _: T, _: SystemTime, _: Vec<KeyValue>) where T: Into<Cow<'static, str>> {}
3630 fn span_context(&self) -> &SpanContext { &SpanContext::NONE }
3631 fn is_recording(&self) -> bool { !self.ended }
3632 fn set_attribute(&mut self, attribute: KeyValue) { self.attributes.push(attribute); }
3633 fn set_status(&mut self, _: Status) {}
3634 fn update_name<T>(&mut self, name: T) where T: Into<Cow<'static, str>> { self.name = name.into().into_owned(); }
3635 fn add_link(&mut self, _: SpanContext, _: Vec<KeyValue>) {}
3636 fn end_with_timestamp(&mut self, _: SystemTime) {
3637 if !self.ended {
3638 self.ended = true;
3639 self.exported.lock().unwrap().push(Exported { name: self.name.clone(), attributes: self.attributes.clone() });
3640 }
3641 }
3642}
3643
3644#[derive(Clone, Copy)]
3645enum BeginMode { Success, Error, Cancelled }
3646#[derive(Clone)]
3647struct ScriptedAdapter { adapter_provider: &'static str, final_usage: bool, overflow: bool, begin_mode: BeginMode }
3648struct ScriptedSession { selected_provider: Option<&'static str>, model: &'static str, final_usage: bool, overflow: bool, begin_mode: BeginMode }
3649struct ScriptedTurn { events: VecDeque<ModelTurnEvent> }
3650#[async_trait]
3651impl ModelAdapter for ScriptedAdapter {
3652 type Session = ScriptedSession;
3653 async fn start_session(&self, _: SessionConfig) -> Result<Self::Session, LoopError> {
3654 Ok(ScriptedSession { selected_provider: Some("before-begin"), model: "before-model", final_usage: self.final_usage, overflow: self.overflow, begin_mode: self.begin_mode })
3655 }
3656 fn provider_name(&self) -> Option<&str> { Some(self.adapter_provider) }
3657}
3658#[async_trait]
3659impl ModelSession for ScriptedSession {
3660 type Turn = ScriptedTurn;
3661 async fn begin_turn(&mut self, _: TurnRequest, _: Option<TurnCancellation>) -> Result<Self::Turn, LoopError> {
3662 match self.begin_mode {
3663 BeginMode::Error => return Err(LoopError::InvalidState("begin failed".into())),
3664 BeginMode::Cancelled => return Err(LoopError::Cancelled),
3665 BeginMode::Success => {}
3666 }
3667 self.selected_provider = if self.final_usage { Some("session-after-begin") } else { None };
3668 self.model = if self.final_usage { "model-after-begin" } else { "fallback-model-after-begin" };
3669 let mut stream_meta = MetadataMap::new();
3670 stream_meta.insert("stream_only".into(), serde_json::json!(true));
3671 stream_meta.insert("shared".into(), serde_json::json!("stream"));
3672 let streamed = Usage {
3673 tokens: Some(TokenUsage::new(if self.overflow { i64::MAX as u64 + 1 } else { 30 }, 40)),
3674 cost: Some(CostUsage::new(0.75, "USD")),
3675 metadata: stream_meta,
3676 };
3677 let final_usage = if self.final_usage {
3678 let mut metadata = MetadataMap::new();
3679 metadata.insert("shared".into(), serde_json::json!("final"));
3680 Some(Usage { tokens: Some(TokenUsage::new(if self.overflow { i64::MAX as u64 + 1 } else { 1 }, 2)), cost: None, metadata })
3681 } else { None };
3682 let mut result_metadata = MetadataMap::new();
3683 set_provider_finish_reasons(&mut result_metadata, ["native", "native", "done"]);
3684 let outputs = vec![
3685 Item::text(ItemKind::Assistant, "first"),
3686 Item::text(ItemKind::Assistant, "second"),
3687 Item::text(ItemKind::Assistant, "third"),
3688 ];
3689 Ok(ScriptedTurn { events: VecDeque::from([
3690 ModelTurnEvent::Usage(streamed),
3691 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()) }),
3692 ]) })
3693 }
3694 fn model_name(&self) -> Option<&str> { Some(self.model) }
3695 fn provider_name(&self) -> Option<&str> { self.selected_provider }
3696}
3697#[async_trait]
3698impl ModelTurn for ScriptedTurn {
3699 async fn next_event(&mut self, _: Option<TurnCancellation>) -> Result<Option<ModelTurnEvent>, LoopError> { Ok(self.events.pop_front()) }
3700}
3701
3702fn attr<'a>(span: &'a Exported, key: &str) -> Option<&'a Value> {
3703 span.attributes.iter().rev().find(|a| a.key.as_str() == key).map(|a| &a.value)
3704}
3705fn operation(span: &Exported) -> Option<&str> {
3706 match attr(span, "gen_ai.operation.name") { Some(Value::String(value)) => Some(value.as_str()), _ => None }
3707}
3708fn json_array(span: &Exported, key: &str) -> Vec<serde_json::Value> {
3709 match attr(span, key) {
3710 Some(Value::Array(Array::String(values))) => values.iter().map(|v| serde_json::from_str(v.as_str()).unwrap()).collect(),
3711 other => panic!("{key} was not Array<String>: {other:?}"),
3712 }
3713}
3714
3715fn run_attempt(adapter: ScriptedAdapter) -> (Vec<Exported>, Result<LoopStep, LoopError>) {
3716 let tracer = MemoryTracer::default();
3717 let subscriber = tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer.clone()));
3718 let runtime = tokio::runtime::Builder::new_current_thread().build().unwrap();
3719 let result = tracing::subscriber::with_default(subscriber, || runtime.block_on(async {
3720 let media = Item::new(ItemKind::User, vec![
3721 Part::media(Modality::Image, "image/png", DataRef::inline_bytes([1, 2, 3])),
3722 Part::media(Modality::Audio, "audio/wav", DataRef::uri("https://secret.invalid/audio")),
3723 ]);
3724 let agent = Agent::builder().model(adapter).transcript(vec![
3725 Item::text(ItemKind::System, "old"), Item::text(ItemKind::User, "middle"), media,
3726 ]).input(vec![Item::text(ItemKind::User, "newest")]).telemetry(
3727 TelemetryConfig::default()
3728 .with_input_messages(MessageCapture::new(3, 100_000).unwrap())
3729 .with_output_messages(MessageCapture::new(2, 100_000).unwrap())
3730 ).build().unwrap();
3731 let mut driver = agent.start(SessionConfig::new("otel-test")).await.unwrap();
3732 driver.next().await
3733 }));
3734 (tracer.exported.lock().unwrap().clone(), result)
3735}
3736fn run(adapter: ScriptedAdapter) -> (Vec<Exported>, agentkit_loop::TurnResult) {
3737 let (spans, result) = run_attempt(adapter);
3738 let result = match result.unwrap() { LoopStep::Finished(result) => result, other => panic!("unexpected step: {other:?}") };
3739 (spans, result)
3740}
3741
3742fn main() {
3743 let (spans, result) = run(ScriptedAdapter { adapter_provider: "adapter", final_usage: true, overflow: true, begin_mode: BeginMode::Success });
3744 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
3745 let agent = spans.iter().find(|s| operation(s) == Some("invoke_agent")).unwrap();
3746 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("session-after-begin".into())));
3747 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("model-after-begin".into())));
3748 assert!(attr(chat, "gen_ai.usage.input_tokens").is_none());
3749 assert_eq!(attr(chat, "gen_ai.usage.output_tokens"), Some(&Value::I64(2)));
3750 assert_eq!(attr(chat, "gen_ai.usage.cost"), Some(&Value::F64(0.75)));
3751 assert_eq!(attr(agent, "gen_ai.provider.name"), Some(&Value::String("session-after-begin".into())));
3752 assert!(attr(agent, "gen_ai.usage.input_tokens").is_none());
3753 assert!(attr(agent, "gen_ai.usage.cost").is_none());
3754 match attr(chat, "gen_ai.response.finish_reasons") {
3755 Some(Value::Array(Array::String(values))) => assert_eq!(values.iter().map(|v| v.as_str()).collect::<Vec<_>>(), ["native", "done"]),
3756 other => panic!("finish reasons were not Array<String>: {other:?}"),
3757 }
3758 assert_eq!(result.usage.as_ref().unwrap().metadata["shared"], "final");
3759 assert_eq!(result.usage.as_ref().unwrap().metadata["stream_only"], true);
3760 let input = json_array(chat, "gen_ai.input.messages");
3761 assert_eq!(input.len(), 3);
3762 assert!(input[0].to_string().contains("middle"));
3763 assert!(input[1].to_string().contains("omitted"));
3764 assert!(input[2].to_string().contains("newest"));
3765 assert!(!input.iter().any(|message| message.to_string().contains("old")));
3766 let output = json_array(chat, "gen_ai.output.messages");
3767 assert_eq!(output.len(), 2);
3768 assert!(output[0].to_string().contains("first"));
3769 assert!(output[1].to_string().contains("second"));
3770 let encoded = format!("{input:?}{output:?}");
3771 assert!(!encoded.contains("secret.invalid"));
3772 assert!(!encoded.contains("[1,2,3]"));
3773
3774 let (spans, result) = run(ScriptedAdapter { adapter_provider: "adapter-fallback", final_usage: false, overflow: false, begin_mode: BeginMode::Success });
3775 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
3776 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("adapter-fallback".into())));
3777 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("fallback-model-after-begin".into())));
3778 assert_eq!(attr(chat, "gen_ai.usage.input_tokens"), Some(&Value::I64(30)));
3779 assert_eq!(attr(chat, "gen_ai.usage.cost"), Some(&Value::F64(0.75)));
3780 assert_eq!(result.usage.unwrap().metadata["stream_only"], true);
3781
3782 for begin_mode in [BeginMode::Error, BeginMode::Cancelled] {
3783 let (spans, result) = run_attempt(ScriptedAdapter { adapter_provider: "adapter-before-error", final_usage: false, overflow: false, begin_mode });
3784 match begin_mode {
3785 BeginMode::Error => assert!(matches!(result, Err(LoopError::InvalidState(_)))),
3786 BeginMode::Cancelled => assert!(matches!(result, Ok(LoopStep::Finished(_)))),
3787 BeginMode::Success => unreachable!(),
3788 }
3789 let chat = spans.iter().find(|s| operation(s) == Some("chat")).unwrap();
3790 let agent = spans.iter().find(|s| operation(s) == Some("invoke_agent")).unwrap();
3791 assert_eq!(attr(chat, "gen_ai.provider.name"), Some(&Value::String("before-begin".into())));
3792 assert_eq!(attr(chat, "gen_ai.request.model"), Some(&Value::String("before-model".into())));
3793 assert_eq!(attr(agent, "gen_ai.provider.name"), Some(&Value::String("before-begin".into())));
3794 }
3795}
3796"#;
3797}
3798
3799fn interrupted_assistant_items() -> Vec<Item> {
3800 vec![Item {
3801 id: None,
3802 kind: ItemKind::Assistant,
3803 parts: vec![Part::Text(TextPart {
3804 text: "Previous assistant response was interrupted by the user before completion."
3805 .into(),
3806 metadata: interrupted_metadata("assistant"),
3807 })],
3808 metadata: interrupted_metadata("assistant"),
3809 usage: None,
3810 finish_reason: None,
3811 created_at: None,
3812 }]
3813}
3814
3815fn unanswered_tool_calls(transcript: &[Item]) -> Vec<ToolCallPart> {
3817 let mut open: Vec<ToolCallPart> = Vec::new();
3818 for part in transcript.iter().flat_map(|item| &item.parts) {
3819 match part {
3820 Part::ToolCall(call) => open.push(call.clone()),
3821 Part::ToolResult(result) => open.retain(|call| call.id != result.call_id),
3822 _ => {}
3823 }
3824 }
3825 open
3826}
3827
3828fn interrupted_tool_result_item(call: ToolCallPart) -> Item {
3833 Item {
3834 id: None,
3835 kind: ItemKind::Tool,
3836 parts: vec![Part::ToolResult(ToolResultPart {
3837 call_id: call.id,
3838 output: ToolOutput::Text("tool call interrupted before it reported a result".into()),
3839 is_error: true,
3840 metadata: interrupted_metadata("tool"),
3841 })],
3842 metadata: interrupted_metadata("tool"),
3843 usage: None,
3844 finish_reason: None,
3845 created_at: None,
3846 }
3847}
3848
3849fn cancelled_approval_item(pending: PendingApprovalToolCall) -> Item {
3850 Item {
3851 id: None,
3852 kind: ItemKind::Tool,
3853 parts: vec![Part::ToolResult(ToolResultPart {
3854 call_id: pending.call.id,
3855 output: ToolOutput::Text("approval cancelled".into()),
3856 is_error: true,
3857 metadata: pending.call.metadata,
3858 })],
3859 metadata: MetadataMap::new(),
3860 usage: None,
3861 finish_reason: None,
3862 created_at: None,
3863 }
3864}
3865
3866fn transcript_has_pending_input(transcript: &[Item]) -> bool {
3872 matches!(
3873 transcript.last().map(|item| item.kind),
3874 Some(ItemKind::User | ItemKind::Tool | ItemKind::Notification)
3875 )
3876}
3877
3878fn extract_tool_calls(items: &[Item]) -> Vec<ToolCallPart> {
3879 let mut calls = Vec::new();
3880 for item in items {
3881 for part in &item.parts {
3882 if let Part::ToolCall(call) = part {
3883 calls.push(call.clone());
3884 }
3885 }
3886 }
3887 calls
3888}
3889
3890fn tool_result_is_error(item: &Item) -> bool {
3891 item.parts
3892 .iter()
3893 .any(|part| matches!(part, Part::ToolResult(result) if result.is_error))
3894}
3895
3896fn tool_result_not_started(item: &Item) -> bool {
3897 item.parts.iter().any(|part| {
3898 matches!(
3899 part,
3900 Part::ToolResult(result)
3901 if result
3902 .metadata
3903 .get(TOOL_RESULT_NOT_STARTED_METADATA_KEY)
3904 .and_then(Value::as_bool)
3905 == Some(true)
3906 )
3907 })
3908}
3909
3910#[derive(Debug, Error)]
3912pub enum LoopError {
3913 #[error("invalid driver state: {0}")]
3915 InvalidState(String),
3916 #[error("turn cancelled")]
3918 Cancelled,
3919 #[error("provider error: {0}")]
3921 Provider(String),
3922 #[error("tool error: {0}")]
3924 Tool(#[from] ToolError),
3925 #[error("mutator error: {0}")]
3927 Mutator(String),
3928 #[error("unsupported operation: {0}")]
3930 Unsupported(String),
3931}
3932
3933struct DriverEmitter<'a> {
3938 session_id: &'a Arc<SessionId>,
3939 observers: &'a [Arc<dyn LoopObserver>],
3940}
3941
3942impl<'a> EventEmitter for DriverEmitter<'a> {
3943 fn emit(&self, event: AgentEvent) {
3944 fan_out_observed_event(self.observers, self.session_id, event);
3945 }
3946}
3947
3948fn fan_out_observed_event(
3949 observers: &[Arc<dyn LoopObserver>],
3950 session_id: &Arc<SessionId>,
3951 event: AgentEvent,
3952) {
3953 if observers.is_empty() {
3954 return;
3955 }
3956 let observed = ObservedEvent {
3957 session_id: Arc::clone(session_id),
3958 event,
3959 };
3960 let last = observers.len() - 1;
3961 for observer in &observers[..last] {
3962 observer.handle_event(observed.clone());
3963 }
3964 observers[last].handle_event(observed);
3965}
3966
3967fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> {
3972 let mut pending: HashSet<ToolCallId> = HashSet::new();
3973 let mut seen_calls: HashSet<ToolCallId> = HashSet::new();
3974 let mut seen_results: HashSet<ToolCallId> = HashSet::new();
3975 for item in transcript {
3976 for part in &item.parts {
3977 match part {
3978 Part::ToolCall(call) => {
3979 if !seen_calls.insert(call.id.clone()) {
3980 return Err(LoopError::Mutator(format!(
3981 "transcript invariant violation: duplicate tool_use: {}",
3982 call.id.0
3983 )));
3984 }
3985 pending.insert(call.id.clone());
3986 }
3987 Part::ToolResult(result) => {
3988 if !pending.remove(&result.call_id) {
3989 let kind = if seen_results.contains(&result.call_id) {
3990 "duplicate"
3991 } else {
3992 "orphaned"
3993 };
3994 return Err(LoopError::Mutator(format!(
3995 "transcript invariant violation: {kind} tool_result: {}",
3996 result.call_id.0
3997 )));
3998 }
3999 seen_results.insert(result.call_id.clone());
4000 }
4001 _ => {}
4002 }
4003 }
4004 }
4005 if !pending.is_empty() {
4006 let missing: Vec<String> = pending.into_iter().map(|id| id.0).collect();
4007 return Err(LoopError::Mutator(format!(
4008 "transcript invariant violation: tool_use(s) without matching tool_result: {}",
4009 missing.join(", ")
4010 )));
4011 }
4012 Ok(())
4013}
4014
4015#[cfg(test)]
4016mod tests {
4017 use std::collections::VecDeque;
4018 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
4019 use std::sync::{Arc as StdArc, Mutex as StdMutex};
4020
4021 use agentkit_core::{
4022 CancellationController, ItemKind, Part, TextPart, ToolCallId, ToolCallPart, ToolOutput,
4023 ToolResultPart,
4024 };
4025 use agentkit_task_manager::{
4026 AsyncTaskManager, RoutingDecision, TaskEvent, TaskManager, TaskManagerHandle,
4027 TaskRoutingPolicy,
4028 };
4029 use agentkit_tools_core::{
4030 FileSystemPermissionRequest, PermissionCode, PermissionDecision, PermissionDenial, Tool,
4031 ToolAnnotations, ToolCatalogEvent, ToolExecutionOutcome, ToolName, ToolRegistry,
4032 ToolResult, ToolSpec,
4033 };
4034 use serde_json::{Value, json};
4035 use tokio::sync::Notify;
4036 use tokio::time::{Duration, timeout};
4037
4038 use super::*;
4039
4040 struct FakeAdapter;
4041 struct SlowAdapter;
4042 struct RecordingAdapter {
4043 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
4044 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
4045 }
4046 struct MultiToolAdapter;
4047 struct DualApprovalAdapter;
4048
4049 struct FakeSession;
4050 struct SlowSession;
4051 struct RecordingSession {
4052 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
4053 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
4054 }
4055 struct MultiToolSession;
4056 struct DualApprovalSession;
4057
4058 struct FakeTurn {
4059 events: VecDeque<ModelTurnEvent>,
4060 }
4061
4062 struct SlowTurn {
4063 emitted: bool,
4064 }
4065
4066 struct RecordingTurn {
4067 emitted: bool,
4068 }
4069 struct MultiToolTurn {
4070 events: VecDeque<ModelTurnEvent>,
4071 }
4072 struct DualApprovalTurn {
4073 events: VecDeque<ModelTurnEvent>,
4074 }
4075
4076 struct DelayedApprovalExecutor {
4077 entered: StdArc<AtomicBool>,
4078 release: StdArc<Notify>,
4079 cancellation: Option<CancellationController>,
4080 spec: ToolSpec,
4081 }
4082
4083 impl DelayedApprovalExecutor {
4084 fn new(entered: StdArc<AtomicBool>, release: StdArc<Notify>) -> Self {
4085 Self {
4086 entered,
4087 release,
4088 cancellation: None,
4089 spec: ToolSpec {
4090 name: ToolName::new("echo"),
4091 description: "delayed approval".into(),
4092 input_schema: json!({
4093 "type": "object",
4094 "properties": {
4095 "value": { "type": "string" }
4096 },
4097 "required": ["value"],
4098 "additionalProperties": false
4099 }),
4100 output_schema: None,
4101 annotations: ToolAnnotations::default(),
4102 metadata: MetadataMap::new(),
4103 },
4104 }
4105 }
4106
4107 fn cancelling_on_approval(mut self, controller: CancellationController) -> Self {
4108 self.cancellation = Some(controller);
4109 self
4110 }
4111 }
4112
4113 #[async_trait]
4114 impl ToolExecutor for DelayedApprovalExecutor {
4115 fn specs(&self) -> Vec<ToolSpec> {
4116 vec![self.spec.clone()]
4117 }
4118
4119 async fn execute(
4120 &self,
4121 request: ToolRequest,
4122 _ctx: &mut ToolContext<'_>,
4123 ) -> ToolExecutionOutcome {
4124 self.entered.store(true, Ordering::SeqCst);
4125 self.release.notified().await;
4126 if let Some(controller) = &self.cancellation {
4127 controller.interrupt();
4128 }
4129 ToolExecutionOutcome::Interrupted(
4130 agentkit_tools_core::ToolInterruption::ApprovalRequired(ApprovalRequest {
4131 task_id: None,
4132 call_id: None,
4133 id: "approval:delayed".into(),
4134 request_kind: "delayed.approval".into(),
4135 reason: agentkit_tools_core::ApprovalReason::PolicyRequiresConfirmation,
4136 summary: "delayed approval".into(),
4137 metadata: request.metadata,
4138 }),
4139 )
4140 }
4141 }
4142
4143 #[async_trait]
4144 impl ModelAdapter for FakeAdapter {
4145 type Session = FakeSession;
4146
4147 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4148 Ok(FakeSession)
4149 }
4150 }
4151
4152 #[async_trait]
4153 impl ModelAdapter for SlowAdapter {
4154 type Session = SlowSession;
4155
4156 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4157 Ok(SlowSession)
4158 }
4159 }
4160
4161 #[async_trait]
4162 impl ModelAdapter for RecordingAdapter {
4163 type Session = RecordingSession;
4164
4165 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4166 Ok(RecordingSession {
4167 seen_descriptions: self.seen_descriptions.clone(),
4168 seen_caches: self.seen_caches.clone(),
4169 })
4170 }
4171 }
4172
4173 #[async_trait]
4174 impl ModelAdapter for MultiToolAdapter {
4175 type Session = MultiToolSession;
4176
4177 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4178 Ok(MultiToolSession)
4179 }
4180 }
4181
4182 #[async_trait]
4183 impl ModelAdapter for DualApprovalAdapter {
4184 type Session = DualApprovalSession;
4185
4186 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4187 Ok(DualApprovalSession)
4188 }
4189 }
4190
4191 #[async_trait]
4192 impl ModelSession for FakeSession {
4193 type Turn = FakeTurn;
4194
4195 async fn begin_turn(
4196 &mut self,
4197 request: TurnRequest,
4198 _cancellation: Option<TurnCancellation>,
4199 ) -> Result<Self::Turn, LoopError> {
4200 let has_tool_result = request.transcript.iter().any(|item| {
4201 item.kind == ItemKind::Tool
4202 && item
4203 .parts
4204 .iter()
4205 .any(|part| matches!(part, Part::ToolResult(_)))
4206 });
4207 let tool_name = request
4208 .available_tools
4209 .first()
4210 .map(|tool| tool.name.0.clone())
4211 .unwrap_or_else(|| "echo".into());
4212
4213 let events = if has_tool_result {
4214 let result_text = request
4215 .transcript
4216 .iter()
4217 .rev()
4218 .find_map(|item| {
4219 item.parts.iter().find_map(|part| match part {
4220 Part::ToolResult(ToolResultPart {
4221 output: ToolOutput::Text(text),
4222 ..
4223 }) => Some(text.clone()),
4224 _ => None,
4225 })
4226 })
4227 .unwrap_or_else(|| "missing".into());
4228
4229 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
4230 model: None,
4231 response_id: None,
4232 finish_reason: FinishReason::Completed,
4233 output_items: vec![Item {
4234 id: None,
4235 kind: ItemKind::Assistant,
4236 parts: vec![Part::Text(TextPart {
4237 text: format!("tool said: {result_text}"),
4238 metadata: MetadataMap::new(),
4239 })],
4240 metadata: MetadataMap::new(),
4241 usage: None,
4242 finish_reason: None,
4243 created_at: None,
4244 }],
4245 usage: None,
4246 metadata: MetadataMap::new(),
4247 })])
4248 } else {
4249 VecDeque::from([
4250 ModelTurnEvent::ToolCall(agentkit_core::ToolCallPart {
4251 id: ToolCallId::new("call-1"),
4252 name: tool_name.clone(),
4253 input: json!({ "value": "pong" }),
4254 metadata: MetadataMap::new(),
4255 }),
4256 ModelTurnEvent::Finished(ModelTurnResult {
4257 model: None,
4258 response_id: None,
4259 finish_reason: FinishReason::ToolCall,
4260 output_items: vec![Item {
4261 id: None,
4262 kind: ItemKind::Assistant,
4263 parts: vec![Part::ToolCall(agentkit_core::ToolCallPart {
4264 id: ToolCallId::new("call-1"),
4265 name: tool_name,
4266 input: json!({ "value": "pong" }),
4267 metadata: MetadataMap::new(),
4268 })],
4269 metadata: MetadataMap::new(),
4270 usage: None,
4271 finish_reason: None,
4272 created_at: None,
4273 }],
4274 usage: None,
4275 metadata: MetadataMap::new(),
4276 }),
4277 ])
4278 };
4279
4280 Ok(FakeTurn { events })
4281 }
4282 }
4283
4284 #[async_trait]
4285 impl ModelSession for SlowSession {
4286 type Turn = SlowTurn;
4287
4288 async fn begin_turn(
4289 &mut self,
4290 request: TurnRequest,
4291 cancellation: Option<TurnCancellation>,
4292 ) -> Result<Self::Turn, LoopError> {
4293 let should_block = request
4294 .transcript
4295 .iter()
4296 .rev()
4297 .find(|item| item.kind == ItemKind::User)
4298 .is_some_and(|item| {
4299 item.parts.iter().any(|part| match part {
4300 Part::Text(text) => text.text == "do the long task",
4301 _ => false,
4302 })
4303 });
4304
4305 if should_block && let Some(cancellation) = cancellation {
4306 cancellation.cancelled().await;
4307 return Err(LoopError::Cancelled);
4308 }
4309
4310 Ok(SlowTurn { emitted: false })
4311 }
4312 }
4313
4314 #[async_trait]
4315 impl ModelSession for RecordingSession {
4316 type Turn = RecordingTurn;
4317
4318 async fn begin_turn(
4319 &mut self,
4320 request: TurnRequest,
4321 _cancellation: Option<TurnCancellation>,
4322 ) -> Result<Self::Turn, LoopError> {
4323 let descriptions = request
4324 .available_tools
4325 .iter()
4326 .map(|tool| tool.description.clone())
4327 .collect::<Vec<_>>();
4328 self.seen_descriptions.lock().unwrap().push(descriptions);
4329 self.seen_caches.lock().unwrap().push(request.cache.clone());
4330
4331 Ok(RecordingTurn { emitted: false })
4332 }
4333 }
4334
4335 #[async_trait]
4336 impl ModelSession for MultiToolSession {
4337 type Turn = MultiToolTurn;
4338
4339 async fn begin_turn(
4340 &mut self,
4341 request: TurnRequest,
4342 _cancellation: Option<TurnCancellation>,
4343 ) -> Result<Self::Turn, LoopError> {
4344 let has_tool_result = request.transcript.iter().any(|item| {
4345 item.kind == ItemKind::Tool
4346 && item
4347 .parts
4348 .iter()
4349 .any(|part| matches!(part, Part::ToolResult(_)))
4350 });
4351
4352 let events = if has_tool_result {
4353 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
4354 model: None,
4355 response_id: None,
4356 finish_reason: FinishReason::Completed,
4357 output_items: vec![Item {
4358 id: None,
4359 kind: ItemKind::Assistant,
4360 parts: vec![Part::Text(TextPart {
4361 text: "mixed tools finished".into(),
4362 metadata: MetadataMap::new(),
4363 })],
4364 metadata: MetadataMap::new(),
4365 usage: None,
4366 finish_reason: None,
4367 created_at: None,
4368 }],
4369 usage: None,
4370 metadata: MetadataMap::new(),
4371 })])
4372 } else {
4373 let foreground = agentkit_core::ToolCallPart {
4374 id: ToolCallId::new("call-foreground"),
4375 name: "foreground-wait".into(),
4376 input: json!({}),
4377 metadata: MetadataMap::new(),
4378 };
4379 let background = agentkit_core::ToolCallPart {
4380 id: ToolCallId::new("call-background"),
4381 name: "background-wait".into(),
4382 input: json!({}),
4383 metadata: MetadataMap::new(),
4384 };
4385 VecDeque::from([
4386 ModelTurnEvent::ToolCall(foreground.clone()),
4387 ModelTurnEvent::ToolCall(background.clone()),
4388 ModelTurnEvent::Finished(ModelTurnResult {
4389 model: None,
4390 response_id: None,
4391 finish_reason: FinishReason::ToolCall,
4392 output_items: vec![Item {
4393 id: None,
4394 kind: ItemKind::Assistant,
4395 parts: vec![Part::ToolCall(foreground), Part::ToolCall(background)],
4396 metadata: MetadataMap::new(),
4397 usage: None,
4398 finish_reason: None,
4399 created_at: None,
4400 }],
4401 usage: None,
4402 metadata: MetadataMap::new(),
4403 }),
4404 ])
4405 };
4406
4407 Ok(MultiToolTurn { events })
4408 }
4409 }
4410
4411 #[async_trait]
4412 impl ModelSession for DualApprovalSession {
4413 type Turn = DualApprovalTurn;
4414
4415 async fn begin_turn(
4416 &mut self,
4417 request: TurnRequest,
4418 _cancellation: Option<TurnCancellation>,
4419 ) -> Result<Self::Turn, LoopError> {
4420 let tool_results = request
4421 .transcript
4422 .iter()
4423 .flat_map(|item| item.parts.iter())
4424 .filter(|part| matches!(part, Part::ToolResult(_)))
4425 .count();
4426
4427 let events = if tool_results >= 2 {
4428 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
4429 model: None,
4430 response_id: None,
4431 finish_reason: FinishReason::Completed,
4432 output_items: vec![Item {
4433 id: None,
4434 kind: ItemKind::Assistant,
4435 parts: vec![Part::Text(TextPart {
4436 text: "both approvals finished".into(),
4437 metadata: MetadataMap::new(),
4438 })],
4439 metadata: MetadataMap::new(),
4440 usage: None,
4441 finish_reason: None,
4442 created_at: None,
4443 }],
4444 usage: None,
4445 metadata: MetadataMap::new(),
4446 })])
4447 } else {
4448 let first = agentkit_core::ToolCallPart {
4449 id: ToolCallId::new("call-1"),
4450 name: "echo".into(),
4451 input: json!({ "value": "first" }),
4452 metadata: MetadataMap::new(),
4453 };
4454 let second = agentkit_core::ToolCallPart {
4455 id: ToolCallId::new("call-2"),
4456 name: "echo".into(),
4457 input: json!({ "value": "second" }),
4458 metadata: MetadataMap::new(),
4459 };
4460 VecDeque::from([
4461 ModelTurnEvent::ToolCall(first.clone()),
4462 ModelTurnEvent::ToolCall(second.clone()),
4463 ModelTurnEvent::Finished(ModelTurnResult {
4464 model: None,
4465 response_id: None,
4466 finish_reason: FinishReason::ToolCall,
4467 output_items: vec![Item {
4468 id: None,
4469 kind: ItemKind::Assistant,
4470 parts: vec![Part::ToolCall(first), Part::ToolCall(second)],
4471 metadata: MetadataMap::new(),
4472 usage: None,
4473 finish_reason: None,
4474 created_at: None,
4475 }],
4476 usage: None,
4477 metadata: MetadataMap::new(),
4478 }),
4479 ])
4480 };
4481
4482 Ok(DualApprovalTurn { events })
4483 }
4484 }
4485
4486 #[async_trait]
4487 impl ModelTurn for FakeTurn {
4488 async fn next_event(
4489 &mut self,
4490 _cancellation: Option<TurnCancellation>,
4491 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4492 Ok(self.events.pop_front())
4493 }
4494 }
4495
4496 #[async_trait]
4497 impl ModelTurn for SlowTurn {
4498 async fn next_event(
4499 &mut self,
4500 cancellation: Option<TurnCancellation>,
4501 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4502 if let Some(cancellation) = cancellation
4503 && cancellation.is_cancelled()
4504 {
4505 return Err(LoopError::Cancelled);
4506 }
4507
4508 if self.emitted {
4509 Ok(None)
4510 } else {
4511 self.emitted = true;
4512 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
4513 model: None,
4514 response_id: None,
4515 finish_reason: FinishReason::Completed,
4516 output_items: vec![Item {
4517 id: None,
4518 kind: ItemKind::Assistant,
4519 parts: vec![Part::Text(TextPart {
4520 text: "done".into(),
4521 metadata: MetadataMap::new(),
4522 })],
4523 metadata: MetadataMap::new(),
4524 usage: None,
4525 finish_reason: None,
4526 created_at: None,
4527 }],
4528 usage: None,
4529 metadata: MetadataMap::new(),
4530 })))
4531 }
4532 }
4533 }
4534
4535 #[async_trait]
4536 impl ModelTurn for RecordingTurn {
4537 async fn next_event(
4538 &mut self,
4539 _cancellation: Option<TurnCancellation>,
4540 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4541 if self.emitted {
4542 Ok(None)
4543 } else {
4544 self.emitted = true;
4545 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
4546 model: None,
4547 response_id: None,
4548 finish_reason: FinishReason::Completed,
4549 output_items: vec![Item {
4550 id: None,
4551 kind: ItemKind::Assistant,
4552 parts: vec![Part::Text(TextPart {
4553 text: "done".into(),
4554 metadata: MetadataMap::new(),
4555 })],
4556 metadata: MetadataMap::new(),
4557 usage: None,
4558 finish_reason: None,
4559 created_at: None,
4560 }],
4561 usage: None,
4562 metadata: MetadataMap::new(),
4563 })))
4564 }
4565 }
4566 }
4567
4568 #[async_trait]
4569 impl ModelTurn for MultiToolTurn {
4570 async fn next_event(
4571 &mut self,
4572 _cancellation: Option<TurnCancellation>,
4573 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4574 Ok(self.events.pop_front())
4575 }
4576 }
4577
4578 #[async_trait]
4579 impl ModelTurn for DualApprovalTurn {
4580 async fn next_event(
4581 &mut self,
4582 _cancellation: Option<TurnCancellation>,
4583 ) -> Result<Option<ModelTurnEvent>, LoopError> {
4584 Ok(self.events.pop_front())
4585 }
4586 }
4587
4588 #[derive(Clone)]
4589 struct EchoTool {
4590 spec: ToolSpec,
4591 }
4592
4593 #[derive(Clone)]
4594 struct FailingTool {
4595 spec: ToolSpec,
4596 }
4597
4598 #[derive(Clone)]
4599 struct RunThenDenyTool {
4600 spec: ToolSpec,
4601 }
4602
4603 impl Default for EchoTool {
4604 fn default() -> Self {
4605 Self {
4606 spec: ToolSpec {
4607 name: ToolName::new("echo"),
4608 description: "Echo back a value".into(),
4609 input_schema: json!({
4610 "type": "object",
4611 "properties": {
4612 "value": { "type": "string" }
4613 },
4614 "required": ["value"],
4615 "additionalProperties": false
4616 }),
4617 output_schema: None,
4618 annotations: ToolAnnotations::default(),
4619 metadata: MetadataMap::new(),
4620 },
4621 }
4622 }
4623 }
4624
4625 impl Default for FailingTool {
4626 fn default() -> Self {
4627 Self {
4628 spec: ToolSpec {
4629 name: ToolName::new("failing"),
4630 description: "Always fails after execution starts".into(),
4631 input_schema: json!({
4632 "type": "object",
4633 "properties": {
4634 "value": { "type": "string" }
4635 },
4636 "additionalProperties": true
4637 }),
4638 output_schema: None,
4639 annotations: ToolAnnotations::default(),
4640 metadata: MetadataMap::new(),
4641 },
4642 }
4643 }
4644 }
4645
4646 impl Default for RunThenDenyTool {
4647 fn default() -> Self {
4648 Self {
4649 spec: ToolSpec {
4650 name: ToolName::new("run_then_deny"),
4651 description: "Runs, then returns a permission-denied error".into(),
4652 input_schema: json!({
4653 "type": "object",
4654 "properties": {
4655 "value": { "type": "string" }
4656 },
4657 "additionalProperties": true
4658 }),
4659 output_schema: None,
4660 annotations: ToolAnnotations::default(),
4661 metadata: MetadataMap::new(),
4662 },
4663 }
4664 }
4665 }
4666
4667 #[derive(Clone)]
4668 struct DynamicSpecTool {
4669 spec: ToolSpec,
4670 version: StdArc<AtomicUsize>,
4671 }
4672
4673 impl DynamicSpecTool {
4674 fn new(version: StdArc<AtomicUsize>) -> Self {
4675 Self {
4676 spec: ToolSpec {
4677 name: ToolName::new("dynamic"),
4678 description: "dynamic version 0".into(),
4679 input_schema: json!({
4680 "type": "object",
4681 "properties": {},
4682 "additionalProperties": false
4683 }),
4684 output_schema: None,
4685 annotations: ToolAnnotations::default(),
4686 metadata: MetadataMap::new(),
4687 },
4688 version,
4689 }
4690 }
4691 }
4692
4693 #[async_trait]
4694 impl Tool for EchoTool {
4695 fn spec(&self) -> &ToolSpec {
4696 &self.spec
4697 }
4698
4699 fn proposed_requests(
4700 &self,
4701 request: &agentkit_tools_core::ToolRequest,
4702 ) -> Result<
4703 Vec<Box<dyn agentkit_tools_core::PermissionRequest>>,
4704 agentkit_tools_core::ToolError,
4705 > {
4706 Ok(vec![Box::new(FileSystemPermissionRequest::Read {
4707 path: "/tmp/echo".into(),
4708 metadata: request.metadata.clone(),
4709 })])
4710 }
4711
4712 async fn invoke(
4713 &self,
4714 request: agentkit_tools_core::ToolRequest,
4715 _ctx: &mut ToolContext<'_>,
4716 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
4717 let value = request
4718 .input
4719 .get("value")
4720 .and_then(Value::as_str)
4721 .ok_or_else(|| {
4722 agentkit_tools_core::ToolError::InvalidInput("missing value".into())
4723 })?;
4724
4725 Ok(ToolResult {
4726 result: ToolResultPart {
4727 call_id: request.call_id,
4728 output: ToolOutput::Text(value.into()),
4729 is_error: false,
4730 metadata: MetadataMap::new(),
4731 },
4732 duration: None,
4733 metadata: MetadataMap::new(),
4734 })
4735 }
4736 }
4737
4738 #[async_trait]
4739 impl Tool for FailingTool {
4740 fn spec(&self) -> &ToolSpec {
4741 &self.spec
4742 }
4743
4744 async fn invoke(
4745 &self,
4746 _request: agentkit_tools_core::ToolRequest,
4747 _ctx: &mut ToolContext<'_>,
4748 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
4749 Err(agentkit_tools_core::ToolError::ExecutionFailed(
4750 "runtime failed".into(),
4751 ))
4752 }
4753 }
4754
4755 #[async_trait]
4756 impl Tool for RunThenDenyTool {
4757 fn spec(&self) -> &ToolSpec {
4758 &self.spec
4759 }
4760
4761 async fn invoke(
4762 &self,
4763 _request: agentkit_tools_core::ToolRequest,
4764 _ctx: &mut ToolContext<'_>,
4765 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
4766 Err(agentkit_tools_core::ToolError::PermissionDenied(
4767 PermissionDenial {
4768 code: PermissionCode::CustomPolicyDenied,
4769 message: "remote 403".into(),
4770 metadata: MetadataMap::new(),
4771 },
4772 ))
4773 }
4774 }
4775
4776 #[async_trait]
4777 impl Tool for DynamicSpecTool {
4778 fn spec(&self) -> &ToolSpec {
4779 &self.spec
4780 }
4781
4782 fn current_spec(&self) -> Option<ToolSpec> {
4783 let mut spec = self.spec.clone();
4784 spec.description = format!("dynamic version {}", self.version.load(Ordering::SeqCst));
4785 Some(spec)
4786 }
4787
4788 async fn invoke(
4789 &self,
4790 request: agentkit_tools_core::ToolRequest,
4791 _ctx: &mut ToolContext<'_>,
4792 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
4793 Ok(ToolResult {
4794 result: ToolResultPart {
4795 call_id: request.call_id,
4796 output: ToolOutput::Text("ok".into()),
4797 is_error: false,
4798 metadata: MetadataMap::new(),
4799 },
4800 duration: None,
4801 metadata: MetadataMap::new(),
4802 })
4803 }
4804 }
4805
4806 struct DenyFsReads;
4807
4808 impl PermissionChecker for DenyFsReads {
4809 fn evaluate(
4810 &self,
4811 request: &dyn agentkit_tools_core::PermissionRequest,
4812 ) -> PermissionDecision {
4813 if request.kind() == "filesystem.read" {
4814 return PermissionDecision::Deny(PermissionDenial {
4815 code: PermissionCode::PathNotAllowed,
4816 message: "reads denied in test".into(),
4817 metadata: MetadataMap::new(),
4818 });
4819 }
4820
4821 PermissionDecision::Allow
4822 }
4823 }
4824
4825 struct ApproveFsReads;
4826
4827 impl PermissionChecker for ApproveFsReads {
4828 fn evaluate(
4829 &self,
4830 request: &dyn agentkit_tools_core::PermissionRequest,
4831 ) -> PermissionDecision {
4832 if request.kind() == "filesystem.read" {
4833 return PermissionDecision::RequireApproval(ApprovalRequest {
4834 task_id: None,
4835 call_id: None,
4836 id: "approval:fs-read".into(),
4837 request_kind: request.kind().into(),
4838 reason: agentkit_tools_core::ApprovalReason::SensitivePath,
4839 summary: request.summary(),
4840 metadata: request.metadata().clone(),
4841 });
4842 }
4843
4844 PermissionDecision::Allow
4845 }
4846 }
4847
4848 struct KeepRecentMutator {
4849 keep: usize,
4850 }
4851
4852 #[async_trait]
4853 impl LoopMutator for KeepRecentMutator {
4854 async fn mutate(
4855 &self,
4856 cursor: &mut TranscriptCursor<'_>,
4857 ctx: LoopCtx<'_>,
4858 ) -> Result<(), LoopError> {
4859 if cursor.len() < 2 {
4860 return Ok(());
4861 }
4862 let drop = cursor.len().saturating_sub(self.keep);
4863 ctx.emitter.emit(AgentEvent::MutationStarted {
4864 session_id: ctx.session_id.clone(),
4865 turn_id: ctx.turn_id.cloned(),
4866 mutator: "keep-recent".into(),
4867 point: ctx.point,
4868 });
4869 cursor.drain(..drop);
4870 ctx.emitter.emit(AgentEvent::MutationFinished {
4871 session_id: ctx.session_id.clone(),
4872 turn_id: ctx.turn_id.cloned(),
4873 mutator: "keep-recent".into(),
4874 dirty: true,
4875 metadata: MetadataMap::new(),
4876 });
4877 Ok(())
4878 }
4879 }
4880
4881 struct PointRecordingMutator {
4884 points: StdArc<StdMutex<Vec<MutationPoint>>>,
4885 }
4886
4887 #[async_trait]
4888 impl LoopMutator for PointRecordingMutator {
4889 async fn mutate(
4890 &self,
4891 _cursor: &mut TranscriptCursor<'_>,
4892 ctx: LoopCtx<'_>,
4893 ) -> Result<(), LoopError> {
4894 self.points.lock().unwrap().push(ctx.point);
4895 Ok(())
4896 }
4897 }
4898
4899 struct RecordingObserver {
4900 events: StdArc<StdMutex<Vec<AgentEvent>>>,
4901 }
4902
4903 impl LoopObserver for RecordingObserver {
4904 fn handle_event(&self, event: ObservedEvent) {
4905 let event = event.event;
4906 self.events.lock().unwrap().push(event);
4907 }
4908 }
4909
4910 struct CatalogExecutor {
4911 version: AtomicUsize,
4912 events: StdMutex<Vec<ToolCatalogEvent>>,
4913 }
4914
4915 impl CatalogExecutor {
4916 fn new() -> Self {
4917 Self {
4918 version: AtomicUsize::new(0),
4919 events: StdMutex::new(Vec::new()),
4920 }
4921 }
4922
4923 fn publish_change(&self, version: usize, event: ToolCatalogEvent) {
4924 self.version.store(version, Ordering::SeqCst);
4925 self.events.lock().unwrap().push(event);
4926 }
4927 }
4928
4929 #[async_trait]
4930 impl ToolExecutor for CatalogExecutor {
4931 fn specs(&self) -> Vec<ToolSpec> {
4932 vec![ToolSpec {
4933 name: ToolName::new("dynamic"),
4934 description: format!("dynamic version {}", self.version.load(Ordering::SeqCst)),
4935 input_schema: json!({
4936 "type": "object",
4937 "properties": {},
4938 "additionalProperties": false
4939 }),
4940 output_schema: None,
4941 annotations: ToolAnnotations::default(),
4942 metadata: MetadataMap::new(),
4943 }]
4944 }
4945
4946 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
4947 std::mem::take(&mut *self.events.lock().unwrap())
4948 }
4949
4950 async fn execute(
4951 &self,
4952 request: ToolRequest,
4953 _ctx: &mut ToolContext<'_>,
4954 ) -> ToolExecutionOutcome {
4955 ToolExecutionOutcome::Completed(ToolResult {
4956 result: ToolResultPart {
4957 call_id: request.call_id,
4958 output: ToolOutput::Text("dynamic-ok".into()),
4959 is_error: false,
4960 metadata: MetadataMap::new(),
4961 },
4962 duration: None,
4963 metadata: MetadataMap::new(),
4964 })
4965 }
4966 }
4967
4968 #[derive(Clone)]
4969 struct BlockingTool {
4970 spec: ToolSpec,
4971 entered: StdArc<AtomicBool>,
4972 release: StdArc<Notify>,
4973 output: &'static str,
4974 }
4975
4976 impl BlockingTool {
4977 fn new(
4978 name: &str,
4979 entered: StdArc<AtomicBool>,
4980 release: StdArc<Notify>,
4981 output: &'static str,
4982 ) -> Self {
4983 Self {
4984 spec: ToolSpec {
4985 name: ToolName::new(name),
4986 description: format!("blocking tool {name}"),
4987 input_schema: json!({
4988 "type": "object",
4989 "properties": {},
4990 "additionalProperties": false
4991 }),
4992 output_schema: None,
4993 annotations: ToolAnnotations::default(),
4994 metadata: MetadataMap::new(),
4995 },
4996 entered,
4997 release,
4998 output,
4999 }
5000 }
5001 }
5002
5003 #[async_trait]
5004 impl Tool for BlockingTool {
5005 fn spec(&self) -> &ToolSpec {
5006 &self.spec
5007 }
5008
5009 async fn invoke(
5010 &self,
5011 request: agentkit_tools_core::ToolRequest,
5012 _ctx: &mut ToolContext<'_>,
5013 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
5014 self.entered.store(true, Ordering::SeqCst);
5015 self.release.notified().await;
5016 Ok(ToolResult {
5017 result: ToolResultPart {
5018 call_id: request.call_id,
5019 output: ToolOutput::Text(self.output.into()),
5020 is_error: false,
5021 metadata: MetadataMap::new(),
5022 },
5023 duration: None,
5024 metadata: MetadataMap::new(),
5025 })
5026 }
5027 }
5028
5029 struct NameRoutingPolicy {
5030 routes: Vec<(String, RoutingDecision)>,
5031 }
5032
5033 impl NameRoutingPolicy {
5034 fn new(routes: impl IntoIterator<Item = (impl Into<String>, RoutingDecision)>) -> Self {
5035 Self {
5036 routes: routes
5037 .into_iter()
5038 .map(|(name, decision)| (name.into(), decision))
5039 .collect(),
5040 }
5041 }
5042 }
5043
5044 impl TaskRoutingPolicy for NameRoutingPolicy {
5045 fn route(&self, request: &ToolRequest) -> RoutingDecision {
5046 self.routes
5047 .iter()
5048 .find(|(name, _)| name == &request.tool_name.0)
5049 .map(|(_, decision)| *decision)
5050 .unwrap_or(RoutingDecision::Foreground)
5051 }
5052 }
5053
5054 async fn wait_for_task_event(handle: &TaskManagerHandle) -> TaskEvent {
5055 timeout(Duration::from_secs(1), handle.next_event())
5056 .await
5057 .expect("timed out waiting for task event")
5058 .expect("task event stream ended unexpectedly")
5059 }
5060
5061 async fn wait_until_entered(flag: &AtomicBool) {
5062 timeout(Duration::from_secs(1), async {
5063 while !flag.load(Ordering::SeqCst) {
5064 tokio::task::yield_now().await;
5065 }
5066 })
5067 .await
5068 .expect("task never entered execution");
5069 }
5070
5071 async fn wait_until_completed(handle: &TaskManagerHandle) {
5072 timeout(Duration::from_secs(1), async {
5073 while handle.list_completed().await.is_empty() {
5074 tokio::task::yield_now().await;
5075 }
5076 })
5077 .await
5078 .expect("task never completed");
5079 }
5080
5081 #[tokio::test]
5082 async fn loop_continues_after_completed_tool_call() {
5083 let tools = ToolRegistry::new().with(EchoTool::default());
5084 let agent = Agent::builder()
5085 .model(FakeAdapter)
5086 .add_tool_source(tools)
5087 .permissions(AllowAllPermissions)
5088 .build()
5089 .unwrap();
5090
5091 let mut driver = agent
5092 .start(SessionConfig {
5093 session_id: SessionId::new("session-1"),
5094 metadata: MetadataMap::new(),
5095 cache: None,
5096 })
5097 .await
5098 .unwrap();
5099
5100 driver
5101 .submit_input(vec![Item {
5102 id: None,
5103 kind: ItemKind::User,
5104 parts: vec![Part::Text(TextPart {
5105 text: "ping".into(),
5106 metadata: MetadataMap::new(),
5107 })],
5108 metadata: MetadataMap::new(),
5109 usage: None,
5110 finish_reason: None,
5111 created_at: None,
5112 }])
5113 .unwrap();
5114
5115 let result = run_until_finished(&mut driver).await;
5116
5117 match result {
5118 LoopStep::Finished(turn) => {
5119 assert_eq!(turn.finish_reason, FinishReason::Completed);
5120 assert_eq!(turn.items.len(), 1);
5121 match &turn.items[0].parts[0] {
5122 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
5123 other => panic!("unexpected part: {other:?}"),
5124 }
5125 }
5126 other => panic!("unexpected loop step: {other:?}"),
5127 }
5128 }
5129
5130 async fn run_until_finished<S: ModelSession + Send>(driver: &mut LoopDriver<S>) -> LoopStep {
5134 loop {
5135 match driver.next().await.unwrap() {
5136 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue,
5137 step => return step,
5138 }
5139 }
5140 }
5141
5142 #[tokio::test]
5149 async fn post_tool_continuation_reports_after_tool_result_mutation_point() {
5150 let points = StdArc::new(StdMutex::new(Vec::<MutationPoint>::new()));
5151 let tools = ToolRegistry::new().with(EchoTool::default());
5152 let agent = Agent::builder()
5153 .model(FakeAdapter)
5154 .add_tool_source(tools)
5155 .permissions(AllowAllPermissions)
5156 .mutator(PointRecordingMutator {
5157 points: points.clone(),
5158 })
5159 .build()
5160 .unwrap();
5161
5162 let mut driver = agent
5163 .start(SessionConfig {
5164 session_id: SessionId::new("session-mutation-point"),
5165 metadata: MetadataMap::new(),
5166 cache: None,
5167 })
5168 .await
5169 .unwrap();
5170
5171 driver
5172 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5173 .unwrap();
5174
5175 let _ = run_until_finished(&mut driver).await;
5177
5178 let recorded = points.lock().unwrap().clone();
5179 assert_eq!(
5180 recorded.first(),
5181 Some(&MutationPoint::AfterTurnEnded),
5182 "first drive of a fresh turn must report AfterTurnEnded, got {recorded:?}"
5183 );
5184 assert!(
5185 recorded.contains(&MutationPoint::AfterToolResult),
5186 "post-tool continuation must report AfterToolResult, got {recorded:?}"
5187 );
5188 }
5189
5190 #[test]
5191 fn pending_input_requires_input_bearing_tail_role() {
5192 assert!(!transcript_has_pending_input(&[]));
5193 assert!(!transcript_has_pending_input(&[Item::text(
5194 ItemKind::System,
5195 "system"
5196 )]));
5197 assert!(!transcript_has_pending_input(&[Item::text(
5198 ItemKind::Developer,
5199 "developer"
5200 )]));
5201 assert!(!transcript_has_pending_input(&[Item::text(
5202 ItemKind::Context,
5203 "context"
5204 )]));
5205 assert!(!transcript_has_pending_input(&[Item::text(
5206 ItemKind::Assistant,
5207 "assistant"
5208 )]));
5209
5210 assert!(transcript_has_pending_input(&[Item::text(
5211 ItemKind::User,
5212 "user"
5213 )]));
5214 assert!(transcript_has_pending_input(&[Item::notification(
5215 "background update"
5216 )]));
5217 assert!(transcript_has_pending_input(&[Item {
5218 id: None,
5219 kind: ItemKind::Tool,
5220 parts: vec![Part::ToolResult(ToolResultPart {
5221 call_id: ToolCallId::new("call-test"),
5222 output: ToolOutput::Text("ok".into()),
5223 is_error: false,
5224 metadata: MetadataMap::new(),
5225 })],
5226 metadata: MetadataMap::new(),
5227 usage: None,
5228 finish_reason: None,
5229 created_at: None,
5230 }]));
5231 }
5232
5233 struct DropTrailingUserMutator;
5239
5240 #[async_trait]
5241 impl LoopMutator for DropTrailingUserMutator {
5242 async fn mutate(
5243 &self,
5244 cursor: &mut TranscriptCursor<'_>,
5245 _ctx: LoopCtx<'_>,
5246 ) -> Result<(), LoopError> {
5247 if cursor.last().map(|item| item.kind) == Some(ItemKind::User) {
5248 cursor.pop();
5249 }
5250 Ok(())
5251 }
5252 }
5253
5254 struct RejectAssistantPrefillAdapter {
5259 saw_assistant_tail: StdArc<AtomicBool>,
5260 }
5261
5262 struct RejectAssistantPrefillSession {
5263 saw_assistant_tail: StdArc<AtomicBool>,
5264 }
5265
5266 #[async_trait]
5267 impl ModelAdapter for RejectAssistantPrefillAdapter {
5268 type Session = RejectAssistantPrefillSession;
5269
5270 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
5271 Ok(RejectAssistantPrefillSession {
5272 saw_assistant_tail: self.saw_assistant_tail.clone(),
5273 })
5274 }
5275 }
5276
5277 #[async_trait]
5278 impl ModelSession for RejectAssistantPrefillSession {
5279 type Turn = FakeTurn;
5280
5281 async fn begin_turn(
5282 &mut self,
5283 request: TurnRequest,
5284 _cancellation: Option<TurnCancellation>,
5285 ) -> Result<Self::Turn, LoopError> {
5286 if request.transcript.last().map(|item| item.kind) == Some(ItemKind::Assistant) {
5287 self.saw_assistant_tail.store(true, Ordering::SeqCst);
5288 return Err(LoopError::Provider(
5289 "conversation must end with a user message".into(),
5290 ));
5291 }
5292 Ok(FakeTurn {
5293 events: VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
5294 model: None,
5295 response_id: None,
5296 finish_reason: FinishReason::Completed,
5297 output_items: vec![Item::text(ItemKind::Assistant, "ok")],
5298 usage: None,
5299 metadata: MetadataMap::new(),
5300 })]),
5301 })
5302 }
5303 }
5304
5305 #[tokio::test]
5313 async fn drive_does_not_dispatch_without_valid_trailing_input() {
5314 let saw_assistant_tail = StdArc::new(AtomicBool::new(false));
5315 let agent = Agent::builder()
5316 .model(RejectAssistantPrefillAdapter {
5317 saw_assistant_tail: saw_assistant_tail.clone(),
5318 })
5319 .mutator(DropTrailingUserMutator)
5320 .transcript(vec![
5323 Item::text(ItemKind::User, "kickoff"),
5324 Item::text(ItemKind::Assistant, "prior reply"),
5325 ])
5326 .build()
5327 .unwrap();
5328
5329 let mut driver = agent
5330 .start(SessionConfig {
5331 session_id: SessionId::new("session-no-valid-input"),
5332 metadata: MetadataMap::new(),
5333 cache: None,
5334 })
5335 .await
5336 .unwrap();
5337
5338 driver
5339 .submit_input(vec![Item::text(ItemKind::User, "follow up")])
5340 .unwrap();
5341
5342 let outcome = driver.next().await;
5344
5345 assert!(
5346 !saw_assistant_tail.load(Ordering::SeqCst),
5347 "loop dispatched a model turn whose transcript ends in an assistant \
5348 message (outcome: {outcome:?}); with no valid trailing input the turn \
5349 must finish instead of driving"
5350 );
5351 }
5352
5353 #[tokio::test]
5354 async fn loop_uses_injected_permission_checker() {
5355 let events = StdArc::new(StdMutex::new(Vec::new()));
5356 let tools = ToolRegistry::new().with(EchoTool::default());
5357 let agent = Agent::builder()
5358 .model(FakeAdapter)
5359 .add_tool_source(tools)
5360 .permissions(DenyFsReads)
5361 .observer(RecordingObserver {
5362 events: events.clone(),
5363 })
5364 .build()
5365 .unwrap();
5366
5367 let mut driver = agent
5368 .start(SessionConfig {
5369 session_id: SessionId::new("session-2"),
5370 metadata: MetadataMap::new(),
5371 cache: None,
5372 })
5373 .await
5374 .unwrap();
5375
5376 driver
5377 .submit_input(vec![Item {
5378 id: None,
5379 kind: ItemKind::User,
5380 parts: vec![Part::Text(TextPart {
5381 text: "ping".into(),
5382 metadata: MetadataMap::new(),
5383 })],
5384 metadata: MetadataMap::new(),
5385 usage: None,
5386 finish_reason: None,
5387 created_at: None,
5388 }])
5389 .unwrap();
5390
5391 let result = run_until_finished(&mut driver).await;
5392
5393 match result {
5394 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
5395 Part::Text(text) => assert!(text.text.contains("tool permission denied")),
5396 other => panic!("unexpected part: {other:?}"),
5397 },
5398 other => panic!("unexpected loop step: {other:?}"),
5399 }
5400
5401 assert!(
5402 events
5403 .lock()
5404 .unwrap()
5405 .iter()
5406 .all(|event| !matches!(event, AgentEvent::ToolExecutionStarted(_))),
5407 "denied tools must not be reported as started"
5408 );
5409 }
5410
5411 #[tokio::test]
5412 async fn failed_tool_execution_still_reports_started() {
5413 let events = StdArc::new(StdMutex::new(Vec::new()));
5414 let tools = ToolRegistry::new().with(FailingTool::default());
5415 let agent = Agent::builder()
5416 .model(FakeAdapter)
5417 .add_tool_source(tools)
5418 .permissions(AllowAllPermissions)
5419 .observer(RecordingObserver {
5420 events: events.clone(),
5421 })
5422 .build()
5423 .unwrap();
5424
5425 let mut driver = agent
5426 .start(SessionConfig {
5427 session_id: SessionId::new("session-failing-start-event"),
5428 metadata: MetadataMap::new(),
5429 cache: None,
5430 })
5431 .await
5432 .unwrap();
5433
5434 driver
5435 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5436 .unwrap();
5437
5438 match run_until_finished(&mut driver).await {
5439 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
5440 other => panic!("unexpected loop step: {other:?}"),
5441 }
5442
5443 let events = events.lock().unwrap();
5444 assert!(events.iter().any(|event| matches!(
5445 event,
5446 AgentEvent::ToolExecutionStarted(call) if call.name == "failing"
5447 )));
5448 assert!(events.iter().any(|event| matches!(
5449 event,
5450 AgentEvent::ToolResultReceived(result) if result.is_error
5451 )));
5452 }
5453
5454 #[tokio::test]
5455 async fn run_then_deny_tool_execution_still_reports_started() {
5456 let events = StdArc::new(StdMutex::new(Vec::new()));
5457 let tools = ToolRegistry::new().with(RunThenDenyTool::default());
5458 let agent = Agent::builder()
5459 .model(FakeAdapter)
5460 .add_tool_source(tools)
5461 .permissions(AllowAllPermissions)
5462 .observer(RecordingObserver {
5463 events: events.clone(),
5464 })
5465 .build()
5466 .unwrap();
5467
5468 let mut driver = agent
5469 .start(SessionConfig {
5470 session_id: SessionId::new("session-run-then-deny-start-event"),
5471 metadata: MetadataMap::new(),
5472 cache: None,
5473 })
5474 .await
5475 .unwrap();
5476
5477 driver
5478 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5479 .unwrap();
5480
5481 match run_until_finished(&mut driver).await {
5482 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
5483 other => panic!("unexpected loop step: {other:?}"),
5484 }
5485
5486 let events = events.lock().unwrap();
5487 assert!(events.iter().any(|event| matches!(
5488 event,
5489 AgentEvent::ToolExecutionStarted(call) if call.name == "run_then_deny"
5490 )));
5491 assert!(events.iter().any(|event| matches!(
5494 event,
5495 AgentEvent::ToolResultReceived(result)
5496 if result.is_error
5497 && result
5498 .metadata
5499 .get(TOOL_RESULT_FAILURE_KIND_METADATA_KEY)
5500 .and_then(Value::as_str)
5501 == Some(TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED)
5502 && result
5503 .metadata
5504 .get(TOOL_RESULT_NOT_STARTED_METADATA_KEY)
5505 .is_none()
5506 )));
5507 }
5508
5509 #[tokio::test]
5510 async fn async_task_manager_background_round_requires_explicit_continue() {
5511 let events = StdArc::new(StdMutex::new(Vec::new()));
5512 let entered = StdArc::new(AtomicBool::new(false));
5513 let release = StdArc::new(Notify::new());
5514 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
5515 "background-wait",
5516 RoutingDecision::Background,
5517 )]));
5518 let handle = task_manager.handle();
5519 let tools = ToolRegistry::new().with(BlockingTool::new(
5520 "background-wait",
5521 entered.clone(),
5522 release.clone(),
5523 "background-done",
5524 ));
5525 let agent = Agent::builder()
5526 .model(FakeAdapter)
5527 .add_tool_source(tools)
5528 .permissions(AllowAllPermissions)
5529 .task_manager(task_manager)
5530 .observer(RecordingObserver {
5531 events: events.clone(),
5532 })
5533 .build()
5534 .unwrap();
5535
5536 let mut driver = agent
5537 .start(SessionConfig {
5538 session_id: SessionId::new("session-background"),
5539 metadata: MetadataMap::new(),
5540 cache: None,
5541 })
5542 .await
5543 .unwrap();
5544
5545 driver
5546 .submit_input(vec![Item {
5547 id: None,
5548 kind: ItemKind::User,
5549 parts: vec![Part::Text(TextPart {
5550 text: "ping".into(),
5551 metadata: MetadataMap::new(),
5552 })],
5553 metadata: MetadataMap::new(),
5554 usage: None,
5555 finish_reason: None,
5556 created_at: None,
5557 }])
5558 .unwrap();
5559
5560 let first = driver.next().await.unwrap();
5561 match first {
5562 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
5563 other => panic!("unexpected first loop step: {other:?}"),
5564 }
5565
5566 match wait_for_task_event(&handle).await {
5567 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "background-wait"),
5568 other => panic!("unexpected task event: {other:?}"),
5569 }
5570 wait_until_entered(entered.as_ref()).await;
5571 release.notify_waiters();
5572
5573 match wait_for_task_event(&handle).await {
5574 TaskEvent::Completed(_, result) => {
5575 assert_eq!(result.output, ToolOutput::Text("background-done".into()))
5576 }
5577 other => panic!("unexpected completion event: {other:?}"),
5578 }
5579
5580 let resumed = driver.next().await.unwrap();
5581 match resumed {
5582 LoopStep::Finished(turn) => {
5583 assert_eq!(turn.finish_reason, FinishReason::Completed);
5584 match &turn.items[0].parts[0] {
5585 Part::Text(text) => assert_eq!(text.text, "tool said: background-done"),
5586 other => panic!("unexpected part after resume: {other:?}"),
5587 }
5588 }
5589 other => panic!("unexpected resumed step: {other:?}"),
5590 }
5591
5592 let events = events.lock().unwrap();
5593 let terminal_results: Vec<_> = events
5594 .iter()
5595 .filter_map(|event| match event {
5596 AgentEvent::ToolResultReceived(result)
5597 if result.call_id == ToolCallId::new("call-1") =>
5598 {
5599 Some(result)
5600 }
5601 _ => None,
5602 })
5603 .collect();
5604 assert_eq!(
5605 terminal_results.len(),
5606 1,
5607 "background completion must emit one terminal result event per call: {events:?}"
5608 );
5609 }
5610
5611 #[tokio::test]
5612 async fn detached_tool_placeholder_is_progress_not_terminal_result() {
5613 let events = StdArc::new(StdMutex::new(Vec::new()));
5614 let entered = StdArc::new(AtomicBool::new(false));
5615 let release = StdArc::new(Notify::new());
5616 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
5617 "detaching-wait",
5618 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10)),
5619 )]));
5620 let handle = task_manager.handle();
5621 let tools = ToolRegistry::new().with(BlockingTool::new(
5622 "detaching-wait",
5623 entered.clone(),
5624 release.clone(),
5625 "detached-done",
5626 ));
5627 let agent = Agent::builder()
5628 .model(FakeAdapter)
5629 .add_tool_source(tools)
5630 .permissions(AllowAllPermissions)
5631 .task_manager(task_manager)
5632 .observer(RecordingObserver {
5633 events: events.clone(),
5634 })
5635 .build()
5636 .unwrap();
5637
5638 let mut driver = agent
5639 .start(SessionConfig {
5640 session_id: SessionId::new("session-detached-progress"),
5641 metadata: MetadataMap::new(),
5642 cache: None,
5643 })
5644 .await
5645 .unwrap();
5646
5647 driver
5648 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5649 .unwrap();
5650
5651 match driver.next().await.unwrap() {
5652 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => {}
5653 other => panic!("unexpected detach step: {other:?}"),
5654 }
5655
5656 match wait_for_task_event(&handle).await {
5657 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "detaching-wait"),
5658 other => panic!("unexpected task event: {other:?}"),
5659 }
5660 match wait_for_task_event(&handle).await {
5661 TaskEvent::Detached(snapshot) => assert_eq!(snapshot.tool_name, "detaching-wait"),
5662 other => panic!("unexpected detach event: {other:?}"),
5663 }
5664 wait_until_entered(entered.as_ref()).await;
5665 release.notify_waiters();
5666
5667 match wait_for_task_event(&handle).await {
5668 TaskEvent::Completed(_, result) => {
5669 assert_eq!(result.output, ToolOutput::Text("detached-done".into()))
5670 }
5671 other => panic!("unexpected completion event: {other:?}"),
5672 }
5673
5674 match driver.next().await.unwrap() {
5675 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
5676 other => panic!("unexpected resumed step: {other:?}"),
5677 }
5678
5679 let events = events.lock().unwrap();
5680 assert!(events.iter().any(|event| matches!(
5681 event,
5682 AgentEvent::ToolExecutionProgress(result)
5683 if result.call_id == ToolCallId::new("call-1") && !result.is_error
5684 )));
5685 let terminal_results: Vec<_> = events
5686 .iter()
5687 .filter_map(|event| match event {
5688 AgentEvent::ToolResultReceived(result)
5689 if result.call_id == ToolCallId::new("call-1") =>
5690 {
5691 Some(result)
5692 }
5693 _ => None,
5694 })
5695 .collect();
5696 assert_eq!(
5697 terminal_results.len(),
5698 1,
5699 "detached call must emit one terminal result event: {events:?}"
5700 );
5701 }
5702
5703 #[tokio::test]
5704 async fn cancelled_background_approval_auto_resolves_when_drained() {
5705 let controller = CancellationController::new();
5706 let events = StdArc::new(StdMutex::new(Vec::new()));
5707 let entered = StdArc::new(AtomicBool::new(false));
5708 let release = StdArc::new(Notify::new());
5709 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
5710 "echo",
5711 RoutingDecision::Background,
5712 )]));
5713 let handle = task_manager.handle();
5714 let agent = Agent::builder()
5715 .model(FakeAdapter)
5716 .tool_executor(DelayedApprovalExecutor::new(
5717 entered.clone(),
5718 release.clone(),
5719 ))
5720 .task_manager(task_manager)
5721 .cancellation(controller.handle())
5722 .observer(RecordingObserver {
5723 events: events.clone(),
5724 })
5725 .build()
5726 .unwrap();
5727
5728 let mut driver = agent
5729 .start(SessionConfig {
5730 session_id: SessionId::new("session-cancel-delayed-background-approval"),
5731 metadata: MetadataMap::new(),
5732 cache: None,
5733 })
5734 .await
5735 .unwrap();
5736
5737 driver
5738 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5739 .unwrap();
5740
5741 match driver.next().await.unwrap() {
5742 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
5743 other => panic!("unexpected first step: {other:?}"),
5744 }
5745
5746 match wait_for_task_event(&handle).await {
5747 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "echo"),
5748 other => panic!("unexpected task event: {other:?}"),
5749 }
5750
5751 wait_until_entered(entered.as_ref()).await;
5752 controller.interrupt();
5753 release.notify_waiters();
5754 wait_until_completed(&handle).await;
5755
5756 match driver.next().await.unwrap() {
5757 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
5758 other => panic!("cancelled background approval should finish cancelled, got {other:?}"),
5759 }
5760
5761 let events = events.lock().unwrap();
5762 assert!(
5763 events
5764 .iter()
5765 .any(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false }))
5766 );
5767 assert!(events.iter().any(|event| matches!(
5768 event,
5769 AgentEvent::ToolResultReceived(result)
5770 if result.call_id == ToolCallId::new("call-1") && result.is_error
5771 )));
5772 }
5773
5774 #[tokio::test]
5775 async fn loop_can_cancel_a_turn_and_continue_after_new_input() {
5776 let controller = CancellationController::new();
5777 let agent = Agent::builder()
5778 .model(SlowAdapter)
5779 .cancellation(controller.handle())
5780 .build()
5781 .unwrap();
5782
5783 let mut driver = agent
5784 .start(SessionConfig {
5785 session_id: SessionId::new("session-cancel"),
5786 metadata: MetadataMap::new(),
5787 cache: None,
5788 })
5789 .await
5790 .unwrap();
5791
5792 driver
5793 .submit_input(vec![Item {
5794 id: None,
5795 kind: ItemKind::User,
5796 parts: vec![Part::Text(TextPart {
5797 text: "do the long task".into(),
5798 metadata: MetadataMap::new(),
5799 })],
5800 metadata: MetadataMap::new(),
5801 usage: None,
5802 finish_reason: None,
5803 created_at: None,
5804 }])
5805 .unwrap();
5806
5807 let cancelled = tokio::join!(async { driver.next().await }, async {
5808 tokio::task::yield_now().await;
5809 controller.interrupt();
5810 })
5811 .0
5812 .unwrap();
5813
5814 match cancelled {
5815 LoopStep::Finished(turn) => {
5816 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
5817 assert_eq!(turn.items.len(), 1);
5818 assert_eq!(turn.items[0].kind, ItemKind::Assistant);
5819 assert_eq!(
5820 turn.items[0].metadata.get(INTERRUPTED_METADATA_KEY),
5821 Some(&Value::Bool(true))
5822 );
5823 }
5824 other => panic!("unexpected loop step: {other:?}"),
5825 }
5826
5827 driver
5828 .submit_input(vec![Item {
5829 id: None,
5830 kind: ItemKind::User,
5831 parts: vec![Part::Text(TextPart {
5832 text: "try again".into(),
5833 metadata: MetadataMap::new(),
5834 })],
5835 metadata: MetadataMap::new(),
5836 usage: None,
5837 finish_reason: None,
5838 created_at: None,
5839 }])
5840 .unwrap();
5841
5842 let result = driver.next().await.unwrap();
5843 match result {
5844 LoopStep::Finished(turn) => {
5845 assert_eq!(turn.finish_reason, FinishReason::Completed);
5846 }
5847 other => panic!("unexpected loop step after retry: {other:?}"),
5848 }
5849 }
5850
5851 #[tokio::test]
5852 async fn loop_interrupt_cancels_foreground_tasks_but_keeps_background_tasks_running() {
5853 let controller = CancellationController::new();
5854 let fg_entered = StdArc::new(AtomicBool::new(false));
5855 let fg_release = StdArc::new(Notify::new());
5856 let bg_entered = StdArc::new(AtomicBool::new(false));
5857 let bg_release = StdArc::new(Notify::new());
5858 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([
5859 ("foreground-wait", RoutingDecision::Foreground),
5860 ("background-wait", RoutingDecision::Background),
5861 ]));
5862 let handle = task_manager.handle();
5863 let tools = ToolRegistry::new()
5864 .with(BlockingTool::new(
5865 "foreground-wait",
5866 fg_entered.clone(),
5867 fg_release,
5868 "foreground-done",
5869 ))
5870 .with(BlockingTool::new(
5871 "background-wait",
5872 bg_entered.clone(),
5873 bg_release.clone(),
5874 "background-done",
5875 ));
5876 let agent = Agent::builder()
5877 .model(MultiToolAdapter)
5878 .add_tool_source(tools)
5879 .permissions(AllowAllPermissions)
5880 .cancellation(controller.handle())
5881 .task_manager(task_manager)
5882 .build()
5883 .unwrap();
5884
5885 let mut driver = agent
5886 .start(SessionConfig {
5887 session_id: SessionId::new("session-mixed-cancel"),
5888 metadata: MetadataMap::new(),
5889 cache: None,
5890 })
5891 .await
5892 .unwrap();
5893
5894 driver
5895 .submit_input(vec![Item {
5896 id: None,
5897 kind: ItemKind::User,
5898 parts: vec![Part::Text(TextPart {
5899 text: "run both".into(),
5900 metadata: MetadataMap::new(),
5901 })],
5902 metadata: MetadataMap::new(),
5903 usage: None,
5904 finish_reason: None,
5905 created_at: None,
5906 }])
5907 .unwrap();
5908
5909 let cancelled = tokio::join!(async { driver.next().await }, async {
5910 let _ = wait_for_task_event(&handle).await;
5911 let _ = wait_for_task_event(&handle).await;
5912 wait_until_entered(fg_entered.as_ref()).await;
5913 wait_until_entered(bg_entered.as_ref()).await;
5914 controller.interrupt();
5915 })
5916 .0
5917 .unwrap();
5918
5919 match cancelled {
5920 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
5921 other => panic!("unexpected loop step after interrupt: {other:?}"),
5922 }
5923
5924 match wait_for_task_event(&handle).await {
5925 TaskEvent::Cancelled(snapshot) => assert_eq!(snapshot.tool_name, "foreground-wait"),
5926 other => panic!("unexpected post-interrupt event: {other:?}"),
5927 }
5928
5929 let running = handle.list_running().await;
5930 assert_eq!(running.len(), 1);
5931 assert_eq!(running[0].tool_name, "background-wait");
5932
5933 bg_release.notify_waiters();
5934 match wait_for_task_event(&handle).await {
5935 TaskEvent::Completed(snapshot, result) => {
5936 assert_eq!(snapshot.tool_name, "background-wait");
5937 assert_eq!(result.output, ToolOutput::Text("background-done".into()));
5938 }
5939 other => panic!("unexpected background completion event: {other:?}"),
5940 }
5941 }
5942
5943 #[tokio::test]
5944 async fn a_cancelled_turn_answers_the_tool_call_it_abandoned() {
5945 let controller = CancellationController::new();
5950 let entered = StdArc::new(AtomicBool::new(false));
5951 let release = StdArc::new(Notify::new());
5952 let items = StdArc::new(StdMutex::new(Vec::<Item>::new()));
5953 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
5954 "wait",
5955 RoutingDecision::Foreground,
5956 )]));
5957 let agent = Agent::builder()
5958 .model(FakeAdapter)
5959 .add_tool_source(ToolRegistry::new().with(BlockingTool::new(
5960 "wait",
5961 entered.clone(),
5962 release,
5963 "done",
5964 )))
5965 .permissions(AllowAllPermissions)
5966 .cancellation(controller.handle())
5967 .task_manager(task_manager)
5968 .transcript_observer(RecordingTranscriptObserver {
5969 items: items.clone(),
5970 })
5971 .build()
5972 .unwrap();
5973
5974 let mut driver = agent
5975 .start(SessionConfig {
5976 session_id: SessionId::new("session-cancel-mid-call"),
5977 metadata: MetadataMap::new(),
5978 cache: None,
5979 })
5980 .await
5981 .unwrap();
5982
5983 driver
5984 .submit_input(vec![Item {
5985 id: None,
5986 kind: ItemKind::User,
5987 parts: vec![Part::Text(TextPart {
5988 text: "run the tool".into(),
5989 metadata: MetadataMap::new(),
5990 })],
5991 metadata: MetadataMap::new(),
5992 usage: None,
5993 finish_reason: None,
5994 created_at: None,
5995 }])
5996 .unwrap();
5997
5998 let cancelled = tokio::join!(async { driver.next().await }, async {
5999 wait_until_entered(entered.as_ref()).await;
6000 controller.interrupt();
6001 })
6002 .0
6003 .unwrap();
6004
6005 match cancelled {
6006 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
6007 other => panic!("unexpected loop step after interrupt: {other:?}"),
6008 }
6009
6010 let transcript = driver.snapshot().transcript;
6011 assert!(
6012 unanswered_tool_calls(&transcript).is_empty(),
6013 "the cancelled turn left a tool call unanswered: {transcript:?}"
6014 );
6015 validate_transcript_invariants(&transcript)
6016 .expect("a cancelled turn must leave a resumable transcript");
6017
6018 let persisted = items.lock().unwrap().clone();
6019 let results: Vec<&ToolResultPart> = persisted
6020 .iter()
6021 .flat_map(|item| &item.parts)
6022 .filter_map(|part| match part {
6023 Part::ToolResult(result) => Some(result),
6024 _ => None,
6025 })
6026 .collect();
6027 assert_eq!(results.len(), 1, "{persisted:?}");
6028 assert_eq!(results[0].call_id, ToolCallId::new("call-1"));
6029 assert!(results[0].is_error);
6030 assert_eq!(
6031 results[0].metadata.get(INTERRUPTED_METADATA_KEY),
6032 Some(&Value::Bool(true))
6033 );
6034 }
6035
6036 #[tokio::test]
6037 async fn regression_cancelled_background_completion_emits_one_terminal_result() {
6038 let events = StdArc::new(StdMutex::new(Vec::new()));
6039 let agent = Agent::builder()
6040 .model(FakeAdapter)
6041 .observer(RecordingObserver {
6042 events: events.clone(),
6043 })
6044 .build()
6045 .unwrap();
6046 let mut driver = agent
6047 .start(SessionConfig::new("session-cancelled-background-event"))
6048 .await
6049 .unwrap();
6050
6051 driver.append_item(Item::new(
6052 ItemKind::Assistant,
6053 vec![Part::ToolCall(ToolCallPart {
6054 id: ToolCallId::new("call-1"),
6055 name: "wait".into(),
6056 input: json!({}),
6057 metadata: MetadataMap::new(),
6058 })],
6059 ));
6060 driver.background_call_ids.insert(ToolCallId::new("call-1"));
6061 driver.close_interrupted_tool_calls();
6062 driver.append_tool_result_item(Item::new(
6063 ItemKind::Tool,
6064 vec![Part::ToolResult(ToolResultPart {
6065 call_id: ToolCallId::new("call-1"),
6066 output: ToolOutput::Text("background-done".into()),
6067 is_error: false,
6068 metadata: MetadataMap::new(),
6069 })],
6070 ));
6071
6072 let events = events.lock().unwrap();
6073 let terminal_results = events
6074 .iter()
6075 .filter(|event| {
6076 matches!(
6077 event,
6078 AgentEvent::ToolResultReceived(result)
6079 if result.call_id == ToolCallId::new("call-1")
6080 )
6081 })
6082 .count();
6083 assert_eq!(
6084 terminal_results, 1,
6085 "a cancelled background call emitted multiple terminal results: {events:?}"
6086 );
6087 }
6088
6089 #[tokio::test]
6090 async fn regression_cancelled_queued_approval_is_answered_once() {
6091 let controller = CancellationController::new();
6092 let entered = StdArc::new(AtomicBool::new(false));
6093 let release = StdArc::new(Notify::new());
6094 release.notify_one();
6095 let events = StdArc::new(StdMutex::new(Vec::new()));
6096 let agent = Agent::builder()
6097 .model(FakeAdapter)
6098 .tool_executor(
6099 DelayedApprovalExecutor::new(entered, release)
6100 .cancelling_on_approval(controller.clone()),
6101 )
6102 .cancellation(controller.handle())
6103 .observer(RecordingObserver {
6104 events: events.clone(),
6105 })
6106 .build()
6107 .unwrap();
6108 let mut driver = agent
6109 .start(SessionConfig::new("session-cancelled-queued-approval"))
6110 .await
6111 .unwrap();
6112 driver
6113 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6114 .unwrap();
6115
6116 match driver.next().await.unwrap() {
6117 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
6118 other => panic!("unexpected first cancellation step: {other:?}"),
6119 }
6120 match driver.next().await.unwrap() {
6121 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
6122 other => panic!("unexpected post-cancellation step: {other:?}"),
6123 }
6124
6125 let events = events.lock().unwrap();
6126 let terminal_results = events
6127 .iter()
6128 .filter(|event| {
6129 matches!(
6130 event,
6131 AgentEvent::ToolResultReceived(result)
6132 if result.call_id == ToolCallId::new("call-1")
6133 )
6134 })
6135 .count();
6136 assert_eq!(
6137 terminal_results, 1,
6138 "a cancelled queued approval was answered more than once: {events:?}"
6139 );
6140 drop(events);
6141
6142 let transcript = driver.snapshot().transcript;
6143 assert!(
6144 !transcript.iter().any(|item| {
6145 item.kind == ItemKind::Notification
6146 && item.parts.iter().any(|part| {
6147 matches!(part, Part::Text(text) if text.text.contains("Background tool call"))
6148 })
6149 }),
6150 "a queued approval was misreported as a background call: {transcript:?}"
6151 );
6152 }
6153
6154 #[tokio::test]
6155 async fn regression_cancelled_unstarted_call_is_not_tracked_as_detached() {
6156 let agent = Agent::builder().model(FakeAdapter).build().unwrap();
6157 let mut driver = agent
6158 .start(SessionConfig::new("session-cancelled-unstarted-call"))
6159 .await
6160 .unwrap();
6161
6162 driver.append_item(Item::new(
6163 ItemKind::Assistant,
6164 vec![Part::ToolCall(ToolCallPart {
6165 id: ToolCallId::new("call-never-started"),
6166 name: "wait".into(),
6167 input: json!({}),
6168 metadata: MetadataMap::new(),
6169 })],
6170 ));
6171 driver
6172 .finish_cancelled(agentkit_core::TurnId::new("turn-cancelled"), Vec::new())
6173 .unwrap();
6174
6175 assert!(
6176 !driver
6177 .detached_call_ids
6178 .contains(&ToolCallId::new("call-never-started")),
6179 "an unstarted call can never deliver a detached result"
6180 );
6181 }
6182
6183 #[tokio::test]
6184 async fn loop_resumes_after_approved_tool_request() {
6185 let tools = ToolRegistry::new().with(EchoTool::default());
6186 let agent = Agent::builder()
6187 .model(FakeAdapter)
6188 .add_tool_source(tools)
6189 .permissions(ApproveFsReads)
6190 .build()
6191 .unwrap();
6192
6193 let mut driver = agent
6194 .start(SessionConfig {
6195 session_id: SessionId::new("session-approval"),
6196 metadata: MetadataMap::new(),
6197 cache: None,
6198 })
6199 .await
6200 .unwrap();
6201
6202 driver
6203 .submit_input(vec![Item {
6204 id: None,
6205 kind: ItemKind::User,
6206 parts: vec![Part::Text(TextPart {
6207 text: "ping".into(),
6208 metadata: MetadataMap::new(),
6209 })],
6210 metadata: MetadataMap::new(),
6211 usage: None,
6212 finish_reason: None,
6213 created_at: None,
6214 }])
6215 .unwrap();
6216
6217 let first = driver.next().await.unwrap();
6218 match first {
6219 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
6220 assert!(pending.request.task_id.is_some());
6221 assert_eq!(pending.request.id.0, "approval:fs-read");
6222 pending.approve(&mut driver).unwrap();
6223 }
6224 other => panic!("unexpected loop step: {other:?}"),
6225 }
6226 let second = driver.next().await.unwrap();
6227 match second {
6228 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
6229 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
6230 other => panic!("unexpected part: {other:?}"),
6231 },
6232 other => panic!("unexpected loop step after approval: {other:?}"),
6233 }
6234 }
6235
6236 #[tokio::test]
6237 async fn approval_gated_tool_does_not_start_before_approval() {
6238 let events = StdArc::new(StdMutex::new(Vec::new()));
6239 let tools = ToolRegistry::new().with(EchoTool::default());
6240 let agent = Agent::builder()
6241 .model(FakeAdapter)
6242 .add_tool_source(tools)
6243 .permissions(ApproveFsReads)
6244 .observer(RecordingObserver {
6245 events: events.clone(),
6246 })
6247 .build()
6248 .unwrap();
6249
6250 let mut driver = agent
6251 .start(SessionConfig {
6252 session_id: SessionId::new("session-approval-start-event"),
6253 metadata: MetadataMap::new(),
6254 cache: None,
6255 })
6256 .await
6257 .unwrap();
6258
6259 driver
6260 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6261 .unwrap();
6262
6263 let pending = match driver.next().await.unwrap() {
6264 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
6265 other => panic!("unexpected loop step: {other:?}"),
6266 };
6267
6268 assert!(
6269 events
6270 .lock()
6271 .unwrap()
6272 .iter()
6273 .all(|event| !matches!(event, AgentEvent::ToolExecutionStarted(_))),
6274 "tool start must not be reported before approval"
6275 );
6276
6277 pending.approve(&mut driver).unwrap();
6278 match driver.next().await.unwrap() {
6279 LoopStep::Finished(_) => {}
6280 other => panic!("unexpected loop step after approval: {other:?}"),
6281 }
6282
6283 let started = events
6284 .lock()
6285 .unwrap()
6286 .iter()
6287 .filter(|event| matches!(event, AgentEvent::ToolExecutionStarted(_)))
6288 .count();
6289 assert_eq!(started, 1);
6290 }
6291
6292 #[tokio::test]
6293 async fn cancelling_pending_approval_resolves_it_and_pairs_tool_result() {
6294 let controller = CancellationController::new();
6295 let events = StdArc::new(StdMutex::new(Vec::new()));
6296 let tools = ToolRegistry::new().with(EchoTool::default());
6297 let agent = Agent::builder()
6298 .model(FakeAdapter)
6299 .add_tool_source(tools)
6300 .permissions(ApproveFsReads)
6301 .cancellation(controller.handle())
6302 .observer(RecordingObserver {
6303 events: events.clone(),
6304 })
6305 .build()
6306 .unwrap();
6307
6308 let mut driver = agent
6309 .start(SessionConfig {
6310 session_id: SessionId::new("session-cancel-pending-approval"),
6311 metadata: MetadataMap::new(),
6312 cache: None,
6313 })
6314 .await
6315 .unwrap();
6316
6317 driver
6318 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6319 .unwrap();
6320
6321 match driver.next().await.unwrap() {
6322 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) => {}
6323 other => panic!("unexpected loop step: {other:?}"),
6324 }
6325
6326 controller.interrupt();
6327
6328 match driver.next().await.unwrap() {
6329 LoopStep::Finished(turn) => {
6330 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
6331 }
6332 other => panic!("unexpected loop step after cancel: {other:?}"),
6333 }
6334
6335 let events = events.lock().unwrap();
6336 assert!(
6337 events
6338 .iter()
6339 .any(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false })),
6340 "pending approval cancellation should close approval UI state"
6341 );
6342 assert!(
6343 events.iter().any(|event| matches!(
6344 event,
6345 AgentEvent::ToolResultReceived(result)
6346 if result.call_id == ToolCallId::new("call-1") && result.is_error
6347 )),
6348 "pending approval cancellation should pair the assistant tool_use"
6349 );
6350 drop(events);
6351
6352 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
6353 }
6354
6355 #[tokio::test]
6356 async fn resolved_approval_runs_even_if_cancellation_also_fired() {
6357 let controller = CancellationController::new();
6358 let tools = ToolRegistry::new().with(EchoTool::default());
6359 let agent = Agent::builder()
6360 .model(FakeAdapter)
6361 .add_tool_source(tools)
6362 .permissions(ApproveFsReads)
6363 .cancellation(controller.handle())
6364 .build()
6365 .unwrap();
6366
6367 let mut driver = agent
6368 .start(SessionConfig {
6369 session_id: SessionId::new("session-resolved-approval-cancel-race"),
6370 metadata: MetadataMap::new(),
6371 cache: None,
6372 })
6373 .await
6374 .unwrap();
6375
6376 driver
6377 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6378 .unwrap();
6379
6380 let pending = match driver.next().await.unwrap() {
6381 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
6382 other => panic!("unexpected loop step: {other:?}"),
6383 };
6384
6385 controller.interrupt();
6386 pending.approve(&mut driver).unwrap();
6387
6388 match driver.next().await.unwrap() {
6389 LoopStep::Finished(turn) => {
6390 assert_eq!(turn.finish_reason, FinishReason::Completed);
6391 match &turn.items[0].parts[0] {
6392 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
6393 other => panic!("unexpected part after approval: {other:?}"),
6394 }
6395 }
6396 other => panic!("unexpected loop step after approved cancel race: {other:?}"),
6397 }
6398 }
6399
6400 #[tokio::test]
6401 async fn loop_resumes_with_patched_input_on_approval() {
6402 let tools = ToolRegistry::new().with(EchoTool::default());
6403 let agent = Agent::builder()
6404 .model(FakeAdapter)
6405 .add_tool_source(tools)
6406 .permissions(ApproveFsReads)
6407 .build()
6408 .unwrap();
6409
6410 let mut driver = agent
6411 .start(SessionConfig {
6412 session_id: SessionId::new("session-approval-patched"),
6413 metadata: MetadataMap::new(),
6414 cache: None,
6415 })
6416 .await
6417 .unwrap();
6418
6419 driver
6420 .submit_input(vec![Item {
6421 id: None,
6422 kind: ItemKind::User,
6423 parts: vec![Part::Text(TextPart {
6424 text: "ping".into(),
6425 metadata: MetadataMap::new(),
6426 })],
6427 metadata: MetadataMap::new(),
6428 usage: None,
6429 finish_reason: None,
6430 created_at: None,
6431 }])
6432 .unwrap();
6433
6434 match driver.next().await.unwrap() {
6435 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
6436 pending
6437 .approve_with_patched_input(&mut driver, json!({ "value": "patched" }))
6438 .unwrap();
6439 }
6440 other => panic!("unexpected loop step: {other:?}"),
6441 }
6442 match driver.next().await.unwrap() {
6443 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
6444 Part::Text(text) => assert_eq!(text.text, "tool said: patched"),
6445 other => panic!("unexpected part: {other:?}"),
6446 },
6447 other => panic!("unexpected loop step after approval: {other:?}"),
6448 }
6449 }
6450
6451 #[tokio::test]
6452 async fn loop_tracks_multiple_pending_approvals_by_call_id() {
6453 let tools = ToolRegistry::new().with(EchoTool::default());
6454 let agent = Agent::builder()
6455 .model(DualApprovalAdapter)
6456 .add_tool_source(tools)
6457 .permissions(ApproveFsReads)
6458 .build()
6459 .unwrap();
6460
6461 let mut driver = agent
6462 .start(SessionConfig {
6463 session_id: SessionId::new("session-dual-approval"),
6464 metadata: MetadataMap::new(),
6465 cache: None,
6466 })
6467 .await
6468 .unwrap();
6469
6470 driver
6471 .submit_input(vec![Item {
6472 id: None,
6473 kind: ItemKind::User,
6474 parts: vec![Part::Text(TextPart {
6475 text: "run both approvals".into(),
6476 metadata: MetadataMap::new(),
6477 })],
6478 metadata: MetadataMap::new(),
6479 usage: None,
6480 finish_reason: None,
6481 created_at: None,
6482 }])
6483 .unwrap();
6484
6485 let pending_first = match driver.next().await.unwrap() {
6486 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
6487 assert_eq!(
6488 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
6489 Some("call-1")
6490 );
6491 pending
6492 }
6493 other => panic!("unexpected first loop step: {other:?}"),
6494 };
6495
6496 let pending_second = match driver.next().await.unwrap() {
6497 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
6498 assert_eq!(
6499 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
6500 Some("call-2")
6501 );
6502 pending
6503 }
6504 other => panic!("unexpected second loop step: {other:?}"),
6505 };
6506
6507 pending_second.approve(&mut driver).unwrap();
6508 match driver.next().await.unwrap() {
6509 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
6510 assert_eq!(
6511 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
6512 Some("call-1")
6513 );
6514 }
6515 other => panic!("unexpected step after approving second request: {other:?}"),
6516 }
6517
6518 pending_first.approve(&mut driver).unwrap();
6519 match driver.next().await.unwrap() {
6520 LoopStep::Finished(turn) => {
6521 assert_eq!(turn.finish_reason, FinishReason::Completed);
6522 match &turn.items[0].parts[0] {
6523 Part::Text(text) => assert_eq!(text.text, "both approvals finished"),
6524 other => panic!("unexpected final part: {other:?}"),
6525 }
6526 }
6527 other => panic!("unexpected final loop step: {other:?}"),
6528 }
6529 }
6530
6531 #[tokio::test]
6532 async fn cancelling_all_pending_approvals_pairs_every_tool_use() {
6533 let events = StdArc::new(StdMutex::new(Vec::new()));
6534 let tools = ToolRegistry::new().with(EchoTool::default());
6535 let agent = Agent::builder()
6536 .model(DualApprovalAdapter)
6537 .add_tool_source(tools)
6538 .permissions(ApproveFsReads)
6539 .observer(RecordingObserver {
6540 events: events.clone(),
6541 })
6542 .build()
6543 .unwrap();
6544
6545 let mut driver = agent
6546 .start(SessionConfig {
6547 session_id: SessionId::new("session-dual-approval-cancel"),
6548 metadata: MetadataMap::new(),
6549 cache: None,
6550 })
6551 .await
6552 .unwrap();
6553
6554 driver
6555 .submit_input(vec![Item {
6556 id: None,
6557 kind: ItemKind::User,
6558 parts: vec![Part::Text(TextPart {
6559 text: "run both approvals".into(),
6560 metadata: MetadataMap::new(),
6561 })],
6562 metadata: MetadataMap::new(),
6563 usage: None,
6564 finish_reason: None,
6565 created_at: None,
6566 }])
6567 .unwrap();
6568
6569 for expected_call in ["call-1", "call-2"] {
6570 match driver.next().await.unwrap() {
6571 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
6572 assert_eq!(
6573 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
6574 Some(expected_call)
6575 );
6576 }
6577 other => panic!("unexpected approval step: {other:?}"),
6578 }
6579 }
6580
6581 match driver.cancel_pending_approvals().await.unwrap() {
6582 Some(LoopStep::Finished(turn)) => {
6583 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
6584 }
6585 other => panic!("unexpected cancellation result: {other:?}"),
6586 }
6587 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
6588
6589 let events = events.lock().unwrap();
6590 let cancelled = events
6591 .iter()
6592 .filter(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false }))
6593 .count();
6594 assert_eq!(cancelled, 2);
6595 assert!(events.iter().any(|event| matches!(
6596 event,
6597 AgentEvent::TurnFinished(turn) if turn.finish_reason == FinishReason::Cancelled
6598 )));
6599 for expected_call in ["call-1", "call-2"] {
6600 assert!(events.iter().any(|event| matches!(
6601 event,
6602 AgentEvent::ToolResultReceived(result)
6603 if result.call_id == ToolCallId::new(expected_call) && result.is_error
6604 )));
6605 }
6606 }
6607
6608 #[tokio::test]
6609 async fn loop_compacts_transcript_before_new_turns() {
6610 let events = StdArc::new(StdMutex::new(Vec::new()));
6611 let agent = Agent::builder()
6612 .model(FakeAdapter)
6613 .mutator(KeepRecentMutator { keep: 1 })
6614 .observer(RecordingObserver {
6615 events: events.clone(),
6616 })
6617 .build()
6618 .unwrap();
6619
6620 let mut driver = agent
6621 .start(SessionConfig {
6622 session_id: SessionId::new("session-4"),
6623 metadata: MetadataMap::new(),
6624 cache: None,
6625 })
6626 .await
6627 .unwrap();
6628
6629 for text in ["first", "second"] {
6630 driver
6631 .submit_input(vec![Item {
6632 id: None,
6633 kind: ItemKind::User,
6634 parts: vec![Part::Text(TextPart {
6635 text: text.into(),
6636 metadata: MetadataMap::new(),
6637 })],
6638 metadata: MetadataMap::new(),
6639 usage: None,
6640 finish_reason: None,
6641 created_at: None,
6642 }])
6643 .unwrap();
6644 let _ = driver.next().await.unwrap();
6645 }
6646
6647 let events = events.lock().unwrap();
6648 assert!(
6649 events
6650 .iter()
6651 .any(|event| matches!(event, AgentEvent::MutationFinished { dirty: true, .. }))
6652 );
6653 }
6654
6655 #[test]
6656 fn transcript_validation_rejects_orphaned_tool_result() {
6657 let transcript = vec![Item {
6658 id: None,
6659 kind: ItemKind::Tool,
6660 parts: vec![Part::ToolResult(ToolResultPart {
6661 call_id: "call-1".into(),
6662 output: ToolOutput::Text("result".into()),
6663 is_error: false,
6664 metadata: MetadataMap::new(),
6665 })],
6666 metadata: MetadataMap::new(),
6667 usage: None,
6668 finish_reason: None,
6669 created_at: None,
6670 }];
6671
6672 let error = validate_transcript_invariants(&transcript).unwrap_err();
6673 assert!(error.to_string().contains("orphaned tool_result"));
6674 }
6675
6676 #[test]
6677 fn transcript_validation_rejects_duplicate_tool_result() {
6678 let transcript = vec![
6679 Item {
6680 id: None,
6681 kind: ItemKind::Assistant,
6682 parts: vec![Part::ToolCall(ToolCallPart {
6683 id: "call-1".into(),
6684 name: "lookup".into(),
6685 input: serde_json::json!({}),
6686 metadata: MetadataMap::new(),
6687 })],
6688 metadata: MetadataMap::new(),
6689 usage: None,
6690 finish_reason: None,
6691 created_at: None,
6692 },
6693 Item {
6694 id: None,
6695 kind: ItemKind::Tool,
6696 parts: vec![Part::ToolResult(ToolResultPart {
6697 call_id: "call-1".into(),
6698 output: ToolOutput::Text("result".into()),
6699 is_error: false,
6700 metadata: MetadataMap::new(),
6701 })],
6702 metadata: MetadataMap::new(),
6703 usage: None,
6704 finish_reason: None,
6705 created_at: None,
6706 },
6707 Item {
6708 id: None,
6709 kind: ItemKind::Tool,
6710 parts: vec![Part::ToolResult(ToolResultPart {
6711 call_id: "call-1".into(),
6712 output: ToolOutput::Text("again".into()),
6713 is_error: false,
6714 metadata: MetadataMap::new(),
6715 })],
6716 metadata: MetadataMap::new(),
6717 usage: None,
6718 finish_reason: None,
6719 created_at: None,
6720 },
6721 ];
6722
6723 let error = validate_transcript_invariants(&transcript).unwrap_err();
6724 assert!(error.to_string().contains("duplicate tool_result"));
6725 }
6726
6727 #[tokio::test]
6728 async fn loop_refreshes_tool_specs_each_turn() {
6729 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
6730 let version = StdArc::new(AtomicUsize::new(1));
6731 let tools = ToolRegistry::new().with(DynamicSpecTool::new(version.clone()));
6732 let agent = Agent::builder()
6733 .model(RecordingAdapter {
6734 seen_descriptions: seen_descriptions.clone(),
6735 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
6736 })
6737 .add_tool_source(tools)
6738 .permissions(AllowAllPermissions)
6739 .build()
6740 .unwrap();
6741
6742 let mut driver = agent
6743 .start(SessionConfig {
6744 session_id: SessionId::new("session-dynamic-tools"),
6745 metadata: MetadataMap::new(),
6746 cache: None,
6747 })
6748 .await
6749 .unwrap();
6750
6751 for text in ["first", "second"] {
6752 driver
6753 .submit_input(vec![Item {
6754 id: None,
6755 kind: ItemKind::User,
6756 parts: vec![Part::Text(TextPart {
6757 text: text.into(),
6758 metadata: MetadataMap::new(),
6759 })],
6760 metadata: MetadataMap::new(),
6761 usage: None,
6762 finish_reason: None,
6763 created_at: None,
6764 }])
6765 .unwrap();
6766
6767 let _ = driver.next().await.unwrap();
6768 if text == "first" {
6769 version.store(2, Ordering::SeqCst);
6770 }
6771 }
6772
6773 let seen_descriptions = seen_descriptions.lock().unwrap();
6774 assert_eq!(seen_descriptions.len(), 2);
6775 assert_eq!(seen_descriptions[0], vec!["dynamic version 1".to_string()]);
6776 assert_eq!(seen_descriptions[1], vec!["dynamic version 2".to_string()]);
6777 }
6778
6779 #[tokio::test]
6780 async fn loop_emits_catalog_change_and_uses_updated_specs_next_turn() {
6781 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
6782 let events = StdArc::new(StdMutex::new(Vec::new()));
6783 let executor = StdArc::new(CatalogExecutor::new());
6784 let executor_for_agent: Arc<dyn ToolExecutor> = executor.clone();
6785 let agent = Agent::builder()
6786 .model(RecordingAdapter {
6787 seen_descriptions: seen_descriptions.clone(),
6788 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
6789 })
6790 .tool_executor(executor_for_agent)
6791 .permissions(AllowAllPermissions)
6792 .observer(RecordingObserver {
6793 events: events.clone(),
6794 })
6795 .build()
6796 .unwrap();
6797
6798 let mut driver = agent
6799 .start(SessionConfig {
6800 session_id: SessionId::new("session-catalog-events"),
6801 metadata: MetadataMap::new(),
6802 cache: None,
6803 })
6804 .await
6805 .unwrap();
6806
6807 driver
6808 .submit_input(vec![Item::text(ItemKind::User, "first")])
6809 .unwrap();
6810 let _ = driver.next().await.unwrap();
6811
6812 executor.publish_change(
6813 1,
6814 ToolCatalogEvent {
6815 source: "mcp:mock".into(),
6816 added: vec!["dynamic".into()],
6817 removed: Vec::new(),
6818 changed: Vec::new(),
6819 },
6820 );
6821
6822 driver
6823 .submit_input(vec![Item::text(ItemKind::User, "second")])
6824 .unwrap();
6825 let _ = driver.next().await.unwrap();
6826
6827 let seen_descriptions = seen_descriptions.lock().unwrap();
6828 assert_eq!(seen_descriptions.len(), 2);
6829 assert_eq!(seen_descriptions[0], vec!["dynamic version 0".to_string()]);
6830 assert_eq!(seen_descriptions[1], vec!["dynamic version 1".to_string()]);
6831
6832 let events = events.lock().unwrap();
6833 assert!(events.iter().any(|event| matches!(
6834 event,
6835 AgentEvent::ToolCatalogChanged(ToolCatalogEvent {
6836 source,
6837 added,
6838 removed,
6839 changed,
6840 }) if source == "mcp:mock"
6841 && added == &vec!["dynamic".to_string()]
6842 && removed.is_empty()
6843 && changed.is_empty()
6844 )));
6845 }
6846
6847 #[tokio::test]
6848 async fn loop_passes_session_default_and_next_turn_cache_requests() {
6849 let seen_caches = StdArc::new(StdMutex::new(Vec::new()));
6850 let agent = Agent::builder()
6851 .model(RecordingAdapter {
6852 seen_descriptions: StdArc::new(StdMutex::new(Vec::new())),
6853 seen_caches: seen_caches.clone(),
6854 })
6855 .permissions(AllowAllPermissions)
6856 .build()
6857 .unwrap();
6858
6859 let default_cache = PromptCacheRequest::best_effort(PromptCacheStrategy::Automatic)
6860 .with_retention(PromptCacheRetention::Short);
6861 let override_cache = PromptCacheRequest::required(PromptCacheStrategy::Explicit {
6862 breakpoints: vec![PromptCacheBreakpoint::TranscriptItemEnd { index: 0 }],
6863 });
6864
6865 let mut driver = agent
6866 .start(SessionConfig {
6867 session_id: SessionId::new("session-cache"),
6868 metadata: MetadataMap::new(),
6869 cache: Some(default_cache.clone()),
6870 })
6871 .await
6872 .unwrap();
6873
6874 driver
6875 .submit_input(vec![Item {
6876 id: None,
6877 kind: ItemKind::User,
6878 parts: vec![Part::Text(TextPart {
6879 text: "first".into(),
6880 metadata: MetadataMap::new(),
6881 })],
6882 metadata: MetadataMap::new(),
6883 usage: None,
6884 finish_reason: None,
6885 created_at: None,
6886 }])
6887 .unwrap();
6888 let _ = driver.next().await.unwrap();
6889
6890 driver
6891 .submit_input_with_cache(
6892 vec![Item {
6893 id: None,
6894 kind: ItemKind::User,
6895 parts: vec![Part::Text(TextPart {
6896 text: "second".into(),
6897 metadata: MetadataMap::new(),
6898 })],
6899 metadata: MetadataMap::new(),
6900 usage: None,
6901 finish_reason: None,
6902 created_at: None,
6903 }],
6904 override_cache.clone(),
6905 )
6906 .unwrap();
6907 let _ = driver.next().await.unwrap();
6908
6909 let seen = seen_caches.lock().unwrap();
6910 assert_eq!(seen.len(), 2);
6911 assert_eq!(seen[0], Some(default_cache));
6912 assert_eq!(seen[1], Some(override_cache));
6913 }
6914
6915 #[tokio::test]
6916 async fn loop_yields_after_tool_result_between_rounds() {
6917 let tools = ToolRegistry::new().with(EchoTool::default());
6918 let agent = Agent::builder()
6919 .model(FakeAdapter)
6920 .add_tool_source(tools)
6921 .permissions(AllowAllPermissions)
6922 .build()
6923 .unwrap();
6924
6925 let mut driver = agent
6926 .start(SessionConfig {
6927 session_id: SessionId::new("yield-session"),
6928 metadata: MetadataMap::new(),
6929 cache: None,
6930 })
6931 .await
6932 .unwrap();
6933
6934 driver
6935 .submit_input(vec![Item::text(ItemKind::User, "ping")])
6936 .unwrap();
6937
6938 let step = driver.next().await.unwrap();
6941 let info = match step {
6942 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info,
6943 other => panic!("expected AfterToolResult, got {other:?}"),
6944 };
6945 assert_eq!(info.session_id, SessionId::new("yield-session"));
6946 assert_eq!(info.transcript_len, 3);
6948
6949 let interrupt = LoopInterrupt::AfterToolResult(info.clone());
6951 assert!(!interrupt.is_blocking());
6952
6953 driver
6955 .submit_input(vec![Item::text(ItemKind::User, "also: report back")])
6956 .unwrap();
6957
6958 let step = driver.next().await.unwrap();
6961 match step {
6962 LoopStep::Finished(turn) => {
6963 assert_eq!(turn.finish_reason, FinishReason::Completed);
6964 }
6965 other => panic!("expected Finished, got {other:?}"),
6966 }
6967
6968 let snapshot = driver.snapshot();
6970 let has_injected_message = snapshot.transcript.iter().any(|item| {
6971 item.kind == ItemKind::User
6972 && item.parts.iter().any(|part| match part {
6973 Part::Text(text) => text.text == "also: report back",
6974 _ => false,
6975 })
6976 });
6977 assert!(
6978 has_injected_message,
6979 "injected user message should be in transcript, got: {:?}",
6980 snapshot.transcript
6981 );
6982 }
6983
6984 struct RecordingTranscriptObserver {
6985 items: StdArc<StdMutex<Vec<Item>>>,
6986 }
6987
6988 impl TranscriptObserver for RecordingTranscriptObserver {
6989 fn on_transcript_event(&self, event: TranscriptEvent<'_>) {
6990 self.items.lock().unwrap().push(event.item.clone());
6991 }
6992 }
6993
6994 #[tokio::test]
6995 async fn observers_see_full_tool_round() {
6996 let events = StdArc::new(StdMutex::new(Vec::<AgentEvent>::new()));
7002 let items = StdArc::new(StdMutex::new(Vec::<Item>::new()));
7003 let agent = Agent::builder()
7004 .model(FakeAdapter)
7005 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
7006 .permissions(AllowAllPermissions)
7007 .observer(RecordingObserver {
7008 events: events.clone(),
7009 })
7010 .transcript_observer(RecordingTranscriptObserver {
7011 items: items.clone(),
7012 })
7013 .build()
7014 .unwrap();
7015
7016 let mut driver = agent
7017 .start(SessionConfig {
7018 session_id: SessionId::new("observer-session"),
7019 metadata: MetadataMap::new(),
7020 cache: None,
7021 })
7022 .await
7023 .unwrap();
7024
7025 driver
7026 .submit_input(vec![Item {
7027 id: None,
7028 kind: ItemKind::User,
7029 parts: vec![Part::Text(TextPart {
7030 text: "ping".into(),
7031 metadata: MetadataMap::new(),
7032 })],
7033 metadata: MetadataMap::new(),
7034 usage: None,
7035 finish_reason: None,
7036 created_at: None,
7037 }])
7038 .unwrap();
7039
7040 let result = run_until_finished(&mut driver).await;
7041 assert!(matches!(result, LoopStep::Finished(_)), "got {result:?}");
7042
7043 let events = events.lock().unwrap().clone();
7046 let tool_call_id = events.iter().find_map(|e| match e {
7047 AgentEvent::ToolCallRequested(c) => Some(c.id.clone()),
7048 _ => None,
7049 });
7050 let tool_results: Vec<_> = events
7051 .iter()
7052 .filter_map(|e| match e {
7053 AgentEvent::ToolResultReceived(r) => Some(r.clone()),
7054 _ => None,
7055 })
7056 .collect();
7057 assert_eq!(tool_results.len(), 1, "events: {events:?}");
7058 assert_eq!(Some(tool_results[0].call_id.clone()), tool_call_id);
7059 assert!(!tool_results[0].is_error);
7060
7061 let items = items.lock().unwrap().clone();
7065 assert_eq!(items.len(), 4, "items: {items:?}");
7066 assert_eq!(items[0].kind, ItemKind::User);
7067 assert_eq!(items[1].kind, ItemKind::Assistant);
7068 assert!(
7069 items[1]
7070 .parts
7071 .iter()
7072 .any(|p| matches!(p, Part::ToolCall(_)))
7073 );
7074 assert_eq!(items[2].kind, ItemKind::Tool);
7075 assert!(
7076 items[2]
7077 .parts
7078 .iter()
7079 .any(|p| matches!(p, Part::ToolResult(_)))
7080 );
7081 assert_eq!(items[3].kind, ItemKind::Assistant);
7082 }
7083
7084 #[test]
7085 fn convenience_cache_builders_construct_expected_defaults() {
7086 let cache = PromptCacheRequest::automatic()
7087 .with_retention(PromptCacheRetention::Short)
7088 .with_key("workspace:demo");
7089 let session = SessionConfig::new("demo").with_cache(cache.clone());
7090
7091 assert_eq!(session.session_id, SessionId::new("demo"));
7092 assert_eq!(session.cache, Some(cache));
7093
7094 let explicit = PromptCacheRequest::explicit([
7095 PromptCacheBreakpoint::tools_end(),
7096 PromptCacheBreakpoint::transcript_item_end(2),
7097 PromptCacheBreakpoint::transcript_part_end(3, 1),
7098 ]);
7099
7100 assert_eq!(explicit.mode, PromptCacheMode::BestEffort);
7101 assert_eq!(
7102 explicit.strategy,
7103 PromptCacheStrategy::Explicit {
7104 breakpoints: vec![
7105 PromptCacheBreakpoint::ToolsEnd,
7106 PromptCacheBreakpoint::TranscriptItemEnd { index: 2 },
7107 PromptCacheBreakpoint::TranscriptPartEnd {
7108 item_index: 3,
7109 part_index: 1,
7110 },
7111 ],
7112 }
7113 );
7114 }
7115}