1use llm::types::IsoString;
2use llm::{ChatMessage, ContentBlock, ToolCallError, ToolCallRequest, ToolCallResult, ToolDefinition};
3use mcp_utils::display_meta::ToolResultMeta;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
9#[serde(tag = "type", rename_all = "snake_case")]
10pub enum ToolEvent {
11 Call { request: ToolCallRequest },
13 CallUpdate { tool_call_id: String, chunk: String },
15 ExecutionStarted { tool_id: String, tool_name: String },
17 Progress { request: ToolCallRequest, progress: f64, total: Option<f64>, message: Option<String> },
19 TaskCreated { request: ToolCallRequest, task_id: String, status_message: Option<String> },
21 TaskStatus { request: ToolCallRequest, task_id: String, status: String, status_message: Option<String> },
23 TaskCompleted {
25 request: ToolCallRequest,
26 task_id: String,
27 result: ToolCallResult,
28 result_meta: Option<ToolResultMeta>,
29 },
30 TaskFailed { request: ToolCallRequest, task_id: String, error: ToolCallError },
32 TaskCancelled { request: ToolCallRequest, task_id: String },
34 Result { result: ToolCallResult, result_meta: Option<ToolResultMeta> },
36 Error { error: ToolCallError },
38 DefinitionsUpdated { tools: Vec<ToolDefinition> },
40}
41
42impl ToolEvent {
43 pub fn task_context_message(&self) -> Option<ChatMessage> {
46 let (request, task_id, status, body) = match self {
47 Self::TaskCompleted { request, task_id, result, .. } => {
48 (request, task_id, "completed", result.result.as_str())
49 }
50 Self::TaskFailed { request, task_id, error } => (request, task_id, "failed", error.error.as_str()),
51 Self::TaskCancelled { request, task_id } => (request, task_id, "cancelled", TASK_CANCELLED_BODY),
52 _ => return None,
53 };
54 Some(task_result_message(request, task_id, status, body))
55 }
56}
57
58#[derive(Debug, Clone, PartialEq)]
59pub struct TaskOutcome {
60 pub request: ToolCallRequest,
61 pub task_id: String,
62 pub state: TaskOutcomeState,
63}
64
65#[derive(Debug, Clone, PartialEq)]
66pub enum TaskOutcomeState {
67 Completed { result: ToolCallResult, result_meta: Option<ToolResultMeta> },
68 Failed { error: ToolCallError },
69 Cancelled,
70}
71
72impl TaskOutcome {
73 pub fn context_message(&self) -> ChatMessage {
74 ChatMessage::User { content: self.content_blocks(), timestamp: IsoString::now() }
75 }
76
77 pub fn content_blocks(&self) -> Vec<ContentBlock> {
78 let (status, body) = match &self.state {
79 TaskOutcomeState::Completed { result, .. } => ("completed", result.result.as_str()),
80 TaskOutcomeState::Failed { error } => ("failed", error.error.as_str()),
81 TaskOutcomeState::Cancelled => ("cancelled", TASK_CANCELLED_BODY),
82 };
83 task_result_content(&self.request, &self.task_id, status, body)
84 }
85}
86
87impl From<TaskOutcome> for ToolEvent {
88 fn from(outcome: TaskOutcome) -> Self {
89 let TaskOutcome { request, task_id, state } = outcome;
90 match state {
91 TaskOutcomeState::Completed { result, result_meta } => {
92 Self::TaskCompleted { request, task_id, result, result_meta }
93 }
94 TaskOutcomeState::Failed { error } => Self::TaskFailed { request, task_id, error },
95 TaskOutcomeState::Cancelled => Self::TaskCancelled { request, task_id },
96 }
97 }
98}
99
100pub fn task_created_result(request: &ToolCallRequest, task_id: &str) -> ToolCallResult {
101 ToolCallResult {
102 id: request.id.clone(),
103 name: request.name.clone(),
104 arguments: request.arguments.clone(),
105 result: format!(
106 "This tool is running as a background task, id: {task_id}. The result will be automatically injected into context when it completes, you may continue working."
107 ),
108 }
109}
110
111const TASK_CANCELLED_BODY: &str = "The background task was cancelled and will not produce a result.";
112
113fn task_result_message(request: &ToolCallRequest, task_id: &str, status: &str, body: &str) -> ChatMessage {
114 ChatMessage::User { content: task_result_content(request, task_id, status, body), timestamp: IsoString::now() }
115}
116
117fn task_result_content(request: &ToolCallRequest, task_id: &str, status: &str, body: &str) -> Vec<ContentBlock> {
118 let content = format!(
119 "<task-result task-id=\"{}\" tool=\"{}\" status=\"{status}\">{}</task-result>",
120 escape_xml(task_id),
121 escape_xml(&request.name),
122 escape_xml(body),
123 );
124 vec![ContentBlock::text(content)]
125}
126
127fn escape_xml(value: &str) -> String {
128 value.replace('&', "&").replace('<', "<").replace('>', ">").replace('"', """).replace('\'', "'")
129}