1use std::collections::{BTreeMap, 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, TaskApproval, TaskLaunchKind, TaskLaunchRequest,
74 TaskManager, TaskResolution, TaskStartContext, TaskStartOutcome, TurnTaskUpdate,
75};
76#[cfg(test)]
77use agentkit_tools_core::ToolContext;
78use agentkit_tools_core::{
79 AllowAllPermissions, ApprovalDecision, ApprovalRequest, BasicToolExecutor, OwnedToolContext,
80 PermissionChecker, ToolCatalogEvent, ToolError, ToolExecutionScope, ToolExecutor, ToolRequest,
81 ToolResources, ToolSource, ToolSpec,
82};
83use async_trait::async_trait;
84use serde::{Deserialize, Serialize};
85use thiserror::Error;
86
87const INTERRUPTED_METADATA_KEY: &str = "agentkit.interrupted";
88const INTERRUPT_REASON_METADATA_KEY: &str = "agentkit.interrupt_reason";
89const INTERRUPT_STAGE_METADATA_KEY: &str = "agentkit.interrupt_stage";
90const USER_CANCELLED_REASON: &str = "user_cancelled";
91
92#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
107pub struct SessionConfig {
108 pub session_id: SessionId,
110 pub metadata: MetadataMap,
112 pub cache: Option<PromptCacheRequest>,
114}
115
116impl SessionConfig {
117 pub fn new(session_id: impl Into<SessionId>) -> Self {
119 Self {
120 session_id: session_id.into(),
121 metadata: MetadataMap::new(),
122 cache: None,
123 }
124 }
125
126 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
128 self.metadata = metadata;
129 self
130 }
131
132 pub fn with_cache(mut self, cache: PromptCacheRequest) -> Self {
134 self.cache = Some(cache);
135 self
136 }
137
138 pub fn without_cache(mut self) -> Self {
140 self.cache = None;
141 self
142 }
143}
144
145#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
151pub enum PromptCacheMode {
152 Disabled,
154 #[default]
156 BestEffort,
157 Required,
159}
160
161#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
167pub enum PromptCacheRetention {
168 Default,
170 Short,
172 Extended,
174}
175
176#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
178pub enum PromptCacheStrategy {
179 #[default]
181 Automatic,
182 Explicit {
184 breakpoints: Vec<PromptCacheBreakpoint>,
186 },
187}
188
189impl PromptCacheStrategy {
190 pub fn automatic() -> Self {
193 Self::Automatic
194 }
195
196 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
198 Self::Explicit {
199 breakpoints: breakpoints.into_iter().collect(),
200 }
201 }
202}
203
204#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
206pub enum PromptCacheBreakpoint {
207 ToolsEnd,
209 TranscriptItemEnd { index: usize },
211 TranscriptPartEnd {
217 item_index: usize,
218 part_index: usize,
219 },
220}
221
222impl PromptCacheBreakpoint {
223 pub fn tools_end() -> Self {
225 Self::ToolsEnd
226 }
227
228 pub fn transcript_item_end(index: usize) -> Self {
230 Self::TranscriptItemEnd { index }
231 }
232
233 pub fn transcript_part_end(item_index: usize, part_index: usize) -> Self {
235 Self::TranscriptPartEnd {
236 item_index,
237 part_index,
238 }
239 }
240}
241
242#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
244pub struct PromptCacheRequest {
245 pub mode: PromptCacheMode,
247 pub strategy: PromptCacheStrategy,
249 pub retention: Option<PromptCacheRetention>,
251 pub key: Option<String>,
253}
254
255impl PromptCacheRequest {
256 pub fn automatic() -> Self {
258 Self::best_effort(PromptCacheStrategy::automatic())
259 }
260
261 pub fn automatic_required() -> Self {
263 Self::required(PromptCacheStrategy::automatic())
264 }
265
266 pub fn explicit(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
268 Self::best_effort(PromptCacheStrategy::explicit(breakpoints))
269 }
270
271 pub fn explicit_required(breakpoints: impl IntoIterator<Item = PromptCacheBreakpoint>) -> Self {
273 Self::required(PromptCacheStrategy::explicit(breakpoints))
274 }
275
276 pub fn disabled() -> Self {
278 Self {
279 mode: PromptCacheMode::Disabled,
280 strategy: PromptCacheStrategy::Automatic,
281 retention: None,
282 key: None,
283 }
284 }
285
286 pub fn best_effort(strategy: PromptCacheStrategy) -> Self {
288 Self {
289 mode: PromptCacheMode::BestEffort,
290 strategy,
291 retention: None,
292 key: None,
293 }
294 }
295
296 pub fn required(strategy: PromptCacheStrategy) -> Self {
298 Self {
299 mode: PromptCacheMode::Required,
300 strategy,
301 retention: None,
302 key: None,
303 }
304 }
305
306 pub fn with_mode(mut self, mode: PromptCacheMode) -> Self {
308 self.mode = mode;
309 self
310 }
311
312 pub fn with_strategy(mut self, strategy: PromptCacheStrategy) -> Self {
314 self.strategy = strategy;
315 self
316 }
317
318 pub fn with_retention(mut self, retention: PromptCacheRetention) -> Self {
320 self.retention = Some(retention);
321 self
322 }
323
324 pub fn with_key(mut self, key: impl Into<String>) -> Self {
326 self.key = Some(key.into());
327 self
328 }
329
330 pub fn without_retention(mut self) -> Self {
332 self.retention = None;
333 self
334 }
335
336 pub fn without_key(mut self) -> Self {
338 self.key = None;
339 self
340 }
341
342 pub fn is_enabled(&self) -> bool {
344 !matches!(self.mode, PromptCacheMode::Disabled)
345 }
346}
347
348#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
354pub struct TurnRequest {
355 pub session_id: SessionId,
357 pub turn_id: agentkit_core::TurnId,
359 pub transcript: Vec<Item>,
361 pub available_tools: Vec<ToolSpec>,
363 pub cache: Option<PromptCacheRequest>,
365 pub metadata: MetadataMap,
367}
368
369#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
374pub struct ModelTurnResult {
375 pub finish_reason: FinishReason,
377 pub output_items: Vec<Item>,
379 pub usage: Option<Usage>,
381 pub metadata: MetadataMap,
383 #[serde(default)]
387 pub model: Option<String>,
388 #[serde(default)]
392 pub response_id: Option<String>,
393}
394
395#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
401pub enum ModelTurnEvent {
402 Delta(Delta),
404 ToolCall(ToolCallPart),
406 Usage(Usage),
408 Finished(ModelTurnResult),
410}
411
412#[async_trait]
449pub trait ModelAdapter: Send + Sync {
450 type Session: ModelSession;
452
453 async fn start_session(&self, config: SessionConfig) -> Result<Self::Session, LoopError>;
459
460 fn provider_name(&self) -> Option<&str> {
467 None
468 }
469}
470
471#[async_trait]
478pub trait ModelSession: Send {
479 type Turn: ModelTurn;
481
482 async fn begin_turn(
495 &mut self,
496 request: TurnRequest,
497 cancellation: Option<TurnCancellation>,
498 ) -> Result<Self::Turn, LoopError>;
499
500 fn model_name(&self) -> Option<&str> {
506 None
507 }
508}
509
510#[async_trait]
516pub trait ModelTurn: Send {
517 async fn next_event(
526 &mut self,
527 cancellation: Option<TurnCancellation>,
528 ) -> Result<Option<ModelTurnEvent>, LoopError>;
529}
530
531pub trait LoopObserver: Send + Sync {
551 fn handle_event(&self, event: AgentEvent);
556}
557
558pub trait TranscriptObserver: Send + Sync {
593 fn on_item_appended(&self, item: &Item);
598}
599
600#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
605#[non_exhaustive]
606pub enum MutationPoint {
607 AfterToolResult,
610 AfterTurnEnded,
613}
614
615pub trait EventEmitter: Send + Sync {
618 fn emit(&self, event: AgentEvent);
620}
621
622#[non_exhaustive]
624pub struct LoopCtx<'a> {
625 pub session_id: &'a SessionId,
627 pub turn_id: Option<&'a agentkit_core::TurnId>,
629 pub point: MutationPoint,
631 pub cancellation: Option<TurnCancellation>,
633 pub emitter: &'a dyn EventEmitter,
635}
636
637pub struct TranscriptCursor<'a> {
645 items: &'a mut Vec<Item>,
646 pub(crate) dirty: bool,
647}
648
649impl<'a> std::ops::Deref for TranscriptCursor<'a> {
650 type Target = Vec<Item>;
651 fn deref(&self) -> &Vec<Item> {
652 self.items
653 }
654}
655
656impl<'a> std::ops::DerefMut for TranscriptCursor<'a> {
657 fn deref_mut(&mut self) -> &mut Vec<Item> {
658 self.dirty = true;
659 self.items
660 }
661}
662
663#[async_trait]
671pub trait LoopMutator: Send + Sync {
672 async fn mutate(
677 &self,
678 cursor: &mut TranscriptCursor<'_>,
679 ctx: LoopCtx<'_>,
680 ) -> Result<(), LoopError> {
681 let _ = (cursor, ctx);
682 Ok(())
683 }
684}
685
686#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
691pub enum AgentEvent {
692 RunStarted { session_id: SessionId },
694 TurnStarted {
696 session_id: SessionId,
697 turn_id: agentkit_core::TurnId,
698 },
699 InputAccepted {
701 session_id: SessionId,
702 items: Vec<Item>,
703 },
704 ContentDelta(Delta),
706 ToolCallRequested(ToolCallPart),
708 ToolResultReceived(ToolResultPart),
720 ApprovalRequired(ApprovalRequest),
722 ApprovalResolved { approved: bool },
724 ToolCatalogChanged(ToolCatalogEvent),
726 MutationStarted {
729 session_id: SessionId,
730 turn_id: Option<agentkit_core::TurnId>,
731 mutator: String,
732 point: MutationPoint,
733 },
734 MutationFinished {
738 session_id: SessionId,
739 turn_id: Option<agentkit_core::TurnId>,
740 mutator: String,
741 dirty: bool,
742 metadata: MetadataMap,
743 },
744 UsageUpdated(Usage),
746 Warning { message: String },
748 RunFailed { message: String },
750 TurnFinished(TurnResult),
752}
753
754#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
776pub struct PendingApproval {
777 pub request: ApprovalRequest,
779}
780
781impl std::ops::Deref for PendingApproval {
782 type Target = ApprovalRequest;
783 fn deref(&self) -> &ApprovalRequest {
784 &self.request
785 }
786}
787
788impl PendingApproval {
789 pub fn approve<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
791 let call_id = self
792 .request
793 .call_id
794 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
795 driver.resolve_approval_for(call_id, ApprovalDecision::Approve)
796 }
797
798 pub fn deny<S: ModelSession>(self, driver: &mut LoopDriver<S>) -> Result<(), LoopError> {
800 let call_id = self
801 .request
802 .call_id
803 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
804 driver.resolve_approval_for(call_id, ApprovalDecision::Deny { reason: None })
805 }
806
807 pub fn deny_with_reason<S: ModelSession>(
809 self,
810 driver: &mut LoopDriver<S>,
811 reason: impl Into<String>,
812 ) -> Result<(), LoopError> {
813 let call_id = self
814 .request
815 .call_id
816 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
817 driver.resolve_approval_for(
818 call_id,
819 ApprovalDecision::Deny {
820 reason: Some(reason.into()),
821 },
822 )
823 }
824
825 pub fn approve_with_patched_input<S: ModelSession>(
835 self,
836 driver: &mut LoopDriver<S>,
837 input: serde_json::Value,
838 ) -> Result<(), LoopError> {
839 let call_id = self
840 .request
841 .call_id
842 .ok_or_else(|| LoopError::InvalidState("pending approval is missing call id".into()))?;
843 driver.resolve_approval_for_with_patched_input(call_id, input)
844 }
845}
846
847#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
872pub struct InputRequest {
873 pub session_id: SessionId,
875 pub reason: String,
877}
878
879impl InputRequest {
880 pub fn submit<S: ModelSession>(
882 self,
883 driver: &mut LoopDriver<S>,
884 items: Vec<Item>,
885 ) -> Result<(), LoopError> {
886 driver.submit_input(items)
887 }
888}
889
890#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
895pub struct TurnResult {
896 pub turn_id: agentkit_core::TurnId,
898 pub finish_reason: FinishReason,
900 pub items: Vec<Item>,
902 pub usage: Option<Usage>,
904 pub metadata: MetadataMap,
906}
907
908#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
944pub enum LoopInterrupt {
945 ApprovalRequest(PendingApproval),
947 AwaitingInput(InputRequest),
949 AfterToolResult(ToolRoundInfo),
959}
960
961impl LoopInterrupt {
962 pub fn is_blocking(&self) -> bool {
968 matches!(self, LoopInterrupt::ApprovalRequest(_))
969 }
970}
971
972#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
975pub struct ToolRoundInfo {
976 pub session_id: SessionId,
978 pub turn_id: agentkit_core::TurnId,
980 pub transcript_len: usize,
982}
983
984impl ToolRoundInfo {
985 pub fn submit<S: ModelSession>(
988 self,
989 driver: &mut LoopDriver<S>,
990 items: Vec<Item>,
991 ) -> Result<(), LoopError> {
992 driver.submit_input(items)
993 }
994}
995
996#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1024pub enum LoopStep {
1025 Interrupt(LoopInterrupt),
1027 Finished(TurnResult),
1029}
1030
1031#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1037pub struct LoopSnapshot {
1038 pub session_id: SessionId,
1040 pub transcript: Vec<Item>,
1042 pub pending_input: Vec<Item>,
1044}
1045
1046#[derive(Clone, Debug)]
1047struct PendingApprovalToolCall {
1048 request: ApprovalRequest,
1049 decision: Option<ApprovalDecision>,
1050 surfaced: bool,
1051 turn_id: agentkit_core::TurnId,
1052 task_id: TaskId,
1053 call: ToolCallPart,
1054 tool_request: ToolRequest,
1055}
1056
1057#[derive(Clone, Debug, Default)]
1058struct ActiveToolRound {
1059 turn_id: agentkit_core::TurnId,
1060 pending_calls: VecDeque<(ToolCallPart, ToolRequest)>,
1061 background_pending: bool,
1062 foreground_progressed: bool,
1063}
1064
1065pub struct Agent<M>
1107where
1108 M: ModelAdapter,
1109{
1110 model: M,
1111 tool_sources: Vec<Arc<dyn ToolSource>>,
1112 tool_executor: Option<Arc<dyn ToolExecutor>>,
1113 task_manager: Arc<dyn TaskManager>,
1114 permissions: Arc<dyn PermissionChecker>,
1115 resources: Arc<dyn ToolResources>,
1116 cancellation: Option<CancellationHandle>,
1117 mutators: Vec<Arc<dyn LoopMutator>>,
1118 observers: Vec<Arc<dyn LoopObserver>>,
1119 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1120 transcript: Vec<Item>,
1121 input: Vec<Item>,
1122}
1123
1124impl<M> Agent<M>
1125where
1126 M: ModelAdapter,
1127{
1128 pub fn builder() -> AgentBuilder<M> {
1130 AgentBuilder::default()
1131 }
1132
1133 pub async fn start(&self, config: SessionConfig) -> Result<LoopDriver<M::Session>, LoopError> {
1149 let session_id = config.session_id.clone();
1150 let default_cache = config.cache.clone();
1151 let session = self.model.start_session(config).await?;
1152 let tool_executor = self
1153 .tool_executor
1154 .clone()
1155 .unwrap_or_else(|| Arc::new(BasicToolExecutor::new(self.tool_sources.clone())));
1156 let driver = LoopDriver {
1157 session_id: session_id.clone(),
1158 provider_name: self.model.provider_name().map(str::to_owned),
1159 default_cache,
1160 next_turn_cache: None,
1161 session: Some(session),
1162 tool_executor,
1163 task_manager: self.task_manager.clone(),
1164 permissions: self.permissions.clone(),
1165 resources: self.resources.clone(),
1166 cancellation: self.cancellation.clone(),
1167 mutators: self.mutators.clone(),
1168 observers: self.observers.clone(),
1169 transcript_observers: self.transcript_observers.clone(),
1170 transcript: self.transcript.clone(),
1171 pending_input: self.input.clone(),
1172 pending_approvals: BTreeMap::new(),
1173 pending_approval_order: VecDeque::new(),
1174 active_tool_round: None,
1175 pending_round_resume: None,
1176 next_turn_index: 1,
1177 detached_call_ids: HashSet::new(),
1178 };
1179 driver.emit(AgentEvent::RunStarted { session_id });
1180 Ok(driver)
1181 }
1182}
1183
1184pub struct AgentBuilder<M>
1190where
1191 M: ModelAdapter,
1192{
1193 model: Option<M>,
1194 tool_sources: Vec<Arc<dyn ToolSource>>,
1195 tool_executor: Option<Arc<dyn ToolExecutor>>,
1196 task_manager: Option<Arc<dyn TaskManager>>,
1197 permissions: Arc<dyn PermissionChecker>,
1198 resources: Arc<dyn ToolResources>,
1199 cancellation: Option<CancellationHandle>,
1200 mutators: Vec<Arc<dyn LoopMutator>>,
1201 observers: Vec<Arc<dyn LoopObserver>>,
1202 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1203 transcript: Vec<Item>,
1204 input: Vec<Item>,
1205}
1206
1207impl<M> Default for AgentBuilder<M>
1208where
1209 M: ModelAdapter,
1210{
1211 fn default() -> Self {
1212 Self {
1213 model: None,
1214 tool_sources: Vec::new(),
1215 tool_executor: None,
1216 task_manager: None,
1217 permissions: Arc::new(AllowAllPermissions),
1218 resources: Arc::new(()),
1219 cancellation: None,
1220 mutators: Vec::new(),
1221 observers: Vec::new(),
1222 transcript_observers: Vec::new(),
1223 transcript: Vec::new(),
1224 input: Vec::new(),
1225 }
1226 }
1227}
1228
1229impl<M> AgentBuilder<M>
1230where
1231 M: ModelAdapter,
1232{
1233 pub fn model(mut self, model: M) -> Self {
1235 self.model = Some(model);
1236 self
1237 }
1238
1239 pub fn add_tool_source<S: ToolSource + 'static>(mut self, source: S) -> Self {
1252 self.tool_sources.push(Arc::new(source));
1253 self
1254 }
1255
1256 pub fn tool_executor(mut self, executor: impl ToolExecutor + 'static) -> Self {
1262 self.tool_executor = Some(Arc::new(executor));
1263 self
1264 }
1265
1266 pub fn task_manager(mut self, manager: impl TaskManager + 'static) -> Self {
1271 self.task_manager = Some(Arc::new(manager));
1272 self
1273 }
1274
1275 pub fn permissions(mut self, permissions: impl PermissionChecker + 'static) -> Self {
1279 self.permissions = Arc::new(permissions);
1280 self
1281 }
1282
1283 pub fn resources(mut self, resources: impl ToolResources + 'static) -> Self {
1285 self.resources = Arc::new(resources);
1286 self
1287 }
1288
1289 pub fn cancellation(mut self, handle: CancellationHandle) -> Self {
1291 self.cancellation = Some(handle);
1292 self
1293 }
1294
1295 pub fn mutator<L: LoopMutator + 'static>(mut self, mutator: L) -> Self {
1303 self.mutators.push(Arc::new(mutator));
1304 self
1305 }
1306
1307 pub fn observer<O: LoopObserver + 'static>(mut self, observer: O) -> Self {
1311 self.observers.push(Arc::new(observer));
1312 self
1313 }
1314
1315 pub fn transcript_observer<O: TranscriptObserver + 'static>(mut self, observer: O) -> Self {
1324 self.transcript_observers.push(Arc::new(observer));
1325 self
1326 }
1327
1328 pub fn transcript(mut self, transcript: Vec<Item>) -> Self {
1336 self.transcript = transcript;
1337 self
1338 }
1339
1340 pub fn input(mut self, input: Vec<Item>) -> Self {
1349 self.input = input;
1350 self
1351 }
1352
1353 pub fn build(self) -> Result<Agent<M>, LoopError> {
1359 let model = self
1360 .model
1361 .ok_or_else(|| LoopError::InvalidState("model adapter is required".into()))?;
1362 Ok(Agent {
1363 model,
1364 tool_sources: self.tool_sources,
1365 tool_executor: self.tool_executor,
1366 task_manager: self
1367 .task_manager
1368 .unwrap_or_else(|| Arc::new(SimpleTaskManager::new())),
1369 permissions: self.permissions,
1370 resources: self.resources,
1371 cancellation: self.cancellation,
1372 mutators: self.mutators,
1373 observers: self.observers,
1374 transcript_observers: self.transcript_observers,
1375 transcript: self.transcript,
1376 input: self.input,
1377 })
1378 }
1379}
1380
1381pub struct LoopDriver<S>
1412where
1413 S: ModelSession,
1414{
1415 session_id: SessionId,
1416 provider_name: Option<String>,
1417 default_cache: Option<PromptCacheRequest>,
1418 next_turn_cache: Option<PromptCacheRequest>,
1419 session: Option<S>,
1420 tool_executor: Arc<dyn ToolExecutor>,
1421 task_manager: Arc<dyn TaskManager>,
1422 permissions: Arc<dyn PermissionChecker>,
1423 resources: Arc<dyn ToolResources>,
1424 cancellation: Option<CancellationHandle>,
1425 mutators: Vec<Arc<dyn LoopMutator>>,
1426 observers: Vec<Arc<dyn LoopObserver>>,
1427 transcript_observers: Vec<Arc<dyn TranscriptObserver>>,
1428 transcript: Vec<Item>,
1429 pending_input: Vec<Item>,
1430 pending_approvals: BTreeMap<ToolCallId, PendingApprovalToolCall>,
1431 pending_approval_order: VecDeque<ToolCallId>,
1432 active_tool_round: Option<ActiveToolRound>,
1433 pending_round_resume: Option<agentkit_core::TurnId>,
1434 next_turn_index: u64,
1435 detached_call_ids: HashSet<ToolCallId>,
1443}
1444
1445impl<S> LoopDriver<S>
1446where
1447 S: ModelSession,
1448{
1449 fn execute_tool_span(
1450 &self,
1451 request: &ToolRequest,
1452 turn_id: &agentkit_core::TurnId,
1453 launch_kind: &'static str,
1454 ) -> tracing::Span {
1455 tracing::info_span!(
1456 "agent.execute_tool",
1457 "otel.name" = %format!("execute_tool {}", request.tool_name),
1458 "gen_ai.operation.name" = "execute_tool",
1459 "gen_ai.tool.name" = %request.tool_name,
1460 "gen_ai.tool.call.id" = %request.call_id,
1461 "gen_ai.conversation.id" = %self.session_id,
1462 "error.type" = tracing::field::Empty,
1463 session.id = %self.session_id,
1464 turn.id = %turn_id,
1465 launch_kind = launch_kind,
1466 )
1467 }
1468
1469 fn start_task_via_manager(
1470 &self,
1471 task_id: Option<TaskId>,
1472 tool_request: ToolRequest,
1473 kind: TaskLaunchKind,
1474 cancellation: Option<TurnCancellation>,
1475 ) -> impl std::future::Future<Output = Result<TaskStartOutcome, LoopError>> + Send + 'static
1476 {
1477 let task_manager = self.task_manager.clone();
1478 let tool_executor = self.tool_executor.clone();
1479 let permissions = self.permissions.clone();
1480 let resources = self.resources.clone();
1481 let session_id = self.session_id.clone();
1482 let turn_id = tool_request.turn_id.clone();
1483 let metadata = tool_request.metadata.clone();
1484
1485 async move {
1486 task_manager
1487 .start_task(
1488 TaskLaunchRequest {
1489 task_id,
1490 request: tool_request.clone(),
1491 kind,
1492 },
1493 TaskStartContext {
1494 executor: tool_executor.clone(),
1495 tool_context: {
1496 let execution_scope = ToolExecutionScope {
1497 executor: tool_executor,
1498 session_id: session_id.clone(),
1499 turn_id: turn_id.clone(),
1500 permissions: permissions.clone(),
1501 resources: resources.clone(),
1502 cancellation: cancellation.clone(),
1503 };
1504 OwnedToolContext {
1505 session_id,
1506 turn_id,
1507 metadata,
1508 permissions,
1509 resources,
1510 cancellation,
1511 execution_scope: Some(execution_scope),
1512 approved_request: None,
1513 }
1514 },
1515 },
1516 )
1517 .await
1518 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))
1519 }
1520 }
1521
1522 fn has_pending_interrupts(&self) -> bool {
1523 !self.pending_approvals.is_empty()
1524 }
1525
1526 fn emit_tool_catalog_events(&mut self, events: Vec<ToolCatalogEvent>) {
1527 for event in events {
1528 self.emit(AgentEvent::ToolCatalogChanged(event));
1529 }
1530 }
1531
1532 fn enqueue_pending_approval(&mut self, turn_id: &agentkit_core::TurnId, task: TaskApproval) {
1533 let call_id = task.tool_request.call_id.clone();
1534 let call = ToolCallPart {
1535 id: call_id.clone(),
1536 name: task.tool_request.tool_name.to_string(),
1537 input: task.tool_request.input.clone(),
1538 metadata: task.tool_request.metadata.clone(),
1539 };
1540 let mut request = task.approval;
1541 request.call_id = Some(call_id.clone());
1542 let pending = PendingApprovalToolCall {
1543 request: request.clone(),
1544 decision: None,
1545 surfaced: false,
1546 turn_id: turn_id.clone(),
1547 task_id: task.task_id,
1548 call,
1549 tool_request: task.tool_request,
1550 };
1551 self.pending_approvals.insert(call_id.clone(), pending);
1552 if !self.pending_approval_order.iter().any(|id| id == &call_id) {
1553 self.pending_approval_order.push_back(call_id);
1554 }
1555 self.emit(AgentEvent::ApprovalRequired(request));
1556 }
1557
1558 fn take_next_unsurfaced_approval_interrupt(&mut self) -> Option<LoopStep> {
1559 for call_id in self.pending_approval_order.clone() {
1560 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
1561 continue;
1562 };
1563 if pending.decision.is_none() && !pending.surfaced {
1564 pending.surfaced = true;
1565 return Some(LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(
1566 PendingApproval {
1567 request: pending.request.clone(),
1568 },
1569 )));
1570 }
1571 }
1572 None
1573 }
1574
1575 fn next_unresolved_approval_interrupt(&self) -> Option<LoopStep> {
1576 self.pending_approval_order.iter().find_map(|call_id| {
1577 self.pending_approvals.get(call_id).and_then(|pending| {
1578 pending.decision.is_none().then(|| {
1579 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(PendingApproval {
1580 request: pending.request.clone(),
1581 }))
1582 })
1583 })
1584 })
1585 }
1586
1587 fn take_next_resolved_approval(&mut self) -> Option<PendingApprovalToolCall> {
1588 let call_id = self.pending_approval_order.iter().find_map(|call_id| {
1589 self.pending_approvals
1590 .get(call_id)
1591 .and_then(|pending| pending.decision.as_ref().map(|_| call_id.clone()))
1592 })?;
1593 self.pending_approval_order.retain(|id| id != &call_id);
1594 self.pending_approvals.remove(&call_id)
1595 }
1596
1597 fn queue_resolution_interrupt(
1598 &mut self,
1599 turn_id: &agentkit_core::TurnId,
1600 resolution: TaskResolution,
1601 ) -> Option<LoopStep> {
1602 match resolution {
1603 TaskResolution::Item(item) => {
1604 self.append_tool_result_item(item);
1605 None
1606 }
1607 TaskResolution::Approval(task) => {
1608 self.enqueue_pending_approval(turn_id, task);
1609 self.take_next_unsurfaced_approval_interrupt()
1610 }
1611 }
1612 }
1613
1614 async fn drain_pending_loop_updates(&mut self) -> Result<(bool, Option<LoopStep>), LoopError> {
1615 let PendingLoopUpdates { mut resolutions } = self
1616 .task_manager
1617 .take_pending_loop_updates()
1618 .await
1619 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
1620 let mut saw_items = false;
1621 while let Some(resolution) = resolutions.pop_front() {
1622 match resolution {
1623 TaskResolution::Item(item) => {
1624 self.append_tool_result_item(item);
1625 saw_items = true;
1626 }
1627 TaskResolution::Approval(task) => {
1628 self.enqueue_pending_approval(&task.tool_request.turn_id.clone(), task);
1629 }
1630 }
1631 }
1632 Ok((saw_items, self.take_next_unsurfaced_approval_interrupt()))
1633 }
1634
1635 async fn run_mutators(
1636 &mut self,
1637 point: MutationPoint,
1638 turn_id: Option<&agentkit_core::TurnId>,
1639 cancellation: Option<TurnCancellation>,
1640 ) -> Result<(), LoopError> {
1641 if self.mutators.is_empty() {
1642 return Ok(());
1643 }
1644 if cancellation
1645 .as_ref()
1646 .is_some_and(TurnCancellation::is_cancelled)
1647 {
1648 return Err(LoopError::Cancelled);
1649 }
1650 let mutators = self.mutators.clone();
1651 let session_id = self.session_id.clone();
1652 let observers = self.observers.clone();
1653 let emitter = DriverEmitter {
1654 observers: &observers,
1655 };
1656 let mut cursor = TranscriptCursor {
1657 items: &mut self.transcript,
1658 dirty: false,
1659 };
1660 for mutator in &mutators {
1661 if cancellation
1662 .as_ref()
1663 .is_some_and(TurnCancellation::is_cancelled)
1664 {
1665 return Err(LoopError::Cancelled);
1666 }
1667 let ctx = LoopCtx {
1668 session_id: &session_id,
1669 turn_id,
1670 point,
1671 cancellation: cancellation.clone(),
1672 emitter: &emitter,
1673 };
1674 mutator.mutate(&mut cursor, ctx).await?;
1675 }
1676 if cursor.dirty {
1677 validate_transcript_invariants(cursor.items)?;
1678 }
1679 Ok(())
1680 }
1681
1682 async fn continue_active_tool_round(&mut self) -> Result<Option<LoopStep>, LoopError> {
1683 let Some(_) = self.active_tool_round.as_ref() else {
1684 return Ok(None);
1685 };
1686 loop {
1687 let cancellation = self
1688 .cancellation
1689 .as_ref()
1690 .map(CancellationHandle::checkpoint);
1691 let turn_id = self
1692 .active_tool_round
1693 .as_ref()
1694 .map(|active| active.turn_id.clone())
1695 .ok_or_else(|| LoopError::InvalidState("missing active tool round".into()))?;
1696
1697 if cancellation
1698 .as_ref()
1699 .is_some_and(TurnCancellation::is_cancelled)
1700 {
1701 self.task_manager
1702 .on_turn_interrupted(&turn_id)
1703 .await
1704 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
1705 self.active_tool_round = None;
1706 return self.finish_cancelled(turn_id, Vec::new()).map(Some);
1707 }
1708
1709 let next_call = self
1710 .active_tool_round
1711 .as_mut()
1712 .and_then(|active| active.pending_calls.pop_front());
1713 if let Some((_call, tool_request)) = next_call {
1714 use tracing::Instrument;
1715 let dispatch_span = self.execute_tool_span(&tool_request, &turn_id, "plain");
1716 match self
1717 .start_task_via_manager(
1718 None,
1719 tool_request.clone(),
1720 TaskLaunchKind::Plain,
1721 cancellation.clone(),
1722 )
1723 .instrument(dispatch_span.clone())
1724 .await?
1725 {
1726 TaskStartOutcome::Ready(resolution) => {
1727 let resolution = *resolution;
1728 match resolution {
1729 TaskResolution::Item(item) => {
1730 if tool_result_is_error(&item) {
1731 dispatch_span.record("error.type", "tool_error");
1732 }
1733 if let Some(active) = self.active_tool_round.as_mut() {
1734 active.foreground_progressed = true;
1735 }
1736 self.append_tool_result_item(item);
1737 }
1738 TaskResolution::Approval(task) => {
1739 self.enqueue_pending_approval(&turn_id, task);
1740 }
1741 }
1742 continue;
1743 }
1744 TaskStartOutcome::Pending { kind, .. } => {
1745 if kind == agentkit_task_manager::TaskKind::Background
1746 && let Some(active) = self.active_tool_round.as_mut()
1747 {
1748 active.background_pending = true;
1749 }
1750 continue;
1751 }
1752 }
1753 }
1754
1755 match self
1756 .task_manager
1757 .wait_for_turn(&turn_id, cancellation.clone())
1758 .await
1759 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?
1760 {
1761 Some(TurnTaskUpdate::Resolution(resolution)) => {
1762 let resolution = *resolution;
1763 match resolution {
1764 TaskResolution::Item(item) => {
1765 if let Some(active) = self.active_tool_round.as_mut() {
1766 active.foreground_progressed = true;
1767 }
1768 self.append_tool_result_item(item);
1769 }
1770 TaskResolution::Approval(task) => {
1771 self.enqueue_pending_approval(&turn_id, task);
1772 }
1773 }
1774 }
1775 Some(TurnTaskUpdate::Detached(snapshot)) => {
1776 let detached_call_id = snapshot.call_id.clone();
1794 self.append_tool_result_item(Item {
1795 id: None,
1796 kind: ItemKind::Tool,
1797 parts: vec![Part::ToolResult(ToolResultPart {
1798 call_id: detached_call_id.clone(),
1799 output: ToolOutput::Text(format!(
1800 "Tool {} is now running in the background. \
1801 The result will be delivered when it completes.",
1802 snapshot.tool_name,
1803 )),
1804 is_error: false,
1805 metadata: MetadataMap::new(),
1806 })],
1807 metadata: MetadataMap::new(),
1808 usage: None,
1809 finish_reason: None,
1810 created_at: None,
1811 });
1812 self.detached_call_ids.insert(detached_call_id);
1813 if let Some(active) = self.active_tool_round.as_mut() {
1814 active.background_pending = true;
1815 active.foreground_progressed = true;
1816 }
1817 }
1818 None => {
1819 if cancellation
1820 .as_ref()
1821 .is_some_and(TurnCancellation::is_cancelled)
1822 {
1823 self.task_manager
1824 .on_turn_interrupted(&turn_id)
1825 .await
1826 .map_err(|error| {
1827 LoopError::Tool(ToolError::Internal(error.to_string()))
1828 })?;
1829 self.active_tool_round = None;
1830 return self.finish_cancelled(turn_id, Vec::new()).map(Some);
1831 }
1832 let active = self.active_tool_round.take().ok_or_else(|| {
1833 LoopError::InvalidState("missing active tool round".into())
1834 })?;
1835 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
1836 return Ok(Some(step));
1837 }
1838 if let Some(step) = self.next_unresolved_approval_interrupt() {
1839 return Ok(Some(step));
1840 }
1841 if active.background_pending && !active.foreground_progressed {
1842 return Ok(None);
1843 }
1844 let info = ToolRoundInfo {
1851 session_id: self.session_id.clone(),
1852 turn_id: turn_id.clone(),
1853 transcript_len: self.transcript.len(),
1854 };
1855 self.pending_round_resume = Some(turn_id);
1856 return Ok(Some(LoopStep::Interrupt(LoopInterrupt::AfterToolResult(
1857 info,
1858 ))));
1859 }
1860 }
1861 }
1862 }
1863
1864 #[tracing::instrument(
1865 name = "agent.turn",
1866 skip_all,
1867 fields(
1868 otel.name = "invoke_agent",
1869 gen_ai.operation.name = "invoke_agent",
1870 gen_ai.conversation.id = %self.session_id,
1871 gen_ai.provider.name = tracing::field::Empty,
1872 gen_ai.usage.input_tokens = tracing::field::Empty,
1873 gen_ai.usage.output_tokens = tracing::field::Empty,
1874 gen_ai.usage.cost = tracing::field::Empty,
1875 session.id = %self.session_id,
1876 turn.id = %turn_id,
1877 transcript.len = self.transcript.len(),
1878 saw_tool_call = tracing::field::Empty,
1879 finish_reason = tracing::field::Empty,
1880 ),
1881 )]
1882 async fn drive_turn(
1883 &mut self,
1884 turn_id: agentkit_core::TurnId,
1885 emit_started: bool,
1886 mutation_point: MutationPoint,
1887 ) -> Result<LoopStep, LoopError> {
1888 if let Some(provider) = &self.provider_name {
1889 tracing::Span::current().record("gen_ai.provider.name", provider.as_str());
1890 }
1891 let cancellation = self
1892 .cancellation
1893 .as_ref()
1894 .map(CancellationHandle::checkpoint);
1895 match self
1896 .run_mutators(mutation_point, Some(&turn_id), cancellation.clone())
1897 .await
1898 {
1899 Ok(()) => {}
1900 Err(LoopError::Cancelled) => {
1901 return self.finish_cancelled(turn_id, interrupted_assistant_items());
1902 }
1903 Err(error) => return Err(error),
1904 }
1905
1906 if !transcript_has_pending_input(&self.transcript) {
1912 let turn_result = TurnResult {
1913 turn_id,
1914 finish_reason: FinishReason::Completed,
1915 items: Vec::new(),
1916 usage: None,
1917 metadata: MetadataMap::new(),
1918 };
1919 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
1920 return Ok(LoopStep::Finished(turn_result));
1921 }
1922
1923 if emit_started {
1924 self.emit(AgentEvent::TurnStarted {
1925 session_id: self.session_id.clone(),
1926 turn_id: turn_id.clone(),
1927 });
1928 }
1929 if cancellation
1930 .as_ref()
1931 .is_some_and(TurnCancellation::is_cancelled)
1932 {
1933 return self.finish_cancelled(turn_id, interrupted_assistant_items());
1934 }
1935
1936 let catalog_events = self.tool_executor.drain_catalog_events();
1937 self.emit_tool_catalog_events(catalog_events);
1938
1939 let request = TurnRequest {
1940 session_id: self.session_id.clone(),
1941 turn_id: turn_id.clone(),
1942 transcript: self.transcript.clone(),
1943 available_tools: self.tool_executor.specs(),
1944 cache: self
1945 .next_turn_cache
1946 .take()
1947 .or_else(|| self.default_cache.clone()),
1948 metadata: MetadataMap::new(),
1949 };
1950
1951 let session = self
1952 .session
1953 .as_mut()
1954 .ok_or_else(|| LoopError::InvalidState("model session is not available".into()))?;
1955
1956 let chat_span = tracing::info_span!(
1963 "chat",
1964 "otel.name" = tracing::field::Empty,
1965 "otel.kind" = "client",
1966 "gen_ai.operation.name" = "chat",
1967 "gen_ai.provider.name" = tracing::field::Empty,
1968 "gen_ai.conversation.id" = %self.session_id,
1969 "gen_ai.request.model" = tracing::field::Empty,
1970 "gen_ai.response.model" = tracing::field::Empty,
1971 "gen_ai.response.id" = tracing::field::Empty,
1972 "gen_ai.response.finish_reasons" = tracing::field::Empty,
1973 "gen_ai.usage.input_tokens" = tracing::field::Empty,
1974 "gen_ai.usage.output_tokens" = tracing::field::Empty,
1975 "gen_ai.usage.cost" = tracing::field::Empty,
1976 );
1977 if let Some(provider) = &self.provider_name {
1978 chat_span.record("gen_ai.provider.name", provider.as_str());
1979 }
1980 match session.model_name() {
1981 Some(model) => {
1982 chat_span.record("gen_ai.request.model", model);
1983 chat_span.record("otel.name", format!("chat {model}").as_str());
1984 }
1985 None => {
1986 chat_span.record("otel.name", "chat");
1987 }
1988 }
1989
1990 use tracing::Instrument;
1991 let mut turn = match session
1992 .begin_turn(request, cancellation.clone())
1993 .instrument(chat_span.clone())
1994 .await
1995 {
1996 Ok(turn) => turn,
1997 Err(LoopError::Cancelled) => {
1998 self.task_manager
1999 .on_turn_interrupted(&turn_id)
2000 .await
2001 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2002 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2003 }
2004 Err(error) => return Err(error),
2005 };
2006 let mut saw_tool_call = false;
2007 let mut finished_result = None;
2008
2009 while let Some(event) = match turn
2010 .next_event(cancellation.clone())
2011 .instrument(chat_span.clone())
2012 .await
2013 {
2014 Ok(event) => event,
2015 Err(LoopError::Cancelled) => {
2016 self.task_manager
2017 .on_turn_interrupted(&turn_id)
2018 .await
2019 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2020 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2021 }
2022 Err(error) => return Err(error),
2023 } {
2024 if cancellation
2025 .as_ref()
2026 .is_some_and(TurnCancellation::is_cancelled)
2027 {
2028 self.task_manager
2029 .on_turn_interrupted(&turn_id)
2030 .await
2031 .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?;
2032 return self.finish_cancelled(turn_id, interrupted_assistant_items());
2033 }
2034 match event {
2035 ModelTurnEvent::Delta(delta) => self.emit(AgentEvent::ContentDelta(delta)),
2036 ModelTurnEvent::Usage(usage) => {
2037 if let Some(tokens) = &usage.tokens {
2038 chat_span.record("gen_ai.usage.input_tokens", tokens.input_tokens);
2039 chat_span.record("gen_ai.usage.output_tokens", tokens.output_tokens);
2040 }
2041 if let Some(cost) = &usage.cost {
2042 chat_span.record("gen_ai.usage.cost", cost.amount);
2043 }
2044 self.emit(AgentEvent::UsageUpdated(usage));
2045 }
2046 ModelTurnEvent::ToolCall(call) => {
2047 saw_tool_call = true;
2048 self.emit(AgentEvent::ToolCallRequested(call.clone()));
2049 }
2050 ModelTurnEvent::Finished(result) => {
2051 finished_result = Some(result);
2052 break;
2053 }
2054 }
2055 }
2056
2057 let mut result = finished_result.ok_or_else(|| {
2058 LoopError::Provider("model turn ended without a Finished event".into())
2059 })?;
2060 if let Some(model) = &result.model {
2061 chat_span.record("gen_ai.response.model", model.as_str());
2062 }
2063 if let Some(id) = &result.response_id {
2064 chat_span.record("gen_ai.response.id", id.as_str());
2065 }
2066 if let Some(tokens) = result
2067 .usage
2068 .as_ref()
2069 .and_then(|usage| usage.tokens.as_ref())
2070 {
2071 chat_span.record("gen_ai.usage.input_tokens", tokens.input_tokens);
2072 chat_span.record("gen_ai.usage.output_tokens", tokens.output_tokens);
2073 }
2074 if let Some(cost) = result.usage.as_ref().and_then(|usage| usage.cost.as_ref()) {
2075 chat_span.record("gen_ai.usage.cost", cost.amount);
2076 }
2077 chat_span.record(
2078 "gen_ai.response.finish_reasons",
2079 tracing::field::debug(&result.finish_reason),
2080 );
2081 drop(chat_span);
2082 tracing::Span::current().record("saw_tool_call", saw_tool_call);
2083 tracing::Span::current().record(
2084 "finish_reason",
2085 tracing::field::debug(&result.finish_reason),
2086 );
2087 if let Some(tokens) = result
2088 .usage
2089 .as_ref()
2090 .and_then(|usage| usage.tokens.as_ref())
2091 {
2092 tracing::Span::current().record("gen_ai.usage.input_tokens", tokens.input_tokens);
2093 tracing::Span::current().record("gen_ai.usage.output_tokens", tokens.output_tokens);
2094 }
2095 if let Some(cost) = result.usage.as_ref().and_then(|usage| usage.cost.as_ref()) {
2096 tracing::Span::current().record("gen_ai.usage.cost", cost.amount);
2097 }
2098 let now = Timestamp::now();
2099 let usage = result.usage.clone();
2100 let finish_reason = result.finish_reason.clone();
2101 let output_items: Vec<Item> = result
2102 .output_items
2103 .drain(..)
2104 .map(|mut item| {
2105 if matches!(item.kind, ItemKind::Assistant) {
2106 if item.usage.is_none() {
2107 item.usage = usage.clone();
2108 }
2109 if item.finish_reason.is_none() {
2110 item.finish_reason = Some(finish_reason.clone());
2111 }
2112 }
2113 if item.created_at.is_none() {
2114 item.created_at = Some(now);
2115 }
2116 item
2117 })
2118 .collect();
2119 self.extend_transcript(output_items.clone());
2120
2121 if saw_tool_call {
2122 let pending_calls = extract_tool_calls(&output_items)
2123 .into_iter()
2124 .map(|call| {
2125 let tool_request = ToolRequest {
2126 call_id: call.id.clone(),
2127 tool_name: agentkit_tools_core::ToolName::new(call.name.clone()),
2128 input: call.input.clone(),
2129 session_id: self.session_id.clone(),
2130 turn_id: turn_id.clone(),
2131 metadata: call.metadata.clone(),
2132 };
2133 (call, tool_request)
2134 })
2135 .collect();
2136 self.active_tool_round = Some(ActiveToolRound {
2137 turn_id: turn_id.clone(),
2138 pending_calls,
2139 background_pending: false,
2140 foreground_progressed: false,
2141 });
2142 if let Some(step) = self.continue_active_tool_round().await? {
2143 return Ok(step);
2144 }
2145 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
2146 InputRequest {
2147 session_id: self.session_id.clone(),
2148 reason: "driver is waiting for input".into(),
2149 },
2150 )));
2151 }
2152
2153 let turn_result = TurnResult {
2154 turn_id,
2155 finish_reason: result.finish_reason,
2156 items: output_items,
2157 usage: result.usage,
2158 metadata: result.metadata,
2159 };
2160 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
2161 Ok(LoopStep::Finished(turn_result))
2162 }
2163
2164 async fn resume_after_approval(
2165 &mut self,
2166 pending: PendingApprovalToolCall,
2167 ) -> Result<LoopStep, LoopError> {
2168 let decision = pending
2169 .decision
2170 .clone()
2171 .ok_or_else(|| LoopError::InvalidState("pending approval has no decision".into()))?;
2172
2173 match decision {
2174 ApprovalDecision::Approve => {
2175 use tracing::Instrument;
2176 let dispatch_span =
2177 self.execute_tool_span(&pending.tool_request, &pending.turn_id, "approved");
2178 match self
2179 .start_task_via_manager(
2180 Some(pending.task_id.clone()),
2181 pending.tool_request.clone(),
2182 TaskLaunchKind::Approved(pending.request.clone()),
2183 self.cancellation
2184 .as_ref()
2185 .map(CancellationHandle::checkpoint),
2186 )
2187 .instrument(dispatch_span.clone())
2188 .await?
2189 {
2190 TaskStartOutcome::Ready(resolution) => {
2191 let resolution = *resolution;
2192 if let TaskResolution::Item(item) = &resolution
2193 && tool_result_is_error(item)
2194 {
2195 dispatch_span.record("error.type", "tool_error");
2196 }
2197 if let Some(step) =
2198 self.queue_resolution_interrupt(&pending.turn_id, resolution)
2199 {
2200 return Ok(step);
2201 }
2202 }
2203 TaskStartOutcome::Pending { .. } => {}
2204 }
2205 }
2206 ApprovalDecision::Deny { reason } => {
2207 self.append_tool_result_item(Item {
2208 id: None,
2209 kind: ItemKind::Tool,
2210 parts: vec![Part::ToolResult(ToolResultPart {
2211 call_id: pending.call.id.clone(),
2212 output: ToolOutput::Text(
2213 reason.unwrap_or_else(|| "approval denied".into()),
2214 ),
2215 is_error: true,
2216 metadata: pending.call.metadata.clone(),
2217 })],
2218 metadata: MetadataMap::new(),
2219 usage: None,
2220 finish_reason: None,
2221 created_at: None,
2222 });
2223 }
2224 }
2225
2226 if let Some(step) = self.continue_active_tool_round().await? {
2227 Ok(step)
2228 } else if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2229 Ok(step)
2230 } else if let Some(step) = self.next_unresolved_approval_interrupt() {
2231 Ok(step)
2232 } else {
2233 self.drive_turn(pending.turn_id, false, MutationPoint::AfterToolResult)
2234 .await
2235 }
2236 }
2237
2238 fn finish_cancelled(
2239 &mut self,
2240 turn_id: agentkit_core::TurnId,
2241 items: Vec<Item>,
2242 ) -> Result<LoopStep, LoopError> {
2243 self.extend_transcript(items.clone());
2244 let turn_result = TurnResult {
2245 turn_id,
2246 finish_reason: FinishReason::Cancelled,
2247 items,
2248 usage: None,
2249 metadata: interrupted_metadata("turn"),
2250 };
2251 self.emit(AgentEvent::TurnFinished(turn_result.clone()));
2252 Ok(LoopStep::Finished(turn_result))
2253 }
2254
2255 pub fn submit_input(&mut self, input: Vec<Item>) -> Result<(), LoopError> {
2265 if self.has_pending_interrupts() {
2266 return Err(LoopError::InvalidState(
2267 "cannot submit input while an interrupt is pending".into(),
2268 ));
2269 }
2270 self.emit(AgentEvent::InputAccepted {
2271 session_id: self.session_id.clone(),
2272 items: input.clone(),
2273 });
2274 self.pending_input.extend(input);
2275 Ok(())
2276 }
2277
2278 pub fn set_next_turn_cache(&mut self, cache: PromptCacheRequest) -> Result<(), LoopError> {
2283 if self.has_pending_interrupts() {
2284 return Err(LoopError::InvalidState(
2285 "cannot update next-turn cache while an interrupt is pending".into(),
2286 ));
2287 }
2288 self.next_turn_cache = Some(cache);
2289 Ok(())
2290 }
2291
2292 #[cfg(test)]
2293 pub(crate) fn submit_input_with_cache(
2294 &mut self,
2295 input: Vec<Item>,
2296 cache: PromptCacheRequest,
2297 ) -> Result<(), LoopError> {
2298 self.set_next_turn_cache(cache)?;
2299 self.submit_input(input)
2300 }
2301
2302 pub fn resolve_approval_for(
2312 &mut self,
2313 call_id: ToolCallId,
2314 decision: ApprovalDecision,
2315 ) -> Result<(), LoopError> {
2316 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
2317 return Err(LoopError::InvalidState(format!(
2318 "no approval request is pending for call {}",
2319 call_id.0
2320 )));
2321 };
2322 pending.decision = Some(decision.clone());
2323 self.emit(AgentEvent::ApprovalResolved {
2324 approved: matches!(decision, ApprovalDecision::Approve),
2325 });
2326 Ok(())
2327 }
2328
2329 pub fn resolve_approval_for_with_patched_input(
2342 &mut self,
2343 call_id: ToolCallId,
2344 input: serde_json::Value,
2345 ) -> Result<(), LoopError> {
2346 let Some(pending) = self.pending_approvals.get_mut(&call_id) else {
2347 return Err(LoopError::InvalidState(format!(
2348 "no approval request is pending for call {}",
2349 call_id.0
2350 )));
2351 };
2352 pending.tool_request.input = input;
2353 self.resolve_approval_for(call_id, ApprovalDecision::Approve)
2354 }
2355
2356 pub fn resolve_approval(&mut self, decision: ApprovalDecision) -> Result<(), LoopError> {
2359 let mut unresolved = self
2360 .pending_approval_order
2361 .iter()
2362 .filter(|call_id| {
2363 self.pending_approvals
2364 .get(*call_id)
2365 .is_some_and(|pending| pending.decision.is_none())
2366 })
2367 .cloned();
2368 let Some(call_id) = unresolved.next() else {
2369 return Err(LoopError::InvalidState(
2370 "no approval request is pending".into(),
2371 ));
2372 };
2373 if unresolved.next().is_some() {
2374 return Err(LoopError::InvalidState(
2375 "multiple approvals are pending; use resolve_approval_for".into(),
2376 ));
2377 }
2378 self.resolve_approval_for(call_id, decision)
2379 }
2380
2381 pub fn snapshot(&self) -> LoopSnapshot {
2383 LoopSnapshot {
2384 session_id: self.session_id.clone(),
2385 transcript: self.transcript.clone(),
2386 pending_input: self.pending_input.clone(),
2387 }
2388 }
2389
2390 pub async fn next(&mut self) -> Result<LoopStep, LoopError> {
2411 if let Some(pending) = self.take_next_resolved_approval() {
2412 return self.resume_after_approval(pending).await;
2413 }
2414
2415 if let Some(step) = self.take_next_unsurfaced_approval_interrupt() {
2416 return Ok(step);
2417 }
2418
2419 if let Some(step) = self.next_unresolved_approval_interrupt() {
2420 return Ok(step);
2421 }
2422
2423 if let Some(step) = self.continue_active_tool_round().await? {
2424 return Ok(step);
2425 }
2426
2427 let (had_loop_updates, loop_step) = self.drain_pending_loop_updates().await?;
2428 if let Some(step) = loop_step {
2429 return Ok(step);
2430 }
2431
2432 if let Some(turn_id) = self.pending_round_resume.take() {
2437 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
2438 self.extend_transcript(drained);
2439 return self
2440 .drive_turn(turn_id, false, MutationPoint::AfterToolResult)
2441 .await;
2442 }
2443
2444 if self.pending_input.is_empty() && !had_loop_updates {
2445 return Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(
2446 InputRequest {
2447 session_id: self.session_id.clone(),
2448 reason: "driver is waiting for input".into(),
2449 },
2450 )));
2451 }
2452
2453 let turn_id = agentkit_core::TurnId::new(format!("turn-{}", self.next_turn_index));
2454 self.next_turn_index += 1;
2455 let drained: Vec<Item> = std::mem::take(&mut self.pending_input);
2456 self.extend_transcript(drained);
2457 self.drive_turn(turn_id, true, MutationPoint::AfterTurnEnded)
2458 .await
2459 }
2460
2461 fn emit(&self, event: AgentEvent) {
2462 for observer in &self.observers {
2463 observer.handle_event(event.clone());
2464 }
2465 }
2466
2467 fn append_item(&mut self, mut item: Item) {
2472 if item.created_at.is_none() {
2473 item.created_at = Some(Timestamp::now());
2474 }
2475 for observer in &self.transcript_observers {
2476 observer.on_item_appended(&item);
2477 }
2478 self.transcript.push(item);
2479 }
2480
2481 fn append_tool_result_item(&mut self, item: Item) {
2494 for part in &item.parts {
2495 if let Part::ToolResult(result) = part {
2496 self.emit(AgentEvent::ToolResultReceived(result.clone()));
2497 }
2498 }
2499 let item = self.maybe_convert_detached(item);
2500 self.append_item(item);
2501 }
2502
2503 fn maybe_convert_detached(&mut self, item: Item) -> Item {
2504 if !matches!(item.kind, ItemKind::Tool) {
2505 return item;
2506 }
2507 let results: Vec<&ToolResultPart> = item
2508 .parts
2509 .iter()
2510 .filter_map(|p| match p {
2511 Part::ToolResult(r) => Some(r),
2512 _ => None,
2513 })
2514 .collect();
2515 if results.is_empty()
2516 || !results
2517 .iter()
2518 .all(|r| self.detached_call_ids.contains(&r.call_id))
2519 {
2520 return item;
2521 }
2522 let mut text = String::new();
2523 for result in &results {
2524 self.detached_call_ids.remove(&result.call_id);
2525 if !text.is_empty() {
2526 text.push_str("\n\n");
2527 }
2528 let label = if result.is_error {
2529 "failed"
2530 } else {
2531 "completed"
2532 };
2533 let body = render_tool_output_brief(&result.output);
2534 text.push_str(&format!(
2535 "Background tool call {} {}: {body}",
2536 result.call_id.0, label
2537 ));
2538 }
2539 Item::notification(text)
2540 }
2541
2542 fn extend_transcript(&mut self, items: impl IntoIterator<Item = Item>) {
2546 let now = Timestamp::now();
2547 for mut item in items {
2548 if item.created_at.is_none() {
2549 item.created_at = Some(now);
2550 }
2551 self.append_item(item);
2552 }
2553 }
2554}
2555
2556fn render_tool_output_brief(output: &ToolOutput) -> String {
2557 match output {
2558 ToolOutput::Text(t) => t.clone(),
2559 ToolOutput::Structured(value) => value.to_string(),
2560 ToolOutput::Parts(parts) => format!("[{} parts]", parts.len()),
2561 ToolOutput::Files(files) => format!("[{} files]", files.len()),
2562 }
2563}
2564
2565fn interrupted_metadata(stage: &str) -> MetadataMap {
2566 let mut metadata = MetadataMap::new();
2567 metadata.insert(INTERRUPTED_METADATA_KEY.into(), true.into());
2568 metadata.insert(
2569 INTERRUPT_REASON_METADATA_KEY.into(),
2570 USER_CANCELLED_REASON.into(),
2571 );
2572 metadata.insert(INTERRUPT_STAGE_METADATA_KEY.into(), stage.into());
2573 metadata
2574}
2575
2576fn interrupted_assistant_items() -> Vec<Item> {
2577 vec![Item {
2578 id: None,
2579 kind: ItemKind::Assistant,
2580 parts: vec![Part::Text(TextPart {
2581 text: "Previous assistant response was interrupted by the user before completion."
2582 .into(),
2583 metadata: interrupted_metadata("assistant"),
2584 })],
2585 metadata: interrupted_metadata("assistant"),
2586 usage: None,
2587 finish_reason: None,
2588 created_at: None,
2589 }]
2590}
2591
2592fn transcript_has_pending_input(transcript: &[Item]) -> bool {
2598 matches!(
2599 transcript.last().map(|item| item.kind),
2600 Some(ItemKind::User | ItemKind::Tool | ItemKind::Notification)
2601 )
2602}
2603
2604fn extract_tool_calls(items: &[Item]) -> Vec<ToolCallPart> {
2605 let mut calls = Vec::new();
2606 for item in items {
2607 for part in &item.parts {
2608 if let Part::ToolCall(call) = part {
2609 calls.push(call.clone());
2610 }
2611 }
2612 }
2613 calls
2614}
2615
2616fn tool_result_is_error(item: &Item) -> bool {
2617 item.parts
2618 .iter()
2619 .any(|part| matches!(part, Part::ToolResult(result) if result.is_error))
2620}
2621
2622#[derive(Debug, Error)]
2624pub enum LoopError {
2625 #[error("invalid driver state: {0}")]
2627 InvalidState(String),
2628 #[error("turn cancelled")]
2630 Cancelled,
2631 #[error("provider error: {0}")]
2633 Provider(String),
2634 #[error("tool error: {0}")]
2636 Tool(#[from] ToolError),
2637 #[error("mutator error: {0}")]
2639 Mutator(String),
2640 #[error("unsupported operation: {0}")]
2642 Unsupported(String),
2643}
2644
2645struct DriverEmitter<'a> {
2650 observers: &'a [Arc<dyn LoopObserver>],
2651}
2652
2653impl<'a> EventEmitter for DriverEmitter<'a> {
2654 fn emit(&self, event: AgentEvent) {
2655 for observer in self.observers {
2656 observer.handle_event(event.clone());
2657 }
2658 }
2659}
2660
2661fn validate_transcript_invariants(transcript: &[Item]) -> Result<(), LoopError> {
2666 let mut pending: HashSet<ToolCallId> = HashSet::new();
2667 let mut seen_calls: HashSet<ToolCallId> = HashSet::new();
2668 let mut seen_results: HashSet<ToolCallId> = HashSet::new();
2669 for item in transcript {
2670 for part in &item.parts {
2671 match part {
2672 Part::ToolCall(call) => {
2673 if !seen_calls.insert(call.id.clone()) {
2674 return Err(LoopError::Mutator(format!(
2675 "transcript invariant violation: duplicate tool_use: {}",
2676 call.id.0
2677 )));
2678 }
2679 pending.insert(call.id.clone());
2680 }
2681 Part::ToolResult(result) => {
2682 if !pending.remove(&result.call_id) {
2683 let kind = if seen_results.contains(&result.call_id) {
2684 "duplicate"
2685 } else {
2686 "orphaned"
2687 };
2688 return Err(LoopError::Mutator(format!(
2689 "transcript invariant violation: {kind} tool_result: {}",
2690 result.call_id.0
2691 )));
2692 }
2693 seen_results.insert(result.call_id.clone());
2694 }
2695 _ => {}
2696 }
2697 }
2698 }
2699 if !pending.is_empty() {
2700 let missing: Vec<String> = pending.into_iter().map(|id| id.0).collect();
2701 return Err(LoopError::Mutator(format!(
2702 "transcript invariant violation: tool_use(s) without matching tool_result: {}",
2703 missing.join(", ")
2704 )));
2705 }
2706 Ok(())
2707}
2708
2709#[cfg(test)]
2710mod tests {
2711 use std::collections::VecDeque;
2712 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
2713 use std::sync::{Arc as StdArc, Mutex as StdMutex};
2714
2715 use agentkit_core::{
2716 CancellationController, ItemKind, Part, TextPart, ToolCallId, ToolCallPart, ToolOutput,
2717 ToolResultPart,
2718 };
2719 use agentkit_task_manager::{
2720 AsyncTaskManager, RoutingDecision, TaskEvent, TaskManager, TaskManagerHandle,
2721 TaskRoutingPolicy,
2722 };
2723 use agentkit_tools_core::{
2724 FileSystemPermissionRequest, PermissionCode, PermissionDecision, PermissionDenial, Tool,
2725 ToolAnnotations, ToolCatalogEvent, ToolExecutionOutcome, ToolName, ToolRegistry,
2726 ToolResult, ToolSpec,
2727 };
2728 use serde_json::{Value, json};
2729 use tokio::sync::Notify;
2730 use tokio::time::{Duration, timeout};
2731
2732 use super::*;
2733
2734 struct FakeAdapter;
2735 struct SlowAdapter;
2736 struct RecordingAdapter {
2737 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
2738 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
2739 }
2740 struct MultiToolAdapter;
2741 struct DualApprovalAdapter;
2742
2743 struct FakeSession;
2744 struct SlowSession;
2745 struct RecordingSession {
2746 seen_descriptions: StdArc<StdMutex<Vec<Vec<String>>>>,
2747 seen_caches: StdArc<StdMutex<Vec<Option<PromptCacheRequest>>>>,
2748 }
2749 struct MultiToolSession;
2750 struct DualApprovalSession;
2751
2752 struct FakeTurn {
2753 events: VecDeque<ModelTurnEvent>,
2754 }
2755
2756 struct SlowTurn {
2757 emitted: bool,
2758 }
2759
2760 struct RecordingTurn {
2761 emitted: bool,
2762 }
2763 struct MultiToolTurn {
2764 events: VecDeque<ModelTurnEvent>,
2765 }
2766 struct DualApprovalTurn {
2767 events: VecDeque<ModelTurnEvent>,
2768 }
2769
2770 #[async_trait]
2771 impl ModelAdapter for FakeAdapter {
2772 type Session = FakeSession;
2773
2774 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
2775 Ok(FakeSession)
2776 }
2777 }
2778
2779 #[async_trait]
2780 impl ModelAdapter for SlowAdapter {
2781 type Session = SlowSession;
2782
2783 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
2784 Ok(SlowSession)
2785 }
2786 }
2787
2788 #[async_trait]
2789 impl ModelAdapter for RecordingAdapter {
2790 type Session = RecordingSession;
2791
2792 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
2793 Ok(RecordingSession {
2794 seen_descriptions: self.seen_descriptions.clone(),
2795 seen_caches: self.seen_caches.clone(),
2796 })
2797 }
2798 }
2799
2800 #[async_trait]
2801 impl ModelAdapter for MultiToolAdapter {
2802 type Session = MultiToolSession;
2803
2804 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
2805 Ok(MultiToolSession)
2806 }
2807 }
2808
2809 #[async_trait]
2810 impl ModelAdapter for DualApprovalAdapter {
2811 type Session = DualApprovalSession;
2812
2813 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
2814 Ok(DualApprovalSession)
2815 }
2816 }
2817
2818 #[async_trait]
2819 impl ModelSession for FakeSession {
2820 type Turn = FakeTurn;
2821
2822 async fn begin_turn(
2823 &mut self,
2824 request: TurnRequest,
2825 _cancellation: Option<TurnCancellation>,
2826 ) -> Result<Self::Turn, LoopError> {
2827 let has_tool_result = request.transcript.iter().any(|item| {
2828 item.kind == ItemKind::Tool
2829 && item
2830 .parts
2831 .iter()
2832 .any(|part| matches!(part, Part::ToolResult(_)))
2833 });
2834 let tool_name = request
2835 .available_tools
2836 .first()
2837 .map(|tool| tool.name.0.clone())
2838 .unwrap_or_else(|| "echo".into());
2839
2840 let events = if has_tool_result {
2841 let result_text = request
2842 .transcript
2843 .iter()
2844 .rev()
2845 .find_map(|item| {
2846 item.parts.iter().find_map(|part| match part {
2847 Part::ToolResult(ToolResultPart {
2848 output: ToolOutput::Text(text),
2849 ..
2850 }) => Some(text.clone()),
2851 _ => None,
2852 })
2853 })
2854 .unwrap_or_else(|| "missing".into());
2855
2856 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
2857 model: None,
2858 response_id: None,
2859 finish_reason: FinishReason::Completed,
2860 output_items: vec![Item {
2861 id: None,
2862 kind: ItemKind::Assistant,
2863 parts: vec![Part::Text(TextPart {
2864 text: format!("tool said: {result_text}"),
2865 metadata: MetadataMap::new(),
2866 })],
2867 metadata: MetadataMap::new(),
2868 usage: None,
2869 finish_reason: None,
2870 created_at: None,
2871 }],
2872 usage: None,
2873 metadata: MetadataMap::new(),
2874 })])
2875 } else {
2876 VecDeque::from([
2877 ModelTurnEvent::ToolCall(agentkit_core::ToolCallPart {
2878 id: ToolCallId::new("call-1"),
2879 name: tool_name.clone(),
2880 input: json!({ "value": "pong" }),
2881 metadata: MetadataMap::new(),
2882 }),
2883 ModelTurnEvent::Finished(ModelTurnResult {
2884 model: None,
2885 response_id: None,
2886 finish_reason: FinishReason::ToolCall,
2887 output_items: vec![Item {
2888 id: None,
2889 kind: ItemKind::Assistant,
2890 parts: vec![Part::ToolCall(agentkit_core::ToolCallPart {
2891 id: ToolCallId::new("call-1"),
2892 name: tool_name,
2893 input: json!({ "value": "pong" }),
2894 metadata: MetadataMap::new(),
2895 })],
2896 metadata: MetadataMap::new(),
2897 usage: None,
2898 finish_reason: None,
2899 created_at: None,
2900 }],
2901 usage: None,
2902 metadata: MetadataMap::new(),
2903 }),
2904 ])
2905 };
2906
2907 Ok(FakeTurn { events })
2908 }
2909 }
2910
2911 #[async_trait]
2912 impl ModelSession for SlowSession {
2913 type Turn = SlowTurn;
2914
2915 async fn begin_turn(
2916 &mut self,
2917 request: TurnRequest,
2918 cancellation: Option<TurnCancellation>,
2919 ) -> Result<Self::Turn, LoopError> {
2920 let should_block = request
2921 .transcript
2922 .iter()
2923 .rev()
2924 .find(|item| item.kind == ItemKind::User)
2925 .is_some_and(|item| {
2926 item.parts.iter().any(|part| match part {
2927 Part::Text(text) => text.text == "do the long task",
2928 _ => false,
2929 })
2930 });
2931
2932 if should_block && let Some(cancellation) = cancellation {
2933 cancellation.cancelled().await;
2934 return Err(LoopError::Cancelled);
2935 }
2936
2937 Ok(SlowTurn { emitted: false })
2938 }
2939 }
2940
2941 #[async_trait]
2942 impl ModelSession for RecordingSession {
2943 type Turn = RecordingTurn;
2944
2945 async fn begin_turn(
2946 &mut self,
2947 request: TurnRequest,
2948 _cancellation: Option<TurnCancellation>,
2949 ) -> Result<Self::Turn, LoopError> {
2950 let descriptions = request
2951 .available_tools
2952 .iter()
2953 .map(|tool| tool.description.clone())
2954 .collect::<Vec<_>>();
2955 self.seen_descriptions.lock().unwrap().push(descriptions);
2956 self.seen_caches.lock().unwrap().push(request.cache.clone());
2957
2958 Ok(RecordingTurn { emitted: false })
2959 }
2960 }
2961
2962 #[async_trait]
2963 impl ModelSession for MultiToolSession {
2964 type Turn = MultiToolTurn;
2965
2966 async fn begin_turn(
2967 &mut self,
2968 request: TurnRequest,
2969 _cancellation: Option<TurnCancellation>,
2970 ) -> Result<Self::Turn, LoopError> {
2971 let has_tool_result = request.transcript.iter().any(|item| {
2972 item.kind == ItemKind::Tool
2973 && item
2974 .parts
2975 .iter()
2976 .any(|part| matches!(part, Part::ToolResult(_)))
2977 });
2978
2979 let events = if has_tool_result {
2980 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
2981 model: None,
2982 response_id: None,
2983 finish_reason: FinishReason::Completed,
2984 output_items: vec![Item {
2985 id: None,
2986 kind: ItemKind::Assistant,
2987 parts: vec![Part::Text(TextPart {
2988 text: "mixed tools finished".into(),
2989 metadata: MetadataMap::new(),
2990 })],
2991 metadata: MetadataMap::new(),
2992 usage: None,
2993 finish_reason: None,
2994 created_at: None,
2995 }],
2996 usage: None,
2997 metadata: MetadataMap::new(),
2998 })])
2999 } else {
3000 let foreground = agentkit_core::ToolCallPart {
3001 id: ToolCallId::new("call-foreground"),
3002 name: "foreground-wait".into(),
3003 input: json!({}),
3004 metadata: MetadataMap::new(),
3005 };
3006 let background = agentkit_core::ToolCallPart {
3007 id: ToolCallId::new("call-background"),
3008 name: "background-wait".into(),
3009 input: json!({}),
3010 metadata: MetadataMap::new(),
3011 };
3012 VecDeque::from([
3013 ModelTurnEvent::ToolCall(foreground.clone()),
3014 ModelTurnEvent::ToolCall(background.clone()),
3015 ModelTurnEvent::Finished(ModelTurnResult {
3016 model: None,
3017 response_id: None,
3018 finish_reason: FinishReason::ToolCall,
3019 output_items: vec![Item {
3020 id: None,
3021 kind: ItemKind::Assistant,
3022 parts: vec![Part::ToolCall(foreground), Part::ToolCall(background)],
3023 metadata: MetadataMap::new(),
3024 usage: None,
3025 finish_reason: None,
3026 created_at: None,
3027 }],
3028 usage: None,
3029 metadata: MetadataMap::new(),
3030 }),
3031 ])
3032 };
3033
3034 Ok(MultiToolTurn { events })
3035 }
3036 }
3037
3038 #[async_trait]
3039 impl ModelSession for DualApprovalSession {
3040 type Turn = DualApprovalTurn;
3041
3042 async fn begin_turn(
3043 &mut self,
3044 request: TurnRequest,
3045 _cancellation: Option<TurnCancellation>,
3046 ) -> Result<Self::Turn, LoopError> {
3047 let tool_results = request
3048 .transcript
3049 .iter()
3050 .flat_map(|item| item.parts.iter())
3051 .filter(|part| matches!(part, Part::ToolResult(_)))
3052 .count();
3053
3054 let events = if tool_results >= 2 {
3055 VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
3056 model: None,
3057 response_id: None,
3058 finish_reason: FinishReason::Completed,
3059 output_items: vec![Item {
3060 id: None,
3061 kind: ItemKind::Assistant,
3062 parts: vec![Part::Text(TextPart {
3063 text: "both approvals finished".into(),
3064 metadata: MetadataMap::new(),
3065 })],
3066 metadata: MetadataMap::new(),
3067 usage: None,
3068 finish_reason: None,
3069 created_at: None,
3070 }],
3071 usage: None,
3072 metadata: MetadataMap::new(),
3073 })])
3074 } else {
3075 let first = agentkit_core::ToolCallPart {
3076 id: ToolCallId::new("call-1"),
3077 name: "echo".into(),
3078 input: json!({ "value": "first" }),
3079 metadata: MetadataMap::new(),
3080 };
3081 let second = agentkit_core::ToolCallPart {
3082 id: ToolCallId::new("call-2"),
3083 name: "echo".into(),
3084 input: json!({ "value": "second" }),
3085 metadata: MetadataMap::new(),
3086 };
3087 VecDeque::from([
3088 ModelTurnEvent::ToolCall(first.clone()),
3089 ModelTurnEvent::ToolCall(second.clone()),
3090 ModelTurnEvent::Finished(ModelTurnResult {
3091 model: None,
3092 response_id: None,
3093 finish_reason: FinishReason::ToolCall,
3094 output_items: vec![Item {
3095 id: None,
3096 kind: ItemKind::Assistant,
3097 parts: vec![Part::ToolCall(first), Part::ToolCall(second)],
3098 metadata: MetadataMap::new(),
3099 usage: None,
3100 finish_reason: None,
3101 created_at: None,
3102 }],
3103 usage: None,
3104 metadata: MetadataMap::new(),
3105 }),
3106 ])
3107 };
3108
3109 Ok(DualApprovalTurn { events })
3110 }
3111 }
3112
3113 #[async_trait]
3114 impl ModelTurn for FakeTurn {
3115 async fn next_event(
3116 &mut self,
3117 _cancellation: Option<TurnCancellation>,
3118 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3119 Ok(self.events.pop_front())
3120 }
3121 }
3122
3123 #[async_trait]
3124 impl ModelTurn for SlowTurn {
3125 async fn next_event(
3126 &mut self,
3127 cancellation: Option<TurnCancellation>,
3128 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3129 if let Some(cancellation) = cancellation
3130 && cancellation.is_cancelled()
3131 {
3132 return Err(LoopError::Cancelled);
3133 }
3134
3135 if self.emitted {
3136 Ok(None)
3137 } else {
3138 self.emitted = true;
3139 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
3140 model: None,
3141 response_id: None,
3142 finish_reason: FinishReason::Completed,
3143 output_items: vec![Item {
3144 id: None,
3145 kind: ItemKind::Assistant,
3146 parts: vec![Part::Text(TextPart {
3147 text: "done".into(),
3148 metadata: MetadataMap::new(),
3149 })],
3150 metadata: MetadataMap::new(),
3151 usage: None,
3152 finish_reason: None,
3153 created_at: None,
3154 }],
3155 usage: None,
3156 metadata: MetadataMap::new(),
3157 })))
3158 }
3159 }
3160 }
3161
3162 #[async_trait]
3163 impl ModelTurn for RecordingTurn {
3164 async fn next_event(
3165 &mut self,
3166 _cancellation: Option<TurnCancellation>,
3167 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3168 if self.emitted {
3169 Ok(None)
3170 } else {
3171 self.emitted = true;
3172 Ok(Some(ModelTurnEvent::Finished(ModelTurnResult {
3173 model: None,
3174 response_id: None,
3175 finish_reason: FinishReason::Completed,
3176 output_items: vec![Item {
3177 id: None,
3178 kind: ItemKind::Assistant,
3179 parts: vec![Part::Text(TextPart {
3180 text: "done".into(),
3181 metadata: MetadataMap::new(),
3182 })],
3183 metadata: MetadataMap::new(),
3184 usage: None,
3185 finish_reason: None,
3186 created_at: None,
3187 }],
3188 usage: None,
3189 metadata: MetadataMap::new(),
3190 })))
3191 }
3192 }
3193 }
3194
3195 #[async_trait]
3196 impl ModelTurn for MultiToolTurn {
3197 async fn next_event(
3198 &mut self,
3199 _cancellation: Option<TurnCancellation>,
3200 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3201 Ok(self.events.pop_front())
3202 }
3203 }
3204
3205 #[async_trait]
3206 impl ModelTurn for DualApprovalTurn {
3207 async fn next_event(
3208 &mut self,
3209 _cancellation: Option<TurnCancellation>,
3210 ) -> Result<Option<ModelTurnEvent>, LoopError> {
3211 Ok(self.events.pop_front())
3212 }
3213 }
3214
3215 #[derive(Clone)]
3216 struct EchoTool {
3217 spec: ToolSpec,
3218 }
3219
3220 impl Default for EchoTool {
3221 fn default() -> Self {
3222 Self {
3223 spec: ToolSpec {
3224 name: ToolName::new("echo"),
3225 description: "Echo back a value".into(),
3226 input_schema: json!({
3227 "type": "object",
3228 "properties": {
3229 "value": { "type": "string" }
3230 },
3231 "required": ["value"],
3232 "additionalProperties": false
3233 }),
3234 output_schema: None,
3235 annotations: ToolAnnotations::default(),
3236 metadata: MetadataMap::new(),
3237 },
3238 }
3239 }
3240 }
3241
3242 #[derive(Clone)]
3243 struct DynamicSpecTool {
3244 spec: ToolSpec,
3245 version: StdArc<AtomicUsize>,
3246 }
3247
3248 impl DynamicSpecTool {
3249 fn new(version: StdArc<AtomicUsize>) -> Self {
3250 Self {
3251 spec: ToolSpec {
3252 name: ToolName::new("dynamic"),
3253 description: "dynamic version 0".into(),
3254 input_schema: json!({
3255 "type": "object",
3256 "properties": {},
3257 "additionalProperties": false
3258 }),
3259 output_schema: None,
3260 annotations: ToolAnnotations::default(),
3261 metadata: MetadataMap::new(),
3262 },
3263 version,
3264 }
3265 }
3266 }
3267
3268 #[async_trait]
3269 impl Tool for EchoTool {
3270 fn spec(&self) -> &ToolSpec {
3271 &self.spec
3272 }
3273
3274 fn proposed_requests(
3275 &self,
3276 request: &agentkit_tools_core::ToolRequest,
3277 ) -> Result<
3278 Vec<Box<dyn agentkit_tools_core::PermissionRequest>>,
3279 agentkit_tools_core::ToolError,
3280 > {
3281 Ok(vec![Box::new(FileSystemPermissionRequest::Read {
3282 path: "/tmp/echo".into(),
3283 metadata: request.metadata.clone(),
3284 })])
3285 }
3286
3287 async fn invoke(
3288 &self,
3289 request: agentkit_tools_core::ToolRequest,
3290 _ctx: &mut ToolContext<'_>,
3291 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
3292 let value = request
3293 .input
3294 .get("value")
3295 .and_then(Value::as_str)
3296 .ok_or_else(|| {
3297 agentkit_tools_core::ToolError::InvalidInput("missing value".into())
3298 })?;
3299
3300 Ok(ToolResult {
3301 result: ToolResultPart {
3302 call_id: request.call_id,
3303 output: ToolOutput::Text(value.into()),
3304 is_error: false,
3305 metadata: MetadataMap::new(),
3306 },
3307 duration: None,
3308 metadata: MetadataMap::new(),
3309 })
3310 }
3311 }
3312
3313 #[async_trait]
3314 impl Tool for DynamicSpecTool {
3315 fn spec(&self) -> &ToolSpec {
3316 &self.spec
3317 }
3318
3319 fn current_spec(&self) -> Option<ToolSpec> {
3320 let mut spec = self.spec.clone();
3321 spec.description = format!("dynamic version {}", self.version.load(Ordering::SeqCst));
3322 Some(spec)
3323 }
3324
3325 async fn invoke(
3326 &self,
3327 request: agentkit_tools_core::ToolRequest,
3328 _ctx: &mut ToolContext<'_>,
3329 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
3330 Ok(ToolResult {
3331 result: ToolResultPart {
3332 call_id: request.call_id,
3333 output: ToolOutput::Text("ok".into()),
3334 is_error: false,
3335 metadata: MetadataMap::new(),
3336 },
3337 duration: None,
3338 metadata: MetadataMap::new(),
3339 })
3340 }
3341 }
3342
3343 struct DenyFsReads;
3344
3345 impl PermissionChecker for DenyFsReads {
3346 fn evaluate(
3347 &self,
3348 request: &dyn agentkit_tools_core::PermissionRequest,
3349 ) -> PermissionDecision {
3350 if request.kind() == "filesystem.read" {
3351 return PermissionDecision::Deny(PermissionDenial {
3352 code: PermissionCode::PathNotAllowed,
3353 message: "reads denied in test".into(),
3354 metadata: MetadataMap::new(),
3355 });
3356 }
3357
3358 PermissionDecision::Allow
3359 }
3360 }
3361
3362 struct ApproveFsReads;
3363
3364 impl PermissionChecker for ApproveFsReads {
3365 fn evaluate(
3366 &self,
3367 request: &dyn agentkit_tools_core::PermissionRequest,
3368 ) -> PermissionDecision {
3369 if request.kind() == "filesystem.read" {
3370 return PermissionDecision::RequireApproval(ApprovalRequest {
3371 task_id: None,
3372 call_id: None,
3373 id: "approval:fs-read".into(),
3374 request_kind: request.kind().into(),
3375 reason: agentkit_tools_core::ApprovalReason::SensitivePath,
3376 summary: request.summary(),
3377 metadata: request.metadata().clone(),
3378 });
3379 }
3380
3381 PermissionDecision::Allow
3382 }
3383 }
3384
3385 struct KeepRecentMutator {
3386 keep: usize,
3387 }
3388
3389 #[async_trait]
3390 impl LoopMutator for KeepRecentMutator {
3391 async fn mutate(
3392 &self,
3393 cursor: &mut TranscriptCursor<'_>,
3394 ctx: LoopCtx<'_>,
3395 ) -> Result<(), LoopError> {
3396 if cursor.len() < 2 {
3397 return Ok(());
3398 }
3399 let drop = cursor.len().saturating_sub(self.keep);
3400 ctx.emitter.emit(AgentEvent::MutationStarted {
3401 session_id: ctx.session_id.clone(),
3402 turn_id: ctx.turn_id.cloned(),
3403 mutator: "keep-recent".into(),
3404 point: ctx.point,
3405 });
3406 cursor.drain(..drop);
3407 ctx.emitter.emit(AgentEvent::MutationFinished {
3408 session_id: ctx.session_id.clone(),
3409 turn_id: ctx.turn_id.cloned(),
3410 mutator: "keep-recent".into(),
3411 dirty: true,
3412 metadata: MetadataMap::new(),
3413 });
3414 Ok(())
3415 }
3416 }
3417
3418 struct PointRecordingMutator {
3421 points: StdArc<StdMutex<Vec<MutationPoint>>>,
3422 }
3423
3424 #[async_trait]
3425 impl LoopMutator for PointRecordingMutator {
3426 async fn mutate(
3427 &self,
3428 _cursor: &mut TranscriptCursor<'_>,
3429 ctx: LoopCtx<'_>,
3430 ) -> Result<(), LoopError> {
3431 self.points.lock().unwrap().push(ctx.point);
3432 Ok(())
3433 }
3434 }
3435
3436 struct RecordingObserver {
3437 events: StdArc<StdMutex<Vec<AgentEvent>>>,
3438 }
3439
3440 impl LoopObserver for RecordingObserver {
3441 fn handle_event(&self, event: AgentEvent) {
3442 self.events.lock().unwrap().push(event);
3443 }
3444 }
3445
3446 struct CatalogExecutor {
3447 version: AtomicUsize,
3448 events: StdMutex<Vec<ToolCatalogEvent>>,
3449 }
3450
3451 impl CatalogExecutor {
3452 fn new() -> Self {
3453 Self {
3454 version: AtomicUsize::new(0),
3455 events: StdMutex::new(Vec::new()),
3456 }
3457 }
3458
3459 fn publish_change(&self, version: usize, event: ToolCatalogEvent) {
3460 self.version.store(version, Ordering::SeqCst);
3461 self.events.lock().unwrap().push(event);
3462 }
3463 }
3464
3465 #[async_trait]
3466 impl ToolExecutor for CatalogExecutor {
3467 fn specs(&self) -> Vec<ToolSpec> {
3468 vec![ToolSpec {
3469 name: ToolName::new("dynamic"),
3470 description: format!("dynamic version {}", self.version.load(Ordering::SeqCst)),
3471 input_schema: json!({
3472 "type": "object",
3473 "properties": {},
3474 "additionalProperties": false
3475 }),
3476 output_schema: None,
3477 annotations: ToolAnnotations::default(),
3478 metadata: MetadataMap::new(),
3479 }]
3480 }
3481
3482 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3483 std::mem::take(&mut *self.events.lock().unwrap())
3484 }
3485
3486 async fn execute(
3487 &self,
3488 request: ToolRequest,
3489 _ctx: &mut ToolContext<'_>,
3490 ) -> ToolExecutionOutcome {
3491 ToolExecutionOutcome::Completed(ToolResult {
3492 result: ToolResultPart {
3493 call_id: request.call_id,
3494 output: ToolOutput::Text("dynamic-ok".into()),
3495 is_error: false,
3496 metadata: MetadataMap::new(),
3497 },
3498 duration: None,
3499 metadata: MetadataMap::new(),
3500 })
3501 }
3502 }
3503
3504 #[derive(Clone)]
3505 struct BlockingTool {
3506 spec: ToolSpec,
3507 entered: StdArc<AtomicBool>,
3508 release: StdArc<Notify>,
3509 output: &'static str,
3510 }
3511
3512 impl BlockingTool {
3513 fn new(
3514 name: &str,
3515 entered: StdArc<AtomicBool>,
3516 release: StdArc<Notify>,
3517 output: &'static str,
3518 ) -> Self {
3519 Self {
3520 spec: ToolSpec {
3521 name: ToolName::new(name),
3522 description: format!("blocking tool {name}"),
3523 input_schema: json!({
3524 "type": "object",
3525 "properties": {},
3526 "additionalProperties": false
3527 }),
3528 output_schema: None,
3529 annotations: ToolAnnotations::default(),
3530 metadata: MetadataMap::new(),
3531 },
3532 entered,
3533 release,
3534 output,
3535 }
3536 }
3537 }
3538
3539 #[async_trait]
3540 impl Tool for BlockingTool {
3541 fn spec(&self) -> &ToolSpec {
3542 &self.spec
3543 }
3544
3545 async fn invoke(
3546 &self,
3547 request: agentkit_tools_core::ToolRequest,
3548 _ctx: &mut ToolContext<'_>,
3549 ) -> Result<ToolResult, agentkit_tools_core::ToolError> {
3550 self.entered.store(true, Ordering::SeqCst);
3551 self.release.notified().await;
3552 Ok(ToolResult {
3553 result: ToolResultPart {
3554 call_id: request.call_id,
3555 output: ToolOutput::Text(self.output.into()),
3556 is_error: false,
3557 metadata: MetadataMap::new(),
3558 },
3559 duration: None,
3560 metadata: MetadataMap::new(),
3561 })
3562 }
3563 }
3564
3565 struct NameRoutingPolicy {
3566 routes: Vec<(String, RoutingDecision)>,
3567 }
3568
3569 impl NameRoutingPolicy {
3570 fn new(routes: impl IntoIterator<Item = (impl Into<String>, RoutingDecision)>) -> Self {
3571 Self {
3572 routes: routes
3573 .into_iter()
3574 .map(|(name, decision)| (name.into(), decision))
3575 .collect(),
3576 }
3577 }
3578 }
3579
3580 impl TaskRoutingPolicy for NameRoutingPolicy {
3581 fn route(&self, request: &ToolRequest) -> RoutingDecision {
3582 self.routes
3583 .iter()
3584 .find(|(name, _)| name == &request.tool_name.0)
3585 .map(|(_, decision)| *decision)
3586 .unwrap_or(RoutingDecision::Foreground)
3587 }
3588 }
3589
3590 async fn wait_for_task_event(handle: &TaskManagerHandle) -> TaskEvent {
3591 timeout(Duration::from_secs(1), handle.next_event())
3592 .await
3593 .expect("timed out waiting for task event")
3594 .expect("task event stream ended unexpectedly")
3595 }
3596
3597 async fn wait_until_entered(flag: &AtomicBool) {
3598 timeout(Duration::from_secs(1), async {
3599 while !flag.load(Ordering::SeqCst) {
3600 tokio::task::yield_now().await;
3601 }
3602 })
3603 .await
3604 .expect("task never entered execution");
3605 }
3606
3607 #[tokio::test]
3608 async fn loop_continues_after_completed_tool_call() {
3609 let tools = ToolRegistry::new().with(EchoTool::default());
3610 let agent = Agent::builder()
3611 .model(FakeAdapter)
3612 .add_tool_source(tools)
3613 .permissions(AllowAllPermissions)
3614 .build()
3615 .unwrap();
3616
3617 let mut driver = agent
3618 .start(SessionConfig {
3619 session_id: SessionId::new("session-1"),
3620 metadata: MetadataMap::new(),
3621 cache: None,
3622 })
3623 .await
3624 .unwrap();
3625
3626 driver
3627 .submit_input(vec![Item {
3628 id: None,
3629 kind: ItemKind::User,
3630 parts: vec![Part::Text(TextPart {
3631 text: "ping".into(),
3632 metadata: MetadataMap::new(),
3633 })],
3634 metadata: MetadataMap::new(),
3635 usage: None,
3636 finish_reason: None,
3637 created_at: None,
3638 }])
3639 .unwrap();
3640
3641 let result = run_until_finished(&mut driver).await;
3642
3643 match result {
3644 LoopStep::Finished(turn) => {
3645 assert_eq!(turn.finish_reason, FinishReason::Completed);
3646 assert_eq!(turn.items.len(), 1);
3647 match &turn.items[0].parts[0] {
3648 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
3649 other => panic!("unexpected part: {other:?}"),
3650 }
3651 }
3652 other => panic!("unexpected loop step: {other:?}"),
3653 }
3654 }
3655
3656 async fn run_until_finished<S: ModelSession + Send>(driver: &mut LoopDriver<S>) -> LoopStep {
3660 loop {
3661 match driver.next().await.unwrap() {
3662 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => continue,
3663 step => return step,
3664 }
3665 }
3666 }
3667
3668 #[tokio::test]
3675 async fn post_tool_continuation_reports_after_tool_result_mutation_point() {
3676 let points = StdArc::new(StdMutex::new(Vec::<MutationPoint>::new()));
3677 let tools = ToolRegistry::new().with(EchoTool::default());
3678 let agent = Agent::builder()
3679 .model(FakeAdapter)
3680 .add_tool_source(tools)
3681 .permissions(AllowAllPermissions)
3682 .mutator(PointRecordingMutator {
3683 points: points.clone(),
3684 })
3685 .build()
3686 .unwrap();
3687
3688 let mut driver = agent
3689 .start(SessionConfig {
3690 session_id: SessionId::new("session-mutation-point"),
3691 metadata: MetadataMap::new(),
3692 cache: None,
3693 })
3694 .await
3695 .unwrap();
3696
3697 driver
3698 .submit_input(vec![Item::text(ItemKind::User, "ping")])
3699 .unwrap();
3700
3701 let _ = run_until_finished(&mut driver).await;
3703
3704 let recorded = points.lock().unwrap().clone();
3705 assert_eq!(
3706 recorded.first(),
3707 Some(&MutationPoint::AfterTurnEnded),
3708 "first drive of a fresh turn must report AfterTurnEnded, got {recorded:?}"
3709 );
3710 assert!(
3711 recorded.contains(&MutationPoint::AfterToolResult),
3712 "post-tool continuation must report AfterToolResult, got {recorded:?}"
3713 );
3714 }
3715
3716 #[test]
3717 fn pending_input_requires_input_bearing_tail_role() {
3718 assert!(!transcript_has_pending_input(&[]));
3719 assert!(!transcript_has_pending_input(&[Item::text(
3720 ItemKind::System,
3721 "system"
3722 )]));
3723 assert!(!transcript_has_pending_input(&[Item::text(
3724 ItemKind::Developer,
3725 "developer"
3726 )]));
3727 assert!(!transcript_has_pending_input(&[Item::text(
3728 ItemKind::Context,
3729 "context"
3730 )]));
3731 assert!(!transcript_has_pending_input(&[Item::text(
3732 ItemKind::Assistant,
3733 "assistant"
3734 )]));
3735
3736 assert!(transcript_has_pending_input(&[Item::text(
3737 ItemKind::User,
3738 "user"
3739 )]));
3740 assert!(transcript_has_pending_input(&[Item::notification(
3741 "background update"
3742 )]));
3743 assert!(transcript_has_pending_input(&[Item {
3744 id: None,
3745 kind: ItemKind::Tool,
3746 parts: vec![Part::ToolResult(ToolResultPart {
3747 call_id: ToolCallId::new("call-test"),
3748 output: ToolOutput::Text("ok".into()),
3749 is_error: false,
3750 metadata: MetadataMap::new(),
3751 })],
3752 metadata: MetadataMap::new(),
3753 usage: None,
3754 finish_reason: None,
3755 created_at: None,
3756 }]));
3757 }
3758
3759 struct DropTrailingUserMutator;
3765
3766 #[async_trait]
3767 impl LoopMutator for DropTrailingUserMutator {
3768 async fn mutate(
3769 &self,
3770 cursor: &mut TranscriptCursor<'_>,
3771 _ctx: LoopCtx<'_>,
3772 ) -> Result<(), LoopError> {
3773 if cursor.last().map(|item| item.kind) == Some(ItemKind::User) {
3774 cursor.pop();
3775 }
3776 Ok(())
3777 }
3778 }
3779
3780 struct RejectAssistantPrefillAdapter {
3785 saw_assistant_tail: StdArc<AtomicBool>,
3786 }
3787
3788 struct RejectAssistantPrefillSession {
3789 saw_assistant_tail: StdArc<AtomicBool>,
3790 }
3791
3792 #[async_trait]
3793 impl ModelAdapter for RejectAssistantPrefillAdapter {
3794 type Session = RejectAssistantPrefillSession;
3795
3796 async fn start_session(&self, _config: SessionConfig) -> Result<Self::Session, LoopError> {
3797 Ok(RejectAssistantPrefillSession {
3798 saw_assistant_tail: self.saw_assistant_tail.clone(),
3799 })
3800 }
3801 }
3802
3803 #[async_trait]
3804 impl ModelSession for RejectAssistantPrefillSession {
3805 type Turn = FakeTurn;
3806
3807 async fn begin_turn(
3808 &mut self,
3809 request: TurnRequest,
3810 _cancellation: Option<TurnCancellation>,
3811 ) -> Result<Self::Turn, LoopError> {
3812 if request.transcript.last().map(|item| item.kind) == Some(ItemKind::Assistant) {
3813 self.saw_assistant_tail.store(true, Ordering::SeqCst);
3814 return Err(LoopError::Provider(
3815 "conversation must end with a user message".into(),
3816 ));
3817 }
3818 Ok(FakeTurn {
3819 events: VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult {
3820 model: None,
3821 response_id: None,
3822 finish_reason: FinishReason::Completed,
3823 output_items: vec![Item::text(ItemKind::Assistant, "ok")],
3824 usage: None,
3825 metadata: MetadataMap::new(),
3826 })]),
3827 })
3828 }
3829 }
3830
3831 #[tokio::test]
3839 async fn drive_does_not_dispatch_without_valid_trailing_input() {
3840 let saw_assistant_tail = StdArc::new(AtomicBool::new(false));
3841 let agent = Agent::builder()
3842 .model(RejectAssistantPrefillAdapter {
3843 saw_assistant_tail: saw_assistant_tail.clone(),
3844 })
3845 .mutator(DropTrailingUserMutator)
3846 .transcript(vec![
3849 Item::text(ItemKind::User, "kickoff"),
3850 Item::text(ItemKind::Assistant, "prior reply"),
3851 ])
3852 .build()
3853 .unwrap();
3854
3855 let mut driver = agent
3856 .start(SessionConfig {
3857 session_id: SessionId::new("session-no-valid-input"),
3858 metadata: MetadataMap::new(),
3859 cache: None,
3860 })
3861 .await
3862 .unwrap();
3863
3864 driver
3865 .submit_input(vec![Item::text(ItemKind::User, "follow up")])
3866 .unwrap();
3867
3868 let outcome = driver.next().await;
3870
3871 assert!(
3872 !saw_assistant_tail.load(Ordering::SeqCst),
3873 "loop dispatched a model turn whose transcript ends in an assistant \
3874 message (outcome: {outcome:?}); with no valid trailing input the turn \
3875 must finish instead of driving"
3876 );
3877 }
3878
3879 #[tokio::test]
3880 async fn loop_uses_injected_permission_checker() {
3881 let tools = ToolRegistry::new().with(EchoTool::default());
3882 let agent = Agent::builder()
3883 .model(FakeAdapter)
3884 .add_tool_source(tools)
3885 .permissions(DenyFsReads)
3886 .build()
3887 .unwrap();
3888
3889 let mut driver = agent
3890 .start(SessionConfig {
3891 session_id: SessionId::new("session-2"),
3892 metadata: MetadataMap::new(),
3893 cache: None,
3894 })
3895 .await
3896 .unwrap();
3897
3898 driver
3899 .submit_input(vec![Item {
3900 id: None,
3901 kind: ItemKind::User,
3902 parts: vec![Part::Text(TextPart {
3903 text: "ping".into(),
3904 metadata: MetadataMap::new(),
3905 })],
3906 metadata: MetadataMap::new(),
3907 usage: None,
3908 finish_reason: None,
3909 created_at: None,
3910 }])
3911 .unwrap();
3912
3913 let result = run_until_finished(&mut driver).await;
3914
3915 match result {
3916 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
3917 Part::Text(text) => assert!(text.text.contains("tool permission denied")),
3918 other => panic!("unexpected part: {other:?}"),
3919 },
3920 other => panic!("unexpected loop step: {other:?}"),
3921 }
3922 }
3923
3924 #[tokio::test]
3925 async fn async_task_manager_background_round_requires_explicit_continue() {
3926 let entered = StdArc::new(AtomicBool::new(false));
3927 let release = StdArc::new(Notify::new());
3928 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
3929 "background-wait",
3930 RoutingDecision::Background,
3931 )]));
3932 let handle = task_manager.handle();
3933 let tools = ToolRegistry::new().with(BlockingTool::new(
3934 "background-wait",
3935 entered.clone(),
3936 release.clone(),
3937 "background-done",
3938 ));
3939 let agent = Agent::builder()
3940 .model(FakeAdapter)
3941 .add_tool_source(tools)
3942 .permissions(AllowAllPermissions)
3943 .task_manager(task_manager)
3944 .build()
3945 .unwrap();
3946
3947 let mut driver = agent
3948 .start(SessionConfig {
3949 session_id: SessionId::new("session-background"),
3950 metadata: MetadataMap::new(),
3951 cache: None,
3952 })
3953 .await
3954 .unwrap();
3955
3956 driver
3957 .submit_input(vec![Item {
3958 id: None,
3959 kind: ItemKind::User,
3960 parts: vec![Part::Text(TextPart {
3961 text: "ping".into(),
3962 metadata: MetadataMap::new(),
3963 })],
3964 metadata: MetadataMap::new(),
3965 usage: None,
3966 finish_reason: None,
3967 created_at: None,
3968 }])
3969 .unwrap();
3970
3971 let first = driver.next().await.unwrap();
3972 match first {
3973 LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => {}
3974 other => panic!("unexpected first loop step: {other:?}"),
3975 }
3976
3977 match wait_for_task_event(&handle).await {
3978 TaskEvent::Started(snapshot) => assert_eq!(snapshot.tool_name, "background-wait"),
3979 other => panic!("unexpected task event: {other:?}"),
3980 }
3981 wait_until_entered(entered.as_ref()).await;
3982 release.notify_waiters();
3983
3984 match wait_for_task_event(&handle).await {
3985 TaskEvent::Completed(_, result) => {
3986 assert_eq!(result.output, ToolOutput::Text("background-done".into()))
3987 }
3988 other => panic!("unexpected completion event: {other:?}"),
3989 }
3990
3991 let resumed = driver.next().await.unwrap();
3992 match resumed {
3993 LoopStep::Finished(turn) => {
3994 assert_eq!(turn.finish_reason, FinishReason::Completed);
3995 match &turn.items[0].parts[0] {
3996 Part::Text(text) => assert_eq!(text.text, "tool said: background-done"),
3997 other => panic!("unexpected part after resume: {other:?}"),
3998 }
3999 }
4000 other => panic!("unexpected resumed step: {other:?}"),
4001 }
4002 }
4003
4004 #[tokio::test]
4005 async fn loop_can_cancel_a_turn_and_continue_after_new_input() {
4006 let controller = CancellationController::new();
4007 let agent = Agent::builder()
4008 .model(SlowAdapter)
4009 .cancellation(controller.handle())
4010 .build()
4011 .unwrap();
4012
4013 let mut driver = agent
4014 .start(SessionConfig {
4015 session_id: SessionId::new("session-cancel"),
4016 metadata: MetadataMap::new(),
4017 cache: None,
4018 })
4019 .await
4020 .unwrap();
4021
4022 driver
4023 .submit_input(vec![Item {
4024 id: None,
4025 kind: ItemKind::User,
4026 parts: vec![Part::Text(TextPart {
4027 text: "do the long task".into(),
4028 metadata: MetadataMap::new(),
4029 })],
4030 metadata: MetadataMap::new(),
4031 usage: None,
4032 finish_reason: None,
4033 created_at: None,
4034 }])
4035 .unwrap();
4036
4037 let cancelled = tokio::join!(async { driver.next().await }, async {
4038 tokio::task::yield_now().await;
4039 controller.interrupt();
4040 })
4041 .0
4042 .unwrap();
4043
4044 match cancelled {
4045 LoopStep::Finished(turn) => {
4046 assert_eq!(turn.finish_reason, FinishReason::Cancelled);
4047 assert_eq!(turn.items.len(), 1);
4048 assert_eq!(turn.items[0].kind, ItemKind::Assistant);
4049 assert_eq!(
4050 turn.items[0].metadata.get(INTERRUPTED_METADATA_KEY),
4051 Some(&Value::Bool(true))
4052 );
4053 }
4054 other => panic!("unexpected loop step: {other:?}"),
4055 }
4056
4057 driver
4058 .submit_input(vec![Item {
4059 id: None,
4060 kind: ItemKind::User,
4061 parts: vec![Part::Text(TextPart {
4062 text: "try again".into(),
4063 metadata: MetadataMap::new(),
4064 })],
4065 metadata: MetadataMap::new(),
4066 usage: None,
4067 finish_reason: None,
4068 created_at: None,
4069 }])
4070 .unwrap();
4071
4072 let result = driver.next().await.unwrap();
4073 match result {
4074 LoopStep::Finished(turn) => {
4075 assert_eq!(turn.finish_reason, FinishReason::Completed);
4076 }
4077 other => panic!("unexpected loop step after retry: {other:?}"),
4078 }
4079 }
4080
4081 #[tokio::test]
4082 async fn loop_interrupt_cancels_foreground_tasks_but_keeps_background_tasks_running() {
4083 let controller = CancellationController::new();
4084 let fg_entered = StdArc::new(AtomicBool::new(false));
4085 let fg_release = StdArc::new(Notify::new());
4086 let bg_entered = StdArc::new(AtomicBool::new(false));
4087 let bg_release = StdArc::new(Notify::new());
4088 let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([
4089 ("foreground-wait", RoutingDecision::Foreground),
4090 ("background-wait", RoutingDecision::Background),
4091 ]));
4092 let handle = task_manager.handle();
4093 let tools = ToolRegistry::new()
4094 .with(BlockingTool::new(
4095 "foreground-wait",
4096 fg_entered.clone(),
4097 fg_release,
4098 "foreground-done",
4099 ))
4100 .with(BlockingTool::new(
4101 "background-wait",
4102 bg_entered.clone(),
4103 bg_release.clone(),
4104 "background-done",
4105 ));
4106 let agent = Agent::builder()
4107 .model(MultiToolAdapter)
4108 .add_tool_source(tools)
4109 .permissions(AllowAllPermissions)
4110 .cancellation(controller.handle())
4111 .task_manager(task_manager)
4112 .build()
4113 .unwrap();
4114
4115 let mut driver = agent
4116 .start(SessionConfig {
4117 session_id: SessionId::new("session-mixed-cancel"),
4118 metadata: MetadataMap::new(),
4119 cache: None,
4120 })
4121 .await
4122 .unwrap();
4123
4124 driver
4125 .submit_input(vec![Item {
4126 id: None,
4127 kind: ItemKind::User,
4128 parts: vec![Part::Text(TextPart {
4129 text: "run both".into(),
4130 metadata: MetadataMap::new(),
4131 })],
4132 metadata: MetadataMap::new(),
4133 usage: None,
4134 finish_reason: None,
4135 created_at: None,
4136 }])
4137 .unwrap();
4138
4139 let cancelled = tokio::join!(async { driver.next().await }, async {
4140 let _ = wait_for_task_event(&handle).await;
4141 let _ = wait_for_task_event(&handle).await;
4142 wait_until_entered(fg_entered.as_ref()).await;
4143 wait_until_entered(bg_entered.as_ref()).await;
4144 controller.interrupt();
4145 })
4146 .0
4147 .unwrap();
4148
4149 match cancelled {
4150 LoopStep::Finished(turn) => assert_eq!(turn.finish_reason, FinishReason::Cancelled),
4151 other => panic!("unexpected loop step after interrupt: {other:?}"),
4152 }
4153
4154 match wait_for_task_event(&handle).await {
4155 TaskEvent::Cancelled(snapshot) => assert_eq!(snapshot.tool_name, "foreground-wait"),
4156 other => panic!("unexpected post-interrupt event: {other:?}"),
4157 }
4158
4159 let running = handle.list_running().await;
4160 assert_eq!(running.len(), 1);
4161 assert_eq!(running[0].tool_name, "background-wait");
4162
4163 bg_release.notify_waiters();
4164 match wait_for_task_event(&handle).await {
4165 TaskEvent::Completed(snapshot, result) => {
4166 assert_eq!(snapshot.tool_name, "background-wait");
4167 assert_eq!(result.output, ToolOutput::Text("background-done".into()));
4168 }
4169 other => panic!("unexpected background completion event: {other:?}"),
4170 }
4171 }
4172
4173 #[tokio::test]
4174 async fn loop_resumes_after_approved_tool_request() {
4175 let tools = ToolRegistry::new().with(EchoTool::default());
4176 let agent = Agent::builder()
4177 .model(FakeAdapter)
4178 .add_tool_source(tools)
4179 .permissions(ApproveFsReads)
4180 .build()
4181 .unwrap();
4182
4183 let mut driver = agent
4184 .start(SessionConfig {
4185 session_id: SessionId::new("session-approval"),
4186 metadata: MetadataMap::new(),
4187 cache: None,
4188 })
4189 .await
4190 .unwrap();
4191
4192 driver
4193 .submit_input(vec![Item {
4194 id: None,
4195 kind: ItemKind::User,
4196 parts: vec![Part::Text(TextPart {
4197 text: "ping".into(),
4198 metadata: MetadataMap::new(),
4199 })],
4200 metadata: MetadataMap::new(),
4201 usage: None,
4202 finish_reason: None,
4203 created_at: None,
4204 }])
4205 .unwrap();
4206
4207 let first = driver.next().await.unwrap();
4208 match first {
4209 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
4210 assert!(pending.request.task_id.is_some());
4211 assert_eq!(pending.request.id.0, "approval:fs-read");
4212 pending.approve(&mut driver).unwrap();
4213 }
4214 other => panic!("unexpected loop step: {other:?}"),
4215 }
4216 let second = driver.next().await.unwrap();
4217 match second {
4218 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
4219 Part::Text(text) => assert_eq!(text.text, "tool said: pong"),
4220 other => panic!("unexpected part: {other:?}"),
4221 },
4222 other => panic!("unexpected loop step after approval: {other:?}"),
4223 }
4224 }
4225
4226 #[tokio::test]
4227 async fn loop_resumes_with_patched_input_on_approval() {
4228 let tools = ToolRegistry::new().with(EchoTool::default());
4229 let agent = Agent::builder()
4230 .model(FakeAdapter)
4231 .add_tool_source(tools)
4232 .permissions(ApproveFsReads)
4233 .build()
4234 .unwrap();
4235
4236 let mut driver = agent
4237 .start(SessionConfig {
4238 session_id: SessionId::new("session-approval-patched"),
4239 metadata: MetadataMap::new(),
4240 cache: None,
4241 })
4242 .await
4243 .unwrap();
4244
4245 driver
4246 .submit_input(vec![Item {
4247 id: None,
4248 kind: ItemKind::User,
4249 parts: vec![Part::Text(TextPart {
4250 text: "ping".into(),
4251 metadata: MetadataMap::new(),
4252 })],
4253 metadata: MetadataMap::new(),
4254 usage: None,
4255 finish_reason: None,
4256 created_at: None,
4257 }])
4258 .unwrap();
4259
4260 match driver.next().await.unwrap() {
4261 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
4262 pending
4263 .approve_with_patched_input(&mut driver, json!({ "value": "patched" }))
4264 .unwrap();
4265 }
4266 other => panic!("unexpected loop step: {other:?}"),
4267 }
4268 match driver.next().await.unwrap() {
4269 LoopStep::Finished(turn) => match &turn.items[0].parts[0] {
4270 Part::Text(text) => assert_eq!(text.text, "tool said: patched"),
4271 other => panic!("unexpected part: {other:?}"),
4272 },
4273 other => panic!("unexpected loop step after approval: {other:?}"),
4274 }
4275 }
4276
4277 #[tokio::test]
4278 async fn loop_tracks_multiple_pending_approvals_by_call_id() {
4279 let tools = ToolRegistry::new().with(EchoTool::default());
4280 let agent = Agent::builder()
4281 .model(DualApprovalAdapter)
4282 .add_tool_source(tools)
4283 .permissions(ApproveFsReads)
4284 .build()
4285 .unwrap();
4286
4287 let mut driver = agent
4288 .start(SessionConfig {
4289 session_id: SessionId::new("session-dual-approval"),
4290 metadata: MetadataMap::new(),
4291 cache: None,
4292 })
4293 .await
4294 .unwrap();
4295
4296 driver
4297 .submit_input(vec![Item {
4298 id: None,
4299 kind: ItemKind::User,
4300 parts: vec![Part::Text(TextPart {
4301 text: "run both approvals".into(),
4302 metadata: MetadataMap::new(),
4303 })],
4304 metadata: MetadataMap::new(),
4305 usage: None,
4306 finish_reason: None,
4307 created_at: None,
4308 }])
4309 .unwrap();
4310
4311 let pending_first = match driver.next().await.unwrap() {
4312 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
4313 assert_eq!(
4314 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
4315 Some("call-1")
4316 );
4317 pending
4318 }
4319 other => panic!("unexpected first loop step: {other:?}"),
4320 };
4321
4322 let pending_second = match driver.next().await.unwrap() {
4323 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
4324 assert_eq!(
4325 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
4326 Some("call-2")
4327 );
4328 pending
4329 }
4330 other => panic!("unexpected second loop step: {other:?}"),
4331 };
4332
4333 pending_second.approve(&mut driver).unwrap();
4334 match driver.next().await.unwrap() {
4335 LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(pending)) => {
4336 assert_eq!(
4337 pending.request.call_id.as_ref().map(|id| id.0.as_str()),
4338 Some("call-1")
4339 );
4340 }
4341 other => panic!("unexpected step after approving second request: {other:?}"),
4342 }
4343
4344 pending_first.approve(&mut driver).unwrap();
4345 match driver.next().await.unwrap() {
4346 LoopStep::Finished(turn) => {
4347 assert_eq!(turn.finish_reason, FinishReason::Completed);
4348 match &turn.items[0].parts[0] {
4349 Part::Text(text) => assert_eq!(text.text, "both approvals finished"),
4350 other => panic!("unexpected final part: {other:?}"),
4351 }
4352 }
4353 other => panic!("unexpected final loop step: {other:?}"),
4354 }
4355 }
4356
4357 #[tokio::test]
4358 async fn loop_compacts_transcript_before_new_turns() {
4359 let events = StdArc::new(StdMutex::new(Vec::new()));
4360 let agent = Agent::builder()
4361 .model(FakeAdapter)
4362 .mutator(KeepRecentMutator { keep: 1 })
4363 .observer(RecordingObserver {
4364 events: events.clone(),
4365 })
4366 .build()
4367 .unwrap();
4368
4369 let mut driver = agent
4370 .start(SessionConfig {
4371 session_id: SessionId::new("session-4"),
4372 metadata: MetadataMap::new(),
4373 cache: None,
4374 })
4375 .await
4376 .unwrap();
4377
4378 for text in ["first", "second"] {
4379 driver
4380 .submit_input(vec![Item {
4381 id: None,
4382 kind: ItemKind::User,
4383 parts: vec![Part::Text(TextPart {
4384 text: text.into(),
4385 metadata: MetadataMap::new(),
4386 })],
4387 metadata: MetadataMap::new(),
4388 usage: None,
4389 finish_reason: None,
4390 created_at: None,
4391 }])
4392 .unwrap();
4393 let _ = driver.next().await.unwrap();
4394 }
4395
4396 let events = events.lock().unwrap();
4397 assert!(
4398 events
4399 .iter()
4400 .any(|event| matches!(event, AgentEvent::MutationFinished { dirty: true, .. }))
4401 );
4402 }
4403
4404 #[test]
4405 fn transcript_validation_rejects_orphaned_tool_result() {
4406 let transcript = vec![Item {
4407 id: None,
4408 kind: ItemKind::Tool,
4409 parts: vec![Part::ToolResult(ToolResultPart {
4410 call_id: "call-1".into(),
4411 output: ToolOutput::Text("result".into()),
4412 is_error: false,
4413 metadata: MetadataMap::new(),
4414 })],
4415 metadata: MetadataMap::new(),
4416 usage: None,
4417 finish_reason: None,
4418 created_at: None,
4419 }];
4420
4421 let error = validate_transcript_invariants(&transcript).unwrap_err();
4422 assert!(error.to_string().contains("orphaned tool_result"));
4423 }
4424
4425 #[test]
4426 fn transcript_validation_rejects_duplicate_tool_result() {
4427 let transcript = vec![
4428 Item {
4429 id: None,
4430 kind: ItemKind::Assistant,
4431 parts: vec![Part::ToolCall(ToolCallPart {
4432 id: "call-1".into(),
4433 name: "lookup".into(),
4434 input: serde_json::json!({}),
4435 metadata: MetadataMap::new(),
4436 })],
4437 metadata: MetadataMap::new(),
4438 usage: None,
4439 finish_reason: None,
4440 created_at: None,
4441 },
4442 Item {
4443 id: None,
4444 kind: ItemKind::Tool,
4445 parts: vec![Part::ToolResult(ToolResultPart {
4446 call_id: "call-1".into(),
4447 output: ToolOutput::Text("result".into()),
4448 is_error: false,
4449 metadata: MetadataMap::new(),
4450 })],
4451 metadata: MetadataMap::new(),
4452 usage: None,
4453 finish_reason: None,
4454 created_at: None,
4455 },
4456 Item {
4457 id: None,
4458 kind: ItemKind::Tool,
4459 parts: vec![Part::ToolResult(ToolResultPart {
4460 call_id: "call-1".into(),
4461 output: ToolOutput::Text("again".into()),
4462 is_error: false,
4463 metadata: MetadataMap::new(),
4464 })],
4465 metadata: MetadataMap::new(),
4466 usage: None,
4467 finish_reason: None,
4468 created_at: None,
4469 },
4470 ];
4471
4472 let error = validate_transcript_invariants(&transcript).unwrap_err();
4473 assert!(error.to_string().contains("duplicate tool_result"));
4474 }
4475
4476 #[tokio::test]
4477 async fn loop_refreshes_tool_specs_each_turn() {
4478 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
4479 let version = StdArc::new(AtomicUsize::new(1));
4480 let tools = ToolRegistry::new().with(DynamicSpecTool::new(version.clone()));
4481 let agent = Agent::builder()
4482 .model(RecordingAdapter {
4483 seen_descriptions: seen_descriptions.clone(),
4484 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
4485 })
4486 .add_tool_source(tools)
4487 .permissions(AllowAllPermissions)
4488 .build()
4489 .unwrap();
4490
4491 let mut driver = agent
4492 .start(SessionConfig {
4493 session_id: SessionId::new("session-dynamic-tools"),
4494 metadata: MetadataMap::new(),
4495 cache: None,
4496 })
4497 .await
4498 .unwrap();
4499
4500 for text in ["first", "second"] {
4501 driver
4502 .submit_input(vec![Item {
4503 id: None,
4504 kind: ItemKind::User,
4505 parts: vec![Part::Text(TextPart {
4506 text: text.into(),
4507 metadata: MetadataMap::new(),
4508 })],
4509 metadata: MetadataMap::new(),
4510 usage: None,
4511 finish_reason: None,
4512 created_at: None,
4513 }])
4514 .unwrap();
4515
4516 let _ = driver.next().await.unwrap();
4517 if text == "first" {
4518 version.store(2, Ordering::SeqCst);
4519 }
4520 }
4521
4522 let seen_descriptions = seen_descriptions.lock().unwrap();
4523 assert_eq!(seen_descriptions.len(), 2);
4524 assert_eq!(seen_descriptions[0], vec!["dynamic version 1".to_string()]);
4525 assert_eq!(seen_descriptions[1], vec!["dynamic version 2".to_string()]);
4526 }
4527
4528 #[tokio::test]
4529 async fn loop_emits_catalog_change_and_uses_updated_specs_next_turn() {
4530 let seen_descriptions = StdArc::new(StdMutex::new(Vec::new()));
4531 let events = StdArc::new(StdMutex::new(Vec::new()));
4532 let executor = StdArc::new(CatalogExecutor::new());
4533 let executor_for_agent: Arc<dyn ToolExecutor> = executor.clone();
4534 let agent = Agent::builder()
4535 .model(RecordingAdapter {
4536 seen_descriptions: seen_descriptions.clone(),
4537 seen_caches: StdArc::new(StdMutex::new(Vec::new())),
4538 })
4539 .tool_executor(executor_for_agent)
4540 .permissions(AllowAllPermissions)
4541 .observer(RecordingObserver {
4542 events: events.clone(),
4543 })
4544 .build()
4545 .unwrap();
4546
4547 let mut driver = agent
4548 .start(SessionConfig {
4549 session_id: SessionId::new("session-catalog-events"),
4550 metadata: MetadataMap::new(),
4551 cache: None,
4552 })
4553 .await
4554 .unwrap();
4555
4556 driver
4557 .submit_input(vec![Item::text(ItemKind::User, "first")])
4558 .unwrap();
4559 let _ = driver.next().await.unwrap();
4560
4561 executor.publish_change(
4562 1,
4563 ToolCatalogEvent {
4564 source: "mcp:mock".into(),
4565 added: vec!["dynamic".into()],
4566 removed: Vec::new(),
4567 changed: Vec::new(),
4568 },
4569 );
4570
4571 driver
4572 .submit_input(vec![Item::text(ItemKind::User, "second")])
4573 .unwrap();
4574 let _ = driver.next().await.unwrap();
4575
4576 let seen_descriptions = seen_descriptions.lock().unwrap();
4577 assert_eq!(seen_descriptions.len(), 2);
4578 assert_eq!(seen_descriptions[0], vec!["dynamic version 0".to_string()]);
4579 assert_eq!(seen_descriptions[1], vec!["dynamic version 1".to_string()]);
4580
4581 let events = events.lock().unwrap();
4582 assert!(events.iter().any(|event| matches!(
4583 event,
4584 AgentEvent::ToolCatalogChanged(ToolCatalogEvent {
4585 source,
4586 added,
4587 removed,
4588 changed,
4589 }) if source == "mcp:mock"
4590 && added == &vec!["dynamic".to_string()]
4591 && removed.is_empty()
4592 && changed.is_empty()
4593 )));
4594 }
4595
4596 #[tokio::test]
4597 async fn loop_passes_session_default_and_next_turn_cache_requests() {
4598 let seen_caches = StdArc::new(StdMutex::new(Vec::new()));
4599 let agent = Agent::builder()
4600 .model(RecordingAdapter {
4601 seen_descriptions: StdArc::new(StdMutex::new(Vec::new())),
4602 seen_caches: seen_caches.clone(),
4603 })
4604 .permissions(AllowAllPermissions)
4605 .build()
4606 .unwrap();
4607
4608 let default_cache = PromptCacheRequest::best_effort(PromptCacheStrategy::Automatic)
4609 .with_retention(PromptCacheRetention::Short);
4610 let override_cache = PromptCacheRequest::required(PromptCacheStrategy::Explicit {
4611 breakpoints: vec![PromptCacheBreakpoint::TranscriptItemEnd { index: 0 }],
4612 });
4613
4614 let mut driver = agent
4615 .start(SessionConfig {
4616 session_id: SessionId::new("session-cache"),
4617 metadata: MetadataMap::new(),
4618 cache: Some(default_cache.clone()),
4619 })
4620 .await
4621 .unwrap();
4622
4623 driver
4624 .submit_input(vec![Item {
4625 id: None,
4626 kind: ItemKind::User,
4627 parts: vec![Part::Text(TextPart {
4628 text: "first".into(),
4629 metadata: MetadataMap::new(),
4630 })],
4631 metadata: MetadataMap::new(),
4632 usage: None,
4633 finish_reason: None,
4634 created_at: None,
4635 }])
4636 .unwrap();
4637 let _ = driver.next().await.unwrap();
4638
4639 driver
4640 .submit_input_with_cache(
4641 vec![Item {
4642 id: None,
4643 kind: ItemKind::User,
4644 parts: vec![Part::Text(TextPart {
4645 text: "second".into(),
4646 metadata: MetadataMap::new(),
4647 })],
4648 metadata: MetadataMap::new(),
4649 usage: None,
4650 finish_reason: None,
4651 created_at: None,
4652 }],
4653 override_cache.clone(),
4654 )
4655 .unwrap();
4656 let _ = driver.next().await.unwrap();
4657
4658 let seen = seen_caches.lock().unwrap();
4659 assert_eq!(seen.len(), 2);
4660 assert_eq!(seen[0], Some(default_cache));
4661 assert_eq!(seen[1], Some(override_cache));
4662 }
4663
4664 #[tokio::test]
4665 async fn loop_yields_after_tool_result_between_rounds() {
4666 let tools = ToolRegistry::new().with(EchoTool::default());
4667 let agent = Agent::builder()
4668 .model(FakeAdapter)
4669 .add_tool_source(tools)
4670 .permissions(AllowAllPermissions)
4671 .build()
4672 .unwrap();
4673
4674 let mut driver = agent
4675 .start(SessionConfig {
4676 session_id: SessionId::new("yield-session"),
4677 metadata: MetadataMap::new(),
4678 cache: None,
4679 })
4680 .await
4681 .unwrap();
4682
4683 driver
4684 .submit_input(vec![Item::text(ItemKind::User, "ping")])
4685 .unwrap();
4686
4687 let step = driver.next().await.unwrap();
4690 let info = match step {
4691 LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info,
4692 other => panic!("expected AfterToolResult, got {other:?}"),
4693 };
4694 assert_eq!(info.session_id, SessionId::new("yield-session"));
4695 assert_eq!(info.transcript_len, 3);
4697
4698 let interrupt = LoopInterrupt::AfterToolResult(info.clone());
4700 assert!(!interrupt.is_blocking());
4701
4702 driver
4704 .submit_input(vec![Item::text(ItemKind::User, "also: report back")])
4705 .unwrap();
4706
4707 let step = driver.next().await.unwrap();
4710 match step {
4711 LoopStep::Finished(turn) => {
4712 assert_eq!(turn.finish_reason, FinishReason::Completed);
4713 }
4714 other => panic!("expected Finished, got {other:?}"),
4715 }
4716
4717 let snapshot = driver.snapshot();
4719 let has_injected_message = snapshot.transcript.iter().any(|item| {
4720 item.kind == ItemKind::User
4721 && item.parts.iter().any(|part| match part {
4722 Part::Text(text) => text.text == "also: report back",
4723 _ => false,
4724 })
4725 });
4726 assert!(
4727 has_injected_message,
4728 "injected user message should be in transcript, got: {:?}",
4729 snapshot.transcript
4730 );
4731 }
4732
4733 struct RecordingTranscriptObserver {
4734 items: StdArc<StdMutex<Vec<Item>>>,
4735 }
4736
4737 impl TranscriptObserver for RecordingTranscriptObserver {
4738 fn on_item_appended(&self, item: &Item) {
4739 self.items.lock().unwrap().push(item.clone());
4740 }
4741 }
4742
4743 #[tokio::test]
4744 async fn observers_see_full_tool_round() {
4745 let events = StdArc::new(StdMutex::new(Vec::<AgentEvent>::new()));
4751 let items = StdArc::new(StdMutex::new(Vec::<Item>::new()));
4752 let agent = Agent::builder()
4753 .model(FakeAdapter)
4754 .add_tool_source(ToolRegistry::new().with(EchoTool::default()))
4755 .permissions(AllowAllPermissions)
4756 .observer(RecordingObserver {
4757 events: events.clone(),
4758 })
4759 .transcript_observer(RecordingTranscriptObserver {
4760 items: items.clone(),
4761 })
4762 .build()
4763 .unwrap();
4764
4765 let mut driver = agent
4766 .start(SessionConfig {
4767 session_id: SessionId::new("observer-session"),
4768 metadata: MetadataMap::new(),
4769 cache: None,
4770 })
4771 .await
4772 .unwrap();
4773
4774 driver
4775 .submit_input(vec![Item {
4776 id: None,
4777 kind: ItemKind::User,
4778 parts: vec![Part::Text(TextPart {
4779 text: "ping".into(),
4780 metadata: MetadataMap::new(),
4781 })],
4782 metadata: MetadataMap::new(),
4783 usage: None,
4784 finish_reason: None,
4785 created_at: None,
4786 }])
4787 .unwrap();
4788
4789 let result = run_until_finished(&mut driver).await;
4790 assert!(matches!(result, LoopStep::Finished(_)), "got {result:?}");
4791
4792 let events = events.lock().unwrap().clone();
4795 let tool_call_id = events.iter().find_map(|e| match e {
4796 AgentEvent::ToolCallRequested(c) => Some(c.id.clone()),
4797 _ => None,
4798 });
4799 let tool_results: Vec<_> = events
4800 .iter()
4801 .filter_map(|e| match e {
4802 AgentEvent::ToolResultReceived(r) => Some(r.clone()),
4803 _ => None,
4804 })
4805 .collect();
4806 assert_eq!(tool_results.len(), 1, "events: {events:?}");
4807 assert_eq!(Some(tool_results[0].call_id.clone()), tool_call_id);
4808 assert!(!tool_results[0].is_error);
4809
4810 let items = items.lock().unwrap().clone();
4814 assert_eq!(items.len(), 4, "items: {items:?}");
4815 assert_eq!(items[0].kind, ItemKind::User);
4816 assert_eq!(items[1].kind, ItemKind::Assistant);
4817 assert!(
4818 items[1]
4819 .parts
4820 .iter()
4821 .any(|p| matches!(p, Part::ToolCall(_)))
4822 );
4823 assert_eq!(items[2].kind, ItemKind::Tool);
4824 assert!(
4825 items[2]
4826 .parts
4827 .iter()
4828 .any(|p| matches!(p, Part::ToolResult(_)))
4829 );
4830 assert_eq!(items[3].kind, ItemKind::Assistant);
4831 }
4832
4833 #[test]
4834 fn convenience_cache_builders_construct_expected_defaults() {
4835 let cache = PromptCacheRequest::automatic()
4836 .with_retention(PromptCacheRetention::Short)
4837 .with_key("workspace:demo");
4838 let session = SessionConfig::new("demo").with_cache(cache.clone());
4839
4840 assert_eq!(session.session_id, SessionId::new("demo"));
4841 assert_eq!(session.cache, Some(cache));
4842
4843 let explicit = PromptCacheRequest::explicit([
4844 PromptCacheBreakpoint::tools_end(),
4845 PromptCacheBreakpoint::transcript_item_end(2),
4846 PromptCacheBreakpoint::transcript_part_end(3, 1),
4847 ]);
4848
4849 assert_eq!(explicit.mode, PromptCacheMode::BestEffort);
4850 assert_eq!(
4851 explicit.strategy,
4852 PromptCacheStrategy::Explicit {
4853 breakpoints: vec![
4854 PromptCacheBreakpoint::ToolsEnd,
4855 PromptCacheBreakpoint::TranscriptItemEnd { index: 2 },
4856 PromptCacheBreakpoint::TranscriptPartEnd {
4857 item_index: 3,
4858 part_index: 1,
4859 },
4860 ],
4861 }
4862 );
4863 }
4864}