1use crate::agent::Agent;
9use crate::harness::Harness;
10use crate::provider::DriverId;
11use crate::session_file::{
12 FileInfo, FileStat, GrepMatch, GrepOptions, GrepSearchResult, InitialFile, SessionFile,
13};
14use crate::tool_types::{ToolCall, ToolDefinition, ToolResult};
15use crate::typed_id::{AgentId, HarnessId, ImageId, ModelId, SessionId, WorkspaceId};
16use async_trait::async_trait;
17use chrono::{DateTime, Utc};
18use std::any::{Any, TypeId};
19use std::collections::{HashMap, HashSet};
20use std::sync::Arc;
21use uuid::Uuid;
22
23fn build_tool_map(tool_defs: &[ToolDefinition]) -> HashMap<&str, &ToolDefinition> {
25 tool_defs.iter().map(|def| (def.name(), def)).collect()
26}
27
28use crate::error::Result;
29
30#[derive(Clone, Default)]
45pub struct ReasoningEffortHandle {
46 inner: Arc<std::sync::RwLock<Option<String>>>,
47}
48
49impl ReasoningEffortHandle {
50 pub fn new() -> Self {
52 Self::default()
53 }
54
55 pub fn with_effort(effort: impl Into<String>) -> Self {
57 Self {
58 inner: Arc::new(std::sync::RwLock::new(Some(effort.into()))),
59 }
60 }
61
62 pub fn set(&self, effort: Option<String>) {
66 let mut guard = self.inner.write().unwrap_or_else(|e| e.into_inner());
70 *guard = effort;
71 }
72
73 pub fn get(&self) -> Option<String> {
75 let guard = self.inner.read().unwrap_or_else(|e| e.into_inner());
78 guard.clone()
79 }
80}
81
82impl std::fmt::Debug for ReasoningEffortHandle {
83 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84 f.debug_struct("ReasoningEffortHandle")
85 .field("effort", &self.get())
86 .finish()
87 }
88}
89
90#[async_trait]
101pub trait AgentStore: Send + Sync {
102 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>>;
104}
105
106#[async_trait]
107impl<T: AgentStore + ?Sized> AgentStore for std::sync::Arc<T> {
108 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>> {
109 (**self).get_agent(agent_id).await
110 }
111}
112
113#[async_trait]
128pub trait HarnessStore: Send + Sync {
129 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>>;
134}
135
136#[async_trait]
137impl<T: HarnessStore + ?Sized> HarnessStore for std::sync::Arc<T> {
138 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>> {
139 (**self).get_harness_chain(harness_id).await
140 }
141}
142
143use crate::leased_resource::{LeasedResource, UpsertLeasedResource};
148use crate::session::Session;
149
150#[async_trait]
156pub trait SessionStore: Send + Sync {
157 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>>;
159}
160
161#[async_trait]
162impl<T: SessionStore + ?Sized> SessionStore for std::sync::Arc<T> {
163 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>> {
164 (**self).get_session(session_id).await
165 }
166}
167
168#[async_trait]
170pub trait SessionMutator: Send + Sync {
171 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session>;
173}
174
175#[async_trait]
176impl<T: SessionMutator + ?Sized> SessionMutator for std::sync::Arc<T> {
177 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session> {
178 (**self).update_session_title(session_id, title).await
179 }
180}
181
182#[derive(Debug, Clone)]
188pub struct ResolvedModel {
189 pub model: String,
191 pub provider_type: DriverId,
193 pub api_key: Option<String>,
195 pub base_url: Option<String>,
197 pub provider_metadata: Option<crate::driver_registry::ProviderMetadata>,
200}
201
202#[async_trait]
212pub trait ProviderStore: Send + Sync {
213 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>>;
218
219 async fn get_default_model(&self) -> Result<Option<ResolvedModel>>;
223}
224
225#[async_trait]
226impl<T: ProviderStore + ?Sized> ProviderStore for std::sync::Arc<T> {
227 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>> {
228 (**self).get_resolved_model(model_id).await
229 }
230
231 async fn get_default_model(&self) -> Result<Option<ResolvedModel>> {
232 (**self).get_default_model().await
233 }
234}
235
236#[derive(Debug, Clone)]
242pub struct StoredImageInfo {
243 pub id: ImageId,
244 pub filename: String,
245 pub content_type: String,
246 pub size_bytes: i64,
247 pub metadata: serde_json::Value,
248 pub created_at: DateTime<Utc>,
249}
250
251#[derive(Debug, Clone)]
253pub struct StoredImage {
254 pub info: StoredImageInfo,
255 pub data: Vec<u8>,
256}
257
258#[derive(Debug, Clone)]
260pub struct CreateStoredImage {
261 pub filename: String,
262 pub content_type: String,
263 pub data: Vec<u8>,
264 pub metadata: serde_json::Value,
265}
266
267#[async_trait]
268pub trait ImageArtifactStore: Send + Sync {
269 async fn create_image(&self, input: CreateStoredImage) -> Result<StoredImageInfo>;
271
272 async fn get_image(&self, image_id: ImageId) -> Result<Option<StoredImage>>;
274
275 async fn get_image_info(&self, image_id: ImageId) -> Result<Option<StoredImageInfo>>;
277}
278
279#[derive(Debug, Clone)]
285pub struct ProviderCredentials {
286 pub api_key: String,
287 pub base_url: Option<String>,
288}
289
290#[async_trait]
291pub trait ProviderCredentialStore: Send + Sync {
292 async fn get_default_provider_credentials(
297 &self,
298 provider_type: &str,
299 ) -> Result<Option<ProviderCredentials>>;
300}
301
302#[async_trait]
313pub trait ToolExecutor: Send + Sync {
314 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult>;
319
320 async fn execute_with_context(
325 &self,
326 tool_call: &ToolCall,
327 tool_def: &ToolDefinition,
328 _context: &ToolContext,
329 ) -> Result<ToolResult> {
330 self.execute(tool_call, tool_def).await
332 }
333
334 async fn execute_batch(
336 &self,
337 tool_calls: &[ToolCall],
338 tool_defs: &[ToolDefinition],
339 ) -> Result<Vec<ToolResult>> {
340 let mut results = Vec::with_capacity(tool_calls.len());
341
342 let tool_map = build_tool_map(tool_defs);
343
344 for tool_call in tool_calls {
345 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
346 crate::error::AgentLoopError::tool(format!(
347 "Tool definition not found: {}",
348 tool_call.name
349 ))
350 })?;
351
352 results.push(self.execute(tool_call, tool_def).await?);
353 }
354
355 Ok(results)
356 }
357
358 async fn execute_parallel(
360 &self,
361 tool_calls: &[ToolCall],
362 tool_defs: &[ToolDefinition],
363 ) -> Result<Vec<ToolResult>>
364 where
365 Self: Sized,
366 {
367 use futures::future::join_all;
368
369 let tool_map = build_tool_map(tool_defs);
370
371 let futures: Vec<_> = tool_calls
372 .iter()
373 .map(|tool_call| async {
374 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
375 crate::error::AgentLoopError::tool(format!(
376 "Tool definition not found: {}",
377 tool_call.name
378 ))
379 })?;
380 self.execute(tool_call, tool_def).await
381 })
382 .collect();
383
384 let results = join_all(futures).await;
385 results.into_iter().collect()
386 }
387}
388
389#[async_trait]
393impl ToolExecutor for std::sync::Arc<dyn ToolExecutor> {
394 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult> {
395 (**self).execute(tool_call, tool_def).await
396 }
397
398 async fn execute_with_context(
399 &self,
400 tool_call: &ToolCall,
401 tool_def: &ToolDefinition,
402 context: &ToolContext,
403 ) -> Result<ToolResult> {
404 (**self)
405 .execute_with_context(tool_call, tool_def, context)
406 .await
407 }
408
409 async fn execute_batch(
410 &self,
411 tool_calls: &[ToolCall],
412 tool_defs: &[ToolDefinition],
413 ) -> Result<Vec<ToolResult>> {
414 (**self).execute_batch(tool_calls, tool_defs).await
415 }
416}
417
418#[async_trait]
430pub trait SessionFileSystem: Send + Sync {
431 fn display_root(&self) -> String {
437 crate::session_path::WORKSPACE_PREFIX.to_string()
438 }
439
440 fn display_path(&self, path: &str) -> String {
446 crate::session_path::to_display_path(path)
447 }
448
449 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(
527 &self,
528 session_id: SessionId,
529 pattern: &str,
530 path_pattern: Option<&str>,
531 ) -> Result<Vec<GrepMatch>>;
532
533 async fn grep_files_with_options(
539 &self,
540 session_id: SessionId,
541 pattern: &str,
542 options: &GrepOptions,
543 ) -> Result<GrepSearchResult> {
544 if options.before_context != 0 || options.after_context != 0 {
545 return Err(crate::error::AgentLoopError::tool(
546 "this file store does not support grep context",
547 ));
548 }
549 let all = self
550 .grep_files(session_id, pattern, options.path_pattern.as_deref())
551 .await?;
552 Ok(crate::session_file::bound_grep_matches(all, options))
553 }
554
555 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
557
558 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
560 if file.is_readonly {
561 return Err(crate::error::AgentLoopError::store(
562 "read-only initial files require a SessionFileSystem-specific seed implementation",
563 ));
564 }
565 self.write_file(session_id, &file.path, &file.content, &file.encoding)
566 .await?;
567 Ok(())
568 }
569}
570
571pub struct WorkspaceScopedFileSystem {
581 inner: Arc<dyn SessionFileSystem>,
582 key: SessionId,
583}
584
585impl WorkspaceScopedFileSystem {
586 pub fn wrap(
588 inner: Arc<dyn SessionFileSystem>,
589 workspace_id: WorkspaceId,
590 ) -> Arc<dyn SessionFileSystem> {
591 Arc::new(Self {
592 inner,
593 key: SessionId::from_uuid(workspace_id.uuid()),
594 })
595 }
596}
597
598#[async_trait]
599impl SessionFileSystem for WorkspaceScopedFileSystem {
600 async fn read_file(&self, _session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
601 self.inner.read_file(self.key, path).await
602 }
603 async fn write_file(
604 &self,
605 _session_id: SessionId,
606 path: &str,
607 content: &str,
608 encoding: &str,
609 ) -> Result<SessionFile> {
610 self.inner
611 .write_file(self.key, path, content, encoding)
612 .await
613 }
614 async fn write_file_if_content_matches(
615 &self,
616 _session_id: SessionId,
617 path: &str,
618 expected_content: &str,
619 expected_encoding: &str,
620 content: &str,
621 encoding: &str,
622 ) -> Result<Option<SessionFile>> {
623 self.inner
624 .write_file_if_content_matches(
625 self.key,
626 path,
627 expected_content,
628 expected_encoding,
629 content,
630 encoding,
631 )
632 .await
633 }
634 async fn delete_file(
635 &self,
636 _session_id: SessionId,
637 path: &str,
638 recursive: bool,
639 ) -> Result<bool> {
640 self.inner.delete_file(self.key, path, recursive).await
641 }
642 async fn list_directory(&self, _session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
643 self.inner.list_directory(self.key, path).await
644 }
645 async fn stat_file(&self, _session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
646 self.inner.stat_file(self.key, path).await
647 }
648 async fn grep_files(
649 &self,
650 _session_id: SessionId,
651 pattern: &str,
652 path_pattern: Option<&str>,
653 ) -> Result<Vec<GrepMatch>> {
654 self.inner.grep_files(self.key, pattern, path_pattern).await
655 }
656 async fn grep_files_with_options(
657 &self,
658 _session_id: SessionId,
659 pattern: &str,
660 options: &GrepOptions,
661 ) -> Result<GrepSearchResult> {
662 self.inner
663 .grep_files_with_options(self.key, pattern, options)
664 .await
665 }
666 async fn create_directory(&self, _session_id: SessionId, path: &str) -> Result<FileInfo> {
667 self.inner.create_directory(self.key, path).await
668 }
669 async fn seed_initial_file(&self, _session_id: SessionId, file: &InitialFile) -> Result<()> {
670 self.inner.seed_initial_file(self.key, file).await
671 }
672
673 fn display_root(&self) -> String {
674 self.inner.display_root()
675 }
676
677 fn display_path(&self, path: &str) -> String {
678 self.inner.display_path(path)
679 }
680
681 fn resolve_path(&self, input: &str) -> String {
682 self.inner.resolve_path(input)
683 }
684
685 fn is_mount_resolver(&self) -> bool {
686 self.inner.is_mount_resolver()
687 }
688}
689
690#[async_trait]
691impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
692 fn display_root(&self) -> String {
693 (**self).display_root()
694 }
695
696 fn display_path(&self, path: &str) -> String {
697 (**self).display_path(path)
698 }
699
700 fn resolve_path(&self, input: &str) -> String {
701 (**self).resolve_path(input)
702 }
703
704 fn is_mount_resolver(&self) -> bool {
705 (**self).is_mount_resolver()
706 }
707
708 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
709 (**self).read_file(session_id, path).await
710 }
711
712 async fn write_file(
713 &self,
714 session_id: SessionId,
715 path: &str,
716 content: &str,
717 encoding: &str,
718 ) -> Result<SessionFile> {
719 (**self)
720 .write_file(session_id, path, content, encoding)
721 .await
722 }
723
724 async fn write_file_if_content_matches(
725 &self,
726 session_id: SessionId,
727 path: &str,
728 expected_content: &str,
729 expected_encoding: &str,
730 content: &str,
731 encoding: &str,
732 ) -> Result<Option<SessionFile>> {
733 (**self)
734 .write_file_if_content_matches(
735 session_id,
736 path,
737 expected_content,
738 expected_encoding,
739 content,
740 encoding,
741 )
742 .await
743 }
744
745 async fn delete_file(
746 &self,
747 session_id: SessionId,
748 path: &str,
749 recursive: bool,
750 ) -> Result<bool> {
751 (**self).delete_file(session_id, path, recursive).await
752 }
753
754 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
755 (**self).list_directory(session_id, path).await
756 }
757
758 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
759 (**self).stat_file(session_id, path).await
760 }
761
762 async fn grep_files(
763 &self,
764 session_id: SessionId,
765 pattern: &str,
766 path_pattern: Option<&str>,
767 ) -> Result<Vec<GrepMatch>> {
768 (**self).grep_files(session_id, pattern, path_pattern).await
769 }
770
771 async fn grep_files_with_options(
772 &self,
773 session_id: SessionId,
774 pattern: &str,
775 options: &GrepOptions,
776 ) -> Result<GrepSearchResult> {
777 (**self)
778 .grep_files_with_options(session_id, pattern, options)
779 .await
780 }
781
782 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
783 (**self).create_directory(session_id, path).await
784 }
785
786 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
787 (**self).seed_initial_file(session_id, file).await
788 }
789}
790
791pub use SessionFileSystem as SessionFileStore;
793
794#[derive(Clone, Default)]
800pub struct SessionFileSystemFactoryContext {
801 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
802}
803
804impl SessionFileSystemFactoryContext {
805 pub fn new() -> Self {
806 Self::default()
807 }
808
809 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
810 let values = Arc::make_mut(&mut self.values);
811 values.insert(TypeId::of::<T>(), value);
812 self
813 }
814
815 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
816 self.values
817 .get(&TypeId::of::<T>())
818 .and_then(|value| value.clone().downcast::<T>().ok())
819 }
820
821 pub fn with_workspace_roots(self, roots: Arc<crate::WorkspaceRootSet>) -> Self {
822 self.with(roots)
823 }
824
825 pub fn workspace_roots(&self) -> Option<Arc<crate::WorkspaceRootSet>> {
826 self.get::<crate::WorkspaceRootSet>()
827 }
828}
829
830#[async_trait]
832pub trait SessionFileSystemFactory: Send + Sync {
833 fn name(&self) -> &'static str {
835 "SessionFileSystemFactory"
836 }
837
838 fn is_disabled(&self) -> bool {
841 false
842 }
843
844 async fn create_session_file_system(
846 &self,
847 context: SessionFileSystemFactoryContext,
848 ) -> Result<Arc<dyn SessionFileSystem>>;
849}
850
851#[derive(Debug, Clone, Default)]
853pub struct DisabledSessionFileSystemFactory;
854
855#[async_trait]
856impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
857 fn name(&self) -> &'static str {
858 "DisabledSessionFileSystemFactory"
859 }
860
861 fn is_disabled(&self) -> bool {
862 true
863 }
864
865 async fn create_session_file_system(
866 &self,
867 _context: SessionFileSystemFactoryContext,
868 ) -> Result<Arc<dyn SessionFileSystem>> {
869 Err(crate::error::AgentLoopError::config(
870 "session filesystem is disabled",
871 ))
872 }
873}
874
875#[derive(Debug, Clone)]
881pub struct KeyInfo {
882 pub key: String,
883 pub created_at: chrono::DateTime<chrono::Utc>,
884 pub updated_at: chrono::DateTime<chrono::Utc>,
885}
886
887#[derive(Debug, Clone)]
889pub struct SecretInfo {
890 pub name: String,
891 pub created_at: chrono::DateTime<chrono::Utc>,
892 pub updated_at: chrono::DateTime<chrono::Utc>,
893}
894
895#[derive(Debug, Clone, serde::Serialize)]
905pub struct KnowledgeSearchHit {
906 pub id: String,
908 pub kb_id: String,
910 pub title: String,
911 pub kind: String,
912 pub tags: Vec<String>,
913 pub snippet: String,
915 pub resource: Option<String>,
917}
918
919#[async_trait]
923pub trait KnowledgeStore: Send + Sync {
924 async fn search_knowledge(
925 &self,
926 org_id: crate::typed_id::OrgId,
927 kb_public_ids: &[String],
928 query: &str,
929 kind: Option<&str>,
930 tags: &[String],
931 limit: usize,
932 ) -> Result<Vec<KnowledgeSearchHit>>;
933}
934
935#[async_trait]
940pub trait SessionStorageStore: Send + Sync {
941 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
945
946 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
948
949 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
951
952 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
954
955 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
959
960 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
962
963 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
965
966 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
968}
969
970use crate::session_schedule::SessionSchedule;
975use crate::typed_id::ScheduleId;
976
977#[async_trait]
981pub trait SessionScheduleStore: Send + Sync {
982 async fn create_schedule(
984 &self,
985 session_id: SessionId,
986 description: String,
987 cron_expression: Option<String>,
988 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
989 timezone: String,
990 ) -> Result<SessionSchedule>;
991
992 async fn create_schedule_enforcing_limits(
996 &self,
997 session_id: SessionId,
998 description: String,
999 cron_expression: Option<String>,
1000 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
1001 timezone: String,
1002 ) -> std::result::Result<SessionSchedule, crate::session_schedule::ScheduleLimitError> {
1003 let per_session = self
1004 .count_active_schedules(session_id)
1005 .await
1006 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
1007 if per_session >= crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION {
1008 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
1009 format!(
1010 "Maximum {} active schedules per session. Cancel an existing schedule first.",
1011 crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION
1012 ),
1013 ));
1014 }
1015
1016 let max_per_org = crate::session_schedule::max_active_schedules_per_org();
1017 let per_org = self
1018 .count_active_org_schedules()
1019 .await
1020 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
1021 if i64::from(per_org) >= max_per_org {
1022 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
1023 format!(
1024 "Maximum {max_per_org} active schedules per org reached. Cancel an existing schedule first."
1025 ),
1026 ));
1027 }
1028
1029 if let Some(cron) = cron_expression.as_deref() {
1030 crate::session_schedule::validate_cron_min_interval(cron)
1031 .map_err(crate::session_schedule::ScheduleLimitError::Rejected)?;
1032 }
1033
1034 self.create_schedule(
1035 session_id,
1036 description,
1037 cron_expression,
1038 scheduled_at,
1039 timezone,
1040 )
1041 .await
1042 .map_err(crate::session_schedule::ScheduleLimitError::Store)
1043 }
1044
1045 async fn cancel_schedule(
1047 &self,
1048 session_id: SessionId,
1049 schedule_id: ScheduleId,
1050 ) -> Result<SessionSchedule>;
1051
1052 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
1054
1055 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
1057
1058 async fn count_active_org_schedules(&self) -> Result<u32>;
1063}
1064
1065#[async_trait]
1075pub trait SessionResourceRegistry: Send + Sync {
1076 async fn register(
1078 &self,
1079 entry: crate::session_resource::RegisterSessionResource,
1080 ) -> Result<crate::session_resource::SessionResourceEntry>;
1081
1082 async fn update_status(
1084 &self,
1085 session_id: SessionId,
1086 resource_id: &str,
1087 status: crate::session_resource::SessionResourceStatus,
1088 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1089
1090 async fn get(
1092 &self,
1093 session_id: SessionId,
1094 resource_id: &str,
1095 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1096
1097 async fn list(
1099 &self,
1100 session_id: SessionId,
1101 filter: Option<&crate::session_resource::SessionResourceFilter>,
1102 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
1103
1104 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
1106}
1107
1108#[async_trait]
1118pub trait LeasedResourceStore: Send + Sync {
1119 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
1125
1126 async fn release_resource(
1132 &self,
1133 session_id: SessionId,
1134 provider: &str,
1135 resource_type: &str,
1136 external_id: &str,
1137 ) -> Result<Option<LeasedResource>>;
1138
1139 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
1144}
1145
1146pub const DEFAULT_MAX_SUBAGENT_DEPTH: u32 = 2;
1153pub const DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS: u32 = 16;
1154pub const DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS: u32 = 200;
1155pub const DEFAULT_MAX_ACTIVE_DETACHED_TASKS: u32 = 8;
1161pub const DEFAULT_MAX_TOTAL_DETACHED_TASKS: u32 = 50;
1162
1163#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1165pub struct SubagentNestingPolicy {
1166 pub platform_default: u32,
1167 pub org_override: Option<u32>,
1168 pub agent_override: Option<u32>,
1169 pub platform_default_max_active_descendant_tasks: u32,
1170 pub org_override_max_active_descendant_tasks: Option<u32>,
1171 pub agent_override_max_active_descendant_tasks: Option<u32>,
1172 pub platform_default_max_total_descendant_tasks: u32,
1173 pub org_override_max_total_descendant_tasks: Option<u32>,
1174 pub agent_override_max_total_descendant_tasks: Option<u32>,
1175 pub platform_default_max_active_detached_tasks: u32,
1176 pub org_override_max_active_detached_tasks: Option<u32>,
1177 pub agent_override_max_active_detached_tasks: Option<u32>,
1178 pub platform_default_max_total_detached_tasks: u32,
1179 pub org_override_max_total_detached_tasks: Option<u32>,
1180 pub agent_override_max_total_detached_tasks: Option<u32>,
1181}
1182
1183impl Default for SubagentNestingPolicy {
1184 fn default() -> Self {
1185 Self {
1186 platform_default: DEFAULT_MAX_SUBAGENT_DEPTH,
1187 org_override: None,
1188 agent_override: None,
1189 platform_default_max_active_descendant_tasks:
1190 DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS,
1191 org_override_max_active_descendant_tasks: None,
1192 agent_override_max_active_descendant_tasks: None,
1193 platform_default_max_total_descendant_tasks:
1194 DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS,
1195 org_override_max_total_descendant_tasks: None,
1196 agent_override_max_total_descendant_tasks: None,
1197 platform_default_max_active_detached_tasks: DEFAULT_MAX_ACTIVE_DETACHED_TASKS,
1198 org_override_max_active_detached_tasks: None,
1199 agent_override_max_active_detached_tasks: None,
1200 platform_default_max_total_detached_tasks: DEFAULT_MAX_TOTAL_DETACHED_TASKS,
1201 org_override_max_total_detached_tasks: None,
1202 agent_override_max_total_detached_tasks: None,
1203 }
1204 }
1205}
1206
1207impl SubagentNestingPolicy {
1208 pub fn max_subagent_depth(self) -> u32 {
1209 self.agent_override
1210 .or(self.org_override)
1211 .unwrap_or(self.platform_default)
1212 }
1213
1214 pub fn max_active_descendant_tasks(self) -> u32 {
1215 self.agent_override_max_active_descendant_tasks
1216 .or(self.org_override_max_active_descendant_tasks)
1217 .unwrap_or(self.platform_default_max_active_descendant_tasks)
1218 }
1219
1220 pub fn max_total_descendant_tasks(self) -> u32 {
1221 self.agent_override_max_total_descendant_tasks
1222 .or(self.org_override_max_total_descendant_tasks)
1223 .unwrap_or(self.platform_default_max_total_descendant_tasks)
1224 }
1225
1226 pub fn max_active_detached_tasks(self) -> u32 {
1227 self.agent_override_max_active_detached_tasks
1228 .or(self.org_override_max_active_detached_tasks)
1229 .unwrap_or(self.platform_default_max_active_detached_tasks)
1230 }
1231
1232 pub fn max_total_detached_tasks(self) -> u32 {
1233 self.agent_override_max_total_detached_tasks
1234 .or(self.org_override_max_total_detached_tasks)
1235 .unwrap_or(self.platform_default_max_total_detached_tasks)
1236 }
1237
1238 pub fn with_platform_default(mut self, depth: u32) -> Self {
1239 self.platform_default = depth;
1240 self
1241 }
1242
1243 pub fn with_org_override(mut self, depth: Option<u32>) -> Self {
1244 self.org_override = depth;
1245 self
1246 }
1247
1248 pub fn with_agent_override(mut self, depth: Option<u32>) -> Self {
1249 self.agent_override = depth;
1250 self
1251 }
1252
1253 pub fn with_agent_task_caps_override(
1254 mut self,
1255 max_active: Option<u32>,
1256 max_total: Option<u32>,
1257 ) -> Self {
1258 self.agent_override_max_active_descendant_tasks = max_active;
1259 self.agent_override_max_total_descendant_tasks = max_total;
1260 self
1261 }
1262
1263 pub fn with_agent_detached_task_caps_override(
1264 mut self,
1265 max_active: Option<u32>,
1266 max_total: Option<u32>,
1267 ) -> Self {
1268 self.agent_override_max_active_detached_tasks = max_active;
1269 self.agent_override_max_total_detached_tasks = max_total;
1270 self
1271 }
1272}
1273
1274pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
1276
1277#[async_trait]
1282pub trait UserConnectionResolver: Send + Sync {
1283 async fn get_connection_token(
1286 &self,
1287 session_id: SessionId,
1288 provider: &str,
1289 ) -> Result<Option<String>>;
1290
1291 async fn get_connection_user(
1296 &self,
1297 _session_id: SessionId,
1298 _provider: &str,
1299 ) -> Result<Option<Uuid>> {
1300 Ok(None)
1301 }
1302
1303 async fn get_connection_token_for_user(
1308 &self,
1309 _user_id: Uuid,
1310 _provider: &str,
1311 ) -> Result<Option<String>> {
1312 Ok(None)
1313 }
1314
1315 async fn get_connection_metadata(
1318 &self,
1319 _session_id: SessionId,
1320 _provider: &str,
1321 ) -> Result<Option<serde_json::Value>> {
1322 Ok(None)
1323 }
1324}
1325
1326#[async_trait]
1336pub trait BudgetChecker: Send + Sync {
1337 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
1339}
1340
1341#[async_trait]
1350pub trait PaymentAuthority: Send + Sync {
1351 async fn execute_machine_payment(
1352 &self,
1353 session_id: SessionId,
1354 request: crate::payment::MachinePaymentRequest,
1355 ) -> Result<crate::payment::MachinePaymentResponse>;
1356}
1357
1358#[async_trait]
1368pub trait SessionCreationAuthority: Send + Sync {
1369 async fn authorize_session_creation(&self, session_id: SessionId) -> Result<SessionId>;
1373}
1374
1375#[async_trait]
1385pub trait OutboundToolRateLimiter: Send + Sync {
1386 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
1388}
1389
1390#[derive(Debug)]
1396pub enum ToolCallClaimResult {
1397 Claimed { claim_token: uuid::Uuid },
1400 AlreadySettled {
1402 result_json: serde_json::Value,
1403 args_fingerprint: String,
1404 },
1405 AlreadyRunning { args_fingerprint: String },
1410 DeterminismViolation {
1414 stored_fingerprint: String,
1415 current_fingerprint: String,
1416 },
1417}
1418
1419#[derive(Debug, Clone)]
1421pub enum DurableToolCallStatus {
1422 Settled { result_json: serde_json::Value },
1424 Interrupted {
1426 result_json: Option<serde_json::Value>,
1427 },
1428 Running,
1430}
1431
1432#[async_trait]
1437pub trait DurableToolResultStore: Send + Sync + 'static {
1438 async fn try_claim_tool_call(
1446 &self,
1447 turn_id: &str,
1448 tool_call_id: &str,
1449 tool_name: &str,
1450 args_fingerprint: &str,
1451 ) -> Result<ToolCallClaimResult>;
1452
1453 async fn settle_tool_call(
1459 &self,
1460 turn_id: &str,
1461 tool_call_id: &str,
1462 result_json: serde_json::Value,
1463 status: &str,
1464 claim_token: uuid::Uuid,
1465 ) -> Result<bool>;
1466
1467 async fn get_tool_call_status(
1472 &self,
1473 turn_id: &str,
1474 tool_call_id: &str,
1475 ) -> Result<Option<DurableToolCallStatus>>;
1476}
1477
1478pub struct NoopDurableToolResultStore;
1481
1482#[async_trait]
1483impl DurableToolResultStore for NoopDurableToolResultStore {
1484 async fn try_claim_tool_call(
1485 &self,
1486 _turn_id: &str,
1487 _tool_call_id: &str,
1488 _tool_name: &str,
1489 _args_fingerprint: &str,
1490 ) -> Result<ToolCallClaimResult> {
1491 Ok(ToolCallClaimResult::Claimed {
1492 claim_token: uuid::Uuid::new_v4(),
1493 })
1494 }
1495
1496 async fn settle_tool_call(
1497 &self,
1498 _turn_id: &str,
1499 _tool_call_id: &str,
1500 _result_json: serde_json::Value,
1501 _status: &str,
1502 _claim_token: uuid::Uuid,
1503 ) -> Result<bool> {
1504 Ok(true)
1505 }
1506
1507 async fn get_tool_call_status(
1508 &self,
1509 _turn_id: &str,
1510 _tool_call_id: &str,
1511 ) -> Result<Option<DurableToolCallStatus>> {
1512 Ok(None)
1513 }
1514}
1515
1516#[derive(Debug, Clone)]
1522pub struct StreamProgress {
1523 pub accumulated_len: usize,
1525 pub last_delta_at: u64,
1527}
1528
1529#[async_trait]
1535pub trait StreamHeartbeater: Send + Sync {
1536 async fn heartbeat(&self, progress: StreamProgress);
1542}
1543
1544pub struct NoopStreamHeartbeater;
1546
1547#[async_trait]
1548impl StreamHeartbeater for NoopStreamHeartbeater {
1549 async fn heartbeat(&self, _progress: StreamProgress) {}
1550}
1551
1552#[derive(Debug, Clone)]
1558pub struct PartialStreamState {
1559 pub accumulated: String,
1562}
1563
1564#[async_trait]
1572pub trait PartialStreamStore: Send + Sync {
1573 async fn get_partial_stream(
1576 &self,
1577 session_id: SessionId,
1578 turn_id: &str,
1579 ) -> Result<Option<PartialStreamState>>;
1580}
1581
1582pub struct NoopPartialStreamStore;
1584
1585#[async_trait]
1586impl PartialStreamStore for NoopPartialStreamStore {
1587 async fn get_partial_stream(
1588 &self,
1589 _session_id: SessionId,
1590 _turn_id: &str,
1591 ) -> Result<Option<PartialStreamState>> {
1592 Ok(None)
1593 }
1594}
1595
1596#[derive(Clone)]
1605pub struct ToolContext {
1606 pub session_id: SessionId,
1608 pub workspace_id: WorkspaceId,
1615
1616 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1618
1619 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1621
1622 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1624
1625 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1627
1628 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1630
1631 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1637
1638 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1640
1641 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1643
1644 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1646
1647 pub session_store: Option<Arc<dyn SessionStore>>,
1649
1650 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1652
1653 pub agent_store: Option<Arc<dyn AgentStore>>,
1655
1656 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1658
1659 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1661
1662 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1664 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1666
1667 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1671
1672 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1674
1675 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1677
1678 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1681
1682 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1685
1686 pub event_context: Option<crate::events::EventContext>,
1689
1690 pub tool_call_id: Option<String>,
1693 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1695
1696 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1699
1700 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1704
1705 pub org_id: Option<crate::typed_id::OrgId>,
1707
1708 pub network_access: Option<crate::network_access::NetworkAccessList>,
1711
1712 pub locale: Option<String>,
1716
1717 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1719
1720 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1722
1723 pub session_creation_authority: Option<Arc<dyn SessionCreationAuthority>>,
1725
1726 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1730
1731 pub subagent_nesting_policy: SubagentNestingPolicy,
1733
1734 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1738}
1739
1740impl ToolContext {
1741 pub fn workspace_fs_key(&self) -> SessionId {
1746 SessionId::from_uuid(self.workspace_id.uuid())
1747 }
1748
1749 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1751 self.workspace_id = workspace_id;
1752 self
1753 }
1754
1755 pub fn new(session_id: SessionId) -> Self {
1757 Self {
1758 session_id,
1759 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1760 file_store: None,
1761 storage_store: None,
1762 image_store: None,
1763 provider_credential_store: None,
1764 utility_llm_service: None,
1765 mcp_invoker: None,
1766 egress_service: None,
1767 sqldb_store: None,
1768 message_retriever: None,
1769 session_store: None,
1770 session_mutator: None,
1771 agent_store: None,
1772 connection_resolver: None,
1773 schedule_store: None,
1774 platform_store: None,
1775 knowledge_store: None,
1776 knowledge_index_search: None,
1777 leased_resource_store: None,
1778 session_resource_registry: None,
1779 session_task_registry: None,
1780 event_emitter: None,
1781 event_context: None,
1782 tool_call_id: None,
1783 capability_registry: None,
1784 tool_registry: None,
1785 visible_tool_names: None,
1786 org_id: None,
1787 network_access: None,
1788 locale: None,
1789 budget_checker: None,
1790 payment_authority: None,
1791 session_creation_authority: None,
1792 subagent_spawn_store: None,
1793 subagent_nesting_policy: SubagentNestingPolicy::default(),
1794 reasoning_effort_handle: None,
1795 }
1796 }
1797
1798 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1800 Self {
1801 session_id,
1802 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1803 file_store: Some(file_store),
1804 storage_store: None,
1805 image_store: None,
1806 provider_credential_store: None,
1807 utility_llm_service: None,
1808 mcp_invoker: None,
1809 egress_service: None,
1810 sqldb_store: None,
1811 message_retriever: None,
1812 session_store: None,
1813 session_mutator: None,
1814 agent_store: None,
1815 connection_resolver: None,
1816 schedule_store: None,
1817 platform_store: None,
1818 knowledge_store: None,
1819 knowledge_index_search: None,
1820 leased_resource_store: None,
1821 session_resource_registry: None,
1822 session_task_registry: None,
1823 event_emitter: None,
1824 event_context: None,
1825 tool_call_id: None,
1826 capability_registry: None,
1827 tool_registry: None,
1828 visible_tool_names: None,
1829 org_id: None,
1830 network_access: None,
1831 locale: None,
1832 budget_checker: None,
1833 payment_authority: None,
1834 session_creation_authority: None,
1835 subagent_spawn_store: None,
1836 subagent_nesting_policy: SubagentNestingPolicy::default(),
1837 reasoning_effort_handle: None,
1838 }
1839 }
1840
1841 pub fn with_storage_store(
1843 session_id: SessionId,
1844 storage_store: Arc<dyn SessionStorageStore>,
1845 ) -> Self {
1846 Self {
1847 session_id,
1848 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1849 file_store: None,
1850 storage_store: Some(storage_store),
1851 image_store: None,
1852 provider_credential_store: None,
1853 utility_llm_service: None,
1854 mcp_invoker: None,
1855 egress_service: None,
1856 sqldb_store: None,
1857 message_retriever: None,
1858 session_store: None,
1859 session_mutator: None,
1860 agent_store: None,
1861 connection_resolver: None,
1862 schedule_store: None,
1863 platform_store: None,
1864 knowledge_store: None,
1865 knowledge_index_search: None,
1866 leased_resource_store: None,
1867 session_resource_registry: None,
1868 session_task_registry: None,
1869 event_emitter: None,
1870 event_context: None,
1871 tool_call_id: None,
1872 capability_registry: None,
1873 tool_registry: None,
1874 visible_tool_names: None,
1875 org_id: None,
1876 network_access: None,
1877 locale: None,
1878 budget_checker: None,
1879 payment_authority: None,
1880 session_creation_authority: None,
1881 subagent_spawn_store: None,
1882 subagent_nesting_policy: SubagentNestingPolicy::default(),
1883 reasoning_effort_handle: None,
1884 }
1885 }
1886
1887 pub fn with_stores(
1889 session_id: SessionId,
1890 file_store: Arc<dyn SessionFileSystem>,
1891 storage_store: Arc<dyn SessionStorageStore>,
1892 ) -> Self {
1893 Self {
1894 session_id,
1895 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1896 file_store: Some(file_store),
1897 storage_store: Some(storage_store),
1898 sqldb_store: None,
1899 image_store: None,
1900 provider_credential_store: None,
1901 utility_llm_service: None,
1902 mcp_invoker: None,
1903 egress_service: None,
1904 message_retriever: None,
1905 session_store: None,
1906 session_mutator: None,
1907 agent_store: None,
1908 connection_resolver: None,
1909 schedule_store: None,
1910 platform_store: None,
1911 knowledge_store: None,
1912 knowledge_index_search: None,
1913 leased_resource_store: None,
1914 session_resource_registry: None,
1915 session_task_registry: None,
1916 event_emitter: None,
1917 event_context: None,
1918 tool_call_id: None,
1919 capability_registry: None,
1920 tool_registry: None,
1921 visible_tool_names: None,
1922 org_id: None,
1923 network_access: None,
1924 locale: None,
1925 budget_checker: None,
1926 payment_authority: None,
1927 session_creation_authority: None,
1928 subagent_spawn_store: None,
1929 subagent_nesting_policy: SubagentNestingPolicy::default(),
1930 reasoning_effort_handle: None,
1931 }
1932 }
1933
1934 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1936 self.sqldb_store = Some(sqldb_store);
1937 self
1938 }
1939
1940 pub fn with_message_retriever(
1942 mut self,
1943 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1944 ) -> Self {
1945 self.message_retriever = Some(retriever);
1946 self
1947 }
1948
1949 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1951 self.session_store = Some(store);
1952 self
1953 }
1954
1955 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1957 self.session_mutator = Some(mutator);
1958 self
1959 }
1960
1961 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
1965 self.reasoning_effort_handle = Some(handle);
1966 self
1967 }
1968
1969 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1971 self.agent_store = Some(store);
1972 self
1973 }
1974
1975 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1977 self.connection_resolver = Some(resolver);
1978 self
1979 }
1980
1981 pub fn with_image_store(
1983 session_id: SessionId,
1984 image_store: Arc<dyn ImageArtifactStore>,
1985 ) -> Self {
1986 Self {
1987 session_id,
1988 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1989 file_store: None,
1990 storage_store: None,
1991 image_store: Some(image_store),
1992 provider_credential_store: None,
1993 utility_llm_service: None,
1994 mcp_invoker: None,
1995 egress_service: None,
1996 sqldb_store: None,
1997 message_retriever: None,
1998 session_store: None,
1999 session_mutator: None,
2000 agent_store: None,
2001 connection_resolver: None,
2002 schedule_store: None,
2003 platform_store: None,
2004 knowledge_store: None,
2005 knowledge_index_search: None,
2006 leased_resource_store: None,
2007 session_resource_registry: None,
2008 session_task_registry: None,
2009 event_emitter: None,
2010 event_context: None,
2011 tool_call_id: None,
2012 capability_registry: None,
2013 tool_registry: None,
2014 visible_tool_names: None,
2015 org_id: None,
2016 network_access: None,
2017 locale: None,
2018 budget_checker: None,
2019 payment_authority: None,
2020 session_creation_authority: None,
2021 subagent_spawn_store: None,
2022 subagent_nesting_policy: SubagentNestingPolicy::default(),
2023 reasoning_effort_handle: None,
2024 }
2025 }
2026
2027 pub fn with_provider_credential_store(
2029 mut self,
2030 store: Arc<dyn ProviderCredentialStore>,
2031 ) -> Self {
2032 self.provider_credential_store = Some(store);
2033 self
2034 }
2035
2036 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
2038 self.utility_llm_service = Some(service);
2039 self
2040 }
2041
2042 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
2044 self.mcp_invoker = Some(invoker);
2045 self
2046 }
2047
2048 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
2050 self.egress_service = Some(service);
2051 self
2052 }
2053
2054 pub fn with_egress_service_opt(
2057 mut self,
2058 service: Option<Arc<dyn crate::EgressService>>,
2059 ) -> Self {
2060 if let Some(service) = service {
2061 self.egress_service = Some(service);
2062 }
2063 self
2064 }
2065
2066 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
2068 self.storage_store = Some(store);
2069 self
2070 }
2071
2072 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
2074 self.schedule_store = Some(store);
2075 self
2076 }
2077
2078 pub fn with_platform_store(
2080 mut self,
2081 store: Arc<dyn crate::platform_store::PlatformStore>,
2082 ) -> Self {
2083 self.platform_store = Some(store);
2084 self
2085 }
2086
2087 pub fn with_knowledge_index_search(
2089 mut self,
2090 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
2091 ) -> Self {
2092 self.knowledge_index_search = Some(search);
2093 self
2094 }
2095
2096 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
2098 self.leased_resource_store = Some(store);
2099 self
2100 }
2101
2102 pub fn with_session_resource_registry(
2104 mut self,
2105 registry: Arc<dyn SessionResourceRegistry>,
2106 ) -> Self {
2107 self.session_resource_registry = Some(registry);
2108 self
2109 }
2110
2111 pub fn with_session_task_registry(
2113 mut self,
2114 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
2115 ) -> Self {
2116 self.session_task_registry = Some(registry);
2117 self
2118 }
2119
2120 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
2122 self.org_id = Some(org_id);
2123 self
2124 }
2125
2126 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
2128 self.tool_registry = Some(registry);
2129 self
2130 }
2131
2132 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
2134 self.visible_tool_names = Some(names);
2135 self
2136 }
2137
2138 pub fn with_network_access(
2140 mut self,
2141 network_access: Option<crate::network_access::NetworkAccessList>,
2142 ) -> Self {
2143 self.network_access = network_access;
2144 self
2145 }
2146
2147 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
2149 self.payment_authority = Some(authority);
2150 self
2151 }
2152
2153 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
2155 self.subagent_spawn_store = Some(store);
2156 self
2157 }
2158
2159 pub fn with_subagent_nesting_policy(mut self, policy: SubagentNestingPolicy) -> Self {
2161 self.subagent_nesting_policy = policy;
2162 self
2163 }
2164
2165 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
2170 let (Some(emitter), Some(ctx), Some(call_id)) =
2171 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2172 else {
2173 return;
2174 };
2175 if let Err(e) = emitter
2176 .emit(EventRequest::new(
2177 self.session_id,
2178 ctx.clone(),
2179 crate::events::ToolProgressData {
2180 tool_call_id: call_id.clone(),
2181 tool_name: tool_name.to_string(),
2182 message: message.to_string(),
2183 display_name: None,
2184 },
2185 ))
2186 .await
2187 {
2188 tracing::debug!(
2189 tool_call_id = call_id,
2190 tool_name,
2191 error = %e,
2192 "Failed to emit tool.progress event"
2193 );
2194 }
2195 }
2196
2197 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
2202 let (Some(emitter), Some(ctx), Some(call_id)) =
2203 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2204 else {
2205 return;
2206 };
2207 if let Err(e) = emitter
2208 .emit(EventRequest::new(
2209 self.session_id,
2210 ctx.clone(),
2211 crate::events::ToolOutputDeltaData {
2212 tool_call_id: call_id.clone(),
2213 tool_name: tool_name.to_string(),
2214 delta: delta.to_string(),
2215 stream: stream.to_string(),
2216 },
2217 ))
2218 .await
2219 {
2220 tracing::debug!(
2221 tool_call_id = call_id,
2222 tool_name,
2223 error = %e,
2224 "Failed to emit tool.output.delta event"
2225 );
2226 }
2227 }
2228}
2229
2230impl std::fmt::Debug for ToolContext {
2231 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2232 f.debug_struct("ToolContext")
2233 .field("session_id", &self.session_id)
2234 .field("file_store", &self.file_store.is_some())
2235 .field("storage_store", &self.storage_store.is_some())
2236 .field("image_store", &self.image_store.is_some())
2237 .field(
2238 "provider_credential_store",
2239 &self.provider_credential_store.is_some(),
2240 )
2241 .field("utility_llm_service", &self.utility_llm_service.is_some())
2242 .field("egress_service", &self.egress_service.is_some())
2243 .field("sqldb_store", &self.sqldb_store.is_some())
2244 .field("message_retriever", &self.message_retriever.is_some())
2245 .field("session_store", &self.session_store.is_some())
2246 .field("session_mutator", &self.session_mutator.is_some())
2247 .field("agent_store", &self.agent_store.is_some())
2248 .field("connection_resolver", &self.connection_resolver.is_some())
2249 .field("schedule_store", &self.schedule_store.is_some())
2250 .field("platform_store", &self.platform_store.is_some())
2251 .field(
2252 "knowledge_index_search",
2253 &self.knowledge_index_search.is_some(),
2254 )
2255 .field(
2256 "leased_resource_store",
2257 &self.leased_resource_store.is_some(),
2258 )
2259 .field("event_emitter", &self.event_emitter.is_some())
2260 .field("tool_registry", &self.tool_registry.is_some())
2261 .field("payment_authority", &self.payment_authority.is_some())
2262 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
2263 .field("subagent_nesting_policy", &self.subagent_nesting_policy)
2264 .field("org_id", &self.org_id)
2265 .finish()
2266 }
2267}
2268
2269use crate::events::{Event, EventRequest};
2274
2275#[async_trait]
2286pub trait EventEmitter: Send + Sync {
2287 async fn emit(&self, request: EventRequest) -> Result<Event>;
2292}
2293
2294#[async_trait]
2296impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
2297 async fn emit(&self, request: EventRequest) -> Result<Event> {
2298 (**self).emit(request).await
2299 }
2300}
2301
2302#[derive(Debug, Clone, Default)]
2306pub struct NoopEventEmitter;
2307
2308#[async_trait]
2309impl EventEmitter for NoopEventEmitter {
2310 async fn emit(&self, request: EventRequest) -> Result<Event> {
2311 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2313 }
2314}
2315
2316#[derive(Debug, Clone)]
2329pub struct ResolvedImage {
2330 pub base64: String,
2332 pub media_type: String,
2334}
2335
2336impl ResolvedImage {
2337 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2339 Self {
2340 base64: base64.into(),
2341 media_type: media_type.into(),
2342 }
2343 }
2344
2345 pub fn to_data_url(&self) -> String {
2349 format!("data:{};base64,{}", self.media_type, self.base64)
2350 }
2351}
2352
2353#[async_trait]
2386pub trait ImageResolver: Send + Sync {
2387 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2391}
2392
2393#[derive(Debug)]
2399pub enum SpawnClaimResult {
2400 Claimed {
2403 spawn_handle_id: uuid::Uuid,
2404 claim_token: uuid::Uuid,
2405 },
2406 ClaimedPendingChild {
2410 spawn_handle_id: uuid::Uuid,
2411 claim_token: uuid::Uuid,
2412 },
2413 AlreadyRunning {
2416 child_session_id: crate::typed_id::SessionId,
2417 claim_token: uuid::Uuid,
2419 },
2420 AlreadySettled {
2423 child_session_id: crate::typed_id::SessionId,
2424 terminal_status: String,
2426 terminal_result: String,
2427 },
2428}
2429
2430#[async_trait]
2438pub trait SubagentSpawnStore: Send + Sync + 'static {
2439 async fn try_claim_spawn(
2444 &self,
2445 parent_session_id: crate::typed_id::SessionId,
2446 tool_call_id: &str,
2447 claim_token: uuid::Uuid,
2448 ) -> Result<SpawnClaimResult>;
2449
2450 async fn register_child_session(
2455 &self,
2456 spawn_handle_id: uuid::Uuid,
2457 claim_token: uuid::Uuid,
2458 child_session_id: crate::typed_id::SessionId,
2459 ) -> Result<()>;
2460
2461 async fn settle_spawn(
2467 &self,
2468 parent_session_id: crate::typed_id::SessionId,
2469 tool_call_id: &str,
2470 claim_token: uuid::Uuid,
2471 terminal_status: &str,
2472 terminal_result: &str,
2473 ) -> Result<()>;
2474}
2475
2476#[async_trait]
2478impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2479 async fn try_claim_spawn(
2480 &self,
2481 parent_session_id: crate::typed_id::SessionId,
2482 tool_call_id: &str,
2483 claim_token: uuid::Uuid,
2484 ) -> Result<SpawnClaimResult> {
2485 (**self)
2486 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2487 .await
2488 }
2489
2490 async fn register_child_session(
2491 &self,
2492 spawn_handle_id: uuid::Uuid,
2493 claim_token: uuid::Uuid,
2494 child_session_id: crate::typed_id::SessionId,
2495 ) -> Result<()> {
2496 (**self)
2497 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2498 .await
2499 }
2500
2501 async fn settle_spawn(
2502 &self,
2503 parent_session_id: crate::typed_id::SessionId,
2504 tool_call_id: &str,
2505 claim_token: uuid::Uuid,
2506 terminal_status: &str,
2507 terminal_result: &str,
2508 ) -> Result<()> {
2509 (**self)
2510 .settle_spawn(
2511 parent_session_id,
2512 tool_call_id,
2513 claim_token,
2514 terminal_status,
2515 terminal_result,
2516 )
2517 .await
2518 }
2519}
2520
2521pub struct NoopSubagentSpawnStore;
2525
2526#[async_trait]
2527impl SubagentSpawnStore for NoopSubagentSpawnStore {
2528 async fn try_claim_spawn(
2529 &self,
2530 _parent_session_id: crate::typed_id::SessionId,
2531 _tool_call_id: &str,
2532 claim_token: uuid::Uuid,
2533 ) -> Result<SpawnClaimResult> {
2534 Ok(SpawnClaimResult::Claimed {
2535 spawn_handle_id: uuid::Uuid::new_v4(),
2536 claim_token,
2537 })
2538 }
2539
2540 async fn register_child_session(
2541 &self,
2542 _spawn_handle_id: uuid::Uuid,
2543 _claim_token: uuid::Uuid,
2544 _child_session_id: crate::typed_id::SessionId,
2545 ) -> Result<()> {
2546 Ok(())
2547 }
2548
2549 async fn settle_spawn(
2550 &self,
2551 _parent_session_id: crate::typed_id::SessionId,
2552 _tool_call_id: &str,
2553 _claim_token: uuid::Uuid,
2554 _terminal_status: &str,
2555 _terminal_result: &str,
2556 ) -> Result<()> {
2557 Ok(())
2558 }
2559}
2560
2561#[cfg(test)]
2566mod tests {
2567 use super::*;
2568
2569 #[test]
2570 fn test_resolved_image_new() {
2571 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2572 assert_eq!(image.base64, "SGVsbG8=");
2573 assert_eq!(image.media_type, "image/png");
2574 }
2575
2576 #[test]
2577 fn test_resolved_image_to_data_url() {
2578 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2579 let data_url = image.to_data_url();
2580 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2581 }
2582
2583 #[test]
2584 fn test_resolved_image_jpeg() {
2585 let image = ResolvedImage::new("base64data", "image/jpeg");
2586 let data_url = image.to_data_url();
2587 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2588 }
2589}