Annotate execute request source code with code location

Cargo.toml

@@ -26,6 +26,7 @@ authors = ["Posit Software, PBC"]
 
 [workspace.dependencies]
 biome_line_index = { git = "https://github.com/biomejs/biome", rev = "c13fc60726883781e4530a4437724273b560c8e0" }
+biome_rowan = { git = "https://github.com/biomejs/biome", rev = "c13fc60726883781e4530a4437724273b560c8e0" }
 aether_lsp_utils = { git = "https://github.com/posit-dev/air", rev = "f959e32eee91" }
 aether_parser = { git = "https://github.com/posit-dev/air", package = "air_r_parser", rev = "f959e32eee91" }
 aether_syntax = { git = "https://github.com/posit-dev/air", package = "air_r_syntax", rev = "f959e32eee91" }

crates/amalthea/Cargo.toml

@@ -34,6 +34,7 @@ serde_with = "3.0.0"
 serde_repr = "0.1.17"
 tracing = "0.1.40"
 assert_matches = "1.5.0"
+url = "2.5.7"
 
 [dev-dependencies]
 env_logger = "0.10.0"

crates/amalthea/src/fixtures/dummy_frontend.rs

@@ -15,6 +15,8 @@ use crate::session::Session;
 use crate::socket::socket::Socket;
 use crate::wire::execute_input::ExecuteInput;
 use crate::wire::execute_request::ExecuteRequest;
+use crate::wire::execute_request::ExecuteRequestPositron;
+use crate::wire::execute_request::JupyterPositronLocation;
 use crate::wire::handshake_reply::HandshakeReply;
 use crate::wire::input_reply::InputReply;
 use crate::wire::jupyter_message::JupyterMessage;
@@ -48,6 +50,7 @@ pub struct DummyFrontend {
 
 pub struct ExecuteRequestOptions {
     pub allow_stdin: bool,
+    pub positron: Option<ExecuteRequestPositron>,
 }
 
 impl DummyConnection {
@@ -233,6 +236,7 @@ impl DummyFrontend {
             user_expressions: serde_json::Value::Null,
             allow_stdin: options.allow_stdin,
             stop_on_error: false,
+            positron: options.positron,
         })
     }
 
@@ -264,7 +268,38 @@ impl DummyFrontend {
     where
         F: FnOnce(String),
     {
-        self.send_execute_request(code, ExecuteRequestOptions::default());
+        self.execute_request_with_options(code, result_check, Default::default())
+    }
+
+    #[track_caller]
+    pub fn execute_request_with_location<F>(
+        &self,
+        code: &str,
+        result_check: F,
+        code_location: JupyterPositronLocation,
+    ) -> u32
+    where
+        F: FnOnce(String),
+    {
+        self.execute_request_with_options(code, result_check, ExecuteRequestOptions {
+            positron: Some(ExecuteRequestPositron {
+                code_location: Some(code_location),
+            }),
+            ..Default::default()
+        })
+    }
+
+    #[track_caller]
+    pub fn execute_request_with_options<F>(
+        &self,
+        code: &str,
+        result_check: F,
+        options: ExecuteRequestOptions,
+    ) -> u32
+    where
+        F: FnOnce(String),
+    {
+        self.send_execute_request(code, options);
         self.recv_iopub_busy();
 
         let input = self.recv_iopub_execute_input();
@@ -663,6 +698,9 @@ impl DummyFrontend {
 
 impl Default for ExecuteRequestOptions {
     fn default() -> Self {
-        Self { allow_stdin: false }
+        Self {
+            allow_stdin: false,
+            positron: None,
+        }
     }
 }

crates/amalthea/src/wire/execute_request.rs

@@ -5,9 +5,11 @@
  *
  */
 
+use anyhow::Context;
 use serde::Deserialize;
 use serde::Serialize;
 use serde_json::Value;
+use url::Url;
 
 use crate::wire::jupyter_message::MessageType;
 
@@ -33,6 +35,157 @@ pub struct ExecuteRequest {
     /// Whether the kernel should discard the execution queue if evaluating the
     /// code results in an error
     pub stop_on_error: bool,
+
+    /// Posit extension
+    pub positron: Option<ExecuteRequestPositron>,
+}
+
+#[serde_with::skip_serializing_none]
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct ExecuteRequestPositron {
+    pub code_location: Option<JupyterPositronLocation>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct JupyterPositronLocation {
+    pub uri: String,
+    pub range: JupyterPositronRange,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct JupyterPositronRange {
+    pub start: JupyterPositronPosition,
+    pub end: JupyterPositronPosition,
+}
+
+/// See https://jupyter-client.readthedocs.io/en/stable/messaging.html#cursor-pos-unicode-note
+/// regarding choice of offset in unicode points
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct JupyterPositronPosition {
+    pub line: u32,
+    /// Column offset in unicode points
+    pub character: u32,
+}
+
+/// Code location with `character` in UTF-8 offset
+#[derive(Debug, Clone)]
+pub struct CodeLocation {
+    pub uri: Url,
+    pub start: Position,
+    pub end: Position,
+}
+
+/// `character` in UTF-8 offset
+#[derive(Debug, Clone)]
+pub struct Position {
+    pub line: u32,
+    pub character: usize,
+}
+
+impl ExecuteRequest {
+    pub fn code_location(&self) -> anyhow::Result<Option<CodeLocation>> {
+        let Some(positron) = &self.positron else {
+            return Ok(None);
+        };
+
+        let Some(location) = &positron.code_location else {
+            return Ok(None);
+        };
+
+        let uri = Url::parse(&location.uri).context("Failed to parse URI from code location")?;
+
+        // The location maps `self.code` to a range in the document. We'll first
+        // do a sanity check that the span dimensions (end - start) match the
+        // code extents.
+        let span_lines = location.range.end.line - location.range.start.line;
+
+        // For multiline code, the last line's expected length is just `end.character`.
+        // For single-line code, the expected length is `end.character - start.character`.
+        let expected_last_line_chars = if span_lines == 0 {
+            location.range.end.character - location.range.start.character
+        } else {
+            location.range.end.character
+        };
+
+        let code_lines: Vec<&str> = self.code.lines().collect();
+        let code_line_count = code_lines.len().saturating_sub(1);
+
+        // Sanity check: `code` conforms exactly to expected number of lines in the span
+        if code_line_count != span_lines as usize {
+            return Err(anyhow::anyhow!(
+                "Line information does not match code line count (expected {}, got {})",
+                code_line_count,
+                span_lines
+            ));
+        }
+
+        let last_line_idx = code_lines.len().saturating_sub(1);
+        let last_line = code_lines.get(last_line_idx).unwrap_or(&"");
+        let last_line = last_line.strip_suffix('\r').unwrap_or(last_line);
+        let last_line_chars = last_line.chars().count() as u32;
+
+        // Sanity check: the last line has exactly the expected number of characters
+        if last_line_chars != expected_last_line_chars {
+            return Err(anyhow::anyhow!(
+                "Expected last line to have {expected} characters, got {actual}",
+                expected = expected_last_line_chars,
+                actual = last_line_chars
+            ));
+        }
+
+        // Convert start character from unicode code points to UTF-8 bytes
+        let character_start =
+            unicode_char_to_utf8_offset(&self.code, 0, location.range.start.character)?;
+
+        // End character is start + last line byte length (for single line)
+        // or just last line byte length (for multiline, since it's on a new line)
+        let last_line_bytes = last_line.len();
+        let character_end = if span_lines == 0 {
+            character_start + last_line_bytes
+        } else {
+            last_line_bytes
+        };
+
+        let start = Position {
+            line: location.range.start.line,
+            character: character_start,
+        };
+        let end = Position {
+            line: location.range.end.line,
+            character: character_end,
+        };
+
+        Ok(Some(CodeLocation { uri, start, end }))
+    }
+}
+
+/// Converts a character position in unicode scalar values to a UTF-8 byte
+/// offset within the specified line.
+fn unicode_char_to_utf8_offset(text: &str, line: u32, character: u32) -> anyhow::Result<usize> {
+    let target_line = text
+        .lines()
+        .nth(line as usize)
+        .ok_or_else(|| anyhow::anyhow!("Line {line} not found in text"))?;
+
+    unicode_char_to_utf8_offset_in_line(target_line, character)
+}
+
+/// Converts a character count in unicode scalar values to a UTF-8 byte count.
+fn unicode_char_to_utf8_offset_in_line(line: &str, character: u32) -> anyhow::Result<usize> {
+    let line_chars = line.chars().count();
+    if character as usize > line_chars {
+        return Err(anyhow::anyhow!(
+            "Character position {character} exceeds line length ({line_chars})"
+        ));
+    }
+
+    let byte_offset = line
+        .char_indices()
+        .nth(character as usize)
+        .map(|(byte_idx, _)| byte_idx)
+        .unwrap_or(line.len());
+
+    Ok(byte_offset)
 }
 
 impl MessageType for ExecuteRequest {

crates/ark/Cargo.toml

@@ -17,6 +17,7 @@ anyhow = "1.0.80"
 async-trait = "0.1.66"
 base64 = "0.21.0"
 biome_line_index.workspace = true
+biome_rowan.workspace = true
 bus = "2.3.0"
 cfg-if = "1.0.0"
 crossbeam = { version = "0.8.2", features = ["crossbeam-channel"] }