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 type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
1090
1091#[async_trait]
1096pub trait UserConnectionResolver: Send + Sync {
1097 async fn get_connection_token(
1100 &self,
1101 session_id: SessionId,
1102 provider: &str,
1103 ) -> Result<Option<String>>;
1104
1105 async fn get_connection_user(
1110 &self,
1111 _session_id: SessionId,
1112 _provider: &str,
1113 ) -> Result<Option<Uuid>> {
1114 Ok(None)
1115 }
1116
1117 async fn get_connection_token_for_user(
1122 &self,
1123 _user_id: Uuid,
1124 _provider: &str,
1125 ) -> Result<Option<String>> {
1126 Ok(None)
1127 }
1128
1129 async fn get_connection_metadata(
1132 &self,
1133 _session_id: SessionId,
1134 _provider: &str,
1135 ) -> Result<Option<serde_json::Value>> {
1136 Ok(None)
1137 }
1138}
1139
1140#[async_trait]
1150pub trait BudgetChecker: Send + Sync {
1151 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
1153}
1154
1155#[async_trait]
1164pub trait PaymentAuthority: Send + Sync {
1165 async fn execute_machine_payment(
1166 &self,
1167 session_id: SessionId,
1168 request: crate::payment::MachinePaymentRequest,
1169 ) -> Result<crate::payment::MachinePaymentResponse>;
1170}
1171
1172#[async_trait]
1182pub trait OutboundToolRateLimiter: Send + Sync {
1183 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
1185}
1186
1187#[derive(Debug)]
1193pub enum ToolCallClaimResult {
1194 Claimed { claim_token: uuid::Uuid },
1197 AlreadySettled {
1199 result_json: serde_json::Value,
1200 args_fingerprint: String,
1201 },
1202 AlreadyRunning { args_fingerprint: String },
1207 DeterminismViolation {
1211 stored_fingerprint: String,
1212 current_fingerprint: String,
1213 },
1214}
1215
1216#[derive(Debug, Clone)]
1218pub enum DurableToolCallStatus {
1219 Settled { result_json: serde_json::Value },
1221 Interrupted {
1223 result_json: Option<serde_json::Value>,
1224 },
1225 Running,
1227}
1228
1229#[async_trait]
1234pub trait DurableToolResultStore: Send + Sync + 'static {
1235 async fn try_claim_tool_call(
1243 &self,
1244 turn_id: &str,
1245 tool_call_id: &str,
1246 tool_name: &str,
1247 args_fingerprint: &str,
1248 ) -> Result<ToolCallClaimResult>;
1249
1250 async fn settle_tool_call(
1256 &self,
1257 turn_id: &str,
1258 tool_call_id: &str,
1259 result_json: serde_json::Value,
1260 status: &str,
1261 claim_token: uuid::Uuid,
1262 ) -> Result<bool>;
1263
1264 async fn get_tool_call_status(
1269 &self,
1270 turn_id: &str,
1271 tool_call_id: &str,
1272 ) -> Result<Option<DurableToolCallStatus>>;
1273}
1274
1275pub struct NoopDurableToolResultStore;
1278
1279#[async_trait]
1280impl DurableToolResultStore for NoopDurableToolResultStore {
1281 async fn try_claim_tool_call(
1282 &self,
1283 _turn_id: &str,
1284 _tool_call_id: &str,
1285 _tool_name: &str,
1286 _args_fingerprint: &str,
1287 ) -> Result<ToolCallClaimResult> {
1288 Ok(ToolCallClaimResult::Claimed {
1289 claim_token: uuid::Uuid::new_v4(),
1290 })
1291 }
1292
1293 async fn settle_tool_call(
1294 &self,
1295 _turn_id: &str,
1296 _tool_call_id: &str,
1297 _result_json: serde_json::Value,
1298 _status: &str,
1299 _claim_token: uuid::Uuid,
1300 ) -> Result<bool> {
1301 Ok(true)
1302 }
1303
1304 async fn get_tool_call_status(
1305 &self,
1306 _turn_id: &str,
1307 _tool_call_id: &str,
1308 ) -> Result<Option<DurableToolCallStatus>> {
1309 Ok(None)
1310 }
1311}
1312
1313#[derive(Debug, Clone)]
1319pub struct StreamProgress {
1320 pub accumulated_len: usize,
1322 pub last_delta_at: u64,
1324}
1325
1326#[async_trait]
1332pub trait StreamHeartbeater: Send + Sync {
1333 async fn heartbeat(&self, progress: StreamProgress);
1339}
1340
1341pub struct NoopStreamHeartbeater;
1343
1344#[async_trait]
1345impl StreamHeartbeater for NoopStreamHeartbeater {
1346 async fn heartbeat(&self, _progress: StreamProgress) {}
1347}
1348
1349#[derive(Debug, Clone)]
1355pub struct PartialStreamState {
1356 pub accumulated: String,
1359}
1360
1361#[async_trait]
1369pub trait PartialStreamStore: Send + Sync {
1370 async fn get_partial_stream(
1373 &self,
1374 session_id: SessionId,
1375 turn_id: &str,
1376 ) -> Result<Option<PartialStreamState>>;
1377}
1378
1379pub struct NoopPartialStreamStore;
1381
1382#[async_trait]
1383impl PartialStreamStore for NoopPartialStreamStore {
1384 async fn get_partial_stream(
1385 &self,
1386 _session_id: SessionId,
1387 _turn_id: &str,
1388 ) -> Result<Option<PartialStreamState>> {
1389 Ok(None)
1390 }
1391}
1392
1393#[derive(Clone)]
1402pub struct ToolContext {
1403 pub session_id: SessionId,
1405 pub workspace_id: WorkspaceId,
1412
1413 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1415
1416 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1418
1419 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1421
1422 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1424
1425 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1427
1428 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1434
1435 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1437
1438 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1440
1441 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1443
1444 pub session_store: Option<Arc<dyn SessionStore>>,
1446
1447 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1449
1450 pub agent_store: Option<Arc<dyn AgentStore>>,
1452
1453 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1455
1456 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1458
1459 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1461 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1463
1464 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1468
1469 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1471
1472 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1474
1475 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1478
1479 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1482
1483 pub event_context: Option<crate::events::EventContext>,
1486
1487 pub tool_call_id: Option<String>,
1490 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1492
1493 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1496
1497 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1501
1502 pub org_id: Option<crate::typed_id::OrgId>,
1504
1505 pub network_access: Option<crate::network_access::NetworkAccessList>,
1508
1509 pub locale: Option<String>,
1513
1514 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1516
1517 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1519
1520 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1524
1525 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1529}
1530
1531impl ToolContext {
1532 pub fn workspace_fs_key(&self) -> SessionId {
1537 SessionId::from_uuid(self.workspace_id.uuid())
1538 }
1539
1540 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1542 self.workspace_id = workspace_id;
1543 self
1544 }
1545
1546 pub fn new(session_id: SessionId) -> Self {
1548 Self {
1549 session_id,
1550 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1551 file_store: None,
1552 storage_store: None,
1553 image_store: None,
1554 provider_credential_store: None,
1555 utility_llm_service: None,
1556 mcp_invoker: None,
1557 egress_service: None,
1558 sqldb_store: None,
1559 message_retriever: None,
1560 session_store: None,
1561 session_mutator: None,
1562 agent_store: None,
1563 connection_resolver: None,
1564 schedule_store: None,
1565 platform_store: None,
1566 knowledge_store: None,
1567 knowledge_index_search: None,
1568 leased_resource_store: None,
1569 session_resource_registry: None,
1570 session_task_registry: None,
1571 event_emitter: None,
1572 event_context: None,
1573 tool_call_id: None,
1574 capability_registry: None,
1575 tool_registry: None,
1576 visible_tool_names: None,
1577 org_id: None,
1578 network_access: None,
1579 locale: None,
1580 budget_checker: None,
1581 payment_authority: None,
1582 subagent_spawn_store: None,
1583 reasoning_effort_handle: None,
1584 }
1585 }
1586
1587 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1589 Self {
1590 session_id,
1591 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1592 file_store: Some(file_store),
1593 storage_store: None,
1594 image_store: None,
1595 provider_credential_store: None,
1596 utility_llm_service: None,
1597 mcp_invoker: None,
1598 egress_service: None,
1599 sqldb_store: None,
1600 message_retriever: None,
1601 session_store: None,
1602 session_mutator: None,
1603 agent_store: None,
1604 connection_resolver: None,
1605 schedule_store: None,
1606 platform_store: None,
1607 knowledge_store: None,
1608 knowledge_index_search: None,
1609 leased_resource_store: None,
1610 session_resource_registry: None,
1611 session_task_registry: None,
1612 event_emitter: None,
1613 event_context: None,
1614 tool_call_id: None,
1615 capability_registry: None,
1616 tool_registry: None,
1617 visible_tool_names: None,
1618 org_id: None,
1619 network_access: None,
1620 locale: None,
1621 budget_checker: None,
1622 payment_authority: None,
1623 subagent_spawn_store: None,
1624 reasoning_effort_handle: None,
1625 }
1626 }
1627
1628 pub fn with_storage_store(
1630 session_id: SessionId,
1631 storage_store: Arc<dyn SessionStorageStore>,
1632 ) -> Self {
1633 Self {
1634 session_id,
1635 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1636 file_store: None,
1637 storage_store: Some(storage_store),
1638 image_store: None,
1639 provider_credential_store: None,
1640 utility_llm_service: None,
1641 mcp_invoker: None,
1642 egress_service: None,
1643 sqldb_store: None,
1644 message_retriever: None,
1645 session_store: None,
1646 session_mutator: None,
1647 agent_store: None,
1648 connection_resolver: None,
1649 schedule_store: None,
1650 platform_store: None,
1651 knowledge_store: None,
1652 knowledge_index_search: None,
1653 leased_resource_store: None,
1654 session_resource_registry: None,
1655 session_task_registry: None,
1656 event_emitter: None,
1657 event_context: None,
1658 tool_call_id: None,
1659 capability_registry: None,
1660 tool_registry: None,
1661 visible_tool_names: None,
1662 org_id: None,
1663 network_access: None,
1664 locale: None,
1665 budget_checker: None,
1666 payment_authority: None,
1667 subagent_spawn_store: None,
1668 reasoning_effort_handle: None,
1669 }
1670 }
1671
1672 pub fn with_stores(
1674 session_id: SessionId,
1675 file_store: Arc<dyn SessionFileSystem>,
1676 storage_store: Arc<dyn SessionStorageStore>,
1677 ) -> Self {
1678 Self {
1679 session_id,
1680 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1681 file_store: Some(file_store),
1682 storage_store: Some(storage_store),
1683 sqldb_store: None,
1684 image_store: None,
1685 provider_credential_store: None,
1686 utility_llm_service: None,
1687 mcp_invoker: None,
1688 egress_service: None,
1689 message_retriever: None,
1690 session_store: None,
1691 session_mutator: None,
1692 agent_store: None,
1693 connection_resolver: None,
1694 schedule_store: None,
1695 platform_store: None,
1696 knowledge_store: None,
1697 knowledge_index_search: None,
1698 leased_resource_store: None,
1699 session_resource_registry: None,
1700 session_task_registry: None,
1701 event_emitter: None,
1702 event_context: None,
1703 tool_call_id: None,
1704 capability_registry: None,
1705 tool_registry: None,
1706 visible_tool_names: None,
1707 org_id: None,
1708 network_access: None,
1709 locale: None,
1710 budget_checker: None,
1711 payment_authority: None,
1712 subagent_spawn_store: None,
1713 reasoning_effort_handle: None,
1714 }
1715 }
1716
1717 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1719 self.sqldb_store = Some(sqldb_store);
1720 self
1721 }
1722
1723 pub fn with_message_retriever(
1725 mut self,
1726 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1727 ) -> Self {
1728 self.message_retriever = Some(retriever);
1729 self
1730 }
1731
1732 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1734 self.session_store = Some(store);
1735 self
1736 }
1737
1738 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1740 self.session_mutator = Some(mutator);
1741 self
1742 }
1743
1744 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
1748 self.reasoning_effort_handle = Some(handle);
1749 self
1750 }
1751
1752 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1754 self.agent_store = Some(store);
1755 self
1756 }
1757
1758 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1760 self.connection_resolver = Some(resolver);
1761 self
1762 }
1763
1764 pub fn with_image_store(
1766 session_id: SessionId,
1767 image_store: Arc<dyn ImageArtifactStore>,
1768 ) -> Self {
1769 Self {
1770 session_id,
1771 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1772 file_store: None,
1773 storage_store: None,
1774 image_store: Some(image_store),
1775 provider_credential_store: None,
1776 utility_llm_service: None,
1777 mcp_invoker: None,
1778 egress_service: None,
1779 sqldb_store: None,
1780 message_retriever: None,
1781 session_store: None,
1782 session_mutator: None,
1783 agent_store: None,
1784 connection_resolver: None,
1785 schedule_store: None,
1786 platform_store: None,
1787 knowledge_store: None,
1788 knowledge_index_search: None,
1789 leased_resource_store: None,
1790 session_resource_registry: None,
1791 session_task_registry: None,
1792 event_emitter: None,
1793 event_context: None,
1794 tool_call_id: None,
1795 capability_registry: None,
1796 tool_registry: None,
1797 visible_tool_names: None,
1798 org_id: None,
1799 network_access: None,
1800 locale: None,
1801 budget_checker: None,
1802 payment_authority: None,
1803 subagent_spawn_store: None,
1804 reasoning_effort_handle: None,
1805 }
1806 }
1807
1808 pub fn with_provider_credential_store(
1810 mut self,
1811 store: Arc<dyn ProviderCredentialStore>,
1812 ) -> Self {
1813 self.provider_credential_store = Some(store);
1814 self
1815 }
1816
1817 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1819 self.utility_llm_service = Some(service);
1820 self
1821 }
1822
1823 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
1825 self.mcp_invoker = Some(invoker);
1826 self
1827 }
1828
1829 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1831 self.egress_service = Some(service);
1832 self
1833 }
1834
1835 pub fn with_egress_service_opt(
1838 mut self,
1839 service: Option<Arc<dyn crate::EgressService>>,
1840 ) -> Self {
1841 if let Some(service) = service {
1842 self.egress_service = Some(service);
1843 }
1844 self
1845 }
1846
1847 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
1849 self.storage_store = Some(store);
1850 self
1851 }
1852
1853 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1855 self.schedule_store = Some(store);
1856 self
1857 }
1858
1859 pub fn with_platform_store(
1861 mut self,
1862 store: Arc<dyn crate::platform_store::PlatformStore>,
1863 ) -> Self {
1864 self.platform_store = Some(store);
1865 self
1866 }
1867
1868 pub fn with_knowledge_index_search(
1870 mut self,
1871 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
1872 ) -> Self {
1873 self.knowledge_index_search = Some(search);
1874 self
1875 }
1876
1877 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1879 self.leased_resource_store = Some(store);
1880 self
1881 }
1882
1883 pub fn with_session_resource_registry(
1885 mut self,
1886 registry: Arc<dyn SessionResourceRegistry>,
1887 ) -> Self {
1888 self.session_resource_registry = Some(registry);
1889 self
1890 }
1891
1892 pub fn with_session_task_registry(
1894 mut self,
1895 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
1896 ) -> Self {
1897 self.session_task_registry = Some(registry);
1898 self
1899 }
1900
1901 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
1903 self.org_id = Some(org_id);
1904 self
1905 }
1906
1907 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
1909 self.tool_registry = Some(registry);
1910 self
1911 }
1912
1913 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
1915 self.visible_tool_names = Some(names);
1916 self
1917 }
1918
1919 pub fn with_network_access(
1921 mut self,
1922 network_access: Option<crate::network_access::NetworkAccessList>,
1923 ) -> Self {
1924 self.network_access = network_access;
1925 self
1926 }
1927
1928 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
1930 self.payment_authority = Some(authority);
1931 self
1932 }
1933
1934 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
1936 self.subagent_spawn_store = Some(store);
1937 self
1938 }
1939
1940 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
1945 let (Some(emitter), Some(ctx), Some(call_id)) =
1946 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1947 else {
1948 return;
1949 };
1950 if let Err(e) = emitter
1951 .emit(EventRequest::new(
1952 self.session_id,
1953 ctx.clone(),
1954 crate::events::ToolProgressData {
1955 tool_call_id: call_id.clone(),
1956 tool_name: tool_name.to_string(),
1957 message: message.to_string(),
1958 display_name: None,
1959 },
1960 ))
1961 .await
1962 {
1963 tracing::debug!(
1964 tool_call_id = call_id,
1965 tool_name,
1966 error = %e,
1967 "Failed to emit tool.progress event"
1968 );
1969 }
1970 }
1971
1972 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
1977 let (Some(emitter), Some(ctx), Some(call_id)) =
1978 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1979 else {
1980 return;
1981 };
1982 if let Err(e) = emitter
1983 .emit(EventRequest::new(
1984 self.session_id,
1985 ctx.clone(),
1986 crate::events::ToolOutputDeltaData {
1987 tool_call_id: call_id.clone(),
1988 tool_name: tool_name.to_string(),
1989 delta: delta.to_string(),
1990 stream: stream.to_string(),
1991 },
1992 ))
1993 .await
1994 {
1995 tracing::debug!(
1996 tool_call_id = call_id,
1997 tool_name,
1998 error = %e,
1999 "Failed to emit tool.output.delta event"
2000 );
2001 }
2002 }
2003}
2004
2005impl std::fmt::Debug for ToolContext {
2006 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2007 f.debug_struct("ToolContext")
2008 .field("session_id", &self.session_id)
2009 .field("file_store", &self.file_store.is_some())
2010 .field("storage_store", &self.storage_store.is_some())
2011 .field("image_store", &self.image_store.is_some())
2012 .field(
2013 "provider_credential_store",
2014 &self.provider_credential_store.is_some(),
2015 )
2016 .field("utility_llm_service", &self.utility_llm_service.is_some())
2017 .field("egress_service", &self.egress_service.is_some())
2018 .field("sqldb_store", &self.sqldb_store.is_some())
2019 .field("message_retriever", &self.message_retriever.is_some())
2020 .field("session_store", &self.session_store.is_some())
2021 .field("session_mutator", &self.session_mutator.is_some())
2022 .field("agent_store", &self.agent_store.is_some())
2023 .field("connection_resolver", &self.connection_resolver.is_some())
2024 .field("schedule_store", &self.schedule_store.is_some())
2025 .field("platform_store", &self.platform_store.is_some())
2026 .field(
2027 "knowledge_index_search",
2028 &self.knowledge_index_search.is_some(),
2029 )
2030 .field(
2031 "leased_resource_store",
2032 &self.leased_resource_store.is_some(),
2033 )
2034 .field("event_emitter", &self.event_emitter.is_some())
2035 .field("tool_registry", &self.tool_registry.is_some())
2036 .field("payment_authority", &self.payment_authority.is_some())
2037 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
2038 .field("org_id", &self.org_id)
2039 .finish()
2040 }
2041}
2042
2043use crate::events::{Event, EventRequest};
2048
2049#[async_trait]
2060pub trait EventEmitter: Send + Sync {
2061 async fn emit(&self, request: EventRequest) -> Result<Event>;
2066}
2067
2068#[async_trait]
2070impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
2071 async fn emit(&self, request: EventRequest) -> Result<Event> {
2072 (**self).emit(request).await
2073 }
2074}
2075
2076#[derive(Debug, Clone, Default)]
2080pub struct NoopEventEmitter;
2081
2082#[async_trait]
2083impl EventEmitter for NoopEventEmitter {
2084 async fn emit(&self, request: EventRequest) -> Result<Event> {
2085 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2087 }
2088}
2089
2090#[derive(Debug, Clone)]
2103pub struct ResolvedImage {
2104 pub base64: String,
2106 pub media_type: String,
2108}
2109
2110impl ResolvedImage {
2111 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2113 Self {
2114 base64: base64.into(),
2115 media_type: media_type.into(),
2116 }
2117 }
2118
2119 pub fn to_data_url(&self) -> String {
2123 format!("data:{};base64,{}", self.media_type, self.base64)
2124 }
2125}
2126
2127#[async_trait]
2160pub trait ImageResolver: Send + Sync {
2161 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2165}
2166
2167#[derive(Debug)]
2173pub enum SpawnClaimResult {
2174 Claimed {
2177 spawn_handle_id: uuid::Uuid,
2178 claim_token: uuid::Uuid,
2179 },
2180 ClaimedPendingChild {
2184 spawn_handle_id: uuid::Uuid,
2185 claim_token: uuid::Uuid,
2186 },
2187 AlreadyRunning {
2190 child_session_id: crate::typed_id::SessionId,
2191 claim_token: uuid::Uuid,
2193 },
2194 AlreadySettled {
2197 child_session_id: crate::typed_id::SessionId,
2198 terminal_status: String,
2200 terminal_result: String,
2201 },
2202}
2203
2204#[async_trait]
2212pub trait SubagentSpawnStore: Send + Sync + 'static {
2213 async fn try_claim_spawn(
2218 &self,
2219 parent_session_id: crate::typed_id::SessionId,
2220 tool_call_id: &str,
2221 claim_token: uuid::Uuid,
2222 ) -> Result<SpawnClaimResult>;
2223
2224 async fn register_child_session(
2229 &self,
2230 spawn_handle_id: uuid::Uuid,
2231 claim_token: uuid::Uuid,
2232 child_session_id: crate::typed_id::SessionId,
2233 ) -> Result<()>;
2234
2235 async fn settle_spawn(
2241 &self,
2242 parent_session_id: crate::typed_id::SessionId,
2243 tool_call_id: &str,
2244 claim_token: uuid::Uuid,
2245 terminal_status: &str,
2246 terminal_result: &str,
2247 ) -> Result<()>;
2248}
2249
2250#[async_trait]
2252impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2253 async fn try_claim_spawn(
2254 &self,
2255 parent_session_id: crate::typed_id::SessionId,
2256 tool_call_id: &str,
2257 claim_token: uuid::Uuid,
2258 ) -> Result<SpawnClaimResult> {
2259 (**self)
2260 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2261 .await
2262 }
2263
2264 async fn register_child_session(
2265 &self,
2266 spawn_handle_id: uuid::Uuid,
2267 claim_token: uuid::Uuid,
2268 child_session_id: crate::typed_id::SessionId,
2269 ) -> Result<()> {
2270 (**self)
2271 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2272 .await
2273 }
2274
2275 async fn settle_spawn(
2276 &self,
2277 parent_session_id: crate::typed_id::SessionId,
2278 tool_call_id: &str,
2279 claim_token: uuid::Uuid,
2280 terminal_status: &str,
2281 terminal_result: &str,
2282 ) -> Result<()> {
2283 (**self)
2284 .settle_spawn(
2285 parent_session_id,
2286 tool_call_id,
2287 claim_token,
2288 terminal_status,
2289 terminal_result,
2290 )
2291 .await
2292 }
2293}
2294
2295pub struct NoopSubagentSpawnStore;
2299
2300#[async_trait]
2301impl SubagentSpawnStore for NoopSubagentSpawnStore {
2302 async fn try_claim_spawn(
2303 &self,
2304 _parent_session_id: crate::typed_id::SessionId,
2305 _tool_call_id: &str,
2306 claim_token: uuid::Uuid,
2307 ) -> Result<SpawnClaimResult> {
2308 Ok(SpawnClaimResult::Claimed {
2309 spawn_handle_id: uuid::Uuid::new_v4(),
2310 claim_token,
2311 })
2312 }
2313
2314 async fn register_child_session(
2315 &self,
2316 _spawn_handle_id: uuid::Uuid,
2317 _claim_token: uuid::Uuid,
2318 _child_session_id: crate::typed_id::SessionId,
2319 ) -> Result<()> {
2320 Ok(())
2321 }
2322
2323 async fn settle_spawn(
2324 &self,
2325 _parent_session_id: crate::typed_id::SessionId,
2326 _tool_call_id: &str,
2327 _claim_token: uuid::Uuid,
2328 _terminal_status: &str,
2329 _terminal_result: &str,
2330 ) -> Result<()> {
2331 Ok(())
2332 }
2333}
2334
2335#[cfg(test)]
2340mod tests {
2341 use super::*;
2342
2343 #[test]
2344 fn test_resolved_image_new() {
2345 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2346 assert_eq!(image.base64, "SGVsbG8=");
2347 assert_eq!(image.media_type, "image/png");
2348 }
2349
2350 #[test]
2351 fn test_resolved_image_to_data_url() {
2352 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2353 let data_url = image.to_data_url();
2354 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2355 }
2356
2357 #[test]
2358 fn test_resolved_image_jpeg() {
2359 let image = ResolvedImage::new("base64data", "image/jpeg");
2360 let data_url = image.to_data_url();
2361 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2362 }
2363}