[ty] Fix outdated version in publish diagnostics after `didChange`

crates/ty_server/src/server/api/notifications/did_change.rs

@@ -26,7 +26,7 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             content_changes,
         } = params;
 
-        let document = session
+        let mut document = session
             .document_handle(&uri)
             .with_failure_code(ErrorCode::InternalError)?;
 

crates/ty_server/src/server/api/notifications/did_change_notebook.rs

@@ -24,7 +24,7 @@ impl SyncNotificationHandler for DidChangeNotebookHandler {
             change: types::NotebookDocumentChangeEvent { cells, metadata },
         }: types::DidChangeNotebookDocumentParams,
     ) -> Result<()> {
-        let document = session
+        let mut document = session
             .document_handle(&uri)
             .with_failure_code(ErrorCode::InternalError)?;
 

crates/ty_server/src/session.rs

@@ -1451,7 +1451,7 @@ impl DocumentHandle {
     }
 
     pub(crate) fn update_text_document(
-        &self,
+        &mut self,
         session: &mut Session,
         content_changes: Vec<TextDocumentContentChangeEvent>,
         new_version: DocumentVersion,
@@ -1471,6 +1471,8 @@ impl DocumentHandle {
             } else {
                 document.apply_changes(content_changes, new_version, position_encoding);
             }
+
+            self.set_version(document.version());
         }
 
         self.update_in_db(session);
@@ -1479,7 +1481,7 @@ impl DocumentHandle {
     }
 
     pub(crate) fn update_notebook_document(
-        &self,
+        &mut self,
         session: &mut Session,
         cells: Option<lsp_types::NotebookDocumentCellChange>,
         metadata: Option<lsp_types::LSPObject>,
@@ -1496,6 +1498,8 @@ impl DocumentHandle {
                 new_version,
                 position_encoding,
             )?;
+
+            self.set_version(new_version);
         }
 
         self.update_in_db(session);
@@ -1516,6 +1520,16 @@ impl DocumentHandle {
         session.apply_changes(path, changes);
     }
 
+    fn set_version(&mut self, version: DocumentVersion) {
+        let self_version = match self {
+            DocumentHandle::Text { version, .. }
+            | DocumentHandle::Notebook { version, .. }
+            | DocumentHandle::Cell { version, .. } => version,
+        };
+
+        *self_version = version;
+    }
+
     /// De-registers a document, specified by its key.
     /// Calling this multiple times for the same document is a logic error.
     ///

crates/ty_server/tests/e2e/publish_diagnostics.rs

@@ -33,6 +33,42 @@ def foo() -> str:
     Ok(())
 }
 
+#[test]
+fn on_did_change() -> Result<()> {
+    let workspace_root = SystemPath::new("src");
+    let foo = SystemPath::new("src/foo.py");
+    let foo_content = "\
+def foo() -> str:
+    return 42
+";
+
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(workspace_root, None)?
+        .with_file(foo, foo_content)?
+        .enable_pull_diagnostics(false)
+        .build()
+        .wait_until_workspaces_are_initialized();
+
+    server.open_text_document(foo, foo_content, 1);
+    let _ = server.await_notification::<PublishDiagnostics>();
+
+    let changes = vec![lsp_types::TextDocumentContentChangeEvent {
+        range: None,
+        range_length: None,
+        text: "def foo() -> int: return 42".to_string(),
+    }];
+
+    server.change_text_document(foo, changes, 2);
+
+    let diagnostics = server.await_notification::<PublishDiagnostics>();
+
+    assert_eq!(diagnostics.version, Some(2));
+
+    insta::assert_debug_snapshot!(diagnostics);
+
+    Ok(())
+}
+
 #[test]
 fn message_without_related_information_support() -> Result<()> {
     let workspace_root = SystemPath::new("src");

crates/ty_server/tests/e2e/snapshots/e2e__publish_diagnostics__on_did_change.snap

@@ -0,0 +1,21 @@
+---
+source: crates/ty_server/tests/e2e/publish_diagnostics.rs
+expression: diagnostics
+---
+PublishDiagnosticsParams {
+    uri: Url {
+        scheme: "file",
+        cannot_be_a_base: false,
+        username: "",
+        password: None,
+        host: None,
+        port: None,
+        path: "<temp_dir>/src/foo.py",
+        query: None,
+        fragment: None,
+    },
+    diagnostics: [],
+    version: Some(
+        2,
+    ),
+}