adk-plugin 0.6.0

//! Plugin callback types
//!
//! Defines the callback function signatures used by plugins.

pub use adk_core::OnToolErrorCallback;
use adk_core::{
    CallbackContext, Content, Event, InvocationContext, LlmRequest, LlmResponse, Result,
};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

/// Callback invoked when a user message is received.
///
/// Can modify the message content before processing.
/// Return `Ok(Some(content))` to replace the message, `Ok(None)` to keep original.
pub type OnUserMessageCallback = Box<
    dyn Fn(
            Arc<dyn InvocationContext>,
            Content,
        ) -> Pin<Box<dyn Future<Output = Result<Option<Content>>> + Send>>
        + Send
        + Sync,
>;

/// Callback invoked for each event generated by the agent.
///
/// Can modify events before they are yielded.
/// Return `Ok(Some(event))` to replace the event, `Ok(None)` to keep original.
pub type OnEventCallback = Box<
    dyn Fn(
            Arc<dyn InvocationContext>,
            Event,
        ) -> Pin<Box<dyn Future<Output = Result<Option<Event>>> + Send>>
        + Send
        + Sync,
>;

/// Callback invoked before the agent run starts.
///
/// Can perform setup, validation, or early exit.
/// Return `Ok(Some(content))` to skip the run and return this content.
/// Return `Ok(None)` to continue with the run.
pub type BeforeRunCallback = Box<
    dyn Fn(
            Arc<dyn InvocationContext>,
        ) -> Pin<Box<dyn Future<Output = Result<Option<Content>>> + Send>>
        + Send
        + Sync,
>;

/// Callback invoked after the agent run completes.
///
/// Used for cleanup, logging, metrics collection.
/// This callback does NOT emit events - it's for side effects only.
pub type AfterRunCallback = Box<
    dyn Fn(Arc<dyn InvocationContext>) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync,
>;

/// Callback invoked when a model error occurs.
///
/// Can handle or transform model errors.
/// Return `Ok(Some(response))` to provide a fallback response.
/// Return `Ok(None)` to propagate the original error.
pub type OnModelErrorCallback = Box<
    dyn Fn(
            Arc<dyn CallbackContext>,
            LlmRequest,
            String, // error message
        ) -> Pin<Box<dyn Future<Output = Result<Option<LlmResponse>>> + Send>>
        + Send
        + Sync,
>;

/// Helper to create a simple logging callback for user messages.
pub fn log_user_messages() -> OnUserMessageCallback {
    Box::new(|_ctx, content| {
        Box::pin(async move {
            tracing::info!(role = %content.role, parts = ?content.parts.len(), "User message received");
            Ok(None)
        })
    })
}

/// Helper to create a simple logging callback for events.
pub fn log_events() -> OnEventCallback {
    Box::new(|_ctx, event| {
        Box::pin(async move {
            tracing::info!(
                id = %event.id,
                author = %event.author,
                partial = event.llm_response.partial,
                "Event generated"
            );
            Ok(None)
        })
    })
}

/// Helper to create a metrics collection callback.
pub fn collect_metrics(
    on_run_start: impl Fn() + Send + Sync + 'static,
    on_run_end: impl Fn() + Send + Sync + 'static,
) -> (BeforeRunCallback, AfterRunCallback) {
    let start_fn = Arc::new(on_run_start);
    let end_fn = Arc::new(on_run_end);

    let before = Box::new(move |_ctx: Arc<dyn InvocationContext>| {
        let f = start_fn.clone();
        Box::pin(async move {
            f();
            Ok(None)
        }) as Pin<Box<dyn Future<Output = Result<Option<Content>>> + Send>>
    });

    let after = Box::new(move |_ctx: Arc<dyn InvocationContext>| {
        let f = end_fn.clone();
        Box::pin(async move {
            f();
        }) as Pin<Box<dyn Future<Output = ()> + Send>>
    });

    (before, after)
}