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(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1652pub enum ToolContextService {
1653 SessionFileSystem,
1654 SessionStorageStore,
1655 ImageArtifactStore,
1656 ProviderCredentialStore,
1657 UtilityLlmService,
1658 McpInvoker,
1659 EgressService,
1660 SessionSqlDbStore,
1661 MessageRetriever,
1662 SessionStore,
1663 SessionMutator,
1664 AgentStore,
1665 ConnectionResolver,
1666 ScheduleStore,
1667 SubagentSessionDelegate,
1668 KnowledgeStore,
1669 KnowledgeIndexSearch,
1670 LeasedResourceStore,
1671 SessionResourceRegistry,
1672 SessionTaskRegistry,
1673 EventEmitter,
1674 CapabilityRegistry,
1675 ToolRegistry,
1676 OrgId,
1677 BudgetChecker,
1678 PaymentAuthority,
1679 SessionCreationAuthority,
1680 SubagentSpawnStore,
1681 ReasoningEffortHandle,
1682}
1683
1684impl ToolContextService {
1685 pub const fn name(self) -> &'static str {
1686 match self {
1687 Self::SessionFileSystem => "SessionFileSystem",
1688 Self::SessionStorageStore => "SessionStorageStore",
1689 Self::ImageArtifactStore => "ImageArtifactStore",
1690 Self::ProviderCredentialStore => "ProviderCredentialStore",
1691 Self::UtilityLlmService => "UtilityLlmService",
1692 Self::McpInvoker => "McpInvoker",
1693 Self::EgressService => "EgressService",
1694 Self::SessionSqlDbStore => "SessionSqlDbStore",
1695 Self::MessageRetriever => "MessageRetriever",
1696 Self::SessionStore => "SessionStore",
1697 Self::SessionMutator => "SessionMutator",
1698 Self::AgentStore => "AgentStore",
1699 Self::ConnectionResolver => "ConnectionResolver",
1700 Self::ScheduleStore => "SessionScheduleStore",
1701 Self::SubagentSessionDelegate => "SubagentSessionDelegate",
1702 Self::KnowledgeStore => "KnowledgeStore",
1703 Self::KnowledgeIndexSearch => "KnowledgeIndexSearch",
1704 Self::LeasedResourceStore => "LeasedResourceStore",
1705 Self::SessionResourceRegistry => "SessionResourceRegistry",
1706 Self::SessionTaskRegistry => "SessionTaskRegistry",
1707 Self::EventEmitter => "EventEmitter",
1708 Self::CapabilityRegistry => "CapabilityRegistry",
1709 Self::ToolRegistry => "ToolRegistry",
1710 Self::OrgId => "OrgId",
1711 Self::BudgetChecker => "BudgetChecker",
1712 Self::PaymentAuthority => "PaymentAuthority",
1713 Self::SessionCreationAuthority => "SessionCreationAuthority",
1714 Self::SubagentSpawnStore => "SubagentSpawnStore",
1715 Self::ReasoningEffortHandle => "ReasoningEffortHandle",
1716 }
1717 }
1718}
1719
1720#[derive(Clone, Default)]
1725pub struct ToolContextServices {
1726 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1727 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1728 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1729 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1730 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1731 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1732 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1733 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1734 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1735 pub session_store: Option<Arc<dyn SessionStore>>,
1736 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1737 pub agent_store: Option<Arc<dyn AgentStore>>,
1738 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1739 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1740 pub subagent_delegate: Option<Arc<dyn crate::subagent_delegation::SubagentSessionDelegate>>,
1741 pub extensions: ToolContextExtensions,
1742 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1743 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1744 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1745 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1746 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1747 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1748 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1749 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1750 pub org_id: Option<crate::typed_id::OrgId>,
1751 pub network_access: Option<crate::network_access::NetworkAccessList>,
1752 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1753 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1754 pub session_creation_authority: Option<Arc<dyn SessionCreationAuthority>>,
1755 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1756 pub subagent_nesting_policy: SubagentNestingPolicy,
1757 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1758}
1759
1760impl ToolContextServices {
1761 pub fn provides(&self, service: ToolContextService) -> bool {
1762 match service {
1763 ToolContextService::SessionFileSystem => self.file_store.is_some(),
1764 ToolContextService::SessionStorageStore => self.storage_store.is_some(),
1765 ToolContextService::ImageArtifactStore => self.image_store.is_some(),
1766 ToolContextService::ProviderCredentialStore => self.provider_credential_store.is_some(),
1767 ToolContextService::UtilityLlmService => self.utility_llm_service.is_some(),
1768 ToolContextService::McpInvoker => self.mcp_invoker.is_some(),
1769 ToolContextService::EgressService => self.egress_service.is_some(),
1770 ToolContextService::SessionSqlDbStore => self.sqldb_store.is_some(),
1771 ToolContextService::MessageRetriever => self.message_retriever.is_some(),
1772 ToolContextService::SessionStore => self.session_store.is_some(),
1773 ToolContextService::SessionMutator => self.session_mutator.is_some(),
1774 ToolContextService::AgentStore => self.agent_store.is_some(),
1775 ToolContextService::ConnectionResolver => self.connection_resolver.is_some(),
1776 ToolContextService::ScheduleStore => self.schedule_store.is_some(),
1777 ToolContextService::SubagentSessionDelegate => self.subagent_delegate.is_some(),
1778 ToolContextService::KnowledgeStore => self.knowledge_store.is_some(),
1779 ToolContextService::KnowledgeIndexSearch => self.knowledge_index_search.is_some(),
1780 ToolContextService::LeasedResourceStore => self.leased_resource_store.is_some(),
1781 ToolContextService::SessionResourceRegistry => self.session_resource_registry.is_some(),
1782 ToolContextService::SessionTaskRegistry => self.session_task_registry.is_some(),
1783 ToolContextService::EventEmitter => self.event_emitter.is_some(),
1784 ToolContextService::CapabilityRegistry => self.capability_registry.is_some(),
1785 ToolContextService::ToolRegistry => self.tool_registry.is_some(),
1786 ToolContextService::OrgId => self.org_id.is_some(),
1787 ToolContextService::BudgetChecker => self.budget_checker.is_some(),
1788 ToolContextService::PaymentAuthority => self.payment_authority.is_some(),
1789 ToolContextService::SessionCreationAuthority => {
1790 self.session_creation_authority.is_some()
1791 }
1792 ToolContextService::SubagentSpawnStore => self.subagent_spawn_store.is_some(),
1793 ToolContextService::ReasoningEffortHandle => self.reasoning_effort_handle.is_some(),
1794 }
1795 }
1796}
1797
1798#[derive(Clone, Default)]
1803pub struct ToolContextExtensions {
1804 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
1805}
1806
1807impl ToolContextExtensions {
1808 pub fn insert<T: Any + Send + Sync>(&mut self, value: Arc<T>) {
1810 Arc::make_mut(&mut self.values).insert(TypeId::of::<T>(), value);
1811 }
1812
1813 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
1815 self.values
1816 .get(&TypeId::of::<T>())
1817 .and_then(|value| value.clone().downcast::<T>().ok())
1818 }
1819
1820 pub fn is_empty(&self) -> bool {
1821 self.values.is_empty()
1822 }
1823}
1824
1825impl std::fmt::Debug for ToolContextExtensions {
1826 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1827 f.debug_struct("ToolContextExtensions")
1828 .field("len", &self.values.len())
1829 .finish()
1830 }
1831}
1832
1833#[derive(Clone)]
1842pub struct ToolContext {
1843 pub session_id: SessionId,
1845 pub workspace_id: WorkspaceId,
1852
1853 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1855
1856 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1858
1859 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1861
1862 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1864
1865 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1867
1868 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1874
1875 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1877
1878 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1880
1881 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1883
1884 pub session_store: Option<Arc<dyn SessionStore>>,
1886
1887 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1889
1890 pub agent_store: Option<Arc<dyn AgentStore>>,
1892
1893 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1895
1896 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1898
1899 pub subagent_delegate: Option<Arc<dyn crate::subagent_delegation::SubagentSessionDelegate>>,
1904 pub extensions: ToolContextExtensions,
1908 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1910
1911 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1915
1916 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1918
1919 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1921
1922 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1925
1926 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1929
1930 pub event_context: Option<crate::events::EventContext>,
1933
1934 pub tool_call_id: Option<String>,
1937 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1939
1940 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1943
1944 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1948
1949 pub org_id: Option<crate::typed_id::OrgId>,
1951
1952 pub network_access: Option<crate::network_access::NetworkAccessList>,
1955
1956 pub locale: Option<String>,
1960
1961 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1963
1964 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1966
1967 pub session_creation_authority: Option<Arc<dyn SessionCreationAuthority>>,
1969
1970 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1974
1975 pub subagent_nesting_policy: SubagentNestingPolicy,
1977
1978 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1982
1983 pub cancellation: Option<tokio_util::sync::CancellationToken>,
1993}
1994
1995impl ToolContext {
1996 pub fn workspace_fs_key(&self) -> SessionId {
2001 SessionId::from_uuid(self.workspace_id.uuid())
2002 }
2003
2004 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
2006 self.workspace_id = workspace_id;
2007 self
2008 }
2009
2010 pub fn new(session_id: SessionId) -> Self {
2012 Self {
2013 session_id,
2014 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
2015 file_store: None,
2016 storage_store: None,
2017 image_store: None,
2018 provider_credential_store: None,
2019 utility_llm_service: None,
2020 mcp_invoker: None,
2021 egress_service: None,
2022 sqldb_store: None,
2023 message_retriever: None,
2024 session_store: None,
2025 session_mutator: None,
2026 agent_store: None,
2027 connection_resolver: None,
2028 schedule_store: None,
2029 subagent_delegate: None,
2030 extensions: ToolContextExtensions::default(),
2031 knowledge_store: None,
2032 knowledge_index_search: None,
2033 leased_resource_store: None,
2034 session_resource_registry: None,
2035 session_task_registry: None,
2036 event_emitter: None,
2037 event_context: None,
2038 tool_call_id: None,
2039 capability_registry: None,
2040 tool_registry: None,
2041 visible_tool_names: None,
2042 org_id: None,
2043 network_access: None,
2044 locale: None,
2045 budget_checker: None,
2046 payment_authority: None,
2047 session_creation_authority: None,
2048 subagent_spawn_store: None,
2049 subagent_nesting_policy: SubagentNestingPolicy::default(),
2050 reasoning_effort_handle: None,
2051 cancellation: None,
2052 }
2053 }
2054
2055 pub fn from_services(session_id: SessionId, services: &ToolContextServices) -> Self {
2057 Self {
2058 session_id,
2059 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
2060 file_store: services.file_store.clone(),
2061 storage_store: services.storage_store.clone(),
2062 image_store: services.image_store.clone(),
2063 provider_credential_store: services.provider_credential_store.clone(),
2064 utility_llm_service: services.utility_llm_service.clone(),
2065 mcp_invoker: services.mcp_invoker.clone(),
2066 egress_service: services.egress_service.clone(),
2067 sqldb_store: services.sqldb_store.clone(),
2068 message_retriever: services.message_retriever.clone(),
2069 session_store: services.session_store.clone(),
2070 session_mutator: services.session_mutator.clone(),
2071 agent_store: services.agent_store.clone(),
2072 connection_resolver: services.connection_resolver.clone(),
2073 schedule_store: services.schedule_store.clone(),
2074 subagent_delegate: services.subagent_delegate.clone(),
2075 extensions: services.extensions.clone(),
2076 knowledge_store: services.knowledge_store.clone(),
2077 knowledge_index_search: services.knowledge_index_search.clone(),
2078 leased_resource_store: services.leased_resource_store.clone(),
2079 session_resource_registry: services.session_resource_registry.clone(),
2080 session_task_registry: services.session_task_registry.clone(),
2081 event_emitter: services.event_emitter.clone(),
2082 event_context: None,
2083 tool_call_id: None,
2084 capability_registry: services.capability_registry.clone(),
2085 tool_registry: services.tool_registry.clone(),
2086 visible_tool_names: None,
2087 org_id: services.org_id,
2088 network_access: services.network_access.clone(),
2089 locale: None,
2090 budget_checker: services.budget_checker.clone(),
2091 payment_authority: services.payment_authority.clone(),
2092 session_creation_authority: services.session_creation_authority.clone(),
2093 subagent_spawn_store: services.subagent_spawn_store.clone(),
2094 subagent_nesting_policy: services.subagent_nesting_policy,
2095 reasoning_effort_handle: services.reasoning_effort_handle.clone(),
2096 cancellation: None,
2097 }
2098 }
2099
2100 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
2102 Self {
2103 session_id,
2104 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
2105 file_store: Some(file_store),
2106 storage_store: None,
2107 image_store: None,
2108 provider_credential_store: None,
2109 utility_llm_service: None,
2110 mcp_invoker: None,
2111 egress_service: None,
2112 sqldb_store: None,
2113 message_retriever: None,
2114 session_store: None,
2115 session_mutator: None,
2116 agent_store: None,
2117 connection_resolver: None,
2118 schedule_store: None,
2119 subagent_delegate: None,
2120 extensions: ToolContextExtensions::default(),
2121 knowledge_store: None,
2122 knowledge_index_search: None,
2123 leased_resource_store: None,
2124 session_resource_registry: None,
2125 session_task_registry: None,
2126 event_emitter: None,
2127 event_context: None,
2128 tool_call_id: None,
2129 capability_registry: None,
2130 tool_registry: None,
2131 visible_tool_names: None,
2132 org_id: None,
2133 network_access: None,
2134 locale: None,
2135 budget_checker: None,
2136 payment_authority: None,
2137 session_creation_authority: None,
2138 subagent_spawn_store: None,
2139 subagent_nesting_policy: SubagentNestingPolicy::default(),
2140 reasoning_effort_handle: None,
2141 cancellation: None,
2142 }
2143 }
2144
2145 pub fn with_storage_store(
2147 session_id: SessionId,
2148 storage_store: Arc<dyn SessionStorageStore>,
2149 ) -> Self {
2150 Self {
2151 session_id,
2152 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
2153 file_store: None,
2154 storage_store: Some(storage_store),
2155 image_store: None,
2156 provider_credential_store: None,
2157 utility_llm_service: None,
2158 mcp_invoker: None,
2159 egress_service: None,
2160 sqldb_store: None,
2161 message_retriever: None,
2162 session_store: None,
2163 session_mutator: None,
2164 agent_store: None,
2165 connection_resolver: None,
2166 schedule_store: None,
2167 subagent_delegate: None,
2168 extensions: ToolContextExtensions::default(),
2169 knowledge_store: None,
2170 knowledge_index_search: None,
2171 leased_resource_store: None,
2172 session_resource_registry: None,
2173 session_task_registry: None,
2174 event_emitter: None,
2175 event_context: None,
2176 tool_call_id: None,
2177 capability_registry: None,
2178 tool_registry: None,
2179 visible_tool_names: None,
2180 org_id: None,
2181 network_access: None,
2182 locale: None,
2183 budget_checker: None,
2184 payment_authority: None,
2185 session_creation_authority: None,
2186 subagent_spawn_store: None,
2187 subagent_nesting_policy: SubagentNestingPolicy::default(),
2188 reasoning_effort_handle: None,
2189 cancellation: None,
2190 }
2191 }
2192
2193 pub fn with_stores(
2195 session_id: SessionId,
2196 file_store: Arc<dyn SessionFileSystem>,
2197 storage_store: Arc<dyn SessionStorageStore>,
2198 ) -> Self {
2199 Self {
2200 session_id,
2201 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
2202 file_store: Some(file_store),
2203 storage_store: Some(storage_store),
2204 sqldb_store: None,
2205 image_store: None,
2206 provider_credential_store: None,
2207 utility_llm_service: None,
2208 mcp_invoker: None,
2209 egress_service: None,
2210 message_retriever: None,
2211 session_store: None,
2212 session_mutator: None,
2213 agent_store: None,
2214 connection_resolver: None,
2215 schedule_store: None,
2216 subagent_delegate: None,
2217 extensions: ToolContextExtensions::default(),
2218 knowledge_store: None,
2219 knowledge_index_search: None,
2220 leased_resource_store: None,
2221 session_resource_registry: None,
2222 session_task_registry: None,
2223 event_emitter: None,
2224 event_context: None,
2225 tool_call_id: None,
2226 capability_registry: None,
2227 tool_registry: None,
2228 visible_tool_names: None,
2229 org_id: None,
2230 network_access: None,
2231 locale: None,
2232 budget_checker: None,
2233 payment_authority: None,
2234 session_creation_authority: None,
2235 subagent_spawn_store: None,
2236 subagent_nesting_policy: SubagentNestingPolicy::default(),
2237 reasoning_effort_handle: None,
2238 cancellation: None,
2239 }
2240 }
2241
2242 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
2244 self.sqldb_store = Some(sqldb_store);
2245 self
2246 }
2247
2248 pub fn with_message_retriever(
2250 mut self,
2251 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
2252 ) -> Self {
2253 self.message_retriever = Some(retriever);
2254 self
2255 }
2256
2257 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
2259 self.session_store = Some(store);
2260 self
2261 }
2262
2263 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
2265 self.session_mutator = Some(mutator);
2266 self
2267 }
2268
2269 pub fn with_cancellation(mut self, token: tokio_util::sync::CancellationToken) -> Self {
2275 self.cancellation = Some(token);
2276 self
2277 }
2278
2279 pub fn is_cancelled(&self) -> bool {
2282 self.cancellation
2283 .as_ref()
2284 .is_some_and(|token| token.is_cancelled())
2285 }
2286
2287 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
2288 self.reasoning_effort_handle = Some(handle);
2289 self
2290 }
2291
2292 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
2294 self.agent_store = Some(store);
2295 self
2296 }
2297
2298 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
2300 self.connection_resolver = Some(resolver);
2301 self
2302 }
2303
2304 pub fn with_image_store(
2306 session_id: SessionId,
2307 image_store: Arc<dyn ImageArtifactStore>,
2308 ) -> Self {
2309 Self {
2310 session_id,
2311 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
2312 file_store: None,
2313 storage_store: None,
2314 image_store: Some(image_store),
2315 provider_credential_store: None,
2316 utility_llm_service: None,
2317 mcp_invoker: None,
2318 egress_service: None,
2319 sqldb_store: None,
2320 message_retriever: None,
2321 session_store: None,
2322 session_mutator: None,
2323 agent_store: None,
2324 connection_resolver: None,
2325 schedule_store: None,
2326 subagent_delegate: None,
2327 extensions: ToolContextExtensions::default(),
2328 knowledge_store: None,
2329 knowledge_index_search: None,
2330 leased_resource_store: None,
2331 session_resource_registry: None,
2332 session_task_registry: None,
2333 event_emitter: None,
2334 event_context: None,
2335 tool_call_id: None,
2336 capability_registry: None,
2337 tool_registry: None,
2338 visible_tool_names: None,
2339 org_id: None,
2340 network_access: None,
2341 locale: None,
2342 budget_checker: None,
2343 payment_authority: None,
2344 session_creation_authority: None,
2345 subagent_spawn_store: None,
2346 subagent_nesting_policy: SubagentNestingPolicy::default(),
2347 reasoning_effort_handle: None,
2348 cancellation: None,
2349 }
2350 }
2351
2352 pub fn with_provider_credential_store(
2354 mut self,
2355 store: Arc<dyn ProviderCredentialStore>,
2356 ) -> Self {
2357 self.provider_credential_store = Some(store);
2358 self
2359 }
2360
2361 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
2363 self.utility_llm_service = Some(service);
2364 self
2365 }
2366
2367 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
2369 self.mcp_invoker = Some(invoker);
2370 self
2371 }
2372
2373 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
2375 self.egress_service = Some(service);
2376 self
2377 }
2378
2379 pub fn with_egress_service_opt(
2382 mut self,
2383 service: Option<Arc<dyn crate::EgressService>>,
2384 ) -> Self {
2385 if let Some(service) = service {
2386 self.egress_service = Some(service);
2387 }
2388 self
2389 }
2390
2391 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
2393 self.storage_store = Some(store);
2394 self
2395 }
2396
2397 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
2399 self.schedule_store = Some(store);
2400 self
2401 }
2402
2403 pub fn with_subagent_delegate(
2405 mut self,
2406 delegate: Arc<dyn crate::subagent_delegation::SubagentSessionDelegate>,
2407 ) -> Self {
2408 self.subagent_delegate = Some(delegate);
2409 self
2410 }
2411
2412 pub fn with_extension<T: std::any::Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
2415 self.extensions.insert(value);
2416 self
2417 }
2418
2419 pub fn extension<T: std::any::Any + Send + Sync>(&self) -> Option<Arc<T>> {
2422 self.extensions.get::<T>()
2423 }
2424
2425 pub fn with_knowledge_index_search(
2427 mut self,
2428 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
2429 ) -> Self {
2430 self.knowledge_index_search = Some(search);
2431 self
2432 }
2433
2434 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
2436 self.leased_resource_store = Some(store);
2437 self
2438 }
2439
2440 pub fn with_session_resource_registry(
2442 mut self,
2443 registry: Arc<dyn SessionResourceRegistry>,
2444 ) -> Self {
2445 self.session_resource_registry = Some(registry);
2446 self
2447 }
2448
2449 pub fn with_session_task_registry(
2451 mut self,
2452 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
2453 ) -> Self {
2454 self.session_task_registry = Some(registry);
2455 self
2456 }
2457
2458 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
2460 self.org_id = Some(org_id);
2461 self
2462 }
2463
2464 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
2466 self.tool_registry = Some(registry);
2467 self
2468 }
2469
2470 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
2472 self.visible_tool_names = Some(names);
2473 self
2474 }
2475
2476 pub fn with_network_access(
2478 mut self,
2479 network_access: Option<crate::network_access::NetworkAccessList>,
2480 ) -> Self {
2481 self.network_access = network_access;
2482 self
2483 }
2484
2485 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
2487 self.payment_authority = Some(authority);
2488 self
2489 }
2490
2491 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
2493 self.subagent_spawn_store = Some(store);
2494 self
2495 }
2496
2497 pub fn with_subagent_nesting_policy(mut self, policy: SubagentNestingPolicy) -> Self {
2499 self.subagent_nesting_policy = policy;
2500 self
2501 }
2502
2503 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
2508 let (Some(emitter), Some(ctx), Some(call_id)) =
2509 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2510 else {
2511 return;
2512 };
2513 if let Err(e) = emitter
2514 .emit(EventRequest::new(
2515 self.session_id,
2516 ctx.clone(),
2517 crate::events::ToolProgressData {
2518 tool_call_id: call_id.clone(),
2519 tool_name: tool_name.to_string(),
2520 message: message.to_string(),
2521 display_name: None,
2522 },
2523 ))
2524 .await
2525 {
2526 tracing::debug!(
2527 tool_call_id = call_id,
2528 tool_name,
2529 error = %e,
2530 "Failed to emit tool.progress event"
2531 );
2532 }
2533 }
2534
2535 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
2540 let (Some(emitter), Some(ctx), Some(call_id)) =
2541 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2542 else {
2543 return;
2544 };
2545 if let Err(e) = emitter
2546 .emit(EventRequest::new(
2547 self.session_id,
2548 ctx.clone(),
2549 crate::events::ToolOutputDeltaData {
2550 tool_call_id: call_id.clone(),
2551 tool_name: tool_name.to_string(),
2552 delta: delta.to_string(),
2553 stream: stream.to_string(),
2554 },
2555 ))
2556 .await
2557 {
2558 tracing::debug!(
2559 tool_call_id = call_id,
2560 tool_name,
2561 error = %e,
2562 "Failed to emit tool.output.delta event"
2563 );
2564 }
2565 }
2566}
2567
2568impl std::fmt::Debug for ToolContext {
2569 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2570 f.debug_struct("ToolContext")
2571 .field("session_id", &self.session_id)
2572 .field("file_store", &self.file_store.is_some())
2573 .field("storage_store", &self.storage_store.is_some())
2574 .field("image_store", &self.image_store.is_some())
2575 .field(
2576 "provider_credential_store",
2577 &self.provider_credential_store.is_some(),
2578 )
2579 .field("utility_llm_service", &self.utility_llm_service.is_some())
2580 .field("egress_service", &self.egress_service.is_some())
2581 .field("sqldb_store", &self.sqldb_store.is_some())
2582 .field("message_retriever", &self.message_retriever.is_some())
2583 .field("session_store", &self.session_store.is_some())
2584 .field("session_mutator", &self.session_mutator.is_some())
2585 .field("agent_store", &self.agent_store.is_some())
2586 .field("connection_resolver", &self.connection_resolver.is_some())
2587 .field("schedule_store", &self.schedule_store.is_some())
2588 .field("subagent_delegate", &self.subagent_delegate.is_some())
2589 .field(
2590 "knowledge_index_search",
2591 &self.knowledge_index_search.is_some(),
2592 )
2593 .field(
2594 "leased_resource_store",
2595 &self.leased_resource_store.is_some(),
2596 )
2597 .field("event_emitter", &self.event_emitter.is_some())
2598 .field("tool_registry", &self.tool_registry.is_some())
2599 .field("payment_authority", &self.payment_authority.is_some())
2600 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
2601 .field("subagent_nesting_policy", &self.subagent_nesting_policy)
2602 .field("org_id", &self.org_id)
2603 .finish()
2604 }
2605}
2606
2607use crate::events::{Event, EventRequest};
2612
2613#[async_trait]
2624pub trait EventEmitter: Send + Sync {
2625 async fn emit(&self, request: EventRequest) -> Result<Event>;
2630}
2631
2632#[async_trait]
2634impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
2635 async fn emit(&self, request: EventRequest) -> Result<Event> {
2636 (**self).emit(request).await
2637 }
2638}
2639
2640#[derive(Debug, Clone, Default)]
2644pub struct NoopEventEmitter;
2645
2646#[async_trait]
2647impl EventEmitter for NoopEventEmitter {
2648 async fn emit(&self, request: EventRequest) -> Result<Event> {
2649 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2651 }
2652}
2653
2654#[derive(Debug, Clone)]
2667pub struct ResolvedImage {
2668 pub base64: String,
2670 pub media_type: String,
2672}
2673
2674impl ResolvedImage {
2675 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2677 Self {
2678 base64: base64.into(),
2679 media_type: media_type.into(),
2680 }
2681 }
2682
2683 pub fn to_data_url(&self) -> String {
2687 format!("data:{};base64,{}", self.media_type, self.base64)
2688 }
2689}
2690
2691#[async_trait]
2724pub trait ImageResolver: Send + Sync {
2725 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2729}
2730
2731#[derive(Debug)]
2737pub enum SpawnClaimResult {
2738 Claimed {
2741 spawn_handle_id: uuid::Uuid,
2742 claim_token: uuid::Uuid,
2743 },
2744 ClaimedPendingChild {
2748 spawn_handle_id: uuid::Uuid,
2749 claim_token: uuid::Uuid,
2750 },
2751 AlreadyRunning {
2754 child_session_id: crate::typed_id::SessionId,
2755 claim_token: uuid::Uuid,
2757 },
2758 AlreadySettled {
2761 child_session_id: crate::typed_id::SessionId,
2762 terminal_status: String,
2764 terminal_result: String,
2765 },
2766}
2767
2768#[async_trait]
2776pub trait SubagentSpawnStore: Send + Sync + 'static {
2777 async fn try_claim_spawn(
2782 &self,
2783 parent_session_id: crate::typed_id::SessionId,
2784 tool_call_id: &str,
2785 claim_token: uuid::Uuid,
2786 ) -> Result<SpawnClaimResult>;
2787
2788 async fn register_child_session(
2793 &self,
2794 spawn_handle_id: uuid::Uuid,
2795 claim_token: uuid::Uuid,
2796 child_session_id: crate::typed_id::SessionId,
2797 ) -> Result<()>;
2798
2799 async fn settle_spawn(
2805 &self,
2806 parent_session_id: crate::typed_id::SessionId,
2807 tool_call_id: &str,
2808 claim_token: uuid::Uuid,
2809 terminal_status: &str,
2810 terminal_result: &str,
2811 ) -> Result<()>;
2812}
2813
2814#[async_trait]
2816impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2817 async fn try_claim_spawn(
2818 &self,
2819 parent_session_id: crate::typed_id::SessionId,
2820 tool_call_id: &str,
2821 claim_token: uuid::Uuid,
2822 ) -> Result<SpawnClaimResult> {
2823 (**self)
2824 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2825 .await
2826 }
2827
2828 async fn register_child_session(
2829 &self,
2830 spawn_handle_id: uuid::Uuid,
2831 claim_token: uuid::Uuid,
2832 child_session_id: crate::typed_id::SessionId,
2833 ) -> Result<()> {
2834 (**self)
2835 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2836 .await
2837 }
2838
2839 async fn settle_spawn(
2840 &self,
2841 parent_session_id: crate::typed_id::SessionId,
2842 tool_call_id: &str,
2843 claim_token: uuid::Uuid,
2844 terminal_status: &str,
2845 terminal_result: &str,
2846 ) -> Result<()> {
2847 (**self)
2848 .settle_spawn(
2849 parent_session_id,
2850 tool_call_id,
2851 claim_token,
2852 terminal_status,
2853 terminal_result,
2854 )
2855 .await
2856 }
2857}
2858
2859pub struct NoopSubagentSpawnStore;
2863
2864#[async_trait]
2865impl SubagentSpawnStore for NoopSubagentSpawnStore {
2866 async fn try_claim_spawn(
2867 &self,
2868 _parent_session_id: crate::typed_id::SessionId,
2869 _tool_call_id: &str,
2870 claim_token: uuid::Uuid,
2871 ) -> Result<SpawnClaimResult> {
2872 Ok(SpawnClaimResult::Claimed {
2873 spawn_handle_id: uuid::Uuid::new_v4(),
2874 claim_token,
2875 })
2876 }
2877
2878 async fn register_child_session(
2879 &self,
2880 _spawn_handle_id: uuid::Uuid,
2881 _claim_token: uuid::Uuid,
2882 _child_session_id: crate::typed_id::SessionId,
2883 ) -> Result<()> {
2884 Ok(())
2885 }
2886
2887 async fn settle_spawn(
2888 &self,
2889 _parent_session_id: crate::typed_id::SessionId,
2890 _tool_call_id: &str,
2891 _claim_token: uuid::Uuid,
2892 _terminal_status: &str,
2893 _terminal_result: &str,
2894 ) -> Result<()> {
2895 Ok(())
2896 }
2897}
2898
2899#[cfg(test)]
2904mod tests {
2905 use super::*;
2906
2907 #[test]
2908 fn test_resolved_image_new() {
2909 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2910 assert_eq!(image.base64, "SGVsbG8=");
2911 assert_eq!(image.media_type, "image/png");
2912 }
2913
2914 #[test]
2915 fn test_resolved_image_to_data_url() {
2916 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2917 let data_url = image.to_data_url();
2918 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2919 }
2920
2921 #[test]
2922 fn test_resolved_image_jpeg() {
2923 let image = ResolvedImage::new("base64data", "image/jpeg");
2924 let data_url = image.to_data_url();
2925 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2926 }
2927}