1use super::SubAgentProgressPayload;
2use llm::types::IsoString;
3use llm::{ChatMessage, ContentBlock, ToolCallError, ToolCallRequest, ToolCallResult, ToolDefinition};
4use mcp_utils::display_meta::ToolResultMeta;
5use schemars::JsonSchema;
6use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
10#[serde(tag = "type", rename_all = "snake_case")]
11pub enum ToolEvent {
12 Call { request: ToolCallRequest },
14 CallUpdate { tool_call_id: String, chunk: String },
16 ExecutionStarted { tool_id: String, tool_name: String },
18 Progress { request: ToolCallRequest, progress: f64, total: Option<f64>, message: Option<String> },
20 SubAgentProgress { request: ToolCallRequest, payload: Box<SubAgentProgressPayload> },
22 DisplayUpdate { request: ToolCallRequest, meta: ToolResultMeta },
24 TaskCreated { request: ToolCallRequest, task_id: String, status_message: Option<String> },
26 TaskStatus { request: ToolCallRequest, task_id: String, status: String, status_message: Option<String> },
28 TaskCompleted {
30 request: ToolCallRequest,
31 task_id: String,
32 result: ToolCallResult,
33 result_meta: Option<ToolResultMeta>,
34 },
35 TaskFailed { request: ToolCallRequest, task_id: String, error: ToolCallError },
37 TaskCancelled { request: ToolCallRequest, task_id: String },
39 Result { result: ToolCallResult, result_meta: Option<ToolResultMeta> },
41 Error { error: ToolCallError },
43 DefinitionsUpdated { tools: Vec<ToolDefinition> },
45}
46
47impl ToolEvent {
48 pub fn task_context_message(&self) -> Option<ChatMessage> {
51 let (request, task_id, status, body) = match self {
52 Self::TaskCompleted { request, task_id, result, .. } => {
53 (request, task_id, "completed", result.result.as_str())
54 }
55 Self::TaskFailed { request, task_id, error } => (request, task_id, "failed", error.error.as_str()),
56 Self::TaskCancelled { request, task_id } => (request, task_id, "cancelled", TASK_CANCELLED_BODY),
57 _ => return None,
58 };
59 Some(task_result_message(request, task_id, status, body))
60 }
61}
62
63#[derive(Debug, Clone, PartialEq)]
64pub struct TaskOutcome {
65 pub request: ToolCallRequest,
66 pub task_id: String,
67 pub state: TaskOutcomeState,
68}
69
70#[derive(Debug, Clone, PartialEq)]
71pub enum TaskOutcomeState {
72 Completed { result: ToolCallResult, result_meta: Option<ToolResultMeta> },
73 Failed { error: ToolCallError },
74 Cancelled,
75}
76
77impl TaskOutcome {
78 pub fn context_message(&self) -> ChatMessage {
79 ChatMessage::User { content: self.content_blocks(), timestamp: IsoString::now() }
80 }
81
82 pub fn content_blocks(&self) -> Vec<ContentBlock> {
83 let (status, body) = match &self.state {
84 TaskOutcomeState::Completed { result, .. } => ("completed", result.result.as_str()),
85 TaskOutcomeState::Failed { error } => ("failed", error.error.as_str()),
86 TaskOutcomeState::Cancelled => ("cancelled", TASK_CANCELLED_BODY),
87 };
88 task_result_content(&self.request, &self.task_id, status, body)
89 }
90}
91
92impl From<TaskOutcome> for ToolEvent {
93 fn from(outcome: TaskOutcome) -> Self {
94 let TaskOutcome { request, task_id, state } = outcome;
95 match state {
96 TaskOutcomeState::Completed { result, result_meta } => {
97 Self::TaskCompleted { request, task_id, result, result_meta }
98 }
99 TaskOutcomeState::Failed { error } => Self::TaskFailed { request, task_id, error },
100 TaskOutcomeState::Cancelled => Self::TaskCancelled { request, task_id },
101 }
102 }
103}
104
105pub fn task_created_result(request: &ToolCallRequest, task_id: &str) -> ToolCallResult {
106 ToolCallResult {
107 id: request.id.clone(),
108 name: request.name.clone(),
109 arguments: request.arguments.clone(),
110 result: format!(
111 "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."
112 ),
113 }
114}
115
116const TASK_CANCELLED_BODY: &str = "The background task was cancelled and will not produce a result.";
117
118fn task_result_message(request: &ToolCallRequest, task_id: &str, status: &str, body: &str) -> ChatMessage {
119 ChatMessage::User { content: task_result_content(request, task_id, status, body), timestamp: IsoString::now() }
120}
121
122fn task_result_content(request: &ToolCallRequest, task_id: &str, status: &str, body: &str) -> Vec<ContentBlock> {
123 let content = format!(
124 "<task-result task-id=\"{}\" tool=\"{}\" status=\"{status}\">{}</task-result>",
125 escape_xml(task_id),
126 escape_xml(&request.name),
127 escape_xml(body),
128 );
129 vec![ContentBlock::text(content)]
130}
131
132fn escape_xml(value: &str) -> String {
133 value.replace('&', "&").replace('<', "<").replace('>', ">").replace('"', """).replace('\'', "'")
134}