adk_plugin/plugin.rs
1//! Plugin definition
2//!
3//! A Plugin bundles related callbacks together for a specific purpose.
4
5use crate::callbacks::*;
6use adk_core::{
7 AfterAgentCallback, AfterModelCallback, AfterToolCallback, BeforeAgentCallback,
8 BeforeModelCallback, BeforeToolCallback,
9};
10use std::future::Future;
11use std::pin::Pin;
12
13/// Type alias for async cleanup functions.
14pub type CloseFn = Box<dyn Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>;
15
16/// Configuration for creating a Plugin.
17///
18/// All callbacks are optional - only set the ones you need.
19///
20/// # Example
21///
22/// ```rust,ignore
23/// let config = PluginConfig {
24/// name: "my-plugin".to_string(),
25/// on_user_message: Some(Box::new(|ctx, content| {
26/// Box::pin(async move {
27/// // Process user message
28/// Ok(None)
29/// })
30/// })),
31/// ..Default::default()
32/// };
33/// ```
34pub struct PluginConfig {
35 /// Unique name for this plugin
36 pub name: String,
37
38 // Run lifecycle callbacks
39 /// Called when a user message is received (can modify)
40 pub on_user_message: Option<OnUserMessageCallback>,
41 /// Called for each event (can modify)
42 pub on_event: Option<OnEventCallback>,
43 /// Called before the run starts (can skip run)
44 pub before_run: Option<BeforeRunCallback>,
45 /// Called after the run completes (cleanup)
46 pub after_run: Option<AfterRunCallback>,
47
48 // Agent callbacks
49 /// Called before agent execution
50 pub before_agent: Option<BeforeAgentCallback>,
51 /// Called after agent execution
52 pub after_agent: Option<AfterAgentCallback>,
53
54 // Model callbacks
55 /// Called before LLM call (can modify request or skip)
56 pub before_model: Option<BeforeModelCallback>,
57 /// Called after LLM call (can modify response)
58 pub after_model: Option<AfterModelCallback>,
59 /// Called when LLM returns an error
60 pub on_model_error: Option<OnModelErrorCallback>,
61
62 // Tool callbacks
63 /// Called before tool execution
64 pub before_tool: Option<BeforeToolCallback>,
65 /// Called after tool execution
66 pub after_tool: Option<AfterToolCallback>,
67 /// Called when tool returns an error
68 pub on_tool_error: Option<OnToolErrorCallback>,
69
70 /// Cleanup function called when plugin is closed
71 pub close_fn: Option<CloseFn>,
72}
73
74impl Default for PluginConfig {
75 fn default() -> Self {
76 Self {
77 name: "unnamed".to_string(),
78 on_user_message: None,
79 on_event: None,
80 before_run: None,
81 after_run: None,
82 before_agent: None,
83 after_agent: None,
84 before_model: None,
85 after_model: None,
86 on_model_error: None,
87 before_tool: None,
88 after_tool: None,
89 on_tool_error: None,
90 close_fn: None,
91 }
92 }
93}
94
95/// A Plugin bundles related callbacks for extending agent behavior.
96///
97/// Plugins are registered with a PluginManager which coordinates
98/// callback execution across all registered plugins.
99///
100/// # Example
101///
102/// ```rust,ignore
103/// use adk_plugin::{Plugin, PluginConfig};
104///
105/// // Create a caching plugin
106/// let cache_plugin = Plugin::new(PluginConfig {
107/// name: "cache".to_string(),
108/// before_model: Some(Box::new(|ctx, request| {
109/// Box::pin(async move {
110/// // Check cache for this request
111/// if let Some(cached) = check_cache(&request).await {
112/// return Ok(BeforeModelResult::Skip(cached));
113/// }
114/// Ok(BeforeModelResult::Continue(request))
115/// })
116/// })),
117/// after_model: Some(Box::new(|ctx, response| {
118/// Box::pin(async move {
119/// // Store response in cache
120/// store_in_cache(&response).await;
121/// Ok(None)
122/// })
123/// })),
124/// ..Default::default()
125/// });
126/// ```
127pub struct Plugin {
128 config: PluginConfig,
129}
130
131impl Plugin {
132 /// Create a new plugin from configuration.
133 pub fn new(config: PluginConfig) -> Self {
134 Self { config }
135 }
136
137 /// Get the plugin name.
138 pub fn name(&self) -> &str {
139 &self.config.name
140 }
141
142 /// Get the on_user_message callback if set.
143 pub fn on_user_message(&self) -> Option<&OnUserMessageCallback> {
144 self.config.on_user_message.as_ref()
145 }
146
147 /// Get the on_event callback if set.
148 pub fn on_event(&self) -> Option<&OnEventCallback> {
149 self.config.on_event.as_ref()
150 }
151
152 /// Get the before_run callback if set.
153 pub fn before_run(&self) -> Option<&BeforeRunCallback> {
154 self.config.before_run.as_ref()
155 }
156
157 /// Get the after_run callback if set.
158 pub fn after_run(&self) -> Option<&AfterRunCallback> {
159 self.config.after_run.as_ref()
160 }
161
162 /// Get the before_agent callback if set.
163 pub fn before_agent(&self) -> Option<&BeforeAgentCallback> {
164 self.config.before_agent.as_ref()
165 }
166
167 /// Get the after_agent callback if set.
168 pub fn after_agent(&self) -> Option<&AfterAgentCallback> {
169 self.config.after_agent.as_ref()
170 }
171
172 /// Get the before_model callback if set.
173 pub fn before_model(&self) -> Option<&BeforeModelCallback> {
174 self.config.before_model.as_ref()
175 }
176
177 /// Get the after_model callback if set.
178 pub fn after_model(&self) -> Option<&AfterModelCallback> {
179 self.config.after_model.as_ref()
180 }
181
182 /// Get the on_model_error callback if set.
183 pub fn on_model_error(&self) -> Option<&OnModelErrorCallback> {
184 self.config.on_model_error.as_ref()
185 }
186
187 /// Get the before_tool callback if set.
188 pub fn before_tool(&self) -> Option<&BeforeToolCallback> {
189 self.config.before_tool.as_ref()
190 }
191
192 /// Get the after_tool callback if set.
193 pub fn after_tool(&self) -> Option<&AfterToolCallback> {
194 self.config.after_tool.as_ref()
195 }
196
197 /// Get the on_tool_error callback if set.
198 pub fn on_tool_error(&self) -> Option<&OnToolErrorCallback> {
199 self.config.on_tool_error.as_ref()
200 }
201
202 /// Close the plugin, running cleanup if configured.
203 pub async fn close(&self) {
204 if let Some(ref close_fn) = self.config.close_fn {
205 close_fn().await;
206 }
207 }
208}
209
210impl std::fmt::Debug for Plugin {
211 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
212 f.debug_struct("Plugin")
213 .field("name", &self.config.name)
214 .field("has_on_user_message", &self.config.on_user_message.is_some())
215 .field("has_on_event", &self.config.on_event.is_some())
216 .field("has_before_run", &self.config.before_run.is_some())
217 .field("has_after_run", &self.config.after_run.is_some())
218 .field("has_before_agent", &self.config.before_agent.is_some())
219 .field("has_after_agent", &self.config.after_agent.is_some())
220 .field("has_before_model", &self.config.before_model.is_some())
221 .field("has_after_model", &self.config.after_model.is_some())
222 .field("has_before_tool", &self.config.before_tool.is_some())
223 .field("has_after_tool", &self.config.after_tool.is_some())
224 .finish()
225 }
226}
227
228/// Builder for creating plugins with a fluent API.
229///
230/// `PluginBuilder` provides a chainable interface for constructing [`Plugin`] instances
231/// without manually filling out a [`PluginConfig`] struct. Each method configures a
232/// single lifecycle callback, and [`build`](Self::build) produces the final `Plugin`.
233///
234/// # End-to-End Example
235///
236/// ```rust,ignore
237/// use adk_plugin::PluginBuilder;
238/// use std::sync::Arc;
239///
240/// let plugin = PluginBuilder::new("audit-trail")
241/// .before_agent(Box::new(|ctx| {
242/// Box::pin(async move {
243/// tracing::info!("agent starting");
244/// Ok(())
245/// })
246/// }))
247/// .after_agent(Box::new(|ctx| {
248/// Box::pin(async move {
249/// tracing::info!("agent finished");
250/// Ok(())
251/// })
252/// }))
253/// .before_tool(Box::new(|ctx| {
254/// Box::pin(async move {
255/// if let Some(name) = ctx.tool_name() {
256/// tracing::info!(tool = name, "tool invoked");
257/// }
258/// Ok(())
259/// })
260/// }))
261/// .after_tool(Box::new(|ctx| {
262/// Box::pin(async move {
263/// tracing::info!("tool completed");
264/// Ok(())
265/// })
266/// }))
267/// .before_model(Box::new(|ctx, request| {
268/// Box::pin(async move {
269/// tracing::debug!(model = ?request.model, "sending request to model");
270/// Ok(request)
271/// })
272/// }))
273/// .after_model(Box::new(|ctx, response| {
274/// Box::pin(async move {
275/// tracing::debug!("model responded");
276/// Ok(None)
277/// })
278/// }))
279/// .close_fn(|| Box::pin(async { tracing::info!("plugin closed") }))
280/// .build();
281///
282/// assert_eq!(plugin.name(), "audit-trail");
283/// ```
284pub struct PluginBuilder {
285 config: PluginConfig,
286}
287
288impl PluginBuilder {
289 /// Create a new plugin builder with the given name.
290 ///
291 /// The name uniquely identifies this plugin within a [`PluginManager`](crate::PluginManager).
292 ///
293 /// ```rust,ignore
294 /// let builder = PluginBuilder::new("my-plugin");
295 /// ```
296 pub fn new(name: impl Into<String>) -> Self {
297 Self { config: PluginConfig { name: name.into(), ..Default::default() } }
298 }
299
300 /// Set the callback invoked when a user message is received.
301 ///
302 /// The callback can inspect or modify the incoming [`Content`](adk_core::Content).
303 /// Return `Ok(Some(content))` to replace the message, or `Ok(None)` to keep the original.
304 ///
305 /// ```rust,ignore
306 /// builder.on_user_message(Box::new(|ctx, content| {
307 /// Box::pin(async move {
308 /// tracing::info!(role = %content.role, "user message received");
309 /// Ok(None) // pass through unchanged
310 /// })
311 /// }))
312 /// ```
313 pub fn on_user_message(mut self, callback: OnUserMessageCallback) -> Self {
314 self.config.on_user_message = Some(callback);
315 self
316 }
317
318 /// Set the callback invoked for each event generated by the agent.
319 ///
320 /// The callback can inspect or modify [`Event`](adk_core::Event) objects before they
321 /// are yielded to the caller. Return `Ok(Some(event))` to replace the event, or
322 /// `Ok(None)` to keep the original.
323 ///
324 /// ```rust,ignore
325 /// builder.on_event(Box::new(|ctx, event| {
326 /// Box::pin(async move {
327 /// tracing::debug!(id = %event.id, "event emitted");
328 /// Ok(None)
329 /// })
330 /// }))
331 /// ```
332 pub fn on_event(mut self, callback: OnEventCallback) -> Self {
333 self.config.on_event = Some(callback);
334 self
335 }
336
337 /// Set the callback invoked before the agent run starts.
338 ///
339 /// Use this for setup, validation, or early exit. Return `Ok(Some(content))` to
340 /// skip the run entirely and return that content, or `Ok(None)` to proceed normally.
341 ///
342 /// ```rust,ignore
343 /// builder.before_run(Box::new(|ctx| {
344 /// Box::pin(async move {
345 /// tracing::info!("run starting");
346 /// Ok(None) // continue with the run
347 /// })
348 /// }))
349 /// ```
350 pub fn before_run(mut self, callback: BeforeRunCallback) -> Self {
351 self.config.before_run = Some(callback);
352 self
353 }
354
355 /// Set the callback invoked after the agent run completes.
356 ///
357 /// Use this for cleanup, logging, or metrics collection. This callback is for
358 /// side effects only and does not emit events.
359 ///
360 /// ```rust,ignore
361 /// builder.after_run(Box::new(|ctx| {
362 /// Box::pin(async move {
363 /// tracing::info!("run completed");
364 /// })
365 /// }))
366 /// ```
367 pub fn after_run(mut self, callback: AfterRunCallback) -> Self {
368 self.config.after_run = Some(callback);
369 self
370 }
371
372 /// Set the callback invoked before agent execution.
373 ///
374 /// Runs each time an agent is about to execute. Useful for tracing, permission
375 /// checks, or injecting state.
376 ///
377 /// ```rust,ignore
378 /// builder.before_agent(Box::new(|ctx| {
379 /// Box::pin(async move {
380 /// tracing::info!("agent starting");
381 /// Ok(())
382 /// })
383 /// }))
384 /// ```
385 pub fn before_agent(mut self, callback: BeforeAgentCallback) -> Self {
386 self.config.before_agent = Some(callback);
387 self
388 }
389
390 /// Set the callback invoked after agent execution.
391 ///
392 /// Runs each time an agent finishes. Useful for logging duration, collecting
393 /// metrics, or post-processing results.
394 ///
395 /// ```rust,ignore
396 /// builder.after_agent(Box::new(|ctx| {
397 /// Box::pin(async move {
398 /// tracing::info!("agent finished");
399 /// Ok(())
400 /// })
401 /// }))
402 /// ```
403 pub fn after_agent(mut self, callback: AfterAgentCallback) -> Self {
404 self.config.after_agent = Some(callback);
405 self
406 }
407
408 /// Set the callback invoked before an LLM call.
409 ///
410 /// The callback receives the [`LlmRequest`](adk_core::LlmRequest) and can modify it
411 /// or skip the call entirely by returning a pre-built response.
412 ///
413 /// ```rust,ignore
414 /// builder.before_model(Box::new(|ctx, request| {
415 /// Box::pin(async move {
416 /// tracing::debug!(model = ?request.model, "calling model");
417 /// Ok(request) // pass through unchanged
418 /// })
419 /// }))
420 /// ```
421 pub fn before_model(mut self, callback: BeforeModelCallback) -> Self {
422 self.config.before_model = Some(callback);
423 self
424 }
425
426 /// Set the callback invoked after an LLM call.
427 ///
428 /// The callback receives the [`LlmResponse`](adk_core::LlmResponse) and can modify
429 /// or replace it. Return `Ok(Some(response))` to replace, or `Ok(None)` to keep
430 /// the original.
431 ///
432 /// ```rust,ignore
433 /// builder.after_model(Box::new(|ctx, response| {
434 /// Box::pin(async move {
435 /// tracing::debug!("model responded");
436 /// Ok(None) // keep original response
437 /// })
438 /// }))
439 /// ```
440 pub fn after_model(mut self, callback: AfterModelCallback) -> Self {
441 self.config.after_model = Some(callback);
442 self
443 }
444
445 /// Set the callback invoked when an LLM call returns an error.
446 ///
447 /// The callback receives the original [`LlmRequest`](adk_core::LlmRequest) and the
448 /// error message. Return `Ok(Some(response))` to provide a fallback response, or
449 /// `Ok(None)` to propagate the error.
450 ///
451 /// ```rust,ignore
452 /// builder.on_model_error(Box::new(|ctx, request, error| {
453 /// Box::pin(async move {
454 /// tracing::warn!(error = %error, "model error, no fallback");
455 /// Ok(None) // propagate the error
456 /// })
457 /// }))
458 /// ```
459 pub fn on_model_error(mut self, callback: OnModelErrorCallback) -> Self {
460 self.config.on_model_error = Some(callback);
461 self
462 }
463
464 /// Set the callback invoked before tool execution.
465 ///
466 /// Runs each time a tool is about to execute. The [`CallbackContext`](adk_core::CallbackContext)
467 /// provides `tool_name()` and `tool_input()` for tool-specific logic.
468 ///
469 /// ```rust,ignore
470 /// builder.before_tool(Box::new(|ctx| {
471 /// Box::pin(async move {
472 /// if let Some(name) = ctx.tool_name() {
473 /// tracing::info!(tool = name, "tool starting");
474 /// }
475 /// Ok(())
476 /// })
477 /// }))
478 /// ```
479 pub fn before_tool(mut self, callback: BeforeToolCallback) -> Self {
480 self.config.before_tool = Some(callback);
481 self
482 }
483
484 /// Set the callback invoked after tool execution.
485 ///
486 /// Runs each time a tool finishes. Useful for logging results, collecting metrics,
487 /// or auditing tool usage.
488 ///
489 /// ```rust,ignore
490 /// builder.after_tool(Box::new(|ctx| {
491 /// Box::pin(async move {
492 /// tracing::info!("tool completed");
493 /// Ok(())
494 /// })
495 /// }))
496 /// ```
497 pub fn after_tool(mut self, callback: AfterToolCallback) -> Self {
498 self.config.after_tool = Some(callback);
499 self
500 }
501
502 /// Set the callback invoked when a tool returns an error.
503 ///
504 /// The callback receives the tool name, the error, and the
505 /// [`CallbackContext`](adk_core::CallbackContext). Use this to log failures or
506 /// provide fallback values.
507 ///
508 /// ```rust,ignore
509 /// builder.on_tool_error(Box::new(|ctx, tool_name, error| {
510 /// Box::pin(async move {
511 /// tracing::warn!(tool = %tool_name, error = %error, "tool failed");
512 /// Ok(None) // propagate the error
513 /// })
514 /// }))
515 /// ```
516 pub fn on_tool_error(mut self, callback: OnToolErrorCallback) -> Self {
517 self.config.on_tool_error = Some(callback);
518 self
519 }
520
521 /// Set the async cleanup function called when the plugin is closed.
522 ///
523 /// Use this to release resources, flush buffers, or perform final cleanup.
524 ///
525 /// ```rust,ignore
526 /// builder.close_fn(|| Box::pin(async {
527 /// tracing::info!("plugin shutting down");
528 /// }))
529 /// ```
530 pub fn close_fn(
531 mut self,
532 f: impl Fn() -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync + 'static,
533 ) -> Self {
534 self.config.close_fn = Some(Box::new(f));
535 self
536 }
537
538 /// Consume the builder and produce a [`Plugin`].
539 ///
540 /// Only callbacks that were explicitly set will be active; all others remain `None`.
541 pub fn build(self) -> Plugin {
542 Plugin::new(self.config)
543 }
544}