use crate::AgentToolResult;
use futures::future::BoxFuture;
use serde_json::Value;
use std::{future::Future, sync::Arc};
use tokio_util::sync::CancellationToken;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ToolProgramChild {
pub root_id: String,
pub root_name: String,
pub parent_id: String,
pub request_id: String,
pub name: String,
pub ordinal: usize,
}
pub trait ToolProgramObserver: Send + Sync {
fn begin<'a>(
&'a self,
child: ToolProgramChild,
arguments: &'a Value,
signal: CancellationToken,
) -> BoxFuture<'a, Result<Arc<dyn ToolProgramObservation>, String>>;
}
pub trait ToolProgramObservation: Send + Sync {
fn id(&self) -> &str;
fn update(&self, partial: Arc<AgentToolResult>) -> Result<(), String>;
fn prepared<'a>(&'a self, arguments: &'a Value) -> BoxFuture<'a, Result<(), String>>;
fn finish<'a>(
&'a self,
result: &'a AgentToolResult,
is_error: bool,
cancelled: bool,
) -> BoxFuture<'a, Result<(), String>>;
fn close<'a>(&'a self, reason: Option<&'a str>) -> BoxFuture<'a, Result<(), String>>;
}
tokio::task_local! { static CURRENT: Arc<dyn ToolProgramObservation>; }
pub fn current_observation() -> Option<Arc<dyn ToolProgramObservation>> {
CURRENT.try_with(Clone::clone).ok()
}
pub async fn scope<F: Future>(
observation: Arc<dyn ToolProgramObservation>,
future: F,
) -> F::Output {
CURRENT.scope(observation, future).await
}