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;
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]
340pub trait SessionFileSystem: Send + Sync {
341 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>>;
343
344 async fn write_file(
346 &self,
347 session_id: SessionId,
348 path: &str,
349 content: &str,
350 encoding: &str,
351 ) -> Result<SessionFile>;
352
353 async fn write_file_if_content_matches(
358 &self,
359 session_id: SessionId,
360 path: &str,
361 expected_content: &str,
362 expected_encoding: &str,
363 content: &str,
364 encoding: &str,
365 ) -> Result<Option<SessionFile>> {
366 let Some(existing) = self.read_file(session_id, path).await? else {
367 return Ok(None);
368 };
369
370 if existing.is_directory {
371 return Ok(None);
372 }
373
374 let current_content = existing.content.unwrap_or_default();
375 if current_content != expected_content || existing.encoding != expected_encoding {
376 return Ok(None);
377 }
378
379 self.write_file(session_id, path, content, encoding)
380 .await
381 .map(Some)
382 }
383
384 async fn delete_file(&self, session_id: SessionId, path: &str, recursive: bool)
386 -> Result<bool>;
387
388 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>>;
390
391 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>>;
393
394 async fn grep_files(
396 &self,
397 session_id: SessionId,
398 pattern: &str,
399 path_pattern: Option<&str>,
400 ) -> Result<Vec<GrepMatch>>;
401
402 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
404
405 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
407 if file.is_readonly {
408 return Err(crate::error::AgentLoopError::store(
409 "read-only initial files require a SessionFileSystem-specific seed implementation",
410 ));
411 }
412 self.write_file(session_id, &file.path, &file.content, &file.encoding)
413 .await?;
414 Ok(())
415 }
416}
417
418#[async_trait]
419impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
420 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
421 (**self).read_file(session_id, path).await
422 }
423
424 async fn write_file(
425 &self,
426 session_id: SessionId,
427 path: &str,
428 content: &str,
429 encoding: &str,
430 ) -> Result<SessionFile> {
431 (**self)
432 .write_file(session_id, path, content, encoding)
433 .await
434 }
435
436 async fn write_file_if_content_matches(
437 &self,
438 session_id: SessionId,
439 path: &str,
440 expected_content: &str,
441 expected_encoding: &str,
442 content: &str,
443 encoding: &str,
444 ) -> Result<Option<SessionFile>> {
445 (**self)
446 .write_file_if_content_matches(
447 session_id,
448 path,
449 expected_content,
450 expected_encoding,
451 content,
452 encoding,
453 )
454 .await
455 }
456
457 async fn delete_file(
458 &self,
459 session_id: SessionId,
460 path: &str,
461 recursive: bool,
462 ) -> Result<bool> {
463 (**self).delete_file(session_id, path, recursive).await
464 }
465
466 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
467 (**self).list_directory(session_id, path).await
468 }
469
470 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
471 (**self).stat_file(session_id, path).await
472 }
473
474 async fn grep_files(
475 &self,
476 session_id: SessionId,
477 pattern: &str,
478 path_pattern: Option<&str>,
479 ) -> Result<Vec<GrepMatch>> {
480 (**self).grep_files(session_id, pattern, path_pattern).await
481 }
482
483 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
484 (**self).create_directory(session_id, path).await
485 }
486
487 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
488 (**self).seed_initial_file(session_id, file).await
489 }
490}
491
492pub use SessionFileSystem as SessionFileStore;
494
495#[derive(Clone, Default)]
501pub struct SessionFileSystemFactoryContext {
502 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
503}
504
505impl SessionFileSystemFactoryContext {
506 pub fn new() -> Self {
507 Self::default()
508 }
509
510 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
511 let values = Arc::make_mut(&mut self.values);
512 values.insert(TypeId::of::<T>(), value);
513 self
514 }
515
516 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
517 self.values
518 .get(&TypeId::of::<T>())
519 .and_then(|value| value.clone().downcast::<T>().ok())
520 }
521}
522
523#[async_trait]
525pub trait SessionFileSystemFactory: Send + Sync {
526 fn name(&self) -> &'static str {
528 "SessionFileSystemFactory"
529 }
530
531 fn is_disabled(&self) -> bool {
534 false
535 }
536
537 async fn create_session_file_system(
539 &self,
540 context: SessionFileSystemFactoryContext,
541 ) -> Result<Arc<dyn SessionFileSystem>>;
542}
543
544#[derive(Debug, Clone, Default)]
546pub struct DisabledSessionFileSystemFactory;
547
548#[async_trait]
549impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
550 fn name(&self) -> &'static str {
551 "DisabledSessionFileSystemFactory"
552 }
553
554 fn is_disabled(&self) -> bool {
555 true
556 }
557
558 async fn create_session_file_system(
559 &self,
560 _context: SessionFileSystemFactoryContext,
561 ) -> Result<Arc<dyn SessionFileSystem>> {
562 Err(crate::error::AgentLoopError::config(
563 "session filesystem is disabled",
564 ))
565 }
566}
567
568#[derive(Debug, Clone)]
574pub struct KeyInfo {
575 pub key: String,
576 pub created_at: chrono::DateTime<chrono::Utc>,
577 pub updated_at: chrono::DateTime<chrono::Utc>,
578}
579
580#[derive(Debug, Clone)]
582pub struct SecretInfo {
583 pub name: String,
584 pub created_at: chrono::DateTime<chrono::Utc>,
585 pub updated_at: chrono::DateTime<chrono::Utc>,
586}
587
588#[async_trait]
598pub trait SessionStorageStore: Send + Sync {
599 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
603
604 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
606
607 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
609
610 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
612
613 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
617
618 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
620
621 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
623
624 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
626}
627
628use crate::session_schedule::SessionSchedule;
633use crate::typed_id::ScheduleId;
634
635#[async_trait]
639pub trait SessionScheduleStore: Send + Sync {
640 async fn create_schedule(
642 &self,
643 session_id: SessionId,
644 description: String,
645 cron_expression: Option<String>,
646 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
647 timezone: String,
648 ) -> Result<SessionSchedule>;
649
650 async fn cancel_schedule(
652 &self,
653 session_id: SessionId,
654 schedule_id: ScheduleId,
655 ) -> Result<SessionSchedule>;
656
657 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
659
660 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
662}
663
664#[async_trait]
674pub trait SessionResourceRegistry: Send + Sync {
675 async fn register(
677 &self,
678 entry: crate::session_resource::RegisterSessionResource,
679 ) -> Result<crate::session_resource::SessionResourceEntry>;
680
681 async fn update_status(
683 &self,
684 session_id: SessionId,
685 resource_id: &str,
686 status: crate::session_resource::SessionResourceStatus,
687 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
688
689 async fn get(
691 &self,
692 session_id: SessionId,
693 resource_id: &str,
694 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
695
696 async fn list(
698 &self,
699 session_id: SessionId,
700 filter: Option<&crate::session_resource::SessionResourceFilter>,
701 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
702
703 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
705}
706
707#[async_trait]
717pub trait LeasedResourceStore: Send + Sync {
718 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
724
725 async fn release_resource(
731 &self,
732 session_id: SessionId,
733 provider: &str,
734 resource_type: &str,
735 external_id: &str,
736 ) -> Result<Option<LeasedResource>>;
737
738 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
743}
744
745pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
751
752#[async_trait]
757pub trait UserConnectionResolver: Send + Sync {
758 async fn get_connection_token(
761 &self,
762 session_id: SessionId,
763 provider: &str,
764 ) -> Result<Option<String>>;
765
766 async fn get_connection_user(
771 &self,
772 _session_id: SessionId,
773 _provider: &str,
774 ) -> Result<Option<Uuid>> {
775 Ok(None)
776 }
777
778 async fn get_connection_token_for_user(
783 &self,
784 _user_id: Uuid,
785 _provider: &str,
786 ) -> Result<Option<String>> {
787 Ok(None)
788 }
789
790 async fn get_connection_metadata(
793 &self,
794 _session_id: SessionId,
795 _provider: &str,
796 ) -> Result<Option<serde_json::Value>> {
797 Ok(None)
798 }
799}
800
801#[async_trait]
811pub trait BudgetChecker: Send + Sync {
812 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
814}
815
816#[async_trait]
825pub trait PaymentAuthority: Send + Sync {
826 async fn execute_machine_payment(
827 &self,
828 session_id: SessionId,
829 request: crate::payment::MachinePaymentRequest,
830 ) -> Result<crate::payment::MachinePaymentResponse>;
831}
832
833#[derive(Clone)]
842pub struct ToolContext {
843 pub session_id: SessionId,
845
846 pub file_store: Option<Arc<dyn SessionFileSystem>>,
848
849 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
851
852 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
854
855 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
857
858 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
860
861 pub egress_service: Option<Arc<dyn crate::EgressService>>,
863
864 pub sqldb_store: Option<SessionSqlDbStoreRef>,
866
867 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
869
870 pub session_store: Option<Arc<dyn SessionStore>>,
872
873 pub session_mutator: Option<Arc<dyn SessionMutator>>,
875
876 pub agent_store: Option<Arc<dyn AgentStore>>,
878
879 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
881
882 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
884
885 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
887 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
889
890 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
892
893 pub event_emitter: Option<Arc<dyn EventEmitter>>,
896
897 pub event_context: Option<crate::events::EventContext>,
900
901 pub tool_call_id: Option<String>,
904 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
906
907 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
910
911 pub memory_store: Option<Arc<dyn crate::memory_store::MemoryStoreBackend>>,
913
914 pub org_id: Option<crate::typed_id::OrgId>,
916
917 pub network_access: Option<crate::network_access::NetworkAccessList>,
920
921 pub locale: Option<String>,
925
926 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
928
929 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
931}
932
933impl ToolContext {
934 pub fn new(session_id: SessionId) -> Self {
936 Self {
937 session_id,
938 file_store: None,
939 storage_store: None,
940 image_store: None,
941 provider_credential_store: None,
942 utility_llm_service: None,
943 egress_service: None,
944 sqldb_store: None,
945 message_retriever: None,
946 session_store: None,
947 session_mutator: None,
948 agent_store: None,
949 connection_resolver: None,
950 schedule_store: None,
951 platform_store: None,
952 leased_resource_store: None,
953 session_resource_registry: None,
954 event_emitter: None,
955 event_context: None,
956 tool_call_id: None,
957 capability_registry: None,
958 tool_registry: None,
959 memory_store: None,
960 org_id: None,
961 network_access: None,
962 locale: None,
963 budget_checker: None,
964 payment_authority: None,
965 }
966 }
967
968 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
970 Self {
971 session_id,
972 file_store: Some(file_store),
973 storage_store: None,
974 image_store: None,
975 provider_credential_store: None,
976 utility_llm_service: None,
977 egress_service: None,
978 sqldb_store: None,
979 message_retriever: None,
980 session_store: None,
981 session_mutator: None,
982 agent_store: None,
983 connection_resolver: None,
984 schedule_store: None,
985 platform_store: None,
986 leased_resource_store: None,
987 session_resource_registry: None,
988 event_emitter: None,
989 event_context: None,
990 tool_call_id: None,
991 capability_registry: None,
992 tool_registry: None,
993 memory_store: None,
994 org_id: None,
995 network_access: None,
996 locale: None,
997 budget_checker: None,
998 payment_authority: None,
999 }
1000 }
1001
1002 pub fn with_storage_store(
1004 session_id: SessionId,
1005 storage_store: Arc<dyn SessionStorageStore>,
1006 ) -> Self {
1007 Self {
1008 session_id,
1009 file_store: None,
1010 storage_store: Some(storage_store),
1011 image_store: None,
1012 provider_credential_store: None,
1013 utility_llm_service: None,
1014 egress_service: None,
1015 sqldb_store: None,
1016 message_retriever: None,
1017 session_store: None,
1018 session_mutator: None,
1019 agent_store: None,
1020 connection_resolver: None,
1021 schedule_store: None,
1022 platform_store: None,
1023 leased_resource_store: None,
1024 session_resource_registry: None,
1025 event_emitter: None,
1026 event_context: None,
1027 tool_call_id: None,
1028 capability_registry: None,
1029 tool_registry: None,
1030 memory_store: None,
1031 org_id: None,
1032 network_access: None,
1033 locale: None,
1034 budget_checker: None,
1035 payment_authority: None,
1036 }
1037 }
1038
1039 pub fn with_stores(
1041 session_id: SessionId,
1042 file_store: Arc<dyn SessionFileSystem>,
1043 storage_store: Arc<dyn SessionStorageStore>,
1044 ) -> Self {
1045 Self {
1046 session_id,
1047 file_store: Some(file_store),
1048 storage_store: Some(storage_store),
1049 sqldb_store: None,
1050 image_store: None,
1051 provider_credential_store: None,
1052 utility_llm_service: None,
1053 egress_service: None,
1054 message_retriever: None,
1055 session_store: None,
1056 session_mutator: None,
1057 agent_store: None,
1058 connection_resolver: None,
1059 schedule_store: None,
1060 platform_store: None,
1061 leased_resource_store: None,
1062 session_resource_registry: None,
1063 event_emitter: None,
1064 event_context: None,
1065 tool_call_id: None,
1066 capability_registry: None,
1067 tool_registry: None,
1068 memory_store: None,
1069 org_id: None,
1070 network_access: None,
1071 locale: None,
1072 budget_checker: None,
1073 payment_authority: None,
1074 }
1075 }
1076
1077 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1079 self.sqldb_store = Some(sqldb_store);
1080 self
1081 }
1082
1083 pub fn with_message_retriever(
1085 mut self,
1086 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1087 ) -> Self {
1088 self.message_retriever = Some(retriever);
1089 self
1090 }
1091
1092 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1094 self.session_store = Some(store);
1095 self
1096 }
1097
1098 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1100 self.session_mutator = Some(mutator);
1101 self
1102 }
1103
1104 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1106 self.agent_store = Some(store);
1107 self
1108 }
1109
1110 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1112 self.connection_resolver = Some(resolver);
1113 self
1114 }
1115
1116 pub fn with_image_store(
1118 session_id: SessionId,
1119 image_store: Arc<dyn ImageArtifactStore>,
1120 ) -> Self {
1121 Self {
1122 session_id,
1123 file_store: None,
1124 storage_store: None,
1125 image_store: Some(image_store),
1126 provider_credential_store: None,
1127 utility_llm_service: None,
1128 egress_service: None,
1129 sqldb_store: None,
1130 message_retriever: None,
1131 session_store: None,
1132 session_mutator: None,
1133 agent_store: None,
1134 connection_resolver: None,
1135 schedule_store: None,
1136 platform_store: None,
1137 leased_resource_store: None,
1138 session_resource_registry: None,
1139 event_emitter: None,
1140 event_context: None,
1141 tool_call_id: None,
1142 capability_registry: None,
1143 tool_registry: None,
1144 memory_store: None,
1145 org_id: None,
1146 network_access: None,
1147 locale: None,
1148 budget_checker: None,
1149 payment_authority: None,
1150 }
1151 }
1152
1153 pub fn with_provider_credential_store(
1155 mut self,
1156 store: Arc<dyn ProviderCredentialStore>,
1157 ) -> Self {
1158 self.provider_credential_store = Some(store);
1159 self
1160 }
1161
1162 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1164 self.utility_llm_service = Some(service);
1165 self
1166 }
1167
1168 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1170 self.egress_service = Some(service);
1171 self
1172 }
1173
1174 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1176 self.schedule_store = Some(store);
1177 self
1178 }
1179
1180 pub fn with_platform_store(
1182 mut self,
1183 store: Arc<dyn crate::platform_store::PlatformStore>,
1184 ) -> Self {
1185 self.platform_store = Some(store);
1186 self
1187 }
1188
1189 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1191 self.leased_resource_store = Some(store);
1192 self
1193 }
1194
1195 pub fn with_session_resource_registry(
1197 mut self,
1198 registry: Arc<dyn SessionResourceRegistry>,
1199 ) -> Self {
1200 self.session_resource_registry = Some(registry);
1201 self
1202 }
1203
1204 pub fn with_memory_store(
1206 mut self,
1207 store: Arc<dyn crate::memory_store::MemoryStoreBackend>,
1208 ) -> Self {
1209 self.memory_store = Some(store);
1210 self
1211 }
1212
1213 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
1215 self.org_id = Some(org_id);
1216 self
1217 }
1218
1219 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
1221 self.tool_registry = Some(registry);
1222 self
1223 }
1224
1225 pub fn with_network_access(
1227 mut self,
1228 network_access: Option<crate::network_access::NetworkAccessList>,
1229 ) -> Self {
1230 self.network_access = network_access;
1231 self
1232 }
1233
1234 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
1236 self.payment_authority = Some(authority);
1237 self
1238 }
1239
1240 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
1245 let (Some(emitter), Some(ctx), Some(call_id)) =
1246 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1247 else {
1248 return;
1249 };
1250 if let Err(e) = emitter
1251 .emit(EventRequest::new(
1252 self.session_id,
1253 ctx.clone(),
1254 crate::events::ToolProgressData {
1255 tool_call_id: call_id.clone(),
1256 tool_name: tool_name.to_string(),
1257 message: message.to_string(),
1258 display_name: None,
1259 },
1260 ))
1261 .await
1262 {
1263 tracing::debug!(
1264 tool_call_id = call_id,
1265 tool_name,
1266 error = %e,
1267 "Failed to emit tool.progress event"
1268 );
1269 }
1270 }
1271
1272 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
1277 let (Some(emitter), Some(ctx), Some(call_id)) =
1278 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1279 else {
1280 return;
1281 };
1282 if let Err(e) = emitter
1283 .emit(EventRequest::new(
1284 self.session_id,
1285 ctx.clone(),
1286 crate::events::ToolOutputDeltaData {
1287 tool_call_id: call_id.clone(),
1288 tool_name: tool_name.to_string(),
1289 delta: delta.to_string(),
1290 stream: stream.to_string(),
1291 },
1292 ))
1293 .await
1294 {
1295 tracing::debug!(
1296 tool_call_id = call_id,
1297 tool_name,
1298 error = %e,
1299 "Failed to emit tool.output.delta event"
1300 );
1301 }
1302 }
1303}
1304
1305impl std::fmt::Debug for ToolContext {
1306 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1307 f.debug_struct("ToolContext")
1308 .field("session_id", &self.session_id)
1309 .field("file_store", &self.file_store.is_some())
1310 .field("storage_store", &self.storage_store.is_some())
1311 .field("image_store", &self.image_store.is_some())
1312 .field(
1313 "provider_credential_store",
1314 &self.provider_credential_store.is_some(),
1315 )
1316 .field("utility_llm_service", &self.utility_llm_service.is_some())
1317 .field("egress_service", &self.egress_service.is_some())
1318 .field("sqldb_store", &self.sqldb_store.is_some())
1319 .field("message_retriever", &self.message_retriever.is_some())
1320 .field("session_store", &self.session_store.is_some())
1321 .field("session_mutator", &self.session_mutator.is_some())
1322 .field("agent_store", &self.agent_store.is_some())
1323 .field("connection_resolver", &self.connection_resolver.is_some())
1324 .field("schedule_store", &self.schedule_store.is_some())
1325 .field("platform_store", &self.platform_store.is_some())
1326 .field(
1327 "leased_resource_store",
1328 &self.leased_resource_store.is_some(),
1329 )
1330 .field("event_emitter", &self.event_emitter.is_some())
1331 .field("tool_registry", &self.tool_registry.is_some())
1332 .field("memory_store", &self.memory_store.is_some())
1333 .field("payment_authority", &self.payment_authority.is_some())
1334 .field("org_id", &self.org_id)
1335 .finish()
1336 }
1337}
1338
1339use crate::events::{Event, EventRequest};
1344
1345#[async_trait]
1356pub trait EventEmitter: Send + Sync {
1357 async fn emit(&self, request: EventRequest) -> Result<Event>;
1362}
1363
1364#[async_trait]
1366impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
1367 async fn emit(&self, request: EventRequest) -> Result<Event> {
1368 (**self).emit(request).await
1369 }
1370}
1371
1372#[derive(Debug, Clone, Default)]
1376pub struct NoopEventEmitter;
1377
1378#[async_trait]
1379impl EventEmitter for NoopEventEmitter {
1380 async fn emit(&self, request: EventRequest) -> Result<Event> {
1381 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
1383 }
1384}
1385
1386#[derive(Debug, Clone)]
1399pub struct ResolvedImage {
1400 pub base64: String,
1402 pub media_type: String,
1404}
1405
1406impl ResolvedImage {
1407 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
1409 Self {
1410 base64: base64.into(),
1411 media_type: media_type.into(),
1412 }
1413 }
1414
1415 pub fn to_data_url(&self) -> String {
1419 format!("data:{};base64,{}", self.media_type, self.base64)
1420 }
1421}
1422
1423#[async_trait]
1456pub trait ImageResolver: Send + Sync {
1457 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
1461}
1462
1463#[cfg(test)]
1468mod tests {
1469 use super::*;
1470
1471 #[test]
1472 fn test_resolved_image_new() {
1473 let image = ResolvedImage::new("SGVsbG8=", "image/png");
1474 assert_eq!(image.base64, "SGVsbG8=");
1475 assert_eq!(image.media_type, "image/png");
1476 }
1477
1478 #[test]
1479 fn test_resolved_image_to_data_url() {
1480 let image = ResolvedImage::new("SGVsbG8=", "image/png");
1481 let data_url = image.to_data_url();
1482 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
1483 }
1484
1485 #[test]
1486 fn test_resolved_image_jpeg() {
1487 let image = ResolvedImage::new("base64data", "image/jpeg");
1488 let data_url = image.to_data_url();
1489 assert!(data_url.starts_with("data:image/jpeg;base64,"));
1490 }
1491}