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(
525 &self,
526 session_id: SessionId,
527 pattern: &str,
528 path_pattern: Option<&str>,
529 ) -> Result<Vec<GrepMatch>>;
530
531 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
533
534 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
536 if file.is_readonly {
537 return Err(crate::error::AgentLoopError::store(
538 "read-only initial files require a SessionFileSystem-specific seed implementation",
539 ));
540 }
541 self.write_file(session_id, &file.path, &file.content, &file.encoding)
542 .await?;
543 Ok(())
544 }
545}
546
547pub struct WorkspaceScopedFileSystem {
557 inner: Arc<dyn SessionFileSystem>,
558 key: SessionId,
559}
560
561impl WorkspaceScopedFileSystem {
562 pub fn wrap(
564 inner: Arc<dyn SessionFileSystem>,
565 workspace_id: WorkspaceId,
566 ) -> Arc<dyn SessionFileSystem> {
567 Arc::new(Self {
568 inner,
569 key: SessionId::from_uuid(workspace_id.uuid()),
570 })
571 }
572}
573
574#[async_trait]
575impl SessionFileSystem for WorkspaceScopedFileSystem {
576 async fn read_file(&self, _session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
577 self.inner.read_file(self.key, path).await
578 }
579 async fn write_file(
580 &self,
581 _session_id: SessionId,
582 path: &str,
583 content: &str,
584 encoding: &str,
585 ) -> Result<SessionFile> {
586 self.inner
587 .write_file(self.key, path, content, encoding)
588 .await
589 }
590 async fn write_file_if_content_matches(
591 &self,
592 _session_id: SessionId,
593 path: &str,
594 expected_content: &str,
595 expected_encoding: &str,
596 content: &str,
597 encoding: &str,
598 ) -> Result<Option<SessionFile>> {
599 self.inner
600 .write_file_if_content_matches(
601 self.key,
602 path,
603 expected_content,
604 expected_encoding,
605 content,
606 encoding,
607 )
608 .await
609 }
610 async fn delete_file(
611 &self,
612 _session_id: SessionId,
613 path: &str,
614 recursive: bool,
615 ) -> Result<bool> {
616 self.inner.delete_file(self.key, path, recursive).await
617 }
618 async fn list_directory(&self, _session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
619 self.inner.list_directory(self.key, path).await
620 }
621 async fn stat_file(&self, _session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
622 self.inner.stat_file(self.key, path).await
623 }
624 async fn grep_files(
625 &self,
626 _session_id: SessionId,
627 pattern: &str,
628 path_pattern: Option<&str>,
629 ) -> Result<Vec<GrepMatch>> {
630 self.inner.grep_files(self.key, pattern, path_pattern).await
631 }
632 async fn create_directory(&self, _session_id: SessionId, path: &str) -> Result<FileInfo> {
633 self.inner.create_directory(self.key, path).await
634 }
635 async fn seed_initial_file(&self, _session_id: SessionId, file: &InitialFile) -> Result<()> {
636 self.inner.seed_initial_file(self.key, file).await
637 }
638
639 fn display_root(&self) -> String {
640 self.inner.display_root()
641 }
642
643 fn display_path(&self, path: &str) -> String {
644 self.inner.display_path(path)
645 }
646
647 fn resolve_path(&self, input: &str) -> String {
648 self.inner.resolve_path(input)
649 }
650
651 fn is_mount_resolver(&self) -> bool {
652 self.inner.is_mount_resolver()
653 }
654}
655
656#[async_trait]
657impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
658 fn display_root(&self) -> String {
659 (**self).display_root()
660 }
661
662 fn display_path(&self, path: &str) -> String {
663 (**self).display_path(path)
664 }
665
666 fn resolve_path(&self, input: &str) -> String {
667 (**self).resolve_path(input)
668 }
669
670 fn is_mount_resolver(&self) -> bool {
671 (**self).is_mount_resolver()
672 }
673
674 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
675 (**self).read_file(session_id, path).await
676 }
677
678 async fn write_file(
679 &self,
680 session_id: SessionId,
681 path: &str,
682 content: &str,
683 encoding: &str,
684 ) -> Result<SessionFile> {
685 (**self)
686 .write_file(session_id, path, content, encoding)
687 .await
688 }
689
690 async fn write_file_if_content_matches(
691 &self,
692 session_id: SessionId,
693 path: &str,
694 expected_content: &str,
695 expected_encoding: &str,
696 content: &str,
697 encoding: &str,
698 ) -> Result<Option<SessionFile>> {
699 (**self)
700 .write_file_if_content_matches(
701 session_id,
702 path,
703 expected_content,
704 expected_encoding,
705 content,
706 encoding,
707 )
708 .await
709 }
710
711 async fn delete_file(
712 &self,
713 session_id: SessionId,
714 path: &str,
715 recursive: bool,
716 ) -> Result<bool> {
717 (**self).delete_file(session_id, path, recursive).await
718 }
719
720 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
721 (**self).list_directory(session_id, path).await
722 }
723
724 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
725 (**self).stat_file(session_id, path).await
726 }
727
728 async fn grep_files(
729 &self,
730 session_id: SessionId,
731 pattern: &str,
732 path_pattern: Option<&str>,
733 ) -> Result<Vec<GrepMatch>> {
734 (**self).grep_files(session_id, pattern, path_pattern).await
735 }
736
737 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
738 (**self).create_directory(session_id, path).await
739 }
740
741 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
742 (**self).seed_initial_file(session_id, file).await
743 }
744}
745
746pub use SessionFileSystem as SessionFileStore;
748
749#[derive(Clone, Default)]
755pub struct SessionFileSystemFactoryContext {
756 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
757}
758
759impl SessionFileSystemFactoryContext {
760 pub fn new() -> Self {
761 Self::default()
762 }
763
764 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
765 let values = Arc::make_mut(&mut self.values);
766 values.insert(TypeId::of::<T>(), value);
767 self
768 }
769
770 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
771 self.values
772 .get(&TypeId::of::<T>())
773 .and_then(|value| value.clone().downcast::<T>().ok())
774 }
775
776 pub fn with_workspace_roots(self, roots: Arc<crate::WorkspaceRootSet>) -> Self {
777 self.with(roots)
778 }
779
780 pub fn workspace_roots(&self) -> Option<Arc<crate::WorkspaceRootSet>> {
781 self.get::<crate::WorkspaceRootSet>()
782 }
783}
784
785#[async_trait]
787pub trait SessionFileSystemFactory: Send + Sync {
788 fn name(&self) -> &'static str {
790 "SessionFileSystemFactory"
791 }
792
793 fn is_disabled(&self) -> bool {
796 false
797 }
798
799 async fn create_session_file_system(
801 &self,
802 context: SessionFileSystemFactoryContext,
803 ) -> Result<Arc<dyn SessionFileSystem>>;
804}
805
806#[derive(Debug, Clone, Default)]
808pub struct DisabledSessionFileSystemFactory;
809
810#[async_trait]
811impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
812 fn name(&self) -> &'static str {
813 "DisabledSessionFileSystemFactory"
814 }
815
816 fn is_disabled(&self) -> bool {
817 true
818 }
819
820 async fn create_session_file_system(
821 &self,
822 _context: SessionFileSystemFactoryContext,
823 ) -> Result<Arc<dyn SessionFileSystem>> {
824 Err(crate::error::AgentLoopError::config(
825 "session filesystem is disabled",
826 ))
827 }
828}
829
830#[derive(Debug, Clone)]
836pub struct KeyInfo {
837 pub key: String,
838 pub created_at: chrono::DateTime<chrono::Utc>,
839 pub updated_at: chrono::DateTime<chrono::Utc>,
840}
841
842#[derive(Debug, Clone)]
844pub struct SecretInfo {
845 pub name: String,
846 pub created_at: chrono::DateTime<chrono::Utc>,
847 pub updated_at: chrono::DateTime<chrono::Utc>,
848}
849
850#[derive(Debug, Clone, serde::Serialize)]
860pub struct KnowledgeSearchHit {
861 pub id: String,
863 pub kb_id: String,
865 pub title: String,
866 pub kind: String,
867 pub tags: Vec<String>,
868 pub snippet: String,
870 pub resource: Option<String>,
872}
873
874#[async_trait]
878pub trait KnowledgeStore: Send + Sync {
879 async fn search_knowledge(
880 &self,
881 org_id: crate::typed_id::OrgId,
882 kb_public_ids: &[String],
883 query: &str,
884 kind: Option<&str>,
885 tags: &[String],
886 limit: usize,
887 ) -> Result<Vec<KnowledgeSearchHit>>;
888}
889
890#[async_trait]
895pub trait SessionStorageStore: Send + Sync {
896 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
900
901 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
903
904 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
906
907 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
909
910 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
914
915 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
917
918 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
920
921 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
923}
924
925use crate::session_schedule::SessionSchedule;
930use crate::typed_id::ScheduleId;
931
932#[async_trait]
936pub trait SessionScheduleStore: Send + Sync {
937 async fn create_schedule(
939 &self,
940 session_id: SessionId,
941 description: String,
942 cron_expression: Option<String>,
943 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
944 timezone: String,
945 ) -> Result<SessionSchedule>;
946
947 async fn create_schedule_enforcing_limits(
951 &self,
952 session_id: SessionId,
953 description: String,
954 cron_expression: Option<String>,
955 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
956 timezone: String,
957 ) -> std::result::Result<SessionSchedule, crate::session_schedule::ScheduleLimitError> {
958 let per_session = self
959 .count_active_schedules(session_id)
960 .await
961 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
962 if per_session >= crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION {
963 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
964 format!(
965 "Maximum {} active schedules per session. Cancel an existing schedule first.",
966 crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION
967 ),
968 ));
969 }
970
971 let max_per_org = crate::session_schedule::max_active_schedules_per_org();
972 let per_org = self
973 .count_active_org_schedules()
974 .await
975 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
976 if i64::from(per_org) >= max_per_org {
977 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
978 format!(
979 "Maximum {max_per_org} active schedules per org reached. Cancel an existing schedule first."
980 ),
981 ));
982 }
983
984 if let Some(cron) = cron_expression.as_deref() {
985 crate::session_schedule::validate_cron_min_interval(cron)
986 .map_err(crate::session_schedule::ScheduleLimitError::Rejected)?;
987 }
988
989 self.create_schedule(
990 session_id,
991 description,
992 cron_expression,
993 scheduled_at,
994 timezone,
995 )
996 .await
997 .map_err(crate::session_schedule::ScheduleLimitError::Store)
998 }
999
1000 async fn cancel_schedule(
1002 &self,
1003 session_id: SessionId,
1004 schedule_id: ScheduleId,
1005 ) -> Result<SessionSchedule>;
1006
1007 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
1009
1010 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
1012
1013 async fn count_active_org_schedules(&self) -> Result<u32>;
1018}
1019
1020#[async_trait]
1030pub trait SessionResourceRegistry: Send + Sync {
1031 async fn register(
1033 &self,
1034 entry: crate::session_resource::RegisterSessionResource,
1035 ) -> Result<crate::session_resource::SessionResourceEntry>;
1036
1037 async fn update_status(
1039 &self,
1040 session_id: SessionId,
1041 resource_id: &str,
1042 status: crate::session_resource::SessionResourceStatus,
1043 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1044
1045 async fn get(
1047 &self,
1048 session_id: SessionId,
1049 resource_id: &str,
1050 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1051
1052 async fn list(
1054 &self,
1055 session_id: SessionId,
1056 filter: Option<&crate::session_resource::SessionResourceFilter>,
1057 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
1058
1059 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
1061}
1062
1063#[async_trait]
1073pub trait LeasedResourceStore: Send + Sync {
1074 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
1080
1081 async fn release_resource(
1087 &self,
1088 session_id: SessionId,
1089 provider: &str,
1090 resource_type: &str,
1091 external_id: &str,
1092 ) -> Result<Option<LeasedResource>>;
1093
1094 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
1099}
1100
1101pub const DEFAULT_MAX_SUBAGENT_DEPTH: u32 = 2;
1108pub const DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS: u32 = 16;
1109pub const DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS: u32 = 200;
1110
1111#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1113pub struct SubagentNestingPolicy {
1114 pub platform_default: u32,
1115 pub org_override: Option<u32>,
1116 pub agent_override: Option<u32>,
1117 pub platform_default_max_active_descendant_tasks: u32,
1118 pub org_override_max_active_descendant_tasks: Option<u32>,
1119 pub agent_override_max_active_descendant_tasks: Option<u32>,
1120 pub platform_default_max_total_descendant_tasks: u32,
1121 pub org_override_max_total_descendant_tasks: Option<u32>,
1122 pub agent_override_max_total_descendant_tasks: Option<u32>,
1123}
1124
1125impl Default for SubagentNestingPolicy {
1126 fn default() -> Self {
1127 Self {
1128 platform_default: DEFAULT_MAX_SUBAGENT_DEPTH,
1129 org_override: None,
1130 agent_override: None,
1131 platform_default_max_active_descendant_tasks:
1132 DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS,
1133 org_override_max_active_descendant_tasks: None,
1134 agent_override_max_active_descendant_tasks: None,
1135 platform_default_max_total_descendant_tasks:
1136 DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS,
1137 org_override_max_total_descendant_tasks: None,
1138 agent_override_max_total_descendant_tasks: None,
1139 }
1140 }
1141}
1142
1143impl SubagentNestingPolicy {
1144 pub fn max_subagent_depth(self) -> u32 {
1145 self.agent_override
1146 .or(self.org_override)
1147 .unwrap_or(self.platform_default)
1148 }
1149
1150 pub fn max_active_descendant_tasks(self) -> u32 {
1151 self.agent_override_max_active_descendant_tasks
1152 .or(self.org_override_max_active_descendant_tasks)
1153 .unwrap_or(self.platform_default_max_active_descendant_tasks)
1154 }
1155
1156 pub fn max_total_descendant_tasks(self) -> u32 {
1157 self.agent_override_max_total_descendant_tasks
1158 .or(self.org_override_max_total_descendant_tasks)
1159 .unwrap_or(self.platform_default_max_total_descendant_tasks)
1160 }
1161
1162 pub fn with_platform_default(mut self, depth: u32) -> Self {
1163 self.platform_default = depth;
1164 self
1165 }
1166
1167 pub fn with_org_override(mut self, depth: Option<u32>) -> Self {
1168 self.org_override = depth;
1169 self
1170 }
1171
1172 pub fn with_agent_override(mut self, depth: Option<u32>) -> Self {
1173 self.agent_override = depth;
1174 self
1175 }
1176
1177 pub fn with_agent_task_caps_override(
1178 mut self,
1179 max_active: Option<u32>,
1180 max_total: Option<u32>,
1181 ) -> Self {
1182 self.agent_override_max_active_descendant_tasks = max_active;
1183 self.agent_override_max_total_descendant_tasks = max_total;
1184 self
1185 }
1186}
1187
1188pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
1190
1191#[async_trait]
1196pub trait UserConnectionResolver: Send + Sync {
1197 async fn get_connection_token(
1200 &self,
1201 session_id: SessionId,
1202 provider: &str,
1203 ) -> Result<Option<String>>;
1204
1205 async fn get_connection_user(
1210 &self,
1211 _session_id: SessionId,
1212 _provider: &str,
1213 ) -> Result<Option<Uuid>> {
1214 Ok(None)
1215 }
1216
1217 async fn get_connection_token_for_user(
1222 &self,
1223 _user_id: Uuid,
1224 _provider: &str,
1225 ) -> Result<Option<String>> {
1226 Ok(None)
1227 }
1228
1229 async fn get_connection_metadata(
1232 &self,
1233 _session_id: SessionId,
1234 _provider: &str,
1235 ) -> Result<Option<serde_json::Value>> {
1236 Ok(None)
1237 }
1238}
1239
1240#[async_trait]
1250pub trait BudgetChecker: Send + Sync {
1251 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
1253}
1254
1255#[async_trait]
1264pub trait PaymentAuthority: Send + Sync {
1265 async fn execute_machine_payment(
1266 &self,
1267 session_id: SessionId,
1268 request: crate::payment::MachinePaymentRequest,
1269 ) -> Result<crate::payment::MachinePaymentResponse>;
1270}
1271
1272#[async_trait]
1282pub trait OutboundToolRateLimiter: Send + Sync {
1283 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
1285}
1286
1287#[derive(Debug)]
1293pub enum ToolCallClaimResult {
1294 Claimed { claim_token: uuid::Uuid },
1297 AlreadySettled {
1299 result_json: serde_json::Value,
1300 args_fingerprint: String,
1301 },
1302 AlreadyRunning { args_fingerprint: String },
1307 DeterminismViolation {
1311 stored_fingerprint: String,
1312 current_fingerprint: String,
1313 },
1314}
1315
1316#[derive(Debug, Clone)]
1318pub enum DurableToolCallStatus {
1319 Settled { result_json: serde_json::Value },
1321 Interrupted {
1323 result_json: Option<serde_json::Value>,
1324 },
1325 Running,
1327}
1328
1329#[async_trait]
1334pub trait DurableToolResultStore: Send + Sync + 'static {
1335 async fn try_claim_tool_call(
1343 &self,
1344 turn_id: &str,
1345 tool_call_id: &str,
1346 tool_name: &str,
1347 args_fingerprint: &str,
1348 ) -> Result<ToolCallClaimResult>;
1349
1350 async fn settle_tool_call(
1356 &self,
1357 turn_id: &str,
1358 tool_call_id: &str,
1359 result_json: serde_json::Value,
1360 status: &str,
1361 claim_token: uuid::Uuid,
1362 ) -> Result<bool>;
1363
1364 async fn get_tool_call_status(
1369 &self,
1370 turn_id: &str,
1371 tool_call_id: &str,
1372 ) -> Result<Option<DurableToolCallStatus>>;
1373}
1374
1375pub struct NoopDurableToolResultStore;
1378
1379#[async_trait]
1380impl DurableToolResultStore for NoopDurableToolResultStore {
1381 async fn try_claim_tool_call(
1382 &self,
1383 _turn_id: &str,
1384 _tool_call_id: &str,
1385 _tool_name: &str,
1386 _args_fingerprint: &str,
1387 ) -> Result<ToolCallClaimResult> {
1388 Ok(ToolCallClaimResult::Claimed {
1389 claim_token: uuid::Uuid::new_v4(),
1390 })
1391 }
1392
1393 async fn settle_tool_call(
1394 &self,
1395 _turn_id: &str,
1396 _tool_call_id: &str,
1397 _result_json: serde_json::Value,
1398 _status: &str,
1399 _claim_token: uuid::Uuid,
1400 ) -> Result<bool> {
1401 Ok(true)
1402 }
1403
1404 async fn get_tool_call_status(
1405 &self,
1406 _turn_id: &str,
1407 _tool_call_id: &str,
1408 ) -> Result<Option<DurableToolCallStatus>> {
1409 Ok(None)
1410 }
1411}
1412
1413#[derive(Debug, Clone)]
1419pub struct StreamProgress {
1420 pub accumulated_len: usize,
1422 pub last_delta_at: u64,
1424}
1425
1426#[async_trait]
1432pub trait StreamHeartbeater: Send + Sync {
1433 async fn heartbeat(&self, progress: StreamProgress);
1439}
1440
1441pub struct NoopStreamHeartbeater;
1443
1444#[async_trait]
1445impl StreamHeartbeater for NoopStreamHeartbeater {
1446 async fn heartbeat(&self, _progress: StreamProgress) {}
1447}
1448
1449#[derive(Debug, Clone)]
1455pub struct PartialStreamState {
1456 pub accumulated: String,
1459}
1460
1461#[async_trait]
1469pub trait PartialStreamStore: Send + Sync {
1470 async fn get_partial_stream(
1473 &self,
1474 session_id: SessionId,
1475 turn_id: &str,
1476 ) -> Result<Option<PartialStreamState>>;
1477}
1478
1479pub struct NoopPartialStreamStore;
1481
1482#[async_trait]
1483impl PartialStreamStore for NoopPartialStreamStore {
1484 async fn get_partial_stream(
1485 &self,
1486 _session_id: SessionId,
1487 _turn_id: &str,
1488 ) -> Result<Option<PartialStreamState>> {
1489 Ok(None)
1490 }
1491}
1492
1493#[derive(Clone)]
1502pub struct ToolContext {
1503 pub session_id: SessionId,
1505 pub workspace_id: WorkspaceId,
1512
1513 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1515
1516 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1518
1519 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1521
1522 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1524
1525 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1527
1528 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1534
1535 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1537
1538 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1540
1541 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1543
1544 pub session_store: Option<Arc<dyn SessionStore>>,
1546
1547 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1549
1550 pub agent_store: Option<Arc<dyn AgentStore>>,
1552
1553 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1555
1556 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1558
1559 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1561 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1563
1564 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1568
1569 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1571
1572 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1574
1575 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1578
1579 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1582
1583 pub event_context: Option<crate::events::EventContext>,
1586
1587 pub tool_call_id: Option<String>,
1590 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1592
1593 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1596
1597 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1601
1602 pub org_id: Option<crate::typed_id::OrgId>,
1604
1605 pub network_access: Option<crate::network_access::NetworkAccessList>,
1608
1609 pub locale: Option<String>,
1613
1614 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1616
1617 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1619
1620 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1624
1625 pub subagent_nesting_policy: SubagentNestingPolicy,
1627
1628 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1632}
1633
1634impl ToolContext {
1635 pub fn workspace_fs_key(&self) -> SessionId {
1640 SessionId::from_uuid(self.workspace_id.uuid())
1641 }
1642
1643 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1645 self.workspace_id = workspace_id;
1646 self
1647 }
1648
1649 pub fn new(session_id: SessionId) -> Self {
1651 Self {
1652 session_id,
1653 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1654 file_store: None,
1655 storage_store: None,
1656 image_store: None,
1657 provider_credential_store: None,
1658 utility_llm_service: None,
1659 mcp_invoker: None,
1660 egress_service: None,
1661 sqldb_store: None,
1662 message_retriever: None,
1663 session_store: None,
1664 session_mutator: None,
1665 agent_store: None,
1666 connection_resolver: None,
1667 schedule_store: None,
1668 platform_store: None,
1669 knowledge_store: None,
1670 knowledge_index_search: None,
1671 leased_resource_store: None,
1672 session_resource_registry: None,
1673 session_task_registry: None,
1674 event_emitter: None,
1675 event_context: None,
1676 tool_call_id: None,
1677 capability_registry: None,
1678 tool_registry: None,
1679 visible_tool_names: None,
1680 org_id: None,
1681 network_access: None,
1682 locale: None,
1683 budget_checker: None,
1684 payment_authority: None,
1685 subagent_spawn_store: None,
1686 subagent_nesting_policy: SubagentNestingPolicy::default(),
1687 reasoning_effort_handle: None,
1688 }
1689 }
1690
1691 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1693 Self {
1694 session_id,
1695 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1696 file_store: Some(file_store),
1697 storage_store: None,
1698 image_store: None,
1699 provider_credential_store: None,
1700 utility_llm_service: None,
1701 mcp_invoker: None,
1702 egress_service: None,
1703 sqldb_store: None,
1704 message_retriever: None,
1705 session_store: None,
1706 session_mutator: None,
1707 agent_store: None,
1708 connection_resolver: None,
1709 schedule_store: None,
1710 platform_store: None,
1711 knowledge_store: None,
1712 knowledge_index_search: None,
1713 leased_resource_store: None,
1714 session_resource_registry: None,
1715 session_task_registry: None,
1716 event_emitter: None,
1717 event_context: None,
1718 tool_call_id: None,
1719 capability_registry: None,
1720 tool_registry: None,
1721 visible_tool_names: None,
1722 org_id: None,
1723 network_access: None,
1724 locale: None,
1725 budget_checker: None,
1726 payment_authority: None,
1727 subagent_spawn_store: None,
1728 subagent_nesting_policy: SubagentNestingPolicy::default(),
1729 reasoning_effort_handle: None,
1730 }
1731 }
1732
1733 pub fn with_storage_store(
1735 session_id: SessionId,
1736 storage_store: Arc<dyn SessionStorageStore>,
1737 ) -> Self {
1738 Self {
1739 session_id,
1740 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1741 file_store: None,
1742 storage_store: Some(storage_store),
1743 image_store: None,
1744 provider_credential_store: None,
1745 utility_llm_service: None,
1746 mcp_invoker: None,
1747 egress_service: None,
1748 sqldb_store: None,
1749 message_retriever: None,
1750 session_store: None,
1751 session_mutator: None,
1752 agent_store: None,
1753 connection_resolver: None,
1754 schedule_store: None,
1755 platform_store: None,
1756 knowledge_store: None,
1757 knowledge_index_search: None,
1758 leased_resource_store: None,
1759 session_resource_registry: None,
1760 session_task_registry: None,
1761 event_emitter: None,
1762 event_context: None,
1763 tool_call_id: None,
1764 capability_registry: None,
1765 tool_registry: None,
1766 visible_tool_names: None,
1767 org_id: None,
1768 network_access: None,
1769 locale: None,
1770 budget_checker: None,
1771 payment_authority: None,
1772 subagent_spawn_store: None,
1773 subagent_nesting_policy: SubagentNestingPolicy::default(),
1774 reasoning_effort_handle: None,
1775 }
1776 }
1777
1778 pub fn with_stores(
1780 session_id: SessionId,
1781 file_store: Arc<dyn SessionFileSystem>,
1782 storage_store: Arc<dyn SessionStorageStore>,
1783 ) -> Self {
1784 Self {
1785 session_id,
1786 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1787 file_store: Some(file_store),
1788 storage_store: Some(storage_store),
1789 sqldb_store: None,
1790 image_store: None,
1791 provider_credential_store: None,
1792 utility_llm_service: None,
1793 mcp_invoker: None,
1794 egress_service: None,
1795 message_retriever: None,
1796 session_store: None,
1797 session_mutator: None,
1798 agent_store: None,
1799 connection_resolver: None,
1800 schedule_store: None,
1801 platform_store: None,
1802 knowledge_store: None,
1803 knowledge_index_search: None,
1804 leased_resource_store: None,
1805 session_resource_registry: None,
1806 session_task_registry: None,
1807 event_emitter: None,
1808 event_context: None,
1809 tool_call_id: None,
1810 capability_registry: None,
1811 tool_registry: None,
1812 visible_tool_names: None,
1813 org_id: None,
1814 network_access: None,
1815 locale: None,
1816 budget_checker: None,
1817 payment_authority: None,
1818 subagent_spawn_store: None,
1819 subagent_nesting_policy: SubagentNestingPolicy::default(),
1820 reasoning_effort_handle: None,
1821 }
1822 }
1823
1824 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1826 self.sqldb_store = Some(sqldb_store);
1827 self
1828 }
1829
1830 pub fn with_message_retriever(
1832 mut self,
1833 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1834 ) -> Self {
1835 self.message_retriever = Some(retriever);
1836 self
1837 }
1838
1839 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1841 self.session_store = Some(store);
1842 self
1843 }
1844
1845 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1847 self.session_mutator = Some(mutator);
1848 self
1849 }
1850
1851 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
1855 self.reasoning_effort_handle = Some(handle);
1856 self
1857 }
1858
1859 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1861 self.agent_store = Some(store);
1862 self
1863 }
1864
1865 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1867 self.connection_resolver = Some(resolver);
1868 self
1869 }
1870
1871 pub fn with_image_store(
1873 session_id: SessionId,
1874 image_store: Arc<dyn ImageArtifactStore>,
1875 ) -> Self {
1876 Self {
1877 session_id,
1878 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1879 file_store: None,
1880 storage_store: None,
1881 image_store: Some(image_store),
1882 provider_credential_store: None,
1883 utility_llm_service: None,
1884 mcp_invoker: None,
1885 egress_service: None,
1886 sqldb_store: None,
1887 message_retriever: None,
1888 session_store: None,
1889 session_mutator: None,
1890 agent_store: None,
1891 connection_resolver: None,
1892 schedule_store: None,
1893 platform_store: None,
1894 knowledge_store: None,
1895 knowledge_index_search: None,
1896 leased_resource_store: None,
1897 session_resource_registry: None,
1898 session_task_registry: None,
1899 event_emitter: None,
1900 event_context: None,
1901 tool_call_id: None,
1902 capability_registry: None,
1903 tool_registry: None,
1904 visible_tool_names: None,
1905 org_id: None,
1906 network_access: None,
1907 locale: None,
1908 budget_checker: None,
1909 payment_authority: None,
1910 subagent_spawn_store: None,
1911 subagent_nesting_policy: SubagentNestingPolicy::default(),
1912 reasoning_effort_handle: None,
1913 }
1914 }
1915
1916 pub fn with_provider_credential_store(
1918 mut self,
1919 store: Arc<dyn ProviderCredentialStore>,
1920 ) -> Self {
1921 self.provider_credential_store = Some(store);
1922 self
1923 }
1924
1925 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1927 self.utility_llm_service = Some(service);
1928 self
1929 }
1930
1931 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
1933 self.mcp_invoker = Some(invoker);
1934 self
1935 }
1936
1937 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1939 self.egress_service = Some(service);
1940 self
1941 }
1942
1943 pub fn with_egress_service_opt(
1946 mut self,
1947 service: Option<Arc<dyn crate::EgressService>>,
1948 ) -> Self {
1949 if let Some(service) = service {
1950 self.egress_service = Some(service);
1951 }
1952 self
1953 }
1954
1955 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
1957 self.storage_store = Some(store);
1958 self
1959 }
1960
1961 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1963 self.schedule_store = Some(store);
1964 self
1965 }
1966
1967 pub fn with_platform_store(
1969 mut self,
1970 store: Arc<dyn crate::platform_store::PlatformStore>,
1971 ) -> Self {
1972 self.platform_store = Some(store);
1973 self
1974 }
1975
1976 pub fn with_knowledge_index_search(
1978 mut self,
1979 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
1980 ) -> Self {
1981 self.knowledge_index_search = Some(search);
1982 self
1983 }
1984
1985 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1987 self.leased_resource_store = Some(store);
1988 self
1989 }
1990
1991 pub fn with_session_resource_registry(
1993 mut self,
1994 registry: Arc<dyn SessionResourceRegistry>,
1995 ) -> Self {
1996 self.session_resource_registry = Some(registry);
1997 self
1998 }
1999
2000 pub fn with_session_task_registry(
2002 mut self,
2003 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
2004 ) -> Self {
2005 self.session_task_registry = Some(registry);
2006 self
2007 }
2008
2009 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
2011 self.org_id = Some(org_id);
2012 self
2013 }
2014
2015 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
2017 self.tool_registry = Some(registry);
2018 self
2019 }
2020
2021 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
2023 self.visible_tool_names = Some(names);
2024 self
2025 }
2026
2027 pub fn with_network_access(
2029 mut self,
2030 network_access: Option<crate::network_access::NetworkAccessList>,
2031 ) -> Self {
2032 self.network_access = network_access;
2033 self
2034 }
2035
2036 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
2038 self.payment_authority = Some(authority);
2039 self
2040 }
2041
2042 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
2044 self.subagent_spawn_store = Some(store);
2045 self
2046 }
2047
2048 pub fn with_subagent_nesting_policy(mut self, policy: SubagentNestingPolicy) -> Self {
2050 self.subagent_nesting_policy = policy;
2051 self
2052 }
2053
2054 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
2059 let (Some(emitter), Some(ctx), Some(call_id)) =
2060 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2061 else {
2062 return;
2063 };
2064 if let Err(e) = emitter
2065 .emit(EventRequest::new(
2066 self.session_id,
2067 ctx.clone(),
2068 crate::events::ToolProgressData {
2069 tool_call_id: call_id.clone(),
2070 tool_name: tool_name.to_string(),
2071 message: message.to_string(),
2072 display_name: None,
2073 },
2074 ))
2075 .await
2076 {
2077 tracing::debug!(
2078 tool_call_id = call_id,
2079 tool_name,
2080 error = %e,
2081 "Failed to emit tool.progress event"
2082 );
2083 }
2084 }
2085
2086 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
2091 let (Some(emitter), Some(ctx), Some(call_id)) =
2092 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2093 else {
2094 return;
2095 };
2096 if let Err(e) = emitter
2097 .emit(EventRequest::new(
2098 self.session_id,
2099 ctx.clone(),
2100 crate::events::ToolOutputDeltaData {
2101 tool_call_id: call_id.clone(),
2102 tool_name: tool_name.to_string(),
2103 delta: delta.to_string(),
2104 stream: stream.to_string(),
2105 },
2106 ))
2107 .await
2108 {
2109 tracing::debug!(
2110 tool_call_id = call_id,
2111 tool_name,
2112 error = %e,
2113 "Failed to emit tool.output.delta event"
2114 );
2115 }
2116 }
2117}
2118
2119impl std::fmt::Debug for ToolContext {
2120 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2121 f.debug_struct("ToolContext")
2122 .field("session_id", &self.session_id)
2123 .field("file_store", &self.file_store.is_some())
2124 .field("storage_store", &self.storage_store.is_some())
2125 .field("image_store", &self.image_store.is_some())
2126 .field(
2127 "provider_credential_store",
2128 &self.provider_credential_store.is_some(),
2129 )
2130 .field("utility_llm_service", &self.utility_llm_service.is_some())
2131 .field("egress_service", &self.egress_service.is_some())
2132 .field("sqldb_store", &self.sqldb_store.is_some())
2133 .field("message_retriever", &self.message_retriever.is_some())
2134 .field("session_store", &self.session_store.is_some())
2135 .field("session_mutator", &self.session_mutator.is_some())
2136 .field("agent_store", &self.agent_store.is_some())
2137 .field("connection_resolver", &self.connection_resolver.is_some())
2138 .field("schedule_store", &self.schedule_store.is_some())
2139 .field("platform_store", &self.platform_store.is_some())
2140 .field(
2141 "knowledge_index_search",
2142 &self.knowledge_index_search.is_some(),
2143 )
2144 .field(
2145 "leased_resource_store",
2146 &self.leased_resource_store.is_some(),
2147 )
2148 .field("event_emitter", &self.event_emitter.is_some())
2149 .field("tool_registry", &self.tool_registry.is_some())
2150 .field("payment_authority", &self.payment_authority.is_some())
2151 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
2152 .field("subagent_nesting_policy", &self.subagent_nesting_policy)
2153 .field("org_id", &self.org_id)
2154 .finish()
2155 }
2156}
2157
2158use crate::events::{Event, EventRequest};
2163
2164#[async_trait]
2175pub trait EventEmitter: Send + Sync {
2176 async fn emit(&self, request: EventRequest) -> Result<Event>;
2181}
2182
2183#[async_trait]
2185impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
2186 async fn emit(&self, request: EventRequest) -> Result<Event> {
2187 (**self).emit(request).await
2188 }
2189}
2190
2191#[derive(Debug, Clone, Default)]
2195pub struct NoopEventEmitter;
2196
2197#[async_trait]
2198impl EventEmitter for NoopEventEmitter {
2199 async fn emit(&self, request: EventRequest) -> Result<Event> {
2200 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2202 }
2203}
2204
2205#[derive(Debug, Clone)]
2218pub struct ResolvedImage {
2219 pub base64: String,
2221 pub media_type: String,
2223}
2224
2225impl ResolvedImage {
2226 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2228 Self {
2229 base64: base64.into(),
2230 media_type: media_type.into(),
2231 }
2232 }
2233
2234 pub fn to_data_url(&self) -> String {
2238 format!("data:{};base64,{}", self.media_type, self.base64)
2239 }
2240}
2241
2242#[async_trait]
2275pub trait ImageResolver: Send + Sync {
2276 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2280}
2281
2282#[derive(Debug)]
2288pub enum SpawnClaimResult {
2289 Claimed {
2292 spawn_handle_id: uuid::Uuid,
2293 claim_token: uuid::Uuid,
2294 },
2295 ClaimedPendingChild {
2299 spawn_handle_id: uuid::Uuid,
2300 claim_token: uuid::Uuid,
2301 },
2302 AlreadyRunning {
2305 child_session_id: crate::typed_id::SessionId,
2306 claim_token: uuid::Uuid,
2308 },
2309 AlreadySettled {
2312 child_session_id: crate::typed_id::SessionId,
2313 terminal_status: String,
2315 terminal_result: String,
2316 },
2317}
2318
2319#[async_trait]
2327pub trait SubagentSpawnStore: Send + Sync + 'static {
2328 async fn try_claim_spawn(
2333 &self,
2334 parent_session_id: crate::typed_id::SessionId,
2335 tool_call_id: &str,
2336 claim_token: uuid::Uuid,
2337 ) -> Result<SpawnClaimResult>;
2338
2339 async fn register_child_session(
2344 &self,
2345 spawn_handle_id: uuid::Uuid,
2346 claim_token: uuid::Uuid,
2347 child_session_id: crate::typed_id::SessionId,
2348 ) -> Result<()>;
2349
2350 async fn settle_spawn(
2356 &self,
2357 parent_session_id: crate::typed_id::SessionId,
2358 tool_call_id: &str,
2359 claim_token: uuid::Uuid,
2360 terminal_status: &str,
2361 terminal_result: &str,
2362 ) -> Result<()>;
2363}
2364
2365#[async_trait]
2367impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2368 async fn try_claim_spawn(
2369 &self,
2370 parent_session_id: crate::typed_id::SessionId,
2371 tool_call_id: &str,
2372 claim_token: uuid::Uuid,
2373 ) -> Result<SpawnClaimResult> {
2374 (**self)
2375 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2376 .await
2377 }
2378
2379 async fn register_child_session(
2380 &self,
2381 spawn_handle_id: uuid::Uuid,
2382 claim_token: uuid::Uuid,
2383 child_session_id: crate::typed_id::SessionId,
2384 ) -> Result<()> {
2385 (**self)
2386 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2387 .await
2388 }
2389
2390 async fn settle_spawn(
2391 &self,
2392 parent_session_id: crate::typed_id::SessionId,
2393 tool_call_id: &str,
2394 claim_token: uuid::Uuid,
2395 terminal_status: &str,
2396 terminal_result: &str,
2397 ) -> Result<()> {
2398 (**self)
2399 .settle_spawn(
2400 parent_session_id,
2401 tool_call_id,
2402 claim_token,
2403 terminal_status,
2404 terminal_result,
2405 )
2406 .await
2407 }
2408}
2409
2410pub struct NoopSubagentSpawnStore;
2414
2415#[async_trait]
2416impl SubagentSpawnStore for NoopSubagentSpawnStore {
2417 async fn try_claim_spawn(
2418 &self,
2419 _parent_session_id: crate::typed_id::SessionId,
2420 _tool_call_id: &str,
2421 claim_token: uuid::Uuid,
2422 ) -> Result<SpawnClaimResult> {
2423 Ok(SpawnClaimResult::Claimed {
2424 spawn_handle_id: uuid::Uuid::new_v4(),
2425 claim_token,
2426 })
2427 }
2428
2429 async fn register_child_session(
2430 &self,
2431 _spawn_handle_id: uuid::Uuid,
2432 _claim_token: uuid::Uuid,
2433 _child_session_id: crate::typed_id::SessionId,
2434 ) -> Result<()> {
2435 Ok(())
2436 }
2437
2438 async fn settle_spawn(
2439 &self,
2440 _parent_session_id: crate::typed_id::SessionId,
2441 _tool_call_id: &str,
2442 _claim_token: uuid::Uuid,
2443 _terminal_status: &str,
2444 _terminal_result: &str,
2445 ) -> Result<()> {
2446 Ok(())
2447 }
2448}
2449
2450#[cfg(test)]
2455mod tests {
2456 use super::*;
2457
2458 #[test]
2459 fn test_resolved_image_new() {
2460 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2461 assert_eq!(image.base64, "SGVsbG8=");
2462 assert_eq!(image.media_type, "image/png");
2463 }
2464
2465 #[test]
2466 fn test_resolved_image_to_data_url() {
2467 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2468 let data_url = image.to_data_url();
2469 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2470 }
2471
2472 #[test]
2473 fn test_resolved_image_jpeg() {
2474 let image = ResolvedImage::new("base64data", "image/jpeg");
2475 let data_url = image.to_data_url();
2476 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2477 }
2478}