From 052fdcfc34a6e19ed6deebee2935896b5c7e188f Mon Sep 17 00:00:00 2001 From: Rohan Kumar Date: Thu, 1 Oct 2026 20:17:54 -0700 Subject: [PATCH 1/3] feat(analyzer): add MCP for collaboration with agents --- Cargo.lock | 89 ++ crates/analyzer/Cargo.toml | 2 + crates/analyzer/src/agent.rs | 1066 ++++++++++++++++++++ crates/analyzer/src/cli.rs | 84 +- crates/analyzer/src/document.rs | 4 + crates/analyzer/src/hook.rs | 125 +++ crates/analyzer/src/lib.rs | 227 ++++- crates/analyzer/src/mcp.rs | 741 ++++++++++++++ crates/analyzer/src/rpc.rs | 83 +- crates/analyzer/src/session.rs | 219 ++++ crates/analyzer/src/transport.rs | 340 +++++++ crates/analyzer/src/watcher.rs | 221 ++++ crates/compiler/src/compile.rs | 38 + crates/gui/src/cli.rs | 334 ++---- crates/gui/src/editor/canvas.rs | 45 + crates/gui/src/editor/mod.rs | 95 +- crates/gui/src/editor/render_regression.rs | 3 + crates/gui/src/lib.rs | 45 +- crates/gui/src/rpc.rs | 239 ++--- crates/tests/Cargo.toml | 1 + crates/tests/fixtures/nvim/full_stack.lua | 141 +++ crates/tests/src/full_stack.rs | 728 +++++++++++-- docs/developers.md | 10 +- docs/docs/gui/shortcuts-config.md | 5 +- docs/docs/tools/agents.md | 89 ++ docs/docs/tools/argone.md | 8 +- docs/docs/tools/neovim.md | 5 +- docs/docs/tools/overview.md | 2 +- docs/sidebars.ts | 1 + lua/argon/agent.lua | 258 +++++ lua/argon/commands/init.lua | 18 + lua/argon/init.lua | 13 + 32 files changed, 4643 insertions(+), 636 deletions(-) create mode 100644 crates/analyzer/src/agent.rs create mode 100644 crates/analyzer/src/hook.rs create mode 100644 crates/analyzer/src/mcp.rs create mode 100644 crates/analyzer/src/session.rs create mode 100644 crates/analyzer/src/transport.rs create mode 100644 crates/analyzer/src/watcher.rs create mode 100644 docs/docs/tools/agents.md create mode 100644 lua/argon/agent.lua diff --git a/Cargo.lock b/Cargo.lock index 7cd86830..dc143b28 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -227,9 +227,11 @@ dependencies = [ "cfgrammar", "clap", "futures", + "getrandom 0.3.4", "homedir", "indexmap", "lsp-document", + "notify", "serde", "serde_json", "tarpc", @@ -250,6 +252,7 @@ dependencies = [ "argon-analyzer", "argonc", "futures", + "serde_json", "tarpc", "tempfile", "tokio", @@ -2907,6 +2910,15 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "fsevent-sys" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76ee7a02da4d231650c7cea31349b889be2f45ddb3ef3032d2ec8185f6313fd2" +dependencies = [ + "libc", +] + [[package]] name = "futf" version = "0.1.5" @@ -4156,6 +4168,26 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inotify" +version = "0.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cc00ea907cab49550b7da656f80ebb97be1b997d931fbcd28d39734e17ce592" +dependencies = [ + "bitflags 2.13.2", + "inotify-sys", + "libc", +] + +[[package]] +name = "inotify-sys" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c033f80b2c113cdf91ab7a33faa9cbc014726dcad99880c8609af2a370edf37d" +dependencies = [ + "libc", +] + [[package]] name = "inout" version = "0.1.4" @@ -4328,6 +4360,26 @@ version = "0.2.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4933f3f57a8e9d9da04db23fb153356ecaf00cbd14aee46279c33dc80925c37" +[[package]] +name = "kqueue" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d763e5b24120b4ddf50de6c92308156765aabfbbccebf401da7cff2d70a41ea" +dependencies = [ + "kqueue-sys", + "libc", +] + +[[package]] +name = "kqueue-sys" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07293a4e297ac234359b510362495713f75ea345d5307140414f20c69ffeb087" +dependencies = [ + "bitflags 2.13.2", + "libc", +] + [[package]] name = "kurbo" version = "0.11.3" @@ -4816,6 +4868,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4b18443e9c262bfe8fa82f51666e2642c53393f7e5c27b3e1aeab922cff5b9d8" dependencies = [ "libc", + "log", "wasi", "windows-sys 0.61.2", ] @@ -5096,6 +5149,33 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8" +[[package]] +name = "notify" +version = "8.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d3d07927151ff8575b7087f245456e549fea62edf0ec4e565a5ee50c8402bc3" +dependencies = [ + "bitflags 2.13.2", + "fsevent-sys", + "inotify", + "kqueue", + "libc", + "log", + "mio", + "notify-types", + "walkdir", + "windows-sys 0.60.2", +] + +[[package]] +name = "notify-types" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42b8cfee0e339a0337359f3c88165702ac6e600dc01c0cc9579a92d62b08477a" +dependencies = [ + "bitflags 2.13.2", +] + [[package]] name = "ntapi" version = "0.4.3" @@ -9561,6 +9641,15 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + [[package]] name = "windows-sys" version = "0.61.2" diff --git a/crates/analyzer/Cargo.toml b/crates/analyzer/Cargo.toml index baae999c..8615cdb7 100644 --- a/crates/analyzer/Cargo.toml +++ b/crates/analyzer/Cargo.toml @@ -34,6 +34,8 @@ homedir = "0.3" clap = { workspace = true } tempfile = { workspace = true } toml = { workspace = true } +getrandom = "0.3" +notify = "8" [dev-dependencies] bincode = "1" diff --git a/crates/analyzer/src/agent.rs b/crates/analyzer/src/agent.rs new file mode 100644 index 00000000..70112af5 --- /dev/null +++ b/crates/analyzer/src/agent.rs @@ -0,0 +1,1066 @@ +//! The service agents use to read and edit the workspace. +//! +//! Agent edits are applied to Neovim's buffers rather than to disk, so the +//! user sees each one as it lands, can undo it, and decides when to save. + +use std::{ + path::{Component, Path, PathBuf}, + sync::atomic::Ordering, + time::Duration, +}; + +use argonc::{ + ast::Span, + compile::{CompileOutput, SolvedValue}, +}; +use indexmap::IndexMap; +use serde::{Deserialize, Serialize}; +use tarpc::context; +use tokio::time::Instant; +use tower_lsp_server::{ + NotCancellable, OngoingProgress, Unbounded, + ls_types::{Diagnostic, ProgressToken, Range, TextEdit, Uri, request::Request}, +}; + +use crate::{ + ApprovalMode, Backend, SourceState, + compiler_worker::{CompileIdentity, CompileRequest}, + document::{Document, PositionEncoding}, +}; + +/// How long the bridge lets one call run, which bounds an approval prompt. +pub const CALL_DEADLINE: Duration = Duration::from_secs(600); + +/// How long Neovim may take to report a change the analyzer asked for. +const SETTLE_TIMEOUT: Duration = Duration::from_secs(10); + +/// How long an edit waits for the compilation that includes it. +const COMPILE_TIMEOUT: Duration = Duration::from_secs(120); + +/// Attempts at an edit whose buffer changed between planning and applying. +const MAX_EDIT_ATTEMPTS: usize = 3; + +pub const FOLLOW_REQUIRED: &str = "The GUI's cell can only be switched while the user has follow mode on (`:Argon follow` in Neovim). Use compile_cell to check a cell without changing the GUI."; + +/// One find-and-replace within a file. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Replacement { + pub old: String, + pub new: String, + /// Replace every occurrence instead of requiring exactly one. + pub replace_all: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EditRequest { + pub path: PathBuf, + /// Matched against the file as it was before any of them applied. + pub edits: Vec, + /// Shown to the user while the edit is applied. + pub label: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CreateRequest { + pub path: PathBuf, + pub contents: String, + pub label: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FileContents { + pub path: PathBuf, + pub contents: String, + /// Whether the text came from a Neovim buffer, which may be unsaved. + pub from_editor: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct AgentDiagnostic { + pub path: PathBuf, + /// One-based. + pub line: u32, + /// One-based, in the editor's position encoding. + pub column: u32, + pub message: String, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum CompileStatus { + Valid, + ExecErrors, + StaticErrors, + ParseErrors, +} + +/// The shape of a compiled cell, so an agent can tell whether an edit did +/// what it meant without seeing the layout. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct CompileSummary { + pub cell: Option, + pub status: CompileStatus, + /// Distinct cells in the compiled hierarchy. + pub cells: usize, + pub rects: usize, + pub polygons: usize, + pub paths: usize, + pub instances: usize, + /// Coordinates the solver left to initial conditions in the top cell. + pub unsolved_vars: usize, + /// `[x0, y0, x1, y1]` of the top cell's layout geometry. + pub bbox: Option<[f64; 4]>, +} + +impl CompileSummary { + pub(crate) fn new(cell: Option<&str>, output: &CompileOutput) -> Self { + let (status, data) = match output { + CompileOutput::Valid(data) => (CompileStatus::Valid, Some(data)), + CompileOutput::ExecErrors(errors) => { + (CompileStatus::ExecErrors, errors.output.as_ref()) + } + CompileOutput::StaticErrors(_) => (CompileStatus::StaticErrors, None), + CompileOutput::FatalParseErrors => (CompileStatus::ParseErrors, None), + }; + let mut summary = Self { + cell: cell.map(str::to_owned), + status, + cells: 0, + rects: 0, + polygons: 0, + paths: 0, + instances: 0, + unsolved_vars: 0, + bbox: None, + }; + let Some(data) = data else { + return summary; + }; + summary.cells = data.cells.len(); + if let Some(top) = data.cells.get(&data.top) { + for object in top.objects.values() { + match object { + SolvedValue::Rect(_) => summary.rects += 1, + SolvedValue::Polygon(_) => summary.polygons += 1, + SolvedValue::Path(_) => summary.paths += 1, + SolvedValue::Instance(_) => summary.instances += 1, + _ => {} + } + } + summary.unsolved_vars = top.unsolved_vars.len(); + } + summary.bbox = data + .layout_bbox(data.top) + .map(|bbox| [bbox.x0, bbox.y0, bbox.x1, bbox.y1]); + summary + } +} + +/// The latest committed compilation and the cell it was for. +#[derive(Debug, Clone)] +pub(crate) struct LatestCompile { + pub(crate) cell: Option, + pub(crate) output: CompileOutput, +} + +/// What the workspace looks like after an agent call. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Report { + /// The source revision the call produced or observed. + pub revision: u64, + /// The revision `summary` and `diagnostics` describe. Older than + /// `revision` when compilation did not finish in time. + pub compiled_revision: u64, + pub summary: Option, + pub diagnostics: Vec, + pub messages: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AgentStatus { + pub root: Option, + pub open_cell: Option, + pub gui_connected: bool, + pub follow_mode: bool, + pub approval: ApprovalMode, + /// Whether any Neovim buffer in the workspace has unsaved changes. + pub workspace_modified: bool, + /// Whether files changed on disk outside Neovim are picked up. + pub watching_files: bool, + pub open_files: Vec, + pub revision: u64, + pub compiled_revision: u64, +} + +#[tarpc::service] +pub trait Agent { + async fn status() -> AgentStatus; + /// The file as the user currently has it, unsaved edits included. + async fn read_file(path: PathBuf) -> Result; + async fn edit_file(request: EditRequest) -> Result; + /// Creates a new source file as an unsaved Neovim buffer. + async fn create_file(request: CreateRequest) -> Result; + async fn diagnostics() -> Report; + /// Compiles a cell without changing what the GUI shows. + async fn compile_cell(cell: String) -> Result; + /// Shows a cell in the GUI. Allowed only in follow mode. + async fn open_cell(cell: String) -> Result; +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct EnsureBufferParams { + uri: Uri, + create: bool, +} + +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct EnsureBufferResult { + ok: bool, + #[serde(default)] + message: Option, +} + +/// Asks Neovim to load a file into a buffer attached to the analyzer. +enum EnsureBuffer {} + +impl Request for EnsureBuffer { + type Params = EnsureBufferParams; + type Result = EnsureBufferResult; + const METHOD: &'static str = "custom/ensureBuffer"; +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct ApplyAgentEditParams { + uri: Uri, + /// The buffer version the edits were planned against. + version: i32, + edits: Vec, + /// Where the replacement text lands once applied. + highlights: Vec, + label: String, + confirm: bool, + /// The whole edited file, sent only when confirming, for the preview. + new_text: Option, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(tag = "status", rename_all = "lowercase")] +pub(crate) enum ApplyAgentEditResult { + Applied { version: i32 }, + Stale { version: i32 }, + Rejected, + Failed { message: String }, +} + +/// Asks Neovim to apply an agent edit if its buffer is still at the planned +/// version, confirming with the user first when `confirm` is set. +enum ApplyAgentEdit {} + +impl Request for ApplyAgentEdit { + type Params = ApplyAgentEditParams; + type Result = ApplyAgentEditResult; + const METHOD: &'static str = "custom/applyAgentEdit"; +} + +#[derive(Debug, Deserialize)] +pub(crate) struct FollowModeParams { + enabled: bool, +} + +/// A resolved edit, ready to send to Neovim. +#[derive(Debug, Clone)] +pub(crate) struct EditPlan { + version: i32, + edits: Vec, + highlights: Vec, + /// Byte spans of the replacement text in the edited file. + spans: Vec, + new_text: String, +} + +/// Every start offset of `pattern` in `text`, overlapping ones included. +fn occurrences(text: &str, pattern: &str) -> Vec { + let mut starts = Vec::new(); + let mut from = 0; + while let Some(found) = text[from..].find(pattern) { + let start = from + found; + starts.push(start); + from = start + text[start..].chars().next().map_or(1, char::len_utf8); + } + starts +} + +/// Resolves find-and-replace edits against `document`. +pub(crate) fn plan_edit( + document: &Document, + replacements: &[Replacement], + encoding: PositionEncoding, +) -> Result { + if replacements.is_empty() { + return Err("No edits were given.".to_owned()); + } + let text = document.contents(); + let mut matches = Vec::new(); + for (index, replacement) in replacements.iter().enumerate() { + let number = index + 1; + if replacement.old.is_empty() { + return Err(format!( + "Edit {number}: old_string is empty. Use create_file to write a new file." + )); + } + if replacement.old == replacement.new { + return Err(format!( + "Edit {number}: old_string and new_string are identical." + )); + } + let starts = occurrences(text, &replacement.old); + match starts.len() { + 0 => { + return Err(format!( + "Edit {number}: old_string was not found. Read the file again; it may have changed." + )); + } + count if count > 1 && !replacement.replace_all => { + return Err(format!( + "Edit {number}: old_string occurs {count} times. Include more surrounding text to make it unique, or set replace_all." + )); + } + _ => {} + } + for start in starts { + matches.push(( + start, + start + replacement.old.len(), + number, + &replacement.new, + )); + } + } + matches.sort_by_key(|&(start, ..)| start); + for pair in matches.windows(2) { + let (_, end, first, _) = pair[0]; + let (start, _, second, _) = pair[1]; + if end > start { + return Err(if first == second { + format!("Edit {first}: old_string matches overlapping text.") + } else { + format!("Edits {first} and {second} overlap.") + }); + } + } + + let mut new_text = String::with_capacity(text.len()); + let mut spans = Vec::with_capacity(matches.len()); + let mut copied = 0; + for &(start, end, _, new) in &matches { + new_text.push_str(&text[copied..start]); + let span_start = new_text.len(); + new_text.push_str(new); + spans.push(cfgrammar::Span::new(span_start, new_text.len())); + copied = end; + } + new_text.push_str(&text[copied..]); + + let edited = Document::new(new_text.clone(), 0, encoding); + Ok(EditPlan { + version: document.version(), + edits: matches + .iter() + .map(|&(start, end, _, new)| TextEdit { + range: Range::new(document.offset_to_pos(start), document.offset_to_pos(end)), + new_text: new.clone(), + }) + .collect(), + highlights: spans + .iter() + .map(|span| { + Range::new( + edited.offset_to_pos(span.start()), + edited.offset_to_pos(span.end()), + ) + }) + .collect(), + spans, + new_text, + }) +} + +/// Fills a blank document, such as a new buffer, with `contents`. +pub(crate) fn plan_insert( + document: &Document, + contents: &str, + encoding: PositionEncoding, +) -> Result { + let existing = document.contents(); + if !existing.trim().is_empty() { + return Err("The file already has contents; use edit_file instead.".to_owned()); + } + let edited = Document::new(contents.to_owned(), 0, encoding); + Ok(EditPlan { + version: document.version(), + edits: vec![TextEdit { + range: Range::new( + document.offset_to_pos(0), + document.offset_to_pos(existing.len()), + ), + new_text: contents.to_owned(), + }], + highlights: vec![Range::new( + edited.offset_to_pos(0), + edited.offset_to_pos(contents.len()), + )], + spans: vec![cfgrammar::Span::new(0, contents.len())], + new_text: contents.to_owned(), + }) +} + +/// Removes `.` and `..` components without touching the file system. +fn normalize(path: &Path) -> PathBuf { + let mut normalized = PathBuf::new(); + for component in path.components() { + match component { + Component::CurDir => {} + Component::ParentDir => { + normalized.pop(); + } + other => normalized.push(other), + } + } + normalized +} + +/// A stable identity for a file that may not exist yet. +fn identity(path: &Path) -> PathBuf { + if let Ok(path) = path.canonicalize() { + return path; + } + match (path.parent().map(Path::canonicalize), path.file_name()) { + (Some(Ok(parent)), Some(name)) => parent.join(name), + _ => path.to_path_buf(), + } +} + +/// `path` relative to `root` when it lies inside it. +fn display_path(root: Option<&Path>, path: &Path) -> PathBuf { + let Some(root) = root else { + return path.to_path_buf(); + }; + if let Ok(relative) = path.strip_prefix(root) { + return relative.to_path_buf(); + } + identity(path) + .strip_prefix(identity(root)) + .map(Path::to_path_buf) + .unwrap_or_else(|_| path.to_path_buf()) +} + +fn agent_diagnostics( + root: Option<&Path>, + diagnostics: &IndexMap>, +) -> Vec { + diagnostics + .iter() + .flat_map(|(uri, diagnostics)| { + let path = uri + .to_file_path() + .map(|path| display_path(root, &path)) + .unwrap_or_else(|| PathBuf::from(uri.as_str())); + diagnostics.iter().map(move |diagnostic| AgentDiagnostic { + path: path.clone(), + line: diagnostic.range.start.line + 1, + column: diagnostic.range.start.character + 1, + message: diagnostic.message.clone(), + }) + }) + .collect() +} + +/// The editor's URI for `path`, if Neovim has it open under any name. +pub(crate) fn open_uri(source: &SourceState, path: &Path) -> Option { + let target = identity(path); + source + .editor_files + .keys() + .find(|uri| { + uri.to_file_path() + .is_some_and(|open| identity(&open) == target) + }) + .cloned() +} + +struct AgentActivity { + progress: OngoingProgress, +} + +impl Backend { + fn root(&self) -> Result { + self.state + .root_dir + .get() + .cloned() + .ok_or_else(|| "The analyzer has not opened a workspace yet.".to_owned()) + } + + /// An agent-supplied path, relative to the workspace root unless absolute. + fn resolve(&self, path: &Path) -> Result { + let root = self.root()?; + Ok(normalize(&if path.is_absolute() { + path.to_path_buf() + } else { + root.join(path) + })) + } + + /// Whether `path` is in the workspace or one of the modules it compiles. + async fn in_workspace(&self, path: &Path) -> bool { + let Ok(root) = self.root() else { + return false; + }; + let target = identity(path); + if target.starts_with(identity(&root)) { + return true; + } + let ast = self.state.published_state.lock().await.ast.clone(); + ast.values().any(|module| identity(&module.path) == target) + } + + async fn check_editable(&self, path: &Path) -> Result<(), String> { + if path.extension().is_none_or(|extension| extension != "ar") { + return Err(format!( + "{} is not an Argon source file. Only .ar files are edited through Neovim.", + path.display() + )); + } + if !self.in_workspace(path).await { + return Err(format!( + "{} is outside this Argon workspace.", + path.display() + )); + } + Ok(()) + } + + /// Waits until `found` returns a value for the source state, which is + /// rechecked each time an editor document changes. + async fn wait_for_source(&self, found: impl Fn(&SourceState) -> Option) -> Option { + let mut changes = self.state.source_changes.subscribe(); + let deadline = Instant::now() + SETTLE_TIMEOUT; + loop { + if let Some(value) = found(&*self.state.source_state.lock().await) { + return Some(value); + } + match tokio::time::timeout_at(deadline, changes.changed()).await { + Ok(Ok(())) => {} + _ => return None, + } + } + } + + async fn wait_for_version(&self, uri: &Uri, version: i32) -> bool { + self.wait_for_source(|source| { + source + .editor_files + .get(uri) + .filter(|document| document.version() >= version) + .map(|_| ()) + }) + .await + .is_some() + } + + async fn wait_for_compiled(&self, revision: u64) -> bool { + let mut compiled = self.state.compiled_revisions.subscribe(); + tokio::time::timeout(COMPILE_TIMEOUT, compiled.wait_for(|&done| done >= revision)) + .await + .is_ok_and(|result| result.is_ok()) + } + + /// Makes sure Neovim has `path` open in a buffer attached to this analyzer. + async fn ensure_buffer(&self, path: &Path, create: bool) -> Result { + if let Some(uri) = open_uri(&*self.state.source_state.lock().await, path) { + return Ok(uri); + } + if create { + if path.exists() { + return Err(format!( + "{} already exists; use edit_file instead.", + path.display() + )); + } + if !path.parent().is_some_and(Path::is_dir) { + return Err(format!( + "The directory for {} does not exist.", + path.display() + )); + } + } else if !path.is_file() { + return Err(format!("{} does not exist.", path.display())); + } + let uri = Uri::from_file_path(path) + .ok_or_else(|| format!("{} is not a valid file path.", path.display()))?; + let response = self + .state + .editor_client + .send_request::(EnsureBufferParams { uri, create }) + .await + .map_err(|error| format!("Could not ask Neovim to open {}: {error}", path.display()))?; + if !response.ok { + return Err(response + .message + .unwrap_or_else(|| format!("Neovim could not open {}.", path.display()))); + } + self.wait_for_source(|source| open_uri(source, path)) + .await + .ok_or_else(|| format!("Neovim opened {} but did not attach it.", path.display())) + } + + /// Applies the plan `plan_for` makes against the current buffer text, + /// replanning if the buffer moved on first. Returns the source revision + /// that includes the edit. + async fn apply_agent_edit( + &self, + uri: &Uri, + label: &str, + plan_for: impl Fn(&Document) -> Result, + ) -> Result<(EditPlan, u64), String> { + for _ in 0..MAX_EDIT_ATTEMPTS { + let plan = { + let source = self.state.source_state.lock().await; + let document = source + .editor_files + .get(uri) + .ok_or_else(|| "Neovim closed the file before the edit.".to_owned())?; + plan_for(document)? + }; + let confirm = self.state.config().agent.approval == ApprovalMode::Always; + self.state.begin_pending_edit(uri).await; + let result = self + .state + .editor_client + .send_request::(ApplyAgentEditParams { + uri: uri.clone(), + version: plan.version, + edits: plan.edits.clone(), + highlights: plan.highlights.clone(), + label: label.to_owned(), + confirm, + new_text: confirm.then(|| plan.new_text.clone()), + }) + .await; + match result { + Ok(ApplyAgentEditResult::Applied { version }) => { + if !self.wait_for_version(uri, version).await { + self.state.end_pending_edit(uri).await; + return Err( + "Neovim applied the edit but never reported the change.".to_owned() + ); + } + let revision = self.state.source_state.lock().await.revision; + return Ok((plan, revision)); + } + Ok(ApplyAgentEditResult::Stale { version }) => { + self.state.end_pending_edit(uri).await; + self.wait_for_version(uri, version).await; + } + Ok(ApplyAgentEditResult::Rejected) => { + self.state.end_pending_edit(uri).await; + return Err("The user rejected this edit in Neovim.".to_owned()); + } + Ok(ApplyAgentEditResult::Failed { message }) => { + self.state.end_pending_edit(uri).await; + return Err(format!("Neovim could not apply the edit: {message}")); + } + Err(error) => { + self.state.end_pending_edit(uri).await; + return Err(format!("Could not reach Neovim: {error}")); + } + } + } + Err( + "The file kept changing while the edit was applied. Read it again and retry." + .to_owned(), + ) + } + + async fn report(&self, revision: u64) -> Report { + let root = self.state.root_dir.get().cloned(); + let published = self.state.published_state.lock().await; + Report { + revision, + compiled_revision: published.compiled_revision, + summary: published + .latest + .as_ref() + .map(|latest| CompileSummary::new(latest.cell.as_deref(), &latest.output)), + diagnostics: agent_diagnostics(root.as_deref(), &published.prev_diagnostics), + messages: published.messages.clone(), + } + } + + /// Waits for the edit's compilation, outlines what it changed in the + /// GUI, and reports the result. + async fn settle(&self, path: &Path, plan: &EditPlan, revision: u64) -> Report { + if self.wait_for_compiled(revision).await { + self.highlight_in_gui(revision, path, &plan.spans).await; + } + self.report(revision).await + } + + async fn highlight_in_gui(&self, revision: u64, path: &Path, spans: &[cfgrammar::Span]) { + let Some(connection) = self.state.gui_connection().await else { + return; + }; + // Compiled spans name files the way the compiler found them. + let target = identity(path); + let path = self + .state + .published_state + .lock() + .await + .ast + .values() + .find(|module| identity(&module.path) == target) + .map_or_else(|| path.to_path_buf(), |module| module.path.clone()); + let spans = spans + .iter() + .map(|&span| Span { + path: path.clone(), + span, + }) + .collect(); + let result = connection + .client + .highlight_spans(context::current(), revision, spans) + .await; + self.handle_gui_result(&connection, result).await; + } + + async fn begin_agent_activity(&self, label: &str) -> AgentActivity { + let id = self + .state + .next_compilation_activity_id + .fetch_add(1, Ordering::Relaxed) + .saturating_add(1); + let token = ProgressToken::String(format!("argon-agent-{id}")); + let _ = self + .state + .editor_client + .create_work_done_progress(token.clone()) + .await; + let progress = self + .state + .editor_client + .progress(token, "Argon agent") + .with_message(label) + .begin() + .await; + self.publish_agent_activity(Some(label.to_owned())).await; + AgentActivity { progress } + } + + async fn finish_agent_activity(&self, activity: AgentActivity) { + activity.progress.finish().await; + self.publish_agent_activity(None).await; + } + + async fn publish_agent_activity(&self, label: Option) { + if let Some(connection) = self.state.gui_connection().await { + let result = connection + .client + .agent_activity(context::current(), label) + .await; + self.handle_gui_result(&connection, result).await; + } + } + + pub(crate) async fn follow_mode(&self, params: FollowModeParams) -> crate::Result<()> { + self.state + .follow_mode + .store(params.enabled, Ordering::Release); + Ok(()) + } + + async fn edit_with( + &self, + path: PathBuf, + label: String, + create: bool, + plan_for: impl Fn(&Document) -> Result, + ) -> Result { + let path = self.resolve(&path)?; + self.check_editable(&path).await?; + // One edit at a time, so each plans against text the last one settled. + let _serial = self.state.agent_edits.lock().await; + let activity = self.begin_agent_activity(&label).await; + let result = async { + let uri = self.ensure_buffer(&path, create).await?; + let (plan, revision) = self.apply_agent_edit(&uri, &label, &plan_for).await?; + Ok(self.settle(&path, &plan, revision).await) + } + .await; + self.finish_agent_activity(activity).await; + result + } +} + +impl Agent for Backend { + async fn status(self, _: context::Context) -> AgentStatus { + let root = self.state.root_dir.get().cloned(); + let gui_connected = self.state.gui_connection().await.is_some(); + let watching_files = self + .state + .watcher + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .is_some(); + let source = self.state.source_state.lock().await; + let compiled_revision = self.state.published_state.lock().await.compiled_revision; + AgentStatus { + open_cell: source.cell.clone(), + gui_connected, + follow_mode: self.state.follow_mode.load(Ordering::Acquire), + approval: self.state.config().agent.approval, + workspace_modified: source.workspace_modified, + watching_files, + open_files: source + .editor_files + .keys() + .filter_map(|uri| uri.to_file_path()) + .map(|path| display_path(root.as_deref(), &path)) + .collect(), + revision: source.revision, + compiled_revision, + root, + } + } + + async fn read_file(self, _: context::Context, path: PathBuf) -> Result { + let path = self.resolve(&path)?; + if !self.in_workspace(&path).await { + return Err(format!( + "{} is outside this Argon workspace.", + path.display() + )); + } + let root = self.state.root_dir.get().cloned(); + let display = display_path(root.as_deref(), &path); + { + let source = self.state.source_state.lock().await; + if let Some(uri) = open_uri(&source, &path) { + return Ok(FileContents { + path: display, + contents: source.editor_files[&uri].contents().to_owned(), + from_editor: true, + }); + } + } + let contents = tokio::fs::read_to_string(&path) + .await + .map_err(|error| format!("Could not read {}: {error}", path.display()))?; + Ok(FileContents { + path: display, + contents, + from_editor: false, + }) + } + + async fn edit_file(self, _: context::Context, request: EditRequest) -> Result { + let encoding = self.state.position_encoding(); + let edits = request.edits; + self.edit_with(request.path, request.label, false, |document| { + plan_edit(document, &edits, encoding) + }) + .await + } + + async fn create_file( + self, + _: context::Context, + request: CreateRequest, + ) -> Result { + let encoding = self.state.position_encoding(); + let contents = request.contents; + self.edit_with(request.path, request.label, true, |document| { + plan_insert(document, &contents, encoding) + }) + .await + } + + async fn diagnostics(self, _: context::Context) -> Report { + let revision = self.state.source_state.lock().await.revision; + self.report(revision).await + } + + async fn compile_cell(self, _: context::Context, cell: String) -> Result { + let root_dir = self.root()?; + let revision = self.state.source_state.lock().await.revision; + let result = self + .state + .compiler + .compile(CompileRequest { + identity: CompileIdentity { + revision, + cell: Some(cell.clone()), + }, + root_dir: root_dir.clone(), + }) + .await + .ok_or_else(|| "The compiler is unavailable.".to_owned())?; + let diagnostics = crate::diagnostics( + &result.root_dir, + &result.ast, + result.output.as_ref(), + self.state.position_encoding(), + ); + Ok(Report { + revision, + compiled_revision: revision, + summary: result + .output + .as_ref() + .map(|output| CompileSummary::new(Some(&cell), output)), + diagnostics: agent_diagnostics(Some(&root_dir), &diagnostics), + messages: result.messages, + }) + } + + async fn open_cell(self, _: context::Context, cell: String) -> Result { + if !self.state.follow_mode.load(Ordering::Acquire) { + return Err(FOLLOW_REQUIRED.to_owned()); + } + let identity = { + let mut source = self.state.source_state.lock().await; + source.cell = Some(cell); + source.compile_identity() + }; + if let Some(connection) = self.update_cell(identity.clone()).await { + let result = connection.client.fit(context::current()).await; + self.handle_gui_result(&connection, result).await; + } + Ok(self.report(identity.revision).await) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn document(text: &str) -> Document { + Document::new(text.to_owned(), 7, PositionEncoding::Utf8) + } + + fn replacement(old: &str, new: &str) -> Replacement { + Replacement { + old: old.to_owned(), + new: new.to_owned(), + replace_all: false, + } + } + + #[test] + fn plans_edits_against_the_original_text() { + let text = "cell top() {\n let a = 1.;\n let b = 2.;\n}\n"; + let plan = plan_edit( + &document(text), + &[ + replacement("let b = 2.;", "let b = 20.;"), + replacement("1.", "10."), + ], + PositionEncoding::Utf8, + ) + .unwrap(); + assert_eq!(plan.version, 7); + assert_eq!( + plan.new_text, + "cell top() {\n let a = 10.;\n let b = 20.;\n}\n" + ); + assert_eq!(plan.edits.len(), 2); + assert_eq!( + plan.edits[0].range.start, + tower_lsp_server::ls_types::Position::new(1, 12) + ); + assert_eq!( + &plan.new_text[plan.spans[0].start()..plan.spans[0].end()], + "10." + ); + assert_eq!( + &plan.new_text[plan.spans[1].start()..plan.spans[1].end()], + "let b = 20.;" + ); + assert_eq!( + plan.highlights[1].start, + tower_lsp_server::ls_types::Position::new(2, 4) + ); + } + + #[test] + fn rejects_missing_ambiguous_and_overlapping_edits() { + let doc = document("aaa bb bb"); + let error = + |edits: &[Replacement]| plan_edit(&doc, edits, PositionEncoding::Utf8).unwrap_err(); + assert!(error(&[replacement("zz", "y")]).contains("not found")); + assert!(error(&[replacement("bb", "c")]).contains("occurs 2 times")); + assert!(error(&[replacement("aa", "c")]).contains("occurs 2 times")); + assert!(error(&[replacement("aaa", "x"), replacement("a b", "y")]).contains("overlap")); + assert!(error(&[replacement("", "x")]).contains("create_file")); + assert!(error(&[]).contains("No edits")); + let all = plan_edit( + &doc, + &[Replacement { + replace_all: true, + ..replacement("bb", "c") + }], + PositionEncoding::Utf8, + ) + .unwrap(); + assert_eq!(all.new_text, "aaa c c"); + assert!( + plan_edit( + &doc, + &[Replacement { + replace_all: true, + ..replacement("aa", "c") + }], + PositionEncoding::Utf8, + ) + .unwrap_err() + .contains("overlapping") + ); + } + + #[test] + fn insertion_requires_a_blank_document() { + let plan = plan_insert(&document("\n"), "cell a() {}\n", PositionEncoding::Utf8).unwrap(); + assert_eq!( + plan.edits[0].range.start, + tower_lsp_server::ls_types::Position::new(0, 0) + ); + assert_eq!(plan.new_text, "cell a() {}\n"); + assert_eq!(plan.spans, vec![cfgrammar::Span::new(0, 12)]); + assert_eq!( + plan.highlights[0].start, + tower_lsp_server::ls_types::Position::new(0, 0) + ); + assert!(plan_insert(&document("x"), "y", PositionEncoding::Utf8).is_err()); + } + + #[test] + fn paths_are_normalized_and_shown_relative_to_the_root() { + assert_eq!( + normalize(Path::new("/a/./b/../c.ar")), + PathBuf::from("/a/c.ar") + ); + assert_eq!( + display_path( + Some(Path::new("/work/lib")), + Path::new("/work/lib/src/a.ar") + ), + PathBuf::from("src/a.ar") + ); + assert_eq!( + display_path(Some(Path::new("/work/lib")), Path::new("/elsewhere/a.ar")), + PathBuf::from("/elsewhere/a.ar") + ); + } +} diff --git a/crates/analyzer/src/cli.rs b/crates/analyzer/src/cli.rs index ab36bee0..56fe34d3 100644 --- a/crates/analyzer/src/cli.rs +++ b/crates/analyzer/src/cli.rs @@ -22,21 +22,56 @@ struct Args { enum Command { /// Coordinate startup with a local Argone SSH session. Relay, + /// Serve the Model Context Protocol, forwarding to the running analyzer + /// whose workspace contains the working directory. + Mcp { + /// Directory to find a session for, instead of the working directory. + #[arg(long)] + root: Option, + }, + /// Run as a Claude Code `PreToolUse` hook that redirects `.ar` file + /// tools to the analyzer while a session is running. + Hook, } #[tokio::main] pub async fn run() { let args = Args::parse(); - if matches!(args.command, Some(Command::Relay)) { - if let Err(error) = run_relay().await { - eprintln!("argon-analyzer relay: {error}"); - std::process::exit(1); + match args.command { + Some(Command::Relay) => { + if let Err(error) = run_relay().await { + eprintln!("argon-analyzer relay: {error}"); + std::process::exit(1); + } } - } else { - crate::main(args.rpc_port, args.relay_socket).await; + Some(Command::Mcp { root }) => { + crate::init_logging(); + if let Err(error) = crate::mcp::run(root).await { + eprintln!("argon-analyzer mcp: {error}"); + std::process::exit(1); + } + } + Some(Command::Hook) => { + if let Err(error) = crate::hook::run().await { + eprintln!("argon-analyzer hook: {error}"); + } + } + None => crate::main(args.rpc_port, args.relay_socket).await, } } +/// Parses the `PORT TOKEN` line the analyzer writes to the relay socket. +fn parse_port_and_token(line: &str) -> Option<(u16, String)> { + let mut fields = line.split_whitespace(); + let port = fields + .next()? + .parse::() + .ok() + .filter(|port| *port != 0)?; + let token = crate::transport::SessionToken::from_hex(fields.next()?)?; + fields.next().is_none().then(|| (port, token.to_hex())) +} + #[cfg(unix)] async fn run_relay() -> io::Result<()> { let directory = tempfile::tempdir()?; @@ -49,16 +84,17 @@ async fn run_relay() -> io::Result<()> { stdout.flush().await?; let (stream, _) = listener.accept().await?; - let mut port = String::new(); - BufReader::new(stream).read_line(&mut port).await?; - let port = port - .trim() - .parse::() - .ok() - .filter(|port| *port != 0) - .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "invalid analyzer RPC port"))?; + let mut line = String::new(); + BufReader::new(stream).read_line(&mut line).await?; + let (port, token) = parse_port_and_token(&line).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "invalid analyzer RPC announcement", + ) + })?; + // Stdout is the SSH channel back to Argone, so the token stays private. stdout - .write_all(format!("ARGON_ANALYZER 1 {port}\n").as_bytes()) + .write_all(format!("ARGON_ANALYZER 2 {port} {token}\n").as_bytes()) .await?; stdout.flush().await } @@ -70,3 +106,21 @@ async fn run_relay() -> io::Result<()> { "relay requires a Unix-like remote host", )) } + +#[cfg(test)] +mod tests { + use super::parse_port_and_token; + + #[test] + fn relay_announcements_carry_a_port_and_a_token() { + let token = "ab".repeat(32); + assert_eq!( + parse_port_and_token(&format!("1234 {token}\n")), + Some((1234, token.clone())) + ); + assert_eq!(parse_port_and_token("1234\n"), None); + assert_eq!(parse_port_and_token(&format!("0 {token}")), None); + assert_eq!(parse_port_and_token("1234 not-a-token"), None); + assert_eq!(parse_port_and_token(&format!("1234 {token} extra")), None); + } +} diff --git a/crates/analyzer/src/document.rs b/crates/analyzer/src/document.rs index beaa2943..44d3baa2 100644 --- a/crates/analyzer/src/document.rs +++ b/crates/analyzer/src/document.rs @@ -186,6 +186,10 @@ impl Document { } } + pub(crate) fn version(&self) -> i32 { + self.version + } + pub(crate) fn contents(&self) -> &str { self.contents.text() } diff --git a/crates/analyzer/src/hook.rs b/crates/analyzer/src/hook.rs new file mode 100644 index 00000000..6253430a --- /dev/null +++ b/crates/analyzer/src/hook.rs @@ -0,0 +1,125 @@ +//! A Claude Code `PreToolUse` hook that steers file tools on `.ar` files to +//! the analyzer while an Argon session is running. +//! +//! Edits are always redirected so they land in Neovim. Reads are redirected +//! only when Neovim holds unsaved changes the file on disk lacks. + +use std::{ + io::{self, Read}, + path::{Path, PathBuf}, + time::Duration, +}; + +use serde_json::{Value, json}; + +use crate::{agent::AgentClient, mcp}; + +/// How long the hook waits for the analyzer before letting the tool run. +const HOOK_TIMEOUT: Duration = Duration::from_secs(3); + +/// The decision for one tool call, or `None` to let it proceed. +async fn decide( + input: &Value, + connect: impl AsyncFnOnce(&Path) -> Option, +) -> Option { + let tool = input.get("tool_name")?.as_str()?; + let file = input.get("tool_input")?.get("file_path")?.as_str()?; + let mut path = PathBuf::from(file); + if path.extension().is_none_or(|extension| extension != "ar") { + return None; + } + if path.is_relative() + && let Some(cwd) = input.get("cwd").and_then(Value::as_str) + { + path = Path::new(cwd).join(path); + } + let directory = path.parent()?.to_path_buf(); + let client = connect(&directory).await?; + if tool == "Read" { + let file = client + .read_file(tarpc::context::current(), path.clone()) + .await + .ok()? + .ok()?; + let disk = std::fs::read_to_string(&path).ok(); + if !file.from_editor || disk.as_deref() == Some(file.contents.as_str()) { + return None; + } + return Some(format!( + "{} has unsaved changes in the user's Neovim buffer, so the file on disk is out of date. Read it with the argon MCP server's read_file tool (mcp__argon__read_file) instead.", + path.display() + )); + } + Some(format!( + "An Argon session is running for {}. Change .ar files with the argon MCP server's edit_file or create_file tools (mcp__argon__edit_file, mcp__argon__create_file) so the change goes through the user's Neovim buffer and shows up live in Neovim and the GUI. Don't write .ar files directly.", + directory.display() + )) +} + +/// Why a `PreToolUse` call should be redirected, or `None` to let it run. +pub async fn decision(input: &Value) -> Option { + tokio::time::timeout( + HOOK_TIMEOUT, + decide(input, async |directory: &Path| { + mcp::connect_session(directory).await.ok() + }), + ) + .await + .ok() + .flatten() +} + +fn deny(reason: &str) -> Value { + json!({ + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": reason, + } + }) +} + +/// Reads the hook input from stdin and prints a decision when the call +/// should be redirected. Any failure lets the tool call proceed. +pub async fn run() -> io::Result<()> { + let mut input = String::new(); + io::stdin().read_to_string(&mut input)?; + let Ok(input) = serde_json::from_str::(&input) else { + return Ok(()); + }; + if let Some(reason) = decision(&input).await { + println!("{}", deny(&reason)); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + async fn no_session(_: &Path) -> Option { + None + } + + #[tokio::test] + async fn ignores_other_files_and_workspaces_without_a_session() { + for input in [ + json!({ "tool_name": "Edit", "tool_input": { "file_path": "/w/notes.md" } }), + json!({ "tool_name": "Bash", "tool_input": { "command": "ls" } }), + json!({ "tool_name": "Edit", "tool_input": { "file_path": "/w/lib.ar" } }), + ] { + assert_eq!(decide(&input, no_session).await, None); + } + } + + #[test] + fn denials_use_the_pre_tool_use_decision_format() { + let output = deny("use the MCP tool"); + assert_eq!(output["hookSpecificOutput"]["hookEventName"], "PreToolUse"); + assert_eq!(output["hookSpecificOutput"]["permissionDecision"], "deny"); + assert_eq!( + output["hookSpecificOutput"]["permissionDecisionReason"], + "use the MCP tool" + ); + } +} diff --git a/crates/analyzer/src/lib.rs b/crates/analyzer/src/lib.rs index 5bd54318..85c27888 100644 --- a/crates/analyzer/src/lib.rs +++ b/crates/analyzer/src/lib.rs @@ -1,9 +1,15 @@ +pub mod agent; mod cell_edit; mod command_completion; mod compiler_worker; pub mod document; +pub mod hook; +pub mod mcp; mod navigation; pub mod rpc; +pub mod session; +pub mod transport; +mod watcher; pub mod cli; @@ -92,10 +98,28 @@ static LOG_RELOAD: OnceLock> = OnceLock::new #[serde(default, deny_unknown_fields)] pub struct ArgonConfig { pub analyzer: AnalyzerConfig, + pub agent: AgentConfig, pub gui: GuiConfig, pub log: LogConfig, } +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct AgentConfig { + pub approval: ApprovalMode, +} + +/// Whether the user confirms each agent edit in Neovim before it is applied. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ApprovalMode { + /// Agent edits apply immediately and Neovim is never asked to confirm. + #[default] + Never, + /// Neovim shows each agent edit and applies it only if the user accepts. + Always, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(default, deny_unknown_fields)] pub struct AnalyzerConfig { @@ -388,6 +412,10 @@ pub(crate) struct PublishedState { pub(crate) nav: Option>, pub(crate) prev_diagnostics: IndexMap>, pub(crate) compiled_revision: u64, + /// Compiler messages that accompanied the latest committed result. + pub(crate) messages: Vec, + /// The latest committed result, summarized on demand for agents. + pub(crate) latest: Option, } #[derive(Clone, Debug)] @@ -629,10 +657,22 @@ pub struct State { gui: Arc>, next_compilation_activity_id: Arc, app_config: Arc>, + /// Secret that every RPC connection must present. + token: transport::SessionToken, + session_file: Arc>>, + /// Bumped whenever an editor document opens, changes, or closes. + source_changes: Arc>, + /// The source revision of the latest committed compilation. + compiled_revisions: Arc>, + /// Whether Neovim has a follow window for agent edits. + follow_mode: Arc, + /// Serializes agent edits so each resolves against settled text. + agent_edits: Arc>, + watcher: Arc>>, } impl State { - fn new(server_addr: SocketAddr, editor_client: Client) -> Self { + fn new(server_addr: SocketAddr, editor_client: Client, token: transport::SessionToken) -> Self { let config = read_config().unwrap_or_else(|error| { error!("{error}; using default configuration"); ArgonConfig::default() @@ -649,9 +689,21 @@ impl State { gui: Default::default(), next_compilation_activity_id: Default::default(), app_config: Arc::new(RwLock::new(config)), + token, + session_file: Default::default(), + source_changes: Arc::new(tokio::sync::watch::Sender::new(0)), + compiled_revisions: Arc::new(tokio::sync::watch::Sender::new(0)), + follow_mode: Default::default(), + agent_edits: Default::default(), + watcher: Default::default(), } } + pub(crate) fn notify_source_changed(&self) { + self.source_changes + .send_modify(|count| *count = count.wrapping_add(1)); + } + fn apply_config(&self, config: ArgonConfig) { *write_unpoisoned(&self.app_config) = config; } @@ -747,6 +799,34 @@ impl State { connection } + /// Records this analyzer in the user's session directory so local + /// clients, such as the agent bridge, can find and authenticate to it. + fn publish_session(&self, root_dir: &Path) { + let info = session::SessionInfo { + pid: std::process::id(), + root: root_dir.to_path_buf(), + port: self.server_addr.port(), + token: self.token.to_hex(), + }; + match session::publish(&info) { + Ok(file) => { + *self.session_file.lock().unwrap_or_else(|p| p.into_inner()) = Some(file); + } + Err(error) => error!("could not publish the analyzer session: {error}"), + } + } + + fn retire_session(&self) { + self.session_file + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + self.watcher + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .take(); + } + async fn take_gui_process(&self) -> Option { self.gui.lock().await.process.take() } @@ -908,6 +988,11 @@ impl Backend { result.output.as_ref(), self.state.position_encoding(), ); + // Cells are shared, so this copy costs one pointer per cell. + let latest = result.output.clone().map(|output| agent::LatestCompile { + cell: identity.cell.clone(), + output, + }); // The compiled GDS can contain millions of geometry values. Move it // into the RPC snapshot after diagnostics have borrowed it instead of // deep-cloning the complete layout immediately before serialization. @@ -938,7 +1023,12 @@ impl Backend { compiled.nav = Some(nav); } compiled.compiled_revision = identity.revision; + compiled.messages = result.messages.clone(); + compiled.latest = latest; } + self.state + .compiled_revisions + .send_replace(identity.revision); for message in result.messages { if !self.state.is_latest_compile_request(&identity).await { @@ -1128,7 +1218,12 @@ impl LanguageServer for Backend { let root_uri = root_uri.or(params.root_uri); let root_dir = root_uri.and_then(|root| root.to_file_path().map(|path| path.into_owned())); if let Some(root_dir) = root_dir { - let _ = self.state.root_dir.set(root_dir); + let _ = self.state.root_dir.set(root_dir.clone()); + self.state.publish_session(&root_dir); + // A recursive watch walks the whole tree on some platforms. + let backend = self.clone(); + let watched = root_dir.clone(); + tokio::task::spawn_blocking(move || backend.start_watcher(&watched)); let connection = self.state.gui_connection().await; self.state.publish_workspace_path(connection).await; } @@ -1189,6 +1284,7 @@ impl LanguageServer for Backend { source.advance_revision(); let identity = source.compile_identity(); drop(source); + self.state.notify_source_changed(); self.compile_after_debounce(identity); } @@ -1242,6 +1338,7 @@ impl LanguageServer for Backend { } let identity = source.compile_identity(); drop(source); + self.state.notify_source_changed(); if analyzer_edit { self.update_cell(identity).await; } else { @@ -1343,10 +1440,12 @@ impl LanguageServer for Backend { source.advance_revision(); let identity = source.compile_identity(); drop(source); + self.state.notify_source_changed(); self.compile_after_debounce(identity); } async fn shutdown(&self) -> Result<()> { + self.state.retire_session(); if let Some(mut gui) = self.state.take_gui_process().await { let _ = gui.kill().await; } @@ -1442,6 +1541,7 @@ impl Backend { match Command::new("argone") .arg("gui") .arg(format!("{}", state.server_addr)) + .env(transport::TOKEN_ENV, state.token.to_hex()) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) @@ -1924,24 +2024,79 @@ async fn spawn(fut: impl Future + Send + 'static) { tokio::spawn(fut); } +/// Most authenticated RPC connections served at once. +const MAX_RPC_CONNECTIONS: usize = 10; + #[cfg(unix)] -fn announce_rpc_port(path: &Path, port: u16) -> io::Result<()> { +fn announce_rpc_port(path: &Path, port: u16, token: &transport::SessionToken) -> io::Result<()> { use std::io::Write; use std::os::unix::net::UnixStream; let mut stream = UnixStream::connect(path)?; - writeln!(stream, "{port}")?; + writeln!(stream, "{port} {}", token.to_hex())?; stream.flush() } #[cfg(not(unix))] -fn announce_rpc_port(_: &Path, _: u16) -> io::Result<()> { +fn announce_rpc_port(_: &Path, _: u16, _: &transport::SessionToken) -> io::Result<()> { Err(io::Error::new( io::ErrorKind::Unsupported, "relay requires a Unix-like remote host", )) } +/// Accepts RPC connections, serving each one the service its role asked for +/// once it has presented the session token. +async fn serve_rpc(listener: tokio::net::TcpListener, state: State) { + let slots = Arc::new(tokio::sync::Semaphore::new(MAX_RPC_CONNECTIONS)); + loop { + let stream = match listener.accept().await { + Ok((stream, _)) => stream, + Err(error) => { + error!("could not accept an RPC connection: {error}"); + tokio::time::sleep(Duration::from_millis(100)).await; + continue; + } + }; + let state = state.clone(); + let slots = slots.clone(); + tokio::spawn(async move { + let (role, stream) = match transport::accept(stream, &state.token).await { + Ok(accepted) => accepted, + Err(error) => { + info!("rejected an RPC connection: {error}"); + return; + } + }; + let Ok(_slot) = slots.acquire_owned().await else { + return; + }; + match role { + transport::Role::Gui => { + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + tarpc::server::BaseChannel::with_defaults(transport) + .execute(state.serve()) + .for_each(spawn) + .await; + } + transport::Role::GuiCallback => { + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + let client = + GuiClient::new(tarpc::client::Config::default(), transport).spawn(); + state.connect_gui(client).await; + } + transport::Role::Agent => { + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + tarpc::server::BaseChannel::with_defaults(transport) + .execute(agent::Agent::serve(Backend { state })) + .for_each(spawn) + .await; + } + } + }); + } +} + pub async fn main(rpc_port: Option, relay_socket: Option) { main_with_io( rpc_port, @@ -1974,15 +2129,23 @@ pub async fn main_with_io( return; } }; - main_with_io_on_listener(listener, relay_socket, stdin, stdout).await; + let token = match transport::SessionToken::generate() { + Ok(token) => token, + Err(error) => { + eprintln!("failed to generate the analyzer session token: {error}"); + return; + } + }; + main_with_io_on_listener(listener, token, relay_socket, stdin, stdout).await; } /// Runs the analyzer using pre-bound GUI RPC and LSP transports. /// -/// Supplying the listener lets tests reserve the RPC address without a -/// bind-release-bind race. +/// Supplying the listener and token lets tests reserve the RPC address +/// without a bind-release-bind race and authenticate their own clients. pub async fn main_with_io_on_listener( rpc_listener: tokio::net::TcpListener, + token: transport::SessionToken, relay_socket: Option, stdin: I, stdout: O, @@ -1990,17 +2153,15 @@ pub async fn main_with_io_on_listener( I: AsyncRead + Unpin, O: AsyncWrite, { - let mut listener = - match tarpc::serde_transport::tcp::listen_on(rpc_listener, Bincode::default).await { - Ok(listener) => listener, - Err(error) => { - eprintln!("failed to configure analyzer RPC listener: {error}"); - return; - } - }; - let server_addr = listener.local_addr(); + let server_addr = match rpc_listener.local_addr() { + Ok(addr) => addr, + Err(error) => { + eprintln!("failed to read the analyzer RPC address: {error}"); + return; + } + }; if let Some(path) = relay_socket - && let Err(error) = announce_rpc_port(&path, server_addr.port()) + && let Err(error) = announce_rpc_port(&path, server_addr.port(), &token) { eprintln!( "failed to announce analyzer RPC port through `{}`: {error}", @@ -2011,7 +2172,7 @@ pub async fn main_with_io_on_listener( let mut ext_state = None; let (service, socket) = LspService::build(|client| { - let state = State::new(server_addr, client); + let state = State::new(server_addr, client, token); ext_state = Some(state.clone()); Backend { state } }) @@ -2025,26 +2186,13 @@ pub async fn main_with_io_on_listener( .custom_method("custom/setConfig", Backend::set_config) .custom_method("custom/saveConfig", Backend::save_config) .custom_method("custom/workspaceModified", Backend::workspace_modified) + .custom_method("custom/followMode", Backend::follow_mode) .finish(); let Some(state) = ext_state else { eprintln!("failed to initialize analyzer state"); return; }; - listener.config_mut().max_frame_length(usize::MAX); - let state_clone = state.clone(); - tokio::spawn(async move { - listener - // Ignore accept errors. - .filter_map(|r| futures::future::ready(r.ok())) - .map(tarpc::server::BaseChannel::with_defaults) - // serve is generated by the service attribute. It takes as input any type implementing - // the generated World trait. - .map(|channel| channel.execute(state_clone.clone().serve()).for_each(spawn)) - // Max 10 channels. - .buffer_unordered(10) - .for_each(|_| async {}) - .await; - }); + tokio::spawn(serve_rpc(rpc_listener, state.clone())); state .editor_client @@ -2058,6 +2206,7 @@ pub async fn main_with_io_on_listener( // Start actual LSP server. Server::new(stdin, stdout, socket).serve(service).await; + state.retire_session(); } #[cfg(test)] @@ -2288,7 +2437,7 @@ mod tests { #[cfg(unix)] #[test] - fn rpc_port_is_announced_through_a_unix_socket() { + fn rpc_port_and_token_are_announced_through_a_unix_socket() { use std::io::Read; let directory = tempfile::tempdir().unwrap(); @@ -2300,7 +2449,11 @@ mod tests { stream.read_to_string(&mut port).unwrap(); port }); - super::announce_rpc_port(&path, 43210).unwrap(); - assert_eq!(receiver.join().unwrap(), "43210\n"); + let token = super::transport::SessionToken::generate().unwrap(); + super::announce_rpc_port(&path, 43210, &token).unwrap(); + assert_eq!( + receiver.join().unwrap(), + format!("43210 {}\n", token.to_hex()) + ); } } diff --git a/crates/analyzer/src/mcp.rs b/crates/analyzer/src/mcp.rs new file mode 100644 index 00000000..23298902 --- /dev/null +++ b/crates/analyzer/src/mcp.rs @@ -0,0 +1,741 @@ +//! A Model Context Protocol server that lets agents work through a running +//! analyzer, so their edits land in the user's Neovim buffers. +//! +//! It speaks newline-delimited JSON-RPC on stdin and stdout and forwards tool +//! calls to the analyzer whose workspace contains its working directory. + +use std::{ + io, + net::{Ipv4Addr, SocketAddr}, + path::{Path, PathBuf}, + sync::Arc, +}; + +use serde_json::{Value, json}; +use tarpc::{context, tokio_serde::formats::Bincode}; +use tokio::{ + io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, + sync::{Mutex, mpsc}, +}; + +use crate::{ + agent::{ + AgentClient, AgentStatus, CALL_DEADLINE, CompileStatus, CreateRequest, EditRequest, + Replacement, Report, + }, + session, transport, +}; + +/// Used when a client asks for no particular protocol version. +const DEFAULT_PROTOCOL_VERSION: &str = "2025-06-18"; + +const INSTRUCTIONS: &str = "This server edits Argon layout sources through the user's running \ +Argon session. Edits go into the user's Neovim buffers, so the user watches them land in Neovim \ +and the layout GUI and can undo them. Buffers may hold unsaved changes, so read .ar files with \ +read_file rather than from disk, edit them with edit_file or create_file rather than writing \ +files, and check cells with compile_cell rather than `arc run`. Never save files; the user does."; + +pub(crate) fn no_session_message(directory: &Path) -> String { + format!( + "No running Argon session was found for {}. Ask the user to open the project with `argone` (or open its lib.ar in Neovim with the Argon plugin), then try again.", + directory.display() + ) +} + +/// Connects to the innermost running analyzer whose workspace contains +/// `directory`, discarding records of analyzers that are gone. +pub(crate) async fn connect_session(directory: &Path) -> Result { + let mut last_error = None; + for (record, info) in session::discover(directory) { + let Some(token) = info.token() else { + continue; + }; + let addr = SocketAddr::from((Ipv4Addr::LOCALHOST, info.port)); + match transport::connect(addr, &token, transport::Role::Agent).await { + Ok(stream) => { + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + return Ok(AgentClient::new(tarpc::client::Config::default(), transport).spawn()); + } + Err(error) => { + if error.kind() == io::ErrorKind::ConnectionRefused { + let _ = std::fs::remove_file(record); + } + last_error = Some(error); + } + } + } + Err(match last_error { + Some(error) => format!( + "Could not connect to the Argon session for {}: {error}", + directory.display() + ), + None => no_session_message(directory), + }) +} + +fn call_context() -> context::Context { + let mut context = context::current(); + context.deadline = std::time::Instant::now() + CALL_DEADLINE; + context +} + +fn is_disconnected(error: &tarpc::client::RpcError) -> bool { + matches!( + error, + tarpc::client::RpcError::Shutdown + | tarpc::client::RpcError::Send(_) + | tarpc::client::RpcError::Channel(_) + ) +} + +/// Answers MCP messages for the session whose workspace contains `directory`. +pub struct Bridge { + directory: PathBuf, + client: Mutex>, +} + +impl Bridge { + pub fn new(directory: PathBuf) -> Self { + Self { + directory, + client: Mutex::new(None), + } + } + + async fn client(&self) -> Result { + let mut client = self.client.lock().await; + if let Some(client) = client.as_ref() { + return Ok(client.clone()); + } + let connected = connect_session(&self.directory).await?; + *client = Some(connected.clone()); + Ok(connected) + } + + /// Runs `call`, reconnecting once if the analyzer went away. + async fn call(&self, call: F) -> Result + where + F: Fn(AgentClient) -> Fut, + Fut: Future>, + { + let client = self.client().await?; + match call(client).await { + Ok(value) => Ok(value), + Err(error) if is_disconnected(&error) => { + *self.client.lock().await = None; + let client = self.client().await?; + call(client).await.map_err(|error| error.to_string()) + } + Err(error) => Err(error.to_string()), + } + } + + /// The response to one JSON-RPC message, or `None` for a notification. + pub async fn handle(&self, message: Value) -> Option { + let id = message.get("id").cloned()?; + let method = message.get("method").and_then(Value::as_str).unwrap_or(""); + let params = message.get("params").cloned().unwrap_or(Value::Null); + let result = match method { + "initialize" => Ok(initialize_result(¶ms)), + "ping" => Ok(json!({})), + "tools/list" => Ok(json!({ "tools": tool_definitions() })), + "tools/call" => Ok(self.call_tool(¶ms).await), + _ => Err((-32601, format!("method not found: {method}"))), + }; + Some(match result { + Ok(result) => json!({ "jsonrpc": "2.0", "id": id, "result": result }), + Err((code, message)) => json!({ + "jsonrpc": "2.0", + "id": id, + "error": { "code": code, "message": message }, + }), + }) + } + + async fn call_tool(&self, params: &Value) -> Value { + let name = params.get("name").and_then(Value::as_str).unwrap_or(""); + let arguments = params + .get("arguments") + .cloned() + .unwrap_or_else(|| json!({})); + let outcome = match self.run_tool(name, &arguments).await { + Ok(text) => (text, false), + Err(text) => (text, true), + }; + json!({ + "content": [{ "type": "text", "text": outcome.0 }], + "isError": outcome.1, + }) + } + + async fn run_tool(&self, name: &str, arguments: &Value) -> Result { + match name { + "status" => { + let status = self + .call(|client| async move { client.status(call_context()).await }) + .await?; + Ok(format_status(&status)) + } + "read_file" => { + let path = PathBuf::from(string_argument(arguments, "path")?); + let offset = arguments.get("offset").and_then(Value::as_u64); + let limit = arguments.get("limit").and_then(Value::as_u64); + let file = self + .call(|client| { + let path = path.clone(); + async move { client.read_file(call_context(), path).await } + }) + .await??; + Ok(format!( + "{} ({})\n{}", + file.path.display(), + if file.from_editor { + "from the Neovim buffer, which may be unsaved" + } else { + "from disk; not open in Neovim" + }, + numbered_lines(&file.contents, offset, limit) + )) + } + "edit_file" => { + let request = EditRequest { + path: PathBuf::from(string_argument(arguments, "path")?), + edits: replacements(arguments)?, + label: label(arguments, "Agent edit"), + }; + let report = self + .call(|client| { + let request = request.clone(); + async move { client.edit_file(call_context(), request).await } + }) + .await??; + Ok(format!( + "Applied the edit to the Neovim buffer for {}; it is unsaved.\n{}", + request.path.display(), + format_report(&report) + )) + } + "create_file" => { + let request = CreateRequest { + path: PathBuf::from(string_argument(arguments, "path")?), + contents: string_argument(arguments, "contents")?.to_owned(), + label: label(arguments, "Agent created a file"), + }; + let report = self + .call(|client| { + let request = request.clone(); + async move { client.create_file(call_context(), request).await } + }) + .await??; + Ok(format!( + "Created {} as an unsaved Neovim buffer.\n{}", + request.path.display(), + format_report(&report) + )) + } + "diagnostics" => { + let report = self + .call(|client| async move { client.diagnostics(call_context()).await }) + .await?; + Ok(format_report(&report)) + } + "compile_cell" => { + let cell = string_argument(arguments, "cell")?.to_owned(); + let report = self + .call(|client| { + let cell = cell.clone(); + async move { client.compile_cell(call_context(), cell).await } + }) + .await??; + Ok(format_report(&report)) + } + "open_cell" => { + let cell = string_argument(arguments, "cell")?.to_owned(); + let report = self + .call(|client| { + let cell = cell.clone(); + async move { client.open_cell(call_context(), cell).await } + }) + .await??; + Ok(format!( + "Opened {cell} in the GUI.\n{}", + format_report(&report) + )) + } + _ => Err(format!("Unknown tool `{name}`.")), + } + } +} + +fn initialize_result(params: &Value) -> Value { + // The tools use only features every protocol revision shares, so the + // client's preferred revision is accepted as is. + let version = params + .get("protocolVersion") + .and_then(Value::as_str) + .unwrap_or(DEFAULT_PROTOCOL_VERSION); + json!({ + "protocolVersion": version, + "capabilities": { "tools": { "listChanged": false } }, + "serverInfo": { "name": "argon", "version": env!("CARGO_PKG_VERSION") }, + "instructions": INSTRUCTIONS, + }) +} + +fn string_argument<'a>(arguments: &'a Value, name: &str) -> Result<&'a str, String> { + arguments + .get(name) + .and_then(Value::as_str) + .ok_or_else(|| format!("Missing string argument `{name}`.")) +} + +fn label(arguments: &Value, default: &str) -> String { + arguments + .get("label") + .and_then(Value::as_str) + .map(str::trim) + .filter(|label| !label.is_empty()) + .unwrap_or(default) + .to_owned() +} + +fn replacements(arguments: &Value) -> Result, String> { + let edits = arguments + .get("edits") + .and_then(Value::as_array) + .ok_or_else(|| "Missing array argument `edits`.".to_owned())?; + edits + .iter() + .enumerate() + .map(|(index, edit)| { + let field = |name| { + edit.get(name) + .and_then(Value::as_str) + .map(str::to_owned) + .ok_or_else(|| format!("Edit {}: missing string `{name}`.", index + 1)) + }; + Ok(Replacement { + old: field("old_string")?, + new: field("new_string")?, + replace_all: edit + .get("replace_all") + .and_then(Value::as_bool) + .unwrap_or(false), + }) + }) + .collect() +} + +/// Lines prefixed with their one-based number and a tab. +fn numbered_lines(contents: &str, offset: Option, limit: Option) -> String { + let first = offset.unwrap_or(1).max(1) as usize; + let count = limit.map_or(usize::MAX, |limit| limit as usize); + contents + .lines() + .enumerate() + .skip(first - 1) + .take(count) + .map(|(index, line)| format!("{:>6}\t{line}", index + 1)) + .collect::>() + .join("\n") +} + +fn format_status(status: &AgentStatus) -> String { + let mut lines = vec![ + format!( + "Workspace: {}", + status.root.as_ref().map_or_else( + || "(not open yet)".to_owned(), + |root| root.display().to_string() + ) + ), + format!( + "Open cell: {}", + status.open_cell.as_deref().unwrap_or("(none)") + ), + format!( + "GUI: {}", + if status.gui_connected { + "connected" + } else { + "not running" + } + ), + format!( + "Follow mode: {}", + if status.follow_mode { + "on (open_cell may switch the GUI's cell)" + } else { + "off (open_cell is unavailable; use compile_cell)" + } + ), + format!( + "Approval: {}", + match status.approval { + crate::ApprovalMode::Never => "edits apply immediately", + crate::ApprovalMode::Always => "the user confirms each edit in Neovim", + } + ), + format!( + "Unsaved changes in Neovim: {}", + if status.workspace_modified { + "yes" + } else { + "no" + } + ), + format!( + "Files changed on disk: {}", + if status.watching_files { + "reloaded automatically" + } else { + "not watched" + } + ), + ]; + if !status.open_files.is_empty() { + lines.push("Open in Neovim:".to_owned()); + lines.extend( + status + .open_files + .iter() + .map(|path| format!(" {}", path.display())), + ); + } + lines.join("\n") +} + +fn format_report(report: &Report) -> String { + let mut lines = Vec::new(); + if report.compiled_revision < report.revision { + lines.push(format!( + "Compilation of revision {} has not finished; the results below are from revision {}.", + report.revision, report.compiled_revision + )); + } + match &report.summary { + Some(summary) => { + let cell = summary.cell.as_deref().unwrap_or("(no cell open)"); + let status = match summary.status { + CompileStatus::Valid => "compiled", + CompileStatus::ExecErrors => "compiled with execution errors", + CompileStatus::StaticErrors => "has static errors", + CompileStatus::ParseErrors => "has parse errors", + }; + lines.push(format!("Cell {cell} {status}.")); + if matches!( + summary.status, + CompileStatus::Valid | CompileStatus::ExecErrors + ) { + lines.push(format!( + "Top cell: {} rects, {} polygons, {} paths, {} instances; {} cells in the hierarchy.", + summary.rects, summary.polygons, summary.paths, summary.instances, summary.cells + )); + if let Some([x0, y0, x1, y1]) = summary.bbox { + lines.push(format!("Bounding box: ({x0}, {y0}) to ({x1}, {y1}).")); + } + if summary.unsolved_vars > 0 { + lines.push(format!( + "{} coordinates are set only by initial conditions.", + summary.unsolved_vars + )); + } + } + } + None => lines.push( + "No cell is open in the GUI, so only the workspace was checked. Use compile_cell to compile one." + .to_owned(), + ), + } + if report.diagnostics.is_empty() { + lines.push("No diagnostics.".to_owned()); + } else { + lines.push("Diagnostics:".to_owned()); + lines.extend(report.diagnostics.iter().map(|diagnostic| { + format!( + " {}:{}:{}: {}", + diagnostic.path.display(), + diagnostic.line, + diagnostic.column, + diagnostic.message + ) + })); + } + if !report.messages.is_empty() { + lines.push("Messages:".to_owned()); + lines.extend(report.messages.iter().map(|message| format!(" {message}"))); + } + lines.join("\n") +} + +fn tool_definitions() -> Value { + let path = json!({ + "type": "string", + "description": "File path, relative to the workspace root or absolute." + }); + let label = json!({ + "type": "string", + "description": "A few words describing the change, shown to the user while it is applied." + }); + let cell = json!({ + "type": "string", + "description": "A cell invocation, such as `top()` or `inverter(2, nf=4)`." + }); + json!([ + { + "name": "status", + "description": "Describe the running Argon session: workspace root, the cell open in the GUI, follow mode, the approval mode, and the files open in Neovim.", + "inputSchema": { "type": "object", "properties": {} } + }, + { + "name": "read_file", + "description": "Read a workspace file as the user currently has it, including unsaved Neovim edits. Use this instead of reading .ar files from disk, which can be out of date. Each line is prefixed with its number and a tab; never include that prefix in edit_file's old_string.", + "inputSchema": { + "type": "object", + "properties": { + "path": path, + "offset": { "type": "integer", "description": "First line to show, one-based." }, + "limit": { "type": "integer", "description": "Most lines to show." } + }, + "required": ["path"] + } + }, + { + "name": "edit_file", + "description": "Edit an Argon source file through the user's Neovim buffer, so the change shows up live in Neovim and the GUI and can be undone. Each old_string must match the current text exactly and be unique unless replace_all is set. All edits in one call are matched against the file before any is applied, must not overlap, and land as one undo step. The buffer is left unsaved. The result reports the recompiled cell's shape and any diagnostics. The user may be asked to approve the edit, and may reject it.", + "inputSchema": { + "type": "object", + "properties": { + "path": path, + "edits": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "properties": { + "old_string": { "type": "string", "description": "Exact text to replace." }, + "new_string": { "type": "string", "description": "Replacement text." }, + "replace_all": { "type": "boolean", "description": "Replace every occurrence. Defaults to false." } + }, + "required": ["old_string", "new_string"] + } + }, + "label": label + }, + "required": ["path", "edits", "label"] + } + }, + { + "name": "create_file", + "description": "Create a new Argon source file as an unsaved Neovim buffer. A module must also be declared from its parent, for example `mod utils;` in lib.ar loads utils.ar. The directory must already exist.", + "inputSchema": { + "type": "object", + "properties": { + "path": path, + "contents": { "type": "string" }, + "label": label + }, + "required": ["path", "contents", "label"] + } + }, + { + "name": "diagnostics", + "description": "List the current errors across the workspace and summarize the cell open in the GUI.", + "inputSchema": { "type": "object", "properties": {} } + }, + { + "name": "compile_cell", + "description": "Compile a cell against the current source, unsaved edits included, without changing what the GUI shows. Reports errors and the cell's shape and bounding box. Use this instead of `arc run`.", + "inputSchema": { + "type": "object", + "properties": { "cell": cell }, + "required": ["cell"] + } + }, + { + "name": "open_cell", + "description": "Show a cell in the user's GUI. Works only while the user has follow mode on (`:Argon follow` in Neovim); otherwise use compile_cell.", + "inputSchema": { + "type": "object", + "properties": { "cell": cell }, + "required": ["cell"] + } + } + ]) +} + +fn parse_error(error: serde_json::Error) -> Value { + json!({ + "jsonrpc": "2.0", + "id": Value::Null, + "error": { "code": -32700, "message": format!("parse error: {error}") }, + }) +} + +/// Serves MCP on stdin and stdout until stdin closes. +pub async fn run(directory: Option) -> io::Result<()> { + let directory = match directory { + Some(directory) => directory, + None => std::env::current_dir()?, + }; + let bridge = Arc::new(Bridge::new(directory)); + let (responses, mut outgoing) = mpsc::unbounded_channel::(); + let writer = tokio::spawn(async move { + let mut stdout = tokio::io::stdout(); + while let Some(response) = outgoing.recv().await { + let mut line = response.to_string(); + line.push('\n'); + stdout.write_all(line.as_bytes()).await?; + stdout.flush().await?; + } + io::Result::Ok(()) + }); + + let mut lines = BufReader::new(tokio::io::stdin()).lines(); + while let Some(line) = lines.next_line().await? { + if line.trim().is_empty() { + continue; + } + let message = match serde_json::from_str::(&line) { + Ok(message) => message, + Err(error) => { + let _ = responses.send(parse_error(error)); + continue; + } + }; + // Requests run concurrently, so a call waiting on the user's + // approval does not hold up the others. + let bridge = bridge.clone(); + let responses = responses.clone(); + tokio::spawn(async move { + if let Some(response) = bridge.handle(message).await { + let _ = responses.send(response); + } + }); + } + drop(responses); + writer.await.map_err(io::Error::other)? +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::{AgentDiagnostic, CompileSummary}; + + fn bridge(directory: &Path) -> Bridge { + Bridge::new(directory.to_path_buf()) + } + + #[tokio::test] + async fn answers_the_handshake_and_lists_tools() { + let directory = tempfile::tempdir().unwrap(); + let bridge = bridge(directory.path()); + let response = bridge + .handle(json!({ + "jsonrpc": "2.0", "id": 1, "method": "initialize", + "params": { "protocolVersion": "2025-03-26", "capabilities": {} } + })) + .await + .unwrap(); + assert_eq!(response["result"]["protocolVersion"], "2025-03-26"); + assert!(response["result"]["capabilities"]["tools"].is_object()); + + assert!( + bridge + .handle(json!({ "jsonrpc": "2.0", "method": "notifications/initialized" })) + .await + .is_none() + ); + + let tools = bridge + .handle(json!({ "jsonrpc": "2.0", "id": 2, "method": "tools/list" })) + .await + .unwrap(); + let names = tools["result"]["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| tool["name"].as_str().unwrap().to_owned()) + .collect::>(); + assert_eq!( + names, + [ + "status", + "read_file", + "edit_file", + "create_file", + "diagnostics", + "compile_cell", + "open_cell" + ] + ); + + let unknown = bridge + .handle(json!({ "jsonrpc": "2.0", "id": 3, "method": "resources/list" })) + .await + .unwrap(); + assert_eq!(unknown["error"]["code"], -32601); + } + + #[tokio::test] + async fn tool_calls_report_errors_as_tool_results() { + let directory = tempfile::tempdir().unwrap(); + let bridge = bridge(directory.path()); + let missing = bridge + .handle(json!({ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": { "name": "edit_file", "arguments": { "path": "lib.ar" } } + })) + .await + .unwrap(); + assert_eq!(missing["result"]["isError"], true); + assert!( + missing["result"]["content"][0]["text"] + .as_str() + .unwrap() + .contains("edits") + ); + } + + #[test] + fn numbers_lines_within_the_requested_window() { + let text = "a\nb\nc\n"; + assert_eq!( + numbered_lines(text, None, None), + " 1\ta\n 2\tb\n 3\tc" + ); + assert_eq!(numbered_lines(text, Some(2), Some(1)), " 2\tb"); + } + + #[test] + fn reports_say_when_compilation_lags_behind() { + let report = Report { + revision: 5, + compiled_revision: 4, + summary: Some(CompileSummary { + cell: Some("top()".to_owned()), + status: CompileStatus::Valid, + cells: 2, + rects: 3, + polygons: 0, + paths: 0, + instances: 1, + unsolved_vars: 0, + bbox: Some([0., 0., 10., 5.]), + }), + diagnostics: vec![AgentDiagnostic { + path: PathBuf::from("lib.ar"), + line: 3, + column: 5, + message: "undeclared variable".to_owned(), + }], + messages: Vec::new(), + }; + let text = format_report(&report); + assert!(text.contains("revision 5 has not finished")); + assert!(text.contains("Cell top() compiled.")); + assert!(text.contains("3 rects")); + assert!(text.contains("Bounding box: (0, 0) to (10, 5).")); + assert!(text.contains("lib.ar:3:5: undeclared variable")); + } +} diff --git a/crates/analyzer/src/rpc.rs b/crates/analyzer/src/rpc.rs index 07708c40..728c7cd8 100644 --- a/crates/analyzer/src/rpc.rs +++ b/crates/analyzer/src/rpc.rs @@ -1,6 +1,6 @@ //! RPC types shared by the analyzer and Argone. -use std::{collections::HashMap, net::SocketAddr, path::PathBuf}; +use std::{collections::HashMap, path::PathBuf}; use argonc::{ ast::Span, @@ -9,7 +9,7 @@ use argonc::{ }; use serde::{Deserialize, Serialize}; -use tarpc::{context, tokio_serde::formats::Bincode}; +use tarpc::context; use tower_lsp_server::ls_types::{ Diagnostic, DiagnosticSeverity, MessageType, Position, Range, ShowDocumentParams, TextEdit, Uri, WorkspaceEdit, @@ -192,7 +192,6 @@ pub struct FocusEditorParams { #[tarpc::service] pub trait LangServer { - async fn register(addr: SocketAddr); async fn select_rect(span: Span); async fn draw_rect( scope_span: Span, @@ -232,6 +231,11 @@ pub trait Gui { async fn place_instance(preview: InstancePreview); async fn configure(config: ArgonConfig); async fn activate(); + /// Briefly outlines objects whose source overlaps `spans`, once the GUI + /// shows `revision` or later. + async fn highlight_spans(revision: u64, spans: Vec); + /// Shows what an agent is doing, or clears it with `None`. + async fn agent_activity(label: Option); } pub(crate) const OUT_OF_SYNC_MESSAGE: &str = "Editor buffer state is inconsistent with GUI state."; @@ -432,20 +436,37 @@ impl State { } } + /// Marks an analyzer edit to `uri` as in flight, so the `didChange` it + /// produces compiles immediately. + pub(crate) async fn begin_pending_edit(&self, uri: &Uri) { + *self + .source_state + .lock() + .await + .pending_workspace_edits + .entry(uri.clone()) + .or_default() += 1; + } + + /// Withdraws an in-flight edit that will not produce a `didChange`. + pub(crate) async fn end_pending_edit(&self, uri: &Uri) { + let mut state = self.source_state.lock().await; + if let Some(count) = state.pending_workspace_edits.get_mut(uri) { + *count = count.saturating_sub(1); + if *count == 0 { + state.pending_workspace_edits.shift_remove(uri); + } + } + } + pub(crate) async fn apply_source_changes( &self, changes: HashMap>, focus: Option, ) -> bool { let pending_uris = changes.keys().cloned().collect::>(); - { - let mut state = self.source_state.lock().await; - for uri in &pending_uris { - *state - .pending_workspace_edits - .entry(uri.clone()) - .or_default() += 1; - } + for uri in &pending_uris { + self.begin_pending_edit(uri).await; } let result: Result<(), String> = async { if let Some(uri) = focus { @@ -480,16 +501,9 @@ impl State { .await; if let Err(error) = result { - let mut state = self.source_state.lock().await; - for uri in pending_uris { - if let Some(count) = state.pending_workspace_edits.get_mut(&uri) { - *count = count.saturating_sub(1); - if *count == 0 { - state.pending_workspace_edits.shift_remove(&uri); - } - } + for uri in &pending_uris { + self.end_pending_edit(uri).await; } - drop(state); self.report_message(MessageType::ERROR, error).await; false } else { @@ -503,22 +517,10 @@ impl State { } } -impl LangServer for State { - async fn register(self, _: tarpc::context::Context, addr: SocketAddr) -> () { - let mut transport = tarpc::serde_transport::tcp::connect(addr, Bincode::default); - transport.config_mut().max_frame_length(usize::MAX); - let gui_client = match transport.await { - Ok(transport) => GuiClient::new(tarpc::client::Config::default(), transport).spawn(), - Err(error) => { - self.editor_client - .show_message( - MessageType::ERROR, - format!("Could not connect to the GUI: {error}"), - ) - .await; - return; - } - }; +impl State { + /// Adopts a GUI that connected back for callbacks, replacing any earlier + /// one, and sends it the session's configuration and current cell. + pub(crate) async fn connect_gui(self, gui_client: GuiClient) { let connection = self.install_gui_connection(gui_client.clone()).await; let modified = self.source_state.lock().await.workspace_modified; if let Err(error) = gui_client @@ -541,14 +543,11 @@ impl LangServer for State { } self.publish_workspace_modified(modified, Some(connection)) .await; - // Acknowledge the connection before compiling/preparing its first - // snapshot. That work may need GUI callbacks or exceed the handshake - // deadline; neither should prevent the window from opening. - tokio::spawn(async move { - Backend { state: self }.open_current().await; - }); + Backend { state: self }.open_current().await; } +} +impl LangServer for State { async fn select_rect(self, _: tarpc::context::Context, span: Span) { // TODO: check that vim file is in sync with GUI file. let compiled = self.published_state.lock().await; diff --git a/crates/analyzer/src/session.rs b/crates/analyzer/src/session.rs new file mode 100644 index 00000000..3dc587d7 --- /dev/null +++ b/crates/analyzer/src/session.rs @@ -0,0 +1,219 @@ +//! Per-user records of running analyzers, so local clients can find them. + +use std::{ + fs, io, + path::{Path, PathBuf}, +}; + +use serde::{Deserialize, Serialize}; + +use crate::transport::SessionToken; + +/// Where one analyzer can be reached, readable only by the user that ran it. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SessionInfo { + pub pid: u32, + pub root: PathBuf, + pub port: u16, + pub token: String, +} + +impl SessionInfo { + pub fn token(&self) -> Option { + SessionToken::from_hex(&self.token) + } +} + +pub fn sessions_dir() -> Option { + crate::argon_state_dir().map(|directory| directory.join("sessions")) +} + +/// A published session record, removed when dropped. +#[derive(Debug)] +pub struct SessionFile { + path: PathBuf, +} + +impl Drop for SessionFile { + fn drop(&mut self) { + let _ = fs::remove_file(&self.path); + } +} + +fn create_private_dir(directory: &Path) -> io::Result<()> { + #[cfg(unix)] + { + use std::os::unix::fs::{DirBuilderExt, PermissionsExt}; + fs::DirBuilder::new() + .recursive(true) + .mode(0o700) + .create(directory)?; + fs::set_permissions(directory, fs::Permissions::from_mode(0o700)) + } + #[cfg(not(unix))] + { + fs::create_dir_all(directory) + } +} + +fn write_private_file(path: &Path, contents: &[u8]) -> io::Result<()> { + use std::io::Write; + let mut options = 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(contents)?; + file.sync_all() +} + +/// Publishes `info` in `directory` under the analyzer's process ID. +pub fn publish_in(directory: &Path, info: &SessionInfo) -> io::Result { + create_private_dir(directory)?; + let path = directory.join(format!("{}.json", info.pid)); + let staging = directory.join(format!("{}.json.tmp", info.pid)); + let _ = fs::remove_file(&staging); + let contents = serde_json::to_vec_pretty(info).map_err(io::Error::other)?; + write_private_file(&staging, &contents)?; + fs::rename(&staging, &path)?; + Ok(SessionFile { path }) +} + +pub fn publish(info: &SessionInfo) -> io::Result { + let directory = sessions_dir() + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "no state directory"))?; + publish_in(&directory, info) +} + +/// Resolves symlinks in the longest prefix of `path` that exists. +fn canonical(path: &Path) -> PathBuf { + for ancestor in path.ancestors() { + if let Ok(resolved) = ancestor.canonicalize() { + let rest = path.strip_prefix(ancestor).unwrap_or(Path::new("")); + return resolved.join(rest); + } + } + path.to_path_buf() +} + +/// Every readable session record in `directory`. +pub fn list_in(directory: &Path) -> Vec<(PathBuf, SessionInfo)> { + let Ok(entries) = fs::read_dir(directory) else { + return Vec::new(); + }; + entries + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| { + path.extension() + .is_some_and(|extension| extension == "json") + }) + .filter_map(|path| { + let info = serde_json::from_slice(&fs::read(&path).ok()?).ok()?; + Some((path, info)) + }) + .collect() +} + +/// Sessions whose workspace contains `directory`, innermost workspace first. +pub fn discover_in(sessions: &Path, directory: &Path) -> Vec<(PathBuf, SessionInfo)> { + let directory = canonical(directory); + let mut matches = list_in(sessions) + .into_iter() + .filter_map(|(path, info)| { + let root = canonical(&info.root); + directory + .starts_with(&root) + .then(|| (root.components().count(), path, info)) + }) + .collect::>(); + matches.sort_by(|(a_depth, a_path, _), (b_depth, b_path, _)| { + let modified = |path: &Path| fs::metadata(path).and_then(|m| m.modified()).ok(); + b_depth + .cmp(a_depth) + .then_with(|| modified(b_path).cmp(&modified(a_path))) + }); + matches + .into_iter() + .map(|(_, path, info)| (path, info)) + .collect() +} + +pub fn discover(directory: &Path) -> Vec<(PathBuf, SessionInfo)> { + sessions_dir() + .map(|sessions| discover_in(&sessions, directory)) + .unwrap_or_default() +} + +/// The session listening on `port`, for clients given only an address. +pub fn find_by_port(port: u16) -> Option { + list_in(&sessions_dir()?) + .into_iter() + .map(|(_, info)| info) + .find(|info| info.port == port) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn info(pid: u32, root: &Path, port: u16) -> SessionInfo { + SessionInfo { + pid, + root: root.to_path_buf(), + port, + token: SessionToken::generate().unwrap().to_hex(), + } + } + + #[test] + fn published_sessions_are_private_and_removed_on_drop() { + let state = tempfile::tempdir().unwrap(); + let sessions = state.path().join("sessions"); + let file = publish_in(&sessions, &info(7, state.path(), 1234)).unwrap(); + let path = sessions.join("7.json"); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + assert_eq!( + fs::metadata(&sessions).unwrap().permissions().mode() & 0o777, + 0o700 + ); + assert_eq!( + fs::metadata(&path).unwrap().permissions().mode() & 0o777, + 0o600 + ); + } + assert_eq!(list_in(&sessions).len(), 1); + drop(file); + assert!(!path.exists()); + } + + #[test] + fn discovery_prefers_the_innermost_workspace() { + let state = tempfile::tempdir().unwrap(); + let sessions = state.path().join("sessions"); + let outer = state.path().join("outer"); + let inner = outer.join("inner"); + let elsewhere = state.path().join("elsewhere"); + for directory in [&inner, &elsewhere] { + fs::create_dir_all(directory).unwrap(); + } + let _outer = publish_in(&sessions, &info(1, &outer, 1)).unwrap(); + let _inner = publish_in(&sessions, &info(2, &inner, 2)).unwrap(); + let _elsewhere = publish_in(&sessions, &info(3, &elsewhere, 3)).unwrap(); + + let ports = |directory: &Path| { + discover_in(&sessions, directory) + .into_iter() + .map(|(_, info)| info.port) + .collect::>() + }; + assert_eq!(ports(&inner.join("src")), vec![2, 1]); + assert_eq!(ports(&outer), vec![1]); + assert_eq!(ports(state.path()), Vec::::new()); + } +} diff --git a/crates/analyzer/src/transport.rs b/crates/analyzer/src/transport.rs new file mode 100644 index 00000000..265594b6 --- /dev/null +++ b/crates/analyzer/src/transport.rs @@ -0,0 +1,340 @@ +//! Authenticated connections to the analyzer's RPC listener. +//! +//! Every connection opens with a JSON `Hello` frame carrying the session token +//! and the client's role. Only after the analyzer accepts it does the stream +//! carry tarpc traffic, so an unauthenticated peer never reaches a service. + +use std::{fmt, io, net::SocketAddr, time::Duration}; + +use bytes::Bytes; +use futures::{SinkExt, StreamExt}; +use serde::{Deserialize, Serialize}; +use tokio::net::TcpStream; +use tokio_util::codec::{Framed, LengthDelimitedCodec}; + +/// Environment variable that hands the session token to a launched client. +pub const TOKEN_ENV: &str = "ARGON_SESSION_TOKEN"; + +/// Version of the handshake and of the RPC services behind it. +pub const PROTOCOL_VERSION: u32 = 1; + +/// How long a peer may take to complete the handshake. +pub const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(5); + +/// Bounds what an unauthenticated peer can make the analyzer buffer. +const MAX_HANDSHAKE_FRAME: usize = 4096; + +const TOKEN_BYTES: usize = 32; + +/// A stream that has completed the handshake, ready for a tarpc transport. +pub type AuthenticatedStream = Framed; + +/// The secret that every RPC connection to one analyzer must present. +#[derive(Clone, PartialEq, Eq)] +pub struct SessionToken([u8; TOKEN_BYTES]); + +impl SessionToken { + pub fn generate() -> io::Result { + let mut bytes = [0; TOKEN_BYTES]; + getrandom::fill(&mut bytes).map_err(|error| io::Error::other(error.to_string()))?; + Ok(Self(bytes)) + } + + pub fn to_hex(&self) -> String { + self.0.iter().map(|byte| format!("{byte:02x}")).collect() + } + + pub fn from_hex(text: &str) -> Option { + let text = text.trim(); + if text.len() != TOKEN_BYTES * 2 || !text.is_ascii() { + return None; + } + let mut bytes = [0; TOKEN_BYTES]; + for (index, byte) in bytes.iter_mut().enumerate() { + *byte = u8::from_str_radix(&text[2 * index..2 * index + 2], 16).ok()?; + } + Some(Self(bytes)) + } + + /// Reads the token from [`TOKEN_ENV`]. + pub fn from_env() -> Option { + std::env::var(TOKEN_ENV) + .ok() + .and_then(|text| Self::from_hex(&text)) + } + + /// Compares in constant time, so response timing reveals nothing about the token. + fn matches(&self, other: &Self) -> bool { + self.0 + .iter() + .zip(other.0.iter()) + .fold(0u8, |difference, (a, b)| difference | (a ^ b)) + == 0 + } +} + +impl fmt::Debug for SessionToken { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("SessionToken(..)") + } +} + +/// What a connection is for, which decides the service it is served. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum Role { + /// The GUI calling the analyzer. + Gui, + /// The analyzer calling the GUI, over a connection the GUI opened. + GuiCallback, + /// An agent bridge calling the analyzer. + Agent, +} + +#[derive(Serialize, Deserialize)] +struct Hello { + version: u32, + token: String, + role: Role, +} + +#[derive(Serialize, Deserialize)] +enum HelloReply { + Accepted, + Rejected(String), +} + +fn handshake_codec() -> LengthDelimitedCodec { + LengthDelimitedCodec::builder() + .max_frame_length(MAX_HANDSHAKE_FRAME) + .new_codec() +} + +fn into_rpc_stream(mut framed: AuthenticatedStream) -> AuthenticatedStream { + framed.codec_mut().set_max_frame_length(usize::MAX); + framed +} + +async fn send_json(framed: &mut AuthenticatedStream, value: &impl Serialize) -> io::Result<()> { + let bytes = serde_json::to_vec(value).map_err(io::Error::other)?; + framed.send(Bytes::from(bytes)).await +} + +async fn receive_json Deserialize<'de>>( + framed: &mut AuthenticatedStream, +) -> io::Result { + let frame = framed.next().await.ok_or_else(|| { + io::Error::new( + io::ErrorKind::UnexpectedEof, + "connection closed during the handshake", + ) + })??; + serde_json::from_slice(&frame) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) +} + +/// Connects to an analyzer and authenticates as `role`. +pub async fn connect( + addr: SocketAddr, + token: &SessionToken, + role: Role, +) -> io::Result { + tokio::time::timeout(HANDSHAKE_TIMEOUT, async { + let stream = TcpStream::connect(addr).await?; + let mut framed = Framed::new(stream, handshake_codec()); + send_json( + &mut framed, + &Hello { + version: PROTOCOL_VERSION, + token: token.to_hex(), + role, + }, + ) + .await?; + match receive_json(&mut framed).await? { + HelloReply::Accepted => Ok(into_rpc_stream(framed)), + HelloReply::Rejected(reason) => Err(io::Error::new( + io::ErrorKind::PermissionDenied, + format!("the analyzer rejected the connection: {reason}"), + )), + } + }) + .await + .map_err(|_| { + io::Error::new( + io::ErrorKind::TimedOut, + format!("timed out connecting to the analyzer at {addr}"), + ) + })? +} + +/// Authenticates an accepted connection, returning the role it asked for. +pub async fn accept( + stream: TcpStream, + token: &SessionToken, +) -> io::Result<(Role, AuthenticatedStream)> { + accept_within(stream, token, HANDSHAKE_TIMEOUT).await +} + +async fn accept_within( + stream: TcpStream, + token: &SessionToken, + limit: Duration, +) -> io::Result<(Role, AuthenticatedStream)> { + let mut framed = Framed::new(stream, handshake_codec()); + let role = tokio::time::timeout(limit, async { + let hello: Hello = receive_json(&mut framed).await?; + let rejection = if !SessionToken::from_hex(&hello.token).is_some_and(|t| t.matches(token)) + { + Some("invalid session token".to_owned()) + } else if hello.version != PROTOCOL_VERSION { + Some(format!( + "protocol version {} is not supported; this analyzer speaks version {PROTOCOL_VERSION}", + hello.version + )) + } else { + None + }; + if let Some(reason) = rejection { + let _ = send_json(&mut framed, &HelloReply::Rejected(reason.clone())).await; + return Err(io::Error::new(io::ErrorKind::PermissionDenied, reason)); + } + send_json(&mut framed, &HelloReply::Accepted).await?; + Ok(hello.role) + }) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "handshake timed out"))??; + Ok((role, into_rpc_stream(framed))) +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use futures::SinkExt; + use tokio::net::TcpListener; + + use super::*; + + async fn listener() -> (TcpListener, SocketAddr) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + (listener, addr) + } + + #[test] + fn tokens_round_trip_through_hex_and_never_print() { + let token = SessionToken::generate().unwrap(); + assert_eq!(SessionToken::from_hex(&token.to_hex()), Some(token.clone())); + assert_ne!(SessionToken::generate().unwrap(), token); + assert!(!format!("{token:?}").contains(&token.to_hex())); + assert_eq!(SessionToken::from_hex("abc"), None); + assert_eq!(SessionToken::from_hex(&"g".repeat(64)), None); + } + + #[tokio::test] + async fn accepts_the_session_token_and_reports_the_role() { + let (listener, addr) = listener().await; + let token = SessionToken::generate().unwrap(); + let server_token = token.clone(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + accept(stream, &server_token).await.map(|(role, _)| role) + }); + connect(addr, &token, Role::Agent).await.unwrap(); + assert_eq!(server.await.unwrap().unwrap(), Role::Agent); + } + + #[tokio::test] + async fn rejects_a_wrong_token() { + let (listener, addr) = listener().await; + let token = SessionToken::generate().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + accept(stream, &token).await.map(|(role, _)| role) + }); + let error = connect(addr, &SessionToken::generate().unwrap(), Role::Gui) + .await + .unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::PermissionDenied); + assert!(error.to_string().contains("invalid session token")); + assert_eq!( + server.await.unwrap().unwrap_err().kind(), + io::ErrorKind::PermissionDenied + ); + } + + #[tokio::test] + async fn rejects_a_peer_that_skips_the_handshake() { + let (listener, addr) = listener().await; + let token = SessionToken::generate().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + accept_within(stream, &token, Duration::from_millis(100)) + .await + .map(|(role, _)| role) + }); + let _silent = TcpStream::connect(addr).await.unwrap(); + assert_eq!( + server.await.unwrap().unwrap_err().kind(), + io::ErrorKind::TimedOut + ); + } + + #[tokio::test] + async fn rejects_oversized_and_malformed_hellos() { + let (listener, addr) = listener().await; + let token = SessionToken::generate().unwrap(); + let server = tokio::spawn(async move { + let mut results = Vec::new(); + for _ in 0..2 { + let (stream, _) = listener.accept().await.unwrap(); + results.push(accept(stream, &token).await.map(|(role, _)| role)); + } + results + }); + let mut oversized = Framed::new( + TcpStream::connect(addr).await.unwrap(), + LengthDelimitedCodec::new(), + ); + let _ = oversized + .send(Bytes::from(vec![b' '; MAX_HANDSHAKE_FRAME + 1])) + .await; + let mut malformed = Framed::new( + TcpStream::connect(addr).await.unwrap(), + LengthDelimitedCodec::new(), + ); + malformed + .send(Bytes::from_static(b"not json")) + .await + .unwrap(); + let results = server.await.unwrap(); + assert!(results.iter().all(Result::is_err)); + } + + #[tokio::test] + async fn rejects_a_mismatched_protocol_version_with_a_reason() { + let (listener, addr) = listener().await; + let token = SessionToken::generate().unwrap(); + let client_token = token.clone(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + accept(stream, &token).await.map(|(role, _)| role) + }); + let mut framed = Framed::new(TcpStream::connect(addr).await.unwrap(), handshake_codec()); + send_json( + &mut framed, + &Hello { + version: PROTOCOL_VERSION + 1, + token: client_token.to_hex(), + role: Role::Gui, + }, + ) + .await + .unwrap(); + let HelloReply::Rejected(reason) = receive_json(&mut framed).await.unwrap() else { + panic!("a mismatched version should be rejected"); + }; + assert!(reason.contains("protocol version")); + assert!(server.await.unwrap().is_err()); + } +} diff --git a/crates/analyzer/src/watcher.rs b/crates/analyzer/src/watcher.rs new file mode 100644 index 00000000..2ac52236 --- /dev/null +++ b/crates/analyzer/src/watcher.rs @@ -0,0 +1,221 @@ +//! Picks up workspace files that change on disk outside Neovim. + +use std::{ + path::{Path, PathBuf}, + time::Duration, +}; + +use indexmap::IndexSet; +use notify::{EventKind, RecommendedWatcher, RecursiveMode, Watcher}; +use serde::{Deserialize, Serialize}; +use tower_lsp_server::ls_types::{MessageType, Uri, request::Request}; +use tracing::error; + +use crate::{Backend, agent}; + +/// How long to gather a burst of events, such as an editor's atomic save. +const DEBOUNCE: Duration = Duration::from_millis(100); + +/// Keeps the operating-system watch alive. +pub(crate) struct WorkspaceWatcher { + _watcher: RecommendedWatcher, +} + +impl std::fmt::Debug for WorkspaceWatcher { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("WorkspaceWatcher") + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct ReloadBufferParams { + uri: Uri, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(tag = "status", rename_all = "lowercase")] +pub(crate) enum ReloadBufferResult { + Reloaded, + Unchanged, + /// The buffer has unsaved changes, so Neovim kept them. + Modified, + Closed, +} + +/// Asks Neovim to reread a file that changed on disk, unless its buffer has +/// unsaved changes. +enum ReloadBuffer {} + +impl Request for ReloadBuffer { + type Params = ReloadBufferParams; + type Result = ReloadBufferResult; + const METHOD: &'static str = "custom/reloadBuffer"; +} + +/// `path` relative to the workspace, which events may name through either +/// the root as given or its canonical form. +fn relative_to<'a>(roots: &[PathBuf], path: &'a Path) -> Option<&'a Path> { + roots.iter().find_map(|root| path.strip_prefix(root).ok()) +} + +/// Whether a change to `path` can affect compilation. +fn is_relevant(roots: &[PathBuf], path: &Path) -> bool { + let Some(relative) = relative_to(roots, path) else { + return false; + }; + let hidden_or_output = relative.components().any(|component| { + let name = component.as_os_str().to_string_lossy(); + name.starts_with('.') || name == "target" + }); + !hidden_or_output + && path + .extension() + .is_some_and(|extension| extension == "ar" || extension == "toml" || extension == "gds") +} + +fn workspace_roots(root: &Path) -> Vec { + let mut roots = vec![root.to_path_buf()]; + if let Ok(canonical) = root.canonicalize() + && canonical != root + { + roots.push(canonical); + } + roots +} + +impl Backend { + pub(crate) fn start_watcher(&self, root: &Path) { + let (events, mut received) = tokio::sync::mpsc::unbounded_channel(); + let watcher = notify::recommended_watcher(move |event| { + let _ = events.send(event); + }) + .and_then(|mut watcher| { + watcher.watch(root, RecursiveMode::Recursive)?; + Ok(watcher) + }); + let watcher = match watcher { + Ok(watcher) => watcher, + Err(error) => { + error!("could not watch {} for changes: {error}", root.display()); + return; + } + }; + *self + .state + .watcher + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = + Some(WorkspaceWatcher { _watcher: watcher }); + + let backend = self.clone(); + let roots = workspace_roots(root); + tokio::spawn(async move { + let collect = |paths: &mut IndexSet, event: notify::Result| { + let Ok(event) = event else { + return; + }; + if matches!(event.kind, EventKind::Access(_)) { + return; + } + paths.extend( + event + .paths + .into_iter() + .filter(|path| is_relevant(&roots, path)), + ); + }; + while let Some(event) = received.recv().await { + let mut paths = IndexSet::new(); + collect(&mut paths, event); + tokio::time::sleep(DEBOUNCE).await; + while let Ok(event) = received.try_recv() { + collect(&mut paths, event); + } + if !paths.is_empty() { + backend.disk_changed(paths).await; + } + } + }); + } + + /// Reloads open buffers whose files changed and recompiles for closed + /// ones, which the compiler rereads from disk. + async fn disk_changed(&self, paths: IndexSet) { + let mut closed_file_changed = false; + for path in paths { + let open = { + let source = self.state.source_state.lock().await; + agent::open_uri(&source, &path) + .map(|uri| (uri.clone(), source.editor_files[&uri].contents().to_owned())) + }; + let Some((uri, buffer)) = open else { + closed_file_changed = true; + continue; + }; + // A missing file is Neovim's to report, and matching text is + // Neovim's own write. + match tokio::fs::read_to_string(&path).await { + Ok(disk) if disk != buffer => {} + _ => continue, + } + match self + .state + .editor_client + .send_request::(ReloadBufferParams { uri }) + .await + { + Ok(ReloadBufferResult::Modified) => { + let roots = self + .state + .root_dir + .get() + .map(|root| workspace_roots(root)) + .unwrap_or_default(); + let name = relative_to(&roots, &path).unwrap_or(&path); + self.state + .report_message( + MessageType::WARNING, + format!( + "{} changed on disk, but Neovim has unsaved changes to it, so it was not reloaded.", + name.display() + ), + ) + .await; + } + Ok(_) => {} + Err(error) => error!("could not ask Neovim to reload {}: {error}", path.display()), + } + } + if closed_file_changed { + let identity = { + let mut source = self.state.source_state.lock().await; + source.advance_revision(); + source.compile_identity() + }; + self.compile_after_debounce(identity); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn only_sources_and_inputs_outside_hidden_and_output_directories_matter() { + let roots = [ + PathBuf::from("/tmp/.hidden/lib"), + PathBuf::from("/private/tmp/.hidden/lib"), + ]; + let relevant = |path: &str| is_relevant(&roots, Path::new(path)); + assert!(relevant("/tmp/.hidden/lib/lib.ar")); + assert!(relevant("/private/tmp/.hidden/lib/src/cells.ar")); + assert!(relevant("/tmp/.hidden/lib/Argon.toml")); + assert!(relevant("/tmp/.hidden/lib/layout/sram.gds")); + assert!(!relevant("/tmp/.hidden/lib/target/argon.gds")); + assert!(!relevant("/tmp/.hidden/lib/.git/index")); + assert!(!relevant("/tmp/.hidden/lib/.lib.ar.swp")); + assert!(!relevant("/tmp/.hidden/lib/notes.md")); + assert!(!relevant("/elsewhere/lib.ar")); + } +} diff --git a/crates/compiler/src/compile.rs b/crates/compiler/src/compile.rs index 45d001d5..92166527 100644 --- a/crates/compiler/src/compile.rs +++ b/crates/compiler/src/compile.rs @@ -13701,6 +13701,44 @@ fn object_id(id: &mut u64) -> ObjectId { ObjectId(next_id) } +impl CompiledData { + /// The bounding box of `cell`'s layout geometry, instances included, in + /// `cell`'s own coordinate frame. + pub fn layout_bbox(&self, cell: CellId) -> Option> { + self.layout_bbox_memoized(cell, &mut HashMap::new()) + } + + fn layout_bbox_memoized( + &self, + cell: CellId, + memo: &mut HashMap>>, + ) -> Option> { + if let Some(bbox) = memo.get(&cell) { + return bbox.clone(); + } + let mut bbox = None; + for object in self.cells.get(&cell)?.objects.values() { + if !object.is_layout() { + continue; + } + match object { + SolvedValue::Rect(r) => bbox = bbox_union(bbox, Some(r.to_float())), + SolvedValue::Polygon(p) => bbox = bbox_union(bbox, p.bbox()), + SolvedValue::Path(p) => bbox = bbox_union(bbox, p.bbox()), + SolvedValue::Instance(i) => { + let child = self + .layout_bbox_memoized(i.cell, memo) + .map(|r| r.transform(i.reflect, i.angle).translate(i.x, i.y)); + bbox = bbox_union(bbox, child); + } + _ => {} + } + } + memo.insert(cell, bbox.clone()); + bbox + } +} + impl CompiledData { /// Resolve one object's source name without allocating names for every /// reachable object. Later bindings and array entries win, as they do in diff --git a/crates/gui/src/cli.rs b/crates/gui/src/cli.rs index 32f9739e..be596246 100644 --- a/crates/gui/src/cli.rs +++ b/crates/gui/src/cli.rs @@ -2,15 +2,16 @@ use std::{ env, ffi::{OsStr, OsString}, fs, - io::{self, BufRead, BufReader, Write}, + io::{self, BufRead, BufReader}, net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener, TcpStream}, path::{Path, PathBuf}, - process::{Child, ChildStdin, Command, ExitCode, ExitStatus, Stdio}, + process::{Child, Command, ExitCode, ExitStatus, Stdio}, sync::mpsc, thread, time::{Duration, Instant}, }; +use analyzer::transport::{SessionToken, TOKEN_ENV}; use anyhow::{Context, Result, anyhow, bail}; use clap::{Args, Parser, Subcommand}; @@ -70,32 +71,12 @@ struct SshArgs { /// Analyzer RPC port on the remote host. #[arg(long)] remote_analyzer_port: Option, - - /// Local GUI callback port. - #[arg(long)] - local_gui_port: Option, - - /// GUI callback port exposed on the remote host. - #[arg(long)] - remote_gui_port: Option, } #[derive(Debug, Args)] struct GuiArgs { /// Analyzer RPC address to connect to. - lang_server_addr: Option, - - /// Local callback port. Omit to let the operating system allocate one. - #[arg(long)] - listen_port: Option, - - /// Callback address advertised to the analyzer, for example through an SSH tunnel. - #[arg(long)] - register_addr: Option, - - /// Coordinate an SSH launch through stdin and stdout. - #[arg(long)] - ssh_control: bool, + lang_server_addr: SocketAddr, } pub fn run() -> ExitCode { @@ -123,31 +104,29 @@ fn execute(cli: Cli) -> Result { } } -fn run_gui(args: GuiArgs) -> Result { - if args.ssh_control { - let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, args.listen_port.unwrap_or(0))) - .context("failed to bind the GUI callback listener")?; - let listen_addr = listener.local_addr()?; - println!("ARGON_GUI 1 {}", listen_addr.port()); - io_flush_stdout()?; - let mut forwarding = String::new(); - std::io::stdin() - .lock() - .read_line(&mut forwarding) - .context("failed to read SSH forwarding configuration")?; - let (lang_server_addr, register_addr) = parse_address_pair(&forwarding)?; - crate::run_with_listener(lang_server_addr, listener, register_addr); - } else { - let lang_server_addr = args.lang_server_addr.ok_or_else(|| { - anyhow!("an analyzer address is required unless --ssh-control is used") - })?; - crate::run_gui(lang_server_addr, args.listen_port, args.register_addr); - } - Ok(success_status()) +/// The analyzer's session token, from the environment or, for a local +/// analyzer, from its session record. +fn session_token(lang_server_addr: SocketAddr) -> Result { + SessionToken::from_env() + .or_else(|| { + lang_server_addr + .ip() + .is_loopback() + .then(|| analyzer::session::find_by_port(lang_server_addr.port())) + .flatten() + .and_then(|session| session.token()) + }) + .ok_or_else(|| { + anyhow!( + "no session token for the analyzer at {lang_server_addr}; set {TOKEN_ENV}, or start the GUI from Neovim with `:Argon gui`" + ) + }) } -fn io_flush_stdout() -> Result<()> { - std::io::stdout().flush().context("failed to flush stdout") +fn run_gui(args: GuiArgs) -> Result { + let token = session_token(args.lang_server_addr)?; + crate::run_gui(args.lang_server_addr, token); + Ok(success_status()) } fn run_nvim(nvim: &OsStr, path: &Path, focus_target: Option<&str>) -> Result { @@ -247,48 +226,28 @@ fn start_forwarded_session( relay: &mut Relay, focus_target: Option<&str>, ) -> Result<(Child, Child)> { - let remote_analyzer_port = wait_for_remote_analyzer(nvim_ssh, relay)?; - let mut gui = launch_forwarded_gui(args.local_gui_port, focus_target)?; - let local_gui_port = gui.port; - - let tunnel = match start_tunnel(args, control, remote_analyzer_port, local_gui_port) { - Ok(tunnel) => tunnel, + let (remote_analyzer_port, token) = wait_for_remote_analyzer(nvim_ssh, relay)?; + let mut tunnel = start_tunnel(args, control, remote_analyzer_port)?; + let analyzer_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, tunnel.local_analyzer_port); + match launch_forwarded_gui(analyzer_addr.into(), &token, focus_target) { + Ok(gui) => Ok((tunnel.child, gui)), Err(error) => { - gui.stop(); - return Err(error); + let _ = tunnel.child.kill(); + let _ = tunnel.child.wait(); + Err(error) } - }; - let analyzer_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, tunnel.local_analyzer_port); - let gui_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, tunnel.remote_gui_port); - if let Err(error) = - writeln!(gui.input, "{analyzer_addr} {gui_addr}").and_then(|_| gui.input.flush()) - { - let mut child = tunnel.child; - let _ = child.kill(); - let _ = child.wait(); - gui.stop(); - return Err(error).context("failed to send forwarding configuration to the GUI"); } - Ok((tunnel.child, gui.child)) } fn validate_port_overrides(args: &SshArgs) -> Result<()> { for (name, port) in [ ("--local-analyzer-port", args.local_analyzer_port), ("--remote-analyzer-port", args.remote_analyzer_port), - ("--local-gui-port", args.local_gui_port), - ("--remote-gui-port", args.remote_gui_port), ] { if port == Some(0) { bail!("{name} must be nonzero; omit it to allocate a port automatically"); } } - if args.local_analyzer_port.is_some() && args.local_analyzer_port == args.local_gui_port { - bail!("local analyzer and GUI ports must be different"); - } - if args.remote_analyzer_port.is_some() && args.remote_analyzer_port == args.remote_gui_port { - bail!("remote analyzer and GUI ports must be different"); - } Ok(()) } @@ -399,10 +358,12 @@ fn parse_relay_announcement(line: &str) -> Option<&str> { .filter(|path| !path.is_empty()) } -fn parse_analyzer_announcement(line: &str) -> Option { - line.strip_prefix("ARGON_ANALYZER 1 ") - .and_then(|port| port.trim().parse().ok()) - .filter(|port| *port != 0) +/// The remote analyzer's port and session token, reported over SSH. +fn parse_analyzer_announcement(line: &str) -> Option<(u16, SessionToken)> { + let mut fields = line.strip_prefix("ARGON_ANALYZER 2 ")?.split_whitespace(); + let port = fields.next()?.parse().ok().filter(|port| *port != 0)?; + let token = SessionToken::from_hex(fields.next()?)?; + fields.next().is_none().then_some((port, token)) } fn ssh_command(args: &SshArgs) -> Command { @@ -525,14 +486,17 @@ fn remote_nvim_command( ) } -fn wait_for_remote_analyzer(nvim_ssh: &mut Child, relay: &mut Relay) -> Result { +fn wait_for_remote_analyzer( + nvim_ssh: &mut Child, + relay: &mut Relay, +) -> Result<(u16, SessionToken)> { let deadline = Instant::now() + STARTUP_TIMEOUT; loop { match relay.lines.recv_timeout(STARTUP_POLL_INTERVAL) { Ok(Ok(line)) => { - if let Some(port) = parse_analyzer_announcement(&line) { + if let Some(announcement) = parse_analyzer_announcement(&line) { let _ = relay.child.wait(); - return Ok(port); + return Ok(announcement); } } Ok(Err(error)) => return Err(error.into()), @@ -558,95 +522,34 @@ fn wait_for_remote_analyzer(nvim_ssh: &mut Child, relay: &mut Relay) -> Result, focus_target: Option<&str>) -> Result { +/// Starts the local GUI against the forwarded analyzer port. +fn launch_forwarded_gui( + analyzer_addr: SocketAddr, + token: &SessionToken, + focus_target: Option<&str>, +) -> Result { let mut command = Command::new(env::current_exe()?); - command.arg("gui").arg("--ssh-control"); - if let Some(port) = gui_port { - command.arg("--listen-port").arg(port.to_string()); - } + command + .arg("gui") + .arg(analyzer_addr.to_string()) + .env(TOKEN_ENV, token.to_hex()) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::inherit()); if let Some(focus_target) = focus_target { command.env(crate::focus::TARGET_ENV, focus_target); } - let mut child = command - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::inherit()) + command .spawn() - .context("failed to start the local Argon GUI")?; - let input = child.stdin.take().expect("GUI stdin was piped"); - let stdout = child.stdout.take().expect("GUI stdout was piped"); - let (line_tx, line_rx) = mpsc::channel(); - thread::spawn(move || { - let mut stdout = BufReader::new(stdout); - let mut line = String::new(); - let result = stdout.read_line(&mut line).map(|_| line); - if line_tx.send(result).is_ok() { - let _ = io::copy(&mut stdout, &mut io::sink()); - } - }); - - let deadline = Instant::now() + STARTUP_TIMEOUT; - loop { - match line_rx.recv_timeout(STARTUP_POLL_INTERVAL) { - Ok(Ok(line)) => { - let port = match parse_gui_announcement(&line) { - Ok(port) => port, - Err(error) => { - let _ = child.kill(); - let _ = child.wait(); - return Err(error).context("the GUI reported an invalid callback port"); - } - }; - return Ok(GuiLaunch { child, input, port }); - } - Ok(Err(error)) => { - let _ = child.kill(); - let _ = child.wait(); - return Err(error.into()); - } - Err(mpsc::RecvTimeoutError::Timeout) => {} - Err(mpsc::RecvTimeoutError::Disconnected) => { - let _ = child.kill(); - let _ = child.wait(); - bail!("GUI closed stdout before reporting its callback port"); - } - } - if let Some(status) = child.try_wait()? { - bail!("GUI exited before startup completed ({status})"); - } - if Instant::now() >= deadline { - let _ = child.kill(); - let _ = child.wait(); - bail!("timed out waiting for the GUI to start"); - } - } + .context("failed to start the local Argon GUI") } struct Tunnel { child: Child, local_analyzer_port: u16, - remote_gui_port: u16, } -fn start_tunnel( - args: &SshArgs, - control: &SshControl, - remote_analyzer_port: u16, - local_gui_port: u16, -) -> Result { +fn start_tunnel(args: &SshArgs, control: &SshControl, remote_analyzer_port: u16) -> Result { let attempts = if args.local_analyzer_port.is_some() { 1 } else { @@ -658,13 +561,7 @@ fn start_tunnel( Some(port) => port, None => available_local_port()?, }; - match start_tunnel_once( - args, - control, - local_analyzer_port, - remote_analyzer_port, - local_gui_port, - ) { + match start_tunnel_once(args, control, local_analyzer_port, remote_analyzer_port) { Ok(tunnel) => return Ok(tunnel), Err(error) if args.local_analyzer_port.is_none() && error.local_collision => { last_collision = Some(error.message); @@ -690,12 +587,11 @@ fn start_tunnel_once( control: &SshControl, local_analyzer_port: u16, remote_analyzer_port: u16, - local_gui_port: u16, ) -> std::result::Result { let local_analyzer_addr = SocketAddrV4::new(Ipv4Addr::LOCALHOST, local_analyzer_port); let mut command = control.command(args); command - // A multiplex control operation does not report the port allocated for `-R 0`. + // Owning the forward in this process removes it when the tunnel stops. .arg("-S") .arg("none") .arg("-o") @@ -706,11 +602,6 @@ fn start_tunnel_once( .arg(format!( "{local_analyzer_addr}:127.0.0.1:{remote_analyzer_port}" )) - .arg("-R") - .arg(format!( - "127.0.0.1:{}:127.0.0.1:{local_gui_port}", - args.remote_gui_port.unwrap_or(0) - )) .arg("-N") .arg(&args.host) .stdin(Stdio::null()) @@ -732,16 +623,10 @@ fn start_tunnel_once( let deadline = Instant::now() + STARTUP_TIMEOUT; let mut lines = Vec::new(); - let mut remote_gui_port = args.remote_gui_port; loop { while let Ok(line) = line_rx.try_recv() { match line { - Ok(line) => { - if remote_gui_port.is_none() { - remote_gui_port = parse_allocated_remote_port(&line); - } - lines.push(line); - } + Ok(line) => lines.push(line), Err(error) => lines.push(error.to_string()), } } @@ -774,9 +659,7 @@ fn start_tunnel_once( message: format!("SSH could not establish forwarding: {detail}"), }); } - if let Some(remote_gui_port) = remote_gui_port.filter(|_| { - TcpStream::connect_timeout(&local_analyzer_addr.into(), STARTUP_POLL_INTERVAL).is_ok() - }) { + if TcpStream::connect_timeout(&local_analyzer_addr.into(), STARTUP_POLL_INTERVAL).is_ok() { thread::spawn(move || { while let Ok(line) = line_rx.recv() { if let Ok(line) = line { @@ -787,7 +670,6 @@ fn start_tunnel_once( return Ok(Tunnel { child, local_analyzer_port, - remote_gui_port, }); } if Instant::now() >= deadline { @@ -807,13 +689,6 @@ fn available_local_port() -> Result { Ok(listener.local_addr()?.port()) } -fn parse_allocated_remote_port(line: &str) -> Option { - let rest = line - .find("Allocated port ") - .map(|start| &line[start + "Allocated port ".len()..])?; - rest.split_whitespace().next()?.parse().ok() -} - fn local_forward_collision(stderr: &str, port: u16) -> bool { let mentions_port = stderr.contains(&port.to_string()); mentions_port @@ -822,37 +697,6 @@ fn local_forward_collision(stderr: &str, port: u16) -> bool { || stderr.contains("Could not request local forwarding")) } -fn parse_port(value: &str) -> Result { - let port = value.trim().parse::()?; - if port == 0 { - bail!("port must be nonzero"); - } - Ok(port) -} - -fn parse_gui_announcement(value: &str) -> Result { - let port = value - .strip_prefix("ARGON_GUI 1 ") - .ok_or_else(|| anyhow!("unexpected GUI startup response"))?; - parse_port(port) -} - -fn parse_address_pair(value: &str) -> Result<(SocketAddr, SocketAddr)> { - let mut fields = value.split_whitespace(); - let analyzer = fields - .next() - .ok_or_else(|| anyhow!("missing analyzer forwarding address"))? - .parse()?; - let gui = fields - .next() - .ok_or_else(|| anyhow!("missing GUI forwarding address"))? - .parse()?; - if fields.next().is_some() { - bail!("forwarding configuration contains extra fields"); - } - Ok((analyzer, gui)) -} - fn shell_quote(value: &str) -> String { format!("'{}'", value.replace('\'', "'\"'\"'")) } @@ -896,8 +740,6 @@ mod tests { ssh_options: Vec::new(), local_analyzer_port: None, remote_analyzer_port: None, - local_gui_port: None, - remote_gui_port: None, } } @@ -915,49 +757,33 @@ mod tests { assert_eq!(remote.path, Path::new("/work/chip")); } - #[test] - fn parses_forwarding_addresses() { - assert_eq!( - parse_address_pair("127.0.0.1:1234 127.0.0.1:5678\n").unwrap(), - ( - "127.0.0.1:1234".parse().unwrap(), - "127.0.0.1:5678".parse().unwrap() - ) - ); - assert!(parse_address_pair("127.0.0.1:1234").is_err()); - assert!(parse_address_pair("127.0.0.1:1 127.0.0.1:2 extra").is_err()); - } - #[test] fn parses_startup_protocol_messages() { - assert_eq!(parse_gui_announcement("ARGON_GUI 1 1234\n").unwrap(), 1234); assert_eq!( parse_relay_announcement("ARGON_RELAY 1 /tmp/session/rpc.sock\n"), Some("/tmp/session/rpc.sock") ); + let token = SessionToken::generate().unwrap(); assert_eq!( - parse_analyzer_announcement("ARGON_ANALYZER 1 5678\n"), - Some(5678) + parse_analyzer_announcement(&format!("ARGON_ANALYZER 2 5678 {}\n", token.to_hex())), + Some((5678, token.clone())) ); - assert!(parse_gui_announcement("1234").is_err()); + // The first protocol carried no token, which this GUI cannot use. + assert_eq!(parse_analyzer_announcement("ARGON_ANALYZER 1 5678"), None); assert_eq!(parse_analyzer_announcement("ARGON_ANALYZER 2 5678"), None); + assert_eq!( + parse_analyzer_announcement(&format!("ARGON_ANALYZER 2 0 {}", token.to_hex())), + None + ); } #[test] - fn parses_openssh_allocated_port_message() { - assert_eq!( - parse_allocated_remote_port( - "Allocated port 45678 for remote forward to 127.0.0.1:1234" - ), - Some(45678) - ); - assert_eq!( - parse_allocated_remote_port( - "debug1: Allocated port 45678 for remote forward to 127.0.0.1:1234" - ), - Some(45678) - ); - assert_eq!(parse_allocated_remote_port("Permission denied"), None); + fn rejects_zero_port_overrides() { + let mut args = ssh_args("ssh"); + args.local_analyzer_port = Some(0); + assert!(validate_port_overrides(&args).is_err()); + args.local_analyzer_port = Some(1234); + assert!(validate_port_overrides(&args).is_ok()); } #[test] diff --git a/crates/gui/src/editor/canvas.rs b/crates/gui/src/editor/canvas.rs index 73c9e7b1..36147053 100644 --- a/crates/gui/src/editor/canvas.rs +++ b/crates/gui/src/editor/canvas.rs @@ -127,6 +127,8 @@ impl ShapeFill { } const SELECT_WIDTH: Pixels = px(3.); +/// Outline color for objects an agent edit just changed. +const AGENT_EDIT_COLOR: u32 = 0x00d8ff; const DEFAULT_BORDER_WIDTH: Pixels = px(2.); /// Above this visible primitive count, retain a viewport raster for interactive /// redraws. Sparse and screen-space-aggregated views stay on the geometry path @@ -5884,6 +5886,7 @@ impl Element for CanvasElement { let state = inner.state.read(cx); let pattern_tiles = inner.pattern_tiles.clone(); let tool = state.tool.read(cx).clone(); + let agent_highlight = state.visible_agent_highlight().map(<[Span]>::to_vec); let coalesced_lod_display = inner.raster_tiles.as_ref().and_then(|tiles| { (raster_tiles_visible_lod_is_scale_safe(tiles, bounds, inner.scale, inner.offset)) .then(|| inner.raster_display_transform_for_current_view()) @@ -7994,6 +7997,48 @@ impl Element for CanvasElement { } } let inner = self.inner.read(cx); + if let Some(spans) = &agent_highlight { + // Outline the objects an agent edit just changed. + let changed = |id: &Option| { + id.as_ref().is_some_and(|id| { + spans.iter().any(|span| { + span.path == id.path + && span.span.start() < id.span.end() + && id.span.start() < span.span.end() + }) + }) + }; + let color = rgb(AGENT_EDIT_COLOR); + let changed_rects = rects + .iter() + .map(|(rect, _)| rect) + .chain(scope_rects.iter().map(|scope| &scope.rect)) + .filter(|rect| changed(&rect.id)); + for rect in changed_rects { + window.paint_quad(get_paint_quad( + get_rect_bounds(rect, bounds, scale, offset), + ShapeFill::Solid, + Rgba { a: 0., ..color }, + color, + Edges::all(SELECT_WIDTH), + rect.border_styles, + )); + } + for (polygon, _) in polygons.iter().filter(|(polygon, _)| changed(&polygon.id)) { + let points = polygon + .points + .iter() + .map(|point| inner.layout_to_px(*point)) + .collect::>(); + paint_polygon_border( + window, + &points, + &polygon.edge_styles, + SELECT_WIDTH, + color, + ); + } + } // highlight hover edges // TODO: reduce repeat code from on_left_mouse_down match tool { diff --git a/crates/gui/src/editor/mod.rs b/crates/gui/src/editor/mod.rs index 3427ddc0..f723b125 100644 --- a/crates/gui/src/editor/mod.rs +++ b/crates/gui/src/editor/mod.rs @@ -34,6 +34,18 @@ pub mod toolbars; pub(crate) const SOURCE_EDIT_REJECTED_MESSAGE: &str = "Could not apply the source edit. Press Ctrl-Shift-M to view the detailed error in Neovim."; +/// How long objects an agent edit changed stay outlined. +const AGENT_HIGHLIGHT_DURATION: Duration = Duration::from_millis(2000); + +/// Source spans an agent edit changed, outlined once the GUI shows +/// `revision`. +#[derive(Clone, Debug)] +pub struct AgentHighlight { + pub revision: u64, + pub spans: Vec, + generation: u64, +} + #[derive(Clone)] pub struct LayerState { pub name: SharedString, @@ -92,6 +104,10 @@ pub struct EditorState { pub fatal_error: Option, pub message: Option, pub connection_error: Option, + /// What an agent is doing right now, if anything. + pub agent_activity: Option, + pub agent_highlight: Option, + agent_highlight_generation: u64, pub solved_cell: Entity>, pub hide_external_geometry: bool, pub layers: Entity, @@ -232,6 +248,12 @@ impl EditorState { } } + /// The agent highlight to draw, once the layout it describes is shown. + pub(crate) fn visible_agent_highlight(&self) -> Option<&[argonc::ast::Span]> { + let highlight = self.agent_highlight.as_ref()?; + (self.compilation_revision? >= highlight.revision).then_some(highlight.spans.as_slice()) + } + pub fn show_message(&mut self, typ: MessageType, message: impl Into) { if typ != MessageType::LOG { self.message = Some(EditorMessage { @@ -426,12 +448,10 @@ impl Editor { cx: &mut Context, window: &mut Window, lang_server_addr: SocketAddr, - gui_listen_port: Option, - gui_listener: Option, - gui_register_addr: Option, + token: analyzer::transport::SessionToken, ) -> Self { let (lang_server_client, mut rx) = - SyncLangServerClient::new(cx.to_async(), lang_server_addr); + SyncLangServerClient::new(cx.to_async(), lang_server_addr, token); let solved_cell = cx.new(|_cx| None); let tool = cx.new(|_cx| ToolState::default()); let layers = cx.new(|_cx| Layers { @@ -459,6 +479,9 @@ impl Editor { fatal_error: None, message: None, connection_error: None, + agent_activity: None, + agent_highlight: None, + agent_highlight_generation: 0, solved_cell, hide_external_geometry: false, tool, @@ -510,7 +533,7 @@ impl Editor { } }) .detach(); - lang_server_client.register_server(gui_listen_port, gui_listener, gui_register_addr); + lang_server_client.serve_callbacks(); editor } @@ -527,6 +550,14 @@ impl Editor { self.state.update(cx, |state, cx| { state.connection_error = None; state.compilation_revision = Some(snapshot.revision); + // A later revision may have moved the source the highlight names. + if state + .agent_highlight + .as_ref() + .is_some_and(|highlight| highlight.revision < snapshot.revision) + { + state.agent_highlight = None; + } if snapshot.prepared_output.is_some() { // Keep the status animation alive between hierarchy preparation // and the first raster worker started by the next paint. @@ -627,6 +658,47 @@ impl Editor { }); } + pub fn highlight_agent_edit(&self, cx: &mut App, revision: u64, spans: Vec) { + let generation = self.state.update(cx, |state, cx| { + state.agent_highlight_generation += 1; + let generation = state.agent_highlight_generation; + state.agent_highlight = Some(AgentHighlight { + revision, + spans, + generation, + }); + cx.notify(); + generation + }); + self.canvas.update(cx, |_, cx| cx.notify()); + let state = self.state.clone(); + let canvas = self.canvas.clone(); + cx.spawn(async move |cx| { + cx.background_executor() + .timer(AGENT_HIGHLIGHT_DURATION) + .await; + let _ = state.update(cx, |state, cx| { + if state + .agent_highlight + .as_ref() + .is_some_and(|highlight| highlight.generation == generation) + { + state.agent_highlight = None; + cx.notify(); + } + }); + let _ = canvas.update(cx, |_, cx| cx.notify()); + }) + .detach(); + } + + pub fn set_agent_activity(&self, cx: &mut App, label: Option) { + self.state.update(cx, |state, cx| { + state.agent_activity = label.map(SharedString::from); + cx.notify(); + }); + } + pub fn set_workspace_modified(&self, cx: &mut App, modified: bool) { self.state.update(cx, |state, cx| { state.workspace_modified = modified; @@ -828,6 +900,7 @@ impl Render for Editor { fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { let theme = self.theme(cx); let font_size = self.state.read(cx).font_size; + let agent_activity = self.state.read(cx).agent_activity.clone(); let (displayed_status, activity_label) = { let state = self.state.read(cx); ( @@ -930,6 +1003,18 @@ impl Render for Editor { if !status_fills_space { status_bar = status_bar.child(div().flex_1()); } + if let Some(agent_activity) = agent_activity { + status_bar = status_bar.child( + div() + .id("agent_activity") + .flex_shrink() + .min_w_0() + .overflow_x_hidden() + .text_ellipsis() + .text_color(theme.subtext) + .child(SharedString::from(format!("Agent: {agent_activity}"))), + ); + } // Keep the activity row's intrinsic height even while idle. Adding // the label used to shrink the canvas, invalidating every raster tile; // completing that raster hid the label and restarted the cycle. diff --git a/crates/gui/src/editor/render_regression.rs b/crates/gui/src/editor/render_regression.rs index 07ab28a1..c9aa18c8 100644 --- a/crates/gui/src/editor/render_regression.rs +++ b/crates/gui/src/editor/render_regression.rs @@ -24,6 +24,9 @@ fn test_canvas(cx: &mut gpui::TestAppContext) -> Entity { fatal_error: None, message: None, connection_error: None, + agent_activity: None, + agent_highlight: None, + agent_highlight_generation: 0, subscriptions: vec![ cx.observe(&solved_cell, |_, _, cx| cx.notify()), cx.observe(&layers, |_, _, cx| cx.notify()), diff --git a/crates/gui/src/lib.rs b/crates/gui/src/lib.rs index 8f5eca69..75b579eb 100644 --- a/crates/gui/src/lib.rs +++ b/crates/gui/src/lib.rs @@ -1,8 +1,6 @@ -use std::{ - borrow::Cow, - collections::BTreeSet, - net::{SocketAddr, TcpListener}, -}; +use std::{borrow::Cow, collections::BTreeSet, net::SocketAddr}; + +use analyzer::transport::SessionToken; use editor::Editor; use gpui::*; @@ -47,33 +45,7 @@ impl AssetSource for Assets { } } -pub fn run_gui( - lang_server_addr: SocketAddr, - gui_listen_port: Option, - gui_register_addr: Option, -) { - run_inner(lang_server_addr, gui_listen_port, None, gui_register_addr); -} - -pub fn run_with_listener( - lang_server_addr: SocketAddr, - gui_listener: TcpListener, - gui_register_addr: SocketAddr, -) { - run_inner( - lang_server_addr, - None, - Some(gui_listener), - Some(gui_register_addr), - ); -} - -fn run_inner( - lang_server_addr: SocketAddr, - gui_listen_port: Option, - gui_listener: Option, - gui_register_addr: Option, -) { +pub fn run_gui(lang_server_addr: SocketAddr, token: SessionToken) { focus::initialize_target(); analyzer::init_logging(); @@ -140,14 +112,7 @@ fn run_inner( cx.open_window(editor_window_options(), |window, cx| { window.replace_root(cx, |window, cx| { - Editor::new( - cx, - window, - lang_server_addr, - gui_listen_port, - gui_listener, - gui_register_addr, - ) + Editor::new(cx, window, lang_server_addr, token.clone()) }) }) .unwrap(); diff --git a/crates/gui/src/rpc.rs b/crates/gui/src/rpc.rs index b7881464..60cf7821 100644 --- a/crates/gui/src/rpc.rs +++ b/crates/gui/src/rpc.rs @@ -1,7 +1,7 @@ use std::{ fmt::Display, future::Future, - net::{Ipv4Addr, SocketAddr, TcpListener}, + net::SocketAddr, sync::{ Arc, Mutex, MutexGuard, atomic::{AtomicU64, Ordering}, @@ -15,6 +15,7 @@ use analyzer::rpc::{ InitialConditionEdit, InstancePreview, LangServerAction, LangServerClient, PathParams, PolygonParams, RectangleEditResult, ValueEdit, }; +use analyzer::transport::{self, Role, SessionToken}; use anyhow::{Result, anyhow}; use argonc::{ast::Span, compile::BasicRect}; use async_compat::CompatExt; @@ -26,11 +27,7 @@ use futures::{ prelude::*, }; use gpui::AsyncApp; -use tarpc::{ - context, - server::{Channel, incoming::Incoming}, - tokio_serde::formats::Bincode, -}; +use tarpc::{context, server::Channel, tokio_serde::formats::Bincode}; use tower_lsp_server::ls_types::MessageType; use tracing::error; @@ -55,6 +52,7 @@ fn lock_unpoisoned(lock: &Mutex) -> MutexGuard<'_, T> { pub struct SyncLangServerClient { app: AsyncApp, lang_server_addr: SocketAddr, + token: SessionToken, client: Arc>, to_exec: UnboundedSender, } @@ -73,20 +71,22 @@ fn is_disconnected(error: &tarpc::client::RpcError) -> bool { ) } -fn connect_client(app: &AsyncApp, lang_server_addr: SocketAddr) -> Result { +async fn connect_lang_server( + lang_server_addr: SocketAddr, + token: SessionToken, +) -> std::io::Result { + let stream = transport::connect(lang_server_addr, &token, Role::Gui).await?; + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + Ok(LangServerClient::new(tarpc::client::Config::default(), transport).spawn()) +} + +fn connect_client( + app: &AsyncApp, + lang_server_addr: SocketAddr, + token: &SessionToken, +) -> Result { app.background_executor() - .block( - async move { - let mut transport = - tarpc::serde_transport::tcp::connect(lang_server_addr, Bincode::default); - transport.config_mut().max_frame_length(usize::MAX); - let transport = transport.await?; - Ok::<_, std::io::Error>( - LangServerClient::new(tarpc::client::Config::default(), transport).spawn(), - ) - } - .compat(), - ) + .block(connect_lang_server(lang_server_addr, token.clone()).compat()) .map_err(Into::into) } @@ -106,6 +106,7 @@ impl SyncLangServerClient { Self { app, lang_server_addr: "127.0.0.1:1".parse().unwrap(), + token: SessionToken::generate().unwrap(), client: Arc::new(Mutex::new(client)), to_exec, } @@ -125,6 +126,7 @@ impl SyncLangServerClient { Self { app, lang_server_addr: "127.0.0.1:1".parse().unwrap(), + token: SessionToken::generate().unwrap(), client: Arc::new(Mutex::new(client.client)), to_exec, }, @@ -132,13 +134,18 @@ impl SyncLangServerClient { ) } - pub fn new(app: AsyncApp, lang_server_addr: SocketAddr) -> (Self, UnboundedReceiver) { - let client = connect_client(&app, lang_server_addr).unwrap(); + pub fn new( + app: AsyncApp, + lang_server_addr: SocketAddr, + token: SessionToken, + ) -> (Self, UnboundedReceiver) { + let client = connect_client(&app, lang_server_addr, &token).unwrap(); let (to_exec, rx) = mpsc::unbounded(); ( Self { app, lang_server_addr, + token, client: Arc::new(Mutex::new(client)), to_exec, }, @@ -189,23 +196,17 @@ impl SyncLangServerClient { let result = match result { Err(RpcCallError::Rpc(error)) if is_disconnected(&error) => { let addr = self.lang_server_addr; + let token = self.token.clone(); let reconnect = self .app .background_executor() .spawn( async move { - let mut transport = - tarpc::serde_transport::tcp::connect(addr, Bincode::default); - transport.config_mut().max_frame_length(usize::MAX); - tokio::time::timeout(LANG_SERVER_CLIENT_TIMEOUT, transport) - .await? - .map(|transport| { - LangServerClient::new( - tarpc::client::Config::default(), - transport, - ) - .spawn() - }) + tokio::time::timeout( + LANG_SERVER_CLIENT_TIMEOUT, + connect_lang_server(addr, token), + ) + .await? } .compat(), ) @@ -271,7 +272,7 @@ impl SyncLangServerClient { } fn reconnect(&self) -> Result { - let client = connect_client(&self.app, self.lang_server_addr)?; + let client = connect_client(&self.app, self.lang_server_addr, &self.token)?; *lock_unpoisoned(&self.client) = client.clone(); Ok(client) } @@ -302,104 +303,49 @@ impl SyncLangServerClient { })); } - pub fn register_server( - &self, - configured_port: Option, - prebound_listener: Option, - register_addr: Option, - ) { + /// Opens the connection the analyzer calls the GUI back over. + pub fn serve_callbacks(&self) { let client = self.clone(); self.app .spawn(async move |_| { - let result = client - .start_server(configured_port, prebound_listener, register_addr) - .await; + let result = client.connect_callbacks().await; if let Err(error) = &result { - error!("Failed to register GUI: {error}"); + error!("Failed to connect the GUI to the analyzer: {error}"); } client.report_connection_result(&result); }) .detach(); } - async fn start_server( - &self, - configured_port: Option, - prebound_listener: Option, - register_addr: Option, - ) -> Result<()> { - let background_executor = self.app.background_executor().clone(); - let mut listener = self + async fn connect_callbacks(&self) -> Result<()> { + let addr = self.lang_server_addr; + let token = self.token.clone(); + let stream = self .app .background_executor() .spawn( - async move { - if let Some(listener) = prebound_listener { - match listener - .set_nonblocking(true) - .and_then(|_| tokio::net::TcpListener::from_std(listener)) - { - Ok(listener) => { - tarpc::serde_transport::tcp::listen_on(listener, Bincode::default) - .await - } - Err(error) => Err(error), - } - } else { - let port = configured_port.unwrap_or(0); - tarpc::serde_transport::tcp::listen( - (Ipv4Addr::LOCALHOST, port), - Bincode::default, - ) - .await - } - } - .compat(), + async move { transport::connect(addr, &token, Role::GuiCallback).await }.compat(), ) .await?; - let server_addr = listener.local_addr(); - let register_addr = register_addr.unwrap_or(server_addr); - let to_exec = self.to_exec.clone(); + let background_executor = self.app.background_executor().clone(); + let server = GuiServer { + to_exec: self.to_exec.clone(), + snapshot: Arc::default(), + }; self.app .background_executor() .spawn( async move { - listener.config_mut().max_frame_length(usize::MAX); - listener - // Ignore accept errors. - .filter_map(|r| futures::future::ready(r.ok())) - .map(tarpc::server::BaseChannel::with_defaults) - // Limit channels to 1 per IP. - .max_channels_per_key(1, |t| t.transport().peer_addr().unwrap().ip()) - // serve is generated by the service attribute. It takes as input any type implementing - // the generated World trait. - .map(|channel| { - let server = GuiServer { - to_exec: to_exec.clone(), - snapshot: Arc::default(), - }; - channel - .execute(server.serve()) - .for_each(|t| background_executor.spawn(t)) - }) - // Max 10 channels. - .buffer_unordered(10) - .for_each(|_| async {}) + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + tarpc::server::BaseChannel::with_defaults(transport) + .execute(server.serve()) + .for_each(|response| background_executor.spawn(response)) .await; } .compat(), ) .detach(); - self.register_with_analyzer(register_addr).await - } - - async fn register_with_analyzer(&self, register_addr: SocketAddr) -> Result<()> { - // Registration can call back into the GUI for its first snapshot. - // The main thread must keep servicing those callbacks while we wait. - self.call_async(move |client| async move { - client.register(context::current(), register_addr).await - }) - .await + Ok(()) } pub fn select_rect(&self, span: Span) -> Result<()> { @@ -714,6 +660,29 @@ impl Gui for GuiServer { .unwrap(); } + async fn highlight_spans( + mut self, + _: context::Context, + revision: u64, + spans: Vec, + ) { + self.to_exec + .send(Box::new(move |editor, cx| { + let _ = cx.update(|cx| editor.highlight_agent_edit(cx, revision, spans)); + })) + .await + .ok(); + } + + async fn agent_activity(mut self, _: context::Context, label: Option) { + self.to_exec + .send(Box::new(move |editor, cx| { + let _ = cx.update(|cx| editor.set_agent_activity(cx, label)); + })) + .await + .ok(); + } + async fn activate(mut self, _context: ::tarpc::context::Context) -> () { self.to_exec .send(Box::new(|editor, cx| { @@ -738,60 +707,6 @@ mod tests { use super::{is_disconnected, lock_unpoisoned}; - #[gpui::test] - fn registration_yields_until_the_analyzer_replies(cx: &mut gpui::TestAppContext) { - use analyzer::rpc::{LangServerRequest, LangServerResponse}; - use futures::{FutureExt, SinkExt, StreamExt}; - use std::sync::atomic::{AtomicBool, Ordering}; - - let (client, mut server) = super::SyncLangServerClient::for_rpc_test(cx.to_async()); - let completed = Arc::new(AtomicBool::new(false)); - let done = completed.clone(); - cx.to_async() - .spawn(async move |_| { - client - .register_with_analyzer("127.0.0.1:12345".parse().unwrap()) - .await - .unwrap(); - done.store(true, Ordering::Release); - }) - .detach(); - let request = loop { - if let Some(message) = server.next().now_or_never() - && let tarpc::ClientMessage::Request(request) = message.unwrap().unwrap() - { - break request; - } - assert!(cx.dispatcher.tick(false)); - }; - assert!(matches!( - request.message, - LangServerRequest::Register { .. } - )); - assert!(!completed.load(Ordering::Acquire)); - // Foreground work can finish while the handshake is unanswered. - let foreground_ran = Arc::new(AtomicBool::new(false)); - let ran = foreground_ran.clone(); - cx.to_async() - .spawn(async move |_| ran.store(true, Ordering::Release)) - .detach(); - while !foreground_ran.load(Ordering::Acquire) { - assert!(cx.dispatcher.tick(false)); - } - assert!(!completed.load(Ordering::Acquire)); - server - .send(tarpc::Response { - request_id: request.id, - message: Ok(LangServerResponse::Register(())), - }) - .now_or_never() - .unwrap() - .unwrap(); - while !completed.load(Ordering::Acquire) { - assert!(cx.dispatcher.tick(false)); - } - } - #[test] fn client_lock_recovers_from_poisoning() { let lock = Arc::new(Mutex::new(1)); diff --git a/crates/tests/Cargo.toml b/crates/tests/Cargo.toml index 7bb4d00c..7c716548 100644 --- a/crates/tests/Cargo.toml +++ b/crates/tests/Cargo.toml @@ -12,6 +12,7 @@ workspace = true analyzer = { package = "argon-analyzer", path = "../analyzer" } argonc = { path = "../compiler" } futures = { workspace = true } +serde_json = { workspace = true } tarpc = { workspace = true } tempfile = { workspace = true } tokio = { version = "1", features = ["full"] } diff --git a/crates/tests/fixtures/nvim/full_stack.lua b/crates/tests/fixtures/nvim/full_stack.lua index 989e689c..cb48e6be 100644 --- a/crates/tests/fixtures/nvim/full_stack.lua +++ b/crates/tests/fixtures/nvim/full_stack.lua @@ -12,6 +12,36 @@ local function wait_for(description, predicate) end local bufnr = vim.api.nvim_get_current_buf() + +local function buffer_text(buf) + return table.concat(vim.api.nvim_buf_get_lines(buf, 0, -1, false), '\n') +end + +local function signal(name) + vim.fn.writefile({ 'ok' }, vim.env.ARGON_TEST_STEPS .. '/' .. name) +end + +local function await(name) + wait_for('test step ' .. name, function() + return vim.uv.fs_stat(vim.env.ARGON_TEST_STEPS .. '/' .. name) ~= nil + end) +end + +local function find_buffer(suffix) + for _, buf in ipairs(vim.api.nvim_list_bufs()) do + if vim.api.nvim_buf_get_name(buf):sub(-#suffix) == suffix then + return buf + end + end +end + +local function argon_client() + return vim.lsp.get_clients({ name = 'argon', bufnr = bufnr })[1] +end + +local function contains(buf, text) + return buffer_text(buf):find(text, 1, true) ~= nil +end local expected_completeopt if vim.env.ARGON_TEST_MODE == 'navigation' then expected_completeopt = 'menu,preview' @@ -340,6 +370,117 @@ elseif vim.env.ARGON_TEST_MODE == 'navigation' then .. ' got ' .. tostring(shifted.range.start.line) ) +elseif vim.env.ARGON_TEST_MODE == 'agent_edit' then + wait_for('the agent edit to reach the buffer', function() + return contains(bufnr, 'let b = rect(') + end) + assert(vim.bo[bufnr].modified, 'an agent edit should leave the buffer modified') + await('undo') + vim.cmd('undo') + assert(not contains(bufnr, 'let b'), 'one undo should revert the whole agent edit') + assert(contains(bufnr, 'let a = rect('), 'undo should keep the original text') + signal('undone') +elseif vim.env.ARGON_TEST_MODE == 'agent_hidden' then + await('edited_hidden') + local utils = find_buffer('/utils.ar') + assert(utils, 'editing a closed file should load it into a buffer') + assert(vim.bo[utils].buflisted, 'the loaded buffer should be listed') + assert(vim.bo[utils].modified, 'the loaded buffer should hold the unsaved edit') + assert(contains(utils, 'x1=17.'), 'the edit should be in the buffer') + assert(vim.lsp.buf_is_attached(utils, argon_client().id), 'the buffer should be attached') + assert( + vim.api.nvim_get_current_buf() == bufnr, + 'loading a file for the agent should not change the current window' + ) + await('created') + local extra = find_buffer('/extra.ar') + assert(extra and contains(extra, 'cell extra'), 'a created file should be a buffer') + assert(vim.bo[extra].modified, 'a created file should be unsaved') + assert( + vim.uv.fs_stat(vim.api.nvim_buf_get_name(extra)) == nil, + 'a created file should not be written to disk' + ) + signal('hidden_checked') +elseif vim.env.ARGON_TEST_MODE == 'agent_approval' then + local agent = require('argon.agent') + local prompts = {} + agent.confirm = function(_, params) + table.insert(prompts, params) + return #prompts > 1 + end + local function set_approval(value) + local response = argon_client():request_sync('custom/setConfig', { + key = 'agent.approval', + value = value, + }, 10000, bufnr) + assert(response and not response.err, 'could not set the approval mode') + end + set_approval('always') + signal('approval_on') + await('approved') + assert(#prompts == 2, 'each edit should be confirmed, got ' .. #prompts) + assert( + prompts[1].newText:find('rejected_rect', 1, true), + 'the prompt should preview the edited file' + ) + assert(not contains(bufnr, 'rejected_rect'), 'a rejected edit should not be applied') + assert(contains(bufnr, 'approved_rect'), 'an accepted edit should be applied') + set_approval('never') + signal('approval_off') + await('unprompted') + assert(#prompts == 2, 'edits must not be confirmed when approval is off') + assert(contains(bufnr, 'unprompted_rect'), 'the unconfirmed edit should be applied') +elseif vim.env.ARGON_TEST_MODE == 'agent_stale' then + local agent = require('argon.agent') + local original = agent.apply_edit + local interrupted = false + agent.apply_edit = function(params, ctx) + if not interrupted then + interrupted = true + vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, { '// typed during the agent edit' }) + end + return original(params, ctx) + end + signal('typing_armed') + wait_for('the agent edit after the interruption', function() + return contains(bufnr, 'let b = rect(') + end) + assert(interrupted, 'the first attempt should have been interrupted') + assert(contains(bufnr, '// typed during'), 'the concurrent change should survive') +elseif vim.env.ARGON_TEST_MODE == 'agent_follow' then + await('follow') + local win = vim.api.nvim_get_current_win() + vim.cmd('Argon follow') + signal('following') + await('edited_followed') + local shown = vim.api.nvim_win_get_buf(win) + assert( + vim.api.nvim_buf_get_name(shown):sub(-#'/utils.ar') == '/utils.ar', + 'the follow window should show the edited file' + ) + assert( + vim.api.nvim_win_get_cursor(win)[1] == 2, + 'the follow window should move to the edit, got line ' + .. vim.api.nvim_win_get_cursor(win)[1] + ) + signal('followed') +elseif vim.env.ARGON_TEST_MODE == 'watcher' then + wait_for('the changed file to reload', function() + return contains(bufnr, 'let c = rect(') + end) + assert(not vim.bo[bufnr].modified, 'reloading should leave the buffer unmodified') + vim.api.nvim_buf_set_lines(bufnr, -1, -1, false, { '// unsaved' }) + signal('buffer_modified') + await('rewritten') + assert(contains(bufnr, '// unsaved'), 'a disk change must not discard unsaved edits') + assert(not contains(bufnr, 'let d = rect('), 'a modified buffer should not reload') +elseif vim.env.ARGON_TEST_MODE == 'mcp' then + wait_for('the MCP edit to reach the buffer', function() + return contains(bufnr, 'let b = rect(') + end) + await('save') + vim.cmd('write') + signal('saved') elseif vim.env.ARGON_TEST_MODE == 'rpc_errors' then -- The Rust test drives the analyzer RPC directly and acknowledges the -- mirrored GUI error after observing it. diff --git a/crates/tests/src/full_stack.rs b/crates/tests/src/full_stack.rs index f5988f3c..1b04d5d4 100644 --- a/crates/tests/src/full_stack.rs +++ b/crates/tests/src/full_stack.rs @@ -4,7 +4,9 @@ use std::{net::Ipv4Addr, path::PathBuf, sync::Arc}; use analyzer::{ ArgonConfig, + agent::AgentClient, rpc::{CompilationSnapshot, CompilationUpdate, Gui, InstancePreview, LangServerClient}, + transport::{self, Role, SessionToken}, }; use argonc::{ ast::Span, @@ -41,6 +43,11 @@ pub enum GuiEvent { Fit, WorkspacePath(Option), WorkspaceModified(bool), + Highlight { + revision: u64, + spans: Vec, + }, + AgentActivity(Option), } #[derive(Clone)] @@ -129,6 +136,18 @@ impl Gui for HeadlessGui { async fn configure(self, _: context::Context, _: ArgonConfig) {} async fn activate(self, _: context::Context) {} + + async fn highlight_spans(self, _: context::Context, revision: u64, spans: Vec) { + self.events + .send(GuiEvent::Highlight { revision, spans }) + .expect("full-stack test should still be receiving GUI events"); + } + + async fn agent_activity(self, _: context::Context, label: Option) { + self.events + .send(GuiEvent::AgentActivity(label)) + .expect("full-stack test should still be receiving GUI events"); + } } fn snapshot_details(output: &CompileOutput) -> (OutputKind, Option, usize) { @@ -161,11 +180,13 @@ pub struct Session { ack: PathBuf, gui_edit_ack: PathBuf, diagnostic_ack: PathBuf, + steps: PathBuf, analyzer_addr: std::net::SocketAddr, analyzer_listener: Option, + token: SessionToken, lsp_port: u16, lsp_listener: Option, - gui_addr: std::net::SocketAddr, + gui: HeadlessGui, events: mpsc::UnboundedReceiver, #[cfg(test)] update_gate: std::sync::Arc, @@ -173,10 +194,20 @@ pub struct Session { impl Session { pub async fn new(source: &str) -> Self { + Self::with_files(source, &[]).await + } + + /// A project whose `lib.ar` is `source`, alongside the other `files`. + pub async fn with_files(source: &str, files: &[(&str, &str)]) -> Self { let directory = tempfile::tempdir().expect("create full-stack test directory"); let project = directory.path().join("project"); std::fs::create_dir(&project).expect("create test project"); std::fs::write(project.join("lib.ar"), source).expect("write test source"); + for (name, contents) in files { + std::fs::write(project.join(name), contents).expect("write test module"); + } + let steps = directory.path().join("steps"); + std::fs::create_dir(&steps).expect("create test step directory"); std::fs::write( project.join("Argon.toml"), "name = \"full-stack-test\"\ntech = \"tech.toml\"\n", @@ -202,36 +233,13 @@ impl Session { .expect("read analyzer LSP address") .port(); let (events_tx, events) = mpsc::unbounded_channel(); - let selected_scope = Arc::new(std::sync::Mutex::new(None)); - let mut listener = - tarpc::serde_transport::tcp::listen((Ipv4Addr::LOCALHOST, 0), Bincode::default) - .await - .expect("bind headless GUI RPC listener"); - listener.config_mut().max_frame_length(usize::MAX); - let gui_addr = listener.local_addr(); let update_gate = std::sync::Arc::new(tokio::sync::Semaphore::new(1)); - let server_update_gate = update_gate.clone(); - tokio::spawn(async move { - listener - .filter_map(|connection| futures::future::ready(connection.ok())) - .map(tarpc::server::BaseChannel::with_defaults) - .map(move |channel| { - let server = HeadlessGui { - events: events_tx.clone(), - snapshot: Default::default(), - update_gate: server_update_gate.clone(), - selected_scope: selected_scope.clone(), - }; - channel - .execute(server.serve()) - .for_each(|response| async move { - tokio::spawn(response); - }) - }) - .buffer_unordered(10) - .for_each(|_| async {}) - .await; - }); + let gui = HeadlessGui { + events: events_tx, + snapshot: Default::default(), + update_gate: update_gate.clone(), + selected_scope: Arc::new(std::sync::Mutex::new(None)), + }; let ack = directory.path().join("gui.ack"); let gui_edit_ack = directory.path().join("gui-edit.ack"); @@ -242,11 +250,13 @@ impl Session { ack, gui_edit_ack, diagnostic_ack, + steps, analyzer_addr, analyzer_listener: Some(analyzer_listener), + token: SessionToken::generate().expect("generate a session token"), lsp_port, lsp_listener: Some(lsp_listener), - gui_addr, + gui, events, #[cfg(test)] update_gate, @@ -262,15 +272,52 @@ impl Session { .analyzer_listener .take() .expect("analyzer should only be started once"); + let token = self.token.clone(); tokio::spawn(async move { let (stream, _) = listener.accept().await.expect("accept Neovim LSP stream"); let (reader, writer) = tokio::io::split(stream); - analyzer::main_with_io_on_listener(rpc_listener, None, reader, writer).await; + analyzer::main_with_io_on_listener(rpc_listener, token, None, reader, writer).await; }); } - pub fn gui_addr(&self) -> std::net::SocketAddr { - self.gui_addr + pub fn project(&self) -> &std::path::Path { + &self.project + } + + async fn connect(&self, role: Role) -> transport::AuthenticatedStream { + time::timeout(TEST_TIMEOUT, async { + loop { + if let Ok(stream) = transport::connect(self.analyzer_addr, &self.token, role).await + { + return stream; + } + time::sleep(std::time::Duration::from_millis(25)).await; + } + }) + .await + .expect("analyzer RPC server did not start") + } + + /// Connects the headless GUI back to the analyzer for callbacks, as the + /// real GUI does once its window is up. + pub async fn connect_gui(&self) { + let stream = self.connect(Role::GuiCallback).await; + let gui = self.gui.clone(); + tokio::spawn(async move { + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + tarpc::server::BaseChannel::with_defaults(transport) + .execute(gui.serve()) + .for_each(|response| async move { + tokio::spawn(response); + }) + .await; + }); + } + + pub async fn connect_agent(&self) -> AgentClient { + let stream = self.connect(Role::Agent).await; + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + AgentClient::new(tarpc::client::Config::default(), transport).spawn() } pub fn spawn_nvim(&self, mode: &str) -> tokio::process::Child { @@ -281,6 +328,7 @@ impl Session { .env("ARGON_TEST_ACK", &self.ack) .env("ARGON_TEST_GUI_EDIT_ACK", &self.gui_edit_ack) .env("ARGON_TEST_DIAGNOSTIC_ACK", &self.diagnostic_ack) + .env("ARGON_TEST_STEPS", &self.steps) .env("ARGON_TEST_MODE", mode) .env("ARGON_TEST_READY", self.project.join("startup.ready")) .arg("--cmd") @@ -299,20 +347,39 @@ impl Session { } pub async fn connect_analyzer(&self) -> LangServerClient { + let stream = self.connect(Role::Gui).await; + let transport = tarpc::serde_transport::new(stream, Bincode::default()); + LangServerClient::new(tarpc::client::Config::default(), transport).spawn() + } + + /// Tells the Neovim fixture that step `name` happened. + pub fn signal(&self, name: &str) { + std::fs::write(self.steps.join(name), "ok\n").expect("signal a test step"); + } + + /// Waits for the Neovim fixture to signal step `name`. + pub async fn await_signal(&self, name: &str) { + let path = self.steps.join(name); time::timeout(TEST_TIMEOUT, async { - loop { - let mut transport = - tarpc::serde_transport::tcp::connect(self.analyzer_addr, Bincode::default); - transport.config_mut().max_frame_length(usize::MAX); - if let Ok(transport) = transport.await { - return LangServerClient::new(tarpc::client::Config::default(), transport) - .spawn(); - } - time::sleep(std::time::Duration::from_millis(25)).await; + while !path.exists() { + time::sleep(std::time::Duration::from_millis(10)).await; } }) .await - .expect("analyzer RPC server did not start") + .unwrap_or_else(|_| panic!("Neovim never signalled {name}")); + } + + /// Waits for a GUI event matching `predicate`, dropping the others. + pub async fn wait_for_event( + &mut self, + mut predicate: impl FnMut(&GuiEvent) -> bool, + ) -> GuiEvent { + loop { + let event = self.next_event().await; + if predicate(&event) { + return event; + } + } } pub async fn next_event(&mut self) -> GuiEvent { @@ -331,6 +398,72 @@ mod tests { use super::*; use crate::finish_nvim; + use analyzer::{ + ApprovalMode, + agent::{CALL_DEADLINE, CompileStatus, CreateRequest, EditRequest, Replacement}, + }; + use serde_json::{Value, json}; + use tower_lsp_server::ls_types::MessageType; + + const FIRST_RECT: &str = "let a = rect(\"met1\", x0=0., y0=0., x1=10., y1=10.);"; + const SECOND_RECT: &str = "let b = rect(\"met1\", x0=20., y0=0., x1=30., y1=10.);"; + const UNIT: &str = + "cell unit() {\n let body = rect(\"met1\", x0=0., y0=0., x1=5., y1=5.);\n}\n"; + + fn one_rect() -> String { + format!("cell top() {{\n {FIRST_RECT}\n}}\n") + } + + fn with_unit() -> String { + format!( + "mod utils;\n\ncell top() {{\n let unit = inst(utils::unit());\n eq(unit.x, 0.);\n eq(unit.y, 0.);\n {FIRST_RECT}\n}}\n" + ) + } + + fn agent_context() -> context::Context { + let mut context = context::current(); + context.deadline = std::time::Instant::now() + CALL_DEADLINE; + context + } + + fn replace(path: &str, old: &str, new: &str, label: &str) -> EditRequest { + EditRequest { + path: PathBuf::from(path), + edits: vec![Replacement { + old: old.to_owned(), + new: new.to_owned(), + replace_all: false, + }], + label: label.to_owned(), + } + } + + fn add_rect(statement: &str, label: &str) -> EditRequest { + replace( + "lib.ar", + FIRST_RECT, + &format!("{FIRST_RECT}\n {statement}"), + label, + ) + } + + impl Session { + async fn wait_for_rects(&mut self, count: usize) -> u64 { + let GuiEvent::UpdateCell { revision, .. } = self + .wait_for_event(|event| { + matches!( + event, + GuiEvent::UpdateCell { kind: OutputKind::Data, rect_count, .. } + if *rect_count == count + ) + }) + .await + else { + unreachable!() + }; + revision + } + } // Process-heavy scenarios share ports and startup deadlines, so keep them // serial even though Cargo runs Rust tests in parallel by default. @@ -350,10 +483,7 @@ mod tests { session.start_analyzer(); let child = session.spawn_nvim("roundtrip"); let analyzer = session.connect_analyzer().await; - analyzer - .register(context::current(), session.gui_addr()) - .await - .expect("register headless GUI"); + session.connect_gui().await; let mut drew_rect = false; let mut saw_editor_update = false; @@ -455,7 +585,7 @@ mod tests { } #[tokio::test(flavor = "multi_thread")] - async fn gui_registration_does_not_wait_for_initial_error_presentation() { + async fn gui_connection_does_not_wait_for_initial_error_presentation() { assert_completes("registering GUI with initial source errors", async { let _guard = FULL_STACK_LOCK.lock().await; let mut session = Session::new("cell top() { missing; }\n").await; @@ -465,17 +595,11 @@ mod tests { while !session.project.join("startup.ready").exists() { time::sleep(std::time::Duration::from_millis(10)).await; } - let analyzer = session.connect_analyzer().await; - // A real GUI cannot service initial snapshot callbacks until its - // constructor returns. Registration must acknowledge independently - // of that first presentation, including when no cell is selected. - time::timeout( - std::time::Duration::from_secs(1), - analyzer.register(context::current(), session.gui_addr()), - ) - .await - .expect("GUI registration waited for presentation") - .expect("register headless GUI"); + // Connecting must not wait on the first presentation, including + // when no cell is selected. + time::timeout(std::time::Duration::from_secs(1), session.connect_gui()) + .await + .expect("GUI connection waited for presentation"); loop { if matches!( session.next_event().await, @@ -501,11 +625,7 @@ mod tests { let mut session = Session::new("cell top() {\n missing;\n}\n").await; session.start_analyzer(); let child = session.spawn_nvim("diagnostics"); - let analyzer = session.connect_analyzer().await; - analyzer - .register(context::current(), session.gui_addr()) - .await - .expect("register headless GUI"); + session.connect_gui().await; let mut saw_errors = false; let mut saw_recovery = false; @@ -543,11 +663,7 @@ mod tests { .await; session.start_analyzer(); let child = session.spawn_nvim("navigation"); - let analyzer = session.connect_analyzer().await; - analyzer - .register(context::current(), session.gui_addr()) - .await - .expect("register headless GUI"); + session.connect_gui().await; std::fs::write(&session.ack, "ok\n").expect("acknowledge navigation"); finish_nvim(child).await; @@ -563,10 +679,7 @@ mod tests { session.start_analyzer(); let child = session.spawn_nvim("rpc_errors"); let analyzer = session.connect_analyzer().await; - analyzer - .register(context::current(), session.gui_addr()) - .await - .expect("register headless GUI"); + session.connect_gui().await; analyzer .open_cell(context::current(), "top(".to_owned()) .await @@ -586,4 +699,471 @@ mod tests { }) .await; } + + #[tokio::test(flavor = "multi_thread")] + async fn agent_edits_land_in_neovim_and_the_gui() { + assert_completes("waiting for an agent edit round trip", async { + let _guard = FULL_STACK_LOCK.lock().await; + let mut session = Session::new(&one_rect()).await; + session.start_analyzer(); + let child = session.spawn_nvim("agent_edit"); + session.connect_gui().await; + session.wait_for_rects(1).await; + let agent = session.connect_agent().await; + + let status = agent.status(agent_context()).await.unwrap(); + assert!(status.gui_connected); + assert!(!status.follow_mode); + assert_eq!(status.approval, ApprovalMode::Never); + assert_eq!(status.open_cell.as_deref(), Some("top()")); + assert_eq!(status.open_files, vec![PathBuf::from("lib.ar")]); + + let report = agent + .edit_file(agent_context(), add_rect(SECOND_RECT, "add a second rect")) + .await + .unwrap() + .expect("the agent edit should apply"); + assert!(report.compiled_revision >= report.revision); + let summary = report.summary.clone().expect("the edit should compile"); + assert_eq!(summary.status, CompileStatus::Valid); + assert_eq!(summary.rects, 2); + assert_eq!(summary.bbox, Some([0., 0., 30., 10.])); + assert!(report.diagnostics.is_empty(), "{:?}", report.diagnostics); + + let file = agent + .read_file(agent_context(), PathBuf::from("lib.ar")) + .await + .unwrap() + .unwrap(); + assert!(file.from_editor); + assert!(file.contents.contains(SECOND_RECT)); + let disk = std::fs::read_to_string(session.project().join("lib.ar")).unwrap(); + assert!(!disk.contains(SECOND_RECT), "agent edits must stay unsaved"); + + let (mut activity, mut geometry, mut highlight, mut idle) = + (false, false, false, false); + while !(activity && geometry && highlight && idle) { + match session.next_event().await { + GuiEvent::AgentActivity(Some(label)) => { + assert_eq!(label, "add a second rect"); + activity = true; + } + GuiEvent::AgentActivity(None) => idle = activity, + GuiEvent::UpdateCell { + kind: OutputKind::Data, + rect_count: 2, + .. + } => geometry = true, + GuiEvent::Highlight { revision, spans } => { + assert_eq!(revision, report.revision); + let [span] = spans.as_slice() else { + panic!("expected one highlighted span, got {spans:?}"); + }; + assert!(span.path.ends_with("lib.ar")); + assert_eq!( + &file.contents[span.span.start()..span.span.end()], + format!("{FIRST_RECT}\n {SECOND_RECT}") + ); + highlight = true; + } + _ => {} + } + } + + session.signal("undo"); + session.await_signal("undone").await; + session.wait_for_rects(1).await; + std::fs::write(&session.ack, "ok\n").unwrap(); + finish_nvim(child).await; + }) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn agent_edits_closed_files_and_reports_problems() { + assert_completes("waiting for agent edits to closed files", async { + let _guard = FULL_STACK_LOCK.lock().await; + let mut session = Session::with_files(&with_unit(), &[("utils.ar", UNIT)]).await; + session.start_analyzer(); + let child = session.spawn_nvim("agent_hidden"); + session.connect_gui().await; + session.wait_for_rects(1).await; + let agent = session.connect_agent().await; + + let report = agent + .edit_file( + agent_context(), + replace("utils.ar", "x1=5.", "x1=17.", "widen the unit"), + ) + .await + .unwrap() + .expect("editing a closed file should load it"); + assert_eq!(report.summary.unwrap().bbox, Some([0., 0., 17., 10.])); + session.signal("edited_hidden"); + + agent + .create_file( + agent_context(), + CreateRequest { + path: PathBuf::from("extra.ar"), + contents: "cell extra() {\n}\n".to_owned(), + label: "add a module".to_owned(), + }, + ) + .await + .unwrap() + .expect("creating a file should open an unsaved buffer"); + session.signal("created"); + session.await_signal("hidden_checked").await; + + let probe = agent + .compile_cell(agent_context(), "utils::unit()".to_owned()) + .await + .unwrap() + .unwrap(); + let summary = probe.summary.unwrap(); + assert_eq!(summary.status, CompileStatus::Valid); + assert_eq!(summary.bbox, Some([0., 0., 17., 5.])); + let status = agent.status(agent_context()).await.unwrap(); + assert_eq!( + status.open_cell.as_deref(), + Some("top()"), + "probing a cell must not change what the GUI shows" + ); + + let broken = agent + .edit_file( + agent_context(), + replace("utils.ar", "x1=17.", "x1=missing", "break it"), + ) + .await + .unwrap() + .expect("an edit that does not compile should still apply"); + assert!( + broken.diagnostics.iter().any(|diagnostic| { + diagnostic.path == std::path::Path::new("utils.ar") + && diagnostic.line == 2 + && diagnostic.message.contains("missing") + }), + "{:?}", + broken.diagnostics + ); + + let error = |result: Result| { + result.expect_err("the edit should be refused") + }; + let refused = |request| { + let agent = agent.clone(); + async move { error(agent.edit_file(agent_context(), request).await.unwrap()) } + }; + assert!( + refused(replace("lib.ar", "0.", "1.", "ambiguous")) + .await + .contains("occurs") + ); + assert!( + refused(replace("lib.ar", "nowhere", "x", "missing")) + .await + .contains("not found") + ); + assert!( + refused(replace("Argon.toml", "name", "x", "manifest")) + .await + .contains("not an Argon source file") + ); + assert!( + refused(replace("/elsewhere/lib.ar", "a", "b", "outside")) + .await + .contains("outside") + ); + let follow = agent + .open_cell(agent_context(), "top()".to_owned()) + .await + .unwrap() + .unwrap_err(); + assert_eq!(follow, analyzer::agent::FOLLOW_REQUIRED); + + std::fs::write(&session.ack, "ok\n").unwrap(); + finish_nvim(child).await; + }) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn approval_mode_asks_before_each_agent_edit() { + assert_completes("waiting for approved agent edits", async { + let _guard = FULL_STACK_LOCK.lock().await; + let mut session = Session::new(&one_rect()).await; + session.start_analyzer(); + let child = session.spawn_nvim("agent_approval"); + session.connect_gui().await; + session.wait_for_rects(1).await; + let agent = session.connect_agent().await; + let rect = |name: &str| { + add_rect( + &format!("let {name} = rect(\"met1\", x0=20., y0=0., x1=30., y1=10.);"), + name, + ) + }; + + session.await_signal("approval_on").await; + assert_eq!( + agent.status(agent_context()).await.unwrap().approval, + ApprovalMode::Always + ); + let rejected = agent + .edit_file(agent_context(), rect("rejected_rect")) + .await + .unwrap(); + assert_eq!( + rejected.unwrap_err(), + "The user rejected this edit in Neovim." + ); + agent + .edit_file(agent_context(), rect("approved_rect")) + .await + .unwrap() + .expect("an approved edit should apply"); + session.signal("approved"); + + session.await_signal("approval_off").await; + agent + .edit_file(agent_context(), rect("unprompted_rect")) + .await + .unwrap() + .expect("an edit should apply without approval"); + session.signal("unprompted"); + + std::fs::write(&session.ack, "ok\n").unwrap(); + finish_nvim(child).await; + }) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn agent_edits_survive_concurrent_typing() { + assert_completes("waiting for an interrupted agent edit", async { + let _guard = FULL_STACK_LOCK.lock().await; + let mut session = Session::new(&one_rect()).await; + session.start_analyzer(); + let child = session.spawn_nvim("agent_stale"); + session.connect_gui().await; + session.wait_for_rects(1).await; + let agent = session.connect_agent().await; + session.await_signal("typing_armed").await; + + agent + .edit_file(agent_context(), add_rect(SECOND_RECT, "add a second rect")) + .await + .unwrap() + .expect("the edit should be replanned after the buffer changed"); + let file = agent + .read_file(agent_context(), PathBuf::from("lib.ar")) + .await + .unwrap() + .unwrap(); + assert!(file.contents.contains(SECOND_RECT)); + assert!(file.contents.contains("// typed during the agent edit")); + + std::fs::write(&session.ack, "ok\n").unwrap(); + finish_nvim(child).await; + }) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn follow_mode_gates_cell_switching_and_tracks_edits() { + assert_completes("waiting for follow mode", async { + let _guard = FULL_STACK_LOCK.lock().await; + let mut session = Session::with_files(&with_unit(), &[("utils.ar", UNIT)]).await; + session.start_analyzer(); + let child = session.spawn_nvim("agent_follow"); + session.connect_gui().await; + session.wait_for_rects(1).await; + let agent = session.connect_agent().await; + + let refused = agent + .open_cell(agent_context(), "utils::unit()".to_owned()) + .await + .unwrap(); + assert!(refused.unwrap_err().contains("follow mode")); + + session.signal("follow"); + session.await_signal("following").await; + while !agent.status(agent_context()).await.unwrap().follow_mode { + time::sleep(std::time::Duration::from_millis(10)).await; + } + let opened = agent + .open_cell(agent_context(), "utils::unit()".to_owned()) + .await + .unwrap() + .expect("follow mode should allow switching the GUI's cell"); + assert_eq!( + opened.summary.unwrap().cell.as_deref(), + Some("utils::unit()") + ); + session + .wait_for_event(|event| matches!(event, GuiEvent::Fit)) + .await; + + agent + .edit_file( + agent_context(), + replace("utils.ar", "x1=5.", "x1=6.", "widen the unit"), + ) + .await + .unwrap() + .unwrap(); + session.signal("edited_followed"); + session.await_signal("followed").await; + + std::fs::write(&session.ack, "ok\n").unwrap(); + finish_nvim(child).await; + }) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn disk_changes_reach_neovim_and_the_gui() { + assert_completes("waiting for disk changes", async { + let _guard = FULL_STACK_LOCK.lock().await; + let mut session = Session::with_files(&with_unit(), &[("utils.ar", UNIT)]).await; + session.start_analyzer(); + let child = session.spawn_nvim("watcher"); + session.connect_gui().await; + session.wait_for_rects(1).await; + let agent = session.connect_agent().await; + while !agent.status(agent_context()).await.unwrap().watching_files { + time::sleep(std::time::Duration::from_millis(10)).await; + } + + // A closed file is reread from disk. + std::fs::write( + session.project().join("utils.ar"), + UNIT.replace("x1=5.", "x1=25."), + ) + .unwrap(); + loop { + let report = agent.diagnostics(agent_context()).await.unwrap(); + if report.summary.and_then(|summary| summary.bbox) == Some([0., 0., 25., 10.]) { + break; + } + time::sleep(std::time::Duration::from_millis(20)).await; + } + + // An open buffer without unsaved changes reloads. + let lib = session.project().join("lib.ar"); + let third = "let c = rect(\"met1\", x0=40., y0=0., x1=50., y1=10.);"; + std::fs::write( + &lib, + with_unit().replace(FIRST_RECT, &format!("{FIRST_RECT}\n {third}")), + ) + .unwrap(); + session.wait_for_rects(2).await; + + // An open buffer with unsaved changes keeps them. + session.await_signal("buffer_modified").await; + let fourth = "let d = rect(\"met1\", x0=60., y0=0., x1=70., y1=10.);"; + std::fs::write( + &lib, + with_unit().replace(FIRST_RECT, &format!("{FIRST_RECT}\n {fourth}")), + ) + .unwrap(); + session + .wait_for_event(|event| { + matches!( + event, + GuiEvent::Message { typ, message } + if *typ == MessageType::WARNING && message.contains("unsaved changes") + ) + }) + .await; + session.signal("rewritten"); + + std::fs::write(&session.ack, "ok\n").unwrap(); + finish_nvim(child).await; + }) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn mcp_bridge_and_hook_use_the_running_session() { + assert_completes("waiting for the MCP bridge", async { + let _guard = FULL_STACK_LOCK.lock().await; + let mut session = Session::new(&one_rect()).await; + session.start_analyzer(); + let child = session.spawn_nvim("mcp"); + session.connect_gui().await; + session.wait_for_rects(1).await; + + let bridge = analyzer::mcp::Bridge::new(session.project().to_path_buf()); + let call = |name: &str, arguments: Value| { + let request = json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { "name": name, "arguments": arguments }, + }); + let bridge = &bridge; + async move { + let response = bridge.handle(request).await.unwrap(); + let result = &response["result"]; + ( + result["content"][0]["text"].as_str().unwrap().to_owned(), + result["isError"].as_bool().unwrap(), + ) + } + }; + + let (status, failed) = call("status", json!({})).await; + assert!(!failed && status.contains("Open cell: top()"), "{status}"); + let (edited, failed) = call( + "edit_file", + json!({ + "path": "lib.ar", + "label": "add a second rect", + "edits": [{ + "old_string": FIRST_RECT, + "new_string": format!("{FIRST_RECT}\n {SECOND_RECT}"), + }], + }), + ) + .await; + assert!(!failed && edited.contains("2 rects"), "{edited}"); + let (read, failed) = call("read_file", json!({ "path": "lib.ar" })).await; + assert!(!failed && read.contains("from the Neovim buffer"), "{read}"); + assert!(read.contains("let b = rect("), "{read}"); + + let lib = session.project().join("lib.ar"); + let hook_input = |tool: &str, path: &std::path::Path| { + json!({ + "hook_event_name": "PreToolUse", + "tool_name": tool, + "tool_input": { "file_path": path }, + "cwd": session.project(), + }) + }; + let edit = analyzer::hook::decision(&hook_input("Edit", &lib)).await; + assert!(edit.is_some_and(|reason| reason.contains("edit_file"))); + let read = analyzer::hook::decision(&hook_input("Read", &lib)).await; + assert!(read.is_some_and(|reason| reason.contains("unsaved"))); + let notes = session.project().join("notes.md"); + assert_eq!( + analyzer::hook::decision(&hook_input("Edit", ¬es)).await, + None + ); + + session.signal("save"); + session.await_signal("saved").await; + assert_eq!( + analyzer::hook::decision(&hook_input("Read", &lib)).await, + None, + "a saved file can be read from disk" + ); + + std::fs::write(&session.ack, "ok\n").unwrap(); + finish_nvim(child).await; + }) + .await; + } } diff --git a/docs/developers.md b/docs/developers.md index d1e71111..0883e2b3 100644 --- a/docs/developers.md +++ b/docs/developers.md @@ -12,8 +12,14 @@ or newer must be available on `PATH` to run them locally. In normal use, Neovim launches the analyzer as an LSP child and communicates with it over standard input and output. The unit tests instead run the real analyzer library in-process and connect Neovim to its LSP stream over TCP using -`vim.lsp.rpc.connect()`. The analyzer's separate GUI-facing RPC port connects -to the headless GUI harness. +`vim.lsp.rpc.connect()`. The headless GUI harness and the agent tests connect +to the analyzer's separate RPC port. + +Every RPC connection opens with a handshake carrying the analyzer's session +token and the client's role (`Gui`, `GuiCallback`, or `Agent`), and is served +only once the token matches. The GUI opens both of its connections, so it never +listens on a port. Tests generate their own token and pass it to +`main_with_io_on_listener`. ## Debugging diff --git a/docs/docs/gui/shortcuts-config.md b/docs/docs/gui/shortcuts-config.md index c5a0e77b..2eafe19d 100644 --- a/docs/docs/gui/shortcuts-config.md +++ b/docs/docs/gui/shortcuts-config.md @@ -44,9 +44,12 @@ font_size = 14 [log] level = "info" + +[agent] +approval = "never" ``` -`font_size` and `icon_size` are in logical pixels, from 1 to 256. After editing the file, run `:Argon reload`. To change a value for this session only, use `:Argon set`; to write the current settings back to disk, use `:Argon saveConfig`. +`font_size` and `icon_size` are in logical pixels, from 1 to 256. `agent.approval` is `"never"` or `"always"`; see [Agents](/tools/agents#approving-edits). After editing the file, run `:Argon reload`. To change a value for this session only, use `:Argon set`; to write the current settings back to disk, use `:Argon saveConfig`. ## Troubleshooting diff --git a/docs/docs/tools/agents.md b/docs/docs/tools/agents.md new file mode 100644 index 00000000..f9558682 --- /dev/null +++ b/docs/docs/tools/agents.md @@ -0,0 +1,89 @@ +--- +title: Agents +description: Let a coding agent edit Argon source through your Neovim buffers while you watch. +sidebar_label: Agents +--- + +# Agents + +A coding agent can edit an Argon project while you have it open. Its edits go into your Neovim buffers, the same way GUI edits do, so you watch each one land in Neovim and in the GUI. Edits stay unsaved until you save them, and each one can be undone with `u`. + +Agents connect through `argon-analyzer mcp`, a [Model Context Protocol](https://modelcontextprotocol.io) server. It finds the analyzer that serves the agent's working directory, so start Argon first (for example with `argone`) and then the agent. + +## Setup + +Register the server with your agent. For Claude Code, run this in the project: + +```bash +claude mcp add --scope project argon -- argon-analyzer mcp +``` + +That writes `.mcp.json`: + +```json +{ + "mcpServers": { + "argon": { + "type": "stdio", + "command": "argon-analyzer", + "args": ["mcp"] + } + } +} +``` + +To search for a session from a directory other than the agent's working directory, pass `--root `. + +## Tools + +| Tool | Description | +| --- | --- | +| `status` | Workspace root, the cell open in the GUI, follow mode, the approval mode, and the files open in Neovim. | +| `read_file` | A file as you currently have it, unsaved edits included. | +| `edit_file` | Find-and-replace edits applied to a Neovim buffer as one undo step. Returns the recompiled cell's shape, bounding box, and diagnostics. | +| `create_file` | A new source file, opened as an unsaved buffer. | +| `diagnostics` | Current errors across the workspace. | +| `compile_cell` | Compiles a cell against the current source without changing what the GUI shows. | +| `open_cell` | Shows a cell in the GUI. Available only in follow mode. | + +The agent can edit `.ar` files in the workspace and in the libraries it depends on. A file that isn't open is loaded into a hidden buffer, which `:ls` lists. The agent never saves. + +## Watching edits + +Each agent edit briefly highlights the text it wrote with the `ArgonAgentEdit` group (linked to `IncSearch` by default). The GUI outlines the objects the edit changed, and its status bar shows what the agent is doing. LSP progress plugins such as Fidget show it in Neovim. + +`:Argon follow` makes the current window follow agent edits: it switches to each edited file and scrolls to the change, unless you are typing in it. `:Argon follow off`, or closing the window, turns it off. The agent can switch the GUI's cell with `open_cell` only while follow mode is on. + +## Approving edits + +By default edits apply immediately. To review each one first, set the approval mode: + +```toml +[agent] +approval = "always" +``` + +Neovim then shows each edit as a diff and asks before applying it. A rejected edit is reported to the agent and isn't retried. With `approval = "never"`, the default, Neovim is never asked. Change it for the current session with `:Argon set agent.approval always`. + +## Redirecting file tools + +An agent can still read and write `.ar` files with its own tools, bypassing your buffers. If a file changes on disk, the analyzer reloads it into Neovim when its buffer has no unsaved changes, and warns otherwise. + +For Claude Code, `argon-analyzer hook` is a `PreToolUse` hook that steers those tools to the MCP server while a session is running. It blocks edits to `.ar` files, and blocks reads when Neovim has unsaved changes the file on disk lacks. Add it to `.claude/settings.json`: + +```json +{ + "hooks": { + "PreToolUse": [ + { + "matcher": "Read|Edit|Write|MultiEdit", + "hooks": [{ "type": "command", "command": "argon-analyzer hook" }] + } + ] + } +} +``` + +## Security + +The analyzer listens on a local port, which the GUI and the MCP server connect to. Every connection must present a token that the analyzer generates at startup, so other users on the same machine can't connect to it. The token reaches the GUI through its environment, or over SSH for `argone ssh`. The MCP server reads it from a session record in `$XDG_STATE_HOME/argon/sessions` (`~/.local/state/argon/sessions` by default), which only you can read. diff --git a/docs/docs/tools/argone.md b/docs/docs/tools/argone.md index a980b4eb..9dbf0e10 100644 --- a/docs/docs/tools/argone.md +++ b/docs/docs/tools/argone.md @@ -38,8 +38,6 @@ argone ssh [PATH] [OPTIONS] | `-o, --ssh-option