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 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>>;
464
465 async fn write_file(
467 &self,
468 session_id: SessionId,
469 path: &str,
470 content: &str,
471 encoding: &str,
472 ) -> Result<SessionFile>;
473
474 async fn write_file_if_content_matches(
479 &self,
480 session_id: SessionId,
481 path: &str,
482 expected_content: &str,
483 expected_encoding: &str,
484 content: &str,
485 encoding: &str,
486 ) -> Result<Option<SessionFile>> {
487 let Some(existing) = self.read_file(session_id, path).await? else {
488 return Ok(None);
489 };
490
491 if existing.is_directory {
492 return Ok(None);
493 }
494
495 let current_content = existing.content.unwrap_or_default();
496 if current_content != expected_content || existing.encoding != expected_encoding {
497 return Ok(None);
498 }
499
500 self.write_file(session_id, path, content, encoding)
501 .await
502 .map(Some)
503 }
504
505 async fn delete_file(&self, session_id: SessionId, path: &str, recursive: bool)
507 -> Result<bool>;
508
509 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>>;
511
512 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>>;
514
515 async fn grep_files(
517 &self,
518 session_id: SessionId,
519 pattern: &str,
520 path_pattern: Option<&str>,
521 ) -> Result<Vec<GrepMatch>>;
522
523 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
525
526 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
528 if file.is_readonly {
529 return Err(crate::error::AgentLoopError::store(
530 "read-only initial files require a SessionFileSystem-specific seed implementation",
531 ));
532 }
533 self.write_file(session_id, &file.path, &file.content, &file.encoding)
534 .await?;
535 Ok(())
536 }
537}
538
539pub struct WorkspaceScopedFileSystem {
549 inner: Arc<dyn SessionFileSystem>,
550 key: SessionId,
551}
552
553impl WorkspaceScopedFileSystem {
554 pub fn wrap(
556 inner: Arc<dyn SessionFileSystem>,
557 workspace_id: WorkspaceId,
558 ) -> Arc<dyn SessionFileSystem> {
559 Arc::new(Self {
560 inner,
561 key: SessionId::from_uuid(workspace_id.uuid()),
562 })
563 }
564}
565
566#[async_trait]
567impl SessionFileSystem for WorkspaceScopedFileSystem {
568 async fn read_file(&self, _session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
569 self.inner.read_file(self.key, path).await
570 }
571 async fn write_file(
572 &self,
573 _session_id: SessionId,
574 path: &str,
575 content: &str,
576 encoding: &str,
577 ) -> Result<SessionFile> {
578 self.inner
579 .write_file(self.key, path, content, encoding)
580 .await
581 }
582 async fn write_file_if_content_matches(
583 &self,
584 _session_id: SessionId,
585 path: &str,
586 expected_content: &str,
587 expected_encoding: &str,
588 content: &str,
589 encoding: &str,
590 ) -> Result<Option<SessionFile>> {
591 self.inner
592 .write_file_if_content_matches(
593 self.key,
594 path,
595 expected_content,
596 expected_encoding,
597 content,
598 encoding,
599 )
600 .await
601 }
602 async fn delete_file(
603 &self,
604 _session_id: SessionId,
605 path: &str,
606 recursive: bool,
607 ) -> Result<bool> {
608 self.inner.delete_file(self.key, path, recursive).await
609 }
610 async fn list_directory(&self, _session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
611 self.inner.list_directory(self.key, path).await
612 }
613 async fn stat_file(&self, _session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
614 self.inner.stat_file(self.key, path).await
615 }
616 async fn grep_files(
617 &self,
618 _session_id: SessionId,
619 pattern: &str,
620 path_pattern: Option<&str>,
621 ) -> Result<Vec<GrepMatch>> {
622 self.inner.grep_files(self.key, pattern, path_pattern).await
623 }
624 async fn create_directory(&self, _session_id: SessionId, path: &str) -> Result<FileInfo> {
625 self.inner.create_directory(self.key, path).await
626 }
627 async fn seed_initial_file(&self, _session_id: SessionId, file: &InitialFile) -> Result<()> {
628 self.inner.seed_initial_file(self.key, file).await
629 }
630
631 fn display_root(&self) -> String {
632 self.inner.display_root()
633 }
634
635 fn display_path(&self, path: &str) -> String {
636 self.inner.display_path(path)
637 }
638
639 fn resolve_path(&self, input: &str) -> String {
640 self.inner.resolve_path(input)
641 }
642}
643
644#[async_trait]
645impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
646 fn display_root(&self) -> String {
647 (**self).display_root()
648 }
649
650 fn display_path(&self, path: &str) -> String {
651 (**self).display_path(path)
652 }
653
654 fn resolve_path(&self, input: &str) -> String {
655 (**self).resolve_path(input)
656 }
657
658 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
659 (**self).read_file(session_id, path).await
660 }
661
662 async fn write_file(
663 &self,
664 session_id: SessionId,
665 path: &str,
666 content: &str,
667 encoding: &str,
668 ) -> Result<SessionFile> {
669 (**self)
670 .write_file(session_id, path, content, encoding)
671 .await
672 }
673
674 async fn write_file_if_content_matches(
675 &self,
676 session_id: SessionId,
677 path: &str,
678 expected_content: &str,
679 expected_encoding: &str,
680 content: &str,
681 encoding: &str,
682 ) -> Result<Option<SessionFile>> {
683 (**self)
684 .write_file_if_content_matches(
685 session_id,
686 path,
687 expected_content,
688 expected_encoding,
689 content,
690 encoding,
691 )
692 .await
693 }
694
695 async fn delete_file(
696 &self,
697 session_id: SessionId,
698 path: &str,
699 recursive: bool,
700 ) -> Result<bool> {
701 (**self).delete_file(session_id, path, recursive).await
702 }
703
704 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
705 (**self).list_directory(session_id, path).await
706 }
707
708 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
709 (**self).stat_file(session_id, path).await
710 }
711
712 async fn grep_files(
713 &self,
714 session_id: SessionId,
715 pattern: &str,
716 path_pattern: Option<&str>,
717 ) -> Result<Vec<GrepMatch>> {
718 (**self).grep_files(session_id, pattern, path_pattern).await
719 }
720
721 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
722 (**self).create_directory(session_id, path).await
723 }
724
725 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
726 (**self).seed_initial_file(session_id, file).await
727 }
728}
729
730pub use SessionFileSystem as SessionFileStore;
732
733#[derive(Clone, Default)]
739pub struct SessionFileSystemFactoryContext {
740 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
741}
742
743impl SessionFileSystemFactoryContext {
744 pub fn new() -> Self {
745 Self::default()
746 }
747
748 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
749 let values = Arc::make_mut(&mut self.values);
750 values.insert(TypeId::of::<T>(), value);
751 self
752 }
753
754 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
755 self.values
756 .get(&TypeId::of::<T>())
757 .and_then(|value| value.clone().downcast::<T>().ok())
758 }
759
760 pub fn with_workspace_roots(self, roots: Arc<crate::WorkspaceRootSet>) -> Self {
761 self.with(roots)
762 }
763
764 pub fn workspace_roots(&self) -> Option<Arc<crate::WorkspaceRootSet>> {
765 self.get::<crate::WorkspaceRootSet>()
766 }
767}
768
769#[async_trait]
771pub trait SessionFileSystemFactory: Send + Sync {
772 fn name(&self) -> &'static str {
774 "SessionFileSystemFactory"
775 }
776
777 fn is_disabled(&self) -> bool {
780 false
781 }
782
783 async fn create_session_file_system(
785 &self,
786 context: SessionFileSystemFactoryContext,
787 ) -> Result<Arc<dyn SessionFileSystem>>;
788}
789
790#[derive(Debug, Clone, Default)]
792pub struct DisabledSessionFileSystemFactory;
793
794#[async_trait]
795impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
796 fn name(&self) -> &'static str {
797 "DisabledSessionFileSystemFactory"
798 }
799
800 fn is_disabled(&self) -> bool {
801 true
802 }
803
804 async fn create_session_file_system(
805 &self,
806 _context: SessionFileSystemFactoryContext,
807 ) -> Result<Arc<dyn SessionFileSystem>> {
808 Err(crate::error::AgentLoopError::config(
809 "session filesystem is disabled",
810 ))
811 }
812}
813
814#[derive(Debug, Clone)]
820pub struct KeyInfo {
821 pub key: String,
822 pub created_at: chrono::DateTime<chrono::Utc>,
823 pub updated_at: chrono::DateTime<chrono::Utc>,
824}
825
826#[derive(Debug, Clone)]
828pub struct SecretInfo {
829 pub name: String,
830 pub created_at: chrono::DateTime<chrono::Utc>,
831 pub updated_at: chrono::DateTime<chrono::Utc>,
832}
833
834#[derive(Debug, Clone, serde::Serialize)]
844pub struct KnowledgeSearchHit {
845 pub id: String,
847 pub kb_id: String,
849 pub title: String,
850 pub kind: String,
851 pub tags: Vec<String>,
852 pub snippet: String,
854 pub resource: Option<String>,
856}
857
858#[async_trait]
862pub trait KnowledgeStore: Send + Sync {
863 async fn search_knowledge(
864 &self,
865 org_id: crate::typed_id::OrgId,
866 kb_public_ids: &[String],
867 query: &str,
868 kind: Option<&str>,
869 tags: &[String],
870 limit: usize,
871 ) -> Result<Vec<KnowledgeSearchHit>>;
872}
873
874#[async_trait]
879pub trait SessionStorageStore: Send + Sync {
880 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
884
885 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
887
888 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
890
891 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
893
894 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
898
899 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
901
902 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
904
905 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
907}
908
909use crate::session_schedule::SessionSchedule;
914use crate::typed_id::ScheduleId;
915
916#[async_trait]
920pub trait SessionScheduleStore: Send + Sync {
921 async fn create_schedule(
923 &self,
924 session_id: SessionId,
925 description: String,
926 cron_expression: Option<String>,
927 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
928 timezone: String,
929 ) -> Result<SessionSchedule>;
930
931 async fn create_schedule_enforcing_limits(
935 &self,
936 session_id: SessionId,
937 description: String,
938 cron_expression: Option<String>,
939 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
940 timezone: String,
941 ) -> std::result::Result<SessionSchedule, crate::session_schedule::ScheduleLimitError> {
942 let per_session = self
943 .count_active_schedules(session_id)
944 .await
945 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
946 if per_session >= crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION {
947 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
948 format!(
949 "Maximum {} active schedules per session. Cancel an existing schedule first.",
950 crate::session_schedule::MAX_ACTIVE_SCHEDULES_PER_SESSION
951 ),
952 ));
953 }
954
955 let max_per_org = crate::session_schedule::max_active_schedules_per_org();
956 let per_org = self
957 .count_active_org_schedules()
958 .await
959 .map_err(crate::session_schedule::ScheduleLimitError::Store)?;
960 if i64::from(per_org) >= max_per_org {
961 return Err(crate::session_schedule::ScheduleLimitError::Rejected(
962 format!(
963 "Maximum {max_per_org} active schedules per org reached. Cancel an existing schedule first."
964 ),
965 ));
966 }
967
968 if let Some(cron) = cron_expression.as_deref() {
969 crate::session_schedule::validate_cron_min_interval(cron)
970 .map_err(crate::session_schedule::ScheduleLimitError::Rejected)?;
971 }
972
973 self.create_schedule(
974 session_id,
975 description,
976 cron_expression,
977 scheduled_at,
978 timezone,
979 )
980 .await
981 .map_err(crate::session_schedule::ScheduleLimitError::Store)
982 }
983
984 async fn cancel_schedule(
986 &self,
987 session_id: SessionId,
988 schedule_id: ScheduleId,
989 ) -> Result<SessionSchedule>;
990
991 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
993
994 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
996
997 async fn count_active_org_schedules(&self) -> Result<u32>;
1002}
1003
1004#[async_trait]
1014pub trait SessionResourceRegistry: Send + Sync {
1015 async fn register(
1017 &self,
1018 entry: crate::session_resource::RegisterSessionResource,
1019 ) -> Result<crate::session_resource::SessionResourceEntry>;
1020
1021 async fn update_status(
1023 &self,
1024 session_id: SessionId,
1025 resource_id: &str,
1026 status: crate::session_resource::SessionResourceStatus,
1027 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1028
1029 async fn get(
1031 &self,
1032 session_id: SessionId,
1033 resource_id: &str,
1034 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
1035
1036 async fn list(
1038 &self,
1039 session_id: SessionId,
1040 filter: Option<&crate::session_resource::SessionResourceFilter>,
1041 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
1042
1043 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
1045}
1046
1047#[async_trait]
1057pub trait LeasedResourceStore: Send + Sync {
1058 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
1064
1065 async fn release_resource(
1071 &self,
1072 session_id: SessionId,
1073 provider: &str,
1074 resource_type: &str,
1075 external_id: &str,
1076 ) -> Result<Option<LeasedResource>>;
1077
1078 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
1083}
1084
1085pub const DEFAULT_MAX_SUBAGENT_DEPTH: u32 = 2;
1092pub const DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS: u32 = 16;
1093pub const DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS: u32 = 200;
1094
1095#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1097pub struct SubagentNestingPolicy {
1098 pub platform_default: u32,
1099 pub org_override: Option<u32>,
1100 pub agent_override: Option<u32>,
1101 pub platform_default_max_active_descendant_tasks: u32,
1102 pub org_override_max_active_descendant_tasks: Option<u32>,
1103 pub agent_override_max_active_descendant_tasks: Option<u32>,
1104 pub platform_default_max_total_descendant_tasks: u32,
1105 pub org_override_max_total_descendant_tasks: Option<u32>,
1106 pub agent_override_max_total_descendant_tasks: Option<u32>,
1107}
1108
1109impl Default for SubagentNestingPolicy {
1110 fn default() -> Self {
1111 Self {
1112 platform_default: DEFAULT_MAX_SUBAGENT_DEPTH,
1113 org_override: None,
1114 agent_override: None,
1115 platform_default_max_active_descendant_tasks:
1116 DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS,
1117 org_override_max_active_descendant_tasks: None,
1118 agent_override_max_active_descendant_tasks: None,
1119 platform_default_max_total_descendant_tasks:
1120 DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS,
1121 org_override_max_total_descendant_tasks: None,
1122 agent_override_max_total_descendant_tasks: None,
1123 }
1124 }
1125}
1126
1127impl SubagentNestingPolicy {
1128 pub fn max_subagent_depth(self) -> u32 {
1129 self.agent_override
1130 .or(self.org_override)
1131 .unwrap_or(self.platform_default)
1132 }
1133
1134 pub fn max_active_descendant_tasks(self) -> u32 {
1135 self.agent_override_max_active_descendant_tasks
1136 .or(self.org_override_max_active_descendant_tasks)
1137 .unwrap_or(self.platform_default_max_active_descendant_tasks)
1138 }
1139
1140 pub fn max_total_descendant_tasks(self) -> u32 {
1141 self.agent_override_max_total_descendant_tasks
1142 .or(self.org_override_max_total_descendant_tasks)
1143 .unwrap_or(self.platform_default_max_total_descendant_tasks)
1144 }
1145
1146 pub fn with_platform_default(mut self, depth: u32) -> Self {
1147 self.platform_default = depth;
1148 self
1149 }
1150
1151 pub fn with_org_override(mut self, depth: Option<u32>) -> Self {
1152 self.org_override = depth;
1153 self
1154 }
1155
1156 pub fn with_agent_override(mut self, depth: Option<u32>) -> Self {
1157 self.agent_override = depth;
1158 self
1159 }
1160
1161 pub fn with_agent_task_caps_override(
1162 mut self,
1163 max_active: Option<u32>,
1164 max_total: Option<u32>,
1165 ) -> Self {
1166 self.agent_override_max_active_descendant_tasks = max_active;
1167 self.agent_override_max_total_descendant_tasks = max_total;
1168 self
1169 }
1170}
1171
1172pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
1174
1175#[async_trait]
1180pub trait UserConnectionResolver: Send + Sync {
1181 async fn get_connection_token(
1184 &self,
1185 session_id: SessionId,
1186 provider: &str,
1187 ) -> Result<Option<String>>;
1188
1189 async fn get_connection_user(
1194 &self,
1195 _session_id: SessionId,
1196 _provider: &str,
1197 ) -> Result<Option<Uuid>> {
1198 Ok(None)
1199 }
1200
1201 async fn get_connection_token_for_user(
1206 &self,
1207 _user_id: Uuid,
1208 _provider: &str,
1209 ) -> Result<Option<String>> {
1210 Ok(None)
1211 }
1212
1213 async fn get_connection_metadata(
1216 &self,
1217 _session_id: SessionId,
1218 _provider: &str,
1219 ) -> Result<Option<serde_json::Value>> {
1220 Ok(None)
1221 }
1222}
1223
1224#[async_trait]
1234pub trait BudgetChecker: Send + Sync {
1235 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
1237}
1238
1239#[async_trait]
1248pub trait PaymentAuthority: Send + Sync {
1249 async fn execute_machine_payment(
1250 &self,
1251 session_id: SessionId,
1252 request: crate::payment::MachinePaymentRequest,
1253 ) -> Result<crate::payment::MachinePaymentResponse>;
1254}
1255
1256#[async_trait]
1266pub trait OutboundToolRateLimiter: Send + Sync {
1267 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
1269}
1270
1271#[derive(Debug)]
1277pub enum ToolCallClaimResult {
1278 Claimed { claim_token: uuid::Uuid },
1281 AlreadySettled {
1283 result_json: serde_json::Value,
1284 args_fingerprint: String,
1285 },
1286 AlreadyRunning { args_fingerprint: String },
1291 DeterminismViolation {
1295 stored_fingerprint: String,
1296 current_fingerprint: String,
1297 },
1298}
1299
1300#[derive(Debug, Clone)]
1302pub enum DurableToolCallStatus {
1303 Settled { result_json: serde_json::Value },
1305 Interrupted {
1307 result_json: Option<serde_json::Value>,
1308 },
1309 Running,
1311}
1312
1313#[async_trait]
1318pub trait DurableToolResultStore: Send + Sync + 'static {
1319 async fn try_claim_tool_call(
1327 &self,
1328 turn_id: &str,
1329 tool_call_id: &str,
1330 tool_name: &str,
1331 args_fingerprint: &str,
1332 ) -> Result<ToolCallClaimResult>;
1333
1334 async fn settle_tool_call(
1340 &self,
1341 turn_id: &str,
1342 tool_call_id: &str,
1343 result_json: serde_json::Value,
1344 status: &str,
1345 claim_token: uuid::Uuid,
1346 ) -> Result<bool>;
1347
1348 async fn get_tool_call_status(
1353 &self,
1354 turn_id: &str,
1355 tool_call_id: &str,
1356 ) -> Result<Option<DurableToolCallStatus>>;
1357}
1358
1359pub struct NoopDurableToolResultStore;
1362
1363#[async_trait]
1364impl DurableToolResultStore for NoopDurableToolResultStore {
1365 async fn try_claim_tool_call(
1366 &self,
1367 _turn_id: &str,
1368 _tool_call_id: &str,
1369 _tool_name: &str,
1370 _args_fingerprint: &str,
1371 ) -> Result<ToolCallClaimResult> {
1372 Ok(ToolCallClaimResult::Claimed {
1373 claim_token: uuid::Uuid::new_v4(),
1374 })
1375 }
1376
1377 async fn settle_tool_call(
1378 &self,
1379 _turn_id: &str,
1380 _tool_call_id: &str,
1381 _result_json: serde_json::Value,
1382 _status: &str,
1383 _claim_token: uuid::Uuid,
1384 ) -> Result<bool> {
1385 Ok(true)
1386 }
1387
1388 async fn get_tool_call_status(
1389 &self,
1390 _turn_id: &str,
1391 _tool_call_id: &str,
1392 ) -> Result<Option<DurableToolCallStatus>> {
1393 Ok(None)
1394 }
1395}
1396
1397#[derive(Debug, Clone)]
1403pub struct StreamProgress {
1404 pub accumulated_len: usize,
1406 pub last_delta_at: u64,
1408}
1409
1410#[async_trait]
1416pub trait StreamHeartbeater: Send + Sync {
1417 async fn heartbeat(&self, progress: StreamProgress);
1423}
1424
1425pub struct NoopStreamHeartbeater;
1427
1428#[async_trait]
1429impl StreamHeartbeater for NoopStreamHeartbeater {
1430 async fn heartbeat(&self, _progress: StreamProgress) {}
1431}
1432
1433#[derive(Debug, Clone)]
1439pub struct PartialStreamState {
1440 pub accumulated: String,
1443}
1444
1445#[async_trait]
1453pub trait PartialStreamStore: Send + Sync {
1454 async fn get_partial_stream(
1457 &self,
1458 session_id: SessionId,
1459 turn_id: &str,
1460 ) -> Result<Option<PartialStreamState>>;
1461}
1462
1463pub struct NoopPartialStreamStore;
1465
1466#[async_trait]
1467impl PartialStreamStore for NoopPartialStreamStore {
1468 async fn get_partial_stream(
1469 &self,
1470 _session_id: SessionId,
1471 _turn_id: &str,
1472 ) -> Result<Option<PartialStreamState>> {
1473 Ok(None)
1474 }
1475}
1476
1477#[derive(Clone)]
1486pub struct ToolContext {
1487 pub session_id: SessionId,
1489 pub workspace_id: WorkspaceId,
1496
1497 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1499
1500 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1502
1503 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1505
1506 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1508
1509 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1511
1512 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1518
1519 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1521
1522 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1524
1525 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1527
1528 pub session_store: Option<Arc<dyn SessionStore>>,
1530
1531 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1533
1534 pub agent_store: Option<Arc<dyn AgentStore>>,
1536
1537 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1539
1540 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1542
1543 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1545 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1547
1548 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1552
1553 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1555
1556 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1558
1559 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1562
1563 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1566
1567 pub event_context: Option<crate::events::EventContext>,
1570
1571 pub tool_call_id: Option<String>,
1574 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1576
1577 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1580
1581 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1585
1586 pub org_id: Option<crate::typed_id::OrgId>,
1588
1589 pub network_access: Option<crate::network_access::NetworkAccessList>,
1592
1593 pub locale: Option<String>,
1597
1598 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1600
1601 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1603
1604 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1608
1609 pub subagent_nesting_policy: SubagentNestingPolicy,
1611
1612 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1616}
1617
1618impl ToolContext {
1619 pub fn workspace_fs_key(&self) -> SessionId {
1624 SessionId::from_uuid(self.workspace_id.uuid())
1625 }
1626
1627 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1629 self.workspace_id = workspace_id;
1630 self
1631 }
1632
1633 pub fn new(session_id: SessionId) -> Self {
1635 Self {
1636 session_id,
1637 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1638 file_store: None,
1639 storage_store: None,
1640 image_store: None,
1641 provider_credential_store: None,
1642 utility_llm_service: None,
1643 mcp_invoker: None,
1644 egress_service: None,
1645 sqldb_store: None,
1646 message_retriever: None,
1647 session_store: None,
1648 session_mutator: None,
1649 agent_store: None,
1650 connection_resolver: None,
1651 schedule_store: None,
1652 platform_store: None,
1653 knowledge_store: None,
1654 knowledge_index_search: None,
1655 leased_resource_store: None,
1656 session_resource_registry: None,
1657 session_task_registry: None,
1658 event_emitter: None,
1659 event_context: None,
1660 tool_call_id: None,
1661 capability_registry: None,
1662 tool_registry: None,
1663 visible_tool_names: None,
1664 org_id: None,
1665 network_access: None,
1666 locale: None,
1667 budget_checker: None,
1668 payment_authority: None,
1669 subagent_spawn_store: None,
1670 subagent_nesting_policy: SubagentNestingPolicy::default(),
1671 reasoning_effort_handle: None,
1672 }
1673 }
1674
1675 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1677 Self {
1678 session_id,
1679 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1680 file_store: Some(file_store),
1681 storage_store: None,
1682 image_store: None,
1683 provider_credential_store: None,
1684 utility_llm_service: None,
1685 mcp_invoker: None,
1686 egress_service: None,
1687 sqldb_store: None,
1688 message_retriever: None,
1689 session_store: None,
1690 session_mutator: None,
1691 agent_store: None,
1692 connection_resolver: None,
1693 schedule_store: None,
1694 platform_store: None,
1695 knowledge_store: None,
1696 knowledge_index_search: None,
1697 leased_resource_store: None,
1698 session_resource_registry: None,
1699 session_task_registry: None,
1700 event_emitter: None,
1701 event_context: None,
1702 tool_call_id: None,
1703 capability_registry: None,
1704 tool_registry: None,
1705 visible_tool_names: None,
1706 org_id: None,
1707 network_access: None,
1708 locale: None,
1709 budget_checker: None,
1710 payment_authority: None,
1711 subagent_spawn_store: None,
1712 subagent_nesting_policy: SubagentNestingPolicy::default(),
1713 reasoning_effort_handle: None,
1714 }
1715 }
1716
1717 pub fn with_storage_store(
1719 session_id: SessionId,
1720 storage_store: Arc<dyn SessionStorageStore>,
1721 ) -> Self {
1722 Self {
1723 session_id,
1724 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1725 file_store: None,
1726 storage_store: Some(storage_store),
1727 image_store: None,
1728 provider_credential_store: None,
1729 utility_llm_service: None,
1730 mcp_invoker: None,
1731 egress_service: None,
1732 sqldb_store: None,
1733 message_retriever: None,
1734 session_store: None,
1735 session_mutator: None,
1736 agent_store: None,
1737 connection_resolver: None,
1738 schedule_store: None,
1739 platform_store: None,
1740 knowledge_store: None,
1741 knowledge_index_search: None,
1742 leased_resource_store: None,
1743 session_resource_registry: None,
1744 session_task_registry: None,
1745 event_emitter: None,
1746 event_context: None,
1747 tool_call_id: None,
1748 capability_registry: None,
1749 tool_registry: None,
1750 visible_tool_names: None,
1751 org_id: None,
1752 network_access: None,
1753 locale: None,
1754 budget_checker: None,
1755 payment_authority: None,
1756 subagent_spawn_store: None,
1757 subagent_nesting_policy: SubagentNestingPolicy::default(),
1758 reasoning_effort_handle: None,
1759 }
1760 }
1761
1762 pub fn with_stores(
1764 session_id: SessionId,
1765 file_store: Arc<dyn SessionFileSystem>,
1766 storage_store: Arc<dyn SessionStorageStore>,
1767 ) -> Self {
1768 Self {
1769 session_id,
1770 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1771 file_store: Some(file_store),
1772 storage_store: Some(storage_store),
1773 sqldb_store: None,
1774 image_store: None,
1775 provider_credential_store: None,
1776 utility_llm_service: None,
1777 mcp_invoker: None,
1778 egress_service: None,
1779 message_retriever: None,
1780 session_store: None,
1781 session_mutator: None,
1782 agent_store: None,
1783 connection_resolver: None,
1784 schedule_store: None,
1785 platform_store: None,
1786 knowledge_store: None,
1787 knowledge_index_search: None,
1788 leased_resource_store: None,
1789 session_resource_registry: None,
1790 session_task_registry: None,
1791 event_emitter: None,
1792 event_context: None,
1793 tool_call_id: None,
1794 capability_registry: None,
1795 tool_registry: None,
1796 visible_tool_names: None,
1797 org_id: None,
1798 network_access: None,
1799 locale: None,
1800 budget_checker: None,
1801 payment_authority: None,
1802 subagent_spawn_store: None,
1803 subagent_nesting_policy: SubagentNestingPolicy::default(),
1804 reasoning_effort_handle: None,
1805 }
1806 }
1807
1808 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1810 self.sqldb_store = Some(sqldb_store);
1811 self
1812 }
1813
1814 pub fn with_message_retriever(
1816 mut self,
1817 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1818 ) -> Self {
1819 self.message_retriever = Some(retriever);
1820 self
1821 }
1822
1823 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1825 self.session_store = Some(store);
1826 self
1827 }
1828
1829 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1831 self.session_mutator = Some(mutator);
1832 self
1833 }
1834
1835 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
1839 self.reasoning_effort_handle = Some(handle);
1840 self
1841 }
1842
1843 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1845 self.agent_store = Some(store);
1846 self
1847 }
1848
1849 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1851 self.connection_resolver = Some(resolver);
1852 self
1853 }
1854
1855 pub fn with_image_store(
1857 session_id: SessionId,
1858 image_store: Arc<dyn ImageArtifactStore>,
1859 ) -> Self {
1860 Self {
1861 session_id,
1862 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1863 file_store: None,
1864 storage_store: None,
1865 image_store: Some(image_store),
1866 provider_credential_store: None,
1867 utility_llm_service: None,
1868 mcp_invoker: None,
1869 egress_service: None,
1870 sqldb_store: None,
1871 message_retriever: None,
1872 session_store: None,
1873 session_mutator: None,
1874 agent_store: None,
1875 connection_resolver: None,
1876 schedule_store: None,
1877 platform_store: None,
1878 knowledge_store: None,
1879 knowledge_index_search: None,
1880 leased_resource_store: None,
1881 session_resource_registry: None,
1882 session_task_registry: None,
1883 event_emitter: None,
1884 event_context: None,
1885 tool_call_id: None,
1886 capability_registry: None,
1887 tool_registry: None,
1888 visible_tool_names: None,
1889 org_id: None,
1890 network_access: None,
1891 locale: None,
1892 budget_checker: None,
1893 payment_authority: None,
1894 subagent_spawn_store: None,
1895 subagent_nesting_policy: SubagentNestingPolicy::default(),
1896 reasoning_effort_handle: None,
1897 }
1898 }
1899
1900 pub fn with_provider_credential_store(
1902 mut self,
1903 store: Arc<dyn ProviderCredentialStore>,
1904 ) -> Self {
1905 self.provider_credential_store = Some(store);
1906 self
1907 }
1908
1909 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1911 self.utility_llm_service = Some(service);
1912 self
1913 }
1914
1915 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
1917 self.mcp_invoker = Some(invoker);
1918 self
1919 }
1920
1921 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1923 self.egress_service = Some(service);
1924 self
1925 }
1926
1927 pub fn with_egress_service_opt(
1930 mut self,
1931 service: Option<Arc<dyn crate::EgressService>>,
1932 ) -> Self {
1933 if let Some(service) = service {
1934 self.egress_service = Some(service);
1935 }
1936 self
1937 }
1938
1939 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
1941 self.storage_store = Some(store);
1942 self
1943 }
1944
1945 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1947 self.schedule_store = Some(store);
1948 self
1949 }
1950
1951 pub fn with_platform_store(
1953 mut self,
1954 store: Arc<dyn crate::platform_store::PlatformStore>,
1955 ) -> Self {
1956 self.platform_store = Some(store);
1957 self
1958 }
1959
1960 pub fn with_knowledge_index_search(
1962 mut self,
1963 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
1964 ) -> Self {
1965 self.knowledge_index_search = Some(search);
1966 self
1967 }
1968
1969 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1971 self.leased_resource_store = Some(store);
1972 self
1973 }
1974
1975 pub fn with_session_resource_registry(
1977 mut self,
1978 registry: Arc<dyn SessionResourceRegistry>,
1979 ) -> Self {
1980 self.session_resource_registry = Some(registry);
1981 self
1982 }
1983
1984 pub fn with_session_task_registry(
1986 mut self,
1987 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
1988 ) -> Self {
1989 self.session_task_registry = Some(registry);
1990 self
1991 }
1992
1993 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
1995 self.org_id = Some(org_id);
1996 self
1997 }
1998
1999 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
2001 self.tool_registry = Some(registry);
2002 self
2003 }
2004
2005 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
2007 self.visible_tool_names = Some(names);
2008 self
2009 }
2010
2011 pub fn with_network_access(
2013 mut self,
2014 network_access: Option<crate::network_access::NetworkAccessList>,
2015 ) -> Self {
2016 self.network_access = network_access;
2017 self
2018 }
2019
2020 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
2022 self.payment_authority = Some(authority);
2023 self
2024 }
2025
2026 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
2028 self.subagent_spawn_store = Some(store);
2029 self
2030 }
2031
2032 pub fn with_subagent_nesting_policy(mut self, policy: SubagentNestingPolicy) -> Self {
2034 self.subagent_nesting_policy = policy;
2035 self
2036 }
2037
2038 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
2043 let (Some(emitter), Some(ctx), Some(call_id)) =
2044 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2045 else {
2046 return;
2047 };
2048 if let Err(e) = emitter
2049 .emit(EventRequest::new(
2050 self.session_id,
2051 ctx.clone(),
2052 crate::events::ToolProgressData {
2053 tool_call_id: call_id.clone(),
2054 tool_name: tool_name.to_string(),
2055 message: message.to_string(),
2056 display_name: None,
2057 },
2058 ))
2059 .await
2060 {
2061 tracing::debug!(
2062 tool_call_id = call_id,
2063 tool_name,
2064 error = %e,
2065 "Failed to emit tool.progress event"
2066 );
2067 }
2068 }
2069
2070 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
2075 let (Some(emitter), Some(ctx), Some(call_id)) =
2076 (&self.event_emitter, &self.event_context, &self.tool_call_id)
2077 else {
2078 return;
2079 };
2080 if let Err(e) = emitter
2081 .emit(EventRequest::new(
2082 self.session_id,
2083 ctx.clone(),
2084 crate::events::ToolOutputDeltaData {
2085 tool_call_id: call_id.clone(),
2086 tool_name: tool_name.to_string(),
2087 delta: delta.to_string(),
2088 stream: stream.to_string(),
2089 },
2090 ))
2091 .await
2092 {
2093 tracing::debug!(
2094 tool_call_id = call_id,
2095 tool_name,
2096 error = %e,
2097 "Failed to emit tool.output.delta event"
2098 );
2099 }
2100 }
2101}
2102
2103impl std::fmt::Debug for ToolContext {
2104 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2105 f.debug_struct("ToolContext")
2106 .field("session_id", &self.session_id)
2107 .field("file_store", &self.file_store.is_some())
2108 .field("storage_store", &self.storage_store.is_some())
2109 .field("image_store", &self.image_store.is_some())
2110 .field(
2111 "provider_credential_store",
2112 &self.provider_credential_store.is_some(),
2113 )
2114 .field("utility_llm_service", &self.utility_llm_service.is_some())
2115 .field("egress_service", &self.egress_service.is_some())
2116 .field("sqldb_store", &self.sqldb_store.is_some())
2117 .field("message_retriever", &self.message_retriever.is_some())
2118 .field("session_store", &self.session_store.is_some())
2119 .field("session_mutator", &self.session_mutator.is_some())
2120 .field("agent_store", &self.agent_store.is_some())
2121 .field("connection_resolver", &self.connection_resolver.is_some())
2122 .field("schedule_store", &self.schedule_store.is_some())
2123 .field("platform_store", &self.platform_store.is_some())
2124 .field(
2125 "knowledge_index_search",
2126 &self.knowledge_index_search.is_some(),
2127 )
2128 .field(
2129 "leased_resource_store",
2130 &self.leased_resource_store.is_some(),
2131 )
2132 .field("event_emitter", &self.event_emitter.is_some())
2133 .field("tool_registry", &self.tool_registry.is_some())
2134 .field("payment_authority", &self.payment_authority.is_some())
2135 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
2136 .field("subagent_nesting_policy", &self.subagent_nesting_policy)
2137 .field("org_id", &self.org_id)
2138 .finish()
2139 }
2140}
2141
2142use crate::events::{Event, EventRequest};
2147
2148#[async_trait]
2159pub trait EventEmitter: Send + Sync {
2160 async fn emit(&self, request: EventRequest) -> Result<Event>;
2165}
2166
2167#[async_trait]
2169impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
2170 async fn emit(&self, request: EventRequest) -> Result<Event> {
2171 (**self).emit(request).await
2172 }
2173}
2174
2175#[derive(Debug, Clone, Default)]
2179pub struct NoopEventEmitter;
2180
2181#[async_trait]
2182impl EventEmitter for NoopEventEmitter {
2183 async fn emit(&self, request: EventRequest) -> Result<Event> {
2184 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2186 }
2187}
2188
2189#[derive(Debug, Clone)]
2202pub struct ResolvedImage {
2203 pub base64: String,
2205 pub media_type: String,
2207}
2208
2209impl ResolvedImage {
2210 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2212 Self {
2213 base64: base64.into(),
2214 media_type: media_type.into(),
2215 }
2216 }
2217
2218 pub fn to_data_url(&self) -> String {
2222 format!("data:{};base64,{}", self.media_type, self.base64)
2223 }
2224}
2225
2226#[async_trait]
2259pub trait ImageResolver: Send + Sync {
2260 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2264}
2265
2266#[derive(Debug)]
2272pub enum SpawnClaimResult {
2273 Claimed {
2276 spawn_handle_id: uuid::Uuid,
2277 claim_token: uuid::Uuid,
2278 },
2279 ClaimedPendingChild {
2283 spawn_handle_id: uuid::Uuid,
2284 claim_token: uuid::Uuid,
2285 },
2286 AlreadyRunning {
2289 child_session_id: crate::typed_id::SessionId,
2290 claim_token: uuid::Uuid,
2292 },
2293 AlreadySettled {
2296 child_session_id: crate::typed_id::SessionId,
2297 terminal_status: String,
2299 terminal_result: String,
2300 },
2301}
2302
2303#[async_trait]
2311pub trait SubagentSpawnStore: Send + Sync + 'static {
2312 async fn try_claim_spawn(
2317 &self,
2318 parent_session_id: crate::typed_id::SessionId,
2319 tool_call_id: &str,
2320 claim_token: uuid::Uuid,
2321 ) -> Result<SpawnClaimResult>;
2322
2323 async fn register_child_session(
2328 &self,
2329 spawn_handle_id: uuid::Uuid,
2330 claim_token: uuid::Uuid,
2331 child_session_id: crate::typed_id::SessionId,
2332 ) -> Result<()>;
2333
2334 async fn settle_spawn(
2340 &self,
2341 parent_session_id: crate::typed_id::SessionId,
2342 tool_call_id: &str,
2343 claim_token: uuid::Uuid,
2344 terminal_status: &str,
2345 terminal_result: &str,
2346 ) -> Result<()>;
2347}
2348
2349#[async_trait]
2351impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2352 async fn try_claim_spawn(
2353 &self,
2354 parent_session_id: crate::typed_id::SessionId,
2355 tool_call_id: &str,
2356 claim_token: uuid::Uuid,
2357 ) -> Result<SpawnClaimResult> {
2358 (**self)
2359 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2360 .await
2361 }
2362
2363 async fn register_child_session(
2364 &self,
2365 spawn_handle_id: uuid::Uuid,
2366 claim_token: uuid::Uuid,
2367 child_session_id: crate::typed_id::SessionId,
2368 ) -> Result<()> {
2369 (**self)
2370 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2371 .await
2372 }
2373
2374 async fn settle_spawn(
2375 &self,
2376 parent_session_id: crate::typed_id::SessionId,
2377 tool_call_id: &str,
2378 claim_token: uuid::Uuid,
2379 terminal_status: &str,
2380 terminal_result: &str,
2381 ) -> Result<()> {
2382 (**self)
2383 .settle_spawn(
2384 parent_session_id,
2385 tool_call_id,
2386 claim_token,
2387 terminal_status,
2388 terminal_result,
2389 )
2390 .await
2391 }
2392}
2393
2394pub struct NoopSubagentSpawnStore;
2398
2399#[async_trait]
2400impl SubagentSpawnStore for NoopSubagentSpawnStore {
2401 async fn try_claim_spawn(
2402 &self,
2403 _parent_session_id: crate::typed_id::SessionId,
2404 _tool_call_id: &str,
2405 claim_token: uuid::Uuid,
2406 ) -> Result<SpawnClaimResult> {
2407 Ok(SpawnClaimResult::Claimed {
2408 spawn_handle_id: uuid::Uuid::new_v4(),
2409 claim_token,
2410 })
2411 }
2412
2413 async fn register_child_session(
2414 &self,
2415 _spawn_handle_id: uuid::Uuid,
2416 _claim_token: uuid::Uuid,
2417 _child_session_id: crate::typed_id::SessionId,
2418 ) -> Result<()> {
2419 Ok(())
2420 }
2421
2422 async fn settle_spawn(
2423 &self,
2424 _parent_session_id: crate::typed_id::SessionId,
2425 _tool_call_id: &str,
2426 _claim_token: uuid::Uuid,
2427 _terminal_status: &str,
2428 _terminal_result: &str,
2429 ) -> Result<()> {
2430 Ok(())
2431 }
2432}
2433
2434#[cfg(test)]
2439mod tests {
2440 use super::*;
2441
2442 #[test]
2443 fn test_resolved_image_new() {
2444 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2445 assert_eq!(image.base64, "SGVsbG8=");
2446 assert_eq!(image.media_type, "image/png");
2447 }
2448
2449 #[test]
2450 fn test_resolved_image_to_data_url() {
2451 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2452 let data_url = image.to_data_url();
2453 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2454 }
2455
2456 #[test]
2457 fn test_resolved_image_jpeg() {
2458 let image = ResolvedImage::new("base64data", "image/jpeg");
2459 let data_url = image.to_data_url();
2460 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2461 }
2462}