diff --git a/Cargo.lock b/Cargo.lock index ed0c4491a..b572dc22b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -327,6 +327,19 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "asynchronous-codec" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a860072022177f903e59730004fb5dc13db9275b79bb2aef7ba8ce831956c233" +dependencies = [ + "bytes", + "futures-sink", + "futures-util", + "memchr", + "pin-project-lite", +] + [[package]] name = "atk" version = "0.18.2" @@ -4199,6 +4212,7 @@ dependencies = [ "ide-jsonrpc", "ide-language", "ide-lsp", + "ide-notebook", "ide-remote", "ide-run", "ide-sdk", @@ -4216,10 +4230,12 @@ dependencies = [ "serde", "serde_json", "specta", + "specta-typescript", "tempfile", "tokio", "tracing", "tracing-subscriber", + "uuid", ] [[package]] @@ -4355,6 +4371,29 @@ dependencies = [ "which", ] +[[package]] +name = "ide-notebook" +version = "0.1.25" +dependencies = [ + "futures", + "ide-core", + "ide-vfs", + "jupyter-protocol", + "jupyter-zmq-client", + "libc", + "parking_lot", + "serde", + "serde_json", + "specta", + "specta-typescript", + "tempfile", + "tokio", + "tracing", + "uuid", + "windows-sys 0.61.2", + "zeromq", +] + [[package]] name = "ide-plugin" version = "0.1.25" @@ -5033,6 +5072,41 @@ dependencies = [ "serde_json", ] +[[package]] +name = "jupyter-protocol" +version = "2.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf199c30e61505042f5c0d79a33b5a71c552ea8b64243ff3498b1f4592d8e692" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures", + "serde", + "serde_json", + "thiserror 2.0.21", + "uuid", +] + +[[package]] +name = "jupyter-zmq-client" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f59bf6241faf6a8c06416c351756be9a79976e8faad220604b9a26f568270ab0" +dependencies = [ + "bytes", + "data-encoding", + "dirs 6.0.0", + "jupyter-protocol", + "ring", + "serde", + "serde_json", + "thiserror 2.0.21", + "tokio", + "uuid", + "zeromq", +] + [[package]] name = "keccak" version = "0.2.2" @@ -6061,6 +6135,7 @@ dependencies = [ "ide-jsonrpc", "ide-language", "ide-lsp", + "ide-notebook", "ide-plugin", "ide-profiler", "ide-remote", @@ -7383,6 +7458,12 @@ dependencies = [ "wait-timeout", ] +[[package]] +name = "saa" +version = "5.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd85992a680de788217d9c70b74aa7685b108cf1aba02f728fb27f1b27199ad1" + [[package]] name = "salsa20" version = "0.11.0" @@ -7402,6 +7483,16 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scc" +version = "3.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d5f8ca67c4f6208caf824037a6a9df0211b4e0937fccc4f543441f294d04511" +dependencies = [ + "saa", + "sdd", +] + [[package]] name = "schannel" version = "0.1.29" @@ -7490,6 +7581,15 @@ dependencies = [ "untrusted", ] +[[package]] +name = "sdd" +version = "4.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31877eee03686b905845993b5d9809a220c9dd5a10e9d6021d7ca67384c9070" +dependencies = [ + "saa", +] + [[package]] name = "sec1" version = "0.8.1" @@ -7783,6 +7883,12 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sha1dc" version = "0.1.5" @@ -9067,7 +9173,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a1446cb4198848d1562301a3340424b4f425ef79f35ef9ee034769a9dd92c10d" dependencies = [ "async-trait", - "asynchronous-codec", + "asynchronous-codec 0.6.2", "byteorder", "bytes", "connection-string", @@ -9146,9 +9252,9 @@ checksum = "fd3ca314f692efd6c868f8408f53fe444634a845f96c028b97d35f6a1f79f0ee" [[package]] name = "tokio" -version = "1.54.0" +version = "1.54.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0c168a474b95bfecfe908a65d4bdffc523144b649a74c09c286afb35d53e0bc" +checksum = "b52fd3255258d26fac7264ab5a76b53c13367136dc011d3d6ad4f8ed070396a3" dependencies = [ "bytes", "libc", @@ -10046,6 +10152,7 @@ dependencies = [ "getrandom 0.4.3", "js-sys", "serde_core", + "sha1_smol", "wasm-bindgen", ] @@ -10651,6 +10758,17 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "626c4bac6755d76ffc12cb01b2eac751db1996b9e0041de9aa02c8c211ddc82c" +[[package]] +name = "win_uds" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd30a1a28a3799479cbf4e17284a220ea9ff6bad098a9d0224543a5d1efe1da" +dependencies = [ + "async-io", + "futures-io", + "socket2", +] + [[package]] name = "winapi" version = "0.3.9" @@ -11389,6 +11507,31 @@ version = "1.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e13084392c5e4bc371903e2935a5eaeed24905a7511356b883835e18a78f6879" +[[package]] +name = "zeromq" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efb2c254fd8f366755335c9e43b865f8484fe3bd717d65ffe7c3f28852863030" +dependencies = [ + "async-trait", + "asynchronous-codec 0.7.0", + "bytes", + "crossbeam-queue", + "futures", + "log", + "num-traits", + "once_cell", + "parking_lot", + "rand 0.9.5", + "regex", + "scc", + "thiserror 1.0.69", + "tokio", + "tokio-util", + "uuid", + "win_uds", +] + [[package]] name = "zerotrie" version = "0.2.5" diff --git a/Cargo.toml b/Cargo.toml index 2dbda21a6..fa4d3c73a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,7 @@ ide-fonts = { path = "crates/ide-fonts" } ide-ssh = { path = "crates/ide-ssh" } ide-remote = { path = "crates/ide-remote" } ide-host = { path = "crates/ide-host" } +ide-notebook = { path = "crates/ide-notebook" } # serialization & IPC typing serde = { version = "1", features = ["derive"] } @@ -85,6 +86,9 @@ redb = "4" tree-sitter = "0.25" # protocols & integrations +jupyter-zmq-client = { version = "=2.0.0", default-features = false, features = ["tokio-runtime", "ring"] } +jupyter-protocol = "=2.0.3" +zeromq = { version = "0.6", default-features = false, features = ["tokio-runtime", "tcp-transport"] } lsp-types = "0.97" gix = { version = "0.89", default-features = false } portable-pty = "0.9" diff --git a/README.md b/README.md index af53c1de2..7200f2339 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,15 @@ pnpm tauri dev Rust 1.96, Node 22, pnpm and the [Tauri prerequisites](https://v2.tauri.app/start/prerequisites/). +Native notebook execution tests are opt-in and need two distinct Python environments +with `ipykernel` installed: + +```sh +OPENIDEA_TEST_PYTHON=/first/env/bin/python \ +OPENIDEA_TEST_PYTHON_SECOND=/second/env/bin/python \ +cargo test -p ide-notebook --test kernel_runtime -- --ignored --test-threads=1 +``` + ## License [Apache 2.0](LICENSE). Bundled third-party assets are listed in [`THIRD_PARTY_LICENSES.md`](THIRD_PARTY_LICENSES.md). diff --git a/crates/ide-core/src/event.rs b/crates/ide-core/src/event.rs index 44231f149..b730a9e13 100644 --- a/crates/ide-core/src/event.rs +++ b/crates/ide-core/src/event.rs @@ -117,6 +117,15 @@ pub enum IdeEvent { AskpassDone { id: u64, }, + /// The kernel of an open notebook changed state (`starting`, `idle`, `busy`, + /// `dead`); `message` says why it failed to start or died. + NotebookKernel { + /// Backend document incarnation; stale queued events cannot target a reopen. + id: String, + path: crate::IpcPath, + state: String, + message: Option, + }, } /// What an askpass prompt asks for. diff --git a/crates/ide-host/Cargo.toml b/crates/ide-host/Cargo.toml index 43df9351e..adc1344ca 100644 --- a/crates/ide-host/Cargo.toml +++ b/crates/ide-host/Cargo.toml @@ -34,17 +34,20 @@ ide-ssh.workspace = true ide-remote.workspace = true ide-jsonrpc.workspace = true ide-agent.workspace = true +ide-notebook.workspace = true ide-tools = { path = "../ide-tools" } serde.workspace = true serde_json.workspace = true specta.workspace = true +specta-typescript.workspace = true tokio = { workspace = true, features = ["io-std"] } tracing.workspace = true tracing-subscriber.workspace = true parking_lot.workspace = true dunce.workspace = true dirs.workspace = true +uuid.workspace = true [dev-dependencies] -tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-util"] } +tokio = { workspace = true, features = ["rt-multi-thread", "macros", "io-util", "test-util"] } tempfile.workspace = true diff --git a/crates/ide-host/src/lib.rs b/crates/ide-host/src/lib.rs index c341b6a9e..294d8a98c 100644 --- a/crates/ide-host/src/lib.rs +++ b/crates/ide-host/src/lib.rs @@ -4,6 +4,7 @@ //! runs one on a remote machine and serves it over JSON-RPC ([`rpc`]). pub mod environment; +pub mod notebook; pub mod ops; mod project; mod proxy; @@ -132,26 +133,75 @@ impl Host { } pub fn close_project(&self, window: &str) { - if let Some(project) = self.projects.write().remove(window) { + let project = self.projects.write().remove(window); + if let Some(project) = project { if !project.light { self.environment.project_closed(); } project.agents.shutdown(); + project.notebooks.stop_admission(); + self.shutdown_notebooks_blocking(vec![project.notebooks.clone()]); let lsp = project.lsp.clone(); self.runtime.spawn(async move { lsp.shutdown_all().await }); } } - /// Stops what this host started: language servers, agents, runs and terminals. + /// Synchronous project/window destruction must finish owned kernel cleanup, + /// not enqueue work which can be lost as the desktop runtime exits. + fn shutdown_notebooks_blocking(&self, notebooks: Vec>) { + for service in ¬ebooks { + service.stop_admission(); + } + let runtime = self.runtime.clone(); + let wait = move || { + let worker = std::thread::spawn(move || { + runtime.block_on(async move { + for service in notebooks { + service.shutdown().await; + } + }); + }); + if let Err(error) = worker.join() { + tracing::error!("notebook shutdown thread failed: {error:?}"); + } + }; + if tokio::runtime::Handle::try_current() + .is_ok_and(|handle| handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread) + { + tokio::task::block_in_place(wait); + } else { + wait(); + } + } + + /// Called before the desktop runtime exits, including update/restart exits. + pub fn shutdown_notebooks_on_exit(&self) { + self.shutdown_notebooks_blocking(self.projects().iter().map(|project| project.notebooks.clone()).collect()); + } + + /// Stops what this host started: notebook kernels, language servers, agents, runs and terminals. /// For a process that exits without dropping the host (the remote server). pub async fn shutdown(&self) { self.environment.shutdown(); self.processes.kill_all(); self.terminals.kill_all(); let projects: Vec<_> = self.projects.write().drain().map(|(_, p)| p).collect(); - for project in projects { + for project in &projects { + project.notebooks.stop_admission(); project.agents.shutdown(); - project.lsp.shutdown_all().await; + } + // Kernel cleanup owns its interrupt/termination/reaping deadlines. An + // outer LSP timeout must never abort it before owned children are reaped. + for project in &projects { + project.notebooks.shutdown().await; + } + let stop_language_servers = async { + for project in &projects { + project.lsp.shutdown_all().await; + } + }; + if tokio::time::timeout(std::time::Duration::from_secs(10), stop_language_servers).await.is_err() { + tracing::warn!("language servers did not shut down in time"); } } diff --git a/crates/ide-host/src/notebook.rs b/crates/ide-host/src/notebook.rs new file mode 100644 index 000000000..63e17d04a --- /dev/null +++ b/crates/ide-host/src/notebook.rs @@ -0,0 +1,796 @@ +//! Project-incarnation notebook ownership and whole-batch execution admission. +use crate::ops::Channel; +use ide_core::error::ErrorKind; +use ide_core::{CancellationToken, IdeError, IdeEvent, IdeResult, SharedSink}; +use ide_notebook::kernel::{KernelManager, KernelOutput}; +use ide_notebook::{NotebookDocument, NotebookEdit, NotebookSnapshot}; +use parking_lot::Mutex; +use serde::Serialize; +use specta::Type; +use std::collections::{HashMap, HashSet}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::{Mutex as AsyncMutex, watch}; + +#[derive(Clone, Debug, Serialize, Type)] +#[serde(rename_all = "camelCase")] +pub struct NotebookOutputEvent { + pub id: String, + pub path: String, + pub cell: String, + pub msg_id: String, + pub kind: String, + #[specta(type = specta_typescript::Unknown)] + pub content: serde_json::Value, + pub version: u32, + pub dirty: bool, + pub changed_cells: Vec, +} + +#[derive(Default)] +struct Registry { + stopped: bool, + paths: HashMap, + documents: HashMap>, +} + +struct ManagedDocument { + id: String, + path: PathBuf, + state: Mutex, + control: AsyncMutex<()>, +} + +struct DocumentState { + document: NotebookDocument, + generation: u64, + transitioning: bool, + closed: bool, + active: Option>, + kernel_status: Option<(String, Option)>, +} + +struct Batch { + cancel: CancellationToken, + done: watch::Sender, +} + +struct BatchGuard { + document: Arc, + batch: Arc, +} + +pub(crate) struct ExecutionAdmission { + guard: BatchGuard, + generation: u64, + plan: Vec<(String, String)>, +} + +impl Drop for BatchGuard { + fn drop(&mut self) { + let mut state = self.document.state.lock(); + if state.active.as_ref().is_some_and(|active| Arc::ptr_eq(active, &self.batch)) { + if let Err(error) = state.document.set_execution_active(false) { + tracing::error!("notebook batch release: {error}"); + } + state.active = None; + } + self.batch.done.send_replace(true); + } +} + +pub struct NotebookService { + registry: Arc>, + kernels: Arc, + sink: SharedSink, +} + +impl NotebookService { + pub fn new(runtime_dir: PathBuf, sink: SharedSink) -> Arc { + let registry = Arc::new(Mutex::new(Registry::default())); + let status_registry = registry.clone(); + let status_sink = sink.clone(); + let kernels = KernelManager::new( + runtime_dir, + Arc::new(move |path, status, message| { + let document = { + let registry = status_registry.lock(); + if registry.stopped { + return; + } + registry.paths.get(path).and_then(|id| registry.documents.get(id)).cloned() + }; + if let Some(document) = document { + let deliver = { + let mut state = document.state.lock(); + if state.closed { + return; + } + state.kernel_status = Some((status.as_str().to_owned(), message.clone())); + !state.transitioning + }; + if deliver { + status_sink.emit(IdeEvent::NotebookKernel { + id: document.id.clone(), + path: path.to_path_buf().into(), + state: status.as_str().to_owned(), + message, + }); + } + } + }), + ); + Arc::new(Self { registry, kernels: Arc::new(kernels), sink }) + } + + pub fn open(&self, path: &Path) -> IdeResult { + let path = dunce::canonicalize(path)?; + // Serialize first load with path insertion: split views share one incarnation. + let mut registry = self.registry.lock(); + if registry.stopped { + return Err(IdeError::cancelled()); + } + if let Some(id) = registry.paths.get(&path) { + return snapshot(®istry.documents[id]); + } + let document = NotebookDocument::load(&path)?; + let initial = document.snapshot(); + let id = initial.id.clone(); + registry.paths.insert(path.clone(), id.clone()); + registry.documents.insert( + id.clone(), + Arc::new(ManagedDocument { + id, + path, + state: Mutex::new(DocumentState { + document, + generation: 0, + transitioning: false, + closed: false, + active: None, + kernel_status: None, + }), + control: AsyncMutex::new(()), + }), + ); + Ok(initial) + } + + fn resolve(&self, id: &str) -> IdeResult> { + let registry = self.registry.lock(); + if registry.stopped { + return Err(IdeError::not_found("notebook project has closed")); + } + registry.documents.get(id).cloned().ok_or_else(|| IdeError::not_found("notebook incarnation is closed")) + } + + pub fn get(&self, id: &str) -> IdeResult { + snapshot(self.resolve(id)?.as_ref()) + } + + pub fn edit(&self, id: &str, version: u32, edit: NotebookEdit) -> IdeResult { + let document = self.resolve(id)?; + let mut state = document.state.lock(); + mutable(&state)?; + state.document.edit(version, edit) + } + + pub fn save(&self, id: &str) -> IdeResult { + let document = self.resolve(id)?; + let mut state = document.state.lock(); + mutable(&state)?; + state.document.save() + } + + pub fn undo(&self, id: &str, redo: bool) -> IdeResult { + let document = self.resolve(id)?; + let mut state = document.state.lock(); + mutable(&state)?; + if state.active.is_some() { + return Err(IdeError::unavailable("stop notebook execution before undo or redo")); + } + if redo { state.document.redo() } else { state.document.undo() } + } + + pub async fn interrupt(&self, id: &str) -> IdeResult<()> { + let document = self.resolve(id)?; + let batch = document.state.lock().active.clone(); + if let Some(batch) = &batch { + batch.cancel.cancel(); + } + // Draining makes prepareClose/save a stable checkpoint, including when + // Python catches KeyboardInterrupt or cancel races with kernel startup. + let interrupt = self.kernels.interrupt(&document.path).await; + self.drain(&document, batch).await?; + interrupt + } + + async fn drain(&self, document: &ManagedDocument, batch: Option>) -> IdeResult<()> { + let Some(batch) = batch else { return Ok(()) }; + let mut done = batch.done.subscribe(); + match tokio::time::timeout(Duration::from_secs(5), done.wait_for(|done| *done)).await { + Ok(Ok(_)) => return Ok(()), + Ok(Err(_)) => return Err(IdeError::internal("notebook execution completion channel closed")), + Err(_) => {} + } + // An uncooperative kernel must not outlive a close/reload authorization. + self.kernels.shutdown_kernel(&document.path).await?; + tokio::time::timeout(Duration::from_secs(5), done.wait_for(|done| *done)) + .await + .map_err(|_| IdeError::unavailable("notebook execution did not stop after kernel shutdown"))? + .map_err(|_| IdeError::internal("notebook execution completion channel closed"))?; + Ok(()) + } + + pub async fn reload(&self, id: &str, restart: bool) -> IdeResult { + let document = self.resolve(id)?; + let _control = document.control.lock().await; + let batch = invalidate(&document)?; + let result = async { + let _ = self.kernels.interrupt(&document.path).await; + self.drain(&document, batch).await?; + if document.state.lock().closed { + return Err(IdeError::cancelled()); + } + if restart { + self.kernels.restart(&document.path).await?; + snapshot(&document) + } else { + self.kernels.shutdown_kernel(&document.path).await?; + document.state.lock().document.reload() + } + } + .await; + let status = { + let mut state = document.state.lock(); + state.transitioning = false; + if state.closed { None } else { state.kernel_status.clone() } + }; + if let Some((state, message)) = status { + self.sink.emit(IdeEvent::NotebookKernel { + id: document.id.clone(), + path: document.path.clone().into(), + state, + message, + }); + } + result + } + + /// Frontend disposal may arrive after project replacement has already retired this id. + pub async fn close(&self, id: &str) -> IdeResult<()> { + let document = match self.resolve(id) { + Ok(document) => document, + Err(error) if error.kind == ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(error), + }; + let _control = document.control.lock().await; + if document.state.lock().closed { + return Ok(()); + } + let batch = invalidate(&document)?; + let result = async { + let _ = self.kernels.interrupt(&document.path).await; + self.drain(&document, batch).await?; + self.kernels.shutdown_kernel(&document.path).await + } + .await; + if let Err(error) = result { + document.state.lock().transitioning = false; + return Err(error); + } + document.state.lock().closed = true; + let mut registry = self.registry.lock(); + registry.documents.remove(id); + registry.paths.remove(&document.path); + Ok(()) + } + + pub(crate) fn admit( + &self, + id: &str, + cells: Vec, + cancel: CancellationToken, + ) -> IdeResult { + admit(self.resolve(id)?, cells, cancel) + } + + pub(crate) async fn execute( + self: &Arc, + admission: ExecutionAdmission, + python: PathBuf, + on_output: Channel, + ) -> IdeResult { + let document = admission.guard.document.clone(); + let service = self.clone(); + // Keep admission and native execution owned if an IPC receiver disappears. + tokio::spawn(async move { + let ExecutionAdmission { guard, generation, plan } = admission; + let result = + service.run_batch(document.clone(), generation, plan, python, guard.batch.clone(), on_output).await; + drop(guard); + match result { + Ok(()) => snapshot(&document), + Err(error) => Err(error), + } + }) + .await + .map_err(|error| IdeError::internal(error.to_string()))? + } + + async fn run_batch( + &self, + document: Arc, + generation: u64, + plan: Vec<(String, String)>, + python: PathBuf, + batch: Arc, + on_output: Channel, + ) -> IdeResult<()> { + for (cell, source) in plan { + { + let mut state = document.state.lock(); + if batch.cancel.is_cancelled() || state.closed || state.generation != generation { + return Err(IdeError::cancelled()); + } + // A user may delete a queued cell while an earlier one executes. + // Never resurrect it from the captured source or a late output. + match state.document.begin_execution(&cell) { + Ok(()) => {} + Err(error) if error.kind == ErrorKind::NotFound => continue, + Err(error) => return Err(error), + } + } + let output_error = Arc::new(Mutex::new(None)); + let callback_document = document.clone(); + let callback_batch = batch.clone(); + let callback_error = output_error.clone(); + let callback_cell = cell.clone(); + let channel = on_output.clone(); + let result = self + .kernels + .execute( + &document.path, + &python, + &source, + batch.cancel.clone(), + Arc::new(move |output: KernelOutput| { + let event = match commit_output( + &callback_document, + &callback_batch, + generation, + &callback_cell, + output, + ) { + Ok(Some(event)) => event, + Ok(None) => return, + Err(error) => { + *callback_error.lock() = Some(error); + callback_batch.cancel.cancel(); + return; + } + }; + let _ = channel.send(event); + }), + ) + .await; + if let Some(error) = output_error.lock().take() { + return Err(error); + } + let result = result?; + if batch.cancel.is_cancelled() { + return Err(IdeError::cancelled()); + } + if result.status != "ok" { + break; + } + } + let state = document.state.lock(); + if state.closed || state.generation != generation || batch.cancel.is_cancelled() { + return Err(IdeError::cancelled()); + } + Ok(()) + } + + /// Synchronously revoke admission before handing cleanup to an async runtime. + pub fn stop_admission(&self) { + let mut registry = self.registry.lock(); + registry.stopped = true; + for document in registry.documents.values() { + let mut state = document.state.lock(); + state.closed = true; + state.generation += 1; + if let Some(batch) = &state.active { + batch.cancel.cancel(); + } + } + } + + pub async fn shutdown(&self) { + self.stop_admission(); + self.kernels.shutdown().await; + let documents: Vec<_> = self.registry.lock().documents.values().cloned().collect(); + for document in documents { + let batch = document.state.lock().active.clone(); + if let Err(error) = self.drain(&document, batch).await { + tracing::error!("notebook cleanup: {error}"); + } + } + let mut registry = self.registry.lock(); + registry.paths.clear(); + registry.documents.clear(); + } +} + +fn mutable(state: &DocumentState) -> IdeResult<()> { + if state.closed { + Err(IdeError::not_found("notebook incarnation is closed")) + } else if state.transitioning { + Err(IdeError::unavailable("notebook lifecycle operation is in progress")) + } else { + Ok(()) + } +} + +fn snapshot(document: &ManagedDocument) -> IdeResult { + let state = document.state.lock(); + if state.closed { + return Err(IdeError::not_found("notebook incarnation is closed")); + } + Ok(state.document.snapshot()) +} + +fn invalidate(document: &ManagedDocument) -> IdeResult>> { + let mut state = document.state.lock(); + mutable(&state)?; + state.transitioning = true; + state.generation += 1; + if let Some(batch) = &state.active { + batch.cancel.cancel(); + } + Ok(state.active.clone()) +} + +fn execution_plan(snapshot: &NotebookSnapshot, cells: Vec) -> IdeResult> { + let mut selected: HashSet<_> = cells.iter().map(String::as_str).collect(); + if selected.is_empty() || selected.len() != cells.len() { + return Err(IdeError::invalid("select distinct code cells to execute")); + } + let mut plan = Vec::with_capacity(selected.len()); + for cell in &snapshot.cells { + if selected.remove(cell.id.as_str()) { + if cell.cell_type != "code" { + return Err(IdeError::invalid("only code cells can execute")); + } + plan.push((cell.id.clone(), cell.source.clone())); + } + } + if !selected.is_empty() { + return Err(IdeError::not_found("an execution cell no longer exists")); + } + Ok(plan) +} + +fn admit( + document: Arc, + cells: Vec, + cancel: CancellationToken, +) -> IdeResult { + let (batch, generation, plan) = { + let mut state = document.state.lock(); + mutable(&state)?; + if state.active.is_some() { + return Err(IdeError::unavailable("notebook execution batch is already running")); + } + if cancel.is_cancelled() { + return Err(IdeError::cancelled()); + } + let plan = execution_plan(&state.document.snapshot(), cells)?; + state.document.set_execution_active(true)?; + let (done, _) = watch::channel(false); + let batch = Arc::new(Batch { cancel, done }); + state.active = Some(batch.clone()); + (batch, state.generation, plan) + }; + Ok(ExecutionAdmission { guard: BatchGuard { document, batch }, generation, plan }) +} + +/// Commit under the short document lock; the caller delivers only after release. +fn commit_output( + document: &ManagedDocument, + batch: &Arc, + generation: u64, + cell: &str, + output: KernelOutput, +) -> IdeResult> { + let mut state = document.state.lock(); + if state.closed + || state.generation != generation + || !state.active.as_ref().is_some_and(|active| Arc::ptr_eq(active, batch)) + { + return Ok(None); + } + let changed_cells = match state.document.record_output(cell, &output) { + Ok(changed) => changed, + Err(error) if error.kind == ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error), + }; + Ok(Some(NotebookOutputEvent { + id: document.id.clone(), + path: document.path.to_string_lossy().into_owned(), + cell: cell.to_owned(), + msg_id: output.msg_id, + kind: output.kind, + content: output.content, + version: state.document.version(), + dirty: state.document.is_dirty(), + changed_cells, + })) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn fixture() -> (tempfile::TempDir, Arc, NotebookSnapshot) { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session.ipynb"); + std::fs::write(&path, serde_json::to_vec(&json!({ + "nbformat": 4, "nbformat_minor": 5, "metadata": {}, + "cells": [ + {"id":"first","cell_type":"code","metadata":{},"source":"x = 1","outputs":[],"execution_count":null}, + {"id":"second","cell_type":"code","metadata":{},"source":"print(x)","outputs":[],"execution_count":null}, + {"id":"prose","cell_type":"markdown","metadata":{},"source":"notes"} + ] + })).unwrap()).unwrap(); + let service = NotebookService::new(dir.path().join("kernels"), Arc::new(ide_core::NullSink)); + let snapshot = service.open(&path).unwrap(); + (dir, service, snapshot) + } + + fn output(kind: &str, content: serde_json::Value) -> KernelOutput { + KernelOutput { kind: kind.into(), content, msg_id: "actual-parent-request".into() } + } + + #[tokio::test] + async fn canonical_split_opens_share_incarnation_but_reopen_and_projects_do_not() { + let (dir, service, initial) = fixture(); + let alternate = dir.path().join(".").join("session.ipynb"); + assert_eq!(service.open(&alternate).unwrap().id, initial.id); + let other = NotebookService::new(dir.path().join("other-runtime"), Arc::new(ide_core::NullSink)); + let independent = other.open(&alternate).unwrap(); + assert_ne!(initial.id, independent.id); + assert_eq!(other.get(&initial.id).unwrap_err().kind, ErrorKind::NotFound); + service.close(&initial.id).await.unwrap(); + let reopened = service.open(&alternate).unwrap(); + assert_ne!(reopened.id, initial.id); + assert_eq!(service.get(&initial.id).unwrap_err().kind, ErrorKind::NotFound); + // A delayed old-incarnation close cannot dispose a new session at that path. + service.close(&initial.id).await.unwrap(); + assert_eq!(service.get(&reopened.id).unwrap().id, reopened.id); + service.shutdown().await; + other.shutdown().await; + } + + #[tokio::test(start_paused = true)] + async fn failed_close_keeps_the_incarnation_recoverable_until_a_retry_succeeds() { + let (dir, service, initial) = fixture(); + // A batch that never drains models a kernel that ignores interrupt and shutdown. + let admission = + service.admit(&initial.id, vec![initial.cells[0].id.clone()], CancellationToken::new()).unwrap(); + let error = service.close(&initial.id).await.unwrap_err(); + assert_eq!(error.kind, ErrorKind::Unavailable); + assert!(admission.guard.batch.cancel.is_cancelled()); + // The backend still owns the incarnation: a reopen shares it instead of + // starting a second kernel, and the failed close left it usable. + let reopened = service.open(&dir.path().join("session.ipynb")).unwrap(); + assert_eq!(reopened.id, initial.id); + assert!(matches!(service.admit(&initial.id, vec![initial.cells[0].id.clone()], CancellationToken::new()), + Err(error) if error.kind == ErrorKind::Unavailable)); + drop(admission); + service.close(&initial.id).await.unwrap(); + assert_eq!(service.get(&initial.id).unwrap_err().kind, ErrorKind::NotFound); + assert_ne!(service.open(&dir.path().join("session.ipynb")).unwrap().id, initial.id); + service.shutdown().await; + } + + #[test] + fn batch_admission_captures_document_order_and_retains_busy_until_drained() { + let (_dir, service, initial) = fixture(); + let document = service.resolve(&initial.id).unwrap(); + let cells = vec![initial.cells[1].id.clone(), initial.cells[0].id.clone()]; + let admission = admit(document.clone(), cells.clone(), CancellationToken::new()).unwrap(); + assert_eq!( + admission.plan, + vec![(initial.cells[0].id.clone(), "x = 1".into()), (initial.cells[1].id.clone(), "print(x)".into()),] + ); + service + .edit( + &initial.id, + initial.version, + NotebookEdit::SetSource { cell: initial.cells[1].id.clone(), source: "print('new source')".into() }, + ) + .unwrap(); + assert_eq!(admission.plan[1].1, "print(x)"); + assert_eq!(service.undo(&initial.id, false).unwrap_err().kind, ErrorKind::Unavailable); + admission.guard.batch.cancel.cancel(); + assert!(matches!(admit(document.clone(), cells.clone(), CancellationToken::new()), Err(error) + if error.kind == ErrorKind::Unavailable)); + drop(admission); + let next = admit(document, cells, CancellationToken::new()).unwrap(); + assert_eq!(next.plan[1].1, "print('new source')"); + } + + #[test] + fn output_commits_revision_and_cross_cell_display_updates_before_delivery() { + let (_dir, service, initial) = fixture(); + let document = service.resolve(&initial.id).unwrap(); + let cells: Vec<_> = initial.cells[..2].iter().map(|cell| cell.id.clone()).collect(); + let admission = admit(document.clone(), cells.clone(), CancellationToken::new()).unwrap(); + let event = commit_output( + &document, + &admission.guard.batch, + admission.generation, + &cells[0], + output( + "display_data", + json!({"data":{"text/plain":"before"},"metadata":{},"transient":{"display_id":"shared"}}), + ), + ) + .unwrap() + .unwrap(); + let committed = service.get(&initial.id).unwrap(); + assert_eq!(event.version, committed.version); + assert!(event.version > initial.version); + assert!(event.dirty); + assert_eq!(event.msg_id, "actual-parent-request"); + assert_eq!(committed.cells[0].outputs[0]["data"]["text/plain"], "before"); + let updated = commit_output( + &document, + &admission.guard.batch, + admission.generation, + &cells[1], + output( + "update_display_data", + json!({"data":{"text/plain":"after"},"metadata":{},"transient":{"display_id":"shared"}}), + ), + ) + .unwrap() + .unwrap(); + assert!(updated.changed_cells.contains(&cells[0])); + let committed = service.get(&initial.id).unwrap(); + assert_eq!(committed.cells[0].outputs[0]["data"]["text/plain"], "after"); + assert!(committed.cells[0].outputs[0].get("transient").is_none()); + assert_eq!(updated.version, committed.version); + } + + #[test] + fn invalidated_or_completed_batches_cannot_commit_and_deleted_cells_stay_deleted() { + let (_dir, service, initial) = fixture(); + let document = service.resolve(&initial.id).unwrap(); + let cell = initial.cells[0].id.clone(); + let admission = admit(document.clone(), vec![cell.clone()], CancellationToken::new()).unwrap(); + let deleted = + service.edit(&initial.id, initial.version, NotebookEdit::DeleteCell { cell: cell.clone() }).unwrap(); + assert!( + commit_output( + &document, + &admission.guard.batch, + admission.generation, + &cell, + output("stream", json!({"name":"stdout","text":"late"})) + ) + .unwrap() + .is_none() + ); + assert_eq!(service.get(&initial.id).unwrap().version, deleted.version); + let other = initial.cells[1].id.clone(); + invalidate(&document).unwrap(); + assert!(admission.guard.batch.cancel.is_cancelled()); + assert!( + commit_output( + &document, + &admission.guard.batch, + admission.generation, + &other, + output("stream", json!({"name":"stdout","text":"old generation"})) + ) + .unwrap() + .is_none() + ); + let batch = admission.guard.batch.clone(); + let generation = document.state.lock().generation; + drop(admission); + assert!( + commit_output( + &document, + &batch, + generation, + &other, + output("stream", json!({"name":"stdout","text":"post idle"})) + ) + .unwrap() + .is_none() + ); + assert_eq!(service.get(&initial.id).unwrap().version, deleted.version); + } + + #[tokio::test] + async fn reload_preserves_incarnation_and_project_shutdown_revokes_all_admission() { + let (_dir, service, initial) = fixture(); + let changed = service + .edit( + &initial.id, + initial.version, + NotebookEdit::SetSource { cell: initial.cells[0].id.clone(), source: "changed".into() }, + ) + .unwrap(); + let reloaded = service.reload(&initial.id, false).await.unwrap(); + assert_eq!(reloaded.id, initial.id); + assert!(reloaded.version > changed.version); + assert_eq!(reloaded.cells[0].source, "x = 1"); + assert!(!reloaded.dirty); + let document = service.resolve(&initial.id).unwrap(); + let admission = admit(document.clone(), vec![initial.cells[0].id.clone()], CancellationToken::new()).unwrap(); + service.stop_admission(); + assert!(admission.guard.batch.cancel.is_cancelled()); + assert_eq!(service.get(&initial.id).unwrap_err().kind, ErrorKind::NotFound); + assert!(matches!(admit(document, vec![initial.cells[0].id.clone()], CancellationToken::new()), Err(error) + if error.kind == ErrorKind::NotFound)); + drop(admission); + service.shutdown().await; + } + + #[test] + fn invalid_batch_selection_does_not_claim_admission() { + let (_dir, service, initial) = fixture(); + let document = service.resolve(&initial.id).unwrap(); + for cells in [ + vec![], + vec![initial.cells[0].id.clone(), initial.cells[0].id.clone()], + vec![initial.cells[2].id.clone()], + vec!["deleted".into()], + ] { + assert!(admit(document.clone(), cells, CancellationToken::new()).is_err()); + } + let cancelled = CancellationToken::new(); + cancelled.cancel(); + assert!(matches!(admit(document.clone(), vec![initial.cells[0].id.clone()], cancelled), Err(error) + if error.kind == ErrorKind::Cancelled)); + let valid = admit(document, vec![initial.cells[0].id.clone()], CancellationToken::new()).unwrap(); + assert_eq!(valid.plan[0].1, "x = 1"); + } + + #[tokio::test] + async fn interrupt_cancels_before_drain_and_does_not_release_a_live_batch() { + let (_dir, service, initial) = fixture(); + service + .edit( + &initial.id, + initial.version, + NotebookEdit::SetSource { cell: initial.cells[0].id.clone(), source: "changed source".into() }, + ) + .unwrap(); + let cancel = CancellationToken::new(); + let admission = service.admit(&initial.id, vec![initial.cells[0].id.clone()], cancel.clone()).unwrap(); + let interrupt_service = service.clone(); + let id = initial.id.clone(); + let interrupt = tokio::spawn(async move { interrupt_service.interrupt(&id).await }); + tokio::time::timeout(Duration::from_secs(1), cancel.cancelled()).await.unwrap(); + assert!(!interrupt.is_finished(), "interrupt must wait for owned execution to drain"); + assert!(matches!(service.admit(&initial.id, vec![initial.cells[0].id.clone()], CancellationToken::new()), + Err(error) if error.kind == ErrorKind::Unavailable)); + let mut done = admission.guard.batch.done.subscribe(); + drop(admission); + assert!(*done.borrow_and_update()); + tokio::time::timeout(Duration::from_secs(1), interrupt).await.unwrap().unwrap().unwrap(); + let next = service.admit(&initial.id, vec![initial.cells[0].id.clone()], CancellationToken::new()).unwrap(); + assert_eq!(next.plan[0].1, "changed source"); + drop(next); + let undone = service.undo(&initial.id, false).unwrap(); + assert_eq!(undone.cells[0].source, "x = 1"); + service.close(&initial.id).await.unwrap(); + } +} diff --git a/crates/ide-host/src/ops/mod.rs b/crates/ide-host/src/ops/mod.rs index f1de3ea33..4baf33cac 100644 --- a/crates/ide-host/src/ops/mod.rs +++ b/crates/ide-host/src/ops/mod.rs @@ -15,6 +15,7 @@ pub mod history; pub mod java; pub mod lsp; pub mod nav; +pub mod notebook; pub mod project; mod python; pub mod refactor; @@ -60,6 +61,7 @@ impl Clone for Channel { pub fn register_all(d: &mut crate::rpc::Dispatcher) { agent::register(d); lsp::register(d); + notebook::register(d); nav::register(d); refactor::register(d); analysis::register(d); diff --git a/crates/ide-host/src/ops/notebook.rs b/crates/ide-host/src/ops/notebook.rs new file mode 100644 index 000000000..409d0ff20 --- /dev/null +++ b/crates/ide-host/src/ops/notebook.rs @@ -0,0 +1,155 @@ +//! Native notebook operations shared by desktop IPC and the remote dispatcher. +use crate::Host; +pub use crate::notebook::NotebookOutputEvent; +use crate::ops::Channel; +use crate::rpc::Dispatcher; +use ide_core::{IdeError, IdeResult, RequestId}; +use ide_notebook::{NotebookEdit, NotebookSnapshot}; +use std::path::PathBuf; + +pub async fn open(host: &Host, window: &str, path: String) -> IdeResult { + let project = host.project(window)?; + let path = PathBuf::from(path); + let path = if path.is_absolute() { path } else { project.root.join(path) }; + let notebooks = project.notebooks.clone(); + host.blocking(move || notebooks.open(&path)).await +} + +pub fn get(host: &Host, window: &str, id: String) -> IdeResult { + host.project(window)?.notebooks.get(&id) +} + +pub async fn edit( + host: &Host, + window: &str, + id: String, + version: u32, + edit: NotebookEdit, +) -> IdeResult { + let notebooks = host.project(window)?.notebooks.clone(); + host.blocking(move || notebooks.edit(&id, version, edit)).await +} + +pub async fn save(host: &Host, window: &str, id: String) -> IdeResult { + let notebooks = host.project(window)?.notebooks.clone(); + host.blocking(move || notebooks.save(&id)).await +} + +pub async fn undo(host: &Host, window: &str, id: String) -> IdeResult { + let notebooks = host.project(window)?.notebooks.clone(); + host.blocking(move || notebooks.undo(&id, false)).await +} + +pub async fn redo(host: &Host, window: &str, id: String) -> IdeResult { + let notebooks = host.project(window)?.notebooks.clone(); + host.blocking(move || notebooks.undo(&id, true)).await +} + +pub async fn reload(host: &Host, window: &str, id: String) -> IdeResult { + let notebooks = host.project(window)?.notebooks.clone(); + // An abandoned IPC future must not leave a lifecycle transition half applied. + tokio::spawn(async move { notebooks.reload(&id, false).await }) + .await + .map_err(|error| IdeError::internal(error.to_string()))? +} + +pub async fn restart(host: &Host, window: &str, id: String) -> IdeResult { + let notebooks = host.project(window)?.notebooks.clone(); + tokio::spawn(async move { notebooks.reload(&id, true).await }) + .await + .map_err(|error| IdeError::internal(error.to_string()))? +} + +pub async fn close(host: &Host, window: &str, id: String) -> IdeResult<()> { + let notebooks = host.project(window)?.notebooks.clone(); + tokio::spawn(async move { notebooks.close(&id).await }) + .await + .map_err(|error| IdeError::internal(error.to_string()))? +} + +pub async fn interrupt(host: &Host, window: &str, id: String) -> IdeResult<()> { + let notebooks = host.project(window)?.notebooks.clone(); + tokio::spawn(async move { notebooks.interrupt(&id).await }) + .await + .map_err(|error| IdeError::internal(error.to_string()))? +} + +#[allow(clippy::too_many_arguments)] +pub async fn execute( + host: &Host, + window: &str, + id: String, + cells: Vec, + python: Option, + request: RequestId, + on_output: Channel, +) -> IdeResult { + let project = host.project(window)?; + let cancel = host.begin_request(request); + let _request = RequestGuard { host, request, cancel: cancel.clone() }; + // Reserve the WHOLE batch and capture source before asynchronous discovery. + let admission = project.notebooks.admit(&id, cells, cancel.clone())?; + let selected = match python { + Some(path) if path.trim().is_empty() => return Err(IdeError::invalid("select a Python interpreter")), + Some(path) => PathBuf::from(path), + None => { + let root = project.root.clone(); + let settings = project.settings.effective(); + tokio::select! { + biased; + _ = cancel.cancelled() => return Err(IdeError::cancelled()), + result = host.blocking(move || { + ide_sdk::apply::python_interpreter(&root, &settings) + .or_else(|| ide_sdk::detect(ide_sdk::SdkKind::Python, Some(&root)) + .into_iter().next().map(|sdk| sdk.path.into_path_buf())) + .ok_or_else(|| IdeError::unavailable("select a Python environment with ipykernel installed")) + }) => result?, + } + } + }; + // Do not canonicalize this path: resolving a venv symlink to base Python + // changes sys.prefix and executes in the wrong environment. + let python = if selected.is_absolute() { selected } else { project.root.join(selected) }; + project.notebooks.execute(admission, python, on_output).await +} + +struct RequestGuard<'a> { + host: &'a Host, + request: RequestId, + cancel: ide_core::CancellationToken, +} + +impl Drop for RequestGuard<'_> { + fn drop(&mut self) { + // If the invocation is dropped, the owned batch still receives cancellation. + self.cancel.cancel(); + self.host.end_request(self.request); + } +} + +pub(crate) fn register(d: &mut Dispatcher) { + d.add("notebook_open", |c| async move { open(&c.host, &c.window, c.arg("path")?).await }); + d.add("notebook_get", |c| async move { get(&c.host, &c.window, c.arg("id")?) }); + d.add("notebook_edit", |c| async move { + edit(&c.host, &c.window, c.arg("id")?, c.arg("version")?, c.arg("edit")?).await + }); + d.add("notebook_save", |c| async move { save(&c.host, &c.window, c.arg("id")?).await }); + d.add("notebook_reload", |c| async move { reload(&c.host, &c.window, c.arg("id")?).await }); + d.add("notebook_undo", |c| async move { undo(&c.host, &c.window, c.arg("id")?).await }); + d.add("notebook_redo", |c| async move { redo(&c.host, &c.window, c.arg("id")?).await }); + d.add("notebook_close", |c| async move { close(&c.host, &c.window, c.arg("id")?).await }); + d.add("notebook_interrupt", |c| async move { interrupt(&c.host, &c.window, c.arg("id")?).await }); + d.add("notebook_restart", |c| async move { restart(&c.host, &c.window, c.arg("id")?).await }); + d.add("notebook_execute", |c| async move { + execute( + &c.host, + &c.window, + c.arg("id")?, + c.arg("cells")?, + c.arg("python")?, + c.arg("request")?, + c.sender("onOutput")?, + ) + .await + }); +} diff --git a/crates/ide-host/src/project.rs b/crates/ide-host/src/project.rs index 12d7c242c..4c84c8d30 100644 --- a/crates/ide-host/src/project.rs +++ b/crates/ide-host/src/project.rs @@ -46,6 +46,8 @@ pub struct Project { pub http_history: Arc, /// Coding agent tasks (Agents tool window). pub agents: Arc, + /// Native notebooks and owned kernels, scoped to this project incarnation. + pub notebooks: Arc, pub(crate) watcher: Watcher, /// Watchers of the Project Structure content roots outside the project folder. pub(crate) content_watchers: Mutex>, @@ -282,6 +284,10 @@ impl Host { tracing::info!("opened {} in window {window} in {} ms {laps:?}", root.display(), started.elapsed().as_millis()); let (db, http_history, agents) = self.project_services(&root, &project_key); + let notebooks = crate::notebook::NotebookService::new( + self.data_dir.join("notebook-runtime").join(uuid::Uuid::new_v4().to_string()), + sink.clone(), + ); ide_agent::registry::maintain(crate::ops::agent::registry_paths(self)); let project = Arc::new(Project { root, @@ -301,6 +307,7 @@ impl Host { db, http_history, agents, + notebooks, watcher, content_watchers: Mutex::new(content_watchers), _scratch_watcher: scratch_watcher, @@ -393,6 +400,10 @@ impl Host { let watcher = light_watcher(&dir, &sink, &documents, &history)?; let project_key = project_key(&dir); let (db, http_history, agents) = self.project_services(&dir, &project_key); + let notebooks = crate::notebook::NotebookService::new( + self.data_dir.join("notebook-runtime").join(uuid::Uuid::new_v4().to_string()), + sink.clone(), + ); let project = Arc::new(Project { sink, languages, @@ -410,6 +421,7 @@ impl Host { db, http_history, agents, + notebooks, watcher, content_watchers: Mutex::new(Vec::new()), _scratch_watcher: None, diff --git a/crates/ide-notebook/Cargo.toml b/crates/ide-notebook/Cargo.toml new file mode 100644 index 000000000..b180560ce --- /dev/null +++ b/crates/ide-notebook/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "ide-notebook" +description = "Lossless native notebook documents and owned Jupyter kernels" +version.workspace = true +edition.workspace = true +license.workspace = true + +[dependencies] +ide-core.workspace = true +ide-vfs.workspace = true +serde.workspace = true +serde_json.workspace = true +specta.workspace = true +specta-typescript.workspace = true +uuid.workspace = true +tokio.workspace = true +futures.workspace = true +parking_lot.workspace = true +tracing.workspace = true +jupyter-zmq-client.workspace = true +jupyter-protocol.workspace = true +zeromq.workspace = true +tempfile.workspace = true + +[target.'cfg(unix)'.dependencies] +libc = "0.2" + +[target.'cfg(windows)'.dependencies] +windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_System_Threading", "Win32_System_JobObjects", "Win32_Security", "Win32_Security_Authorization", "Win32_System_Memory"] } + diff --git a/crates/ide-notebook/src/document.rs b/crates/ide-notebook/src/document.rs new file mode 100644 index 000000000..9eb328422 --- /dev/null +++ b/crates/ide-notebook/src/document.rs @@ -0,0 +1,988 @@ +//! Lossless nbformat 4 documents. Runtime identities and execution bookkeeping +//! never leak into existing notebook JSON. Only explicit mutations alter bytes. +mod raw; + +use std::{ + collections::{HashMap, HashSet}, + path::{Path, PathBuf}, + sync::Arc, + time::{Duration, Instant}, +}; + +use ide_core::{IdeError, IdeResult, error::ErrorKind}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use specta::Type; +use uuid::Uuid; + +use self::raw::{Layout, Member, Node}; +use crate::kernel::KernelOutput; + +#[derive(Clone, Debug, Serialize, Deserialize, Type)] +#[serde(rename_all = "camelCase")] +pub struct NotebookSnapshot { + pub id: String, + pub path: String, + pub version: u32, + pub dirty: bool, + pub can_undo: bool, + pub can_redo: bool, + pub nbformat: u32, + pub nbformat_minor: u32, + #[specta(type = specta_typescript::Unknown)] + pub metadata: Value, + pub cells: Vec, +} + +#[derive(Clone, Debug, Serialize, Deserialize, Type)] +#[serde(rename_all = "camelCase")] +pub struct NotebookCell { + pub id: String, + pub cell_type: String, + pub source: String, + #[specta(type = specta_typescript::Unknown)] + pub metadata: Value, + #[specta(type = Vec)] + pub outputs: Vec, + pub execution_count: Option, + #[specta(type = specta_typescript::Unknown)] + pub attachments: Value, +} + +#[derive(Clone, Debug, Serialize, Deserialize, Type)] +#[serde(tag = "type", rename_all = "camelCase", rename_all_fields = "camelCase")] +pub enum NotebookEdit { + SetSource { + cell: String, + source: String, + }, + InsertCell { + after: Option, + cell_type: String, + }, + DeleteCell { + cell: String, + }, + MoveCell { + cell: String, + before: Option, + }, + SetCellType { + cell: String, + cell_type: String, + }, + SetOutputs { + cell: String, + #[specta(type = Vec)] + outputs: Vec, + execution_count: Option, + }, +} + +#[derive(Clone)] +struct CellState { + id: String, + display_ids: Vec>, + pending_clear: bool, +} + +#[derive(Clone)] +struct Fields { + values: Vec<(String, Option)>, + order: Vec, + layout: Layout, +} + +impl Fields { + fn capture(node: &Node, keys: &[&str]) -> Self { + let members = node.members().expect("validated cell object"); + Self { + values: keys.iter().map(|key| ((*key).into(), members.iter().find(|m| m.key == *key).cloned())).collect(), + order: members.iter().map(|m| m.key.clone()).collect(), + layout: node.layout(), + } + } + + fn restore(&self, node: &mut Node) { + node.restore_fields(&self.values, &self.order, &self.layout); + } +} + +enum Change { + Fields { cell: String, before: Fields, after: Fields }, + Insert { index: usize, node: Node, state: CellState, before: Layout, after: Layout }, + Delete { index: usize, node: Node, state: CellState, before: Layout, after: Layout }, + Move { from: usize, to: usize, before: Layout, after: Layout }, +} + +struct History { + change: Change, + typing_cell: Option, + time: Instant, +} + +/// The caller serializes short document operations, not kernel awaits. Batch +/// admission belongs to the host; set_execution_active protects backend history. +pub struct NotebookDocument { + id: String, + path: PathBuf, + root: Node, + prefix: String, + suffix: String, + baseline: Node, + baseline_text: Arc, + cells: Vec, + version: u32, + last_user_revision: u32, + undo: Vec, + redo: Vec, + execution_active: bool, +} + +impl NotebookDocument { + pub fn load(path: &Path) -> IdeResult { + let text = std::fs::read_to_string(path)?; + Self::from_text(path.to_path_buf(), text) + } + + pub fn from_text(path: PathBuf, text: String) -> IdeResult { + let text: Arc = text.into(); + let (root, prefix, suffix) = Node::parse(text.clone())?; + validate_notebook(&root)?; + let nodes = root.get("cells").and_then(Node::array).expect("validated cells"); + let mut persisted = HashMap::::new(); + for node in nodes { + if let Some(id) = string_field(node, "id").filter(|id| valid_id(id)) { + *persisted.entry(id.to_owned()).or_default() += 1; + } + } + let mut used: HashSet = persisted.keys().cloned().collect(); + let cells = nodes + .iter() + .map(|node| { + let id = string_field(node, "id") + .filter(|id| valid_id(id) && persisted.get(*id) == Some(&1)) + .map(str::to_owned) + .unwrap_or_else(|| unique_id(&mut used)); + let outputs = node.get("outputs").and_then(Node::array).map_or(0, <[Node]>::len); + CellState { id, display_ids: vec![None; outputs], pending_clear: false } + }) + .collect(); + Ok(Self { + id: Uuid::new_v4().to_string(), + path, + baseline: root.clone(), + root, + prefix, + suffix, + baseline_text: text, + cells, + version: 1, + last_user_revision: 1, + undo: Vec::new(), + redo: Vec::new(), + execution_active: false, + }) + } + + pub fn snapshot(&self) -> NotebookSnapshot { + NotebookSnapshot { + id: self.id.clone(), + path: self.path.to_string_lossy().into_owned(), + version: self.version, + dirty: self.is_dirty(), + can_undo: !self.execution_active && !self.undo.is_empty(), + can_redo: !self.execution_active && !self.redo.is_empty(), + nbformat: uint_field(&self.root, "nbformat").expect("validated format"), + nbformat_minor: uint_field(&self.root, "nbformat_minor").expect("validated minor"), + metadata: self.root.get("metadata").expect("validated metadata").value(), + cells: self + .nodes() + .iter() + .zip(&self.cells) + .map(|(node, state)| NotebookCell { + id: state.id.clone(), + cell_type: string_field(node, "cell_type").expect("validated type").into(), + source: source(node).expect("validated source"), + metadata: node.get("metadata").map_or_else(|| json!({}), Node::value), + outputs: node + .get("outputs") + .and_then(Node::array) + .map_or_else(Vec::new, |items| items.iter().map(Node::value).collect()), + execution_count: uint_field(node, "execution_count"), + attachments: node.get("attachments").map_or_else(|| json!({}), Node::value), + }) + .collect(), + } + } + + /// Byte equality, not a sticky flag. Unchanged raw siblings short-circuit + /// without serializing the notebook on each source edit or output chunk. + pub fn is_dirty(&self) -> bool { + !self.root.same_bytes(&self.baseline) + } + + /// The host uses these cheap fields for high-volume output events. + pub fn version(&self) -> u32 { + self.version + } + pub fn id(&self) -> &str { + &self.id + } + pub fn path(&self) -> &Path { + &self.path + } + + /// Intended for save/export and behavioral tests, not per-chunk IPC. + pub fn to_text(&self) -> String { + let mut text = String::with_capacity(self.baseline_text.len()); + text.push_str(&self.prefix); + self.root.write(&mut text); + text.push_str(&self.suffix); + text + } + + pub fn save(&mut self) -> IdeResult { + let disk = std::fs::read(&self.path).map_err(|error| { + if error.kind() == std::io::ErrorKind::NotFound { + out_of_sync("Notebook file was removed; reload or save a separate copy before overwriting") + } else { + error.into() + } + })?; + if disk != self.baseline_text.as_bytes() { + return Err(out_of_sync("Notebook changed on disk; reload or save a separate copy before overwriting")); + } + if self.is_dirty() { + let text = self.to_text(); + ide_vfs::write_atomic_strict(&self.path, text.as_bytes())?; + self.baseline_text = text.into(); + self.baseline = self.root.clone(); + } + // A save is a typing-group boundary; undo after it must become dirty. + self.break_typing_group(); + Ok(self.snapshot()) + } + + pub fn reload(&mut self) -> IdeResult { + self.require_inactive()?; + self.ensure_revision()?; + let mut replacement = Self::load(&self.path)?; + replacement.id.clone_from(&self.id); + replacement.version = self.version + 1; + replacement.last_user_revision = replacement.version; + *self = replacement; + Ok(self.snapshot()) + } + + pub fn set_execution_active(&mut self, active: bool) -> IdeResult<()> { + if active && self.execution_active { + return Err(IdeError::invalid("Notebook already has an active execution batch")); + } + self.execution_active = active; + self.break_typing_group(); + Ok(()) + } + + fn nodes(&self) -> &[Node] { + self.root.get("cells").and_then(Node::array).expect("validated cells") + } + fn array_mut(&mut self) -> &mut Node { + self.root.get_mut("cells").expect("validated cells") + } + fn cell_mut(&mut self, index: usize) -> &mut Node { + self.array_mut().item_mut(index) + } + fn index(&self, id: &str) -> IdeResult { + self.cells + .iter() + .position(|cell| cell.id == id) + .ok_or_else(|| IdeError::not_found(format!("Notebook cell {id} no longer exists"))) + } + fn ensure_revision(&self) -> IdeResult<()> { + if self.version == u32::MAX { + Err(IdeError::unavailable("Notebook revision limit reached; save and reopen the document")) + } else { + Ok(()) + } + } + fn require_inactive(&self) -> IdeResult<()> { + if self.execution_active { + Err(IdeError::invalid("Stop and drain notebook execution before undo, redo, or reload")) + } else { + Ok(()) + } + } + fn break_typing_group(&mut self) { + if let Some(last) = self.undo.last_mut() { + last.typing_cell = None; + } + } +} + +fn out_of_sync(message: &str) -> IdeError { + IdeError::new(ErrorKind::OutOfSync, message) +} +fn string_field<'a>(node: &'a Node, key: &str) -> Option<&'a str> { + node.get(key)?.string() +} +fn uint_field(node: &Node, key: &str) -> Option { + node.get(key)?.uint() +} +fn valid_id(id: &str) -> bool { + !id.is_empty() && id.len() <= 64 && id.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_') +} +fn unique_id(used: &mut HashSet) -> String { + loop { + let id = Uuid::new_v4().to_string(); + if used.insert(id.clone()) { + return id; + } + } +} + +fn text_value(value: &Value) -> IdeResult { + match value { + Value::String(text) => Ok(text.clone()), + Value::Array(parts) => { + let mut text = String::new(); + for part in parts { + text.push_str( + part.as_str().ok_or_else(|| IdeError::invalid("Notebook text arrays must contain only strings"))?, + ); + } + Ok(text) + } + _ => Err(IdeError::invalid("Notebook source/text must be a string or string array")), + } +} + +fn source(node: &Node) -> IdeResult { + match node.get("source") { + Some(value) => node_text(value), + None if !matches!(string_field(node, "cell_type"), Some("code" | "markdown" | "raw")) => Ok(String::new()), + None => Err(IdeError::invalid("Notebook cell source is missing")), + } +} + +fn validate_notebook(root: &Node) -> IdeResult<()> { + if root.members().is_none() + || uint_field(root, "nbformat") != Some(4) + || uint_field(root, "nbformat_minor").is_none() + { + return Err(IdeError::invalid( + "Only nbformat 4 notebooks with a nonnegative integer minor version are supported; the file has not been modified", + )); + } + if root.get("metadata").is_none_or(|n| n.members().is_none()) { + return Err(IdeError::invalid("Notebook metadata must be an object")); + } + let cells = + root.get("cells").and_then(Node::array).ok_or_else(|| IdeError::invalid("Notebook cells must be an array"))?; + for node in cells { + if node.members().is_none() || string_field(node, "cell_type").is_none() { + return Err(IdeError::invalid("Notebook cells must be objects with a string cell_type")); + } + let known = matches!(string_field(node, "cell_type"), Some("code" | "markdown" | "raw")); + match node.get("source") { + Some(value) => validate_text_node(value)?, + None if known => return Err(IdeError::invalid("Notebook cell source is missing")), + None => {} + } + if known && node.get("metadata").is_none() { + return Err(IdeError::invalid("Notebook cell metadata is missing")); + } + if let Some(metadata) = node.get("metadata") + && metadata.members().is_none() + { + return Err(IdeError::invalid("Cell metadata must be an object")); + } + if let Some(attachments) = node.get("attachments") + && attachments.members().is_none() + { + return Err(IdeError::invalid("Cell attachments must be an object")); + } + if let Some(outputs) = node.get("outputs") { + let outputs = outputs.array().ok_or_else(|| IdeError::invalid("Cell outputs must be an array"))?; + for output in outputs { + validate_output_node(output)?; + } + } + if let Some(count) = node.get("execution_count") + && !count.is_null() + && uint_field(node, "execution_count").is_none() + { + return Err(IdeError::invalid("Cell execution_count must be null or a nonnegative 32-bit integer")); + } + if string_field(node, "cell_type") == Some("code") + && (node.get("outputs").is_none() || node.get("execution_count").is_none()) + { + return Err(IdeError::invalid("Code cells require outputs and execution_count")); + } + } + Ok(()) +} + +impl NotebookDocument { + pub fn edit(&mut self, base_version: u32, edit: NotebookEdit) -> IdeResult { + if base_version < self.last_user_revision || base_version > self.version { + return Err(out_of_sync("Notebook was edited elsewhere; refresh and rebase pending edits")); + } + self.ensure_revision()?; + if let NotebookEdit::SetOutputs { cell, outputs, execution_count } = edit { + let index = self.code_index(&cell)?; + let outputs = outputs.into_iter().map(persisted_output).collect::>>()?; + let node = Node::new(Value::Array(outputs)); + for output in node.array().expect("outputs array") { + validate_output_node(output)?; + } + let count = Node::new(json!(execution_count)); + let changed = !self.nodes()[index].get("outputs").expect("code outputs").same_bytes(&node) + || !self.nodes()[index].get("execution_count").expect("code count").same_bytes(&count); + if changed { + let len = node.array().expect("outputs array").len(); + self.cell_mut(index).set("outputs", Some(node)); + self.cell_mut(index).set("execution_count", Some(count)); + self.cells[index].display_ids = vec![None; len]; + self.cells[index].pending_clear = false; + self.version += 1; + } + return Ok(self.snapshot()); + } + let typing_cell = match &edit { + NotebookEdit::SetSource { cell, .. } => Some(cell.clone()), + _ => None, + }; + let change = match edit { + NotebookEdit::SetSource { cell, source: text } => { + let index = self.index(&cell)?; + if self.nodes()[index].get("source").is_some_and(|old| text_matches(old, &text)) { + return Ok(self.snapshot()); + } + let before = Fields::capture(&self.nodes()[index], &["source"]); + let value = rewritten_source(self.nodes()[index].get("source"), text); + self.cell_mut(index).set("source", Some(value)); + let after = Fields::capture(&self.nodes()[index], &["source"]); + Change::Fields { cell, before, after } + } + NotebookEdit::InsertCell { after, cell_type } => { + validate_cell_type(&cell_type)?; + let index = after.map_or(Ok(0), |id| self.index(&id).map(|index| index + 1))?; + let id = unique_id(&mut self.cells.iter().map(|state| state.id.clone()).collect()); + let mut value = json!({ "cell_type": cell_type, "metadata": {}, "source": [] }); + if uint_field(&self.root, "nbformat_minor").expect("minor") >= 5 { + value["id"] = json!(id); + } + if value["cell_type"] == "code" { + value["outputs"] = json!([]); + value["execution_count"] = Value::Null; + } + let node = Node::new(value); + let state = CellState { id, display_ids: Vec::new(), pending_clear: false }; + let before = self.root.get("cells").expect("cells").layout(); + self.array_mut().insert(index, node.clone()); + self.cells.insert(index, state.clone()); + let after = self.root.get("cells").expect("cells").layout(); + Change::Insert { index, node, state, before, after } + } + NotebookEdit::DeleteCell { cell } => { + let index = self.index(&cell)?; + let before = self.root.get("cells").expect("cells").layout(); + let node = self.array_mut().remove(index); + let state = self.cells.remove(index); + let after = self.root.get("cells").expect("cells").layout(); + Change::Delete { index, node, state, before, after } + } + NotebookEdit::MoveCell { cell, before: target } => { + let from = self.index(&cell)?; + let target = target.map_or(Ok(self.cells.len()), |id| self.index(&id))?; + let to = if target > from { target - 1 } else { target }; + if from == to { + return Ok(self.snapshot()); + } + let before = self.root.get("cells").expect("cells").layout(); + self.move_cell(from, to); + let after = self.root.get("cells").expect("cells").layout(); + Change::Move { from, to, before, after } + } + NotebookEdit::SetCellType { cell, cell_type } => { + validate_cell_type(&cell_type)?; + let index = self.index(&cell)?; + if string_field(&self.nodes()[index], "cell_type") == Some(cell_type.as_str()) { + return Ok(self.snapshot()); + } + let keys = ["cell_type", "outputs", "execution_count"]; + let before = Fields::capture(&self.nodes()[index], &keys); + let node = self.cell_mut(index); + node.set("cell_type", Some(Node::new(json!(cell_type)))); + if cell_type == "code" { + node.set("outputs", Some(Node::new(json!([])))); + node.set("execution_count", Some(Node::new(Value::Null))); + } else { + node.set("outputs", None); + node.set("execution_count", None); + } + self.cells[index].display_ids.clear(); + self.cells[index].pending_clear = false; + let after = Fields::capture(&self.nodes()[index], &keys); + Change::Fields { cell, before, after } + } + NotebookEdit::SetOutputs { .. } => unreachable!("output edits handled above"), + }; + self.version += 1; + self.last_user_revision = self.version; + self.redo.clear(); + let now = Instant::now(); + let grouped = typing_cell.as_ref().is_some_and(|cell| { + self.undo.last().is_some_and(|last| { + last.typing_cell.as_ref() == Some(cell) && now.duration_since(last.time) < Duration::from_millis(750) + }) + }); + if grouped { + if let ( + Some(History { change: Change::Fields { after: previous, .. }, time, .. }), + Change::Fields { after, .. }, + ) = (self.undo.last_mut(), change) + { + *previous = after; + *time = now; + } + } else { + self.undo.push(History { change, typing_cell, time: now }); + // Bounded operation count; streaming outputs never enter history. + if self.undo.len() > 200 { + self.undo.remove(0); + } + } + Ok(self.snapshot()) + } + + pub fn undo(&mut self) -> IdeResult { + self.require_inactive()?; + self.ensure_revision()?; + if let Some(mut history) = self.undo.pop() { + if let Err(error) = self.apply_history(&mut history.change, false) { + self.undo.push(history); + return Err(error); + } + history.typing_cell = None; + self.redo.push(history); + self.version += 1; + self.last_user_revision = self.version; + } + self.break_typing_group(); + Ok(self.snapshot()) + } + + pub fn redo(&mut self) -> IdeResult { + self.require_inactive()?; + self.ensure_revision()?; + if let Some(mut history) = self.redo.pop() { + if let Err(error) = self.apply_history(&mut history.change, true) { + self.redo.push(history); + return Err(error); + } + history.typing_cell = None; + self.undo.push(history); + self.version += 1; + self.last_user_revision = self.version; + } + Ok(self.snapshot()) + } + + fn apply_history(&mut self, change: &mut Change, forward: bool) -> IdeResult<()> { + let insertion_change = matches!(change, Change::Insert { .. }); + match change { + Change::Fields { cell, before, after } => { + let index = self.index(cell)?; + let keys: Vec<&str> = before.values.iter().map(|(key, _)| key.as_str()).collect(); + let current = Fields::capture(&self.nodes()[index], &keys); + if forward { + *before = current; + } else { + *after = current; + } + let fields = if forward { after } else { before }; + fields.restore(self.cell_mut(index)); + if fields.values.iter().any(|(key, _)| key == "outputs") { + let len = self.nodes()[index].get("outputs").and_then(Node::array).map_or(0, <[Node]>::len); + self.cells[index].display_ids = vec![None; len]; + self.cells[index].pending_clear = false; + } + } + Change::Insert { index, node, state, before, after } + | Change::Delete { index, node, state, before, after } => { + let insert = insertion_change == forward; + if insert { + self.array_mut().insert(*index, node.clone()); + self.cells.insert(*index, state.clone()); + } else { + *node = self.array_mut().remove(*index); + *state = self.cells.remove(*index); + } + self.array_mut().restore_layout(if forward { after } else { before }.clone()); + } + Change::Move { from, to, before, after } => { + if forward { + self.move_cell(*from, *to); + } else { + self.move_cell(*to, *from); + } + self.array_mut().restore_layout(if forward { after } else { before }.clone()); + } + } + Ok(()) + } + + fn move_cell(&mut self, from: usize, to: usize) { + let node = self.array_mut().remove(from); + self.array_mut().insert(to, node); + let state = self.cells.remove(from); + self.cells.insert(to, state); + } + + fn code_index(&self, cell: &str) -> IdeResult { + let index = self.index(cell)?; + if string_field(&self.nodes()[index], "cell_type") != Some("code") { + return Err(IdeError::invalid("Only code cells accept execution output")); + } + Ok(index) + } +} + +fn validate_cell_type(cell_type: &str) -> IdeResult<()> { + if matches!(cell_type, "code" | "markdown" | "raw") { + Ok(()) + } else { + Err(IdeError::invalid("New/converted cell type must be code, markdown, or raw")) + } +} + +fn persisted_output(mut value: Value) -> IdeResult { + let object = value.as_object_mut().ok_or_else(|| IdeError::invalid("Notebook outputs must be objects"))?; + if !object.get("output_type").is_some_and(Value::is_string) { + return Err(IdeError::invalid("Notebook outputs require a string output_type")); + } + object.remove("transient"); + Ok(value) +} + +impl NotebookDocument { + pub fn begin_execution(&mut self, cell: &str) -> IdeResult<()> { + let index = self.code_index(cell)?; + self.ensure_revision()?; + let changed = self.clear_outputs(index); + let count_changed = !self.nodes()[index].get("execution_count").expect("code count").is_null(); + if count_changed { + self.cell_mut(index).set("execution_count", Some(Node::new(Value::Null))); + } + if changed || count_changed { + self.version += 1; + } + Ok(()) + } + + /// Commits one IOPub event and returns all cells whose persisted bytes + /// changed. Display identifiers and deferred clears exist only in memory. + pub fn record_output(&mut self, cell: &str, event: &KernelOutput) -> IdeResult> { + let index = self.code_index(cell)?; + self.ensure_revision()?; + let content = &event.content; + let mut changed = Vec::new(); + match event.kind.as_str() { + "execute_input" => { + let count = content + .get("execution_count") + .and_then(Value::as_u64) + .and_then(|n| u32::try_from(n).ok()) + .ok_or_else(|| IdeError::invalid("Kernel execute_input lacks a valid execution_count"))?; + if uint_field(&self.nodes()[index], "execution_count") != Some(count) { + self.cell_mut(index).set("execution_count", Some(Node::new(json!(count)))); + changed.push(cell.to_owned()); + } + } + "clear_output" => { + if content.get("wait").and_then(Value::as_bool).unwrap_or(false) { + self.cells[index].pending_clear = true; + } else if self.clear_outputs(index) { + changed.push(cell.to_owned()); + } + } + "stream" => { + let name = content + .get("name") + .and_then(Value::as_str) + .ok_or_else(|| IdeError::invalid("Kernel stream lacks a name"))?; + let text = + text_value(content.get("text").ok_or_else(|| IdeError::invalid("Kernel stream lacks text"))?)?; + if self.clear_pending(index) { + changed.push(cell.to_owned()); + } + let outputs = self.nodes()[index].get("outputs").and_then(Node::array).expect("code outputs"); + let last = outputs.len().checked_sub(1).filter(|last| { + string_field(&outputs[*last], "output_type") == Some("stream") + && string_field(&outputs[*last], "name") == Some(name) + }); + if let Some(last) = last { + if !text.is_empty() { + let old = outputs[last] + .get("text") + .ok_or_else(|| IdeError::invalid("Stored stream output lacks text"))?; + let combined = if old.string().is_some() { None } else { Some(node_text(old)?) }; + let stream = self.cell_mut(index).get_mut("outputs").expect("outputs").item_mut(last); + if let Some(mut combined) = combined { + combined.push_str(&text); + stream.set("text", Some(Node::new(Value::String(combined)))); + } else { + stream.get_mut("text").expect("stream text").append_string(&text); + } + changed.push(cell.to_owned()); + } + } else { + let mut value = content + .as_object() + .cloned() + .ok_or_else(|| IdeError::invalid("Kernel stream content must be an object"))?; + value.insert("output_type".into(), json!("stream")); + value.insert("text".into(), Value::String(text)); + value.remove("transient"); + self.append_output(index, Node::new(Value::Object(value)), None); + changed.push(cell.to_owned()); + } + } + "display_data" | "execute_result" | "error" => { + // Validate before applying a pending clear: malformed messages + // must not erase existing output as a side effect of failure. + let mut value = content + .as_object() + .cloned() + .ok_or_else(|| IdeError::invalid("Kernel output content must be an object"))?; + if event.kind == "error" { + if !value.get("ename").is_some_and(Value::is_string) + || !value.get("evalue").is_some_and(Value::is_string) + || !value + .get("traceback") + .and_then(Value::as_array) + .is_some_and(|lines| lines.iter().all(Value::is_string)) + { + return Err(IdeError::invalid("Kernel error output is malformed")); + } + } else if !value.get("data").is_some_and(Value::is_object) + || !value.get("metadata").is_some_and(Value::is_object) + { + return Err(IdeError::invalid("Kernel display output lacks a MIME bundle or metadata")); + } + if event.kind == "execute_result" + && value.get("execution_count").and_then(Value::as_u64).is_none_or(|n| u32::try_from(n).is_err()) + { + return Err(IdeError::invalid("Kernel execute_result lacks a valid execution_count")); + } + let display_id = display_id(content); + value.remove("transient"); + value.insert("output_type".into(), json!(event.kind)); + if self.clear_pending(index) { + changed.push(cell.to_owned()); + } + self.append_output(index, Node::new(Value::Object(value)), display_id); + changed.push(cell.to_owned()); + } + "update_display_data" => { + let Some(id) = display_id(content) else { + return Ok(Vec::new()); + }; + let data = content + .get("data") + .filter(|v| v.is_object()) + .ok_or_else(|| IdeError::invalid("Display update lacks a MIME bundle"))?; + let metadata = content + .get("metadata") + .filter(|v| v.is_object()) + .ok_or_else(|| IdeError::invalid("Display update lacks metadata"))?; + if self.clear_pending(index) { + changed.push(cell.to_owned()); + } + for current in 0..self.cells.len() { + let matches: Vec = self.cells[current] + .display_ids + .iter() + .enumerate() + .filter_map(|(output, display)| (display.as_deref() == Some(id.as_str())).then_some(output)) + .collect(); + let mut cell_changed = false; + for output in matches { + let Some(old) = self.nodes()[current] + .get("outputs") + .and_then(Node::array) + .and_then(|outputs| outputs.get(output)) + else { + continue; + }; + let data_node = Node::new(data.clone()); + let metadata_node = Node::new(metadata.clone()); + let replace_data = old.get("data").is_none_or(|old| !old.same_bytes(&data_node)); + let replace_metadata = old.get("metadata").is_none_or(|old| !old.same_bytes(&metadata_node)); + if replace_data || replace_metadata { + let output = + self.cell_mut(current).get_mut("outputs").expect("mapped outputs").item_mut(output); + if replace_data { + output.set("data", Some(data_node)); + } + if replace_metadata { + output.set("metadata", Some(metadata_node)); + } + cell_changed = true; + } + } + if cell_changed { + changed.push(self.cells[current].id.clone()); + } + } + } + _ => {} + } + let mut seen = HashSet::new(); + changed.retain(|id| seen.insert(id.clone())); + if !changed.is_empty() { + self.version += 1; + } + Ok(changed) + } + + fn append_output(&mut self, index: usize, node: Node, display_id: Option) { + let len = self.nodes()[index].get("outputs").and_then(Node::array).expect("code outputs").len(); + self.cell_mut(index).get_mut("outputs").expect("outputs").insert(len, node); + self.cells[index].display_ids.push(display_id); + } + + fn clear_pending(&mut self, index: usize) -> bool { + if self.cells[index].pending_clear { self.clear_outputs(index) } else { false } + } + + fn clear_outputs(&mut self, index: usize) -> bool { + self.cells[index].pending_clear = false; + self.cells[index].display_ids.clear(); + let changed = !self.nodes()[index].get("outputs").and_then(Node::array).expect("code outputs").is_empty(); + if changed { + self.cell_mut(index).set("outputs", Some(Node::new(json!([])))); + } + changed + } +} + +fn display_id(content: &Value) -> Option { + content.get("transient")?.get("display_id")?.as_str().map(str::to_owned) +} + +fn node_text(node: &Node) -> IdeResult { + if let Some(text) = node.string() { + return Ok(text.into()); + } + if let Some(parts) = node.array() { + let mut text = String::new(); + for part in parts { + text.push_str( + part.string().ok_or_else(|| IdeError::invalid("Notebook text arrays must contain only strings"))?, + ); + } + return Ok(text); + } + Err(IdeError::invalid("Notebook source/text must be a string or string array")) +} + +fn rewritten_source(old: Option<&Node>, text: String) -> Node { + let Some(old) = old.filter(|node| node.array().is_some()) else { + return Node::new(Value::String(text)); + }; + let parts = old.array().expect("source list"); + let mut first = 0; + let mut start = 0; + while first < parts.len() { + let part = parts[first].string().expect("validated source text"); + if !text[start..].starts_with(part) { + break; + } + start += part.len(); + first += 1; + } + let mut last = parts.len(); + let mut end = text.len(); + while last > first { + let part = parts[last - 1].string().expect("validated source text"); + if end - start < part.len() || !text[start..end].ends_with(part) { + break; + } + end -= part.len(); + last -= 1; + } + let mut updated = old.clone(); + for index in (first..last).rev() { + updated.remove(index); + } + for (offset, part) in text[start..end].split_inclusive('\n').enumerate() { + updated.insert(first + offset, Node::new(Value::String(part.into()))); + } + updated +} + +fn validate_text_node(node: &Node) -> IdeResult<()> { + if node.string().is_some() || node.array().is_some_and(|items| items.iter().all(|part| part.string().is_some())) { + Ok(()) + } else { + Err(IdeError::invalid("Notebook source/text must be a string or string array")) + } +} + +fn text_matches(node: &Node, text: &str) -> bool { + if let Some(old) = node.string() { + return old == text; + } + let Some(parts) = node.array() else { + return false; + }; + let mut remaining = text; + for part in parts { + let Some(part) = part.string() else { + return false; + }; + let Some(rest) = remaining.strip_prefix(part) else { + return false; + }; + remaining = rest; + } + remaining.is_empty() +} + +fn validate_output_node(node: &Node) -> IdeResult<()> { + let kind = string_field(node, "output_type") + .ok_or_else(|| IdeError::invalid("Notebook outputs must be objects with a string output_type"))?; + match kind { + "stream" => { + if string_field(node, "name").is_none() { + return Err(IdeError::invalid("Stored stream output lacks a name")); + } + validate_text_node(node.get("text").ok_or_else(|| IdeError::invalid("Stored stream output lacks text"))?)?; + } + "display_data" | "execute_result" => { + if node.get("data").is_none_or(|n| n.members().is_none()) + || node.get("metadata").is_none_or(|n| n.members().is_none()) + { + return Err(IdeError::invalid("Stored display output requires data and metadata objects")); + } + if kind == "execute_result" && uint_field(node, "execution_count").is_none() { + return Err(IdeError::invalid("Stored execute_result lacks a valid execution_count")); + } + } + "error" + if (string_field(node, "ename").is_none() + || string_field(node, "evalue").is_none() + || !node + .get("traceback") + .and_then(Node::array) + .is_some_and(|lines| lines.iter().all(|line| line.string().is_some()))) => + { + return Err(IdeError::invalid("Stored error output is malformed")); + } + _ => {} + } + Ok(()) +} diff --git a/crates/ide-notebook/src/document/raw.rs b/crates/ide-notebook/src/document/raw.rs new file mode 100644 index 000000000..0aee0a413 --- /dev/null +++ b/crates/ide-notebook/src/document/raw.rs @@ -0,0 +1,408 @@ +//! JSON syntax retained as shared original spans. Mutating one value never +//! normalizes its siblings, property spelling, whitespace, or container order. +use std::{ops::Range, sync::Arc}; + +use ide_core::{IdeError, IdeResult}; +use serde_json::Value; + +#[derive(Clone, Debug)] +struct Raw { + text: Arc, + span: Range, +} + +impl Raw { + fn text(&self) -> &str { + &self.text[self.span.clone()] + } +} + +#[derive(Clone, Debug)] +pub(super) struct Node { + raw: Option, + kind: Kind, +} + +#[derive(Clone, Debug)] +enum Kind { + Scalar(Arc), + Object(Vec, Layout), + Array(Vec, Layout), +} + +#[derive(Clone, Debug)] +pub(super) struct Member { + pub key: String, + head: String, + pub value: Node, +} + +#[derive(Clone, Debug)] +pub(super) struct Layout { + prefix: String, + separators: Vec, + suffix: String, +} + +impl Layout { + fn compact() -> Self { + Self { prefix: String::new(), separators: Vec::new(), suffix: String::new() } + } + + fn separator(&self, index: usize) -> &str { + self.separators.get(index).map(String::as_str).unwrap_or(",") + } + + fn remove(&mut self, index: usize) { + if !self.separators.is_empty() { + self.separators.remove(index.saturating_sub(1).min(self.separators.len() - 1)); + } + } + + fn insert(&mut self, index: usize, previous_len: usize) { + if previous_len > 0 { + let separator = self.separators.first().cloned().unwrap_or_else(|| ",".into()); + self.separators.insert(index.saturating_sub(1).min(self.separators.len()), separator); + } + } +} + +impl Node { + pub fn parse(text: Arc) -> IdeResult<(Self, String, String)> { + // serde supplies strict JSON/UTF-8 string and numeric validation. The + // scanner below only locates the already validated syntax spans. + let start = usize::from(text.starts_with('\u{feff}')) * 3; + serde_json::from_str::(&text[start..]) + .map_err(|e| IdeError::invalid(format!("Invalid notebook JSON: {e}")))?; + let mut parser = Parser { text: text.clone(), pos: start }; + parser.space(); + let begin = parser.pos; + let root = parser.node()?; + Ok((root, text[..begin].into(), text[parser.pos..].into())) + } + + pub fn new(value: Value) -> Self { + match value { + Value::Object(map) => Self { + raw: None, + kind: Kind::Object( + map.into_iter() + .map(|(key, value)| Member { + head: format!("{}:", serde_json::to_string(&key).expect("JSON key")), + key, + value: Self::new(value), + }) + .collect(), + Layout::compact(), + ), + }, + Value::Array(values) => { + Self { raw: None, kind: Kind::Array(values.into_iter().map(Self::new).collect(), Layout::compact()) } + } + value => Self { raw: None, kind: Kind::Scalar(Arc::new(value)) }, + } + } + + pub fn value(&self) -> Value { + match &self.kind { + Kind::Scalar(value) => value.as_ref().clone(), + Kind::Array(items, _) => Value::Array(items.iter().map(Self::value).collect()), + Kind::Object(items, _) => Value::Object(items.iter().map(|m| (m.key.clone(), m.value.value())).collect()), + } + } + + pub fn get(&self, key: &str) -> Option<&Self> { + match &self.kind { + Kind::Object(items, _) => items.iter().find(|m| m.key == key).map(|m| &m.value), + _ => None, + } + } + + pub fn string(&self) -> Option<&str> { + match &self.kind { + Kind::Scalar(value) => value.as_str(), + _ => None, + } + } + + pub fn uint(&self) -> Option { + match &self.kind { + Kind::Scalar(value) => value.as_u64().and_then(|n| u32::try_from(n).ok()), + _ => None, + } + } + + pub fn is_null(&self) -> bool { + matches!(&self.kind, Kind::Scalar(value) if value.is_null()) + } + + pub fn append_string(&mut self, text: &str) { + self.raw = None; + let Kind::Scalar(value) = &mut self.kind else { unreachable!("validated string") }; + let Value::String(current) = Arc::make_mut(value) else { unreachable!("validated string") }; + current.push_str(text); + } + + pub fn get_mut(&mut self, key: &str) -> Option<&mut Self> { + self.raw = None; + match &mut self.kind { + Kind::Object(items, _) => items.iter_mut().find(|m| m.key == key).map(|m| &mut m.value), + _ => None, + } + } + + pub fn members(&self) -> Option<&[Member]> { + match &self.kind { + Kind::Object(items, _) => Some(items), + _ => None, + } + } + + pub fn array(&self) -> Option<&[Self]> { + match &self.kind { + Kind::Array(items, _) => Some(items), + _ => None, + } + } + + pub fn item_mut(&mut self, index: usize) -> &mut Self { + self.raw = None; + match &mut self.kind { + Kind::Array(items, _) => &mut items[index], + _ => unreachable!("validated array"), + } + } + + pub fn layout(&self) -> Layout { + match &self.kind { + Kind::Array(_, layout) | Kind::Object(_, layout) => layout.clone(), + _ => unreachable!("container layout"), + } + } + + pub fn restore_layout(&mut self, layout: Layout) { + self.raw = None; + match &mut self.kind { + Kind::Array(_, old) | Kind::Object(_, old) => *old = layout, + _ => unreachable!("container layout"), + } + } + + pub fn set(&mut self, key: &str, value: Option) { + self.raw = None; + let Kind::Object(items, layout) = &mut self.kind else { unreachable!("validated object") }; + match (items.iter().position(|m| m.key == key), value) { + (Some(index), Some(value)) => items[index].value = value, + (Some(index), None) => { + items.remove(index); + layout.remove(index); + } + (None, Some(value)) => { + layout.insert(items.len(), items.len()); + items.push(Member { + key: key.into(), + head: format!("{}:", serde_json::to_string(key).expect("JSON key")), + value, + }); + } + (None, None) => {} + } + } + + pub fn restore_fields(&mut self, fields: &[(String, Option)], order: &[String], layout: &Layout) { + for (key, field) in fields { + self.set(key, field.as_ref().map(|m| m.value.clone())); + } + let Kind::Object(items, old_layout) = &mut self.kind else { unreachable!("validated object") }; + for (key, field) in fields { + if let (Some(saved), Some(current)) = (field, items.iter_mut().find(|m| &m.key == key)) { + current.head.clone_from(&saved.head); + } + } + items.sort_by_key(|m| order.iter().position(|key| key == &m.key).unwrap_or(usize::MAX)); + if items.len() == order.len() { + *old_layout = layout.clone(); + } + } + + pub fn insert(&mut self, index: usize, node: Self) { + self.raw = None; + let Kind::Array(items, layout) = &mut self.kind else { unreachable!("validated array") }; + layout.insert(index, items.len()); + items.insert(index, node); + } + + pub fn remove(&mut self, index: usize) -> Self { + self.raw = None; + let Kind::Array(items, layout) = &mut self.kind else { unreachable!("validated array") }; + layout.remove(index); + items.remove(index) + } + + pub fn same_bytes(&self, other: &Self) -> bool { + if let (Some(a), Some(b)) = (&self.raw, &other.raw) { + return (Arc::ptr_eq(&a.text, &b.text) && a.span == b.span) || a.text() == b.text(); + } + match (&self.kind, &other.kind) { + (Kind::Scalar(a), Kind::Scalar(b)) => { + if a != b { + return false; + } + match (&self.raw, &other.raw) { + (Some(raw), None) | (None, Some(raw)) => { + raw.text() == serde_json::to_string(a.as_ref()).expect("JSON value") + } + _ => true, + } + } + (Kind::Object(a, al), Kind::Object(b, bl)) => { + a.len() == b.len() + && same_layout(al, bl, a.len()) + && a.iter().zip(b).all(|(a, b)| a.key == b.key && a.head == b.head && a.value.same_bytes(&b.value)) + } + (Kind::Array(a, al), Kind::Array(b, bl)) => { + a.len() == b.len() && same_layout(al, bl, a.len()) && a.iter().zip(b).all(|(a, b)| a.same_bytes(b)) + } + _ => false, + } + } + + pub fn write(&self, output: &mut String) { + if let Some(raw) = &self.raw { + output.push_str(raw.text()); + return; + } + match &self.kind { + Kind::Scalar(value) => output.push_str(&serde_json::to_string(value.as_ref()).expect("JSON value")), + Kind::Object(items, layout) => { + output.push('{'); + output.push_str(&layout.prefix); + for (index, member) in items.iter().enumerate() { + if index > 0 { + output.push_str(layout.separator(index - 1)); + } + output.push_str(&member.head); + member.value.write(output); + } + output.push_str(&layout.suffix); + output.push('}'); + } + Kind::Array(items, layout) => { + output.push('['); + output.push_str(&layout.prefix); + for (index, item) in items.iter().enumerate() { + if index > 0 { + output.push_str(layout.separator(index - 1)); + } + item.write(output); + } + output.push_str(&layout.suffix); + output.push(']'); + } + } + } +} + +fn same_layout(a: &Layout, b: &Layout, count: usize) -> bool { + a.prefix == b.prefix + && a.suffix == b.suffix + && (0..count.saturating_sub(1)).all(|index| a.separator(index) == b.separator(index)) +} + +struct Parser { + text: Arc, + pos: usize, +} + +impl Parser { + fn space(&mut self) { + while self.text.as_bytes().get(self.pos).is_some_and(u8::is_ascii_whitespace) { + self.pos += 1; + } + } + + fn string_end(&mut self) { + self.pos += 1; + loop { + match self.text.as_bytes()[self.pos] { + b'\\' => self.pos += 2, + b'"' => { + self.pos += 1; + break; + } + _ => self.pos += 1, + } + } + } + + fn node(&mut self) -> IdeResult { + let start = self.pos; + let kind = match self.text.as_bytes()[self.pos] { + b'{' | b'[' => { + let object = self.text.as_bytes()[self.pos] == b'{'; + let close = if object { b'}' } else { b']' }; + self.pos += 1; + let gap = self.pos; + self.space(); + let mut layout = + Layout { prefix: self.text[gap..self.pos].into(), separators: Vec::new(), suffix: String::new() }; + let mut members = Vec::new(); + let mut values = Vec::new(); + if self.text.as_bytes()[self.pos] != close { + loop { + if object { + let key_start = self.pos; + self.string_end(); + let key: String = serde_json::from_str(&self.text[key_start..self.pos]) + .map_err(|e| IdeError::invalid(e.to_string()))?; + if members.iter().any(|m: &Member| m.key == key) { + return Err(IdeError::invalid( + "Duplicate JSON object keys are ambiguous; notebook was not opened", + )); + } + self.space(); + self.pos += 1; + self.space(); + let head = self.text[key_start..self.pos].into(); + let value = self.node()?; + members.push(Member { key, head, value }); + } else { + values.push(self.node()?); + } + let end = self.pos; + self.space(); + if self.text.as_bytes()[self.pos] == close { + layout.suffix = self.text[end..self.pos].into(); + break; + } + self.pos += 1; + self.space(); + layout.separators.push(self.text[end..self.pos].into()); + } + } + self.pos += 1; + if object { Kind::Object(members, layout) } else { Kind::Array(values, layout) } + } + b'"' => { + self.string_end(); + Kind::Scalar(Arc::new( + serde_json::from_str(&self.text[start..self.pos]).map_err(|e| IdeError::invalid(e.to_string()))?, + )) + } + _ => { + while self + .text + .as_bytes() + .get(self.pos) + .is_some_and(|b| !b.is_ascii_whitespace() && !b",]}".contains(b)) + { + self.pos += 1; + } + Kind::Scalar(Arc::new( + serde_json::from_str(&self.text[start..self.pos]).map_err(|e| IdeError::invalid(e.to_string()))?, + )) + } + }; + Ok(Node { raw: Some(Raw { text: self.text.clone(), span: start..self.pos }), kind }) + } +} diff --git a/crates/ide-notebook/src/kernel.rs b/crates/ide-notebook/src/kernel.rs new file mode 100644 index 000000000..752990421 --- /dev/null +++ b/crates/ide-notebook/src/kernel.rs @@ -0,0 +1,284 @@ +//! Native, app-owned Jupyter kernels. One supervised process per notebook path. +//! +//! Output routing ends at the matching shell reply AND IOPub idle. Readers remain +//! alive and drain later background output, but never attribute it to another run. + +mod process; +mod transport; + +use ide_core::{CancellationToken, IdeError, IdeResult}; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use specta::Type; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use tokio::sync::Mutex as AsyncMutex; + +pub(super) type StatusCallback = Arc) + Send + Sync>; +pub(super) type OutputCallback = Arc; + +#[derive(Clone, Debug, Serialize, Deserialize, Type)] +#[serde(rename_all = "camelCase")] +pub struct KernelOutput { + pub kind: String, + #[specta(type = specta_typescript::Unknown)] + pub content: Value, + /// Actual Jupyter parent request ID, not a UI request ID. + pub msg_id: String, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, Type)] +#[serde(rename_all = "camelCase")] +pub enum KernelState { + Starting, + Idle, + Busy, + Dead, +} + +impl KernelState { + pub fn as_str(self) -> &'static str { + match self { + Self::Starting => "starting", + Self::Idle => "idle", + Self::Busy => "busy", + Self::Dead => "dead", + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, Type)] +#[serde(rename_all = "camelCase")] +pub struct ExecutionResult { + pub status: String, + pub execution_count: Option, +} + +struct Slot { + operation: AsyncMutex>, + interrupt: Mutex>, + stop: Mutex>, + closed: AtomicBool, +} + +impl Slot { + fn new() -> Self { + Self { + operation: AsyncMutex::new(None), + interrupt: Mutex::new(None), + stop: Mutex::new(None), + closed: AtomicBool::new(false), + } + } + + fn cancel(&self) { + if let Some(token) = self.interrupt.lock().as_ref() { + token.cancel(); + } + } + + fn stop(&self) { + self.cancel(); + if let Some(token) = self.stop.lock().as_ref() { + token.cancel(); + } + } +} + +struct ActiveExecution<'a>(&'a Slot); +impl Drop for ActiveExecution<'_> { + fn drop(&mut self) { + if let Some(token) = self.0.interrupt.lock().take() { + token.cancel(); + } + } +} + +pub struct KernelManager { + runtime_dir: PathBuf, + on_status: StatusCallback, + slots: Mutex>>, + closed: AtomicBool, +} + +impl KernelManager { + pub fn new(runtime_dir: PathBuf, on_status: StatusCallback) -> Self { + Self { runtime_dir, on_status, slots: Mutex::new(HashMap::new()), closed: AtomicBool::new(false) } + } + + pub async fn execute( + &self, + notebook: &Path, + python: &Path, + code: &str, + cancel: CancellationToken, + on_output: OutputCallback, + ) -> IdeResult { + if cancel.is_cancelled() { + return Err(IdeError::cancelled()); + } + let slot = { + let mut slots = self.slots.lock(); + if self.closed.load(Ordering::Acquire) { + return Err(IdeError::unavailable("notebook kernel manager is shut down")); + } + slots.entry(notebook.to_path_buf()).or_insert_with(|| Arc::new(Slot::new())).clone() + }; + // Register cancellation before taking the operation lock: interrupt can + // otherwise race the tiny lock-to-token gap and miss a newly admitted run. + let interrupt = CancellationToken::new(); + { + let mut active = slot.interrupt.lock(); + if active.is_some() { + return Err(IdeError::unavailable("notebook kernel is busy")); + } + *active = Some(interrupt.clone()); + } + let _active = ActiveExecution(&slot); + // A second caller must not sit in a queue and execute stale source later. + let mut running = slot.operation.try_lock().map_err(|_| IdeError::unavailable("notebook kernel is busy"))?; + if slot.closed.load(Ordering::Acquire) { + return Err(IdeError::unavailable("notebook kernel was closed")); + } + if running.as_ref().is_some_and(|r| r.python != python || r.is_finished()) + && let Some(old) = running.take() + { + old.shutdown().await; + } + if cancel.is_cancelled() || interrupt.is_cancelled() || slot.closed.load(Ordering::Acquire) { + return Err(IdeError::cancelled()); + } + if running.is_none() { + let (new, ready) = transport::Running::spawn( + self.runtime_dir.clone(), + notebook.to_path_buf(), + python.to_path_buf(), + self.on_status.clone(), + ); + *slot.stop.lock() = Some(new.stop.clone()); + *running = Some(new); + if slot.closed.load(Ordering::Acquire) { + slot.stop(); + } + let result = tokio::select! { + biased; + _ = cancel.cancelled() => Err(IdeError::cancelled()), + _ = interrupt.cancelled() => Err(IdeError::cancelled()), + result = ready => result.unwrap_or_else(|_| Err(IdeError::unavailable("kernel startup task ended"))), + }; + if let Err(error) = result { + if let Some(failed) = running.take() { + failed.shutdown().await; + } + slot.stop.lock().take(); + return Err(error); + } + } + let result = running.as_ref().expect("initialized kernel").execute(code, cancel, interrupt, on_output).await; + if running.as_ref().is_some_and(transport::Running::is_finished) { + if let Some(dead) = running.take() { + dead.shutdown().await; + } + slot.stop.lock().take(); + } + result + } + + /// Cancels the active request, including startup. The execution future drains + /// reply+idle before returning, or stops the owned process on interrupt timeout. + pub async fn interrupt(&self, notebook: &Path) -> IdeResult<()> { + let slot = self.slots.lock().get(notebook).cloned(); + if let Some(slot) = slot { + slot.cancel(); + // Wait for native interrupt delivery/drain, not just token mutation. + let _drained = slot.operation.lock().await; + } + Ok(()) + } + + pub async fn restart(&self, notebook: &Path) -> IdeResult<()> { + let slot = self.slots.lock().get(notebook).cloned(); + let Some(slot) = slot else { return Ok(()) }; + slot.stop(); + let mut running = slot.operation.lock().await; + if slot.closed.load(Ordering::Acquire) { + return Err(IdeError::unavailable("notebook kernel was closed")); + } + let Some(old) = running.take() else { return Ok(()) }; + let python = old.python.clone(); + old.shutdown().await; + let (new, ready) = + transport::Running::spawn(self.runtime_dir.clone(), notebook.to_path_buf(), python, self.on_status.clone()); + *slot.stop.lock() = Some(new.stop.clone()); + *running = Some(new); + if slot.closed.load(Ordering::Acquire) { + slot.stop(); + } + let result = ready.await.unwrap_or_else(|_| Err(IdeError::unavailable("kernel restart task ended"))); + if result.is_err() { + if let Some(failed) = running.take() { + failed.shutdown().await; + } + slot.stop.lock().take(); + } + result + } + + pub async fn shutdown_kernel(&self, notebook: &Path) -> IdeResult<()> { + let slot = { + let slots = self.slots.lock(); + // Keep a closed tombstone until old readers/process have drained. + // Reopening this path cannot overlap old-generation status events. + let slot = slots.get(notebook).cloned(); + if let Some(slot) = &slot { + slot.closed.store(true, Ordering::Release); + slot.stop(); + } + slot + }; + if let Some(slot) = slot { + if let Some(running) = slot.operation.lock().await.take() { + running.shutdown().await; + } + slot.stop.lock().take(); + let mut slots = self.slots.lock(); + if slots.get(notebook).is_some_and(|current| Arc::ptr_eq(current, &slot)) { + slots.remove(notebook); + } + } + Ok(()) + } + + pub async fn shutdown(&self) { + let slots = { + let mut slots = self.slots.lock(); + self.closed.store(true, Ordering::Release); + slots + .drain() + .map(|(_, slot)| { + slot.closed.store(true, Ordering::Release); + slot.stop(); + slot + }) + .collect::>() + }; + // All stop tokens are signalled before awaiting any process. + for slot in slots { + if let Some(running) = slot.operation.lock().await.take() { + running.shutdown().await; + } + } + } +} + +impl Drop for KernelManager { + fn drop(&mut self) { + for slot in self.slots.get_mut().values() { + slot.closed.store(true, Ordering::Release); + slot.stop(); + } + } +} diff --git a/crates/ide-notebook/src/kernel/process.rs b/crates/ide-notebook/src/kernel/process.rs new file mode 100644 index 000000000..791939e0b --- /dev/null +++ b/crates/ide-notebook/src/kernel/process.rs @@ -0,0 +1,444 @@ +use ide_core::{IdeError, IdeResult}; +use jupyter_protocol::ConnectionInfo; +use parking_lot::Mutex; +use std::collections::VecDeque; +use std::io::Write; +use std::path::Path; +use std::process::Stdio; +use std::sync::Arc; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::process::{Child, Command}; +use tokio::task::JoinHandle; +use uuid::Uuid; + +const STDERR_BYTES: usize = 16 * 1024; + +pub(super) struct Process { + child: Child, + #[cfg(unix)] + pid: Option, + stderr: Arc>>, + drains: Vec>, + // Kept until the process and sockets have stopped; contains the HMAC secret. + _directory: tempfile::TempDir, + #[cfg(windows)] + windows: windows::Resources, +} + +impl Process { + pub async fn spawn(runtime: &Path, notebook: &Path, python: &Path) -> IdeResult<(Self, ConnectionInfo)> { + if !python.is_absolute() { + return Err(IdeError::invalid("select an absolute Python executable path for this notebook")); + } + std::fs::create_dir_all(runtime)?; + let directory = tempfile::Builder::new().prefix("kernel-").tempdir_in(runtime)?; + make_private(directory.path())?; + let ip = std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST); + let (ports, listeners) = jupyter_zmq_client::peek_ports_with_listeners(ip, 5).await.map_err(protocol_error)?; + let info = ConnectionInfo { + transport: jupyter_protocol::connection_info::Transport::TCP, + ip: ip.to_string(), + stdin_port: ports[0], + control_port: ports[1], + hb_port: ports[2], + shell_port: ports[3], + iopub_port: ports[4], + signature_scheme: "hmac-sha256".into(), + key: format!("{}{}", Uuid::new_v4().simple(), Uuid::new_v4().simple()), + kernel_name: Some("openidea-python".into()), + }; + let connection = directory.path().join("connection.json"); + write_private(&connection, &serde_json::to_vec(&info).map_err(protocol_error)?)?; + #[cfg(windows)] + let windows = windows::Resources::new()?; + let mut argv = vec![python.as_os_str().to_os_string()]; + #[cfg(not(windows))] + argv.extend(["-m".into(), "ipykernel_launcher".into(), "-f".into(), "{connection_file}".into()]); + #[cfg(windows)] + argv.extend(["-c".into(), windows::BOOTSTRAP.into(), "-f".into(), "{connection_file}".into()]); + // App-owned spec, not registered in user/global Jupyter directories. The + // argv array is also the launch source of truth, without shell quoting. + let spec = serde_json::json!({ + "argv": argv.iter().map(|v| v.to_string_lossy().into_owned()).collect::>(), + "display_name": "OpenIDEA Python", "language": "python", "interrupt_mode": "signal" + }); + write_private(&directory.path().join("kernel.json"), &serde_json::to_vec(&spec).map_err(protocol_error)?)?; + let mut command = Command::new(python); + for arg in &argv[1..] { + if arg == "{connection_file}" { + command.arg(&connection); + } else { + command.arg(arg); + } + } + command + .current_dir(notebook.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or_else(|| Path::new("."))) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .env_remove("JPY_INTERRUPT_EVENT") + .env_remove("JPY_PARENT_PID") + .env("PYTHONUNBUFFERED", "1"); + ide_core::process::prepare(command.as_std_mut()); + #[cfg(unix)] + { + use std::os::unix::process::CommandExt; + command.as_std_mut().process_group(0); + command.env("JPY_PARENT_PID", std::process::id().to_string()); + } + #[cfg(windows)] + windows.configure(&mut command); + let child = command.spawn().map_err(|e| { + IdeError::unavailable(format!( + "cannot start selected Python {}: {e}. Select an executable environment with ipykernel installed", + python.display() + )) + })?; + drop(listeners); + #[cfg(unix)] + let pid = child.id(); + let mut process = Self { + child, + #[cfg(unix)] + pid, + stderr: Arc::new(Mutex::new(VecDeque::with_capacity(STDERR_BYTES))), + drains: Vec::new(), + _directory: directory, + #[cfg(windows)] + windows, + }; + if let Some(mut stderr) = process.child.stderr.take() { + let tail = process.stderr.clone(); + process.drains.push(tokio::spawn(async move { + let mut buffer = [0_u8; 4096]; + while let Ok(count) = stderr.read(&mut buffer).await { + if count == 0 { + break; + } + let mut tail = tail.lock(); + for byte in &buffer[..count] { + if tail.len() == STDERR_BYTES { + tail.pop_front(); + } + tail.push_back(*byte); + } + } + })); + } + if let Some(mut stdout) = process.child.stdout.take() { + process.drains.push(tokio::spawn(async move { + let mut buffer = [0_u8; 4096]; + while let Ok(count) = stdout.read(&mut buffer).await { + if count == 0 { + break; + } + } + })); + } + #[cfg(windows)] + process.windows.assign_and_release(&process.child)?; + Ok((process, info)) + } + + pub fn stderr(&self) -> String { + let bytes = self.stderr.lock().iter().copied().collect::>(); + String::from_utf8_lossy(&bytes).into_owned() + } + + pub fn death(&mut self) -> IdeResult> { + match self.child.try_wait()? { + None => Ok(None), + Some(status) => { + self.kill_tree(); + Ok(Some(format!("kernel exited ({status}); {}", self.stderr()))) + } + } + } + + pub fn interrupt(&self) -> IdeResult<()> { + #[cfg(unix)] + { + let pid = self.pid.ok_or_else(|| IdeError::unavailable("kernel is no longer running"))?; + // SAFETY: this child was created in its own process group. Never + // signal the host group or a discovered/external kernel. + if unsafe { libc::kill(-(pid as i32), libc::SIGINT) } != 0 { + return Err(std::io::Error::last_os_error().into()); + } + Ok(()) + } + #[cfg(windows)] + { + self.windows.interrupt() + } + #[cfg(not(any(unix, windows)))] + { + Err(IdeError::unavailable("native kernel interruption is unsupported on this platform")) + } + } + + fn kill_tree(&mut self) { + #[cfg(unix)] + if let Some(pid) = self.pid.take() { + // SAFETY: group ID belongs exclusively to the owned kernel. + unsafe { + libc::kill(-(pid as i32), libc::SIGKILL); + } + } + #[cfg(windows)] + self.windows.kill_tree(); + } + + pub async fn shutdown(&mut self) { + self.kill_tree(); + let _ = self.child.start_kill(); + let _ = tokio::time::timeout(Duration::from_secs(3), self.child.wait()).await; + for task in self.drains.drain(..) { + task.abort(); + } + } +} + +impl Drop for Process { + fn drop(&mut self) { + self.kill_tree(); + let _ = self.child.start_kill(); + for task in &self.drains { + task.abort(); + } + // tokio Child's kill_on_drop also kills/reaps if this future is aborted. + } +} + +fn protocol_error(error: impl std::fmt::Display) -> IdeError { + IdeError::unavailable(format!("Jupyter kernel setup: {error}")) +} + +fn write_private(path: &Path, bytes: &[u8]) -> IdeResult<()> { + let mut options = std::fs::OpenOptions::new(); + options.write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600); + } + let mut file = options.open(path)?; + file.write_all(bytes)?; + Ok(()) +} + +fn make_private(path: &Path) -> IdeResult<()> { + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))?; + } + #[cfg(windows)] + windows::private_directory(path)?; + Ok(()) +} + +#[cfg(windows)] +mod windows { + use super::*; + use std::os::windows::ffi::OsStrExt; + use windows_sys::Win32::Foundation::{CloseHandle, ERROR_ALREADY_EXISTS, GetLastError, HANDLE, LocalFree}; + use windows_sys::Win32::Security::Authorization::{ + ConvertStringSecurityDescriptorToSecurityDescriptorW, SDDL_REVISION_1, SE_FILE_OBJECT, SetNamedSecurityInfoW, + }; + use windows_sys::Win32::Security::{ + DACL_SECURITY_INFORMATION, GetSecurityDescriptorDacl, PROTECTED_DACL_SECURITY_INFORMATION, SECURITY_ATTRIBUTES, + }; + use windows_sys::Win32::System::JobObjects::{ + AssignProcessToJobObject, CreateJobObjectW, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JobObjectExtendedLimitInformation, SetInformationJobObject, + TerminateJobObject, + }; + use windows_sys::Win32::System::Threading::{CreateEventW, SetEvent}; + + // Launcher glue only. All Jupyter traffic is authenticated native ZeroMQ. + // Child opens its own handles: no assumption about Rust HANDLE inheritance. + // The gate prevents any ipykernel imports/descendants before job assignment. + pub const BOOTSTRAP: &str = r#"import ctypes, os, runpy, sys +k = ctypes.WinDLL('kernel32', use_last_error=True) +k.OpenEventW.argtypes = [ctypes.c_uint32, ctypes.c_int, ctypes.c_wchar_p] +k.OpenEventW.restype = ctypes.c_void_p +k.WaitForSingleObject.argtypes = [ctypes.c_void_p, ctypes.c_uint32] +k.WaitForSingleObject.restype = ctypes.c_uint32 +k.CloseHandle.argtypes = [ctypes.c_void_p] +k.CloseHandle.restype = ctypes.c_int +interrupt = k.OpenEventW(0x00100002, 0, os.environ.pop('OPENIDEA_INTERRUPT_EVENT')) +gate = k.OpenEventW(0x00100000, 0, os.environ.pop('OPENIDEA_LAUNCH_EVENT')) +if not interrupt or not gate: + raise ctypes.WinError(ctypes.get_last_error()) +if k.WaitForSingleObject(gate, 30000) != 0: + raise RuntimeError('OpenIDEA kernel launch gate failed') +k.CloseHandle(gate) +os.environ['JPY_INTERRUPT_EVENT'] = str(interrupt) +try: + runpy.run_module('ipykernel_launcher', run_name='__main__') +finally: + k.CloseHandle(interrupt) +"#; + + struct Handle(usize); + impl Handle { + fn new(raw: HANDLE) -> IdeResult { + if raw.is_null() { Err(std::io::Error::last_os_error().into()) } else { Ok(Self(raw as usize)) } + } + fn raw(&self) -> HANDLE { + self.0 as HANDLE + } + fn signal(&self) -> IdeResult<()> { + // SAFETY: owned live event handle, retained for this process lifetime. + if unsafe { SetEvent(self.raw()) } == 0 { Err(std::io::Error::last_os_error().into()) } else { Ok(()) } + } + } + impl Drop for Handle { + fn drop(&mut self) { + unsafe { + CloseHandle(self.raw()); + } + } + } + + struct Descriptor(*mut std::ffi::c_void); + impl Descriptor { + fn new() -> IdeResult { + // Protected ACL: owner and SYSTEM only, inherited by runtime files. + let sddl = "D:P(A;OICI;FA;;;OW)(A;OICI;FA;;;SY)\0".encode_utf16().collect::>(); + let mut raw = std::ptr::null_mut(); + if unsafe { + ConvertStringSecurityDescriptorToSecurityDescriptorW( + sddl.as_ptr(), + SDDL_REVISION_1, + &mut raw, + std::ptr::null_mut(), + ) + } == 0 + { + return Err(std::io::Error::last_os_error().into()); + } + Ok(Self(raw)) + } + } + impl Drop for Descriptor { + fn drop(&mut self) { + unsafe { + LocalFree(self.0); + } + } + } + + pub fn private_directory(path: &Path) -> IdeResult<()> { + let descriptor = Descriptor::new()?; + let mut present = 0; + let mut defaulted = 0; + let mut dacl = std::ptr::null_mut(); + if unsafe { GetSecurityDescriptorDacl(descriptor.0, &mut present, &mut dacl, &mut defaulted) } == 0 + || present == 0 + { + return Err(std::io::Error::last_os_error().into()); + } + let path = path.as_os_str().encode_wide().chain(Some(0)).collect::>(); + let error = unsafe { + SetNamedSecurityInfoW( + path.as_ptr(), + SE_FILE_OBJECT, + DACL_SECURITY_INFORMATION | PROTECTED_DACL_SECURITY_INFORMATION, + std::ptr::null_mut(), + std::ptr::null_mut(), + dacl, + std::ptr::null_mut(), + ) + }; + if error != 0 { + return Err(std::io::Error::from_raw_os_error(error as i32).into()); + } + Ok(()) + } + + pub struct Resources { + interrupt: Handle, + gate: Handle, + job: Handle, + interrupt_name: String, + gate_name: String, + } + impl Resources { + pub fn new() -> IdeResult { + let interrupt_name = format!("Local\\OpenIDEA-interrupt-{}", Uuid::new_v4()); + let gate_name = format!("Local\\OpenIDEA-launch-{}", Uuid::new_v4()); + fn event(name: &str) -> IdeResult { + let name = name.encode_utf16().chain(Some(0)).collect::>(); + // Event-specific rights; owner and SYSTEM only. File inheritance + // flags/FA are not used for kernel synchronization objects. + let sddl = "D:P(A;;GA;;;OW)(A;;GA;;;SY)\0".encode_utf16().collect::>(); + let mut raw = std::ptr::null_mut(); + if unsafe { + ConvertStringSecurityDescriptorToSecurityDescriptorW( + sddl.as_ptr(), + SDDL_REVISION_1, + &mut raw, + std::ptr::null_mut(), + ) + } == 0 + { + return Err(std::io::Error::last_os_error().into()); + } + let descriptor = Descriptor(raw); + let attrs = SECURITY_ATTRIBUTES { + nLength: std::mem::size_of::() as u32, + lpSecurityDescriptor: descriptor.0, + bInheritHandle: 0, + }; + let handle = Handle::new(unsafe { CreateEventW(&attrs, 0, 0, name.as_ptr()) })?; + if unsafe { GetLastError() } == ERROR_ALREADY_EXISTS { + return Err(IdeError::unavailable("kernel event name collision")); + } + Ok(handle) + } + let interrupt = event(&interrupt_name)?; + let gate = event(&gate_name)?; + let job = Handle::new(unsafe { CreateJobObjectW(std::ptr::null(), std::ptr::null()) })?; + let mut info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); + info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + if unsafe { + SetInformationJobObject( + job.raw(), + JobObjectExtendedLimitInformation, + &info as *const _ as *const _, + std::mem::size_of_val(&info) as u32, + ) + } == 0 + { + return Err(std::io::Error::last_os_error().into()); + } + Ok(Self { interrupt, gate, job, interrupt_name, gate_name }) + } + pub fn configure(&self, command: &mut Command) { + command.env("OPENIDEA_INTERRUPT_EVENT", &self.interrupt_name).env("OPENIDEA_LAUNCH_EVENT", &self.gate_name); + } + pub fn assign_and_release(&self, child: &Child) -> IdeResult<()> { + let handle = child + .raw_handle() + .ok_or_else(|| IdeError::unavailable("kernel exited before Windows job assignment"))?; + if unsafe { AssignProcessToJobObject(self.job.raw(), handle as HANDLE) } == 0 { + return Err(IdeError::unavailable(format!( + "cannot contain kernel process in owned Windows job: {}", + std::io::Error::last_os_error() + ))); + } + self.gate.signal() + } + pub fn interrupt(&self) -> IdeResult<()> { + self.interrupt.signal() + } + pub fn kill_tree(&self) { + unsafe { + TerminateJobObject(self.job.raw(), 1); + } + } + } +} diff --git a/crates/ide-notebook/src/kernel/transport.rs b/crates/ide-notebook/src/kernel/transport.rs new file mode 100644 index 000000000..46a814b8b --- /dev/null +++ b/crates/ide-notebook/src/kernel/transport.rs @@ -0,0 +1,555 @@ +use super::process::Process; +use super::{ExecutionResult, KernelOutput, KernelState, OutputCallback, StatusCallback}; +use ide_core::{CancellationToken, IdeError, IdeResult}; +use jupyter_protocol::{ExecuteRequest, JupyterMessage, KernelInfoRequest, ShutdownRequest}; +use jupyter_zmq_client::{ClientControlConnection, Connection, DealerSendConnection, RawMessage}; +use parking_lot::Mutex; +use serde_json::Value; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot}; +use tokio::task::JoinHandle; +use tokio::time::{Instant, MissedTickBehavior}; +use uuid::Uuid; +use zeromq::SocketRecv; + +const STARTUP_TIMEOUT: Duration = Duration::from_secs(30); +const INTERRUPT_TIMEOUT: Duration = Duration::from_secs(5); +const SEND_TIMEOUT: Duration = Duration::from_secs(3); +const MAX_MESSAGE_BYTES: usize = 8 * 1024 * 1024; +const MAX_OUTPUT_BYTES: usize = 32 * 1024 * 1024; +const MAX_OUTPUT_MESSAGES: usize = 10_000; +const QUEUED_MESSAGES: usize = 64; +const QUEUED_BYTES: usize = 16 * 1024 * 1024; + +struct Admission(Arc); +impl Drop for Admission { + fn drop(&mut self) { + self.0.store(false, Ordering::Release); + } +} + +struct Execution { + code: String, + cancel: CancellationToken, + interrupt: CancellationToken, + output: OutputCallback, + reply: oneshot::Sender>, + _admission: Admission, +} + +pub(super) struct Running { + pub python: PathBuf, + pub stop: CancellationToken, + requests: mpsc::Sender, + task: Option>, + busy: Arc, +} + +impl Running { + pub fn spawn( + runtime: PathBuf, + notebook: PathBuf, + python: PathBuf, + status: StatusCallback, + ) -> (Self, oneshot::Receiver>) { + let stop = CancellationToken::new(); + let (requests, receiver) = mpsc::channel(1); + let (ready, ready_rx) = oneshot::channel(); + let task = tokio::spawn(supervise(runtime, notebook, python.clone(), status, stop.clone(), receiver, ready)); + (Self { python, stop, requests, task: Some(task), busy: Arc::new(AtomicBool::new(false)) }, ready_rx) + } + + pub fn is_finished(&self) -> bool { + self.task.as_ref().is_none_or(JoinHandle::is_finished) + } + + pub async fn execute( + &self, + code: &str, + cancel: CancellationToken, + interrupt: CancellationToken, + output: OutputCallback, + ) -> IdeResult { + if self.busy.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire).is_err() { + return Err(IdeError::unavailable("previous kernel execution has not drained")); + } + let (reply, result) = oneshot::channel(); + let execution = Execution { + code: code.to_owned(), + cancel, + interrupt, + output, + reply, + _admission: Admission(self.busy.clone()), + }; + self.requests.try_send(execution).map_err(|_| IdeError::unavailable("kernel is busy or has exited"))?; + result.await.unwrap_or_else(|_| Err(IdeError::unavailable("kernel execution task ended"))) + } + + pub async fn shutdown(mut self) { + self.stop.cancel(); + if let Some(mut task) = self.task.take() + && tokio::time::timeout(Duration::from_secs(8), &mut task).await.is_err() + { + task.abort(); + let _ = task.await; + } + } +} + +impl Drop for Running { + fn drop(&mut self) { + self.stop.cancel(); + // Do not abort: supervisor owns graceful shutdown and process reaping. + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Channel { + Shell, + Iopub, +} + +struct Frame { + channel: Channel, + parent: String, + kind: String, + content: Value, + bytes: usize, + _permit: Option, +} + +#[derive(Clone)] +struct Failure { + token: CancellationToken, + error: Arc>>, +} +impl Failure { + fn new() -> Self { + Self { token: CancellationToken::new(), error: Arc::new(Mutex::new(None)) } + } + fn fail(&self, error: IdeError) { + self.error.lock().get_or_insert(error); + self.token.cancel(); + } + fn error(&self) -> IdeError { + self.error.lock().clone().unwrap_or_else(|| IdeError::unavailable("kernel reader stopped")) + } +} + +struct Readers(Vec>); +impl Drop for Readers { + fn drop(&mut self) { + for task in &self.0 { + task.abort(); + } + } +} + +fn reader( + mut connection: Connection, + channel: Channel, + tx: mpsc::Sender, + failure: Failure, + queued: Arc, +) -> JoinHandle<()> { + tokio::spawn(async move { + loop { + let result = receive(&mut connection, channel, &queued).await; + match result { + Ok(frame) => { + if tx.try_send(frame).is_err() { + failure.fail(IdeError::unavailable( + "kernel output queue overflow; execution stopped to bound memory", + )); + break; + } + } + Err(error) => { + failure.fail(error); + break; + } + } + } + }) +} + +async fn receive( + connection: &mut Connection, + channel: Channel, + queued: &Arc, +) -> IdeResult { + let multipart = connection.socket.recv().await.map_err(protocol)?; + let bytes = multipart + .iter() + .try_fold(0_usize, |sum, part| sum.checked_add(part.len())) + .ok_or_else(|| IdeError::unavailable("kernel message size overflow"))?; + if bytes > MAX_MESSAGE_BYTES { + return Err(IdeError::unavailable("kernel message exceeds 8 MiB limit; execution stopped")); + } + let permit = queued + .clone() + .try_acquire_many_owned(bytes as u32) + .map_err(|_| IdeError::unavailable("kernel output queue exceeds 16 MiB; execution stopped"))?; + // Upstream RawMessage indexes four signed parts without checking length. + // Validate framing before calling it, including messages from a bad kernel. + let delimiter = multipart + .iter() + .position(|part| part.as_ref() == b"") + .ok_or_else(|| protocol("missing Jupyter message delimiter"))?; + if multipart.len() < delimiter + 6 { + return Err(protocol("incomplete Jupyter message frames")); + } + let raw = RawMessage::from_multipart(multipart, &connection.mac).map_err(protocol)?; + let header: Value = serde_json::from_slice(&raw.jparts[0]).map_err(protocol)?; + let parent: Value = serde_json::from_slice(&raw.jparts[1]).map_err(protocol)?; + let kind = header + .get("msg_type") + .and_then(Value::as_str) + .ok_or_else(|| protocol("missing Jupyter message type"))? + .to_owned(); + let parent = parent.get("msg_id").and_then(Value::as_str).unwrap_or("").to_owned(); + let content = serde_json::from_slice(&raw.jparts[3]).map_err(protocol)?; + Ok(Frame { channel, parent, kind, content, bytes, _permit: Some(permit) }) +} + +fn protocol(error: impl std::fmt::Display) -> IdeError { + IdeError::new(ide_core::error::ErrorKind::Protocol, format!("Jupyter protocol: {error}")) +} + +struct Session { + process: Process, + shell: DealerSendConnection, + control: ClientControlConnection, + frames: mpsc::Receiver, + failure: Failure, + _readers: Readers, +} + +impl Session { + async fn connect( + runtime: &std::path::Path, + notebook: &std::path::Path, + python: &std::path::Path, + ) -> IdeResult { + let (process, info) = Process::spawn(runtime, notebook, python).await?; + let session = Uuid::new_v4().to_string(); + let identity = jupyter_zmq_client::peer_identity_for_session(&session).map_err(protocol)?; + let (iopub, shell, control) = tokio::try_join!( + jupyter_zmq_client::create_client_iopub_connection(&info, "", &session), + jupyter_zmq_client::create_client_shell_connection_with_identity(&info, &session, identity), + jupyter_zmq_client::create_client_control_connection(&info, &session), + ) + .map_err(protocol)?; + let (shell, shell_recv) = shell.split(); + let (tx, frames) = mpsc::channel(QUEUED_MESSAGES); + let failure = Failure::new(); + let queued = Arc::new(Semaphore::new(QUEUED_BYTES)); + let readers = Readers(vec![ + reader(iopub, Channel::Iopub, tx.clone(), failure.clone(), queued.clone()), + reader(shell_recv, Channel::Shell, tx, failure.clone(), queued), + ]); + Ok(Self { process, shell, control, frames, failure, _readers: readers }) + } + + async fn ready(&mut self, stop: &CancellationToken) -> IdeResult<()> { + let deadline = Instant::now() + STARTUP_TIMEOUT; + let mut probe = tokio::time::interval(Duration::from_millis(250)); + probe.set_missed_tick_behavior(MissedTickBehavior::Delay); + let mut request = String::new(); + let mut reply = false; + let mut idle = false; + loop { + tokio::select! { + biased; + _ = stop.cancelled() => return Err(IdeError::cancelled()), + _ = self.failure.token.cancelled() => return Err(self.failure.error()), + _ = tokio::time::sleep_until(deadline) => return Err(IdeError::unavailable(format!( + "selected Python kernel was not ready within 30 seconds; {}", self.process.stderr()))), + _ = probe.tick() => { + if let Some(message) = self.process.death()? { + return Err(IdeError::unavailable(message)); + } + // Only readiness probes are retried, never user execution. + // Once the shell replies, retain its ID to await IOPub idle. + if !reply { + let message: JupyterMessage = KernelInfoRequest {}.into(); + request = message.header.msg_id.clone(); + idle = false; + send(&mut self.shell, message).await?; + } else if !idle { + // A slow subscriber can miss the first idle. Probe again + // after a reply, proving BOTH live shell and subscription. + reply = false; + } + }, + frame = self.frames.recv() => { + let frame = frame.ok_or_else(|| IdeError::unavailable("kernel readers closed during startup"))?; + if frame.parent != request { continue; } + if frame.channel == Channel::Shell && frame.kind == "kernel_info_reply" { reply = true; } + if frame.channel == Channel::Iopub && is_idle(&frame) { idle = true; } + if reply && idle { return Ok(()); } + } + } + } + } + + async fn run( + &mut self, + execution: &mut Execution, + stop: &CancellationToken, + notebook: &std::path::Path, + status: &StatusCallback, + ) -> IdeResult { + if execution.cancel.is_cancelled() || execution.interrupt.is_cancelled() || execution.reply.is_closed() { + return Err(IdeError::cancelled()); + } + let request: JupyterMessage = ExecuteRequest { + code: std::mem::take(&mut execution.code), + allow_stdin: false, + stop_on_error: true, + ..Default::default() + } + .into(); + let request_id = request.header.msg_id.clone(); + // If sending fails/times out, delivery is uncertain. Supervisor closes + // the process; no caller may retry user code automatically. + send(&mut self.shell, request).await?; + status(notebook, KernelState::Busy, None); + let mut completion = Completion::default(); + let mut budget = Budget::default(); + let mut interrupted = false; + let mut interrupt_deadline = None; + let mut death = tokio::time::interval(Duration::from_millis(100)); + death.set_missed_tick_behavior(MissedTickBehavior::Delay); + loop { + tokio::select! { + biased; + _ = stop.cancelled() => return Err(IdeError::cancelled()), + _ = self.failure.token.cancelled() => return Err(self.failure.error()), + _ = execution.cancel.cancelled(), if !interrupted => { + self.process.interrupt()?; + interrupted = true; + interrupt_deadline = Some(Instant::now() + INTERRUPT_TIMEOUT); + }, + _ = execution.interrupt.cancelled(), if !interrupted => { + self.process.interrupt()?; + interrupted = true; + interrupt_deadline = Some(Instant::now() + INTERRUPT_TIMEOUT); + }, + _ = death.tick() => { + if execution.reply.is_closed() && !interrupted { + self.process.interrupt()?; + interrupted = true; + interrupt_deadline = Some(Instant::now() + INTERRUPT_TIMEOUT); + } + if interrupt_deadline.is_some_and(|deadline| Instant::now() >= deadline) { + return Err(IdeError::unavailable("kernel did not drain within 5 seconds of interrupt; owned kernel stopped")); + } + if let Some(message) = self.process.death()? { return Err(IdeError::unavailable(message)); } + }, + frame = self.frames.recv() => { + let frame = frame.ok_or_else(|| IdeError::unavailable("kernel readers closed during execution"))?; + if frame.parent != request_id { continue; } + budget.add(frame.bytes)?; + completion.observe(&frame, &request_id)?; + if frame.channel == Channel::Iopub && is_output(&frame.kind) { + (execution.output)(KernelOutput { kind: frame.kind, content: frame.content, msg_id: frame.parent }); + } + if completion.complete() { + status(notebook, KernelState::Idle, None); + if interrupted || execution.cancel.is_cancelled() || execution.interrupt.is_cancelled() { + return Err(IdeError::cancelled()); + } + return Ok(completion.reply.take().expect("complete reply")); + } + } + } + } + } + + async fn shutdown(&mut self) { + // Stop is separate from user interrupt: graceful shutdown, then kill only + // our process group/job. It must remain bounded even for wedged Python. + let request: JupyterMessage = ShutdownRequest { restart: false }.into(); + let _ = tokio::time::timeout(Duration::from_millis(500), self.control.send(request)).await; + let deadline = Instant::now() + Duration::from_millis(750); + while Instant::now() < deadline { + match self.process.death() { + Ok(None) => {} + _ => break, + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + self.process.shutdown().await; + } +} + +async fn send(connection: &mut DealerSendConnection, message: JupyterMessage) -> IdeResult<()> { + tokio::time::timeout(SEND_TIMEOUT, connection.send(message)) + .await + .map_err(|_| IdeError::unavailable("kernel request send timed out; delivery uncertain, not retried"))? + .map_err(protocol) +} + +fn is_idle(frame: &Frame) -> bool { + frame.kind == "status" && frame.content.get("execution_state").and_then(Value::as_str) == Some("idle") +} + +fn is_output(kind: &str) -> bool { + matches!( + kind, + "stream" + | "display_data" + | "execute_result" + | "error" + | "execute_input" + | "clear_output" + | "update_display_data" + ) +} + +#[derive(Default)] +struct Completion { + reply: Option, + idle: bool, +} +impl Completion { + fn observe(&mut self, frame: &Frame, request_id: &str) -> IdeResult<()> { + if frame.parent != request_id { + return Ok(()); + } + if frame.channel == Channel::Iopub && is_idle(frame) { + self.idle = true; + } + if frame.channel == Channel::Shell && frame.kind == "execute_reply" { + let status = frame + .content + .get("status") + .and_then(Value::as_str) + .ok_or_else(|| protocol("execute_reply missing status"))?; + if !matches!(status, "ok" | "error" | "abort" | "aborted") { + return Err(protocol("invalid execute_reply status")); + } + let count = match frame.content.get("execution_count") { + None | Some(Value::Null) => None, + Some(value) => Some( + value + .as_u64() + .and_then(|n| u32::try_from(n).ok()) + .ok_or_else(|| protocol("invalid execution count"))?, + ), + }; + self.reply = Some(ExecutionResult { status: status.to_owned(), execution_count: count }); + } + Ok(()) + } + fn complete(&self) -> bool { + self.idle && self.reply.is_some() + } +} + +#[derive(Default)] +struct Budget { + bytes: usize, + messages: usize, +} +impl Budget { + fn add(&mut self, bytes: usize) -> IdeResult<()> { + self.bytes = self.bytes.saturating_add(bytes); + self.messages = self.messages.saturating_add(1); + if self.bytes > MAX_OUTPUT_BYTES || self.messages > MAX_OUTPUT_MESSAGES { + return Err(IdeError::unavailable("kernel output exceeded 32 MiB or 10000 messages; execution stopped")); + } + Ok(()) + } +} + +async fn supervise( + runtime: PathBuf, + notebook: PathBuf, + python: PathBuf, + status: StatusCallback, + stop: CancellationToken, + mut requests: mpsc::Receiver, + ready: oneshot::Sender>, +) { + status(¬ebook, KernelState::Starting, None); + let startup = async { + let mut session = Session::connect(&runtime, ¬ebook, &python).await?; + session.ready(&stop).await?; + IdeResult::Ok(session) + }; + let result = tokio::select! { + biased; + _ = stop.cancelled() => Err(IdeError::cancelled()), + result = tokio::time::timeout(STARTUP_TIMEOUT, startup) => result.unwrap_or_else(|_| Err(IdeError::unavailable("selected Python kernel startup timed out"))), + }; + let mut session = match result { + Ok(session) => session, + Err(error) => { + let message = format!( + "{}; selected Python: {}. Ensure ipykernel is installed in this exact environment ({} -m pip install ipykernel); no fallback interpreter was used", + error.message, + python.display(), + python.display() + ); + let error = IdeError::new(error.kind, message); + status(¬ebook, KernelState::Dead, Some(error.message.clone())); + let _ = ready.send(Err(error)); + return; + } + }; + status(¬ebook, KernelState::Idle, None); + if ready.send(Ok(())).is_err() { + stop.cancel(); + } + let mut death = tokio::time::interval(Duration::from_millis(100)); + death.set_missed_tick_behavior(MissedTickBehavior::Delay); + let failure = session.failure.clone(); + let mut reason = None; + loop { + tokio::select! { + biased; + _ = stop.cancelled() => break, + _ = failure.token.cancelled() => { reason = Some(failure.error().message); break; }, + _ = death.tick() => match session.process.death() { + Ok(None) => {}, + Ok(Some(message)) => { reason = Some(message); break; }, + Err(error) => { reason = Some(error.message); break; }, + }, + request = requests.recv() => { + let Some(mut request) = request else { break }; + let result = session.run(&mut request, &stop, ¬ebook, &status).await; + // Cancelled with reply+idle is reusable. Stop cancellation or any + // transport/death/overflow error invalidates this generation. + let fatal = result.as_ref().err().is_some_and(|e| e.kind != ide_core::error::ErrorKind::Cancelled) + || stop.is_cancelled(); + if fatal { reason = result.as_ref().err().map(|e| e.message.clone()); } + let Execution { reply, _admission, .. } = request; + drop(_admission); + let _ = reply.send(result); + if fatal { break; } + }, + frame = session.frames.recv() => { + // Post-idle/background messages are deliberately NOT forwarded. + // Keeping this drain alive also prevents stale IDs leaking into + // later executions or another process generation. + if frame.is_none() { reason = Some("kernel readers closed".into()); break; } + } + } + } + requests.close(); + while let Ok(request) = requests.try_recv() { + let _ = request.reply.send(Err(IdeError::unavailable("kernel stopped before execution admission"))); + } + session.shutdown().await; + status(¬ebook, KernelState::Dead, reason); +} + +#[cfg(test)] +#[path = "transport_tests.rs"] +mod tests; diff --git a/crates/ide-notebook/src/kernel/transport_tests.rs b/crates/ide-notebook/src/kernel/transport_tests.rs new file mode 100644 index 000000000..292c3114e --- /dev/null +++ b/crates/ide-notebook/src/kernel/transport_tests.rs @@ -0,0 +1,71 @@ +use super::*; +use serde_json::json; + +fn frame(channel: Channel, parent: &str, kind: &str, content: Value) -> Frame { + Frame { channel, parent: parent.into(), kind: kind.into(), content, bytes: 1, _permit: None } +} + +#[test] +fn completion_requires_both_channels_in_either_order_and_ignores_old_requests() { + for idle_first in [false, true] { + let mut completion = Completion::default(); + let idle = frame(Channel::Iopub, "current", "status", json!({"execution_state":"idle"})); + let reply = frame(Channel::Shell, "current", "execute_reply", json!({"status":"error","execution_count":7})); + let old_idle = frame(Channel::Iopub, "old-generation", "status", json!({"execution_state":"idle"})); + let old_reply = + frame(Channel::Shell, "old-generation", "execute_reply", json!({"status":"ok","execution_count":99})); + completion.observe(&old_reply, "current").unwrap(); + completion.observe(&old_idle, "current").unwrap(); + assert!(!completion.complete()); + let (first, second) = if idle_first { (&idle, &reply) } else { (&reply, &idle) }; + completion.observe(first, "current").unwrap(); + assert!(!completion.complete(), "one channel cannot finalize execution"); + completion.observe(second, "current").unwrap(); + assert!(completion.complete()); + let result = completion.reply.unwrap(); + assert_eq!(result.status, "error"); + assert_eq!(result.execution_count, Some(7)); + } +} + +#[test] +fn wrong_channel_cannot_forge_execution_completion() { + let mut completion = Completion::default(); + completion + .observe(&frame(Channel::Iopub, "run", "execute_reply", json!({"status":"ok","execution_count":2})), "run") + .unwrap(); + completion.observe(&frame(Channel::Shell, "run", "status", json!({"execution_state":"idle"})), "run").unwrap(); + assert!(!completion.complete()); + assert!(completion.reply.is_none()); + assert!(!completion.idle); +} + +#[test] +fn malformed_reply_is_rejected_instead_of_reported_as_success() { + for content in [ + json!({"status":"unexpected"}), + json!({"execution_count":1}), + json!({"status":"ok","execution_count":-1}), + json!({"status":"ok","execution_count":4294967296_u64}), + ] { + let mut completion = Completion::default(); + let error = completion.observe(&frame(Channel::Shell, "run", "execute_reply", content), "run").unwrap_err(); + assert_eq!(error.kind, ide_core::error::ErrorKind::Protocol); + assert!(!completion.complete()); + } +} + +#[test] +fn output_budget_stops_at_byte_or_message_boundary_without_wrapping() { + let mut bytes = Budget::default(); + bytes.add(MAX_OUTPUT_BYTES).unwrap(); + assert_eq!(bytes.add(1).unwrap_err().kind, ide_core::error::ErrorKind::Unavailable); + let mut messages = Budget::default(); + for _ in 0..MAX_OUTPUT_MESSAGES { + messages.add(0).unwrap(); + } + assert_eq!(messages.add(0).unwrap_err().kind, ide_core::error::ErrorKind::Unavailable); + let mut overflow = Budget::default(); + assert_eq!(overflow.add(usize::MAX).unwrap_err().kind, ide_core::error::ErrorKind::Unavailable); + assert_eq!(overflow.add(1).unwrap_err().kind, ide_core::error::ErrorKind::Unavailable); +} diff --git a/crates/ide-notebook/src/lib.rs b/crates/ide-notebook/src/lib.rs new file mode 100644 index 000000000..67935b6f9 --- /dev/null +++ b/crates/ide-notebook/src/lib.rs @@ -0,0 +1,5 @@ +//! Native notebook documents and owned Jupyter kernel sessions. +pub mod document; +pub mod kernel; + +pub use document::{NotebookCell, NotebookDocument, NotebookEdit, NotebookSnapshot}; diff --git a/crates/ide-notebook/tests/document.rs b/crates/ide-notebook/tests/document.rs new file mode 100644 index 000000000..da320e590 --- /dev/null +++ b/crates/ide-notebook/tests/document.rs @@ -0,0 +1,511 @@ +use std::{collections::HashSet, path::PathBuf}; + +use ide_core::error::ErrorKind; +use ide_notebook::{NotebookDocument, NotebookEdit, kernel::KernelOutput}; +use serde_json::{Value, json}; + +const ORIGINAL: &str = r##"{ + "unknownRoot": { "escaped": "\u0061", "number": 1e+02 }, + "nbformat_minor": 5, + "metadata": {"custom": [true, {"x": "\u03bb"}]}, + "cells": [ + { "unknownCell": {"keep": "\u0041"}, "source": ["\u0061\n", "b\n", "\u0063"], "id": "first", "metadata": {}, "cell_type": "code", "execution_count": null, "outputs": [] }, + { "cell_type": "markdown", "source": "# title", "metadata": {"custom": 3}, "id": "second", "attachments": {"pic.png": {"image/png": "aW1hZ2U="}}, "future": [1, 2] }, + { "cell_type": "vendor.future", "source": ["opaque"], "metadata": {}, "id": "third", "vendor": {"raw": "\u0042"} } + ], + "nbformat": 4 +}"##; + +fn document(text: &str) -> NotebookDocument { + NotebookDocument::from_text(PathBuf::from("notebook.ipynb"), text.into()).unwrap() +} + +fn apply(doc: &mut NotebookDocument, edit: NotebookEdit) { + let version = doc.version(); + doc.edit(version, edit).unwrap(); +} + +fn output(kind: &str, content: Value) -> KernelOutput { + KernelOutput { kind: kind.into(), content, msg_id: "actual-parent-id".into() } +} + +fn emit(doc: &mut NotebookDocument, cell: &str, kind: &str, content: Value) -> Vec { + doc.record_output(cell, &output(kind, content)).unwrap() +} + +fn two_code_cells() -> NotebookDocument { + document(&json!({ + "nbformat": 4, "nbformat_minor": 5, "metadata": {}, + "cells": [ + {"id": "a", "cell_type": "code", "metadata": {}, "source": "a()", "outputs": [], "execution_count": null}, + {"id": "b", "cell_type": "code", "metadata": {}, "source": ["b()"], "outputs": [], "execution_count": null} + ] + }).to_string()) +} + +#[test] +fn noop_save_preserves_all_bytes_for_utf8_bom_crlf_and_no_final_newline() { + let dir = tempfile::tempdir().unwrap(); + let variants = [ + ORIGINAL.to_owned(), + format!("{}\n", ORIGINAL), + ORIGINAL.replace('\n', "\r\n"), + format!("\u{feff}{}\r\n", ORIGINAL.replace('\n', "\r\n")), + ]; + for (index, text) in variants.iter().enumerate() { + let path = dir.path().join(format!("{index}.ipynb")); + std::fs::write(&path, text).unwrap(); + let mut doc = NotebookDocument::load(&path).unwrap(); + assert!(!doc.snapshot().dirty); + assert_eq!(doc.to_text(), *text); + let before = doc.version(); + doc.save().unwrap(); + assert_eq!(std::fs::read(&path).unwrap(), text.as_bytes()); + assert_eq!(doc.version(), before); + } +} + +#[test] +fn source_list_edit_preserves_unchanged_elements_and_every_other_byte() { + let mut doc = document(ORIGINAL); + apply(&mut doc, NotebookEdit::SetSource { cell: "first".into(), source: "a\nB\nc".into() }); + let expected = ORIGINAL.replace(r#"["\u0061\n", "b\n", "\u0063"]"#, r#"["\u0061\n", "B\n", "\u0063"]"#); + assert_eq!(doc.to_text(), expected); + assert!(doc.snapshot().dirty); + let edited_version = doc.version(); + doc.undo().unwrap(); + assert_eq!(doc.to_text(), ORIGINAL); + assert!(!doc.snapshot().dirty); + assert!(doc.version() > edited_version); + doc.redo().unwrap(); + assert_eq!(doc.to_text(), expected); +} + +#[test] +fn string_source_stays_a_string_and_attachments_and_unknown_cells_survive() { + let mut doc = document(ORIGINAL); + apply(&mut doc, NotebookEdit::SetSource { cell: "second".into(), source: "new\n标题".into() }); + assert_eq!(doc.to_text(), ORIGINAL.replace(r##""# title""##, r#""new\n标题""#)); + let parsed: Value = serde_json::from_str(&doc.to_text()).unwrap(); + assert!(parsed["cells"][1]["source"].is_string()); + assert_eq!(parsed["cells"][1]["attachments"]["pic.png"]["image/png"], "aW1hZ2U="); + assert_eq!(parsed["cells"][2]["vendor"]["raw"], "B"); +} + +#[test] +fn old_minor_and_invalid_or_duplicate_persisted_ids_are_not_repaired_on_load_or_edit() { + let text = r#"{"nbformat":4,"nbformat_minor":2,"metadata":{},"cells":[{"cell_type":"raw","metadata":{},"source":"one","id":"duplicate"},{"cell_type":"raw","metadata":{},"source":"two","id":"duplicate"},{"cell_type":"raw","metadata":{},"source":"three","id":"bad id"},{"cell_type":"raw","metadata":{},"source":"four"}]}"#; + let mut doc = document(text); + let before = doc.snapshot(); + assert_eq!(before.cells.iter().map(|cell| &cell.id).collect::>().len(), 4); + assert_eq!(doc.to_text(), text); + apply(&mut doc, NotebookEdit::SetSource { cell: before.cells[3].id.clone(), source: "changed".into() }); + let after = doc.snapshot(); + assert_eq!( + before.cells.iter().map(|cell| &cell.id).collect::>(), + after.cells.iter().map(|cell| &cell.id).collect::>() + ); + apply(&mut doc, NotebookEdit::InsertCell { after: None, cell_type: "code".into() }); + let parsed: Value = serde_json::from_str(&doc.to_text()).unwrap(); + assert_eq!(parsed["nbformat_minor"], 2); + assert!(parsed["cells"][0].get("id").is_none()); + assert_eq!(parsed["cells"][1]["id"], "duplicate"); + assert_eq!(parsed["cells"][2]["id"], "duplicate"); + assert_eq!(parsed["cells"][3]["id"], "bad id"); + assert!(parsed["cells"][4].get("id").is_none()); +} + +#[test] +fn new_cells_get_valid_ids_only_for_minor_five_and_keep_identity_through_history() { + let mut doc = document(ORIGINAL); + apply(&mut doc, NotebookEdit::InsertCell { after: Some("first".into()), cell_type: "markdown".into() }); + let inserted = doc.snapshot().cells[1].id.clone(); + let parsed: Value = serde_json::from_str(&doc.to_text()).unwrap(); + assert_eq!(parsed["cells"][1]["id"], inserted); + assert!(inserted.len() <= 64 && inserted.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-' || b == b'_')); + doc.undo().unwrap(); + assert_eq!(doc.to_text(), ORIGINAL); + doc.redo().unwrap(); + assert_eq!(doc.snapshot().cells[1].id, inserted); +} + +#[test] +fn move_and_delete_restore_exact_raw_subtrees_and_container_punctuation() { + let mut doc = document(ORIGINAL); + apply(&mut doc, NotebookEdit::MoveCell { cell: "third".into(), before: Some("first".into()) }); + assert_eq!( + doc.snapshot().cells.iter().map(|cell| cell.id.as_str()).collect::>(), + ["third", "first", "second"] + ); + assert!(doc.to_text().contains(r#""vendor": {"raw": "\u0042"}"#)); + apply(&mut doc, NotebookEdit::DeleteCell { cell: "first".into() }); + assert_eq!(doc.snapshot().cells.iter().map(|cell| cell.id.as_str()).collect::>(), ["third", "second"]); + doc.undo().unwrap(); + doc.undo().unwrap(); + assert_eq!(doc.to_text(), ORIGINAL); + doc.redo().unwrap(); + doc.redo().unwrap(); + assert_eq!(doc.snapshot().cells.iter().map(|cell| cell.id.as_str()).collect::>(), ["third", "second"]); +} + +#[test] +fn delete_last_cell_and_reinsert_into_empty_notebook_are_supported() { + let text = r#"{"nbformat":4,"nbformat_minor":5,"metadata":{},"cells":[ {"id":"only","cell_type":"raw","source":"x","metadata":{}} ]}"#; + let mut doc = document(text); + apply(&mut doc, NotebookEdit::DeleteCell { cell: "only".into() }); + assert_eq!(serde_json::from_str::(&doc.to_text()).unwrap()["cells"], json!([])); + doc.undo().unwrap(); + assert_eq!(doc.to_text(), text); + apply(&mut doc, NotebookEdit::DeleteCell { cell: "only".into() }); + apply(&mut doc, NotebookEdit::InsertCell { after: None, cell_type: "code".into() }); + assert_eq!(doc.snapshot().cells[0].cell_type, "code"); + assert_eq!(doc.snapshot().cells[0].outputs, Vec::::new()); +} + +#[test] +fn type_conversion_keeps_unknown_fields_and_restores_removed_fields_exactly() { + let text = + ORIGINAL.replace("\"outputs\": []", r#""outputs": [{"output_type":"vendor.output","payload":{"x":"\u0061"}}]"#); + let mut doc = document(&text); + apply(&mut doc, NotebookEdit::SetCellType { cell: "first".into(), cell_type: "markdown".into() }); + let converted: Value = serde_json::from_str(&doc.to_text()).unwrap(); + assert!(converted["cells"][0].get("outputs").is_none()); + assert!(converted["cells"][0].get("execution_count").is_none()); + assert_eq!(converted["cells"][0]["unknownCell"]["keep"], "A"); + doc.undo().unwrap(); + assert_eq!(doc.to_text(), text); + apply(&mut doc, NotebookEdit::SetCellType { cell: "second".into(), cell_type: "code".into() }); + let converted: Value = serde_json::from_str(&doc.to_text()).unwrap(); + assert_eq!(converted["cells"][1]["outputs"], json!([])); + assert_eq!(converted["cells"][1]["execution_count"], Value::Null); + assert_eq!(converted["cells"][1]["attachments"]["pic.png"]["image/png"], "aW1hZ2U="); +} + +#[test] +fn output_only_revisions_accept_pending_source_edits_but_stale_user_edits_are_rejected() { + let mut doc = two_code_cells(); + let base = doc.version(); + emit(&mut doc, "a", "stream", json!({"name":"stdout", "text":"hello"})); + assert!(doc.version() > base); + doc.edit(base, NotebookEdit::SetSource { cell: "a".into(), source: "new source".into() }).unwrap(); + assert_eq!(doc.snapshot().cells[0].outputs[0]["text"], "hello"); + let current = doc.version(); + let stale = doc.edit(base, NotebookEdit::SetSource { cell: "b".into(), source: "stale".into() }).unwrap_err(); + assert_eq!(stale.kind, ErrorKind::OutOfSync); + assert_eq!(doc.version(), current); + let future = doc.edit(current + 1, NotebookEdit::DeleteCell { cell: "a".into() }).unwrap_err(); + assert_eq!(future.kind, ErrorKind::OutOfSync); + assert_eq!(doc.snapshot().cells[1].source, "b()"); +} + +#[test] +fn source_undo_preserves_execution_outputs_and_output_does_not_clear_redo() { + let mut doc = two_code_cells(); + let initial_version = doc.version(); + apply(&mut doc, NotebookEdit::SetSource { cell: "a".into(), source: "print(1)".into() }); + doc.begin_execution("a").unwrap(); + emit(&mut doc, "a", "execute_input", json!({"execution_count": 7})); + emit(&mut doc, "a", "stream", json!({"name":"stdout", "text":"1\n"})); + let output_version = doc.version(); + doc.undo().unwrap(); + let state = doc.snapshot(); + assert_eq!(state.cells[0].source, "a()"); + assert_eq!(state.cells[0].execution_count, Some(7)); + assert_eq!(state.cells[0].outputs[0]["text"], "1\n"); + assert!(state.version > output_version && output_version > initial_version); + emit(&mut doc, "a", "stream", json!({"name":"stderr", "text":"later"})); + assert!(doc.snapshot().can_redo); + doc.redo().unwrap(); + let state = doc.snapshot(); + assert_eq!(state.cells[0].source, "print(1)"); + assert_eq!(state.cells[0].outputs[1]["text"], "later"); +} + +#[test] +fn history_is_rejected_until_execution_batch_is_drained() { + let mut doc = two_code_cells(); + apply(&mut doc, NotebookEdit::SetSource { cell: "a".into(), source: "changed".into() }); + doc.set_execution_active(true).unwrap(); + assert!(!doc.snapshot().can_undo); + assert_eq!(doc.undo().unwrap_err().kind, ErrorKind::InvalidInput); + assert_eq!(doc.redo().unwrap_err().kind, ErrorKind::InvalidInput); + assert_eq!(doc.reload().unwrap_err().kind, ErrorKind::InvalidInput); + assert!(doc.set_execution_active(true).is_err()); + doc.set_execution_active(false).unwrap(); + doc.undo().unwrap(); + assert_eq!(doc.snapshot().cells[0].source, "a()"); +} + +#[test] +fn inserted_cell_redo_retains_outputs_produced_after_its_creation() { + let mut doc = two_code_cells(); + apply(&mut doc, NotebookEdit::InsertCell { after: None, cell_type: "code".into() }); + let id = doc.snapshot().cells[0].id.clone(); + emit(&mut doc, &id, "stream", json!({"name":"stdout", "text":"keep"})); + doc.undo().unwrap(); + doc.redo().unwrap(); + assert_eq!(doc.snapshot().cells[0].id, id); + assert_eq!(doc.snapshot().cells[0].outputs[0]["text"], "keep"); +} + +#[test] +fn display_updates_reach_earlier_cells_and_never_persist_transient_ids() { + let mut doc = two_code_cells(); + let display = json!({"data":{"text/plain":"old", "vendor/mime":{"arbitrary":true}}, "metadata":{"custom":1}, "transient":{"display_id":"shared"}, "vendorField": 9}); + emit(&mut doc, "a", "display_data", display.clone()); + emit(&mut doc, "b", "execute_result", { + let mut value = display; + value["execution_count"] = json!(3); + value + }); + let changed = emit( + &mut doc, + "b", + "update_display_data", + json!({"data":{"text/plain":"new"}, "metadata":{"custom":2}, "transient":{"display_id":"shared"}}), + ); + assert_eq!(changed.into_iter().collect::>(), HashSet::from(["a".to_owned(), "b".to_owned()])); + let state = doc.snapshot(); + assert_eq!(state.cells[0].outputs[0]["data"]["text/plain"], "new"); + assert_eq!(state.cells[1].outputs[0]["data"]["text/plain"], "new"); + assert_eq!(state.cells[0].outputs[0]["vendorField"], 9); + assert_eq!(state.cells[1].outputs[0]["execution_count"], 3); + assert!(!doc.to_text().contains("display_id")); + assert!(!doc.to_text().contains("transient")); +} + +#[test] +fn deferred_clear_is_applied_only_to_next_output_and_removes_display_targets() { + let mut doc = two_code_cells(); + emit( + &mut doc, + "a", + "display_data", + json!({"data":{"text/plain":"old"}, "metadata":{}, "transient":{"display_id":"d"}}), + ); + let version = doc.version(); + assert!(emit(&mut doc, "a", "clear_output", json!({"wait":true})).is_empty()); + assert_eq!(doc.version(), version); + assert_eq!(doc.snapshot().cells[0].outputs[0]["data"]["text/plain"], "old"); + emit(&mut doc, "a", "execute_input", json!({"execution_count":2})); + assert_eq!(doc.snapshot().cells[0].outputs[0]["data"]["text/plain"], "old"); + assert_eq!(emit(&mut doc, "a", "stream", json!({"name":"stdout","text":"replacement"})), ["a"]); + let state = doc.snapshot(); + assert_eq!(state.cells[0].outputs, vec![json!({"output_type":"stream", "name":"stdout", "text":"replacement"})]); + assert!( + emit( + &mut doc, + "b", + "update_display_data", + json!({"data":{"text/plain":"not resurrected"}, "metadata":{}, "transient":{"display_id":"d"}}) + ) + .is_empty() + ); + emit(&mut doc, "a", "clear_output", json!({"wait":false})); + assert_eq!(doc.snapshot().cells[0].outputs, Vec::::new()); + assert_eq!(doc.snapshot().cells[0].execution_count, Some(2)); +} + +#[test] +fn streams_merge_only_adjacent_same_names_and_errors_and_mime_data_are_retained() { + let mut doc = two_code_cells(); + emit(&mut doc, "a", "stream", json!({"name":"stdout", "text":"one"})); + emit(&mut doc, "a", "stream", json!({"name":"stdout", "text":" two"})); + emit(&mut doc, "a", "stream", json!({"name":"stderr", "text":"warning"})); + emit(&mut doc, "a", "stream", json!({"name":"stdout", "text":"three"})); + emit( + &mut doc, + "a", + "error", + json!({"ename":"ValueError", "evalue":"bad", "traceback":["line one", "line two"], "vendor":true}), + ); + emit( + &mut doc, + "a", + "display_data", + json!({"data":{"image/png":"aW1hZ2U=", "application/vnd.vendor+json":{"x":1}}, "metadata":{"width":100}}), + ); + let state = doc.snapshot(); + assert_eq!(state.cells[0].outputs[0]["text"], "one two"); + assert_eq!(state.cells[0].outputs[1]["text"], "warning"); + assert_eq!(state.cells[0].outputs[2]["text"], "three"); + assert_eq!(state.cells[0].outputs[3]["traceback"], json!(["line one", "line two"])); + assert_eq!(state.cells[0].outputs[3]["vendor"], true); + assert_eq!(state.cells[0].outputs[4]["data"]["application/vnd.vendor+json"]["x"], 1); +} + +#[test] +fn save_undo_and_reload_track_the_exact_disk_baseline_and_document_incarnation() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("document.ipynb"); + std::fs::write(&path, ORIGINAL).unwrap(); + let mut doc = NotebookDocument::load(&path).unwrap(); + let id = doc.id().to_owned(); + apply(&mut doc, NotebookEdit::SetSource { cell: "second".into(), source: "saved".into() }); + let saved = doc.to_text(); + assert!(!doc.save().unwrap().dirty); + assert_eq!(std::fs::read_to_string(&path).unwrap(), saved); + assert!(doc.undo().unwrap().dirty); + assert_eq!(doc.to_text(), ORIGINAL); + let version = doc.version(); + let reloaded = doc.reload().unwrap(); + assert_eq!(reloaded.id, id); + assert!(reloaded.version > version); + assert!(!reloaded.dirty && !reloaded.can_undo && !reloaded.can_redo); + assert_eq!(doc.to_text(), saved); + assert_ne!(NotebookDocument::load(&path).unwrap().id(), id); +} + +#[test] +fn external_conflict_is_detected_even_when_clean_and_failed_save_retains_history() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("conflict.ipynb"); + std::fs::write(&path, ORIGINAL).unwrap(); + let mut doc = NotebookDocument::load(&path).unwrap(); + std::fs::write(&path, format!("{ORIGINAL}\n")).unwrap(); + assert_eq!(doc.save().unwrap_err().kind, ErrorKind::OutOfSync); + apply(&mut doc, NotebookEdit::SetSource { cell: "second".into(), source: "unsaved".into() }); + let unsaved = doc.to_text(); + let version = doc.version(); + assert_eq!(doc.save().unwrap_err().kind, ErrorKind::OutOfSync); + assert_eq!(doc.to_text(), unsaved); + assert_eq!(doc.version(), version); + assert!(doc.snapshot().dirty && doc.snapshot().can_undo); + assert_eq!(std::fs::read_to_string(&path).unwrap(), format!("{ORIGINAL}\n")); + std::fs::remove_file(&path).unwrap(); + assert_eq!(doc.save().unwrap_err().kind, ErrorKind::OutOfSync); + assert!(!path.exists()); + doc.undo().unwrap(); + assert_eq!(doc.to_text(), ORIGINAL); +} + +#[test] +fn malformed_source_duplicate_json_keys_unsupported_format_and_invalid_utf8_are_rejected() { + for text in [ + ORIGINAL.replace(r#"["\u0061\n", "b\n", "\u0063"]"#, "[123]"), + ORIGINAL.replace("\"nbformat\": 4", "\"nbformat\": 3"), + ORIGINAL.replace("\"nbformat\": 4", "\"nbformat\": 4, \"nbformat\": 4"), + ] { + assert!(NotebookDocument::from_text(PathBuf::from("invalid.ipynb"), text).is_err()); + } + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("invalid.ipynb"); + std::fs::write(&path, [0xff, 0xfe, 0x80]).unwrap(); + assert!(NotebookDocument::load(&path).is_err()); + assert_eq!(std::fs::read(&path).unwrap(), [0xff, 0xfe, 0x80]); +} + +#[test] +fn malformed_kernel_output_does_not_apply_a_pending_clear_or_advance_revision() { + let mut doc = two_code_cells(); + emit(&mut doc, "a", "stream", json!({"name":"stdout", "text":"keep"})); + emit(&mut doc, "a", "clear_output", json!({"wait":true})); + let before = doc.to_text(); + let version = doc.version(); + let error = doc.record_output("a", &output("display_data", json!({"data": {}}))).unwrap_err(); + assert_eq!(error.kind, ErrorKind::InvalidInput); + assert_eq!(doc.to_text(), before); + assert_eq!(doc.version(), version); + emit(&mut doc, "a", "stream", json!({"name":"stdout", "text":"replacement"})); + assert_eq!(doc.snapshot().cells[0].outputs[0]["text"], "replacement"); +} + +#[test] +fn late_output_never_recreates_deleted_cells() { + let mut doc = two_code_cells(); + apply(&mut doc, NotebookEdit::DeleteCell { cell: "a".into() }); + let before = doc.to_text(); + let error = doc.record_output("a", &output("stream", json!({"name":"stdout", "text":"late"}))).unwrap_err(); + assert_eq!(error.kind, ErrorKind::NotFound); + assert_eq!(doc.to_text(), before); + assert_eq!(doc.snapshot().cells[0].id, "b"); +} + +#[test] +fn minor_five_existing_duplicate_and_missing_ids_remain_untouched_after_mutation() { + let text = ORIGINAL.replace("\"id\": \"second\"", "\"id\": \"first\"").replace(", \"id\": \"third\"", ""); + let mut doc = document(&text); + let state = doc.snapshot(); + assert_eq!(state.cells.iter().map(|cell| &cell.id).collect::>().len(), 3); + assert_eq!(doc.to_text(), text); + apply(&mut doc, NotebookEdit::SetSource { cell: state.cells[0].id.clone(), source: "changed".into() }); + apply(&mut doc, NotebookEdit::InsertCell { after: Some(state.cells[2].id.clone()), cell_type: "raw".into() }); + let parsed: Value = serde_json::from_str(&doc.to_text()).unwrap(); + assert_eq!(parsed["cells"][0]["id"], "first"); + assert_eq!(parsed["cells"][1]["id"], "first"); + assert!(parsed["cells"][2].get("id").is_none()); + assert_eq!(parsed["cells"][3]["id"], doc.snapshot().cells[3].id); +} + +#[test] +fn edited_crlf_bom_notebook_preserves_every_unaffected_byte() { + let text = format!("\u{feff}{}\r\n", ORIGINAL.replace('\n', "\r\n")); + let mut doc = document(&text); + apply(&mut doc, NotebookEdit::SetSource { cell: "second".into(), source: "edited".into() }); + assert_eq!(doc.to_text(), text.replace(r##""# title""##, r#""edited""#)); + doc.undo().unwrap(); + assert_eq!(doc.to_text(), text); + assert!(!doc.is_dirty()); +} + +#[test] +fn ineffective_edits_do_not_create_dirty_state_history_or_revisions() { + let mut doc = two_code_cells(); + let version = doc.version(); + apply(&mut doc, NotebookEdit::SetSource { cell: "a".into(), source: "a()".into() }); + apply(&mut doc, NotebookEdit::SetCellType { cell: "a".into(), cell_type: "code".into() }); + apply(&mut doc, NotebookEdit::MoveCell { cell: "a".into(), before: Some("a".into()) }); + apply(&mut doc, NotebookEdit::MoveCell { cell: "b".into(), before: None }); + doc.begin_execution("a").unwrap(); + assert!(emit(&mut doc, "a", "clear_output", json!({"wait":false})).is_empty()); + let state = doc.snapshot(); + assert_eq!(state.version, version); + assert!(!state.dirty && !state.can_undo && !state.can_redo); +} + +#[test] +fn explicit_output_replacement_strips_transient_fields_without_discarding_user_redo() { + let mut doc = two_code_cells(); + apply(&mut doc, NotebookEdit::SetSource { cell: "a".into(), source: "changed".into() }); + doc.undo().unwrap(); + let base = doc.version(); + apply( + &mut doc, + NotebookEdit::SetOutputs { + cell: "a".into(), + execution_count: Some(4), + outputs: vec![ + json!({"output_type":"vendor.output", "payload":{"keep":true}, "transient":{"display_id":"not persisted"}}), + ], + }, + ); + assert!(doc.version() > base); + assert!(doc.snapshot().can_redo); + assert!(!doc.to_text().contains("transient")); + doc.edit(base, NotebookEdit::SetSource { cell: "b".into(), source: "accepted across output-only revision".into() }) + .unwrap(); + assert_eq!(doc.snapshot().cells[0].outputs[0]["payload"]["keep"], true); +} + +#[test] +fn malformed_stored_outputs_and_failed_output_edits_never_modify_the_document() { + let text = + ORIGINAL.replace("\"outputs\": []", r#""outputs": [{"output_type":"stream","name":"stdout","text":[123]}]"#); + assert!(NotebookDocument::from_text(PathBuf::from("invalid.ipynb"), text).is_err()); + let mut doc = two_code_cells(); + let before = doc.to_text(); + let version = doc.version(); + let error = doc + .edit( + version, + NotebookEdit::SetOutputs { + cell: "a".into(), + execution_count: Some(1), + outputs: vec![json!({"output_type":"error","ename":"E","evalue":"bad","traceback":[42]})], + }, + ) + .unwrap_err(); + assert_eq!(error.kind, ErrorKind::InvalidInput); + assert_eq!(doc.to_text(), before); + assert_eq!(doc.version(), version); +} diff --git a/crates/ide-notebook/tests/kernel_runtime.rs b/crates/ide-notebook/tests/kernel_runtime.rs new file mode 100644 index 000000000..eff49c6e2 --- /dev/null +++ b/crates/ide-notebook/tests/kernel_runtime.rs @@ -0,0 +1,322 @@ +//! Opt-in native acceptance tests. Set OPENIDEA_TEST_PYTHON to an absolute +//! ipykernel-enabled interpreter and run this target with --ignored. +//! These tests execute real Python over authenticated native ZeroMQ, no server. + +use ide_core::{CancellationToken, IdeResult, error::ErrorKind}; +use ide_notebook::kernel::{ExecutionResult, KernelManager, KernelOutput, KernelState}; +use serde_json::Value; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::sync::oneshot; + +type Outputs = Arc>>; + +fn python() -> PathBuf { + let path = PathBuf::from( + std::env::var_os("OPENIDEA_TEST_PYTHON").expect("set OPENIDEA_TEST_PYTHON to a real ipykernel Python"), + ); + assert!(path.is_absolute()); + path +} + +fn manager(root: &Path) -> Arc { + Arc::new(KernelManager::new(root.join("runtime"), Arc::new(|_, _, _| {}))) +} + +async fn run( + manager: &KernelManager, + notebook: &Path, + python: &Path, + code: &str, +) -> (IdeResult, Vec) { + let outputs: Outputs = Arc::new(Mutex::new(Vec::new())); + let sink = outputs.clone(); + let result = tokio::time::timeout( + Duration::from_secs(45), + manager.execute( + notebook, + python, + code, + CancellationToken::new(), + Arc::new(move |output| sink.lock().unwrap().push(output)), + ), + ) + .await + .expect("real kernel request exceeded test deadline"); + let outputs = std::mem::take(&mut *outputs.lock().unwrap()); + (result, outputs) +} + +fn json_display(outputs: &[KernelOutput]) -> &Value { + &outputs + .iter() + .find(|output| output.kind == "display_data" && output.content["data"]["application/json"].is_object()) + .expect("expected native JSON display output") + .content["data"]["application/json"] +} + +#[tokio::test] +async fn cancelled_before_admission_cannot_become_an_interpreter_error() { + let root = tempfile::tempdir().unwrap(); + let manager = manager(root.path()); + let cancel = CancellationToken::new(); + cancel.cancel(); + let result = manager + .execute( + &root.path().join("cancelled.ipynb"), + &root.path().join("does-not-exist"), + "raise RuntimeError('must never execute')", + cancel, + Arc::new(|_| panic!("cancelled code emitted output")), + ) + .await; + assert_eq!(result.unwrap_err().kind, ErrorKind::Cancelled); + manager.shutdown().await; +} + +#[tokio::test] +#[ignore = "requires OPENIDEA_TEST_PYTHON with real ipykernel"] +async fn real_kernel_preserves_state_interpreter_mime_and_request_identity() { + let root = tempfile::tempdir().unwrap(); + let statuses = Arc::new(Mutex::new(Vec::new())); + let sink = statuses.clone(); + let manager = + KernelManager::new(root.path().join("runtime"), Arc::new(move |_, state, _| sink.lock().unwrap().push(state))); + let notebook = root.path().join("native.ipynb"); + let python = python(); + let (result, outputs) = run(&manager, ¬ebook, &python, r#" +import os, sys +from IPython.display import display, update_display, clear_output +persisted_value = 41 +print('native stdout', flush=True) +print('native stderr', file=sys.stderr, flush=True) +display({'application/x-openidea-test': {'opaque': [1, {'unknown': 'kept'}]}, 'text/plain': 'custom'}, raw=True, display_id='stable-display') +update_display({'text/plain': 'updated'}, raw=True, display_id='stable-display') +clear_output(wait=True) +display({'text/html': '
42
', 'image/svg+xml': ''}, raw=True) +persisted_value + 1 +"#).await; + let result = result.unwrap(); + assert_eq!(result.status, "ok"); + let request = &outputs.iter().find(|o| o.kind == "execute_input").unwrap().msg_id; + assert!(!request.is_empty(), "Jupyter must supply a real parent request ID"); + assert!(outputs.iter().all(|o| &o.msg_id == request && o.kind != "status")); + assert!(outputs.iter().any(|o| o.kind == "stream" + && o.content["name"] == "stdout" + && o.content["text"].as_str().unwrap().contains("native stdout"))); + assert!(outputs.iter().any(|o| o.kind == "stream" + && o.content["name"] == "stderr" + && o.content["text"].as_str().unwrap().contains("native stderr"))); + assert!(outputs.iter().any(|o| o.content["data"]["application/x-openidea-test"]["opaque"][1]["unknown"] == "kept")); + assert!( + outputs + .iter() + .any(|o| o.kind == "update_display_data" && o.content["transient"]["display_id"] == "stable-display") + ); + assert!(outputs.iter().any(|o| o.kind == "clear_output" && o.content["wait"] == true)); + assert!(outputs.iter().any(|o| o.kind == "execute_result" && o.content["data"]["text/plain"] == "42")); + let (result, next) = run(&manager, ¬ebook, &python, + "display({'application/json': {'value': persisted_value, 'executable': sys.executable, 'cwd': os.getcwd()}}, raw=True)").await; + assert_eq!(result.unwrap().status, "ok"); + let value = json_display(&next); + assert_eq!(value["value"], 41); + assert_eq!( + std::fs::canonicalize(value["executable"].as_str().unwrap()).unwrap(), + std::fs::canonicalize(&python).unwrap() + ); + assert_eq!( + std::fs::canonicalize(value["cwd"].as_str().unwrap()).unwrap(), + std::fs::canonicalize(root.path()).unwrap() + ); + assert!(next.iter().all(|o| o.msg_id != *request)); + let (result, errors) = run(&manager, ¬ebook, &python, "input('must not hang')").await; + assert_eq!(result.unwrap().status, "error"); + assert!(errors.iter().any(|o| o.kind == "error" && o.content["ename"] == "StdinNotImplementedError")); + let states = statuses.lock().unwrap().clone(); + assert_eq!(states.first(), Some(&KernelState::Starting)); + assert!(states.windows(2).any(|pair| pair == [KernelState::Busy, KernelState::Idle])); + manager.shutdown().await; +} + +async fn long_run( + manager: Arc, + notebook: PathBuf, + python: PathBuf, + cancel: CancellationToken, + catches_interrupt: bool, +) -> (tokio::task::JoinHandle>, oneshot::Receiver<()>) { + let (ready, received) = oneshot::channel(); + let ready = Mutex::new(Some(ready)); + let code = if catches_interrupt { + "import time, os\nkept_after_interrupt = 123\ntry:\n print('INTERRUPT-READY', flush=True)\n while True: time.sleep(0.01)\nexcept KeyboardInterrupt:\n caught_interrupt = True\n" + } else { + "import time, os\nkept_after_interrupt = 123\nprint('INTERRUPT-READY', flush=True)\nwhile True: time.sleep(0.01)\n" + }; + let task = tokio::spawn(async move { + manager + .execute( + ¬ebook, + &python, + code, + cancel, + Arc::new(move |output| { + if output.kind == "stream" + && output.content["text"].as_str().is_some_and(|s| s.contains("INTERRUPT-READY")) + && let Some(ready) = ready.lock().unwrap().take() + { + let _ = ready.send(()); + } + }), + ) + .await + }); + (task, received) +} + +#[tokio::test] +#[ignore = "requires real ipykernel; Windows must exercise JPY_INTERRUPT_EVENT"] +async fn real_interrupt_drains_and_keeps_interpreter_state() { + let root = tempfile::tempdir().unwrap(); + let manager = manager(root.path()); + let notebook = root.path().join("interrupt.ipynb"); + let python = python(); + let (task, ready) = + long_run(manager.clone(), notebook.clone(), python.clone(), CancellationToken::new(), false).await; + tokio::time::timeout(Duration::from_secs(45), ready).await.unwrap().unwrap(); + tokio::time::timeout(Duration::from_secs(10), manager.interrupt(¬ebook)).await.unwrap().unwrap(); + assert_eq!(task.await.unwrap().unwrap_err().kind, ErrorKind::Cancelled); + let (result, outputs) = run( + &manager, + ¬ebook, + &python, + "from IPython.display import display\ndisplay({'application/json': {'kept': kept_after_interrupt}}, raw=True)", + ) + .await; + assert_eq!(result.unwrap().status, "ok"); + assert_eq!(json_display(&outputs)["kept"], 123, "interrupt must not emulate success by restarting Python"); + manager.shutdown().await; +} + +#[tokio::test] +#[ignore = "requires OPENIDEA_TEST_PYTHON with real ipykernel"] +async fn token_cancellation_is_not_success_when_python_catches_keyboard_interrupt() { + let root = tempfile::tempdir().unwrap(); + let manager = manager(root.path()); + let notebook = root.path().join("catch.ipynb"); + let python = python(); + let cancel = CancellationToken::new(); + let (task, ready) = long_run(manager.clone(), notebook.clone(), python.clone(), cancel.clone(), true).await; + tokio::time::timeout(Duration::from_secs(45), ready).await.unwrap().unwrap(); + cancel.cancel(); + let error = tokio::time::timeout(Duration::from_secs(10), task).await.unwrap().unwrap().unwrap_err(); + assert_eq!(error.kind, ErrorKind::Cancelled); + let (result, outputs) = run(&manager, ¬ebook, &python, + "from IPython.display import display\ndisplay({'application/json': {'caught': caught_interrupt, 'kept': kept_after_interrupt}}, raw=True)").await; + assert_eq!(result.unwrap().status, "ok"); + assert_eq!(json_display(&outputs)["caught"], true); + assert_eq!(json_display(&outputs)["kept"], 123); + manager.shutdown().await; +} + +#[tokio::test] +#[ignore = "requires OPENIDEA_TEST_PYTHON with real ipykernel"] +async fn stale_parent_output_is_dropped_and_restart_invalidates_state() { + let root = tempfile::tempdir().unwrap(); + let manager = manager(root.path()); + let notebook = root.path().join("generation.ipynb"); + let python = python(); + let (result, _) = run( + &manager, + ¬ebook, + &python, + r#" +import threading, os +from IPython.display import display +kernel = get_ipython().kernel +old_parent = kernel.get_parent() +old_release = threading.Event() +old_done = threading.Event() +old_state = 99 +def old_worker(): + old_release.wait() + kernel.session.send(kernel.iopub_socket, 'stream', + {'name': 'stdout', 'text': 'OLD-GENERATION\n'}, + parent=old_parent) + old_done.set() +threading.Thread(target=old_worker, daemon=True).start() +"#, + ) + .await; + assert_eq!(result.unwrap().status, "ok"); + // Send an explicitly old protocol parent during the next execution. + // Ordinary thread print() can inherit ipykernel's current global parent. + let (result, outputs) = run(&manager, ¬ebook, &python, + "old_release.set()\nassert old_done.wait(5)\nprint('CURRENT-GENERATION', flush=True)\ndisplay({'application/json': {'pid': os.getpid()}}, raw=True)").await; + assert_eq!(result.unwrap().status, "ok"); + assert!(!outputs.iter().any(|o| o.content["text"].as_str().is_some_and(|s| s.contains("OLD-GENERATION")))); + assert!(outputs.iter().any(|o| o.content["text"].as_str().is_some_and(|s| s.contains("CURRENT-GENERATION")))); + let old_pid = json_display(&outputs)["pid"].as_u64().unwrap(); + manager.restart(¬ebook).await.unwrap(); + let (result, outputs) = run(&manager, ¬ebook, &python, + "import os\nfrom IPython.display import display\ndisplay({'application/json': {'pid': os.getpid(), 'old': 'old_state' in globals()}}, raw=True)").await; + assert_eq!(result.unwrap().status, "ok"); + assert_ne!(json_display(&outputs)["pid"], old_pid); + assert_eq!(json_display(&outputs)["old"], false); + manager.shutdown().await; +} + +#[tokio::test] +#[ignore = "requires OPENIDEA_TEST_PYTHON with real ipykernel"] +async fn process_death_and_output_overflow_fail_explicitly_and_can_restart() { + let root = tempfile::tempdir().unwrap(); + let manager = manager(root.path()); + let notebook = root.path().join("death.ipynb"); + let python = python(); + let (result, _) = run(&manager, ¬ebook, &python, "import os\nos._exit(23)").await; + assert!(result.is_err(), "owned process death cannot report execute success"); + manager.restart(¬ebook).await.unwrap(); + let (result, _) = run( + &manager, + ¬ebook, + &python, + "from IPython.display import display\ndisplay({'text/plain': 'x' * (9 * 1024 * 1024)}, raw=True)", + ) + .await; + let error = result.unwrap_err(); + assert_eq!(error.kind, ErrorKind::Unavailable); + assert!(error.message.contains("limit") || error.message.contains("overflow") || error.message.contains("exceed")); + manager.restart(¬ebook).await.unwrap(); + let (result, outputs) = run(&manager, ¬ebook, &python, "print('RECOVERED', flush=True)").await; + assert_eq!(result.unwrap().status, "ok"); + assert!(outputs.iter().any(|o| o.content["text"].as_str().is_some_and(|s| s.contains("RECOVERED")))); + manager.shutdown().await; +} + +#[tokio::test] +#[ignore = "requires OPENIDEA_TEST_PYTHON and OPENIDEA_TEST_PYTHON_SECOND"] +async fn changing_selected_python_recreates_kernel_instead_of_reusing_old_environment() { + let root = tempfile::tempdir().unwrap(); + let manager = manager(root.path()); + let notebook = root.path().join("selection.ipynb"); + let first = python(); + let second = PathBuf::from( + std::env::var_os("OPENIDEA_TEST_PYTHON_SECOND").expect("set a second absolute ipykernel interpreter path"), + ); + assert_ne!(first, second); + let (result, outputs) = run(&manager, ¬ebook, &first, + "import os\nfrom IPython.display import display\nold_environment_state = 1\ndisplay({'application/json': {'pid': os.getpid()}}, raw=True)").await; + assert_eq!(result.unwrap().status, "ok"); + let first_pid = json_display(&outputs)["pid"].as_u64().unwrap(); + let (result, outputs) = run(&manager, ¬ebook, &second, + "import os, sys\nfrom IPython.display import display\ndisplay({'application/json': {'pid': os.getpid(), 'executable': sys.executable, 'old': 'old_environment_state' in globals()}}, raw=True)").await; + assert_eq!(result.unwrap().status, "ok"); + let data = json_display(&outputs); + assert_ne!(data["pid"], first_pid); + assert_eq!(data["old"], false); + assert_eq!( + std::fs::canonicalize(data["executable"].as_str().unwrap()).unwrap(), + std::fs::canonicalize(second).unwrap() + ); + manager.shutdown().await; +} diff --git a/crates/ide-vfs/src/encoding.rs b/crates/ide-vfs/src/encoding.rs index 6a3d1d8d4..9153abded 100644 --- a/crates/ide-vfs/src/encoding.rs +++ b/crates/ide-vfs/src/encoding.rs @@ -226,6 +226,17 @@ pub(crate) fn is_atomic_temp(path: &Path) -> bool { /// writes through symlinks. Falls back to an in-place write when the directory /// does not allow creating files. Returns the new mtime in ms. pub fn write_atomic(path: &Path, bytes: &[u8]) -> IdeResult { + write_atomic_impl(path, bytes, true) +} + +/// Replaces a file atomically without falling back to truncating it in place. +/// Preserves existing permissions and writes through symlinks; if a sibling +/// temporary file cannot be created, leaves the existing file untouched. +pub fn write_atomic_strict(path: &Path, bytes: &[u8]) -> IdeResult { + write_atomic_impl(path, bytes, false) +} + +fn write_atomic_impl(path: &Path, bytes: &[u8], allow_in_place: bool) -> IdeResult { let target = match fs::symlink_metadata(path) { Ok(m) if m.file_type().is_symlink() => fs::canonicalize(path)?, _ => path.to_path_buf(), @@ -235,10 +246,19 @@ pub fn write_atomic(path: &Path, bytes: &[u8]) -> IdeResult { return Err(IdeError::invalid(format!("{} is a directory", target.display()))); } let tmp = temp_path_for(&target); - let file = fs::OpenOptions::new().write(true).create_new(true).open(&tmp); + let mut options = fs::OpenOptions::new(); + options.write(true).create_new(true); + #[cfg(unix)] + if existing.is_some() || !allow_in_place { + use std::os::unix::fs::OpenOptionsExt; + // Replacement bytes must not become readable before the final + // permissions are restored, even when the directory has a loose umask. + options.mode(0o600); + } + let file = options.open(&tmp); let mut file = match file { Ok(f) => f, - Err(e) if e.kind() == std::io::ErrorKind::PermissionDenied && existing.is_some() => { + Err(e) if allow_in_place && e.kind() == std::io::ErrorKind::PermissionDenied && existing.is_some() => { tracing::debug!("cannot create temp file next to {}: {e}; writing in place", target.display()); fs::write(&target, bytes)?; return Ok(mtime_ms(&fs::metadata(&target)?)); @@ -265,6 +285,61 @@ pub fn write_atomic(path: &Path, bytes: &[u8]) -> IdeResult { mod tests { use super::*; + #[test] + fn strict_write_replaces_complete_contents() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("notebook.ipynb"); + fs::write(&path, b"original, longer contents").unwrap(); + + write_atomic_strict(&path, b"{\"cells\":[]}").unwrap(); + + assert_eq!(fs::read(&path).unwrap(), b"{\"cells\":[]}"); + } + + #[cfg(unix)] + #[test] + fn strict_write_preserves_symlink_and_target_permissions() { + use std::os::unix::fs::{PermissionsExt, symlink}; + + let dir = tempfile::tempdir().unwrap(); + let target = dir.path().join("target.ipynb"); + let link = dir.path().join("link.ipynb"); + fs::write(&target, b"old").unwrap(); + fs::set_permissions(&target, fs::Permissions::from_mode(0o640)).unwrap(); + symlink(&target, &link).unwrap(); + + write_atomic_strict(&link, b"new notebook").unwrap(); + + assert!(fs::symlink_metadata(&link).unwrap().file_type().is_symlink()); + assert_eq!(fs::read(&target).unwrap(), b"new notebook"); + assert_eq!(fs::metadata(&target).unwrap().permissions().mode() & 0o777, 0o640); + } + + #[cfg(unix)] + #[test] + fn strict_write_does_not_truncate_when_sibling_creation_is_denied() { + use std::os::unix::fs::PermissionsExt; + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("notebook.ipynb"); + fs::write(&path, b"original notebook").unwrap(); + fs::set_permissions(&path, fs::Permissions::from_mode(0o666)).unwrap(); + fs::set_permissions(dir.path(), fs::Permissions::from_mode(0o555)).unwrap(); + // Privileged test runners can bypass directory permissions. + let probe = dir.path().join("probe"); + if fs::File::create(&probe).is_ok() { + fs::set_permissions(dir.path(), fs::Permissions::from_mode(0o755)).unwrap(); + fs::remove_file(probe).unwrap(); + return; + } + + let result = write_atomic_strict(&path, b"replacement"); + fs::set_permissions(dir.path(), fs::Permissions::from_mode(0o755)).unwrap(); + + assert_eq!(result.unwrap_err().kind, ide_core::error::ErrorKind::PermissionDenied); + assert_eq!(fs::read(&path).unwrap(), b"original notebook"); + } + #[test] fn decode_variants() { let (t, e, bom, bin) = decode(b"\xEF\xBB\xBFhi"); diff --git a/crates/ide-vfs/src/lib.rs b/crates/ide-vfs/src/lib.rs index baa0c7962..24c882e70 100644 --- a/crates/ide-vfs/src/lib.rs +++ b/crates/ide-vfs/src/lib.rs @@ -21,7 +21,9 @@ pub use archive::{ }; pub use bytes::{ByteRange, FileProbe, MAX_RANGE_LEN, base64, image_mime, probe, read_data_url, read_range}; use encoding::SNIFF_LEN; -pub use encoding::{LARGE_FILE_THRESHOLD, MAX_OPEN_SIZE, READ_ONLY_THRESHOLD, encode, write_atomic}; +pub use encoding::{ + LARGE_FILE_THRESHOLD, MAX_OPEN_SIZE, READ_ONLY_THRESHOLD, encode, write_atomic, write_atomic_strict, +}; pub use rules::ExcludeMatcher; pub use watch::{WatchMode, Watcher}; diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index d012e6533..d292831dd 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -63,6 +63,7 @@ ide-spell.workspace = true ide-ssh.workspace = true ide-remote.workspace = true ide-host.workspace = true +ide-notebook.workspace = true ide-jsonrpc.workspace = true ide-tools = { path = "../crates/ide-tools" } ide-plugin = { path = "../crates/ide-plugin" } diff --git a/src-tauri/src/commands/mod.rs b/src-tauri/src/commands/mod.rs index d743891f6..d6262d3a1 100644 --- a/src-tauri/src/commands/mod.rs +++ b/src-tauri/src/commands/mod.rs @@ -28,6 +28,7 @@ pub mod http; pub mod java; pub mod lsp; pub mod nav; +pub mod notebook; pub mod perf; pub mod plugin; pub mod profiler; @@ -96,6 +97,17 @@ pub fn builder() -> tauri_specta::Builder { agent::agent_delete, agent::agent_changes, agent::agent_base_text, + notebook::notebook_open, + notebook::notebook_get, + notebook::notebook_edit, + notebook::notebook_save, + notebook::notebook_reload, + notebook::notebook_undo, + notebook::notebook_redo, + notebook::notebook_close, + notebook::notebook_execute, + notebook::notebook_interrupt, + notebook::notebook_restart, fs::fs_read_dir, bytes::fs_probe, bytes::fs_read_range, diff --git a/src-tauri/src/commands/notebook.rs b/src-tauri/src/commands/notebook.rs new file mode 100644 index 000000000..d49d67489 --- /dev/null +++ b/src-tauri/src/commands/notebook.rs @@ -0,0 +1,88 @@ +//! Thin desktop adapters for the shared native notebook service. +use crate::state::AppState; +use ide_core::{IdeResult, RequestId}; +use ide_host::ops::notebook::{self, NotebookOutputEvent}; +use ide_notebook::{NotebookEdit, NotebookSnapshot}; +use tauri::ipc::Channel; +use tauri::{State, Window}; + +#[tauri::command] +#[specta::specta] +pub async fn notebook_open(window: Window, state: State<'_, AppState>, path: String) -> IdeResult { + notebook::open(&state.host, window.label(), path).await +} + +#[tauri::command] +#[specta::specta] +pub fn notebook_get(window: Window, state: State<'_, AppState>, id: String) -> IdeResult { + notebook::get(&state.host, window.label(), id) +} + +#[tauri::command] +#[specta::specta] +pub async fn notebook_edit( + window: Window, + state: State<'_, AppState>, + id: String, + version: u32, + edit: NotebookEdit, +) -> IdeResult { + notebook::edit(&state.host, window.label(), id, version, edit).await +} + +#[tauri::command] +#[specta::specta] +pub async fn notebook_save(window: Window, state: State<'_, AppState>, id: String) -> IdeResult { + notebook::save(&state.host, window.label(), id).await +} + +#[tauri::command] +#[specta::specta] +pub async fn notebook_reload(window: Window, state: State<'_, AppState>, id: String) -> IdeResult { + notebook::reload(&state.host, window.label(), id).await +} + +#[tauri::command] +#[specta::specta] +pub async fn notebook_undo(window: Window, state: State<'_, AppState>, id: String) -> IdeResult { + notebook::undo(&state.host, window.label(), id).await +} + +#[tauri::command] +#[specta::specta] +pub async fn notebook_redo(window: Window, state: State<'_, AppState>, id: String) -> IdeResult { + notebook::redo(&state.host, window.label(), id).await +} + +#[tauri::command] +#[specta::specta] +pub async fn notebook_close(window: Window, state: State<'_, AppState>, id: String) -> IdeResult<()> { + notebook::close(&state.host, window.label(), id).await +} + +#[tauri::command] +#[specta::specta] +pub async fn notebook_interrupt(window: Window, state: State<'_, AppState>, id: String) -> IdeResult<()> { + notebook::interrupt(&state.host, window.label(), id).await +} + +#[tauri::command] +#[specta::specta] +pub async fn notebook_restart(window: Window, state: State<'_, AppState>, id: String) -> IdeResult { + notebook::restart(&state.host, window.label(), id).await +} + +#[tauri::command] +#[specta::specta] +#[allow(clippy::too_many_arguments)] +pub async fn notebook_execute( + window: Window, + state: State<'_, AppState>, + id: String, + cells: Vec, + python: Option, + request: RequestId, + on_output: Channel, +) -> IdeResult { + notebook::execute(&state.host, window.label(), id, cells, python, request, super::sender(on_output)).await +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index e714953c2..099c29497 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -139,11 +139,14 @@ pub fn run() { } #[cfg(not(target_os = "macos"))] let _ = &early_opens; - // Quit after installing an update (`app_install_update`): relaunch. - if let tauri::RunEvent::Exit = event - && app.try_state::().is_some_and(|w| w.restart_requested()) - { - tauri::process::restart(&app.env()); + if let tauri::RunEvent::Exit = event { + if let Some(state) = app.try_state::() { + state.host.shutdown_notebooks_on_exit(); + } + // Cleanup must finish before update/restart replaces this process. + if app.try_state::().is_some_and(|w| w.restart_requested()) { + tauri::process::restart(&app.env()); + } } }); } diff --git a/src/bindings.ts b/src/bindings.ts index 60c7596f2..7bf4a4797 100644 --- a/src/bindings.ts +++ b/src/bindings.ts @@ -116,6 +116,17 @@ export const commands = { agentChanges: (task: string) => typedError(__TAURI_INVOKE("agent_changes", { task })), /** Content of `path` when the task started (left side of its diff). */ agentBaseText: (task: string, path: IpcPath) => typedError(__TAURI_INVOKE("agent_base_text", { task, path })), + notebookOpen: (path: string) => typedError(__TAURI_INVOKE("notebook_open", { path })), + notebookGet: (id: string) => typedError(__TAURI_INVOKE("notebook_get", { id })), + notebookEdit: (id: string, version: number, edit: NotebookEdit) => typedError(__TAURI_INVOKE("notebook_edit", { id, version, edit })), + notebookSave: (id: string) => typedError(__TAURI_INVOKE("notebook_save", { id })), + notebookReload: (id: string) => typedError(__TAURI_INVOKE("notebook_reload", { id })), + notebookUndo: (id: string) => typedError(__TAURI_INVOKE("notebook_undo", { id })), + notebookRedo: (id: string) => typedError(__TAURI_INVOKE("notebook_redo", { id })), + notebookClose: (id: string) => typedError(__TAURI_INVOKE("notebook_close", { id })), + notebookExecute: (id: string, cells: string[], python: string | null, request: RequestId, onOutput: Channel) => typedError(__TAURI_INVOKE("notebook_execute", { id, cells, python, request, onOutput })), + notebookInterrupt: (id: string) => typedError(__TAURI_INVOKE("notebook_interrupt", { id })), + notebookRestart: (id: string) => typedError(__TAURI_INVOKE("notebook_restart", { id })), fsReadDir: (path: string) => typedError(__TAURI_INVOKE("fs_read_dir", { path })), /** Size and text/binary classification, from the file's head only. */ fsProbe: (path: string) => typedError(__TAURI_INVOKE("fs_probe", { path })), @@ -3327,7 +3338,14 @@ cancellable: boolean } | /** A child process (`ssh`, `git`) asks for a credential; answered with `ssh_askpass_answer`. */ { type: "askpass"; id: number; prompt: string; kind: AskpassKind } | /** Prompt `id` was answered (in any window) or gave up. */ -{ type: "askpassDone"; id: number }; +{ type: "askpassDone"; id: number } | +/** + * The kernel of an open notebook changed state (`starting`, `idle`, `busy`, + * `dead`); `message` says why it failed to start or died. + */ +{ type: "notebookKernel"; +/** Backend document incarnation; stale queued events cannot target a reopen. */ +id: string; path: IpcPath; state: string; message: string | null }; /** An identifier-like token whose text equals the searched name. */ export type IdentifierOccurrence = { @@ -4168,6 +4186,43 @@ export type NotLocalReason = */ "unresolved"; +export type NotebookCell = { + id: string, + cellType: string, + source: string, + metadata: unknown, + outputs: unknown[], + executionCount: number | null, + attachments: unknown, +}; + +export type NotebookEdit = { type: "setSource"; cell: string; source: string } | { type: "insertCell"; after: string | null; cellType: string } | { type: "deleteCell"; cell: string } | { type: "moveCell"; cell: string; before: string | null } | { type: "setCellType"; cell: string; cellType: string } | { type: "setOutputs"; cell: string; outputs: unknown[]; executionCount: number | null }; + +export type NotebookOutputEvent = { + id: string, + path: string, + cell: string, + msgId: string, + kind: string, + content: unknown, + version: number, + dirty: boolean, + changedCells: string[], +}; + +export type NotebookSnapshot = { + id: string, + path: string, + version: number, + dirty: boolean, + canUndo: boolean, + canRedo: boolean, + nbformat: number, + nbformatMinor: number, + metadata: unknown, + cells: NotebookCell[], +}; + export type NotificationGroupSettings = { /** `balloon` | `sticky` | `none` */ display?: string, diff --git a/src/editor/core/actions.ts b/src/editor/core/actions.ts index 2c4043083..8de51c9e9 100644 --- a/src/editor/core/actions.ts +++ b/src/editor/core/actions.ts @@ -805,16 +805,18 @@ export function registerEditorActions(): () => void { category: "Window", icon: "close", when: "(editorTabsOpen || editorSplit) && !toolWindowFocus", - run: () => ctrl.closeActiveTab(), + run: async () => { + await ctrl.closeActiveTab(); + }, }, { id: "editor.closeOtherTabs", title: "Close Other Tabs", category: "Window", when: TABS_WHEN, - run: () => { + run: async () => { const t = activePaneTab(); - if (t) ctrl.closeOthers(t.paneId, t.tabId); + if (t) await ctrl.closeOthers(t.paneId, t.tabId); }, }, { @@ -822,7 +824,9 @@ export function registerEditorActions(): () => void { title: "Close All Tabs", category: "Window", when: "editorTabsOpen", - run: () => ctrl.closeEverything(), + run: async () => { + await ctrl.closeEverything(); + }, }, { id: "editor.pinTab", @@ -896,14 +900,16 @@ export function registerEditorActions(): () => void { title: "Unsplit", category: "Window", when: "editorSplit", - run: () => ctrl.unsplit(useTabs.getState().activePane), + run: async () => { + await ctrl.unsplit(useTabs.getState().activePane); + }, }, { id: "editor.unsplitAll", title: "Unsplit All", category: "Window", when: "editorSplit", - run: () => ctrl.unsplitAll(useTabs.getState().activePane), + run: () => void ctrl.unsplitAll(useTabs.getState().activePane), }, { id: "editor.changeSplitOrientation", diff --git a/src/editor/core/controller.test.ts b/src/editor/core/controller.test.ts index 031c37f62..761bb64c5 100644 --- a/src/editor/core/controller.test.ts +++ b/src/editor/core/controller.test.ts @@ -166,7 +166,7 @@ describe("EditorController", () => { expect(pane().tabs[2]!.preview).toBe(false); useDocMeta.setState({ "/p/b.ts": { ...useDocMeta.getState()["/p/b.ts"]!, readOnly: true } }); expect(labels()).toContain("Close All Read-Only Tabs"); - ctrl.closeReadOnly(pane().id); + await ctrl.closeReadOnly(pane().id); expect(pane().tabs.map((t) => t.input.path)).toEqual(["/p/a.ts", "/p/c.ts"]); }); @@ -180,12 +180,13 @@ describe("EditorController", () => { useProject.setState({ lightEdit: false }); expect(openFilePaths(useTabs.getState())).toEqual([]); useProject.setState({ lightEdit: true }); - expect(openFilePaths(useTabs.getState())).toEqual(["/x/a.txt", "/y/b.md"]); - // Closing every tab leaves nothing to reopen. - ctrl.closeEverything(); + await vi.waitFor(() => expect(openFilePaths(useTabs.getState())).toEqual(["/x/a.txt", "/y/b.md"])); + // Clearing the session must not resurrect its other files on the next explicit open. + await ctrl.closeEverything(); useProject.setState({ lightEdit: false }); useProject.setState({ lightEdit: true }); - expect(openFilePaths(useTabs.getState())).toEqual([]); + await ctrl.openFile("/x/a.txt", { focus: false }); + expect(openFilePaths(useTabs.getState())).toEqual(["/x/a.txt"]); useProject.setState({ lightEdit: false }); }); @@ -197,7 +198,7 @@ describe("EditorController", () => { const key = Object.values(useTabs.getState().panes) .flatMap((p) => p.tabs) .find((t) => t.input.path === "/p/b.ts")!.input.key; - ctrl.closeInput(key); + await ctrl.closeInput(key); expect(openFilePaths(useTabs.getState())).toEqual(["/p/a.ts"]); }); @@ -223,31 +224,6 @@ describe("EditorController", () => { const s = useTabs.getState(); const pane = s.panes[s.activePane]!; const items = resolveMenuItems(ctrl.tabMenu(pane.id, pane.tabs[0]!)); - expect(items.map((i) => (i.kind === "separator" ? "-" : i.label))).toEqual([ - "Close", - "Close Other Tabs", - "Close Tabs to the Left", - "Close Tabs to the Right", - "Close All Tabs", - "Close Unmodified Tabs", - "-", - "Pin Tab", - "-", - "Split and Move Right", - "Split and Move Down", - "Split Right", - "Split Down", - "-", - "Open In", - "-", - "Local History", - "Git", - "-", - "Copy Path/Reference…", - "Rename File…", - "-", - "Configure Editor Tabs…", - ]); expect(items.find((i) => i.label === "Close Tabs to the Left")?.disabled).toBe(true); expect(items.find((i) => i.label === "Close Tabs to the Right")?.disabled).toBe(false); dispose(); diff --git a/src/editor/core/controller.ts b/src/editor/core/controller.ts index 40cbba52b..6c85fcf13 100644 --- a/src/editor/core/controller.ts +++ b/src/editor/core/controller.ts @@ -22,11 +22,15 @@ import { notifications, resolveFileEditor, useContributions, + type EditorCloseSession, type EditorInput, type OpenOptions, + type FileEditorResolution, type OpenEditorTab, } from "@/workbench/contributions"; import { useDistractionFree } from "@/workbench/distractionFree"; +import { ask } from "@/workbench/messages"; +import { t } from "@/platform/i18n"; import { fileIconName, showContextMenu, type MenuItem } from "@/ui"; import { usePlatform } from "@/platform/platformInfo"; import { @@ -46,6 +50,7 @@ import { moveTab, moveTabToNewSplit, neighbourPane, + openKeys, openTab, paneIds, paneOf, @@ -64,6 +69,16 @@ import { type TabsState, } from "./model/tabs"; import { DocumentManager, useDocMeta } from "./model/document"; +import { + closedDocuments, + combinedError, + customDocuments, + documentDirty, + documentLifecycle, + failures, + inputDirty, + type CustomDocument, +} from "./model/customDocuments"; import { EditorSlot, type SlotHost } from "./model/slot"; import { NavigationHistory } from "./model/navHistory"; import { EditHistory } from "./model/editHistory"; @@ -194,6 +209,16 @@ export class EditorController implements SlotHost { /** When each input key was last the active tab (a counter, across panes). */ private lastActive = new Map(); private activeClock = 0; + /** Serializes close decisions without retaining a stale tabs snapshot. */ + private removalTail: Promise = Promise.resolve(); + private closingDocuments = new Map>(); + private preparedDocuments = new Map(); + private closeSession: EditorCloseSession | null = null; + private closeCommitted = false; + private projectGeneration = 0; + private admissionGeneration = 0; + private closeRequests = 0; + private projectTransition: Promise = Promise.resolve(); constructor(backend: DocBackend = docBackend(), scheduler?: Scheduler) { this.manager = new DocumentManager( @@ -224,7 +249,7 @@ export class EditorController implements SlotHost { let prevRoot = useProject.getState().project?.root ?? null; let prevLight = useProject.getState().lightEdit; this.disposers.push( - useProject.subscribe((s) => { + useProject.subscribe((s, previous) => { if (s.settings !== prevSettings) { prevSettings = s.settings; tabOptions.activateOnClose = editorPrefs(s.settings).tabs @@ -233,16 +258,33 @@ export class EditorController implements SlotHost { } const root = s.project?.root ?? null; // Entering or leaving LightEdit replaces the backend's documents too. - if (root !== prevRoot || s.lightEdit !== prevLight) { + if (root !== prevRoot || s.lightEdit !== prevLight || (s.opening && s.project !== previous.project)) { prevRoot = root; prevLight = s.lightEdit; this.saveSession(); this.sessionRoot = null; - this.closeEverything(); + const generation = ++this.projectGeneration; + const documents = customDocuments(useTabs.getState(), useContributions.getState().editorKinds); + // The guard decided already. Disposal must never prompt or autosave + // custom documents after the backend has replaced their project. + useTabs.setState(initialTabsState(), true); + this.closeSession = null; + this.closeCommitted = false; + this.preparedDocuments.clear(); this.nav.clear(); this.edits.clear(); - if (root) this.restoreSession(root); - else if (s.lightEdit) this.restoreSession(LIGHT_EDIT_SESSION); + this.restoredTabs.clear(); + this.restoredViews.clear(); + this.projectTransition = this.releaseDocuments(documents.values()) + .catch((e) => this.documentError(e)) + .then(() => { + if (generation !== this.projectGeneration) return; + if (root) this.restoreSession(root); + else if (s.lightEdit) this.restoreSession(LIGHT_EDIT_SESSION); + }); + } else if (previous.opening && !s.opening && this.closeCommitted) { + // A failed project open returns to the same project. + this.closeSession?.cancel(); } }), ); @@ -252,7 +294,10 @@ export class EditorController implements SlotHost { if (this.manager.backend.onFsChanged) this.disposers.push(this.manager.backend.onFsChanged((e) => this.onFilesRemoved(e.removed))); const onWindowBlur = () => { - if (this.autoSave()) void this.manager.saveAll("auto"); + if (this.autoSave() && !this.closeSession) { + void this.manager.saveAll("auto"); + void this.saveCustomDocuments(true).catch((e) => this.documentError(e)); + } }; const onPageHide = () => this.saveSession(); if (typeof window !== "undefined") { @@ -271,11 +316,18 @@ export class EditorController implements SlotHost { stop() { this.saveSession(); this.sessionRoot = null; + ++this.projectGeneration; clearTimeout(this.sessionTimer); + const documents = customDocuments(useTabs.getState(), useContributions.getState().editorKinds); this.disposers.splice(0).forEach((d) => d()); + this.closeSession = null; + this.closeCommitted = false; + this.preparedDocuments.clear(); for (const slot of this.slots.values()) slot.dispose(); this.slots.clear(); this.manager.stop(); + useTabs.setState(initialTabsState(), true); + void this.releaseDocuments(documents.values()).catch((e) => this.documentError(e)); } autoSave(): boolean { @@ -302,13 +354,23 @@ export class EditorController implements SlotHost { return n; }, true); - // Autosave on tab switch (IDEA). - if (this.autoSave()) { + // Save only a document that remains open. A discarded/closing document must + // never be saved by the tab removal that follows its close decision. + if (this.autoSave() && !this.closeSession) { + const saved = new Set(); + const open = openKeys(next); for (const [paneId, pane] of Object.entries(prev.panes)) { const now = next.panes[paneId]; if (pane.active && now && now.active !== pane.active) { const slot = this.slots.get(pane.active); if (slot) void slot.model.save("auto").catch(() => {}); + const input = pane.tabs.find((t) => t.id === pane.active)?.input; + if (input && open.has(input.key) && !saved.has(input.key) && !this.preparedDocuments.has(input.key)) { + saved.add(input.key); + const lifecycle = documentLifecycle(input, useContributions.getState().editorKinds); + if (lifecycle && documentDirty({ input, lifecycle })) + void lifecycle.save(input).catch((e) => this.documentError(e)); + } } } } @@ -356,6 +418,9 @@ export class EditorController implements SlotHost { const session = loadSession(root); if (!session) return; const state = restoreTabsState(session, fileKey); + const kinds = useContributions.getState().editorKinds; + for (const pane of Object.values(state.panes)) + pane.tabs = pane.tabs.map((tab) => (documentLifecycle(tab.input, kinds) ? { ...tab, preview: false } : tab)); this.restoredViews = new Map(Object.entries(session.files)); this.restoredTabs = new Set(Object.values(state.panes).flatMap((p) => p.tabs.map((t) => t.id))); useTabs.setState(state, true); @@ -767,18 +832,43 @@ export class EditorController implements SlotHost { // --- opening --------------------------------------------------------------------------- + private openingGeneration() { + const generation = { project: this.projectGeneration, admission: this.admissionGeneration }; + this.checkOpening(generation); + return generation; + } + + private checkOpening(generation: { project: number; admission: number }) { + if (generation.project !== this.projectGeneration) throw new Error("The project changed while opening the editor"); + if (generation.admission !== this.admissionGeneration || this.closeSession || this.closeRequests) + throw new Error("The project is closing"); + } + async openFile(path: string, o: OpenOptions = {}) { + const generation = this.openingGeneration(); if (isVirtualUri(path)) return this.openVirtual(path, o); + await this.projectTransition; + this.checkOpening(generation); // Images, binary and huge files open in a contributed viewer unless already open as text. if (!o.asText && !findTab(useTabs.getState(), fileKey(path))) { const viewer = await resolveFileEditor(path); - if (viewer) return this.openInput(viewer, o); + if (viewer) { + try { + await this.admitInput(viewer.input, o, generation, viewer); + } finally { + await viewer.rollback(); + } + return; + } } - await this.openInput(fileInput(path), o); + await this.admitInput(fileInput(path), o, generation); } /** Opens the library class (or other server-provided document) behind `uri` read-only. */ async openVirtual(uri: string, o: OpenOptions = {}) { + const generation = this.openingGeneration(); + await this.projectTransition; + this.checkOpening(generation); let vd: VirtualDocument; try { vd = await virtualDocument(uri); @@ -791,10 +881,23 @@ export class EditorController implements SlotHost { }); return; } - await this.openInput(virtualInput(vd), o); + await this.admitInput(virtualInput(vd), o, generation); } async openInput(input: EditorInput, o: OpenOptions = {}) { + return this.admitInput(input, o, this.openingGeneration()); + } + + private async admitInput( + input: EditorInput, + o: OpenOptions, + generation: { project: number; admission: number }, + resolution?: FileEditorResolution, + ) { + await this.projectTransition; + this.checkOpening(generation); + await this.closingDocuments.get(input.key); + this.checkOpening(generation); this.navSuppress++; try { let s = useTabs.getState(); @@ -816,13 +919,17 @@ export class EditorController implements SlotHost { if (found && (!(o.side || o.pane) || found.paneId === paneId)) { paneId = found.paneId; s = activateTab(setActivePane(s, paneId), paneId, found.tab.id); - if (!o.preview) s = promoteTab(s, paneId, found.tab.id); + if (!o.preview || documentLifecycle(input, useContributions.getState().editorKinds)) + s = promoteTab(s, paneId, found.tab.id); tabId = found.tab.id; if (found.tab.input.data !== input.data && input.data !== undefined) - s = openTab(s, paneId, input, { preview: o.preview }).state; + s = openTab(s, paneId, input, { + preview: o.preview && !documentLifecycle(input, useContributions.getState().editorKinds), + }).state; } else { const tabs = editorPrefs().tabs; - const preview = !!o.preview && tabs.previewTab; + const preview = + !!o.preview && tabs.previewTab && !documentLifecycle(input, useContributions.getState().editorKinds); const r = openTab(s, paneId, input, { preview, ...(tabs.openAtEnd ? { index: s.panes[paneId]?.tabs.length ?? 0 } : {}), @@ -830,16 +937,27 @@ export class EditorController implements SlotHost { s = r.state; tabId = r.tabId; if (tabs.sortAlphabetically) s = sortTabs(s, paneId, (t) => tabSortKey(t.input)); - s = enforceTabLimit( - s, - paneId, - tabs.tabLimit, - tabs.closingPolicy, - (t) => !!(t.input.path && useDocMeta.getState()[t.input.path]?.dirty), - tabId, - ); } + resolution?.admit(); useTabs.setState(s, true); + const tabPrefs = editorPrefs().tabs; + await this.removeGuarded( + (current) => + current.panes[paneId] + ? enforceTabLimit( + current, + paneId, + tabPrefs.tabLimit, + tabPrefs.closingPolicy, + (tab) => + inputDirty(tab.input, useContributions.getState().editorKinds) ?? + !!(tab.input.path && useDocMeta.getState()[tab.input.path]?.dirty), + tabId, + (tab) => inputDirty(tab.input, useContributions.getState().editorKinds) === true, + ) + : current, + { protect: (tab) => inputDirty(tab.input, useContributions.getState().editorKinds) === true }, + ); if (!isTextInput(input)) { if (o.focus !== false) requestAnimationFrame(() => this.focus()); return; @@ -893,32 +1011,171 @@ export class EditorController implements SlotHost { useTabs.setState((s) => f(s), true); } - closeTab(paneId: string, tabId: string) { - this.apply((s) => (s.panes[paneId] ? closeTab(s, paneId, tabId) : s)); + private documentError(error: unknown) { + notifications().notify({ + level: "error", + title: t("app.unsaved.saveFailed"), + body: error instanceof Error ? error.message : String(error), + }); } - closeOthers(paneId: string, tabId: string) { - this.apply((s) => closeOthers(s, paneId, tabId)); + + private enqueueRemoval(run: () => Promise): Promise { + const result = this.removalTail.then(run); + this.removalTail = result.catch(() => {}); + return result; } - closeAll(paneId: string, keepPinned = false) { - this.apply((s) => closeAll(s, paneId, keepPinned)); + + private releaseDocuments(documents: Iterable): Promise { + const pending = [...documents].map((d) => { + const key = d.input.key; + const previous = this.closingDocuments.get(key) ?? Promise.resolve(); + const close = previous.then(() => d.lifecycle.close(d.input)); + this.closingDocuments.set(key, close); + return close.finally(() => { + if (this.closingDocuments.get(key) === close) this.closingDocuments.delete(key); + }); + }); + return Promise.allSettled(pending).then((results) => { + const errors = failures(results); + if (errors.length) throw combinedError(errors); + }); } - closeToRight(paneId: string, tabId: string) { - this.apply((s) => closeToRight(s, paneId, tabId)); + + /** Close exactly the originally selected tabs, evaluating references after every await. */ + private removeGuarded( + select: (s: TabsState) => TabsState, + options: { protect?: (tab: Tab) => boolean; clearHistory?: boolean } = {}, + ): Promise { + const generation = this.projectGeneration; + const initial = useTabs.getState(); + const planned = select(initial); + const remainingIds = new Set(Object.values(planned.panes).flatMap((p) => p.tabs.map((t) => t.id))); + const targets = new Set( + Object.values(initial.panes).flatMap((p) => p.tabs.filter((t) => !remainingIds.has(t.id)).map((t) => t.id)), + ); + const emptyPanes = paneIds(initial).filter((id) => !planned.panes[id] && !initial.panes[id]?.tabs.length); + return this.enqueueRemoval(async () => { + if (generation !== this.projectGeneration || this.closeSession) return false; + const prepared: CustomDocument[] = []; + const currentRemoval = (s: TabsState) => { + let next = s; + for (const [paneId, pane] of Object.entries(s.panes)) + for (const tab of pane.tabs) + if (targets.has(tab.id) && next.panes[paneId] && !options.protect?.(tab)) + next = closeTab(next, paneId, tab.id); + for (const id of emptyPanes) + if (next.panes[id]?.tabs.length === 0 && paneIds(next).length > 1) next = removePane(next, id); + return next; + }; + const stillLast = (d: CustomDocument) => { + if (generation !== this.projectGeneration) return false; + const current = useTabs.getState(); + return openKeys(current).has(d.input.key) && !openKeys(currentRemoval(current)).has(d.input.key); + }; + try { + const authorized = new Set(); + const discards = new Map(); + for (;;) { + const current = useTabs.getState(); + const d = closedDocuments(current, currentRemoval(current), useContributions.getState().editorKinds).find( + (document) => !authorized.has(document.input.key), + ); + if (!d) { + // A later Cancel/save failure must not discard earlier documents. + const discard = [...discards.values()].find(stillLast); + if (!discard) break; + discards.delete(discard.input.key); + await discard.lifecycle.discard(discard.input); + // The await may have changed which documents lose their last tab. + if (generation !== this.projectGeneration) return false; + continue; + } + if (!prepared.some((document) => document.input.key === d.input.key)) { + prepared.push(d); + this.preparedDocuments.set(d.input.key, d); + await d.lifecycle.prepareClose(d.input); + } + if (generation !== this.projectGeneration) return false; + if (!stillLast(d)) continue; + if (documentDirty(d)) { + const decision = await ask({ + title: t("app.unsaved.title"), + body: t("app.unsaved.body", { count: 1, project: d.input.path ?? d.input.key }), + buttons: [ + { id: "discard", label: t("app.unsaved.discard") }, + { id: "save", label: t("app.unsaved.save"), primary: true }, + ], + icon: "warning", + }); + if (!stillLast(d)) continue; + if (decision.button === "save") await d.lifecycle.save(d.input); + else if (decision.button === "discard") discards.set(d.input.key, d); + else return false; + } + authorized.add(d.input.key); + } + // Commit only the original target ids against the latest tab state. + if (generation !== this.projectGeneration) return false; + const before = useTabs.getState(); + const next = currentRemoval(before); + const closed = closedDocuments(before, next, useContributions.getState().editorKinds); + const closedKeys = new Set(closed.map((d) => d.input.key)); + for (const d of prepared) if (!closedKeys.has(d.input.key)) d.lifecycle.cancelClose(d.input); + // Remove frozen markers only after the tab reaction, to prevent autosave. + useTabs.setState(options.clearHistory ? { ...next, closed: [] } : next, true); + for (const d of prepared) this.preparedDocuments.delete(d.input.key); + prepared.length = 0; + await this.releaseDocuments(closed); + return true; + } catch (e) { + this.documentError(e); + return false; + } finally { + for (const d of prepared) { + if (this.preparedDocuments.get(d.input.key) === d) this.preparedDocuments.delete(d.input.key); + if (generation === this.projectGeneration) d.lifecycle.cancelClose(d.input); + } + } + }); } - closeToLeft(paneId: string, tabId: string) { - this.apply((s) => closeToLeft(s, paneId, tabId)); + + closeTab(paneId: string, tabId: string): Promise { + return this.removeGuarded((s) => (s.panes[paneId] ? closeTab(s, paneId, tabId) : s)); + } + closeOthers(paneId: string, tabId: string): Promise { + return this.removeGuarded((s) => (s.panes[paneId] ? closeOthers(s, paneId, tabId) : s), { + protect: (tab) => tab.pinned, + }); + } + closeAll(paneId: string, keepPinned = false): Promise { + return this.removeGuarded((s) => (s.panes[paneId] ? closeAll(s, paneId, keepPinned) : s), { + protect: (tab) => keepPinned && tab.pinned, + }); + } + closeToRight(paneId: string, tabId: string): Promise { + return this.removeGuarded((s) => (s.panes[paneId] ? closeToRight(s, paneId, tabId) : s), { + protect: (tab) => tab.pinned, + }); + } + closeToLeft(paneId: string, tabId: string): Promise { + return this.removeGuarded((s) => (s.panes[paneId] ? closeToLeft(s, paneId, tabId) : s), { + protect: (tab) => tab.pinned, + }); } /** Keeps tabs with unsaved changes or a VCS status (IDEA: tabs shown as modified). */ - closeUnmodified(paneId: string) { - const modified = (t: Tab) => { - const p = t.input.path; - return !!p && (!!useDocMeta.getState()[p]?.dirty || !!getFileDecoration(p)?.color); - }; - this.apply((s) => closeUnmodified(s, paneId, modified)); + closeUnmodified(paneId: string): Promise { + const modified = (tab: Tab) => + (inputDirty(tab.input, useContributions.getState().editorKinds) ?? false) || + !!(tab.input.path && (useDocMeta.getState()[tab.input.path]?.dirty || getFileDecoration(tab.input.path)?.color)); + return this.removeGuarded((s) => (s.panes[paneId] ? closeUnmodified(s, paneId, modified) : s), { + protect: (tab) => tab.pinned || modified(tab), + }); } - closeReadOnly(paneId: string) { + closeReadOnly(paneId: string): Promise { const writable = (t: Tab) => !t.input.path || !useDocMeta.getState()[t.input.path]?.readOnly; - this.apply((s) => closeUnmodified(s, paneId, writable)); + return this.removeGuarded((s) => (s.panes[paneId] ? closeUnmodified(s, paneId, writable) : s), { + protect: (tab) => tab.pinned || writable(tab), + }); } /** Reload from Disk: re-reads the open, unmodified documents of `paths`. */ async reloadFromDisk(paths: string[]) { @@ -991,7 +1248,11 @@ export class EditorController implements SlotHost { this.apply((s) => cycleTab(s, delta)); this.focus(); } - reopenClosed() { + async reopenClosed(): Promise { + const generation = this.projectGeneration; + await this.projectTransition; + await this.removalTail; + if (generation !== this.projectGeneration || this.closeSession) return; const r = reopenClosed(useTabs.getState()); if (!r) return; useTabs.setState(r.state, true); @@ -1013,12 +1274,16 @@ export class EditorController implements SlotHost { } } - unsplit(paneId: string) { - this.apply((s) => removePane(s, paneId)); + unsplit(paneId: string): Promise { + return this.removeGuarded((s) => (s.panes[paneId] ? removePane(s, paneId) : s)); } - unsplitAll(paneId: string) { - this.apply((s) => unsplitAll(s, paneId)); - this.focus(); + async unsplitAll(paneId: string): Promise { + const closed = await this.removeGuarded((s) => unsplitAll(s, paneId)); + if (closed) { + this.apply((s) => setActivePane(s, paneId)); + this.focus(); + } + return closed; } changeSplitOrientation(paneId: string) { this.apply((s) => changeSplitOrientation(s, paneId)); @@ -1041,35 +1306,46 @@ export class EditorController implements SlotHost { }); } - closeActiveTab() { + closeActiveTab(): Promise { const s = useTabs.getState(); const pane = s.panes[s.activePane]; - if (pane?.active) this.closeTab(pane.id, pane.active); - else if (paneIds(s).length > 1) this.unsplit(s.activePane); + if (pane?.active) return this.closeTab(pane.id, pane.active); + if (paneIds(s).length > 1) return this.unsplit(s.activePane); + return Promise.resolve(true); } - closeEverything() { - this.apply((s) => { - let next = s; - for (const id of paneIds(s)) if (next.panes[id]) next = closeAll(next, id); - return { ...next, closed: [] }; - }); + closeEverything(): Promise { + return this.removeGuarded( + (s) => { + let next = s; + for (const id of paneIds(s)) if (next.panes[id]) next = closeAll(next, id); + return next; + }, + { clearHistory: true }, + ); } private onFilesRemoved(removed: string[]) { if (!removed.length) return; - for (const path of removed) { - const isUnder = (k: string) => k === fileKey(path) || k.startsWith(fileKey(path) + "/"); - const s = useTabs.getState(); - const keys = new Set(); - for (const pane of Object.values(s.panes)) - for (const t of pane.tabs) if (isUnder(t.input.key)) keys.add(t.input.key); - for (const key of keys) { - const model = this.manager.get(key.slice(5)); - if (model?.getMeta().dirty) continue; - this.apply((st) => closeKey(st, key)); - this.nav.removePath(key.slice(5)); - this.edits.removePath(key.slice(5)); + const kinds = useContributions.getState().editorKinds; + const inputs = new Map(); + for (const pane of Object.values(useTabs.getState().panes)) + for (const tab of pane.tabs) { + const path = tab.input.path; + if (path && removed.some((root) => path === root || path.startsWith(`${root}/`))) + inputs.set(tab.input.key, tab.input); + } + for (const input of inputs.values()) { + if (inputDirty(input, kinds) === true || (input.path && this.manager.get(input.path)?.getMeta().dirty)) continue; + // A pending output/edit can make the document dirty before this queued + // removal runs. Recheck then; external deletion never discards edits. + void this.removeGuarded((s) => closeKey(s, input.key), { + protect: (tab) => + inputDirty(tab.input, kinds) === true || !!(tab.input.path && useDocMeta.getState()[tab.input.path]?.dirty), + }); + if (input.path) { + this.nav.removePath(input.path); + this.edits.removePath(input.path); } } } @@ -1125,8 +1401,8 @@ export class EditorController implements SlotHost { return [...out]; } - closePath(path: string) { - this.apply((s) => { + closePath(path: string): Promise { + return this.removeGuarded((s) => { const keys = new Set(); for (const pane of Object.values(s.panes)) for (const t of pane.tabs) if (t.input.path === path) keys.add(t.input.key); @@ -1136,8 +1412,8 @@ export class EditorController implements SlotHost { }); } - closeInput(key: string) { - this.apply((s) => closeKey(s, key)); + closeInput(key: string): Promise { + return this.removeGuarded((s) => closeKey(s, key)); } activeInput(): EditorInput | null { @@ -1164,8 +1440,96 @@ export class EditorController implements SlotHost { return { doc: slot.model.doc, path: slot.model.path, ...slot.caret() }; } - saveAll() { - return this.manager.saveAll("explicit"); + private async saveCustomDocuments(auto = false): Promise { + const documents = customDocuments(useTabs.getState(), useContributions.getState().editorKinds); + const results = await Promise.allSettled( + [...documents.values()] + .filter((d) => (!auto || !this.preparedDocuments.has(d.input.key)) && documentDirty(d)) + .map((d) => Promise.resolve().then(() => d.lifecycle.save(d.input))), + ); + const errors = failures(results); + if (errors.length) throw combinedError(errors); + } + + async saveAll(): Promise { + const results = await Promise.allSettled([this.manager.saveAll("explicit"), this.saveCustomDocuments()]); + const errors = failures(results); + if (errors.length) throw combinedError(errors); + } + + /** Hold every open custom document still until the project/window decision. */ + prepareClose(): Promise { + ++this.admissionGeneration; + ++this.closeRequests; + return this.enqueueRemoval(async () => { + if (this.closeSession) throw new Error("An editor close decision is already pending"); + const generation = this.projectGeneration; + const documents = [...customDocuments(useTabs.getState(), useContributions.getState().editorKinds).values()]; + const prepared: CustomDocument[] = []; + let settled = false; + const session: EditorCloseSession = { + dirty: 0, + discard: async () => { + if (settled || this.closeSession !== session || generation !== this.projectGeneration) + throw new Error("The editor close decision is no longer current"); + const results = await Promise.allSettled( + documents.filter(documentDirty).map((d) => Promise.resolve().then(() => d.lifecycle.discard(d.input))), + ); + const errors = failures(results); + if (errors.length) throw combinedError(errors); + }, + cancel: () => { + if (settled) return; + settled = true; + if (this.closeSession !== session || generation !== this.projectGeneration) return; + this.closeSession = null; + this.closeCommitted = false; + for (const d of prepared) { + this.preparedDocuments.delete(d.input.key); + d.lifecycle.cancelClose(d.input); + } + }, + commit: () => { + if (settled || this.closeSession !== session || generation !== this.projectGeneration) + throw new Error("The editor close decision is no longer current"); + this.closeCommitted = true; + }, + }; + this.closeSession = session; + try { + for (const d of documents) { + prepared.push(d); + this.preparedDocuments.set(d.input.key, d); + await d.lifecycle.prepareClose(d.input); + if (generation !== this.projectGeneration) throw new Error("The project changed during close"); + if (this.closeSession !== session) { + d.lifecycle.cancelClose(d.input); + throw new Error("The editor close decision was cancelled"); + } + } + session.dirty = documents.filter(documentDirty).length; + return session; + } catch (e) { + session.cancel(); + throw e; + } + }).finally(() => { + --this.closeRequests; + }); + } + + cancelClose(): void { + this.closeSession?.cancel(); + } + + documentHistory(command: "undo" | "redo"): boolean { + const input = this.activeTabOf()?.input; + if (!input) return false; + const lifecycle = documentLifecycle(input, useContributions.getState().editorKinds); + const run = lifecycle?.[command]; + if (!run) return false; + void run.call(lifecycle, input).catch((e) => this.documentError(e)); + return true; } /** Context menu of a tab (IDEA's EditorTabPopupMenu). */ @@ -1181,21 +1545,21 @@ export class EditorController implements SlotHost { run(); }; return [ - { label: "Close", shortcut: shortcutLabel("editor.closeTab"), run: () => this.closeTab(paneId, tab.id) }, - { label: "Close Other Tabs", disabled: tabs.length < 2, run: () => this.closeOthers(paneId, tab.id) }, - { label: "Close Tabs to the Left", disabled: index <= 0, run: () => this.closeToLeft(paneId, tab.id) }, + { label: "Close", shortcut: shortcutLabel("editor.closeTab"), run: () => void this.closeTab(paneId, tab.id) }, + { label: "Close Other Tabs", disabled: tabs.length < 2, run: () => void this.closeOthers(paneId, tab.id) }, + { label: "Close Tabs to the Left", disabled: index <= 0, run: () => void this.closeToLeft(paneId, tab.id) }, { label: "Close Tabs to the Right", disabled: index < 0 || index >= tabs.length - 1, - run: () => this.closeToRight(paneId, tab.id), + run: () => void this.closeToRight(paneId, tab.id), }, - { label: "Close All Tabs", run: () => this.closeAll(paneId) }, - { label: "Close Unmodified Tabs", run: () => this.closeUnmodified(paneId) }, + { label: "Close All Tabs", run: () => void this.closeAll(paneId) }, + { label: "Close Unmodified Tabs", run: () => void this.closeUnmodified(paneId) }, ...(tabs.some((t) => t.pinned) - ? ([{ label: "Close All but Pinned", run: () => this.closeAll(paneId, true) }] as MenuItem[]) + ? ([{ label: "Close All but Pinned", run: () => void this.closeAll(paneId, true) }] as MenuItem[]) : []), ...(tabs.some((t) => !t.pinned && !!t.input.path && useDocMeta.getState()[t.input.path]?.readOnly) - ? ([{ label: "Close All Read-Only Tabs", run: () => this.closeReadOnly(paneId) }] as MenuItem[]) + ? ([{ label: "Close All Read-Only Tabs", run: () => void this.closeReadOnly(paneId) }] as MenuItem[]) : []), { kind: "separator" }, { label: tab.pinned ? "Unpin Tab" : "Pin Tab", icon: "pin", run: () => this.togglePin(paneId, tab.id) }, @@ -1220,8 +1584,8 @@ export class EditorController implements SlotHost { { label: "Move to Opposite Group", run: () => this.moveToOpposite(paneId, tab.id) }, { label: "Open in Opposite Group", run: () => this.openInOpposite(paneId, tab.id) }, { label: "Change Splitter Orientation", run: () => this.changeSplitOrientation(paneId) }, - { label: "Unsplit", run: () => this.unsplit(paneId) }, - { label: "Unsplit All", run: () => this.unsplitAll(paneId) }, + { label: "Unsplit", run: () => void this.unsplit(paneId) }, + { label: "Unsplit All", run: () => void this.unsplitAll(paneId) }, ] as MenuItem[]) : []), { kind: "separator" }, diff --git a/src/editor/core/customDocuments.test.ts b/src/editor/core/customDocuments.test.ts new file mode 100644 index 000000000..b9f8a6993 --- /dev/null +++ b/src/editor/core/customDocuments.test.ts @@ -0,0 +1,660 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { EditorController, useTabs } from "./controller"; +import { createMemoryBackend } from "./backend"; +import { findTab, initialTabsState, paneIds } from "./model/tabs"; +import { useDocMeta } from "./model/document"; +import { useProject } from "@/platform/project"; +import { ask } from "@/workbench/messages"; +import type * as Messages from "@/workbench/messages"; +import * as virtualDocs from "@/platform/virtualDocs"; +import type { VirtualDocument } from "@/platform/ipc"; +import { + registerEditorKind, + registerFileEditorProvider, + type EditorDocumentLifecycle, + type EditorInput, +} from "@/workbench/contributions"; + +vi.mock("@/workbench/messages", async (original) => ({ + ...(await original()), + ask: vi.fn(), +})); + +class TestDocument implements EditorDocumentLifecycle { + saved = "original"; + text = "edited"; + pending = ""; + running = false; + frozen = false; + closed = false; + saves = 0; + releases = 0; + saveError: Error | null = null; + prepareGate: Promise | null = null; + prepareError: Error | null = null; + + isDirty() { + return this.text !== this.saved || this.pending !== ""; + } + async flush() { + this.text += this.pending; + this.pending = ""; + } + async save() { + await this.flush(); + if (this.saveError) throw this.saveError; + this.saves++; + this.saved = this.text; + } + async prepareClose() { + this.frozen = true; + await this.flush(); + if (this.prepareGate) await this.prepareGate; + if (this.running) this.text += "\nfinal output"; + this.running = false; + if (this.prepareError) throw this.prepareError; + } + cancelClose() { + this.frozen = false; + } + async discard() { + this.pending = ""; + this.text = this.saved; + } + async close() { + this.closed = true; + this.running = false; + this.releases++; + } +} + +const input = (name = "a"): EditorInput => ({ + kind: "lifecycle-test", + key: `custom:${name}`, + path: `/p/${name}.ipynb`, +}); +let ctrl: EditorController; +let documents: Map; +let unregister: () => void; + +beforeEach(() => { + useTabs.setState(initialTabsState(), true); + useProject.setState({ project: null, settings: { files: { autoSave: false } } as never }); + localStorage.clear(); + documents = new Map(); + const sessions = documents; + const lifecycle: EditorDocumentLifecycle = { + isDirty: (i) => sessions.get(i.key)!.isDirty(), + flush: (i) => sessions.get(i.key)!.flush(), + save: (i) => sessions.get(i.key)!.save(), + prepareClose: (i) => sessions.get(i.key)!.prepareClose(), + cancelClose: (i) => sessions.get(i.key)!.cancelClose(), + discard: (i) => sessions.get(i.key)!.discard(), + close: (i) => sessions.get(i.key)!.close(), + }; + unregister = registerEditorKind({ + kind: "lifecycle-test", + component: () => null, + title: (i) => i.key, + document: lifecycle, + }); + ctrl = new EditorController(createMemoryBackend()); + ctrl.start(); + vi.mocked(ask).mockReset(); + vi.mocked(ask).mockResolvedValue({ button: "cancel", checked: false }); +}); + +afterEach(() => { + ctrl.stop(); + unregister(); + useDocMeta.setState({}, true); + useProject.setState({ project: null, settings: null }); + localStorage.clear(); +}); + +async function open(name = "a") { + const i = input(name); + const doc = new TestDocument(); + documents.set(i.key, doc); + await ctrl.openInput(i, { preview: true, focus: false }); + return doc; +} + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +describe("custom editor document lifetime", () => { + it("rejects a delayed virtual document from the previous project", async () => { + const started = deferred(); + const response = deferred(); + const probe = vi.spyOn(virtualDocs, "virtualDocument").mockImplementationOnce(() => { + started.resolve(); + return response.promise; + }); + try { + const opening = ctrl.openVirtual("jdt://old/Class", { focus: false }); + await started.promise; + useProject.setState({ project: { root: "/next", name: "next", settings: {} } as never }); + const rejected = expect(opening).rejects.toThrow("project changed"); + response.resolve({ + uri: "jdt://old/Class", + text: "class Old {}", + languageId: "java", + title: "Class", + description: null, + decompiled: true, + }); + await rejected; + expect(findTab(useTabs.getState(), "virtual:jdt://old/Class")).toBeNull(); + } finally { + probe.mockRestore(); + } + }); + + it.each(["pending", "cancelled", "replaced"] as const)( + "rolls back delayed provider acquisition after a %s project close", + async (route) => { + const started = deferred(); + const response = deferred(); + let acquired = false; + let installed = false; + const dispose = registerFileEditorProvider({ + id: "delayed-lifecycle-test", + order: -100, + resolve: async (_path, context) => { + started.resolve(); + await response.promise; + acquired = true; + context.onAdmission( + () => { + installed = true; + }, + async () => { + acquired = false; + }, + ); + return input(); + }, + }); + try { + const opening = ctrl.openFile(input().path!, { focus: false }); + await started.promise; + const session = await ctrl.prepareClose(); + if (route === "cancelled") session.cancel(); + if (route === "replaced") { + session.commit(); + useProject.setState({ project: { root: "/replacement", name: "replacement", settings: {} } as never }); + } + const rejected = expect(opening).rejects.toThrow(route === "replaced" ? "project changed" : "closing"); + response.resolve(); + await rejected; + session.cancel(); + expect(installed).toBe(false); + expect(acquired).toBe(false); + expect(findTab(useTabs.getState(), input().key)).toBeNull(); + } finally { + dispose(); + } + }, + ); + + it("rejects a new open before invoking a provider while close admission is pending", async () => { + let probes = 0; + const tracked = registerFileEditorProvider({ + id: "tracked-probe", + order: -101, + resolve: () => { + ++probes; + return input(); + }, + }); + try { + const closing = ctrl.prepareClose(); + await expect(ctrl.openFile(input().path!, { focus: false })).rejects.toThrow("closing"); + (await closing).cancel(); + expect(probes).toBe(0); + } finally { + tracked(); + } + }); + + it("retains pending changes and the last tab on Cancel", async () => { + const doc = await open(); + doc.pending = " pending"; + await expect(ctrl.closeInput(input().key)).resolves.toBe(false); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(doc.text).toBe("edited pending"); + expect(doc.saved).toBe("original"); + expect(doc.frozen).toBe(false); + expect(doc.closed).toBe(false); + }); + + it("saves a drained execution checkpoint before removing the last tab", async () => { + const doc = await open(); + doc.running = true; + vi.mocked(ask).mockResolvedValue({ button: "save", checked: false }); + await expect(ctrl.closePath(input().path!)).resolves.toBe(true); + expect(doc.saved).toBe("edited\nfinal output"); + expect(doc.running).toBe(false); + expect(doc.closed).toBe(true); + expect(findTab(useTabs.getState(), input().key)).toBeNull(); + }); + + it("discards without a removal-triggered autosave", async () => { + const doc = await open(); + useProject.setState({ settings: { files: { autoSave: true } } as never }); + doc.pending = " unsaved"; + vi.mocked(ask).mockResolvedValue({ button: "discard", checked: false }); + await ctrl.closeEverything(); + expect(doc.saved).toBe("original"); + expect(doc.text).toBe("original"); + expect(doc.saves).toBe(0); + expect(doc.closed).toBe(true); + }); + + it("keeps dirty state and tabs when save fails", async () => { + const doc = await open(); + doc.saveError = new Error("external modification conflict"); + vi.mocked(ask).mockResolvedValue({ button: "save", checked: false }); + await expect(ctrl.closeActiveTab()).resolves.toBe(false); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(doc.isDirty()).toBe(true); + expect(doc.frozen).toBe(false); + expect(doc.closed).toBe(false); + }); + + it("closing one split neither prompts nor stops its shared execution", async () => { + const doc = await open(); + doc.running = true; + const first = useTabs.getState().activePane; + ctrl.split(first, "row"); + await ctrl.unsplit(first); + expect(doc.running).toBe(true); + expect(doc.closed).toBe(false); + expect(doc.frozen).toBe(false); + expect(ask).not.toHaveBeenCalled(); + vi.mocked(ask).mockResolvedValue({ button: "discard", checked: false }); + await ctrl.closeInput(input().key); + expect(doc.releases).toBe(1); + expect(doc.running).toBe(false); + }); + + it("does not discard shared state when a split opens while a close prompt awaits", async () => { + const doc = await open(); + const first = useTabs.getState().activePane; + let answer!: (value: { button: string; checked: boolean }) => void; + const asked = new Promise((resolve) => { + vi.mocked(ask).mockImplementationOnce(() => { + resolve(); + return new Promise((r) => { + answer = r; + }); + }); + }); + const closing = ctrl.closeAll(first); + await asked; + ctrl.split(first, "row"); + answer({ button: "discard", checked: false }); + await closing; + expect(paneIds(useTabs.getState())).toHaveLength(1); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(doc.text).toBe("edited"); + expect(doc.closed).toBe(false); + expect(doc.frozen).toBe(false); + }); + + it("saves a split document once and reports failures after trying every document", async () => { + const a = await open(); + ctrl.split(useTabs.getState().activePane, "row"); + const b = await open("b"); + a.saveError = new Error("disk full"); + await expect(ctrl.saveAll()).rejects.toThrow("disk full"); + expect(a.saved).toBe("original"); + expect(a.isDirty()).toBe(true); + expect(b.saved).toBe("edited"); + a.saveError = null; + await ctrl.saveAll(); + expect(a.saves).toBe(1); + expect(b.saves).toBe(1); + }); + + it("never previews or evicts a dirty custom editor under LRU pressure", async () => { + useProject.setState({ + settings: { files: { autoSave: false }, editor: { tabs: { tabLimit: 1, closingPolicy: "lru" } } } as never, + }); + const doc = await open(); + await ctrl.openInput({ kind: "read-only-test", key: "other" }, { preview: true, focus: false }); + const tab = findTab(useTabs.getState(), input().key)?.tab; + expect(tab?.preview).toBe(false); + expect(doc.closed).toBe(false); + expect(doc.text).toBe("edited"); + expect(findTab(useTabs.getState(), "other")).not.toBeNull(); + }); + + it("Close Unmodified retains pending custom changes", async () => { + const doc = await open(); + doc.text = doc.saved; + doc.pending = "new typing"; + await ctrl.openInput({ kind: "read-only-test", key: "other" }, { focus: false }); + await ctrl.closeUnmodified(useTabs.getState().activePane); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(findTab(useTabs.getState(), "other")).toBeNull(); + expect(doc.pending).toBe("new typing"); + }); + + it("Close Read-Only retains a dirty custom document on Cancel and releases it after Discard", async () => { + const doc = await open(); + const pane = useTabs.getState().activePane; + useDocMeta.setState({ [input().path!]: { readOnly: true } as never }); + const cancelled = await ctrl.closeReadOnly(pane); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(doc.text).toBe("edited"); + expect(doc.closed).toBe(false); + expect(cancelled).toBe(false); + + vi.mocked(ask).mockResolvedValueOnce({ button: "discard", checked: false }); + expect(await ctrl.closeReadOnly(pane)).toBe(true); + expect(findTab(useTabs.getState(), input().key)).toBeNull(); + expect(doc.saved).toBe("original"); + expect(doc.closed).toBe(true); + }); + + it("Unsplit All preserves cancelled documents and only releases their last removed view", async () => { + const kept = await open(); + const keepPane = useTabs.getState().activePane; + ctrl.split(keepPane, "row"); + const removedPane = useTabs.getState().activePane; + const removed = await open("removed"); + + const cancelled = await ctrl.unsplitAll(keepPane); + expect(paneIds(useTabs.getState())).toContain(removedPane); + expect(findTab(useTabs.getState(), input("removed").key)).not.toBeNull(); + expect(removed.text).toBe("edited"); + expect(removed.closed).toBe(false); + expect(cancelled).toBe(false); + + vi.mocked(ask).mockResolvedValueOnce({ button: "discard", checked: false }); + expect(await ctrl.unsplitAll(keepPane)).toBe(true); + expect(paneIds(useTabs.getState())).toEqual([keepPane]); + expect(removed.saved).toBe("original"); + expect(removed.closed).toBe(true); + expect(kept.text).toBe("edited"); + expect(kept.closed).toBe(false); + }); + + it("prepares hidden documents, deduplicates split state, and resumes after a failed project transition", async () => { + const a = await open(); + ctrl.split(useTabs.getState().activePane, "row"); + const b = await open("b"); + b.running = true; + const session = await ctrl.prepareClose(); + expect(session.dirty).toBe(2); + expect(a.frozen).toBe(true); + expect(b.running).toBe(false); + await ctrl.saveAll(); + expect(b.saved).toBe("edited\nfinal output"); + session.commit(); + ctrl.cancelClose(); + expect(a.frozen).toBe(false); + expect(b.frozen).toBe(false); + expect(a.closed).toBe(false); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + }); + + it("project replacement releases old split documents once, without saving discarded changes", async () => { + const doc = await open(); + ctrl.split(useTabs.getState().activePane, "row"); + const session = await ctrl.prepareClose(); + await session.discard(); + session.commit(); + useProject.setState({ project: { root: "/new-project", name: "new", settings: {} } as never }); + await ctrl.openInput({ kind: "read-only-test", key: "new-project-tab" }, { focus: false }); + expect(doc.saved).toBe("original"); + expect(doc.saves).toBe(0); + expect(doc.releases).toBe(1); + expect(findTab(useTabs.getState(), input().key)).toBeNull(); + expect(findTab(useTabs.getState(), "new-project-tab")).not.toBeNull(); + }); + + it.each(["all", "others", "left", "right", "unsplit"] as const)( + "%s respects a last-document Cancel without removing other tabs", + async (route) => { + const marker: EditorInput = { kind: "read-only-test", key: "marker" }; + if (route !== "left") await ctrl.openInput(marker, { focus: false }); + const doc = await open(); + const pane = useTabs.getState().activePane; + if (route === "left") await ctrl.openInput(marker, { focus: false }); + if (route === "unsplit") await ctrl.openInput(marker, { side: true, focus: false }); + const markerTab = findTab(useTabs.getState(), marker.key)!.tab.id; + let result: boolean; + switch (route) { + case "all": + result = await ctrl.closeAll(pane); + break; + case "others": + result = await ctrl.closeOthers(pane, markerTab); + break; + case "left": + result = await ctrl.closeToLeft(pane, markerTab); + break; + case "right": + result = await ctrl.closeToRight(pane, markerTab); + break; + case "unsplit": + result = await ctrl.unsplit(pane); + break; + } + expect(result).toBe(false); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(findTab(useTabs.getState(), marker.key)).not.toBeNull(); + expect(doc.text).toBe("edited"); + expect(doc.closed).toBe(false); + expect(doc.frozen).toBe(false); + }, + ); + + it("keeps newly opened tabs when a bulk close completes", async () => { + await open(); + let answer!: (value: { button: string; checked: boolean }) => void; + let prompted!: () => void; + const asked = new Promise((resolve) => { + prompted = resolve; + }); + vi.mocked(ask).mockImplementationOnce(() => { + prompted(); + return new Promise((resolve) => { + answer = resolve; + }); + }); + const closing = ctrl.closeEverything(); + await asked; + const opening = ctrl.openInput({ kind: "read-only-test", key: "new-tab" }, { focus: false }); + await vi.waitFor(() => expect(findTab(useTabs.getState(), "new-tab")).not.toBeNull()); + answer({ button: "discard", checked: false }); + await Promise.all([closing, opening]); + expect(findTab(useTabs.getState(), "new-tab")).not.toBeNull(); + expect(findTab(useTabs.getState(), input().key)).toBeNull(); + }); + + it("does not prompt or blur-autosave a document until execution has drained", async () => { + const doc = await open(); + useProject.setState({ settings: { files: { autoSave: true } } as never }); + let drain!: () => void; + doc.prepareGate = new Promise((resolve) => { + drain = resolve; + }); + doc.running = true; + const closing = ctrl.closeInput(input().key); + await vi.waitFor(() => expect(doc.frozen).toBe(true)); + window.dispatchEvent(new Event("blur")); + expect(doc.saved).toBe("original"); + expect(ask).not.toHaveBeenCalled(); + drain(); + await closing; + expect(doc.text).toBe("edited\nfinal output"); + expect(doc.saved).toBe("original"); + expect(doc.frozen).toBe(false); + expect(doc.closed).toBe(false); + }); + + it("unfreezes every prepared document when another preparation fails", async () => { + const a = await open(); + const b = await open("b"); + b.prepareError = new Error("pending source could not be flushed"); + await expect(ctrl.prepareClose()).rejects.toThrow("could not be flushed"); + expect(a.frozen).toBe(false); + expect(b.frozen).toBe(false); + expect(a.text).toBe("edited"); + expect(b.text).toBe("edited"); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(findTab(useTabs.getState(), input("b").key)).not.toBeNull(); + }); + + it("evicts a clean custom document through its last-reference lifecycle", async () => { + useProject.setState({ settings: { files: { autoSave: false }, editor: { tabs: { tabLimit: 1 } } } as never }); + const doc = await open(); + doc.text = doc.saved; + await ctrl.openInput({ kind: "read-only-test", key: "new-tab" }, { focus: false }); + expect(findTab(useTabs.getState(), input().key)).toBeNull(); + expect(findTab(useTabs.getState(), "new-tab")).not.toBeNull(); + expect(doc.releases).toBe(1); + expect(doc.saved).toBe("original"); + expect(ask).not.toHaveBeenCalled(); + }); + + it("releases old documents when a replacement project has the same root", async () => { + useProject.setState({ project: { root: "/p", name: "first", settings: {} } as never }); + const doc = await open(); + const session = await ctrl.prepareClose(); + await session.discard(); + session.commit(); + useProject.setState({ opening: true }); + useProject.setState({ project: { root: "/p", name: "replacement", settings: {} } as never }); + useProject.setState({ opening: false }); + await ctrl.openInput({ kind: "read-only-test", key: "replacement-tab" }, { focus: false }); + expect(doc.releases).toBe(1); + expect(findTab(useTabs.getState(), input().key)).toBeNull(); + expect(findTab(useTabs.getState(), "replacement-tab")).not.toBeNull(); + }); + + it.each(["cancel", "save"] as const)( + "does not apply an earlier bulk discard when a later %s aborts closing", + async (button) => { + const a = await open(); + const b = await open("b"); + b.saveError = new Error("disk full"); + vi.mocked(ask) + .mockResolvedValueOnce({ button: "discard", checked: false }) + .mockResolvedValueOnce({ button, checked: false }); + await expect(ctrl.closeEverything()).resolves.toBe(false); + expect(a.text).toBe("edited"); + expect(a.saved).toBe("original"); + expect(b.text).toBe("edited"); + expect(a.frozen).toBe(false); + expect(b.frozen).toBe(false); + expect(a.closed).toBe(false); + expect(b.closed).toBe(false); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(findTab(useTabs.getState(), input("b").key)).not.toBeNull(); + }, + ); + + it("guards a document that becomes the last split reference during another close prompt", async () => { + const a = await open(); + const first = useTabs.getState().activePane; + const firstTab = findTab(useTabs.getState(), input().key, first)!.tab.id; + ctrl.split(first, "row"); + const second = useTabs.getState().activePane; + const b = await open("b"); + await ctrl.openInput({ kind: "read-only-test", key: "marker" }, { focus: false }); + const marker = findTab(useTabs.getState(), "marker")!.tab.id; + let answer!: (value: { button: string; checked: boolean }) => void; + const asked = new Promise((resolve) => { + vi.mocked(ask).mockImplementationOnce(() => { + resolve(); + return new Promise((r) => { + answer = r; + }); + }); + }); + const closing = ctrl.closeOthers(second, marker); + await asked; + // Moving onto the existing split merges its duplicate, leaving the + // originally targeted tab as the only remaining reference. + ctrl.moveTab(first, firstTab, second, 0); + answer({ button: "discard", checked: false }); + await expect(closing).resolves.toBe(false); + expect(a.text).toBe("edited"); + expect(b.text).toBe("edited"); + expect(a.closed).toBe(false); + expect(b.closed).toBe(false); + expect(a.frozen).toBe(false); + expect(b.frozen).toBe(false); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + }); + + it("retains a tab that becomes dirty while a tab-limit close drains execution", async () => { + useProject.setState({ settings: { files: { autoSave: false }, editor: { tabs: { tabLimit: 1 } } } as never }); + const doc = await open(); + doc.text = doc.saved; + doc.running = true; + let drain!: () => void; + doc.prepareGate = new Promise((resolve) => { + drain = resolve; + }); + const opening = ctrl.openInput({ kind: "read-only-test", key: "other" }, { focus: false }); + await vi.waitFor(() => expect(doc.frozen).toBe(true)); + drain(); + await opening; + expect(doc.text).toBe("original\nfinal output"); + expect(doc.saved).toBe("original"); + expect(doc.closed).toBe(false); + expect(doc.frozen).toBe(false); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + expect(ask).not.toHaveBeenCalled(); + }); + + it("attempts remaining saves when a lifecycle callback throws synchronously", async () => { + const a = await open(); + const b = await open("b"); + a.save = () => { + throw new Error("session unavailable"); + }; + await expect(ctrl.saveAll()).rejects.toThrow("session unavailable"); + expect(a.saved).toBe("original"); + expect(a.text).toBe("edited"); + expect(b.saved).toBe("edited"); + }); + + it("rejects authorization from a cancelled project close session", async () => { + const doc = await open(); + const session = await ctrl.prepareClose(); + session.cancel(); + expect(() => session.commit()).toThrow(); + expect(doc.frozen).toBe(false); + expect(doc.text).toBe("edited"); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + }); + + it("does not return an authorized session after preparation was cancelled in flight", async () => { + const doc = await open(); + let drain!: () => void; + doc.prepareGate = new Promise((resolve) => { + drain = resolve; + }); + const preparing = ctrl.prepareClose(); + await vi.waitFor(() => expect(doc.frozen).toBe(true)); + ctrl.cancelClose(); + drain(); + await expect(preparing).rejects.toThrow("cancelled"); + expect(doc.frozen).toBe(false); + expect(doc.text).toBe("edited"); + expect(doc.closed).toBe(false); + expect(findTab(useTabs.getState(), input().key)).not.toBeNull(); + }); +}); diff --git a/src/editor/core/index.ts b/src/editor/core/index.ts index a83a1843c..7330eed8f 100644 --- a/src/editor/core/index.ts +++ b/src/editor/core/index.ts @@ -47,6 +47,9 @@ export function activate(): void | (() => void) { onActiveInputChange: (fn) => ctrl.onActiveInputChange(fn), closePath: (path) => ctrl.closePath(path), closeInput: (key) => ctrl.closeInput(key), + prepareClose: () => ctrl.prepareClose(), + cancelClose: () => ctrl.cancelClose(), + documentHistory: (command) => ctrl.documentHistory(command), reformat: async (path) => { const handle = editorsFor(path)[0]; return !!handle && (await formatDocument(handle.view, { explicit: false, selectionOnly: false })).applied; diff --git a/src/editor/core/model/customDocuments.ts b/src/editor/core/model/customDocuments.ts new file mode 100644 index 000000000..4d6e09fb5 --- /dev/null +++ b/src/editor/core/model/customDocuments.ts @@ -0,0 +1,67 @@ +/** + * Documents of custom editor kinds (`EditorKindContribution.document`, e.g. + * notebooks): one per input key, however many tabs and splits show it. The + * editor controller keeps them alive while any tab shows the key and closes + * them after the last one goes. + */ +import type { EditorDocumentLifecycle, EditorInput, EditorKindContribution } from "@/workbench/contributions"; +import { openKeys, paneIds, type TabsState } from "./tabs"; + +export interface CustomDocument { + input: EditorInput; + lifecycle: EditorDocumentLifecycle; +} + +type Kinds = ReadonlyMap; + +export function documentLifecycle(input: EditorInput, kinds: Kinds): EditorDocumentLifecycle | undefined { + return kinds.get(input.kind)?.document; +} + +/** Open custom documents by input key, in visual tab order (the first tab's input). */ +export function customDocuments(s: TabsState, kinds: Kinds): Map { + const out = new Map(); + for (const id of paneIds(s)) + for (const t of s.panes[id]?.tabs ?? []) { + if (out.has(t.input.key)) continue; + const lifecycle = documentLifecycle(t.input, kinds); + if (lifecycle) out.set(t.input.key, { input: t.input, lifecycle }); + } + return out; +} + +/** Custom documents open in `prev` that no tab of `next` shows any more. */ +export function closedDocuments(prev: TabsState, next: TabsState, kinds: Kinds): CustomDocument[] { + const open = openKeys(next); + return [...customDocuments(prev, kinds).values()].filter((d) => !open.has(d.input.key)); +} + +/** Unsaved changes; a failing check counts as unsaved so nothing is dropped on its word. */ +export function documentDirty(d: CustomDocument): boolean { + try { + return d.lifecycle.isDirty(d.input); + } catch (e) { + console.warn(`editor document ${d.input.key}: isDirty failed`, e); + return true; + } +} + +/** Whether the tab's input is a custom document with unsaved changes; null when it is no custom document. */ +export function inputDirty(input: EditorInput, kinds: Kinds): boolean | null { + const lifecycle = documentLifecycle(input, kinds); + return lifecycle ? documentDirty({ input, lifecycle }) : null; +} + +/** The rejection reasons of `results` (an AggregateError contributes each of its errors). */ +export function failures(results: readonly PromiseSettledResult[]): unknown[] { + return results.flatMap((r) => + r.status === "fulfilled" ? [] : r.reason instanceof AggregateError ? r.reason.errors : [r.reason], + ); +} + +/** One error for `errors` (non-empty): itself when single, else an AggregateError with the first message. */ +export function combinedError(errors: readonly unknown[]): unknown { + if (errors.length === 1) return errors[0]; + const first = errors[0]; + return new AggregateError(errors, first instanceof Error ? first.message : String(first)); +} diff --git a/src/editor/core/model/document.test.ts b/src/editor/core/model/document.test.ts index f21b62e60..cc5f7e3d4 100644 --- a/src/editor/core/model/document.test.ts +++ b/src/editor/core/model/document.test.ts @@ -100,6 +100,28 @@ describe("DocumentManager + slots", () => { expect(backend.docs.size).toBe(0); expect(manager.get("/p/a.txt")).toBeUndefined(); }); + + it("Save All rejects a failed document save after saving the remaining dirty documents", async () => { + const { backend, manager } = setup({ "/p/a.txt": "a\n", "/p/b.txt": "b\n" }, { files: { autoSave: false } }); + try { + const a = await manager.acquire("/p/a.txt"); + const b = await manager.acquire("/p/b.txt"); + backend.patch(a.doc, [{ from: 0, to: 0, insert: "changed " }]); + backend.patch(b.doc, [{ from: 0, to: 0, insert: "changed " }]); + const save = backend.save.bind(backend); + vi.spyOn(backend, "save").mockImplementation(async (id) => { + if (id === a.doc) throw new Error("disk full"); + await save(id); + }); + await expect(manager.saveAll("explicit")).rejects.toThrow("disk full"); + expect(backend.files.get("/p/a.txt")).toBe("a\n"); + expect(backend.files.get("/p/b.txt")).toBe("changed b\n"); + expect(a.getMeta().dirty).toBe(true); + expect(b.getMeta().dirty).toBe(false); + } finally { + manager.stop(); + } + }); }); describe("save cleanup", () => { diff --git a/src/editor/core/model/document.ts b/src/editor/core/model/document.ts index db0f7ba6f..2a458aaaf 100644 --- a/src/editor/core/model/document.ts +++ b/src/editor/core/model/document.ts @@ -735,6 +735,7 @@ export class DocumentManager { for (const m of this.byPath.values()) m.applySettings(); } + /** Saves every document; failures are notified, and an explicit Save All also rejects with them. */ async saveAll(reason: "explicit" | "auto") { const errors: unknown[] = []; await Promise.all( @@ -747,6 +748,8 @@ export class DocumentManager { if (errors.length) { const e = errors[0]; this.notify({ level: "error", title: "Could not save", body: e instanceof Error ? e.message : String(e) }); + if (reason === "explicit") + throw errors.length === 1 ? e : new AggregateError(errors, e instanceof Error ? e.message : String(e)); } } } diff --git a/src/editor/core/model/tabs.ts b/src/editor/core/model/tabs.ts index dc58c9c47..269e4b206 100644 --- a/src/editor/core/model/tabs.ts +++ b/src/editor/core/model/tabs.ts @@ -310,9 +310,9 @@ export function reopenClosed(s: TabsState): { state: TabsState; tabId: string; p } /** - * Tab limit: closes tabs of `paneId` beyond `limit` (pinned and `keep` tabs - * excluded). `unchanged` closes the least recently used unmodified tab first, - * `lru` the least recently used one. + * Tab limit: closes tabs of `paneId` beyond `limit` (pinned, `keep` and + * `isProtected` tabs excluded). `unchanged` closes the least recently used + * unmodified tab first, `lru` the least recently used one. */ export function enforceTabLimit( s: TabsState, @@ -321,12 +321,13 @@ export function enforceTabLimit( policy: string, isDirty: (tab: Tab) => boolean, keep?: string, + isProtected?: (tab: Tab) => boolean, ): TabsState { let next = s; for (;;) { const pane = next.panes[paneId]; if (!pane || limit <= 0) return next; - const closable = pane.tabs.filter((t) => !t.pinned && t.id !== keep && t.id !== pane.active); + const closable = pane.tabs.filter((t) => !t.pinned && t.id !== keep && t.id !== pane.active && !isProtected?.(t)); if (pane.tabs.filter((t) => !t.pinned).length <= limit || !closable.length) return next; const lruOrder = [...closable].sort((a, b) => rank(pane, b) - rank(pane, a)); const victim = (policy === "unchanged" ? lruOrder.find((t) => !isDirty(t)) : undefined) ?? lruOrder[0]!; diff --git a/src/editor/core/view/EditorArea.tsx b/src/editor/core/view/EditorArea.tsx index 71899f5e0..f948b322e 100644 --- a/src/editor/core/view/EditorArea.tsx +++ b/src/editor/core/view/EditorArea.tsx @@ -17,6 +17,7 @@ import { textEditorWrapperFor, useContributions, useFileDecorations, + useEditorDocumentRevision, useTextEditorWrappers, type EditorInput, } from "@/workbench/contributions"; @@ -39,6 +40,7 @@ import { } from "@/ui"; import { editorController, isTextInput, tabIcon, tabTitle, useTabs, type VirtualInputData } from "../controller"; import { useDocMeta } from "../model/document"; +import { inputDirty } from "../model/customDocuments"; import type { LayoutNode, SplitEdge, Tab } from "../model/tabs"; import { disambiguateTitles } from "../model/tabTitles"; import { TextEditorHost } from "./TextEditorHost"; @@ -130,7 +132,9 @@ function PaneView({ paneId }: { paneId: string }) { const kinds = useContributions((s) => s.editorKinds); const tabs = useMemo(() => pane?.tabs ?? [], [pane]); const paths = tabs.map((t) => t.input.path ?? ""); - const dirty = useDocMeta(useShallow((s) => paths.map((p) => !!s[p]?.dirty))); + const textDirty = useDocMeta(useShallow((s) => paths.map((p) => !!s[p]?.dirty))); + useEditorDocumentRevision(); + const dirty = tabs.map((tab, i) => inputDirty(tab.input, kinds) ?? textDirty[i] ?? false); const colors = useFileDecorations(useShallow(() => paths.map((p) => (p ? (getFileDecoration(p)?.color ?? "") : "")))); const tints = useFileDecorations(useShallow(() => paths.map((p) => (p ? (getFileDecoration(p)?.tint ?? "") : "")))); @@ -257,8 +261,8 @@ function PaneView({ paneId }: { paneId: string }) { layout={tabPrefs.multipleRows ? "wrap" : tabPrefs.oneRowOverflow === "squeeze" ? "squeeze" : "scroll"} closePosition={tabPrefs.closeButtonPosition === "left" ? "left" : "right"} onSelect={(id) => ctrl.activate(paneId, id)} - onClose={(id) => ctrl.closeTab(paneId, id)} - onCloseOthers={(id) => ctrl.closeOthers(paneId, id)} + onClose={(id) => void ctrl.closeTab(paneId, id)} + onCloseOthers={(id) => void ctrl.closeOthers(paneId, id)} onDoubleClick={(id) => { // IDEA: a preview tab becomes a regular one; any other tab maximizes the editor. if (pane.tabs.find((t) => t.id === id)?.preview) ctrl.keepTab(paneId, id); diff --git a/src/features/agents/store.ts b/src/features/agents/store.ts index b0d04ed22..2f89cd723 100644 --- a/src/features/agents/store.ts +++ b/src/features/agents/store.ts @@ -180,7 +180,7 @@ export async function connect(): Promise { for (const t of editors.openTabs?.() ?? []) { const task = (t.input.data as { task?: string } | undefined)?.task; if (t.input.kind === SESSION_KIND && task && !useAgents.getState().tasks.has(task)) - editors.closeInput?.(t.input.key); + await editors.closeInput?.(t.input.key); } } @@ -219,7 +219,7 @@ function onEvent(e: AgentEvent) { const unread = new Set(s.unread); unread.delete(e.id); useAgents.setState({ tasks, composers, unread }); - editorService().closeInput?.(sessionKey(e.id)); + void editorService().closeInput?.(sessionKey(e.id)); break; } case "truncated": { @@ -366,8 +366,7 @@ export async function startTask(): Promise { })); remember(composer.text.trim()); updateComposer(HOME, emptyDraft()); - editorService().closeInput?.(HOME_KEY); - openTask(task.id); + if ((await editorService().closeInput?.(HOME_KEY)) !== false) openTask(task.id); } function remember(text: string) { diff --git a/src/features/analysis/cleanup.ts b/src/features/analysis/cleanup.ts index f1b971cf7..f6d3b691f 100644 --- a/src/features/analysis/cleanup.ts +++ b/src/features/analysis/cleanup.ts @@ -126,7 +126,10 @@ export async function runCodeCleanup(choice: ScopeChoice, options: CleanupOption }); await executeAction("code.cleanupFile", { path, ...options }); await editors?.saveAll(); - if (!wasOpen.has(path)) editors?.closePath?.(path); + if (!wasOpen.has(path) && (await editors?.closePath?.(path)) === false) { + run.cancelled = true; + break; + } done++; } } finally { diff --git a/src/features/app/__tests__/closeGuard.test.ts b/src/features/app/__tests__/closeGuard.test.ts index 1cd3a17c6..596365ac5 100644 --- a/src/features/app/__tests__/closeGuard.test.ts +++ b/src/features/app/__tests__/closeGuard.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it, vi } from "vitest"; import { confirmExitEnabled, decideClose, type CloseDeps, type CloseInput } from "../closeGuard"; import type { MessageRequest as PromptRequest, MessageResult as PromptResult } from "@/workbench/messages"; +import type { EditorCloseSession } from "@/workbench/contributions"; function deps(answers: PromptResult[]) { const asked: PromptRequest[] = []; @@ -26,9 +27,8 @@ describe("decideClose", () => { }); it("confirms exit for the last window when enabled", async () => { - const { d, asked } = deps([{ button: "exit", checked: false }]); + const { d } = deps([{ button: "exit", checked: false }]); await expect(decideClose({ ...base, confirmExit: true }, d)).resolves.toBe(true); - expect(asked[0]?.title).toBe("Confirm Exit"); expect(d.disableConfirmExit).not.toHaveBeenCalled(); }); @@ -51,9 +51,8 @@ describe("decideClose", () => { }); it("saves all on Save", async () => { - const { d, asked } = deps([{ button: "save", checked: false }]); + const { d } = deps([{ button: "save", checked: false }]); await expect(decideClose({ ...base, dirty: 2 }, d)).resolves.toBe(true); - expect(asked[0]?.body).toBe("2 files have unsaved changes in openidea. Save before closing?"); expect(d.saveAll).toHaveBeenCalled(); }); @@ -79,4 +78,83 @@ describe("decideClose", () => { expect(confirmExitEnabled({ general: {} })).toBe(false); expect(confirmExitEnabled(null)).toBe(false); }); + + it("includes dirty changes produced while preparing a custom document", async () => { + const { d } = deps([{ button: "save", checked: false }]); + let text = "old output"; + let saved = ""; + let frozen = false; + const session: EditorCloseSession = { + dirty: 1, + discard: async () => { + text = ""; + }, + cancel: () => { + frozen = false; + }, + commit: () => { + frozen = true; + }, + }; + d.prepareClose = async () => { + frozen = true; + text = "final execution output"; + return session; + }; + d.saveAll = async () => { + expect(frozen).toBe(true); + saved = text; + }; + await expect(decideClose(base, d)).resolves.toBe(true); + expect(saved).toBe("final execution output"); + expect(frozen).toBe(true); + }); + + it.each(["cancel", "save", "discard"])( + "resumes a prepared document when %s cannot authorize closing", + async (button) => { + const { d } = deps([{ button, checked: false }]); + let frozen = true; + let text = "unsaved source"; + const session: EditorCloseSession = { + dirty: 1, + discard: async () => { + throw new Error("reload failed"); + }, + cancel: () => { + frozen = false; + }, + commit: () => { + text = "closed"; + }, + }; + d.prepareClose = async () => session; + d.saveAll = async () => { + throw new Error("disk full"); + }; + await expect(decideClose(base, d)).resolves.toBe(false); + expect(frozen).toBe(false); + expect(text).toBe("unsaved source"); + }, + ); + + it("Don't Save discards prepared custom edits without saving them", async () => { + const { d } = deps([{ button: "discard", checked: false }]); + let text = "unsaved"; + let saved = "disk baseline"; + d.saveAll = async () => { + saved = text; + }; + d.prepareClose = async () => ({ + dirty: 1, + discard: async () => { + text = saved; + }, + cancel: () => {}, + commit: () => {}, + }); + await expect(decideClose(base, d)).resolves.toBe(true); + expect(saved).toBe("disk baseline"); + expect(text).toBe("disk baseline"); + }); }); diff --git a/src/features/app/__tests__/editCommand.test.ts b/src/features/app/__tests__/editCommand.test.ts index 6d1460cf9..4fa8a2477 100644 --- a/src/features/app/__tests__/editCommand.test.ts +++ b/src/features/app/__tests__/editCommand.test.ts @@ -1,14 +1,16 @@ -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { history } from "@codemirror/commands"; import { EditorState } from "@codemirror/state"; import { EditorView } from "@codemirror/view"; import { editCommand } from "../appActions"; +import * as contributions from "@/workbench/contributions"; describe("Edit menu Undo / Redo", () => { let view: EditorView | null = null; afterEach(() => { view?.destroy(); view = null; + vi.restoreAllMocks(); }); it("use the focused editor's history, including programmatic edits", () => { @@ -20,4 +22,32 @@ describe("Edit menu Undo / Redo", () => { editCommand("redo")(); expect(view.state.doc.toString()).toBe("ab"); }); + + it("routes custom document history before CodeMirror's independent local history", () => { + const parent = document.createElement("div"); + parent.dataset.paneId = "custom-pane"; + document.body.appendChild(parent); + view = new EditorView({ state: EditorState.create({ doc: "local", extensions: history() }), parent }); + view.dispatch({ changes: { from: 5, insert: " edit" }, userEvent: "input" }); + let source = "backend edited"; + vi.spyOn(contributions, "maybeEditorService").mockReturnValue({ + openFile: async () => {}, + openInput: async () => {}, + activePath: () => "/p/notebook.ipynb", + activeCursor: () => null, + saveAll: async () => {}, + documentHistory: (command) => { + source = command === "undo" ? "backend original" : "backend edited"; + return true; + }, + }); + view.contentDOM.focus(); + editCommand("undo")(); + expect(source).toBe("backend original"); + expect(view.state.doc.toString()).toBe("local edit"); + editCommand("redo")(); + expect(source).toBe("backend edited"); + expect(view.state.doc.toString()).toBe("local edit"); + parent.remove(); + }); }); diff --git a/src/features/app/appActions.ts b/src/features/app/appActions.ts index 19d7e4ce7..39f33ce9f 100644 --- a/src/features/app/appActions.ts +++ b/src/features/app/appActions.ts @@ -3,7 +3,7 @@ import { EditorView } from "@codemirror/view"; import type { Action } from "@/platform/actions"; import { t } from "@/platform/i18n"; import { call, commands, isTauri } from "@/platform/ipc"; -import { notifications } from "@/workbench/contributions"; +import { maybeEditorService, notifications } from "@/workbench/contributions"; import { showAbout } from "./AboutDialog"; import { quitApp } from "./closeGuard"; import { guessOs } from "./host"; @@ -35,6 +35,12 @@ export function editCommand(command: "undo" | "redo" | "cut" | "copy" | "paste" return () => { if (command === "undo" || command === "redo") { const focused = document.activeElement; + if ( + focused instanceof Element && + focused.closest("[data-editor-area], [data-pane-id]") && + maybeEditorService()?.documentHistory?.(command) + ) + return; const view = focused instanceof HTMLElement ? EditorView.findFromDOM(focused) : null; if (view) { (command === "undo" ? undo : redo)(view); diff --git a/src/features/app/closeGuard.ts b/src/features/app/closeGuard.ts index 28f894eba..29739ab4a 100644 --- a/src/features/app/closeGuard.ts +++ b/src/features/app/closeGuard.ts @@ -10,16 +10,12 @@ import { commands, isTauri, onIdeEvent, type Settings } from "@/platform/ipc"; import { t } from "@/platform/i18n"; import { closeProject, provideLeaveProjectGuard, useProject } from "@/platform/project"; -import { maybeEditorService, notifications } from "@/workbench/contributions"; +import { maybeEditorService, notifications, type EditorCloseSession } from "@/workbench/contributions"; import { ask, type Ask } from "@/workbench/messages"; /** Document ids the backend reports dirty in this window. */ const dirtyDocs = new Set(); -export function dirtyCount(): number { - return dirtyDocs.size; -} - export function confirmExitEnabled(settings: Settings | null | undefined): boolean { return settings?.general?.confirmExit === true; } @@ -37,6 +33,8 @@ export interface CloseDeps { saveAll: () => Promise; disableConfirmExit: () => Promise; onSaveFailed: (e: unknown) => void; + /** Stops custom document edits/execution before reading their final dirty state. */ + prepareClose?: () => Promise; } /** Asks whatever must be asked; resolves true when the window may close. */ @@ -51,27 +49,34 @@ export async function decideClose(input: CloseInput, deps: CloseDeps): Promise {}); } - if (input.dirty > 0) { - const r = await deps.ask({ - title: t("app.unsaved.title"), - body: t("app.unsaved.body", { count: input.dirty, project: input.project ?? t("app.title.default") }), - buttons: [ - { id: "discard", label: t("app.unsaved.discard") }, - { id: "save", label: t("app.unsaved.save"), primary: true }, - ], - icon: "warning", - }); - if (r.button === "cancel") return false; - if (r.button === "save") { - try { - await deps.saveAll(); - } catch (e) { - deps.onSaveFailed(e); - return false; - } + let session: EditorCloseSession | undefined; + let authorized = false; + try { + session = await deps.prepareClose?.(); + const dirty = input.dirty + (session?.dirty ?? 0); + if (dirty > 0) { + const r = await deps.ask({ + title: t("app.unsaved.title"), + body: t("app.unsaved.body", { count: dirty, project: input.project ?? t("app.title.default") }), + buttons: [ + { id: "discard", label: t("app.unsaved.discard") }, + { id: "save", label: t("app.unsaved.save"), primary: true }, + ], + icon: "warning", + }); + if (r.button === "save") await deps.saveAll(); + else if (r.button === "discard") await session?.discard(); + else return false; } + session?.commit(); + authorized = true; + return true; + } catch (e) { + deps.onSaveFailed(e); + return false; + } finally { + if (!authorized) session?.cancel(); } - return true; } const defaultDeps: CloseDeps = { @@ -79,6 +84,7 @@ const defaultDeps: CloseDeps = { saveAll: async () => { await maybeEditorService()?.saveAll(); }, + prepareClose: async () => maybeEditorService()?.prepareClose?.(), disableConfirmExit: async () => { const r = await commands.settingsSet("global", "general.confirmExit", "false"); if (r.status === "ok") useProject.setState({ settings: r.data }); @@ -110,13 +116,17 @@ async function onCloseRequested(deps: CloseDeps): Promise { quitting: ctx.quitting, windows: ctx.windows, confirmExit: confirmExitEnabled(s.settings), - dirty: dirtyCount(), + dirty: dirtyDocs.size, project: s.project?.name ?? null, }, deps, ); if (!ok && ctx.quitting) await commands.appCancelQuit(); return ok; + } catch (e) { + maybeEditorService()?.cancelClose?.(); + deps.onSaveFailed(e); + return false; } finally { closing = false; } @@ -126,7 +136,7 @@ async function onCloseRequested(deps: CloseDeps): Promise { export function confirmLeave(deps: CloseDeps = defaultDeps): Promise { const s = useProject.getState(); return decideClose( - { quitting: false, windows: 2, confirmExit: false, dirty: dirtyCount(), project: s.project?.name ?? null }, + { quitting: false, windows: 2, confirmExit: false, dirty: dirtyDocs.size, project: s.project?.name ?? null }, deps, ); } @@ -135,7 +145,10 @@ export function confirmLeave(deps: CloseDeps = defaultDeps): Promise { export async function quitApp(deps: CloseDeps = defaultDeps): Promise { if (!isTauri()) return; if (confirmExitEnabled(useProject.getState().settings)) { - const ok = await decideClose({ quitting: false, windows: 1, confirmExit: true, dirty: 0, project: null }, deps); + const ok = await decideClose( + { quitting: false, windows: 1, confirmExit: true, dirty: 0, project: null }, + { ...deps, prepareClose: undefined }, + ); if (!ok) return; } await commands.appQuit(); diff --git a/src/features/index.ts b/src/features/index.ts index b66a461b8..ae77994eb 100644 --- a/src/features/index.ts +++ b/src/features/index.ts @@ -1,7 +1,8 @@ /** * Feature activation. Each feature exports `activate(): void | (() => void)` * that registers actions, tool windows, status items and services, returning - * a disposer. The editor goes first because it provides the editor service. + * a disposer. The editor provides the editor service; app must install its launch + * barrier before welcome can restore a remembered project. */ import { activate as editor } from "@/editor/core"; import { activate as editorLsp } from "@/editor/lsp"; @@ -29,6 +30,7 @@ import { activate as tools, activateSettingsPages as toolsPages } from "./tools" import { activate as scratch } from "./scratch"; import { activate as libraries } from "./libraries"; import { activate as viewers } from "./viewers"; +import { activate as notebook } from "./notebook"; import { activate as web } from "./web"; import { activate as analysis } from "./analysis"; import { activate as services } from "./services"; @@ -46,6 +48,7 @@ type Activate = () => void | (() => void); const features: Array<[string, Activate]> = [ ["editor", editor], ["editor-lsp", editorLsp], + ["app", app], ["welcome", welcome], ["project-tree", projectTree], ["search", search], @@ -63,12 +66,12 @@ const features: Array<[string, Activate]> = [ ["templates", templates], ["todo", todo], ["settings", settings], - ["app", app], ["perf", perf], ["plugins", plugins], ["tools", tools], ["scratch", scratch], ["libraries", libraries], + ["notebook", notebook], ["viewers", viewers], ["web", web], ["analysis", analysis], diff --git a/src/features/notebook/CellEditor.tsx b/src/features/notebook/CellEditor.tsx new file mode 100644 index 000000000..627ea58ff --- /dev/null +++ b/src/features/notebook/CellEditor.tsx @@ -0,0 +1,129 @@ +import { useLayoutEffect, useRef } from "react"; +import { Annotation, Compartment, EditorState, Prec, StateEffect } from "@codemirror/state"; +import { + EditorView, + drawSelection, + highlightActiveLine, + highlightSpecialChars, + keymap, + lineNumbers, +} from "@codemirror/view"; +import { defaultKeymap, indentWithTab } from "@codemirror/commands"; +import { bracketMatching, indentOnInput, syntaxHighlighting } from "@codemirror/language"; +import { openideaTheme } from "@/editor/core/cm/theme"; +import { openideaHighlightStyle } from "@/editor/core/cm/highlight"; +import { findLanguageDescription, loadSupport } from "@/editor/core/cm/languages"; +import { tx } from "@/platform/i18n"; +import type { NotebookSession } from "./session"; +import styles from "./Notebook.module.css"; + +const external = Annotation.define(); +export function CellEditor({ + session, + id, + source, + type, + frozen, +}: { + session: NotebookSession; + id: string; + source: string; + type: string; + frozen: boolean; +}) { + const host = useRef(null); + const view = useRef(null); + const readonly = useRef(new Compartment()); + const initialSource = useRef(source); + useLayoutEffect(() => { + if (!host.current) return; + const language = new Compartment(); + const configuration = [ + openideaTheme, + syntaxHighlighting(openideaHighlightStyle), + lineNumbers(), + drawSelection(), + highlightActiveLine(), + highlightSpecialChars(), + bracketMatching(), + indentOnInput(), + Prec.highest( + keymap.of([ + { + key: "Mod-z", + run: () => { + void session.undo().catch((e) => session.report(e)); + return true; + }, + }, + { + key: "Mod-Shift-z", + run: () => { + void session.redo().catch((e) => session.report(e)); + return true; + }, + }, + { + key: "Mod-y", + run: () => { + void session.redo().catch((e) => session.report(e)); + return true; + }, + }, + ]), + ), + keymap.of([indentWithTab, ...defaultKeymap]), + language.of([]), + readonly.current.of(EditorState.readOnly.of(session.getState().frozen)), + EditorView.lineWrapping, + EditorView.theme({ "&": { height: "auto" }, ".cm-scroller": { maxHeight: "640px", overflow: "auto" } }), + EditorView.contentAttributes.of({ + "aria-label": tx(type === "code" ? "Python cell source" : "Markdown cell source"), + }), + EditorState.transactionFilter.of((transaction) => + transaction.docChanged && session.getState().frozen && !transaction.annotation(external) ? [] : transaction, + ), + EditorView.updateListener.of((update) => { + session.editors.set(id, update.state); + if (update.docChanged && update.transactions.some((tr) => !tr.annotation(external))) { + session.edit({ type: "setSource", cell: id, source: update.state.doc.toString() }); + } + }), + ]; + // Reconfigure saved state to replace stale view closures; retain selection. + const saved = session.editors.get(id); + const state = saved + ? saved.update({ effects: StateEffect.reconfigure.of(configuration) }).state + : EditorState.create({ doc: initialSource.current, extensions: configuration }); + const editor = new EditorView({ state, parent: host.current }); + view.current = editor; + let alive = true; + const description = findLanguageDescription( + type === "code" ? "Python" : "Markdown", + type === "code" ? "cell.py" : "cell.md", + ); + if (description) + void loadSupport(description).then((support) => { + if (alive && support) editor.dispatch({ effects: language.reconfigure(support) }); + }); + return () => { + alive = false; + session.editors.set(id, editor.state); + view.current = null; + editor.destroy(); + }; + }, [session, id, type]); + useLayoutEffect(() => { + const editor = view.current; + if (!editor) return; + if (editor.state.doc.toString() !== source) + editor.dispatch({ + changes: { from: 0, to: editor.state.doc.length, insert: source }, + annotations: external.of(true), + }); + editor.dispatch({ effects: readonly.current.reconfigure(EditorState.readOnly.of(frozen)) }); + }, [source, frozen]); + return
; +} + +// No CodeMirror history extension: source and structural undo share backend history. diff --git a/src/features/notebook/Notebook.module.css b/src/features/notebook/Notebook.module.css new file mode 100644 index 000000000..c5bb59f5d --- /dev/null +++ b/src/features/notebook/Notebook.module.css @@ -0,0 +1,220 @@ +.root { + display: flex; + flex: 1; + flex-direction: column; + min-height: 0; + min-width: 0; + background: var(--bg-editor, var(--bg-app)); + color: var(--fg); + outline: none; +} +.toolbar { + display: flex; + flex: none; + align-items: center; + flex-wrap: wrap; + gap: 3px; + min-height: 36px; + padding: 4px 8px; + border-bottom: 1px solid var(--border-editor-tabs, var(--border)); + font-size: var(--fs-small); +} +.separator { + width: 1px; + height: 16px; + margin: 0 4px; + background: var(--border); +} +.spacer { + flex: 1; +} +.interpreter { + min-width: 120px; + max-width: 260px; +} +.kernel { + display: inline-flex; + align-items: center; + gap: 5px; + padding: 0 6px; + color: var(--fg-muted); +} +.error { + display: flex; + align-items: flex-start; + gap: 12px; + padding: 8px 12px; + color: var(--danger, var(--fg)); + border-bottom: 1px solid var(--border); + overflow-wrap: anywhere; + user-select: text; +} +.errorMessage { + flex: 1; + min-width: 0; + white-space: pre-wrap; +} +.errorMessage pre { + max-height: 200px; + overflow: auto; +} +.kernelHelp { + display: block; + margin-top: 6px; +} +.scroll { + flex: 1; + min-height: 0; + min-width: 0; + overflow: auto; + overflow-anchor: none; +} +.cellList { + position: relative; + width: 100%; +} +.virtualCell { + position: absolute; + top: 0; + left: 0; + width: 100%; + padding: 6px 12px; + box-sizing: border-box; +} +.cell { + border: 1px solid var(--border); + border-radius: var(--radius-small, 4px); + min-width: 0; + overflow: hidden; +} +.selected { + border-color: var(--accent); +} +.cellToolbar { + display: flex; + align-items: center; + gap: 2px; + min-height: 28px; + padding: 2px 6px; + background: var(--bg-elev, var(--bg-app)); + border-bottom: 1px solid var(--border); +} +.executionCount { + min-width: 32px; + text-align: right; + font: var(--fs-small) var(--font-mono); + color: var(--fg-muted); + margin-right: 6px; +} +.cellEditor { + min-height: 48px; +} +.outputs { + border-top: 1px solid var(--border); + padding: 8px 12px; + overflow: auto; + max-height: 640px; + user-select: text; +} +.textOutput, +.errorOutput, +.rawCell pre { + white-space: pre-wrap; + overflow-wrap: anywhere; + font: var(--fs-code) / var(--lh-code, 1.5) var(--font-mono); + margin: 0; + user-select: text; +} +.errorOutput { + color: var(--danger, var(--fg)); +} +.outputImage { + display: block; + max-width: 100%; + height: auto; + margin: 4px 0; +} +.outputFrame { + display: block; + width: 100%; + height: 280px; + min-height: 80px; + max-height: 600px; + border: 1px solid var(--border); + background: white; + resize: vertical; + box-sizing: border-box; +} +.notice { + color: var(--fg-muted); + font-size: var(--fs-small); + padding: 4px 0; + margin: 0; + user-select: text; +} +.rawCell { + padding: 8px 12px; + max-height: 480px; + overflow: auto; +} +.markdown { + padding: 8px 14px; + min-height: 32px; + overflow-wrap: anywhere; + line-height: 1.6; + user-select: text; +} +.markdown > :first-child { + margin-top: 0; +} +.markdown > :last-child { + margin-bottom: 0; +} +.markdown img { + max-width: 100%; + height: auto; +} +.markdown pre { + overflow: auto; + padding: 8px; + background: var(--bg-elev); + font-family: var(--font-mono); +} +.markdown code { + font-family: var(--font-mono); +} +.markdown a { + color: var(--link); + cursor: pointer; +} +.markdown table { + border-collapse: collapse; + display: block; + overflow: auto; +} +.markdown th, +.markdown td { + border: 1px solid var(--border); + padding: 4px 8px; +} +.addBar { + display: flex; + justify-content: center; + gap: 8px; + flex-wrap: wrap; + padding: 12px 16px 32px; +} +@media (max-width: 600px) { + .virtualCell { + padding: 5px; + } + .toolbar { + gap: 2px; + } + .interpreter { + max-width: 180px; + } + .cellToolbar { + flex-wrap: wrap; + } +} diff --git a/src/features/notebook/NotebookEditor.tsx b/src/features/notebook/NotebookEditor.tsx new file mode 100644 index 000000000..339fa9d6e --- /dev/null +++ b/src/features/notebook/NotebookEditor.tsx @@ -0,0 +1,436 @@ +import { useCallback, useEffect, useLayoutEffect, useRef, useState, useSyncExternalStore } from "react"; +import { useVirtualizer } from "@tanstack/react-virtual"; +import type { NotebookCell } from "@/bindings"; +import { tx } from "@/platform/i18n"; +import { Button, Dialog, EmptyState, IconButton, Select, Spinner } from "@/ui"; +import type { EditorInput } from "@/workbench/contributions"; +import { CellEditor } from "./CellEditor"; +import { CellOutputs, MarkdownCell } from "./OutputView"; +import { notebookSession, type NotebookSession } from "./session"; +import styles from "./Notebook.module.css"; + +export function NotebookEditor({ input }: { input: EditorInput }) { + if (!input.path) return ; + return ; +} + +function NotebookDocument({ session }: { session: NotebookSession }) { + const state = useSyncExternalStore(session.subscribe, session.getState); + const root = useRef(null); + const scroll = useRef(null); + const [selected, select] = useState(session.selectedCell); + const [reloadOpen, setReloadOpen] = useState(false); + const cells = session.cells(); + const active = cells.find((cell) => cell.id === selected) ?? cells.find((cell) => cell.cellType === "code"); + const invoke = useCallback( + (work: Promise) => { + void work.catch((error) => session.report(error)); + }, + [session], + ); + const choose = (id: string) => { + session.selectedCell = id; + select(id); + }; + const virtualizer = useVirtualizer({ + count: cells.length, + getScrollElement: () => scroll.current, + getItemKey: (index) => cells[index]!.id, + estimateSize: () => 220, + overscan: 3, + initialOffset: session.scrollTop, + }); + useEffect(() => { + invoke(session.load()); + }, [session, invoke]); + useLayoutEffect(() => { + const element = root.current; + if (!element) return; + const history = (event: InputEvent) => { + if (event.inputType !== "historyUndo" && event.inputType !== "historyRedo") return; + event.preventDefault(); + event.stopImmediatePropagation(); + invoke(event.inputType === "historyUndo" ? session.undo() : session.redo()); + }; + element.addEventListener("beforeinput", history, true); + return () => element.removeEventListener("beforeinput", history, true); + }, [session, invoke]); + + const disabled = state.frozen || !state.snapshot; + const historyDisabled = disabled || state.running; + const options = state.environments.map((sdk) => ({ value: sdk.path, label: sdk.name, detail: sdk.path })); + if (state.python && !options.some((option) => option.value === state.python)) + options.unshift({ + value: state.python, + label: state.python, + detail: state.python, + }); + return ( +
{ + if (event.nativeEvent.isComposing || event.defaultPrevented) return; + if (event.metaKey || event.ctrlKey) { + const key = event.key.toLowerCase(); + if (key === "s") { + event.preventDefault(); + event.stopPropagation(); + invoke(session.save()); + } + if (key === "z" || key === "y") { + event.preventDefault(); + event.stopPropagation(); + invoke(event.shiftKey || key === "y" ? session.redo() : session.undo()); + } + } + if (event.key === "Enter" && (event.shiftKey || event.ctrlKey) && active?.cellType === "code") { + event.preventDefault(); + event.stopPropagation(); + invoke(session.execute([active.id])); + } + }} + > +
+ + + invoke(session.interrupt())} + /> + invoke(session.restart())} + /> + + session.edit({ type: "setCellType", cell: cell.id, cellType: value })} + /> + {cell.cellType === "code" && ( + { + void session.execute([cell.id]).catch((error) => session.report(error)); + }} + /> + )} + {cell.cellType === "markdown" && ( + { + if (editing) session.markdownEditing.delete(cell.id); + else session.markdownEditing.add(cell.id); + setEditing(!editing); + }} + /> + )} + + session.edit({ type: "moveCell", cell: cell.id, before: cells[index - 1]!.id })} + /> + session.edit({ type: "moveCell", cell: cell.id, before: cells[index + 2]?.id ?? null })} + /> + session.edit({ type: "insertCell", after: cell.id, cellType: "code" })} + /> + session.edit({ type: "deleteCell", cell: cell.id })} + /> +
+ {!known ? ( +
+

{tx("Raw or unknown cell; original content is preserved read-only.")}

+ 
{cell.source}
+
+ ) : cell.cellType === "markdown" && !editing ? ( +
{ + if (!frozen) { + session.markdownEditing.add(cell.id); + setEditing(true); + } + }} + > + {cell.source ? ( + + ) : ( + + )} +
+ ) : ( + + )} + {cell.cellType === "code" && } + + ); +} diff --git a/src/features/notebook/OutputView.test.tsx b/src/features/notebook/OutputView.test.tsx new file mode 100644 index 000000000..f4e9992db --- /dev/null +++ b/src/features/notebook/OutputView.test.tsx @@ -0,0 +1,119 @@ +import { describe, expect, it, vi } from "vitest"; +import { createEvent, fireEvent, render } from "@testing-library/react"; + +const openExternalUrl = vi.hoisted(() => vi.fn(async () => {})); +const sanitizedSvgUrl = vi.hoisted(() => vi.fn((_source: string) => "data:image/svg+xml,clean")); +vi.mock("@/platform/openUrl", () => ({ openExternalUrl })); +vi.mock("./security", async (original) => { + const actual = await original(); + return { + ...actual, + imageMimeUrl: (bundle: Record) => + typeof bundle["image/svg+xml"] === "string" && bundle["image/svg+xml"] + ? sanitizedSvgUrl(bundle["image/svg+xml"]) + : actual.imageMimeUrl(bundle), + }; +}); + +const { CellOutputs, MarkdownCell } = await import("./OutputView"); +const PNG = "iVBORw0KGgo="; + +describe("notebook Markdown preview", () => { + it("leaves native interactions such as
alone and only claims links", () => { + const { container } = render( + Morebody
"} attachments={{}} path="/p/a.ipynb" />, + ); + const summary = container.querySelector("summary")!; + const click = createEvent.click(summary); + fireEvent(summary, click); + expect(click.defaultPrevented).toBe(false); + }); + + it("opens web links with the system handler", () => { + openExternalUrl.mockClear(); + const { container } = render( + , + ); + const link = container.querySelector("a[data-notebook-href]")!; + const click = createEvent.click(link); + fireEvent(link, click); + expect(click.defaultPrevented).toBe(true); + expect(openExternalUrl).toHaveBeenCalledWith("https://example.com/docs"); + }); + + it("keeps the rendered preview across snapshots with equal attachments", () => { + const attachments = { "a.png": { "image/png": PNG } }; + const source = "text ![a](attachment:a.png)"; + const { container, rerender } = render( + , + ); + const paragraph = container.querySelector("p")!; + rerender(); + expect(container.querySelector("p")).toBe(paragraph); + rerender(); + expect(container.querySelector("p")).not.toBe(paragraph); + expect(container.querySelector("img")!.getAttribute("src")).toBe("data:image/png;base64,AAAA"); + }); + + it.each([ + ["Markdown destination", "![a](attachment:a.png)"], + ["reference definition", "![a][r]\n\n[r]: attachment:a.png"], + ["raw HTML", 'a'], + ])("resolves an entity-encoded attachment in a %s", (_kind, source) => { + const { container } = render( + , + ); + expect(container.querySelector("img")?.getAttribute("src")).toBe(`data:image/png;base64,${PNG}`); + }); + + it.each(["__proto__", "constructor"])("resolves an attachment named %s from parsed notebook JSON", (name) => { + const attachments = JSON.parse(`{${JSON.stringify(name)}:{"image/png":"${PNG}"}}`); + const { container } = render( + , + ); + expect(container.querySelector("img")?.getAttribute("src")).toBe(`data:image/png;base64,${PNG}`); + }); + + it("does not resolve inherited members as attachments", () => { + const { container } = render( + , + ); + expect(container.querySelector("img")).toBeNull(); + expect(container.textContent?.trim()).toBe("a"); + }); + + it("does not rebuild when only an unreferenced attachment changes", () => { + const source = "![a](attachment:a.png)"; + const { container, rerender } = render( + , + ); + const image = container.querySelector("img")!; + rerender( + , + ); + expect(container.querySelector("img")).toBe(image); + }); + + it("ignores attachment identity when the source references none", () => { + const { container, rerender } = render(); + const paragraph = container.querySelector("p")!; + rerender(); + expect(container.querySelector("p")).toBe(paragraph); + }); +}); + +describe("notebook image outputs", () => { + it("does not re-sanitize an unchanged SVG when a fresh snapshot arrives", () => { + sanitizedSvgUrl.mockClear(); + const output = () => [{ output_type: "display_data", data: { "image/svg+xml": "" }, metadata: {} }]; + const { rerender } = render(); + rerender(); + expect(sanitizedSvgUrl).toHaveBeenCalledTimes(1); + rerender(" } }]} />); + expect(sanitizedSvgUrl).toHaveBeenCalledTimes(2); + }); +}); diff --git a/src/features/notebook/OutputView.tsx b/src/features/notebook/OutputView.tsx new file mode 100644 index 000000000..19f2c01af --- /dev/null +++ b/src/features/notebook/OutputView.tsx @@ -0,0 +1,203 @@ +import { useLayoutEffect, useMemo, useRef } from "react"; +import Anser from "anser"; +import { tx } from "@/platform/i18n"; +import { openExternalUrl } from "@/platform/openUrl"; +import { editorService, notifications } from "@/workbench/contributions"; +import { resolveLink } from "@/features/viewers/markdown/render"; +import { + imageMimeUrl, + markdownAttachmentNames, + markdownRoot, + resolveMarkdownAttachments, + staticOutputDocument, +} from "./security"; +import { mimeText, object, own } from "./outputs"; +import styles from "./Notebook.module.css"; + +const PREVIEW_LIMIT = 200_000; +function AnsiText({ text, error = false }: { text: string; error?: boolean }) { + const entries = useMemo(() => Anser.ansiToJson(text.slice(0, PREVIEW_LIMIT), { json: true }), [text]); + return ( + <> + 
+        {entries.map((entry, index) => (
+          
+            {entry.content}
+          
+        ))}
+
+ {text.length > PREVIEW_LIMIT && ( +

{tx("Output preview truncated; full output is retained in the notebook.")}

+ )} + + ); +} + +function StaticHtml({ html, title }: { html: string; title: string }) { + const srcDoc = useMemo(() => staticOutputDocument(html), [html]); + return ( + + link + +
safe text
+
namevalue
alpha42
+ `), + ); + expect(document.querySelector("tbody tr")?.textContent).toBe("alpha42"); + expect(document.querySelector("script,base,form,input,button,iframe,object")).toBeNull(); + expect(document.querySelector("meta[http-equiv=refresh]")).toBeNull(); + expect(document.querySelector("[href],[target],[srcset],[onerror],[ping]")).toBeNull(); + expect(document.querySelector("img")?.hasAttribute("src")).toBe(false); + expect(document.querySelector("body div")?.hasAttribute("style")).toBe(false); + }); + + it("retains safe SVG chart paint, local definitions and raster images but removes active/external SVG", () => { + const url = sanitizedSvgUrl(` +
host markup
 + + + + + + + + `); + const xml = decodeURIComponent(url.slice(url.indexOf(",") + 1)); + const document = new DOMParser().parseFromString(xml, "image/svg+xml"); + expect(document.querySelector("script,foreignObject,animate,style")).toBeNull(); + expect(document.querySelector('[href^="https:"],[href^="http:"]')).toBeNull(); + expect(document.querySelector("#chart")?.getAttribute("stroke")).toBe("#123456"); + expect(document.querySelector("#chart")?.getAttribute("fill")).toBe("none"); + expect(document.querySelector("rect")?.getAttribute("fill")).toBe("url(#paint)"); + expect(document.querySelector("#local-marker")?.getAttribute("href")).toBe("#chart"); + expect(document.querySelector("#external")?.hasAttribute("fill")).toBe(false); + expect(document.querySelector("image")?.getAttribute("href")).toBe("data:image/png;base64,aGVsbG8="); + expect(document.querySelector("[onclick],[style]")).toBeNull(); + }); + + it("resolves only embedded image attachments and makes Markdown file links editor-controlled", () => { + const root = resolveMarkdownAttachments( + markdownRoot( + `![chart](attachment:chart.png)\n\n![remote](https://attacker.invalid/a.png)\n\n[local](./other.py#L8)\n\n
text
`, + ), + { + "chart.png": { "image/png": "aGVsbG8=" }, + }, + ); + const images = root.querySelectorAll("img"); + expect(images[0]?.getAttribute("src")).toBe("data:image/png;base64,aGVsbG8="); + expect(images[1]?.hasAttribute("src")).toBe(false); + const link = root.querySelector("a"); + expect(link?.hasAttribute("href")).toBe(false); + expect(link?.getAttribute("data-notebook-href")).toBe("./other.py#L8"); + expect(root.querySelector("script,[style]")).toBeNull(); + }); +}); diff --git a/src/features/notebook/security.ts b/src/features/notebook/security.ts new file mode 100644 index 000000000..fb3fbb1b0 --- /dev/null +++ b/src/features/notebook/security.ts @@ -0,0 +1,285 @@ +import DOMPurify from "dompurify"; +import { Marked } from "marked"; +import { mimeText, object, own } from "./outputs"; + +export const OUTPUT_CSP = + "default-src 'none'; img-src data: blob:; style-src 'unsafe-inline'; script-src 'none'; connect-src 'none'; form-action 'none'; base-uri 'none'; frame-src 'none'; object-src 'none'; font-src 'none'; media-src 'none'; navigate-to 'none'"; +const HTML_TAGS = [ + "div", + "span", + "p", + "br", + "hr", + "pre", + "code", + "samp", + "kbd", + "b", + "strong", + "i", + "em", + "u", + "s", + "small", + "sub", + "sup", + "blockquote", + "ul", + "ol", + "li", + "dl", + "dt", + "dd", + "h1", + "h2", + "h3", + "h4", + "h5", + "h6", + "table", + "caption", + "thead", + "tbody", + "tfoot", + "tr", + "th", + "td", + "colgroup", + "col", + "img", + "a", + "details", + "summary", +]; +const RASTER = /^data:image\/(?:png|jpeg);base64,[a-z0-9+/=\s]+$/i; +const SVG_TAGS: Record = { + svg: true, + g: true, + defs: true, + path: true, + rect: true, + circle: true, + ellipse: true, + line: true, + polyline: true, + polygon: true, + text: true, + tspan: true, + title: true, + desc: true, + clippath: true, + mask: true, + lineargradient: true, + radialgradient: true, + stop: true, + use: true, + symbol: true, + pattern: true, + image: true, +}; +const SVG_ATTRS: Record = { + xmlns: true, + viewbox: true, + width: true, + height: true, + x: true, + y: true, + x1: true, + x2: true, + y1: true, + y2: true, + cx: true, + cy: true, + r: true, + rx: true, + ry: true, + d: true, + points: true, + transform: true, + id: true, + fill: true, + "fill-opacity": true, + "fill-rule": true, + stroke: true, + "stroke-width": true, + "stroke-opacity": true, + "stroke-linecap": true, + "stroke-linejoin": true, + "stroke-dasharray": true, + "stroke-dashoffset": true, + opacity: true, + "clip-path": true, + "clip-rule": true, + clippathunits: true, + mask: true, + maskunits: true, + gradientunits: true, + gradienttransform: true, + offset: true, + "stop-color": true, + "stop-opacity": true, + "font-family": true, + "font-size": true, + "font-weight": true, + "font-style": true, + "text-anchor": true, + "dominant-baseline": true, + dx: true, + dy: true, + preserveaspectratio: true, + patternunits: true, + patterntransform: true, + href: true, + "xlink:href": true, +}; + +/** Image-only SVG: no foreign HTML, animation, scripts, CSS, or external references. */ +export function sanitizedSvgUrl(source: string): string { + const clean = DOMPurify.sanitize(source, { + USE_PROFILES: { svg: true, svgFilters: false }, + ADD_TAGS: ["use"], + RETURN_DOM_FRAGMENT: true, + }); + const svg = clean.querySelector("svg"); + if (!svg) return ""; + for (const node of [svg, ...svg.querySelectorAll("*")]) { + if (!Object.hasOwn(SVG_TAGS, node.localName.toLowerCase())) { + node.remove(); + continue; + } + // Matplotlib uses presentation CSS. Convert only safe paint/text properties + // to SVG attributes before dropping style itself. + const presentation = node.getAttribute("style"); + if (presentation) + for (const declaration of presentation.split(";")) { + const colon = declaration.indexOf(":"); + if (colon < 0) continue; + const name = declaration.slice(0, colon).trim().toLowerCase(); + const value = declaration.slice(colon + 1).trim(); + if ( + /^(?:fill(?:-opacity|-rule)?|stroke(?:-width|-opacity|-linecap|-linejoin|-dasharray|-dashoffset)?|opacity|font-size|font-weight|font-style|font-family|text-anchor|stop-color|stop-opacity)$/.test( + name, + ) + ) + node.setAttribute(name, value); + } + for (const attr of [...node.attributes]) { + const name = attr.name.toLowerCase(); + const value = attr.value.trim(); + if ( + !Object.hasOwn(SVG_ATTRS, name) || + /[\p{Cc}\\]/u.test(value) || + ((name === "href" || name === "xlink:href") && + !/^#[\w.-]+$/.test(value) && + !(node.localName.toLowerCase() === "image" && RASTER.test(value))) || + (/url\s*\(/i.test(value) && !/^url\(\s*#[\w.-]+\s*\)$/i.test(value)) + ) + node.removeAttribute(attr.name); + } + } + svg.setAttribute("xmlns", "http://www.w3.org/2000/svg"); + return `data:image/svg+xml;charset=utf-8,${encodeURIComponent(new XMLSerializer().serializeToString(svg))}`; +} + +function htmlRoot(html: string, markdown = false): HTMLElement { + const root = DOMPurify.sanitize(html, { + ALLOWED_TAGS: HTML_TAGS, + ALLOWED_ATTR: [ + "src", + "alt", + "title", + "width", + "height", + "colspan", + "rowspan", + "scope", + "align", + ...(markdown ? ["href"] : ["style"]), + ], + // Attachment references are resolved by us after sanitization. No arbitrary + // URI survives the explicit src allowlist below. + ADD_URI_SAFE_ATTR: ["src"], + ALLOW_DATA_ATTR: false, + ALLOW_ARIA_ATTR: false, + RETURN_DOM: true, + }) as HTMLElement; + for (const node of root.querySelectorAll("*")) { + // Style attributes are iframe-only. Strip URL-bearing declarations even + // though the child CSP also blocks their network requests. + const style = node.getAttribute("style"); + if (style && /url\s*\(|image-set\s*\(|src\s*\(|https?:|\/\/|@|\\|expression\s*\(|behavior|binding/i.test(style)) + node.removeAttribute("style"); + const src = node.getAttribute("src"); + if (src && !RASTER.test(src) && !(markdown && src.startsWith("attachment:"))) node.removeAttribute("src"); + if (!markdown) node.removeAttribute("href"); + } + return root; +} + +export function staticOutputDocument(html: string): string { + const root = htmlRoot(html); + return `${root.innerHTML}`; +} + +const markdown = new Marked({ gfm: true }); +/** + * Sanitized Markdown with inert links; link clicks are handled by the editor and + * `attachment:` images are left for `resolveMarkdownAttachments`. + */ +export function markdownRoot(source: string): HTMLElement { + const root = htmlRoot(markdown.parse(source, { async: false }) as string, true); + for (const link of root.querySelectorAll("a[href]")) { + const href = link.getAttribute("href")!; + link.removeAttribute("href"); + link.setAttribute("data-notebook-href", href); + link.setAttribute("role", "link"); + link.setAttribute("tabindex", "0"); + } + return root; +} + +function attachmentName(img: Element): string | null | undefined { + const src = img.getAttribute("src"); + if (!src?.startsWith("attachment:")) return undefined; + try { + return decodeURIComponent(src.slice("attachment:".length)); + } catch { + return null; + } +} + +/** Names referenced after Markdown entity and URI decoding, which a source-text search would miss. */ +export function markdownAttachmentNames(root: HTMLElement): string[] { + const names = new Set(); + for (const img of root.querySelectorAll("img")) { + const name = attachmentName(img); + if (name) names.add(name); + } + return [...names]; +} + +/** Host Markdown has no arbitrary CSS or remote images: only embedded attachments resolve. */ +export function resolveMarkdownAttachments(root: HTMLElement, attachments: unknown): HTMLElement { + for (const img of root.querySelectorAll("img")) { + const name = attachmentName(img); + if (name === undefined) continue; + if (name === null) { + img.removeAttribute("src"); + continue; + } + const data = object(own(attachments, name)); + const url = imageMimeUrl(data); + if (url) img.setAttribute("src", url); + else img.replaceWith(document.createTextNode(img.getAttribute("alt") || name)); + } + return root; +} + +export function imageMimeUrl(bundle: Record): string { + for (const mime of ["image/png", "image/jpeg"]) { + const value = mimeText(bundle[mime]).replace(/\s/g, ""); + if (value && /^[a-z0-9+/=]+$/i.test(value)) return `data:${mime};base64,${value}`; + } + const svg = mimeText(bundle["image/svg+xml"]); + return svg ? sanitizedSvgUrl(svg) : ""; +} diff --git a/src/features/notebook/session.test.ts b/src/features/notebook/session.test.ts new file mode 100644 index 000000000..94d28b341 --- /dev/null +++ b/src/features/notebook/session.test.ts @@ -0,0 +1,494 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import type * as TauriCore from "@tauri-apps/api/core"; +import type { IdeEvent, NotebookSnapshot } from "@/bindings"; +import { IpcError } from "@/platform/ipc"; +import type * as Ipc from "@/platform/ipc"; +import { NotebookSession, notebookBackend } from "./session"; + +type KernelEvent = Extract; +const kernelListeners = vi.hoisted(() => new Set<(event: KernelEvent) => void>()); +vi.mock("@/platform/ipc", async (original) => { + const actual = await original(); + return { + ...actual, + onIdeEvent: (_type: string, listener: (event: KernelEvent) => void) => { + kernelListeners.add(listener); + return () => kernelListeners.delete(listener); + }, + }; +}); +vi.mock("@tauri-apps/api/core", async (original) => ({ + ...(await original()), + Channel: class { + onmessage = (_event: unknown) => {}; + }, +})); +afterEach(() => kernelListeners.clear()); + +function kernelFailure(message: string) { + for (const listener of kernelListeners) + listener({ type: "notebookKernel", id: "document-a", path: "/p/a.ipynb", state: "dead", message }); +} + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (error: unknown) => void; + const promise = new Promise((yes, no) => { + resolve = yes; + reject = no; + }); + return { promise, resolve, reject }; +} +function snapshot(): NotebookSnapshot { + return { + id: "document-a", + path: "/p/a.ipynb", + version: 1, + dirty: false, + nbformat: 4, + nbformatMinor: 5, + metadata: {}, + canUndo: false, + canRedo: false, + cells: [ + { + id: "cell-a", + cellType: "code", + source: "original", + metadata: {}, + attachments: {}, + executionCount: null, + outputs: [], + }, + ], + }; +} +function model() { + let document = snapshot(); + const backend: typeof notebookBackend = { + ...notebookBackend, + open: async () => structuredClone(document), + get: async () => structuredClone(document), + environments: async () => [[], null], + close: async () => null, + interrupt: async () => null, + cancel: async () => {}, + edit: async (_id, version, edit) => { + if (version !== document.version) throw new IpcError({ kind: "outOfSync", message: "The document changed" }); + if (edit.type !== "setSource") throw new Error("This regression model accepts source edits only"); + document = { + ...document, + version: document.version + 1, + dirty: true, + canUndo: true, + cells: document.cells.map((cell) => (cell.id === edit.cell ? { ...cell, source: edit.source } : cell)), + }; + return structuredClone(document); + }, + save: async () => { + document = { ...document, dirty: false }; + return structuredClone(document); + }, + reload: async () => { + document = snapshot(); + return structuredClone(document); + }, + }; + return { + backend, + read: () => document, + replace: (next: NotebookSnapshot) => { + document = next; + }, + }; +} + +describe("shared notebook mutation session", () => { + it("retires event listeners even when the backend project disappears during close", async () => { + const m = model(); + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + const kernel = session.getState().kernel; + const failure = new IpcError({ kind: "notFound", message: "The backend project was replaced" }); + m.backend.close = async () => { + throw failure; + }; + await expect(session.close()).rejects.toBe(failure); + kernelFailure("Late failure from the retired kernel"); + expect(session.getState().error).toBeNull(); + expect(session.getState().kernel).toBe(kernel); + }); + + it("still releases the backend document when interrupting a running batch fails during close", async () => { + const m = model(); + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + const running = deferred(); + const retired = new IpcError({ kind: "notFound", message: "notebook project has closed" }); + const started = deferred(); + m.backend.execute = () => { + started.resolve(); + return running.promise; + }; + m.backend.interrupt = async () => { + running.reject(retired); + throw retired; + }; + const closed = vi.fn(async () => null); + m.backend.close = closed; + const execution = session.execute().catch(() => {}); + await started.promise; + await session.close(); + await execution; + expect(closed).toHaveBeenCalledWith("document-a"); + }); + + it("retains rapid typing over an older acknowledgement and serializes the latest intent", async () => { + const m = model(); + const started = deferred(); + const release = deferred(); + const edit = m.backend.edit; + let first = true; + m.backend.edit = async (...args) => { + if (first) { + first = false; + started.resolve(); + await release.promise; + } + return edit(...args); + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + session.edit({ type: "setSource", cell: "cell-a", source: "first" }); + await started.promise; + session.edit({ type: "setSource", cell: "cell-a", source: "second" }); + session.edit({ type: "setSource", cell: "cell-a", source: "latest typing" }); + const visible: string[] = []; + const unsubscribe = session.subscribe(() => visible.push(session.cells()[0]!.source)); + expect(session.cells()[0]!.source).toBe("latest typing"); + expect(session.isDirty()).toBe(true); + release.resolve(); + await session.flush(); + expect(visible.every((source) => source === "latest typing")).toBe(true); + expect(m.read().cells[0]!.source).toBe("latest typing"); + expect(m.read().version).toBe(3); + await session.save(); + expect(session.isDirty()).toBe(false); + unsubscribe(); + await session.close(); + }); + + it("rebases stable-cell source intent after an OutOfSync response", async () => { + const m = model(); + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + m.replace({ ...snapshot(), version: 4, cells: [{ ...snapshot().cells[0]!, source: "other split changed this" }] }); + session.edit({ type: "setSource", cell: "cell-a", source: "retained user intent" }); + await session.flush(); + expect(m.read().cells[0]!.source).toBe("retained user intent"); + expect(session.getState().snapshot?.version).toBe(5); + expect(session.getState().pending).toBe(0); + await session.close(); + }); + + it("orders typing after an in-flight undo so its response cannot overwrite the new source", async () => { + const m = model(); + const started = deferred(); + const release = deferred(); + m.backend.undo = async () => { + started.resolve(); + await release.promise; + const current = m.read(); + const undone = { + ...current, + version: current.version + 1, + canRedo: true, + cells: current.cells.map((cell) => ({ ...cell, source: "original" })), + }; + m.replace(undone); + return structuredClone(undone); + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + session.edit({ type: "setSource", cell: "cell-a", source: "before undo" }); + const undo = session.undo(); + session.edit({ type: "setSource", cell: "cell-a", source: "immediately after undo" }); + await started.promise; + session.edit({ type: "setSource", cell: "cell-a", source: "typing after undo" }); + expect(session.cells()[0]!.source).toBe("typing after undo"); + release.resolve(); + await undo; + await session.flush(); + expect(m.read().cells[0]!.source).toBe("typing after undo"); + expect(session.cells()[0]!.source).toBe("typing after undo"); + expect(m.read().version).toBe(4); + await session.close(); + }); + + it("keeps failed and subsequent source visible/dirty until an explicit retry succeeds", async () => { + const m = model(); + const edit = m.backend.edit; + let unavailable = true; + m.backend.edit = async (...args) => { + if (unavailable) throw new Error("Backend disconnected"); + return edit(...args); + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + session.edit({ type: "setSource", cell: "cell-a", source: "unsent draft" }); + await expect(session.flush()).rejects.toThrow("Backend disconnected"); + session.edit({ type: "setSource", cell: "cell-a", source: "more typing after failure" }); + expect(session.cells()[0]!.source).toBe("more typing after failure"); + expect(session.pendingSources()).toEqual([{ cell: "cell-a", source: "more typing after failure" }]); + expect(session.isDirty()).toBe(true); + expect(m.read().cells[0]!.source).toBe("original"); + unavailable = false; + await session.save(); + expect(m.read().cells[0]!.source).toBe("more typing after failure"); + expect(session.isDirty()).toBe(false); + await session.close(); + }); + + it("freezes edit admission during a close decision and resumes only after cancellation", async () => { + const m = model(); + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + session.edit({ type: "setSource", cell: "cell-a", source: "checkpoint" }); + await session.prepareClose(); + session.edit({ type: "setSource", cell: "cell-a", source: "must not enter closing document" }); + expect(session.cells()[0]!.source).toBe("checkpoint"); + session.cancelClose(); + session.edit({ type: "setSource", cell: "cell-a", source: "resumed" }); + await session.flush(); + expect(m.read().cells[0]!.source).toBe("resumed"); + await session.close(); + }); + + it("retains a kernel failure arriving before a late source acknowledgement and subsequent undo", async () => { + const m = model(); + const started = deferred(); + const release = deferred(); + const edit = m.backend.edit; + m.backend.edit = async (...args) => { + started.resolve(); + await release.promise; + return edit(...args); + }; + m.backend.undo = async () => { + const current = m.read(); + const undone = { + ...current, + version: current.version + 1, + cells: current.cells.map((cell) => ({ ...cell, source: "original" })), + }; + m.replace(undone); + return structuredClone(undone); + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + session.edit({ type: "setSource", cell: "cell-a", source: "pending draft" }); + await started.promise; + const failure = "Kernel process exited"; + kernelFailure(failure); + release.resolve(); + await session.flush(); + expect(session.cells()[0]!.source).toBe("pending draft"); + expect(session.getState().pending).toBe(0); + expect(session.getState().error).toBe(failure); + await session.undo(); + expect(session.cells()[0]!.source).toBe("original"); + expect(session.getState().error).toBe(failure); + await session.close(); + }); + + it("retains an external save conflict through editing until save succeeds", async () => { + const m = model(); + const save = m.backend.save; + const conflict = new Error("File changed on disk"); + m.backend.save = async () => { + throw conflict; + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + await expect(session.save()).rejects.toBe(conflict); + session.report(conflict); // The editor's promise catch must preserve ownership. + session.edit({ type: "setSource", cell: "cell-a", source: "draft after conflict" }); + await session.flush(); + expect(m.read().cells[0]!.source).toBe("draft after conflict"); + expect(session.isDirty()).toBe(true); + expect(session.getState().error).toBe(conflict.message); + m.backend.save = save; + await session.save(); + expect(session.isDirty()).toBe(false); + expect(session.getState().error).toBeNull(); + await session.close(); + }); + + it("clears a recovered mutation failure on source retry without requiring save", async () => { + const m = model(); + const edit = m.backend.edit; + const failure = new Error("Edit transport unavailable"); + m.backend.edit = async () => { + throw failure; + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + session.edit({ type: "setSource", cell: "cell-a", source: "retained draft" }); + await expect(session.save()).rejects.toBe(failure); + session.report(failure); + expect(session.getState().error).toBe(failure.message); + expect(session.getState().pending).toBe(1); + expect(session.cells()[0]!.source).toBe("retained draft"); + m.backend.edit = edit; + await session.flush(); + expect(m.read().cells[0]!.source).toBe("retained draft"); + expect(session.getState().pending).toBe(0); + expect(session.isDirty()).toBe(true); + expect(session.getState().error).toBeNull(); + await session.close(); + }); + + it("reveals a retained save conflict after a newer mutation failure recovers", async () => { + const m = model(); + const conflict = new Error("External save conflict"); + m.backend.save = async () => { + throw conflict; + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + await expect(session.save()).rejects.toBe(conflict); + const edit = m.backend.edit; + const mutation = new Error("Edit rejected"); + m.backend.edit = async () => { + throw mutation; + }; + session.edit({ type: "setSource", cell: "cell-a", source: "draft" }); + await expect(session.flush()).rejects.toBe(mutation); + expect(session.getState().error).toBe(mutation.message); + m.backend.edit = edit; + await session.flush(); + expect(session.cells()[0]!.source).toBe("draft"); + expect(session.getState().error).toBe(conflict.message); + await session.close(); + }); + + it("retains execution failure across editing and save, then clears it after successful execution", async () => { + const m = model(); + const failure = new Error("Selected interpreter lacks ipykernel"); + m.backend.execute = async () => { + throw failure; + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + await expect(session.execute()).rejects.toBe(failure); + session.report(failure); + session.edit({ type: "setSource", cell: "cell-a", source: "print(42)" }); + await session.save(); + expect(m.read().cells[0]!.source).toBe("print(42)"); + expect(session.getState().error).toBe(failure.message); + m.backend.execute = async () => structuredClone(m.read()); + await session.execute(); + expect(session.getState().running).toBe(false); + expect(session.getState().error).toBeNull(); + await session.close(); + }); + + it("recovers kernel failure on restart without erasing a save conflict", async () => { + const m = model(); + const conflict = new Error("External save conflict"); + m.backend.save = async () => { + throw conflict; + }; + m.backend.restart = async () => structuredClone(m.read()); + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + await expect(session.save()).rejects.toBe(conflict); + kernelFailure("Kernel died"); + await session.restart(); + expect(session.getState().kernel).toBe("idle"); + expect(session.getState().frozen).toBe(false); + expect(session.getState().error).toBe(conflict.message); + await session.close(); + }); + + it("retains environment discovery failure through editing, save, execution and restart until discovery recovers", async () => { + const m = model(); + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + const environments = m.backend.environments; + const failure = new Error("Environment discovery unavailable"); + m.backend.environments = async () => { + throw failure; + }; + await session.detectEnvironments(); + session.edit({ type: "setSource", cell: "cell-a", source: "draft" }); + await session.save(); + m.backend.execute = async () => structuredClone(m.read()); + await session.execute(); + m.backend.restart = async () => structuredClone(m.read()); + await session.restart(); + expect(session.getState().error).toBe(failure.message); + m.backend.environments = environments; + await session.detectEnvironments(); + expect(session.getState().error).toBeNull(); + await session.close(); + }); + + it("recovers reload and document failures after discard without clearing a kernel failure", async () => { + const m = model(); + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + const conflict = new Error("External save conflict"); + m.backend.save = async () => { + throw conflict; + }; + await expect(session.save()).rejects.toBe(conflict); + const edit = m.backend.edit; + m.backend.edit = async () => { + throw new Error("Edit unavailable"); + }; + session.edit({ type: "setSource", cell: "cell-a", source: "discarded draft" }); + await expect(session.flush()).rejects.toThrow("Edit unavailable"); + const reload = m.backend.reload; + const failure = new Error("Reload unavailable"); + m.backend.reload = async () => { + throw failure; + }; + await expect(session.discard()).rejects.toBe(failure); + session.report(failure); + kernelFailure("Kernel is still dead"); + m.backend.reload = reload; + m.backend.edit = edit; + await session.discard(); + expect(session.cells()[0]!.source).toBe("original"); + expect(session.getState().pending).toBe(0); + expect(session.isDirty()).toBe(false); + expect(session.getState().error).toBe("Kernel is still dead"); + session.cancelClose(); + m.backend.restart = async () => structuredClone(m.read()); + await session.restart(); + expect(session.getState().error).toBeNull(); + await session.close(); + }); + + it("retains restart failure across editing and clears it only after a successful restart", async () => { + const m = model(); + const failure = new Error("Kernel restart failed"); + m.backend.restart = async () => { + throw failure; + }; + const session = new NotebookSession("/p/a.ipynb", m.backend, true); + await session.load(); + await expect(session.restart()).rejects.toBe(failure); + session.report(failure); + expect(session.getState().frozen).toBe(false); + session.edit({ type: "setSource", cell: "cell-a", source: "editable after failed restart" }); + await session.flush(); + expect(session.cells()[0]!.source).toBe("editable after failed restart"); + expect(session.getState().error).toBe(failure.message); + m.backend.restart = async () => structuredClone(m.read()); + await session.restart(); + expect(session.getState().error).toBeNull(); + await session.close(); + }); +}); diff --git a/src/features/notebook/session.ts b/src/features/notebook/session.ts new file mode 100644 index 000000000..26c78892d --- /dev/null +++ b/src/features/notebook/session.ts @@ -0,0 +1,509 @@ +import { Channel } from "@tauri-apps/api/core"; +import type { EditorState } from "@codemirror/state"; +import type { NotebookSnapshot, NotebookEdit, NotebookOutputEvent, Sdk, RequestId } from "@/bindings"; +import { call, commands, IpcError, isTauri, nextRequestId, onIdeEvent } from "@/platform/ipc"; +import { notifyEditorDocumentChanged } from "@/workbench/contributions"; +import { OutputReducer } from "./outputs"; + +export const notebookBackend = { + open: (path: string) => call(commands.notebookOpen(path)), + get: (id: string) => call(commands.notebookGet(id)), + edit: (id: string, version: number, edit: NotebookEdit) => call(commands.notebookEdit(id, version, edit)), + save: (id: string) => call(commands.notebookSave(id)), + reload: (id: string) => call(commands.notebookReload(id)), + undo: (id: string) => call(commands.notebookUndo(id)), + redo: (id: string) => call(commands.notebookRedo(id)), + close: (id: string) => call(commands.notebookClose(id)), + interrupt: (id: string) => call(commands.notebookInterrupt(id)), + restart: (id: string) => call(commands.notebookRestart(id)), + execute: ( + id: string, + cells: string[], + python: string | null, + request: RequestId, + output: Channel, + ) => call(commands.notebookExecute(id, cells, python, request, output)), + environments: () => Promise.all([call(commands.sdkDetect("python")), call(commands.sdkPythonEffective())]), + cancel: (request: RequestId) => commands.cancelRequest(request), +}; + +type Pending = { edit: NotebookEdit | { type: "undo" } | { type: "redo" }; committed: boolean }; +type ErrorOwner = "load" | "mutation" | "save" | "execution" | "kernel" | "environment" | "reload" | "other"; + +const PYTHON_KEY = "openidea.notebook.python:"; +function savedPython(path: string): string | null { + try { + return globalThis.localStorage?.getItem(PYTHON_KEY + path) || null; + } catch { + return null; + } +} +export interface NotebookViewState { + snapshot: NotebookSnapshot | null; + loading: boolean; + error: string | null; + running: boolean; + frozen: boolean; + kernel: string; + python: string | null; + environments: Sdk[]; + pending: number; +} + +/** One session/queue per file; tabs and mounted React views do not own its lifetime. */ +export class NotebookSession { + private state: NotebookViewState = { + snapshot: null, + loading: false, + error: null, + running: false, + frozen: false, + kernel: "idle", + python: null, + environments: [], + pending: 0, + }; + private listeners = new Set<() => void>(); + private loading: Promise | null = null; + private pending: Pending[] = []; + private pumping: Promise | null = null; + private blocked = false; + private disposed = false; + private authoritativeVersion = -1; + private executionCancelled = false; + private generation = 0; + private execution: Promise | null = null; + private request: RequestId | null = null; + private refreshing: Promise | null = null; + private refreshAgain = false; + private output = new OutputReducer(); + private stopEvents: (() => void) | null = null; + private errors = new Map(); + // EditorState contains selection and syntax state, but deliberately no local history. + readonly editors = new Map(); + readonly markdownEditing = new Set(); + selectedCell: string | null = null; + scrollTop = 0; + + constructor( + readonly path: string, + private backend = notebookBackend, + private desktop = isTauri(), + ) { + this.state = { ...this.state, python: savedPython(path) }; + } + subscribe = (listener: () => void) => { + this.listeners.add(listener); + return () => { + this.listeners.delete(listener); + }; + }; + getState = () => this.state; + isDirty = () => this.pending.length > 0 || !!this.state.snapshot?.dirty; + private publish(patch: Partial = {}) { + const error = [...this.errors.values()].at(-1)?.message ?? null; + this.state = { ...this.state, ...patch, error, pending: this.pending.length }; + this.listeners.forEach((listener) => listener()); + notifyEditorDocumentChanged(); + } + report(error: unknown) { + this.reportOperation(error, "other"); + } + private reportOperation(error: unknown, owner: ErrorOwner) { + // UI promise catches and operations waiting on flush re-report the same + // failure. Keep its original owner so only the real recovery clears it. + if ([...this.errors.values()].some((entry) => entry.cause === error)) return; + this.errors.delete(owner); + this.errors.set(owner, { cause: error, message: error instanceof Error ? error.message : String(error) }); + this.publish(); + } + private recover(...owners: ErrorOwner[]) { + for (const owner of owners) this.errors.delete(owner); + this.publish(); + } + private accept(snapshot: NotebookSnapshot) { + if (this.disposed) return; + const current = this.state.snapshot; + if (current && (current.id !== snapshot.id || snapshot.version < this.authoritativeVersion)) return; + this.authoritativeVersion = snapshot.version; + if (current && current.version > snapshot.version) { + // Only Channel output reductions can be newer than the latest full + // snapshot. Take acknowledged user fields/order, retain newer outputs. + const outputCells = new Map(current.cells.map((cell) => [cell.id, cell])); + snapshot = { + ...snapshot, + version: current.version, + dirty: current.dirty, + cells: snapshot.cells.map((cell) => { + const newer = outputCells.get(cell.id); + return newer && cell.cellType === "code" + ? { ...cell, outputs: newer.outputs, executionCount: newer.executionCount } + : cell; + }), + }; + } + this.publish({ snapshot }); + const live = new Set(snapshot.cells.map((cell) => cell.id)); + for (const id of this.editors.keys()) if (!live.has(id)) this.editors.delete(id); + } + load(opened?: NotebookSnapshot): Promise { + if (this.loading) return this.loading; + if (this.disposed) return Promise.reject(new Error("Notebook is closed")); + this.publish({ loading: true }); + this.loading = (async () => { + if (!this.desktop) throw new Error("Notebook editing and execution require the desktop app."); + const snapshot = opened ?? (await this.backend.open(this.path)); + // Closing a still-loading tab must release the real backend document too. + if (this.disposed) { + await this.backend.close(snapshot.id); + return; + } + this.accept(snapshot); + this.recover("load"); + this.stopEvents = onIdeEvent("notebookKernel", (event) => { + if (this.disposed || event.path !== snapshot.path || event.id !== snapshot.id) return; + this.publish({ kernel: event.state }); + if (event.message) this.reportOperation(event.message, "kernel"); + }); + void this.detectEnvironments(); + })() + .catch((error) => { + this.reportOperation(error, "load"); + this.loading = null; + throw error; + }) + .finally(() => this.publish({ loading: false })); + return this.loading; + } + async detectEnvironments() { + try { + const [environments, effective] = await this.backend.environments(); + if (!this.disposed) { + this.publish({ environments, python: this.state.python ?? effective }); + this.recover("environment"); + } + } catch (error) { + if (!this.disposed) this.reportOperation(error, "environment"); + } + } + selectPython(python: string) { + if (this.state.running || this.state.frozen) return; + try { + globalThis.localStorage?.setItem(PYTHON_KEY + this.path, python); + } catch { + /* Interpreter preference is optional when storage is unavailable. */ + } + this.publish({ python }); + } + pendingSources() { + const sources = new Map(); + for (const { edit } of this.pending) if (edit.type === "setSource") sources.set(edit.cell, edit.source); + return [...sources].map(([cell, source]) => ({ cell, source })); + } + + /** Pending source is overlaid, never replaced by an older IPC response. */ + cells() { + const cells = this.state.snapshot?.cells ?? []; + return cells.map((cell) => { + let source = cell.source; + for (const item of this.pending) + if (item.edit.type === "setSource" && item.edit.cell === cell.id) source = item.edit.source; + return source === cell.source ? cell : { ...cell, source }; + }); + } + edit(edit: NotebookEdit) { + if (this.disposed || this.state.frozen || !this.state.snapshot) return; + if (this.state.running && edit.type !== "setSource") { + this.reportOperation(new Error("Stop execution before changing notebook structure."), "mutation"); + return; + } + const last = this.pending.at(-1); + // Never mutate an in-flight intent; only coalesce an unsent tail. + if ( + edit.type === "setSource" && + last && + !last.committed && + last !== this.inflight && + last.edit.type === "setSource" && + last.edit.cell === edit.cell + ) + last.edit = edit; + else this.pending.push({ edit, committed: false }); + this.publish(); + if (!this.blocked) void this.pump().catch(() => {}); + } + private inflight: Pending | null = null; + private pump(): Promise { + if (this.pumping) return this.pumping; + this.pumping = Promise.resolve() + .then(async () => { + while (this.pending.length && !this.disposed) { + const item = this.pending[0]!; + this.inflight = item; + let snapshot = this.state.snapshot!; + if (!item.committed) { + try { + const result = + item.edit.type === "undo" || item.edit.type === "redo" + ? await this.backend[item.edit.type](snapshot.id) + : await this.backend.edit(snapshot.id, snapshot.version, item.edit); + item.committed = true; + this.accept(result); + } catch (error) { + if ( + !(error instanceof IpcError) || + error.kind !== "outOfSync" || + item.edit.type === "undo" || + item.edit.type === "redo" + ) + throw error; + snapshot = await this.backend.get(snapshot.id); + this.accept(snapshot); + // Semantic stable-id rebase, once. Missing/deleted targets remain pending + // with an actionable backend error rather than silently dropping user intent. + const result = await this.backend.edit(snapshot.id, snapshot.version, item.edit); + item.committed = true; + this.accept(result); + } + } + this.pending.shift(); + this.recover("mutation"); + } + }) + .catch((error) => { + this.blocked = true; + this.reportOperation(error, "mutation"); + throw error; + }) + .finally(() => { + this.inflight = null; + this.pumping = null; + }); + return this.pumping; + } + async flush() { + await this.load(); + this.blocked = false; + await this.pump(); + if (this.pending.length) + throw new Error("Notebook edits have not been committed. Retry edits or discard explicitly."); + } + async save() { + try { + await this.flush(); + this.accept(await this.backend.save(this.state.snapshot!.id)); + this.recover("save"); + } catch (error) { + this.reportOperation(error, "save"); + throw error; + } + } + private async history(direction: "undo" | "redo") { + if (this.state.running || this.state.frozen) throw new Error("Stop execution before undo or redo."); + if (!this.state.snapshot) await this.load(); + if (this.state.running || this.state.frozen || this.disposed) throw new Error("Notebook is busy or closing."); + // History is a queue barrier: later typing must commit AFTER undo/redo, not + // race it and get overwritten by the history response. + this.pending.push({ edit: { type: direction }, committed: false }); + this.publish(); + this.blocked = false; + await this.pump(); + } + undo = () => this.history("undo"); + redo = () => this.history("redo"); + + private refresh(): Promise { + this.refreshAgain = true; + if (this.refreshing) return this.refreshing; + this.refreshing = (async () => { + while (this.refreshAgain && !this.disposed) { + this.refreshAgain = false; + const id = this.state.snapshot!.id; + this.accept(await this.backend.get(id)); + } + })().finally(() => { + this.refreshing = null; + }); + return this.refreshing; + } + execute(cells?: string[]): Promise { + if (this.state.running || this.state.frozen || this.disposed) + return Promise.reject(new Error("Notebook is busy or closing.")); + this.publish({ running: true }); + this.executionCancelled = false; + const generation = ++this.generation; + this.output.reset(); + this.execution = (async () => { + await this.flush(); + if (this.state.frozen || this.executionCancelled || generation !== this.generation) return; + const snapshot = this.state.snapshot!; + const ids = cells ?? snapshot.cells.filter((cell) => cell.cellType === "code").map((cell) => cell.id); + if (!ids.length) return; + const request = nextRequestId(); + this.request = request; + const channel = new Channel(); + let buffered: NotebookOutputEvent[] = []; + let frame: number | null = null; + let overflow = false; + channel.onmessage = (event) => { + if (this.disposed || generation !== this.generation || event.id !== snapshot.id) return; + // Bound retained Channel work. A flood recovers from backend state instead + // of retaining an unbounded browser queue or losing authoritative output. + if (buffered.length < 256 && !overflow) buffered.push(event); + else { + overflow = true; + buffered = []; + this.output.recover(); + } + if (frame !== null) return; + frame = requestAnimationFrame(() => { + frame = null; + if (this.disposed || generation !== this.generation) { + buffered = []; + return; + } + let current = this.state.snapshot!; + let refresh = overflow; + overflow = false; + for (const item of buffered) { + const result = this.output.reduce(current, item); + current = result.snapshot; + refresh ||= result.refresh; + } + buffered = []; + if (current !== this.state.snapshot) this.publish({ snapshot: current }); + if (refresh) void this.refresh().catch((error) => this.reportOperation(error, "execution")); + }); + }; + try { + const result = await this.backend.execute(snapshot.id, ids, this.state.python, request, channel); + if (generation === this.generation) this.accept(result); + } finally { + if (generation === this.generation && !this.disposed) await this.refresh(); + } + if (generation === this.generation && !this.executionCancelled) this.recover("execution", "kernel"); + })() + .catch((error) => { + this.reportOperation(error, "execution"); + throw error; + }) + .finally(() => { + this.request = null; + this.execution = null; + this.publish({ running: false }); + }); + return this.execution; + } + async interrupt() { + this.executionCancelled = true; + if (!this.state.snapshot) return; + // Batch cancellation precedes interrupt in the host; cancelRequest is an + // additional route for requests still waiting on kernel admission/startup. + try { + if (this.request !== null) await this.backend.cancel(this.request); + } finally { + await this.backend.interrupt(this.state.snapshot.id); + } + } + private async drain() { + const execution = this.execution; + if (!execution) return; + try { + await this.interrupt(); + } finally { + await execution.catch(() => {}); + } + } + async restart() { + if (this.state.frozen) throw new Error("Notebook is closing."); + this.publish({ frozen: true }); + ++this.generation; + try { + await this.drain(); + await this.flush(); + this.output.reset(); + this.accept(await this.backend.restart(this.state.snapshot!.id)); + this.publish({ kernel: "idle" }); + this.recover("execution", "kernel"); + } catch (error) { + this.reportOperation(error, "kernel"); + throw error; + } finally { + this.publish({ frozen: false }); + } + } + async prepareClose() { + this.publish({ frozen: true }); + if (this.loading) await this.loading.catch(() => {}); + // An invalid/unavailable file never acquired a backend document. + if (!this.state.snapshot) return; + try { + await this.flush(); + } finally { + await this.drain(); + } + } + cancelClose() { + if (!this.disposed) this.publish({ frozen: false }); + } + async discard() { + this.publish({ frozen: true }); + ++this.generation; + try { + await this.drain(); + if (this.pumping) await this.pumping.catch(() => {}); + const snapshot = this.state.snapshot; + if (!snapshot) return; + const reloaded = await this.backend.reload(snapshot.id); + this.pending = []; + this.blocked = false; + this.editors.clear(); + this.output.reset(); + this.accept(reloaded); + this.recover("mutation", "save", "reload"); + } catch (error) { + this.reportOperation(error, "reload"); + throw error; + } + } + async close() { + this.publish({ frozen: true }); + ++this.generation; + try { + // The host close interrupts, drains and shuts down authoritatively, and a + // retired project already rejects interrupt. Never let that skip release. + await this.drain().catch(() => {}); + if (this.pumping) await this.pumping.catch(() => {}); + if (this.loading) await this.loading.catch(() => {}); + const snapshot = this.state.snapshot; + if (snapshot) await this.backend.close(snapshot.id); + } finally { + this.disposed = true; + this.stopEvents?.(); + this.editors.clear(); + this.listeners.clear(); + } + } +} + +const sessions = new Map(); +export function notebookSession(path: string): NotebookSession { + let session = sessions.get(path); + if (!session) { + session = new NotebookSession(path); + sessions.set(path, session); + } + return session; +} +export function existingNotebookSession(path: string) { + return sessions.get(path); +} +export async function closeNotebook(path: string) { + const session = sessions.get(path); + if (!session) return; + try { + await session.close(); + } finally { + if (sessions.get(path) === session) sessions.delete(path); + } +} diff --git a/src/features/remote-dev/index.ts b/src/features/remote-dev/index.ts index f58905968..607a41652 100644 --- a/src/features/remote-dev/index.ts +++ b/src/features/remote-dev/index.ts @@ -5,7 +5,8 @@ */ import { registerActions } from "@/platform/actions"; import { disconnectRemote } from "@/platform/remote"; -import { registerOverlay } from "@/workbench/contributions"; +import { maybeEditorService, registerOverlay } from "@/workbench/contributions"; +import { confirmLeaveProject } from "@/platform/project"; import { ConnectHost } from "./ConnectDialog"; import { reloadEnvironment } from "./environment"; import { closeConnectDialog, openConnectDialog } from "./model"; @@ -29,7 +30,15 @@ export function activate(): () => void { category: "File", keywords: ["ssh", "remote", "close"], when: "projectOpen", - run: () => disconnectRemote(), + run: async () => { + if (!(await confirmLeaveProject())) return; + try { + await disconnectRemote(); + } catch (e) { + maybeEditorService()?.cancelClose?.(); + throw e; + } + }, }, { id: "remote.reloadEnvironment", diff --git a/src/features/search/recent/RecentPopups.tsx b/src/features/search/recent/RecentPopups.tsx index f6688c18f..20220c24d 100644 --- a/src/features/search/recent/RecentPopups.tsx +++ b/src/features/search/recent/RecentPopups.tsx @@ -127,7 +127,7 @@ function RecentFilesDialog() { }; // IDEA: Delete closes the selected editor (and drops it from the list) or hides the tool window. - const removeSelected = () => { + const removeSelected = async () => { if (effectiveColumn === "tools") { const tw = tools[tIndex]?.tw; if (tw) toolWindowService().hide(tw.id); @@ -135,8 +135,10 @@ function RecentFilesDialog() { } const item = files[fIndex]; if (!item) return; - if (item.editor) maybeEditorService()?.closeInput?.(item.path); - else maybeEditorService()?.closePath?.(item.path); + const closed = item.editor + ? await maybeEditorService()?.closeInput?.(item.path) + : await maybeEditorService()?.closePath?.(item.path); + if (closed === false) return; if (!switcher) void call(commands.removeRecent(item.path)).catch(() => {}); setRemoved((r) => new Set(r).add(item.path)); setFileIndex(fIndex); @@ -192,7 +194,7 @@ function RecentFilesDialog() { // In Recent Files Backspace edits the speed search text first. if (e.key === "Delete" || (e.key === "Backspace" && !query)) { e.preventDefault(); - removeSelected(); + void removeSelected(); return; } if (e.key === "Tab" || ((e.key === "ArrowLeft" || e.key === "ArrowRight") && !query)) { diff --git a/src/features/vcs/MergeEditor.tsx b/src/features/vcs/MergeEditor.tsx index b9960c869..2d955dae1 100644 --- a/src/features/vcs/MergeEditor.tsx +++ b/src/features/vcs/MergeEditor.tsx @@ -280,7 +280,7 @@ export function MergeEditor({ input }: { input: EditorInput }) { setResolved(true); await refreshNow(); // Like IDEA, applying the result closes the merge editor. - editorService().closeInput?.(input.key); + await editorService().closeInput?.(input.key); } }; diff --git a/src/features/welcome/session.test.ts b/src/features/welcome/session.test.ts index 45c464d37..b18239a56 100644 --- a/src/features/welcome/session.test.ts +++ b/src/features/welcome/session.test.ts @@ -21,7 +21,7 @@ const { provideStartupLaunch, useProject } = await import("@/platform/project"); beforeEach(() => { openProject.mockReset(); - useProject.setState({ project: null }); + useProject.setState({ project: null, lightEdit: false }); sessionStorage.clear(); localStorage.setItem("openidea.session.openProject", "/p/last"); _resetSessionForTests(); @@ -46,12 +46,21 @@ describe("reopenLastProject", () => { expect(openProject).not.toHaveBeenCalled(); }); - it("leaves files opened at launch in LightEdit", async () => { - provideStartupLaunch(Promise.resolve().then(() => useProject.setState({ lightEdit: true }))); - await reopenLastProject(); - expect(openProject).not.toHaveBeenCalled(); - useProject.setState({ lightEdit: false }); - }); + it.each(["last session", "window session"])( + "keeps launched loose files instead of restoring the %s project", + async (remembered) => { + if (remembered === "window session") sessionStorage.setItem("openidea.window.project", "/p/window"); + provideStartupLaunch( + Promise.resolve().then(() => { + useProject.setState({ lightEdit: true }); + }), + ); + await reopenLastProject(); + expect(openProject).not.toHaveBeenCalled(); + expect(useProject.getState().lightEdit).toBe(true); + expect(useProject.getState().project).toBeNull(); + }, + ); it("reopens the last project when the launch named none", async () => { provideStartupLaunch(Promise.resolve()); diff --git a/src/features/welcome/session.ts b/src/features/welcome/session.ts index 65c4b98a5..c772a1d56 100644 --- a/src/features/welcome/session.ts +++ b/src/features/welcome/session.ts @@ -87,7 +87,10 @@ let attempted = false; export async function reopenLastProject(): Promise { if (attempted || !isTauri()) return; attempted = true; - if (useProject.getState().project) return; + // An explicit launch wins over both this window's and the last session's project. + await startupLaunchHandled(); + const startup = useProject.getState(); + if (startup.project || startup.lightEdit) return; // A reloaded webview gets back its own project, whatever window it is. const own = windowProject(); if (own) { @@ -96,7 +99,6 @@ export async function reopenLastProject(): Promise { } // Secondary windows open the project they were created for (features/app/windows.ts). if (getCurrentWindow().label !== "main") return; - if (useProject.getState().project) return; const path = lastSessionProject(); if (!path) return; let settings: Settings | null = null; @@ -106,9 +108,8 @@ export async function reopenLastProject(): Promise { /* use the local preference */ } if (!reopenEnabled(settings, useSession.getState().reopenPref)) return; - // A project or files passed on the command line (or from Finder) win over the last session's. - await startupLaunchHandled(); - if (useProject.getState().project || useProject.getState().lightEdit) return; + const current = useProject.getState(); + if (current.project || current.lightEdit) return; await reopen(path); } diff --git a/src/platform/locales/text/zh.json b/src/platform/locales/text/zh.json index 3f89ea982..aba694d8c 100644 --- a/src/platform/locales/text/zh.json +++ b/src/platform/locales/text/zh.json @@ -4765,5 +4765,51 @@ "No matching fonts. Enter a family name manually.": "没有匹配的字体。可手动输入字体家族名称。", "This name was not found on this desktop. The saved preference is kept.": "此桌面未发现该字体，已保留保存的设置。", "No system fonts were returned. Bundled fonts and manual input are available.": "未发现系统字体。可选择内置字体或手动输入。", - "Type a family name or search installed fonts.": "输入字体家族名称或搜索已安装字体。" + "Type a family name or search installed fonts.": "输入字体家族名称或搜索已安装字体。", + "Add Code Cell": "添加代码单元格", + "Add Markdown Cell": "添加 Markdown 单元格", + "Cell": "单元格", + "Cell outputs": "单元格输出", + "Cell type": "单元格类型", + "Delete Cell": "删除单元格", + "Discard Changes and Reload": "放弃更改并重新加载", + "Discard all unsaved notebook changes and stop execution? This cannot be undone.": "放弃笔记本的所有未保存更改并停止执行？此操作无法撤消。", + "Edit Markdown": "编辑 Markdown", + "Inert JavaScript source": "不执行的 JavaScript 源代码", + "Insert Code Cell Below": "在下方插入代码单元格", + "Install ipykernel in the selected Python environment, then restart the kernel. No packages are installed automatically.": "请在所选 Python 环境中安装 ipykernel，然后重启内核。不会自动安装任何软件包。", + "Interrupt Kernel": "中断内核", + "JavaScript output is inert. Interactive widgets are not supported.": "JavaScript 输出不会执行。不支持交互式组件。", + "Loading notebook…": "正在加载笔记本…", + "Markdown cell source": "Markdown 单元格源代码", + "Move Cell Down": "下移单元格", + "Move Cell Up": "上移单元格", + "Notebook Python interpreter": "笔记本 Python 解释器", + "Notebook editing and execution require the desktop app.": "编辑和执行笔记本需要使用桌面应用。", + "Notebook image output": "笔记本图像输出", + "Notebook path is missing": "缺少笔记本路径", + "Notebook toolbar": "笔记本工具栏", + "Notebook unavailable": "笔记本不可用", + "Output preview truncated; full output is retained in the notebook.": "输出预览已截断；完整输出仍保留在笔记本中。", + "Preview Markdown": "预览 Markdown", + "Python cell source": "Python 单元格源代码", + "Raw": "原始内容", + "Raw or unknown cell; original content is preserved read-only.": "原始或未知类型的单元格；原始内容以只读方式保留。", + "Redo": "重做", + "Refresh Python interpreters": "刷新 Python 解释器", + "Reload Notebook from Disk": "从磁盘重新加载笔记本", + "Restart Kernel": "重启内核", + "Retry Pending Edits": "重试待提交更改", + "Run All": "运行全部", + "Run Cell": "运行单元格", + "Select Python interpreter": "选择 Python 解释器", + "Static notebook HTML output": "笔记本静态 HTML 输出", + "This notebook has no cells": "此笔记本没有单元格", + "Uncommitted source edits": "未提交的源代码更改", + "Unsupported output MIME type; original data is preserved.": "不支持此输出 MIME 类型；原始数据已保留。", + "Unsupported output type; original data is preserved.": "不支持此输出类型；原始数据已保留。", + "busy": "忙碌", + "dead": "已终止", + "idle": "空闲", + "starting": "正在启动" } diff --git a/src/workbench/contributions.ts b/src/workbench/contributions.ts index b7bdeb2e9..b3fa65d61 100644 --- a/src/workbench/contributions.ts +++ b/src/workbench/contributions.ts @@ -118,6 +118,53 @@ export interface EditorKindContribution { * `data` must be plain JSON). The editor closes tabs whose target is gone itself. */ restorable?: boolean; + /** + * The editor owns a document with unsaved state (a notebook): it takes part in + * the dirty mark, Save All, autosave and the close prompts. Tabs with equal + * `input.key` (splits) share one document. + */ + document?: EditorDocumentLifecycle; +} + +/** + * Document lifecycle of a custom editor kind, keyed by its tabs' input. The + * workbench calls it whatever React has mounted: hidden tabs keep their + * document. Call `notifyEditorDocumentChanged` whenever `isDirty` may change. + */ +export interface EditorDocumentLifecycle { + /** Unsaved changes, pending edits not committed yet included. */ + isDirty(input: EditorInput): boolean; + /** Commits pending edits; rejects when they can't be. */ + flush(input: EditorInput): Promise; + /** Flushes and saves; rejects on failure (the document stays dirty and open). */ + save(input: EditorInput): Promise; + /** + * Before a close decision: stops admitting edits and runs, flushes, and + * interrupts and drains execution, so `isDirty` holds until `cancelClose`. + */ + prepareClose(input: EditorInput): Promise; + /** The close didn't happen (Cancel, a failed save, a tab of it opened again): edits and runs resume. */ + cancelClose(input: EditorInput): void; + /** Don't Save: drops unsaved changes without saving them. */ + discard(input: EditorInput): Promise; + /** Releases the document (and its kernel…) once no tab shows it any more. Called once per open. */ + close(input: EditorInput): Promise; + /** Edit › Undo / Redo while the editor has focus. */ + undo?(input: EditorInput): Promise; + redo?(input: EditorInput): Promise; +} + +/** Bumped by `notifyEditorDocumentChanged`; selectors reading `isDirty` re-run on it. */ +export const useEditorDocuments = create<{ revision: number }>(() => ({ revision: 0 })); + +/** A custom editor document's dirty state (or pending edits) may have changed: tab marks and counts refresh. */ +export function notifyEditorDocumentChanged() { + useEditorDocuments.setState((s) => ({ revision: s.revision + 1 })); +} + +/** Changes with every `notifyEditorDocumentChanged`, for components that read `isDirty`. */ +export function useEditorDocumentRevision(): number { + return useEditorDocuments((s) => s.revision); } /** @@ -307,11 +354,23 @@ export const provideProjectPopup = (component: ComponentType) // editor kind (images, hex view) before the text editor. A provider may probe // the file (size, binary content); the first non-null answer wins. +/** Resources acquired while resolving belong to the request until its tab is admitted. */ +export interface FileEditorResolutionContext { + /** Commit is synchronous: no close/project transition may interleave with tab insertion. */ + onAdmission(commit: () => void, rollback: () => Promise): void; +} + +export interface FileEditorResolution { + input: EditorInput; + admit(): void; + rollback(): Promise; +} + export interface FileEditorProvider { id: string; /** Lower asks first. Default 100. */ order?: number; - resolve(path: string): EditorInput | null | Promise; + resolve(path: string, context: FileEditorResolutionContext): EditorInput | null | Promise; } const fileEditorProviders = new Map(); @@ -324,15 +383,39 @@ export function registerFileEditorProvider(p: FileEditorProvider): () => void { } /** The non-text editor input for `path`, or null for the text editor. Provider failures fall through. */ -export async function resolveFileEditor(path: string): Promise { +export async function resolveFileEditor(path: string): Promise { const providers = [...fileEditorProviders.values()].sort((a, b) => (a.order ?? 100) - (b.order ?? 100)); for (const p of providers) { + const acquisitions: Array<{ commit: () => void; rollback: () => Promise }> = []; + let admitted = false; + const rollback = async () => { + if (admitted) return; + const results = await Promise.allSettled(acquisitions.splice(0).map((a) => Promise.resolve().then(a.rollback))); + const errors = results.filter((r): r is PromiseRejectedResult => r.status === "rejected"); + if (errors.length) + throw new AggregateError( + errors.map((r) => r.reason), + "File editor acquisition cleanup failed", + ); + }; try { - const input = await p.resolve(path); - if (input) return input; + const input = await p.resolve(path, { + onAdmission: (commit, release) => acquisitions.push({ commit, rollback: release }), + }); + if (input) + return { + input, + admit: () => { + for (const acquisition of acquisitions) acquisition.commit(); + admitted = true; + acquisitions.length = 0; + }, + rollback, + }; } catch (e) { console.warn(`file editor provider ${p.id} failed`, e); } + await rollback(); } return null; } @@ -412,7 +495,7 @@ export interface EditorService { activePath(): string | null; /** Doc id + cursor of the focused text editor. */ activeCursor(): { doc: number; path: string; line: number; col: number; offset: number } | null; - /** Save all dirty documents. */ + /** Save all dirty documents; rejects with the failures once every save was tried. */ saveAll(): Promise; /** Move keyboard focus to the active editor (used by Escape from tool windows). */ focus?(): void; @@ -424,12 +507,37 @@ export interface EditorService { activeInput?(): EditorInput | null; /** Calls `fn` whenever the active tab changes; returns the unsubscribe. */ onActiveInputChange?(fn: (input: EditorInput | null) => void): () => void; - /** Closes every tab showing `path`. */ - closePath?(path: string): void; - /** Closes the tab of the input with this `key` (e.g. a merge editor once applied). */ - closeInput?(key: string): void; + /** + * Closes every tab showing `path`. Tabs without a custom document close at + * once; unsaved custom documents are asked about first. Resolves false when + * the user kept them (Cancel, a failed save). + */ + closePath?(path: string): Promise; + /** Closes the tab of the input with this `key` (e.g. a merge editor once applied), like `closePath`. */ + closeInput?(key: string): Promise; /** Reformat Code on the whole open document at `path` (quietly); whether anything changed. */ reformat?(path: string): Promise; + /** + * Window / project close: prepares every open custom editor document for the + * Save / Don't Save / Cancel decision (see `EditorCloseSession`). + */ + prepareClose?(): Promise; + /** Backend project/window transition failed after authorization: resume its prepared documents. */ + cancelClose?(): void; + /** Edit › Undo / Redo for a focused custom editor that has its own history; false when none takes it. */ + documentHistory?(command: "undo" | "redo"): boolean; +} + +/** Custom editor documents held still for a window / project close decision. */ +export interface EditorCloseSession { + /** Documents with unsaved changes. */ + dirty: number; + /** Don't Save: drops their unsaved changes. */ + discard(): Promise; + /** The close didn't happen (Cancel, a failed save): the documents resume. */ + cancel(): void; + /** The close goes ahead: they stay held until their tabs close (they resume if the project stays open). */ + commit(): void; } export interface ToolWindowService { diff --git a/src/workbench/fileEditors.test.ts b/src/workbench/fileEditors.test.ts index d14275aa4..aa613a31d 100644 --- a/src/workbench/fileEditors.test.ts +++ b/src/workbench/fileEditors.test.ts @@ -27,12 +27,52 @@ describe("file editor providers", () => { resolve: (p) => (p.endsWith(".png") ? { kind: "image", key: p, path: p } : null), }), ]; - expect((await resolveFileEditor("/a.png"))?.kind).toBe("image"); - expect((await resolveFileEditor("/a.txt"))?.kind).toBe("late"); + expect((await resolveFileEditor("/a.png"))?.input.kind).toBe("image"); + expect((await resolveFileEditor("/a.txt"))?.input.kind).toBe("late"); disposers.forEach((d) => d()); expect(await resolveFileEditor("/a.png")).toBeNull(); }); + it("releases a failed provider's acquisition before trying the next provider", async () => { + const live = new Set(["failed", "winner"]); + const disposers = [ + registerFileEditorProvider({ + id: "acquired-failure", + order: 1, + resolve: (_path, context) => { + context.onAdmission( + () => {}, + async () => { + live.delete("failed"); + }, + ); + throw new Error("probe failed after acquisition"); + }, + }), + registerFileEditorProvider({ + id: "acquired-winner", + order: 2, + resolve: (_path, context) => { + context.onAdmission( + () => {}, + async () => { + live.delete("winner"); + }, + ); + return { kind: "viewer", key: "winner" }; + }, + }), + ]; + try { + const resolution = (await resolveFileEditor("/a"))!; + expect(live).toEqual(new Set(["winner"])); + await resolution.rollback(); + expect(live).toEqual(new Set()); + } finally { + disposers.forEach((dispose) => dispose()); + } + }); + it("finds the text editor wrapper of a path", () => { const dispose = registerTextEditorWrapper({ id: "md", matches: (p) => p.endsWith(".md"), component: () => null }); const wrappers = useTextEditorWrappers.getState().wrappers; diff --git a/src/workbench/projectWindow.ts b/src/workbench/projectWindow.ts index 2d83c1149..701147f64 100644 --- a/src/workbench/projectWindow.ts +++ b/src/workbench/projectWindow.ts @@ -4,6 +4,7 @@ import { call, commands } from "@/platform/ipc"; import { confirmLeaveProject, useProject } from "@/platform/project"; import { connectRemote } from "@/platform/remote"; import { ask } from "./messages"; +import { maybeEditorService } from "./contributions"; /** "Open project in": `new` | `same` | `ask`. */ export function openPolicy(): "ask" | "new" | "same" { @@ -62,6 +63,11 @@ export async function openRemoteProject( return true; } if (!(await confirmLeaveProject())) return false; - await here(); - return true; + try { + await here(); + return true; + } catch (e) { + maybeEditorService()?.cancelClose?.(); + throw e; + } } diff --git a/src/workbench/providerAcquisition.test.ts b/src/workbench/providerAcquisition.test.ts new file mode 100644 index 000000000..fe3c816c2 --- /dev/null +++ b/src/workbench/providerAcquisition.test.ts @@ -0,0 +1,91 @@ +import { describe, expect, it } from "vitest"; +import { ProviderAcquisitions } from "./providerAcquisition"; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +function fixture() { + const owned = new Set(); + const closed: string[] = []; + const pool = new ProviderAcquisitions<{ id: string }>( + (value) => owned.has(value.id), + async (value) => { + closed.push(value.id); + }, + ); + return { owned, closed, pool }; +} + +describe("provider incarnation ownership", () => { + it("does not close an admitted alias when another resolution is abandoned", async () => { + const { pool, owned, closed } = fixture(); + const first = await pool.acquire(async () => ({ id: "shared" })); + const response = deferred<{ id: string }>(); + const started = deferred(); + const opening = pool.acquire(() => { + started.resolve(); + return response.promise; + }); + await started.promise; + const rollback = first.rollback(); + response.resolve({ id: "shared" }); + const second = await opening; + // The unresolved acquisition now owns a lease, even before tab admission. + await rollback; + expect(closed).toEqual([]); + owned.add("shared"); + second.admit(); + expect(closed).toEqual([]); + }); + + it("releases a shared backend once when every alias is abandoned", async () => { + const { pool, closed } = fixture(); + const a = await pool.acquire(async () => ({ id: "shared" })); + const b = await pool.acquire(async () => ({ id: "shared" })); + await Promise.all([a.rollback(), b.rollback()]); + await a.rollback(); + expect(closed).toEqual(["shared"]); + }); + + it("never releases an existing session or a replacement incarnation", async () => { + const { pool, owned, closed } = fixture(); + owned.add("existing"); + const existing = await pool.acquire(async () => ({ id: "existing" })); + await existing.rollback(); + const old = await pool.acquire(async () => ({ id: "old-project" })); + owned.add("new-project"); + const replacement = await pool.acquire(async () => ({ id: "new-project" })); + replacement.admit(); + await old.rollback(); + expect(closed).toEqual(["old-project"]); + expect(owned).toEqual(new Set(["existing", "new-project"])); + }); + + it("waits for abandoned cleanup before reopening the backend path", async () => { + const closed = deferred(); + const closing = deferred(); + let incarnation = "old"; + const pool = new ProviderAcquisitions<{ id: string }>( + () => false, + async () => { + closing.resolve(); + await closed.promise; + incarnation = "new"; + }, + ); + const old = await pool.acquire(async () => ({ id: incarnation })); + const rollback = old.rollback(); + await closing.promise; + const reopening = pool.acquire(async () => ({ id: incarnation })); + closed.resolve(); + await rollback; + const fresh = await reopening; + expect(fresh.value.id).toBe("new"); + fresh.admit(); + }); +}); diff --git a/src/workbench/providerAcquisition.ts b/src/workbench/providerAcquisition.ts new file mode 100644 index 000000000..c0a1ee3a3 --- /dev/null +++ b/src/workbench/providerAcquisition.ts @@ -0,0 +1,63 @@ +/** Resolver leases protect shared backend incarnations before their tabs own them. */ +export class ProviderAcquisitions { + private pending = 0; + private idle: Array<() => void> = []; + private resources = new Map }>(); + private closing: Promise = Promise.resolve(); + + constructor( + private owned: (value: T) => boolean, + private close: (value: T) => Promise, + ) {} + + async acquire(open: () => Promise) { + // Do not let a new open reuse an incarnation whose rollback is in flight. + let closing: Promise; + do { + closing = this.closing; + await closing; + } while (closing !== this.closing); + ++this.pending; + let value: T; + try { + value = await open(); + const resource = this.resources.get(value.id); + if (resource) ++resource.leases; + else this.resources.set(value.id, { value, leases: 1 }); + } finally { + if (--this.pending === 0) this.idle.splice(0).forEach((resolve) => resolve()); + } + let settled = false; + const release = async (admitted: boolean) => { + if (settled) return; + settled = true; + const resource = this.resources.get(value.id)!; + --resource.leases; + if (!admitted && this.pending) await new Promise((resolve) => this.idle.push(resolve)); + // An unresolved alias may have acquired this same id; an admitted session + // (including one that predated this request) owns it independently. + if (resource.leases || this.owned(value)) { + if (!resource.leases) this.resources.delete(value.id); + return; + } + if (admitted) { + this.resources.delete(value.id); + return; + } + if (!resource.closing) { + resource.closing = this.close(value).finally(() => { + if (this.resources.get(value.id) === resource) this.resources.delete(value.id); + }); + this.closing = Promise.all([this.closing, resource.closing]).catch(() => {}); + } + await resource.closing; + }; + return { + value, + admit: () => { + void release(true); + }, + rollback: () => release(false), + }; + } +} diff --git a/src/workbench/shellActions.ts b/src/workbench/shellActions.ts index 51b294e53..8ef0762d9 100644 --- a/src/workbench/shellActions.ts +++ b/src/workbench/shellActions.ts @@ -9,7 +9,7 @@ import { NOTIFICATIONS_TOOL_WINDOW } from "./Notifications"; import { activateToolWindow } from "./toolwindows/store"; import { toggleTheme } from "./theme"; import { cancelTask, canCancel, useTasks } from "./tasks"; -import { showStatusInfo } from "./contributions"; +import { maybeEditorService, showStatusInfo } from "./contributions"; export async function openProjectDialog() { if (!isTauri()) { @@ -30,6 +30,7 @@ export async function openProjectDialog() { await executeAction("file.openProjectPath", path); return; } + if (!(await confirmLeaveProject())) return; try { await openProject(path); } catch (e) { @@ -66,7 +67,13 @@ export function registerShellActions(): () => void { category: "File", when: "projectOpen", run: async () => { - if (useProject.getState().project && (await confirmLeaveProject())) await closeProject(); + if (!useProject.getState().project || !(await confirmLeaveProject())) return; + try { + await closeProject(); + } catch (e) { + maybeEditorService()?.cancelClose?.(); + throw e; + } }, }, {