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, InvocationId,
7 Memory, 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
55pub struct RunnerInvocation {
60 invocation_id: InvocationId,
61 events: EventStream,
62}
63
64impl RunnerInvocation {
65 pub fn invocation_id(&self) -> &InvocationId {
67 &self.invocation_id
68 }
69
70 pub fn into_events(self) -> EventStream {
72 self.events
73 }
74}
75
76impl std::fmt::Debug for RunnerInvocation {
77 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78 formatter
79 .debug_struct("RunnerInvocation")
80 .field("invocation_id", &self.invocation_id)
81 .finish_non_exhaustive()
82 }
83}
84
85impl Drop for ActiveRunCleanup {
86 fn drop(&mut self) {
87 let mut runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
88 runs.remove(&self.run_id);
89 }
90}
91use tracing::Instrument;
92
93pub struct RunnerConfig {
97 pub app_name: String,
99 pub agent: Arc<dyn Agent>,
101 pub session_service: Arc<dyn SessionService>,
103 #[cfg(feature = "artifacts")]
104 pub artifact_service: Option<Arc<dyn ArtifactService>>,
106 pub memory_service: Option<Arc<dyn Memory>>,
108 #[cfg(feature = "plugins")]
109 pub plugin_manager: Option<Arc<PluginManager>>,
111 #[allow(dead_code)]
114 pub run_config: Option<RunConfig>,
115 pub compaction_config: Option<adk_core::EventsCompactionConfig>,
119 pub context_cache_config: Option<ContextCacheConfig>,
127 pub cache_capable: Option<Arc<dyn CacheCapable>>,
130 pub request_context: Option<adk_core::RequestContext>,
134 pub cancellation_token: Option<CancellationToken>,
139 pub intra_compaction_config: Option<adk_core::IntraCompactionConfig>,
143 pub intra_compaction_summarizer: Option<Arc<dyn adk_core::BaseEventsSummarizer>>,
146 #[cfg(feature = "context-compaction")]
154 pub context_compaction: Option<crate::compaction::CompactionConfig>,
155}
156
157pub struct Runner {
163 app_name: String,
164 root_agent: Arc<dyn Agent>,
165 session_service: Arc<dyn SessionService>,
166 #[cfg(feature = "artifacts")]
167 artifact_service: Option<Arc<dyn ArtifactService>>,
168 memory_service: Option<Arc<dyn Memory>>,
169 #[cfg(feature = "plugins")]
170 plugin_manager: Option<Arc<PluginManager>>,
171 #[cfg(feature = "skills")]
172 skill_injector: Option<Arc<SkillInjector>>,
173 run_config: RunConfig,
174 compaction_config: Option<adk_core::EventsCompactionConfig>,
175 context_cache_config: Option<ContextCacheConfig>,
176 cache_capable: Option<Arc<dyn CacheCapable>>,
177 cache_manager: Option<Arc<tokio::sync::Mutex<CacheManager>>>,
178 request_context: Option<adk_core::RequestContext>,
179 cancellation_token: Option<CancellationToken>,
180 intra_compactor: Option<Arc<crate::intra_compaction::IntraInvocationCompactor>>,
181 #[cfg(feature = "context-compaction")]
183 context_compaction: Option<Arc<crate::compaction::CompactionConfig>>,
184 active_runs: ActiveRuns,
187 next_run_id: Arc<std::sync::atomic::AtomicU64>,
188 pub(crate) external_session_locks: Arc<
191 std::sync::Mutex<
192 std::collections::HashMap<String, std::sync::Weak<tokio::sync::Mutex<()>>>,
193 >,
194 >,
195}
196
197impl Runner {
198 pub fn builder() -> crate::builder::RunnerConfigBuilder<
214 crate::builder::NoAppName,
215 crate::builder::NoAgent,
216 crate::builder::NoSessionService,
217 > {
218 crate::builder::RunnerConfigBuilder::new()
219 }
220
221 pub fn new(config: RunnerConfig) -> Result<Self> {
225 let run_config = config.run_config.unwrap_or_default();
226
227 let effective_cache_config = config
230 .context_cache_config
231 .or_else(|| config.cache_capable.as_ref().map(|_| ContextCacheConfig::default()));
232
233 let cache_manager = effective_cache_config
234 .as_ref()
235 .map(|c| Arc::new(tokio::sync::Mutex::new(CacheManager::new(c.clone()))));
236
237 let intra_compactor = config.intra_compaction_config.as_ref().and_then(|ic_config| {
238 config.intra_compaction_summarizer.as_ref().map(|summarizer| {
239 Arc::new(crate::intra_compaction::IntraInvocationCompactor::new(
240 ic_config.clone(),
241 summarizer.clone(),
242 ))
243 })
244 });
245
246 Ok(Self {
247 app_name: config.app_name,
248 root_agent: config.agent,
249 session_service: config.session_service,
250 #[cfg(feature = "artifacts")]
251 artifact_service: config.artifact_service,
252 memory_service: config.memory_service,
253 #[cfg(feature = "plugins")]
254 plugin_manager: config.plugin_manager,
255 #[cfg(feature = "skills")]
256 skill_injector: None,
257 run_config,
258 compaction_config: config.compaction_config,
259 context_cache_config: effective_cache_config,
260 cache_capable: config.cache_capable,
261 cache_manager,
262 request_context: config.request_context,
263 cancellation_token: config.cancellation_token,
264 intra_compactor,
265 #[cfg(feature = "context-compaction")]
266 context_compaction: config.context_compaction.map(Arc::new),
267 active_runs: Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
268 next_run_id: Arc::new(std::sync::atomic::AtomicU64::new(1)),
269 external_session_locks: Arc::new(std::sync::Mutex::new(
270 std::collections::HashMap::new(),
271 )),
272 })
273 }
274
275 pub(crate) fn root_agent(&self) -> Arc<dyn Agent> {
277 Arc::clone(&self.root_agent)
278 }
279
280 #[cfg(feature = "skills")]
284 pub fn with_skill_injector(mut self, injector: SkillInjector) -> Self {
285 self.skill_injector = Some(Arc::new(injector));
286 self
287 }
288
289 #[cfg(feature = "skills")]
291 #[deprecated(note = "Use with_auto_skills_mut instead")]
292 pub fn with_auto_skills(
293 mut self,
294 root: impl AsRef<std::path::Path>,
295 config: SkillInjectorConfig,
296 ) -> adk_skill::SkillResult<Self> {
297 self.with_auto_skills_mut(root, config)?;
298 Ok(self)
299 }
300
301 #[cfg(feature = "skills")]
307 pub fn with_auto_skills_mut(
308 &mut self,
309 root: impl AsRef<std::path::Path>,
310 config: SkillInjectorConfig,
311 ) -> adk_skill::SkillResult<()> {
312 let injector = SkillInjector::from_root(root, config)?;
313 self.skill_injector = Some(Arc::new(injector));
314 Ok(())
315 }
316
317 pub async fn run(
322 &self,
323 user_id: UserId,
324 session_id: SessionId,
325 user_content: Content,
326 ) -> Result<EventStream> {
327 self.run_with_config(user_id, session_id, user_content, None).await
328 }
329
330 pub fn app_name(&self) -> &str {
332 &self.app_name
333 }
334
335 pub fn session_service(&self) -> &Arc<dyn adk_session::SessionService> {
337 &self.session_service
338 }
339
340 pub fn run_config(&self) -> &RunConfig {
344 &self.run_config
345 }
346
347 pub async fn run_with_config(
358 &self,
359 user_id: UserId,
360 session_id: SessionId,
361 user_content: Content,
362 run_config: Option<RunConfig>,
363 ) -> Result<EventStream> {
364 Ok(self
365 .start_with_config(user_id, session_id, user_content, run_config)
366 .await?
367 .into_events())
368 }
369
370 pub async fn start_with_config(
381 &self,
382 user_id: UserId,
383 session_id: SessionId,
384 user_content: Content,
385 run_config: Option<RunConfig>,
386 ) -> Result<RunnerInvocation> {
387 let invocation_id = InvocationId::try_from(format!("inv-{}", uuid::Uuid::new_v4()))?;
388 let stream_invocation_id = invocation_id.to_string();
389 let app_name = self.app_name.clone();
390 let typed_app_name = AppName::try_from(app_name.clone())?;
391 let session_service = self.session_service.clone();
392 let root_agent = self.root_agent.clone();
393 #[cfg(feature = "artifacts")]
394 let artifact_service = self.artifact_service.clone();
395 let memory_service = self.memory_service.clone();
396 #[cfg(feature = "plugins")]
397 let plugin_manager = self.plugin_manager.clone();
398 #[cfg(feature = "skills")]
399 let skill_injector = self.skill_injector.clone();
400 let mut run_config = run_config.unwrap_or_else(|| self.run_config.clone());
401 let compaction_config = self.compaction_config.clone();
402 let context_cache_config = self.context_cache_config.clone();
403 let cache_capable = self.cache_capable.clone();
404 let cache_manager_ref = self.cache_manager.clone();
405 let request_context = self.request_context.clone();
406 let intra_compactor = self.intra_compactor.clone();
407 #[cfg(feature = "context-compaction")]
408 let context_compaction = self.context_compaction.clone();
409
410 let identity =
413 adk_core::AdkIdentity::new(typed_app_name.clone(), user_id.clone(), session_id.clone());
414
415 let session_token = match self.cancellation_token.as_ref() {
424 Some(global) => global.child_token(),
425 None => CancellationToken::new(),
426 };
427 let run_id = self.next_run_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
428 {
429 let mut runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
430 runs.insert(
431 run_id,
432 ActiveRun { identity: identity.clone(), token: session_token.clone() },
433 );
434 }
435 let active_runs = self.active_runs.clone();
436 let effective_token = Some(session_token);
437
438 let cleanup = ActiveRunCleanup { active_runs: active_runs.clone(), run_id };
443
444 let s = stream! {
445 let _cleanup = cleanup;
446
447 let cancellation_token = effective_token;
449 let session = match session_service
451 .get(adk_session::GetRequest {
452 app_name: app_name.clone(),
453 user_id: user_id.to_string(),
454 session_id: session_id.to_string(),
455 num_recent_events: run_config.history_max_events,
456 after: None,
457 })
458 .await
459 {
460 Ok(s) => s,
461 Err(e) => {
462 yield Err(e);
463 return;
464 }
465 };
466
467 let agent_to_run = Self::find_agent_to_run(&root_agent, session.as_ref());
469
470 root_agent.configure_run(agent_to_run.name(), &mut run_config);
473 if let Some(targets) = root_agent.transfer_targets_for(agent_to_run.name()) {
474 run_config.transfer_targets = targets;
475 run_config.parent_agent = None;
476 }
477
478 #[cfg(feature = "artifacts")]
480 let artifact_service_clone = artifact_service.clone();
481 let memory_service_clone = memory_service.clone();
482
483 let invocation_id = stream_invocation_id;
485 #[cfg(any(feature = "skills", feature = "plugins"))]
486 let mut effective_user_content = user_content.clone();
487 #[cfg(not(any(feature = "skills", feature = "plugins")))]
488 let effective_user_content = user_content.clone();
489 #[cfg(feature = "skills")]
490 let mut selected_skill_name = String::new();
491 #[cfg(not(feature = "skills"))]
492 let selected_skill_name = String::new();
493 #[cfg(feature = "skills")]
494 let mut selected_skill_id = String::new();
495 #[cfg(not(feature = "skills"))]
496 let selected_skill_id = String::new();
497
498 #[cfg(feature = "skills")]
499 if let Some(injector) = skill_injector.as_ref()
500 && let Some(matched) = adk_skill::apply_skill_injection(
501 &mut effective_user_content,
502 injector.index(),
503 injector.policy(),
504 injector.max_injected_chars(),
505 ) {
506 selected_skill_name = matched.skill.name;
507 selected_skill_id = matched.skill.id;
508 }
509
510 let mut invocation_ctx = match InvocationContext::new_typed(
511 invocation_id.clone(),
512 agent_to_run.clone(),
513 user_id.clone(),
514 typed_app_name.clone(),
515 session_id.clone(),
516 effective_user_content.clone(),
517 Arc::from(session),
518 ) {
519 Ok(ctx) => ctx,
520 Err(e) => {
521 yield Err(e);
522 return;
523 }
524 };
525
526 #[cfg(feature = "artifacts")]
528 if let Some(service) = artifact_service {
529 let scoped = adk_artifact::ScopedArtifacts::new(
531 service,
532 app_name.clone(),
533 user_id.to_string(),
534 session_id.to_string(),
535 );
536 invocation_ctx = invocation_ctx.with_artifacts(Arc::new(scoped));
537 }
538 if let Some(memory) = memory_service {
539 invocation_ctx = invocation_ctx.with_memory(memory);
540 }
541
542 invocation_ctx = invocation_ctx.with_run_config(run_config.clone());
544
545 if let Some(rc) = request_context.clone() {
547 invocation_ctx = invocation_ctx.with_request_context(rc);
548 }
549
550 if let Some(token) = cancellation_token.as_ref() {
552 invocation_ctx = invocation_ctx.with_cancellation_token(token.clone());
553 }
554
555 let mut ctx = Arc::new(invocation_ctx);
556
557 #[cfg(feature = "plugins")]
558 if let Some(manager) = plugin_manager.as_ref() {
559 match manager
560 .run_before_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>)
561 .await
562 {
563 Ok(Some(content)) => {
564 let mut early_event = adk_core::Event::new(ctx.invocation_id());
565 early_event.author = agent_to_run.name().to_string();
566 early_event.llm_response.content = Some(content);
567
568 ctx.mutable_session().append_event(early_event.clone());
569 if let Err(e) = session_service
570 .append_event_for_identity(adk_session::AppendEventRequest {
571 identity: identity.clone(),
572 event: early_event.clone(),
573 })
574 .await {
575 yield Err(e);
576 return;
577 }
578
579 yield Ok(early_event);
580 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
581 return;
582 }
583 Ok(None) => {}
584 Err(e) => {
585 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
586 yield Err(e);
587 return;
588 }
589 }
590
591 match manager
592 .run_on_user_message(
593 ctx.clone() as Arc<dyn adk_core::InvocationContext>,
594 effective_user_content.clone(),
595 )
596 .await
597 {
598 Ok(Some(modified)) => {
599 effective_user_content = modified;
600
601 let mut refreshed_ctx = match InvocationContext::with_mutable_session(
602 ctx.invocation_id().to_string(),
603 agent_to_run.clone(),
604 ctx.user_id().to_string(),
605 ctx.app_name().to_string(),
606 ctx.session_id().to_string(),
607 effective_user_content.clone(),
608 ctx.mutable_session().clone(),
609 ) {
610 Ok(ctx) => ctx,
611 Err(e) => {
612 yield Err(e);
613 return;
614 }
615 };
616 refreshed_ctx = refreshed_ctx.with_orchestration_root_invocation_id(
617 adk_core::InvocationContext::orchestration_root_invocation_id(
618 ctx.as_ref(),
619 )
620 .to_string(),
621 );
622
623 #[cfg(feature = "artifacts")]
624 if let Some(service) = artifact_service_clone.clone() {
625 let scoped = adk_artifact::ScopedArtifacts::new(
626 service,
627 ctx.app_name().to_string(),
628 ctx.user_id().to_string(),
629 ctx.session_id().to_string(),
630 );
631 refreshed_ctx = refreshed_ctx.with_artifacts(Arc::new(scoped));
632 }
633 if let Some(memory) = memory_service_clone.clone() {
634 refreshed_ctx = refreshed_ctx.with_memory(memory);
635 }
636 refreshed_ctx = refreshed_ctx.with_run_config(run_config.clone());
637 if let Some(rc) = request_context.clone() {
638 refreshed_ctx = refreshed_ctx.with_request_context(rc);
639 }
640 if let Some(token) = cancellation_token.as_ref() {
641 refreshed_ctx = refreshed_ctx.with_cancellation_token(token.clone());
642 }
643 ctx = Arc::new(refreshed_ctx);
644 }
645 Ok(None) => {}
646 Err(e) => {
647 if let Some(manager) = plugin_manager.as_ref() {
648 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
649 }
650 yield Err(e);
651 return;
652 }
653 }
654 }
655
656 let mut user_event = adk_core::Event::new(ctx.invocation_id());
658 user_event.author = "user".to_string();
659 user_event.llm_response.content = Some(effective_user_content.clone());
660
661 ctx.mutable_session().append_event(user_event.clone());
664
665 if let Err(e) = session_service
666 .append_event_for_identity(adk_session::AppendEventRequest {
667 identity: identity.clone(),
668 event: user_event,
669 })
670 .await {
671 #[cfg(feature = "plugins")]
672 if let Some(manager) = plugin_manager.as_ref() {
673 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
674 }
675 yield Err(e);
676 return;
677 }
678
679 if let (Some(cm_mutex), Some(cache_model)) = (&cache_manager_ref, &cache_capable) {
684 let should_refresh_cache = {
685 let cm = cm_mutex.lock().await;
686 cm.is_enabled() && (cm.active_cache_name().is_none() || cm.needs_refresh())
687 };
688
689 if should_refresh_cache {
690 let system_instruction = agent_to_run.description().to_string();
694 let tools = std::collections::HashMap::new();
695 let ttl = context_cache_config.as_ref().map_or(600, |c| c.ttl_seconds);
696
697 match cache_model.create_cache(&system_instruction, &tools, ttl).await {
698 Ok(name) => {
699 let old_cache = {
700 let mut cm = cm_mutex.lock().await;
701 let old = cm.clear_active_cache();
702 cm.set_active_cache(name);
703 old
704 };
705
706 if let Some(old) = old_cache
707 && let Err(e) = cache_model.delete_cache(&old).await {
708 tracing::warn!(
709 old_cache = %old,
710 error = %e,
711 "failed to delete old cache, proceeding with new cache"
712 );
713 }
714 }
715 Err(e) => {
716 tracing::warn!(
717 error = %e,
718 "cache creation failed, proceeding without cache"
719 );
720 }
721 }
722 }
723
724 let cache_name = {
726 let mut cm = cm_mutex.lock().await;
727 if cm.is_enabled() {
728 cm.record_invocation().map(str::to_string)
729 } else {
730 None
731 }
732 };
733
734 if let Some(cache_name) = cache_name {
735 run_config.cached_content = Some(cache_name);
736 let mut refreshed_ctx = match InvocationContext::with_mutable_session(
738 ctx.invocation_id().to_string(),
739 agent_to_run.clone(),
740 ctx.user_id().to_string(),
741 ctx.app_name().to_string(),
742 ctx.session_id().to_string(),
743 effective_user_content.clone(),
744 ctx.mutable_session().clone(),
745 ) {
746 Ok(ctx) => ctx,
747 Err(e) => {
748 yield Err(e);
749 return;
750 }
751 };
752 refreshed_ctx = refreshed_ctx.with_orchestration_root_invocation_id(
753 adk_core::InvocationContext::orchestration_root_invocation_id(ctx.as_ref())
754 .to_string(),
755 );
756 #[cfg(feature = "artifacts")]
757 if let Some(service) = artifact_service_clone.clone() {
758 let scoped = adk_artifact::ScopedArtifacts::new(
759 service,
760 ctx.app_name().to_string(),
761 ctx.user_id().to_string(),
762 ctx.session_id().to_string(),
763 );
764 refreshed_ctx = refreshed_ctx.with_artifacts(Arc::new(scoped));
765 }
766 if let Some(memory) = memory_service_clone.clone() {
767 refreshed_ctx = refreshed_ctx.with_memory(memory);
768 }
769 refreshed_ctx = refreshed_ctx.with_run_config(run_config.clone());
770 if let Some(rc) = request_context.clone() {
771 refreshed_ctx = refreshed_ctx.with_request_context(rc);
772 }
773 if let Some(token) = cancellation_token.as_ref() {
774 refreshed_ctx = refreshed_ctx.with_cancellation_token(token.clone());
775 }
776 ctx = Arc::new(refreshed_ctx);
777 }
778 }
779
780 if let Some(ref compactor) = intra_compactor {
784 compactor.reset_cycle();
785 let session_events = ctx.mutable_session().as_ref().events_snapshot();
786 match compactor.maybe_compact(&session_events).await {
787 Ok(Some(compacted_events)) => {
788 ctx.mutable_session().replace_events(compacted_events);
789 tracing::info!("intra-invocation compaction applied before agent execution");
790 }
791 Ok(None) => {} Err(e) => {
793 tracing::warn!(error = %e, "intra-invocation compaction check failed");
794 }
795 }
796 }
797
798 #[cfg(feature = "context-compaction")]
803 if let Some(ref cc_config) = context_compaction {
804 let session_events = ctx.mutable_session().events_snapshot();
805 let estimated = crate::compaction::estimate_event_tokens(&session_events);
806 if estimated > cc_config.context_budget {
807 tracing::info!(
808 estimated_tokens = estimated,
809 budget = cc_config.context_budget,
810 "context exceeds budget, applying proactive compaction"
811 );
812 match crate::compaction::apply_compaction_with_retry(cc_config, session_events).await {
813 Ok(compacted) => {
814 ctx.mutable_session().replace_events(compacted);
815 tracing::info!("proactive context compaction succeeded");
816 }
817 Err(e) => {
818 tracing::warn!(error = %e, "proactive context compaction failed, proceeding with full context");
821 }
822 }
823 }
824 }
825
826 let agent_span = tracing::info_span!(
828 "agent.execute",
829 "gcp.vertex.agent.invocation_id" = ctx.invocation_id(),
830 "gcp.vertex.agent.session_id" = ctx.session_id(),
831 "gcp.vertex.agent.event_id" = ctx.invocation_id(), "gen_ai.conversation.id" = ctx.session_id(),
833 "adk.app_name" = ctx.app_name(),
834 "adk.user_id" = ctx.user_id(),
835 "agent.name" = %agent_to_run.name(),
836 "adk.skills.selected_name" = %selected_skill_name,
837 "adk.skills.selected_id" = %selected_skill_id
838 );
839
840 let mut agent_stream = match agent_to_run.run(ctx.clone()).instrument(agent_span.clone()).await {
841 Ok(s) => s,
842 #[cfg(feature = "context-compaction")]
843 Err(e) if context_compaction.is_some() && crate::compaction::is_token_limit_error(&e) => {
844 let cc_config = context_compaction.as_ref().unwrap();
846 tracing::warn!(
847 error = %e,
848 "agent execution failed with token limit error, attempting compaction"
849 );
850 let session_events = ctx.mutable_session().events_snapshot();
851 match crate::compaction::apply_compaction_with_retry(cc_config, session_events).await {
852 Ok(compacted) => {
853 ctx.mutable_session().replace_events(compacted);
854 tracing::info!("context compaction succeeded after token limit error, retrying agent");
855 match agent_to_run.run(ctx.clone()).instrument(agent_span.clone()).await {
857 Ok(s) => s,
858 Err(retry_err) => {
859 #[cfg(feature = "plugins")]
860 if let Some(manager) = plugin_manager.as_ref() {
861 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
862 }
863 yield Err(retry_err);
864 return;
865 }
866 }
867 }
868 Err(compaction_err) => {
869 #[cfg(feature = "plugins")]
870 if let Some(manager) = plugin_manager.as_ref() {
871 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
872 }
873 yield Err(compaction_err);
874 return;
875 }
876 }
877 }
878 Err(e) => {
879 #[cfg(feature = "plugins")]
880 if let Some(manager) = plugin_manager.as_ref() {
881 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
882 }
883 yield Err(e);
884 return;
885 }
886 };
887
888 use futures::StreamExt;
890 let mut transfer_target: Option<(String, String)> = None;
891 let mut streamed_content = HashMap::new();
892
893 while let Some(result) = {
894 match cancellation_token.as_ref() {
899 Some(token) => {
900 tokio::select! {
901 biased;
902 _ = token.cancelled() => {
903 tracing::info!("cancellation fired during agent stream await");
904 #[cfg(feature = "plugins")]
905 if let Some(manager) = plugin_manager.as_ref() {
906 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
907 }
908 return;
909 }
910 result = agent_stream.next().instrument(agent_span.clone()) => result,
916 }
917 }
918 None => agent_stream.next().instrument(agent_span.clone()).await,
919 }
920 } {
921 match result {
922 Ok(mut event) => {
923 #[cfg(feature = "plugins")]
924 if let Some(manager) = plugin_manager.as_ref() {
925 match manager
926 .run_on_event(
927 ctx.clone() as Arc<dyn adk_core::InvocationContext>,
928 event.clone(),
929 )
930 .await
931 {
932 Ok(Some(modified)) => {
933 event = modified;
934 }
935 Ok(None) => {}
936 Err(e) => {
937 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
938 yield Err(e);
939 return;
940 }
941 }
942 }
943
944 preserve_streamed_content(&mut streamed_content, &mut event);
945
946 if let Some(target) = &event.actions.transfer_to_agent {
948 let source = if event.author.is_empty() {
949 agent_to_run.name()
950 } else {
951 &event.author
952 };
953 if let Some(allowed) = root_agent.transfer_targets_for(source)
954 && !allowed.contains(target)
955 {
956 if root_agent.strict_transfer_policy() {
957 yield Err(adk_core::AdkError::new(
958 adk_core::ErrorComponent::Agent,
959 adk_core::ErrorCategory::Forbidden,
960 "agent.transfer.target_forbidden",
961 format!(
962 "agent '{source}' cannot hand off to '{target}'; allowed targets: {}",
963 allowed.join(", ")
964 ),
965 ));
966 return;
967 }
968 tracing::warn!(source, target, "handoff target rejected by root policy");
969 } else {
970 transfer_target = Some((source.to_string(), target.clone()));
971 }
972 }
973
974 if !event.actions.state_delta.is_empty() {
980 ctx.mutable_session().apply_state_delta(&event.actions.state_delta);
981 }
982
983 ctx.mutable_session().append_event(event.clone());
985
986 if !event.llm_response.partial
993 && let Err(e) = session_service
994 .append_event_for_identity(adk_session::AppendEventRequest {
995 identity: identity.clone(),
996 event: event.clone(),
997 })
998 .await {
999 #[cfg(feature = "plugins")]
1000 if let Some(manager) = plugin_manager.as_ref() {
1001 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1002 }
1003 yield Err(e);
1004 return;
1005 }
1006 yield Ok(event);
1007 }
1008 Err(e) => {
1009 #[cfg(feature = "plugins")]
1010 if let Some(manager) = plugin_manager.as_ref() {
1011 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1012 }
1013 yield Err(e);
1014 return;
1015 }
1016 }
1017 }
1018
1019 const DEFAULT_MAX_TRANSFER_DEPTH: u32 = 10;
1026 let max_depth = run_config.max_transfer_depth.unwrap_or(DEFAULT_MAX_TRANSFER_DEPTH);
1027 let mut transfer_depth: u32 = 0;
1028 let mut current_transfer_target = transfer_target;
1029
1030 while let Some((transfer_source, target_name)) = current_transfer_target.take() {
1031 transfer_depth += 1;
1032 if transfer_depth > max_depth {
1033 tracing::warn!(
1034 depth = transfer_depth,
1035 target = %target_name,
1036 "max transfer depth exceeded, stopping transfer chain"
1037 );
1038 if root_agent.strict_transfer_policy() {
1039 yield Err(adk_core::AdkError::new(
1040 adk_core::ErrorComponent::Agent,
1041 adk_core::ErrorCategory::InvalidInput,
1042 "agent.transfer.depth_exceeded",
1043 format!(
1044 "maximum handoff depth {max_depth} exceeded while transferring to '{target_name}'"
1045 ),
1046 ));
1047 return;
1048 }
1049 break;
1050 }
1051
1052 let governance = adk_core::AgentTransferRequest {
1053 invocation_id: invocation_id.clone(),
1054 from: transfer_source.clone(),
1055 to: target_name.clone(),
1056 depth: transfer_depth,
1057 };
1058 match root_agent.govern_transfer(&governance).await {
1059 Ok(adk_core::AgentTransferDecision::Allow) => {}
1060 Ok(adk_core::AgentTransferDecision::Deny { reason }) => {
1061 yield Err(adk_core::AdkError::new(
1062 adk_core::ErrorComponent::Agent,
1063 adk_core::ErrorCategory::Forbidden,
1064 "agent.transfer.denied",
1065 format!(
1066 "handoff from '{transfer_source}' to '{target_name}' was denied: {reason}"
1067 ),
1068 ));
1069 return;
1070 }
1071 Err(error) => {
1072 yield Err(error);
1073 return;
1074 }
1075 }
1076
1077 let target_agent = match Self::find_agent(&root_agent, &target_name) {
1078 Some(a) => a,
1079 None => {
1080 tracing::warn!(target = %target_name, "transfer target not found in agent tree");
1081 if root_agent.strict_transfer_policy() {
1082 yield Err(adk_core::AdkError::new(
1083 adk_core::ErrorComponent::Agent,
1084 adk_core::ErrorCategory::NotFound,
1085 "agent.transfer.target_not_found",
1086 format!(
1087 "handoff target '{target_name}' was not found in the agent tree"
1088 ),
1089 ));
1090 return;
1091 }
1092 break;
1093 }
1094 };
1095
1096 let mut transfer_run_config = run_config.clone();
1101 if let Some(targets) = root_agent.transfer_targets_for(&target_name) {
1102 transfer_run_config.transfer_targets = targets;
1103 transfer_run_config.parent_agent = None;
1104 } else {
1105 let (parent_name, peer_names) =
1106 Self::compute_transfer_context(&root_agent, &target_name);
1107 let mut targets = Vec::new();
1108 if let Some(ref parent) = parent_name {
1109 targets.push(parent.clone());
1110 }
1111 targets.extend(peer_names);
1112 transfer_run_config.transfer_targets = targets;
1113 transfer_run_config.parent_agent = parent_name;
1114 }
1115 root_agent.configure_run(&target_name, &mut transfer_run_config);
1116
1117 let transfer_invocation_id = format!("inv-{}", uuid::Uuid::new_v4());
1119 let mut transfer_ctx = match InvocationContext::with_mutable_session(
1120 transfer_invocation_id.clone(),
1121 target_agent.clone(),
1122 ctx.user_id().to_string(),
1123 ctx.app_name().to_string(),
1124 ctx.session_id().to_string(),
1125 effective_user_content.clone(),
1126 ctx.mutable_session().clone(),
1127 ) {
1128 Ok(ctx) => ctx,
1129 Err(e) => {
1130 yield Err(e);
1131 return;
1132 }
1133 };
1134 transfer_ctx = transfer_ctx.with_orchestration_root_invocation_id(
1135 adk_core::InvocationContext::orchestration_root_invocation_id(ctx.as_ref())
1136 .to_string(),
1137 );
1138
1139 #[cfg(feature = "artifacts")]
1140 if let Some(ref service) = artifact_service_clone {
1141 let scoped = adk_artifact::ScopedArtifacts::new(
1142 service.clone(),
1143 ctx.app_name().to_string(),
1144 ctx.user_id().to_string(),
1145 ctx.session_id().to_string(),
1146 );
1147 transfer_ctx = transfer_ctx.with_artifacts(Arc::new(scoped));
1148 }
1149 if let Some(ref memory) = memory_service_clone {
1150 transfer_ctx = transfer_ctx.with_memory(memory.clone());
1151 }
1152 transfer_ctx = transfer_ctx.with_run_config(transfer_run_config);
1153 if let Some(rc) = request_context.clone() {
1154 transfer_ctx = transfer_ctx.with_request_context(rc);
1155 }
1156 if let Some(token) = cancellation_token.as_ref() {
1157 transfer_ctx = transfer_ctx.with_cancellation_token(token.clone());
1158 }
1159 if let Some(shared_state) = adk_core::CallbackContext::shared_state(ctx.as_ref()) {
1160 transfer_ctx = transfer_ctx.with_shared_state(shared_state);
1161 }
1162
1163 let transfer_ctx = Arc::new(transfer_ctx);
1164
1165 let mut transfer_stream = match target_agent.run(transfer_ctx.clone()).await {
1167 Ok(s) => s,
1168 Err(e) => {
1169 #[cfg(feature = "plugins")]
1170 if let Some(manager) = plugin_manager.as_ref() {
1171 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1172 }
1173 yield Err(e);
1174 return;
1175 }
1176 };
1177
1178 while let Some(result) = {
1180 match cancellation_token.as_ref() {
1183 Some(token) => {
1184 tokio::select! {
1185 biased;
1186 _ = token.cancelled() => {
1187 tracing::info!("cancellation fired during transferred agent stream await");
1188 #[cfg(feature = "plugins")]
1189 if let Some(manager) = plugin_manager.as_ref() {
1190 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1191 }
1192 return;
1193 }
1194 result = transfer_stream.next() => result,
1195 }
1196 }
1197 None => transfer_stream.next().await,
1198 }
1199 } {
1200 match result {
1201 Ok(mut event) => {
1202 #[cfg(feature = "plugins")]
1203 if let Some(manager) = plugin_manager.as_ref() {
1204 match manager
1205 .run_on_event(
1206 transfer_ctx.clone() as Arc<dyn adk_core::InvocationContext>,
1207 event.clone(),
1208 )
1209 .await
1210 {
1211 Ok(Some(modified)) => {
1212 event = modified;
1213 }
1214 Ok(None) => {}
1215 Err(e) => {
1216 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1217 yield Err(e);
1218 return;
1219 }
1220 }
1221 }
1222
1223 preserve_streamed_content(&mut streamed_content, &mut event);
1224
1225 if let Some(target) = &event.actions.transfer_to_agent {
1227 let source = if event.author.is_empty() {
1228 target_agent.name()
1229 } else {
1230 &event.author
1231 };
1232 if let Some(allowed) = root_agent.transfer_targets_for(source)
1233 && !allowed.contains(target)
1234 {
1235 if root_agent.strict_transfer_policy() {
1236 yield Err(adk_core::AdkError::new(
1237 adk_core::ErrorComponent::Agent,
1238 adk_core::ErrorCategory::Forbidden,
1239 "agent.transfer.target_forbidden",
1240 format!(
1241 "agent '{source}' cannot hand off to '{target}'; allowed targets: {}",
1242 allowed.join(", ")
1243 ),
1244 ));
1245 return;
1246 }
1247 tracing::warn!(source, target, "handoff target rejected by root policy");
1248 } else {
1249 current_transfer_target =
1250 Some((source.to_string(), target.clone()));
1251 }
1252 }
1253
1254 if !event.actions.state_delta.is_empty() {
1256 transfer_ctx.mutable_session().apply_state_delta(&event.actions.state_delta);
1257 }
1258
1259 transfer_ctx.mutable_session().append_event(event.clone());
1261
1262 if !event.llm_response.partial
1263 && let Err(e) = session_service
1264 .append_event_for_identity(adk_session::AppendEventRequest {
1265 identity: identity.clone(),
1266 event: event.clone(),
1267 })
1268 .await {
1269 #[cfg(feature = "plugins")]
1270 if let Some(manager) = plugin_manager.as_ref() {
1271 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1272 }
1273 yield Err(e);
1274 return;
1275 }
1276 yield Ok(event);
1277 }
1278 Err(e) => {
1279 #[cfg(feature = "plugins")]
1280 if let Some(manager) = plugin_manager.as_ref() {
1281 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1282 }
1283 yield Err(e);
1284 return;
1285 }
1286 }
1287 }
1288 }
1289
1290 if let Some(ref compaction_cfg) = compaction_config {
1294 let event_count = ctx.mutable_session().as_ref().events_len();
1295
1296 if event_count > 0 {
1297 let all_events = ctx.mutable_session().as_ref().events_snapshot();
1298 let invocation_count = all_events.iter().filter(|e| e.author == "user").count()
1299 as u32;
1300
1301 if invocation_count > 0
1302 && invocation_count.is_multiple_of(compaction_cfg.compaction_interval)
1303 {
1304 let overlap = compaction_cfg.overlap_size as usize;
1307
1308 let user_msg_indices: Vec<usize> = all_events.iter()
1310 .enumerate()
1311 .filter(|(_, e)| e.author == "user")
1312 .map(|(i, _)| i)
1313 .collect();
1314
1315 let compact_up_to = if overlap == 0 {
1318 all_events.len()
1319 } else if user_msg_indices.len() > overlap {
1320 user_msg_indices[user_msg_indices.len() - overlap]
1322 } else {
1323 0
1325 };
1326
1327 if compact_up_to > 0 {
1328 let events_to_compact = &all_events[..compact_up_to];
1329
1330 match compaction_cfg.summarizer.summarize_events(events_to_compact).await {
1331 Ok(Some(compaction_event)) => {
1332 if let Err(e) = session_service
1334 .append_event_for_identity(adk_session::AppendEventRequest {
1335 identity: identity.clone(),
1336 event: compaction_event.clone(),
1337 })
1338 .await {
1339 tracing::warn!(error = %e, "Failed to persist compaction event");
1340 } else {
1341 tracing::info!(
1342 compacted_events = compact_up_to,
1343 "Context compaction completed"
1344 );
1345 }
1346 }
1347 Ok(None) => {
1348 tracing::debug!("Compaction summarizer returned no result");
1349 }
1350 Err(e) => {
1351 tracing::warn!(error = %e, "Context compaction failed");
1353 }
1354 }
1355 }
1356 }
1357 }
1358 }
1359
1360 #[cfg(feature = "plugins")]
1361 if let Some(manager) = plugin_manager.as_ref() {
1362 manager.run_after_run(ctx.clone() as Arc<dyn adk_core::InvocationContext>).await;
1363 }
1364 };
1365
1366 Ok(RunnerInvocation { invocation_id, events: Box::pin(s) })
1367 }
1368
1369 pub async fn run_str(
1379 &self,
1380 user_id: &str,
1381 session_id: &str,
1382 user_content: Content,
1383 ) -> Result<EventStream> {
1384 let user_id = UserId::try_from(user_id)?;
1385 let session_id = SessionId::try_from(session_id)?;
1386 self.run(user_id, session_id, user_content).await
1387 }
1388
1389 pub fn interrupt(&self, session_id: &str) -> bool {
1414 let runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
1415 let matching: Vec<&ActiveRun> =
1416 runs.values().filter(|run| run.identity.session_id.as_ref() == session_id).collect();
1417 if matching.is_empty() {
1418 tracing::debug!(session.id = session_id, "no active run to interrupt");
1419 return false;
1420 }
1421 tracing::info!(
1422 session.id = session_id,
1423 run.count = matching.len(),
1424 "interrupting running agent"
1425 );
1426 for run in matching {
1427 run.token.cancel();
1428 }
1429 true
1430 }
1431
1432 pub fn interrupt_identity(&self, app_name: &str, user_id: &str, session_id: &str) -> bool {
1444 let runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
1445 let mut cancelled = false;
1446 for run in runs.values() {
1447 if run.identity.app_name.as_ref() == app_name
1448 && run.identity.user_id.as_ref() == user_id
1449 && run.identity.session_id.as_ref() == session_id
1450 {
1451 run.token.cancel();
1452 cancelled = true;
1453 }
1454 }
1455 if !cancelled {
1456 tracing::debug!(
1457 app.name = app_name,
1458 user.id = user_id,
1459 session.id = session_id,
1460 "no active run to interrupt"
1461 );
1462 }
1463 cancelled
1464 }
1465
1466 pub fn active_runs(&self) -> Vec<adk_core::AdkIdentity> {
1471 let runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
1472 runs.values().map(|run| run.identity.clone()).collect()
1473 }
1474
1475 pub fn active_session_ids(&self) -> Vec<String> {
1480 let runs = self.active_runs.lock().unwrap_or_else(|e| e.into_inner());
1481 let mut ids: Vec<String> =
1482 runs.values().map(|run| run.identity.session_id.as_ref().to_string()).collect();
1483 ids.sort();
1484 ids.dedup();
1485 ids
1486 }
1487
1488 #[cfg(feature = "context-compaction")]
1493 pub fn context_compaction(&self) -> Option<&crate::compaction::CompactionConfig> {
1494 self.context_compaction.as_deref()
1495 }
1496
1497 pub fn find_agent_to_run(
1499 root_agent: &Arc<dyn Agent>,
1500 session: &dyn adk_session::Session,
1501 ) -> Arc<dyn Agent> {
1502 let events = session.events();
1504 for i in (0..events.len()).rev() {
1505 if let Some(event) = events.at(i) {
1506 if let Some(target_name) = &event.actions.transfer_to_agent
1508 && let Some(agent) = Self::find_agent(root_agent, target_name)
1509 {
1510 return agent;
1511 }
1512
1513 if event.author == "user" {
1514 continue;
1515 }
1516
1517 if let Some(agent) = Self::find_agent(root_agent, &event.author) {
1519 if Self::is_transferable(root_agent, &agent) {
1521 return agent;
1522 }
1523 }
1524 }
1525 }
1526
1527 root_agent.clone()
1529 }
1530
1531 fn is_transferable(root_agent: &Arc<dyn Agent>, agent: &Arc<dyn Agent>) -> bool {
1547 fn path_supports_transfer(current: &Arc<dyn Agent>, target: &str) -> Option<bool> {
1553 if current.name() == target {
1554 return Some(current.supports_agent_transfer());
1555 }
1556 for sub in current.sub_agents() {
1557 if let Some(sub_ok) = path_supports_transfer(sub, target) {
1558 return Some(current.supports_agent_transfer() && sub_ok);
1559 }
1560 }
1561 None
1562 }
1563
1564 path_supports_transfer(root_agent, agent.name()).unwrap_or(true)
1565 }
1566
1567 pub fn find_agent(current: &Arc<dyn Agent>, target_name: &str) -> Option<Arc<dyn Agent>> {
1569 if current.name() == target_name {
1570 return Some(current.clone());
1571 }
1572
1573 for sub_agent in current.sub_agents() {
1574 if let Some(found) = Self::find_agent(sub_agent, target_name) {
1575 return Some(found);
1576 }
1577 }
1578
1579 None
1580 }
1581
1582 pub fn compute_transfer_context(
1588 root: &Arc<dyn Agent>,
1589 target_name: &str,
1590 ) -> (Option<String>, Vec<String>) {
1591 if root.name() == target_name {
1593 return (None, Vec::new());
1594 }
1595
1596 fn find_parent(current: &Arc<dyn Agent>, target: &str) -> Option<Arc<dyn Agent>> {
1598 for sub in current.sub_agents() {
1599 if sub.name() == target {
1600 return Some(current.clone());
1601 }
1602 if let Some(found) = find_parent(sub, target) {
1603 return Some(found);
1604 }
1605 }
1606 None
1607 }
1608
1609 match find_parent(root, target_name) {
1610 Some(parent) => {
1611 let parent_name = parent.name().to_string();
1612 let peers: Vec<String> = parent
1613 .sub_agents()
1614 .iter()
1615 .filter(|a| a.name() != target_name)
1616 .map(|a| a.name().to_string())
1617 .collect();
1618 (Some(parent_name), peers)
1619 }
1620 None => (None, Vec::new()),
1621 }
1622 }
1623}
1624
1625#[cfg(test)]
1626mod streamed_content_tests {
1627 use super::preserve_streamed_content;
1628 use adk_core::{Content, Event, Part};
1629 use std::collections::HashMap;
1630
1631 fn text(event: &Event) -> String {
1632 event
1633 .content()
1634 .map(|content| content.parts.iter().filter_map(Part::text).collect())
1635 .unwrap_or_default()
1636 }
1637
1638 #[test]
1639 fn final_empty_event_preserves_streamed_text_for_persistence() {
1640 let mut accumulated = HashMap::new();
1641 let mut first = Event::with_id("response-1", "inv-1");
1642 first.llm_response.partial = true;
1643 first.llm_response.content = Some(Content::new("model").with_text("Verify "));
1644 preserve_streamed_content(&mut accumulated, &mut first);
1645
1646 let mut second = Event::with_id("response-1", "inv-1");
1647 second.llm_response.partial = true;
1648 second.llm_response.content = Some(Content::new("model").with_text("the invoice."));
1649 preserve_streamed_content(&mut accumulated, &mut second);
1650
1651 let mut final_event = Event::with_id("response-1", "inv-1");
1652 final_event.llm_response.turn_complete = true;
1653 preserve_streamed_content(&mut accumulated, &mut final_event);
1654
1655 assert_eq!(text(&final_event), "Verify the invoice.");
1656 assert!(accumulated.is_empty());
1657 }
1658}