1use crate::InvocationContext;
2use crate::cache::CacheManager;
3#[cfg(feature = "artifacts")]
4use adk_artifact::ArtifactService;
5use adk_core::{
6 Agent, AppName, CacheCapable, Content, ContextCacheConfig, Event, EventStream, Memory,
7 ReadonlyContext, Result, RunConfig, SessionId, UserId,
8};
9#[cfg(feature = "plugins")]
10use adk_plugin::PluginManager;
11use adk_session::SessionService;
12#[cfg(feature = "skills")]
13use adk_skill::{SkillInjector, SkillInjectorConfig};
14use async_stream::stream;
15use std::{collections::HashMap, sync::Arc};
16use tokio_util::sync::CancellationToken;
17
18#[derive(Debug, Clone)]
23struct ActiveRun {
24 identity: adk_core::AdkIdentity,
25 token: CancellationToken,
26}
27
28type ActiveRuns = Arc<std::sync::Mutex<std::collections::HashMap<u64, ActiveRun>>>;
30
31fn preserve_streamed_content(accumulated: &mut HashMap<String, Content>, event: &mut Event) {
32 if event.llm_response.partial {
33 if let Some(chunk) = &event.llm_response.content {
34 accumulated
35 .entry(event.id.clone())
36 .and_modify(|content| content.parts.extend(chunk.parts.clone()))
37 .or_insert_with(|| chunk.clone());
38 }
39 } else if event.llm_response.content.is_none() {
40 event.llm_response.content = accumulated.remove(&event.id);
41 } else {
42 accumulated.remove(&event.id);
43 }
44}
45
46struct ActiveRunCleanup {
51 active_runs: ActiveRuns,
52 run_id: u64,
53}
54
55impl Drop for ActiveRunCleanup {
56 fn drop(&mut self) {
57 let mut runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
58 runs.remove(&self.run_id);
59 }
60}
61use tracing::Instrument;
62
63pub struct RunnerConfig {
67 pub app_name: String,
69 pub agent: Arc<dyn Agent>,
71 pub session_service: Arc<dyn SessionService>,
73 #[cfg(feature = "artifacts")]
74 pub artifact_service: Option<Arc<dyn ArtifactService>>,
76 pub memory_service: Option<Arc<dyn Memory>>,
78 #[cfg(feature = "plugins")]
79 pub plugin_manager: Option<Arc<PluginManager>>,
81 #[allow(dead_code)]
84 pub run_config: Option<RunConfig>,
85 pub compaction_config: Option<adk_core::EventsCompactionConfig>,
89 pub context_cache_config: Option<ContextCacheConfig>,
97 pub cache_capable: Option<Arc<dyn CacheCapable>>,
100 pub request_context: Option<adk_core::RequestContext>,
104 pub cancellation_token: Option<CancellationToken>,
106 pub intra_compaction_config: Option<adk_core::IntraCompactionConfig>,
110 pub intra_compaction_summarizer: Option<Arc<dyn adk_core::BaseEventsSummarizer>>,
113 #[cfg(feature = "context-compaction")]
121 pub context_compaction: Option<crate::compaction::CompactionConfig>,
122}
123
124pub struct Runner {
130 app_name: String,
131 root_agent: Arc<dyn Agent>,
132 session_service: Arc<dyn SessionService>,
133 #[cfg(feature = "artifacts")]
134 artifact_service: Option<Arc<dyn ArtifactService>>,
135 memory_service: Option<Arc<dyn Memory>>,
136 #[cfg(feature = "plugins")]
137 plugin_manager: Option<Arc<PluginManager>>,
138 #[cfg(feature = "skills")]
139 skill_injector: Option<Arc<SkillInjector>>,
140 run_config: RunConfig,
141 compaction_config: Option<adk_core::EventsCompactionConfig>,
142 context_cache_config: Option<ContextCacheConfig>,
143 cache_capable: Option<Arc<dyn CacheCapable>>,
144 cache_manager: Option<Arc<tokio::sync::Mutex<CacheManager>>>,
145 request_context: Option<adk_core::RequestContext>,
146 cancellation_token: Option<CancellationToken>,
147 intra_compactor: Option<Arc<crate::intra_compaction::IntraInvocationCompactor>>,
148 #[cfg(feature = "context-compaction")]
150 context_compaction: Option<Arc<crate::compaction::CompactionConfig>>,
151 active_runs: ActiveRuns,
154 next_run_id: Arc<std::sync::atomic::AtomicU64>,
155 pub(crate) external_session_locks: Arc<
158 std::sync::Mutex<
159 std::collections::HashMap<String, std::sync::Weak<tokio::sync::Mutex<()>>>,
160 >,
161 >,
162}
163
164impl Runner {
165 pub fn builder() -> crate::builder::RunnerConfigBuilder<
181 crate::builder::NoAppName,
182 crate::builder::NoAgent,
183 crate::builder::NoSessionService,
184 > {
185 crate::builder::RunnerConfigBuilder::new()
186 }
187
188 pub fn new(config: RunnerConfig) -> Result<Self> {
192 let run_config = config.run_config.unwrap_or_default();
193
194 let effective_cache_config = config
197 .context_cache_config
198 .or_else(|| config.cache_capable.as_ref().map(|_| ContextCacheConfig::default()));
199
200 let cache_manager = effective_cache_config
201 .as_ref()
202 .map(|c| Arc::new(tokio::sync::Mutex::new(CacheManager::new(c.clone()))));
203
204 let intra_compactor = config.intra_compaction_config.as_ref().and_then(|ic_config| {
205 config.intra_compaction_summarizer.as_ref().map(|summarizer| {
206 Arc::new(crate::intra_compaction::IntraInvocationCompactor::new(
207 ic_config.clone(),
208 summarizer.clone(),
209 ))
210 })
211 });
212
213 Ok(Self {
214 app_name: config.app_name,
215 root_agent: config.agent,
216 session_service: config.session_service,
217 #[cfg(feature = "artifacts")]
218 artifact_service: config.artifact_service,
219 memory_service: config.memory_service,
220 #[cfg(feature = "plugins")]
221 plugin_manager: config.plugin_manager,
222 #[cfg(feature = "skills")]
223 skill_injector: None,
224 run_config,
225 compaction_config: config.compaction_config,
226 context_cache_config: effective_cache_config,
227 cache_capable: config.cache_capable,
228 cache_manager,
229 request_context: config.request_context,
230 cancellation_token: config.cancellation_token,
231 intra_compactor,
232 #[cfg(feature = "context-compaction")]
233 context_compaction: config.context_compaction.map(Arc::new),
234 active_runs: Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
235 next_run_id: Arc::new(std::sync::atomic::AtomicU64::new(1)),
236 external_session_locks: Arc::new(std::sync::Mutex::new(
237 std::collections::HashMap::new(),
238 )),
239 })
240 }
241
242 pub(crate) fn root_agent(&self) -> Arc<dyn Agent> {
244 Arc::clone(&self.root_agent)
245 }
246
247 #[cfg(feature = "skills")]
251 pub fn with_skill_injector(mut self, injector: SkillInjector) -> Self {
252 self.skill_injector = Some(Arc::new(injector));
253 self
254 }
255
256 #[cfg(feature = "skills")]
258 #[deprecated(note = "Use with_auto_skills_mut instead")]
259 pub fn with_auto_skills(
260 mut self,
261 root: impl AsRef<std::path::Path>,
262 config: SkillInjectorConfig,
263 ) -> adk_skill::SkillResult<Self> {
264 self.with_auto_skills_mut(root, config)?;
265 Ok(self)
266 }
267
268 #[cfg(feature = "skills")]
274 pub fn with_auto_skills_mut(
275 &mut self,
276 root: impl AsRef<std::path::Path>,
277 config: SkillInjectorConfig,
278 ) -> adk_skill::SkillResult<()> {
279 let injector = SkillInjector::from_root(root, config)?;
280 self.skill_injector = Some(Arc::new(injector));
281 Ok(())
282 }
283
284 pub async fn run(
289 &self,
290 user_id: UserId,
291 session_id: SessionId,
292 user_content: Content,
293 ) -> Result<EventStream> {
294 self.run_with_config(user_id, session_id, user_content, None).await
295 }
296
297 pub fn app_name(&self) -> &str {
299 &self.app_name
300 }
301
302 pub fn session_service(&self) -> &Arc<dyn adk_session::SessionService> {
304 &self.session_service
305 }
306
307 pub fn run_config(&self) -> &RunConfig {
311 &self.run_config
312 }
313
314 pub async fn run_with_config(
325 &self,
326 user_id: UserId,
327 session_id: SessionId,
328 user_content: Content,
329 run_config: Option<RunConfig>,
330 ) -> Result<EventStream> {
331 let app_name = self.app_name.clone();
332 let typed_app_name = AppName::try_from(app_name.clone())?;
333 let session_service = self.session_service.clone();
334 let root_agent = self.root_agent.clone();
335 #[cfg(feature = "artifacts")]
336 let artifact_service = self.artifact_service.clone();
337 let memory_service = self.memory_service.clone();
338 #[cfg(feature = "plugins")]
339 let plugin_manager = self.plugin_manager.clone();
340 #[cfg(feature = "skills")]
341 let skill_injector = self.skill_injector.clone();
342 let mut run_config = run_config.unwrap_or_else(|| self.run_config.clone());
343 let compaction_config = self.compaction_config.clone();
344 let context_cache_config = self.context_cache_config.clone();
345 let cache_capable = self.cache_capable.clone();
346 let cache_manager_ref = self.cache_manager.clone();
347 let request_context = self.request_context.clone();
348 let cancellation_token = self.cancellation_token.clone();
349 let intra_compactor = self.intra_compactor.clone();
350 #[cfg(feature = "context-compaction")]
351 let context_compaction = self.context_compaction.clone();
352
353 let identity =
356 adk_core::AdkIdentity::new(typed_app_name.clone(), user_id.clone(), session_id.clone());
357
358 let session_token = CancellationToken::new();
363 let run_id = self.next_run_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
364 {
365 let mut runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
366 runs.insert(
367 run_id,
368 ActiveRun { identity: identity.clone(), token: session_token.clone() },
369 );
370 }
371 let active_runs = self.active_runs.clone();
372 let effective_token = if let Some(ref global) = cancellation_token {
374 let combined = CancellationToken::new();
375 let combined_clone = combined.clone();
376 let global_clone = global.clone();
377 let session_clone = session_token.clone();
378 let combined_for_global = combined_clone.clone();
380 tokio::spawn(async move {
381 global_clone.cancelled().await;
382 combined_for_global.cancel();
383 });
384 let combined_for_session = combined_clone;
385 tokio::spawn(async move {
386 session_clone.cancelled().await;
387 combined_for_session.cancel();
388 });
389 Some(combined)
390 } else {
391 Some(session_token.clone())
392 };
393
394 let cleanup = ActiveRunCleanup { active_runs: active_runs.clone(), run_id };
399
400 let s = stream! {
401 let _cleanup = cleanup;
402
403 let cancellation_token = effective_token;
405 let session = match session_service
407 .get(adk_session::GetRequest {
408 app_name: app_name.clone(),
409 user_id: user_id.to_string(),
410 session_id: session_id.to_string(),
411 num_recent_events: run_config.history_max_events,
412 after: None,
413 })
414 .await
415 {
416 Ok(s) => s,
417 Err(e) => {
418 yield Err(e);
419 return;
420 }
421 };
422
423 let agent_to_run = Self::find_agent_to_run(&root_agent, session.as_ref());
425
426 root_agent.configure_run(agent_to_run.name(), &mut run_config);
429 if let Some(targets) = root_agent.transfer_targets_for(agent_to_run.name()) {
430 run_config.transfer_targets = targets;
431 run_config.parent_agent = None;
432 }
433
434 #[cfg(feature = "artifacts")]
436 let artifact_service_clone = artifact_service.clone();
437 let memory_service_clone = memory_service.clone();
438
439 let invocation_id = format!("inv-{}", uuid::Uuid::new_v4());
441 #[cfg(any(feature = "skills", feature = "plugins"))]
442 let mut effective_user_content = user_content.clone();
443 #[cfg(not(any(feature = "skills", feature = "plugins")))]
444 let effective_user_content = user_content.clone();
445 #[cfg(feature = "skills")]
446 let mut selected_skill_name = String::new();
447 #[cfg(not(feature = "skills"))]
448 let selected_skill_name = String::new();
449 #[cfg(feature = "skills")]
450 let mut selected_skill_id = String::new();
451 #[cfg(not(feature = "skills"))]
452 let selected_skill_id = String::new();
453
454 #[cfg(feature = "skills")]
455 if let Some(injector) = skill_injector.as_ref()
456 && let Some(matched) = adk_skill::apply_skill_injection(
457 &mut effective_user_content,
458 injector.index(),
459 injector.policy(),
460 injector.max_injected_chars(),
461 ) {
462 selected_skill_name = matched.skill.name;
463 selected_skill_id = matched.skill.id;
464 }
465
466 let mut invocation_ctx = match InvocationContext::new_typed(
467 invocation_id.clone(),
468 agent_to_run.clone(),
469 user_id.clone(),
470 typed_app_name.clone(),
471 session_id.clone(),
472 effective_user_content.clone(),
473 Arc::from(session),
474 ) {
475 Ok(ctx) => ctx,
476 Err(e) => {
477 yield Err(e);
478 return;
479 }
480 };
481
482 #[cfg(feature = "artifacts")]
484 if let Some(service) = artifact_service {
485 let scoped = adk_artifact::ScopedArtifacts::new(
487 service,
488 app_name.clone(),
489 user_id.to_string(),
490 session_id.to_string(),
491 );
492 invocation_ctx = invocation_ctx.with_artifacts(Arc::new(scoped));
493 }
494 if let Some(memory) = memory_service {
495 invocation_ctx = invocation_ctx.with_memory(memory);
496 }
497
498 invocation_ctx = invocation_ctx.with_run_config(run_config.clone());
500
501 if let Some(rc) = request_context.clone() {
503 invocation_ctx = invocation_ctx.with_request_context(rc);
504 }
505
506 if let Some(token) = cancellation_token.as_ref() {
508 invocation_ctx = invocation_ctx.with_cancellation_token(token.clone());
509 }
510
511 let mut ctx = Arc::new(invocation_ctx);
512
513 #[cfg(feature = "plugins")]
514 if let Some(manager) = plugin_manager.as_ref() {
515 match manager
516 .run_before_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>)
517 .await
518 {
519 Ok(Some(content)) => {
520 let mut early_event = adk_core::Event::new(ctx.invocation_id());
521 early_event.author = agent_to_run.name().to_string();
522 early_event.llm_response.content = Some(content);
523
524 ctx.mutable_session().append_event(early_event.clone());
525 if let Err(e) = session_service
526 .append_event_for_identity(adk_session::AppendEventRequest {
527 identity: identity.clone(),
528 event: early_event.clone(),
529 })
530 .await {
531 yield Err(e);
532 return;
533 }
534
535 yield Ok(early_event);
536 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
537 return;
538 }
539 Ok(None) => {}
540 Err(e) => {
541 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
542 yield Err(e);
543 return;
544 }
545 }
546
547 match manager
548 .run_on_user_message(
549 ctx.clone() as Arc<dyn adk_core::InvocationContext>,
550 effective_user_content.clone(),
551 )
552 .await
553 {
554 Ok(Some(modified)) => {
555 effective_user_content = modified;
556
557 let mut refreshed_ctx = match InvocationContext::with_mutable_session(
558 ctx.invocation_id().to_string(),
559 agent_to_run.clone(),
560 ctx.user_id().to_string(),
561 ctx.app_name().to_string(),
562 ctx.session_id().to_string(),
563 effective_user_content.clone(),
564 ctx.mutable_session().clone(),
565 ) {
566 Ok(ctx) => ctx,
567 Err(e) => {
568 yield Err(e);
569 return;
570 }
571 };
572 refreshed_ctx = refreshed_ctx.with_orchestration_root_invocation_id(
573 adk_core::InvocationContext::orchestration_root_invocation_id(
574 ctx.as_ref(),
575 )
576 .to_string(),
577 );
578
579 #[cfg(feature = "artifacts")]
580 if let Some(service) = artifact_service_clone.clone() {
581 let scoped = adk_artifact::ScopedArtifacts::new(
582 service,
583 ctx.app_name().to_string(),
584 ctx.user_id().to_string(),
585 ctx.session_id().to_string(),
586 );
587 refreshed_ctx = refreshed_ctx.with_artifacts(Arc::new(scoped));
588 }
589 if let Some(memory) = memory_service_clone.clone() {
590 refreshed_ctx = refreshed_ctx.with_memory(memory);
591 }
592 refreshed_ctx = refreshed_ctx.with_run_config(run_config.clone());
593 if let Some(rc) = request_context.clone() {
594 refreshed_ctx = refreshed_ctx.with_request_context(rc);
595 }
596 if let Some(token) = cancellation_token.as_ref() {
597 refreshed_ctx = refreshed_ctx.with_cancellation_token(token.clone());
598 }
599 ctx = Arc::new(refreshed_ctx);
600 }
601 Ok(None) => {}
602 Err(e) => {
603 if let Some(manager) = plugin_manager.as_ref() {
604 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
605 }
606 yield Err(e);
607 return;
608 }
609 }
610 }
611
612 let mut user_event = adk_core::Event::new(ctx.invocation_id());
614 user_event.author = "user".to_string();
615 user_event.llm_response.content = Some(effective_user_content.clone());
616
617 ctx.mutable_session().append_event(user_event.clone());
620
621 if let Err(e) = session_service
622 .append_event_for_identity(adk_session::AppendEventRequest {
623 identity: identity.clone(),
624 event: user_event,
625 })
626 .await {
627 #[cfg(feature = "plugins")]
628 if let Some(manager) = plugin_manager.as_ref() {
629 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
630 }
631 yield Err(e);
632 return;
633 }
634
635 if let (Some(cm_mutex), Some(cache_model)) = (&cache_manager_ref, &cache_capable) {
640 let should_refresh_cache = {
641 let cm = cm_mutex.lock().await;
642 cm.is_enabled() && (cm.active_cache_name().is_none() || cm.needs_refresh())
643 };
644
645 if should_refresh_cache {
646 let system_instruction = agent_to_run.description().to_string();
650 let tools = std::collections::HashMap::new();
651 let ttl = context_cache_config.as_ref().map_or(600, |c| c.ttl_seconds);
652
653 match cache_model.create_cache(&system_instruction, &tools, ttl).await {
654 Ok(name) => {
655 let old_cache = {
656 let mut cm = cm_mutex.lock().await;
657 let old = cm.clear_active_cache();
658 cm.set_active_cache(name);
659 old
660 };
661
662 if let Some(old) = old_cache
663 && let Err(e) = cache_model.delete_cache(&old).await {
664 tracing::warn!(
665 old_cache = %old,
666 error = %e,
667 "failed to delete old cache, proceeding with new cache"
668 );
669 }
670 }
671 Err(e) => {
672 tracing::warn!(
673 error = %e,
674 "cache creation failed, proceeding without cache"
675 );
676 }
677 }
678 }
679
680 let cache_name = {
682 let mut cm = cm_mutex.lock().await;
683 if cm.is_enabled() {
684 cm.record_invocation().map(str::to_string)
685 } else {
686 None
687 }
688 };
689
690 if let Some(cache_name) = cache_name {
691 run_config.cached_content = Some(cache_name);
692 let mut refreshed_ctx = match InvocationContext::with_mutable_session(
694 ctx.invocation_id().to_string(),
695 agent_to_run.clone(),
696 ctx.user_id().to_string(),
697 ctx.app_name().to_string(),
698 ctx.session_id().to_string(),
699 effective_user_content.clone(),
700 ctx.mutable_session().clone(),
701 ) {
702 Ok(ctx) => ctx,
703 Err(e) => {
704 yield Err(e);
705 return;
706 }
707 };
708 refreshed_ctx = refreshed_ctx.with_orchestration_root_invocation_id(
709 adk_core::InvocationContext::orchestration_root_invocation_id(ctx.as_ref())
710 .to_string(),
711 );
712 #[cfg(feature = "artifacts")]
713 if let Some(service) = artifact_service_clone.clone() {
714 let scoped = adk_artifact::ScopedArtifacts::new(
715 service,
716 ctx.app_name().to_string(),
717 ctx.user_id().to_string(),
718 ctx.session_id().to_string(),
719 );
720 refreshed_ctx = refreshed_ctx.with_artifacts(Arc::new(scoped));
721 }
722 if let Some(memory) = memory_service_clone.clone() {
723 refreshed_ctx = refreshed_ctx.with_memory(memory);
724 }
725 refreshed_ctx = refreshed_ctx.with_run_config(run_config.clone());
726 if let Some(rc) = request_context.clone() {
727 refreshed_ctx = refreshed_ctx.with_request_context(rc);
728 }
729 if let Some(token) = cancellation_token.as_ref() {
730 refreshed_ctx = refreshed_ctx.with_cancellation_token(token.clone());
731 }
732 ctx = Arc::new(refreshed_ctx);
733 }
734 }
735
736 if let Some(ref compactor) = intra_compactor {
740 compactor.reset_cycle();
741 let session_events = ctx.mutable_session().as_ref().events_snapshot();
742 match compactor.maybe_compact(&session_events).await {
743 Ok(Some(compacted_events)) => {
744 ctx.mutable_session().replace_events(compacted_events);
745 tracing::info!("intra-invocation compaction applied before agent execution");
746 }
747 Ok(None) => {} Err(e) => {
749 tracing::warn!(error = %e, "intra-invocation compaction check failed");
750 }
751 }
752 }
753
754 #[cfg(feature = "context-compaction")]
759 if let Some(ref cc_config) = context_compaction {
760 let session_events = ctx.mutable_session().events_snapshot();
761 let estimated = crate::compaction::estimate_event_tokens(&session_events);
762 if estimated > cc_config.context_budget {
763 tracing::info!(
764 estimated_tokens = estimated,
765 budget = cc_config.context_budget,
766 "context exceeds budget, applying proactive compaction"
767 );
768 match crate::compaction::apply_compaction_with_retry(cc_config, session_events).await {
769 Ok(compacted) => {
770 ctx.mutable_session().replace_events(compacted);
771 tracing::info!("proactive context compaction succeeded");
772 }
773 Err(e) => {
774 tracing::warn!(error = %e, "proactive context compaction failed, proceeding with full context");
777 }
778 }
779 }
780 }
781
782 let agent_span = tracing::info_span!(
784 "agent.execute",
785 "gcp.vertex.agent.invocation_id" = ctx.invocation_id(),
786 "gcp.vertex.agent.session_id" = ctx.session_id(),
787 "gcp.vertex.agent.event_id" = ctx.invocation_id(), "gen_ai.conversation.id" = ctx.session_id(),
789 "adk.app_name" = ctx.app_name(),
790 "adk.user_id" = ctx.user_id(),
791 "agent.name" = %agent_to_run.name(),
792 "adk.skills.selected_name" = %selected_skill_name,
793 "adk.skills.selected_id" = %selected_skill_id
794 );
795
796 let mut agent_stream = match agent_to_run.run(ctx.clone()).instrument(agent_span.clone()).await {
797 Ok(s) => s,
798 #[cfg(feature = "context-compaction")]
799 Err(e) if context_compaction.is_some() && crate::compaction::is_token_limit_error(&e) => {
800 let cc_config = context_compaction.as_ref().unwrap();
802 tracing::warn!(
803 error = %e,
804 "agent execution failed with token limit error, attempting compaction"
805 );
806 let session_events = ctx.mutable_session().events_snapshot();
807 match crate::compaction::apply_compaction_with_retry(cc_config, session_events).await {
808 Ok(compacted) => {
809 ctx.mutable_session().replace_events(compacted);
810 tracing::info!("context compaction succeeded after token limit error, retrying agent");
811 match agent_to_run.run(ctx.clone()).instrument(agent_span).await {
813 Ok(s) => s,
814 Err(retry_err) => {
815 #[cfg(feature = "plugins")]
816 if let Some(manager) = plugin_manager.as_ref() {
817 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
818 }
819 yield Err(retry_err);
820 return;
821 }
822 }
823 }
824 Err(compaction_err) => {
825 #[cfg(feature = "plugins")]
826 if let Some(manager) = plugin_manager.as_ref() {
827 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
828 }
829 yield Err(compaction_err);
830 return;
831 }
832 }
833 }
834 Err(e) => {
835 #[cfg(feature = "plugins")]
836 if let Some(manager) = plugin_manager.as_ref() {
837 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
838 }
839 yield Err(e);
840 return;
841 }
842 };
843
844 use futures::StreamExt;
846 let mut transfer_target: Option<(String, String)> = None;
847 let mut streamed_content = HashMap::new();
848
849 while let Some(result) = {
850 match cancellation_token.as_ref() {
855 Some(token) => {
856 tokio::select! {
857 biased;
858 _ = token.cancelled() => {
859 tracing::info!("cancellation fired during agent stream await");
860 #[cfg(feature = "plugins")]
861 if let Some(manager) = plugin_manager.as_ref() {
862 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
863 }
864 return;
865 }
866 result = agent_stream.next() => result,
867 }
868 }
869 None => agent_stream.next().await,
870 }
871 } {
872 match result {
873 Ok(mut event) => {
874 #[cfg(feature = "plugins")]
875 if let Some(manager) = plugin_manager.as_ref() {
876 match manager
877 .run_on_event(
878 ctx.clone() as Arc<dyn adk_core::InvocationContext>,
879 event.clone(),
880 )
881 .await
882 {
883 Ok(Some(modified)) => {
884 event = modified;
885 }
886 Ok(None) => {}
887 Err(e) => {
888 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
889 yield Err(e);
890 return;
891 }
892 }
893 }
894
895 preserve_streamed_content(&mut streamed_content, &mut event);
896
897 if let Some(target) = &event.actions.transfer_to_agent {
899 let source = if event.author.is_empty() {
900 agent_to_run.name()
901 } else {
902 &event.author
903 };
904 if let Some(allowed) = root_agent.transfer_targets_for(source)
905 && !allowed.contains(target)
906 {
907 if root_agent.strict_transfer_policy() {
908 yield Err(adk_core::AdkError::new(
909 adk_core::ErrorComponent::Agent,
910 adk_core::ErrorCategory::Forbidden,
911 "agent.transfer.target_forbidden",
912 format!(
913 "agent '{source}' cannot hand off to '{target}'; allowed targets: {}",
914 allowed.join(", ")
915 ),
916 ));
917 return;
918 }
919 tracing::warn!(source, target, "handoff target rejected by root policy");
920 } else {
921 transfer_target = Some((source.to_string(), target.clone()));
922 }
923 }
924
925 if !event.actions.state_delta.is_empty() {
931 ctx.mutable_session().apply_state_delta(&event.actions.state_delta);
932 }
933
934 ctx.mutable_session().append_event(event.clone());
936
937 if !event.llm_response.partial
944 && let Err(e) = session_service
945 .append_event_for_identity(adk_session::AppendEventRequest {
946 identity: identity.clone(),
947 event: event.clone(),
948 })
949 .await {
950 #[cfg(feature = "plugins")]
951 if let Some(manager) = plugin_manager.as_ref() {
952 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
953 }
954 yield Err(e);
955 return;
956 }
957 yield Ok(event);
958 }
959 Err(e) => {
960 #[cfg(feature = "plugins")]
961 if let Some(manager) = plugin_manager.as_ref() {
962 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
963 }
964 yield Err(e);
965 return;
966 }
967 }
968 }
969
970 const DEFAULT_MAX_TRANSFER_DEPTH: u32 = 10;
977 let max_depth = run_config.max_transfer_depth.unwrap_or(DEFAULT_MAX_TRANSFER_DEPTH);
978 let mut transfer_depth: u32 = 0;
979 let mut current_transfer_target = transfer_target;
980
981 while let Some((transfer_source, target_name)) = current_transfer_target.take() {
982 transfer_depth += 1;
983 if transfer_depth > max_depth {
984 tracing::warn!(
985 depth = transfer_depth,
986 target = %target_name,
987 "max transfer depth exceeded, stopping transfer chain"
988 );
989 if root_agent.strict_transfer_policy() {
990 yield Err(adk_core::AdkError::new(
991 adk_core::ErrorComponent::Agent,
992 adk_core::ErrorCategory::InvalidInput,
993 "agent.transfer.depth_exceeded",
994 format!(
995 "maximum handoff depth {max_depth} exceeded while transferring to '{target_name}'"
996 ),
997 ));
998 return;
999 }
1000 break;
1001 }
1002
1003 let governance = adk_core::AgentTransferRequest {
1004 invocation_id: invocation_id.clone(),
1005 from: transfer_source.clone(),
1006 to: target_name.clone(),
1007 depth: transfer_depth,
1008 };
1009 match root_agent.govern_transfer(&governance).await {
1010 Ok(adk_core::AgentTransferDecision::Allow) => {}
1011 Ok(adk_core::AgentTransferDecision::Deny { reason }) => {
1012 yield Err(adk_core::AdkError::new(
1013 adk_core::ErrorComponent::Agent,
1014 adk_core::ErrorCategory::Forbidden,
1015 "agent.transfer.denied",
1016 format!(
1017 "handoff from '{transfer_source}' to '{target_name}' was denied: {reason}"
1018 ),
1019 ));
1020 return;
1021 }
1022 Err(error) => {
1023 yield Err(error);
1024 return;
1025 }
1026 }
1027
1028 let target_agent = match Self::find_agent(&root_agent, &target_name) {
1029 Some(a) => a,
1030 None => {
1031 tracing::warn!(target = %target_name, "transfer target not found in agent tree");
1032 if root_agent.strict_transfer_policy() {
1033 yield Err(adk_core::AdkError::new(
1034 adk_core::ErrorComponent::Agent,
1035 adk_core::ErrorCategory::NotFound,
1036 "agent.transfer.target_not_found",
1037 format!(
1038 "handoff target '{target_name}' was not found in the agent tree"
1039 ),
1040 ));
1041 return;
1042 }
1043 break;
1044 }
1045 };
1046
1047 let mut transfer_run_config = run_config.clone();
1052 if let Some(targets) = root_agent.transfer_targets_for(&target_name) {
1053 transfer_run_config.transfer_targets = targets;
1054 transfer_run_config.parent_agent = None;
1055 } else {
1056 let (parent_name, peer_names) =
1057 Self::compute_transfer_context(&root_agent, &target_name);
1058 let mut targets = Vec::new();
1059 if let Some(ref parent) = parent_name {
1060 targets.push(parent.clone());
1061 }
1062 targets.extend(peer_names);
1063 transfer_run_config.transfer_targets = targets;
1064 transfer_run_config.parent_agent = parent_name;
1065 }
1066 root_agent.configure_run(&target_name, &mut transfer_run_config);
1067
1068 let transfer_invocation_id = format!("inv-{}", uuid::Uuid::new_v4());
1070 let mut transfer_ctx = match InvocationContext::with_mutable_session(
1071 transfer_invocation_id.clone(),
1072 target_agent.clone(),
1073 ctx.user_id().to_string(),
1074 ctx.app_name().to_string(),
1075 ctx.session_id().to_string(),
1076 effective_user_content.clone(),
1077 ctx.mutable_session().clone(),
1078 ) {
1079 Ok(ctx) => ctx,
1080 Err(e) => {
1081 yield Err(e);
1082 return;
1083 }
1084 };
1085 transfer_ctx = transfer_ctx.with_orchestration_root_invocation_id(
1086 adk_core::InvocationContext::orchestration_root_invocation_id(ctx.as_ref())
1087 .to_string(),
1088 );
1089
1090 #[cfg(feature = "artifacts")]
1091 if let Some(ref service) = artifact_service_clone {
1092 let scoped = adk_artifact::ScopedArtifacts::new(
1093 service.clone(),
1094 ctx.app_name().to_string(),
1095 ctx.user_id().to_string(),
1096 ctx.session_id().to_string(),
1097 );
1098 transfer_ctx = transfer_ctx.with_artifacts(Arc::new(scoped));
1099 }
1100 if let Some(ref memory) = memory_service_clone {
1101 transfer_ctx = transfer_ctx.with_memory(memory.clone());
1102 }
1103 transfer_ctx = transfer_ctx.with_run_config(transfer_run_config);
1104 if let Some(rc) = request_context.clone() {
1105 transfer_ctx = transfer_ctx.with_request_context(rc);
1106 }
1107 if let Some(token) = cancellation_token.as_ref() {
1108 transfer_ctx = transfer_ctx.with_cancellation_token(token.clone());
1109 }
1110 if let Some(shared_state) = adk_core::CallbackContext::shared_state(ctx.as_ref()) {
1111 transfer_ctx = transfer_ctx.with_shared_state(shared_state);
1112 }
1113
1114 let transfer_ctx = Arc::new(transfer_ctx);
1115
1116 let mut transfer_stream = match target_agent.run(transfer_ctx.clone()).await {
1118 Ok(s) => s,
1119 Err(e) => {
1120 #[cfg(feature = "plugins")]
1121 if let Some(manager) = plugin_manager.as_ref() {
1122 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1123 }
1124 yield Err(e);
1125 return;
1126 }
1127 };
1128
1129 while let Some(result) = {
1131 match cancellation_token.as_ref() {
1134 Some(token) => {
1135 tokio::select! {
1136 biased;
1137 _ = token.cancelled() => {
1138 tracing::info!("cancellation fired during transferred agent stream await");
1139 #[cfg(feature = "plugins")]
1140 if let Some(manager) = plugin_manager.as_ref() {
1141 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1142 }
1143 return;
1144 }
1145 result = transfer_stream.next() => result,
1146 }
1147 }
1148 None => transfer_stream.next().await,
1149 }
1150 } {
1151 match result {
1152 Ok(mut event) => {
1153 #[cfg(feature = "plugins")]
1154 if let Some(manager) = plugin_manager.as_ref() {
1155 match manager
1156 .run_on_event(
1157 transfer_ctx.clone() as Arc<dyn adk_core::InvocationContext>,
1158 event.clone(),
1159 )
1160 .await
1161 {
1162 Ok(Some(modified)) => {
1163 event = modified;
1164 }
1165 Ok(None) => {}
1166 Err(e) => {
1167 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1168 yield Err(e);
1169 return;
1170 }
1171 }
1172 }
1173
1174 preserve_streamed_content(&mut streamed_content, &mut event);
1175
1176 if let Some(target) = &event.actions.transfer_to_agent {
1178 let source = if event.author.is_empty() {
1179 target_agent.name()
1180 } else {
1181 &event.author
1182 };
1183 if let Some(allowed) = root_agent.transfer_targets_for(source)
1184 && !allowed.contains(target)
1185 {
1186 if root_agent.strict_transfer_policy() {
1187 yield Err(adk_core::AdkError::new(
1188 adk_core::ErrorComponent::Agent,
1189 adk_core::ErrorCategory::Forbidden,
1190 "agent.transfer.target_forbidden",
1191 format!(
1192 "agent '{source}' cannot hand off to '{target}'; allowed targets: {}",
1193 allowed.join(", ")
1194 ),
1195 ));
1196 return;
1197 }
1198 tracing::warn!(source, target, "handoff target rejected by root policy");
1199 } else {
1200 current_transfer_target =
1201 Some((source.to_string(), target.clone()));
1202 }
1203 }
1204
1205 if !event.actions.state_delta.is_empty() {
1207 transfer_ctx.mutable_session().apply_state_delta(&event.actions.state_delta);
1208 }
1209
1210 transfer_ctx.mutable_session().append_event(event.clone());
1212
1213 if !event.llm_response.partial
1214 && let Err(e) = session_service
1215 .append_event_for_identity(adk_session::AppendEventRequest {
1216 identity: identity.clone(),
1217 event: event.clone(),
1218 })
1219 .await {
1220 #[cfg(feature = "plugins")]
1221 if let Some(manager) = plugin_manager.as_ref() {
1222 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1223 }
1224 yield Err(e);
1225 return;
1226 }
1227 yield Ok(event);
1228 }
1229 Err(e) => {
1230 #[cfg(feature = "plugins")]
1231 if let Some(manager) = plugin_manager.as_ref() {
1232 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1233 }
1234 yield Err(e);
1235 return;
1236 }
1237 }
1238 }
1239 }
1240
1241 if let Some(ref compaction_cfg) = compaction_config {
1245 let event_count = ctx.mutable_session().as_ref().events_len();
1246
1247 if event_count > 0 {
1248 let all_events = ctx.mutable_session().as_ref().events_snapshot();
1249 let invocation_count = all_events.iter().filter(|e| e.author == "user").count()
1250 as u32;
1251
1252 if invocation_count > 0
1253 && invocation_count.is_multiple_of(compaction_cfg.compaction_interval)
1254 {
1255 let overlap = compaction_cfg.overlap_size as usize;
1258
1259 let user_msg_indices: Vec<usize> = all_events.iter()
1261 .enumerate()
1262 .filter(|(_, e)| e.author == "user")
1263 .map(|(i, _)| i)
1264 .collect();
1265
1266 let compact_up_to = if overlap == 0 {
1269 all_events.len()
1270 } else if user_msg_indices.len() > overlap {
1271 user_msg_indices[user_msg_indices.len() - overlap]
1273 } else {
1274 0
1276 };
1277
1278 if compact_up_to > 0 {
1279 let events_to_compact = &all_events[..compact_up_to];
1280
1281 match compaction_cfg.summarizer.summarize_events(events_to_compact).await {
1282 Ok(Some(compaction_event)) => {
1283 if let Err(e) = session_service
1285 .append_event_for_identity(adk_session::AppendEventRequest {
1286 identity: identity.clone(),
1287 event: compaction_event.clone(),
1288 })
1289 .await {
1290 tracing::warn!(error = %e, "Failed to persist compaction event");
1291 } else {
1292 tracing::info!(
1293 compacted_events = compact_up_to,
1294 "Context compaction completed"
1295 );
1296 }
1297 }
1298 Ok(None) => {
1299 tracing::debug!("Compaction summarizer returned no result");
1300 }
1301 Err(e) => {
1302 tracing::warn!(error = %e, "Context compaction failed");
1304 }
1305 }
1306 }
1307 }
1308 }
1309 }
1310
1311 #[cfg(feature = "plugins")]
1312 if let Some(manager) = plugin_manager.as_ref() {
1313 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1314 }
1315 };
1316
1317 Ok(Box::pin(s))
1318 }
1319
1320 pub async fn run_str(
1330 &self,
1331 user_id: &str,
1332 session_id: &str,
1333 user_content: Content,
1334 ) -> Result<EventStream> {
1335 let user_id = UserId::try_from(user_id)?;
1336 let session_id = SessionId::try_from(session_id)?;
1337 self.run(user_id, session_id, user_content).await
1338 }
1339
1340 pub fn interrupt(&self, session_id: &str) -> bool {
1365 let runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
1366 let matching: Vec<&ActiveRun> =
1367 runs.values().filter(|run| run.identity.session_id.as_ref() == session_id).collect();
1368 if matching.is_empty() {
1369 tracing::debug!(session.id = session_id, "no active run to interrupt");
1370 return false;
1371 }
1372 tracing::info!(
1373 session.id = session_id,
1374 run.count = matching.len(),
1375 "interrupting running agent"
1376 );
1377 for run in matching {
1378 run.token.cancel();
1379 }
1380 true
1381 }
1382
1383 pub fn interrupt_identity(&self, app_name: &str, user_id: &str, session_id: &str) -> bool {
1395 let runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
1396 let mut cancelled = false;
1397 for run in runs.values() {
1398 if run.identity.app_name.as_ref() == app_name
1399 && run.identity.user_id.as_ref() == user_id
1400 && run.identity.session_id.as_ref() == session_id
1401 {
1402 run.token.cancel();
1403 cancelled = true;
1404 }
1405 }
1406 if !cancelled {
1407 tracing::debug!(
1408 app.name = app_name,
1409 user.id = user_id,
1410 session.id = session_id,
1411 "no active run to interrupt"
1412 );
1413 }
1414 cancelled
1415 }
1416
1417 pub fn active_runs(&self) -> Vec<adk_core::AdkIdentity> {
1422 let runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
1423 runs.values().map(|run| run.identity.clone()).collect()
1424 }
1425
1426 pub fn active_session_ids(&self) -> Vec<String> {
1431 let runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
1432 let mut ids: Vec<String> =
1433 runs.values().map(|run| run.identity.session_id.as_ref().to_string()).collect();
1434 ids.sort();
1435 ids.dedup();
1436 ids
1437 }
1438
1439 #[cfg(feature = "context-compaction")]
1444 pub fn context_compaction(&self) -> Option<&crate::compaction::CompactionConfig> {
1445 self.context_compaction.as_deref()
1446 }
1447
1448 pub fn find_agent_to_run(
1450 root_agent: &Arc<dyn Agent>,
1451 session: &dyn adk_session::Session,
1452 ) -> Arc<dyn Agent> {
1453 let events = session.events();
1455 for i in (0..events.len()).rev() {
1456 if let Some(event) = events.at(i) {
1457 if let Some(target_name) = &event.actions.transfer_to_agent
1459 && let Some(agent) = Self::find_agent(root_agent, target_name)
1460 {
1461 return agent;
1462 }
1463
1464 if event.author == "user" {
1465 continue;
1466 }
1467
1468 if let Some(agent) = Self::find_agent(root_agent, &event.author) {
1470 if Self::is_transferable(root_agent, &agent) {
1472 return agent;
1473 }
1474 }
1475 }
1476 }
1477
1478 root_agent.clone()
1480 }
1481
1482 fn is_transferable(root_agent: &Arc<dyn Agent>, agent: &Arc<dyn Agent>) -> bool {
1498 fn path_supports_transfer(current: &Arc<dyn Agent>, target: &str) -> Option<bool> {
1504 if current.name() == target {
1505 return Some(current.supports_agent_transfer());
1506 }
1507 for sub in current.sub_agents() {
1508 if let Some(sub_ok) = path_supports_transfer(sub, target) {
1509 return Some(current.supports_agent_transfer() && sub_ok);
1510 }
1511 }
1512 None
1513 }
1514
1515 path_supports_transfer(root_agent, agent.name()).unwrap_or(true)
1516 }
1517
1518 pub fn find_agent(current: &Arc<dyn Agent>, target_name: &str) -> Option<Arc<dyn Agent>> {
1520 if current.name() == target_name {
1521 return Some(current.clone());
1522 }
1523
1524 for sub_agent in current.sub_agents() {
1525 if let Some(found) = Self::find_agent(sub_agent, target_name) {
1526 return Some(found);
1527 }
1528 }
1529
1530 None
1531 }
1532
1533 pub fn compute_transfer_context(
1539 root: &Arc<dyn Agent>,
1540 target_name: &str,
1541 ) -> (Option<String>, Vec<String>) {
1542 if root.name() == target_name {
1544 return (None, Vec::new());
1545 }
1546
1547 fn find_parent(current: &Arc<dyn Agent>, target: &str) -> Option<Arc<dyn Agent>> {
1549 for sub in current.sub_agents() {
1550 if sub.name() == target {
1551 return Some(current.clone());
1552 }
1553 if let Some(found) = find_parent(sub, target) {
1554 return Some(found);
1555 }
1556 }
1557 None
1558 }
1559
1560 match find_parent(root, target_name) {
1561 Some(parent) => {
1562 let parent_name = parent.name().to_string();
1563 let peers: Vec<String> = parent
1564 .sub_agents()
1565 .iter()
1566 .filter(|a| a.name() != target_name)
1567 .map(|a| a.name().to_string())
1568 .collect();
1569 (Some(parent_name), peers)
1570 }
1571 None => (None, Vec::new()),
1572 }
1573 }
1574}
1575
1576#[cfg(test)]
1577mod streamed_content_tests {
1578 use super::preserve_streamed_content;
1579 use adk_core::{Content, Event, Part};
1580 use std::collections::HashMap;
1581
1582 fn text(event: &Event) -> String {
1583 event
1584 .content()
1585 .map(|content| content.parts.iter().filter_map(Part::text).collect())
1586 .unwrap_or_default()
1587 }
1588
1589 #[test]
1590 fn final_empty_event_preserves_streamed_text_for_persistence() {
1591 let mut accumulated = HashMap::new();
1592 let mut first = Event::with_id("response-1", "inv-1");
1593 first.llm_response.partial = true;
1594 first.llm_response.content = Some(Content::new("model").with_text("Verify "));
1595 preserve_streamed_content(&mut accumulated, &mut first);
1596
1597 let mut second = Event::with_id("response-1", "inv-1");
1598 second.llm_response.partial = true;
1599 second.llm_response.content = Some(Content::new("model").with_text("the invoice."));
1600 preserve_streamed_content(&mut accumulated, &mut second);
1601
1602 let mut final_event = Event::with_id("response-1", "inv-1");
1603 final_event.llm_response.turn_complete = true;
1604 preserve_streamed_content(&mut accumulated, &mut final_event);
1605
1606 assert_eq!(text(&final_event), "Verify the invoice.");
1607 assert!(accumulated.is_empty());
1608 }
1609}