Callbacks & plugins

Lifecycle hooks in adk-rs: the callback type aliases and their contexts, and the plugin system the Runner invokes around every invocation and event.

adk-rs has two hook surfaces, both fully wired. Callbacks are typed async closures aimed at specific lifecycle points — before/after the agent, the model, and each tool — registered on one agent via its builder. Plugins are trait objects the Runner calls around every invocation and for every event, hooking the whole runner.

§Callback contexts

Two lightweight wrappers around Arc<InvocationContext> accompany the callback types (both in src/core/callback.rs):

pub struct ReadonlyContext { pub invocation: Arc<InvocationContext> }

Read-only view of the invocation. This is what dynamic instruction providers receive — they can read session state, services and run config but are not expected to mutate.

pub struct CallbackContext { pub invocation: Arc<InvocationContext> }

Mutable callback context. Same shape today; the docs in source note a future revision may add per-callback mutation helpers.

§Callback types

Each callback is a type alias for an Arc<dyn Fn(...) -> BoxFuture<Result<Option<T>>>> — cheap to clone and store. The Option return is the control signal: None means “continue normally”, Some(value) short-circuits or rewrites, depending on the hook.

type BeforeAgentCallback = Arc<dyn Fn(&mut CallbackContext) -> BoxFuture<Result<Option<Content>>> + Send + Sync>

Runs before an agent. Returning Some(content) short-circuits the agent entirely and uses that content as its sole response.

type AfterAgentCallback = BeforeAgentCallback

Runs after an agent. Returning Some(content) appends one more event carrying that content after the agent’s final response.

type BeforeModelCallback = Arc<dyn Fn(&mut CallbackContext, &mut LlmRequest) -> BoxFuture<Result<Option<LlmResponse>>> + Send + Sync>

Runs before each model call. May mutate the outgoing LlmRequest in place (system text, contents, config) or short-circuit the call by returning a synthetic LlmResponse.

type AfterModelCallback = Arc<dyn Fn(&mut CallbackContext, &mut LlmResponse) -> BoxFuture<Result<Option<LlmResponse>>> + Send + Sync>

Runs after each model call; may edit the response in place or return a replacement.

type OnModelErrorCallback = Arc<dyn Fn(&mut CallbackContext, &mut LlmRequest, &Error) -> BoxFuture<Result<Option<LlmResponse>>> + Send + Sync>

Recovery hook for model failures: inspect the error and request, optionally return a substitute LlmResponse instead of propagating the error.

type BeforeToolCallback = Arc<dyn Fn(&mut ToolContext, &Arc<dyn DynTool>, &mut Value) -> BoxFuture<Result<Option<Value>>> + Send + Sync>

Runs before a tool executes, with mutable access to the JSON args. Returning Some(result) skips the tool and uses that value as its response.

type AfterToolCallback = Arc<dyn Fn(&mut ToolContext, &Arc<dyn DynTool>, &Value, &mut Value) -> BoxFuture<Result<Option<Value>>> + Send + Sync>

Runs after a tool, seeing the original args and the mutable result; may rewrite the result in place or return a replacement.

type OnToolErrorCallback = Arc<dyn Fn(&mut ToolContext, &Arc<dyn DynTool>, &Value, &Error) -> BoxFuture<Result<Option<Value>>> + Send + Sync>

Recovery hook for tool failures: optionally return a substitute result value.

The crate also exports a before_agent_callback! macro that turns an ordinary async closure into a BeforeAgentCallback-shaped value, sparing you the Arc/BoxFuture boilerplate.

use std::sync::Arc;
use adk_rs::before_agent_callback;
use adk_rs::core::{BeforeAgentCallback, BeforeModelCallback, CallbackContext};

// Via the macro: an async closure over &mut CallbackContext.
let gate: BeforeAgentCallback = before_agent_callback!(|_ctx| async move {
    // Return Some(content) to short-circuit the agent.
    Ok(None)
});

// By hand: a before-model callback that rewrites the outgoing request.
let guard: BeforeModelCallback = Arc::new(|_ctx: &mut CallbackContext, req| {
    Box::pin(async move {
        req.append_system_text("Never reveal internal tool names.");
        Ok(None) // None = proceed with the (mutated) request
    })
});

§The BasePlugin trait

Plugins live in src/runner/plugin.rs. All hooks have safe defaults, so an implementation overrides only what it needs:

async fn on_register(&self) -> Result<()>

Called once when the plugin is registered with the manager.

async fn before_run(&self, ctx: &InvocationContext) -> Result<()>

Called before each invocation begins (after the user event is persisted).

async fn on_event(&self, ctx: &InvocationContext, event: &Event) -> Result<()>

Called for every event the runner yields — model turns, tool responses, checkpoints.

async fn after_run(&self, ctx: &InvocationContext, err: Option<&Error>) -> Result<()>

Called when the invocation finishes, gracefully or with an error (the error is passed in).

§PluginManager and registration

PluginManager holds an ordered Vec<Arc<dyn BasePlugin>> and fans every hook out to each plugin in registration order. You rarely touch it directly — Runner::builder().plugin(...) registers into the builder’s manager for you (note that plugin is async because it awaits on_register).

pub fn new() -> PluginManager

Constructs an empty manager.

pub async fn register(&mut self, p: Arc<dyn BasePlugin>) -> Result<()>

Calls the plugin’s on_register and adds it to the fan-out list.

pub async fn plugin(self, p: Arc<dyn BasePlugin>) -> Result<RunnerBuilder>

On RunnerBuilder: registers a plugin with the runner being built.

§LoggingPlugin and a custom example

The crate ships one plugin out of the box: LoggingPlugin, which logs every event at INFO under the adk::event target via the tracing facade (author, invocation id, concatenated text). It pairs well with the telemetry feature.

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use adk_rs::Runner;
use adk_rs::core::{Event, InvocationContext};
use adk_rs::error::Result;
use adk_rs::runner::{BasePlugin, LoggingPlugin};
use async_trait::async_trait;

#[derive(Debug, Default)]
struct EventCounter {
    seen: AtomicUsize,
}

#[async_trait]
impl BasePlugin for EventCounter {
    async fn on_event(&self, _ctx: &InvocationContext, _event: &Event) -> Result<()> {
        self.seen.fetch_add(1, Ordering::Relaxed);
        Ok(())
    }
}

let runner = Runner::builder()
    .app_name("hello")
    .agent(agent)
    .session_service(sessions)
    .plugin(Arc::new(LoggingPlugin)).await?
    .plugin(Arc::new(EventCounter::default())).await?
    .build()?;

§Related pages

• Runner — where plugin hooks fire in the invocation lifecycle.
• Events — the payload on_event observes.
• LlmAgent — the builder’s *_callback registration methods and dynamic instruction providers.
• Telemetry — structured tracing to complement LoggingPlugin.