Skip to main content

adk_runner/
builder.rs

1//! Typestate builder for [`RunnerConfig`] / [`Runner`].
2//!
3//! The builder enforces at compile time that the three required fields
4//! (`app_name`, `agent`, `session_service`) are set before `build()` is
5//! callable.
6//!
7//! # Example
8//!
9//! ```rust,ignore
10//! let runner = Runner::builder()
11//!     .app_name("my-app")
12//!     .agent(agent)
13//!     .session_service(session_service)
14//!     .memory_service(memory)
15//!     .build()?;
16//! ```
17
18use std::marker::PhantomData;
19use std::sync::Arc;
20
21use adk_artifact::ArtifactService;
22use adk_core::{Agent, CacheCapable, ContextCacheConfig, Memory, Result, RunConfig};
23use adk_plugin::PluginManager;
24use adk_session::SessionService;
25use tokio_util::sync::CancellationToken;
26
27use crate::runner::{Runner, RunnerConfig};
28
29// ---------------------------------------------------------------------------
30// Typestate marker types
31// ---------------------------------------------------------------------------
32
33/// Marker: `app_name` has not been set.
34pub struct NoAppName;
35/// Marker: `app_name` has been set.
36pub struct HasAppName;
37/// Marker: `agent` has not been set.
38pub struct NoAgent;
39/// Marker: `agent` has been set.
40pub struct HasAgent;
41/// Marker: `session_service` has not been set.
42pub struct NoSessionService;
43/// Marker: `session_service` has been set.
44pub struct HasSessionService;
45
46// ---------------------------------------------------------------------------
47// Builder
48// ---------------------------------------------------------------------------
49
50/// A typestate builder for constructing a [`Runner`].
51///
52/// The three type parameters track whether the required fields have been
53/// provided. `build()` is only available when all three are `Has*`.
54pub struct RunnerConfigBuilder<A, G, S> {
55    app_name: Option<String>,
56    agent: Option<Arc<dyn Agent>>,
57    session_service: Option<Arc<dyn SessionService>>,
58    artifact_service: Option<Arc<dyn ArtifactService>>,
59    memory_service: Option<Arc<dyn Memory>>,
60    plugin_manager: Option<Arc<PluginManager>>,
61    run_config: Option<RunConfig>,
62    compaction_config: Option<adk_core::EventsCompactionConfig>,
63    context_cache_config: Option<ContextCacheConfig>,
64    cache_capable: Option<Arc<dyn CacheCapable>>,
65    request_context: Option<adk_core::RequestContext>,
66    cancellation_token: Option<CancellationToken>,
67    _marker: PhantomData<(A, G, S)>,
68}
69
70impl RunnerConfigBuilder<NoAppName, NoAgent, NoSessionService> {
71    /// Create a new builder with all fields unset and defaults applied.
72    pub fn new() -> Self {
73        Self {
74            app_name: None,
75            agent: None,
76            session_service: None,
77            artifact_service: None,
78            memory_service: None,
79            plugin_manager: None,
80            run_config: None,
81            compaction_config: None,
82            context_cache_config: None,
83            cache_capable: None,
84            request_context: None,
85            cancellation_token: None,
86            _marker: PhantomData,
87        }
88    }
89}
90
91impl Default for RunnerConfigBuilder<NoAppName, NoAgent, NoSessionService> {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97// ---------------------------------------------------------------------------
98// Required-field setters (transition type state)
99// ---------------------------------------------------------------------------
100
101impl<A, G, S> RunnerConfigBuilder<A, G, S> {
102    /// Set the application name (required).
103    pub fn app_name(self, name: impl Into<String>) -> RunnerConfigBuilder<HasAppName, G, S> {
104        RunnerConfigBuilder {
105            app_name: Some(name.into()),
106            agent: self.agent,
107            session_service: self.session_service,
108            artifact_service: self.artifact_service,
109            memory_service: self.memory_service,
110            plugin_manager: self.plugin_manager,
111            run_config: self.run_config,
112            compaction_config: self.compaction_config,
113            context_cache_config: self.context_cache_config,
114            cache_capable: self.cache_capable,
115            request_context: self.request_context,
116            cancellation_token: self.cancellation_token,
117            _marker: PhantomData,
118        }
119    }
120
121    /// Set the root agent (required).
122    pub fn agent(self, agent: Arc<dyn Agent>) -> RunnerConfigBuilder<A, HasAgent, S> {
123        RunnerConfigBuilder {
124            app_name: self.app_name,
125            agent: Some(agent),
126            session_service: self.session_service,
127            artifact_service: self.artifact_service,
128            memory_service: self.memory_service,
129            plugin_manager: self.plugin_manager,
130            run_config: self.run_config,
131            compaction_config: self.compaction_config,
132            context_cache_config: self.context_cache_config,
133            cache_capable: self.cache_capable,
134            request_context: self.request_context,
135            cancellation_token: self.cancellation_token,
136            _marker: PhantomData,
137        }
138    }
139
140    /// Set the session service (required).
141    pub fn session_service(
142        self,
143        service: Arc<dyn SessionService>,
144    ) -> RunnerConfigBuilder<A, G, HasSessionService> {
145        RunnerConfigBuilder {
146            app_name: self.app_name,
147            agent: self.agent,
148            session_service: Some(service),
149            artifact_service: self.artifact_service,
150            memory_service: self.memory_service,
151            plugin_manager: self.plugin_manager,
152            run_config: self.run_config,
153            compaction_config: self.compaction_config,
154            context_cache_config: self.context_cache_config,
155            cache_capable: self.cache_capable,
156            request_context: self.request_context,
157            cancellation_token: self.cancellation_token,
158            _marker: PhantomData,
159        }
160    }
161}
162
163// ---------------------------------------------------------------------------
164// Optional-field setters (no type-state change)
165// ---------------------------------------------------------------------------
166
167impl<A, G, S> RunnerConfigBuilder<A, G, S> {
168    /// Set the artifact service (optional).
169    pub fn artifact_service(mut self, service: Arc<dyn ArtifactService>) -> Self {
170        self.artifact_service = Some(service);
171        self
172    }
173
174    /// Set the memory service (optional).
175    pub fn memory_service(mut self, service: Arc<dyn Memory>) -> Self {
176        self.memory_service = Some(service);
177        self
178    }
179
180    /// Set the plugin manager (optional).
181    pub fn plugin_manager(mut self, manager: Arc<PluginManager>) -> Self {
182        self.plugin_manager = Some(manager);
183        self
184    }
185
186    /// Set the run configuration (optional).
187    pub fn run_config(mut self, config: RunConfig) -> Self {
188        self.run_config = Some(config);
189        self
190    }
191
192    /// Set the events compaction configuration (optional).
193    pub fn compaction_config(mut self, config: adk_core::EventsCompactionConfig) -> Self {
194        self.compaction_config = Some(config);
195        self
196    }
197
198    /// Set the context cache configuration (optional).
199    pub fn context_cache_config(mut self, config: ContextCacheConfig) -> Self {
200        self.context_cache_config = Some(config);
201        self
202    }
203
204    /// Set the cache-capable model reference (optional).
205    pub fn cache_capable(mut self, model: Arc<dyn CacheCapable>) -> Self {
206        self.cache_capable = Some(model);
207        self
208    }
209
210    /// Set the request context from auth middleware (optional).
211    pub fn request_context(mut self, ctx: adk_core::RequestContext) -> Self {
212        self.request_context = Some(ctx);
213        self
214    }
215
216    /// Set a cooperative cancellation token (optional).
217    pub fn cancellation_token(mut self, token: CancellationToken) -> Self {
218        self.cancellation_token = Some(token);
219        self
220    }
221}
222
223// ---------------------------------------------------------------------------
224// build() — only available when all required fields are set
225// ---------------------------------------------------------------------------
226
227impl RunnerConfigBuilder<HasAppName, HasAgent, HasSessionService> {
228    /// Consume the builder and create a [`Runner`].
229    ///
230    /// Delegates to [`Runner::new()`] internally.
231    ///
232    /// # Errors
233    ///
234    /// Returns an error if `Runner::new()` fails (e.g. invalid `app_name`).
235    pub fn build(self) -> Result<Runner> {
236        let config = RunnerConfig {
237            // SAFETY: typestate guarantees these are `Some`.
238            app_name: self.app_name.expect("typestate guarantees app_name is set"),
239            agent: self.agent.expect("typestate guarantees agent is set"),
240            session_service: self
241                .session_service
242                .expect("typestate guarantees session_service is set"),
243            artifact_service: self.artifact_service,
244            memory_service: self.memory_service,
245            plugin_manager: self.plugin_manager,
246            run_config: self.run_config,
247            compaction_config: self.compaction_config,
248            context_cache_config: self.context_cache_config,
249            cache_capable: self.cache_capable,
250            request_context: self.request_context,
251            cancellation_token: self.cancellation_token,
252        };
253        Runner::new(config)
254    }
255}