use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::mpsc;
use crate::tools::{BashCompletionSink, ToolMutability, ToolResult, ToolSchema};
use crate::{AgentEvent, PendingQuestion};
pub enum ToolOutcome {
Completed(ToolResult),
Running(RunningHandle),
NeedsHuman {
question: PendingQuestion,
result: ToolResult,
},
}
impl std::fmt::Debug for ToolOutcome {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ToolOutcome::Completed(result) => f.debug_tuple("Completed").field(result).finish(),
ToolOutcome::Running(handle) => {
write!(f, "Running(tool_call_id={:?})", handle.tool_call_id)
}
ToolOutcome::NeedsHuman { question, .. } => {
write!(f, "NeedsHuman(tool_call_id={:?})", question.tool_call_id)
}
}
}
}
impl ToolOutcome {
pub fn into_tool_result(self) -> ToolResult {
match self {
ToolOutcome::Completed(result) => result,
ToolOutcome::Running(handle) => handle.ack,
ToolOutcome::NeedsHuman { result, .. } => result,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ToolClass {
pub mutability: ToolMutability,
pub parallel_safe: bool,
pub promotable: bool,
}
impl ToolClass {
pub const MUTATING_SERIAL: Self = Self {
mutability: ToolMutability::Mutating,
parallel_safe: false,
promotable: false,
};
pub const READONLY_PARALLEL: Self = Self {
mutability: ToolMutability::ReadOnly,
parallel_safe: true,
promotable: false,
};
pub const fn promotable(mut self) -> Self {
self.promotable = true;
self
}
}
impl Default for ToolClass {
fn default() -> Self {
Self::MUTATING_SERIAL
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AsyncWaitKind {
AsyncTools,
Children,
}
pub type ToolResultFuture = Pin<Box<dyn Future<Output = ToolResult> + Send + 'static>>;
pub enum RunningCompletion {
Detached,
Driven(ToolResultFuture),
}
pub struct RunningHandle {
pub tool_call_id: String,
pub ack: ToolResult,
pub completion: RunningCompletion,
pub wait_kind: AsyncWaitKind,
pub kill: Box<dyn FnOnce() + Send>,
}
#[derive(Clone)]
pub struct ToolCtx {
pub session_id: Option<Arc<str>>,
pub tool_call_id: Arc<str>,
pub event_tx: Option<mpsc::Sender<AgentEvent>>,
pub available_tool_schemas: Arc<[ToolSchema]>,
pub bypass_permissions: bool,
pub can_async_resume: bool,
pub async_completion_sink: Option<Arc<dyn AsyncToolCompletionSink>>,
pub bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
}
impl ToolCtx {
pub fn none(tool_call_id: impl Into<Arc<str>>) -> Self {
Self {
session_id: None,
tool_call_id: tool_call_id.into(),
event_tx: None,
available_tool_schemas: Arc::from(Vec::new()),
bypass_permissions: false,
can_async_resume: false,
async_completion_sink: None,
bash_completion_sink: None,
}
}
pub fn cloned_bash_completion_sink(&self) -> Option<Arc<dyn BashCompletionSink>> {
self.bash_completion_sink.clone()
}
pub fn session_id(&self) -> Option<&str> {
self.session_id.as_deref()
}
pub fn cloned_sender(&self) -> Option<mpsc::Sender<AgentEvent>> {
self.event_tx.clone()
}
pub fn cloned_async_completion_sink(&self) -> Option<Arc<dyn AsyncToolCompletionSink>> {
self.async_completion_sink.clone()
}
pub async fn emit(&self, event: AgentEvent) {
if let Some(tx) = &self.event_tx {
let event = match event {
AgentEvent::Token { content } => AgentEvent::ToolToken {
tool_call_id: self.tool_call_id.to_string(),
content,
},
other => other,
};
let _ = tx.try_send(event);
}
}
pub async fn emit_tool_token(&self, content: impl Into<String>) {
self.emit(AgentEvent::ToolToken {
tool_call_id: self.tool_call_id.to_string(),
content: content.into(),
})
.await;
}
}
pub struct AsyncToolCompletionInfo {
pub session_id: String,
pub tool_call_id: String,
pub result: ToolResult,
}
pub trait AsyncToolCompletionSink: Send + Sync {
fn on_tool_completed(&self, info: AsyncToolCompletionInfo);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn tool_class_defaults_and_builders() {
assert_eq!(ToolClass::default(), ToolClass::MUTATING_SERIAL);
assert_eq!(ToolClass::MUTATING_SERIAL.mutability, ToolMutability::Mutating);
assert!(!ToolClass::MUTATING_SERIAL.parallel_safe);
assert!(!ToolClass::MUTATING_SERIAL.promotable);
assert_eq!(ToolClass::READONLY_PARALLEL.mutability, ToolMutability::ReadOnly);
assert!(ToolClass::READONLY_PARALLEL.parallel_safe);
assert!(ToolClass::READONLY_PARALLEL.promotable().promotable);
}
#[test]
fn tool_ctx_none_and_accessors() {
let ctx = ToolCtx::none("call-1");
assert_eq!(ctx.tool_call_id.as_ref(), "call-1");
assert!(ctx.session_id().is_none());
assert!(ctx.cloned_sender().is_none());
assert!(ctx.cloned_async_completion_sink().is_none());
}
}