1mod messages;
2mod supermaven_completion_provider;
3
4pub use supermaven_completion_provider::*;
5
6use anyhow::{Context as _, Result};
7#[allow(unused_imports)]
8use client::{Client, proto};
9use collections::BTreeMap;
10
11use futures::{AsyncBufReadExt, StreamExt, channel::mpsc, io::BufReader};
12use gpui::{App, AsyncApp, Context, Entity, EntityId, Global, Task, WeakEntity, actions};
13use language::{
14 Anchor, Buffer, BufferSnapshot, ToOffset, language_settings::all_language_settings,
15};
16use messages::*;
17use postage::watch;
18use serde::{Deserialize, Serialize};
19use settings::SettingsStore;
20use smol::{
21 io::AsyncWriteExt,
22 process::{Child, ChildStdin, ChildStdout},
23};
24use std::{path::PathBuf, process::Stdio, sync::Arc};
25use ui::prelude::*;
26use util::ResultExt;
27
28actions!(
29 supermaven,
30 [
31 /// Signs out of Supermaven.
32 SignOut
33 ]
34);
35
36pub fn init(client: Arc<Client>, cx: &mut App) {
37 let supermaven = cx.new(|_| Supermaven::Starting);
38 Supermaven::set_global(supermaven.clone(), cx);
39
40 let mut provider = all_language_settings(None, cx).edit_predictions.provider;
41 if provider == language::language_settings::EditPredictionProvider::Supermaven {
42 supermaven.update(cx, |supermaven, cx| supermaven.start(client.clone(), cx));
43 }
44
45 cx.observe_global::<SettingsStore>(move |cx| {
46 let new_provider = all_language_settings(None, cx).edit_predictions.provider;
47 if new_provider != provider {
48 provider = new_provider;
49 if provider == language::language_settings::EditPredictionProvider::Supermaven {
50 supermaven.update(cx, |supermaven, cx| supermaven.start(client.clone(), cx));
51 } else {
52 supermaven.update(cx, |supermaven, _cx| supermaven.stop());
53 }
54 }
55 })
56 .detach();
57
58 cx.on_action(|_: &SignOut, cx| {
59 if let Some(supermaven) = Supermaven::global(cx) {
60 supermaven.update(cx, |supermaven, _cx| supermaven.sign_out());
61 }
62 });
63}
64
65pub enum Supermaven {
66 Starting,
67 FailedDownload { error: anyhow::Error },
68 Spawned(SupermavenAgent),
69 Error { error: anyhow::Error },
70}
71
72#[derive(Clone)]
73pub enum AccountStatus {
74 Unknown,
75 NeedsActivation { activate_url: String },
76 Ready,
77}
78
79#[derive(Clone)]
80struct SupermavenGlobal(Entity<Supermaven>);
81
82impl Global for SupermavenGlobal {}
83
84impl Supermaven {
85 pub fn global(cx: &App) -> Option<Entity<Self>> {
86 cx.try_global::<SupermavenGlobal>()
87 .map(|model| model.0.clone())
88 }
89
90 pub fn set_global(supermaven: Entity<Self>, cx: &mut App) {
91 cx.set_global(SupermavenGlobal(supermaven));
92 }
93
94 pub fn start(&mut self, client: Arc<Client>, cx: &mut Context<Self>) {
95 if let Self::Starting = self {
96 cx.spawn(async move |this, cx| {
97 let binary_path =
98 supermaven_api::get_supermaven_agent_path(client.http_client()).await?;
99
100 this.update(cx, |this, cx| {
101 if let Self::Starting = this {
102 *this =
103 Self::Spawned(SupermavenAgent::new(binary_path, client.clone(), cx)?);
104 }
105 anyhow::Ok(())
106 })
107 })
108 .detach_and_log_err(cx)
109 }
110 }
111
112 pub fn stop(&mut self) {
113 *self = Self::Starting;
114 }
115
116 pub fn is_enabled(&self) -> bool {
117 matches!(self, Self::Spawned { .. })
118 }
119
120 pub fn complete(
121 &mut self,
122 buffer: &Entity<Buffer>,
123 cursor_position: Anchor,
124 cx: &App,
125 ) -> Option<SupermavenCompletion> {
126 if let Self::Spawned(agent) = self {
127 let buffer_id = buffer.entity_id();
128 let buffer = buffer.read(cx);
129 let path = buffer
130 .file()
131 .and_then(|file| Some(file.as_local()?.abs_path(cx)))
132 .unwrap_or_else(|| PathBuf::from("untitled"))
133 .to_string_lossy()
134 .to_string();
135 let content = buffer.text();
136 let offset = cursor_position.to_offset(buffer);
137 let state_id = agent.next_state_id;
138 agent.next_state_id.0 += 1;
139
140 let (updates_tx, mut updates_rx) = watch::channel();
141 postage::stream::Stream::try_recv(&mut updates_rx).unwrap();
142
143 agent.states.insert(
144 state_id,
145 SupermavenCompletionState {
146 buffer_id,
147 prefix_anchor: cursor_position,
148 prefix_offset: offset,
149 text: String::new(),
150 dedent: String::new(),
151 updates_tx,
152 },
153 );
154 // ensure the states map is max 1000 elements
155 if agent.states.len() > 1000 {
156 // state id is monotonic so it's sufficient to remove the first element
157 agent
158 .states
159 .remove(&agent.states.keys().next().unwrap().clone());
160 }
161
162 let _ = agent
163 .outgoing_tx
164 .unbounded_send(OutboundMessage::StateUpdate(StateUpdateMessage {
165 new_id: state_id.0.to_string(),
166 updates: vec![
167 StateUpdate::FileUpdate(FileUpdateMessage {
168 path: path.clone(),
169 content,
170 }),
171 StateUpdate::CursorUpdate(CursorPositionUpdateMessage { path, offset }),
172 ],
173 }));
174
175 Some(SupermavenCompletion {
176 id: state_id,
177 updates: updates_rx,
178 })
179 } else {
180 None
181 }
182 }
183
184 pub fn completion(
185 &self,
186 buffer: &Entity<Buffer>,
187 cursor_position: Anchor,
188 cx: &App,
189 ) -> Option<&str> {
190 if let Self::Spawned(agent) = self {
191 find_relevant_completion(
192 &agent.states,
193 buffer.entity_id(),
194 &buffer.read(cx).snapshot(),
195 cursor_position,
196 )
197 } else {
198 None
199 }
200 }
201
202 pub fn sign_out(&mut self) {
203 if let Self::Spawned(agent) = self {
204 agent
205 .outgoing_tx
206 .unbounded_send(OutboundMessage::Logout)
207 .ok();
208 // The account status will get set to RequiresActivation or Ready when the next
209 // message from the agent comes in. Until that happens, set the status to Unknown
210 // to disable the button.
211 agent.account_status = AccountStatus::Unknown;
212 }
213 }
214}
215
216fn find_relevant_completion<'a>(
217 states: &'a BTreeMap<SupermavenCompletionStateId, SupermavenCompletionState>,
218 buffer_id: EntityId,
219 buffer: &BufferSnapshot,
220 cursor_position: Anchor,
221) -> Option<&'a str> {
222 let mut best_completion: Option<&str> = None;
223 'completions: for state in states.values() {
224 if state.buffer_id != buffer_id {
225 continue;
226 }
227 let Some(state_completion) = state.text.strip_prefix(&state.dedent) else {
228 continue;
229 };
230
231 let current_cursor_offset = cursor_position.to_offset(buffer);
232 if current_cursor_offset < state.prefix_offset {
233 continue;
234 }
235
236 let original_cursor_offset = buffer.clip_offset(state.prefix_offset, text::Bias::Left);
237 let text_inserted_since_completion_request =
238 buffer.text_for_range(original_cursor_offset..current_cursor_offset);
239 let mut trimmed_completion = state_completion;
240 for chunk in text_inserted_since_completion_request {
241 if let Some(suffix) = trimmed_completion.strip_prefix(chunk) {
242 trimmed_completion = suffix;
243 } else {
244 continue 'completions;
245 }
246 }
247
248 if best_completion.map_or(false, |best| best.len() > trimmed_completion.len()) {
249 continue;
250 }
251
252 best_completion = Some(trimmed_completion);
253 }
254 best_completion
255}
256
257pub struct SupermavenAgent {
258 _process: Child,
259 next_state_id: SupermavenCompletionStateId,
260 states: BTreeMap<SupermavenCompletionStateId, SupermavenCompletionState>,
261 outgoing_tx: mpsc::UnboundedSender<OutboundMessage>,
262 _handle_outgoing_messages: Task<Result<()>>,
263 _handle_incoming_messages: Task<Result<()>>,
264 pub account_status: AccountStatus,
265 service_tier: Option<ServiceTier>,
266 #[allow(dead_code)]
267 client: Arc<Client>,
268}
269
270impl SupermavenAgent {
271 fn new(
272 binary_path: PathBuf,
273 client: Arc<Client>,
274 cx: &mut Context<Supermaven>,
275 ) -> Result<Self> {
276 let mut process = util::command::new_smol_command(&binary_path)
277 .arg("stdio")
278 .stdin(Stdio::piped())
279 .stdout(Stdio::piped())
280 .stderr(Stdio::piped())
281 .kill_on_drop(true)
282 .spawn()
283 .context("failed to start the binary")?;
284
285 let stdin = process
286 .stdin
287 .take()
288 .context("failed to get stdin for process")?;
289 let stdout = process
290 .stdout
291 .take()
292 .context("failed to get stdout for process")?;
293
294 let (outgoing_tx, outgoing_rx) = mpsc::unbounded();
295
296 cx.spawn({
297 let client = client.clone();
298 let outgoing_tx = outgoing_tx.clone();
299 async move |this, cx| {
300 let mut status = client.status();
301 while let Some(status) = status.next().await {
302 if status.is_connected() {
303 let api_key = client.request(proto::GetSupermavenApiKey {}).await?.api_key;
304 outgoing_tx
305 .unbounded_send(OutboundMessage::SetApiKey(SetApiKey { api_key }))
306 .ok();
307 this.update(cx, |this, cx| {
308 if let Supermaven::Spawned(this) = this {
309 this.account_status = AccountStatus::Ready;
310 cx.notify();
311 }
312 })?;
313 break;
314 }
315 }
316 anyhow::Ok(())
317 }
318 })
319 .detach();
320
321 Ok(Self {
322 _process: process,
323 next_state_id: SupermavenCompletionStateId::default(),
324 states: BTreeMap::default(),
325 outgoing_tx,
326 _handle_outgoing_messages: cx.spawn(async move |_, _cx| {
327 Self::handle_outgoing_messages(outgoing_rx, stdin).await
328 }),
329 _handle_incoming_messages: cx.spawn(async move |this, cx| {
330 Self::handle_incoming_messages(this, stdout, cx).await
331 }),
332 account_status: AccountStatus::Unknown,
333 service_tier: None,
334 client,
335 })
336 }
337
338 async fn handle_outgoing_messages(
339 mut outgoing: mpsc::UnboundedReceiver<OutboundMessage>,
340 mut stdin: ChildStdin,
341 ) -> Result<()> {
342 while let Some(message) = outgoing.next().await {
343 let bytes = serde_json::to_vec(&message)?;
344 stdin.write_all(&bytes).await?;
345 stdin.write_all(&[b'\n']).await?;
346 }
347 Ok(())
348 }
349
350 async fn handle_incoming_messages(
351 this: WeakEntity<Supermaven>,
352 stdout: ChildStdout,
353 cx: &mut AsyncApp,
354 ) -> Result<()> {
355 const MESSAGE_PREFIX: &str = "SM-MESSAGE ";
356
357 let stdout = BufReader::new(stdout);
358 let mut lines = stdout.lines();
359 while let Some(line) = lines.next().await {
360 let Some(line) = line.context("failed to read line from stdout").log_err() else {
361 continue;
362 };
363 let Some(line) = line.strip_prefix(MESSAGE_PREFIX) else {
364 continue;
365 };
366 let Some(message) = serde_json::from_str::<SupermavenMessage>(line)
367 .with_context(|| format!("failed to deserialize line from stdout: {:?}", line))
368 .log_err()
369 else {
370 continue;
371 };
372
373 this.update(cx, |this, _cx| {
374 if let Supermaven::Spawned(this) = this {
375 this.handle_message(message);
376 }
377 Task::ready(anyhow::Ok(()))
378 })?
379 .await?;
380 }
381
382 Ok(())
383 }
384
385 fn handle_message(&mut self, message: SupermavenMessage) {
386 match message {
387 SupermavenMessage::ActivationRequest(request) => {
388 self.account_status = match request.activate_url {
389 Some(activate_url) => AccountStatus::NeedsActivation {
390 activate_url: activate_url.clone(),
391 },
392 None => AccountStatus::Ready,
393 };
394 }
395 SupermavenMessage::ActivationSuccess => {
396 self.account_status = AccountStatus::Ready;
397 }
398 SupermavenMessage::ServiceTier { service_tier } => {
399 self.account_status = AccountStatus::Ready;
400 self.service_tier = Some(service_tier);
401 }
402 SupermavenMessage::Response(response) => {
403 let state_id = SupermavenCompletionStateId(response.state_id.parse().unwrap());
404 if let Some(state) = self.states.get_mut(&state_id) {
405 for item in &response.items {
406 match item {
407 ResponseItem::Text { text } => state.text.push_str(text),
408 ResponseItem::Dedent { text } => state.dedent.push_str(text),
409 _ => {}
410 }
411 }
412 *state.updates_tx.borrow_mut() = ();
413 }
414 }
415 SupermavenMessage::Passthrough { passthrough } => self.handle_message(*passthrough),
416 _ => {
417 log::warn!("unhandled message: {:?}", message);
418 }
419 }
420 }
421}
422
423#[derive(Copy, Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
424pub struct SupermavenCompletionStateId(usize);
425
426#[allow(dead_code)]
427pub struct SupermavenCompletionState {
428 buffer_id: EntityId,
429 prefix_anchor: Anchor,
430 // prefix_offset is tracked independently because the anchor biases left which
431 // doesn't allow us to determine if the prior text has been deleted.
432 prefix_offset: usize,
433 text: String,
434 dedent: String,
435 updates_tx: watch::Sender<()>,
436}
437
438pub struct SupermavenCompletion {
439 pub id: SupermavenCompletionStateId,
440 pub updates: watch::Receiver<()>,
441}