Skip to main content

ferrin_tool/
execute.rs

1//! Execution contract: [`ToolExecute`], [`ToolOutput`], [`ToolContext`].
2
3use std::fmt;
4use std::sync::Arc;
5
6use ferrin_message::Message;
7use ferrin_spec::BoxStream;
8use ferrin_spec::JsonValue;
9use ferrin_spec::ToolCallId;
10use futures_util::StreamExt;
11use tokio_util::sync::CancellationToken;
12
13use crate::error::ToolError;
14
15/// Stream of outputs produced by one execution.
16pub type ToolOutputStream = BoxStream<'static, Result<ToolOutput, ToolError>>;
17
18/// One output of a tool execution.
19///
20/// Streaming tools emit any number of [`ToolOutput::Preliminary`] values
21/// followed by exactly one [`ToolOutput::Final`]; single-value tools emit
22/// only the final value.
23#[derive(Debug, Clone, PartialEq, Eq)]
24#[non_exhaustive]
25pub enum ToolOutput {
26    /// An intermediate result (`preliminary: true` in the reference model).
27    Preliminary(JsonValue),
28    /// The final result.
29    Final(JsonValue),
30}
31
32impl ToolOutput {
33    /// Returns `true` for the final output.
34    #[must_use]
35    pub fn is_final(&self) -> bool {
36        matches!(self, Self::Final(_))
37    }
38
39    /// The JSON value.
40    #[must_use]
41    pub fn value(&self) -> &JsonValue {
42        match self {
43            Self::Preliminary(value) | Self::Final(value) => value,
44        }
45    }
46
47    /// Consumes the output, returning the JSON value.
48    #[must_use]
49    pub fn into_value(self) -> JsonValue {
50        match self {
51            Self::Preliminary(value) | Self::Final(value) => value,
52        }
53    }
54}
55
56/// Executes a tool.
57///
58/// Implement this for custom execution strategies (remote execution,
59/// recording); ordinary tools use the closure adapters on [`crate::ToolBuilder`].
60/// Implementations receive input that already passed the tool's input
61/// schema and must honour `ctx.cancellation`.
62pub trait ToolExecute: Send + Sync {
63    /// Starts an execution.
64    fn execute(&self, input: JsonValue, ctx: ToolContext) -> ToolOutputStream;
65}
66
67/// Per-call information handed to a tool execution.
68#[derive(Clone)]
69pub struct ToolContext {
70    /// Id of the tool call being executed.
71    pub tool_call_id: ToolCallId,
72    /// Messages sent to the model for the step that produced the call
73    /// (without the system prompt and without the assistant response).
74    pub messages: Arc<[Message]>,
75    /// Cancels the execution.
76    pub cancellation: CancellationToken,
77    /// This tool's selected context, validated when it declares a context schema.
78    pub tools_context: Option<JsonValue>,
79    /// Sandbox the tool operates in.
80    #[cfg(feature = "sandbox")]
81    pub sandbox: Option<Arc<dyn crate::sandbox::Sandbox>>,
82}
83
84impl fmt::Debug for ToolContext {
85    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
86        let mut debug = f.debug_struct("ToolContext");
87        debug
88            .field("tool_call_id", &self.tool_call_id)
89            .field("messages", &self.messages.len())
90            .field("cancelled", &self.cancellation.is_cancelled())
91            .field("tools_context", &self.tools_context);
92        #[cfg(feature = "sandbox")]
93        debug.field(
94            "sandbox",
95            &self.sandbox.as_ref().map(|sandbox| sandbox.description()),
96        );
97        debug.finish()
98    }
99}
100
101impl ToolContext {
102    /// Creates a context with no messages, a fresh token and no tool context.
103    #[must_use]
104    pub fn new(tool_call_id: impl Into<ToolCallId>) -> Self {
105        Self {
106            tool_call_id: tool_call_id.into(),
107            messages: Arc::from(Vec::new()),
108            cancellation: CancellationToken::new(),
109            tools_context: None,
110            #[cfg(feature = "sandbox")]
111            sandbox: None,
112        }
113    }
114
115    /// Sets the step messages.
116    #[must_use]
117    pub fn with_messages(mut self, messages: impl Into<Arc<[Message]>>) -> Self {
118        self.messages = messages.into();
119        self
120    }
121
122    /// Sets the cancellation token.
123    #[must_use]
124    pub fn with_cancellation(mut self, cancellation: CancellationToken) -> Self {
125        self.cancellation = cancellation;
126        self
127    }
128
129    /// Sets the validated tool context.
130    #[must_use]
131    pub fn with_tools_context(mut self, tools_context: Option<JsonValue>) -> Self {
132        self.tools_context = tools_context;
133        self
134    }
135
136    /// Sets the sandbox.
137    #[cfg(feature = "sandbox")]
138    #[must_use]
139    pub fn with_sandbox(mut self, sandbox: Arc<dyn crate::sandbox::Sandbox>) -> Self {
140        self.sandbox = Some(sandbox);
141        self
142    }
143}
144
145/// Drives an output stream to its end, forwarding preliminary values to
146/// `on_preliminary` and returning the final value.
147///
148/// # Errors
149///
150/// Returns the first error produced by the stream, or a
151/// [`ToolError::Message`] when the stream ends without a final output.
152pub async fn execute_to_completion(
153    mut stream: ToolOutputStream,
154    mut on_preliminary: impl FnMut(JsonValue),
155) -> Result<JsonValue, ToolError> {
156    while let Some(item) = stream.next().await {
157        match item? {
158            ToolOutput::Preliminary(value) => on_preliminary(value),
159            ToolOutput::Final(value) => return Ok(value),
160        }
161    }
162    Err(ToolError::message("tool produced no final output"))
163}