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#[async_trait]
39pub trait AgentStore: Send + Sync {
40 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>>;
42}
43
44#[async_trait]
45impl<T: AgentStore + ?Sized> AgentStore for std::sync::Arc<T> {
46 async fn get_agent(&self, agent_id: AgentId) -> Result<Option<Agent>> {
47 (**self).get_agent(agent_id).await
48 }
49}
50
51#[async_trait]
66pub trait HarnessStore: Send + Sync {
67 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>>;
72}
73
74#[async_trait]
75impl<T: HarnessStore + ?Sized> HarnessStore for std::sync::Arc<T> {
76 async fn get_harness_chain(&self, harness_id: HarnessId) -> Result<Vec<Harness>> {
77 (**self).get_harness_chain(harness_id).await
78 }
79}
80
81use crate::leased_resource::{LeasedResource, UpsertLeasedResource};
86use crate::session::Session;
87
88#[async_trait]
94pub trait SessionStore: Send + Sync {
95 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>>;
97}
98
99#[async_trait]
100impl<T: SessionStore + ?Sized> SessionStore for std::sync::Arc<T> {
101 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>> {
102 (**self).get_session(session_id).await
103 }
104}
105
106#[async_trait]
108pub trait SessionMutator: Send + Sync {
109 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session>;
111}
112
113#[async_trait]
114impl<T: SessionMutator + ?Sized> SessionMutator for std::sync::Arc<T> {
115 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session> {
116 (**self).update_session_title(session_id, title).await
117 }
118}
119
120#[derive(Debug, Clone)]
126pub struct ResolvedModel {
127 pub model: String,
129 pub provider_type: DriverId,
131 pub api_key: Option<String>,
133 pub base_url: Option<String>,
135 pub provider_metadata: Option<crate::llm_driver_registry::ProviderMetadata>,
138}
139
140#[async_trait]
150pub trait ProviderStore: Send + Sync {
151 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>>;
156
157 async fn get_default_model(&self) -> Result<Option<ResolvedModel>>;
161}
162
163#[async_trait]
164impl<T: ProviderStore + ?Sized> ProviderStore for std::sync::Arc<T> {
165 async fn get_resolved_model(&self, model_id: ModelId) -> Result<Option<ResolvedModel>> {
166 (**self).get_resolved_model(model_id).await
167 }
168
169 async fn get_default_model(&self) -> Result<Option<ResolvedModel>> {
170 (**self).get_default_model().await
171 }
172}
173
174#[derive(Debug, Clone)]
180pub struct StoredImageInfo {
181 pub id: ImageId,
182 pub filename: String,
183 pub content_type: String,
184 pub size_bytes: i64,
185 pub metadata: serde_json::Value,
186 pub created_at: DateTime<Utc>,
187}
188
189#[derive(Debug, Clone)]
191pub struct StoredImage {
192 pub info: StoredImageInfo,
193 pub data: Vec<u8>,
194}
195
196#[derive(Debug, Clone)]
198pub struct CreateStoredImage {
199 pub filename: String,
200 pub content_type: String,
201 pub data: Vec<u8>,
202 pub metadata: serde_json::Value,
203}
204
205#[async_trait]
206pub trait ImageArtifactStore: Send + Sync {
207 async fn create_image(&self, input: CreateStoredImage) -> Result<StoredImageInfo>;
209
210 async fn get_image(&self, image_id: ImageId) -> Result<Option<StoredImage>>;
212
213 async fn get_image_info(&self, image_id: ImageId) -> Result<Option<StoredImageInfo>>;
215}
216
217#[derive(Debug, Clone)]
223pub struct ProviderCredentials {
224 pub api_key: String,
225 pub base_url: Option<String>,
226}
227
228#[async_trait]
229pub trait ProviderCredentialStore: Send + Sync {
230 async fn get_default_provider_credentials(
235 &self,
236 provider_type: &str,
237 ) -> Result<Option<ProviderCredentials>>;
238}
239
240#[async_trait]
251pub trait ToolExecutor: Send + Sync {
252 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult>;
257
258 async fn execute_with_context(
263 &self,
264 tool_call: &ToolCall,
265 tool_def: &ToolDefinition,
266 _context: &ToolContext,
267 ) -> Result<ToolResult> {
268 self.execute(tool_call, tool_def).await
270 }
271
272 async fn execute_batch(
274 &self,
275 tool_calls: &[ToolCall],
276 tool_defs: &[ToolDefinition],
277 ) -> Result<Vec<ToolResult>> {
278 let mut results = Vec::with_capacity(tool_calls.len());
279
280 let tool_map = build_tool_map(tool_defs);
281
282 for tool_call in tool_calls {
283 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
284 crate::error::AgentLoopError::tool(format!(
285 "Tool definition not found: {}",
286 tool_call.name
287 ))
288 })?;
289
290 results.push(self.execute(tool_call, tool_def).await?);
291 }
292
293 Ok(results)
294 }
295
296 async fn execute_parallel(
298 &self,
299 tool_calls: &[ToolCall],
300 tool_defs: &[ToolDefinition],
301 ) -> Result<Vec<ToolResult>>
302 where
303 Self: Sized,
304 {
305 use futures::future::join_all;
306
307 let tool_map = build_tool_map(tool_defs);
308
309 let futures: Vec<_> = tool_calls
310 .iter()
311 .map(|tool_call| async {
312 let tool_def = tool_map.get(tool_call.name.as_str()).ok_or_else(|| {
313 crate::error::AgentLoopError::tool(format!(
314 "Tool definition not found: {}",
315 tool_call.name
316 ))
317 })?;
318 self.execute(tool_call, tool_def).await
319 })
320 .collect();
321
322 let results = join_all(futures).await;
323 results.into_iter().collect()
324 }
325}
326
327#[async_trait]
331impl ToolExecutor for std::sync::Arc<dyn ToolExecutor> {
332 async fn execute(&self, tool_call: &ToolCall, tool_def: &ToolDefinition) -> Result<ToolResult> {
333 (**self).execute(tool_call, tool_def).await
334 }
335
336 async fn execute_with_context(
337 &self,
338 tool_call: &ToolCall,
339 tool_def: &ToolDefinition,
340 context: &ToolContext,
341 ) -> Result<ToolResult> {
342 (**self)
343 .execute_with_context(tool_call, tool_def, context)
344 .await
345 }
346
347 async fn execute_batch(
348 &self,
349 tool_calls: &[ToolCall],
350 tool_defs: &[ToolDefinition],
351 ) -> Result<Vec<ToolResult>> {
352 (**self).execute_batch(tool_calls, tool_defs).await
353 }
354}
355
356#[async_trait]
368pub trait SessionFileSystem: Send + Sync {
369 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>>;
371
372 async fn write_file(
374 &self,
375 session_id: SessionId,
376 path: &str,
377 content: &str,
378 encoding: &str,
379 ) -> Result<SessionFile>;
380
381 async fn write_file_if_content_matches(
386 &self,
387 session_id: SessionId,
388 path: &str,
389 expected_content: &str,
390 expected_encoding: &str,
391 content: &str,
392 encoding: &str,
393 ) -> Result<Option<SessionFile>> {
394 let Some(existing) = self.read_file(session_id, path).await? else {
395 return Ok(None);
396 };
397
398 if existing.is_directory {
399 return Ok(None);
400 }
401
402 let current_content = existing.content.unwrap_or_default();
403 if current_content != expected_content || existing.encoding != expected_encoding {
404 return Ok(None);
405 }
406
407 self.write_file(session_id, path, content, encoding)
408 .await
409 .map(Some)
410 }
411
412 async fn delete_file(&self, session_id: SessionId, path: &str, recursive: bool)
414 -> Result<bool>;
415
416 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>>;
418
419 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>>;
421
422 async fn grep_files(
424 &self,
425 session_id: SessionId,
426 pattern: &str,
427 path_pattern: Option<&str>,
428 ) -> Result<Vec<GrepMatch>>;
429
430 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo>;
432
433 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
435 if file.is_readonly {
436 return Err(crate::error::AgentLoopError::store(
437 "read-only initial files require a SessionFileSystem-specific seed implementation",
438 ));
439 }
440 self.write_file(session_id, &file.path, &file.content, &file.encoding)
441 .await?;
442 Ok(())
443 }
444}
445
446pub struct WorkspaceScopedFileSystem {
456 inner: Arc<dyn SessionFileSystem>,
457 key: SessionId,
458}
459
460impl WorkspaceScopedFileSystem {
461 pub fn wrap(
463 inner: Arc<dyn SessionFileSystem>,
464 workspace_id: WorkspaceId,
465 ) -> Arc<dyn SessionFileSystem> {
466 Arc::new(Self {
467 inner,
468 key: SessionId::from_uuid(workspace_id.uuid()),
469 })
470 }
471}
472
473#[async_trait]
474impl SessionFileSystem for WorkspaceScopedFileSystem {
475 async fn read_file(&self, _session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
476 self.inner.read_file(self.key, path).await
477 }
478 async fn write_file(
479 &self,
480 _session_id: SessionId,
481 path: &str,
482 content: &str,
483 encoding: &str,
484 ) -> Result<SessionFile> {
485 self.inner
486 .write_file(self.key, path, content, encoding)
487 .await
488 }
489 async fn write_file_if_content_matches(
490 &self,
491 _session_id: SessionId,
492 path: &str,
493 expected_content: &str,
494 expected_encoding: &str,
495 content: &str,
496 encoding: &str,
497 ) -> Result<Option<SessionFile>> {
498 self.inner
499 .write_file_if_content_matches(
500 self.key,
501 path,
502 expected_content,
503 expected_encoding,
504 content,
505 encoding,
506 )
507 .await
508 }
509 async fn delete_file(
510 &self,
511 _session_id: SessionId,
512 path: &str,
513 recursive: bool,
514 ) -> Result<bool> {
515 self.inner.delete_file(self.key, path, recursive).await
516 }
517 async fn list_directory(&self, _session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
518 self.inner.list_directory(self.key, path).await
519 }
520 async fn stat_file(&self, _session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
521 self.inner.stat_file(self.key, path).await
522 }
523 async fn grep_files(
524 &self,
525 _session_id: SessionId,
526 pattern: &str,
527 path_pattern: Option<&str>,
528 ) -> Result<Vec<GrepMatch>> {
529 self.inner.grep_files(self.key, pattern, path_pattern).await
530 }
531 async fn create_directory(&self, _session_id: SessionId, path: &str) -> Result<FileInfo> {
532 self.inner.create_directory(self.key, path).await
533 }
534 async fn seed_initial_file(&self, _session_id: SessionId, file: &InitialFile) -> Result<()> {
535 self.inner.seed_initial_file(self.key, file).await
536 }
537}
538
539#[async_trait]
540impl<T: SessionFileSystem + ?Sized> SessionFileSystem for std::sync::Arc<T> {
541 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
542 (**self).read_file(session_id, path).await
543 }
544
545 async fn write_file(
546 &self,
547 session_id: SessionId,
548 path: &str,
549 content: &str,
550 encoding: &str,
551 ) -> Result<SessionFile> {
552 (**self)
553 .write_file(session_id, path, content, encoding)
554 .await
555 }
556
557 async fn write_file_if_content_matches(
558 &self,
559 session_id: SessionId,
560 path: &str,
561 expected_content: &str,
562 expected_encoding: &str,
563 content: &str,
564 encoding: &str,
565 ) -> Result<Option<SessionFile>> {
566 (**self)
567 .write_file_if_content_matches(
568 session_id,
569 path,
570 expected_content,
571 expected_encoding,
572 content,
573 encoding,
574 )
575 .await
576 }
577
578 async fn delete_file(
579 &self,
580 session_id: SessionId,
581 path: &str,
582 recursive: bool,
583 ) -> Result<bool> {
584 (**self).delete_file(session_id, path, recursive).await
585 }
586
587 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
588 (**self).list_directory(session_id, path).await
589 }
590
591 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
592 (**self).stat_file(session_id, path).await
593 }
594
595 async fn grep_files(
596 &self,
597 session_id: SessionId,
598 pattern: &str,
599 path_pattern: Option<&str>,
600 ) -> Result<Vec<GrepMatch>> {
601 (**self).grep_files(session_id, pattern, path_pattern).await
602 }
603
604 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
605 (**self).create_directory(session_id, path).await
606 }
607
608 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
609 (**self).seed_initial_file(session_id, file).await
610 }
611}
612
613pub use SessionFileSystem as SessionFileStore;
615
616#[derive(Clone, Default)]
622pub struct SessionFileSystemFactoryContext {
623 values: Arc<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>,
624}
625
626impl SessionFileSystemFactoryContext {
627 pub fn new() -> Self {
628 Self::default()
629 }
630
631 pub fn with<T: Any + Send + Sync>(mut self, value: Arc<T>) -> Self {
632 let values = Arc::make_mut(&mut self.values);
633 values.insert(TypeId::of::<T>(), value);
634 self
635 }
636
637 pub fn get<T: Any + Send + Sync>(&self) -> Option<Arc<T>> {
638 self.values
639 .get(&TypeId::of::<T>())
640 .and_then(|value| value.clone().downcast::<T>().ok())
641 }
642}
643
644#[async_trait]
646pub trait SessionFileSystemFactory: Send + Sync {
647 fn name(&self) -> &'static str {
649 "SessionFileSystemFactory"
650 }
651
652 fn is_disabled(&self) -> bool {
655 false
656 }
657
658 async fn create_session_file_system(
660 &self,
661 context: SessionFileSystemFactoryContext,
662 ) -> Result<Arc<dyn SessionFileSystem>>;
663}
664
665#[derive(Debug, Clone, Default)]
667pub struct DisabledSessionFileSystemFactory;
668
669#[async_trait]
670impl SessionFileSystemFactory for DisabledSessionFileSystemFactory {
671 fn name(&self) -> &'static str {
672 "DisabledSessionFileSystemFactory"
673 }
674
675 fn is_disabled(&self) -> bool {
676 true
677 }
678
679 async fn create_session_file_system(
680 &self,
681 _context: SessionFileSystemFactoryContext,
682 ) -> Result<Arc<dyn SessionFileSystem>> {
683 Err(crate::error::AgentLoopError::config(
684 "session filesystem is disabled",
685 ))
686 }
687}
688
689#[derive(Debug, Clone)]
695pub struct KeyInfo {
696 pub key: String,
697 pub created_at: chrono::DateTime<chrono::Utc>,
698 pub updated_at: chrono::DateTime<chrono::Utc>,
699}
700
701#[derive(Debug, Clone)]
703pub struct SecretInfo {
704 pub name: String,
705 pub created_at: chrono::DateTime<chrono::Utc>,
706 pub updated_at: chrono::DateTime<chrono::Utc>,
707}
708
709#[async_trait]
719pub trait SessionStorageStore: Send + Sync {
720 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()>;
724
725 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>>;
727
728 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool>;
730
731 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>>;
733
734 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()>;
738
739 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>>;
741
742 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool>;
744
745 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>>;
747}
748
749use crate::session_schedule::SessionSchedule;
754use crate::typed_id::ScheduleId;
755
756#[async_trait]
760pub trait SessionScheduleStore: Send + Sync {
761 async fn create_schedule(
763 &self,
764 session_id: SessionId,
765 description: String,
766 cron_expression: Option<String>,
767 scheduled_at: Option<chrono::DateTime<chrono::Utc>>,
768 timezone: String,
769 ) -> Result<SessionSchedule>;
770
771 async fn cancel_schedule(
773 &self,
774 session_id: SessionId,
775 schedule_id: ScheduleId,
776 ) -> Result<SessionSchedule>;
777
778 async fn list_schedules(&self, session_id: SessionId) -> Result<Vec<SessionSchedule>>;
780
781 async fn count_active_schedules(&self, session_id: SessionId) -> Result<u32>;
783}
784
785#[async_trait]
795pub trait SessionResourceRegistry: Send + Sync {
796 async fn register(
798 &self,
799 entry: crate::session_resource::RegisterSessionResource,
800 ) -> Result<crate::session_resource::SessionResourceEntry>;
801
802 async fn update_status(
804 &self,
805 session_id: SessionId,
806 resource_id: &str,
807 status: crate::session_resource::SessionResourceStatus,
808 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
809
810 async fn get(
812 &self,
813 session_id: SessionId,
814 resource_id: &str,
815 ) -> Result<Option<crate::session_resource::SessionResourceEntry>>;
816
817 async fn list(
819 &self,
820 session_id: SessionId,
821 filter: Option<&crate::session_resource::SessionResourceFilter>,
822 ) -> Result<Vec<crate::session_resource::SessionResourceEntry>>;
823
824 async fn deregister(&self, session_id: SessionId, resource_id: &str) -> Result<bool>;
826}
827
828#[async_trait]
838pub trait LeasedResourceStore: Send + Sync {
839 async fn upsert_resource(&self, input: UpsertLeasedResource) -> Result<LeasedResource>;
845
846 async fn release_resource(
852 &self,
853 session_id: SessionId,
854 provider: &str,
855 resource_type: &str,
856 external_id: &str,
857 ) -> Result<Option<LeasedResource>>;
858
859 async fn list_resources(&self, session_id: SessionId) -> Result<Vec<LeasedResource>>;
864}
865
866pub type SessionSqlDbStoreRef = Arc<dyn crate::session_sqldb::SessionSqlDbStore>;
872
873#[async_trait]
878pub trait UserConnectionResolver: Send + Sync {
879 async fn get_connection_token(
882 &self,
883 session_id: SessionId,
884 provider: &str,
885 ) -> Result<Option<String>>;
886
887 async fn get_connection_user(
892 &self,
893 _session_id: SessionId,
894 _provider: &str,
895 ) -> Result<Option<Uuid>> {
896 Ok(None)
897 }
898
899 async fn get_connection_token_for_user(
904 &self,
905 _user_id: Uuid,
906 _provider: &str,
907 ) -> Result<Option<String>> {
908 Ok(None)
909 }
910
911 async fn get_connection_metadata(
914 &self,
915 _session_id: SessionId,
916 _provider: &str,
917 ) -> Result<Option<serde_json::Value>> {
918 Ok(None)
919 }
920}
921
922#[async_trait]
932pub trait BudgetChecker: Send + Sync {
933 async fn check_budgets(&self, session_id: &str) -> Result<crate::budget::BudgetToolResponse>;
935}
936
937#[async_trait]
946pub trait PaymentAuthority: Send + Sync {
947 async fn execute_machine_payment(
948 &self,
949 session_id: SessionId,
950 request: crate::payment::MachinePaymentRequest,
951 ) -> Result<crate::payment::MachinePaymentResponse>;
952}
953
954#[async_trait]
964pub trait OutboundToolRateLimiter: Send + Sync {
965 async fn check_org(&self, org_id: &crate::typed_id::OrgId) -> bool;
967}
968
969#[derive(Debug)]
975pub enum ToolCallClaimResult {
976 Claimed { claim_token: uuid::Uuid },
979 AlreadySettled {
981 result_json: serde_json::Value,
982 args_fingerprint: String,
983 },
984 AlreadyRunning { args_fingerprint: String },
989 DeterminismViolation {
993 stored_fingerprint: String,
994 current_fingerprint: String,
995 },
996}
997
998#[derive(Debug, Clone)]
1000pub enum DurableToolCallStatus {
1001 Settled { result_json: serde_json::Value },
1003 Interrupted {
1005 result_json: Option<serde_json::Value>,
1006 },
1007 Running,
1009}
1010
1011#[async_trait]
1016pub trait DurableToolResultStore: Send + Sync + 'static {
1017 async fn try_claim_tool_call(
1025 &self,
1026 turn_id: &str,
1027 tool_call_id: &str,
1028 tool_name: &str,
1029 args_fingerprint: &str,
1030 ) -> Result<ToolCallClaimResult>;
1031
1032 async fn settle_tool_call(
1038 &self,
1039 turn_id: &str,
1040 tool_call_id: &str,
1041 result_json: serde_json::Value,
1042 status: &str,
1043 claim_token: uuid::Uuid,
1044 ) -> Result<bool>;
1045
1046 async fn get_tool_call_status(
1051 &self,
1052 turn_id: &str,
1053 tool_call_id: &str,
1054 ) -> Result<Option<DurableToolCallStatus>>;
1055}
1056
1057pub struct NoopDurableToolResultStore;
1060
1061#[async_trait]
1062impl DurableToolResultStore for NoopDurableToolResultStore {
1063 async fn try_claim_tool_call(
1064 &self,
1065 _turn_id: &str,
1066 _tool_call_id: &str,
1067 _tool_name: &str,
1068 _args_fingerprint: &str,
1069 ) -> Result<ToolCallClaimResult> {
1070 Ok(ToolCallClaimResult::Claimed {
1071 claim_token: uuid::Uuid::new_v4(),
1072 })
1073 }
1074
1075 async fn settle_tool_call(
1076 &self,
1077 _turn_id: &str,
1078 _tool_call_id: &str,
1079 _result_json: serde_json::Value,
1080 _status: &str,
1081 _claim_token: uuid::Uuid,
1082 ) -> Result<bool> {
1083 Ok(true)
1084 }
1085
1086 async fn get_tool_call_status(
1087 &self,
1088 _turn_id: &str,
1089 _tool_call_id: &str,
1090 ) -> Result<Option<DurableToolCallStatus>> {
1091 Ok(None)
1092 }
1093}
1094
1095#[derive(Debug, Clone)]
1101pub struct StreamProgress {
1102 pub accumulated_len: usize,
1104 pub last_delta_at: u64,
1106}
1107
1108#[async_trait]
1114pub trait StreamHeartbeater: Send + Sync {
1115 async fn heartbeat(&self, progress: StreamProgress);
1121}
1122
1123pub struct NoopStreamHeartbeater;
1125
1126#[async_trait]
1127impl StreamHeartbeater for NoopStreamHeartbeater {
1128 async fn heartbeat(&self, _progress: StreamProgress) {}
1129}
1130
1131#[derive(Debug, Clone)]
1137pub struct PartialStreamState {
1138 pub accumulated: String,
1141}
1142
1143#[async_trait]
1151pub trait PartialStreamStore: Send + Sync {
1152 async fn get_partial_stream(
1155 &self,
1156 session_id: SessionId,
1157 turn_id: &str,
1158 ) -> Result<Option<PartialStreamState>>;
1159}
1160
1161pub struct NoopPartialStreamStore;
1163
1164#[async_trait]
1165impl PartialStreamStore for NoopPartialStreamStore {
1166 async fn get_partial_stream(
1167 &self,
1168 _session_id: SessionId,
1169 _turn_id: &str,
1170 ) -> Result<Option<PartialStreamState>> {
1171 Ok(None)
1172 }
1173}
1174
1175#[derive(Clone)]
1184pub struct ToolContext {
1185 pub session_id: SessionId,
1187 pub workspace_id: WorkspaceId,
1194
1195 pub file_store: Option<Arc<dyn SessionFileSystem>>,
1197
1198 pub storage_store: Option<Arc<dyn SessionStorageStore>>,
1200
1201 pub image_store: Option<Arc<dyn ImageArtifactStore>>,
1203
1204 pub provider_credential_store: Option<Arc<dyn ProviderCredentialStore>>,
1206
1207 pub utility_llm_service: Option<Arc<dyn crate::UtilityLlmService>>,
1209
1210 pub egress_service: Option<Arc<dyn crate::EgressService>>,
1212
1213 pub sqldb_store: Option<SessionSqlDbStoreRef>,
1215
1216 pub message_retriever: Option<Arc<dyn crate::message_retriever::MessageRetriever>>,
1218
1219 pub session_store: Option<Arc<dyn SessionStore>>,
1221
1222 pub session_mutator: Option<Arc<dyn SessionMutator>>,
1224
1225 pub agent_store: Option<Arc<dyn AgentStore>>,
1227
1228 pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
1230
1231 pub schedule_store: Option<Arc<dyn SessionScheduleStore>>,
1233
1234 pub platform_store: Option<Arc<dyn crate::platform_store::PlatformStore>>,
1236 pub leased_resource_store: Option<Arc<dyn LeasedResourceStore>>,
1238
1239 pub session_resource_registry: Option<Arc<dyn SessionResourceRegistry>>,
1241
1242 pub session_task_registry: Option<Arc<dyn crate::session_task::SessionTaskRegistry>>,
1245
1246 pub event_emitter: Option<Arc<dyn EventEmitter>>,
1249
1250 pub event_context: Option<crate::events::EventContext>,
1253
1254 pub tool_call_id: Option<String>,
1257 pub capability_registry: Option<crate::capabilities::CapabilityRegistry>,
1259
1260 pub tool_registry: Option<Arc<crate::tools::ToolRegistry>>,
1263
1264 pub visible_tool_names: Option<Arc<HashSet<String>>>,
1268
1269 pub org_id: Option<crate::typed_id::OrgId>,
1271
1272 pub network_access: Option<crate::network_access::NetworkAccessList>,
1275
1276 pub locale: Option<String>,
1280
1281 pub budget_checker: Option<Arc<dyn BudgetChecker>>,
1283
1284 pub payment_authority: Option<Arc<dyn PaymentAuthority>>,
1286
1287 pub subagent_spawn_store: Option<Arc<dyn SubagentSpawnStore>>,
1291}
1292
1293impl ToolContext {
1294 pub fn workspace_fs_key(&self) -> SessionId {
1299 SessionId::from_uuid(self.workspace_id.uuid())
1300 }
1301
1302 pub fn with_workspace_id(mut self, workspace_id: WorkspaceId) -> Self {
1304 self.workspace_id = workspace_id;
1305 self
1306 }
1307
1308 pub fn new(session_id: SessionId) -> Self {
1310 Self {
1311 session_id,
1312 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1313 file_store: None,
1314 storage_store: None,
1315 image_store: None,
1316 provider_credential_store: None,
1317 utility_llm_service: None,
1318 egress_service: None,
1319 sqldb_store: None,
1320 message_retriever: None,
1321 session_store: None,
1322 session_mutator: None,
1323 agent_store: None,
1324 connection_resolver: None,
1325 schedule_store: None,
1326 platform_store: None,
1327 leased_resource_store: None,
1328 session_resource_registry: None,
1329 session_task_registry: None,
1330 event_emitter: None,
1331 event_context: None,
1332 tool_call_id: None,
1333 capability_registry: None,
1334 tool_registry: None,
1335 visible_tool_names: None,
1336 org_id: None,
1337 network_access: None,
1338 locale: None,
1339 budget_checker: None,
1340 payment_authority: None,
1341 subagent_spawn_store: None,
1342 }
1343 }
1344
1345 pub fn with_file_store(session_id: SessionId, file_store: Arc<dyn SessionFileSystem>) -> Self {
1347 Self {
1348 session_id,
1349 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1350 file_store: Some(file_store),
1351 storage_store: None,
1352 image_store: None,
1353 provider_credential_store: None,
1354 utility_llm_service: None,
1355 egress_service: None,
1356 sqldb_store: None,
1357 message_retriever: None,
1358 session_store: None,
1359 session_mutator: None,
1360 agent_store: None,
1361 connection_resolver: None,
1362 schedule_store: None,
1363 platform_store: None,
1364 leased_resource_store: None,
1365 session_resource_registry: None,
1366 session_task_registry: None,
1367 event_emitter: None,
1368 event_context: None,
1369 tool_call_id: None,
1370 capability_registry: None,
1371 tool_registry: None,
1372 visible_tool_names: None,
1373 org_id: None,
1374 network_access: None,
1375 locale: None,
1376 budget_checker: None,
1377 payment_authority: None,
1378 subagent_spawn_store: None,
1379 }
1380 }
1381
1382 pub fn with_storage_store(
1384 session_id: SessionId,
1385 storage_store: Arc<dyn SessionStorageStore>,
1386 ) -> Self {
1387 Self {
1388 session_id,
1389 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1390 file_store: None,
1391 storage_store: Some(storage_store),
1392 image_store: None,
1393 provider_credential_store: None,
1394 utility_llm_service: None,
1395 egress_service: None,
1396 sqldb_store: None,
1397 message_retriever: None,
1398 session_store: None,
1399 session_mutator: None,
1400 agent_store: None,
1401 connection_resolver: None,
1402 schedule_store: None,
1403 platform_store: None,
1404 leased_resource_store: None,
1405 session_resource_registry: None,
1406 session_task_registry: None,
1407 event_emitter: None,
1408 event_context: None,
1409 tool_call_id: None,
1410 capability_registry: None,
1411 tool_registry: None,
1412 visible_tool_names: None,
1413 org_id: None,
1414 network_access: None,
1415 locale: None,
1416 budget_checker: None,
1417 payment_authority: None,
1418 subagent_spawn_store: None,
1419 }
1420 }
1421
1422 pub fn with_stores(
1424 session_id: SessionId,
1425 file_store: Arc<dyn SessionFileSystem>,
1426 storage_store: Arc<dyn SessionStorageStore>,
1427 ) -> Self {
1428 Self {
1429 session_id,
1430 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1431 file_store: Some(file_store),
1432 storage_store: Some(storage_store),
1433 sqldb_store: None,
1434 image_store: None,
1435 provider_credential_store: None,
1436 utility_llm_service: None,
1437 egress_service: None,
1438 message_retriever: None,
1439 session_store: None,
1440 session_mutator: None,
1441 agent_store: None,
1442 connection_resolver: None,
1443 schedule_store: None,
1444 platform_store: None,
1445 leased_resource_store: None,
1446 session_resource_registry: None,
1447 session_task_registry: None,
1448 event_emitter: None,
1449 event_context: None,
1450 tool_call_id: None,
1451 capability_registry: None,
1452 tool_registry: None,
1453 visible_tool_names: None,
1454 org_id: None,
1455 network_access: None,
1456 locale: None,
1457 budget_checker: None,
1458 payment_authority: None,
1459 subagent_spawn_store: None,
1460 }
1461 }
1462
1463 pub fn with_sqldb_store(mut self, sqldb_store: SessionSqlDbStoreRef) -> Self {
1465 self.sqldb_store = Some(sqldb_store);
1466 self
1467 }
1468
1469 pub fn with_message_retriever(
1471 mut self,
1472 retriever: Arc<dyn crate::message_retriever::MessageRetriever>,
1473 ) -> Self {
1474 self.message_retriever = Some(retriever);
1475 self
1476 }
1477
1478 pub fn with_session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
1480 self.session_store = Some(store);
1481 self
1482 }
1483
1484 pub fn with_session_mutator(mut self, mutator: Arc<dyn SessionMutator>) -> Self {
1486 self.session_mutator = Some(mutator);
1487 self
1488 }
1489
1490 pub fn with_agent_store(mut self, store: Arc<dyn AgentStore>) -> Self {
1492 self.agent_store = Some(store);
1493 self
1494 }
1495
1496 pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
1498 self.connection_resolver = Some(resolver);
1499 self
1500 }
1501
1502 pub fn with_image_store(
1504 session_id: SessionId,
1505 image_store: Arc<dyn ImageArtifactStore>,
1506 ) -> Self {
1507 Self {
1508 session_id,
1509 workspace_id: WorkspaceId::from_uuid(session_id.uuid()),
1510 file_store: None,
1511 storage_store: None,
1512 image_store: Some(image_store),
1513 provider_credential_store: None,
1514 utility_llm_service: 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 leased_resource_store: None,
1525 session_resource_registry: None,
1526 session_task_registry: None,
1527 event_emitter: None,
1528 event_context: None,
1529 tool_call_id: None,
1530 capability_registry: None,
1531 tool_registry: None,
1532 visible_tool_names: None,
1533 org_id: None,
1534 network_access: None,
1535 locale: None,
1536 budget_checker: None,
1537 payment_authority: None,
1538 subagent_spawn_store: None,
1539 }
1540 }
1541
1542 pub fn with_provider_credential_store(
1544 mut self,
1545 store: Arc<dyn ProviderCredentialStore>,
1546 ) -> Self {
1547 self.provider_credential_store = Some(store);
1548 self
1549 }
1550
1551 pub fn with_utility_llm_service(mut self, service: Arc<dyn crate::UtilityLlmService>) -> Self {
1553 self.utility_llm_service = Some(service);
1554 self
1555 }
1556
1557 pub fn with_egress_service(mut self, service: Arc<dyn crate::EgressService>) -> Self {
1559 self.egress_service = Some(service);
1560 self
1561 }
1562
1563 pub fn with_egress_service_opt(
1566 mut self,
1567 service: Option<Arc<dyn crate::EgressService>>,
1568 ) -> Self {
1569 if let Some(service) = service {
1570 self.egress_service = Some(service);
1571 }
1572 self
1573 }
1574
1575 pub fn with_storage_store_arc(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
1577 self.storage_store = Some(store);
1578 self
1579 }
1580
1581 pub fn with_schedule_store(mut self, store: Arc<dyn SessionScheduleStore>) -> Self {
1583 self.schedule_store = Some(store);
1584 self
1585 }
1586
1587 pub fn with_platform_store(
1589 mut self,
1590 store: Arc<dyn crate::platform_store::PlatformStore>,
1591 ) -> Self {
1592 self.platform_store = Some(store);
1593 self
1594 }
1595
1596 pub fn with_leased_resource_store(mut self, store: Arc<dyn LeasedResourceStore>) -> Self {
1598 self.leased_resource_store = Some(store);
1599 self
1600 }
1601
1602 pub fn with_session_resource_registry(
1604 mut self,
1605 registry: Arc<dyn SessionResourceRegistry>,
1606 ) -> Self {
1607 self.session_resource_registry = Some(registry);
1608 self
1609 }
1610
1611 pub fn with_session_task_registry(
1613 mut self,
1614 registry: Arc<dyn crate::session_task::SessionTaskRegistry>,
1615 ) -> Self {
1616 self.session_task_registry = Some(registry);
1617 self
1618 }
1619
1620 pub fn with_org_id(mut self, org_id: crate::typed_id::OrgId) -> Self {
1622 self.org_id = Some(org_id);
1623 self
1624 }
1625
1626 pub fn with_tool_registry(mut self, registry: Arc<crate::tools::ToolRegistry>) -> Self {
1628 self.tool_registry = Some(registry);
1629 self
1630 }
1631
1632 pub fn with_visible_tool_names(mut self, names: Arc<HashSet<String>>) -> Self {
1634 self.visible_tool_names = Some(names);
1635 self
1636 }
1637
1638 pub fn with_network_access(
1640 mut self,
1641 network_access: Option<crate::network_access::NetworkAccessList>,
1642 ) -> Self {
1643 self.network_access = network_access;
1644 self
1645 }
1646
1647 pub fn with_payment_authority(mut self, authority: Arc<dyn PaymentAuthority>) -> Self {
1649 self.payment_authority = Some(authority);
1650 self
1651 }
1652
1653 pub fn with_subagent_spawn_store(mut self, store: Arc<dyn SubagentSpawnStore>) -> Self {
1655 self.subagent_spawn_store = Some(store);
1656 self
1657 }
1658
1659 pub async fn emit_progress(&self, tool_name: &str, message: &str) {
1664 let (Some(emitter), Some(ctx), Some(call_id)) =
1665 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1666 else {
1667 return;
1668 };
1669 if let Err(e) = emitter
1670 .emit(EventRequest::new(
1671 self.session_id,
1672 ctx.clone(),
1673 crate::events::ToolProgressData {
1674 tool_call_id: call_id.clone(),
1675 tool_name: tool_name.to_string(),
1676 message: message.to_string(),
1677 display_name: None,
1678 },
1679 ))
1680 .await
1681 {
1682 tracing::debug!(
1683 tool_call_id = call_id,
1684 tool_name,
1685 error = %e,
1686 "Failed to emit tool.progress event"
1687 );
1688 }
1689 }
1690
1691 pub async fn emit_tool_output(&self, tool_name: &str, delta: &str, stream: &str) {
1696 let (Some(emitter), Some(ctx), Some(call_id)) =
1697 (&self.event_emitter, &self.event_context, &self.tool_call_id)
1698 else {
1699 return;
1700 };
1701 if let Err(e) = emitter
1702 .emit(EventRequest::new(
1703 self.session_id,
1704 ctx.clone(),
1705 crate::events::ToolOutputDeltaData {
1706 tool_call_id: call_id.clone(),
1707 tool_name: tool_name.to_string(),
1708 delta: delta.to_string(),
1709 stream: stream.to_string(),
1710 },
1711 ))
1712 .await
1713 {
1714 tracing::debug!(
1715 tool_call_id = call_id,
1716 tool_name,
1717 error = %e,
1718 "Failed to emit tool.output.delta event"
1719 );
1720 }
1721 }
1722}
1723
1724impl std::fmt::Debug for ToolContext {
1725 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1726 f.debug_struct("ToolContext")
1727 .field("session_id", &self.session_id)
1728 .field("file_store", &self.file_store.is_some())
1729 .field("storage_store", &self.storage_store.is_some())
1730 .field("image_store", &self.image_store.is_some())
1731 .field(
1732 "provider_credential_store",
1733 &self.provider_credential_store.is_some(),
1734 )
1735 .field("utility_llm_service", &self.utility_llm_service.is_some())
1736 .field("egress_service", &self.egress_service.is_some())
1737 .field("sqldb_store", &self.sqldb_store.is_some())
1738 .field("message_retriever", &self.message_retriever.is_some())
1739 .field("session_store", &self.session_store.is_some())
1740 .field("session_mutator", &self.session_mutator.is_some())
1741 .field("agent_store", &self.agent_store.is_some())
1742 .field("connection_resolver", &self.connection_resolver.is_some())
1743 .field("schedule_store", &self.schedule_store.is_some())
1744 .field("platform_store", &self.platform_store.is_some())
1745 .field(
1746 "leased_resource_store",
1747 &self.leased_resource_store.is_some(),
1748 )
1749 .field("event_emitter", &self.event_emitter.is_some())
1750 .field("tool_registry", &self.tool_registry.is_some())
1751 .field("payment_authority", &self.payment_authority.is_some())
1752 .field("subagent_spawn_store", &self.subagent_spawn_store.is_some())
1753 .field("org_id", &self.org_id)
1754 .finish()
1755 }
1756}
1757
1758use crate::events::{Event, EventRequest};
1763
1764#[async_trait]
1775pub trait EventEmitter: Send + Sync {
1776 async fn emit(&self, request: EventRequest) -> Result<Event>;
1781}
1782
1783#[async_trait]
1785impl<E: EventEmitter + ?Sized> EventEmitter for Arc<E> {
1786 async fn emit(&self, request: EventRequest) -> Result<Event> {
1787 (**self).emit(request).await
1788 }
1789}
1790
1791#[derive(Debug, Clone, Default)]
1795pub struct NoopEventEmitter;
1796
1797#[async_trait]
1798impl EventEmitter for NoopEventEmitter {
1799 async fn emit(&self, request: EventRequest) -> Result<Event> {
1800 Ok(request.into_event(crate::typed_id::EventId::new(), 0))
1802 }
1803}
1804
1805#[derive(Debug, Clone)]
1818pub struct ResolvedImage {
1819 pub base64: String,
1821 pub media_type: String,
1823}
1824
1825impl ResolvedImage {
1826 pub fn new(base64: impl Into<String>, media_type: impl Into<String>) -> Self {
1828 Self {
1829 base64: base64.into(),
1830 media_type: media_type.into(),
1831 }
1832 }
1833
1834 pub fn to_data_url(&self) -> String {
1838 format!("data:{};base64,{}", self.media_type, self.base64)
1839 }
1840}
1841
1842#[async_trait]
1875pub trait ImageResolver: Send + Sync {
1876 async fn resolve_image(&self, image_id: Uuid) -> Result<Option<ResolvedImage>>;
1880}
1881
1882#[derive(Debug)]
1888pub enum SpawnClaimResult {
1889 Claimed {
1892 spawn_handle_id: uuid::Uuid,
1893 claim_token: uuid::Uuid,
1894 },
1895 ClaimedPendingChild {
1899 spawn_handle_id: uuid::Uuid,
1900 claim_token: uuid::Uuid,
1901 },
1902 AlreadyRunning {
1905 child_session_id: crate::typed_id::SessionId,
1906 claim_token: uuid::Uuid,
1908 },
1909 AlreadySettled {
1912 child_session_id: crate::typed_id::SessionId,
1913 terminal_status: String,
1915 terminal_result: String,
1916 },
1917}
1918
1919#[async_trait]
1927pub trait SubagentSpawnStore: Send + Sync + 'static {
1928 async fn try_claim_spawn(
1933 &self,
1934 parent_session_id: crate::typed_id::SessionId,
1935 tool_call_id: &str,
1936 claim_token: uuid::Uuid,
1937 ) -> Result<SpawnClaimResult>;
1938
1939 async fn register_child_session(
1944 &self,
1945 spawn_handle_id: uuid::Uuid,
1946 claim_token: uuid::Uuid,
1947 child_session_id: crate::typed_id::SessionId,
1948 ) -> Result<()>;
1949
1950 async fn settle_spawn(
1956 &self,
1957 parent_session_id: crate::typed_id::SessionId,
1958 tool_call_id: &str,
1959 claim_token: uuid::Uuid,
1960 terminal_status: &str,
1961 terminal_result: &str,
1962 ) -> Result<()>;
1963}
1964
1965#[async_trait]
1967impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
1968 async fn try_claim_spawn(
1969 &self,
1970 parent_session_id: crate::typed_id::SessionId,
1971 tool_call_id: &str,
1972 claim_token: uuid::Uuid,
1973 ) -> Result<SpawnClaimResult> {
1974 (**self)
1975 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
1976 .await
1977 }
1978
1979 async fn register_child_session(
1980 &self,
1981 spawn_handle_id: uuid::Uuid,
1982 claim_token: uuid::Uuid,
1983 child_session_id: crate::typed_id::SessionId,
1984 ) -> Result<()> {
1985 (**self)
1986 .register_child_session(spawn_handle_id, claim_token, child_session_id)
1987 .await
1988 }
1989
1990 async fn settle_spawn(
1991 &self,
1992 parent_session_id: crate::typed_id::SessionId,
1993 tool_call_id: &str,
1994 claim_token: uuid::Uuid,
1995 terminal_status: &str,
1996 terminal_result: &str,
1997 ) -> Result<()> {
1998 (**self)
1999 .settle_spawn(
2000 parent_session_id,
2001 tool_call_id,
2002 claim_token,
2003 terminal_status,
2004 terminal_result,
2005 )
2006 .await
2007 }
2008}
2009
2010pub struct NoopSubagentSpawnStore;
2014
2015#[async_trait]
2016impl SubagentSpawnStore for NoopSubagentSpawnStore {
2017 async fn try_claim_spawn(
2018 &self,
2019 _parent_session_id: crate::typed_id::SessionId,
2020 _tool_call_id: &str,
2021 claim_token: uuid::Uuid,
2022 ) -> Result<SpawnClaimResult> {
2023 Ok(SpawnClaimResult::Claimed {
2024 spawn_handle_id: uuid::Uuid::new_v4(),
2025 claim_token,
2026 })
2027 }
2028
2029 async fn register_child_session(
2030 &self,
2031 _spawn_handle_id: uuid::Uuid,
2032 _claim_token: uuid::Uuid,
2033 _child_session_id: crate::typed_id::SessionId,
2034 ) -> Result<()> {
2035 Ok(())
2036 }
2037
2038 async fn settle_spawn(
2039 &self,
2040 _parent_session_id: crate::typed_id::SessionId,
2041 _tool_call_id: &str,
2042 _claim_token: uuid::Uuid,
2043 _terminal_status: &str,
2044 _terminal_result: &str,
2045 ) -> Result<()> {
2046 Ok(())
2047 }
2048}
2049
2050#[cfg(test)]
2055mod tests {
2056 use super::*;
2057
2058 #[test]
2059 fn test_resolved_image_new() {
2060 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2061 assert_eq!(image.base64, "SGVsbG8=");
2062 assert_eq!(image.media_type, "image/png");
2063 }
2064
2065 #[test]
2066 fn test_resolved_image_to_data_url() {
2067 let image = ResolvedImage::new("SGVsbG8=", "image/png");
2068 let data_url = image.to_data_url();
2069 assert_eq!(data_url, "data:image/png;base64,SGVsbG8=");
2070 }
2071
2072 #[test]
2073 fn test_resolved_image_jpeg() {
2074 let image = ResolvedImage::new("base64data", "image/jpeg");
2075 let data_url = image.to_data_url();
2076 assert!(data_url.starts_with("data:image/jpeg;base64,"));
2077 }
2078}