Skip to main content

rskit_tool/
context.rs

1//! Execution context for tool calls.
2
3use tokio_util::sync::CancellationToken;
4
5use crate::io::{ToolMetadata, ToolOutput};
6
7/// Carries per-request metadata and cancellation through tool execution.
8#[derive(Clone)]
9pub struct Context {
10    /// Unique request identifier.
11    pub request_id: String,
12    /// Tool-use identifier (links back to the LLM tool call).
13    pub tool_use_id: String,
14    /// Maximum size (in bytes) for the result content.
15    pub max_result_size: usize,
16    metadata: ToolMetadata,
17    cancel_token: CancellationToken,
18}
19
20impl Context {
21    /// Create a context with empty identifiers, no result-size limit, and a fresh cancellation token.
22    pub fn new() -> Self {
23        Self {
24            request_id: String::new(),
25            tool_use_id: String::new(),
26            max_result_size: 0,
27            metadata: ToolMetadata::new(),
28            cancel_token: CancellationToken::new(),
29        }
30    }
31
32    /// Create a context that observes the provided cancellation token.
33    pub fn with_cancellation(token: CancellationToken) -> Self {
34        Self {
35            cancel_token: token,
36            ..Self::new()
37        }
38    }
39
40    /// Store metadata for the current tool call under `key`.
41    pub fn set(&mut self, key: &str, value: ToolOutput) {
42        self.metadata.insert(key.to_string(), value);
43    }
44
45    /// Return metadata previously stored under `key`, if present.
46    pub fn get(&self, key: &str) -> Option<&ToolOutput> {
47        self.metadata.get(key)
48    }
49
50    /// Return the cancellation token associated with this tool call.
51    pub const fn cancel_token(&self) -> &CancellationToken {
52        &self.cancel_token
53    }
54
55    /// Return whether cancellation has been requested for this tool call.
56    pub fn is_cancelled(&self) -> bool {
57        self.cancel_token.is_cancelled()
58    }
59}
60
61impl Default for Context {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl std::fmt::Debug for Context {
68    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69        f.debug_struct("Context")
70            .field("request_id", &self.request_id)
71            .field("tool_use_id", &self.tool_use_id)
72            .field("max_result_size", &self.max_result_size)
73            .field("metadata_keys", &self.metadata.keys().collect::<Vec<_>>())
74            .field("cancelled", &self.is_cancelled())
75            .finish_non_exhaustive()
76    }
77}