adk_runner/
agent_invoker.rs

1//! [`AgentInvoker`] for [`Runner`].
2//!
3//! [`Runner::run`] resolves an *existing* session and yields `session.not_found` through the
4//! stream when there is none. A caller driving an agent from an external event — a cron tick, a
5//! webhook, a queue message — has no opportunity to register a session first, so every such
6//! caller ended up writing the same create-then-run dance. This implementation owns it.
7
8use std::sync::{Arc, Weak};
9
10use adk_core::{Agent, AgentInvoker, Content, EventStream, Result};
11use async_trait::async_trait;
12use futures::StreamExt;
13use tokio::sync::{Mutex, OwnedMutexGuard};
14
15use crate::Runner;
16
17type SessionLockRegistry =
18    Arc<std::sync::Mutex<std::collections::HashMap<String, Weak<Mutex<()>>>>>;
19
20/// Keeps one session lock held for the lifetime of the returned event stream and removes an idle
21/// weak registry entry when the stream completes or is dropped.
22struct SessionInvocationLease {
23    key: String,
24    registry: SessionLockRegistry,
25    _guard: OwnedMutexGuard<()>,
26}
27
28impl Drop for SessionInvocationLease {
29    fn drop(&mut self) {
30        let mut registry = self.registry.lock().unwrap_or_else(|error| error.into_inner());
31        if registry.get(&self.key).is_some_and(|lock| lock.strong_count() == 1) {
32            registry.remove(&self.key);
33        }
34    }
35}
36
37#[async_trait]
38impl AgentInvoker for Runner {
39    fn agent(&self) -> Option<Arc<dyn Agent>> {
40        Some(self.root_agent())
41    }
42
43    async fn invoke(
44        &self,
45        user_id: &str,
46        session_id: &str,
47        content: Content,
48    ) -> Result<EventStream> {
49        let lock_key = format!("{}\0{user_id}\0{session_id}", self.app_name());
50        let session_lock = {
51            let mut registry =
52                self.external_session_locks.lock().unwrap_or_else(|error| error.into_inner());
53            if let Some(lock) = registry.get(&lock_key).and_then(Weak::upgrade) {
54                lock
55            } else {
56                let lock = Arc::new(Mutex::new(()));
57                registry.insert(lock_key.clone(), Arc::downgrade(&lock));
58                lock
59            }
60        };
61        let guard = session_lock.lock_owned().await;
62        let lease = SessionInvocationLease {
63            key: lock_key,
64            registry: Arc::clone(&self.external_session_locks),
65            _guard: guard,
66        };
67
68        let existing = self
69            .session_service()
70            .get(adk_session::GetRequest {
71                app_name: self.app_name().to_string(),
72                user_id: user_id.to_string(),
73                session_id: session_id.to_string(),
74                num_recent_events: None,
75                after: None,
76            })
77            .await;
78
79        match existing {
80            Ok(_) => {}
81            // Only a genuine absence is created through. A backend or transport failure must
82            // surface, not be papered over with a fresh session that silently discards history.
83            Err(error) if error.is_not_found() => {
84                let created = self
85                    .session_service()
86                    .create(adk_session::CreateRequest {
87                        app_name: self.app_name().to_string(),
88                        user_id: user_id.to_string(),
89                        session_id: Some(session_id.to_string()),
90                        state: std::collections::HashMap::new(),
91                    })
92                    .await;
93                if let Err(create_error) = created {
94                    // Another runner or process may have won the create race. Only suppress the
95                    // error when a fresh lookup proves the requested session now exists.
96                    if self
97                        .session_service()
98                        .get(adk_session::GetRequest {
99                            app_name: self.app_name().to_string(),
100                            user_id: user_id.to_string(),
101                            session_id: session_id.to_string(),
102                            num_recent_events: None,
103                            after: None,
104                        })
105                        .await
106                        .is_err()
107                    {
108                        return Err(create_error);
109                    }
110                }
111                tracing::debug!(
112                    user_id,
113                    session_id,
114                    "session is ready for an externally triggered invocation"
115                );
116            }
117            Err(error) => return Err(error),
118        }
119
120        let mut events = self.run_str(user_id, session_id, content).await?;
121
122        Ok(Box::pin(async_stream::stream! {
123            let _lease = lease;
124            while let Some(event) = events.next().await {
125                yield event;
126            }
127        }))
128    }
129}
adk_runner/agent_invoker.rs

adk_runner/
agent_invoker.rs
