1use crate::agent::Agent;
9use crate::harness::Harness;
10use crate::provider::DriverId;
11use crate::session_file::{FileInfo, FileStat, GrepMatch, InitialFile, SessionFile};
12use crate::tool_types::{ToolCall, ToolDefinition, ToolResult};
13use crate::typed_id::{AgentId, HarnessId, ImageId, ModelId, SessionId, WorkspaceId};
14use async_trait::async_trait;
15use chrono::{DateTime, Utc};
16use std::any::{Any, TypeId};
17use std::collections::{HashMap, HashSet};
18use std::sync::Arc;
19use uuid::Uuid;
20
21fn build_tool_map(tool_defs: &[ToolDefinition]) -> HashMap<&str, &ToolDefinition> {
23 tool_defs.iter().map(|def| (def.name(), def)).collect()
24}
25
26use crate::error::Result;
27
28#[derive(Clone, Default)]
43pub struct ReasoningEffortHandle {
44 inner: Arc<std::sync::RwLock<Option<String>>>,
45}
46
47impl ReasoningEffortHandle {
48 pub fn new() -> Self {
50 Self::default()
51 }
52
53 pub fn with_effort(effort: impl Into<String>) -> Self {
55 Self {
56 inner: Arc::new(std::sync::RwLock::new(Some(effort.into()))),
57 }
58 }
59
60 pub fn set(&self, effort: Option<String>) {
64 let mut guard = self.inner.write().unwrap_or_else(|e| e.into_inner());
68 *guard = effort;
69 }
70
71 pub fn get(&self) -> Option<String> {
73 let guard = self.inner.read().unwrap_or_else(|e| e.into_inner());
76 guard.clone()
77 }
78}
79
80impl std::fmt::Debug for ReasoningEffortHandle {
81 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82 f.debug_struct("ReasoningEffortHandle")
83 .field("effort", &self.get())
84 .finish()
85 }
86}
87
88#[async_trait]
99pub trait AgentStore: Send + Sync {
100 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>>;
102}
103
104#[async_trait]
105impl<T: AgentStore + ?Sized> AgentStore for std::sync::Arc<T> {
106 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>> {
107 (**self).get_agent(agent_id).await
108 }
109}
110
111#[async_trait]
126pub trait HarnessStore: Send + Sync {
127 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>>;
132}
133
134#[async_trait]
135impl<T: HarnessStore + ?Sized> HarnessStore for std::sync::Arc<T> {
136 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>> {
137 (**self).get_harness_chain(harness_id).await
138 }
139}
140
141use crate::leased_resource::{LeasedResource, UpsertLeasedResource};
146use crate::session::Session;
147
148#[async_trait]
154pub trait SessionStore: Send + Sync {
155 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>>;
157}
158
159#[async_trait]
160impl<T: SessionStore + ?Sized> SessionStore for std::sync::Arc<T> {
161 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>> {
162 (**self).get_session(session_id).await
163 }
164}
165
166#[async_trait]
168pub trait SessionMutator: Send + Sync {
169 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session>;
171}
172
173#[async_trait]
174impl<T: SessionMutator + ?Sized> SessionMutator for std::sync::Arc<T> {
175 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session> {
176 (**self).update_session_title(session_id, title).await
177 }
178}
179
180#[derive(Debug, Clone)]
186pub struct ResolvedModel {
187 pub model: String,
189 pub provider_type: DriverId,
191 pub api_key: Option<String>,
193 pub base_url: Option<String>,
195 pub provider_metadata: Option<crate::driver_registry::ProviderMetadata>,
198}
199
200#[async_trait]
210pub trait ProviderStore: Send + Sync {
211 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>>;
216
217 async fn get_default_model(&self) -> Result<Option<ResolvedModel>>;
221}
222
223#[async_trait]
224impl<T: ProviderStore + ?Sized> ProviderStore for std::sync::Arc<T> {
225 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>> {
226 (**self).get_resolved_model(model_id).await
227 }
228
229 async fn get_default_model(&self) -> Result<Option<ResolvedModel>> {
230 (**self).get_default_model().await
231 }
232}
233
234#[derive(Debug, Clone)]
240pub struct StoredImageInfo {
241 pub id: ImageId,
242 pub filename: String,
243 pub content_type: String,
244 pub size_bytes: i64,
245 pub metadata: serde_json::Value,
246 pub created_at: DateTime<Utc>,
247}
248
249#[derive(Debug, Clone)]
251pub struct StoredImage {
252 pub info: StoredImageInfo,
253 pub data: Vec<u8>,
254}
255
256#[derive(Debug, Clone)]
258pub struct CreateStoredImage {
259 pub filename: String,
260 pub content_type: String,
261 pub data: Vec<u8>,
262 pub metadata: serde_json::Value,
263}
264
265#[async_trait]
266pub trait ImageArtifactStore: Send + Sync {
267 async fn create_image(&self, input: CreateStoredImage) -> Result<StoredImageInfo>;
269
270 async fn get_image(&self, image_id: ImageId) -> Result<Option<StoredImage>>;
272
273 async fn get_image_info(&self, image_id: ImageId) -> Result<Option<StoredImageInfo>>;
275}
276
277#[derive(Debug, Clone)]
283pub struct ProviderCredentials {
284 pub api_key: String,
285 pub base_url: Option<String>,
286}
287
288#[async_trait]
289pub trait ProviderCredentialStore: Send + Sync {
290 async fn get_default_provider_credentials(
295 &self,
296 provider_type: &str,
297 ) -> Result<Option<ProviderCredentials>>;
298}
299
300#[async_trait]
311pub trait ToolExecutor: Send + Sync {
312 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult>;
317
318 async fn execute_with_context(
323 &self,
324 tool_call: &ToolCall,
325 tool_def: &ToolDefinition,
326 _context: &ToolContext,
327 ) -> Result<ToolResult> {
328 self.execute(tool_call, tool_def).await
330 }
331
332 async fn execute_batch(
334 &self,
335 tool_calls: &[ToolCall],
336 tool_defs: &[ToolDefinition],
337 ) -> Result<Vec<ToolResult>> {
338 let mut results = Vec::with_capacity(tool_calls.len());
339
340 let tool_map = build_tool_map(tool_defs);
341
342 for tool_call in tool_calls {
343 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
344 crate::error::AgentLoopError::tool(format!(
345 "Tool definition not found: {}",
346 tool_call.name
347 ))
348 })?;
349
350 results.push(self.execute(tool_call, tool_def).await?);
351 }
352
353 Ok(results)
354 }
355
356 async fn execute_parallel(
358 &self,
359 tool_calls: &[ToolCall],
360 tool_defs: &[ToolDefinition],
361 ) -> Result<Vec<ToolResult>>
362 where
363 Self: Sized,
364 {
365 use futures::future::join_all;
366
367 let tool_map = build_tool_map(tool_defs);
368
369 let futures: Vec<_> = tool_calls
370 .iter()
371 .map(|tool_call| async {
372 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
373 crate::error::AgentLoopError::tool(format!(
374 "Tool definition not found: {}",
375 tool_call.name
376 ))
377 })?;
378 self.execute(tool_call, tool_def).await
379 })
380 .collect();
381
382 let results = join_all(futures).await;
383 results.into_iter().collect()
384 }
385}
386
387#[async_trait]
391impl ToolExecutor for std::sync::Arc<dyn ToolExecutor> {
392 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult> {
393 (**self).execute(tool_call, tool_def).await
394 }
395
396 async fn execute_with_context(
397 &self,
398 tool_call: &ToolCall,
399 tool_def: &ToolDefinition,
400 context: &ToolContext,
401 ) -> Result<ToolResult> {
402 (**self)
403 .execute_with_context(tool_call, tool_def, context)
404 .await
405 }
406
407 async fn execute_batch(
408 &self,
409 tool_calls: &[ToolCall],
410 tool_defs: &[ToolDefinition],
411 ) -> Result<Vec<ToolResult>> {
412 (**self).execute_batch(tool_calls, tool_defs).await
413 }
414}
415
416#[async_trait]
428pub trait SessionFileSystem: Send + Sync {
429 fn display_root(&self) -> String {
436 crate::session_path::WORKSPACE_PREFIX.to_string()
437 }
438
439 fn display_path(&self, path: &str) -> String {
445 crate::session_path::to_display_path(path)
446 }
447
448 fn resolve_path(&self, input: &str) -> String {
459 crate::session_path::to_session_path(input)
460 }
461
462 fn is_mount_resolver(&self) -> bool;
466
467 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>>;
469
470 async fn write_file(
472 &self,
473 session_id: SessionId,
474 path: &str,
475 content: &str,
476 encoding: &str,
477 ) -> Result<SessionFile>;
478
479 async fn write_file_if_content_matches(
484 &self,
485 session_id: SessionId,
486 path: &str,
487 expected_content: &str,
488 expected_encoding: &str,
489 content: &str,
490 encoding: &str,
491 ) -> Result<Option<SessionFile>> {
492 let Some(existing) = self.read_file(session_id, path).await? else {
493 return Ok(None);
494 };
495
496 if existing.is_directory {
497 return Ok(None);
498 }
499
500 let current_content = existing.content.unwrap_or_default();
501 if current_content != expected_content || existing.encoding != expected_encoding {
502 return Ok(None);
503 }
504
505 self.write_file(session_id, path, content, encoding)
506 .await
507 .map(Some)
508 }
509
510 async fn delete_file(&self, session_id: SessionId, path: &str, recursive: bool)
512 -> Result<bool>;
513
514 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>>;
516
517 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>>;
519
520 async fn grep_files(
522 &self,
523 session_id: SessionId,
524 pattern: &str,
525 path_pattern: Option<&str>,
526 ) -> Result<Vec<GrepMatch>>;
527
528 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
530
531 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
533 if file.is_readonly {
534 return Err(crate::error::AgentLoopError::store(
535 "read-only initial files require a SessionFileSystem-specific seed implementation",
536 ));
537 }
538 self.write_file(session_id, &file.path, &file.content, &file.encoding)
539 .await?;
540 Ok(())
541 }
542}
543
544pub struct WorkspaceScopedFileSystem {
554 inner: Arc<dyn SessionFileSystem>,
555 key: SessionId,
556}
557
558impl WorkspaceScopedFileSystem {
559 pub fn wrap(
561 inner: Arc<dyn SessionFileSystem>,
562 workspace_id: WorkspaceId,
563 ) -> Arc<dyn SessionFileSystem> {
564 Arc::new(Self {
565 inner,
566 key: SessionId::from_uuid(workspace_id.uuid()),
567 })
568 }
569}
570
571#[async_trait]
572impl SessionFileSystem for WorkspaceScopedFileSystem {
573 async fn read_file(&self, _session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
574 self.inner.read_file(self.key, path).await
575 }
576 async fn write_file(
577 &self,
578 _session_id: SessionId,
579 path: &str,
580 content: &str,
581 encoding: &str,
582 ) -> Result<SessionFile> {
583 self.inner
584 .write_file(self.key, path, content, encoding)
585 .await
586 }
587 async fn write_file_if_content_matches(
588 &self,
589 _session_id: SessionId,
590 path: &str,
591 expected_content: &str,
592 expected_encoding: &str,
593 content: &str,
594 encoding: &str,
595 ) -> Result<Option<SessionFile>> {
596 self.inner
597 .write_file_if_content_matches(
598 self.key,
599 path,
600 expected_content,
601 expected_encoding,
602 content,
603 encoding,
604 )
605 .await
606 }
607 async fn delete_file(
608 &self,
609 _session_id: SessionId,
610 path: &str,
611 recursive: bool,
612 ) -> Result<bool> {
613 self.inner.delete_file(self.key, path, recursive).await
614 }
615 async fn list_directory(&self, _session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
616 self.inner.list_directory(self.key, path).await
617 }
618 async fn stat_file(&self, _session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
619 self.inner.stat_file(self.key, path).await
620 }
621 async fn grep_files(
622 &self,
623 _session_id: SessionId,
624 pattern: &str,
625 path_pattern: Option<&str>,
626 ) -> Result<Vec<GrepMatch>> {
627 self.inner.grep_files(self.key, pattern, path_pattern).await
628 }
629 async fn create_directory(&self, _session_id: SessionId, path: &str) -> Result<FileInfo> {
630 self.inner.create_directory(self.key, path).await
631 }
632 async fn seed_initial_file(&self, _session_id: SessionId, file: &InitialFile) -> Result<()> {
633 self.inner.seed_initial_file(self.key, file).await
634 }
635
636 fn display_root(&self) -> String {
637 self.inner.display_root()
638 }
639
640 fn display_path(&self, path: &str) -> String {
641 self.inner.display_path(path)
642 }
643
644 fn resolve_path(&self, input: &str) -> String {
645 self.inner.resolve_path(input)
646 }
647
648 fn is_mount_resolver(&self) -> bool {
649 self.inner.is_mount_resolver()
650 }
651}
652
653#[async_trait]
654impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
655 fn display_root(&self) -> String {
656 (**self).display_root()
657 }
658
659 fn display_path(&self, path: &str) -> String {
660 (**self).display_path(path)
661 }
662
663 fn resolve_path(&self, input: &str) -> String {
664 (**self).resolve_path(input)
665 }
666
667 fn is_mount_resolver(&self) -> bool {
668 (**self).is_mount_resolver()
669 }
670
671 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
672 (**self).read_file(session_id, path).await
673 }
674
675 async fn write_file(
676 &self,
677 session_id: SessionId,
678 path: &str,
679 content: &str,
680 encoding: &str,
681 ) -> Result<SessionFile> {
682 (**self)
683 .write_file(session_id, path, content, encoding)
684 .await
685 }
686
687 async fn write_file_if_content_matches(
688 &self,
689 session_id: SessionId,
690 path: &str,
691 expected_content: &str,
692 expected_encoding: &str,
693 content: &str,
694 encoding: &str,
695 ) -> Result<Option<SessionFile>> {
696 (**self)
697 .write_file_if_content_matches(
698 session_id,
699 path,
700 expected_content,
701 expected_encoding,
702 content,
703 encoding,
704 )
705 .await
706 }
707
708 async fn delete_file(
709 &self,
710 session_id: SessionId,
711 path: &str,
712 recursive: bool,
713 ) -> Result<bool> {
714 (**self).delete_file(session_id, path, recursive).await
715 }
716
717 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
718 (**self).list_directory(session_id, path).await
719 }
720
721 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
722 (**self).stat_file(session_id, path).await
723 }
724
725 async fn grep_files(
726 &self,
727 session_id: SessionId,
728 pattern: &str,
729 path_pattern: Option<&str>,
730 ) -> Result<Vec<GrepMatch>> {
731 (**self).grep_files(session_id, pattern, path_pattern).await
732 }
733
734 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
735 (**self).create_directory(session_id, path).await
736 }
737
738 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
739 (**self).seed_initial_file(session_id, file).await
740 }
741}
742
743pub use SessionFileSystem as SessionFileStore;
745
746#[derive(Clone, Default)]
752pub struct SessionFileSystemFactoryContext {
753 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
754}
755
756impl SessionFileSystemFactoryContext {
757 pub fn new() -> Self {
758 Self::default()
759 }
760
761 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
762 let values = Arc::make_mut(&mut self.values);
763 values.insert(TypeId::of::<T>(), value);
764 self
765 }
766
767 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
768 self.values
769 .get(&TypeId::of::<T>())
770 .and_then(|value| value.clone().downcast::<T>().ok())
771 }
772
773 pub fn with_workspace_roots(self, roots: Arc<crate::WorkspaceRootSet>) -> Self {
774 self.with(roots)
775 }
776
777 pub fn workspace_roots(&self) -> Option<Arc<crate::WorkspaceRootSet>> {
778 self.get::<crate::WorkspaceRootSet>()
779 }
780}
781
782#[async_trait]
784pub trait SessionFileSystemFactory: Send + Sync {
785 fn name(&self) -> &'static str {
787 "SessionFileSystemFactory"
788 }
789
790 fn is_disabled(&self) -> bool {
793 false
794 }
795
796 async fn create_session_file_system(
798 &self,
799 context: SessionFileSystemFactoryContext,
800 ) -> Result<Arc<dyn SessionFileSystem>>;
801}
802
803#[derive(Debug, Clone, Default)]
805pub struct DisabledSessionFileSystemFactory;
806
807#[async_trait]
808impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
809 fn name(&self) -> &'static str {
810 "DisabledSessionFileSystemFactory"
811 }
812
813 fn is_disabled(&self) -> bool {
814 true
815 }
816
817 async fn create_session_file_system(
818 &self,
819 _context: SessionFileSystemFactoryContext,
820 ) -> Result<Arc<dyn SessionFileSystem>> {
821 Err(crate::error::AgentLoopError::config(
822 "session filesystem is disabled",
823 ))
824 }
825}
826
827#[derive(Debug, Clone)]
833pub struct KeyInfo {
834 pub key: String,
835 pub created_at: chrono::DateTime<chrono::Utc>,
836 pub updated_at: chrono::DateTime<chrono::Utc>,
837}
838
839#[derive(Debug, Clone)]
841pub struct SecretInfo {
842 pub name: String,
843 pub created_at: chrono::DateTime<chrono::Utc>,
844 pub updated_at: chrono::DateTime<chrono::Utc>,
845}
846
847#[derive(Debug, Clone, serde::Serialize)]
857pub struct KnowledgeSearchHit {
858 pub id: String,
860 pub kb_id: String,
862 pub title: String,
863 pub kind: String,
864 pub tags: Vec<String>,
865 pub snippet: String,
867 pub resource: Option<String>,
869}
870
871#[async_trait]
875pub trait KnowledgeStore: Send + Sync {
876 async fn search_knowledge(
877 &self,
878 org_id: crate::typed_id::OrgId,
879 kb_public_ids: &[String],
880 query: &str,
881 kind: Option<&str>,
882 tags: &[String],
883 limit: usize,
884 ) -> Result<Vec<KnowledgeSearchHit>>;
885}
886
887#[async_trait]
892pub trait SessionStorageStore: Send + Sync {
893 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
897
898 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
900
901 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
903
904 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
906
907 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
911
912 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
914
915 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
917
918 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
920}
921
922use crate::session_schedule::SessionSchedule;
927use crate::typed_id::ScheduleId;
928
929#[async_trait]
933pub trait SessionScheduleStore: Send + Sync {
934 async fn create_schedule(
936 &self,
937 session_id: SessionId,
938 description: String,
939 cron_expression: Option<String>,
940 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
941 timezone: String,
942 ) -> Result<SessionSchedule>;
943
944 async fn create_schedule_enforcing_limits(
948 &self,
949 session_id: SessionId,
950 description: String,
951 cron_expression: Option<String>,
952 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
953 timezone: String,
954 ) -> std::result::Result<SessionSchedule, crate::session_schedule::ScheduleLimitError> {
955 let per_session = self
956 .count_active_schedules(session_id)
957 .await
958 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
959 if per_session >= crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION {
960 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
961 format!(
962 "Maximum {} active schedules per session. Cancel an existing schedule first.",
963 crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION
964 ),
965 ));
966 }
967
968 let max_per_org = crate::session_schedule::max_active_schedules_per_org();
969 let per_org = self
970 .count_active_org_schedules()
971 .await
972 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
973 if i64::from(per_org) >= max_per_org {
974 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
975 format!(
976 "Maximum {max_per_org} active schedules per org reached. Cancel an existing schedule first."
977 ),
978 ));
979 }
980
981 if let Some(cron) = cron_expression.as_deref() {
982 crate::session_schedule::validate_cron_min_interval(cron)
983 .map_err(crate::session_schedule::ScheduleLimitError::Rejected)?;
984 }
985
986 self.create_schedule(
987 session_id,
988 description,
989 cron_expression,
990 scheduled_at,
991 timezone,
992 )
993 .await
994 .map_err(crate::session_schedule::ScheduleLimitError::Store)
995 }
996
997 async fn cancel_schedule(
999 &self,
1000 session_id: SessionId,
1001 schedule_id: ScheduleId,
1002 ) -> Result<SessionSchedule>;
1003
1004 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
1006
1007 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
1009
1010 async fn count_active_org_schedules(&self) -> Result<u32>;
1015}
1016
1017#[async_trait]
1027pub trait SessionResourceRegistry: Send + Sync {
1028 async fn register(
1030 &self,
1031 entry: crate::session_resource::RegisterSessionResource,
1032 ) -> Result<crate::session_resource::SessionResourceEntry>;
1033
1034 async fn update_status(
1036 &self,
1037 session_id: SessionId,
1038 resource_id: &str,
1039 status: crate::session_resource::SessionResourceStatus,
1040 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1041
1042 async fn get(
1044 &self,
1045 session_id: SessionId,
1046 resource_id: &str,
1047 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1048
1049 async fn list(
1051 &self,
1052 session_id: SessionId,
1053 filter: Option<&crate::session_resource::SessionResourceFilter>,
1054 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
1055
1056 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
1058}
1059
1060#[async_trait]
1070pub trait LeasedResourceStore: Send + Sync {
1071 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
1077
1078 async fn release_resource(
1084 &self,
1085 session_id: SessionId,
1086 provider: &str,
1087 resource_type: &str,
1088 external_id: &str,
1089 ) -> Result<Option<LeasedResource>>;
1090
1091 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
1096}
1097
1098pub const DEFAULT_MAX_SUBAGENT_DEPTH: u32 = 2;
1105pub const DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS: u32 = 16;
1106pub const DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS: u32 = 200;
1107
1108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1110pub struct SubagentNestingPolicy {
1111 pub platform_default: u32,
1112 pub org_override: Option<u32>,
1113 pub agent_override: Option<u32>,
1114 pub platform_default_max_active_descendant_tasks: u32,
1115 pub org_override_max_active_descendant_tasks: Option<u32>,
1116 pub agent_override_max_active_descendant_tasks: Option<u32>,
1117 pub platform_default_max_total_descendant_tasks: u32,
1118 pub org_override_max_total_descendant_tasks: Option<u32>,
1119 pub agent_override_max_total_descendant_tasks: Option<u32>,
1120}
1121
1122impl Default for SubagentNestingPolicy {
1123 fn default() -> Self {
1124 Self {
1125 platform_default: DEFAULT_MAX_SUBAGENT_DEPTH,
1126 org_override: None,
1127 agent_override: None,
1128 platform_default_max_active_descendant_tasks:
1129 DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS,
1130 org_override_max_active_descendant_tasks: None,
1131 agent_override_max_active_descendant_tasks: None,
1132 platform_default_max_total_descendant_tasks:
1133 DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS,
1134 org_override_max_total_descendant_tasks: None,
1135 agent_override_max_total_descendant_tasks: None,
1136 }
1137 }
1138}
1139
1140impl SubagentNestingPolicy {
1141 pub fn max_subagent_depth(self) -> u32 {
1142 self.agent_override
1143 .or(self.org_override)
1144 .unwrap_or(self.platform_default)
1145 }
1146
1147 pub fn max_active_descendant_tasks(self) -> u32 {
1148 self.agent_override_max_active_descendant_tasks
1149 .or(self.org_override_max_active_descendant_tasks)
1150 .unwrap_or(self.platform_default_max_active_descendant_tasks)
1151 }
1152
1153 pub fn max_total_descendant_tasks(self) -> u32 {
1154 self.agent_override_max_total_descendant_tasks
1155 .or(self.org_override_max_total_descendant_tasks)
1156 .unwrap_or(self.platform_default_max_total_descendant_tasks)
1157 }
1158
1159 pub fn with_platform_default(mut self, depth: u32) -> Self {
1160 self.platform_default = depth;
1161 self
1162 }
1163
1164 pub fn with_org_override(mut self, depth: Option<u32>) -> Self {
1165 self.org_override = depth;
1166 self
1167 }
1168
1169 pub fn with_agent_override(mut self, depth: Option<u32>) -> Self {
1170 self.agent_override = depth;
1171 self
1172 }
1173
1174 pub fn with_agent_task_caps_override(
1175 mut self,
1176 max_active: Option<u32>,
1177 max_total: Option<u32>,
1178 ) -> Self {
1179 self.agent_override_max_active_descendant_tasks = max_active;
1180 self.agent_override_max_total_descendant_tasks = max_total;
1181 self
1182 }
1183}
1184
1185pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
1187
1188#[async_trait]
1193pub trait UserConnectionResolver: Send + Sync {
1194 async fn get_connection_token(
1197 &self,
1198 session_id: SessionId,
1199 provider: &str,
1200 ) -> Result<Option<String>>;
1201
1202 async fn get_connection_user(
1207 &self,
1208 _session_id: SessionId,
1209 _provider: &str,
1210 ) -> Result<Option<Uuid>> {
1211 Ok(None)
1212 }
1213
1214 async fn get_connection_token_for_user(
1219 &self,
1220 _user_id: Uuid,
1221 _provider: &str,
1222 ) -> Result<Option<String>> {
1223 Ok(None)
1224 }
1225
1226 async fn get_connection_metadata(
1229 &self,
1230 _session_id: SessionId,
1231 _provider: &str,
1232 ) -> Result<Option<serde_json::Value>> {
1233 Ok(None)
1234 }
1235}
1236
1237#[async_trait]
1247pub trait BudgetChecker: Send + Sync {
1248 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
1250}
1251
1252#[async_trait]
1261pub trait PaymentAuthority: Send + Sync {
1262 async fn execute_machine_payment(
1263 &self,
1264 session_id: SessionId,
1265 request: crate::payment::MachinePaymentRequest,
1266 ) -> Result<crate::payment::MachinePaymentResponse>;
1267}
1268
1269#[async_trait]
1279pub trait OutboundToolRateLimiter: Send + Sync {
1280 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
1282}
1283
1284#[derive(Debug)]
1290pub enum ToolCallClaimResult {
1291 Claimed { claim_token: uuid::Uuid },
1294 AlreadySettled {
1296 result_json: serde_json::Value,
1297 args_fingerprint: String,
1298 },
1299 AlreadyRunning { args_fingerprint: String },
1304 DeterminismViolation {
1308 stored_fingerprint: String,
1309 current_fingerprint: String,
1310 },
1311}
1312
1313#[derive(Debug, Clone)]
1315pub enum DurableToolCallStatus {
1316 Settled { result_json: serde_json::Value },
1318 Interrupted {
1320 result_json: Option<serde_json::Value>,
1321 },
1322 Running,
1324}
1325
1326#[async_trait]
1331pub trait DurableToolResultStore: Send + Sync + 'static {
1332 async fn try_claim_tool_call(
1340 &self,
1341 turn_id: &str,
1342 tool_call_id: &str,
1343 tool_name: &str,
1344 args_fingerprint: &str,
1345 ) -> Result<ToolCallClaimResult>;
1346
1347 async fn settle_tool_call(
1353 &self,
1354 turn_id: &str,
1355 tool_call_id: &str,
1356 result_json: serde_json::Value,
1357 status: &str,
1358 claim_token: uuid::Uuid,
1359 ) -> Result<bool>;
1360
1361 async fn get_tool_call_status(
1366 &self,
1367 turn_id: &str,
1368 tool_call_id: &str,
1369 ) -> Result<Option<DurableToolCallStatus>>;
1370}
1371
1372pub struct NoopDurableToolResultStore;
1375
1376#[async_trait]
1377impl DurableToolResultStore for NoopDurableToolResultStore {
1378 async fn try_claim_tool_call(
1379 &self,
1380 _turn_id: &str,
1381 _tool_call_id: &str,
1382 _tool_name: &str,
1383 _args_fingerprint: &str,
1384 ) -> Result<ToolCallClaimResult> {
1385 Ok(ToolCallClaimResult::Claimed {
1386 claim_token: uuid::Uuid::new_v4(),
1387 })
1388 }
1389
1390 async fn settle_tool_call(
1391 &self,
1392 _turn_id: &str,
1393 _tool_call_id: &str,
1394 _result_json: serde_json::Value,
1395 _status: &str,
1396 _claim_token: uuid::Uuid,
1397 ) -> Result<bool> {
1398 Ok(true)
1399 }
1400
1401 async fn get_tool_call_status(
1402 &self,
1403 _turn_id: &str,
1404 _tool_call_id: &str,
1405 ) -> Result<Option<DurableToolCallStatus>> {
1406 Ok(None)
1407 }
1408}
1409
1410#[derive(Debug, Clone)]
1416pub struct StreamProgress {
1417 pub accumulated_len: usize,
1419 pub last_delta_at: u64,
1421}
1422
1423#[async_trait]
1429pub trait StreamHeartbeater: Send + Sync {
1430 async fn heartbeat(&self, progress: StreamProgress);
1436}
1437
1438pub struct NoopStreamHeartbeater;
1440
1441#[async_trait]
1442impl StreamHeartbeater for NoopStreamHeartbeater {
1443 async fn heartbeat(&self, _progress: StreamProgress) {}
1444}
1445
1446#[derive(Debug, Clone)]
1452pub struct PartialStreamState {
1453 pub accumulated: String,
1456}
1457
1458#[async_trait]
1466pub trait PartialStreamStore: Send + Sync {
1467 async fn get_partial_stream(
1470 &self,
1471 session_id: SessionId,
1472 turn_id: &str,
1473 ) -> Result<Option<PartialStreamState>>;
1474}
1475
1476pub struct NoopPartialStreamStore;
1478
1479#[async_trait]
1480impl PartialStreamStore for NoopPartialStreamStore {
1481 async fn get_partial_stream(
1482 &self,
1483 _session_id: SessionId,
1484 _turn_id: &str,
1485 ) -> Result<Option<PartialStreamState>> {
1486 Ok(None)
1487 }
1488}
1489
1490#[derive(Clone)]
1499pub struct ToolContext {
1500 pub session_id: SessionId,
1502 pub workspace_id: WorkspaceId,
1509
1510 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1512
1513 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1515
1516 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1518
1519 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1521
1522 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1524
1525 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1531
1532 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1534
1535 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1537
1538 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1540
1541 pub session_store: Option<Arc<dyn SessionStore>>,
1543
1544 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1546
1547 pub agent_store: Option<Arc<dyn AgentStore>>,
1549
1550 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1552
1553 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1555
1556 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1558 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1560
1561 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1565
1566 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1568
1569 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1571
1572 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1575
1576 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1579
1580 pub event_context: Option<crate::events::EventContext>,
1583
1584 pub tool_call_id: Option<String>,
1587 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1589
1590 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1593
1594 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1598
1599 pub org_id: Option<crate::typed_id::OrgId>,
1601
1602 pub network_access: Option<crate::network_access::NetworkAccessList>,
1605
1606 pub locale: Option<String>,
1610
1611 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1613
1614 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1616
1617 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1621
1622 pub subagent_nesting_policy: SubagentNestingPolicy,
1624
1625 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1629}
1630
1631impl ToolContext {
1632 pub fn workspace_fs_key(&self) -> SessionId {
1637 SessionId::from_uuid(self.workspace_id.uuid())
1638 }
1639
1640 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1642 self.workspace_id = workspace_id;
1643 self
1644 }
1645
1646 pub fn new(session_id: SessionId) -> Self {
1648 Self {
1649 session_id,
1650 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1651 file_store: None,
1652 storage_store: None,
1653 image_store: None,
1654 provider_credential_store: None,
1655 utility_llm_service: None,
1656 mcp_invoker: None,
1657 egress_service: None,
1658 sqldb_store: None,
1659 message_retriever: None,
1660 session_store: None,
1661 session_mutator: None,
1662 agent_store: None,
1663 connection_resolver: None,
1664 schedule_store: None,
1665 platform_store: None,
1666 knowledge_store: None,
1667 knowledge_index_search: None,
1668 leased_resource_store: None,
1669 session_resource_registry: None,
1670 session_task_registry: None,
1671 event_emitter: None,
1672 event_context: None,
1673 tool_call_id: None,
1674 capability_registry: None,
1675 tool_registry: None,
1676 visible_tool_names: None,
1677 org_id: None,
1678 network_access: None,
1679 locale: None,
1680 budget_checker: None,
1681 payment_authority: None,
1682 subagent_spawn_store: None,
1683 subagent_nesting_policy: SubagentNestingPolicy::default(),
1684 reasoning_effort_handle: None,
1685 }
1686 }
1687
1688 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1690 Self {
1691 session_id,
1692 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1693 file_store: Some(file_store),
1694 storage_store: None,
1695 image_store: None,
1696 provider_credential_store: None,
1697 utility_llm_service: None,
1698 mcp_invoker: None,
1699 egress_service: None,
1700 sqldb_store: None,
1701 message_retriever: None,
1702 session_store: None,
1703 session_mutator: None,
1704 agent_store: None,
1705 connection_resolver: None,
1706 schedule_store: None,
1707 platform_store: None,
1708 knowledge_store: None,
1709 knowledge_index_search: None,
1710 leased_resource_store: None,
1711 session_resource_registry: None,
1712 session_task_registry: None,
1713 event_emitter: None,
1714 event_context: None,
1715 tool_call_id: None,
1716 capability_registry: None,
1717 tool_registry: None,
1718 visible_tool_names: None,
1719 org_id: None,
1720 network_access: None,
1721 locale: None,
1722 budget_checker: None,
1723 payment_authority: None,
1724 subagent_spawn_store: None,
1725 subagent_nesting_policy: SubagentNestingPolicy::default(),
1726 reasoning_effort_handle: None,
1727 }
1728 }
1729
1730 pub fn with_storage_store(
1732 session_id: SessionId,
1733 storage_store: Arc<dyn SessionStorageStore>,
1734 ) -> Self {
1735 Self {
1736 session_id,
1737 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1738 file_store: None,
1739 storage_store: Some(storage_store),
1740 image_store: None,
1741 provider_credential_store: None,
1742 utility_llm_service: None,
1743 mcp_invoker: None,
1744 egress_service: None,
1745 sqldb_store: None,
1746 message_retriever: None,
1747 session_store: None,
1748 session_mutator: None,
1749 agent_store: None,
1750 connection_resolver: None,
1751 schedule_store: None,
1752 platform_store: None,
1753 knowledge_store: None,
1754 knowledge_index_search: None,
1755 leased_resource_store: None,
1756 session_resource_registry: None,
1757 session_task_registry: None,
1758 event_emitter: None,
1759 event_context: None,
1760 tool_call_id: None,
1761 capability_registry: None,
1762 tool_registry: None,
1763 visible_tool_names: None,
1764 org_id: None,
1765 network_access: None,
1766 locale: None,
1767 budget_checker: None,
1768 payment_authority: None,
1769 subagent_spawn_store: None,
1770 subagent_nesting_policy: SubagentNestingPolicy::default(),
1771 reasoning_effort_handle: None,
1772 }
1773 }
1774
1775 pub fn with_stores(
1777 session_id: SessionId,
1778 file_store: Arc<dyn SessionFileSystem>,
1779 storage_store: Arc<dyn SessionStorageStore>,
1780 ) -> Self {
1781 Self {
1782 session_id,
1783 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1784 file_store: Some(file_store),
1785 storage_store: Some(storage_store),
1786 sqldb_store: None,
1787 image_store: None,
1788 provider_credential_store: None,
1789 utility_llm_service: None,
1790 mcp_invoker: None,
1791 egress_service: None,
1792 message_retriever: None,
1793 session_store: None,
1794 session_mutator: None,
1795 agent_store: None,
1796 connection_resolver: None,
1797 schedule_store: None,
1798 platform_store: None,
1799 knowledge_store: None,
1800 knowledge_index_search: None,
1801 leased_resource_store: None,
1802 session_resource_registry: None,
1803 session_task_registry: None,
1804 event_emitter: None,
1805 event_context: None,
1806 tool_call_id: None,
1807 capability_registry: None,
1808 tool_registry: None,
1809 visible_tool_names: None,
1810 org_id: None,
1811 network_access: None,
1812 locale: None,
1813 budget_checker: None,
1814 payment_authority: None,
1815 subagent_spawn_store: None,
1816 subagent_nesting_policy: SubagentNestingPolicy::default(),
1817 reasoning_effort_handle: None,
1818 }
1819 }
1820
1821 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1823 self.sqldb_store = Some(sqldb_store);
1824 self
1825 }
1826
1827 pub fn with_message_retriever(
1829 mut self,
1830 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1831 ) -> Self {
1832 self.message_retriever = Some(retriever);
1833 self
1834 }
1835
1836 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1838 self.session_store = Some(store);
1839 self
1840 }
1841
1842 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1844 self.session_mutator = Some(mutator);
1845 self
1846 }
1847
1848 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
1852 self.reasoning_effort_handle = Some(handle);
1853 self
1854 }
1855
1856 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1858 self.agent_store = Some(store);
1859 self
1860 }
1861
1862 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1864 self.connection_resolver = Some(resolver);
1865 self
1866 }
1867
1868 pub fn with_image_store(
1870 session_id: SessionId,
1871 image_store: Arc<dyn ImageArtifactStore>,
1872 ) -> Self {
1873 Self {
1874 session_id,
1875 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1876 file_store: None,
1877 storage_store: None,
1878 image_store: Some(image_store),
1879 provider_credential_store: None,
1880 utility_llm_service: None,
1881 mcp_invoker: None,
1882 egress_service: None,
1883 sqldb_store: None,
1884 message_retriever: None,
1885 session_store: None,
1886 session_mutator: None,
1887 agent_store: None,
1888 connection_resolver: None,
1889 schedule_store: None,
1890 platform_store: None,
1891 knowledge_store: None,
1892 knowledge_index_search: None,
1893 leased_resource_store: None,
1894 session_resource_registry: None,
1895 session_task_registry: None,
1896 event_emitter: None,
1897 event_context: None,
1898 tool_call_id: None,
1899 capability_registry: None,
1900 tool_registry: None,
1901 visible_tool_names: None,
1902 org_id: None,
1903 network_access: None,
1904 locale: None,
1905 budget_checker: None,
1906 payment_authority: None,
1907 subagent_spawn_store: None,
1908 subagent_nesting_policy: SubagentNestingPolicy::default(),
1909 reasoning_effort_handle: None,
1910 }
1911 }
1912
1913 pub fn with_provider_credential_store(
1915 mut self,
1916 store: Arc<dyn ProviderCredentialStore>,
1917 ) -> Self {
1918 self.provider_credential_store = Some(store);
1919 self
1920 }
1921
1922 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1924 self.utility_llm_service = Some(service);
1925 self
1926 }
1927
1928 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
1930 self.mcp_invoker = Some(invoker);
1931 self
1932 }
1933
1934 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1936 self.egress_service = Some(service);
1937 self
1938 }
1939
1940 pub fn with_egress_service_opt(
1943 mut self,
1944 service: Option<Arc<dyn crate::EgressService>>,
1945 ) -> Self {
1946 if let Some(service) = service {
1947 self.egress_service = Some(service);
1948 }
1949 self
1950 }
1951
1952 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
1954 self.storage_store = Some(store);
1955 self
1956 }
1957
1958 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1960 self.schedule_store = Some(store);
1961 self
1962 }
1963
1964 pub fn with_platform_store(
1966 mut self,
1967 store: Arc<dyn crate::platform_store::PlatformStore>,
1968 ) -> Self {
1969 self.platform_store = Some(store);
1970 self
1971 }
1972
1973 pub fn with_knowledge_index_search(
1975 mut self,
1976 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
1977 ) -> Self {
1978 self.knowledge_index_search = Some(search);
1979 self
1980 }
1981
1982 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1984 self.leased_resource_store = Some(store);
1985 self
1986 }
1987
1988 pub fn with_session_resource_registry(
1990 mut self,
1991 registry: Arc<dyn SessionResourceRegistry>,
1992 ) -> Self {
1993 self.session_resource_registry = Some(registry);
1994 self
1995 }
1996
1997 pub fn with_session_task_registry(
1999 mut self,
2000 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
2001 ) -> Self {
2002 self.session_task_registry = Some(registry);
2003 self
2004 }
2005
2006 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
2008 self.org_id = Some(org_id);
2009 self
2010 }
2011
2012 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
2014 self.tool_registry = Some(registry);
2015 self
2016 }
2017
2018 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
2020 self.visible_tool_names = Some(names);
2021 self
2022 }
2023
2024 pub fn with_network_access(
2026 mut self,
2027 network_access: Option<crate::network_access::NetworkAccessList>,
2028 ) -> Self {
2029 self.network_access = network_access;
2030 self
2031 }
2032
2033 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
2035 self.payment_authority = Some(authority);
2036 self
2037 }
2038
2039 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
2041 self.subagent_spawn_store = Some(store);
2042 self
2043 }
2044
2045 pub fn with_subagent_nesting_policy(mut self, policy: SubagentNestingPolicy) -> Self {
2047 self.subagent_nesting_policy = policy;
2048 self
2049 }
2050
2051 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
2056 let (Some(emitter), Some(ctx), Some(call_id)) =
2057 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2058 else {
2059 return;
2060 };
2061 if let Err(e) = emitter
2062 .emit(EventRequest::new(
2063 self.session_id,
2064 ctx.clone(),
2065 crate::events::ToolProgressData {
2066 tool_call_id: call_id.clone(),
2067 tool_name: tool_name.to_string(),
2068 message: message.to_string(),
2069 display_name: None,
2070 },
2071 ))
2072 .await
2073 {
2074 tracing::debug!(
2075 tool_call_id = call_id,
2076 tool_name,
2077 error = %e,
2078 "Failed to emit tool.progress event"
2079 );
2080 }
2081 }
2082
2083 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
2088 let (Some(emitter), Some(ctx), Some(call_id)) =
2089 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2090 else {
2091 return;
2092 };
2093 if let Err(e) = emitter
2094 .emit(EventRequest::new(
2095 self.session_id,
2096 ctx.clone(),
2097 crate::events::ToolOutputDeltaData {
2098 tool_call_id: call_id.clone(),
2099 tool_name: tool_name.to_string(),
2100 delta: delta.to_string(),
2101 stream: stream.to_string(),
2102 },
2103 ))
2104 .await
2105 {
2106 tracing::debug!(
2107 tool_call_id = call_id,
2108 tool_name,
2109 error = %e,
2110 "Failed to emit tool.output.delta event"
2111 );
2112 }
2113 }
2114}
2115
2116impl std::fmt::Debug for ToolContext {
2117 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2118 f.debug_struct("ToolContext")
2119 .field("session_id", &self.session_id)
2120 .field("file_store", &self.file_store.is_some())
2121 .field("storage_store", &self.storage_store.is_some())
2122 .field("image_store", &self.image_store.is_some())
2123 .field(
2124 "provider_credential_store",
2125 &self.provider_credential_store.is_some(),
2126 )
2127 .field("utility_llm_service", &self.utility_llm_service.is_some())
2128 .field("egress_service", &self.egress_service.is_some())
2129 .field("sqldb_store", &self.sqldb_store.is_some())
2130 .field("message_retriever", &self.message_retriever.is_some())
2131 .field("session_store", &self.session_store.is_some())
2132 .field("session_mutator", &self.session_mutator.is_some())
2133 .field("agent_store", &self.agent_store.is_some())
2134 .field("connection_resolver", &self.connection_resolver.is_some())
2135 .field("schedule_store", &self.schedule_store.is_some())
2136 .field("platform_store", &self.platform_store.is_some())
2137 .field(
2138 "knowledge_index_search",
2139 &self.knowledge_index_search.is_some(),
2140 )
2141 .field(
2142 "leased_resource_store",
2143 &self.leased_resource_store.is_some(),
2144 )
2145 .field("event_emitter", &self.event_emitter.is_some())
2146 .field("tool_registry", &self.tool_registry.is_some())
2147 .field("payment_authority", &self.payment_authority.is_some())
2148 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
2149 .field("subagent_nesting_policy", &self.subagent_nesting_policy)
2150 .field("org_id", &self.org_id)
2151 .finish()
2152 }
2153}
2154
2155use crate::events::{Event, EventRequest};
2160
2161#[async_trait]
2172pub trait EventEmitter: Send + Sync {
2173 async fn emit(&self, request: EventRequest) -> Result<Event>;
2178}
2179
2180#[async_trait]
2182impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
2183 async fn emit(&self, request: EventRequest) -> Result<Event> {
2184 (**self).emit(request).await
2185 }
2186}
2187
2188#[derive(Debug, Clone, Default)]
2192pub struct NoopEventEmitter;
2193
2194#[async_trait]
2195impl EventEmitter for NoopEventEmitter {
2196 async fn emit(&self, request: EventRequest) -> Result<Event> {
2197 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2199 }
2200}
2201
2202#[derive(Debug, Clone)]
2215pub struct ResolvedImage {
2216 pub base64: String,
2218 pub media_type: String,
2220}
2221
2222impl ResolvedImage {
2223 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2225 Self {
2226 base64: base64.into(),
2227 media_type: media_type.into(),
2228 }
2229 }
2230
2231 pub fn to_data_url(&self) -> String {
2235 format!("data:{};base64,{}", self.media_type, self.base64)
2236 }
2237}
2238
2239#[async_trait]
2272pub trait ImageResolver: Send + Sync {
2273 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2277}
2278
2279#[derive(Debug)]
2285pub enum SpawnClaimResult {
2286 Claimed {
2289 spawn_handle_id: uuid::Uuid,
2290 claim_token: uuid::Uuid,
2291 },
2292 ClaimedPendingChild {
2296 spawn_handle_id: uuid::Uuid,
2297 claim_token: uuid::Uuid,
2298 },
2299 AlreadyRunning {
2302 child_session_id: crate::typed_id::SessionId,
2303 claim_token: uuid::Uuid,
2305 },
2306 AlreadySettled {
2309 child_session_id: crate::typed_id::SessionId,
2310 terminal_status: String,
2312 terminal_result: String,
2313 },
2314}
2315
2316#[async_trait]
2324pub trait SubagentSpawnStore: Send + Sync + 'static {
2325 async fn try_claim_spawn(
2330 &self,
2331 parent_session_id: crate::typed_id::SessionId,
2332 tool_call_id: &str,
2333 claim_token: uuid::Uuid,
2334 ) -> Result<SpawnClaimResult>;
2335
2336 async fn register_child_session(
2341 &self,
2342 spawn_handle_id: uuid::Uuid,
2343 claim_token: uuid::Uuid,
2344 child_session_id: crate::typed_id::SessionId,
2345 ) -> Result<()>;
2346
2347 async fn settle_spawn(
2353 &self,
2354 parent_session_id: crate::typed_id::SessionId,
2355 tool_call_id: &str,
2356 claim_token: uuid::Uuid,
2357 terminal_status: &str,
2358 terminal_result: &str,
2359 ) -> Result<()>;
2360}
2361
2362#[async_trait]
2364impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2365 async fn try_claim_spawn(
2366 &self,
2367 parent_session_id: crate::typed_id::SessionId,
2368 tool_call_id: &str,
2369 claim_token: uuid::Uuid,
2370 ) -> Result<SpawnClaimResult> {
2371 (**self)
2372 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2373 .await
2374 }
2375
2376 async fn register_child_session(
2377 &self,
2378 spawn_handle_id: uuid::Uuid,
2379 claim_token: uuid::Uuid,
2380 child_session_id: crate::typed_id::SessionId,
2381 ) -> Result<()> {
2382 (**self)
2383 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2384 .await
2385 }
2386
2387 async fn settle_spawn(
2388 &self,
2389 parent_session_id: crate::typed_id::SessionId,
2390 tool_call_id: &str,
2391 claim_token: uuid::Uuid,
2392 terminal_status: &str,
2393 terminal_result: &str,
2394 ) -> Result<()> {
2395 (**self)
2396 .settle_spawn(
2397 parent_session_id,
2398 tool_call_id,
2399 claim_token,
2400 terminal_status,
2401 terminal_result,
2402 )
2403 .await
2404 }
2405}
2406
2407pub struct NoopSubagentSpawnStore;
2411
2412#[async_trait]
2413impl SubagentSpawnStore for NoopSubagentSpawnStore {
2414 async fn try_claim_spawn(
2415 &self,
2416 _parent_session_id: crate::typed_id::SessionId,
2417 _tool_call_id: &str,
2418 claim_token: uuid::Uuid,
2419 ) -> Result<SpawnClaimResult> {
2420 Ok(SpawnClaimResult::Claimed {
2421 spawn_handle_id: uuid::Uuid::new_v4(),
2422 claim_token,
2423 })
2424 }
2425
2426 async fn register_child_session(
2427 &self,
2428 _spawn_handle_id: uuid::Uuid,
2429 _claim_token: uuid::Uuid,
2430 _child_session_id: crate::typed_id::SessionId,
2431 ) -> Result<()> {
2432 Ok(())
2433 }
2434
2435 async fn settle_spawn(
2436 &self,
2437 _parent_session_id: crate::typed_id::SessionId,
2438 _tool_call_id: &str,
2439 _claim_token: uuid::Uuid,
2440 _terminal_status: &str,
2441 _terminal_result: &str,
2442 ) -> Result<()> {
2443 Ok(())
2444 }
2445}
2446
2447#[cfg(test)]
2452mod tests {
2453 use super::*;
2454
2455 #[test]
2456 fn test_resolved_image_new() {
2457 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2458 assert_eq!(image.base64, "SGVsbG8=");
2459 assert_eq!(image.media_type, "image/png");
2460 }
2461
2462 #[test]
2463 fn test_resolved_image_to_data_url() {
2464 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2465 let data_url = image.to_data_url();
2466 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2467 }
2468
2469 #[test]
2470 fn test_resolved_image_jpeg() {
2471 let image = ResolvedImage::new("base64data", "image/jpeg");
2472 let data_url = image.to_data_url();
2473 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2474 }
2475}