diff --git a/changelog.d/unreleased/1736.fixed.md b/changelog.d/unreleased/1736.fixed.md new file mode 100644 index 0000000000..7225183b05 --- /dev/null +++ b/changelog.d/unreleased/1736.fixed.md @@ -0,0 +1,19 @@ +--- +category: fixed +issues: + - 1736 + - 1737 + - 1738 + - 1742 +affected: + - src/CodeIndex/Mcp/McpServer.cs + - tests/CodeIndex.Tests/McpServerTests.cs +--- + +## English + +- **MCP server shutdown and request handling now have stronger concurrency guards (#1736, #1737, #1738, #1742)** — request processing now has a bounded timeout response, shared text-writer output is serialized, EOF drains in-flight stdio work with a grace period, and the run-loop flag uses volatile visibility. + +## 日本語 + +- **MCP server の shutdown と request handling の concurrency guard を強化しました (#1736, #1737, #1738, #1742)** — request 処理は bounded timeout response を返し、共有 TextWriter 出力を直列化し、EOF 時は stdio の in-flight 処理を猶予付きで drain し、run-loop flag は volatile visibility を使うようになりました。 diff --git a/src/CodeIndex/Database/DbWriter.cs b/src/CodeIndex/Database/DbWriter.cs index 2c7e2f85ca..49f8a37f2e 100644 --- a/src/CodeIndex/Database/DbWriter.cs +++ b/src/CodeIndex/Database/DbWriter.cs @@ -667,7 +667,10 @@ ON CONFLICT(path) DO UPDATE SET cmd.Parameters["@checksum"].Value = (object?)file.Checksum ?? DBNull.Value; cmd.Parameters["@modified"].Value = file.Modified; cmd.Parameters["@generated"].Value = file.Generated ? 1 : 0; - return (long)cmd.ExecuteScalar()!; + using var reader = cmd.ExecuteReader(); + if (!reader.Read()) + throw new InvalidOperationException("SQLite RETURNING id produced no row for file upsert."); + return reader.GetInt64(0); } finally { diff --git a/src/CodeIndex/Mcp/McpServer.cs b/src/CodeIndex/Mcp/McpServer.cs index e51d07146e..2c99472305 100644 --- a/src/CodeIndex/Mcp/McpServer.cs +++ b/src/CodeIndex/Mcp/McpServer.cs @@ -1,4 +1,5 @@ using System.Collections.Concurrent; +using System.Diagnostics; using System.Runtime.InteropServices; using System.Text; using System.Text.Json; @@ -59,10 +60,11 @@ public partial class McpServer : IDisposable // 直後にリセットする。`WithDbReader` が `DbReader` にライブな cancellation token // を渡せるようにするため (#1567)。 private readonly AsyncLocal _currentRequestToken = new(); + private readonly AsyncLocal _isolateDbForCurrentRequest = new(); private readonly AsyncLocal?> _currentOutOfBandFrameWriter = new(); private readonly AsyncLocal?> _deferredFrameLogs = new(); private static readonly AsyncLocal CurrentCorrelationContext = new(); - private bool _running = true; + private volatile bool _running = true; // Per-session DbContext reused across MCP tool calls. Holding the connection open // avoids reopening SQLite, reapplying pragmas, and re-registering every SQL function // on each invocation (issue #1494). @@ -84,6 +86,8 @@ public partial class McpServer : IDisposable // ツール呼び出し監査ログ (#1562)。`--audit-log` 未指定時は null。AuditLogSink のライフサイクル // (ファイルハンドル / rotation) は ProgramRunner 側で所有する。 private readonly AuditLogSink? _auditLog; + private readonly TimeSpan _requestTimeout; + private readonly SemaphoreSlim _textWriterGate = new(1, 1); // `initialize.clientInfo` echoed into every audit record so the trail can answer // "which client issued this call?" without a second log source. Updated on every // `initialize` so a single-session reconnection picks up the new caller identity. @@ -153,6 +157,9 @@ public partial class McpServer : IDisposable // tool calls wedge the SQLite reader lock or balloon memory (#1567). // 同時 in-flight ツール呼び出し数の既定上限 (#1567)。 internal const int DefaultMaxConcurrency = 8; + internal static readonly TimeSpan DefaultRequestTimeout = TimeSpan.FromSeconds(60); + internal static readonly TimeSpan DefaultEofDrainTimeout = TimeSpan.FromSeconds(5); + internal static readonly TimeSpan DefaultEofPostCancelDrainTimeout = TimeSpan.FromSeconds(5); public McpServer(string dbPath, string version, bool dbPathExplicit = false) : this(dbPath, version, dbPathExplicit, null, null, null, null, DefaultMaxConcurrency) @@ -235,6 +242,7 @@ internal McpServer(string dbPath, string version, bool dbPathExplicit, Func @@ -275,6 +283,7 @@ internal void OverrideRateLimiterForTests(RateLimiter limiter) internal string CurrentSessionId => _sessionId; internal Action? RequestRegisteredForTests { get; set; } + internal Func? RequestDelayForTests { get; set; } /// /// Cap configured for concurrent in-flight tool calls (#1567). Surfaced for tests so @@ -283,6 +292,14 @@ internal void OverrideRateLimiterForTests(RateLimiter limiter) /// internal int MaxConcurrency { get; } + internal TimeSpan RequestTimeout + { + get => _requestTimeout; + init => _requestTimeout = value <= TimeSpan.Zero + ? throw new ArgumentOutOfRangeException(nameof(value), value, "MCP request timeout must be greater than zero.") + : value; + } + /// /// Run the MCP server loop on the default stdio transport. Kept as a thin wrapper around /// so existing callers stay @@ -489,8 +506,8 @@ internal async Task RunAsync(IMcpTransport transport, CancellationToken cancella private async Task RunConcurrentFrameLoopAsync(IMcpTransport transport, CancellationToken loopToken) { - using var writeGate = new SemaphoreSlim(1, 1); - using var normalFrameGate = new SemaphoreSlim(1, 1); + var writeGate = new SemaphoreSlim(1, 1); + var normalFrameGate = new SemaphoreSlim(1, 1); var tasks = new List(); while (_running) @@ -591,10 +608,60 @@ private async Task RunConcurrentFrameLoopAsync(IMcpTransport transport, Cancella SpinWait.SpinUntil(() => !_running || _activeRequests.Count > 0, TimeSpan.FromMilliseconds(50)); } - await Task.WhenAll(tasks).ConfigureAwait(false); + await DrainInFlightTasksAsync(tasks, DefaultEofDrainTimeout, DefaultEofPostCancelDrainTimeout).ConfigureAwait(false); Console.Error.WriteLine("[cdidx-mcp] Server stopped. Restart `cdidx mcp` when your client reconnects."); } + private async Task DrainInFlightTasksAsync(List tasks, TimeSpan gracePeriod, TimeSpan postCancelGracePeriod) + { + tasks.RemoveAll(task => task.IsCompleted); + if (tasks.Count == 0) + return; + + var allTasks = Task.WhenAll(tasks); + var completed = await Task.WhenAny(allTasks, Task.Delay(gracePeriod)).ConfigureAwait(false); + if (completed == allTasks) + { + await ObserveInFlightTasksAsync(allTasks).ConfigureAwait(false); + return; + } + + Console.Error.WriteLine($"[cdidx-mcp] EOF reached with {tasks.Count} in-flight request(s); cancelling after {gracePeriod.TotalMilliseconds:0}ms grace period."); + try + { + if (!_shutdownCts.IsCancellationRequested) + _shutdownCts.Cancel(); + } + catch (ObjectDisposedException) + { + // Disposal raced EOF drain; no further action is possible. + } + + completed = await Task.WhenAny(allTasks, Task.Delay(postCancelGracePeriod)).ConfigureAwait(false); + if (completed == allTasks) + { + await ObserveInFlightTasksAsync(allTasks).ConfigureAwait(false); + return; + } + + _ = allTasks.ContinueWith(task => + { + _ = task.Exception; + }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default); + } + + private static async Task ObserveInFlightTasksAsync(Task tasks) + { + try + { + await tasks.ConfigureAwait(false); + } + catch (Exception ex) + { + Console.Error.WriteLine($"[cdidx-mcp] In-flight request ended during EOF drain ({ex.GetType().Name})."); + } + } + /// /// Process one MCP JSON-RPC line and write any response to the provided writer. Kept as a /// thin wrapper around so existing tests that drive a @@ -610,8 +677,16 @@ internal async Task ProcessLineAsync(string line, TextWriter writer) { try { - await WriteJsonLineAsync(writer, response).ConfigureAwait(false); - FlushDeferredFrameLogs(); + await _textWriterGate.WaitAsync().ConfigureAwait(false); + try + { + await WriteJsonLineAsync(writer, response).ConfigureAwait(false); + FlushDeferredFrameLogs(); + } + finally + { + _textWriterGate.Release(); + } } catch (Exception ex) when (ex is IOException or ObjectDisposedException or OperationCanceledException) { @@ -698,7 +773,7 @@ internal static string BuildInvalidUtf8ErrorLog(string detail) ExtractResponseId(request, out responseHasId, out responseId); if (responseHasId && CurrentCorrelationContext.Value is null) frameCorrelationScope = BeginRequestCorrelation(responseId); - var response = await HandleMessageAsync(request).ConfigureAwait(false); + var response = await HandleMessageAsync(request, isolateRequestDb: true).ConfigureAwait(false); return response != null ? SerializeResponseOrFallback(response, responseHasId, responseId) : null; } catch (JsonException ex) @@ -874,12 +949,15 @@ private static void AppendMinimalCorrelationData(StringBuilder builder) /// JSON-RPCメッセージを適切なハンドラにルーティング。 /// internal JsonNode? HandleMessage(JsonNode request) - => HandleMessageAsync(request).GetAwaiter().GetResult(); + => HandleMessageAsync(request, isolateRequestDb: false).GetAwaiter().GetResult(); - internal async Task HandleMessageAsync(JsonNode request) + internal Task HandleMessageAsync(JsonNode request) + => HandleMessageAsync(request, isolateRequestDb: false); + + private async Task HandleMessageAsync(JsonNode request, bool isolateRequestDb) { if (request is JsonArray batch) - return HandleBatchMessage(batch); + return await HandleBatchMessageAsync(batch, isolateRequestDb).ConfigureAwait(false); if (request is not JsonObject obj) return CreateErrorResponse(hasId: false, id: null, code: -32600, message: "Invalid request: expected JSON object", @@ -977,7 +1055,7 @@ private static void AppendMinimalCorrelationData(StringBuilder builder) retrySafe: false); } - return await DispatchWithRequestCancellationAsync(id, () => method switch + return await DispatchWithRequestCancellationAsync(id, isolateRequestDb, () => method switch { "initialize" => Task.FromResult(HandleInitialize(id, request["params"])), "tools/list" => Task.FromResult(HandleToolsList(id)), @@ -995,7 +1073,7 @@ private static void AppendMinimalCorrelationData(StringBuilder builder) }).ConfigureAwait(false); } - private JsonNode? HandleBatchMessage(JsonArray batch) + private async Task HandleBatchMessageAsync(JsonArray batch, bool isolateRequestDb) { if (batch.Count == 0) return CreateErrorResponse(hasId: true, id: null, code: -32600, message: "Invalid request: empty batch", @@ -1036,7 +1114,7 @@ private static void AppendMinimalCorrelationData(StringBuilder builder) } else { - response = HandleMessage(item); + response = await HandleMessageAsync(item, isolateRequestDb).ConfigureAwait(false); } if (response != null) @@ -1047,17 +1125,18 @@ private static void AppendMinimalCorrelationData(StringBuilder builder) } private JsonNode DispatchWithRequestCancellation(JsonNode? id, Func action) - => DispatchWithRequestCancellationAsync(id, () => Task.FromResult(action())).GetAwaiter().GetResult(); + => DispatchWithRequestCancellationAsync(id, isolateRequestDb: false, () => Task.FromResult(action())).GetAwaiter().GetResult(); - private async Task DispatchWithRequestCancellationAsync(JsonNode? id, Func> action) + private async Task DispatchWithRequestCancellationAsync(JsonNode? id, bool isolateRequestDb, Func> action) { var requestKey = SerializeRequestId(id); if (requestKey == null) return await action().ConfigureAwait(false); - using var requestCts = CancellationTokenSource.CreateLinkedTokenSource(_currentRequestToken.Value, _shutdownCts.Token); + var requestCts = CancellationTokenSource.CreateLinkedTokenSource(_currentRequestToken.Value, _shutdownCts.Token); if (!_activeRequests.TryAdd(requestKey, requestCts)) { + requestCts.Dispose(); return CreateErrorResponse(hasId: true, id: id, code: -32600, message: "Duplicate in-flight request id", category: McpErrorEnvelope.CategoryInvalidRequest, suggestion: "JSON-RPC request ids must be unique while a previous request with the same id is still running.", @@ -1066,23 +1145,80 @@ private async Task DispatchWithRequestCancellationAsync(JsonNode? id, RequestRegisteredForTests?.Invoke(id); var previousToken = _currentRequestToken.Value; + var stopwatch = Stopwatch.StartNew(); + var cleanupNow = true; try { _currentRequestToken.Value = requestCts.Token; + requestCts.CancelAfter(_requestTimeout); requestCts.Token.ThrowIfCancellationRequested(); - return await action().ConfigureAwait(false); + if (!isolateRequestDb) + { + if (RequestDelayForTests is { } delay) + await delay(requestCts.Token).ConfigureAwait(false); + return await action().ConfigureAwait(false); + } + + var actionTask = Task.Run(async () => + { + var previousIsolation = _isolateDbForCurrentRequest.Value; + _isolateDbForCurrentRequest.Value = isolateRequestDb; + try + { + if (RequestDelayForTests is { } delay) + await delay(requestCts.Token).ConfigureAwait(false); + return await action().ConfigureAwait(false); + } + finally + { + _isolateDbForCurrentRequest.Value = previousIsolation; + } + }, requestCts.Token); + var completed = await Task.WhenAny(actionTask, Task.Delay(_requestTimeout)).ConfigureAwait(false); + if (completed != actionTask) + { + try { requestCts.Cancel(); } + catch (ObjectDisposedException) { /* completed while timeout cancellation was being delivered. */ } + cleanupNow = false; + _ = actionTask.ContinueWith(task => + { + _ = task.Exception; + _activeRequests.TryRemove(requestKey, out _); + requestCts.Dispose(); + }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default); + return CreateRequestTimeoutResponse(id, stopwatch.Elapsed); + } + + return await actionTask.ConfigureAwait(false); } catch (OperationCanceledException) when (requestCts.IsCancellationRequested) { + if (!previousToken.IsCancellationRequested && !_shutdownCts.IsCancellationRequested && stopwatch.Elapsed >= _requestTimeout) + return CreateRequestTimeoutResponse(id, stopwatch.Elapsed); return CreateCancelledResponse(id); } finally { _currentRequestToken.Value = previousToken; - _activeRequests.TryRemove(requestKey, out _); + if (cleanupNow) + { + _activeRequests.TryRemove(requestKey, out _); + requestCts.Dispose(); + } } } + private static JsonObject CreateRequestTimeoutResponse(JsonNode? id, TimeSpan elapsed) + => CreateErrorResponse(hasId: true, id: id, code: -32603, message: "Request timed out", + category: McpErrorEnvelope.CategoryInternalError, + suggestion: "Retry with a narrower query, refresh the index if it is degraded, or increase the MCP request timeout before retrying.", + retrySafe: true, + extraData: new JsonObject + { + ["reason"] = "timeout", + ["elapsed_ms"] = (long)Math.Ceiling(elapsed.TotalMilliseconds), + }); + private static IDisposable BeginRequestCorrelation(JsonNode? id) { var previous = CurrentCorrelationContext.Value; @@ -2279,14 +2415,25 @@ private JsonNode WithDbReader(JsonNode? id, JsonNode? args, Func() ?? false; + return isolatedReader.RunWithGeneratedScope(() => action(isolatedReader)); + } + var db = GetOrOpenSharedDb(); if (!_sharedDbReadMigrated) { db.TryMigrateForRead(); _sharedDbReadMigrated = true; } - // Reuse the connection-scoped schema cache so each MCP tool call no longer - // re-runs PRAGMA table_info / PRAGMA index_list per DbReader (issue #1565), + // Reuse the connection-scoped schema cache for single-threaded direct callers so each + // call no longer re-runs PRAGMA table_info / PRAGMA index_list per DbReader (issue #1565), // and hand the per-request cancellation token to the reader so SQLite work // the tool kicks off can observe shutdown / client-disconnect cancellation // (#1567). The token is `CancellationToken.None` outside an in-flight request, @@ -2295,21 +2442,18 @@ private JsonNode WithDbReader(JsonNode? id, JsonNode? args, Func() ?? false; return reader.RunWithGeneratedScope(() => action(reader)); } /// - /// Open the per-session DbContext on first use and reuse it on every subsequent call. + /// Open the per-session DbContext on first use and reuse it on every subsequent direct call. /// Centralising the open lets us pay the connection setup, pragma application, and SQL - /// function registration once per MCP session instead of once per tool invocation - /// (#1494). The MCP loop is single-threaded, so no locking is required. - /// MCP セッション初回呼び出し時に DbContext を開き、以後は再利用する。接続セットアップや - /// PRAGMA・SQL 関数登録のコストを毎ツール呼び出しごとに払わないようにする（#1494）。 - /// MCP ループは単一スレッドのためロック不要。 + /// function registration once per direct session instead of once per tool invocation + /// (#1494). Transport requests that may time out independently use isolated DB contexts. + /// 直接呼び出しセッション初回に DbContext を開き、以後は再利用する。timeout 後も独立して + /// 継続し得る transport リクエストは、共有接続を避けるためリクエスト単位の DB context を使う。 /// internal DbContext GetOrOpenSharedDb() { diff --git a/tests/CodeIndex.Tests/McpServerTests.cs b/tests/CodeIndex.Tests/McpServerTests.cs index d8336391ef..37598c09da 100644 --- a/tests/CodeIndex.Tests/McpServerTests.cs +++ b/tests/CodeIndex.Tests/McpServerTests.cs @@ -151,24 +151,24 @@ public async Task ProcessLineAsync_ToolCallEmitsInvocationTelemetry() using var writer = new StringWriter(); using var error = new StringWriter(); - Monitor.Enter(TestConsoleLock.Gate); - try + await Task.Run(() => { - var previousError = Console.Error; - try - { - Console.SetError(error); - await _server.ProcessLineAsync("""{"jsonrpc":"2.0","id":123,"method":"tools/call","params":{"name":"ping","arguments":{}}}""", writer); - } - finally + lock (TestConsoleLock.Gate) { - Console.SetError(previousError); + var previousError = Console.Error; + try + { + Console.SetError(error); +#pragma warning disable xUnit1031 + _server.ProcessLineAsync("""{"jsonrpc":"2.0","id":123,"method":"tools/call","params":{"name":"ping","arguments":{}}}""", writer).GetAwaiter().GetResult(); +#pragma warning restore xUnit1031 + } + finally + { + Console.SetError(previousError); + } } - } - finally - { - Monitor.Exit(TestConsoleLock.Gate); - } + }); var line = error.ToString() .Split('\n', StringSplitOptions.RemoveEmptyEntries) @@ -191,24 +191,24 @@ public async Task ProcessLineAsync_FallbackErrorIncludesCorrelationData() using var writer = new StringWriter(); using var error = new StringWriter(); - Monitor.Enter(TestConsoleLock.Gate); - try + await Task.Run(() => { - var previousError = Console.Error; - try + lock (TestConsoleLock.Gate) { - Console.SetError(error); - await _server.ProcessLineAsync("""{"jsonrpc":"2.0","id":321,"method":"tools/call","params":{"name":42,"arguments":{}}}""", writer); - } - finally - { - Console.SetError(previousError); + var previousError = Console.Error; + try + { + Console.SetError(error); +#pragma warning disable xUnit1031 + _server.ProcessLineAsync("""{"jsonrpc":"2.0","id":321,"method":"tools/call","params":{"name":42,"arguments":{}}}""", writer).GetAwaiter().GetResult(); +#pragma warning restore xUnit1031 + } + finally + { + Console.SetError(previousError); + } } - } - finally - { - Monitor.Exit(TestConsoleLock.Gate); - } + }); var response = JsonNode.Parse(writer.ToString())!; var data = response["error"]!["data"]!; @@ -8623,6 +8623,73 @@ public void ToolsCall_Definition_AcceptsKindClassCaseInsensitively_Issue199() Assert.Equal(0, structuredWrongKind["count"]!.GetValue()); } + [Fact] + public async Task ProcessFrameAsync_RequestTimeout_ReturnsStructuredTimeoutError() + { + using var server = new McpServer(_dbPath, "1.0", dbPathExplicit: false) + { + RequestTimeout = TimeSpan.FromMilliseconds(20), + }; + server.RequestDelayForTests = _ => + { + Thread.Sleep(TimeSpan.FromMilliseconds(200)); + return Task.CompletedTask; + }; + + var responseText = await server.ProcessFrameAsync( + """{"jsonrpc":"2.0","id":123,"method":"tools/call","params":{"name":"status"}}"""); + + var response = JsonNode.Parse(responseText!)!; + var error = response["error"]!; + Assert.Equal(-32603, error["code"]!.GetValue()); + Assert.Equal("Request timed out", error["message"]!.GetValue()); + Assert.Equal("timeout", error["data"]!["reason"]!.GetValue()); + Assert.True(error["data"]!["elapsed_ms"]!.GetValue() >= 1); + Assert.Equal("internal_error", error["data"]!["category"]!.GetValue()); + Assert.True(error["data"]!["retry_safe"]!.GetValue()); + Assert.Equal(123, response["id"]!.GetValue()); + } + + [Fact] + public async Task ProcessFrameAsync_BatchRequestTimeout_ReturnsStructuredTimeoutError() + { + using var server = new McpServer(_dbPath, "1.0", dbPathExplicit: false) + { + RequestTimeout = TimeSpan.FromMilliseconds(20), + }; + server.RequestDelayForTests = _ => + { + Thread.Sleep(TimeSpan.FromMilliseconds(200)); + return Task.CompletedTask; + }; + + var responseText = await server.ProcessFrameAsync( + """[{"jsonrpc":"2.0","id":123,"method":"tools/call","params":{"name":"status"}}]"""); + + var response = JsonNode.Parse(responseText!)!.AsArray().Single()!; + var error = response["error"]!; + Assert.Equal(-32603, error["code"]!.GetValue()); + Assert.Equal("Request timed out", error["message"]!.GetValue()); + Assert.Equal("timeout", error["data"]!["reason"]!.GetValue()); + Assert.Equal(123, response["id"]!.GetValue()); + } + + [Fact] + public async Task RunAsync_StdioEofDrainsInFlightRequestBeforeReturning() + { + using var server = new McpServer(_dbPath, "1.0", dbPathExplicit: false); + server.RequestRegisteredForTests = _ => { }; + var transport = new QueuedFrameTransport( + """{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"status"}}"""); + + await server.RunAsync(transport, CancellationToken.None); + + Assert.Single(transport.WrittenFrames); + var response = JsonNode.Parse(transport.WrittenFrames[0]!)!; + Assert.Equal(1, response["id"]!.GetValue()); + Assert.Null(response["error"]); + } + private static string CreateLegacyDbWithoutIndexedAt() { var dbPath = Path.Combine(Path.GetTempPath(), $"cdidx_mcp_legacy_{Guid.NewGuid():N}.db"); @@ -9436,6 +9503,31 @@ public async Task RunAsync_CancelledBeforeAnyResponseDoesNotWriteFrame() Assert.Equal(0, transport.WriteCalls); } + private sealed class QueuedFrameTransport : IMcpTransport + { + private readonly Queue _frames; + + public QueuedFrameTransport(params string[] frames) + { + _frames = new Queue(frames.Cast().Append(null)); + } + + public string Name => "stdio"; + public string Endpoint => "memory://queued"; + public List WrittenFrames { get; } = []; + + public Task ReadFrameAsync(CancellationToken cancellationToken) + => Task.FromResult(_frames.Count == 0 ? null : _frames.Dequeue()); + + public Task WriteFrameAsync(string? frame, CancellationToken cancellationToken) + { + WrittenFrames.Add(frame); + return Task.CompletedTask; + } + + public ValueTask DisposeAsync() => ValueTask.CompletedTask; + } + /// /// In-memory IMcpTransport whose ReadFrameAsync blocks until the supplied CancellationToken /// trips. Records read/write counts so tests can assert the loop actually entered the read