From 2eddfbd7bf4c47f054a8b5a317521d7e4a836fe8 Mon Sep 17 00:00:00 2001
From: DavidLoftus <43099493+DavidLoftus@users.noreply.github.com>
Date: Tue, 11 Feb 2025 00:39:49 +0100
Subject: [PATCH] Add restart LSP command to VSCode extension (#4916)

This PR adds a new command to the Carbon VSCode extension "carbon:
Restart language server" which acts similar to clangd / bazel
extensions. While not useful during regular working of the extension
this helps when LSP either crashes or the toolchain is updated.

In order to support graceful restart, I needed to add a Call handler for
"shutdown" which is noop for now. Lsp clients always call shutdown
before sending exit.

---------

Co-authored-by: jonmeow <jperkins@google.com>
---
 toolchain/language_server/handle.h            |  6 ++++
 toolchain/language_server/handle_shutdown.cpp | 21 +++++++++++
 .../language_server/incoming_messages.cpp     |  1 +
 .../basics/fail_shutdown_without_exit.carbon  | 35 +++++++++++++++++++
 .../testdata/basics/initialize.carbon         |  7 ++++
 .../testdata/basics/notify_parse_error.carbon | 21 ++++++++++-
 .../testdata/document_symbol/basics.carbon    | 27 +++++++++++---
 .../testdata/document_symbol/language.carbon  | 21 ++++++++++-
 .../testdata/document_symbol/unknown.carbon   | 21 ++++++++++-
 .../text_document/change_count.carbon         | 21 ++++++++++-
 .../text_document/change_unknown.carbon       | 21 ++++++++++-
 .../text_document/close_unknown.carbon        | 21 ++++++++++-
 .../testdata/text_document/diagnostics.carbon | 21 ++++++++++-
 .../text_document/open_change_close.carbon    | 21 ++++++++++-
 .../text_document/open_duplicate.carbon       | 21 ++++++++++-
 utils/vscode/package.json                     |  8 ++++-
 utils/vscode/src/extension.ts                 |  8 ++++-
 17 files changed, 287 insertions(+), 15 deletions(-)
 create mode 100644 toolchain/language_server/handle_shutdown.cpp
 create mode 100644 toolchain/language_server/testdata/basics/fail_shutdown_without_exit.carbon

diff --git a/toolchain/language_server/handle.h b/toolchain/language_server/handle.h
index 4161f20f4d06..08ba21d2d208 100644
--- a/toolchain/language_server/handle.h
+++ b/toolchain/language_server/handle.h
@@ -39,6 +39,12 @@ auto HandleInitialize(
     llvm::function_ref<void(llvm::Expected<llvm::json::Object>)> on_done)
     -> void;
 
+// Prepares LSP for shutdown.
+auto HandleShutdown(
+    Context& /*context*/,
+    const clang::clangd::NoParams& /*client_capabilities*/,
+    llvm::function_ref<void(llvm::Expected<std::nullptr_t>)> on_done) -> void;
+
 }  // namespace Carbon::LanguageServer
 
 #endif  // CARBON_TOOLCHAIN_LANGUAGE_SERVER_HANDLE_H_
diff --git a/toolchain/language_server/handle_shutdown.cpp b/toolchain/language_server/handle_shutdown.cpp
new file mode 100644
index 000000000000..269bd76c2e4d
--- /dev/null
+++ b/toolchain/language_server/handle_shutdown.cpp
@@ -0,0 +1,21 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "toolchain/language_server/handle.h"
+
+namespace Carbon::LanguageServer {
+
+auto HandleShutdown(
+    Context& /*context*/,
+    const clang::clangd::NoParams& /*client_capabilities*/,
+    llvm::function_ref<void(llvm::Expected<std::nullptr_t>)> on_done) -> void {
+  // TODO: Track that `shutdown` was called, and:
+  // - Warn on duplicate calls.
+  // - Make `exit` return `1` if `shutdown` wasn't called.
+  //   https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#exit
+  // - Error on other post-`shutdown` calls.
+  on_done(nullptr);
+}
+
+}  // namespace Carbon::LanguageServer
diff --git a/toolchain/language_server/incoming_messages.cpp b/toolchain/language_server/incoming_messages.cpp
index 70d404671f0a..42a733a5bf08 100644
--- a/toolchain/language_server/incoming_messages.cpp
+++ b/toolchain/language_server/incoming_messages.cpp
@@ -74,6 +74,7 @@ IncomingMessages::IncomingMessages(clang::clangd::Transport* transport,
     : transport_(transport), context_(context) {
   AddCallHandler("textDocument/documentSymbol", &HandleDocumentSymbol);
   AddCallHandler("initialize", &HandleInitialize);
+  AddCallHandler("shutdown", &HandleShutdown);
   AddNotificationHandler("textDocument/didChange",
                          &HandleDidChangeTextDocument);
   AddNotificationHandler("textDocument/didClose", &HandleDidCloseTextDocument);
diff --git a/toolchain/language_server/testdata/basics/fail_shutdown_without_exit.carbon b/toolchain/language_server/testdata/basics/fail_shutdown_without_exit.carbon
new file mode 100644
index 000000000000..b9e4d21de627
--- /dev/null
+++ b/toolchain/language_server/testdata/basics/fail_shutdown_without_exit.carbon
@@ -0,0 +1,35 @@
+// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
+// Exceptions. See /LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// AUTOUPDATE
+// TIP: To test this file alone, run:
+// TIP:   bazel test //toolchain/testing:file_test --test_arg=--file_tests=toolchain/language_server/testdata/basics/fail_shutdown_without_exit.carbon
+// TIP: To dump output, run:
+// TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/basics/fail_shutdown_without_exit.carbon
+
+// --- STDIN
+[[@LSP-CALL:initialize]]
+[[@LSP-CALL:shutdown]]
+// --- AUTOUPDATE-SPLIT
+
+// CHECK:STDERR: error: Input/output error [LanguageServerTransportError]
+// CHECK:STDERR:
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
+// CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/basics/initialize.carbon b/toolchain/language_server/testdata/basics/initialize.carbon
index 4129c012b990..899ccb163f69 100644
--- a/toolchain/language_server/testdata/basics/initialize.carbon
+++ b/toolchain/language_server/testdata/basics/initialize.carbon
@@ -10,6 +10,7 @@
 
 // --- STDIN
 [[@LSP-CALL:initialize]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
@@ -25,4 +26,10 @@
 // CHECK:STDOUT:       "textDocumentSync": 1
 // CHECK:STDOUT:     }
 // CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
 // CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/basics/notify_parse_error.carbon b/toolchain/language_server/testdata/basics/notify_parse_error.carbon
index 088a963a5baf..0ee79ac29449 100644
--- a/toolchain/language_server/testdata/basics/notify_parse_error.carbon
+++ b/toolchain/language_server/testdata/basics/notify_parse_error.carbon
@@ -10,11 +10,30 @@
 
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-NOTIFY:textDocument/didOpen:"nope": "bad"]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
 // CHECK:STDERR: warning: -32602: in call to `textDocument/didOpen`, JSON parse failed: missing value at (root).textDocument [LanguageServerNotificationParseError]
 // CHECK:STDERR:
-// CHECK:STDOUT:
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
+// CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/document_symbol/basics.carbon b/toolchain/language_server/testdata/document_symbol/basics.carbon
index 3f566439f957..af54bd5becdb 100644
--- a/toolchain/language_server/testdata/document_symbol/basics.carbon
+++ b/toolchain/language_server/testdata/document_symbol/basics.carbon
@@ -9,6 +9,7 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/document_symbol/basics.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-NOTIFY:textDocument/didOpen:
   "textDocument": {"uri": "file:/empty.carbon", "languageId": "carbon",
                    "text": ""}
@@ -30,11 +31,23 @@
 [[@LSP-CALL:textDocument/documentSymbol:
   "textDocument": {"uri": "file:/fn.carbon"}
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
-// CHECK:STDOUT: Content-Length: 145{{\r}}
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 145{{\r}}
 // CHECK:STDOUT: {{\r}}
 // CHECK:STDOUT: {
 // CHECK:STDOUT:   "jsonrpc": "2.0",
@@ -46,7 +59,7 @@
 // CHECK:STDOUT: }Content-Length: 49{{\r}}
 // CHECK:STDOUT: {{\r}}
 // CHECK:STDOUT: {
-// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "id": 2,
 // CHECK:STDOUT:   "jsonrpc": "2.0",
 // CHECK:STDOUT:   "result": []
 // CHECK:STDOUT: }Content-Length: 147{{\r}}
@@ -61,7 +74,7 @@
 // CHECK:STDOUT: }Content-Length: 49{{\r}}
 // CHECK:STDOUT: {{\r}}
 // CHECK:STDOUT: {
-// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "id": 3,
 // CHECK:STDOUT:   "jsonrpc": "2.0",
 // CHECK:STDOUT:   "result": []
 // CHECK:STDOUT: }Content-Length: 142{{\r}}
@@ -76,7 +89,7 @@
 // CHECK:STDOUT: }Content-Length: 459{{\r}}
 // CHECK:STDOUT: {{\r}}
 // CHECK:STDOUT: {
-// CHECK:STDOUT:   "id": 3,
+// CHECK:STDOUT:   "id": 4,
 // CHECK:STDOUT:   "jsonrpc": "2.0",
 // CHECK:STDOUT:   "result": [
 // CHECK:STDOUT:     {
@@ -104,4 +117,10 @@
 // CHECK:STDOUT:       }
 // CHECK:STDOUT:     }
 // CHECK:STDOUT:   ]
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 5,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
 // CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/document_symbol/language.carbon b/toolchain/language_server/testdata/document_symbol/language.carbon
index 0be2b96e7735..00ed58d25c70 100644
--- a/toolchain/language_server/testdata/document_symbol/language.carbon
+++ b/toolchain/language_server/testdata/document_symbol/language.carbon
@@ -9,13 +9,32 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/document_symbol/language.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-CALL:textDocument/documentSymbol:
   "textDocument": {"uri": "file:/test.cpp"}
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
 // CHECK:STDERR: /test.cpp: warning: non-Carbon file requested [LanguageServerFileUnsupported]
 // CHECK:STDERR:
-// CHECK:STDOUT:
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 3,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
+// CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/document_symbol/unknown.carbon b/toolchain/language_server/testdata/document_symbol/unknown.carbon
index 2fca57d2e68c..017efd915a55 100644
--- a/toolchain/language_server/testdata/document_symbol/unknown.carbon
+++ b/toolchain/language_server/testdata/document_symbol/unknown.carbon
@@ -9,13 +9,32 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/document_symbol/unknown.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-CALL:textDocument/documentSymbol:
   "textDocument": {"uri": "file:/test.carbon"}
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
 // CHECK:STDERR: /test.carbon: warning: unknown file requested [LanguageServerFileUnknown]
 // CHECK:STDERR:
-// CHECK:STDOUT:
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 3,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
+// CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/text_document/change_count.carbon b/toolchain/language_server/testdata/text_document/change_count.carbon
index d1fbd895a56c..79d11917b8cf 100644
--- a/toolchain/language_server/testdata/text_document/change_count.carbon
+++ b/toolchain/language_server/testdata/text_document/change_count.carbon
@@ -9,6 +9,7 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/text_document/change_count.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-NOTIFY:textDocument/didOpen:
   "textDocument": {"uri": "file:/test.carbon", "languageId": "carbon",
                    "text": "// Empty"}
@@ -28,6 +29,7 @@
      "text": "a"}
   ]
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
@@ -36,7 +38,18 @@
 // CHECK:STDERR:
 // CHECK:STDERR: /test.carbon: warning: received unsupported contentChanges count: 2 [LanguageServerUnsupportedChanges]
 // CHECK:STDERR:
-// CHECK:STDOUT: Content-Length: 144{{\r}}
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 144{{\r}}
 // CHECK:STDOUT: {{\r}}
 // CHECK:STDOUT: {
 // CHECK:STDOUT:   "jsonrpc": "2.0",
@@ -45,4 +58,10 @@
 // CHECK:STDOUT:     "diagnostics": [],
 // CHECK:STDOUT:     "uri": "file:///test.carbon"
 // CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
 // CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/text_document/change_unknown.carbon b/toolchain/language_server/testdata/text_document/change_unknown.carbon
index edf2b3bdafe7..b397d7e6640c 100644
--- a/toolchain/language_server/testdata/text_document/change_unknown.carbon
+++ b/toolchain/language_server/testdata/text_document/change_unknown.carbon
@@ -9,14 +9,33 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/text_document/change_unknown.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-NOTIFY:textDocument/didChange:
   "textDocument": {"uri": "file:/test.carbon"},
   "contentChanges": [{"text": "new content"}]
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
 // CHECK:STDERR: /test.carbon: warning: unknown file requested [LanguageServerFileUnknown]
 // CHECK:STDERR:
-// CHECK:STDOUT:
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
+// CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/text_document/close_unknown.carbon b/toolchain/language_server/testdata/text_document/close_unknown.carbon
index 2729b85efec7..9397e3959df7 100644
--- a/toolchain/language_server/testdata/text_document/close_unknown.carbon
+++ b/toolchain/language_server/testdata/text_document/close_unknown.carbon
@@ -9,13 +9,32 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/text_document/close_unknown.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-NOTIFY:textDocument/didClose:
   "textDocument": {"uri": "file:/test.carbon"}
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
 // CHECK:STDERR: /test.carbon: warning: tried closing unknown file; ignoring request [LanguageServerCloseUnknownFile]
 // CHECK:STDERR:
-// CHECK:STDOUT:
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
+// CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/text_document/diagnostics.carbon b/toolchain/language_server/testdata/text_document/diagnostics.carbon
index 96251d03b152..d570d2ebd463 100644
--- a/toolchain/language_server/testdata/text_document/diagnostics.carbon
+++ b/toolchain/language_server/testdata/text_document/diagnostics.carbon
@@ -9,6 +9,7 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/text_document/diagnostics.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-NOTIFY:textDocument/didOpen:
   "textDocument": {"uri": "file:/test.carbon", "languageId": "carbon",
                    "version": 1, "text": "{"}
@@ -20,11 +21,23 @@
 [[@LSP-NOTIFY:textDocument/didClose:
   "textDocument": {"uri": "file:/test.carbon"}
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
-// CHECK:STDOUT: Content-Length: 1164{{\r}}
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 1164{{\r}}
 // CHECK:STDOUT: {{\r}}
 // CHECK:STDOUT: {
 // CHECK:STDOUT:   "jsonrpc": "2.0",
@@ -106,4 +119,10 @@
 // CHECK:STDOUT:     "uri": "file:///test.carbon",
 // CHECK:STDOUT:     "version": 2
 // CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
 // CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/text_document/open_change_close.carbon b/toolchain/language_server/testdata/text_document/open_change_close.carbon
index cf8b55ef1acb..6ba033d81785 100644
--- a/toolchain/language_server/testdata/text_document/open_change_close.carbon
+++ b/toolchain/language_server/testdata/text_document/open_change_close.carbon
@@ -9,6 +9,7 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/text_document/open_change_close.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-NOTIFY:textDocument/didOpen:
   "textDocument": {"uri": "file:/test.carbon", "languageId": "carbon",
                    "text": "// Empty"}
@@ -20,11 +21,23 @@
 [[@LSP-NOTIFY:textDocument/didClose:
   "textDocument": {"uri": "file:/test.carbon"}
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
-// CHECK:STDOUT: Content-Length: 144{{\r}}
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 144{{\r}}
 // CHECK:STDOUT: {{\r}}
 // CHECK:STDOUT: {
 // CHECK:STDOUT:   "jsonrpc": "2.0",
@@ -42,4 +55,10 @@
 // CHECK:STDOUT:     "diagnostics": [],
 // CHECK:STDOUT:     "uri": "file:///test.carbon"
 // CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
 // CHECK:STDOUT: }
diff --git a/toolchain/language_server/testdata/text_document/open_duplicate.carbon b/toolchain/language_server/testdata/text_document/open_duplicate.carbon
index 90817231162f..4f48a05babe5 100644
--- a/toolchain/language_server/testdata/text_document/open_duplicate.carbon
+++ b/toolchain/language_server/testdata/text_document/open_duplicate.carbon
@@ -9,6 +9,7 @@
 // TIP:   bazel run //toolchain/testing:file_test -- --dump_output --file_tests=toolchain/language_server/testdata/text_document/open_duplicate.carbon
 
 // --- STDIN
+[[@LSP-CALL:initialize]]
 [[@LSP-NOTIFY:textDocument/didOpen:
   "textDocument": {"uri": "file:/test.carbon", "languageId": "carbon",
                    "text": "// Empty"}
@@ -17,13 +18,25 @@
   "textDocument": {"uri": "file:/test.carbon", "languageId": "carbon",
                    "text": "// Empty"}
 ]]
+[[@LSP-CALL:shutdown]]
 [[@LSP-NOTIFY:exit]]
 
 // --- AUTOUPDATE-SPLIT
 
 // CHECK:STDERR: /test.carbon: warning: duplicate open file request; updating content [LanguageServerOpenDuplicateFile]
 // CHECK:STDERR:
-// CHECK:STDOUT: Content-Length: 144{{\r}}
+// CHECK:STDOUT: Content-Length: 146{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 1,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": {
+// CHECK:STDOUT:     "capabilities": {
+// CHECK:STDOUT:       "documentSymbolProvider": true,
+// CHECK:STDOUT:       "textDocumentSync": 1
+// CHECK:STDOUT:     }
+// CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 144{{\r}}
 // CHECK:STDOUT: {{\r}}
 // CHECK:STDOUT: {
 // CHECK:STDOUT:   "jsonrpc": "2.0",
@@ -41,4 +54,10 @@
 // CHECK:STDOUT:     "diagnostics": [],
 // CHECK:STDOUT:     "uri": "file:///test.carbon"
 // CHECK:STDOUT:   }
+// CHECK:STDOUT: }Content-Length: 51{{\r}}
+// CHECK:STDOUT: {{\r}}
+// CHECK:STDOUT: {
+// CHECK:STDOUT:   "id": 2,
+// CHECK:STDOUT:   "jsonrpc": "2.0",
+// CHECK:STDOUT:   "result": null
 // CHECK:STDOUT: }
diff --git a/utils/vscode/package.json b/utils/vscode/package.json
index 0865f2a45172..af778bfbbcd6 100644
--- a/utils/vscode/package.json
+++ b/utils/vscode/package.json
@@ -46,7 +46,13 @@
           "default": "./bazel-bin/toolchain/carbon"
         }
       }
-    }
+    },
+    "commands": [
+      {
+        "command": "carbon.lsp.restart",
+        "title": "carbon: Restart language server"
+      }
+    ]
   },
   "scripts": {
     "vscode:prepublish": "npm run package",
diff --git a/utils/vscode/src/extension.ts b/utils/vscode/src/extension.ts
index c58773b33cb7..04f9805e841d 100644
--- a/utils/vscode/src/extension.ts
+++ b/utils/vscode/src/extension.ts
@@ -8,7 +8,7 @@
  * This is the main launcher for the LSP extension.
  */
 
-import { workspace, ExtensionContext } from 'vscode';
+import { workspace, ExtensionContext, commands } from 'vscode';
 
 import {
   LanguageClient,
@@ -43,6 +43,12 @@ export function activate(context: ExtensionContext) {
     clientOptions
   );
   client.start();
+
+  context.subscriptions.push(
+    commands.registerCommand('carbon.lsp.restart', () => {
+      client.restart();
+    })
+  );
 }
 
 export function deactivate(): Thenable<void> | undefined {