1use crate::{AgentTool, ToolCallEventStream, ToolInput};
2use agent_client_protocol as acp;
3use gpui::{App, SharedString, Task};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use std::sync::Arc;
7
8#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
9#[serde(rename_all = "snake_case")]
10#[schemars(inline)]
11pub enum PlanEntryStatus {
12 /// The task has not started yet.
13 Pending,
14 /// The task is currently being worked on.
15 InProgress,
16 /// The task has been successfully completed.
17 Completed,
18}
19
20impl From<PlanEntryStatus> for acp::PlanEntryStatus {
21 fn from(value: PlanEntryStatus) -> Self {
22 match value {
23 PlanEntryStatus::Pending => acp::PlanEntryStatus::Pending,
24 PlanEntryStatus::InProgress => acp::PlanEntryStatus::InProgress,
25 PlanEntryStatus::Completed => acp::PlanEntryStatus::Completed,
26 }
27 }
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
31pub struct PlanItem {
32 /// Human-readable description of what this task aims to accomplish.
33 pub step: String,
34 /// The current status of this task.
35 pub status: PlanEntryStatus,
36}
37
38impl From<PlanItem> for acp::PlanEntry {
39 fn from(value: PlanItem) -> Self {
40 acp::PlanEntry::new(
41 value.step,
42 acp::PlanEntryPriority::Medium,
43 value.status.into(),
44 )
45 }
46}
47
48/// Updates the task plan.
49///
50/// Provide a list of plan entries, each with a step and status.
51#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
52pub struct UpdatePlanToolInput {
53 /// The list of plan entries and their current statuses.
54 pub plan: Vec<PlanItem>,
55}
56
57pub struct UpdatePlanTool;
58
59impl UpdatePlanTool {
60 fn to_plan(input: UpdatePlanToolInput) -> acp::Plan {
61 acp::Plan::new(input.plan.into_iter().map(Into::into).collect())
62 }
63}
64
65impl AgentTool for UpdatePlanTool {
66 type Input = UpdatePlanToolInput;
67 type Output = String;
68
69 const NAME: &'static str = "update_plan";
70
71 fn kind() -> acp::ToolKind {
72 acp::ToolKind::Think
73 }
74
75 fn initial_title(
76 &self,
77 input: Result<Self::Input, serde_json::Value>,
78 _cx: &mut App,
79 ) -> SharedString {
80 match input {
81 Ok(input) if input.plan.is_empty() => "Clear plan".into(),
82 Ok(_) | Err(_) => "Update plan".into(),
83 }
84 }
85
86 fn run(
87 self: Arc<Self>,
88 input: ToolInput<Self::Input>,
89 event_stream: ToolCallEventStream,
90 cx: &mut App,
91 ) -> Task<Result<Self::Output, Self::Output>> {
92 cx.spawn(async move |_cx| {
93 let input = input
94 .recv()
95 .await
96 .map_err(|e| format!("Failed to receive tool input: {e}"))?;
97
98 event_stream.update_plan(Self::to_plan(input));
99
100 Ok("Plan updated".to_string())
101 })
102 }
103
104 fn replay(
105 &self,
106 input: Self::Input,
107 _output: Self::Output,
108 event_stream: ToolCallEventStream,
109 _cx: &mut App,
110 ) -> anyhow::Result<()> {
111 event_stream.update_plan(Self::to_plan(input));
112 Ok(())
113 }
114}
115
116#[cfg(test)]
117mod tests {
118 use super::*;
119 use crate::ToolCallEventStream;
120 use gpui::TestAppContext;
121 use pretty_assertions::assert_eq;
122
123 fn sample_input() -> UpdatePlanToolInput {
124 UpdatePlanToolInput {
125 plan: vec![
126 PlanItem {
127 step: "Inspect the existing tool wiring".to_string(),
128 status: PlanEntryStatus::Completed,
129 },
130 PlanItem {
131 step: "Implement the update_plan tool".to_string(),
132 status: PlanEntryStatus::InProgress,
133 },
134 PlanItem {
135 step: "Add tests".to_string(),
136 status: PlanEntryStatus::Pending,
137 },
138 ],
139 }
140 }
141
142 #[gpui::test]
143 async fn test_run_emits_plan_event(cx: &mut TestAppContext) {
144 let tool = Arc::new(UpdatePlanTool);
145 let (event_stream, mut event_rx) = ToolCallEventStream::test();
146
147 let input = sample_input();
148 let result = cx
149 .update(|cx| tool.run(ToolInput::resolved(input.clone()), event_stream, cx))
150 .await
151 .expect("tool should succeed");
152
153 assert_eq!(result, "Plan updated".to_string());
154
155 let plan = event_rx.expect_plan().await;
156 assert_eq!(
157 plan,
158 acp::Plan::new(vec![
159 acp::PlanEntry::new(
160 "Inspect the existing tool wiring",
161 acp::PlanEntryPriority::Medium,
162 acp::PlanEntryStatus::Completed,
163 ),
164 acp::PlanEntry::new(
165 "Implement the update_plan tool",
166 acp::PlanEntryPriority::Medium,
167 acp::PlanEntryStatus::InProgress,
168 ),
169 acp::PlanEntry::new(
170 "Add tests",
171 acp::PlanEntryPriority::Medium,
172 acp::PlanEntryStatus::Pending,
173 ),
174 ])
175 );
176 }
177
178 #[gpui::test]
179 async fn test_replay_emits_plan_event(cx: &mut TestAppContext) {
180 let tool = UpdatePlanTool;
181 let (event_stream, mut event_rx) = ToolCallEventStream::test();
182
183 let input = sample_input();
184
185 cx.update(|cx| {
186 tool.replay(input.clone(), "Plan updated".to_string(), event_stream, cx)
187 .expect("replay should succeed");
188 });
189
190 let plan = event_rx.expect_plan().await;
191 assert_eq!(
192 plan,
193 acp::Plan::new(vec![
194 acp::PlanEntry::new(
195 "Inspect the existing tool wiring",
196 acp::PlanEntryPriority::Medium,
197 acp::PlanEntryStatus::Completed,
198 ),
199 acp::PlanEntry::new(
200 "Implement the update_plan tool",
201 acp::PlanEntryPriority::Medium,
202 acp::PlanEntryStatus::InProgress,
203 ),
204 acp::PlanEntry::new(
205 "Add tests",
206 acp::PlanEntryPriority::Medium,
207 acp::PlanEntryStatus::Pending,
208 ),
209 ])
210 );
211 }
212
213 #[gpui::test]
214 async fn test_initial_title(cx: &mut TestAppContext) {
215 let tool = UpdatePlanTool;
216
217 let title = cx.update(|cx| tool.initial_title(Ok(sample_input()), cx));
218 assert_eq!(title, SharedString::from("Update plan"));
219
220 let title =
221 cx.update(|cx| tool.initial_title(Ok(UpdatePlanToolInput { plan: Vec::new() }), cx));
222 assert_eq!(title, SharedString::from("Clear plan"));
223 }
224}