Fix LSP malformed JSON and lifecycle exit handling

changelog.d/unreleased/2829.fixed.md

@@ -0,0 +1,16 @@
+---
+category: fixed
+issues:
+  - 2829
+affected:
+  - src/CodeIndex/Lsp/LspServer.cs
+  - tests/CodeIndex.Tests/LspServerTests.cs
+---
+
+## English
+
+- **LSP malformed JSON frames no longer terminate the server loop (#2829)** — malformed JSON-RPC payloads now return a parse error response, allowing the stdio LSP loop to continue processing later valid frames.
+
+## 日本語
+
+- **LSP の malformed JSON frame で server loop が終了しなくなりました (#2829)** — 不正な JSON-RPC payload は parse error response として扱われ、stdio LSP loop が後続の有効な frame を処理し続けられるようになりました。

changelog.d/unreleased/2830.fixed.md

@@ -0,0 +1,17 @@
+---
+category: fixed
+issues:
+  - 2830
+affected:
+  - src/CodeIndex/Cli/ProgramRunner.cs
+  - src/CodeIndex/Lsp/LspServer.cs
+  - tests/CodeIndex.Tests/LspServerTests.cs
+---
+
+## English
+
+- **LSP shutdown and exit now terminate the server loop cleanly (#2830)** — the stdio LSP loop now observes `exit` notifications, returns after the lifecycle terminates, avoids processing later frames, and exits non-zero when `exit` arrives before `shutdown`.
+
+## 日本語
+
+- **LSP の shutdown と exit で server loop が正常終了するようになりました (#2830)** — stdio LSP loop は `exit` 通知を監視して lifecycle 終了後に戻り、後続の frame を処理せず、`shutdown` 前に `exit` を受け取った場合は非 0 で終了します。

src/CodeIndex/Cli/ProgramRunner.cs

@@ -1594,8 +1594,7 @@ private static int RunLsp(string[] cmdArgs, string appVersion, JsonSerializerOpt
             }
 
             using var server = new LspServer(new DbReader(db), appVersion, jsonOptions, indexedProjectRoot);
-            server.Run(Console.OpenStandardInput(), Console.OpenStandardOutput());
-            return CommandExitCodes.Success;
+            return server.Run(Console.OpenStandardInput(), Console.OpenStandardOutput());
         }
         catch (OperationCanceledException)
         {

src/CodeIndex/Lsp/LspServer.cs

@@ -21,6 +21,8 @@ internal sealed class LspServer : IDisposable
     private readonly string? _projectRoot;
     private readonly StringComparison _pathStringComparison;
     private bool _shutdownRequested;
+    private bool _exitRequested;
+    private bool _exitRequestedBeforeShutdown;
 
     public LspServer(DbReader reader, string version, JsonSerializerOptions jsonOptions, string? projectRoot = null)
     {
@@ -31,45 +33,67 @@ public LspServer(DbReader reader, string version, JsonSerializerOptions jsonOpti
         _pathStringComparison = PathCasing.ComparisonFor(_projectRoot ?? Environment.CurrentDirectory);
     }
 
-    public void Run(Stream input, Stream output)
+    public int Run(Stream input, Stream output)
     {
         while (TryReadMessage(input, out var payload))
         {
             var response = HandleMessage(payload);
             if (response != null)
                 WriteMessage(output, response.ToJsonString(_jsonOptions));
+            if (_exitRequested)
+                break;
         }
+
+        return _exitRequestedBeforeShutdown ? CommandExitCodes.UsageError : CommandExitCodes.Success;
     }
 
     internal JsonObject? HandleMessage(string payload)
     {
-        using var document = JsonDocument.Parse(payload);
-        var root = document.RootElement;
-        var method = root.TryGetProperty("method", out var methodElement) ? methodElement.GetString() : null;
-        var hasId = root.TryGetProperty("id", out var idElement);
-        JsonNode? id = hasId ? JsonNode.Parse(idElement.GetRawText()) : null;
-
-        if (method == null)
-            return hasId ? Error(id, -32600, "Invalid Request") : null;
-
+        JsonDocument document;
         try
         {
-            return method switch
-            {
-                "initialize" => Result(id, BuildInitializeResult()),
-                "initialized" => null,
-                "shutdown" => HandleShutdown(id),
-                "exit" => null,
-                "workspace/symbol" => Result(id, WorkspaceSymbol(root)),
-                "textDocument/documentSymbol" => Result(id, DocumentSymbol(root)),
-                "textDocument/definition" => Result(id, Definition(root)),
-                "textDocument/references" => Result(id, References(root)),
-                _ => hasId ? Error(id, -32601, $"Method not found: {method}") : null,
-            };
+            document = JsonDocument.Parse(payload);
         }
-        catch (Exception ex) when (ex is ArgumentException or InvalidOperationException or JsonException or IOException)
+        catch (JsonException)
         {
-            return hasId ? Error(id, -32602, ex.Message) : null;
+            return Error(null, -32700, "Parse error");
+        }
+
+        using (document)
+        {
+            JsonNode? id = null;
+            var hasId = false;
+
+            try
+            {
+                var root = document.RootElement;
+                if (root.ValueKind != JsonValueKind.Object)
+                    return Error(null, -32600, "Invalid Request");
+
+                var method = root.TryGetProperty("method", out var methodElement) ? methodElement.GetString() : null;
+                hasId = root.TryGetProperty("id", out var idElement);
+                id = hasId ? JsonNode.Parse(idElement.GetRawText()) : null;
+
+                if (method == null)
+                    return hasId ? Error(id, -32600, "Invalid Request") : null;
+
+                return method switch
+                {
+                    "initialize" => Result(id, BuildInitializeResult()),
+                    "initialized" => null,
+                    "shutdown" => HandleShutdown(id),
+                    "exit" => HandleExit(),
+                    "workspace/symbol" => Result(id, WorkspaceSymbol(root)),
+                    "textDocument/documentSymbol" => Result(id, DocumentSymbol(root)),
+                    "textDocument/definition" => Result(id, Definition(root)),
+                    "textDocument/references" => Result(id, References(root)),
+                    _ => hasId ? Error(id, -32601, $"Method not found: {method}") : null,
+                };
+            }
+            catch (Exception ex) when (ex is ArgumentException or InvalidOperationException or JsonException or IOException)
+            {
+                return hasId ? Error(id, -32602, ex.Message) : null;
+            }
         }
     }
 
@@ -79,6 +103,13 @@ private JsonObject HandleShutdown(JsonNode? id)
         return Result(id, null);
     }
 
+    private JsonObject? HandleExit()
+    {
+        _exitRequestedBeforeShutdown = !_shutdownRequested;
+        _exitRequested = true;
+        return null;
+    }
+
     private JsonObject BuildInitializeResult() => new()
     {
         ["capabilities"] = new JsonObject

tests/CodeIndex.Tests/LspServerTests.cs

@@ -83,6 +83,97 @@ public void HandleMessage_Initialize_AdvertisesCoreCapabilities()
         }
     }
 
+    [Fact]
+    public void Run_MalformedJsonFrame_WritesParseErrorAndContinues()
+    {
+        var projectRoot = TestProjectHelper.CreateTempProject("cdidx_lsp_malformed_json");
+        try
+        {
+            var dbPath = TestProjectHelper.CreateProjectDb(projectRoot);
+            using var db = new DbContext(dbPath);
+            using var server = new LspServer(new DbReader(db), "1.2.3", ProgramRunner.CreateDefaultJsonOptions(), projectRoot);
+            const string initializeRequest = "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{}}";
+            using var input = new MemoryStream(Encoding.UTF8.GetBytes(Frame("{") + Frame(initializeRequest)));
+            using var output = new MemoryStream();
+
+            var exitCode = server.Run(input, output);
+
+            Assert.Equal(CommandExitCodes.Success, exitCode);
+            output.Position = 0;
+            Assert.True(LspServer.TryReadMessage(output, out var parseErrorPayload));
+            using var parseError = JsonDocument.Parse(parseErrorPayload);
+            Assert.Equal(-32700, parseError.RootElement.GetProperty("error").GetProperty("code").GetInt32());
+            Assert.Equal(JsonValueKind.Null, parseError.RootElement.GetProperty("id").ValueKind);
+
+            Assert.True(LspServer.TryReadMessage(output, out var initializePayload));
+            using var initialize = JsonDocument.Parse(initializePayload);
+            Assert.True(initialize.RootElement.GetProperty("result").GetProperty("capabilities").GetProperty("definitionProvider").GetBoolean());
+            Assert.False(LspServer.TryReadMessage(output, out _));
+        }
+        finally
+        {
+            TestProjectHelper.DeleteDirectory(projectRoot);
+        }
+    }
+
+    [Fact]
+    public void Run_ShutdownThenExit_StopsBeforeLaterFrames()
+    {
+        var projectRoot = TestProjectHelper.CreateTempProject("cdidx_lsp_shutdown_exit");
+        try
+        {
+            var dbPath = TestProjectHelper.CreateProjectDb(projectRoot);
+            using var db = new DbContext(dbPath);
+            using var server = new LspServer(new DbReader(db), "1.2.3", ProgramRunner.CreateDefaultJsonOptions(), projectRoot);
+            const string shutdownRequest = "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"shutdown\"}";
+            const string exitNotification = "{\"jsonrpc\":\"2.0\",\"method\":\"exit\"}";
+            const string initializeRequest = "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"initialize\",\"params\":{}}";
+            using var input = new MemoryStream(Encoding.UTF8.GetBytes(
+                Frame(shutdownRequest) + Frame(exitNotification) + Frame(initializeRequest)));
+            using var output = new MemoryStream();
+
+            var exitCode = server.Run(input, output);
+
+            Assert.Equal(CommandExitCodes.Success, exitCode);
+            output.Position = 0;
+            Assert.True(LspServer.TryReadMessage(output, out var shutdownPayload));
+            using var shutdown = JsonDocument.Parse(shutdownPayload);
+            Assert.Equal(2, shutdown.RootElement.GetProperty("id").GetInt32());
+            Assert.Equal(JsonValueKind.Null, shutdown.RootElement.GetProperty("result").ValueKind);
+            Assert.False(LspServer.TryReadMessage(output, out _));
+        }
+        finally
+        {
+            TestProjectHelper.DeleteDirectory(projectRoot);
+        }
+    }
+
+    [Fact]
+    public void Run_ExitBeforeShutdown_ReturnsUsageError()
+    {
+        var projectRoot = TestProjectHelper.CreateTempProject("cdidx_lsp_exit_without_shutdown");
+        try
+        {
+            var dbPath = TestProjectHelper.CreateProjectDb(projectRoot);
+            using var db = new DbContext(dbPath);
+            using var server = new LspServer(new DbReader(db), "1.2.3", ProgramRunner.CreateDefaultJsonOptions(), projectRoot);
+            const string exitNotification = "{\"jsonrpc\":\"2.0\",\"method\":\"exit\"}";
+            const string initializeRequest = "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"initialize\",\"params\":{}}";
+            using var input = new MemoryStream(Encoding.UTF8.GetBytes(Frame(exitNotification) + Frame(initializeRequest)));
+            using var output = new MemoryStream();
+
+            var exitCode = server.Run(input, output);
+
+            Assert.Equal(CommandExitCodes.UsageError, exitCode);
+            output.Position = 0;
+            Assert.False(LspServer.TryReadMessage(output, out _));
+        }
+        finally
+        {
+            TestProjectHelper.DeleteDirectory(projectRoot);
+        }
+    }
+
     [Fact]
     public void HandleMessage_DocumentSymbol_ReturnsIndexedSymbols()
     {
@@ -419,4 +510,7 @@ private static string CreateDefinitionRequest(string sourcePath, int id, int lin
                 position = new { line, character },
             },
         });
+
+    private static string Frame(string payload) =>
+        $"Content-Length: {Encoding.UTF8.GetByteCount(payload)}\r\n\r\n{payload}";
 }