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