1use crate::agent::Agent;
9use crate::harness::Harness;
10use crate::provider::DriverId;
11use crate::session_file::{
12 FileInfo, FileStat, GrepMatch, GrepOptions, GrepSearchResult, InitialFile, SessionFile,
13};
14use crate::tool_types::{ToolCall, ToolDefinition, ToolResult};
15use crate::typed_id::{AgentId, HarnessId, ImageId, MessageId, ModelId, SessionId, WorkspaceId};
16use async_trait::async_trait;
17use chrono::{DateTime, Utc};
18use std::any::{Any, TypeId};
19use std::collections::{HashMap, HashSet};
20use std::sync::Arc;
21use uuid::Uuid;
22
23fn build_tool_map(tool_defs: &[ToolDefinition]) -> HashMap<&str, &ToolDefinition> {
25 tool_defs.iter().map(|def| (def.name(), def)).collect()
26}
27
28use crate::error::Result;
29
30#[derive(Clone, Default)]
45pub struct ReasoningEffortHandle {
46 inner: Arc<std::sync::RwLock<Option<String>>>,
47}
48
49impl ReasoningEffortHandle {
50 pub fn new() -> Self {
52 Self::default()
53 }
54
55 pub fn with_effort(effort: impl Into<String>) -> Self {
57 Self {
58 inner: Arc::new(std::sync::RwLock::new(Some(effort.into()))),
59 }
60 }
61
62 pub fn set(&self, effort: Option<String>) {
66 let mut guard = self.inner.write().unwrap_or_else(|e| e.into_inner());
70 *guard = effort;
71 }
72
73 pub fn get(&self) -> Option<String> {
75 let guard = self.inner.read().unwrap_or_else(|e| e.into_inner());
78 guard.clone()
79 }
80}
81
82impl std::fmt::Debug for ReasoningEffortHandle {
83 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84 f.debug_struct("ReasoningEffortHandle")
85 .field("effort", &self.get())
86 .finish()
87 }
88}
89
90#[async_trait]
101pub trait AgentStore: Send + Sync {
102 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>>;
104}
105
106#[async_trait]
107impl<T: AgentStore + ?Sized> AgentStore for std::sync::Arc<T> {
108 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>> {
109 (**self).get_agent(agent_id).await
110 }
111}
112
113#[async_trait]
128pub trait HarnessStore: Send + Sync {
129 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>>;
134}
135
136#[async_trait]
137impl<T: HarnessStore + ?Sized> HarnessStore for std::sync::Arc<T> {
138 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>> {
139 (**self).get_harness_chain(harness_id).await
140 }
141}
142
143use crate::capability_types::AgentCapabilityConfig;
148use crate::leased_resource::{LeasedResource, UpsertLeasedResource};
149use crate::session::Session;
150
151#[async_trait]
157pub trait SessionStore: Send + Sync {
158 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>>;
160}
161
162#[async_trait]
163impl<T: SessionStore + ?Sized> SessionStore for std::sync::Arc<T> {
164 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>> {
165 (**self).get_session(session_id).await
166 }
167}
168
169#[async_trait]
171pub trait SessionMutator: Send + Sync {
172 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session>;
174
175 async fn upsert_session_capability(
181 &self,
182 _session_id: SessionId,
183 _capability: AgentCapabilityConfig,
184 ) -> Result<Session> {
185 Err(crate::error::AgentLoopError::store(
186 "session backend does not support live capability reconfiguration",
187 ))
188 }
189
190 async fn remove_session_capability(
192 &self,
193 _session_id: SessionId,
194 _capability_id: &str,
195 ) -> Result<Session> {
196 Err(crate::error::AgentLoopError::store(
197 "session backend does not support live capability reconfiguration",
198 ))
199 }
200}
201
202#[async_trait]
203impl<T: SessionMutator + ?Sized> SessionMutator for std::sync::Arc<T> {
204 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session> {
205 (**self).update_session_title(session_id, title).await
206 }
207
208 async fn upsert_session_capability(
209 &self,
210 session_id: SessionId,
211 capability: AgentCapabilityConfig,
212 ) -> Result<Session> {
213 (**self)
214 .upsert_session_capability(session_id, capability)
215 .await
216 }
217
218 async fn remove_session_capability(
219 &self,
220 session_id: SessionId,
221 capability_id: &str,
222 ) -> Result<Session> {
223 (**self)
224 .remove_session_capability(session_id, capability_id)
225 .await
226 }
227}
228
229#[derive(Debug, Clone)]
235pub struct ResolvedModel {
236 pub model: String,
238 pub provider_type: DriverId,
240 pub api_key: Option<String>,
242 pub base_url: Option<String>,
244 pub provider_metadata: Option<crate::driver_registry::ProviderMetadata>,
247}
248
249#[async_trait]
259pub trait ProviderStore: Send + Sync {
260 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>>;
265
266 async fn get_default_model(&self) -> Result<Option<ResolvedModel>>;
270}
271
272#[async_trait]
273impl<T: ProviderStore + ?Sized> ProviderStore for std::sync::Arc<T> {
274 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>> {
275 (**self).get_resolved_model(model_id).await
276 }
277
278 async fn get_default_model(&self) -> Result<Option<ResolvedModel>> {
279 (**self).get_default_model().await
280 }
281}
282
283#[derive(Debug, Clone)]
289pub struct StoredImageInfo {
290 pub id: ImageId,
291 pub filename: String,
292 pub content_type: String,
293 pub size_bytes: i64,
294 pub metadata: serde_json::Value,
295 pub created_at: DateTime<Utc>,
296}
297
298#[derive(Debug, Clone)]
300pub struct StoredImage {
301 pub info: StoredImageInfo,
302 pub data: Vec<u8>,
303}
304
305#[derive(Debug, Clone)]
307pub struct CreateStoredImage {
308 pub filename: String,
309 pub content_type: String,
310 pub data: Vec<u8>,
311 pub metadata: serde_json::Value,
312}
313
314#[async_trait]
315pub trait ImageArtifactStore: Send + Sync {
316 async fn create_image(&self, input: CreateStoredImage) -> Result<StoredImageInfo>;
318
319 async fn get_image(&self, image_id: ImageId) -> Result<Option<StoredImage>>;
321
322 async fn get_image_info(&self, image_id: ImageId) -> Result<Option<StoredImageInfo>>;
324}
325
326#[derive(Debug, Clone)]
332pub struct ProviderCredentials {
333 pub api_key: String,
334 pub base_url: Option<String>,
335}
336
337#[async_trait]
338pub trait ProviderCredentialStore: Send + Sync {
339 async fn get_default_provider_credentials(
344 &self,
345 provider_type: &str,
346 ) -> Result<Option<ProviderCredentials>>;
347}
348
349#[async_trait]
360pub trait ToolExecutor: Send + Sync {
361 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult>;
366
367 async fn execute_with_context(
372 &self,
373 tool_call: &ToolCall,
374 tool_def: &ToolDefinition,
375 _context: &ToolContext,
376 ) -> Result<ToolResult> {
377 self.execute(tool_call, tool_def).await
379 }
380
381 async fn execute_batch(
383 &self,
384 tool_calls: &[ToolCall],
385 tool_defs: &[ToolDefinition],
386 ) -> Result<Vec<ToolResult>> {
387 let mut results = Vec::with_capacity(tool_calls.len());
388
389 let tool_map = build_tool_map(tool_defs);
390
391 for tool_call in tool_calls {
392 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
393 crate::error::AgentLoopError::tool(format!(
394 "Tool definition not found: {}",
395 tool_call.name
396 ))
397 })?;
398
399 results.push(self.execute(tool_call, tool_def).await?);
400 }
401
402 Ok(results)
403 }
404
405 async fn execute_parallel(
407 &self,
408 tool_calls: &[ToolCall],
409 tool_defs: &[ToolDefinition],
410 ) -> Result<Vec<ToolResult>>
411 where
412 Self: Sized,
413 {
414 use futures::future::join_all;
415
416 let tool_map = build_tool_map(tool_defs);
417
418 let futures: Vec<_> = tool_calls
419 .iter()
420 .map(|tool_call| async {
421 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
422 crate::error::AgentLoopError::tool(format!(
423 "Tool definition not found: {}",
424 tool_call.name
425 ))
426 })?;
427 self.execute(tool_call, tool_def).await
428 })
429 .collect();
430
431 let results = join_all(futures).await;
432 results.into_iter().collect()
433 }
434}
435
436#[async_trait]
440impl ToolExecutor for std::sync::Arc<dyn ToolExecutor> {
441 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult> {
442 (**self).execute(tool_call, tool_def).await
443 }
444
445 async fn execute_with_context(
446 &self,
447 tool_call: &ToolCall,
448 tool_def: &ToolDefinition,
449 context: &ToolContext,
450 ) -> Result<ToolResult> {
451 (**self)
452 .execute_with_context(tool_call, tool_def, context)
453 .await
454 }
455
456 async fn execute_batch(
457 &self,
458 tool_calls: &[ToolCall],
459 tool_defs: &[ToolDefinition],
460 ) -> Result<Vec<ToolResult>> {
461 (**self).execute_batch(tool_calls, tool_defs).await
462 }
463}
464
465#[async_trait]
477pub trait SessionFileSystem: Send + Sync {
478 fn display_root(&self) -> String {
484 crate::session_path::WORKSPACE_PREFIX.to_string()
485 }
486
487 fn display_path(&self, path: &str) -> String {
493 crate::session_path::to_display_path(path)
494 }
495
496 fn resolve_path(&self, input: &str) -> String {
506 crate::session_path::to_session_path(input)
507 }
508
509 fn is_mount_resolver(&self) -> bool;
513
514 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>>;
516
517 async fn write_file(
519 &self,
520 session_id: SessionId,
521 path: &str,
522 content: &str,
523 encoding: &str,
524 ) -> Result<SessionFile>;
525
526 async fn write_file_if_content_matches(
531 &self,
532 session_id: SessionId,
533 path: &str,
534 expected_content: &str,
535 expected_encoding: &str,
536 content: &str,
537 encoding: &str,
538 ) -> Result<Option<SessionFile>> {
539 let Some(existing) = self.read_file(session_id, path).await? else {
540 return Ok(None);
541 };
542
543 if existing.is_directory {
544 return Ok(None);
545 }
546
547 let current_content = existing.content.unwrap_or_default();
548 if current_content != expected_content || existing.encoding != expected_encoding {
549 return Ok(None);
550 }
551
552 self.write_file(session_id, path, content, encoding)
553 .await
554 .map(Some)
555 }
556
557 async fn delete_file(&self, session_id: SessionId, path: &str, recursive: bool)
559 -> Result<bool>;
560
561 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>>;
563
564 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>>;
566
567 async fn grep_files(
574 &self,
575 session_id: SessionId,
576 pattern: &str,
577 path_pattern: Option<&str>,
578 ) -> Result<Vec<GrepMatch>>;
579
580 async fn grep_files_with_options(
586 &self,
587 session_id: SessionId,
588 pattern: &str,
589 options: &GrepOptions,
590 ) -> Result<GrepSearchResult> {
591 if options.before_context != 0 || options.after_context != 0 {
592 return Err(crate::error::AgentLoopError::tool(
593 "this file store does not support grep context",
594 ));
595 }
596 let all = self
597 .grep_files(session_id, pattern, options.path_pattern.as_deref())
598 .await?;
599 Ok(crate::session_file::bound_grep_matches(all, options))
600 }
601
602 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
604
605 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
607 if file.is_readonly {
608 return Err(crate::error::AgentLoopError::store(
609 "read-only initial files require a SessionFileSystem-specific seed implementation",
610 ));
611 }
612 self.write_file(session_id, &file.path, &file.content, &file.encoding)
613 .await?;
614 Ok(())
615 }
616}
617
618pub struct WorkspaceScopedFileSystem {
628 inner: Arc<dyn SessionFileSystem>,
629 key: SessionId,
630}
631
632impl WorkspaceScopedFileSystem {
633 pub fn wrap(
635 inner: Arc<dyn SessionFileSystem>,
636 workspace_id: WorkspaceId,
637 ) -> Arc<dyn SessionFileSystem> {
638 Arc::new(Self {
639 inner,
640 key: SessionId::from_uuid(workspace_id.uuid()),
641 })
642 }
643}
644
645#[async_trait]
646impl SessionFileSystem for WorkspaceScopedFileSystem {
647 async fn read_file(&self, _session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
648 self.inner.read_file(self.key, path).await
649 }
650 async fn write_file(
651 &self,
652 _session_id: SessionId,
653 path: &str,
654 content: &str,
655 encoding: &str,
656 ) -> Result<SessionFile> {
657 self.inner
658 .write_file(self.key, path, content, encoding)
659 .await
660 }
661 async fn write_file_if_content_matches(
662 &self,
663 _session_id: SessionId,
664 path: &str,
665 expected_content: &str,
666 expected_encoding: &str,
667 content: &str,
668 encoding: &str,
669 ) -> Result<Option<SessionFile>> {
670 self.inner
671 .write_file_if_content_matches(
672 self.key,
673 path,
674 expected_content,
675 expected_encoding,
676 content,
677 encoding,
678 )
679 .await
680 }
681 async fn delete_file(
682 &self,
683 _session_id: SessionId,
684 path: &str,
685 recursive: bool,
686 ) -> Result<bool> {
687 self.inner.delete_file(self.key, path, recursive).await
688 }
689 async fn list_directory(&self, _session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
690 self.inner.list_directory(self.key, path).await
691 }
692 async fn stat_file(&self, _session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
693 self.inner.stat_file(self.key, path).await
694 }
695 async fn grep_files(
696 &self,
697 _session_id: SessionId,
698 pattern: &str,
699 path_pattern: Option<&str>,
700 ) -> Result<Vec<GrepMatch>> {
701 self.inner.grep_files(self.key, pattern, path_pattern).await
702 }
703 async fn grep_files_with_options(
704 &self,
705 _session_id: SessionId,
706 pattern: &str,
707 options: &GrepOptions,
708 ) -> Result<GrepSearchResult> {
709 self.inner
710 .grep_files_with_options(self.key, pattern, options)
711 .await
712 }
713 async fn create_directory(&self, _session_id: SessionId, path: &str) -> Result<FileInfo> {
714 self.inner.create_directory(self.key, path).await
715 }
716 async fn seed_initial_file(&self, _session_id: SessionId, file: &InitialFile) -> Result<()> {
717 self.inner.seed_initial_file(self.key, file).await
718 }
719
720 fn display_root(&self) -> String {
721 self.inner.display_root()
722 }
723
724 fn display_path(&self, path: &str) -> String {
725 self.inner.display_path(path)
726 }
727
728 fn resolve_path(&self, input: &str) -> String {
729 self.inner.resolve_path(input)
730 }
731
732 fn is_mount_resolver(&self) -> bool {
733 self.inner.is_mount_resolver()
734 }
735}
736
737#[async_trait]
738impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
739 fn display_root(&self) -> String {
740 (**self).display_root()
741 }
742
743 fn display_path(&self, path: &str) -> String {
744 (**self).display_path(path)
745 }
746
747 fn resolve_path(&self, input: &str) -> String {
748 (**self).resolve_path(input)
749 }
750
751 fn is_mount_resolver(&self) -> bool {
752 (**self).is_mount_resolver()
753 }
754
755 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
756 (**self).read_file(session_id, path).await
757 }
758
759 async fn write_file(
760 &self,
761 session_id: SessionId,
762 path: &str,
763 content: &str,
764 encoding: &str,
765 ) -> Result<SessionFile> {
766 (**self)
767 .write_file(session_id, path, content, encoding)
768 .await
769 }
770
771 async fn write_file_if_content_matches(
772 &self,
773 session_id: SessionId,
774 path: &str,
775 expected_content: &str,
776 expected_encoding: &str,
777 content: &str,
778 encoding: &str,
779 ) -> Result<Option<SessionFile>> {
780 (**self)
781 .write_file_if_content_matches(
782 session_id,
783 path,
784 expected_content,
785 expected_encoding,
786 content,
787 encoding,
788 )
789 .await
790 }
791
792 async fn delete_file(
793 &self,
794 session_id: SessionId,
795 path: &str,
796 recursive: bool,
797 ) -> Result<bool> {
798 (**self).delete_file(session_id, path, recursive).await
799 }
800
801 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
802 (**self).list_directory(session_id, path).await
803 }
804
805 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
806 (**self).stat_file(session_id, path).await
807 }
808
809 async fn grep_files(
810 &self,
811 session_id: SessionId,
812 pattern: &str,
813 path_pattern: Option<&str>,
814 ) -> Result<Vec<GrepMatch>> {
815 (**self).grep_files(session_id, pattern, path_pattern).await
816 }
817
818 async fn grep_files_with_options(
819 &self,
820 session_id: SessionId,
821 pattern: &str,
822 options: &GrepOptions,
823 ) -> Result<GrepSearchResult> {
824 (**self)
825 .grep_files_with_options(session_id, pattern, options)
826 .await
827 }
828
829 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
830 (**self).create_directory(session_id, path).await
831 }
832
833 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
834 (**self).seed_initial_file(session_id, file).await
835 }
836}
837
838pub use SessionFileSystem as SessionFileStore;
840
841#[derive(Clone, Default)]
847pub struct SessionFileSystemFactoryContext {
848 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
849}
850
851impl SessionFileSystemFactoryContext {
852 pub fn new() -> Self {
853 Self::default()
854 }
855
856 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
857 let values = Arc::make_mut(&mut self.values);
858 values.insert(TypeId::of::<T>(), value);
859 self
860 }
861
862 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
863 self.values
864 .get(&TypeId::of::<T>())
865 .and_then(|value| value.clone().downcast::<T>().ok())
866 }
867
868 pub fn with_workspace_roots(self, roots: Arc<crate::WorkspaceRootSet>) -> Self {
869 self.with(roots)
870 }
871
872 pub fn workspace_roots(&self) -> Option<Arc<crate::WorkspaceRootSet>> {
873 self.get::<crate::WorkspaceRootSet>()
874 }
875}
876
877#[async_trait]
879pub trait SessionFileSystemFactory: Send + Sync {
880 fn name(&self) -> &'static str {
882 "SessionFileSystemFactory"
883 }
884
885 fn is_disabled(&self) -> bool {
888 false
889 }
890
891 async fn create_session_file_system(
893 &self,
894 context: SessionFileSystemFactoryContext,
895 ) -> Result<Arc<dyn SessionFileSystem>>;
896}
897
898#[derive(Debug, Clone, Default)]
900pub struct DisabledSessionFileSystemFactory;
901
902#[async_trait]
903impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
904 fn name(&self) -> &'static str {
905 "DisabledSessionFileSystemFactory"
906 }
907
908 fn is_disabled(&self) -> bool {
909 true
910 }
911
912 async fn create_session_file_system(
913 &self,
914 _context: SessionFileSystemFactoryContext,
915 ) -> Result<Arc<dyn SessionFileSystem>> {
916 Err(crate::error::AgentLoopError::config(
917 "session filesystem is disabled",
918 ))
919 }
920}
921
922#[derive(Debug, Clone)]
928pub struct KeyInfo {
929 pub key: String,
930 pub created_at: chrono::DateTime<chrono::Utc>,
931 pub updated_at: chrono::DateTime<chrono::Utc>,
932}
933
934#[derive(Debug, Clone)]
936pub struct SecretInfo {
937 pub name: String,
938 pub created_at: chrono::DateTime<chrono::Utc>,
939 pub updated_at: chrono::DateTime<chrono::Utc>,
940}
941
942#[derive(Debug, Clone, serde::Serialize)]
952pub struct KnowledgeSearchHit {
953 pub id: String,
955 pub kb_id: String,
957 pub title: String,
958 pub kind: String,
959 pub tags: Vec<String>,
960 pub snippet: String,
962 pub resource: Option<String>,
964}
965
966#[async_trait]
970pub trait KnowledgeStore: Send + Sync {
971 async fn search_knowledge(
972 &self,
973 org_id: crate::typed_id::OrgId,
974 kb_public_ids: &[String],
975 query: &str,
976 kind: Option<&str>,
977 tags: &[String],
978 limit: usize,
979 ) -> Result<Vec<KnowledgeSearchHit>>;
980}
981
982#[async_trait]
987pub trait SessionStorageStore: Send + Sync {
988 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
992
993 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
995
996 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
998
999 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
1001
1002 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
1006
1007 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
1009
1010 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
1012
1013 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
1015}
1016
1017use crate::session_schedule::SessionSchedule;
1022use crate::typed_id::ScheduleId;
1023
1024#[async_trait]
1028pub trait SessionScheduleStore: Send + Sync {
1029 async fn create_schedule(
1031 &self,
1032 session_id: SessionId,
1033 description: String,
1034 cron_expression: Option<String>,
1035 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
1036 timezone: String,
1037 ) -> Result<SessionSchedule>;
1038
1039 async fn create_schedule_enforcing_limits(
1043 &self,
1044 session_id: SessionId,
1045 description: String,
1046 cron_expression: Option<String>,
1047 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
1048 timezone: String,
1049 ) -> std::result::Result<SessionSchedule, crate::session_schedule::ScheduleLimitError> {
1050 let per_session = self
1051 .count_active_schedules(session_id)
1052 .await
1053 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
1054 if per_session >= crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION {
1055 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
1056 format!(
1057 "Maximum {} active schedules per session. Cancel an existing schedule first.",
1058 crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION
1059 ),
1060 ));
1061 }
1062
1063 let max_per_org = crate::session_schedule::max_active_schedules_per_org();
1064 let per_org = self
1065 .count_active_org_schedules()
1066 .await
1067 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
1068 if i64::from(per_org) >= max_per_org {
1069 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
1070 format!(
1071 "Maximum {max_per_org} active schedules per org reached. Cancel an existing schedule first."
1072 ),
1073 ));
1074 }
1075
1076 if let Some(cron) = cron_expression.as_deref() {
1077 crate::session_schedule::validate_cron_min_interval(cron)
1078 .map_err(crate::session_schedule::ScheduleLimitError::Rejected)?;
1079 }
1080
1081 self.create_schedule(
1082 session_id,
1083 description,
1084 cron_expression,
1085 scheduled_at,
1086 timezone,
1087 )
1088 .await
1089 .map_err(crate::session_schedule::ScheduleLimitError::Store)
1090 }
1091
1092 async fn cancel_schedule(
1094 &self,
1095 session_id: SessionId,
1096 schedule_id: ScheduleId,
1097 ) -> Result<SessionSchedule>;
1098
1099 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
1101
1102 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
1104
1105 async fn count_active_org_schedules(&self) -> Result<u32>;
1110}
1111
1112#[async_trait]
1122pub trait SessionResourceRegistry: Send + Sync {
1123 async fn register(
1125 &self,
1126 entry: crate::session_resource::RegisterSessionResource,
1127 ) -> Result<crate::session_resource::SessionResourceEntry>;
1128
1129 async fn update_status(
1131 &self,
1132 session_id: SessionId,
1133 resource_id: &str,
1134 status: crate::session_resource::SessionResourceStatus,
1135 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1136
1137 async fn get(
1139 &self,
1140 session_id: SessionId,
1141 resource_id: &str,
1142 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1143
1144 async fn list(
1146 &self,
1147 session_id: SessionId,
1148 filter: Option<&crate::session_resource::SessionResourceFilter>,
1149 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
1150
1151 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
1153}
1154
1155#[async_trait]
1165pub trait LeasedResourceStore: Send + Sync {
1166 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
1172
1173 async fn release_resource(
1179 &self,
1180 session_id: SessionId,
1181 provider: &str,
1182 resource_type: &str,
1183 external_id: &str,
1184 ) -> Result<Option<LeasedResource>>;
1185
1186 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
1191}
1192
1193pub const DEFAULT_MAX_SUBAGENT_DEPTH: u32 = 2;
1200pub const DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS: u32 = 16;
1201pub const DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS: u32 = 200;
1202pub const DEFAULT_MAX_ACTIVE_DETACHED_TASKS: u32 = 8;
1208pub const DEFAULT_MAX_TOTAL_DETACHED_TASKS: u32 = 50;
1209
1210#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1212pub struct SubagentNestingPolicy {
1213 pub platform_default: u32,
1214 pub org_override: Option<u32>,
1215 pub agent_override: Option<u32>,
1216 pub platform_default_max_active_descendant_tasks: u32,
1217 pub org_override_max_active_descendant_tasks: Option<u32>,
1218 pub agent_override_max_active_descendant_tasks: Option<u32>,
1219 pub platform_default_max_total_descendant_tasks: u32,
1220 pub org_override_max_total_descendant_tasks: Option<u32>,
1221 pub agent_override_max_total_descendant_tasks: Option<u32>,
1222 pub platform_default_max_active_detached_tasks: u32,
1223 pub org_override_max_active_detached_tasks: Option<u32>,
1224 pub agent_override_max_active_detached_tasks: Option<u32>,
1225 pub platform_default_max_total_detached_tasks: u32,
1226 pub org_override_max_total_detached_tasks: Option<u32>,
1227 pub agent_override_max_total_detached_tasks: Option<u32>,
1228}
1229
1230impl Default for SubagentNestingPolicy {
1231 fn default() -> Self {
1232 Self {
1233 platform_default: DEFAULT_MAX_SUBAGENT_DEPTH,
1234 org_override: None,
1235 agent_override: None,
1236 platform_default_max_active_descendant_tasks:
1237 DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS,
1238 org_override_max_active_descendant_tasks: None,
1239 agent_override_max_active_descendant_tasks: None,
1240 platform_default_max_total_descendant_tasks:
1241 DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS,
1242 org_override_max_total_descendant_tasks: None,
1243 agent_override_max_total_descendant_tasks: None,
1244 platform_default_max_active_detached_tasks: DEFAULT_MAX_ACTIVE_DETACHED_TASKS,
1245 org_override_max_active_detached_tasks: None,
1246 agent_override_max_active_detached_tasks: None,
1247 platform_default_max_total_detached_tasks: DEFAULT_MAX_TOTAL_DETACHED_TASKS,
1248 org_override_max_total_detached_tasks: None,
1249 agent_override_max_total_detached_tasks: None,
1250 }
1251 }
1252}
1253
1254impl SubagentNestingPolicy {
1255 pub fn max_subagent_depth(self) -> u32 {
1256 self.agent_override
1257 .or(self.org_override)
1258 .unwrap_or(self.platform_default)
1259 }
1260
1261 pub fn max_active_descendant_tasks(self) -> u32 {
1262 self.agent_override_max_active_descendant_tasks
1263 .or(self.org_override_max_active_descendant_tasks)
1264 .unwrap_or(self.platform_default_max_active_descendant_tasks)
1265 }
1266
1267 pub fn max_total_descendant_tasks(self) -> u32 {
1268 self.agent_override_max_total_descendant_tasks
1269 .or(self.org_override_max_total_descendant_tasks)
1270 .unwrap_or(self.platform_default_max_total_descendant_tasks)
1271 }
1272
1273 pub fn max_active_detached_tasks(self) -> u32 {
1274 self.agent_override_max_active_detached_tasks
1275 .or(self.org_override_max_active_detached_tasks)
1276 .unwrap_or(self.platform_default_max_active_detached_tasks)
1277 }
1278
1279 pub fn max_total_detached_tasks(self) -> u32 {
1280 self.agent_override_max_total_detached_tasks
1281 .or(self.org_override_max_total_detached_tasks)
1282 .unwrap_or(self.platform_default_max_total_detached_tasks)
1283 }
1284
1285 pub fn with_platform_default(mut self, depth: u32) -> Self {
1286 self.platform_default = depth;
1287 self
1288 }
1289
1290 pub fn with_org_override(mut self, depth: Option<u32>) -> Self {
1291 self.org_override = depth;
1292 self
1293 }
1294
1295 pub fn with_agent_override(mut self, depth: Option<u32>) -> Self {
1296 self.agent_override = depth;
1297 self
1298 }
1299
1300 pub fn with_agent_task_caps_override(
1301 mut self,
1302 max_active: Option<u32>,
1303 max_total: Option<u32>,
1304 ) -> Self {
1305 self.agent_override_max_active_descendant_tasks = max_active;
1306 self.agent_override_max_total_descendant_tasks = max_total;
1307 self
1308 }
1309
1310 pub fn with_agent_detached_task_caps_override(
1311 mut self,
1312 max_active: Option<u32>,
1313 max_total: Option<u32>,
1314 ) -> Self {
1315 self.agent_override_max_active_detached_tasks = max_active;
1316 self.agent_override_max_total_detached_tasks = max_total;
1317 self
1318 }
1319}
1320
1321pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
1323
1324#[async_trait]
1329pub trait UserConnectionResolver: Send + Sync {
1330 async fn get_connection_token(
1333 &self,
1334 session_id: SessionId,
1335 provider: &str,
1336 ) -> Result<Option<String>>;
1337
1338 async fn get_connection_user(
1343 &self,
1344 _session_id: SessionId,
1345 _provider: &str,
1346 ) -> Result<Option<Uuid>> {
1347 Ok(None)
1348 }
1349
1350 async fn get_connection_token_for_user(
1355 &self,
1356 _user_id: Uuid,
1357 _provider: &str,
1358 ) -> Result<Option<String>> {
1359 Ok(None)
1360 }
1361
1362 async fn get_connection_metadata(
1365 &self,
1366 _session_id: SessionId,
1367 _provider: &str,
1368 ) -> Result<Option<serde_json::Value>> {
1369 Ok(None)
1370 }
1371}
1372
1373#[async_trait]
1383pub trait BudgetChecker: Send + Sync {
1384 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
1386}
1387
1388#[async_trait]
1397pub trait PaymentAuthority: Send + Sync {
1398 async fn execute_machine_payment(
1399 &self,
1400 session_id: SessionId,
1401 request: crate::payment::MachinePaymentRequest,
1402 ) -> Result<crate::payment::MachinePaymentResponse>;
1403}
1404
1405#[async_trait]
1415pub trait SessionCreationAuthority: Send + Sync {
1416 async fn authorize_session_creation(&self, session_id: SessionId) -> Result<SessionId>;
1420}
1421
1422#[async_trait]
1432pub trait OutboundToolRateLimiter: Send + Sync {
1433 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
1435}
1436
1437#[derive(Debug)]
1443pub enum ToolCallClaimResult {
1444 Claimed { claim_token: uuid::Uuid },
1447 AlreadySettled {
1449 result_json: serde_json::Value,
1450 args_fingerprint: String,
1451 },
1452 AlreadyRunning { args_fingerprint: String },
1457 DeterminismViolation {
1461 stored_fingerprint: String,
1462 current_fingerprint: String,
1463 },
1464}
1465
1466#[derive(Debug, Clone)]
1468pub enum DurableToolCallStatus {
1469 Settled { result_json: serde_json::Value },
1471 Interrupted {
1473 result_json: Option<serde_json::Value>,
1474 },
1475 Running,
1477}
1478
1479#[async_trait]
1484pub trait DurableToolResultStore: Send + Sync + 'static {
1485 async fn try_claim_tool_call(
1493 &self,
1494 turn_id: &str,
1495 tool_call_id: &str,
1496 tool_name: &str,
1497 args_fingerprint: &str,
1498 ) -> Result<ToolCallClaimResult>;
1499
1500 async fn settle_tool_call(
1506 &self,
1507 turn_id: &str,
1508 tool_call_id: &str,
1509 result_json: serde_json::Value,
1510 status: &str,
1511 claim_token: uuid::Uuid,
1512 ) -> Result<bool>;
1513
1514 async fn get_tool_call_status(
1519 &self,
1520 turn_id: &str,
1521 tool_call_id: &str,
1522 ) -> Result<Option<DurableToolCallStatus>>;
1523}
1524
1525pub struct NoopDurableToolResultStore;
1528
1529#[async_trait]
1530impl DurableToolResultStore for NoopDurableToolResultStore {
1531 async fn try_claim_tool_call(
1532 &self,
1533 _turn_id: &str,
1534 _tool_call_id: &str,
1535 _tool_name: &str,
1536 _args_fingerprint: &str,
1537 ) -> Result<ToolCallClaimResult> {
1538 Ok(ToolCallClaimResult::Claimed {
1539 claim_token: uuid::Uuid::new_v4(),
1540 })
1541 }
1542
1543 async fn settle_tool_call(
1544 &self,
1545 _turn_id: &str,
1546 _tool_call_id: &str,
1547 _result_json: serde_json::Value,
1548 _status: &str,
1549 _claim_token: uuid::Uuid,
1550 ) -> Result<bool> {
1551 Ok(true)
1552 }
1553
1554 async fn get_tool_call_status(
1555 &self,
1556 _turn_id: &str,
1557 _tool_call_id: &str,
1558 ) -> Result<Option<DurableToolCallStatus>> {
1559 Ok(None)
1560 }
1561}
1562
1563#[derive(Debug, Clone)]
1569pub struct StreamProgress {
1570 pub accumulated_len: usize,
1572 pub last_delta_at: u64,
1574}
1575
1576#[async_trait]
1582pub trait StreamHeartbeater: Send + Sync {
1583 async fn heartbeat(&self, progress: StreamProgress);
1589}
1590
1591pub struct NoopStreamHeartbeater;
1593
1594#[async_trait]
1595impl StreamHeartbeater for NoopStreamHeartbeater {
1596 async fn heartbeat(&self, _progress: StreamProgress) {}
1597}
1598
1599#[derive(Debug, Clone)]
1605pub struct PartialStreamState {
1606 pub message_id: MessageId,
1608
1609 pub accumulated: String,
1612}
1613
1614#[async_trait]
1622pub trait PartialStreamStore: Send + Sync {
1623 async fn get_partial_stream(
1626 &self,
1627 session_id: SessionId,
1628 turn_id: &str,
1629 ) -> Result<Option<PartialStreamState>>;
1630}
1631
1632pub struct NoopPartialStreamStore;
1634
1635#[async_trait]
1636impl PartialStreamStore for NoopPartialStreamStore {
1637 async fn get_partial_stream(
1638 &self,
1639 _session_id: SessionId,
1640 _turn_id: &str,
1641 ) -> Result<Option<PartialStreamState>> {
1642 Ok(None)
1643 }
1644}
1645
1646#[derive(Clone)]
1655pub struct ToolContext {
1656 pub session_id: SessionId,
1658 pub workspace_id: WorkspaceId,
1665
1666 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1668
1669 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1671
1672 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1674
1675 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1677
1678 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1680
1681 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1687
1688 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1690
1691 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1693
1694 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1696
1697 pub session_store: Option<Arc<dyn SessionStore>>,
1699
1700 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1702
1703 pub agent_store: Option<Arc<dyn AgentStore>>,
1705
1706 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1708
1709 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1711
1712 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1714 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1716
1717 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1721
1722 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1724
1725 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1727
1728 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1731
1732 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1735
1736 pub event_context: Option<crate::events::EventContext>,
1739
1740 pub tool_call_id: Option<String>,
1743 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1745
1746 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1749
1750 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1754
1755 pub org_id: Option<crate::typed_id::OrgId>,
1757
1758 pub network_access: Option<crate::network_access::NetworkAccessList>,
1761
1762 pub locale: Option<String>,
1766
1767 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1769
1770 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1772
1773 pub session_creation_authority: Option<Arc<dyn SessionCreationAuthority>>,
1775
1776 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1780
1781 pub subagent_nesting_policy: SubagentNestingPolicy,
1783
1784 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1788}
1789
1790impl ToolContext {
1791 pub fn workspace_fs_key(&self) -> SessionId {
1796 SessionId::from_uuid(self.workspace_id.uuid())
1797 }
1798
1799 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1801 self.workspace_id = workspace_id;
1802 self
1803 }
1804
1805 pub fn new(session_id: SessionId) -> Self {
1807 Self {
1808 session_id,
1809 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1810 file_store: None,
1811 storage_store: None,
1812 image_store: None,
1813 provider_credential_store: None,
1814 utility_llm_service: None,
1815 mcp_invoker: None,
1816 egress_service: None,
1817 sqldb_store: None,
1818 message_retriever: None,
1819 session_store: None,
1820 session_mutator: None,
1821 agent_store: None,
1822 connection_resolver: None,
1823 schedule_store: None,
1824 platform_store: None,
1825 knowledge_store: None,
1826 knowledge_index_search: None,
1827 leased_resource_store: None,
1828 session_resource_registry: None,
1829 session_task_registry: None,
1830 event_emitter: None,
1831 event_context: None,
1832 tool_call_id: None,
1833 capability_registry: None,
1834 tool_registry: None,
1835 visible_tool_names: None,
1836 org_id: None,
1837 network_access: None,
1838 locale: None,
1839 budget_checker: None,
1840 payment_authority: None,
1841 session_creation_authority: None,
1842 subagent_spawn_store: None,
1843 subagent_nesting_policy: SubagentNestingPolicy::default(),
1844 reasoning_effort_handle: None,
1845 }
1846 }
1847
1848 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1850 Self {
1851 session_id,
1852 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1853 file_store: Some(file_store),
1854 storage_store: None,
1855 image_store: None,
1856 provider_credential_store: None,
1857 utility_llm_service: None,
1858 mcp_invoker: None,
1859 egress_service: None,
1860 sqldb_store: None,
1861 message_retriever: None,
1862 session_store: None,
1863 session_mutator: None,
1864 agent_store: None,
1865 connection_resolver: None,
1866 schedule_store: None,
1867 platform_store: None,
1868 knowledge_store: None,
1869 knowledge_index_search: None,
1870 leased_resource_store: None,
1871 session_resource_registry: None,
1872 session_task_registry: None,
1873 event_emitter: None,
1874 event_context: None,
1875 tool_call_id: None,
1876 capability_registry: None,
1877 tool_registry: None,
1878 visible_tool_names: None,
1879 org_id: None,
1880 network_access: None,
1881 locale: None,
1882 budget_checker: None,
1883 payment_authority: None,
1884 session_creation_authority: None,
1885 subagent_spawn_store: None,
1886 subagent_nesting_policy: SubagentNestingPolicy::default(),
1887 reasoning_effort_handle: None,
1888 }
1889 }
1890
1891 pub fn with_storage_store(
1893 session_id: SessionId,
1894 storage_store: Arc<dyn SessionStorageStore>,
1895 ) -> Self {
1896 Self {
1897 session_id,
1898 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1899 file_store: None,
1900 storage_store: Some(storage_store),
1901 image_store: None,
1902 provider_credential_store: None,
1903 utility_llm_service: None,
1904 mcp_invoker: None,
1905 egress_service: None,
1906 sqldb_store: None,
1907 message_retriever: None,
1908 session_store: None,
1909 session_mutator: None,
1910 agent_store: None,
1911 connection_resolver: None,
1912 schedule_store: None,
1913 platform_store: None,
1914 knowledge_store: None,
1915 knowledge_index_search: None,
1916 leased_resource_store: None,
1917 session_resource_registry: None,
1918 session_task_registry: None,
1919 event_emitter: None,
1920 event_context: None,
1921 tool_call_id: None,
1922 capability_registry: None,
1923 tool_registry: None,
1924 visible_tool_names: None,
1925 org_id: None,
1926 network_access: None,
1927 locale: None,
1928 budget_checker: None,
1929 payment_authority: None,
1930 session_creation_authority: None,
1931 subagent_spawn_store: None,
1932 subagent_nesting_policy: SubagentNestingPolicy::default(),
1933 reasoning_effort_handle: None,
1934 }
1935 }
1936
1937 pub fn with_stores(
1939 session_id: SessionId,
1940 file_store: Arc<dyn SessionFileSystem>,
1941 storage_store: Arc<dyn SessionStorageStore>,
1942 ) -> Self {
1943 Self {
1944 session_id,
1945 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1946 file_store: Some(file_store),
1947 storage_store: Some(storage_store),
1948 sqldb_store: None,
1949 image_store: None,
1950 provider_credential_store: None,
1951 utility_llm_service: None,
1952 mcp_invoker: None,
1953 egress_service: None,
1954 message_retriever: None,
1955 session_store: None,
1956 session_mutator: None,
1957 agent_store: None,
1958 connection_resolver: None,
1959 schedule_store: None,
1960 platform_store: None,
1961 knowledge_store: None,
1962 knowledge_index_search: None,
1963 leased_resource_store: None,
1964 session_resource_registry: None,
1965 session_task_registry: None,
1966 event_emitter: None,
1967 event_context: None,
1968 tool_call_id: None,
1969 capability_registry: None,
1970 tool_registry: None,
1971 visible_tool_names: None,
1972 org_id: None,
1973 network_access: None,
1974 locale: None,
1975 budget_checker: None,
1976 payment_authority: None,
1977 session_creation_authority: None,
1978 subagent_spawn_store: None,
1979 subagent_nesting_policy: SubagentNestingPolicy::default(),
1980 reasoning_effort_handle: None,
1981 }
1982 }
1983
1984 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1986 self.sqldb_store = Some(sqldb_store);
1987 self
1988 }
1989
1990 pub fn with_message_retriever(
1992 mut self,
1993 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1994 ) -> Self {
1995 self.message_retriever = Some(retriever);
1996 self
1997 }
1998
1999 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
2001 self.session_store = Some(store);
2002 self
2003 }
2004
2005 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
2007 self.session_mutator = Some(mutator);
2008 self
2009 }
2010
2011 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
2015 self.reasoning_effort_handle = Some(handle);
2016 self
2017 }
2018
2019 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
2021 self.agent_store = Some(store);
2022 self
2023 }
2024
2025 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
2027 self.connection_resolver = Some(resolver);
2028 self
2029 }
2030
2031 pub fn with_image_store(
2033 session_id: SessionId,
2034 image_store: Arc<dyn ImageArtifactStore>,
2035 ) -> Self {
2036 Self {
2037 session_id,
2038 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
2039 file_store: None,
2040 storage_store: None,
2041 image_store: Some(image_store),
2042 provider_credential_store: None,
2043 utility_llm_service: None,
2044 mcp_invoker: None,
2045 egress_service: None,
2046 sqldb_store: None,
2047 message_retriever: None,
2048 session_store: None,
2049 session_mutator: None,
2050 agent_store: None,
2051 connection_resolver: None,
2052 schedule_store: None,
2053 platform_store: None,
2054 knowledge_store: None,
2055 knowledge_index_search: None,
2056 leased_resource_store: None,
2057 session_resource_registry: None,
2058 session_task_registry: None,
2059 event_emitter: None,
2060 event_context: None,
2061 tool_call_id: None,
2062 capability_registry: None,
2063 tool_registry: None,
2064 visible_tool_names: None,
2065 org_id: None,
2066 network_access: None,
2067 locale: None,
2068 budget_checker: None,
2069 payment_authority: None,
2070 session_creation_authority: None,
2071 subagent_spawn_store: None,
2072 subagent_nesting_policy: SubagentNestingPolicy::default(),
2073 reasoning_effort_handle: None,
2074 }
2075 }
2076
2077 pub fn with_provider_credential_store(
2079 mut self,
2080 store: Arc<dyn ProviderCredentialStore>,
2081 ) -> Self {
2082 self.provider_credential_store = Some(store);
2083 self
2084 }
2085
2086 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
2088 self.utility_llm_service = Some(service);
2089 self
2090 }
2091
2092 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
2094 self.mcp_invoker = Some(invoker);
2095 self
2096 }
2097
2098 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
2100 self.egress_service = Some(service);
2101 self
2102 }
2103
2104 pub fn with_egress_service_opt(
2107 mut self,
2108 service: Option<Arc<dyn crate::EgressService>>,
2109 ) -> Self {
2110 if let Some(service) = service {
2111 self.egress_service = Some(service);
2112 }
2113 self
2114 }
2115
2116 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
2118 self.storage_store = Some(store);
2119 self
2120 }
2121
2122 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
2124 self.schedule_store = Some(store);
2125 self
2126 }
2127
2128 pub fn with_platform_store(
2130 mut self,
2131 store: Arc<dyn crate::platform_store::PlatformStore>,
2132 ) -> Self {
2133 self.platform_store = Some(store);
2134 self
2135 }
2136
2137 pub fn with_knowledge_index_search(
2139 mut self,
2140 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
2141 ) -> Self {
2142 self.knowledge_index_search = Some(search);
2143 self
2144 }
2145
2146 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
2148 self.leased_resource_store = Some(store);
2149 self
2150 }
2151
2152 pub fn with_session_resource_registry(
2154 mut self,
2155 registry: Arc<dyn SessionResourceRegistry>,
2156 ) -> Self {
2157 self.session_resource_registry = Some(registry);
2158 self
2159 }
2160
2161 pub fn with_session_task_registry(
2163 mut self,
2164 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
2165 ) -> Self {
2166 self.session_task_registry = Some(registry);
2167 self
2168 }
2169
2170 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
2172 self.org_id = Some(org_id);
2173 self
2174 }
2175
2176 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
2178 self.tool_registry = Some(registry);
2179 self
2180 }
2181
2182 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
2184 self.visible_tool_names = Some(names);
2185 self
2186 }
2187
2188 pub fn with_network_access(
2190 mut self,
2191 network_access: Option<crate::network_access::NetworkAccessList>,
2192 ) -> Self {
2193 self.network_access = network_access;
2194 self
2195 }
2196
2197 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
2199 self.payment_authority = Some(authority);
2200 self
2201 }
2202
2203 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
2205 self.subagent_spawn_store = Some(store);
2206 self
2207 }
2208
2209 pub fn with_subagent_nesting_policy(mut self, policy: SubagentNestingPolicy) -> Self {
2211 self.subagent_nesting_policy = policy;
2212 self
2213 }
2214
2215 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
2220 let (Some(emitter), Some(ctx), Some(call_id)) =
2221 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2222 else {
2223 return;
2224 };
2225 if let Err(e) = emitter
2226 .emit(EventRequest::new(
2227 self.session_id,
2228 ctx.clone(),
2229 crate::events::ToolProgressData {
2230 tool_call_id: call_id.clone(),
2231 tool_name: tool_name.to_string(),
2232 message: message.to_string(),
2233 display_name: None,
2234 },
2235 ))
2236 .await
2237 {
2238 tracing::debug!(
2239 tool_call_id = call_id,
2240 tool_name,
2241 error = %e,
2242 "Failed to emit tool.progress event"
2243 );
2244 }
2245 }
2246
2247 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
2252 let (Some(emitter), Some(ctx), Some(call_id)) =
2253 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2254 else {
2255 return;
2256 };
2257 if let Err(e) = emitter
2258 .emit(EventRequest::new(
2259 self.session_id,
2260 ctx.clone(),
2261 crate::events::ToolOutputDeltaData {
2262 tool_call_id: call_id.clone(),
2263 tool_name: tool_name.to_string(),
2264 delta: delta.to_string(),
2265 stream: stream.to_string(),
2266 },
2267 ))
2268 .await
2269 {
2270 tracing::debug!(
2271 tool_call_id = call_id,
2272 tool_name,
2273 error = %e,
2274 "Failed to emit tool.output.delta event"
2275 );
2276 }
2277 }
2278}
2279
2280impl std::fmt::Debug for ToolContext {
2281 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2282 f.debug_struct("ToolContext")
2283 .field("session_id", &self.session_id)
2284 .field("file_store", &self.file_store.is_some())
2285 .field("storage_store", &self.storage_store.is_some())
2286 .field("image_store", &self.image_store.is_some())
2287 .field(
2288 "provider_credential_store",
2289 &self.provider_credential_store.is_some(),
2290 )
2291 .field("utility_llm_service", &self.utility_llm_service.is_some())
2292 .field("egress_service", &self.egress_service.is_some())
2293 .field("sqldb_store", &self.sqldb_store.is_some())
2294 .field("message_retriever", &self.message_retriever.is_some())
2295 .field("session_store", &self.session_store.is_some())
2296 .field("session_mutator", &self.session_mutator.is_some())
2297 .field("agent_store", &self.agent_store.is_some())
2298 .field("connection_resolver", &self.connection_resolver.is_some())
2299 .field("schedule_store", &self.schedule_store.is_some())
2300 .field("platform_store", &self.platform_store.is_some())
2301 .field(
2302 "knowledge_index_search",
2303 &self.knowledge_index_search.is_some(),
2304 )
2305 .field(
2306 "leased_resource_store",
2307 &self.leased_resource_store.is_some(),
2308 )
2309 .field("event_emitter", &self.event_emitter.is_some())
2310 .field("tool_registry", &self.tool_registry.is_some())
2311 .field("payment_authority", &self.payment_authority.is_some())
2312 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
2313 .field("subagent_nesting_policy", &self.subagent_nesting_policy)
2314 .field("org_id", &self.org_id)
2315 .finish()
2316 }
2317}
2318
2319use crate::events::{Event, EventRequest};
2324
2325#[async_trait]
2336pub trait EventEmitter: Send + Sync {
2337 async fn emit(&self, request: EventRequest) -> Result<Event>;
2342}
2343
2344#[async_trait]
2346impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
2347 async fn emit(&self, request: EventRequest) -> Result<Event> {
2348 (**self).emit(request).await
2349 }
2350}
2351
2352#[derive(Debug, Clone, Default)]
2356pub struct NoopEventEmitter;
2357
2358#[async_trait]
2359impl EventEmitter for NoopEventEmitter {
2360 async fn emit(&self, request: EventRequest) -> Result<Event> {
2361 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2363 }
2364}
2365
2366#[derive(Debug, Clone)]
2379pub struct ResolvedImage {
2380 pub base64: String,
2382 pub media_type: String,
2384}
2385
2386impl ResolvedImage {
2387 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2389 Self {
2390 base64: base64.into(),
2391 media_type: media_type.into(),
2392 }
2393 }
2394
2395 pub fn to_data_url(&self) -> String {
2399 format!("data:{};base64,{}", self.media_type, self.base64)
2400 }
2401}
2402
2403#[async_trait]
2436pub trait ImageResolver: Send + Sync {
2437 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2441}
2442
2443#[derive(Debug)]
2449pub enum SpawnClaimResult {
2450 Claimed {
2453 spawn_handle_id: uuid::Uuid,
2454 claim_token: uuid::Uuid,
2455 },
2456 ClaimedPendingChild {
2460 spawn_handle_id: uuid::Uuid,
2461 claim_token: uuid::Uuid,
2462 },
2463 AlreadyRunning {
2466 child_session_id: crate::typed_id::SessionId,
2467 claim_token: uuid::Uuid,
2469 },
2470 AlreadySettled {
2473 child_session_id: crate::typed_id::SessionId,
2474 terminal_status: String,
2476 terminal_result: String,
2477 },
2478}
2479
2480#[async_trait]
2488pub trait SubagentSpawnStore: Send + Sync + 'static {
2489 async fn try_claim_spawn(
2494 &self,
2495 parent_session_id: crate::typed_id::SessionId,
2496 tool_call_id: &str,
2497 claim_token: uuid::Uuid,
2498 ) -> Result<SpawnClaimResult>;
2499
2500 async fn register_child_session(
2505 &self,
2506 spawn_handle_id: uuid::Uuid,
2507 claim_token: uuid::Uuid,
2508 child_session_id: crate::typed_id::SessionId,
2509 ) -> Result<()>;
2510
2511 async fn settle_spawn(
2517 &self,
2518 parent_session_id: crate::typed_id::SessionId,
2519 tool_call_id: &str,
2520 claim_token: uuid::Uuid,
2521 terminal_status: &str,
2522 terminal_result: &str,
2523 ) -> Result<()>;
2524}
2525
2526#[async_trait]
2528impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2529 async fn try_claim_spawn(
2530 &self,
2531 parent_session_id: crate::typed_id::SessionId,
2532 tool_call_id: &str,
2533 claim_token: uuid::Uuid,
2534 ) -> Result<SpawnClaimResult> {
2535 (**self)
2536 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2537 .await
2538 }
2539
2540 async fn register_child_session(
2541 &self,
2542 spawn_handle_id: uuid::Uuid,
2543 claim_token: uuid::Uuid,
2544 child_session_id: crate::typed_id::SessionId,
2545 ) -> Result<()> {
2546 (**self)
2547 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2548 .await
2549 }
2550
2551 async fn settle_spawn(
2552 &self,
2553 parent_session_id: crate::typed_id::SessionId,
2554 tool_call_id: &str,
2555 claim_token: uuid::Uuid,
2556 terminal_status: &str,
2557 terminal_result: &str,
2558 ) -> Result<()> {
2559 (**self)
2560 .settle_spawn(
2561 parent_session_id,
2562 tool_call_id,
2563 claim_token,
2564 terminal_status,
2565 terminal_result,
2566 )
2567 .await
2568 }
2569}
2570
2571pub struct NoopSubagentSpawnStore;
2575
2576#[async_trait]
2577impl SubagentSpawnStore for NoopSubagentSpawnStore {
2578 async fn try_claim_spawn(
2579 &self,
2580 _parent_session_id: crate::typed_id::SessionId,
2581 _tool_call_id: &str,
2582 claim_token: uuid::Uuid,
2583 ) -> Result<SpawnClaimResult> {
2584 Ok(SpawnClaimResult::Claimed {
2585 spawn_handle_id: uuid::Uuid::new_v4(),
2586 claim_token,
2587 })
2588 }
2589
2590 async fn register_child_session(
2591 &self,
2592 _spawn_handle_id: uuid::Uuid,
2593 _claim_token: uuid::Uuid,
2594 _child_session_id: crate::typed_id::SessionId,
2595 ) -> Result<()> {
2596 Ok(())
2597 }
2598
2599 async fn settle_spawn(
2600 &self,
2601 _parent_session_id: crate::typed_id::SessionId,
2602 _tool_call_id: &str,
2603 _claim_token: uuid::Uuid,
2604 _terminal_status: &str,
2605 _terminal_result: &str,
2606 ) -> Result<()> {
2607 Ok(())
2608 }
2609}
2610
2611#[cfg(test)]
2616mod tests {
2617 use super::*;
2618
2619 #[test]
2620 fn test_resolved_image_new() {
2621 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2622 assert_eq!(image.base64, "SGVsbG8=");
2623 assert_eq!(image.media_type, "image/png");
2624 }
2625
2626 #[test]
2627 fn test_resolved_image_to_data_url() {
2628 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2629 let data_url = image.to_data_url();
2630 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2631 }
2632
2633 #[test]
2634 fn test_resolved_image_jpeg() {
2635 let image = ResolvedImage::new("base64data", "image/jpeg");
2636 let data_url = image.to_data_url();
2637 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2638 }
2639}