Fix types/provide-type resolving the wrong token in notebook cells

pyrefly/lib/lsp/non_wasm/server.rs

@@ -2853,7 +2853,8 @@ impl Server {
     ) -> Option<ProvideTypeResponse> {
         let uri = &params.text_document.uri;
         let handle = self.make_handle_if_enabled(uri, None).ok()?;
-        provide_type(transaction, &handle, params.positions)
+        let notebook_cell = self.maybe_get_code_cell_index(uri);
+        provide_type(transaction, &handle, params.positions, notebook_cell)
     }
 
     fn type_error_display_status(&self, path: &Path) -> TypeErrorDisplayStatus {

pyrefly/lib/lsp/wasm/provide_type.rs

@@ -48,6 +48,7 @@ pub fn provide_type(
     transaction: &mut Transaction<'_>,
     handle: &Handle,
     positions: Vec<Position>,
+    notebook_cell: Option<usize>,
 ) -> Option<ProvideTypeResponse> {
     // This LSP method works for unopened files.
     // Check if the file is already loaded in memory. If not, load it.
@@ -59,7 +60,7 @@ pub fn provide_type(
     let mut contents = Vec::new();
 
     for position in positions {
-        let text_size = info.from_lsp_position(position, None);
+        let text_size = info.from_lsp_position(position, notebook_cell);
         if let Some(ty) = transaction.get_result_type_at_for_display(handle, text_size) {
             let mut c = TypeDisplayContext::new(&[&ty]);
             c.set_lsp_display_mode(LspDisplayMode::ProvideType);

pyrefly/lib/test/lsp/lsp_interaction/notebook_provide_type.rs

@@ -37,3 +37,33 @@ fn test_notebook_provide_type() {
 
     interaction.shutdown().unwrap();
 }
+
+#[test]
+fn test_notebook_provide_type_second_cell() {
+    let root = get_test_files_root();
+    let mut interaction = LspInteraction::new();
+    interaction.set_root(root.path().to_path_buf());
+    interaction
+        .initialize(InitializeSettings {
+            configuration: Some(None),
+            ..Default::default()
+        })
+        .unwrap();
+    interaction.open_notebook(
+        "notebook.ipynb",
+        vec!["foo = \"bar\"\nfoo", "x = 1\nbar = x"],
+    );
+
+    // `bar` is at cell2 line 1; its type is the value of `x`.
+    interaction
+        .provide_type_cell("notebook.ipynb", "cell2", 1, 0)
+        .expect_response(json!({
+            "contents": [{
+                "kind": "plaintext",
+                "value": "typing.Literal[1]",
+            }]
+        }))
+        .unwrap();
+
+    interaction.shutdown().unwrap();
+}