diff --git a/crates/repl/src/jupyter_completion_provider.rs b/crates/repl/src/jupyter_completion_provider.rs new file mode 100644 index 00000000000..10bdcaf6aff --- /dev/null +++ b/crates/repl/src/jupyter_completion_provider.rs @@ -0,0 +1,392 @@ +use std::cell::RefCell; +use std::ops::Range; +use std::rc::Rc; +use std::time::Duration; + +use anyhow::Result; +use editor::{CompletionContext, CompletionProvider, Editor}; +use gpui::{Context, Entity, Task, WeakEntity, Window}; +use language::{self, CodeLabel, ToOffset, ToPoint}; +use project::{ + Completion, CompletionDisplayOptions, CompletionResponse, CompletionSource, Project, + lsp_store::CompletionDocumentation, +}; +use runtimelib::media::MediaType; + +use crate::repl_editor::{CompletionChunk, completion_chunk}; +use crate::session::Session; + +pub(crate) struct JupyterCompletionProvider { + project: Entity, + session: WeakEntity, +} + +impl JupyterCompletionProvider { + pub fn new(project: Entity, session: WeakEntity) -> Self { + Self { project, session } + } +} + +impl CompletionProvider for JupyterCompletionProvider { + fn completions( + &self, + buffer: &Entity, + buffer_position: language::Anchor, + trigger: CompletionContext, + window: &mut Window, + cx: &mut Context, + ) -> Task>> { + let lsp_completions = + self.project + .completions(buffer, buffer_position, trigger, window, cx); + + let snapshot = buffer.read(cx).snapshot(); + let cursor_point = buffer_position.to_point(&snapshot); + let chunk = completion_chunk(&snapshot, cursor_point, cx); + + let jupyter_receiver = chunk.as_ref().and_then(|chunk| { + let session = self.session.upgrade()?; + session.update(cx, |session, cx| { + session.request_completions(chunk.code.clone(), chunk.cursor_pos, cx) + }) + }); + + let timer = cx.background_executor().timer(Duration::from_secs(5)); + cx.spawn(async move |_, _cx| { + let mut responses = lsp_completions.await.unwrap_or_default(); + + if let (Some(rx), Some(chunk)) = (jupyter_receiver, chunk) { + let reply = match futures::future::select(std::pin::pin!(rx), std::pin::pin!(timer)) + .await + { + futures::future::Either::Left((Ok(reply), _)) => Some(reply), + _ => None, + }; + + if let Some(reply) = reply + && reply.status == runtimelib::ReplyStatus::Ok + { + if let Some(response) = jupyter_reply_to_completion_response( + &reply, + &chunk, + &snapshot, + buffer_position, + ) { + responses.push(response); + } + } + } + + Ok(responses) + }) + } + + fn resolve_completions( + &self, + buffer: Entity, + completion_indices: Vec, + completions: Rc>>, + cx: &mut Context, + ) -> Task> { + let mut lsp_indices = Vec::new(); + let mut jupyter_indices = Vec::new(); + + for &index in &completion_indices { + match &completions.borrow()[index].source { + CompletionSource::Custom => jupyter_indices.push(index), + _ => lsp_indices.push(index), + } + } + + let lsp_task = if lsp_indices.is_empty() { + Task::ready(Ok(false)) + } else { + self.project + .resolve_completions(buffer, lsp_indices, completions.clone(), cx) + }; + + let jupyter_receivers: Vec<_> = jupyter_indices + .iter() + .filter_map(|&index| { + let session = self.session.upgrade()?; + let new_text = completions.borrow()[index].new_text.clone(); + let cursor_pos = new_text.chars().count(); + let rx = session.update(cx, |session, cx| { + session.request_inspect(new_text, cursor_pos, cx) + })?; + Some((index, rx)) + }) + .collect(); + + let executor = cx.background_executor().clone(); + cx.spawn(async move |_, _cx| { + let lsp_resolved = lsp_task.await.unwrap_or(false); + + let mut jupyter_resolved = false; + for (index, rx) in jupyter_receivers { + let timeout = executor.timer(Duration::from_secs(5)); + let reply = match futures::future::select( + std::pin::pin!(rx), + std::pin::pin!(timeout), + ) + .await + { + futures::future::Either::Left((Ok(reply), _)) => Some(reply), + _ => None, + }; + + if let Some(reply) = reply + && reply.found + && reply.status == runtimelib::ReplyStatus::Ok + { + if let Some(documentation) = inspect_reply_to_documentation(&reply) { + completions.borrow_mut()[index].documentation = Some(documentation); + jupyter_resolved = true; + } + } + } + + Ok(lsp_resolved || jupyter_resolved) + }) + } + + fn apply_additional_edits_for_completion( + &self, + buffer: Entity, + completions: Rc>>, + completion_index: usize, + push_to_history: bool, + all_commit_ranges: Vec>, + cx: &mut Context, + ) -> Task>> { + self.project.apply_additional_edits_for_completion( + buffer, + completions, + completion_index, + push_to_history, + all_commit_ranges, + cx, + ) + } + + fn is_completion_trigger( + &self, + buffer: &Entity, + position: language::Anchor, + text: &str, + trigger_in_words: bool, + cx: &mut Context, + ) -> bool { + self.project + .is_completion_trigger(buffer, position, text, trigger_in_words, cx) + } + + fn sort_completions(&self) -> bool { + true + } + + fn filter_completions(&self) -> bool { + true + } + + fn show_snippets(&self) -> bool { + true + } +} + +fn char_offset_to_byte_offset(text: &str, char_offset: usize) -> usize { + text.char_indices() + .nth(char_offset) + .map(|(byte_idx, _)| byte_idx) + .unwrap_or(text.len()) +} + +fn jupyter_reply_to_completion_response( + reply: &runtimelib::CompleteReply, + chunk: &CompletionChunk, + snapshot: &language::BufferSnapshot, + buffer_position: language::Anchor, +) -> Option { + if reply.matches.is_empty() { + return None; + } + + let start_byte = chunk.start_byte + char_offset_to_byte_offset(&chunk.code, reply.cursor_start); + let end_byte = chunk.start_byte + char_offset_to_byte_offset(&chunk.code, reply.cursor_end); + + let replace_start = snapshot.anchor_after(start_byte); + let replace_end = snapshot.anchor_before(end_byte); + let replace_range = replace_start..replace_end; + + let byte_offset = buffer_position.to_offset(snapshot); + let match_start_byte = start_byte.min(byte_offset); + let match_start = snapshot.anchor_after(match_start_byte); + + let completions = reply + .matches + .iter() + .map(|match_text| Completion { + replace_range: replace_range.clone(), + new_text: match_text.clone(), + label: CodeLabel::filtered( + format!("{match_text} Jupyter"), + match_text.len(), + None, + vec![], + ), + documentation: None, + source: CompletionSource::Custom, + icon_path: None, + match_start: Some(match_start), + snippet_deduplication_key: None, + insert_text_mode: None, + confirm: None, + group: None, + }) + .collect(); + + Some(CompletionResponse { + completions, + display_options: CompletionDisplayOptions::default(), + is_incomplete: false, + }) +} + +fn inspect_reply_to_documentation( + reply: &runtimelib::InspectReply, +) -> Option { + for media_type in &reply.data.content { + if let MediaType::Markdown(text) = media_type { + let trimmed = text.trim(); + if !trimmed.is_empty() { + return Some(CompletionDocumentation::MultiLineMarkdown( + trimmed.to_string().into(), + )); + } + } + } + + for media_type in &reply.data.content { + if let MediaType::Plain(text) = media_type { + let trimmed = text.trim(); + if !trimmed.is_empty() { + return Some(CompletionDocumentation::MultiLinePlainText( + trimmed.to_string().into(), + )); + } + } + } + + None +} + +#[cfg(test)] +mod tests { + use super::*; + use gpui::{App, AppContext as _}; + use language::Buffer; + + fn make_chunk(code: &str, start_byte: usize) -> CompletionChunk { + CompletionChunk { + code: code.into(), + cursor_pos: code.chars().count(), + start_byte, + } + } + + #[gpui::test] + fn test_jupyter_reply_empty_matches(cx: &mut App) { + let buffer = cx.new(|cx| Buffer::local("foo bar", cx)); + let snapshot = buffer.read(cx).snapshot(); + let buffer_position = snapshot.anchor_after(4); + let chunk = make_chunk("foo bar", 0); + let reply = runtimelib::CompleteReply::default(); + + assert!( + jupyter_reply_to_completion_response(&reply, &chunk, &snapshot, buffer_position) + .is_none() + ); + } + + #[gpui::test] + fn test_jupyter_reply_translates_offsets(cx: &mut App) { + // Buffer matches the chunk one-to-one; chunk starts at byte 0. + let text = "import numpy as np\nnp.ar"; + let buffer = cx.new(|cx| Buffer::local(text, cx)); + let snapshot = buffer.read(cx).snapshot(); + let buffer_position = snapshot.anchor_after(text.len()); + let chunk = make_chunk(text, 0); + let reply = runtimelib::CompleteReply { + matches: vec!["np.array".into(), "np.arange".into()], + cursor_start: 19, // start of "np.ar" (after "import numpy as np\n") + cursor_end: 24, // end of "np.ar" + ..Default::default() + }; + + let response = + jupyter_reply_to_completion_response(&reply, &chunk, &snapshot, buffer_position) + .unwrap(); + assert_eq!(response.completions.len(), 2); + let completion = &response.completions[0]; + assert_eq!(completion.replace_range.start.to_offset(&snapshot), 19); + assert_eq!(completion.replace_range.end.to_offset(&snapshot), 24); + assert_eq!(completion.new_text, "np.array"); + // In the common case, `match_start` equals the replace-range start. + assert_eq!(completion.match_start.unwrap().to_offset(&snapshot), 19); + } + + #[gpui::test] + fn test_jupyter_reply_chunk_offset_into_buffer(cx: &mut App) { + // The chunk lives at byte 12 of the buffer; replace_range must be + // shifted by chunk.start_byte to land in the right buffer span. + let prelude = "x = 1\ny = 2\n"; // 12 bytes + let chunk_text = "z.fo"; + let mut text = String::from(prelude); + text.push_str(chunk_text); + let buffer = cx.new({ + let text = text.clone(); + |cx| Buffer::local(text, cx) + }); + let snapshot = buffer.read(cx).snapshot(); + let buffer_position = snapshot.anchor_after(text.len()); + let chunk = make_chunk(chunk_text, 12); + let reply = runtimelib::CompleteReply { + matches: vec!["z.foo".into()], + cursor_start: 0, + cursor_end: 4, + ..Default::default() + }; + + let response = + jupyter_reply_to_completion_response(&reply, &chunk, &snapshot, buffer_position) + .unwrap(); + let completion = &response.completions[0]; + assert_eq!(completion.replace_range.start.to_offset(&snapshot), 12); + assert_eq!(completion.replace_range.end.to_offset(&snapshot), 16); + } + + #[gpui::test] + fn test_jupyter_reply_multibyte_offsets(cx: &mut App) { + // "π" is 2 bytes / 1 char, so char-offset and byte-offset diverge. + let text = "x = π."; + let buffer = cx.new(|cx| Buffer::local(text, cx)); + let snapshot = buffer.read(cx).snapshot(); + let buffer_position = snapshot.anchor_after(text.len()); + let chunk = make_chunk(text, 0); + let reply = runtimelib::CompleteReply { + matches: vec!["π.bit_length".into()], + cursor_start: 4, // char offset of "π" + cursor_end: 6, // char offset after "." + ..Default::default() + }; + + let response = + jupyter_reply_to_completion_response(&reply, &chunk, &snapshot, buffer_position) + .unwrap(); + let completion = &response.completions[0]; + // char offset 4 → byte 4 ("x = " = 4 bytes) + // char offset 6 → byte 7 ("x = π." = 4 + 2 + 1 = 7 bytes) + assert_eq!(completion.replace_range.start.to_offset(&snapshot), 4); + assert_eq!(completion.replace_range.end.to_offset(&snapshot), 7); + } +} diff --git a/crates/repl/src/repl.rs b/crates/repl/src/repl.rs index 8c3d15a2ad2..89fd88344c4 100644 --- a/crates/repl/src/repl.rs +++ b/crates/repl/src/repl.rs @@ -1,4 +1,5 @@ pub mod components; +mod jupyter_completion_provider; mod jupyter_settings; pub mod kernels; pub mod notebook; diff --git a/crates/repl/src/repl_editor.rs b/crates/repl/src/repl_editor.rs index 61bed513a16..7b2c48350a1 100644 --- a/crates/repl/src/repl_editor.rs +++ b/crates/repl/src/repl_editor.rs @@ -1,15 +1,17 @@ //! REPL operations on an [`Editor`]. use std::ops::Range; +use std::rc::Rc; use std::sync::Arc; use anyhow::{Context as _, Result}; use editor::{Editor, MultiBufferOffset}; use gpui::{App, Entity, WeakEntity, Window, prelude::*}; -use language::{BufferSnapshot, Language, LanguageName, Point}; +use language::{BufferSnapshot, Language, LanguageName, Point, ToOffset as _}; use project::{ProjectItem as _, WorktreeId}; use workspace::{Workspace, notifications::NotificationId}; +use crate::jupyter_completion_provider::JupyterCompletionProvider; use crate::kernels::PythonEnvKernelSpecification; use crate::repl_store::ReplStore; use crate::session::SessionEvent; @@ -51,15 +53,24 @@ pub fn assign_kernelspec( cx.new(|cx| Session::new(weak_editor.clone(), fs, kernel_specification, window, cx)); weak_editor - .update(cx, |_editor, cx| { + .update(cx, |editor, cx| { + if let Some(project) = editor.project().cloned() { + let provider = + Rc::new(JupyterCompletionProvider::new(project, session.downgrade())); + editor.set_completion_provider(Some(provider)); + } + cx.notify(); cx.subscribe(&session, { let store = store.clone(); - move |_this, _session, event, cx| match event { - SessionEvent::Shutdown(shutdown_event) => { + move |this, _session, event, cx| match event { + SessionEvent::Shutdown(shutdown_editor) => { + if let Some(project) = this.project().cloned() { + this.set_completion_provider(Some(Rc::new(project))); + } store.update(cx, |store, _cx| { - store.remove_session(shutdown_event.entity_id()); + store.remove_session(shutdown_editor.entity_id()); }); } } @@ -253,15 +264,24 @@ pub fn run( let session = cx.new(|cx| Session::new(weak_editor, fs, kernel_specification, window, cx)); - editor.update(cx, |_editor, cx| { + editor.update(cx, |editor, cx| { + if let Some(project) = editor.project().cloned() { + let provider = + Rc::new(JupyterCompletionProvider::new(project, session.downgrade())); + editor.set_completion_provider(Some(provider)); + } + cx.notify(); cx.subscribe(&session, { let store = store.clone(); - move |_this, _session, event, cx| match event { - SessionEvent::Shutdown(shutdown_event) => { + move |this, _session, event, cx| match event { + SessionEvent::Shutdown(shutdown_editor) => { + if let Some(project) = this.project().cloned() { + this.set_completion_provider(Some(Rc::new(project))); + } store.update(cx, |store, _cx| { - store.remove_session(shutdown_event.entity_id()); + store.remove_session(shutdown_editor.entity_id()); }); } } @@ -503,6 +523,42 @@ pub fn setup_editor_session_actions(editor: &mut Editor, editor_handle: WeakEnti .detach(); } +pub(crate) struct CompletionChunk { + pub code: String, + /// Char offset of the cursor within `code`. + pub cursor_pos: usize, + /// Byte offset within the buffer where `code` starts. + pub start_byte: usize, +} + +/// Returns the runnable chunk surrounding `cursor` along with the cursor's char +/// offset within that chunk. The chunk is the same unit that would be sent to +/// the kernel if the user invoked `run` at this position, so completion and +/// inspection requests share the kernel's mental model of "the current cell". +pub(crate) fn completion_chunk( + buffer: &BufferSnapshot, + cursor: Point, + cx: &mut App, +) -> Option { + let (ranges, _) = runnable_ranges(buffer, cursor..cursor, cx); + // `runnable_ranges` may skip forward to the next cell when the cursor is on + // a blank line; that range is not useful as completion context, so only + // accept a chunk that actually contains the cursor. + let range = ranges + .into_iter() + .find(|r| r.start <= cursor && cursor <= r.end)?; + let start_byte = range.start.to_offset(buffer); + let cursor_byte = cursor.to_offset(buffer); + let code: String = buffer.text_for_range(range).collect(); + let cursor_byte_in_chunk = cursor_byte.saturating_sub(start_byte).min(code.len()); + let cursor_pos = code[..cursor_byte_in_chunk].chars().count(); + Some(CompletionChunk { + code, + cursor_pos, + start_byte, + }) +} + fn cell_range(buffer: &BufferSnapshot, start_row: u32, end_row: u32) -> Range { let mut snippet_end_row = end_row; while buffer.is_line_blank(snippet_end_row) && snippet_end_row > start_row { @@ -1081,4 +1137,131 @@ mod tests { let (snippets, _) = runnable_ranges(&snapshot, Point::new(1, 0)..Point::new(1, 0), cx); assert!(snippets.is_empty()); } + + #[gpui::test] + fn test_completion_chunk_basic(cx: &mut App) { + let test_language = Arc::new(Language::new( + LanguageConfig { + name: "TestLang".into(), + line_comments: vec!["# ".into()], + ..Default::default() + }, + None, + )); + + let buffer = cx.new(|cx| { + Buffer::local( + indoc! { r#" + print(1 + 1) + print(2 + 2) + "# }, + cx, + ) + .with_language(test_language, cx) + }); + let snapshot = buffer.read(cx).snapshot(); + + let chunk = completion_chunk(&snapshot, Point::new(0, 6), cx).unwrap(); + assert_eq!(chunk.code, "print(1 + 1)"); + assert_eq!(chunk.cursor_pos, 6); + assert_eq!(chunk.start_byte, 0); + + // Second line: chunk starts after "print(1 + 1)\n" (13 bytes). + let chunk = completion_chunk(&snapshot, Point::new(1, 4), cx).unwrap(); + assert_eq!(chunk.code, "print(2 + 2)"); + assert_eq!(chunk.cursor_pos, 4); + assert_eq!(chunk.start_byte, 13); + } + + #[gpui::test] + fn test_completion_chunk_jupytext(cx: &mut App) { + let test_language = Arc::new(Language::new( + LanguageConfig { + name: "TestLang".into(), + line_comments: vec!["# ".into()], + ..Default::default() + }, + None, + )); + + let buffer = cx.new(|cx| { + Buffer::local( + indoc! { r#" + # %% + x = 1 + y = 2 + # %% + z = 3 + "# }, + cx, + ) + .with_language(test_language, cx) + }); + let snapshot = buffer.read(cx).snapshot(); + + // Cursor inside the first cell. + let chunk = completion_chunk(&snapshot, Point::new(2, 4), cx).unwrap(); + assert_eq!(chunk.code, "# %%\nx = 1\ny = 2"); + assert_eq!(chunk.cursor_pos, 15); + assert_eq!(chunk.start_byte, 0); + + // Cursor inside the second cell; chunk starts at row 3 (byte 17). + let chunk = completion_chunk(&snapshot, Point::new(4, 2), cx).unwrap(); + assert_eq!(chunk.code, "# %%\nz = 3"); + assert_eq!(chunk.cursor_pos, 7); + assert_eq!(chunk.start_byte, 17); + } + + #[gpui::test] + fn test_completion_chunk_blank_line_returns_none(cx: &mut App) { + let test_language = Arc::new(Language::new( + LanguageConfig { + name: "TestLang".into(), + line_comments: vec!["# ".into()], + ..Default::default() + }, + None, + )); + + let buffer = cx.new(|cx| { + Buffer::local( + indoc! { r#" + print(1 + 1) + + print(2 + 2) + "# }, + cx, + ) + .with_language(test_language, cx) + }); + let snapshot = buffer.read(cx).snapshot(); + + // `runnable_ranges` skips forward to the next non-blank cell, but the + // returned range does not contain the cursor on row 1, so + // `completion_chunk` rejects it. + assert!(completion_chunk(&snapshot, Point::new(1, 0), cx).is_none()); + } + + #[gpui::test] + fn test_completion_chunk_multibyte(cx: &mut App) { + let test_language = Arc::new(Language::new( + LanguageConfig { + name: "TestLang".into(), + line_comments: vec!["# ".into()], + ..Default::default() + }, + None, + )); + + // "π" is 2 bytes, so byte and char counts diverge inside the chunk. + let buffer = cx.new(|cx| Buffer::local("x = π * 2", cx).with_language(test_language, cx)); + let snapshot = buffer.read(cx).snapshot(); + + // Byte column 6 = right after "π" (which occupies bytes 4..6). + let chunk = completion_chunk(&snapshot, Point::new(0, 6), cx).unwrap(); + assert_eq!(chunk.code, "x = π * 2"); + // chunk.code[..6] = "x = π" → 5 chars + assert_eq!(chunk.cursor_pos, 5); + assert_eq!(chunk.start_byte, 0); + } } diff --git a/crates/repl/src/session.rs b/crates/repl/src/session.rs index 9b7bd759504..6aa77843b74 100644 --- a/crates/repl/src/session.rs +++ b/crates/repl/src/session.rs @@ -35,8 +35,8 @@ use gpui::{ use language::Point; use project::Fs; use runtimelib::{ - ExecuteRequest, ExecutionState, InputReply, InterruptRequest, JupyterMessage, - JupyterMessageContent, KernelInfoRequest, ReplyStatus, ShutdownRequest, + CompleteRequest, ExecuteRequest, ExecutionState, InputReply, InspectRequest, InterruptRequest, + JupyterMessage, JupyterMessageContent, KernelInfoRequest, ReplyStatus, ShutdownRequest, }; use settings::Settings as _; use std::{env::temp_dir, ops::Range, sync::Arc, time::Duration}; @@ -53,6 +53,10 @@ pub struct Session { blocks: HashMap, result_inlays: HashMap, usize)>, next_inlay_id: usize, + pending_completion_replies: + HashMap>, + pending_inspect_replies: + HashMap>, _subscriptions: Vec, } @@ -257,6 +261,8 @@ impl Session { blocks: HashMap::default(), result_inlays: HashMap::default(), next_inlay_id: 0, + pending_completion_replies: HashMap::default(), + pending_inspect_replies: HashMap::default(), kernel_specification, _subscriptions: vec![subscription], }; @@ -486,6 +492,54 @@ impl Session { anyhow::Ok(()) } + pub fn request_completions( + &mut self, + code: String, + cursor_pos: usize, + cx: &mut Context, + ) -> Option> { + if !matches!(&self.kernel, Kernel::RunningKernel(_)) { + return None; + } + + self.pending_completion_replies.clear(); + + let request = CompleteRequest { code, cursor_pos }; + let message: JupyterMessage = JupyterMessageContent::CompleteRequest(request).into(); + let msg_id = message.header.msg_id.clone(); + + let (tx, rx) = futures::channel::oneshot::channel(); + self.pending_completion_replies.insert(msg_id, tx); + self.send(message, cx).log_err(); + + Some(rx) + } + + pub fn request_inspect( + &mut self, + code: String, + cursor_pos: usize, + cx: &mut Context, + ) -> Option> { + if !matches!(&self.kernel, Kernel::RunningKernel(_)) { + return None; + } + + let request = InspectRequest { + code, + cursor_pos, + detail_level: Some(0), + }; + let message: JupyterMessage = JupyterMessageContent::InspectRequest(request).into(); + let msg_id = message.header.msg_id.clone(); + + let (tx, rx) = futures::channel::oneshot::channel(); + self.pending_inspect_replies.insert(msg_id, tx); + self.send(message, cx).log_err(); + + Some(rx) + } + fn send_stdin_reply( &mut self, value: String, @@ -999,6 +1053,18 @@ impl KernelSession for Session { self.kernel.set_kernel_info(reply); cx.notify(); } + JupyterMessageContent::CompleteReply(reply) => { + if let Some(tx) = self.pending_completion_replies.remove(parent_message_id) { + tx.send(reply.clone()).ok(); + } + return; + } + JupyterMessageContent::InspectReply(reply) => { + if let Some(tx) = self.pending_inspect_replies.remove(parent_message_id) { + tx.send(reply.clone()).ok(); + } + return; + } JupyterMessageContent::UpdateDisplayData(update) => { let display_id = if let Some(display_id) = update.transient.display_id.clone() { display_id