1use crate::{AgentTool, ToolCallEventStream, ToolInput};
2use agent_client_protocol::schema 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.recv().await.map_err(|e| e.to_string())?;
94
95 event_stream.update_plan(Self::to_plan(input));
96
97 Ok("Plan updated".to_string())
98 })
99 }
100
101 fn replay(
102 &self,
103 input: Self::Input,
104 _output: Self::Output,
105 event_stream: ToolCallEventStream,
106 _cx: &mut App,
107 ) -> anyhow::Result<()> {
108 event_stream.update_plan(Self::to_plan(input));
109 Ok(())
110 }
111}
112
113#[cfg(test)]
114mod tests {
115 use super::*;
116 use crate::ToolCallEventStream;
117 use gpui::TestAppContext;
118 use pretty_assertions::assert_eq;
119
120 fn sample_input() -> UpdatePlanToolInput {
121 UpdatePlanToolInput {
122 plan: vec![
123 PlanItem {
124 step: "Inspect the existing tool wiring".to_string(),
125 status: PlanEntryStatus::Completed,
126 },
127 PlanItem {
128 step: "Implement the update_plan tool".to_string(),
129 status: PlanEntryStatus::InProgress,
130 },
131 PlanItem {
132 step: "Add tests".to_string(),
133 status: PlanEntryStatus::Pending,
134 },
135 ],
136 }
137 }
138
139 #[gpui::test]
140 async fn test_run_emits_plan_event(cx: &mut TestAppContext) {
141 let tool = Arc::new(UpdatePlanTool);
142 let (event_stream, mut event_rx) = ToolCallEventStream::test();
143
144 let input = sample_input();
145 let result = cx
146 .update(|cx| tool.run(ToolInput::resolved(input.clone()), event_stream, cx))
147 .await
148 .expect("tool should succeed");
149
150 assert_eq!(result, "Plan updated".to_string());
151
152 let plan = event_rx.expect_plan().await;
153 assert_eq!(
154 plan,
155 acp::Plan::new(vec![
156 acp::PlanEntry::new(
157 "Inspect the existing tool wiring",
158 acp::PlanEntryPriority::Medium,
159 acp::PlanEntryStatus::Completed,
160 ),
161 acp::PlanEntry::new(
162 "Implement the update_plan tool",
163 acp::PlanEntryPriority::Medium,
164 acp::PlanEntryStatus::InProgress,
165 ),
166 acp::PlanEntry::new(
167 "Add tests",
168 acp::PlanEntryPriority::Medium,
169 acp::PlanEntryStatus::Pending,
170 ),
171 ])
172 );
173 }
174
175 #[gpui::test]
176 async fn test_replay_emits_plan_event(cx: &mut TestAppContext) {
177 let tool = UpdatePlanTool;
178 let (event_stream, mut event_rx) = ToolCallEventStream::test();
179
180 let input = sample_input();
181
182 cx.update(|cx| {
183 tool.replay(input.clone(), "Plan updated".to_string(), event_stream, cx)
184 .expect("replay should succeed");
185 });
186
187 let plan = event_rx.expect_plan().await;
188 assert_eq!(
189 plan,
190 acp::Plan::new(vec![
191 acp::PlanEntry::new(
192 "Inspect the existing tool wiring",
193 acp::PlanEntryPriority::Medium,
194 acp::PlanEntryStatus::Completed,
195 ),
196 acp::PlanEntry::new(
197 "Implement the update_plan tool",
198 acp::PlanEntryPriority::Medium,
199 acp::PlanEntryStatus::InProgress,
200 ),
201 acp::PlanEntry::new(
202 "Add tests",
203 acp::PlanEntryPriority::Medium,
204 acp::PlanEntryStatus::Pending,
205 ),
206 ])
207 );
208 }
209
210 #[gpui::test]
211 async fn test_initial_title(cx: &mut TestAppContext) {
212 let tool = UpdatePlanTool;
213
214 let title = cx.update(|cx| tool.initial_title(Ok(sample_input()), cx));
215 assert_eq!(title, SharedString::from("Update plan"));
216
217 let title =
218 cx.update(|cx| tool.initial_title(Ok(UpdatePlanToolInput { plan: Vec::new() }), cx));
219 assert_eq!(title, SharedString::from("Clear plan"));
220 }
221}