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 workspace_display_path(path: &str) -> String {
22 if path == "/" {
23 "/workspace".to_string()
24 } else if path.starts_with('/') {
25 format!("/workspace{path}")
26 } else {
27 format!("/workspace/{path}")
28 }
29}
30
31fn build_tool_map(tool_defs: &[ToolDefinition]) -> HashMap<&str, &ToolDefinition> {
33 tool_defs.iter().map(|def| (def.name(), def)).collect()
34}
35
36use crate::error::Result;
37
38#[derive(Clone, Default)]
53pub struct ReasoningEffortHandle {
54 inner: Arc<std::sync::RwLock<Option<String>>>,
55}
56
57impl ReasoningEffortHandle {
58 pub fn new() -> Self {
60 Self::default()
61 }
62
63 pub fn with_effort(effort: impl Into<String>) -> Self {
65 Self {
66 inner: Arc::new(std::sync::RwLock::new(Some(effort.into()))),
67 }
68 }
69
70 pub fn set(&self, effort: Option<String>) {
74 let mut guard = self.inner.write().unwrap_or_else(|e| e.into_inner());
78 *guard = effort;
79 }
80
81 pub fn get(&self) -> Option<String> {
83 let guard = self.inner.read().unwrap_or_else(|e| e.into_inner());
86 guard.clone()
87 }
88}
89
90impl std::fmt::Debug for ReasoningEffortHandle {
91 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92 f.debug_struct("ReasoningEffortHandle")
93 .field("effort", &self.get())
94 .finish()
95 }
96}
97
98#[async_trait]
109pub trait AgentStore: Send + Sync {
110 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>>;
112}
113
114#[async_trait]
115impl<T: AgentStore + ?Sized> AgentStore for std::sync::Arc<T> {
116 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>> {
117 (**self).get_agent(agent_id).await
118 }
119}
120
121#[async_trait]
136pub trait HarnessStore: Send + Sync {
137 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>>;
142}
143
144#[async_trait]
145impl<T: HarnessStore + ?Sized> HarnessStore for std::sync::Arc<T> {
146 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>> {
147 (**self).get_harness_chain(harness_id).await
148 }
149}
150
151use crate::leased_resource::{LeasedResource, UpsertLeasedResource};
156use crate::session::Session;
157
158#[async_trait]
164pub trait SessionStore: Send + Sync {
165 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>>;
167}
168
169#[async_trait]
170impl<T: SessionStore + ?Sized> SessionStore for std::sync::Arc<T> {
171 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>> {
172 (**self).get_session(session_id).await
173 }
174}
175
176#[async_trait]
178pub trait SessionMutator: Send + Sync {
179 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session>;
181}
182
183#[async_trait]
184impl<T: SessionMutator + ?Sized> SessionMutator for std::sync::Arc<T> {
185 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session> {
186 (**self).update_session_title(session_id, title).await
187 }
188}
189
190#[derive(Debug, Clone)]
196pub struct ResolvedModel {
197 pub model: String,
199 pub provider_type: DriverId,
201 pub api_key: Option<String>,
203 pub base_url: Option<String>,
205 pub provider_metadata: Option<crate::driver_registry::ProviderMetadata>,
208}
209
210#[async_trait]
220pub trait ProviderStore: Send + Sync {
221 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>>;
226
227 async fn get_default_model(&self) -> Result<Option<ResolvedModel>>;
231}
232
233#[async_trait]
234impl<T: ProviderStore + ?Sized> ProviderStore for std::sync::Arc<T> {
235 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>> {
236 (**self).get_resolved_model(model_id).await
237 }
238
239 async fn get_default_model(&self) -> Result<Option<ResolvedModel>> {
240 (**self).get_default_model().await
241 }
242}
243
244#[derive(Debug, Clone)]
250pub struct StoredImageInfo {
251 pub id: ImageId,
252 pub filename: String,
253 pub content_type: String,
254 pub size_bytes: i64,
255 pub metadata: serde_json::Value,
256 pub created_at: DateTime<Utc>,
257}
258
259#[derive(Debug, Clone)]
261pub struct StoredImage {
262 pub info: StoredImageInfo,
263 pub data: Vec<u8>,
264}
265
266#[derive(Debug, Clone)]
268pub struct CreateStoredImage {
269 pub filename: String,
270 pub content_type: String,
271 pub data: Vec<u8>,
272 pub metadata: serde_json::Value,
273}
274
275#[async_trait]
276pub trait ImageArtifactStore: Send + Sync {
277 async fn create_image(&self, input: CreateStoredImage) -> Result<StoredImageInfo>;
279
280 async fn get_image(&self, image_id: ImageId) -> Result<Option<StoredImage>>;
282
283 async fn get_image_info(&self, image_id: ImageId) -> Result<Option<StoredImageInfo>>;
285}
286
287#[derive(Debug, Clone)]
293pub struct ProviderCredentials {
294 pub api_key: String,
295 pub base_url: Option<String>,
296}
297
298#[async_trait]
299pub trait ProviderCredentialStore: Send + Sync {
300 async fn get_default_provider_credentials(
305 &self,
306 provider_type: &str,
307 ) -> Result<Option<ProviderCredentials>>;
308}
309
310#[async_trait]
321pub trait ToolExecutor: Send + Sync {
322 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult>;
327
328 async fn execute_with_context(
333 &self,
334 tool_call: &ToolCall,
335 tool_def: &ToolDefinition,
336 _context: &ToolContext,
337 ) -> Result<ToolResult> {
338 self.execute(tool_call, tool_def).await
340 }
341
342 async fn execute_batch(
344 &self,
345 tool_calls: &[ToolCall],
346 tool_defs: &[ToolDefinition],
347 ) -> Result<Vec<ToolResult>> {
348 let mut results = Vec::with_capacity(tool_calls.len());
349
350 let tool_map = build_tool_map(tool_defs);
351
352 for tool_call in tool_calls {
353 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
354 crate::error::AgentLoopError::tool(format!(
355 "Tool definition not found: {}",
356 tool_call.name
357 ))
358 })?;
359
360 results.push(self.execute(tool_call, tool_def).await?);
361 }
362
363 Ok(results)
364 }
365
366 async fn execute_parallel(
368 &self,
369 tool_calls: &[ToolCall],
370 tool_defs: &[ToolDefinition],
371 ) -> Result<Vec<ToolResult>>
372 where
373 Self: Sized,
374 {
375 use futures::future::join_all;
376
377 let tool_map = build_tool_map(tool_defs);
378
379 let futures: Vec<_> = tool_calls
380 .iter()
381 .map(|tool_call| async {
382 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
383 crate::error::AgentLoopError::tool(format!(
384 "Tool definition not found: {}",
385 tool_call.name
386 ))
387 })?;
388 self.execute(tool_call, tool_def).await
389 })
390 .collect();
391
392 let results = join_all(futures).await;
393 results.into_iter().collect()
394 }
395}
396
397#[async_trait]
401impl ToolExecutor for std::sync::Arc<dyn ToolExecutor> {
402 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult> {
403 (**self).execute(tool_call, tool_def).await
404 }
405
406 async fn execute_with_context(
407 &self,
408 tool_call: &ToolCall,
409 tool_def: &ToolDefinition,
410 context: &ToolContext,
411 ) -> Result<ToolResult> {
412 (**self)
413 .execute_with_context(tool_call, tool_def, context)
414 .await
415 }
416
417 async fn execute_batch(
418 &self,
419 tool_calls: &[ToolCall],
420 tool_defs: &[ToolDefinition],
421 ) -> Result<Vec<ToolResult>> {
422 (**self).execute_batch(tool_calls, tool_defs).await
423 }
424}
425
426#[async_trait]
438pub trait SessionFileSystem: Send + Sync {
439 fn display_root(&self) -> String {
445 "/workspace".to_string()
446 }
447
448 fn display_path(&self, path: &str) -> String {
450 workspace_display_path(path)
451 }
452
453 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>>;
455
456 async fn write_file(
458 &self,
459 session_id: SessionId,
460 path: &str,
461 content: &str,
462 encoding: &str,
463 ) -> Result<SessionFile>;
464
465 async fn write_file_if_content_matches(
470 &self,
471 session_id: SessionId,
472 path: &str,
473 expected_content: &str,
474 expected_encoding: &str,
475 content: &str,
476 encoding: &str,
477 ) -> Result<Option<SessionFile>> {
478 let Some(existing) = self.read_file(session_id, path).await? else {
479 return Ok(None);
480 };
481
482 if existing.is_directory {
483 return Ok(None);
484 }
485
486 let current_content = existing.content.unwrap_or_default();
487 if current_content != expected_content || existing.encoding != expected_encoding {
488 return Ok(None);
489 }
490
491 self.write_file(session_id, path, content, encoding)
492 .await
493 .map(Some)
494 }
495
496 async fn delete_file(&self, session_id: SessionId, path: &str, recursive: bool)
498 -> Result<bool>;
499
500 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>>;
502
503 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>>;
505
506 async fn grep_files(
508 &self,
509 session_id: SessionId,
510 pattern: &str,
511 path_pattern: Option<&str>,
512 ) -> Result<Vec<GrepMatch>>;
513
514 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
516
517 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
519 if file.is_readonly {
520 return Err(crate::error::AgentLoopError::store(
521 "read-only initial files require a SessionFileSystem-specific seed implementation",
522 ));
523 }
524 self.write_file(session_id, &file.path, &file.content, &file.encoding)
525 .await?;
526 Ok(())
527 }
528}
529
530pub struct WorkspaceScopedFileSystem {
540 inner: Arc<dyn SessionFileSystem>,
541 key: SessionId,
542}
543
544impl WorkspaceScopedFileSystem {
545 pub fn wrap(
547 inner: Arc<dyn SessionFileSystem>,
548 workspace_id: WorkspaceId,
549 ) -> Arc<dyn SessionFileSystem> {
550 Arc::new(Self {
551 inner,
552 key: SessionId::from_uuid(workspace_id.uuid()),
553 })
554 }
555}
556
557#[async_trait]
558impl SessionFileSystem for WorkspaceScopedFileSystem {
559 async fn read_file(&self, _session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
560 self.inner.read_file(self.key, path).await
561 }
562 async fn write_file(
563 &self,
564 _session_id: SessionId,
565 path: &str,
566 content: &str,
567 encoding: &str,
568 ) -> Result<SessionFile> {
569 self.inner
570 .write_file(self.key, path, content, encoding)
571 .await
572 }
573 async fn write_file_if_content_matches(
574 &self,
575 _session_id: SessionId,
576 path: &str,
577 expected_content: &str,
578 expected_encoding: &str,
579 content: &str,
580 encoding: &str,
581 ) -> Result<Option<SessionFile>> {
582 self.inner
583 .write_file_if_content_matches(
584 self.key,
585 path,
586 expected_content,
587 expected_encoding,
588 content,
589 encoding,
590 )
591 .await
592 }
593 async fn delete_file(
594 &self,
595 _session_id: SessionId,
596 path: &str,
597 recursive: bool,
598 ) -> Result<bool> {
599 self.inner.delete_file(self.key, path, recursive).await
600 }
601 async fn list_directory(&self, _session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
602 self.inner.list_directory(self.key, path).await
603 }
604 async fn stat_file(&self, _session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
605 self.inner.stat_file(self.key, path).await
606 }
607 async fn grep_files(
608 &self,
609 _session_id: SessionId,
610 pattern: &str,
611 path_pattern: Option<&str>,
612 ) -> Result<Vec<GrepMatch>> {
613 self.inner.grep_files(self.key, pattern, path_pattern).await
614 }
615 async fn create_directory(&self, _session_id: SessionId, path: &str) -> Result<FileInfo> {
616 self.inner.create_directory(self.key, path).await
617 }
618 async fn seed_initial_file(&self, _session_id: SessionId, file: &InitialFile) -> Result<()> {
619 self.inner.seed_initial_file(self.key, file).await
620 }
621
622 fn display_root(&self) -> String {
623 self.inner.display_root()
624 }
625
626 fn display_path(&self, path: &str) -> String {
627 self.inner.display_path(path)
628 }
629}
630
631#[async_trait]
632impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
633 fn display_root(&self) -> String {
634 (**self).display_root()
635 }
636
637 fn display_path(&self, path: &str) -> String {
638 (**self).display_path(path)
639 }
640
641 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
642 (**self).read_file(session_id, path).await
643 }
644
645 async fn write_file(
646 &self,
647 session_id: SessionId,
648 path: &str,
649 content: &str,
650 encoding: &str,
651 ) -> Result<SessionFile> {
652 (**self)
653 .write_file(session_id, path, content, encoding)
654 .await
655 }
656
657 async fn write_file_if_content_matches(
658 &self,
659 session_id: SessionId,
660 path: &str,
661 expected_content: &str,
662 expected_encoding: &str,
663 content: &str,
664 encoding: &str,
665 ) -> Result<Option<SessionFile>> {
666 (**self)
667 .write_file_if_content_matches(
668 session_id,
669 path,
670 expected_content,
671 expected_encoding,
672 content,
673 encoding,
674 )
675 .await
676 }
677
678 async fn delete_file(
679 &self,
680 session_id: SessionId,
681 path: &str,
682 recursive: bool,
683 ) -> Result<bool> {
684 (**self).delete_file(session_id, path, recursive).await
685 }
686
687 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
688 (**self).list_directory(session_id, path).await
689 }
690
691 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
692 (**self).stat_file(session_id, path).await
693 }
694
695 async fn grep_files(
696 &self,
697 session_id: SessionId,
698 pattern: &str,
699 path_pattern: Option<&str>,
700 ) -> Result<Vec<GrepMatch>> {
701 (**self).grep_files(session_id, pattern, path_pattern).await
702 }
703
704 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
705 (**self).create_directory(session_id, path).await
706 }
707
708 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
709 (**self).seed_initial_file(session_id, file).await
710 }
711}
712
713pub use SessionFileSystem as SessionFileStore;
715
716#[derive(Clone, Default)]
722pub struct SessionFileSystemFactoryContext {
723 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
724}
725
726impl SessionFileSystemFactoryContext {
727 pub fn new() -> Self {
728 Self::default()
729 }
730
731 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
732 let values = Arc::make_mut(&mut self.values);
733 values.insert(TypeId::of::<T>(), value);
734 self
735 }
736
737 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
738 self.values
739 .get(&TypeId::of::<T>())
740 .and_then(|value| value.clone().downcast::<T>().ok())
741 }
742}
743
744#[async_trait]
746pub trait SessionFileSystemFactory: Send + Sync {
747 fn name(&self) -> &'static str {
749 "SessionFileSystemFactory"
750 }
751
752 fn is_disabled(&self) -> bool {
755 false
756 }
757
758 async fn create_session_file_system(
760 &self,
761 context: SessionFileSystemFactoryContext,
762 ) -> Result<Arc<dyn SessionFileSystem>>;
763}
764
765#[derive(Debug, Clone, Default)]
767pub struct DisabledSessionFileSystemFactory;
768
769#[async_trait]
770impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
771 fn name(&self) -> &'static str {
772 "DisabledSessionFileSystemFactory"
773 }
774
775 fn is_disabled(&self) -> bool {
776 true
777 }
778
779 async fn create_session_file_system(
780 &self,
781 _context: SessionFileSystemFactoryContext,
782 ) -> Result<Arc<dyn SessionFileSystem>> {
783 Err(crate::error::AgentLoopError::config(
784 "session filesystem is disabled",
785 ))
786 }
787}
788
789#[derive(Debug, Clone)]
795pub struct KeyInfo {
796 pub key: String,
797 pub created_at: chrono::DateTime<chrono::Utc>,
798 pub updated_at: chrono::DateTime<chrono::Utc>,
799}
800
801#[derive(Debug, Clone)]
803pub struct SecretInfo {
804 pub name: String,
805 pub created_at: chrono::DateTime<chrono::Utc>,
806 pub updated_at: chrono::DateTime<chrono::Utc>,
807}
808
809#[derive(Debug, Clone, serde::Serialize)]
819pub struct KnowledgeSearchHit {
820 pub id: String,
822 pub kb_id: String,
824 pub title: String,
825 pub kind: String,
826 pub tags: Vec<String>,
827 pub snippet: String,
829 pub resource: Option<String>,
831}
832
833#[async_trait]
837pub trait KnowledgeStore: Send + Sync {
838 async fn search_knowledge(
839 &self,
840 org_id: crate::typed_id::OrgId,
841 kb_public_ids: &[String],
842 query: &str,
843 kind: Option<&str>,
844 tags: &[String],
845 limit: usize,
846 ) -> Result<Vec<KnowledgeSearchHit>>;
847}
848
849#[async_trait]
854pub trait SessionStorageStore: Send + Sync {
855 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
859
860 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
862
863 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
865
866 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
868
869 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
873
874 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
876
877 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
879
880 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
882}
883
884use crate::session_schedule::SessionSchedule;
889use crate::typed_id::ScheduleId;
890
891#[async_trait]
895pub trait SessionScheduleStore: Send + Sync {
896 async fn create_schedule(
898 &self,
899 session_id: SessionId,
900 description: String,
901 cron_expression: Option<String>,
902 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
903 timezone: String,
904 ) -> Result<SessionSchedule>;
905
906 async fn cancel_schedule(
908 &self,
909 session_id: SessionId,
910 schedule_id: ScheduleId,
911 ) -> Result<SessionSchedule>;
912
913 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
915
916 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
918}
919
920#[async_trait]
930pub trait SessionResourceRegistry: Send + Sync {
931 async fn register(
933 &self,
934 entry: crate::session_resource::RegisterSessionResource,
935 ) -> Result<crate::session_resource::SessionResourceEntry>;
936
937 async fn update_status(
939 &self,
940 session_id: SessionId,
941 resource_id: &str,
942 status: crate::session_resource::SessionResourceStatus,
943 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
944
945 async fn get(
947 &self,
948 session_id: SessionId,
949 resource_id: &str,
950 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
951
952 async fn list(
954 &self,
955 session_id: SessionId,
956 filter: Option<&crate::session_resource::SessionResourceFilter>,
957 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
958
959 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
961}
962
963#[async_trait]
973pub trait LeasedResourceStore: Send + Sync {
974 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
980
981 async fn release_resource(
987 &self,
988 session_id: SessionId,
989 provider: &str,
990 resource_type: &str,
991 external_id: &str,
992 ) -> Result<Option<LeasedResource>>;
993
994 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
999}
1000
1001pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
1007
1008#[async_trait]
1013pub trait UserConnectionResolver: Send + Sync {
1014 async fn get_connection_token(
1017 &self,
1018 session_id: SessionId,
1019 provider: &str,
1020 ) -> Result<Option<String>>;
1021
1022 async fn get_connection_user(
1027 &self,
1028 _session_id: SessionId,
1029 _provider: &str,
1030 ) -> Result<Option<Uuid>> {
1031 Ok(None)
1032 }
1033
1034 async fn get_connection_token_for_user(
1039 &self,
1040 _user_id: Uuid,
1041 _provider: &str,
1042 ) -> Result<Option<String>> {
1043 Ok(None)
1044 }
1045
1046 async fn get_connection_metadata(
1049 &self,
1050 _session_id: SessionId,
1051 _provider: &str,
1052 ) -> Result<Option<serde_json::Value>> {
1053 Ok(None)
1054 }
1055}
1056
1057#[async_trait]
1067pub trait BudgetChecker: Send + Sync {
1068 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
1070}
1071
1072#[async_trait]
1081pub trait PaymentAuthority: Send + Sync {
1082 async fn execute_machine_payment(
1083 &self,
1084 session_id: SessionId,
1085 request: crate::payment::MachinePaymentRequest,
1086 ) -> Result<crate::payment::MachinePaymentResponse>;
1087}
1088
1089#[async_trait]
1099pub trait OutboundToolRateLimiter: Send + Sync {
1100 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
1102}
1103
1104#[derive(Debug)]
1110pub enum ToolCallClaimResult {
1111 Claimed { claim_token: uuid::Uuid },
1114 AlreadySettled {
1116 result_json: serde_json::Value,
1117 args_fingerprint: String,
1118 },
1119 AlreadyRunning { args_fingerprint: String },
1124 DeterminismViolation {
1128 stored_fingerprint: String,
1129 current_fingerprint: String,
1130 },
1131}
1132
1133#[derive(Debug, Clone)]
1135pub enum DurableToolCallStatus {
1136 Settled { result_json: serde_json::Value },
1138 Interrupted {
1140 result_json: Option<serde_json::Value>,
1141 },
1142 Running,
1144}
1145
1146#[async_trait]
1151pub trait DurableToolResultStore: Send + Sync + 'static {
1152 async fn try_claim_tool_call(
1160 &self,
1161 turn_id: &str,
1162 tool_call_id: &str,
1163 tool_name: &str,
1164 args_fingerprint: &str,
1165 ) -> Result<ToolCallClaimResult>;
1166
1167 async fn settle_tool_call(
1173 &self,
1174 turn_id: &str,
1175 tool_call_id: &str,
1176 result_json: serde_json::Value,
1177 status: &str,
1178 claim_token: uuid::Uuid,
1179 ) -> Result<bool>;
1180
1181 async fn get_tool_call_status(
1186 &self,
1187 turn_id: &str,
1188 tool_call_id: &str,
1189 ) -> Result<Option<DurableToolCallStatus>>;
1190}
1191
1192pub struct NoopDurableToolResultStore;
1195
1196#[async_trait]
1197impl DurableToolResultStore for NoopDurableToolResultStore {
1198 async fn try_claim_tool_call(
1199 &self,
1200 _turn_id: &str,
1201 _tool_call_id: &str,
1202 _tool_name: &str,
1203 _args_fingerprint: &str,
1204 ) -> Result<ToolCallClaimResult> {
1205 Ok(ToolCallClaimResult::Claimed {
1206 claim_token: uuid::Uuid::new_v4(),
1207 })
1208 }
1209
1210 async fn settle_tool_call(
1211 &self,
1212 _turn_id: &str,
1213 _tool_call_id: &str,
1214 _result_json: serde_json::Value,
1215 _status: &str,
1216 _claim_token: uuid::Uuid,
1217 ) -> Result<bool> {
1218 Ok(true)
1219 }
1220
1221 async fn get_tool_call_status(
1222 &self,
1223 _turn_id: &str,
1224 _tool_call_id: &str,
1225 ) -> Result<Option<DurableToolCallStatus>> {
1226 Ok(None)
1227 }
1228}
1229
1230#[derive(Debug, Clone)]
1236pub struct StreamProgress {
1237 pub accumulated_len: usize,
1239 pub last_delta_at: u64,
1241}
1242
1243#[async_trait]
1249pub trait StreamHeartbeater: Send + Sync {
1250 async fn heartbeat(&self, progress: StreamProgress);
1256}
1257
1258pub struct NoopStreamHeartbeater;
1260
1261#[async_trait]
1262impl StreamHeartbeater for NoopStreamHeartbeater {
1263 async fn heartbeat(&self, _progress: StreamProgress) {}
1264}
1265
1266#[derive(Debug, Clone)]
1272pub struct PartialStreamState {
1273 pub accumulated: String,
1276}
1277
1278#[async_trait]
1286pub trait PartialStreamStore: Send + Sync {
1287 async fn get_partial_stream(
1290 &self,
1291 session_id: SessionId,
1292 turn_id: &str,
1293 ) -> Result<Option<PartialStreamState>>;
1294}
1295
1296pub struct NoopPartialStreamStore;
1298
1299#[async_trait]
1300impl PartialStreamStore for NoopPartialStreamStore {
1301 async fn get_partial_stream(
1302 &self,
1303 _session_id: SessionId,
1304 _turn_id: &str,
1305 ) -> Result<Option<PartialStreamState>> {
1306 Ok(None)
1307 }
1308}
1309
1310#[derive(Clone)]
1319pub struct ToolContext {
1320 pub session_id: SessionId,
1322 pub workspace_id: WorkspaceId,
1329
1330 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1332
1333 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1335
1336 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1338
1339 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1341
1342 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1344
1345 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1351
1352 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1354
1355 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1357
1358 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1360
1361 pub session_store: Option<Arc<dyn SessionStore>>,
1363
1364 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1366
1367 pub agent_store: Option<Arc<dyn AgentStore>>,
1369
1370 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1372
1373 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1375
1376 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1378 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1380
1381 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1385
1386 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1388
1389 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1391
1392 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1395
1396 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1399
1400 pub event_context: Option<crate::events::EventContext>,
1403
1404 pub tool_call_id: Option<String>,
1407 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1409
1410 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1413
1414 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1418
1419 pub org_id: Option<crate::typed_id::OrgId>,
1421
1422 pub network_access: Option<crate::network_access::NetworkAccessList>,
1425
1426 pub locale: Option<String>,
1430
1431 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1433
1434 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1436
1437 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1441
1442 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1446}
1447
1448impl ToolContext {
1449 pub fn workspace_fs_key(&self) -> SessionId {
1454 SessionId::from_uuid(self.workspace_id.uuid())
1455 }
1456
1457 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1459 self.workspace_id = workspace_id;
1460 self
1461 }
1462
1463 pub fn new(session_id: SessionId) -> Self {
1465 Self {
1466 session_id,
1467 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1468 file_store: None,
1469 storage_store: None,
1470 image_store: None,
1471 provider_credential_store: None,
1472 utility_llm_service: None,
1473 mcp_invoker: None,
1474 egress_service: None,
1475 sqldb_store: None,
1476 message_retriever: None,
1477 session_store: None,
1478 session_mutator: None,
1479 agent_store: None,
1480 connection_resolver: None,
1481 schedule_store: None,
1482 platform_store: None,
1483 knowledge_store: None,
1484 knowledge_index_search: None,
1485 leased_resource_store: None,
1486 session_resource_registry: None,
1487 session_task_registry: None,
1488 event_emitter: None,
1489 event_context: None,
1490 tool_call_id: None,
1491 capability_registry: None,
1492 tool_registry: None,
1493 visible_tool_names: None,
1494 org_id: None,
1495 network_access: None,
1496 locale: None,
1497 budget_checker: None,
1498 payment_authority: None,
1499 subagent_spawn_store: None,
1500 reasoning_effort_handle: None,
1501 }
1502 }
1503
1504 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1506 Self {
1507 session_id,
1508 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1509 file_store: Some(file_store),
1510 storage_store: None,
1511 image_store: None,
1512 provider_credential_store: None,
1513 utility_llm_service: None,
1514 mcp_invoker: None,
1515 egress_service: None,
1516 sqldb_store: None,
1517 message_retriever: None,
1518 session_store: None,
1519 session_mutator: None,
1520 agent_store: None,
1521 connection_resolver: None,
1522 schedule_store: None,
1523 platform_store: None,
1524 knowledge_store: None,
1525 knowledge_index_search: None,
1526 leased_resource_store: None,
1527 session_resource_registry: None,
1528 session_task_registry: None,
1529 event_emitter: None,
1530 event_context: None,
1531 tool_call_id: None,
1532 capability_registry: None,
1533 tool_registry: None,
1534 visible_tool_names: None,
1535 org_id: None,
1536 network_access: None,
1537 locale: None,
1538 budget_checker: None,
1539 payment_authority: None,
1540 subagent_spawn_store: None,
1541 reasoning_effort_handle: None,
1542 }
1543 }
1544
1545 pub fn with_storage_store(
1547 session_id: SessionId,
1548 storage_store: Arc<dyn SessionStorageStore>,
1549 ) -> Self {
1550 Self {
1551 session_id,
1552 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1553 file_store: None,
1554 storage_store: Some(storage_store),
1555 image_store: None,
1556 provider_credential_store: None,
1557 utility_llm_service: None,
1558 mcp_invoker: None,
1559 egress_service: None,
1560 sqldb_store: None,
1561 message_retriever: None,
1562 session_store: None,
1563 session_mutator: None,
1564 agent_store: None,
1565 connection_resolver: None,
1566 schedule_store: None,
1567 platform_store: None,
1568 knowledge_store: None,
1569 knowledge_index_search: None,
1570 leased_resource_store: None,
1571 session_resource_registry: None,
1572 session_task_registry: None,
1573 event_emitter: None,
1574 event_context: None,
1575 tool_call_id: None,
1576 capability_registry: None,
1577 tool_registry: None,
1578 visible_tool_names: None,
1579 org_id: None,
1580 network_access: None,
1581 locale: None,
1582 budget_checker: None,
1583 payment_authority: None,
1584 subagent_spawn_store: None,
1585 reasoning_effort_handle: None,
1586 }
1587 }
1588
1589 pub fn with_stores(
1591 session_id: SessionId,
1592 file_store: Arc<dyn SessionFileSystem>,
1593 storage_store: Arc<dyn SessionStorageStore>,
1594 ) -> Self {
1595 Self {
1596 session_id,
1597 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1598 file_store: Some(file_store),
1599 storage_store: Some(storage_store),
1600 sqldb_store: None,
1601 image_store: None,
1602 provider_credential_store: None,
1603 utility_llm_service: None,
1604 mcp_invoker: None,
1605 egress_service: None,
1606 message_retriever: None,
1607 session_store: None,
1608 session_mutator: None,
1609 agent_store: None,
1610 connection_resolver: None,
1611 schedule_store: None,
1612 platform_store: None,
1613 knowledge_store: None,
1614 knowledge_index_search: None,
1615 leased_resource_store: None,
1616 session_resource_registry: None,
1617 session_task_registry: None,
1618 event_emitter: None,
1619 event_context: None,
1620 tool_call_id: None,
1621 capability_registry: None,
1622 tool_registry: None,
1623 visible_tool_names: None,
1624 org_id: None,
1625 network_access: None,
1626 locale: None,
1627 budget_checker: None,
1628 payment_authority: None,
1629 subagent_spawn_store: None,
1630 reasoning_effort_handle: None,
1631 }
1632 }
1633
1634 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1636 self.sqldb_store = Some(sqldb_store);
1637 self
1638 }
1639
1640 pub fn with_message_retriever(
1642 mut self,
1643 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1644 ) -> Self {
1645 self.message_retriever = Some(retriever);
1646 self
1647 }
1648
1649 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1651 self.session_store = Some(store);
1652 self
1653 }
1654
1655 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1657 self.session_mutator = Some(mutator);
1658 self
1659 }
1660
1661 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
1665 self.reasoning_effort_handle = Some(handle);
1666 self
1667 }
1668
1669 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1671 self.agent_store = Some(store);
1672 self
1673 }
1674
1675 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1677 self.connection_resolver = Some(resolver);
1678 self
1679 }
1680
1681 pub fn with_image_store(
1683 session_id: SessionId,
1684 image_store: Arc<dyn ImageArtifactStore>,
1685 ) -> Self {
1686 Self {
1687 session_id,
1688 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1689 file_store: None,
1690 storage_store: None,
1691 image_store: Some(image_store),
1692 provider_credential_store: None,
1693 utility_llm_service: None,
1694 mcp_invoker: None,
1695 egress_service: None,
1696 sqldb_store: None,
1697 message_retriever: None,
1698 session_store: None,
1699 session_mutator: None,
1700 agent_store: None,
1701 connection_resolver: None,
1702 schedule_store: None,
1703 platform_store: None,
1704 knowledge_store: None,
1705 knowledge_index_search: None,
1706 leased_resource_store: None,
1707 session_resource_registry: None,
1708 session_task_registry: None,
1709 event_emitter: None,
1710 event_context: None,
1711 tool_call_id: None,
1712 capability_registry: None,
1713 tool_registry: None,
1714 visible_tool_names: None,
1715 org_id: None,
1716 network_access: None,
1717 locale: None,
1718 budget_checker: None,
1719 payment_authority: None,
1720 subagent_spawn_store: None,
1721 reasoning_effort_handle: None,
1722 }
1723 }
1724
1725 pub fn with_provider_credential_store(
1727 mut self,
1728 store: Arc<dyn ProviderCredentialStore>,
1729 ) -> Self {
1730 self.provider_credential_store = Some(store);
1731 self
1732 }
1733
1734 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1736 self.utility_llm_service = Some(service);
1737 self
1738 }
1739
1740 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
1742 self.mcp_invoker = Some(invoker);
1743 self
1744 }
1745
1746 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1748 self.egress_service = Some(service);
1749 self
1750 }
1751
1752 pub fn with_egress_service_opt(
1755 mut self,
1756 service: Option<Arc<dyn crate::EgressService>>,
1757 ) -> Self {
1758 if let Some(service) = service {
1759 self.egress_service = Some(service);
1760 }
1761 self
1762 }
1763
1764 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
1766 self.storage_store = Some(store);
1767 self
1768 }
1769
1770 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1772 self.schedule_store = Some(store);
1773 self
1774 }
1775
1776 pub fn with_platform_store(
1778 mut self,
1779 store: Arc<dyn crate::platform_store::PlatformStore>,
1780 ) -> Self {
1781 self.platform_store = Some(store);
1782 self
1783 }
1784
1785 pub fn with_knowledge_index_search(
1787 mut self,
1788 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
1789 ) -> Self {
1790 self.knowledge_index_search = Some(search);
1791 self
1792 }
1793
1794 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1796 self.leased_resource_store = Some(store);
1797 self
1798 }
1799
1800 pub fn with_session_resource_registry(
1802 mut self,
1803 registry: Arc<dyn SessionResourceRegistry>,
1804 ) -> Self {
1805 self.session_resource_registry = Some(registry);
1806 self
1807 }
1808
1809 pub fn with_session_task_registry(
1811 mut self,
1812 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
1813 ) -> Self {
1814 self.session_task_registry = Some(registry);
1815 self
1816 }
1817
1818 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
1820 self.org_id = Some(org_id);
1821 self
1822 }
1823
1824 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
1826 self.tool_registry = Some(registry);
1827 self
1828 }
1829
1830 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
1832 self.visible_tool_names = Some(names);
1833 self
1834 }
1835
1836 pub fn with_network_access(
1838 mut self,
1839 network_access: Option<crate::network_access::NetworkAccessList>,
1840 ) -> Self {
1841 self.network_access = network_access;
1842 self
1843 }
1844
1845 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
1847 self.payment_authority = Some(authority);
1848 self
1849 }
1850
1851 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
1853 self.subagent_spawn_store = Some(store);
1854 self
1855 }
1856
1857 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
1862 let (Some(emitter), Some(ctx), Some(call_id)) =
1863 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1864 else {
1865 return;
1866 };
1867 if let Err(e) = emitter
1868 .emit(EventRequest::new(
1869 self.session_id,
1870 ctx.clone(),
1871 crate::events::ToolProgressData {
1872 tool_call_id: call_id.clone(),
1873 tool_name: tool_name.to_string(),
1874 message: message.to_string(),
1875 display_name: None,
1876 },
1877 ))
1878 .await
1879 {
1880 tracing::debug!(
1881 tool_call_id = call_id,
1882 tool_name,
1883 error = %e,
1884 "Failed to emit tool.progress event"
1885 );
1886 }
1887 }
1888
1889 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
1894 let (Some(emitter), Some(ctx), Some(call_id)) =
1895 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1896 else {
1897 return;
1898 };
1899 if let Err(e) = emitter
1900 .emit(EventRequest::new(
1901 self.session_id,
1902 ctx.clone(),
1903 crate::events::ToolOutputDeltaData {
1904 tool_call_id: call_id.clone(),
1905 tool_name: tool_name.to_string(),
1906 delta: delta.to_string(),
1907 stream: stream.to_string(),
1908 },
1909 ))
1910 .await
1911 {
1912 tracing::debug!(
1913 tool_call_id = call_id,
1914 tool_name,
1915 error = %e,
1916 "Failed to emit tool.output.delta event"
1917 );
1918 }
1919 }
1920}
1921
1922impl std::fmt::Debug for ToolContext {
1923 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1924 f.debug_struct("ToolContext")
1925 .field("session_id", &self.session_id)
1926 .field("file_store", &self.file_store.is_some())
1927 .field("storage_store", &self.storage_store.is_some())
1928 .field("image_store", &self.image_store.is_some())
1929 .field(
1930 "provider_credential_store",
1931 &self.provider_credential_store.is_some(),
1932 )
1933 .field("utility_llm_service", &self.utility_llm_service.is_some())
1934 .field("egress_service", &self.egress_service.is_some())
1935 .field("sqldb_store", &self.sqldb_store.is_some())
1936 .field("message_retriever", &self.message_retriever.is_some())
1937 .field("session_store", &self.session_store.is_some())
1938 .field("session_mutator", &self.session_mutator.is_some())
1939 .field("agent_store", &self.agent_store.is_some())
1940 .field("connection_resolver", &self.connection_resolver.is_some())
1941 .field("schedule_store", &self.schedule_store.is_some())
1942 .field("platform_store", &self.platform_store.is_some())
1943 .field(
1944 "knowledge_index_search",
1945 &self.knowledge_index_search.is_some(),
1946 )
1947 .field(
1948 "leased_resource_store",
1949 &self.leased_resource_store.is_some(),
1950 )
1951 .field("event_emitter", &self.event_emitter.is_some())
1952 .field("tool_registry", &self.tool_registry.is_some())
1953 .field("payment_authority", &self.payment_authority.is_some())
1954 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
1955 .field("org_id", &self.org_id)
1956 .finish()
1957 }
1958}
1959
1960use crate::events::{Event, EventRequest};
1965
1966#[async_trait]
1977pub trait EventEmitter: Send + Sync {
1978 async fn emit(&self, request: EventRequest) -> Result<Event>;
1983}
1984
1985#[async_trait]
1987impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
1988 async fn emit(&self, request: EventRequest) -> Result<Event> {
1989 (**self).emit(request).await
1990 }
1991}
1992
1993#[derive(Debug, Clone, Default)]
1997pub struct NoopEventEmitter;
1998
1999#[async_trait]
2000impl EventEmitter for NoopEventEmitter {
2001 async fn emit(&self, request: EventRequest) -> Result<Event> {
2002 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2004 }
2005}
2006
2007#[derive(Debug, Clone)]
2020pub struct ResolvedImage {
2021 pub base64: String,
2023 pub media_type: String,
2025}
2026
2027impl ResolvedImage {
2028 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2030 Self {
2031 base64: base64.into(),
2032 media_type: media_type.into(),
2033 }
2034 }
2035
2036 pub fn to_data_url(&self) -> String {
2040 format!("data:{};base64,{}", self.media_type, self.base64)
2041 }
2042}
2043
2044#[async_trait]
2077pub trait ImageResolver: Send + Sync {
2078 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2082}
2083
2084#[derive(Debug)]
2090pub enum SpawnClaimResult {
2091 Claimed {
2094 spawn_handle_id: uuid::Uuid,
2095 claim_token: uuid::Uuid,
2096 },
2097 ClaimedPendingChild {
2101 spawn_handle_id: uuid::Uuid,
2102 claim_token: uuid::Uuid,
2103 },
2104 AlreadyRunning {
2107 child_session_id: crate::typed_id::SessionId,
2108 claim_token: uuid::Uuid,
2110 },
2111 AlreadySettled {
2114 child_session_id: crate::typed_id::SessionId,
2115 terminal_status: String,
2117 terminal_result: String,
2118 },
2119}
2120
2121#[async_trait]
2129pub trait SubagentSpawnStore: Send + Sync + 'static {
2130 async fn try_claim_spawn(
2135 &self,
2136 parent_session_id: crate::typed_id::SessionId,
2137 tool_call_id: &str,
2138 claim_token: uuid::Uuid,
2139 ) -> Result<SpawnClaimResult>;
2140
2141 async fn register_child_session(
2146 &self,
2147 spawn_handle_id: uuid::Uuid,
2148 claim_token: uuid::Uuid,
2149 child_session_id: crate::typed_id::SessionId,
2150 ) -> Result<()>;
2151
2152 async fn settle_spawn(
2158 &self,
2159 parent_session_id: crate::typed_id::SessionId,
2160 tool_call_id: &str,
2161 claim_token: uuid::Uuid,
2162 terminal_status: &str,
2163 terminal_result: &str,
2164 ) -> Result<()>;
2165}
2166
2167#[async_trait]
2169impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2170 async fn try_claim_spawn(
2171 &self,
2172 parent_session_id: crate::typed_id::SessionId,
2173 tool_call_id: &str,
2174 claim_token: uuid::Uuid,
2175 ) -> Result<SpawnClaimResult> {
2176 (**self)
2177 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2178 .await
2179 }
2180
2181 async fn register_child_session(
2182 &self,
2183 spawn_handle_id: uuid::Uuid,
2184 claim_token: uuid::Uuid,
2185 child_session_id: crate::typed_id::SessionId,
2186 ) -> Result<()> {
2187 (**self)
2188 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2189 .await
2190 }
2191
2192 async fn settle_spawn(
2193 &self,
2194 parent_session_id: crate::typed_id::SessionId,
2195 tool_call_id: &str,
2196 claim_token: uuid::Uuid,
2197 terminal_status: &str,
2198 terminal_result: &str,
2199 ) -> Result<()> {
2200 (**self)
2201 .settle_spawn(
2202 parent_session_id,
2203 tool_call_id,
2204 claim_token,
2205 terminal_status,
2206 terminal_result,
2207 )
2208 .await
2209 }
2210}
2211
2212pub struct NoopSubagentSpawnStore;
2216
2217#[async_trait]
2218impl SubagentSpawnStore for NoopSubagentSpawnStore {
2219 async fn try_claim_spawn(
2220 &self,
2221 _parent_session_id: crate::typed_id::SessionId,
2222 _tool_call_id: &str,
2223 claim_token: uuid::Uuid,
2224 ) -> Result<SpawnClaimResult> {
2225 Ok(SpawnClaimResult::Claimed {
2226 spawn_handle_id: uuid::Uuid::new_v4(),
2227 claim_token,
2228 })
2229 }
2230
2231 async fn register_child_session(
2232 &self,
2233 _spawn_handle_id: uuid::Uuid,
2234 _claim_token: uuid::Uuid,
2235 _child_session_id: crate::typed_id::SessionId,
2236 ) -> Result<()> {
2237 Ok(())
2238 }
2239
2240 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 Ok(())
2249 }
2250}
2251
2252#[cfg(test)]
2257mod tests {
2258 use super::*;
2259
2260 #[test]
2261 fn test_resolved_image_new() {
2262 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2263 assert_eq!(image.base64, "SGVsbG8=");
2264 assert_eq!(image.media_type, "image/png");
2265 }
2266
2267 #[test]
2268 fn test_resolved_image_to_data_url() {
2269 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2270 let data_url = image.to_data_url();
2271 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2272 }
2273
2274 #[test]
2275 fn test_resolved_image_jpeg() {
2276 let image = ResolvedImage::new("base64data", "image/jpeg");
2277 let data_url = image.to_data_url();
2278 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2279 }
2280}