1use crate::error::Result;
4use crate::typed_id::SessionId;
5use async_trait::async_trait;
6use std::sync::Arc;
7
8pub const DEFAULT_MAX_SUBAGENT_DEPTH: u32 = 2;
11pub const DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS: u32 = 16;
12pub const DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS: u32 = 200;
13pub const DEFAULT_MAX_ACTIVE_DETACHED_TASKS: u32 = 8;
19pub const DEFAULT_MAX_TOTAL_DETACHED_TASKS: u32 = 50;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct SubagentNestingPolicy {
24 pub platform_default: u32,
25 pub org_override: Option<u32>,
26 pub agent_override: Option<u32>,
27 pub platform_default_max_active_descendant_tasks: u32,
28 pub org_override_max_active_descendant_tasks: Option<u32>,
29 pub agent_override_max_active_descendant_tasks: Option<u32>,
30 pub platform_default_max_total_descendant_tasks: u32,
31 pub org_override_max_total_descendant_tasks: Option<u32>,
32 pub agent_override_max_total_descendant_tasks: Option<u32>,
33 pub platform_default_max_active_detached_tasks: u32,
34 pub org_override_max_active_detached_tasks: Option<u32>,
35 pub agent_override_max_active_detached_tasks: Option<u32>,
36 pub platform_default_max_total_detached_tasks: u32,
37 pub org_override_max_total_detached_tasks: Option<u32>,
38 pub agent_override_max_total_detached_tasks: Option<u32>,
39}
40
41impl Default for SubagentNestingPolicy {
42 fn default() -> Self {
43 Self {
44 platform_default: DEFAULT_MAX_SUBAGENT_DEPTH,
45 org_override: None,
46 agent_override: None,
47 platform_default_max_active_descendant_tasks:
48 DEFAULT_MAX_ACTIVE_DESCENDANT_SUBAGENT_TASKS,
49 org_override_max_active_descendant_tasks: None,
50 agent_override_max_active_descendant_tasks: None,
51 platform_default_max_total_descendant_tasks:
52 DEFAULT_MAX_TOTAL_DESCENDANT_SUBAGENT_TASKS,
53 org_override_max_total_descendant_tasks: None,
54 agent_override_max_total_descendant_tasks: None,
55 platform_default_max_active_detached_tasks: DEFAULT_MAX_ACTIVE_DETACHED_TASKS,
56 org_override_max_active_detached_tasks: None,
57 agent_override_max_active_detached_tasks: None,
58 platform_default_max_total_detached_tasks: DEFAULT_MAX_TOTAL_DETACHED_TASKS,
59 org_override_max_total_detached_tasks: None,
60 agent_override_max_total_detached_tasks: None,
61 }
62 }
63}
64
65impl SubagentNestingPolicy {
66 pub fn max_subagent_depth(self) -> u32 {
67 self.agent_override
68 .or(self.org_override)
69 .unwrap_or(self.platform_default)
70 }
71
72 pub fn max_active_descendant_tasks(self) -> u32 {
73 self.agent_override_max_active_descendant_tasks
74 .or(self.org_override_max_active_descendant_tasks)
75 .unwrap_or(self.platform_default_max_active_descendant_tasks)
76 }
77
78 pub fn max_total_descendant_tasks(self) -> u32 {
79 self.agent_override_max_total_descendant_tasks
80 .or(self.org_override_max_total_descendant_tasks)
81 .unwrap_or(self.platform_default_max_total_descendant_tasks)
82 }
83
84 pub fn max_active_detached_tasks(self) -> u32 {
85 self.agent_override_max_active_detached_tasks
86 .or(self.org_override_max_active_detached_tasks)
87 .unwrap_or(self.platform_default_max_active_detached_tasks)
88 }
89
90 pub fn max_total_detached_tasks(self) -> u32 {
91 self.agent_override_max_total_detached_tasks
92 .or(self.org_override_max_total_detached_tasks)
93 .unwrap_or(self.platform_default_max_total_detached_tasks)
94 }
95
96 pub fn with_platform_default(mut self, depth: u32) -> Self {
97 self.platform_default = depth;
98 self
99 }
100
101 pub fn with_org_override(mut self, depth: Option<u32>) -> Self {
102 self.org_override = depth;
103 self
104 }
105
106 pub fn with_agent_override(mut self, depth: Option<u32>) -> Self {
107 self.agent_override = depth;
108 self
109 }
110
111 pub fn with_agent_task_caps_override(
112 mut self,
113 max_active: Option<u32>,
114 max_total: Option<u32>,
115 ) -> Self {
116 self.agent_override_max_active_descendant_tasks = max_active;
117 self.agent_override_max_total_descendant_tasks = max_total;
118 self
119 }
120
121 pub fn with_agent_detached_task_caps_override(
122 mut self,
123 max_active: Option<u32>,
124 max_total: Option<u32>,
125 ) -> Self {
126 self.agent_override_max_active_detached_tasks = max_active;
127 self.agent_override_max_total_detached_tasks = max_total;
128 self
129 }
130}
131
132#[async_trait]
138pub trait SessionCreationAuthority: Send + Sync {
139 async fn authorize_session_creation(&self, session_id: SessionId) -> Result<SessionId>;
143}
144
145#[derive(Debug)]
147pub enum SpawnClaimResult {
148 Claimed {
151 spawn_handle_id: uuid::Uuid,
152 claim_token: uuid::Uuid,
153 },
154 ClaimedPendingChild {
158 spawn_handle_id: uuid::Uuid,
159 claim_token: uuid::Uuid,
160 },
161 AlreadyRunning {
164 child_session_id: crate::typed_id::SessionId,
165 claim_token: uuid::Uuid,
167 },
168 AlreadySettled {
171 child_session_id: crate::typed_id::SessionId,
172 terminal_status: String,
174 terminal_result: String,
175 },
176}
177
178#[async_trait]
186pub trait SubagentSpawnStore: Send + Sync + 'static {
187 async fn try_claim_spawn(
192 &self,
193 parent_session_id: crate::typed_id::SessionId,
194 tool_call_id: &str,
195 claim_token: uuid::Uuid,
196 ) -> Result<SpawnClaimResult>;
197
198 async fn register_child_session(
203 &self,
204 spawn_handle_id: uuid::Uuid,
205 claim_token: uuid::Uuid,
206 child_session_id: crate::typed_id::SessionId,
207 ) -> Result<()>;
208
209 async fn settle_spawn(
215 &self,
216 parent_session_id: crate::typed_id::SessionId,
217 tool_call_id: &str,
218 claim_token: uuid::Uuid,
219 terminal_status: &str,
220 terminal_result: &str,
221 ) -> Result<()>;
222}
223
224#[async_trait]
226impl<S: SubagentSpawnStore + ?Sized> SubagentSpawnStore for Arc<S> {
227 async fn try_claim_spawn(
228 &self,
229 parent_session_id: crate::typed_id::SessionId,
230 tool_call_id: &str,
231 claim_token: uuid::Uuid,
232 ) -> Result<SpawnClaimResult> {
233 (**self)
234 .try_claim_spawn(parent_session_id, tool_call_id, claim_token)
235 .await
236 }
237
238 async fn register_child_session(
239 &self,
240 spawn_handle_id: uuid::Uuid,
241 claim_token: uuid::Uuid,
242 child_session_id: crate::typed_id::SessionId,
243 ) -> Result<()> {
244 (**self)
245 .register_child_session(spawn_handle_id, claim_token, child_session_id)
246 .await
247 }
248
249 async fn settle_spawn(
250 &self,
251 parent_session_id: crate::typed_id::SessionId,
252 tool_call_id: &str,
253 claim_token: uuid::Uuid,
254 terminal_status: &str,
255 terminal_result: &str,
256 ) -> Result<()> {
257 (**self)
258 .settle_spawn(
259 parent_session_id,
260 tool_call_id,
261 claim_token,
262 terminal_status,
263 terminal_result,
264 )
265 .await
266 }
267}