1use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
65use std::sync::Arc;
66
67use agentkit_core::{
68 CancellationHandle, Delta, FinishReason, Item, ItemKind, MetadataMap, Part, SessionId, TaskId,
69 TextPart, Timestamp, ToolCallId, ToolCallPart, ToolOutput, ToolResultPart, TurnCancellation,
70 Usage,
71};
72use agentkit_task_manager::{
73 PendingLoopUpdates, SimpleTaskManager, TOOL_RESULT_FAILURE_KIND_METADATA_KEY,
74 TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED, TaskApproval, TaskLaunchKind, TaskLaunchRequest,
75 TaskManager, TaskResolution, TaskStartContext, TaskStartOutcome, TurnTaskUpdate,
76};
77#[cfg(test)]
78use agentkit_tools_core::ToolContext;
79use agentkit_tools_core::{
80 AllowAllPermissions, ApprovalDecision, ApprovalRequest, BasicToolExecutor, OwnedToolContext,
81 PermissionChecker, ToolCatalogEvent, ToolError, ToolExecutionScope, ToolExecutor, ToolRequest,
82 ToolResources, ToolSource, ToolSpec,
83};
84use async_trait::async_trait;
85use serde::{Deserialize, Serialize};
86use serde_json::Value;
87use thiserror::Error;
88
89const INTERRUPTED_METADATA_KEY: &str = "agentkit.interrupted";
90const INTERRUPT_REASON_METADATA_KEY: &str = "agentkit.interrupt_reason";
91const INTERRUPT_STAGE_METADATA_KEY: &str = "agentkit.interrupt_stage";
92const USER_CANCELLED_REASON: &str = "user_cancelled";
93
94#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
109pub struct SessionConfig {
110 pub session_id: SessionId,
112 pub metadata: MetadataMap,
114 pub cache: Option<PromptCacheRequest>,
116}
117
118impl SessionConfig {
119 pub fn new(session_id: impl Into<SessionId>) -> Self {
121 Self {
122 session_id: session_id.into(),
123 metadata: MetadataMap::new(),
124 cache: None,
125 }
126 }
127
128 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
130 self.metadata = metadata;
131 self
132 }
133
134 pub fn with_cache(mut self, cache: PromptCacheRequest) -> Self {
136 self.cache = Some(cache);
137 self
138 }
139
140 pub fn without_cache(mut self) -> Self {
142 self.cache = None;
143 self
144 }
145}
146
147#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
153pub enum PromptCacheMode {
154 Disabled,
156 #[default]
158 BestEffort,
159 Required,
161}
162
163#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
169pub enum PromptCacheRetention {
170 Default,
172 Short,
174 Extended,
176}
177
178#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
180pub enum PromptCacheStrategy {
181 #[default]
183 Automatic,
184 Explicit {
186 breakpoints: Vec<PromptCacheBreakpoint>,
188 },
189}
190
191impl PromptCacheStrategy {
192 pub fn automatic() -> Self {
195 Self::Automatic
196 }
197
198 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
200 Self::Explicit {
201 breakpoints: breakpoints.into_iter().collect(),
202 }
203 }
204}
205
206#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
208pub enum PromptCacheBreakpoint {
209 ToolsEnd,
211 TranscriptItemEnd { index: usize },
213 TranscriptPartEnd {
219 item_index: usize,
220 part_index: usize,
221 },
222}
223
224impl PromptCacheBreakpoint {
225 pub fn tools_end() -> Self {
227 Self::ToolsEnd
228 }
229
230 pub fn transcript_item_end(index: usize) -> Self {
232 Self::TranscriptItemEnd { index }
233 }
234
235 pub fn transcript_part_end(item_index: usize, part_index: usize) -> Self {
237 Self::TranscriptPartEnd {
238 item_index,
239 part_index,
240 }
241 }
242}
243
244#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
246pub struct PromptCacheRequest {
247 pub mode: PromptCacheMode,
249 pub strategy: PromptCacheStrategy,
251 pub retention: Option<PromptCacheRetention>,
253 pub key: Option<String>,
255}
256
257impl PromptCacheRequest {
258 pub fn automatic() -> Self {
260 Self::best_effort(PromptCacheStrategy::automatic())
261 }
262
263 pub fn automatic_required() -> Self {
265 Self::required(PromptCacheStrategy::automatic())
266 }
267
268 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
270 Self::best_effort(PromptCacheStrategy::explicit(breakpoints))
271 }
272
273 pub fn explicit_required(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
275 Self::required(PromptCacheStrategy::explicit(breakpoints))
276 }
277
278 pub fn disabled() -> Self {
280 Self {
281 mode: PromptCacheMode::Disabled,
282 strategy: PromptCacheStrategy::Automatic,
283 retention: None,
284 key: None,
285 }
286 }
287
288 pub fn best_effort(strategy: PromptCacheStrategy) -> Self {
290 Self {
291 mode: PromptCacheMode::BestEffort,
292 strategy,
293 retention: None,
294 key: None,
295 }
296 }
297
298 pub fn required(strategy: PromptCacheStrategy) -> Self {
300 Self {
301 mode: PromptCacheMode::Required,
302 strategy,
303 retention: None,
304 key: None,
305 }
306 }
307
308 pub fn with_mode(mut self, mode: PromptCacheMode) -> Self {
310 self.mode = mode;
311 self
312 }
313
314 pub fn with_strategy(mut self, strategy: PromptCacheStrategy) -> Self {
316 self.strategy = strategy;
317 self
318 }
319
320 pub fn with_retention(mut self, retention: PromptCacheRetention) -> Self {
322 self.retention = Some(retention);
323 self
324 }
325
326 pub fn with_key(mut self, key: impl Into<String>) -> Self {
328 self.key = Some(key.into());
329 self
330 }
331
332 pub fn without_retention(mut self) -> Self {
334 self.retention = None;
335 self
336 }
337
338 pub fn without_key(mut self) -> Self {
340 self.key = None;
341 self
342 }
343
344 pub fn is_enabled(&self) -> bool {
346 !matches!(self.mode, PromptCacheMode::Disabled)
347 }
348}
349
350#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
356pub struct TurnRequest {
357 pub session_id: SessionId,
359 pub turn_id: agentkit_core::TurnId,
361 pub transcript: Vec<Item>,
363 pub available_tools: Vec<ToolSpec>,
365 pub cache: Option<PromptCacheRequest>,
367 pub metadata: MetadataMap,
369}
370
371#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
376pub struct ModelTurnResult {
377 pub finish_reason: FinishReason,
379 pub output_items: Vec<Item>,
381 pub usage: Option<Usage>,
383 pub metadata: MetadataMap,
385 #[serde(default)]
389 pub model: Option<String>,
390 #[serde(default)]
394 pub response_id: Option<String>,
395}
396
397#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
403pub enum ModelTurnEvent {
404 Delta(Delta),
406 ToolCall(ToolCallPart),
408 Usage(Usage),
410 Finished(ModelTurnResult),
412}
413
414#[async_trait]
451pub trait ModelAdapter: Send + Sync {
452 type Session: ModelSession;
454
455 async fn start_session(&self, config: SessionConfig) -> Result<Self::Session, LoopError>;
461
462 fn provider_name(&self) -> Option<&str> {
469 None
470 }
471}
472
473#[async_trait]
480pub trait ModelSession: Send {
481 type Turn: ModelTurn;
483
484 async fn begin_turn(
497 &mut self,
498 request: TurnRequest,
499 cancellation: Option<TurnCancellation>,
500 ) -> Result<Self::Turn, LoopError>;
501
502 fn model_name(&self) -> Option<&str> {
508 None
509 }
510}
511
512#[async_trait]
518pub trait ModelTurn: Send {
519 async fn next_event(
528 &mut self,
529 cancellation: Option<TurnCancellation>,
530 ) -> Result<Option<ModelTurnEvent>, LoopError>;
531}
532
533pub trait LoopObserver: Send + Sync {
553 fn handle_event(&self, event: ObservedEvent);
558}
559
560#[derive(Clone, Debug, PartialEq)]
566pub struct ObservedEvent {
567 pub session_id: Arc<SessionId>,
569 pub event: AgentEvent,
571}
572
573pub trait TranscriptObserver: Send + Sync {
608 fn on_transcript_event(&self, event: TranscriptEvent<'_>);
613}
614
615#[derive(Clone, Debug)]
618pub struct TranscriptEvent<'a> {
619 pub session_id: &'a SessionId,
621 pub item: &'a Item,
623}
624
625#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
630#[non_exhaustive]
631pub enum MutationPoint {
632 AfterToolResult,
635 AfterTurnEnded,
638}
639
640pub trait EventEmitter: Send + Sync {
643 fn emit(&self, event: AgentEvent);
645}
646
647#[non_exhaustive]
649pub struct LoopCtx<'a> {
650 pub session_id: &'a SessionId,
652 pub turn_id: Option<&'a agentkit_core::TurnId>,
654 pub point: MutationPoint,
656 pub cancellation: Option<TurnCancellation>,
658 pub emitter: &'a dyn EventEmitter,
660}
661
662pub struct TranscriptCursor<'a> {
670 items: &'a mut Vec<Item>,
671 pub(crate) dirty: bool,
672}
673
674impl<'a> std::ops::Deref for TranscriptCursor<'a> {
675 type Target = Vec<Item>;
676 fn deref(&self) -> &Vec<Item> {
677 self.items
678 }
679}
680
681impl<'a> std::ops::DerefMut for TranscriptCursor<'a> {
682 fn deref_mut(&mut self) -> &mut Vec<Item> {
683 self.dirty = true;
684 self.items
685 }
686}
687
688#[async_trait]
696pub trait LoopMutator: Send + Sync {
697 async fn mutate(
702 &self,
703 cursor: &mut TranscriptCursor<'_>,
704 ctx: LoopCtx<'_>,
705 ) -> Result<(), LoopError> {
706 let _ = (cursor, ctx);
707 Ok(())
708 }
709}
710
711#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
716#[non_exhaustive]
717pub enum AgentEvent {
718 RunStarted { session_id: SessionId },
720 TurnStarted {
722 session_id: SessionId,
723 turn_id: agentkit_core::TurnId,
724 },
725 InputAccepted {
727 session_id: SessionId,
728 items: Vec<Item>,
729 },
730 ContentDelta(Delta),
732 ToolCallRequested(ToolCallPart),
734 ToolExecutionStarted(ToolCallPart),
736 ToolExecutionProgress(ToolResultPart),
742 ToolResultReceived(ToolResultPart),
751 ApprovalRequired(ApprovalRequest),
753 ApprovalResolved { approved: bool },
755 ToolCatalogChanged(ToolCatalogEvent),
757 MutationStarted {
760 session_id: SessionId,
761 turn_id: Option<agentkit_core::TurnId>,
762 mutator: String,
763 point: MutationPoint,
764 },
765 MutationFinished {
769 session_id: SessionId,
770 turn_id: Option<agentkit_core::TurnId>,
771 mutator: String,
772 dirty: bool,
773 metadata: MetadataMap,
774 },
775 UsageUpdated(Usage),
777 Warning { message: String },
779 RunFailed { message: String },
781 TurnFinished(TurnResult),
783}
784
785#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
807pub struct PendingApproval {
808 pub request: ApprovalRequest,
810}
811
812impl std::ops::Deref for PendingApproval {
813 type Target = ApprovalRequest;
814 fn deref(&self) -> &ApprovalRequest {
815 &self.request
816 }
817}
818
819impl PendingApproval {
820 pub fn approve<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
822 let call_id = self
823 .request
824 .call_id
825 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
826 driver.resolve_approval_for(call_id, ApprovalDecision::Approve)
827 }
828
829 pub fn deny<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
831 let call_id = self
832 .request
833 .call_id
834 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
835 driver.resolve_approval_for(call_id, ApprovalDecision::Deny { reason: None })
836 }
837
838 pub fn deny_with_reason<S: ModelSession>(
840 self,
841 driver: &mut LoopDriver<S>,
842 reason: impl Into<String>,
843 ) -> Result<(), LoopError> {
844 let call_id = self
845 .request
846 .call_id
847 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
848 driver.resolve_approval_for(
849 call_id,
850 ApprovalDecision::Deny {
851 reason: Some(reason.into()),
852 },
853 )
854 }
855
856 pub fn approve_with_patched_input<S: ModelSession>(
866 self,
867 driver: &mut LoopDriver<S>,
868 input: serde_json::Value,
869 ) -> Result<(), LoopError> {
870 let call_id = self
871 .request
872 .call_id
873 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
874 driver.resolve_approval_for_with_patched_input(call_id, input)
875 }
876}
877
878#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
903pub struct InputRequest {
904 pub session_id: SessionId,
906 pub reason: String,
908}
909
910impl InputRequest {
911 pub fn submit<S: ModelSession>(
913 self,
914 driver: &mut LoopDriver<S>,
915 items: Vec<Item>,
916 ) -> Result<(), LoopError> {
917 driver.submit_input(items)
918 }
919}
920
921#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
926pub struct TurnResult {
927 pub turn_id: agentkit_core::TurnId,
929 pub finish_reason: FinishReason,
931 pub items: Vec<Item>,
933 pub usage: Option<Usage>,
935 pub metadata: MetadataMap,
937}
938
939#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
975pub enum LoopInterrupt {
976 ApprovalRequest(PendingApproval),
978 AwaitingInput(InputRequest),
980 AfterToolResult(ToolRoundInfo),
990}
991
992impl LoopInterrupt {
993 pub fn is_blocking(&self) -> bool {
999 matches!(self, LoopInterrupt::ApprovalRequest(_))
1000 }
1001}
1002
1003#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1006pub struct ToolRoundInfo {
1007 pub session_id: SessionId,
1009 pub turn_id: agentkit_core::TurnId,
1011 pub transcript_len: usize,
1013}
1014
1015impl ToolRoundInfo {
1016 pub fn submit<S: ModelSession>(
1019 self,
1020 driver: &mut LoopDriver<S>,
1021 items: Vec<Item>,
1022 ) -> Result<(), LoopError> {
1023 driver.submit_input(items)
1024 }
1025}
1026
1027#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1055pub enum LoopStep {
1056 Interrupt(LoopInterrupt),
1058 Finished(TurnResult),
1060}
1061
1062#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1068pub struct LoopSnapshot {
1069 pub session_id: SessionId,
1071 pub transcript: Vec<Item>,
1073 pub pending_input: Vec<Item>,
1075}
1076
1077#[derive(Clone)]
1078struct PendingApprovalToolCall {
1079 request: ApprovalRequest,
1080 decision: Option<ApprovalDecision>,
1081 surfaced: bool,
1082 turn_id: agentkit_core::TurnId,
1083 task_id: TaskId,
1084 call: ToolCallPart,
1085 tool_request: ToolRequest,
1086 cancellation: Option<TurnCancellation>,
1087}
1088
1089#[derive(Clone, Default)]
1090struct ActiveToolRound {
1091 turn_id: agentkit_core::TurnId,
1092 pending_calls: VecDeque<(ToolCallPart, ToolRequest)>,
1093 cancellation: Option<TurnCancellation>,
1094 background_pending: bool,
1095 foreground_progressed: bool,
1096}
1097
1098pub struct Agent<M>
1140where
1141 M: ModelAdapter,
1142{
1143 model: M,
1144 tool_sources: Vec<Arc<dyn ToolSource>>,
1145 tool_executor: Option<Arc<dyn ToolExecutor>>,
1146 task_manager: Arc<dyn TaskManager>,
1147 permissions: Arc<dyn PermissionChecker>,
1148 resources: Arc<dyn ToolResources>,
1149 cancellation: Option<CancellationHandle>,
1150 mutators: Vec<Arc<dyn LoopMutator>>,
1151 observers: Vec<Arc<dyn LoopObserver>>,
1152 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1153 transcript: Vec<Item>,
1154 input: Vec<Item>,
1155}
1156
1157impl<M> Agent<M>
1158where
1159 M: ModelAdapter,
1160{
1161 pub fn builder() -> AgentBuilder<M> {
1163 AgentBuilder::default()
1164 }
1165
1166 pub async fn start(&self, config: SessionConfig) -> Result<LoopDriver<M::Session>, LoopError> {
1182 let session_id = config.session_id.clone();
1183 let default_cache = config.cache.clone();
1184 let session = self.model.start_session(config).await?;
1185 let tool_executor = self
1186 .tool_executor
1187 .clone()
1188 .unwrap_or_else(|| Arc::new(BasicToolExecutor::new(self.tool_sources.clone())));
1189 let driver = LoopDriver {
1190 session_id: session_id.clone(),
1191 observed_session_id: Arc::new(session_id.clone()),
1192 provider_name: self.model.provider_name().map(str::to_owned),
1193 default_cache,
1194 next_turn_cache: None,
1195 session: Some(session),
1196 tool_executor,
1197 task_manager: self.task_manager.clone(),
1198 permissions: self.permissions.clone(),
1199 resources: self.resources.clone(),
1200 cancellation: self.cancellation.clone(),
1201 mutators: self.mutators.clone(),
1202 observers: self.observers.clone(),
1203 transcript_observers: self.transcript_observers.clone(),
1204 transcript: self.transcript.clone(),
1205 pending_input: self.input.clone(),
1206 pending_approvals: BTreeMap::new(),
1207 pending_approval_order: VecDeque::new(),
1208 active_tool_round: None,
1209 pending_round_resume: None,
1210 next_turn_index: 1,
1211 detached_call_ids: HashSet::new(),
1212 tool_cancellations: HashMap::new(),
1213 };
1214 driver.emit(AgentEvent::RunStarted { session_id });
1215 Ok(driver)
1216 }
1217}
1218
1219pub struct AgentBuilder<M>
1225where
1226 M: ModelAdapter,
1227{
1228 model: Option<M>,
1229 tool_sources: Vec<Arc<dyn ToolSource>>,
1230 tool_executor: Option<Arc<dyn ToolExecutor>>,
1231 task_manager: Option<Arc<dyn TaskManager>>,
1232 permissions: Arc<dyn PermissionChecker>,
1233 resources: Arc<dyn ToolResources>,
1234 cancellation: Option<CancellationHandle>,
1235 mutators: Vec<Arc<dyn LoopMutator>>,
1236 observers: Vec<Arc<dyn LoopObserver>>,
1237 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1238 transcript: Vec<Item>,
1239 input: Vec<Item>,
1240}
1241
1242impl<M> Default for AgentBuilder<M>
1243where
1244 M: ModelAdapter,
1245{
1246 fn default() -> Self {
1247 Self {
1248 model: None,
1249 tool_sources: Vec::new(),
1250 tool_executor: None,
1251 task_manager: None,
1252 permissions: Arc::new(AllowAllPermissions),
1253 resources: Arc::new(()),
1254 cancellation: None,
1255 mutators: Vec::new(),
1256 observers: Vec::new(),
1257 transcript_observers: Vec::new(),
1258 transcript: Vec::new(),
1259 input: Vec::new(),
1260 }
1261 }
1262}
1263
1264impl<M> AgentBuilder<M>
1265where
1266 M: ModelAdapter,
1267{
1268 pub fn model(mut self, model: M) -> Self {
1270 self.model = Some(model);
1271 self
1272 }
1273
1274 pub fn add_tool_source<S: ToolSource + 'static>(mut self, source: S) -> Self {
1287 self.tool_sources.push(Arc::new(source));
1288 self
1289 }
1290
1291 pub fn tool_executor(mut self, executor: impl ToolExecutor + 'static) -> Self {
1297 self.tool_executor = Some(Arc::new(executor));
1298 self
1299 }
1300
1301 pub fn task_manager(mut self, manager: impl TaskManager + 'static) -> Self {
1306 self.task_manager = Some(Arc::new(manager));
1307 self
1308 }
1309
1310 pub fn permissions(mut self, permissions: impl PermissionChecker + 'static) -> Self {
1314 self.permissions = Arc::new(permissions);
1315 self
1316 }
1317
1318 pub fn resources(mut self, resources: impl ToolResources + 'static) -> Self {
1320 self.resources = Arc::new(resources);
1321 self
1322 }
1323
1324 pub fn cancellation(mut self, handle: CancellationHandle) -> Self {
1326 self.cancellation = Some(handle);
1327 self
1328 }
1329
1330 pub fn mutator<L: LoopMutator + 'static>(mut self, mutator: L) -> Self {
1338 self.mutators.push(Arc::new(mutator));
1339 self
1340 }
1341
1342 pub fn observer<O: LoopObserver + 'static>(mut self, observer: O) -> Self {
1346 self.observers.push(Arc::new(observer));
1347 self
1348 }
1349
1350 pub fn transcript_observer<O: TranscriptObserver + 'static>(mut self, observer: O) -> Self {
1359 self.transcript_observers.push(Arc::new(observer));
1360 self
1361 }
1362
1363 pub fn transcript(mut self, transcript: Vec<Item>) -> Self {
1371 self.transcript = transcript;
1372 self
1373 }
1374
1375 pub fn input(mut self, input: Vec<Item>) -> Self {
1384 self.input = input;
1385 self
1386 }
1387
1388 pub fn build(self) -> Result<Agent<M>, LoopError> {
1394 let model = self
1395 .model
1396 .ok_or_else(|| LoopError::InvalidState("model adapter is required".into()))?;
1397 Ok(Agent {
1398 model,
1399 tool_sources: self.tool_sources,
1400 tool_executor: self.tool_executor,
1401 task_manager: self
1402 .task_manager
1403 .unwrap_or_else(|| Arc::new(SimpleTaskManager::new())),
1404 permissions: self.permissions,
1405 resources: self.resources,
1406 cancellation: self.cancellation,
1407 mutators: self.mutators,
1408 observers: self.observers,
1409 transcript_observers: self.transcript_observers,
1410 transcript: self.transcript,
1411 input: self.input,
1412 })
1413 }
1414}
1415
1416pub struct LoopDriver<S>
1447where
1448 S: ModelSession,
1449{
1450 session_id: SessionId,
1451 observed_session_id: Arc<SessionId>,
1452 provider_name: Option<String>,
1453 default_cache: Option<PromptCacheRequest>,
1454 next_turn_cache: Option<PromptCacheRequest>,
1455 session: Option<S>,
1456 tool_executor: Arc<dyn ToolExecutor>,
1457 task_manager: Arc<dyn TaskManager>,
1458 permissions: Arc<dyn PermissionChecker>,
1459 resources: Arc<dyn ToolResources>,
1460 cancellation: Option<CancellationHandle>,
1461 mutators: Vec<Arc<dyn LoopMutator>>,
1462 observers: Vec<Arc<dyn LoopObserver>>,
1463 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1464 transcript: Vec<Item>,
1465 pending_input: Vec<Item>,
1466 pending_approvals: BTreeMap<ToolCallId, PendingApprovalToolCall>,
1467 pending_approval_order: VecDeque<ToolCallId>,
1468 active_tool_round: Option<ActiveToolRound>,
1469 pending_round_resume: Option<agentkit_core::TurnId>,
1470 next_turn_index: u64,
1471 detached_call_ids: HashSet<ToolCallId>,
1479 tool_cancellations: HashMap<ToolCallId, TurnCancellation>,
1480}
1481
1482impl<S> LoopDriver<S>
1483where
1484 S: ModelSession,
1485{
1486 fn execute_tool_span(
1487 &self,
1488 request: &ToolRequest,
1489 turn_id: &agentkit_core::TurnId,
1490 launch_kind: &'static str,
1491 ) -> tracing::Span {
1492 tracing::info_span!(
1493 "agent.execute_tool",
1494 "otel.name" = %format!("execute_tool {}", request.tool_name),
1495 "gen_ai.operation.name" = "execute_tool",
1496 "gen_ai.tool.name" = %request.tool_name,
1497 "gen_ai.tool.call.id" = %request.call_id,
1498 "gen_ai.conversation.id" = %self.session_id,
1499 "error.type" = tracing::field::Empty,
1500 session.id = %self.session_id,
1501 turn.id = %turn_id,
1502 launch_kind = launch_kind,
1503 )
1504 }
1505
1506 fn start_task_via_manager(
1507 &self,
1508 task_id: Option<TaskId>,
1509 tool_request: ToolRequest,
1510 kind: TaskLaunchKind,
1511 cancellation: Option<TurnCancellation>,
1512 ) -> impl std::future::Future<Output = Result<TaskStartOutcome, LoopError>> + Send + 'static
1513 {
1514 let task_manager = self.task_manager.clone();
1515 let tool_executor = self.tool_executor.clone();
1516 let permissions = self.permissions.clone();
1517 let resources = self.resources.clone();
1518 let session_id = self.session_id.clone();
1519 let turn_id = tool_request.turn_id.clone();
1520 let metadata = tool_request.metadata.clone();
1521
1522 async move {
1523 task_manager
1524 .start_task(
1525 TaskLaunchRequest {
1526 task_id,
1527 request: tool_request.clone(),
1528 kind,
1529 },
1530 TaskStartContext {
1531 executor: tool_executor.clone(),
1532 tool_context: {
1533 let execution_scope = ToolExecutionScope {
1534 executor: tool_executor,
1535 session_id: session_id.clone(),
1536 turn_id: turn_id.clone(),
1537 permissions: permissions.clone(),
1538 resources: resources.clone(),
1539 cancellation: cancellation.clone(),
1540 };
1541 OwnedToolContext {
1542 session_id,
1543 turn_id,
1544 metadata,
1545 permissions,
1546 resources,
1547 cancellation,
1548 execution_scope: Some(execution_scope),
1549 approved_request: None,
1550 }
1551 },
1552 },
1553 )
1554 .await
1555 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))
1556 }
1557 }
1558
1559 fn register_tool_cancellation(
1560 &mut self,
1561 call_id: &ToolCallId,
1562 cancellation: Option<TurnCancellation>,
1563 ) {
1564 if let Some(cancellation) = cancellation {
1565 self.tool_cancellations
1566 .insert(call_id.clone(), cancellation);
1567 }
1568 }
1569
1570 fn tool_cancellation_for(
1571 &mut self,
1572 call_id: &ToolCallId,
1573 fallback: Option<TurnCancellation>,
1574 ) -> Option<TurnCancellation> {
1575 self.tool_cancellations.get(call_id).cloned().or(fallback)
1576 }
1577
1578 fn clear_tool_cancellation(&mut self, call_id: &ToolCallId) {
1579 self.tool_cancellations.remove(call_id);
1580 }
1581
1582 fn has_pending_interrupts(&self) -> bool {
1583 !self.pending_approvals.is_empty()
1584 }
1585
1586 fn emit_tool_catalog_events(&mut self, events: Vec<ToolCatalogEvent>) {
1587 for event in events {
1588 self.emit(AgentEvent::ToolCatalogChanged(event));
1589 }
1590 }
1591
1592 fn enqueue_pending_approval(
1593 &mut self,
1594 turn_id: &agentkit_core::TurnId,
1595 task: TaskApproval,
1596 cancellation: Option<TurnCancellation>,
1597 ) {
1598 let call_id = task.tool_request.call_id.clone();
1599 let cancellation = self.tool_cancellation_for(&call_id, cancellation);
1600 let call = ToolCallPart {
1601 id: call_id.clone(),
1602 name: task.tool_request.tool_name.to_string(),
1603 input: task.tool_request.input.clone(),
1604 metadata: task.tool_request.metadata.clone(),
1605 };
1606 let mut request = task.approval;
1607 request.call_id = Some(call_id.clone());
1608 let pending = PendingApprovalToolCall {
1609 request: request.clone(),
1610 decision: None,
1611 surfaced: false,
1612 turn_id: turn_id.clone(),
1613 task_id: task.task_id,
1614 call,
1615 tool_request: task.tool_request,
1616 cancellation,
1617 };
1618 self.pending_approvals.insert(call_id.clone(), pending);
1619 if !self.pending_approval_order.iter().any(|id| id == &call_id) {
1620 self.pending_approval_order.push_back(call_id);
1621 }
1622 self.emit(AgentEvent::ApprovalRequired(request));
1623 }
1624
1625 fn take_next_unsurfaced_approval_interrupt(&mut self) -> Option<LoopStep> {
1626 for call_id in self.pending_approval_order.clone() {
1627 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
1628 continue;
1629 };
1630 if pending.decision.is_none() && !pending.surfaced {
1631 pending.surfaced = true;
1632 return Some(LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(
1633 PendingApproval {
1634 request: pending.request.clone(),
1635 },
1636 )));
1637 }
1638 }
1639 None
1640 }
1641
1642 fn next_unresolved_approval_interrupt(&self) -> Option<LoopStep> {
1643 self.pending_approval_order.iter().find_map(|call_id| {
1644 self.pending_approvals.get(call_id).and_then(|pending| {
1645 pending.decision.is_none().then(|| {
1646 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(PendingApproval {
1647 request: pending.request.clone(),
1648 }))
1649 })
1650 })
1651 })
1652 }
1653
1654 fn take_next_resolved_approval(&mut self) -> Option<PendingApprovalToolCall> {
1655 let call_id = self.pending_approval_order.iter().find_map(|call_id| {
1656 self.pending_approvals
1657 .get(call_id)
1658 .and_then(|pending| pending.decision.as_ref().map(|_| call_id.clone()))
1659 })?;
1660 self.pending_approval_order.retain(|id| id != &call_id);
1661 self.pending_approvals.remove(&call_id)
1662 }
1663
1664 fn queue_resolution_interrupt(
1665 &mut self,
1666 turn_id: &agentkit_core::TurnId,
1667 resolution: TaskResolution,
1668 cancellation: Option<TurnCancellation>,
1669 ) -> Option<LoopStep> {
1670 match resolution {
1671 TaskResolution::Item(item) => {
1672 self.append_tool_result_item(item);
1673 None
1674 }
1675 TaskResolution::Approval(task) => {
1676 self.enqueue_pending_approval(turn_id, task, cancellation);
1677 self.take_next_unsurfaced_approval_interrupt()
1678 }
1679 }
1680 }
1681
1682 async fn drain_pending_loop_updates(&mut self) -> Result<(bool, Option<LoopStep>), LoopError> {
1683 let PendingLoopUpdates { mut resolutions } = self
1684 .task_manager
1685 .take_pending_loop_updates()
1686 .await
1687 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
1688 let mut saw_items = false;
1689 while let Some(resolution) = resolutions.pop_front() {
1690 match resolution {
1691 TaskResolution::Item(item) => {
1692 self.append_tool_result_item(item);
1693 saw_items = true;
1694 }
1695 TaskResolution::Approval(task) => {
1696 self.enqueue_pending_approval(&task.tool_request.turn_id.clone(), task, None);
1697 }
1698 }
1699 }
1700 if let Some(step) = self.finish_cancelled_pending_approval().await? {
1701 return Ok((saw_items, Some(step)));
1702 }
1703 Ok((saw_items, self.take_next_unsurfaced_approval_interrupt()))
1704 }
1705
1706 async fn finish_cancelled_pending_approval(&mut self) -> Result<Option<LoopStep>, LoopError> {
1707 if self.pending_approvals.is_empty() {
1708 return Ok(None);
1709 }
1710 if !self.pending_approvals.values().any(|pending| {
1711 pending
1712 .cancellation
1713 .as_ref()
1714 .is_some_and(TurnCancellation::is_cancelled)
1715 }) {
1716 return Ok(None);
1717 }
1718 self.cancel_pending_approvals().await
1719 }
1720
1721 async fn run_mutators(
1722 &mut self,
1723 point: MutationPoint,
1724 turn_id: Option<&agentkit_core::TurnId>,
1725 cancellation: Option<TurnCancellation>,
1726 ) -> Result<(), LoopError> {
1727 if self.mutators.is_empty() {
1728 return Ok(());
1729 }
1730 if cancellation
1731 .as_ref()
1732 .is_some_and(TurnCancellation::is_cancelled)
1733 {
1734 return Err(LoopError::Cancelled);
1735 }
1736 let mutators = self.mutators.clone();
1737 let session_id = self.session_id.clone();
1738 let observed_session_id = Arc::clone(&self.observed_session_id);
1739 let observers = self.observers.clone();
1740 let emitter = DriverEmitter {
1741 session_id: &observed_session_id,
1742 observers: &observers,
1743 };
1744 let mut cursor = TranscriptCursor {
1745 items: &mut self.transcript,
1746 dirty: false,
1747 };
1748 for mutator in &mutators {
1749 if cancellation
1750 .as_ref()
1751 .is_some_and(TurnCancellation::is_cancelled)
1752 {
1753 return Err(LoopError::Cancelled);
1754 }
1755 let ctx = LoopCtx {
1756 session_id: &session_id,
1757 turn_id,
1758 point,
1759 cancellation: cancellation.clone(),
1760 emitter: &emitter,
1761 };
1762 mutator.mutate(&mut cursor, ctx).await?;
1763 }
1764 if cursor.dirty {
1765 validate_transcript_invariants(cursor.items)?;
1766 }
1767 Ok(())
1768 }
1769
1770 async fn continue_active_tool_round(&mut self) -> Result<Option<LoopStep>, LoopError> {
1771 let Some(_) = self.active_tool_round.as_ref() else {
1772 return Ok(None);
1773 };
1774 loop {
1775 let turn_id = self
1776 .active_tool_round
1777 .as_ref()
1778 .map(|active| active.turn_id.clone())
1779 .ok_or_else(|| LoopError::InvalidState("missing active tool round".into()))?;
1780 let cancellation = self
1781 .active_tool_round
1782 .as_ref()
1783 .and_then(|active| active.cancellation.clone());
1784
1785 if cancellation
1786 .as_ref()
1787 .is_some_and(TurnCancellation::is_cancelled)
1788 {
1789 self.task_manager
1790 .on_turn_interrupted(&turn_id)
1791 .await
1792 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
1793 self.active_tool_round = None;
1794 return self.finish_cancelled(turn_id, Vec::new()).map(Some);
1795 }
1796
1797 let next_call = self
1798 .active_tool_round
1799 .as_mut()
1800 .and_then(|active| active.pending_calls.pop_front());
1801 if let Some((call, tool_request)) = next_call {
1802 use tracing::Instrument;
1803 self.register_tool_cancellation(&call.id, cancellation.clone());
1804 let dispatch_span = self.execute_tool_span(&tool_request, &turn_id, "plain");
1805 match self
1806 .start_task_via_manager(
1807 None,
1808 tool_request.clone(),
1809 TaskLaunchKind::Plain,
1810 cancellation.clone(),
1811 )
1812 .instrument(dispatch_span.clone())
1813 .await?
1814 {
1815 TaskStartOutcome::Ready(resolution) => {
1816 let resolution = *resolution;
1817 match resolution {
1818 TaskResolution::Item(item) => {
1819 if !tool_result_is_permission_denied(&item) {
1820 self.emit(AgentEvent::ToolExecutionStarted(call.clone()));
1821 }
1822 if tool_result_is_error(&item) {
1823 dispatch_span.record("error.type", "tool_error");
1824 }
1825 if let Some(active) = self.active_tool_round.as_mut() {
1826 active.foreground_progressed = true;
1827 }
1828 self.append_tool_result_item(item);
1829 }
1830 TaskResolution::Approval(task) => {
1831 self.enqueue_pending_approval(&turn_id, task, cancellation.clone());
1832 }
1833 }
1834 continue;
1835 }
1836 TaskStartOutcome::Pending { kind, .. } => {
1837 self.emit(AgentEvent::ToolExecutionStarted(call.clone()));
1838 if kind == agentkit_task_manager::TaskKind::Background
1839 && let Some(active) = self.active_tool_round.as_mut()
1840 {
1841 active.background_pending = true;
1842 }
1843 continue;
1844 }
1845 }
1846 }
1847
1848 match self
1849 .task_manager
1850 .wait_for_turn(&turn_id, cancellation.clone())
1851 .await
1852 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?
1853 {
1854 Some(TurnTaskUpdate::Resolution(resolution)) => {
1855 let resolution = *resolution;
1856 match resolution {
1857 TaskResolution::Item(item) => {
1858 if let Some(active) = self.active_tool_round.as_mut() {
1859 active.foreground_progressed = true;
1860 }
1861 self.append_tool_result_item(item);
1862 }
1863 TaskResolution::Approval(task) => {
1864 self.enqueue_pending_approval(&turn_id, task, cancellation.clone());
1865 }
1866 }
1867 }
1868 Some(TurnTaskUpdate::Detached(snapshot)) => {
1869 let detached_call_id = snapshot.call_id.clone();
1887 let detached_result = ToolResultPart {
1888 call_id: detached_call_id.clone(),
1889 output: ToolOutput::Text(format!(
1890 "Tool {} is now running in the background. \
1891 The result will be delivered when it completes.",
1892 snapshot.tool_name,
1893 )),
1894 is_error: false,
1895 metadata: MetadataMap::new(),
1896 };
1897 self.emit(AgentEvent::ToolExecutionProgress(detached_result.clone()));
1898 self.append_item(Item {
1899 id: None,
1900 kind: ItemKind::Tool,
1901 parts: vec![Part::ToolResult(detached_result)],
1902 metadata: MetadataMap::new(),
1903 usage: None,
1904 finish_reason: None,
1905 created_at: None,
1906 });
1907 self.detached_call_ids.insert(detached_call_id);
1908 if let Some(active) = self.active_tool_round.as_mut() {
1909 active.background_pending = true;
1910 active.foreground_progressed = true;
1911 }
1912 }
1913 None => {
1914 if cancellation
1915 .as_ref()
1916 .is_some_and(TurnCancellation::is_cancelled)
1917 {
1918 self.task_manager
1919 .on_turn_interrupted(&turn_id)
1920 .await
1921 .map_err(|error| {
1922 LoopError::Tool(ToolError::Internal(error.to_string()))
1923 })?;
1924 self.active_tool_round = None;
1925 return self.finish_cancelled(turn_id, Vec::new()).map(Some);
1926 }
1927 let active = self.active_tool_round.take().ok_or_else(|| {
1928 LoopError::InvalidState("missing active tool round".into())
1929 })?;
1930 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
1931 return Ok(Some(step));
1932 }
1933 if let Some(step) = self.next_unresolved_approval_interrupt() {
1934 return Ok(Some(step));
1935 }
1936 if active.background_pending && !active.foreground_progressed {
1937 return Ok(None);
1938 }
1939 let info = ToolRoundInfo {
1946 session_id: self.session_id.clone(),
1947 turn_id: turn_id.clone(),
1948 transcript_len: self.transcript.len(),
1949 };
1950 self.pending_round_resume = Some(turn_id);
1951 return Ok(Some(LoopStep::Interrupt(LoopInterrupt::AfterToolResult(
1952 info,
1953 ))));
1954 }
1955 }
1956 }
1957 }
1958
1959 #[tracing::instrument(
1960 name = "agent.turn",
1961 skip_all,
1962 fields(
1963 otel.name = "invoke_agent",
1964 gen_ai.operation.name = "invoke_agent",
1965 gen_ai.conversation.id = %self.session_id,
1966 gen_ai.provider.name = tracing::field::Empty,
1967 gen_ai.usage.input_tokens = tracing::field::Empty,
1968 gen_ai.usage.output_tokens = tracing::field::Empty,
1969 gen_ai.usage.cost = tracing::field::Empty,
1970 session.id = %self.session_id,
1971 turn.id = %turn_id,
1972 transcript.len = self.transcript.len(),
1973 saw_tool_call = tracing::field::Empty,
1974 finish_reason = tracing::field::Empty,
1975 ),
1976 )]
1977 async fn drive_turn(
1978 &mut self,
1979 turn_id: agentkit_core::TurnId,
1980 emit_started: bool,
1981 mutation_point: MutationPoint,
1982 ) -> Result<LoopStep, LoopError> {
1983 if let Some(provider) = &self.provider_name {
1984 tracing::Span::current().record("gen_ai.provider.name", provider.as_str());
1985 }
1986 let cancellation = self
1987 .cancellation
1988 .as_ref()
1989 .map(CancellationHandle::checkpoint);
1990 match self
1991 .run_mutators(mutation_point, Some(&turn_id), cancellation.clone())
1992 .await
1993 {
1994 Ok(()) => {}
1995 Err(LoopError::Cancelled) => {
1996 return self.finish_cancelled(turn_id, interrupted_assistant_items());
1997 }
1998 Err(error) => return Err(error),
1999 }
2000
2001 if !transcript_has_pending_input(&self.transcript) {
2007 let turn_result = TurnResult {
2008 turn_id,
2009 finish_reason: FinishReason::Completed,
2010 items: Vec::new(),
2011 usage: None,
2012 metadata: MetadataMap::new(),
2013 };
2014 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
2015 return Ok(LoopStep::Finished(turn_result));
2016 }
2017
2018 if emit_started {
2019 self.emit(AgentEvent::TurnStarted {
2020 session_id: self.session_id.clone(),
2021 turn_id: turn_id.clone(),
2022 });
2023 }
2024 if cancellation
2025 .as_ref()
2026 .is_some_and(TurnCancellation::is_cancelled)
2027 {
2028 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2029 }
2030
2031 let catalog_events = self.tool_executor.drain_catalog_events();
2032 self.emit_tool_catalog_events(catalog_events);
2033
2034 let request = TurnRequest {
2035 session_id: self.session_id.clone(),
2036 turn_id: turn_id.clone(),
2037 transcript: self.transcript.clone(),
2038 available_tools: self.tool_executor.specs(),
2039 cache: self
2040 .next_turn_cache
2041 .take()
2042 .or_else(|| self.default_cache.clone()),
2043 metadata: MetadataMap::new(),
2044 };
2045
2046 let session = self
2047 .session
2048 .as_mut()
2049 .ok_or_else(|| LoopError::InvalidState("model session is not available".into()))?;
2050
2051 let chat_span = tracing::info_span!(
2058 "chat",
2059 "otel.name" = tracing::field::Empty,
2060 "otel.kind" = "client",
2061 "gen_ai.operation.name" = "chat",
2062 "gen_ai.provider.name" = tracing::field::Empty,
2063 "gen_ai.conversation.id" = %self.session_id,
2064 "gen_ai.request.model" = tracing::field::Empty,
2065 "gen_ai.response.model" = tracing::field::Empty,
2066 "gen_ai.response.id" = tracing::field::Empty,
2067 "gen_ai.response.finish_reasons" = tracing::field::Empty,
2068 "gen_ai.usage.input_tokens" = tracing::field::Empty,
2069 "gen_ai.usage.output_tokens" = tracing::field::Empty,
2070 "gen_ai.usage.cost" = tracing::field::Empty,
2071 );
2072 if let Some(provider) = &self.provider_name {
2073 chat_span.record("gen_ai.provider.name", provider.as_str());
2074 }
2075 match session.model_name() {
2076 Some(model) => {
2077 chat_span.record("gen_ai.request.model", model);
2078 chat_span.record("otel.name", format!("chat {model}").as_str());
2079 }
2080 None => {
2081 chat_span.record("otel.name", "chat");
2082 }
2083 }
2084
2085 use tracing::Instrument;
2086 let mut turn = match session
2087 .begin_turn(request, cancellation.clone())
2088 .instrument(chat_span.clone())
2089 .await
2090 {
2091 Ok(turn) => turn,
2092 Err(LoopError::Cancelled) => {
2093 self.task_manager
2094 .on_turn_interrupted(&turn_id)
2095 .await
2096 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2097 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2098 }
2099 Err(error) => return Err(error),
2100 };
2101 let mut saw_tool_call = false;
2102 let mut finished_result = None;
2103
2104 while let Some(event) = match turn
2105 .next_event(cancellation.clone())
2106 .instrument(chat_span.clone())
2107 .await
2108 {
2109 Ok(event) => event,
2110 Err(LoopError::Cancelled) => {
2111 self.task_manager
2112 .on_turn_interrupted(&turn_id)
2113 .await
2114 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2115 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2116 }
2117 Err(error) => return Err(error),
2118 } {
2119 if cancellation
2120 .as_ref()
2121 .is_some_and(TurnCancellation::is_cancelled)
2122 {
2123 self.task_manager
2124 .on_turn_interrupted(&turn_id)
2125 .await
2126 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2127 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2128 }
2129 match event {
2130 ModelTurnEvent::Delta(delta) => self.emit(AgentEvent::ContentDelta(delta)),
2131 ModelTurnEvent::Usage(usage) => {
2132 if let Some(tokens) = &usage.tokens {
2133 chat_span.record("gen_ai.usage.input_tokens", tokens.input_tokens);
2134 chat_span.record("gen_ai.usage.output_tokens", tokens.output_tokens);
2135 }
2136 if let Some(cost) = &usage.cost {
2137 chat_span.record("gen_ai.usage.cost", cost.amount);
2138 }
2139 self.emit(AgentEvent::UsageUpdated(usage));
2140 }
2141 ModelTurnEvent::ToolCall(call) => {
2142 saw_tool_call = true;
2143 self.emit(AgentEvent::ToolCallRequested(call.clone()));
2144 }
2145 ModelTurnEvent::Finished(result) => {
2146 finished_result = Some(result);
2147 break;
2148 }
2149 }
2150 }
2151
2152 let mut result = finished_result.ok_or_else(|| {
2153 LoopError::Provider("model turn ended without a Finished event".into())
2154 })?;
2155 if let Some(model) = &result.model {
2156 chat_span.record("gen_ai.response.model", model.as_str());
2157 }
2158 if let Some(id) = &result.response_id {
2159 chat_span.record("gen_ai.response.id", id.as_str());
2160 }
2161 if let Some(tokens) = result
2162 .usage
2163 .as_ref()
2164 .and_then(|usage| usage.tokens.as_ref())
2165 {
2166 chat_span.record("gen_ai.usage.input_tokens", tokens.input_tokens);
2167 chat_span.record("gen_ai.usage.output_tokens", tokens.output_tokens);
2168 }
2169 if let Some(cost) = result.usage.as_ref().and_then(|usage| usage.cost.as_ref()) {
2170 chat_span.record("gen_ai.usage.cost", cost.amount);
2171 }
2172 chat_span.record(
2173 "gen_ai.response.finish_reasons",
2174 tracing::field::debug(&result.finish_reason),
2175 );
2176 drop(chat_span);
2177 tracing::Span::current().record("saw_tool_call", saw_tool_call);
2178 tracing::Span::current().record(
2179 "finish_reason",
2180 tracing::field::debug(&result.finish_reason),
2181 );
2182 if let Some(tokens) = result
2183 .usage
2184 .as_ref()
2185 .and_then(|usage| usage.tokens.as_ref())
2186 {
2187 tracing::Span::current().record("gen_ai.usage.input_tokens", tokens.input_tokens);
2188 tracing::Span::current().record("gen_ai.usage.output_tokens", tokens.output_tokens);
2189 }
2190 if let Some(cost) = result.usage.as_ref().and_then(|usage| usage.cost.as_ref()) {
2191 tracing::Span::current().record("gen_ai.usage.cost", cost.amount);
2192 }
2193 let now = Timestamp::now();
2194 let usage = result.usage.clone();
2195 let finish_reason = result.finish_reason.clone();
2196 let output_items: Vec<Item> = result
2197 .output_items
2198 .drain(..)
2199 .map(|mut item| {
2200 if matches!(item.kind, ItemKind::Assistant) {
2201 if item.usage.is_none() {
2202 item.usage = usage.clone();
2203 }
2204 if item.finish_reason.is_none() {
2205 item.finish_reason = Some(finish_reason.clone());
2206 }
2207 }
2208 if item.created_at.is_none() {
2209 item.created_at = Some(now);
2210 }
2211 item
2212 })
2213 .collect();
2214 self.extend_transcript(output_items.clone());
2215
2216 if saw_tool_call {
2217 let pending_calls = extract_tool_calls(&output_items)
2218 .into_iter()
2219 .map(|call| {
2220 let tool_request = ToolRequest {
2221 call_id: call.id.clone(),
2222 tool_name: agentkit_tools_core::ToolName::new(call.name.clone()),
2223 input: call.input.clone(),
2224 session_id: self.session_id.clone(),
2225 turn_id: turn_id.clone(),
2226 metadata: call.metadata.clone(),
2227 };
2228 (call, tool_request)
2229 })
2230 .collect();
2231 self.active_tool_round = Some(ActiveToolRound {
2232 turn_id: turn_id.clone(),
2233 pending_calls,
2234 cancellation: cancellation.clone(),
2235 background_pending: false,
2236 foreground_progressed: false,
2237 });
2238 if let Some(step) = self.continue_active_tool_round().await? {
2239 return Ok(step);
2240 }
2241 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
2242 InputRequest {
2243 session_id: self.session_id.clone(),
2244 reason: "driver is waiting for input".into(),
2245 },
2246 )));
2247 }
2248
2249 let turn_result = TurnResult {
2250 turn_id,
2251 finish_reason: result.finish_reason,
2252 items: output_items,
2253 usage: result.usage,
2254 metadata: result.metadata,
2255 };
2256 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
2257 Ok(LoopStep::Finished(turn_result))
2258 }
2259
2260 async fn resume_after_approval(
2261 &mut self,
2262 pending: PendingApprovalToolCall,
2263 ) -> Result<LoopStep, LoopError> {
2264 let decision = pending
2265 .decision
2266 .clone()
2267 .ok_or_else(|| LoopError::InvalidState("pending approval has no decision".into()))?;
2268
2269 match decision {
2270 ApprovalDecision::Approve => {
2271 use tracing::Instrument;
2272 self.emit(AgentEvent::ToolExecutionStarted(pending.call.clone()));
2273 let dispatch_span =
2274 self.execute_tool_span(&pending.tool_request, &pending.turn_id, "approved");
2275 let cancellation = self
2276 .cancellation
2277 .as_ref()
2278 .map(CancellationHandle::checkpoint);
2279 self.register_tool_cancellation(&pending.call.id, cancellation.clone());
2280 match self
2281 .start_task_via_manager(
2282 Some(pending.task_id.clone()),
2283 pending.tool_request.clone(),
2284 TaskLaunchKind::Approved(pending.request.clone()),
2285 cancellation.clone(),
2286 )
2287 .instrument(dispatch_span.clone())
2288 .await?
2289 {
2290 TaskStartOutcome::Ready(resolution) => {
2291 let resolution = *resolution;
2292 if let TaskResolution::Item(item) = &resolution
2293 && tool_result_is_error(item)
2294 {
2295 dispatch_span.record("error.type", "tool_error");
2296 }
2297 if let Some(step) = self.queue_resolution_interrupt(
2298 &pending.turn_id,
2299 resolution,
2300 cancellation,
2301 ) {
2302 return Ok(step);
2303 }
2304 }
2305 TaskStartOutcome::Pending { .. } => {}
2306 }
2307 }
2308 ApprovalDecision::Deny { reason } => {
2309 self.append_tool_result_item(Item {
2310 id: None,
2311 kind: ItemKind::Tool,
2312 parts: vec![Part::ToolResult(ToolResultPart {
2313 call_id: pending.call.id.clone(),
2314 output: ToolOutput::Text(
2315 reason.unwrap_or_else(|| "approval denied".into()),
2316 ),
2317 is_error: true,
2318 metadata: pending.call.metadata.clone(),
2319 })],
2320 metadata: MetadataMap::new(),
2321 usage: None,
2322 finish_reason: None,
2323 created_at: None,
2324 });
2325 }
2326 }
2327
2328 if let Some(step) = self.continue_active_tool_round().await? {
2329 Ok(step)
2330 } else if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2331 Ok(step)
2332 } else if let Some(step) = self.next_unresolved_approval_interrupt() {
2333 Ok(step)
2334 } else {
2335 self.drive_turn(pending.turn_id, false, MutationPoint::AfterToolResult)
2336 .await
2337 }
2338 }
2339
2340 fn finish_cancelled(
2341 &mut self,
2342 turn_id: agentkit_core::TurnId,
2343 items: Vec<Item>,
2344 ) -> Result<LoopStep, LoopError> {
2345 self.extend_transcript(items.clone());
2346 let turn_result = TurnResult {
2347 turn_id,
2348 finish_reason: FinishReason::Cancelled,
2349 items,
2350 usage: None,
2351 metadata: interrupted_metadata("turn"),
2352 };
2353 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
2354 Ok(LoopStep::Finished(turn_result))
2355 }
2356
2357 pub fn submit_input(&mut self, input: Vec<Item>) -> Result<(), LoopError> {
2367 if self.has_pending_interrupts() {
2368 return Err(LoopError::InvalidState(
2369 "cannot submit input while an interrupt is pending".into(),
2370 ));
2371 }
2372 self.emit(AgentEvent::InputAccepted {
2373 session_id: self.session_id.clone(),
2374 items: input.clone(),
2375 });
2376 self.pending_input.extend(input);
2377 Ok(())
2378 }
2379
2380 pub fn set_next_turn_cache(&mut self, cache: PromptCacheRequest) -> Result<(), LoopError> {
2385 if self.has_pending_interrupts() {
2386 return Err(LoopError::InvalidState(
2387 "cannot update next-turn cache while an interrupt is pending".into(),
2388 ));
2389 }
2390 self.next_turn_cache = Some(cache);
2391 Ok(())
2392 }
2393
2394 #[cfg(test)]
2395 pub(crate) fn submit_input_with_cache(
2396 &mut self,
2397 input: Vec<Item>,
2398 cache: PromptCacheRequest,
2399 ) -> Result<(), LoopError> {
2400 self.set_next_turn_cache(cache)?;
2401 self.submit_input(input)
2402 }
2403
2404 pub fn resolve_approval_for(
2414 &mut self,
2415 call_id: ToolCallId,
2416 decision: ApprovalDecision,
2417 ) -> Result<(), LoopError> {
2418 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
2419 return Err(LoopError::InvalidState(format!(
2420 "no approval request is pending for call {}",
2421 call_id.0
2422 )));
2423 };
2424 pending.decision = Some(decision.clone());
2425 self.emit(AgentEvent::ApprovalResolved {
2426 approved: matches!(decision, ApprovalDecision::Approve),
2427 });
2428 Ok(())
2429 }
2430
2431 pub fn resolve_approval_for_with_patched_input(
2444 &mut self,
2445 call_id: ToolCallId,
2446 input: serde_json::Value,
2447 ) -> Result<(), LoopError> {
2448 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
2449 return Err(LoopError::InvalidState(format!(
2450 "no approval request is pending for call {}",
2451 call_id.0
2452 )));
2453 };
2454 pending.tool_request.input = input;
2455 self.resolve_approval_for(call_id, ApprovalDecision::Approve)
2456 }
2457
2458 pub fn resolve_approval(&mut self, decision: ApprovalDecision) -> Result<(), LoopError> {
2461 let mut unresolved = self
2462 .pending_approval_order
2463 .iter()
2464 .filter(|call_id| {
2465 self.pending_approvals
2466 .get(*call_id)
2467 .is_some_and(|pending| pending.decision.is_none())
2468 })
2469 .cloned();
2470 let Some(call_id) = unresolved.next() else {
2471 return Err(LoopError::InvalidState(
2472 "no approval request is pending".into(),
2473 ));
2474 };
2475 if unresolved.next().is_some() {
2476 return Err(LoopError::InvalidState(
2477 "multiple approvals are pending; use resolve_approval_for".into(),
2478 ));
2479 }
2480 self.resolve_approval_for(call_id, decision)
2481 }
2482
2483 pub fn cancel_pending_approval_for(&mut self, call_id: ToolCallId) -> Result<(), LoopError> {
2488 let Some(pending) = self.drain_pending_approval_for(&call_id) else {
2489 return Err(LoopError::InvalidState(format!(
2490 "no approval request is pending for call {}",
2491 call_id.0
2492 )));
2493 };
2494 self.reject_drained_approvals(vec![pending]);
2495 Ok(())
2496 }
2497
2498 pub async fn cancel_pending_approvals(&mut self) -> Result<Option<LoopStep>, LoopError> {
2504 if self.pending_approvals.is_empty() {
2505 return Ok(None);
2506 }
2507 let Some(turn_id) = self
2508 .pending_approval_order
2509 .iter()
2510 .find_map(|call_id| self.pending_approvals.get(call_id))
2511 .map(|pending| pending.turn_id.clone())
2512 else {
2513 return Ok(None);
2514 };
2515 let pending = self.drain_pending_approval_items();
2516 self.active_tool_round = None;
2517 self.task_manager
2518 .on_turn_interrupted(&turn_id)
2519 .await
2520 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2521 self.reject_drained_approvals(pending);
2522 self.finish_cancelled(turn_id, Vec::new()).map(Some)
2523 }
2524
2525 pub fn snapshot(&self) -> LoopSnapshot {
2527 LoopSnapshot {
2528 session_id: self.session_id.clone(),
2529 transcript: self.transcript.clone(),
2530 pending_input: self.pending_input.clone(),
2531 }
2532 }
2533
2534 pub async fn next(&mut self) -> Result<LoopStep, LoopError> {
2555 if let Some(pending) = self.take_next_resolved_approval() {
2556 return self.resume_after_approval(pending).await;
2557 }
2558
2559 if let Some(step) = self.finish_cancelled_pending_approval().await? {
2560 return Ok(step);
2561 }
2562
2563 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2564 return Ok(step);
2565 }
2566
2567 if let Some(step) = self.next_unresolved_approval_interrupt() {
2568 return Ok(step);
2569 }
2570
2571 if let Some(step) = self.continue_active_tool_round().await? {
2572 return Ok(step);
2573 }
2574
2575 let (had_loop_updates, loop_step) = self.drain_pending_loop_updates().await?;
2576 if let Some(step) = loop_step {
2577 return Ok(step);
2578 }
2579
2580 if let Some(turn_id) = self.pending_round_resume.take() {
2585 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
2586 self.extend_transcript(drained);
2587 return self
2588 .drive_turn(turn_id, false, MutationPoint::AfterToolResult)
2589 .await;
2590 }
2591
2592 if self.pending_input.is_empty() && !had_loop_updates {
2593 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
2594 InputRequest {
2595 session_id: self.session_id.clone(),
2596 reason: "driver is waiting for input".into(),
2597 },
2598 )));
2599 }
2600
2601 let turn_id = agentkit_core::TurnId::new(format!("turn-{}", self.next_turn_index));
2602 self.next_turn_index += 1;
2603 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
2604 self.extend_transcript(drained);
2605 self.drive_turn(turn_id, true, MutationPoint::AfterTurnEnded)
2606 .await
2607 }
2608
2609 fn emit(&self, event: AgentEvent) {
2610 fan_out_observed_event(&self.observers, &self.observed_session_id, event);
2611 }
2612
2613 fn append_item(&mut self, mut item: Item) {
2618 if item.created_at.is_none() {
2619 item.created_at = Some(Timestamp::now());
2620 }
2621 for observer in &self.transcript_observers {
2622 observer.on_transcript_event(TranscriptEvent {
2623 session_id: &self.session_id,
2624 item: &item,
2625 });
2626 }
2627 self.transcript.push(item);
2628 }
2629
2630 fn append_tool_result_item(&mut self, item: Item) {
2644 for part in &item.parts {
2645 if let Part::ToolResult(result) = part {
2646 self.emit(AgentEvent::ToolResultReceived(result.clone()));
2647 self.clear_tool_cancellation(&result.call_id);
2648 }
2649 }
2650 let item = self.maybe_convert_detached(item);
2651 self.append_item(item);
2652 }
2653
2654 fn drain_pending_approval_for(
2655 &mut self,
2656 call_id: &ToolCallId,
2657 ) -> Option<PendingApprovalToolCall> {
2658 let pending = self.pending_approvals.remove(call_id)?;
2659 self.pending_approval_order.retain(|id| id != call_id);
2660 self.clear_tool_cancellation(call_id);
2661 Some(pending)
2662 }
2663
2664 fn drain_pending_approval_items(&mut self) -> Vec<PendingApprovalToolCall> {
2665 let order = std::mem::take(&mut self.pending_approval_order);
2666 let pending = order
2667 .iter()
2668 .filter_map(|call_id| {
2669 let pending = self.pending_approvals.remove(call_id);
2670 self.clear_tool_cancellation(call_id);
2671 pending
2672 })
2673 .collect();
2674 self.pending_approvals.clear();
2675 pending
2676 }
2677
2678 fn reject_drained_approvals(&mut self, pending: Vec<PendingApprovalToolCall>) {
2679 for pending in pending {
2680 self.emit(AgentEvent::ApprovalResolved { approved: false });
2681 self.append_tool_result_item(cancelled_approval_item(pending));
2682 }
2683 }
2684
2685 fn maybe_convert_detached(&mut self, item: Item) -> Item {
2686 if !matches!(item.kind, ItemKind::Tool) {
2687 return item;
2688 }
2689 let results: Vec<&ToolResultPart> = item
2690 .parts
2691 .iter()
2692 .filter_map(|p| match p {
2693 Part::ToolResult(r) => Some(r),
2694 _ => None,
2695 })
2696 .collect();
2697 if results.is_empty()
2698 || !results
2699 .iter()
2700 .all(|r| self.detached_call_ids.contains(&r.call_id))
2701 {
2702 return item;
2703 }
2704 let mut text = String::new();
2705 for result in &results {
2706 self.detached_call_ids.remove(&result.call_id);
2707 if !text.is_empty() {
2708 text.push_str("\n\n");
2709 }
2710 let label = if result.is_error {
2711 "failed"
2712 } else {
2713 "completed"
2714 };
2715 let body = render_tool_output_brief(&result.output);
2716 text.push_str(&format!(
2717 "Background tool call {} {}: {body}",
2718 result.call_id.0, label
2719 ));
2720 }
2721 Item::notification(text)
2722 }
2723
2724 fn extend_transcript(&mut self, items: impl IntoIterator<Item = Item>) {
2728 let now = Timestamp::now();
2729 for mut item in items {
2730 if item.created_at.is_none() {
2731 item.created_at = Some(now);
2732 }
2733 self.append_item(item);
2734 }
2735 }
2736}
2737
2738fn render_tool_output_brief(output: &ToolOutput) -> String {
2739 match output {
2740 ToolOutput::Text(t) => t.clone(),
2741 ToolOutput::Structured(value) => value.to_string(),
2742 ToolOutput::Parts(parts) => format!("[{} parts]", parts.len()),
2743 ToolOutput::Files(files) => format!("[{} files]", files.len()),
2744 }
2745}
2746
2747fn interrupted_metadata(stage: &str) -> MetadataMap {
2748 let mut metadata = MetadataMap::new();
2749 metadata.insert(INTERRUPTED_METADATA_KEY.into(), true.into());
2750 metadata.insert(
2751 INTERRUPT_REASON_METADATA_KEY.into(),
2752 USER_CANCELLED_REASON.into(),
2753 );
2754 metadata.insert(INTERRUPT_STAGE_METADATA_KEY.into(), stage.into());
2755 metadata
2756}
2757
2758fn interrupted_assistant_items() -> Vec<Item> {
2759 vec![Item {
2760 id: None,
2761 kind: ItemKind::Assistant,
2762 parts: vec![Part::Text(TextPart {
2763 text: "Previous assistant response was interrupted by the user before completion."
2764 .into(),
2765 metadata: interrupted_metadata("assistant"),
2766 })],
2767 metadata: interrupted_metadata("assistant"),
2768 usage: None,
2769 finish_reason: None,
2770 created_at: None,
2771 }]
2772}
2773
2774fn cancelled_approval_item(pending: PendingApprovalToolCall) -> Item {
2775 Item {
2776 id: None,
2777 kind: ItemKind::Tool,
2778 parts: vec![Part::ToolResult(ToolResultPart {
2779 call_id: pending.call.id,
2780 output: ToolOutput::Text("approval cancelled".into()),
2781 is_error: true,
2782 metadata: pending.call.metadata,
2783 })],
2784 metadata: MetadataMap::new(),
2785 usage: None,
2786 finish_reason: None,
2787 created_at: None,
2788 }
2789}
2790
2791fn transcript_has_pending_input(transcript: &[Item]) -> bool {
2797 matches!(
2798 transcript.last().map(|item| item.kind),
2799 Some(ItemKind::User | ItemKind::Tool | ItemKind::Notification)
2800 )
2801}
2802
2803fn extract_tool_calls(items: &[Item]) -> Vec<ToolCallPart> {
2804 let mut calls = Vec::new();
2805 for item in items {
2806 for part in &item.parts {
2807 if let Part::ToolCall(call) = part {
2808 calls.push(call.clone());
2809 }
2810 }
2811 }
2812 calls
2813}
2814
2815fn tool_result_is_error(item: &Item) -> bool {
2816 item.parts
2817 .iter()
2818 .any(|part| matches!(part, Part::ToolResult(result) if result.is_error))
2819}
2820
2821fn tool_result_is_permission_denied(item: &Item) -> bool {
2822 item.parts.iter().any(|part| {
2823 matches!(
2824 part,
2825 Part::ToolResult(result)
2826 if result
2827 .metadata
2828 .get(TOOL_RESULT_FAILURE_KIND_METADATA_KEY)
2829 .and_then(Value::as_str)
2830 == Some(TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED)
2831 )
2832 })
2833}
2834
2835#[derive(Debug, Error)]
2837pub enum LoopError {
2838 #[error("invalid driver state: {0}")]
2840 InvalidState(String),
2841 #[error("turn cancelled")]
2843 Cancelled,
2844 #[error("provider error: {0}")]
2846 Provider(String),
2847 #[error("tool error: {0}")]
2849 Tool(#[from] ToolError),
2850 #[error("mutator error: {0}")]
2852 Mutator(String),
2853 #[error("unsupported operation: {0}")]
2855 Unsupported(String),
2856}
2857
2858struct DriverEmitter<'a> {
2863 session_id: &'a Arc<SessionId>,
2864 observers: &'a [Arc<dyn LoopObserver>],
2865}
2866
2867impl<'a> EventEmitter for DriverEmitter<'a> {
2868 fn emit(&self, event: AgentEvent) {
2869 fan_out_observed_event(self.observers, self.session_id, event);
2870 }
2871}
2872
2873fn fan_out_observed_event(
2874 observers: &[Arc<dyn LoopObserver>],
2875 session_id: &Arc<SessionId>,
2876 event: AgentEvent,
2877) {
2878 if observers.is_empty() {
2879 return;
2880 }
2881 let observed = ObservedEvent {
2882 session_id: Arc::clone(session_id),
2883 event,
2884 };
2885 let last = observers.len() - 1;
2886 for observer in &observers[..last] {
2887 observer.handle_event(observed.clone());
2888 }
2889 observers[last].handle_event(observed);
2890}
2891
2892fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> {
2897 let mut pending: HashSet<ToolCallId> = HashSet::new();
2898 let mut seen_calls: HashSet<ToolCallId> = HashSet::new();
2899 let mut seen_results: HashSet<ToolCallId> = HashSet::new();
2900 for item in transcript {
2901 for part in &item.parts {
2902 match part {
2903 Part::ToolCall(call) => {
2904 if !seen_calls.insert(call.id.clone()) {
2905 return Err(LoopError::Mutator(format!(
2906 "transcript invariant violation: duplicate tool_use: {}",
2907 call.id.0
2908 )));
2909 }
2910 pending.insert(call.id.clone());
2911 }
2912 Part::ToolResult(result) => {
2913 if !pending.remove(&result.call_id) {
2914 let kind = if seen_results.contains(&result.call_id) {
2915 "duplicate"
2916 } else {
2917 "orphaned"
2918 };
2919 return Err(LoopError::Mutator(format!(
2920 "transcript invariant violation: {kind} tool_result: {}",
2921 result.call_id.0
2922 )));
2923 }
2924 seen_results.insert(result.call_id.clone());
2925 }
2926 _ => {}
2927 }
2928 }
2929 }
2930 if !pending.is_empty() {
2931 let missing: Vec<String> = pending.into_iter().map(|id| id.0).collect();
2932 return Err(LoopError::Mutator(format!(
2933 "transcript invariant violation: tool_use(s) without matching tool_result: {}",
2934 missing.join(", ")
2935 )));
2936 }
2937 Ok(())
2938}
2939
2940#[cfg(test)]
2941mod tests {
2942 use std::collections::VecDeque;
2943 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
2944 use std::sync::{Arc as StdArc, Mutex as StdMutex};
2945
2946 use agentkit_core::{
2947 CancellationController, ItemKind, Part, TextPart, ToolCallId, ToolCallPart, ToolOutput,
2948 ToolResultPart,
2949 };
2950 use agentkit_task_manager::{
2951 AsyncTaskManager, RoutingDecision, TaskEvent, TaskManager, TaskManagerHandle,
2952 TaskRoutingPolicy,
2953 };
2954 use agentkit_tools_core::{
2955 FileSystemPermissionRequest, PermissionCode, PermissionDecision, PermissionDenial, Tool,
2956 ToolAnnotations, ToolCatalogEvent, ToolExecutionOutcome, ToolName, ToolRegistry,
2957 ToolResult, ToolSpec,
2958 };
2959 use serde_json::{Value, json};
2960 use tokio::sync::Notify;
2961 use tokio::time::{Duration, timeout};
2962
2963 use super::*;
2964
2965 struct FakeAdapter;
2966 struct SlowAdapter;
2967 struct RecordingAdapter {
2968 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
2969 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
2970 }
2971 struct MultiToolAdapter;
2972 struct DualApprovalAdapter;
2973
2974 struct FakeSession;
2975 struct SlowSession;
2976 struct RecordingSession {
2977 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
2978 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
2979 }
2980 struct MultiToolSession;
2981 struct DualApprovalSession;
2982
2983 struct FakeTurn {
2984 events: VecDeque<ModelTurnEvent>,
2985 }
2986
2987 struct SlowTurn {
2988 emitted: bool,
2989 }
2990
2991 struct RecordingTurn {
2992 emitted: bool,
2993 }
2994 struct MultiToolTurn {
2995 events: VecDeque<ModelTurnEvent>,
2996 }
2997 struct DualApprovalTurn {
2998 events: VecDeque<ModelTurnEvent>,
2999 }
3000
3001 struct DelayedApprovalExecutor {
3002 entered: StdArc<AtomicBool>,
3003 release: StdArc<Notify>,
3004 spec: ToolSpec,
3005 }
3006
3007 impl DelayedApprovalExecutor {
3008 fn new(entered: StdArc<AtomicBool>, release: StdArc<Notify>) -> Self {
3009 Self {
3010 entered,
3011 release,
3012 spec: ToolSpec {
3013 name: ToolName::new("echo"),
3014 description: "delayed approval".into(),
3015 input_schema: json!({
3016 "type": "object",
3017 "properties": {
3018 "value": { "type": "string" }
3019 },
3020 "required": ["value"],
3021 "additionalProperties": false
3022 }),
3023 output_schema: None,
3024 annotations: ToolAnnotations::default(),
3025 metadata: MetadataMap::new(),
3026 },
3027 }
3028 }
3029 }
3030
3031 #[async_trait]
3032 impl ToolExecutor for DelayedApprovalExecutor {
3033 fn specs(&self) -> Vec<ToolSpec> {
3034 vec![self.spec.clone()]
3035 }
3036
3037 async fn execute(
3038 &self,
3039 request: ToolRequest,
3040 _ctx: &mut ToolContext<'_>,
3041 ) -> ToolExecutionOutcome {
3042 self.entered.store(true, Ordering::SeqCst);
3043 self.release.notified().await;
3044 ToolExecutionOutcome::Interrupted(
3045 agentkit_tools_core::ToolInterruption::ApprovalRequired(ApprovalRequest {
3046 task_id: None,
3047 call_id: None,
3048 id: "approval:delayed".into(),
3049 request_kind: "delayed.approval".into(),
3050 reason: agentkit_tools_core::ApprovalReason::PolicyRequiresConfirmation,
3051 summary: "delayed approval".into(),
3052 metadata: request.metadata,
3053 }),
3054 )
3055 }
3056 }
3057
3058 #[async_trait]
3059 impl ModelAdapter for FakeAdapter {
3060 type Session = FakeSession;
3061
3062 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
3063 Ok(FakeSession)
3064 }
3065 }
3066
3067 #[async_trait]
3068 impl ModelAdapter for SlowAdapter {
3069 type Session = SlowSession;
3070
3071 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
3072 Ok(SlowSession)
3073 }
3074 }
3075
3076 #[async_trait]
3077 impl ModelAdapter for RecordingAdapter {
3078 type Session = RecordingSession;
3079
3080 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
3081 Ok(RecordingSession {
3082 seen_descriptions: self.seen_descriptions.clone(),
3083 seen_caches: self.seen_caches.clone(),
3084 })
3085 }
3086 }
3087
3088 #[async_trait]
3089 impl ModelAdapter for MultiToolAdapter {
3090 type Session = MultiToolSession;
3091
3092 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
3093 Ok(MultiToolSession)
3094 }
3095 }
3096
3097 #[async_trait]
3098 impl ModelAdapter for DualApprovalAdapter {
3099 type Session = DualApprovalSession;
3100
3101 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
3102 Ok(DualApprovalSession)
3103 }
3104 }
3105
3106 #[async_trait]
3107 impl ModelSession for FakeSession {
3108 type Turn = FakeTurn;
3109
3110 async fn begin_turn(
3111 &mut self,
3112 request: TurnRequest,
3113 _cancellation: Option<TurnCancellation>,
3114 ) -> Result<Self::Turn, LoopError> {
3115 let has_tool_result = request.transcript.iter().any(|item| {
3116 item.kind == ItemKind::Tool
3117 && item
3118 .parts
3119 .iter()
3120 .any(|part| matches!(part, Part::ToolResult(_)))
3121 });
3122 let tool_name = request
3123 .available_tools
3124 .first()
3125 .map(|tool| tool.name.0.clone())
3126 .unwrap_or_else(|| "echo".into());
3127
3128 let events = if has_tool_result {
3129 let result_text = request
3130 .transcript
3131 .iter()
3132 .rev()
3133 .find_map(|item| {
3134 item.parts.iter().find_map(|part| match part {
3135 Part::ToolResult(ToolResultPart {
3136 output: ToolOutput::Text(text),
3137 ..
3138 }) => Some(text.clone()),
3139 _ => None,
3140 })
3141 })
3142 .unwrap_or_else(|| "missing".into());
3143
3144 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
3145 model: None,
3146 response_id: None,
3147 finish_reason: FinishReason::Completed,
3148 output_items: vec![Item {
3149 id: None,
3150 kind: ItemKind::Assistant,
3151 parts: vec![Part::Text(TextPart {
3152 text: format!("tool said: {result_text}"),
3153 metadata: MetadataMap::new(),
3154 })],
3155 metadata: MetadataMap::new(),
3156 usage: None,
3157 finish_reason: None,
3158 created_at: None,
3159 }],
3160 usage: None,
3161 metadata: MetadataMap::new(),
3162 })])
3163 } else {
3164 VecDeque::from([
3165 ModelTurnEvent::ToolCall(agentkit_core::ToolCallPart {
3166 id: ToolCallId::new("call-1"),
3167 name: tool_name.clone(),
3168 input: json!({ "value": "pong" }),
3169 metadata: MetadataMap::new(),
3170 }),
3171 ModelTurnEvent::Finished(ModelTurnResult {
3172 model: None,
3173 response_id: None,
3174 finish_reason: FinishReason::ToolCall,
3175 output_items: vec![Item {
3176 id: None,
3177 kind: ItemKind::Assistant,
3178 parts: vec![Part::ToolCall(agentkit_core::ToolCallPart {
3179 id: ToolCallId::new("call-1"),
3180 name: tool_name,
3181 input: json!({ "value": "pong" }),
3182 metadata: MetadataMap::new(),
3183 })],
3184 metadata: MetadataMap::new(),
3185 usage: None,
3186 finish_reason: None,
3187 created_at: None,
3188 }],
3189 usage: None,
3190 metadata: MetadataMap::new(),
3191 }),
3192 ])
3193 };
3194
3195 Ok(FakeTurn { events })
3196 }
3197 }
3198
3199 #[async_trait]
3200 impl ModelSession for SlowSession {
3201 type Turn = SlowTurn;
3202
3203 async fn begin_turn(
3204 &mut self,
3205 request: TurnRequest,
3206 cancellation: Option<TurnCancellation>,
3207 ) -> Result<Self::Turn, LoopError> {
3208 let should_block = request
3209 .transcript
3210 .iter()
3211 .rev()
3212 .find(|item| item.kind == ItemKind::User)
3213 .is_some_and(|item| {
3214 item.parts.iter().any(|part| match part {
3215 Part::Text(text) => text.text == "do the long task",
3216 _ => false,
3217 })
3218 });
3219
3220 if should_block && let Some(cancellation) = cancellation {
3221 cancellation.cancelled().await;
3222 return Err(LoopError::Cancelled);
3223 }
3224
3225 Ok(SlowTurn { emitted: false })
3226 }
3227 }
3228
3229 #[async_trait]
3230 impl ModelSession for RecordingSession {
3231 type Turn = RecordingTurn;
3232
3233 async fn begin_turn(
3234 &mut self,
3235 request: TurnRequest,
3236 _cancellation: Option<TurnCancellation>,
3237 ) -> Result<Self::Turn, LoopError> {
3238 let descriptions = request
3239 .available_tools
3240 .iter()
3241 .map(|tool| tool.description.clone())
3242 .collect::<Vec<_>>();
3243 self.seen_descriptions.lock().unwrap().push(descriptions);
3244 self.seen_caches.lock().unwrap().push(request.cache.clone());
3245
3246 Ok(RecordingTurn { emitted: false })
3247 }
3248 }
3249
3250 #[async_trait]
3251 impl ModelSession for MultiToolSession {
3252 type Turn = MultiToolTurn;
3253
3254 async fn begin_turn(
3255 &mut self,
3256 request: TurnRequest,
3257 _cancellation: Option<TurnCancellation>,
3258 ) -> Result<Self::Turn, LoopError> {
3259 let has_tool_result = request.transcript.iter().any(|item| {
3260 item.kind == ItemKind::Tool
3261 && item
3262 .parts
3263 .iter()
3264 .any(|part| matches!(part, Part::ToolResult(_)))
3265 });
3266
3267 let events = if has_tool_result {
3268 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
3269 model: None,
3270 response_id: None,
3271 finish_reason: FinishReason::Completed,
3272 output_items: vec![Item {
3273 id: None,
3274 kind: ItemKind::Assistant,
3275 parts: vec![Part::Text(TextPart {
3276 text: "mixed tools finished".into(),
3277 metadata: MetadataMap::new(),
3278 })],
3279 metadata: MetadataMap::new(),
3280 usage: None,
3281 finish_reason: None,
3282 created_at: None,
3283 }],
3284 usage: None,
3285 metadata: MetadataMap::new(),
3286 })])
3287 } else {
3288 let foreground = agentkit_core::ToolCallPart {
3289 id: ToolCallId::new("call-foreground"),
3290 name: "foreground-wait".into(),
3291 input: json!({}),
3292 metadata: MetadataMap::new(),
3293 };
3294 let background = agentkit_core::ToolCallPart {
3295 id: ToolCallId::new("call-background"),
3296 name: "background-wait".into(),
3297 input: json!({}),
3298 metadata: MetadataMap::new(),
3299 };
3300 VecDeque::from([
3301 ModelTurnEvent::ToolCall(foreground.clone()),
3302 ModelTurnEvent::ToolCall(background.clone()),
3303 ModelTurnEvent::Finished(ModelTurnResult {
3304 model: None,
3305 response_id: None,
3306 finish_reason: FinishReason::ToolCall,
3307 output_items: vec![Item {
3308 id: None,
3309 kind: ItemKind::Assistant,
3310 parts: vec![Part::ToolCall(foreground), Part::ToolCall(background)],
3311 metadata: MetadataMap::new(),
3312 usage: None,
3313 finish_reason: None,
3314 created_at: None,
3315 }],
3316 usage: None,
3317 metadata: MetadataMap::new(),
3318 }),
3319 ])
3320 };
3321
3322 Ok(MultiToolTurn { events })
3323 }
3324 }
3325
3326 #[async_trait]
3327 impl ModelSession for DualApprovalSession {
3328 type Turn = DualApprovalTurn;
3329
3330 async fn begin_turn(
3331 &mut self,
3332 request: TurnRequest,
3333 _cancellation: Option<TurnCancellation>,
3334 ) -> Result<Self::Turn, LoopError> {
3335 let tool_results = request
3336 .transcript
3337 .iter()
3338 .flat_map(|item| item.parts.iter())
3339 .filter(|part| matches!(part, Part::ToolResult(_)))
3340 .count();
3341
3342 let events = if tool_results >= 2 {
3343 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
3344 model: None,
3345 response_id: None,
3346 finish_reason: FinishReason::Completed,
3347 output_items: vec![Item {
3348 id: None,
3349 kind: ItemKind::Assistant,
3350 parts: vec![Part::Text(TextPart {
3351 text: "both approvals finished".into(),
3352 metadata: MetadataMap::new(),
3353 })],
3354 metadata: MetadataMap::new(),
3355 usage: None,
3356 finish_reason: None,
3357 created_at: None,
3358 }],
3359 usage: None,
3360 metadata: MetadataMap::new(),
3361 })])
3362 } else {
3363 let first = agentkit_core::ToolCallPart {
3364 id: ToolCallId::new("call-1"),
3365 name: "echo".into(),
3366 input: json!({ "value": "first" }),
3367 metadata: MetadataMap::new(),
3368 };
3369 let second = agentkit_core::ToolCallPart {
3370 id: ToolCallId::new("call-2"),
3371 name: "echo".into(),
3372 input: json!({ "value": "second" }),
3373 metadata: MetadataMap::new(),
3374 };
3375 VecDeque::from([
3376 ModelTurnEvent::ToolCall(first.clone()),
3377 ModelTurnEvent::ToolCall(second.clone()),
3378 ModelTurnEvent::Finished(ModelTurnResult {
3379 model: None,
3380 response_id: None,
3381 finish_reason: FinishReason::ToolCall,
3382 output_items: vec![Item {
3383 id: None,
3384 kind: ItemKind::Assistant,
3385 parts: vec![Part::ToolCall(first), Part::ToolCall(second)],
3386 metadata: MetadataMap::new(),
3387 usage: None,
3388 finish_reason: None,
3389 created_at: None,
3390 }],
3391 usage: None,
3392 metadata: MetadataMap::new(),
3393 }),
3394 ])
3395 };
3396
3397 Ok(DualApprovalTurn { events })
3398 }
3399 }
3400
3401 #[async_trait]
3402 impl ModelTurn for FakeTurn {
3403 async fn next_event(
3404 &mut self,
3405 _cancellation: Option<TurnCancellation>,
3406 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3407 Ok(self.events.pop_front())
3408 }
3409 }
3410
3411 #[async_trait]
3412 impl ModelTurn for SlowTurn {
3413 async fn next_event(
3414 &mut self,
3415 cancellation: Option<TurnCancellation>,
3416 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3417 if let Some(cancellation) = cancellation
3418 && cancellation.is_cancelled()
3419 {
3420 return Err(LoopError::Cancelled);
3421 }
3422
3423 if self.emitted {
3424 Ok(None)
3425 } else {
3426 self.emitted = true;
3427 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
3428 model: None,
3429 response_id: None,
3430 finish_reason: FinishReason::Completed,
3431 output_items: vec![Item {
3432 id: None,
3433 kind: ItemKind::Assistant,
3434 parts: vec![Part::Text(TextPart {
3435 text: "done".into(),
3436 metadata: MetadataMap::new(),
3437 })],
3438 metadata: MetadataMap::new(),
3439 usage: None,
3440 finish_reason: None,
3441 created_at: None,
3442 }],
3443 usage: None,
3444 metadata: MetadataMap::new(),
3445 })))
3446 }
3447 }
3448 }
3449
3450 #[async_trait]
3451 impl ModelTurn for RecordingTurn {
3452 async fn next_event(
3453 &mut self,
3454 _cancellation: Option<TurnCancellation>,
3455 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3456 if self.emitted {
3457 Ok(None)
3458 } else {
3459 self.emitted = true;
3460 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
3461 model: None,
3462 response_id: None,
3463 finish_reason: FinishReason::Completed,
3464 output_items: vec![Item {
3465 id: None,
3466 kind: ItemKind::Assistant,
3467 parts: vec![Part::Text(TextPart {
3468 text: "done".into(),
3469 metadata: MetadataMap::new(),
3470 })],
3471 metadata: MetadataMap::new(),
3472 usage: None,
3473 finish_reason: None,
3474 created_at: None,
3475 }],
3476 usage: None,
3477 metadata: MetadataMap::new(),
3478 })))
3479 }
3480 }
3481 }
3482
3483 #[async_trait]
3484 impl ModelTurn for MultiToolTurn {
3485 async fn next_event(
3486 &mut self,
3487 _cancellation: Option<TurnCancellation>,
3488 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3489 Ok(self.events.pop_front())
3490 }
3491 }
3492
3493 #[async_trait]
3494 impl ModelTurn for DualApprovalTurn {
3495 async fn next_event(
3496 &mut self,
3497 _cancellation: Option<TurnCancellation>,
3498 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3499 Ok(self.events.pop_front())
3500 }
3501 }
3502
3503 #[derive(Clone)]
3504 struct EchoTool {
3505 spec: ToolSpec,
3506 }
3507
3508 #[derive(Clone)]
3509 struct FailingTool {
3510 spec: ToolSpec,
3511 }
3512
3513 #[derive(Clone)]
3514 struct RunThenDenyTool {
3515 spec: ToolSpec,
3516 }
3517
3518 impl Default for EchoTool {
3519 fn default() -> Self {
3520 Self {
3521 spec: ToolSpec {
3522 name: ToolName::new("echo"),
3523 description: "Echo back a value".into(),
3524 input_schema: json!({
3525 "type": "object",
3526 "properties": {
3527 "value": { "type": "string" }
3528 },
3529 "required": ["value"],
3530 "additionalProperties": false
3531 }),
3532 output_schema: None,
3533 annotations: ToolAnnotations::default(),
3534 metadata: MetadataMap::new(),
3535 },
3536 }
3537 }
3538 }
3539
3540 impl Default for FailingTool {
3541 fn default() -> Self {
3542 Self {
3543 spec: ToolSpec {
3544 name: ToolName::new("failing"),
3545 description: "Always fails after execution starts".into(),
3546 input_schema: json!({
3547 "type": "object",
3548 "properties": {
3549 "value": { "type": "string" }
3550 },
3551 "additionalProperties": true
3552 }),
3553 output_schema: None,
3554 annotations: ToolAnnotations::default(),
3555 metadata: MetadataMap::new(),
3556 },
3557 }
3558 }
3559 }
3560
3561 impl Default for RunThenDenyTool {
3562 fn default() -> Self {
3563 Self {
3564 spec: ToolSpec {
3565 name: ToolName::new("run_then_deny"),
3566 description: "Runs, then returns a permission-denied error".into(),
3567 input_schema: json!({
3568 "type": "object",
3569 "properties": {
3570 "value": { "type": "string" }
3571 },
3572 "additionalProperties": true
3573 }),
3574 output_schema: None,
3575 annotations: ToolAnnotations::default(),
3576 metadata: MetadataMap::new(),
3577 },
3578 }
3579 }
3580 }
3581
3582 #[derive(Clone)]
3583 struct DynamicSpecTool {
3584 spec: ToolSpec,
3585 version: StdArc<AtomicUsize>,
3586 }
3587
3588 impl DynamicSpecTool {
3589 fn new(version: StdArc<AtomicUsize>) -> Self {
3590 Self {
3591 spec: ToolSpec {
3592 name: ToolName::new("dynamic"),
3593 description: "dynamic version 0".into(),
3594 input_schema: json!({
3595 "type": "object",
3596 "properties": {},
3597 "additionalProperties": false
3598 }),
3599 output_schema: None,
3600 annotations: ToolAnnotations::default(),
3601 metadata: MetadataMap::new(),
3602 },
3603 version,
3604 }
3605 }
3606 }
3607
3608 #[async_trait]
3609 impl Tool for EchoTool {
3610 fn spec(&self) -> &ToolSpec {
3611 &self.spec
3612 }
3613
3614 fn proposed_requests(
3615 &self,
3616 request: &agentkit_tools_core::ToolRequest,
3617 ) -> Result<
3618 Vec<Box<dyn agentkit_tools_core::PermissionRequest>>,
3619 agentkit_tools_core::ToolError,
3620 > {
3621 Ok(vec![Box::new(FileSystemPermissionRequest::Read {
3622 path: "/tmp/echo".into(),
3623 metadata: request.metadata.clone(),
3624 })])
3625 }
3626
3627 async fn invoke(
3628 &self,
3629 request: agentkit_tools_core::ToolRequest,
3630 _ctx: &mut ToolContext<'_>,
3631 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
3632 let value = request
3633 .input
3634 .get("value")
3635 .and_then(Value::as_str)
3636 .ok_or_else(|| {
3637 agentkit_tools_core::ToolError::InvalidInput("missing value".into())
3638 })?;
3639
3640 Ok(ToolResult {
3641 result: ToolResultPart {
3642 call_id: request.call_id,
3643 output: ToolOutput::Text(value.into()),
3644 is_error: false,
3645 metadata: MetadataMap::new(),
3646 },
3647 duration: None,
3648 metadata: MetadataMap::new(),
3649 })
3650 }
3651 }
3652
3653 #[async_trait]
3654 impl Tool for FailingTool {
3655 fn spec(&self) -> &ToolSpec {
3656 &self.spec
3657 }
3658
3659 async fn invoke(
3660 &self,
3661 _request: agentkit_tools_core::ToolRequest,
3662 _ctx: &mut ToolContext<'_>,
3663 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
3664 Err(agentkit_tools_core::ToolError::ExecutionFailed(
3665 "runtime failed".into(),
3666 ))
3667 }
3668 }
3669
3670 #[async_trait]
3671 impl Tool for RunThenDenyTool {
3672 fn spec(&self) -> &ToolSpec {
3673 &self.spec
3674 }
3675
3676 async fn invoke(
3677 &self,
3678 _request: agentkit_tools_core::ToolRequest,
3679 _ctx: &mut ToolContext<'_>,
3680 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
3681 Err(agentkit_tools_core::ToolError::PermissionDenied(
3682 PermissionDenial {
3683 code: PermissionCode::CustomPolicyDenied,
3684 message: "remote 403".into(),
3685 metadata: MetadataMap::new(),
3686 },
3687 ))
3688 }
3689 }
3690
3691 #[async_trait]
3692 impl Tool for DynamicSpecTool {
3693 fn spec(&self) -> &ToolSpec {
3694 &self.spec
3695 }
3696
3697 fn current_spec(&self) -> Option<ToolSpec> {
3698 let mut spec = self.spec.clone();
3699 spec.description = format!("dynamic version {}", self.version.load(Ordering::SeqCst));
3700 Some(spec)
3701 }
3702
3703 async fn invoke(
3704 &self,
3705 request: agentkit_tools_core::ToolRequest,
3706 _ctx: &mut ToolContext<'_>,
3707 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
3708 Ok(ToolResult {
3709 result: ToolResultPart {
3710 call_id: request.call_id,
3711 output: ToolOutput::Text("ok".into()),
3712 is_error: false,
3713 metadata: MetadataMap::new(),
3714 },
3715 duration: None,
3716 metadata: MetadataMap::new(),
3717 })
3718 }
3719 }
3720
3721 struct DenyFsReads;
3722
3723 impl PermissionChecker for DenyFsReads {
3724 fn evaluate(
3725 &self,
3726 request: &dyn agentkit_tools_core::PermissionRequest,
3727 ) -> PermissionDecision {
3728 if request.kind() == "filesystem.read" {
3729 return PermissionDecision::Deny(PermissionDenial {
3730 code: PermissionCode::PathNotAllowed,
3731 message: "reads denied in test".into(),
3732 metadata: MetadataMap::new(),
3733 });
3734 }
3735
3736 PermissionDecision::Allow
3737 }
3738 }
3739
3740 struct ApproveFsReads;
3741
3742 impl PermissionChecker for ApproveFsReads {
3743 fn evaluate(
3744 &self,
3745 request: &dyn agentkit_tools_core::PermissionRequest,
3746 ) -> PermissionDecision {
3747 if request.kind() == "filesystem.read" {
3748 return PermissionDecision::RequireApproval(ApprovalRequest {
3749 task_id: None,
3750 call_id: None,
3751 id: "approval:fs-read".into(),
3752 request_kind: request.kind().into(),
3753 reason: agentkit_tools_core::ApprovalReason::SensitivePath,
3754 summary: request.summary(),
3755 metadata: request.metadata().clone(),
3756 });
3757 }
3758
3759 PermissionDecision::Allow
3760 }
3761 }
3762
3763 struct KeepRecentMutator {
3764 keep: usize,
3765 }
3766
3767 #[async_trait]
3768 impl LoopMutator for KeepRecentMutator {
3769 async fn mutate(
3770 &self,
3771 cursor: &mut TranscriptCursor<'_>,
3772 ctx: LoopCtx<'_>,
3773 ) -> Result<(), LoopError> {
3774 if cursor.len() < 2 {
3775 return Ok(());
3776 }
3777 let drop = cursor.len().saturating_sub(self.keep);
3778 ctx.emitter.emit(AgentEvent::MutationStarted {
3779 session_id: ctx.session_id.clone(),
3780 turn_id: ctx.turn_id.cloned(),
3781 mutator: "keep-recent".into(),
3782 point: ctx.point,
3783 });
3784 cursor.drain(..drop);
3785 ctx.emitter.emit(AgentEvent::MutationFinished {
3786 session_id: ctx.session_id.clone(),
3787 turn_id: ctx.turn_id.cloned(),
3788 mutator: "keep-recent".into(),
3789 dirty: true,
3790 metadata: MetadataMap::new(),
3791 });
3792 Ok(())
3793 }
3794 }
3795
3796 struct PointRecordingMutator {
3799 points: StdArc<StdMutex<Vec<MutationPoint>>>,
3800 }
3801
3802 #[async_trait]
3803 impl LoopMutator for PointRecordingMutator {
3804 async fn mutate(
3805 &self,
3806 _cursor: &mut TranscriptCursor<'_>,
3807 ctx: LoopCtx<'_>,
3808 ) -> Result<(), LoopError> {
3809 self.points.lock().unwrap().push(ctx.point);
3810 Ok(())
3811 }
3812 }
3813
3814 struct RecordingObserver {
3815 events: StdArc<StdMutex<Vec<AgentEvent>>>,
3816 }
3817
3818 impl LoopObserver for RecordingObserver {
3819 fn handle_event(&self, event: ObservedEvent) {
3820 let event = event.event;
3821 self.events.lock().unwrap().push(event);
3822 }
3823 }
3824
3825 struct CatalogExecutor {
3826 version: AtomicUsize,
3827 events: StdMutex<Vec<ToolCatalogEvent>>,
3828 }
3829
3830 impl CatalogExecutor {
3831 fn new() -> Self {
3832 Self {
3833 version: AtomicUsize::new(0),
3834 events: StdMutex::new(Vec::new()),
3835 }
3836 }
3837
3838 fn publish_change(&self, version: usize, event: ToolCatalogEvent) {
3839 self.version.store(version, Ordering::SeqCst);
3840 self.events.lock().unwrap().push(event);
3841 }
3842 }
3843
3844 #[async_trait]
3845 impl ToolExecutor for CatalogExecutor {
3846 fn specs(&self) -> Vec<ToolSpec> {
3847 vec![ToolSpec {
3848 name: ToolName::new("dynamic"),
3849 description: format!("dynamic version {}", self.version.load(Ordering::SeqCst)),
3850 input_schema: json!({
3851 "type": "object",
3852 "properties": {},
3853 "additionalProperties": false
3854 }),
3855 output_schema: None,
3856 annotations: ToolAnnotations::default(),
3857 metadata: MetadataMap::new(),
3858 }]
3859 }
3860
3861 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3862 std::mem::take(&mut *self.events.lock().unwrap())
3863 }
3864
3865 async fn execute(
3866 &self,
3867 request: ToolRequest,
3868 _ctx: &mut ToolContext<'_>,
3869 ) -> ToolExecutionOutcome {
3870 ToolExecutionOutcome::Completed(ToolResult {
3871 result: ToolResultPart {
3872 call_id: request.call_id,
3873 output: ToolOutput::Text("dynamic-ok".into()),
3874 is_error: false,
3875 metadata: MetadataMap::new(),
3876 },
3877 duration: None,
3878 metadata: MetadataMap::new(),
3879 })
3880 }
3881 }
3882
3883 #[derive(Clone)]
3884 struct BlockingTool {
3885 spec: ToolSpec,
3886 entered: StdArc<AtomicBool>,
3887 release: StdArc<Notify>,
3888 output: &'static str,
3889 }
3890
3891 impl BlockingTool {
3892 fn new(
3893 name: &str,
3894 entered: StdArc<AtomicBool>,
3895 release: StdArc<Notify>,
3896 output: &'static str,
3897 ) -> Self {
3898 Self {
3899 spec: ToolSpec {
3900 name: ToolName::new(name),
3901 description: format!("blocking tool {name}"),
3902 input_schema: json!({
3903 "type": "object",
3904 "properties": {},
3905 "additionalProperties": false
3906 }),
3907 output_schema: None,
3908 annotations: ToolAnnotations::default(),
3909 metadata: MetadataMap::new(),
3910 },
3911 entered,
3912 release,
3913 output,
3914 }
3915 }
3916 }
3917
3918 #[async_trait]
3919 impl Tool for BlockingTool {
3920 fn spec(&self) -> &ToolSpec {
3921 &self.spec
3922 }
3923
3924 async fn invoke(
3925 &self,
3926 request: agentkit_tools_core::ToolRequest,
3927 _ctx: &mut ToolContext<'_>,
3928 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
3929 self.entered.store(true, Ordering::SeqCst);
3930 self.release.notified().await;
3931 Ok(ToolResult {
3932 result: ToolResultPart {
3933 call_id: request.call_id,
3934 output: ToolOutput::Text(self.output.into()),
3935 is_error: false,
3936 metadata: MetadataMap::new(),
3937 },
3938 duration: None,
3939 metadata: MetadataMap::new(),
3940 })
3941 }
3942 }
3943
3944 struct NameRoutingPolicy {
3945 routes: Vec<(String, RoutingDecision)>,
3946 }
3947
3948 impl NameRoutingPolicy {
3949 fn new(routes: impl IntoIterator<Item = (impl Into<String>, RoutingDecision)>) -> Self {
3950 Self {
3951 routes: routes
3952 .into_iter()
3953 .map(|(name, decision)| (name.into(), decision))
3954 .collect(),
3955 }
3956 }
3957 }
3958
3959 impl TaskRoutingPolicy for NameRoutingPolicy {
3960 fn route(&self, request: &ToolRequest) -> RoutingDecision {
3961 self.routes
3962 .iter()
3963 .find(|(name, _)| name == &request.tool_name.0)
3964 .map(|(_, decision)| *decision)
3965 .unwrap_or(RoutingDecision::Foreground)
3966 }
3967 }
3968
3969 async fn wait_for_task_event(handle: &TaskManagerHandle) -> TaskEvent {
3970 timeout(Duration::from_secs(1), handle.next_event())
3971 .await
3972 .expect("timed out waiting for task event")
3973 .expect("task event stream ended unexpectedly")
3974 }
3975
3976 async fn wait_until_entered(flag: &AtomicBool) {
3977 timeout(Duration::from_secs(1), async {
3978 while !flag.load(Ordering::SeqCst) {
3979 tokio::task::yield_now().await;
3980 }
3981 })
3982 .await
3983 .expect("task never entered execution");
3984 }
3985
3986 async fn wait_until_completed(handle: &TaskManagerHandle) {
3987 timeout(Duration::from_secs(1), async {
3988 while handle.list_completed().await.is_empty() {
3989 tokio::task::yield_now().await;
3990 }
3991 })
3992 .await
3993 .expect("task never completed");
3994 }
3995
3996 #[tokio::test]
3997 async fn loop_continues_after_completed_tool_call() {
3998 let tools = ToolRegistry::new().with(EchoTool::default());
3999 let agent = Agent::builder()
4000 .model(FakeAdapter)
4001 .add_tool_source(tools)
4002 .permissions(AllowAllPermissions)
4003 .build()
4004 .unwrap();
4005
4006 let mut driver = agent
4007 .start(SessionConfig {
4008 session_id: SessionId::new("session-1"),
4009 metadata: MetadataMap::new(),
4010 cache: None,
4011 })
4012 .await
4013 .unwrap();
4014
4015 driver
4016 .submit_input(vec![Item {
4017 id: None,
4018 kind: ItemKind::User,
4019 parts: vec![Part::Text(TextPart {
4020 text: "ping".into(),
4021 metadata: MetadataMap::new(),
4022 })],
4023 metadata: MetadataMap::new(),
4024 usage: None,
4025 finish_reason: None,
4026 created_at: None,
4027 }])
4028 .unwrap();
4029
4030 let result = run_until_finished(&mut driver).await;
4031
4032 match result {
4033 LoopStep::Finished(turn) => {
4034 assert_eq!(turn.finish_reason, FinishReason::Completed);
4035 assert_eq!(turn.items.len(), 1);
4036 match &turn.items[0].parts[0] {
4037 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
4038 other => panic!("unexpected part: {other:?}"),
4039 }
4040 }
4041 other => panic!("unexpected loop step: {other:?}"),
4042 }
4043 }
4044
4045 async fn run_until_finished<S: ModelSession + Send>(driver: &mut LoopDriver<S>) -> LoopStep {
4049 loop {
4050 match driver.next().await.unwrap() {
4051 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue,
4052 step => return step,
4053 }
4054 }
4055 }
4056
4057 #[tokio::test]
4064 async fn post_tool_continuation_reports_after_tool_result_mutation_point() {
4065 let points = StdArc::new(StdMutex::new(Vec::<MutationPoint>::new()));
4066 let tools = ToolRegistry::new().with(EchoTool::default());
4067 let agent = Agent::builder()
4068 .model(FakeAdapter)
4069 .add_tool_source(tools)
4070 .permissions(AllowAllPermissions)
4071 .mutator(PointRecordingMutator {
4072 points: points.clone(),
4073 })
4074 .build()
4075 .unwrap();
4076
4077 let mut driver = agent
4078 .start(SessionConfig {
4079 session_id: SessionId::new("session-mutation-point"),
4080 metadata: MetadataMap::new(),
4081 cache: None,
4082 })
4083 .await
4084 .unwrap();
4085
4086 driver
4087 .submit_input(vec![Item::text(ItemKind::User, "ping")])
4088 .unwrap();
4089
4090 let _ = run_until_finished(&mut driver).await;
4092
4093 let recorded = points.lock().unwrap().clone();
4094 assert_eq!(
4095 recorded.first(),
4096 Some(&MutationPoint::AfterTurnEnded),
4097 "first drive of a fresh turn must report AfterTurnEnded, got {recorded:?}"
4098 );
4099 assert!(
4100 recorded.contains(&MutationPoint::AfterToolResult),
4101 "post-tool continuation must report AfterToolResult, got {recorded:?}"
4102 );
4103 }
4104
4105 #[test]
4106 fn pending_input_requires_input_bearing_tail_role() {
4107 assert!(!transcript_has_pending_input(&[]));
4108 assert!(!transcript_has_pending_input(&[Item::text(
4109 ItemKind::System,
4110 "system"
4111 )]));
4112 assert!(!transcript_has_pending_input(&[Item::text(
4113 ItemKind::Developer,
4114 "developer"
4115 )]));
4116 assert!(!transcript_has_pending_input(&[Item::text(
4117 ItemKind::Context,
4118 "context"
4119 )]));
4120 assert!(!transcript_has_pending_input(&[Item::text(
4121 ItemKind::Assistant,
4122 "assistant"
4123 )]));
4124
4125 assert!(transcript_has_pending_input(&[Item::text(
4126 ItemKind::User,
4127 "user"
4128 )]));
4129 assert!(transcript_has_pending_input(&[Item::notification(
4130 "background update"
4131 )]));
4132 assert!(transcript_has_pending_input(&[Item {
4133 id: None,
4134 kind: ItemKind::Tool,
4135 parts: vec![Part::ToolResult(ToolResultPart {
4136 call_id: ToolCallId::new("call-test"),
4137 output: ToolOutput::Text("ok".into()),
4138 is_error: false,
4139 metadata: MetadataMap::new(),
4140 })],
4141 metadata: MetadataMap::new(),
4142 usage: None,
4143 finish_reason: None,
4144 created_at: None,
4145 }]));
4146 }
4147
4148 struct DropTrailingUserMutator;
4154
4155 #[async_trait]
4156 impl LoopMutator for DropTrailingUserMutator {
4157 async fn mutate(
4158 &self,
4159 cursor: &mut TranscriptCursor<'_>,
4160 _ctx: LoopCtx<'_>,
4161 ) -> Result<(), LoopError> {
4162 if cursor.last().map(|item| item.kind) == Some(ItemKind::User) {
4163 cursor.pop();
4164 }
4165 Ok(())
4166 }
4167 }
4168
4169 struct RejectAssistantPrefillAdapter {
4174 saw_assistant_tail: StdArc<AtomicBool>,
4175 }
4176
4177 struct RejectAssistantPrefillSession {
4178 saw_assistant_tail: StdArc<AtomicBool>,
4179 }
4180
4181 #[async_trait]
4182 impl ModelAdapter for RejectAssistantPrefillAdapter {
4183 type Session = RejectAssistantPrefillSession;
4184
4185 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
4186 Ok(RejectAssistantPrefillSession {
4187 saw_assistant_tail: self.saw_assistant_tail.clone(),
4188 })
4189 }
4190 }
4191
4192 #[async_trait]
4193 impl ModelSession for RejectAssistantPrefillSession {
4194 type Turn = FakeTurn;
4195
4196 async fn begin_turn(
4197 &mut self,
4198 request: TurnRequest,
4199 _cancellation: Option<TurnCancellation>,
4200 ) -> Result<Self::Turn, LoopError> {
4201 if request.transcript.last().map(|item| item.kind) == Some(ItemKind::Assistant) {
4202 self.saw_assistant_tail.store(true, Ordering::SeqCst);
4203 return Err(LoopError::Provider(
4204 "conversation must end with a user message".into(),
4205 ));
4206 }
4207 Ok(FakeTurn {
4208 events: VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
4209 model: None,
4210 response_id: None,
4211 finish_reason: FinishReason::Completed,
4212 output_items: vec![Item::text(ItemKind::Assistant, "ok")],
4213 usage: None,
4214 metadata: MetadataMap::new(),
4215 })]),
4216 })
4217 }
4218 }
4219
4220 #[tokio::test]
4228 async fn drive_does_not_dispatch_without_valid_trailing_input() {
4229 let saw_assistant_tail = StdArc::new(AtomicBool::new(false));
4230 let agent = Agent::builder()
4231 .model(RejectAssistantPrefillAdapter {
4232 saw_assistant_tail: saw_assistant_tail.clone(),
4233 })
4234 .mutator(DropTrailingUserMutator)
4235 .transcript(vec![
4238 Item::text(ItemKind::User, "kickoff"),
4239 Item::text(ItemKind::Assistant, "prior reply"),
4240 ])
4241 .build()
4242 .unwrap();
4243
4244 let mut driver = agent
4245 .start(SessionConfig {
4246 session_id: SessionId::new("session-no-valid-input"),
4247 metadata: MetadataMap::new(),
4248 cache: None,
4249 })
4250 .await
4251 .unwrap();
4252
4253 driver
4254 .submit_input(vec![Item::text(ItemKind::User, "follow up")])
4255 .unwrap();
4256
4257 let outcome = driver.next().await;
4259
4260 assert!(
4261 !saw_assistant_tail.load(Ordering::SeqCst),
4262 "loop dispatched a model turn whose transcript ends in an assistant \
4263 message (outcome: {outcome:?}); with no valid trailing input the turn \
4264 must finish instead of driving"
4265 );
4266 }
4267
4268 #[tokio::test]
4269 async fn loop_uses_injected_permission_checker() {
4270 let events = StdArc::new(StdMutex::new(Vec::new()));
4271 let tools = ToolRegistry::new().with(EchoTool::default());
4272 let agent = Agent::builder()
4273 .model(FakeAdapter)
4274 .add_tool_source(tools)
4275 .permissions(DenyFsReads)
4276 .observer(RecordingObserver {
4277 events: events.clone(),
4278 })
4279 .build()
4280 .unwrap();
4281
4282 let mut driver = agent
4283 .start(SessionConfig {
4284 session_id: SessionId::new("session-2"),
4285 metadata: MetadataMap::new(),
4286 cache: None,
4287 })
4288 .await
4289 .unwrap();
4290
4291 driver
4292 .submit_input(vec![Item {
4293 id: None,
4294 kind: ItemKind::User,
4295 parts: vec![Part::Text(TextPart {
4296 text: "ping".into(),
4297 metadata: MetadataMap::new(),
4298 })],
4299 metadata: MetadataMap::new(),
4300 usage: None,
4301 finish_reason: None,
4302 created_at: None,
4303 }])
4304 .unwrap();
4305
4306 let result = run_until_finished(&mut driver).await;
4307
4308 match result {
4309 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
4310 Part::Text(text) => assert!(text.text.contains("tool permission denied")),
4311 other => panic!("unexpected part: {other:?}"),
4312 },
4313 other => panic!("unexpected loop step: {other:?}"),
4314 }
4315
4316 assert!(
4317 events
4318 .lock()
4319 .unwrap()
4320 .iter()
4321 .all(|event| !matches!(event, AgentEvent::ToolExecutionStarted(_))),
4322 "denied tools must not be reported as started"
4323 );
4324 }
4325
4326 #[tokio::test]
4327 async fn failed_tool_execution_still_reports_started() {
4328 let events = StdArc::new(StdMutex::new(Vec::new()));
4329 let tools = ToolRegistry::new().with(FailingTool::default());
4330 let agent = Agent::builder()
4331 .model(FakeAdapter)
4332 .add_tool_source(tools)
4333 .permissions(AllowAllPermissions)
4334 .observer(RecordingObserver {
4335 events: events.clone(),
4336 })
4337 .build()
4338 .unwrap();
4339
4340 let mut driver = agent
4341 .start(SessionConfig {
4342 session_id: SessionId::new("session-failing-start-event"),
4343 metadata: MetadataMap::new(),
4344 cache: None,
4345 })
4346 .await
4347 .unwrap();
4348
4349 driver
4350 .submit_input(vec![Item::text(ItemKind::User, "ping")])
4351 .unwrap();
4352
4353 match run_until_finished(&mut driver).await {
4354 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
4355 other => panic!("unexpected loop step: {other:?}"),
4356 }
4357
4358 let events = events.lock().unwrap();
4359 assert!(events.iter().any(|event| matches!(
4360 event,
4361 AgentEvent::ToolExecutionStarted(call) if call.name == "failing"
4362 )));
4363 assert!(events.iter().any(|event| matches!(
4364 event,
4365 AgentEvent::ToolResultReceived(result) if result.is_error
4366 )));
4367 }
4368
4369 #[tokio::test]
4370 async fn run_then_deny_tool_execution_still_reports_started() {
4371 let events = StdArc::new(StdMutex::new(Vec::new()));
4372 let tools = ToolRegistry::new().with(RunThenDenyTool::default());
4373 let agent = Agent::builder()
4374 .model(FakeAdapter)
4375 .add_tool_source(tools)
4376 .permissions(AllowAllPermissions)
4377 .observer(RecordingObserver {
4378 events: events.clone(),
4379 })
4380 .build()
4381 .unwrap();
4382
4383 let mut driver = agent
4384 .start(SessionConfig {
4385 session_id: SessionId::new("session-run-then-deny-start-event"),
4386 metadata: MetadataMap::new(),
4387 cache: None,
4388 })
4389 .await
4390 .unwrap();
4391
4392 driver
4393 .submit_input(vec![Item::text(ItemKind::User, "ping")])
4394 .unwrap();
4395
4396 match run_until_finished(&mut driver).await {
4397 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
4398 other => panic!("unexpected loop step: {other:?}"),
4399 }
4400
4401 let events = events.lock().unwrap();
4402 assert!(events.iter().any(|event| matches!(
4403 event,
4404 AgentEvent::ToolExecutionStarted(call) if call.name == "run_then_deny"
4405 )));
4406 assert!(events.iter().any(|event| matches!(
4407 event,
4408 AgentEvent::ToolResultReceived(result)
4409 if result.is_error
4410 && result
4411 .metadata
4412 .get(TOOL_RESULT_FAILURE_KIND_METADATA_KEY)
4413 .is_none()
4414 )));
4415 }
4416
4417 #[tokio::test]
4418 async fn async_task_manager_background_round_requires_explicit_continue() {
4419 let events = StdArc::new(StdMutex::new(Vec::new()));
4420 let entered = StdArc::new(AtomicBool::new(false));
4421 let release = StdArc::new(Notify::new());
4422 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
4423 "background-wait",
4424 RoutingDecision::Background,
4425 )]));
4426 let handle = task_manager.handle();
4427 let tools = ToolRegistry::new().with(BlockingTool::new(
4428 "background-wait",
4429 entered.clone(),
4430 release.clone(),
4431 "background-done",
4432 ));
4433 let agent = Agent::builder()
4434 .model(FakeAdapter)
4435 .add_tool_source(tools)
4436 .permissions(AllowAllPermissions)
4437 .task_manager(task_manager)
4438 .observer(RecordingObserver {
4439 events: events.clone(),
4440 })
4441 .build()
4442 .unwrap();
4443
4444 let mut driver = agent
4445 .start(SessionConfig {
4446 session_id: SessionId::new("session-background"),
4447 metadata: MetadataMap::new(),
4448 cache: None,
4449 })
4450 .await
4451 .unwrap();
4452
4453 driver
4454 .submit_input(vec![Item {
4455 id: None,
4456 kind: ItemKind::User,
4457 parts: vec![Part::Text(TextPart {
4458 text: "ping".into(),
4459 metadata: MetadataMap::new(),
4460 })],
4461 metadata: MetadataMap::new(),
4462 usage: None,
4463 finish_reason: None,
4464 created_at: None,
4465 }])
4466 .unwrap();
4467
4468 let first = driver.next().await.unwrap();
4469 match first {
4470 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
4471 other => panic!("unexpected first loop step: {other:?}"),
4472 }
4473
4474 match wait_for_task_event(&handle).await {
4475 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "background-wait"),
4476 other => panic!("unexpected task event: {other:?}"),
4477 }
4478 wait_until_entered(entered.as_ref()).await;
4479 release.notify_waiters();
4480
4481 match wait_for_task_event(&handle).await {
4482 TaskEvent::Completed(_, result) => {
4483 assert_eq!(result.output, ToolOutput::Text("background-done".into()))
4484 }
4485 other => panic!("unexpected completion event: {other:?}"),
4486 }
4487
4488 let resumed = driver.next().await.unwrap();
4489 match resumed {
4490 LoopStep::Finished(turn) => {
4491 assert_eq!(turn.finish_reason, FinishReason::Completed);
4492 match &turn.items[0].parts[0] {
4493 Part::Text(text) => assert_eq!(text.text, "tool said: background-done"),
4494 other => panic!("unexpected part after resume: {other:?}"),
4495 }
4496 }
4497 other => panic!("unexpected resumed step: {other:?}"),
4498 }
4499
4500 let events = events.lock().unwrap();
4501 let terminal_results: Vec<_> = events
4502 .iter()
4503 .filter_map(|event| match event {
4504 AgentEvent::ToolResultReceived(result)
4505 if result.call_id == ToolCallId::new("call-1") =>
4506 {
4507 Some(result)
4508 }
4509 _ => None,
4510 })
4511 .collect();
4512 assert_eq!(
4513 terminal_results.len(),
4514 1,
4515 "background completion must emit one terminal result event per call: {events:?}"
4516 );
4517 }
4518
4519 #[tokio::test]
4520 async fn detached_tool_placeholder_is_progress_not_terminal_result() {
4521 let events = StdArc::new(StdMutex::new(Vec::new()));
4522 let entered = StdArc::new(AtomicBool::new(false));
4523 let release = StdArc::new(Notify::new());
4524 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
4525 "detaching-wait",
4526 RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10)),
4527 )]));
4528 let handle = task_manager.handle();
4529 let tools = ToolRegistry::new().with(BlockingTool::new(
4530 "detaching-wait",
4531 entered.clone(),
4532 release.clone(),
4533 "detached-done",
4534 ));
4535 let agent = Agent::builder()
4536 .model(FakeAdapter)
4537 .add_tool_source(tools)
4538 .permissions(AllowAllPermissions)
4539 .task_manager(task_manager)
4540 .observer(RecordingObserver {
4541 events: events.clone(),
4542 })
4543 .build()
4544 .unwrap();
4545
4546 let mut driver = agent
4547 .start(SessionConfig {
4548 session_id: SessionId::new("session-detached-progress"),
4549 metadata: MetadataMap::new(),
4550 cache: None,
4551 })
4552 .await
4553 .unwrap();
4554
4555 driver
4556 .submit_input(vec![Item::text(ItemKind::User, "ping")])
4557 .unwrap();
4558
4559 match driver.next().await.unwrap() {
4560 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => {}
4561 other => panic!("unexpected detach step: {other:?}"),
4562 }
4563
4564 match wait_for_task_event(&handle).await {
4565 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "detaching-wait"),
4566 other => panic!("unexpected task event: {other:?}"),
4567 }
4568 match wait_for_task_event(&handle).await {
4569 TaskEvent::Detached(snapshot) => assert_eq!(snapshot.tool_name, "detaching-wait"),
4570 other => panic!("unexpected detach event: {other:?}"),
4571 }
4572 wait_until_entered(entered.as_ref()).await;
4573 release.notify_waiters();
4574
4575 match wait_for_task_event(&handle).await {
4576 TaskEvent::Completed(_, result) => {
4577 assert_eq!(result.output, ToolOutput::Text("detached-done".into()))
4578 }
4579 other => panic!("unexpected completion event: {other:?}"),
4580 }
4581
4582 match driver.next().await.unwrap() {
4583 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Completed),
4584 other => panic!("unexpected resumed step: {other:?}"),
4585 }
4586
4587 let events = events.lock().unwrap();
4588 assert!(events.iter().any(|event| matches!(
4589 event,
4590 AgentEvent::ToolExecutionProgress(result)
4591 if result.call_id == ToolCallId::new("call-1") && !result.is_error
4592 )));
4593 let terminal_results: Vec<_> = events
4594 .iter()
4595 .filter_map(|event| match event {
4596 AgentEvent::ToolResultReceived(result)
4597 if result.call_id == ToolCallId::new("call-1") =>
4598 {
4599 Some(result)
4600 }
4601 _ => None,
4602 })
4603 .collect();
4604 assert_eq!(
4605 terminal_results.len(),
4606 1,
4607 "detached call must emit one terminal result event: {events:?}"
4608 );
4609 }
4610
4611 #[tokio::test]
4612 async fn cancelled_background_approval_auto_resolves_when_drained() {
4613 let controller = CancellationController::new();
4614 let events = StdArc::new(StdMutex::new(Vec::new()));
4615 let entered = StdArc::new(AtomicBool::new(false));
4616 let release = StdArc::new(Notify::new());
4617 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
4618 "echo",
4619 RoutingDecision::Background,
4620 )]));
4621 let handle = task_manager.handle();
4622 let agent = Agent::builder()
4623 .model(FakeAdapter)
4624 .tool_executor(DelayedApprovalExecutor::new(
4625 entered.clone(),
4626 release.clone(),
4627 ))
4628 .task_manager(task_manager)
4629 .cancellation(controller.handle())
4630 .observer(RecordingObserver {
4631 events: events.clone(),
4632 })
4633 .build()
4634 .unwrap();
4635
4636 let mut driver = agent
4637 .start(SessionConfig {
4638 session_id: SessionId::new("session-cancel-delayed-background-approval"),
4639 metadata: MetadataMap::new(),
4640 cache: None,
4641 })
4642 .await
4643 .unwrap();
4644
4645 driver
4646 .submit_input(vec![Item::text(ItemKind::User, "ping")])
4647 .unwrap();
4648
4649 match driver.next().await.unwrap() {
4650 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
4651 other => panic!("unexpected first step: {other:?}"),
4652 }
4653
4654 match wait_for_task_event(&handle).await {
4655 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "echo"),
4656 other => panic!("unexpected task event: {other:?}"),
4657 }
4658
4659 wait_until_entered(entered.as_ref()).await;
4660 controller.interrupt();
4661 release.notify_waiters();
4662 wait_until_completed(&handle).await;
4663
4664 match driver.next().await.unwrap() {
4665 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
4666 other => panic!("cancelled background approval should finish cancelled, got {other:?}"),
4667 }
4668
4669 let events = events.lock().unwrap();
4670 assert!(
4671 events
4672 .iter()
4673 .any(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false }))
4674 );
4675 assert!(events.iter().any(|event| matches!(
4676 event,
4677 AgentEvent::ToolResultReceived(result)
4678 if result.call_id == ToolCallId::new("call-1") && result.is_error
4679 )));
4680 }
4681
4682 #[tokio::test]
4683 async fn loop_can_cancel_a_turn_and_continue_after_new_input() {
4684 let controller = CancellationController::new();
4685 let agent = Agent::builder()
4686 .model(SlowAdapter)
4687 .cancellation(controller.handle())
4688 .build()
4689 .unwrap();
4690
4691 let mut driver = agent
4692 .start(SessionConfig {
4693 session_id: SessionId::new("session-cancel"),
4694 metadata: MetadataMap::new(),
4695 cache: None,
4696 })
4697 .await
4698 .unwrap();
4699
4700 driver
4701 .submit_input(vec![Item {
4702 id: None,
4703 kind: ItemKind::User,
4704 parts: vec![Part::Text(TextPart {
4705 text: "do the long task".into(),
4706 metadata: MetadataMap::new(),
4707 })],
4708 metadata: MetadataMap::new(),
4709 usage: None,
4710 finish_reason: None,
4711 created_at: None,
4712 }])
4713 .unwrap();
4714
4715 let cancelled = tokio::join!(async { driver.next().await }, async {
4716 tokio::task::yield_now().await;
4717 controller.interrupt();
4718 })
4719 .0
4720 .unwrap();
4721
4722 match cancelled {
4723 LoopStep::Finished(turn) => {
4724 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
4725 assert_eq!(turn.items.len(), 1);
4726 assert_eq!(turn.items[0].kind, ItemKind::Assistant);
4727 assert_eq!(
4728 turn.items[0].metadata.get(INTERRUPTED_METADATA_KEY),
4729 Some(&Value::Bool(true))
4730 );
4731 }
4732 other => panic!("unexpected loop step: {other:?}"),
4733 }
4734
4735 driver
4736 .submit_input(vec![Item {
4737 id: None,
4738 kind: ItemKind::User,
4739 parts: vec![Part::Text(TextPart {
4740 text: "try again".into(),
4741 metadata: MetadataMap::new(),
4742 })],
4743 metadata: MetadataMap::new(),
4744 usage: None,
4745 finish_reason: None,
4746 created_at: None,
4747 }])
4748 .unwrap();
4749
4750 let result = driver.next().await.unwrap();
4751 match result {
4752 LoopStep::Finished(turn) => {
4753 assert_eq!(turn.finish_reason, FinishReason::Completed);
4754 }
4755 other => panic!("unexpected loop step after retry: {other:?}"),
4756 }
4757 }
4758
4759 #[tokio::test]
4760 async fn loop_interrupt_cancels_foreground_tasks_but_keeps_background_tasks_running() {
4761 let controller = CancellationController::new();
4762 let fg_entered = StdArc::new(AtomicBool::new(false));
4763 let fg_release = StdArc::new(Notify::new());
4764 let bg_entered = StdArc::new(AtomicBool::new(false));
4765 let bg_release = StdArc::new(Notify::new());
4766 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([
4767 ("foreground-wait", RoutingDecision::Foreground),
4768 ("background-wait", RoutingDecision::Background),
4769 ]));
4770 let handle = task_manager.handle();
4771 let tools = ToolRegistry::new()
4772 .with(BlockingTool::new(
4773 "foreground-wait",
4774 fg_entered.clone(),
4775 fg_release,
4776 "foreground-done",
4777 ))
4778 .with(BlockingTool::new(
4779 "background-wait",
4780 bg_entered.clone(),
4781 bg_release.clone(),
4782 "background-done",
4783 ));
4784 let agent = Agent::builder()
4785 .model(MultiToolAdapter)
4786 .add_tool_source(tools)
4787 .permissions(AllowAllPermissions)
4788 .cancellation(controller.handle())
4789 .task_manager(task_manager)
4790 .build()
4791 .unwrap();
4792
4793 let mut driver = agent
4794 .start(SessionConfig {
4795 session_id: SessionId::new("session-mixed-cancel"),
4796 metadata: MetadataMap::new(),
4797 cache: None,
4798 })
4799 .await
4800 .unwrap();
4801
4802 driver
4803 .submit_input(vec![Item {
4804 id: None,
4805 kind: ItemKind::User,
4806 parts: vec![Part::Text(TextPart {
4807 text: "run both".into(),
4808 metadata: MetadataMap::new(),
4809 })],
4810 metadata: MetadataMap::new(),
4811 usage: None,
4812 finish_reason: None,
4813 created_at: None,
4814 }])
4815 .unwrap();
4816
4817 let cancelled = tokio::join!(async { driver.next().await }, async {
4818 let _ = wait_for_task_event(&handle).await;
4819 let _ = wait_for_task_event(&handle).await;
4820 wait_until_entered(fg_entered.as_ref()).await;
4821 wait_until_entered(bg_entered.as_ref()).await;
4822 controller.interrupt();
4823 })
4824 .0
4825 .unwrap();
4826
4827 match cancelled {
4828 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
4829 other => panic!("unexpected loop step after interrupt: {other:?}"),
4830 }
4831
4832 match wait_for_task_event(&handle).await {
4833 TaskEvent::Cancelled(snapshot) => assert_eq!(snapshot.tool_name, "foreground-wait"),
4834 other => panic!("unexpected post-interrupt event: {other:?}"),
4835 }
4836
4837 let running = handle.list_running().await;
4838 assert_eq!(running.len(), 1);
4839 assert_eq!(running[0].tool_name, "background-wait");
4840
4841 bg_release.notify_waiters();
4842 match wait_for_task_event(&handle).await {
4843 TaskEvent::Completed(snapshot, result) => {
4844 assert_eq!(snapshot.tool_name, "background-wait");
4845 assert_eq!(result.output, ToolOutput::Text("background-done".into()));
4846 }
4847 other => panic!("unexpected background completion event: {other:?}"),
4848 }
4849 }
4850
4851 #[tokio::test]
4852 async fn loop_resumes_after_approved_tool_request() {
4853 let tools = ToolRegistry::new().with(EchoTool::default());
4854 let agent = Agent::builder()
4855 .model(FakeAdapter)
4856 .add_tool_source(tools)
4857 .permissions(ApproveFsReads)
4858 .build()
4859 .unwrap();
4860
4861 let mut driver = agent
4862 .start(SessionConfig {
4863 session_id: SessionId::new("session-approval"),
4864 metadata: MetadataMap::new(),
4865 cache: None,
4866 })
4867 .await
4868 .unwrap();
4869
4870 driver
4871 .submit_input(vec![Item {
4872 id: None,
4873 kind: ItemKind::User,
4874 parts: vec![Part::Text(TextPart {
4875 text: "ping".into(),
4876 metadata: MetadataMap::new(),
4877 })],
4878 metadata: MetadataMap::new(),
4879 usage: None,
4880 finish_reason: None,
4881 created_at: None,
4882 }])
4883 .unwrap();
4884
4885 let first = driver.next().await.unwrap();
4886 match first {
4887 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
4888 assert!(pending.request.task_id.is_some());
4889 assert_eq!(pending.request.id.0, "approval:fs-read");
4890 pending.approve(&mut driver).unwrap();
4891 }
4892 other => panic!("unexpected loop step: {other:?}"),
4893 }
4894 let second = driver.next().await.unwrap();
4895 match second {
4896 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
4897 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
4898 other => panic!("unexpected part: {other:?}"),
4899 },
4900 other => panic!("unexpected loop step after approval: {other:?}"),
4901 }
4902 }
4903
4904 #[tokio::test]
4905 async fn approval_gated_tool_does_not_start_before_approval() {
4906 let events = StdArc::new(StdMutex::new(Vec::new()));
4907 let tools = ToolRegistry::new().with(EchoTool::default());
4908 let agent = Agent::builder()
4909 .model(FakeAdapter)
4910 .add_tool_source(tools)
4911 .permissions(ApproveFsReads)
4912 .observer(RecordingObserver {
4913 events: events.clone(),
4914 })
4915 .build()
4916 .unwrap();
4917
4918 let mut driver = agent
4919 .start(SessionConfig {
4920 session_id: SessionId::new("session-approval-start-event"),
4921 metadata: MetadataMap::new(),
4922 cache: None,
4923 })
4924 .await
4925 .unwrap();
4926
4927 driver
4928 .submit_input(vec![Item::text(ItemKind::User, "ping")])
4929 .unwrap();
4930
4931 let pending = match driver.next().await.unwrap() {
4932 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
4933 other => panic!("unexpected loop step: {other:?}"),
4934 };
4935
4936 assert!(
4937 events
4938 .lock()
4939 .unwrap()
4940 .iter()
4941 .all(|event| !matches!(event, AgentEvent::ToolExecutionStarted(_))),
4942 "tool start must not be reported before approval"
4943 );
4944
4945 pending.approve(&mut driver).unwrap();
4946 match driver.next().await.unwrap() {
4947 LoopStep::Finished(_) => {}
4948 other => panic!("unexpected loop step after approval: {other:?}"),
4949 }
4950
4951 let started = events
4952 .lock()
4953 .unwrap()
4954 .iter()
4955 .filter(|event| matches!(event, AgentEvent::ToolExecutionStarted(_)))
4956 .count();
4957 assert_eq!(started, 1);
4958 }
4959
4960 #[tokio::test]
4961 async fn cancelling_pending_approval_resolves_it_and_pairs_tool_result() {
4962 let controller = CancellationController::new();
4963 let events = StdArc::new(StdMutex::new(Vec::new()));
4964 let tools = ToolRegistry::new().with(EchoTool::default());
4965 let agent = Agent::builder()
4966 .model(FakeAdapter)
4967 .add_tool_source(tools)
4968 .permissions(ApproveFsReads)
4969 .cancellation(controller.handle())
4970 .observer(RecordingObserver {
4971 events: events.clone(),
4972 })
4973 .build()
4974 .unwrap();
4975
4976 let mut driver = agent
4977 .start(SessionConfig {
4978 session_id: SessionId::new("session-cancel-pending-approval"),
4979 metadata: MetadataMap::new(),
4980 cache: None,
4981 })
4982 .await
4983 .unwrap();
4984
4985 driver
4986 .submit_input(vec![Item::text(ItemKind::User, "ping")])
4987 .unwrap();
4988
4989 match driver.next().await.unwrap() {
4990 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) => {}
4991 other => panic!("unexpected loop step: {other:?}"),
4992 }
4993
4994 controller.interrupt();
4995
4996 match driver.next().await.unwrap() {
4997 LoopStep::Finished(turn) => {
4998 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
4999 }
5000 other => panic!("unexpected loop step after cancel: {other:?}"),
5001 }
5002
5003 let events = events.lock().unwrap();
5004 assert!(
5005 events
5006 .iter()
5007 .any(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false })),
5008 "pending approval cancellation should close approval UI state"
5009 );
5010 assert!(
5011 events.iter().any(|event| matches!(
5012 event,
5013 AgentEvent::ToolResultReceived(result)
5014 if result.call_id == ToolCallId::new("call-1") && result.is_error
5015 )),
5016 "pending approval cancellation should pair the assistant tool_use"
5017 );
5018 drop(events);
5019
5020 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
5021 }
5022
5023 #[tokio::test]
5024 async fn resolved_approval_runs_even_if_cancellation_also_fired() {
5025 let controller = CancellationController::new();
5026 let tools = ToolRegistry::new().with(EchoTool::default());
5027 let agent = Agent::builder()
5028 .model(FakeAdapter)
5029 .add_tool_source(tools)
5030 .permissions(ApproveFsReads)
5031 .cancellation(controller.handle())
5032 .build()
5033 .unwrap();
5034
5035 let mut driver = agent
5036 .start(SessionConfig {
5037 session_id: SessionId::new("session-resolved-approval-cancel-race"),
5038 metadata: MetadataMap::new(),
5039 cache: None,
5040 })
5041 .await
5042 .unwrap();
5043
5044 driver
5045 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5046 .unwrap();
5047
5048 let pending = match driver.next().await.unwrap() {
5049 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => pending,
5050 other => panic!("unexpected loop step: {other:?}"),
5051 };
5052
5053 controller.interrupt();
5054 pending.approve(&mut driver).unwrap();
5055
5056 match driver.next().await.unwrap() {
5057 LoopStep::Finished(turn) => {
5058 assert_eq!(turn.finish_reason, FinishReason::Completed);
5059 match &turn.items[0].parts[0] {
5060 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
5061 other => panic!("unexpected part after approval: {other:?}"),
5062 }
5063 }
5064 other => panic!("unexpected loop step after approved cancel race: {other:?}"),
5065 }
5066 }
5067
5068 #[tokio::test]
5069 async fn loop_resumes_with_patched_input_on_approval() {
5070 let tools = ToolRegistry::new().with(EchoTool::default());
5071 let agent = Agent::builder()
5072 .model(FakeAdapter)
5073 .add_tool_source(tools)
5074 .permissions(ApproveFsReads)
5075 .build()
5076 .unwrap();
5077
5078 let mut driver = agent
5079 .start(SessionConfig {
5080 session_id: SessionId::new("session-approval-patched"),
5081 metadata: MetadataMap::new(),
5082 cache: None,
5083 })
5084 .await
5085 .unwrap();
5086
5087 driver
5088 .submit_input(vec![Item {
5089 id: None,
5090 kind: ItemKind::User,
5091 parts: vec![Part::Text(TextPart {
5092 text: "ping".into(),
5093 metadata: MetadataMap::new(),
5094 })],
5095 metadata: MetadataMap::new(),
5096 usage: None,
5097 finish_reason: None,
5098 created_at: None,
5099 }])
5100 .unwrap();
5101
5102 match driver.next().await.unwrap() {
5103 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
5104 pending
5105 .approve_with_patched_input(&mut driver, json!({ "value": "patched" }))
5106 .unwrap();
5107 }
5108 other => panic!("unexpected loop step: {other:?}"),
5109 }
5110 match driver.next().await.unwrap() {
5111 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
5112 Part::Text(text) => assert_eq!(text.text, "tool said: patched"),
5113 other => panic!("unexpected part: {other:?}"),
5114 },
5115 other => panic!("unexpected loop step after approval: {other:?}"),
5116 }
5117 }
5118
5119 #[tokio::test]
5120 async fn loop_tracks_multiple_pending_approvals_by_call_id() {
5121 let tools = ToolRegistry::new().with(EchoTool::default());
5122 let agent = Agent::builder()
5123 .model(DualApprovalAdapter)
5124 .add_tool_source(tools)
5125 .permissions(ApproveFsReads)
5126 .build()
5127 .unwrap();
5128
5129 let mut driver = agent
5130 .start(SessionConfig {
5131 session_id: SessionId::new("session-dual-approval"),
5132 metadata: MetadataMap::new(),
5133 cache: None,
5134 })
5135 .await
5136 .unwrap();
5137
5138 driver
5139 .submit_input(vec![Item {
5140 id: None,
5141 kind: ItemKind::User,
5142 parts: vec![Part::Text(TextPart {
5143 text: "run both approvals".into(),
5144 metadata: MetadataMap::new(),
5145 })],
5146 metadata: MetadataMap::new(),
5147 usage: None,
5148 finish_reason: None,
5149 created_at: None,
5150 }])
5151 .unwrap();
5152
5153 let pending_first = match driver.next().await.unwrap() {
5154 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
5155 assert_eq!(
5156 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
5157 Some("call-1")
5158 );
5159 pending
5160 }
5161 other => panic!("unexpected first loop step: {other:?}"),
5162 };
5163
5164 let pending_second = match driver.next().await.unwrap() {
5165 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
5166 assert_eq!(
5167 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
5168 Some("call-2")
5169 );
5170 pending
5171 }
5172 other => panic!("unexpected second loop step: {other:?}"),
5173 };
5174
5175 pending_second.approve(&mut driver).unwrap();
5176 match driver.next().await.unwrap() {
5177 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
5178 assert_eq!(
5179 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
5180 Some("call-1")
5181 );
5182 }
5183 other => panic!("unexpected step after approving second request: {other:?}"),
5184 }
5185
5186 pending_first.approve(&mut driver).unwrap();
5187 match driver.next().await.unwrap() {
5188 LoopStep::Finished(turn) => {
5189 assert_eq!(turn.finish_reason, FinishReason::Completed);
5190 match &turn.items[0].parts[0] {
5191 Part::Text(text) => assert_eq!(text.text, "both approvals finished"),
5192 other => panic!("unexpected final part: {other:?}"),
5193 }
5194 }
5195 other => panic!("unexpected final loop step: {other:?}"),
5196 }
5197 }
5198
5199 #[tokio::test]
5200 async fn cancelling_all_pending_approvals_pairs_every_tool_use() {
5201 let events = StdArc::new(StdMutex::new(Vec::new()));
5202 let tools = ToolRegistry::new().with(EchoTool::default());
5203 let agent = Agent::builder()
5204 .model(DualApprovalAdapter)
5205 .add_tool_source(tools)
5206 .permissions(ApproveFsReads)
5207 .observer(RecordingObserver {
5208 events: events.clone(),
5209 })
5210 .build()
5211 .unwrap();
5212
5213 let mut driver = agent
5214 .start(SessionConfig {
5215 session_id: SessionId::new("session-dual-approval-cancel"),
5216 metadata: MetadataMap::new(),
5217 cache: None,
5218 })
5219 .await
5220 .unwrap();
5221
5222 driver
5223 .submit_input(vec![Item {
5224 id: None,
5225 kind: ItemKind::User,
5226 parts: vec![Part::Text(TextPart {
5227 text: "run both approvals".into(),
5228 metadata: MetadataMap::new(),
5229 })],
5230 metadata: MetadataMap::new(),
5231 usage: None,
5232 finish_reason: None,
5233 created_at: None,
5234 }])
5235 .unwrap();
5236
5237 for expected_call in ["call-1", "call-2"] {
5238 match driver.next().await.unwrap() {
5239 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
5240 assert_eq!(
5241 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
5242 Some(expected_call)
5243 );
5244 }
5245 other => panic!("unexpected approval step: {other:?}"),
5246 }
5247 }
5248
5249 match driver.cancel_pending_approvals().await.unwrap() {
5250 Some(LoopStep::Finished(turn)) => {
5251 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
5252 }
5253 other => panic!("unexpected cancellation result: {other:?}"),
5254 }
5255 validate_transcript_invariants(&driver.snapshot().transcript).unwrap();
5256
5257 let events = events.lock().unwrap();
5258 let cancelled = events
5259 .iter()
5260 .filter(|event| matches!(event, AgentEvent::ApprovalResolved { approved: false }))
5261 .count();
5262 assert_eq!(cancelled, 2);
5263 assert!(events.iter().any(|event| matches!(
5264 event,
5265 AgentEvent::TurnFinished(turn) if turn.finish_reason == FinishReason::Cancelled
5266 )));
5267 for expected_call in ["call-1", "call-2"] {
5268 assert!(events.iter().any(|event| matches!(
5269 event,
5270 AgentEvent::ToolResultReceived(result)
5271 if result.call_id == ToolCallId::new(expected_call) && result.is_error
5272 )));
5273 }
5274 }
5275
5276 #[tokio::test]
5277 async fn loop_compacts_transcript_before_new_turns() {
5278 let events = StdArc::new(StdMutex::new(Vec::new()));
5279 let agent = Agent::builder()
5280 .model(FakeAdapter)
5281 .mutator(KeepRecentMutator { keep: 1 })
5282 .observer(RecordingObserver {
5283 events: events.clone(),
5284 })
5285 .build()
5286 .unwrap();
5287
5288 let mut driver = agent
5289 .start(SessionConfig {
5290 session_id: SessionId::new("session-4"),
5291 metadata: MetadataMap::new(),
5292 cache: None,
5293 })
5294 .await
5295 .unwrap();
5296
5297 for text in ["first", "second"] {
5298 driver
5299 .submit_input(vec![Item {
5300 id: None,
5301 kind: ItemKind::User,
5302 parts: vec![Part::Text(TextPart {
5303 text: text.into(),
5304 metadata: MetadataMap::new(),
5305 })],
5306 metadata: MetadataMap::new(),
5307 usage: None,
5308 finish_reason: None,
5309 created_at: None,
5310 }])
5311 .unwrap();
5312 let _ = driver.next().await.unwrap();
5313 }
5314
5315 let events = events.lock().unwrap();
5316 assert!(
5317 events
5318 .iter()
5319 .any(|event| matches!(event, AgentEvent::MutationFinished { dirty: true, .. }))
5320 );
5321 }
5322
5323 #[test]
5324 fn transcript_validation_rejects_orphaned_tool_result() {
5325 let transcript = vec![Item {
5326 id: None,
5327 kind: ItemKind::Tool,
5328 parts: vec![Part::ToolResult(ToolResultPart {
5329 call_id: "call-1".into(),
5330 output: ToolOutput::Text("result".into()),
5331 is_error: false,
5332 metadata: MetadataMap::new(),
5333 })],
5334 metadata: MetadataMap::new(),
5335 usage: None,
5336 finish_reason: None,
5337 created_at: None,
5338 }];
5339
5340 let error = validate_transcript_invariants(&transcript).unwrap_err();
5341 assert!(error.to_string().contains("orphaned tool_result"));
5342 }
5343
5344 #[test]
5345 fn transcript_validation_rejects_duplicate_tool_result() {
5346 let transcript = vec![
5347 Item {
5348 id: None,
5349 kind: ItemKind::Assistant,
5350 parts: vec![Part::ToolCall(ToolCallPart {
5351 id: "call-1".into(),
5352 name: "lookup".into(),
5353 input: serde_json::json!({}),
5354 metadata: MetadataMap::new(),
5355 })],
5356 metadata: MetadataMap::new(),
5357 usage: None,
5358 finish_reason: None,
5359 created_at: None,
5360 },
5361 Item {
5362 id: None,
5363 kind: ItemKind::Tool,
5364 parts: vec![Part::ToolResult(ToolResultPart {
5365 call_id: "call-1".into(),
5366 output: ToolOutput::Text("result".into()),
5367 is_error: false,
5368 metadata: MetadataMap::new(),
5369 })],
5370 metadata: MetadataMap::new(),
5371 usage: None,
5372 finish_reason: None,
5373 created_at: None,
5374 },
5375 Item {
5376 id: None,
5377 kind: ItemKind::Tool,
5378 parts: vec![Part::ToolResult(ToolResultPart {
5379 call_id: "call-1".into(),
5380 output: ToolOutput::Text("again".into()),
5381 is_error: false,
5382 metadata: MetadataMap::new(),
5383 })],
5384 metadata: MetadataMap::new(),
5385 usage: None,
5386 finish_reason: None,
5387 created_at: None,
5388 },
5389 ];
5390
5391 let error = validate_transcript_invariants(&transcript).unwrap_err();
5392 assert!(error.to_string().contains("duplicate tool_result"));
5393 }
5394
5395 #[tokio::test]
5396 async fn loop_refreshes_tool_specs_each_turn() {
5397 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
5398 let version = StdArc::new(AtomicUsize::new(1));
5399 let tools = ToolRegistry::new().with(DynamicSpecTool::new(version.clone()));
5400 let agent = Agent::builder()
5401 .model(RecordingAdapter {
5402 seen_descriptions: seen_descriptions.clone(),
5403 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
5404 })
5405 .add_tool_source(tools)
5406 .permissions(AllowAllPermissions)
5407 .build()
5408 .unwrap();
5409
5410 let mut driver = agent
5411 .start(SessionConfig {
5412 session_id: SessionId::new("session-dynamic-tools"),
5413 metadata: MetadataMap::new(),
5414 cache: None,
5415 })
5416 .await
5417 .unwrap();
5418
5419 for text in ["first", "second"] {
5420 driver
5421 .submit_input(vec![Item {
5422 id: None,
5423 kind: ItemKind::User,
5424 parts: vec![Part::Text(TextPart {
5425 text: text.into(),
5426 metadata: MetadataMap::new(),
5427 })],
5428 metadata: MetadataMap::new(),
5429 usage: None,
5430 finish_reason: None,
5431 created_at: None,
5432 }])
5433 .unwrap();
5434
5435 let _ = driver.next().await.unwrap();
5436 if text == "first" {
5437 version.store(2, Ordering::SeqCst);
5438 }
5439 }
5440
5441 let seen_descriptions = seen_descriptions.lock().unwrap();
5442 assert_eq!(seen_descriptions.len(), 2);
5443 assert_eq!(seen_descriptions[0], vec!["dynamic version 1".to_string()]);
5444 assert_eq!(seen_descriptions[1], vec!["dynamic version 2".to_string()]);
5445 }
5446
5447 #[tokio::test]
5448 async fn loop_emits_catalog_change_and_uses_updated_specs_next_turn() {
5449 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
5450 let events = StdArc::new(StdMutex::new(Vec::new()));
5451 let executor = StdArc::new(CatalogExecutor::new());
5452 let executor_for_agent: Arc<dyn ToolExecutor> = executor.clone();
5453 let agent = Agent::builder()
5454 .model(RecordingAdapter {
5455 seen_descriptions: seen_descriptions.clone(),
5456 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
5457 })
5458 .tool_executor(executor_for_agent)
5459 .permissions(AllowAllPermissions)
5460 .observer(RecordingObserver {
5461 events: events.clone(),
5462 })
5463 .build()
5464 .unwrap();
5465
5466 let mut driver = agent
5467 .start(SessionConfig {
5468 session_id: SessionId::new("session-catalog-events"),
5469 metadata: MetadataMap::new(),
5470 cache: None,
5471 })
5472 .await
5473 .unwrap();
5474
5475 driver
5476 .submit_input(vec![Item::text(ItemKind::User, "first")])
5477 .unwrap();
5478 let _ = driver.next().await.unwrap();
5479
5480 executor.publish_change(
5481 1,
5482 ToolCatalogEvent {
5483 source: "mcp:mock".into(),
5484 added: vec!["dynamic".into()],
5485 removed: Vec::new(),
5486 changed: Vec::new(),
5487 },
5488 );
5489
5490 driver
5491 .submit_input(vec![Item::text(ItemKind::User, "second")])
5492 .unwrap();
5493 let _ = driver.next().await.unwrap();
5494
5495 let seen_descriptions = seen_descriptions.lock().unwrap();
5496 assert_eq!(seen_descriptions.len(), 2);
5497 assert_eq!(seen_descriptions[0], vec!["dynamic version 0".to_string()]);
5498 assert_eq!(seen_descriptions[1], vec!["dynamic version 1".to_string()]);
5499
5500 let events = events.lock().unwrap();
5501 assert!(events.iter().any(|event| matches!(
5502 event,
5503 AgentEvent::ToolCatalogChanged(ToolCatalogEvent {
5504 source,
5505 added,
5506 removed,
5507 changed,
5508 }) if source == "mcp:mock"
5509 && added == &vec!["dynamic".to_string()]
5510 && removed.is_empty()
5511 && changed.is_empty()
5512 )));
5513 }
5514
5515 #[tokio::test]
5516 async fn loop_passes_session_default_and_next_turn_cache_requests() {
5517 let seen_caches = StdArc::new(StdMutex::new(Vec::new()));
5518 let agent = Agent::builder()
5519 .model(RecordingAdapter {
5520 seen_descriptions: StdArc::new(StdMutex::new(Vec::new())),
5521 seen_caches: seen_caches.clone(),
5522 })
5523 .permissions(AllowAllPermissions)
5524 .build()
5525 .unwrap();
5526
5527 let default_cache = PromptCacheRequest::best_effort(PromptCacheStrategy::Automatic)
5528 .with_retention(PromptCacheRetention::Short);
5529 let override_cache = PromptCacheRequest::required(PromptCacheStrategy::Explicit {
5530 breakpoints: vec![PromptCacheBreakpoint::TranscriptItemEnd { index: 0 }],
5531 });
5532
5533 let mut driver = agent
5534 .start(SessionConfig {
5535 session_id: SessionId::new("session-cache"),
5536 metadata: MetadataMap::new(),
5537 cache: Some(default_cache.clone()),
5538 })
5539 .await
5540 .unwrap();
5541
5542 driver
5543 .submit_input(vec![Item {
5544 id: None,
5545 kind: ItemKind::User,
5546 parts: vec![Part::Text(TextPart {
5547 text: "first".into(),
5548 metadata: MetadataMap::new(),
5549 })],
5550 metadata: MetadataMap::new(),
5551 usage: None,
5552 finish_reason: None,
5553 created_at: None,
5554 }])
5555 .unwrap();
5556 let _ = driver.next().await.unwrap();
5557
5558 driver
5559 .submit_input_with_cache(
5560 vec![Item {
5561 id: None,
5562 kind: ItemKind::User,
5563 parts: vec![Part::Text(TextPart {
5564 text: "second".into(),
5565 metadata: MetadataMap::new(),
5566 })],
5567 metadata: MetadataMap::new(),
5568 usage: None,
5569 finish_reason: None,
5570 created_at: None,
5571 }],
5572 override_cache.clone(),
5573 )
5574 .unwrap();
5575 let _ = driver.next().await.unwrap();
5576
5577 let seen = seen_caches.lock().unwrap();
5578 assert_eq!(seen.len(), 2);
5579 assert_eq!(seen[0], Some(default_cache));
5580 assert_eq!(seen[1], Some(override_cache));
5581 }
5582
5583 #[tokio::test]
5584 async fn loop_yields_after_tool_result_between_rounds() {
5585 let tools = ToolRegistry::new().with(EchoTool::default());
5586 let agent = Agent::builder()
5587 .model(FakeAdapter)
5588 .add_tool_source(tools)
5589 .permissions(AllowAllPermissions)
5590 .build()
5591 .unwrap();
5592
5593 let mut driver = agent
5594 .start(SessionConfig {
5595 session_id: SessionId::new("yield-session"),
5596 metadata: MetadataMap::new(),
5597 cache: None,
5598 })
5599 .await
5600 .unwrap();
5601
5602 driver
5603 .submit_input(vec![Item::text(ItemKind::User, "ping")])
5604 .unwrap();
5605
5606 let step = driver.next().await.unwrap();
5609 let info = match step {
5610 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info,
5611 other => panic!("expected AfterToolResult, got {other:?}"),
5612 };
5613 assert_eq!(info.session_id, SessionId::new("yield-session"));
5614 assert_eq!(info.transcript_len, 3);
5616
5617 let interrupt = LoopInterrupt::AfterToolResult(info.clone());
5619 assert!(!interrupt.is_blocking());
5620
5621 driver
5623 .submit_input(vec![Item::text(ItemKind::User, "also: report back")])
5624 .unwrap();
5625
5626 let step = driver.next().await.unwrap();
5629 match step {
5630 LoopStep::Finished(turn) => {
5631 assert_eq!(turn.finish_reason, FinishReason::Completed);
5632 }
5633 other => panic!("expected Finished, got {other:?}"),
5634 }
5635
5636 let snapshot = driver.snapshot();
5638 let has_injected_message = snapshot.transcript.iter().any(|item| {
5639 item.kind == ItemKind::User
5640 && item.parts.iter().any(|part| match part {
5641 Part::Text(text) => text.text == "also: report back",
5642 _ => false,
5643 })
5644 });
5645 assert!(
5646 has_injected_message,
5647 "injected user message should be in transcript, got: {:?}",
5648 snapshot.transcript
5649 );
5650 }
5651
5652 struct RecordingTranscriptObserver {
5653 items: StdArc<StdMutex<Vec<Item>>>,
5654 }
5655
5656 impl TranscriptObserver for RecordingTranscriptObserver {
5657 fn on_transcript_event(&self, event: TranscriptEvent<'_>) {
5658 self.items.lock().unwrap().push(event.item.clone());
5659 }
5660 }
5661
5662 #[tokio::test]
5663 async fn observers_see_full_tool_round() {
5664 let events = StdArc::new(StdMutex::new(Vec::<AgentEvent>::new()));
5670 let items = StdArc::new(StdMutex::new(Vec::<Item>::new()));
5671 let agent = Agent::builder()
5672 .model(FakeAdapter)
5673 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
5674 .permissions(AllowAllPermissions)
5675 .observer(RecordingObserver {
5676 events: events.clone(),
5677 })
5678 .transcript_observer(RecordingTranscriptObserver {
5679 items: items.clone(),
5680 })
5681 .build()
5682 .unwrap();
5683
5684 let mut driver = agent
5685 .start(SessionConfig {
5686 session_id: SessionId::new("observer-session"),
5687 metadata: MetadataMap::new(),
5688 cache: None,
5689 })
5690 .await
5691 .unwrap();
5692
5693 driver
5694 .submit_input(vec![Item {
5695 id: None,
5696 kind: ItemKind::User,
5697 parts: vec![Part::Text(TextPart {
5698 text: "ping".into(),
5699 metadata: MetadataMap::new(),
5700 })],
5701 metadata: MetadataMap::new(),
5702 usage: None,
5703 finish_reason: None,
5704 created_at: None,
5705 }])
5706 .unwrap();
5707
5708 let result = run_until_finished(&mut driver).await;
5709 assert!(matches!(result, LoopStep::Finished(_)), "got {result:?}");
5710
5711 let events = events.lock().unwrap().clone();
5714 let tool_call_id = events.iter().find_map(|e| match e {
5715 AgentEvent::ToolCallRequested(c) => Some(c.id.clone()),
5716 _ => None,
5717 });
5718 let tool_results: Vec<_> = events
5719 .iter()
5720 .filter_map(|e| match e {
5721 AgentEvent::ToolResultReceived(r) => Some(r.clone()),
5722 _ => None,
5723 })
5724 .collect();
5725 assert_eq!(tool_results.len(), 1, "events: {events:?}");
5726 assert_eq!(Some(tool_results[0].call_id.clone()), tool_call_id);
5727 assert!(!tool_results[0].is_error);
5728
5729 let items = items.lock().unwrap().clone();
5733 assert_eq!(items.len(), 4, "items: {items:?}");
5734 assert_eq!(items[0].kind, ItemKind::User);
5735 assert_eq!(items[1].kind, ItemKind::Assistant);
5736 assert!(
5737 items[1]
5738 .parts
5739 .iter()
5740 .any(|p| matches!(p, Part::ToolCall(_)))
5741 );
5742 assert_eq!(items[2].kind, ItemKind::Tool);
5743 assert!(
5744 items[2]
5745 .parts
5746 .iter()
5747 .any(|p| matches!(p, Part::ToolResult(_)))
5748 );
5749 assert_eq!(items[3].kind, ItemKind::Assistant);
5750 }
5751
5752 #[test]
5753 fn convenience_cache_builders_construct_expected_defaults() {
5754 let cache = PromptCacheRequest::automatic()
5755 .with_retention(PromptCacheRetention::Short)
5756 .with_key("workspace:demo");
5757 let session = SessionConfig::new("demo").with_cache(cache.clone());
5758
5759 assert_eq!(session.session_id, SessionId::new("demo"));
5760 assert_eq!(session.cache, Some(cache));
5761
5762 let explicit = PromptCacheRequest::explicit([
5763 PromptCacheBreakpoint::tools_end(),
5764 PromptCacheBreakpoint::transcript_item_end(2),
5765 PromptCacheBreakpoint::transcript_part_end(3, 1),
5766 ]);
5767
5768 assert_eq!(explicit.mode, PromptCacheMode::BestEffort);
5769 assert_eq!(
5770 explicit.strategy,
5771 PromptCacheStrategy::Explicit {
5772 breakpoints: vec![
5773 PromptCacheBreakpoint::ToolsEnd,
5774 PromptCacheBreakpoint::TranscriptItemEnd { index: 2 },
5775 PromptCacheBreakpoint::TranscriptPartEnd {
5776 item_index: 3,
5777 part_index: 1,
5778 },
5779 ],
5780 }
5781 );
5782 }
5783}