1use std::collections::HashMap;
2use std::sync::{Arc, Mutex, RwLock};
3use std::time::{SystemTime, UNIX_EPOCH};
4
5use adk_core::{Event, Part, Result};
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9use super::{
10 RelationshipKind, TeamBudget, TeamError, TeamLifecycleContext, TeamLifecycleDecision,
11 TeamLifecycleHook, TeamLifecycleOutcome, TeamRuntimeError, TeamTerminationPolicy,
12 lifecycle::TeamLifecycleManager,
13};
14
15pub const TEAM_EXECUTION_STATE_KEY: &str = "__adk_team_execution_v1";
17pub const TEAM_ROOT_INVOCATION_KEY: &str = "adk.team.root_invocation_id";
19pub const TEAM_EDGE_ID_KEY: &str = "adk.team.edge_id";
21
22#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
24#[serde(rename_all = "camelCase")]
25pub struct TeamExecutionUsage {
26 pub events: u64,
28 pub model_requests: u64,
30 pub tool_calls: u64,
32 pub tokens: u64,
34 pub cost_microusd: u64,
36 pub delegations: u64,
38 pub handoffs: u64,
40}
41
42#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
44#[serde(rename_all = "camelCase")]
45pub enum TeamExecutionStatus {
46 Running,
48 Completed,
50 Terminated,
52 Failed,
54 BudgetExceeded,
56}
57
58#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
60#[serde(rename_all = "camelCase")]
61pub struct TeamEdgeExecution {
62 pub id: String,
64 #[serde(default, skip_serializing_if = "Option::is_none")]
66 pub parent_id: Option<String>,
67 pub from: String,
69 pub to: String,
71 pub kind: RelationshipKind,
73 pub attempt: u32,
75 pub started_at_ms: u64,
77 #[serde(default, skip_serializing_if = "Option::is_none")]
79 pub finished_at_ms: Option<u64>,
80 pub status: TeamExecutionStatus,
82 #[serde(default, skip_serializing_if = "Option::is_none")]
84 pub error: Option<String>,
85}
86
87#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
89#[serde(rename_all = "camelCase")]
90pub struct TeamExecutionSnapshot {
91 pub team: String,
93 pub invocation_id: String,
95 pub roster: Vec<ResolvedTeamMember>,
97 pub started_at_ms: u64,
99 pub updated_at_ms: u64,
101 pub status: TeamExecutionStatus,
103 pub usage: TeamExecutionUsage,
105 pub edges: Vec<TeamEdgeExecution>,
107 #[serde(default, skip_serializing_if = "Option::is_none")]
109 pub reason: Option<String>,
110}
111
112#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
114#[serde(rename_all = "camelCase")]
115pub struct ResolvedTeamMember {
116 pub member: String,
118 pub binding: String,
120 #[serde(default, skip_serializing_if = "Vec::is_empty")]
122 pub capabilities: Vec<String>,
123 #[serde(default, skip_serializing_if = "Option::is_none")]
125 pub version: Option<String>,
126 #[serde(default, skip_serializing_if = "Option::is_none")]
128 pub digest: Option<String>,
129 #[serde(default, skip_serializing_if = "Vec::is_empty")]
131 pub trust_labels: Vec<String>,
132}
133
134#[derive(Debug)]
135pub(crate) enum EventDisposition {
136 Continue,
137 Terminate,
138}
139
140pub(crate) struct TeamEdgeStart<'a> {
141 pub(crate) execution_id: Option<String>,
142 pub(crate) parent_id: Option<String>,
143 pub(crate) from: &'a str,
144 pub(crate) to: &'a str,
145 pub(crate) kind: RelationshipKind,
146 pub(crate) attempt: u32,
147}
148
149pub(crate) struct TeamRuntimeRegistry {
150 team: String,
151 roster: Vec<ResolvedTeamMember>,
152 budget: TeamBudget,
153 termination: TeamTerminationPolicy,
154 lifecycle: TeamLifecycleManager,
155 invocations: RwLock<HashMap<String, Arc<Mutex<TeamExecutionSnapshot>>>>,
156}
157
158impl TeamRuntimeRegistry {
159 pub(crate) fn team_name(&self) -> &str {
160 &self.team
161 }
162
163 pub(crate) fn new(
164 team: String,
165 roster: Vec<ResolvedTeamMember>,
166 budget: TeamBudget,
167 termination: TeamTerminationPolicy,
168 hooks: Vec<Arc<dyn TeamLifecycleHook>>,
169 ) -> Self {
170 Self {
171 team,
172 roster,
173 budget,
174 termination,
175 lifecycle: TeamLifecycleManager::new(hooks),
176 invocations: RwLock::new(HashMap::new()),
177 }
178 }
179
180 pub(crate) async fn before_lifecycle(
181 &self,
182 context: &TeamLifecycleContext,
183 ) -> Result<TeamLifecycleDecision> {
184 self.lifecycle.before(context).await
185 }
186
187 pub(crate) async fn after_lifecycle(
188 &self,
189 context: &TeamLifecycleContext,
190 outcome: &TeamLifecycleOutcome,
191 ) -> Result<()> {
192 self.lifecycle.after(context, outcome).await
193 }
194
195 pub(crate) fn snapshot(&self, invocation_id: &str) -> Option<TeamExecutionSnapshot> {
196 let ledger = self
197 .invocations
198 .read()
199 .unwrap_or_else(|error| error.into_inner())
200 .get(invocation_id)
201 .cloned()?;
202 let snapshot = ledger.lock().unwrap_or_else(|error| error.into_inner()).clone();
203 Some(snapshot)
204 }
205
206 pub(crate) fn snapshots(&self) -> Vec<TeamExecutionSnapshot> {
207 self.invocations
208 .read()
209 .unwrap_or_else(|error| error.into_inner())
210 .values()
211 .map(|ledger| ledger.lock().unwrap_or_else(|error| error.into_inner()).clone())
212 .collect()
213 }
214
215 pub(crate) fn restore(
216 &self,
217 snapshot: TeamExecutionSnapshot,
218 ) -> std::result::Result<(), TeamError> {
219 if snapshot.team != self.team {
220 return Err(TeamError::IncompatibleExecutionSnapshot(format!(
221 "snapshot belongs to team '{}', expected '{}'",
222 snapshot.team, self.team
223 )));
224 }
225 if snapshot.roster != self.roster {
226 return Err(TeamError::IncompatibleExecutionSnapshot(
227 "the frozen member roster does not match".to_string(),
228 ));
229 }
230 if snapshot.invocation_id.trim().is_empty() {
231 return Err(TeamError::IncompatibleExecutionSnapshot(
232 "invocationId must not be empty".to_string(),
233 ));
234 }
235 self.invocations
236 .write()
237 .unwrap_or_else(|error| error.into_inner())
238 .insert(snapshot.invocation_id.clone(), Arc::new(Mutex::new(snapshot)));
239 Ok(())
240 }
241
242 pub(crate) fn resume_handoff_target(&self, invocation_id: &str) -> Result<Option<String>> {
243 let Some(snapshot) = self.snapshot(invocation_id) else {
244 return Ok(None);
245 };
246 let Some(active) =
247 snapshot.edges.iter().rev().find(|edge| edge.status == TeamExecutionStatus::Running)
248 else {
249 return Ok(None);
250 };
251 match active.kind {
252 RelationshipKind::Handoff => Ok(Some(active.to.clone())),
253 RelationshipKind::Delegate => Err(TeamRuntimeError::UnsafeResume(format!(
254 "cannot replay unresolved delegation '{}' from '{}' to '{}'; inspect CompiledTeam::resume_plan and use a checkpoint-aware durable host",
255 active.id, active.from, active.to
256 ))
257 .into()),
258 }
259 }
260
261 fn ledger(&self, invocation_id: &str) -> Arc<Mutex<TeamExecutionSnapshot>> {
262 if let Some(ledger) = self
263 .invocations
264 .read()
265 .unwrap_or_else(|error| error.into_inner())
266 .get(invocation_id)
267 .cloned()
268 {
269 return ledger;
270 }
271 let now = now_ms();
272 let ledger = Arc::new(Mutex::new(TeamExecutionSnapshot {
273 team: self.team.clone(),
274 invocation_id: invocation_id.to_string(),
275 roster: self.roster.clone(),
276 started_at_ms: now,
277 updated_at_ms: now,
278 status: TeamExecutionStatus::Running,
279 usage: TeamExecutionUsage::default(),
280 edges: Vec::new(),
281 reason: None,
282 }));
283 self.invocations
284 .write()
285 .unwrap_or_else(|error| error.into_inner())
286 .entry(invocation_id.to_string())
287 .or_insert_with(|| ledger.clone())
288 .clone()
289 }
290
291 pub(crate) fn check_budget(&self, invocation_id: &str) -> Result<()> {
292 let ledger = self.ledger(invocation_id);
293 let mut snapshot = ledger.lock().unwrap_or_else(|error| error.into_inner());
294 if let Some(reason) = self.budget_violation(&snapshot) {
295 snapshot.status = TeamExecutionStatus::BudgetExceeded;
296 snapshot.reason = Some(reason.clone());
297 snapshot.updated_at_ms = now_ms();
298 return Err(TeamRuntimeError::BudgetExceeded(reason).into());
299 }
300 Ok(())
301 }
302
303 pub(crate) fn start_edge(
304 &self,
305 invocation_id: &str,
306 start: TeamEdgeStart<'_>,
307 ) -> Result<String> {
308 self.check_budget(invocation_id)?;
309 let ledger = self.ledger(invocation_id);
310 let mut snapshot = ledger.lock().unwrap_or_else(|error| error.into_inner());
311 snapshot.status = TeamExecutionStatus::Running;
312 snapshot.reason = None;
313 match start.kind {
314 RelationshipKind::Delegate => snapshot.usage.delegations += 1,
315 RelationshipKind::Handoff => snapshot.usage.handoffs += 1,
316 }
317 let id = start.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
318 snapshot.edges.push(TeamEdgeExecution {
319 id: id.clone(),
320 parent_id: start.parent_id,
321 from: start.from.to_string(),
322 to: start.to.to_string(),
323 kind: start.kind,
324 attempt: start.attempt,
325 started_at_ms: now_ms(),
326 finished_at_ms: None,
327 status: TeamExecutionStatus::Running,
328 error: None,
329 });
330 snapshot.updated_at_ms = now_ms();
331 if let Some(reason) = self.budget_violation(&snapshot) {
332 snapshot.status = TeamExecutionStatus::BudgetExceeded;
333 snapshot.reason = Some(reason.clone());
334 return Err(TeamRuntimeError::BudgetExceeded(reason).into());
335 }
336 Ok(id)
337 }
338
339 pub(crate) fn finish_edge(&self, invocation_id: &str, id: &str, error: Option<String>) {
340 let ledger = self.ledger(invocation_id);
341 let mut snapshot = ledger.lock().unwrap_or_else(|failure| failure.into_inner());
342 if let Some(edge) = snapshot.edges.iter_mut().find(|edge| edge.id == id) {
343 edge.finished_at_ms = Some(now_ms());
344 edge.status = if error.is_some() {
345 TeamExecutionStatus::Failed
346 } else {
347 TeamExecutionStatus::Completed
348 };
349 edge.error = error;
350 }
351 snapshot.updated_at_ms = now_ms();
352 }
353
354 pub(crate) fn record_event(
355 &self,
356 invocation_id: &str,
357 edge_id: Option<&str>,
358 event: &mut Event,
359 ) -> Result<EventDisposition> {
360 let ledger = self.ledger(invocation_id);
361 let mut snapshot = ledger.lock().unwrap_or_else(|error| error.into_inner());
362 let already_recorded = event
363 .provider_metadata
364 .get(TEAM_ROOT_INVOCATION_KEY)
365 .is_some_and(|root| root == invocation_id);
366 if !already_recorded {
367 snapshot.usage.events += 1;
368 snapshot.usage.tool_calls += event.tool_calls().len() as u64;
369 if let Some(usage) = &event.llm_response.usage_metadata {
370 snapshot.usage.model_requests += 1;
371 snapshot.usage.tokens += u64::try_from(usage.total_token_count.max(0)).unwrap_or(0);
372 if let Some(cost) = usage.cost
373 && cost.is_finite()
374 && cost > 0.0
375 {
376 snapshot.usage.cost_microusd = snapshot
377 .usage
378 .cost_microusd
379 .saturating_add((cost * 1_000_000.0).round() as u64);
380 }
381 }
382 }
383 snapshot.updated_at_ms = now_ms();
384 event
385 .provider_metadata
386 .insert(TEAM_ROOT_INVOCATION_KEY.to_string(), invocation_id.to_string());
387 if let Some(edge_id) = edge_id {
388 event.provider_metadata.insert(TEAM_EDGE_ID_KEY.to_string(), edge_id.to_string());
389 }
390
391 if let Some(reason) = self.budget_violation(&snapshot) {
392 snapshot.status = TeamExecutionStatus::BudgetExceeded;
393 snapshot.reason = Some(reason.clone());
394 persist_snapshot_if_absent(event, &snapshot);
395 return Err(TeamRuntimeError::BudgetExceeded(reason).into());
396 }
397
398 let termination = self.termination_reason(event);
399 if let Some(reason) = &termination {
400 snapshot.status = TeamExecutionStatus::Terminated;
401 snapshot.reason = Some(reason.clone());
402 } else if event.is_final_response()
403 && edge_id.is_none_or(|id| {
404 snapshot
405 .edges
406 .iter()
407 .find(|edge| edge.id == id)
408 .is_none_or(|edge| edge.kind != RelationshipKind::Delegate)
409 })
410 {
411 snapshot.status = TeamExecutionStatus::Completed;
412 }
413 persist_snapshot_if_absent(event, &snapshot);
414 Ok(termination.map_or(EventDisposition::Continue, |_| EventDisposition::Terminate))
415 }
416
417 pub(crate) fn fail(&self, invocation_id: &str, reason: String) {
418 let ledger = self.ledger(invocation_id);
419 let mut snapshot = ledger.lock().unwrap_or_else(|error| error.into_inner());
420 snapshot.status = TeamExecutionStatus::Failed;
421 snapshot.reason = Some(reason);
422 snapshot.updated_at_ms = now_ms();
423 }
424
425 fn termination_reason(&self, event: &Event) -> Option<String> {
426 if self.termination.stop_on_escalation && event.actions.escalate {
427 return Some("team terminated on escalation".to_string());
428 }
429 if event.is_final_response()
430 && self.termination.final_authors.iter().any(|author| author == &event.author)
431 {
432 return Some(format!("team terminated after final response from '{}'", event.author));
433 }
434 for marker in &self.termination.text_markers {
435 let matched = event
436 .content()
437 .into_iter()
438 .flat_map(|content| &content.parts)
439 .any(|part| matches!(part, Part::Text { text } if text.contains(marker)));
440 if matched {
441 return Some(format!("team termination marker matched: {marker}"));
442 }
443 }
444 None
445 }
446
447 fn budget_violation(&self, snapshot: &TeamExecutionSnapshot) -> Option<String> {
448 let usage = &snapshot.usage;
449 let limits = [
450 ("events", usage.events, self.budget.max_events),
451 ("model requests", usage.model_requests, self.budget.max_model_requests),
452 ("tool calls", usage.tool_calls, self.budget.max_tool_calls),
453 ("tokens", usage.tokens, self.budget.max_tokens),
454 ("costMicrousd", usage.cost_microusd, self.budget.max_cost_microusd),
455 ("delegations", usage.delegations, self.budget.max_delegations),
456 ("handoffs", usage.handoffs, self.budget.max_handoffs),
457 ];
458 if let Some((name, used, limit)) =
459 limits.into_iter().find(|(_, used, limit)| limit.is_some_and(|max| *used > max))
460 {
461 return Some(format!(
462 "team budget exceeded for {name}: used {used}, maximum {}",
463 limit.unwrap_or_default()
464 ));
465 }
466 if let Some(max_wall_time_ms) = self.budget.max_wall_time_ms {
467 let elapsed = now_ms().saturating_sub(snapshot.started_at_ms);
468 if elapsed > max_wall_time_ms {
469 return Some(format!(
470 "team wall-time budget exceeded: elapsed {elapsed}ms, maximum {max_wall_time_ms}ms"
471 ));
472 }
473 }
474 None
475 }
476}
477
478fn persist_snapshot(event: &mut Event, snapshot: &TeamExecutionSnapshot) {
479 if let Ok(value) = serde_json::to_value(snapshot) {
480 event.actions.state_delta.insert(TEAM_EXECUTION_STATE_KEY.to_string(), value);
481 }
482}
483
484fn persist_snapshot_if_absent(event: &mut Event, snapshot: &TeamExecutionSnapshot) {
485 if !event.actions.state_delta.contains_key(TEAM_EXECUTION_STATE_KEY) {
486 persist_snapshot(event, snapshot);
487 }
488}
489
490fn now_ms() -> u64 {
491 SystemTime::now()
492 .duration_since(UNIX_EPOCH)
493 .unwrap_or_default()
494 .as_millis()
495 .try_into()
496 .unwrap_or(u64::MAX)
497}
498
499#[cfg(test)]
500mod tests {
501 use super::*;
502 use adk_core::Content;
503
504 fn registry(budget: TeamBudget, termination: TeamTerminationPolicy) -> TeamRuntimeRegistry {
505 TeamRuntimeRegistry::new(
506 "portable_team".to_string(),
507 vec![ResolvedTeamMember {
508 member: "root".to_string(),
509 binding: "root-v1".to_string(),
510 capabilities: vec!["route".to_string()],
511 version: None,
512 digest: None,
513 trust_labels: Vec::new(),
514 }],
515 budget,
516 termination,
517 Vec::new(),
518 )
519 }
520
521 #[test]
522 fn enforces_aggregate_event_budget_and_persists_receipt() {
523 let runtime = registry(
524 TeamBudget { max_events: Some(1), ..TeamBudget::default() },
525 TeamTerminationPolicy::default(),
526 );
527 let mut first = Event::new("invocation");
528 assert!(matches!(
529 runtime.record_event("invocation", None, &mut first),
530 Ok(EventDisposition::Continue)
531 ));
532 assert!(first.actions.state_delta.contains_key(TEAM_EXECUTION_STATE_KEY));
533
534 let mut second = Event::new("invocation");
535 let error = runtime.record_event("invocation", None, &mut second).unwrap_err();
536 assert!(error.to_string().contains("events"));
537 assert_eq!(error.code, "agent.team.budget_exceeded");
538 assert_eq!(
539 runtime.snapshot("invocation").unwrap().status,
540 TeamExecutionStatus::BudgetExceeded
541 );
542 }
543
544 #[test]
545 fn records_causal_edges_and_terminates_on_marker() {
546 let runtime = registry(
547 TeamBudget::default(),
548 TeamTerminationPolicy {
549 text_markers: vec!["APPROVED".to_string()],
550 ..TeamTerminationPolicy::default()
551 },
552 );
553 let parent = runtime
554 .start_edge(
555 "invocation",
556 TeamEdgeStart {
557 execution_id: Some("edge-1".to_string()),
558 parent_id: None,
559 from: "root",
560 to: "a",
561 kind: RelationshipKind::Delegate,
562 attempt: 1,
563 },
564 )
565 .unwrap();
566 let child = runtime
567 .start_edge(
568 "invocation",
569 TeamEdgeStart {
570 execution_id: Some("edge-2".to_string()),
571 parent_id: Some(parent.clone()),
572 from: "a",
573 to: "b",
574 kind: RelationshipKind::Handoff,
575 attempt: 1,
576 },
577 )
578 .unwrap();
579 runtime.finish_edge("invocation", &child, None);
580 runtime.finish_edge("invocation", &parent, None);
581 let mut event = Event::new("invocation");
582 event.author = "b".to_string();
583 event.llm_response.content = Some(Content::new("model").with_text("APPROVED"));
584 assert!(matches!(
585 runtime.record_event("invocation", Some(&child), &mut event),
586 Ok(EventDisposition::Terminate)
587 ));
588 let snapshot = runtime.snapshot("invocation").unwrap();
589 assert_eq!(snapshot.status, TeamExecutionStatus::Terminated);
590 assert_eq!(snapshot.edges[1].parent_id.as_deref(), Some("edge-1"));
591 assert_eq!(event.provider_metadata.get(TEAM_EDGE_ID_KEY), Some(&child));
592 }
593
594 #[test]
595 fn restores_only_matching_frozen_rosters() {
596 let source = registry(TeamBudget::default(), TeamTerminationPolicy::default());
597 source.check_budget("resume-me").unwrap();
598 let snapshot = source.snapshot("resume-me").unwrap();
599
600 let restored = registry(TeamBudget::default(), TeamTerminationPolicy::default());
601 restored.restore(snapshot.clone()).unwrap();
602 assert_eq!(restored.snapshot("resume-me"), Some(snapshot.clone()));
603
604 let other = TeamRuntimeRegistry::new(
605 "portable_team".to_string(),
606 vec![ResolvedTeamMember {
607 member: "root".to_string(),
608 binding: "root-v2".to_string(),
609 capabilities: vec!["route".to_string()],
610 version: None,
611 digest: None,
612 trust_labels: Vec::new(),
613 }],
614 TeamBudget::default(),
615 TeamTerminationPolicy::default(),
616 Vec::new(),
617 );
618 assert!(matches!(
619 other.restore(snapshot),
620 Err(TeamError::IncompatibleExecutionSnapshot(_))
621 ));
622 }
623
624 #[test]
625 fn resumes_handoffs_but_refuses_unsafe_delegate_replay() {
626 let handoff = registry(TeamBudget::default(), TeamTerminationPolicy::default());
627 handoff
628 .start_edge(
629 "handoff-run",
630 TeamEdgeStart {
631 execution_id: Some("handoff-edge".to_string()),
632 parent_id: None,
633 from: "root",
634 to: "specialist",
635 kind: RelationshipKind::Handoff,
636 attempt: 1,
637 },
638 )
639 .unwrap();
640 assert_eq!(
641 handoff.resume_handoff_target("handoff-run").unwrap().as_deref(),
642 Some("specialist")
643 );
644
645 let delegate = registry(TeamBudget::default(), TeamTerminationPolicy::default());
646 delegate
647 .start_edge(
648 "delegate-run",
649 TeamEdgeStart {
650 execution_id: Some("delegate-edge".to_string()),
651 parent_id: None,
652 from: "root",
653 to: "worker",
654 kind: RelationshipKind::Delegate,
655 attempt: 1,
656 },
657 )
658 .unwrap();
659 let error = delegate.resume_handoff_target("delegate-run").unwrap_err();
660 assert!(error.to_string().contains("cannot replay unresolved delegation"));
661 assert_eq!(error.code, "agent.team.resume_unsafe");
662 }
663}