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