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