Skip to main content

adk_runner/
runner.rs

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/// A run currently in flight, tracked so it can be interrupted.
19///
20/// Keyed by run ID rather than session ID: a session ID is only unique within an
21/// app and user, and one identity may have several runs in flight at once.
22#[derive(Debug, Clone)]
23struct ActiveRun {
24    identity: adk_core::AdkIdentity,
25    token: CancellationToken,
26}
27
28/// Registry of in-flight runs, keyed by run ID.
29type 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
46/// Deregisters a run when its event stream is dropped.
47///
48/// Removal is by run ID, so a finishing run can never deregister a different run
49/// that happens to share the same session ID.
50struct ActiveRunCleanup {
51    active_runs: ActiveRuns,
52    run_id: u64,
53}
54
55/// A newly registered runner invocation and its event stream.
56///
57/// The invocation identity is available before the stream is polled, so hosts
58/// can durably correlate failures that happen before the first agent event.
59pub struct RunnerInvocation {
60    invocation_id: InvocationId,
61    events: EventStream,
62}
63
64impl RunnerInvocation {
65    /// Returns the identity assigned to this invocation.
66    pub fn invocation_id(&self) -> &InvocationId {
67        &self.invocation_id
68    }
69
70    /// Consumes the invocation and returns its event stream.
71    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
93/// Configuration for constructing a [`Runner`].
94///
95/// Use [`Runner::builder()`] for a compile-time-safe way to construct this.
96pub struct RunnerConfig {
97    /// Application name used for session scoping.
98    pub app_name: String,
99    /// The root agent to execute.
100    pub agent: Arc<dyn Agent>,
101    /// Session persistence backend.
102    pub session_service: Arc<dyn SessionService>,
103    #[cfg(feature = "artifacts")]
104    /// Optional artifact storage service.
105    pub artifact_service: Option<Arc<dyn ArtifactService>>,
106    /// Optional memory/RAG service.
107    pub memory_service: Option<Arc<dyn Memory>>,
108    #[cfg(feature = "plugins")]
109    /// Optional plugin manager for lifecycle hooks.
110    pub plugin_manager: Option<Arc<PluginManager>>,
111    /// Optional run configuration (streaming mode, etc.)
112    /// If not provided, uses default (SSE streaming)
113    #[allow(dead_code)]
114    pub run_config: Option<RunConfig>,
115    /// Optional context compaction configuration.
116    /// When set, the runner will periodically summarize older events
117    /// to reduce context size sent to the LLM.
118    pub compaction_config: Option<adk_core::EventsCompactionConfig>,
119    /// Optional context cache configuration for automatic prompt caching lifecycle.
120    /// When set alongside `cache_capable`, the runner will automatically create and
121    /// manage cached content resources for supported providers.
122    ///
123    /// When `cache_capable` is set but this field is `None`, the runner
124    /// automatically uses [`ContextCacheConfig::default()`] (4096 min tokens,
125    /// 600s TTL, refresh every 3 invocations).
126    pub context_cache_config: Option<ContextCacheConfig>,
127    /// Optional cache-capable model reference for automatic cache management.
128    /// Set this to the same model used by the agent if it supports caching.
129    pub cache_capable: Option<Arc<dyn CacheCapable>>,
130    /// Optional request context from the server's auth middleware bridge.
131    /// When set, the runner passes it to `InvocationContext` so that
132    /// `user_scopes()` and `user_id()` reflect the authenticated identity.
133    pub request_context: Option<adk_core::RequestContext>,
134    /// Optional cooperative cancellation token for externally managed runs.
135    ///
136    /// Each run derives a child of this token, so cancelling it stops every in-flight
137    /// run, while [`Runner::interrupt`] cancels only the targeted run's child.
138    pub cancellation_token: Option<CancellationToken>,
139    /// Optional intra-invocation compaction configuration.
140    /// When set, the runner estimates token count before each agent run
141    /// and triggers mid-invocation summarization when the threshold is exceeded.
142    pub intra_compaction_config: Option<adk_core::IntraCompactionConfig>,
143    /// Optional summarizer for intra-invocation compaction.
144    /// Required when `intra_compaction_config` is set.
145    pub intra_compaction_summarizer: Option<Arc<dyn adk_core::BaseEventsSummarizer>>,
146    /// Optional context compaction configuration for token-budget overflow handling.
147    ///
148    /// When set, the runner applies the configured [`CompactionStrategy`](crate::compaction::CompactionStrategy)
149    /// to shrink the event history when the context exceeds the token budget,
150    /// retrying the model request up to `max_retries` times.
151    ///
152    /// This field is only available when the `context-compaction` feature is enabled.
153    #[cfg(feature = "context-compaction")]
154    pub context_compaction: Option<crate::compaction::CompactionConfig>,
155}
156
157/// Agent execution runtime.
158///
159/// Orchestrates session retrieval, agent dispatch, event streaming, context
160/// caching, and compaction. Construct via [`Runner::builder()`] or
161/// [`Runner::new()`].
162pub 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    /// Optional context compaction configuration for token-budget overflow handling.
182    #[cfg(feature = "context-compaction")]
183    context_compaction: Option<Arc<crate::compaction::CompactionConfig>>,
184    /// Per-session cancellation tokens for the interrupt API.
185    /// Each `run()` call registers a token here; `interrupt()` cancels it.
186    active_runs: ActiveRuns,
187    next_run_id: Arc<std::sync::atomic::AtomicU64>,
188    /// Serializes externally triggered invocations per session. The weak values prevent the
189    /// per-trigger session policy from retaining one lock forever for every completed event.
190    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    /// Create a typestate builder for constructing a `Runner`.
199    ///
200    /// The builder enforces at compile time that the three required fields
201    /// (`app_name`, `agent`, `session_service`) are set before `build()` is
202    /// callable.
203    ///
204    /// # Example
205    ///
206    /// ```rust,ignore
207    /// let runner = Runner::builder()
208    ///     .app_name("my-app")
209    ///     .agent(agent)
210    ///     .session_service(session_service)
211    ///     .build()?;
212    /// ```
213    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    /// Create a new runner from a [`RunnerConfig`].
222    ///
223    /// Prefer [`Runner::builder()`] for a compile-time-safe construction API.
224    pub fn new(config: RunnerConfig) -> Result<Self> {
225        let run_config = config.run_config.unwrap_or_default();
226
227        // When a cache-capable model is provided but no explicit cache config,
228        // use the default ContextCacheConfig to enable caching automatically.
229        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    /// The executable root owned by this runner.
276    pub(crate) fn root_agent(&self) -> Arc<dyn Agent> {
277        Arc::clone(&self.root_agent)
278    }
279
280    /// Enable skill injection using a pre-built injector.
281    ///
282    /// Skill injection runs before plugin `on_user_message` callbacks.
283    #[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    /// Enable skill injection by auto-loading `.skills/` from the given root path.
290    #[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    /// Enable skill injection by auto-loading `.skills/` from the given root path.
302    ///
303    /// Unlike [`with_auto_skills`](Self::with_auto_skills), this method borrows
304    /// the Runner mutably instead of consuming it. On error, the Runner remains
305    /// valid with no skill injector configured.
306    #[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    /// Execute the root agent for the given user and session, returning an event stream.
318    ///
319    /// Retrieves the existing session, resolves the target agent, runs plugins and skills, and
320    /// streams events as the agent executes.
321    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    /// Returns the runner's application name.
331    pub fn app_name(&self) -> &str {
332        &self.app_name
333    }
334
335    /// Returns the runner's session service.
336    pub fn session_service(&self) -> &Arc<dyn adk_session::SessionService> {
337        &self.session_service
338    }
339
340    /// Returns the runner's configured [`RunConfig`].
341    ///
342    /// Useful as a base to clone and adjust for [`Self::run_with_config`].
343    pub fn run_config(&self) -> &RunConfig {
344        &self.run_config
345    }
346
347    /// Runs the agent with a per-invocation [`RunConfig`] override.
348    ///
349    /// Passing `None` uses the runner's configured `RunConfig`. Supply one to vary a single
350    /// invocation — injecting `runtime_toolsets` for tools that only exist for the duration of
351    /// that run, for example, as `SandboxRunner` does with tools bound to a live sandbox session.
352    ///
353    /// # Errors
354    ///
355    /// Returns an error if invocation setup fails before the stream is created. Session lookup
356    /// and agent execution failures are yielded by the returned stream.
357    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    /// Registers one invocation and returns its identity with the event stream.
371    ///
372    /// Prefer this method when a host must durably associate its own turn or
373    /// operation record before the first agent event can be produced.
374    ///
375    /// # Errors
376    ///
377    /// Returns an error if invocation setup fails before the stream is created.
378    /// Session lookup and agent execution failures are yielded by the returned
379    /// stream.
380    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        // Built once and used for every persistence write, so a backend with a
411        // composite natural key can bind each event to its tenant.
412        let identity =
413            adk_core::AdkIdentity::new(typed_app_name.clone(), user_id.clone(), session_id.clone());
414
415        // Register this run for the interrupt API. The registration is keyed by a
416        // unique run ID rather than the raw session ID: two identities can share a
417        // session ID, and one identity can have two runs in flight, and both cases
418        // previously overwrote each other's token.
419        //
420        // A child of the global token fires when either the global token or this run's
421        // interrupt does, with no watcher task to outlive the run; dropping the run's
422        // last handle detaches the child from the global token.
423        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        // Built here rather than inside the generator: registration happens as soon
439        // as `run` is called, so deregistration must also survive a stream that is
440        // dropped before it is ever polled. Moving the guard into the generator
441        // keeps it alive exactly as long as the stream.
442        let cleanup = ActiveRunCleanup { active_runs: active_runs.clone(), run_id };
443
444        let s = stream! {
445            let _cleanup = cleanup;
446
447            // Use the effective token (combines global + per-session)
448            let cancellation_token = effective_token;
449            // Resolve the existing session.
450            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            // Find which agent should handle this request
468            let agent_to_run = Self::find_agent_to_run(&root_agent, session.as_ref());
469
470            // Let validated composite roots apply policy for the concrete agent
471            // selected for this turn. Ordinary agents keep the default no-op.
472            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            // Clone services for potential reuse in transfer
479            #[cfg(feature = "artifacts")]
480            let artifact_service_clone = artifact_service.clone();
481            let memory_service_clone = memory_service.clone();
482
483            // Create invocation context with MutableSession
484            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            // Add optional services
527            #[cfg(feature = "artifacts")]
528            if let Some(service) = artifact_service {
529                // Wrap service with ScopedArtifacts to bind session context
530                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            // Apply run config (streaming mode, etc.)
543            invocation_ctx = invocation_ctx.with_run_config(run_config.clone());
544
545            // Apply request context from auth middleware bridge if present
546            if let Some(rc) = request_context.clone() {
547                invocation_ctx = invocation_ctx.with_request_context(rc);
548            }
549
550            // Expose cooperative cancellation to the agent/tools.
551            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            // Append user message to session service (persistent storage)
657            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            // Also add to mutable session for immediate visibility
662            // Note: adk_session::Event is a re-export of adk_core::Event, so we can use it directly
663            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            // ===== CONTEXT CACHE LIFECYCLE =====
680            // If context caching is configured and a cache-capable model is available,
681            // create or refresh the cached content before agent execution.
682            // Cache failures are non-fatal — log a warning and proceed without cache.
683            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                    // Gather system instruction from the agent's description
691                    // (the full instruction is resolved inside the agent, but the
692                    // description provides a reasonable proxy for cache keying).
693                    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                // Attach cache name to run config so agents can use it.
725                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                    // Rebuild the invocation context with the updated run config.
737                    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            // ===== INTRA-INVOCATION COMPACTION =====
781            // If intra-compaction is configured, check if the session events
782            // exceed the token threshold and compact them before the agent runs.
783            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) => {} // No compaction needed
792                    Err(e) => {
793                        tracing::warn!(error = %e, "intra-invocation compaction check failed");
794                    }
795                }
796            }
797
798            // ===== CONTEXT COMPACTION (TOKEN BUDGET) =====
799            // If context-compaction is configured, proactively check the estimated
800            // token count before calling the agent. If it exceeds the budget,
801            // apply compaction to bring it under the limit.
802            #[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                            // Proactive compaction failed — proceed anyway and let the
819                            // model reject the request if it's truly too large.
820                            tracing::warn!(error = %e, "proactive context compaction failed, proceeding with full context");
821                        }
822                    }
823                }
824            }
825
826            // Run the agent with instrumentation (ADK-Go style attributes)
827            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(), // Use invocation_id as event_id for agent spans
832                "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                    // Token limit error on agent.run() — apply compaction and retry
845                    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                            // Retry the agent call with compacted context
856                            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            // Stream events and check for transfers
889            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                // Race the next event against cancellation so an in-flight
895                // await (LLM streaming, tool I/O) is interrupted promptly
896                // rather than only at poll boundaries. Dropping the stream on
897                // cancellation releases the underlying provider connection.
898                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                            // Instrument the poll, not just the construction of the
911                            // stream: `agent_to_run.run(..)` merely builds the
912                            // stream, so without this every span created while the
913                            // stream is drained — i.e. the whole agent execution —
914                            // is created outside `agent_span`.
915                            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                        // Check for transfer action
947                        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                        // CRITICAL: Apply state_delta to the mutable session immediately.
975                        // This is the key fix for state propagation between sequential agents.
976                        // When an agent sets output_key, it emits an event with state_delta.
977                        // We must apply this to the mutable session so downstream agents
978                        // can read the value via ctx.session().state().get().
979                        if !event.actions.state_delta.is_empty() {
980                            ctx.mutable_session().apply_state_delta(&event.actions.state_delta);
981                        }
982
983                        // Also add the event to the mutable session's event list
984                        ctx.mutable_session().append_event(event.clone());
985
986                        // Append event to session service (persistent storage)
987                        // Skip partial streaming chunks — only persist the final
988                        // event. Streaming chunks share the same event ID, so
989                        // persisting each one would violate the primary key
990                        // constraint. The final chunk (partial=false) carries the
991                        // complete accumulated content.
992                        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            // ===== TRANSFER LOOP =====
1020            // Support multi-hop transfers with a max-depth guard.
1021            // When an agent emits transfer_to_agent, the runner resolves the
1022            // target from the root agent tree, computes transfer_targets
1023            // (parent + peers) for the new agent, and runs it. This repeats
1024            // until no further transfer is requested or the depth limit is hit.
1025            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                // Compute transfer_targets for the target agent:
1097                // - parent: the agent that transferred to it (or root if applicable)
1098                // - peers: siblings in the agent tree
1099                // - children: handled by the agent itself via sub_agents()
1100                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                // For transfers, we reuse the same mutable session to preserve state
1118                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                // Run the transferred agent
1166                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                // Stream events from the transferred agent, capturing any further transfer
1179                while let Some(result) = {
1180                    // Race the next event against cancellation for prompt
1181                    // mid-await interruption of the transferred agent.
1182                    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                            // Capture further transfer requests
1226                            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                            // Apply state delta for transferred agent too
1255                            if !event.actions.state_delta.is_empty() {
1256                                transfer_ctx.mutable_session().apply_state_delta(&event.actions.state_delta);
1257                            }
1258
1259                            // Add to mutable session
1260                            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            // ===== CONTEXT COMPACTION =====
1291            // After all events have been processed, check if compaction should trigger.
1292            // This runs in the background after the invocation completes.
1293            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                        // Determine the window of events to compact
1305                        // We compact all events except the most recent overlap_size invocations
1306                        let overlap = compaction_cfg.overlap_size as usize;
1307
1308                        // Find the boundary: keep the last `overlap` user messages and everything after
1309                        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                        // Keep the last `overlap` user messages intact.
1316                        // When overlap is 0, compact everything.
1317                        let compact_up_to = if overlap == 0 {
1318                            all_events.len()
1319                        } else if user_msg_indices.len() > overlap {
1320                            // Compact up to (but not including) the overlap-th-from-last user message
1321                            user_msg_indices[user_msg_indices.len() - overlap]
1322                        } else {
1323                            // Not enough user messages to satisfy overlap — skip compaction
1324                            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                                    // Persist the compaction event
1333                                    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                                    // Compaction failure is non-fatal — log and continue
1352                                    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    /// Convenience method that accepts string arguments.
1370    ///
1371    /// Converts `user_id` and `session_id` to their typed equivalents
1372    /// and delegates to [`run()`](Self::run).
1373    ///
1374    /// # Errors
1375    ///
1376    /// Returns an error if either string fails identity validation
1377    /// (empty, contains null bytes, or exceeds length limit).
1378    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    /// Interrupt a running agent for the given session.
1390    ///
1391    /// Cancels the agent's current execution within the event loop. Events
1392    /// already produced and appended to the session are preserved — only
1393    /// future events are stopped. The caller can then issue a new `run()`
1394    /// call with a different instruction to redirect the agent.
1395    ///
1396    /// Returns `true` if a running session was found and interrupted,
1397    /// `false` if no active run exists for that session ID.
1398    ///
1399    /// # Example
1400    ///
1401    /// ```rust,ignore
1402    /// // Start a run in the background
1403    /// let mut stream = runner.run(user_id, session_id, content).await?;
1404    /// tokio::spawn(async move { while stream.next().await.is_some() {} });
1405    ///
1406    /// // Later, interrupt it
1407    /// let was_running = runner.interrupt("session-1");
1408    /// assert!(was_running);
1409    ///
1410    /// // Redirect with a new instruction
1411    /// let mut stream = runner.run(user_id, session_id, new_content).await?;
1412    /// ```
1413    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    /// Interrupts runs for one exact identity.
1433    ///
1434    /// A session ID is only unique within an app and user, so this is the precise
1435    /// form of [`Runner::interrupt`] for a `Runner` shared across tenants.
1436    /// Returns `true` when at least one run was cancelled.
1437    ///
1438    /// # Example
1439    ///
1440    /// ```rust,ignore
1441    /// let cancelled = runner.interrupt_identity("my-app", "user-1", "session-1");
1442    /// ```
1443    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    /// Returns the identity of every run currently in flight.
1467    ///
1468    /// One identity appears once per in-flight run, so a repeated entry means that
1469    /// identity has concurrent runs.
1470    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    /// Returns the session IDs of all currently running agent executions.
1476    ///
1477    /// Session IDs are deduplicated. Use [`Runner::active_runs`] when the app and
1478    /// user dimensions matter.
1479    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    /// Returns a reference to the context compaction configuration, if set.
1489    ///
1490    /// This is used by the runner's generate_content loop to detect token limit
1491    /// errors and apply compaction strategies before retrying.
1492    #[cfg(feature = "context-compaction")]
1493    pub fn context_compaction(&self) -> Option<&crate::compaction::CompactionConfig> {
1494        self.context_compaction.as_deref()
1495    }
1496
1497    /// Find which agent should handle the request based on session history
1498    pub fn find_agent_to_run(
1499        root_agent: &Arc<dyn Agent>,
1500        session: &dyn adk_session::Session,
1501    ) -> Arc<dyn Agent> {
1502        // Look at recent events to find last agent that responded
1503        let events = session.events();
1504        for i in (0..events.len()).rev() {
1505            if let Some(event) = events.at(i) {
1506                // Check for explicit transfer
1507                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                // Try to find this agent in the tree
1518                if let Some(agent) = Self::find_agent(root_agent, &event.author) {
1519                    // Check if agent allows transfer up the tree
1520                    if Self::is_transferable(root_agent, &agent) {
1521                        return agent;
1522                    }
1523                }
1524            }
1525        }
1526
1527        // Default to root agent
1528        root_agent.clone()
1529    }
1530
1531    /// Check if an agent found in session history can be resumed directly for
1532    /// the next user message.
1533    ///
1534    /// An agent is a valid direct-resumption target only if it *and* every
1535    /// ancestor up to the root permit agent transfer
1536    /// ([`Agent::supports_agent_transfer`]). A deterministic workflow agent
1537    /// (sequential, parallel, loop, conditional) anywhere on that path returns
1538    /// `false`, which forces resumption to restart from the workflow root
1539    /// rather than a single sub-agent that happened to respond last. This
1540    /// mirrors Google ADK's `_is_transferable_across_agent_tree`.
1541    ///
1542    /// LLM-driven transfer-policy enforcement
1543    /// (`disallow_transfer_to_parent` / `disallow_transfer_to_peers`) is still
1544    /// handled inside `LlmAgent::run()` when it builds the `transfer_to_agent`
1545    /// tool's valid-target list; this check only governs cross-turn resumption.
1546    fn is_transferable(root_agent: &Arc<dyn Agent>, agent: &Arc<dyn Agent>) -> bool {
1547        // Walk the tree from the root down to the target. Every agent on that
1548        // path (root through target, inclusive) must support transfer for the
1549        // target to be a valid direct-resumption point. `Some(true)` = found
1550        // and fully transferable, `Some(false)` = found but a workflow agent
1551        // sits on the path, `None` = target not present in this subtree.
1552        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    /// Recursively search agent tree for agent with given name
1568    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    /// Compute the parent name and peer names for a given agent in the tree.
1583    /// Returns `(parent_name, peer_names)`.
1584    ///
1585    /// Walks the agent tree to find the parent of `target_name`, then collects
1586    /// the parent's name and the sibling agent names (excluding the target itself).
1587    pub fn compute_transfer_context(
1588        root: &Arc<dyn Agent>,
1589        target_name: &str,
1590    ) -> (Option<String>, Vec<String>) {
1591        // If the target is the root itself, there's no parent or peers
1592        if root.name() == target_name {
1593            return (None, Vec::new());
1594        }
1595
1596        // BFS/DFS to find the parent of target_name
1597        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}