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 async fn count_active_org_schedules(&self) -> Result<u32>;
924}
925
926#[async_trait]
936pub trait SessionResourceRegistry: Send + Sync {
937 async fn register(
939 &self,
940 entry: crate::session_resource::RegisterSessionResource,
941 ) -> Result<crate::session_resource::SessionResourceEntry>;
942
943 async fn update_status(
945 &self,
946 session_id: SessionId,
947 resource_id: &str,
948 status: crate::session_resource::SessionResourceStatus,
949 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
950
951 async fn get(
953 &self,
954 session_id: SessionId,
955 resource_id: &str,
956 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
957
958 async fn list(
960 &self,
961 session_id: SessionId,
962 filter: Option<&crate::session_resource::SessionResourceFilter>,
963 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
964
965 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
967}
968
969#[async_trait]
979pub trait LeasedResourceStore: Send + Sync {
980 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
986
987 async fn release_resource(
993 &self,
994 session_id: SessionId,
995 provider: &str,
996 resource_type: &str,
997 external_id: &str,
998 ) -> Result<Option<LeasedResource>>;
999
1000 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
1005}
1006
1007pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
1013
1014#[async_trait]
1019pub trait UserConnectionResolver: Send + Sync {
1020 async fn get_connection_token(
1023 &self,
1024 session_id: SessionId,
1025 provider: &str,
1026 ) -> Result<Option<String>>;
1027
1028 async fn get_connection_user(
1033 &self,
1034 _session_id: SessionId,
1035 _provider: &str,
1036 ) -> Result<Option<Uuid>> {
1037 Ok(None)
1038 }
1039
1040 async fn get_connection_token_for_user(
1045 &self,
1046 _user_id: Uuid,
1047 _provider: &str,
1048 ) -> Result<Option<String>> {
1049 Ok(None)
1050 }
1051
1052 async fn get_connection_metadata(
1055 &self,
1056 _session_id: SessionId,
1057 _provider: &str,
1058 ) -> Result<Option<serde_json::Value>> {
1059 Ok(None)
1060 }
1061}
1062
1063#[async_trait]
1073pub trait BudgetChecker: Send + Sync {
1074 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
1076}
1077
1078#[async_trait]
1087pub trait PaymentAuthority: Send + Sync {
1088 async fn execute_machine_payment(
1089 &self,
1090 session_id: SessionId,
1091 request: crate::payment::MachinePaymentRequest,
1092 ) -> Result<crate::payment::MachinePaymentResponse>;
1093}
1094
1095#[async_trait]
1105pub trait OutboundToolRateLimiter: Send + Sync {
1106 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
1108}
1109
1110#[derive(Debug)]
1116pub enum ToolCallClaimResult {
1117 Claimed { claim_token: uuid::Uuid },
1120 AlreadySettled {
1122 result_json: serde_json::Value,
1123 args_fingerprint: String,
1124 },
1125 AlreadyRunning { args_fingerprint: String },
1130 DeterminismViolation {
1134 stored_fingerprint: String,
1135 current_fingerprint: String,
1136 },
1137}
1138
1139#[derive(Debug, Clone)]
1141pub enum DurableToolCallStatus {
1142 Settled { result_json: serde_json::Value },
1144 Interrupted {
1146 result_json: Option<serde_json::Value>,
1147 },
1148 Running,
1150}
1151
1152#[async_trait]
1157pub trait DurableToolResultStore: Send + Sync + 'static {
1158 async fn try_claim_tool_call(
1166 &self,
1167 turn_id: &str,
1168 tool_call_id: &str,
1169 tool_name: &str,
1170 args_fingerprint: &str,
1171 ) -> Result<ToolCallClaimResult>;
1172
1173 async fn settle_tool_call(
1179 &self,
1180 turn_id: &str,
1181 tool_call_id: &str,
1182 result_json: serde_json::Value,
1183 status: &str,
1184 claim_token: uuid::Uuid,
1185 ) -> Result<bool>;
1186
1187 async fn get_tool_call_status(
1192 &self,
1193 turn_id: &str,
1194 tool_call_id: &str,
1195 ) -> Result<Option<DurableToolCallStatus>>;
1196}
1197
1198pub struct NoopDurableToolResultStore;
1201
1202#[async_trait]
1203impl DurableToolResultStore for NoopDurableToolResultStore {
1204 async fn try_claim_tool_call(
1205 &self,
1206 _turn_id: &str,
1207 _tool_call_id: &str,
1208 _tool_name: &str,
1209 _args_fingerprint: &str,
1210 ) -> Result<ToolCallClaimResult> {
1211 Ok(ToolCallClaimResult::Claimed {
1212 claim_token: uuid::Uuid::new_v4(),
1213 })
1214 }
1215
1216 async fn settle_tool_call(
1217 &self,
1218 _turn_id: &str,
1219 _tool_call_id: &str,
1220 _result_json: serde_json::Value,
1221 _status: &str,
1222 _claim_token: uuid::Uuid,
1223 ) -> Result<bool> {
1224 Ok(true)
1225 }
1226
1227 async fn get_tool_call_status(
1228 &self,
1229 _turn_id: &str,
1230 _tool_call_id: &str,
1231 ) -> Result<Option<DurableToolCallStatus>> {
1232 Ok(None)
1233 }
1234}
1235
1236#[derive(Debug, Clone)]
1242pub struct StreamProgress {
1243 pub accumulated_len: usize,
1245 pub last_delta_at: u64,
1247}
1248
1249#[async_trait]
1255pub trait StreamHeartbeater: Send + Sync {
1256 async fn heartbeat(&self, progress: StreamProgress);
1262}
1263
1264pub struct NoopStreamHeartbeater;
1266
1267#[async_trait]
1268impl StreamHeartbeater for NoopStreamHeartbeater {
1269 async fn heartbeat(&self, _progress: StreamProgress) {}
1270}
1271
1272#[derive(Debug, Clone)]
1278pub struct PartialStreamState {
1279 pub accumulated: String,
1282}
1283
1284#[async_trait]
1292pub trait PartialStreamStore: Send + Sync {
1293 async fn get_partial_stream(
1296 &self,
1297 session_id: SessionId,
1298 turn_id: &str,
1299 ) -> Result<Option<PartialStreamState>>;
1300}
1301
1302pub struct NoopPartialStreamStore;
1304
1305#[async_trait]
1306impl PartialStreamStore for NoopPartialStreamStore {
1307 async fn get_partial_stream(
1308 &self,
1309 _session_id: SessionId,
1310 _turn_id: &str,
1311 ) -> Result<Option<PartialStreamState>> {
1312 Ok(None)
1313 }
1314}
1315
1316#[derive(Clone)]
1325pub struct ToolContext {
1326 pub session_id: SessionId,
1328 pub workspace_id: WorkspaceId,
1335
1336 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1338
1339 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1341
1342 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1344
1345 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1347
1348 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1350
1351 pub mcp_invoker: Option<Arc<dyn crate::McpToolInvoker>>,
1357
1358 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1360
1361 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1363
1364 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1366
1367 pub session_store: Option<Arc<dyn SessionStore>>,
1369
1370 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1372
1373 pub agent_store: Option<Arc<dyn AgentStore>>,
1375
1376 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1378
1379 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1381
1382 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1384 pub knowledge_store: Option<Arc<dyn KnowledgeStore>>,
1386
1387 pub knowledge_index_search: Option<Arc<dyn crate::vector_store::KnowledgeIndexSearch>>,
1391
1392 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1394
1395 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1397
1398 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1401
1402 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1405
1406 pub event_context: Option<crate::events::EventContext>,
1409
1410 pub tool_call_id: Option<String>,
1413 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1415
1416 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1419
1420 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1424
1425 pub org_id: Option<crate::typed_id::OrgId>,
1427
1428 pub network_access: Option<crate::network_access::NetworkAccessList>,
1431
1432 pub locale: Option<String>,
1436
1437 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1439
1440 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1442
1443 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1447
1448 pub reasoning_effort_handle: Option<ReasoningEffortHandle>,
1452}
1453
1454impl ToolContext {
1455 pub fn workspace_fs_key(&self) -> SessionId {
1460 SessionId::from_uuid(self.workspace_id.uuid())
1461 }
1462
1463 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1465 self.workspace_id = workspace_id;
1466 self
1467 }
1468
1469 pub fn new(session_id: SessionId) -> Self {
1471 Self {
1472 session_id,
1473 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1474 file_store: None,
1475 storage_store: None,
1476 image_store: None,
1477 provider_credential_store: None,
1478 utility_llm_service: None,
1479 mcp_invoker: None,
1480 egress_service: None,
1481 sqldb_store: None,
1482 message_retriever: None,
1483 session_store: None,
1484 session_mutator: None,
1485 agent_store: None,
1486 connection_resolver: None,
1487 schedule_store: None,
1488 platform_store: None,
1489 knowledge_store: None,
1490 knowledge_index_search: None,
1491 leased_resource_store: None,
1492 session_resource_registry: None,
1493 session_task_registry: None,
1494 event_emitter: None,
1495 event_context: None,
1496 tool_call_id: None,
1497 capability_registry: None,
1498 tool_registry: None,
1499 visible_tool_names: None,
1500 org_id: None,
1501 network_access: None,
1502 locale: None,
1503 budget_checker: None,
1504 payment_authority: None,
1505 subagent_spawn_store: None,
1506 reasoning_effort_handle: None,
1507 }
1508 }
1509
1510 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1512 Self {
1513 session_id,
1514 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1515 file_store: Some(file_store),
1516 storage_store: None,
1517 image_store: None,
1518 provider_credential_store: None,
1519 utility_llm_service: None,
1520 mcp_invoker: None,
1521 egress_service: None,
1522 sqldb_store: None,
1523 message_retriever: None,
1524 session_store: None,
1525 session_mutator: None,
1526 agent_store: None,
1527 connection_resolver: None,
1528 schedule_store: None,
1529 platform_store: None,
1530 knowledge_store: None,
1531 knowledge_index_search: None,
1532 leased_resource_store: None,
1533 session_resource_registry: None,
1534 session_task_registry: None,
1535 event_emitter: None,
1536 event_context: None,
1537 tool_call_id: None,
1538 capability_registry: None,
1539 tool_registry: None,
1540 visible_tool_names: None,
1541 org_id: None,
1542 network_access: None,
1543 locale: None,
1544 budget_checker: None,
1545 payment_authority: None,
1546 subagent_spawn_store: None,
1547 reasoning_effort_handle: None,
1548 }
1549 }
1550
1551 pub fn with_storage_store(
1553 session_id: SessionId,
1554 storage_store: Arc<dyn SessionStorageStore>,
1555 ) -> Self {
1556 Self {
1557 session_id,
1558 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1559 file_store: None,
1560 storage_store: Some(storage_store),
1561 image_store: None,
1562 provider_credential_store: None,
1563 utility_llm_service: None,
1564 mcp_invoker: None,
1565 egress_service: None,
1566 sqldb_store: None,
1567 message_retriever: None,
1568 session_store: None,
1569 session_mutator: None,
1570 agent_store: None,
1571 connection_resolver: None,
1572 schedule_store: None,
1573 platform_store: None,
1574 knowledge_store: None,
1575 knowledge_index_search: None,
1576 leased_resource_store: None,
1577 session_resource_registry: None,
1578 session_task_registry: None,
1579 event_emitter: None,
1580 event_context: None,
1581 tool_call_id: None,
1582 capability_registry: None,
1583 tool_registry: None,
1584 visible_tool_names: None,
1585 org_id: None,
1586 network_access: None,
1587 locale: None,
1588 budget_checker: None,
1589 payment_authority: None,
1590 subagent_spawn_store: None,
1591 reasoning_effort_handle: None,
1592 }
1593 }
1594
1595 pub fn with_stores(
1597 session_id: SessionId,
1598 file_store: Arc<dyn SessionFileSystem>,
1599 storage_store: Arc<dyn SessionStorageStore>,
1600 ) -> Self {
1601 Self {
1602 session_id,
1603 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1604 file_store: Some(file_store),
1605 storage_store: Some(storage_store),
1606 sqldb_store: None,
1607 image_store: None,
1608 provider_credential_store: None,
1609 utility_llm_service: None,
1610 mcp_invoker: None,
1611 egress_service: None,
1612 message_retriever: None,
1613 session_store: None,
1614 session_mutator: None,
1615 agent_store: None,
1616 connection_resolver: None,
1617 schedule_store: None,
1618 platform_store: None,
1619 knowledge_store: None,
1620 knowledge_index_search: None,
1621 leased_resource_store: None,
1622 session_resource_registry: None,
1623 session_task_registry: None,
1624 event_emitter: None,
1625 event_context: None,
1626 tool_call_id: None,
1627 capability_registry: None,
1628 tool_registry: None,
1629 visible_tool_names: None,
1630 org_id: None,
1631 network_access: None,
1632 locale: None,
1633 budget_checker: None,
1634 payment_authority: None,
1635 subagent_spawn_store: None,
1636 reasoning_effort_handle: None,
1637 }
1638 }
1639
1640 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1642 self.sqldb_store = Some(sqldb_store);
1643 self
1644 }
1645
1646 pub fn with_message_retriever(
1648 mut self,
1649 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1650 ) -> Self {
1651 self.message_retriever = Some(retriever);
1652 self
1653 }
1654
1655 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1657 self.session_store = Some(store);
1658 self
1659 }
1660
1661 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1663 self.session_mutator = Some(mutator);
1664 self
1665 }
1666
1667 pub fn with_reasoning_effort_handle(mut self, handle: ReasoningEffortHandle) -> Self {
1671 self.reasoning_effort_handle = Some(handle);
1672 self
1673 }
1674
1675 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1677 self.agent_store = Some(store);
1678 self
1679 }
1680
1681 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1683 self.connection_resolver = Some(resolver);
1684 self
1685 }
1686
1687 pub fn with_image_store(
1689 session_id: SessionId,
1690 image_store: Arc<dyn ImageArtifactStore>,
1691 ) -> Self {
1692 Self {
1693 session_id,
1694 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1695 file_store: None,
1696 storage_store: None,
1697 image_store: Some(image_store),
1698 provider_credential_store: None,
1699 utility_llm_service: None,
1700 mcp_invoker: None,
1701 egress_service: None,
1702 sqldb_store: None,
1703 message_retriever: None,
1704 session_store: None,
1705 session_mutator: None,
1706 agent_store: None,
1707 connection_resolver: None,
1708 schedule_store: None,
1709 platform_store: None,
1710 knowledge_store: None,
1711 knowledge_index_search: None,
1712 leased_resource_store: None,
1713 session_resource_registry: None,
1714 session_task_registry: None,
1715 event_emitter: None,
1716 event_context: None,
1717 tool_call_id: None,
1718 capability_registry: None,
1719 tool_registry: None,
1720 visible_tool_names: None,
1721 org_id: None,
1722 network_access: None,
1723 locale: None,
1724 budget_checker: None,
1725 payment_authority: None,
1726 subagent_spawn_store: None,
1727 reasoning_effort_handle: None,
1728 }
1729 }
1730
1731 pub fn with_provider_credential_store(
1733 mut self,
1734 store: Arc<dyn ProviderCredentialStore>,
1735 ) -> Self {
1736 self.provider_credential_store = Some(store);
1737 self
1738 }
1739
1740 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1742 self.utility_llm_service = Some(service);
1743 self
1744 }
1745
1746 pub fn with_mcp_invoker(mut self, invoker: Arc<dyn crate::McpToolInvoker>) -> Self {
1748 self.mcp_invoker = Some(invoker);
1749 self
1750 }
1751
1752 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1754 self.egress_service = Some(service);
1755 self
1756 }
1757
1758 pub fn with_egress_service_opt(
1761 mut self,
1762 service: Option<Arc<dyn crate::EgressService>>,
1763 ) -> Self {
1764 if let Some(service) = service {
1765 self.egress_service = Some(service);
1766 }
1767 self
1768 }
1769
1770 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
1772 self.storage_store = Some(store);
1773 self
1774 }
1775
1776 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1778 self.schedule_store = Some(store);
1779 self
1780 }
1781
1782 pub fn with_platform_store(
1784 mut self,
1785 store: Arc<dyn crate::platform_store::PlatformStore>,
1786 ) -> Self {
1787 self.platform_store = Some(store);
1788 self
1789 }
1790
1791 pub fn with_knowledge_index_search(
1793 mut self,
1794 search: Arc<dyn crate::vector_store::KnowledgeIndexSearch>,
1795 ) -> Self {
1796 self.knowledge_index_search = Some(search);
1797 self
1798 }
1799
1800 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1802 self.leased_resource_store = Some(store);
1803 self
1804 }
1805
1806 pub fn with_session_resource_registry(
1808 mut self,
1809 registry: Arc<dyn SessionResourceRegistry>,
1810 ) -> Self {
1811 self.session_resource_registry = Some(registry);
1812 self
1813 }
1814
1815 pub fn with_session_task_registry(
1817 mut self,
1818 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
1819 ) -> Self {
1820 self.session_task_registry = Some(registry);
1821 self
1822 }
1823
1824 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
1826 self.org_id = Some(org_id);
1827 self
1828 }
1829
1830 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
1832 self.tool_registry = Some(registry);
1833 self
1834 }
1835
1836 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
1838 self.visible_tool_names = Some(names);
1839 self
1840 }
1841
1842 pub fn with_network_access(
1844 mut self,
1845 network_access: Option<crate::network_access::NetworkAccessList>,
1846 ) -> Self {
1847 self.network_access = network_access;
1848 self
1849 }
1850
1851 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
1853 self.payment_authority = Some(authority);
1854 self
1855 }
1856
1857 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
1859 self.subagent_spawn_store = Some(store);
1860 self
1861 }
1862
1863 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
1868 let (Some(emitter), Some(ctx), Some(call_id)) =
1869 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1870 else {
1871 return;
1872 };
1873 if let Err(e) = emitter
1874 .emit(EventRequest::new(
1875 self.session_id,
1876 ctx.clone(),
1877 crate::events::ToolProgressData {
1878 tool_call_id: call_id.clone(),
1879 tool_name: tool_name.to_string(),
1880 message: message.to_string(),
1881 display_name: None,
1882 },
1883 ))
1884 .await
1885 {
1886 tracing::debug!(
1887 tool_call_id = call_id,
1888 tool_name,
1889 error = %e,
1890 "Failed to emit tool.progress event"
1891 );
1892 }
1893 }
1894
1895 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
1900 let (Some(emitter), Some(ctx), Some(call_id)) =
1901 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1902 else {
1903 return;
1904 };
1905 if let Err(e) = emitter
1906 .emit(EventRequest::new(
1907 self.session_id,
1908 ctx.clone(),
1909 crate::events::ToolOutputDeltaData {
1910 tool_call_id: call_id.clone(),
1911 tool_name: tool_name.to_string(),
1912 delta: delta.to_string(),
1913 stream: stream.to_string(),
1914 },
1915 ))
1916 .await
1917 {
1918 tracing::debug!(
1919 tool_call_id = call_id,
1920 tool_name,
1921 error = %e,
1922 "Failed to emit tool.output.delta event"
1923 );
1924 }
1925 }
1926}
1927
1928impl std::fmt::Debug for ToolContext {
1929 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1930 f.debug_struct("ToolContext")
1931 .field("session_id", &self.session_id)
1932 .field("file_store", &self.file_store.is_some())
1933 .field("storage_store", &self.storage_store.is_some())
1934 .field("image_store", &self.image_store.is_some())
1935 .field(
1936 "provider_credential_store",
1937 &self.provider_credential_store.is_some(),
1938 )
1939 .field("utility_llm_service", &self.utility_llm_service.is_some())
1940 .field("egress_service", &self.egress_service.is_some())
1941 .field("sqldb_store", &self.sqldb_store.is_some())
1942 .field("message_retriever", &self.message_retriever.is_some())
1943 .field("session_store", &self.session_store.is_some())
1944 .field("session_mutator", &self.session_mutator.is_some())
1945 .field("agent_store", &self.agent_store.is_some())
1946 .field("connection_resolver", &self.connection_resolver.is_some())
1947 .field("schedule_store", &self.schedule_store.is_some())
1948 .field("platform_store", &self.platform_store.is_some())
1949 .field(
1950 "knowledge_index_search",
1951 &self.knowledge_index_search.is_some(),
1952 )
1953 .field(
1954 "leased_resource_store",
1955 &self.leased_resource_store.is_some(),
1956 )
1957 .field("event_emitter", &self.event_emitter.is_some())
1958 .field("tool_registry", &self.tool_registry.is_some())
1959 .field("payment_authority", &self.payment_authority.is_some())
1960 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
1961 .field("org_id", &self.org_id)
1962 .finish()
1963 }
1964}
1965
1966use crate::events::{Event, EventRequest};
1971
1972#[async_trait]
1983pub trait EventEmitter: Send + Sync {
1984 async fn emit(&self, request: EventRequest) -> Result<Event>;
1989}
1990
1991#[async_trait]
1993impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
1994 async fn emit(&self, request: EventRequest) -> Result<Event> {
1995 (**self).emit(request).await
1996 }
1997}
1998
1999#[derive(Debug, Clone, Default)]
2003pub struct NoopEventEmitter;
2004
2005#[async_trait]
2006impl EventEmitter for NoopEventEmitter {
2007 async fn emit(&self, request: EventRequest) -> Result<Event> {
2008 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
2010 }
2011}
2012
2013#[derive(Debug, Clone)]
2026pub struct ResolvedImage {
2027 pub base64: String,
2029 pub media_type: String,
2031}
2032
2033impl ResolvedImage {
2034 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
2036 Self {
2037 base64: base64.into(),
2038 media_type: media_type.into(),
2039 }
2040 }
2041
2042 pub fn to_data_url(&self) -> String {
2046 format!("data:{};base64,{}", self.media_type, self.base64)
2047 }
2048}
2049
2050#[async_trait]
2083pub trait ImageResolver: Send + Sync {
2084 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
2088}
2089
2090#[derive(Debug)]
2096pub enum SpawnClaimResult {
2097 Claimed {
2100 spawn_handle_id: uuid::Uuid,
2101 claim_token: uuid::Uuid,
2102 },
2103 ClaimedPendingChild {
2107 spawn_handle_id: uuid::Uuid,
2108 claim_token: uuid::Uuid,
2109 },
2110 AlreadyRunning {
2113 child_session_id: crate::typed_id::SessionId,
2114 claim_token: uuid::Uuid,
2116 },
2117 AlreadySettled {
2120 child_session_id: crate::typed_id::SessionId,
2121 terminal_status: String,
2123 terminal_result: String,
2124 },
2125}
2126
2127#[async_trait]
2135pub trait SubagentSpawnStore: Send + Sync + 'static {
2136 async fn try_claim_spawn(
2141 &self,
2142 parent_session_id: crate::typed_id::SessionId,
2143 tool_call_id: &str,
2144 claim_token: uuid::Uuid,
2145 ) -> Result<SpawnClaimResult>;
2146
2147 async fn register_child_session(
2152 &self,
2153 spawn_handle_id: uuid::Uuid,
2154 claim_token: uuid::Uuid,
2155 child_session_id: crate::typed_id::SessionId,
2156 ) -> Result<()>;
2157
2158 async fn settle_spawn(
2164 &self,
2165 parent_session_id: crate::typed_id::SessionId,
2166 tool_call_id: &str,
2167 claim_token: uuid::Uuid,
2168 terminal_status: &str,
2169 terminal_result: &str,
2170 ) -> Result<()>;
2171}
2172
2173#[async_trait]
2175impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
2176 async fn try_claim_spawn(
2177 &self,
2178 parent_session_id: crate::typed_id::SessionId,
2179 tool_call_id: &str,
2180 claim_token: uuid::Uuid,
2181 ) -> Result<SpawnClaimResult> {
2182 (**self)
2183 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
2184 .await
2185 }
2186
2187 async fn register_child_session(
2188 &self,
2189 spawn_handle_id: uuid::Uuid,
2190 claim_token: uuid::Uuid,
2191 child_session_id: crate::typed_id::SessionId,
2192 ) -> Result<()> {
2193 (**self)
2194 .register_child_session(spawn_handle_id, claim_token, child_session_id)
2195 .await
2196 }
2197
2198 async fn settle_spawn(
2199 &self,
2200 parent_session_id: crate::typed_id::SessionId,
2201 tool_call_id: &str,
2202 claim_token: uuid::Uuid,
2203 terminal_status: &str,
2204 terminal_result: &str,
2205 ) -> Result<()> {
2206 (**self)
2207 .settle_spawn(
2208 parent_session_id,
2209 tool_call_id,
2210 claim_token,
2211 terminal_status,
2212 terminal_result,
2213 )
2214 .await
2215 }
2216}
2217
2218pub struct NoopSubagentSpawnStore;
2222
2223#[async_trait]
2224impl SubagentSpawnStore for NoopSubagentSpawnStore {
2225 async fn try_claim_spawn(
2226 &self,
2227 _parent_session_id: crate::typed_id::SessionId,
2228 _tool_call_id: &str,
2229 claim_token: uuid::Uuid,
2230 ) -> Result<SpawnClaimResult> {
2231 Ok(SpawnClaimResult::Claimed {
2232 spawn_handle_id: uuid::Uuid::new_v4(),
2233 claim_token,
2234 })
2235 }
2236
2237 async fn register_child_session(
2238 &self,
2239 _spawn_handle_id: uuid::Uuid,
2240 _claim_token: uuid::Uuid,
2241 _child_session_id: crate::typed_id::SessionId,
2242 ) -> Result<()> {
2243 Ok(())
2244 }
2245
2246 async fn settle_spawn(
2247 &self,
2248 _parent_session_id: crate::typed_id::SessionId,
2249 _tool_call_id: &str,
2250 _claim_token: uuid::Uuid,
2251 _terminal_status: &str,
2252 _terminal_result: &str,
2253 ) -> Result<()> {
2254 Ok(())
2255 }
2256}
2257
2258#[cfg(test)]
2263mod tests {
2264 use super::*;
2265
2266 #[test]
2267 fn test_resolved_image_new() {
2268 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2269 assert_eq!(image.base64, "SGVsbG8=");
2270 assert_eq!(image.media_type, "image/png");
2271 }
2272
2273 #[test]
2274 fn test_resolved_image_to_data_url() {
2275 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2276 let data_url = image.to_data_url();
2277 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2278 }
2279
2280 #[test]
2281 fn test_resolved_image_jpeg() {
2282 let image = ResolvedImage::new("base64data", "image/jpeg");
2283 let data_url = image.to_data_url();
2284 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2285 }
2286}