1use crate::agent::Agent;
9use crate::harness::Harness;
10use crate::llm_models::LlmProviderType;
11use crate::session_file::{FileInfo, FileStat, GrepMatch, InitialFile, SessionFile};
12use crate::tool_types::{ToolCall, ToolDefinition, ToolResult};
13use crate::typed_id::{AgentId, HarnessId, ImageId, ModelId, SessionId};
14use async_trait::async_trait;
15use chrono::{DateTime, Utc};
16use std::any::{Any, TypeId};
17use std::collections::{HashMap, HashSet};
18use std::sync::Arc;
19use uuid::Uuid;
20
21fn build_tool_map(tool_defs: &[ToolDefinition]) -> HashMap<&str, &ToolDefinition> {
23 tool_defs.iter().map(|def| (def.name(), def)).collect()
24}
25
26use crate::error::Result;
27
28#[async_trait]
39pub trait AgentStore: Send + Sync {
40 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>>;
42}
43
44#[async_trait]
45impl<T: AgentStore + ?Sized> AgentStore for std::sync::Arc<T> {
46 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>> {
47 (**self).get_agent(agent_id).await
48 }
49}
50
51#[async_trait]
66pub trait HarnessStore: Send + Sync {
67 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>>;
72}
73
74#[async_trait]
75impl<T: HarnessStore + ?Sized> HarnessStore for std::sync::Arc<T> {
76 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>> {
77 (**self).get_harness_chain(harness_id).await
78 }
79}
80
81use crate::leased_resource::{LeasedResource, UpsertLeasedResource};
86use crate::session::Session;
87
88#[async_trait]
94pub trait SessionStore: Send + Sync {
95 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>>;
97}
98
99#[async_trait]
100impl<T: SessionStore + ?Sized> SessionStore for std::sync::Arc<T> {
101 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>> {
102 (**self).get_session(session_id).await
103 }
104}
105
106#[async_trait]
108pub trait SessionMutator: Send + Sync {
109 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session>;
111}
112
113#[async_trait]
114impl<T: SessionMutator + ?Sized> SessionMutator for std::sync::Arc<T> {
115 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session> {
116 (**self).update_session_title(session_id, title).await
117 }
118}
119
120#[derive(Debug, Clone)]
126pub struct ModelWithProvider {
127 pub model: String,
129 pub provider_type: LlmProviderType,
131 pub api_key: Option<String>,
133 pub base_url: Option<String>,
135}
136
137#[async_trait]
147pub trait LlmProviderStore: Send + Sync {
148 async fn get_model_with_provider(&self, model_id: ModelId)
153 -> Result<Option<ModelWithProvider>>;
154
155 async fn get_default_model(&self) -> Result<Option<ModelWithProvider>>;
159}
160
161#[async_trait]
162impl<T: LlmProviderStore + ?Sized> LlmProviderStore for std::sync::Arc<T> {
163 async fn get_model_with_provider(
164 &self,
165 model_id: ModelId,
166 ) -> Result<Option<ModelWithProvider>> {
167 (**self).get_model_with_provider(model_id).await
168 }
169
170 async fn get_default_model(&self) -> Result<Option<ModelWithProvider>> {
171 (**self).get_default_model().await
172 }
173}
174
175#[derive(Debug, Clone)]
181pub struct StoredImageInfo {
182 pub id: ImageId,
183 pub filename: String,
184 pub content_type: String,
185 pub size_bytes: i64,
186 pub metadata: serde_json::Value,
187 pub created_at: DateTime<Utc>,
188}
189
190#[derive(Debug, Clone)]
192pub struct StoredImage {
193 pub info: StoredImageInfo,
194 pub data: Vec<u8>,
195}
196
197#[derive(Debug, Clone)]
199pub struct CreateStoredImage {
200 pub filename: String,
201 pub content_type: String,
202 pub data: Vec<u8>,
203 pub metadata: serde_json::Value,
204}
205
206#[async_trait]
207pub trait ImageArtifactStore: Send + Sync {
208 async fn create_image(&self, input: CreateStoredImage) -> Result<StoredImageInfo>;
210
211 async fn get_image(&self, image_id: ImageId) -> Result<Option<StoredImage>>;
213
214 async fn get_image_info(&self, image_id: ImageId) -> Result<Option<StoredImageInfo>>;
216}
217
218#[derive(Debug, Clone)]
224pub struct ProviderCredentials {
225 pub api_key: String,
226 pub base_url: Option<String>,
227}
228
229#[async_trait]
230pub trait ProviderCredentialStore: Send + Sync {
231 async fn get_default_provider_credentials(
236 &self,
237 provider_type: &str,
238 ) -> Result<Option<ProviderCredentials>>;
239}
240
241#[async_trait]
252pub trait ToolExecutor: Send + Sync {
253 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult>;
258
259 async fn execute_with_context(
264 &self,
265 tool_call: &ToolCall,
266 tool_def: &ToolDefinition,
267 _context: &ToolContext,
268 ) -> Result<ToolResult> {
269 self.execute(tool_call, tool_def).await
271 }
272
273 async fn execute_batch(
275 &self,
276 tool_calls: &[ToolCall],
277 tool_defs: &[ToolDefinition],
278 ) -> Result<Vec<ToolResult>> {
279 let mut results = Vec::with_capacity(tool_calls.len());
280
281 let tool_map = build_tool_map(tool_defs);
282
283 for tool_call in tool_calls {
284 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
285 crate::error::AgentLoopError::tool(format!(
286 "Tool definition not found: {}",
287 tool_call.name
288 ))
289 })?;
290
291 results.push(self.execute(tool_call, tool_def).await?);
292 }
293
294 Ok(results)
295 }
296
297 async fn execute_parallel(
299 &self,
300 tool_calls: &[ToolCall],
301 tool_defs: &[ToolDefinition],
302 ) -> Result<Vec<ToolResult>>
303 where
304 Self: Sized,
305 {
306 use futures::future::join_all;
307
308 let tool_map = build_tool_map(tool_defs);
309
310 let futures: Vec<_> = tool_calls
311 .iter()
312 .map(|tool_call| async {
313 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
314 crate::error::AgentLoopError::tool(format!(
315 "Tool definition not found: {}",
316 tool_call.name
317 ))
318 })?;
319 self.execute(tool_call, tool_def).await
320 })
321 .collect();
322
323 let results = join_all(futures).await;
324 results.into_iter().collect()
325 }
326}
327
328#[async_trait]
332impl ToolExecutor for std::sync::Arc<dyn ToolExecutor> {
333 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult> {
334 (**self).execute(tool_call, tool_def).await
335 }
336
337 async fn execute_with_context(
338 &self,
339 tool_call: &ToolCall,
340 tool_def: &ToolDefinition,
341 context: &ToolContext,
342 ) -> Result<ToolResult> {
343 (**self)
344 .execute_with_context(tool_call, tool_def, context)
345 .await
346 }
347
348 async fn execute_batch(
349 &self,
350 tool_calls: &[ToolCall],
351 tool_defs: &[ToolDefinition],
352 ) -> Result<Vec<ToolResult>> {
353 (**self).execute_batch(tool_calls, tool_defs).await
354 }
355}
356
357#[async_trait]
369pub trait SessionFileSystem: Send + Sync {
370 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>>;
372
373 async fn write_file(
375 &self,
376 session_id: SessionId,
377 path: &str,
378 content: &str,
379 encoding: &str,
380 ) -> Result<SessionFile>;
381
382 async fn write_file_if_content_matches(
387 &self,
388 session_id: SessionId,
389 path: &str,
390 expected_content: &str,
391 expected_encoding: &str,
392 content: &str,
393 encoding: &str,
394 ) -> Result<Option<SessionFile>> {
395 let Some(existing) = self.read_file(session_id, path).await? else {
396 return Ok(None);
397 };
398
399 if existing.is_directory {
400 return Ok(None);
401 }
402
403 let current_content = existing.content.unwrap_or_default();
404 if current_content != expected_content || existing.encoding != expected_encoding {
405 return Ok(None);
406 }
407
408 self.write_file(session_id, path, content, encoding)
409 .await
410 .map(Some)
411 }
412
413 async fn delete_file(&self, session_id: SessionId, path: &str, recursive: bool)
415 -> Result<bool>;
416
417 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>>;
419
420 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>>;
422
423 async fn grep_files(
425 &self,
426 session_id: SessionId,
427 pattern: &str,
428 path_pattern: Option<&str>,
429 ) -> Result<Vec<GrepMatch>>;
430
431 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
433
434 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
436 if file.is_readonly {
437 return Err(crate::error::AgentLoopError::store(
438 "read-only initial files require a SessionFileSystem-specific seed implementation",
439 ));
440 }
441 self.write_file(session_id, &file.path, &file.content, &file.encoding)
442 .await?;
443 Ok(())
444 }
445}
446
447#[async_trait]
448impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
449 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
450 (**self).read_file(session_id, path).await
451 }
452
453 async fn write_file(
454 &self,
455 session_id: SessionId,
456 path: &str,
457 content: &str,
458 encoding: &str,
459 ) -> Result<SessionFile> {
460 (**self)
461 .write_file(session_id, path, content, encoding)
462 .await
463 }
464
465 async fn write_file_if_content_matches(
466 &self,
467 session_id: SessionId,
468 path: &str,
469 expected_content: &str,
470 expected_encoding: &str,
471 content: &str,
472 encoding: &str,
473 ) -> Result<Option<SessionFile>> {
474 (**self)
475 .write_file_if_content_matches(
476 session_id,
477 path,
478 expected_content,
479 expected_encoding,
480 content,
481 encoding,
482 )
483 .await
484 }
485
486 async fn delete_file(
487 &self,
488 session_id: SessionId,
489 path: &str,
490 recursive: bool,
491 ) -> Result<bool> {
492 (**self).delete_file(session_id, path, recursive).await
493 }
494
495 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
496 (**self).list_directory(session_id, path).await
497 }
498
499 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
500 (**self).stat_file(session_id, path).await
501 }
502
503 async fn grep_files(
504 &self,
505 session_id: SessionId,
506 pattern: &str,
507 path_pattern: Option<&str>,
508 ) -> Result<Vec<GrepMatch>> {
509 (**self).grep_files(session_id, pattern, path_pattern).await
510 }
511
512 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
513 (**self).create_directory(session_id, path).await
514 }
515
516 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
517 (**self).seed_initial_file(session_id, file).await
518 }
519}
520
521pub use SessionFileSystem as SessionFileStore;
523
524#[derive(Clone, Default)]
530pub struct SessionFileSystemFactoryContext {
531 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
532}
533
534impl SessionFileSystemFactoryContext {
535 pub fn new() -> Self {
536 Self::default()
537 }
538
539 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
540 let values = Arc::make_mut(&mut self.values);
541 values.insert(TypeId::of::<T>(), value);
542 self
543 }
544
545 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
546 self.values
547 .get(&TypeId::of::<T>())
548 .and_then(|value| value.clone().downcast::<T>().ok())
549 }
550}
551
552#[async_trait]
554pub trait SessionFileSystemFactory: Send + Sync {
555 fn name(&self) -> &'static str {
557 "SessionFileSystemFactory"
558 }
559
560 fn is_disabled(&self) -> bool {
563 false
564 }
565
566 async fn create_session_file_system(
568 &self,
569 context: SessionFileSystemFactoryContext,
570 ) -> Result<Arc<dyn SessionFileSystem>>;
571}
572
573#[derive(Debug, Clone, Default)]
575pub struct DisabledSessionFileSystemFactory;
576
577#[async_trait]
578impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
579 fn name(&self) -> &'static str {
580 "DisabledSessionFileSystemFactory"
581 }
582
583 fn is_disabled(&self) -> bool {
584 true
585 }
586
587 async fn create_session_file_system(
588 &self,
589 _context: SessionFileSystemFactoryContext,
590 ) -> Result<Arc<dyn SessionFileSystem>> {
591 Err(crate::error::AgentLoopError::config(
592 "session filesystem is disabled",
593 ))
594 }
595}
596
597#[derive(Debug, Clone)]
603pub struct KeyInfo {
604 pub key: String,
605 pub created_at: chrono::DateTime<chrono::Utc>,
606 pub updated_at: chrono::DateTime<chrono::Utc>,
607}
608
609#[derive(Debug, Clone)]
611pub struct SecretInfo {
612 pub name: String,
613 pub created_at: chrono::DateTime<chrono::Utc>,
614 pub updated_at: chrono::DateTime<chrono::Utc>,
615}
616
617#[async_trait]
627pub trait SessionStorageStore: Send + Sync {
628 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
632
633 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
635
636 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
638
639 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
641
642 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
646
647 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
649
650 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
652
653 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
655}
656
657use crate::session_schedule::SessionSchedule;
662use crate::typed_id::ScheduleId;
663
664#[async_trait]
668pub trait SessionScheduleStore: Send + Sync {
669 async fn create_schedule(
671 &self,
672 session_id: SessionId,
673 description: String,
674 cron_expression: Option<String>,
675 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
676 timezone: String,
677 ) -> Result<SessionSchedule>;
678
679 async fn cancel_schedule(
681 &self,
682 session_id: SessionId,
683 schedule_id: ScheduleId,
684 ) -> Result<SessionSchedule>;
685
686 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
688
689 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
691}
692
693#[async_trait]
703pub trait SessionResourceRegistry: Send + Sync {
704 async fn register(
706 &self,
707 entry: crate::session_resource::RegisterSessionResource,
708 ) -> Result<crate::session_resource::SessionResourceEntry>;
709
710 async fn update_status(
712 &self,
713 session_id: SessionId,
714 resource_id: &str,
715 status: crate::session_resource::SessionResourceStatus,
716 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
717
718 async fn get(
720 &self,
721 session_id: SessionId,
722 resource_id: &str,
723 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
724
725 async fn list(
727 &self,
728 session_id: SessionId,
729 filter: Option<&crate::session_resource::SessionResourceFilter>,
730 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
731
732 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
734}
735
736#[async_trait]
746pub trait LeasedResourceStore: Send + Sync {
747 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
753
754 async fn release_resource(
760 &self,
761 session_id: SessionId,
762 provider: &str,
763 resource_type: &str,
764 external_id: &str,
765 ) -> Result<Option<LeasedResource>>;
766
767 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
772}
773
774pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
780
781#[async_trait]
786pub trait UserConnectionResolver: Send + Sync {
787 async fn get_connection_token(
790 &self,
791 session_id: SessionId,
792 provider: &str,
793 ) -> Result<Option<String>>;
794
795 async fn get_connection_user(
800 &self,
801 _session_id: SessionId,
802 _provider: &str,
803 ) -> Result<Option<Uuid>> {
804 Ok(None)
805 }
806
807 async fn get_connection_token_for_user(
812 &self,
813 _user_id: Uuid,
814 _provider: &str,
815 ) -> Result<Option<String>> {
816 Ok(None)
817 }
818
819 async fn get_connection_metadata(
822 &self,
823 _session_id: SessionId,
824 _provider: &str,
825 ) -> Result<Option<serde_json::Value>> {
826 Ok(None)
827 }
828}
829
830#[async_trait]
840pub trait BudgetChecker: Send + Sync {
841 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
843}
844
845#[async_trait]
854pub trait PaymentAuthority: Send + Sync {
855 async fn execute_machine_payment(
856 &self,
857 session_id: SessionId,
858 request: crate::payment::MachinePaymentRequest,
859 ) -> Result<crate::payment::MachinePaymentResponse>;
860}
861
862#[async_trait]
872pub trait OutboundToolRateLimiter: Send + Sync {
873 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
875}
876
877#[derive(Debug)]
883pub enum ToolCallClaimResult {
884 Claimed { claim_token: uuid::Uuid },
887 AlreadySettled {
889 result_json: serde_json::Value,
890 args_fingerprint: String,
891 },
892 AlreadyRunning { args_fingerprint: String },
897 DeterminismViolation {
901 stored_fingerprint: String,
902 current_fingerprint: String,
903 },
904}
905
906#[derive(Debug, Clone)]
908pub enum DurableToolCallStatus {
909 Settled { result_json: serde_json::Value },
911 Interrupted {
913 result_json: Option<serde_json::Value>,
914 },
915 Running,
917}
918
919#[async_trait]
924pub trait DurableToolResultStore: Send + Sync + 'static {
925 async fn try_claim_tool_call(
933 &self,
934 turn_id: &str,
935 tool_call_id: &str,
936 tool_name: &str,
937 args_fingerprint: &str,
938 ) -> Result<ToolCallClaimResult>;
939
940 async fn settle_tool_call(
946 &self,
947 turn_id: &str,
948 tool_call_id: &str,
949 result_json: serde_json::Value,
950 status: &str,
951 claim_token: uuid::Uuid,
952 ) -> Result<bool>;
953
954 async fn get_tool_call_status(
959 &self,
960 turn_id: &str,
961 tool_call_id: &str,
962 ) -> Result<Option<DurableToolCallStatus>>;
963}
964
965pub struct NoopDurableToolResultStore;
968
969#[async_trait]
970impl DurableToolResultStore for NoopDurableToolResultStore {
971 async fn try_claim_tool_call(
972 &self,
973 _turn_id: &str,
974 _tool_call_id: &str,
975 _tool_name: &str,
976 _args_fingerprint: &str,
977 ) -> Result<ToolCallClaimResult> {
978 Ok(ToolCallClaimResult::Claimed {
979 claim_token: uuid::Uuid::new_v4(),
980 })
981 }
982
983 async fn settle_tool_call(
984 &self,
985 _turn_id: &str,
986 _tool_call_id: &str,
987 _result_json: serde_json::Value,
988 _status: &str,
989 _claim_token: uuid::Uuid,
990 ) -> Result<bool> {
991 Ok(true)
992 }
993
994 async fn get_tool_call_status(
995 &self,
996 _turn_id: &str,
997 _tool_call_id: &str,
998 ) -> Result<Option<DurableToolCallStatus>> {
999 Ok(None)
1000 }
1001}
1002
1003#[derive(Debug, Clone)]
1009pub struct StreamProgress {
1010 pub accumulated_len: usize,
1012 pub last_delta_at: u64,
1014}
1015
1016#[async_trait]
1022pub trait StreamHeartbeater: Send + Sync {
1023 async fn heartbeat(&self, progress: StreamProgress);
1029}
1030
1031pub struct NoopStreamHeartbeater;
1033
1034#[async_trait]
1035impl StreamHeartbeater for NoopStreamHeartbeater {
1036 async fn heartbeat(&self, _progress: StreamProgress) {}
1037}
1038
1039#[derive(Debug, Clone)]
1045pub struct PartialStreamState {
1046 pub accumulated: String,
1049}
1050
1051#[async_trait]
1059pub trait PartialStreamStore: Send + Sync {
1060 async fn get_partial_stream(
1063 &self,
1064 session_id: SessionId,
1065 turn_id: &str,
1066 ) -> Result<Option<PartialStreamState>>;
1067}
1068
1069pub struct NoopPartialStreamStore;
1071
1072#[async_trait]
1073impl PartialStreamStore for NoopPartialStreamStore {
1074 async fn get_partial_stream(
1075 &self,
1076 _session_id: SessionId,
1077 _turn_id: &str,
1078 ) -> Result<Option<PartialStreamState>> {
1079 Ok(None)
1080 }
1081}
1082
1083#[derive(Clone)]
1092pub struct ToolContext {
1093 pub session_id: SessionId,
1095
1096 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1098
1099 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1101
1102 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1104
1105 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1107
1108 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1110
1111 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1113
1114 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1116
1117 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1119
1120 pub session_store: Option<Arc<dyn SessionStore>>,
1122
1123 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1125
1126 pub agent_store: Option<Arc<dyn AgentStore>>,
1128
1129 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1131
1132 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1134
1135 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1137 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1139
1140 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1142
1143 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1146
1147 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1150
1151 pub event_context: Option<crate::events::EventContext>,
1154
1155 pub tool_call_id: Option<String>,
1158 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1160
1161 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1164
1165 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1169
1170 pub org_id: Option<crate::typed_id::OrgId>,
1172
1173 pub network_access: Option<crate::network_access::NetworkAccessList>,
1176
1177 pub locale: Option<String>,
1181
1182 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1184
1185 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1187
1188 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1192}
1193
1194impl ToolContext {
1195 pub fn new(session_id: SessionId) -> Self {
1197 Self {
1198 session_id,
1199 file_store: None,
1200 storage_store: None,
1201 image_store: None,
1202 provider_credential_store: None,
1203 utility_llm_service: None,
1204 egress_service: None,
1205 sqldb_store: None,
1206 message_retriever: None,
1207 session_store: None,
1208 session_mutator: None,
1209 agent_store: None,
1210 connection_resolver: None,
1211 schedule_store: None,
1212 platform_store: None,
1213 leased_resource_store: None,
1214 session_resource_registry: None,
1215 session_task_registry: None,
1216 event_emitter: None,
1217 event_context: None,
1218 tool_call_id: None,
1219 capability_registry: None,
1220 tool_registry: None,
1221 visible_tool_names: None,
1222 org_id: None,
1223 network_access: None,
1224 locale: None,
1225 budget_checker: None,
1226 payment_authority: None,
1227 subagent_spawn_store: None,
1228 }
1229 }
1230
1231 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1233 Self {
1234 session_id,
1235 file_store: Some(file_store),
1236 storage_store: None,
1237 image_store: None,
1238 provider_credential_store: None,
1239 utility_llm_service: None,
1240 egress_service: None,
1241 sqldb_store: None,
1242 message_retriever: None,
1243 session_store: None,
1244 session_mutator: None,
1245 agent_store: None,
1246 connection_resolver: None,
1247 schedule_store: None,
1248 platform_store: None,
1249 leased_resource_store: None,
1250 session_resource_registry: None,
1251 session_task_registry: None,
1252 event_emitter: None,
1253 event_context: None,
1254 tool_call_id: None,
1255 capability_registry: None,
1256 tool_registry: None,
1257 visible_tool_names: None,
1258 org_id: None,
1259 network_access: None,
1260 locale: None,
1261 budget_checker: None,
1262 payment_authority: None,
1263 subagent_spawn_store: None,
1264 }
1265 }
1266
1267 pub fn with_storage_store(
1269 session_id: SessionId,
1270 storage_store: Arc<dyn SessionStorageStore>,
1271 ) -> Self {
1272 Self {
1273 session_id,
1274 file_store: None,
1275 storage_store: Some(storage_store),
1276 image_store: None,
1277 provider_credential_store: None,
1278 utility_llm_service: None,
1279 egress_service: None,
1280 sqldb_store: None,
1281 message_retriever: None,
1282 session_store: None,
1283 session_mutator: None,
1284 agent_store: None,
1285 connection_resolver: None,
1286 schedule_store: None,
1287 platform_store: None,
1288 leased_resource_store: None,
1289 session_resource_registry: None,
1290 session_task_registry: None,
1291 event_emitter: None,
1292 event_context: None,
1293 tool_call_id: None,
1294 capability_registry: None,
1295 tool_registry: None,
1296 visible_tool_names: None,
1297 org_id: None,
1298 network_access: None,
1299 locale: None,
1300 budget_checker: None,
1301 payment_authority: None,
1302 subagent_spawn_store: None,
1303 }
1304 }
1305
1306 pub fn with_stores(
1308 session_id: SessionId,
1309 file_store: Arc<dyn SessionFileSystem>,
1310 storage_store: Arc<dyn SessionStorageStore>,
1311 ) -> Self {
1312 Self {
1313 session_id,
1314 file_store: Some(file_store),
1315 storage_store: Some(storage_store),
1316 sqldb_store: None,
1317 image_store: None,
1318 provider_credential_store: None,
1319 utility_llm_service: None,
1320 egress_service: None,
1321 message_retriever: None,
1322 session_store: None,
1323 session_mutator: None,
1324 agent_store: None,
1325 connection_resolver: None,
1326 schedule_store: None,
1327 platform_store: None,
1328 leased_resource_store: None,
1329 session_resource_registry: None,
1330 session_task_registry: None,
1331 event_emitter: None,
1332 event_context: None,
1333 tool_call_id: None,
1334 capability_registry: None,
1335 tool_registry: None,
1336 visible_tool_names: None,
1337 org_id: None,
1338 network_access: None,
1339 locale: None,
1340 budget_checker: None,
1341 payment_authority: None,
1342 subagent_spawn_store: None,
1343 }
1344 }
1345
1346 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1348 self.sqldb_store = Some(sqldb_store);
1349 self
1350 }
1351
1352 pub fn with_message_retriever(
1354 mut self,
1355 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1356 ) -> Self {
1357 self.message_retriever = Some(retriever);
1358 self
1359 }
1360
1361 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1363 self.session_store = Some(store);
1364 self
1365 }
1366
1367 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1369 self.session_mutator = Some(mutator);
1370 self
1371 }
1372
1373 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1375 self.agent_store = Some(store);
1376 self
1377 }
1378
1379 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1381 self.connection_resolver = Some(resolver);
1382 self
1383 }
1384
1385 pub fn with_image_store(
1387 session_id: SessionId,
1388 image_store: Arc<dyn ImageArtifactStore>,
1389 ) -> Self {
1390 Self {
1391 session_id,
1392 file_store: None,
1393 storage_store: None,
1394 image_store: Some(image_store),
1395 provider_credential_store: None,
1396 utility_llm_service: None,
1397 egress_service: None,
1398 sqldb_store: None,
1399 message_retriever: None,
1400 session_store: None,
1401 session_mutator: None,
1402 agent_store: None,
1403 connection_resolver: None,
1404 schedule_store: None,
1405 platform_store: None,
1406 leased_resource_store: None,
1407 session_resource_registry: None,
1408 session_task_registry: None,
1409 event_emitter: None,
1410 event_context: None,
1411 tool_call_id: None,
1412 capability_registry: None,
1413 tool_registry: None,
1414 visible_tool_names: None,
1415 org_id: None,
1416 network_access: None,
1417 locale: None,
1418 budget_checker: None,
1419 payment_authority: None,
1420 subagent_spawn_store: None,
1421 }
1422 }
1423
1424 pub fn with_provider_credential_store(
1426 mut self,
1427 store: Arc<dyn ProviderCredentialStore>,
1428 ) -> Self {
1429 self.provider_credential_store = Some(store);
1430 self
1431 }
1432
1433 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1435 self.utility_llm_service = Some(service);
1436 self
1437 }
1438
1439 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1441 self.egress_service = Some(service);
1442 self
1443 }
1444
1445 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1447 self.schedule_store = Some(store);
1448 self
1449 }
1450
1451 pub fn with_platform_store(
1453 mut self,
1454 store: Arc<dyn crate::platform_store::PlatformStore>,
1455 ) -> Self {
1456 self.platform_store = Some(store);
1457 self
1458 }
1459
1460 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1462 self.leased_resource_store = Some(store);
1463 self
1464 }
1465
1466 pub fn with_session_resource_registry(
1468 mut self,
1469 registry: Arc<dyn SessionResourceRegistry>,
1470 ) -> Self {
1471 self.session_resource_registry = Some(registry);
1472 self
1473 }
1474
1475 pub fn with_session_task_registry(
1477 mut self,
1478 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
1479 ) -> Self {
1480 self.session_task_registry = Some(registry);
1481 self
1482 }
1483
1484 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
1486 self.org_id = Some(org_id);
1487 self
1488 }
1489
1490 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
1492 self.tool_registry = Some(registry);
1493 self
1494 }
1495
1496 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
1498 self.visible_tool_names = Some(names);
1499 self
1500 }
1501
1502 pub fn with_network_access(
1504 mut self,
1505 network_access: Option<crate::network_access::NetworkAccessList>,
1506 ) -> Self {
1507 self.network_access = network_access;
1508 self
1509 }
1510
1511 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
1513 self.payment_authority = Some(authority);
1514 self
1515 }
1516
1517 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
1519 self.subagent_spawn_store = Some(store);
1520 self
1521 }
1522
1523 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
1528 let (Some(emitter), Some(ctx), Some(call_id)) =
1529 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1530 else {
1531 return;
1532 };
1533 if let Err(e) = emitter
1534 .emit(EventRequest::new(
1535 self.session_id,
1536 ctx.clone(),
1537 crate::events::ToolProgressData {
1538 tool_call_id: call_id.clone(),
1539 tool_name: tool_name.to_string(),
1540 message: message.to_string(),
1541 display_name: None,
1542 },
1543 ))
1544 .await
1545 {
1546 tracing::debug!(
1547 tool_call_id = call_id,
1548 tool_name,
1549 error = %e,
1550 "Failed to emit tool.progress event"
1551 );
1552 }
1553 }
1554
1555 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
1560 let (Some(emitter), Some(ctx), Some(call_id)) =
1561 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1562 else {
1563 return;
1564 };
1565 if let Err(e) = emitter
1566 .emit(EventRequest::new(
1567 self.session_id,
1568 ctx.clone(),
1569 crate::events::ToolOutputDeltaData {
1570 tool_call_id: call_id.clone(),
1571 tool_name: tool_name.to_string(),
1572 delta: delta.to_string(),
1573 stream: stream.to_string(),
1574 },
1575 ))
1576 .await
1577 {
1578 tracing::debug!(
1579 tool_call_id = call_id,
1580 tool_name,
1581 error = %e,
1582 "Failed to emit tool.output.delta event"
1583 );
1584 }
1585 }
1586}
1587
1588impl std::fmt::Debug for ToolContext {
1589 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1590 f.debug_struct("ToolContext")
1591 .field("session_id", &self.session_id)
1592 .field("file_store", &self.file_store.is_some())
1593 .field("storage_store", &self.storage_store.is_some())
1594 .field("image_store", &self.image_store.is_some())
1595 .field(
1596 "provider_credential_store",
1597 &self.provider_credential_store.is_some(),
1598 )
1599 .field("utility_llm_service", &self.utility_llm_service.is_some())
1600 .field("egress_service", &self.egress_service.is_some())
1601 .field("sqldb_store", &self.sqldb_store.is_some())
1602 .field("message_retriever", &self.message_retriever.is_some())
1603 .field("session_store", &self.session_store.is_some())
1604 .field("session_mutator", &self.session_mutator.is_some())
1605 .field("agent_store", &self.agent_store.is_some())
1606 .field("connection_resolver", &self.connection_resolver.is_some())
1607 .field("schedule_store", &self.schedule_store.is_some())
1608 .field("platform_store", &self.platform_store.is_some())
1609 .field(
1610 "leased_resource_store",
1611 &self.leased_resource_store.is_some(),
1612 )
1613 .field("event_emitter", &self.event_emitter.is_some())
1614 .field("tool_registry", &self.tool_registry.is_some())
1615 .field("payment_authority", &self.payment_authority.is_some())
1616 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
1617 .field("org_id", &self.org_id)
1618 .finish()
1619 }
1620}
1621
1622use crate::events::{Event, EventRequest};
1627
1628#[async_trait]
1639pub trait EventEmitter: Send + Sync {
1640 async fn emit(&self, request: EventRequest) -> Result<Event>;
1645}
1646
1647#[async_trait]
1649impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
1650 async fn emit(&self, request: EventRequest) -> Result<Event> {
1651 (**self).emit(request).await
1652 }
1653}
1654
1655#[derive(Debug, Clone, Default)]
1659pub struct NoopEventEmitter;
1660
1661#[async_trait]
1662impl EventEmitter for NoopEventEmitter {
1663 async fn emit(&self, request: EventRequest) -> Result<Event> {
1664 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
1666 }
1667}
1668
1669#[derive(Debug, Clone)]
1682pub struct ResolvedImage {
1683 pub base64: String,
1685 pub media_type: String,
1687}
1688
1689impl ResolvedImage {
1690 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
1692 Self {
1693 base64: base64.into(),
1694 media_type: media_type.into(),
1695 }
1696 }
1697
1698 pub fn to_data_url(&self) -> String {
1702 format!("data:{};base64,{}", self.media_type, self.base64)
1703 }
1704}
1705
1706#[async_trait]
1739pub trait ImageResolver: Send + Sync {
1740 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
1744}
1745
1746#[derive(Debug)]
1752pub enum SpawnClaimResult {
1753 Claimed {
1756 spawn_handle_id: uuid::Uuid,
1757 claim_token: uuid::Uuid,
1758 },
1759 ClaimedPendingChild {
1763 spawn_handle_id: uuid::Uuid,
1764 claim_token: uuid::Uuid,
1765 },
1766 AlreadyRunning {
1769 child_session_id: crate::typed_id::SessionId,
1770 claim_token: uuid::Uuid,
1772 },
1773 AlreadySettled {
1776 child_session_id: crate::typed_id::SessionId,
1777 terminal_status: String,
1779 terminal_result: String,
1780 },
1781}
1782
1783#[async_trait]
1791pub trait SubagentSpawnStore: Send + Sync + 'static {
1792 async fn try_claim_spawn(
1797 &self,
1798 parent_session_id: crate::typed_id::SessionId,
1799 tool_call_id: &str,
1800 subagent_name: &str,
1801 subagent_task: &str,
1802 claim_token: uuid::Uuid,
1803 ) -> Result<SpawnClaimResult>;
1804
1805 async fn register_child_session(
1810 &self,
1811 spawn_handle_id: uuid::Uuid,
1812 claim_token: uuid::Uuid,
1813 child_session_id: crate::typed_id::SessionId,
1814 ) -> Result<()>;
1815
1816 async fn settle_spawn(
1822 &self,
1823 parent_session_id: crate::typed_id::SessionId,
1824 tool_call_id: &str,
1825 claim_token: uuid::Uuid,
1826 terminal_status: &str,
1827 terminal_result: &str,
1828 ) -> Result<()>;
1829}
1830
1831#[async_trait]
1833impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
1834 async fn try_claim_spawn(
1835 &self,
1836 parent_session_id: crate::typed_id::SessionId,
1837 tool_call_id: &str,
1838 subagent_name: &str,
1839 subagent_task: &str,
1840 claim_token: uuid::Uuid,
1841 ) -> Result<SpawnClaimResult> {
1842 (**self)
1843 .try_claim_spawn(
1844 parent_session_id,
1845 tool_call_id,
1846 subagent_name,
1847 subagent_task,
1848 claim_token,
1849 )
1850 .await
1851 }
1852
1853 async fn register_child_session(
1854 &self,
1855 spawn_handle_id: uuid::Uuid,
1856 claim_token: uuid::Uuid,
1857 child_session_id: crate::typed_id::SessionId,
1858 ) -> Result<()> {
1859 (**self)
1860 .register_child_session(spawn_handle_id, claim_token, child_session_id)
1861 .await
1862 }
1863
1864 async fn settle_spawn(
1865 &self,
1866 parent_session_id: crate::typed_id::SessionId,
1867 tool_call_id: &str,
1868 claim_token: uuid::Uuid,
1869 terminal_status: &str,
1870 terminal_result: &str,
1871 ) -> Result<()> {
1872 (**self)
1873 .settle_spawn(
1874 parent_session_id,
1875 tool_call_id,
1876 claim_token,
1877 terminal_status,
1878 terminal_result,
1879 )
1880 .await
1881 }
1882}
1883
1884pub struct NoopSubagentSpawnStore;
1888
1889#[async_trait]
1890impl SubagentSpawnStore for NoopSubagentSpawnStore {
1891 async fn try_claim_spawn(
1892 &self,
1893 _parent_session_id: crate::typed_id::SessionId,
1894 _tool_call_id: &str,
1895 _subagent_name: &str,
1896 _subagent_task: &str,
1897 claim_token: uuid::Uuid,
1898 ) -> Result<SpawnClaimResult> {
1899 Ok(SpawnClaimResult::Claimed {
1900 spawn_handle_id: uuid::Uuid::new_v4(),
1901 claim_token,
1902 })
1903 }
1904
1905 async fn register_child_session(
1906 &self,
1907 _spawn_handle_id: uuid::Uuid,
1908 _claim_token: uuid::Uuid,
1909 _child_session_id: crate::typed_id::SessionId,
1910 ) -> Result<()> {
1911 Ok(())
1912 }
1913
1914 async fn settle_spawn(
1915 &self,
1916 _parent_session_id: crate::typed_id::SessionId,
1917 _tool_call_id: &str,
1918 _claim_token: uuid::Uuid,
1919 _terminal_status: &str,
1920 _terminal_result: &str,
1921 ) -> Result<()> {
1922 Ok(())
1923 }
1924}
1925
1926#[cfg(test)]
1931mod tests {
1932 use super::*;
1933
1934 #[test]
1935 fn test_resolved_image_new() {
1936 let image = ResolvedImage::new("SGVsbG8=", "image/png");
1937 assert_eq!(image.base64, "SGVsbG8=");
1938 assert_eq!(image.media_type, "image/png");
1939 }
1940
1941 #[test]
1942 fn test_resolved_image_to_data_url() {
1943 let image = ResolvedImage::new("SGVsbG8=", "image/png");
1944 let data_url = image.to_data_url();
1945 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
1946 }
1947
1948 #[test]
1949 fn test_resolved_image_jpeg() {
1950 let image = ResolvedImage::new("base64data", "image/jpeg");
1951 let data_url = image.to_data_url();
1952 assert!(data_url.starts_with("data:image/jpeg;base64,"));
1953 }
1954}