feat: add code lens support for base image scanning

src/app/commands.rs

@@ -35,8 +35,8 @@ where
     C: LSPClient,
 {
     pub async fn update_document_with_text(&self, uri: &str, text: &str) {
-        self.document_database.remove_document(uri).await;
         self.document_database.write_document_text(uri, text).await;
+        self.document_database.remove_diagnostics(uri).await;
         let _ = self.publish_all_diagnostics().await;
     }
 

src/app/lsp_server.rs

@@ -6,10 +6,10 @@ use tower_lsp::LanguageServer;
 use tower_lsp::jsonrpc::{Error, ErrorCode, Result};
 use tower_lsp::lsp_types::{
     CodeActionOrCommand, CodeActionParams, CodeActionProviderCapability, CodeActionResponse,
-    Command, DidChangeConfigurationParams, DidChangeTextDocumentParams, DidOpenTextDocumentParams,
-    ExecuteCommandOptions, ExecuteCommandParams, InitializeParams, InitializeResult,
-    InitializedParams, MessageType, ServerCapabilities, TextDocumentSyncCapability,
-    TextDocumentSyncKind,
+    CodeLens, CodeLensOptions, CodeLensParams, Command, DidChangeConfigurationParams,
+    DidChangeTextDocumentParams, DidOpenTextDocumentParams, ExecuteCommandOptions,
+    ExecuteCommandParams, InitializeParams, InitializeResult, InitializedParams, MessageType,
+    Position, Range, ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind,
 };
 use tracing::{debug, info};
 
@@ -106,6 +106,9 @@ where
                     TextDocumentSyncKind::FULL,
                 )),
                 code_action_provider: Some(CodeActionProviderCapability::Simple(true)),
+                code_lens_provider: Some(CodeLensOptions {
+                    resolve_provider: Some(false),
+                }),
                 execute_command_provider: Some(ExecuteCommandOptions {
                     commands: vec![SupportedCommands::ExecuteScan.to_string()],
                     ..Default::default()
@@ -180,7 +183,7 @@ where
 
         if last_line_starting_with_from_statement == line_selected_as_usize {
             let action = Command {
-                title: "Scan Image".to_string(),
+                title: "Scan base image".to_string(),
                 command: SupportedCommands::ExecuteScan.to_string(),
                 arguments: Some(vec![
                     json!(params.text_document.uri),
@@ -194,6 +197,58 @@ where
         return Ok(None);
     }
 
+    async fn code_lens(&self, params: CodeLensParams) -> Result<Option<Vec<CodeLens>>> {
+        info!("{}", format!("received code lens params: {params:?}"));
+
+        let Some(content) = self
+            .query_executor
+            .get_document_text(params.text_document.uri.as_str())
+            .await
+        else {
+            return Err(lsp_error(
+                ErrorCode::InternalError,
+                format!(
+                    "unable to extract document content for document: {}",
+                    &params.text_document.uri
+                ),
+            ));
+        };
+
+        let Some(last_line_starting_with_from_statement) = content
+            .lines()
+            .enumerate()
+            .filter(|(_, line)| line.trim_start().starts_with("FROM "))
+            .map(|(line_num, _)| line_num)
+            .last()
+        else {
+            return Ok(None);
+        };
+
+        let scan_base_image_lens = CodeLens {
+            range: Range {
+                start: Position {
+                    line: last_line_starting_with_from_statement as u32,
+                    character: 0,
+                },
+                end: Position {
+                    line: last_line_starting_with_from_statement as u32,
+                    character: 0,
+                },
+            },
+            command: Some(Command {
+                title: "Scan base image".to_string(),
+                command: SupportedCommands::ExecuteScan.to_string(),
+                arguments: Some(vec![
+                    json!(params.text_document.uri),
+                    json!(last_line_starting_with_from_statement),
+                ]),
+            }),
+            data: None,
+        };
+
+        Ok(Some(vec![scan_base_image_lens]))
+    }
+
     async fn execute_command(&self, params: ExecuteCommandParams) -> Result<Option<Value>> {
         let command: SupportedCommands = params.command.as_str().try_into().map_err(|e| {
             lsp_error(

tests/general.rs

@@ -1,5 +1,5 @@
 use serde_json::json;
-use tower_lsp::lsp_types::{CodeActionOrCommand, Command, MessageType};
+use tower_lsp::lsp_types::{CodeActionOrCommand, CodeLens, Command, MessageType, Position, Range};
 
 mod test;
 
@@ -34,7 +34,7 @@ async fn when_the_client_asks_for_the_existing_code_actions_it_receives_the_avai
     assert_eq!(
         response.unwrap(),
         vec![CodeActionOrCommand::Command(Command {
-            title: "Scan Image".to_string(),
+            title: "Scan base image".to_string(),
             command: "sysdig-lsp.execute-scan".to_string(),
             arguments: Some(vec![json!("file://dockerfile/"), json!(0)])
         })]
@@ -62,9 +62,82 @@ async fn when_the_client_asks_for_the_existing_code_actions_but_the_dockerfile_c
     assert_eq!(
         response_for_second_line.unwrap(),
         vec![CodeActionOrCommand::Command(Command {
-            title: "Scan Image".to_string(),
+            title: "Scan base image".to_string(),
             command: "sysdig-lsp.execute-scan".to_string(),
             arguments: Some(vec![json!("file://dockerfile/"), json!(1)])
         })]
     );
 }
+
+#[tokio::test]
+async fn when_the_client_asks_for_the_existing_code_lens_it_receives_the_available_code_lens() {
+    let mut client = test::TestClient::new_initialized().await;
+
+    // Open a Dockerfile containing a single "FROM" statement.
+    client
+        .open_file_with_contents("Dockerfile", "FROM alpine")
+        .await;
+
+    // Request code lens on the line with the FROM statement (line 0).
+    let response = client
+        .request_available_code_lens_in_file("Dockerfile")
+        .await;
+
+    // Expect a CodeLens with the appropriate command.
+    assert_eq!(
+        response.unwrap(),
+        vec![CodeLens {
+            range: Range {
+                start: Position {
+                    line: 0,
+                    character: 0
+                },
+                end: Position {
+                    line: 0,
+                    character: 0
+                }
+            },
+            command: Some(Command {
+                title: "Scan base image".to_string(),
+                command: "sysdig-lsp.execute-scan".to_string(),
+                arguments: Some(vec![json!("file://dockerfile/"), json!(0)])
+            }),
+            data: None
+        }]
+    );
+}
+
+#[tokio::test]
+async fn when_the_client_asks_for_the_existing_code_lens_but_the_dockerfile_contains_multiple_froms_it_only_returns_the_latest()
+ {
+    let mut client = test::TestClient::new_initialized().await;
+    client
+        .open_file_with_contents("Dockerfile", "FROM alpine\nFROM ubuntu")
+        .await;
+
+    let response = client
+        .request_available_code_lens_in_file("Dockerfile")
+        .await;
+
+    assert_eq!(
+        response.unwrap(),
+        vec![CodeLens {
+            range: Range {
+                start: Position {
+                    line: 1,
+                    character: 0
+                },
+                end: Position {
+                    line: 1,
+                    character: 0
+                }
+            },
+            command: Some(Command {
+                title: "Scan base image".to_string(),
+                command: "sysdig-lsp.execute-scan".to_string(),
+                arguments: Some(vec![json!("file://dockerfile/"), json!(1)])
+            }),
+            data: None
+        }]
+    );
+}

tests/test.rs

@@ -8,9 +8,9 @@ use sysdig_lsp::app::{LSPClient, LSPServer};
 use tokio::sync::Mutex;
 use tower_lsp::LanguageServer;
 use tower_lsp::lsp_types::{
-    CodeActionOrCommand, CodeActionParams, Diagnostic, DidOpenTextDocumentParams, InitializeParams,
-    InitializeResult, InitializedParams, MessageType, Position, Range, TextDocumentIdentifier,
-    TextDocumentItem, Url,
+    CodeActionOrCommand, CodeActionParams, CodeLens, CodeLensParams, Diagnostic,
+    DidOpenTextDocumentParams, InitializeParams, InitializeResult, InitializedParams, MessageType,
+    Position, Range, TextDocumentIdentifier, TextDocumentItem, Url,
 };
 
 pub struct TestClient {
@@ -86,11 +86,24 @@ impl TestClient {
             .await
             .unwrap_or_else(|_| {
                 panic!(
-                    "unable to send code action for filename {} in line number {}",
-                    filename, line_number
+                    "unable to send code action for filename {filename} in line number {line_number}",
                 )
             })
     }
+
+    pub async fn request_available_code_lens_in_file(
+        &mut self,
+        filename: &str,
+    ) -> Option<Vec<CodeLens>> {
+        self.server
+            .code_lens(CodeLensParams {
+                text_document: TextDocumentIdentifier::new(url_from(filename)),
+                work_done_progress_params: Default::default(),
+                partial_result_params: Default::default(),
+            })
+            .await
+            .unwrap_or_else(|_| panic!("unable to send code lens for filename {filename}"))
+    }
 }
 
 fn url_from(filename: &str) -> Url {