1use tokio_util::sync::CancellationToken;
4
5use crate::io::{ToolMetadata, ToolOutput};
6
7#[derive(Clone)]
9pub struct Context {
10 pub request_id: String,
12 pub tool_use_id: String,
14 pub max_result_size: usize,
16 metadata: ToolMetadata,
17 cancel_token: CancellationToken,
18}
19
20impl Context {
21 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 pub fn with_cancellation(token: CancellationToken) -> Self {
34 Self {
35 cancel_token: token,
36 ..Self::new()
37 }
38 }
39
40 pub fn set(&mut self, key: &str, value: ToolOutput) {
42 self.metadata.insert(key.to_string(), value);
43 }
44
45 pub fn get(&self, key: &str) -> Option<&ToolOutput> {
47 self.metadata.get(key)
48 }
49
50 pub const fn cancel_token(&self) -> &CancellationToken {
52 &self.cancel_token
53 }
54
55 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}