Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions changelog.d/unreleased/1684.fixed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
---
category: fixed
issues:
- 1684
affected:
- src/CodeIndex/Mcp/McpServer.cs
- src/CodeIndex/Mcp/HttpMcpTransport.cs
- src/CodeIndex/Mcp/IMcpTransport.cs
- src/CodeIndex/Mcp/McpToolHandlers.cs
- src/CodeIndex/Mcp/McpToolDefinitions.cs
- tests/CodeIndex.Tests/HttpMcpTransportTests.cs
- tests/CodeIndex.Tests/McpServerTests.cs
---

## English

- **MCP write tools now emit progress notifications (#1684)** — `index` and `backfill_fold` now honor `tools/call.params._meta.progressToken` and send `notifications/progress` during long-running work over stdio and HTTP `/events`, so clients can keep connections alive without changing the final tool result shape.

## 日本語

- **MCP の書き込みツールが progress notification を送るようになりました (#1684)** — `index` と `backfill_fold` は `tools/call.params._meta.progressToken` を受け取り、stdio と HTTP `/events` で長時間処理中に `notifications/progress` を送るため、最終 tool result の形を変えずにクライアントが接続を維持しやすくなりました。
54 changes: 53 additions & 1 deletion src/CodeIndex/Mcp/HttpMcpTransport.cs
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,13 @@ namespace CodeIndex.Mcp;
/// SSE / マルチクライアント対応は将来作業として切り出す(現サーバーは自発的なサーバー→クライアント
/// メッセージを発生させないため、最小単位として POST/response で十分)。
/// </summary>
internal sealed class HttpMcpTransport : IMcpTransport
internal sealed class HttpMcpTransport : IMcpTransport, IOutOfBandMcpTransport
{
private readonly HttpListener _listener;
private readonly string _endpoint;
private readonly Action<HttpRequestLogRecord>? _requestLogger;
private readonly object _requestLoggerGate = new();
private readonly ConcurrentDictionary<Guid, EventStream> _eventStreams = new();
private readonly ConcurrentBag<Task> _sseStreams = new();
private readonly CancellationTokenSource _acceptCts = new();
private readonly Channel<PendingRequest> _requestQueue = Channel.CreateUnbounded<PendingRequest>();
Expand Down Expand Up @@ -350,6 +351,25 @@ public async Task WriteFrameAsync(string? frame, CancellationToken cancellationT
}
}

public async Task WriteOutOfBandFrameAsync(string frame, CancellationToken cancellationToken)
{
if (_eventStreams.IsEmpty)
return;

foreach (var (id, stream) in _eventStreams)
{
try
{
await stream.WriteJsonRpcEventAsync(frame, cancellationToken).ConfigureAwait(false);
}
catch
{
_eventStreams.TryRemove(id, out _);
try { stream.Response.Abort(); } catch { /* ignore */ }
}
}
}

private async Task<bool> TryAuthorizeAsync(PendingRequest request)
{
var context = request.Context;
Expand Down Expand Up @@ -418,13 +438,16 @@ private static bool IsEventsPath(string? path)
private async Task RunEventStreamAsync(PendingRequest request, CancellationToken cancellationToken)
{
var context = request.Context;
var streamId = Guid.NewGuid();
var stream = new EventStream(context.Response);
try
{
context.Response.StatusCode = (int)HttpStatusCode.OK;
context.Response.ContentType = "text/event-stream; charset=utf-8";
context.Response.SendChunked = true;
context.Response.AddHeader("Cache-Control", "no-cache");
context.Response.AddHeader("Connection", "keep-alive");
_eventStreams[streamId] = stream;

var prelude = Encoding.UTF8.GetBytes(": cdidx mcp event stream ready\n\n");
await context.Response.OutputStream.WriteAsync(prelude.AsMemory(), cancellationToken).ConfigureAwait(false);
Expand All @@ -445,11 +468,40 @@ private async Task RunEventStreamAsync(PendingRequest request, CancellationToken
}
finally
{
_eventStreams.TryRemove(streamId, out _);
LogRequest(request, (int)HttpStatusCode.OK);
try { context.Response.Close(); } catch { /* ignore */ }
}
}

private sealed class EventStream(HttpListenerResponse response)
{
private readonly SemaphoreSlim _writeGate = new(1, 1);

public HttpListenerResponse Response { get; } = response;

public async Task WriteJsonRpcEventAsync(string frame, CancellationToken cancellationToken)
{
var builder = new StringBuilder();
builder.Append("event: message\n");
foreach (var line in frame.Replace("\r\n", "\n").Replace('\r', '\n').Split('\n'))
builder.Append("data: ").Append(line).Append('\n');
builder.Append('\n');
var bytes = Encoding.UTF8.GetBytes(builder.ToString());

await _writeGate.WaitAsync(cancellationToken).ConfigureAwait(false);
try
{
await Response.OutputStream.WriteAsync(bytes.AsMemory(), cancellationToken).ConfigureAwait(false);
await Response.OutputStream.FlushAsync(cancellationToken).ConfigureAwait(false);
}
finally
{
_writeGate.Release();
}
}
}

private bool HashEqualsConfiguredToken(string provided)
{
// Hash only the attacker-supplied input and compare to the pre-computed configured-token
Expand Down
5 changes: 5 additions & 0 deletions src/CodeIndex/Mcp/IMcpTransport.cs
Original file line number Diff line number Diff line change
Expand Up @@ -34,3 +34,8 @@ internal interface IMcpTransport : IAsyncDisposable
/// </summary>
Task WriteFrameAsync(string? frame, CancellationToken cancellationToken);
}

internal interface IOutOfBandMcpTransport
{
Task WriteOutOfBandFrameAsync(string frame, CancellationToken cancellationToken);
}
53 changes: 51 additions & 2 deletions src/CodeIndex/Mcp/McpServer.cs
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ public partial class McpServer : IDisposable
// 直後にリセットする。`WithDbReader` が `DbReader` にライブな cancellation token
// を渡せるようにするため (#1567)。
private readonly AsyncLocal<CancellationToken> _currentRequestToken = new();
private readonly AsyncLocal<Action<string>?> _currentOutOfBandFrameWriter = new();
private 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
Expand Down Expand Up @@ -433,11 +434,15 @@ internal async Task RunAsync(IMcpTransport transport, CancellationToken cancella
// ツールが起動する SQLite 作業が shutdown / 切断を観測できるよう per-request
// token を `WithDbReader` に渡す (#1567)。
_currentRequestToken.Value = loopToken;
_currentOutOfBandFrameWriter.Value = transport is IOutOfBandMcpTransport outOfBandTransport
? frameToWrite => outOfBandTransport.WriteOutOfBandFrameAsync(frameToWrite, loopToken).GetAwaiter().GetResult()
: null;
response = ProcessFrame(frame);
}
finally
{
_currentRequestToken.Value = CancellationToken.None;
_currentOutOfBandFrameWriter.Value = null;
_concurrencyGate.Release();
}

Expand Down Expand Up @@ -529,11 +534,24 @@ private async Task RunConcurrentFrameLoopAsync(IMcpTransport transport, Cancella
try
{
_currentRequestToken.Value = loopToken;
_currentOutOfBandFrameWriter.Value = frameToWrite =>
{
writeGate.Wait(loopToken);
try
{
transport.WriteFrameAsync(frameToWrite, loopToken).GetAwaiter().GetResult();
}
finally
{
writeGate.Release();
}
};
response = ProcessFrame(frame);
}
finally
{
_currentRequestToken.Value = CancellationToken.None;
_currentOutOfBandFrameWriter.Value = null;
normalFrameGate.Release();
}

Expand Down Expand Up @@ -1458,6 +1476,7 @@ private JsonNode HandleToolsCall(JsonNode? id, JsonNode? callParams)
{
var toolName = callParams?["name"]?.GetValue<string>();
var args = callParams?["arguments"];
var progressToken = TryReadProgressToken(callParams);

if (toolName == null)
{
Expand Down Expand Up @@ -1548,8 +1567,8 @@ private JsonNode HandleToolsCall(JsonNode? id, JsonNode? callParams)
"unused_symbols" => ExecuteUnusedSymbols(id, args),
"symbol_hotspots" => ExecuteSymbolHotspots(id, args),
"ping" => ExecutePing(id),
"index" => ExecuteIndex(id, args),
"backfill_fold" => ExecuteBackfillFold(id),
"index" => ExecuteIndex(id, args, progressToken),
"backfill_fold" => ExecuteBackfillFold(id, progressToken),
"suggest_improvement" => ExecuteSuggestImprovement(id, args),
_ => CreateErrorResponse(hasId: true, id: id, code: -32602, message: $"Unknown tool: {toolName}",
category: McpErrorEnvelope.CategoryToolUnknown,
Expand Down Expand Up @@ -1618,6 +1637,36 @@ private JsonNode HandleToolsCall(JsonNode? id, JsonNode? callParams)
return response;
}

private static JsonNode? TryReadProgressToken(JsonNode? callParams)
{
var token = callParams?["_meta"]?["progressToken"];
return token is null ? null : JsonNode.Parse(token.ToJsonString());
}

private void EmitProgressNotification(JsonNode? progressToken, long progress, long? total, string? message = null)
{
if (progressToken is null || _currentOutOfBandFrameWriter.Value is not { } writer)
return;

var parameters = new JsonObject
{
["progressToken"] = JsonNode.Parse(progressToken.ToJsonString()),
["progress"] = progress,
};
if (total.HasValue)
parameters["total"] = total.Value;
if (!string.IsNullOrWhiteSpace(message))
parameters["message"] = message;

var notification = new JsonObject
{
["jsonrpc"] = "2.0",
["method"] = "notifications/progress",
["params"] = parameters,
};
writer(notification.ToJsonString(_jsonOptions));
}

/// <summary>
/// Emit a single audit record for the just-executed tool call. Inspects the wire
/// response to derive the result count and error code so the audit trail matches what
Expand Down
4 changes: 2 additions & 2 deletions src/CodeIndex/Mcp/McpToolDefinitions.cs
Original file line number Diff line number Diff line change
Expand Up @@ -385,7 +385,7 @@ private JsonNode HandleToolsList(JsonNode? id)
ReadOnlyAnnotations()),
CreateToolDefinition(
"index",
"Index or re-index a project directory. Scans source files, extracts symbols, and builds FTS5 search index. / プロジェクトディレクトリをインデックス(再インデックス)。ソースファイルをスキャンし、シンボルを抽出してFTS5検索インデックスを構築。",
"Index or re-index a project directory. Scans source files, extracts symbols, and builds FTS5 search index. On transports that can carry out-of-band server messages (stdio, and HTTP clients connected to `/events`), when the tools/call request includes `_meta.progressToken`, this tool emits `notifications/progress` with that token while scanning, indexing, and finalizing. / プロジェクトディレクトリをインデックス(再インデックス)。ソースファイルをスキャンし、シンボルを抽出してFTS5検索インデックスを構築。out-of-band のサーバーメッセージを送れる transport(stdio、および `/events` に接続した HTTP クライアント)では、tools/call リクエストに `_meta.progressToken` が含まれる場合、スキャン・インデックス・finalize 中に同じ token の `notifications/progress` を送信する。",
new JsonObject
{
["type"] = "object",
Expand All @@ -400,7 +400,7 @@ private JsonNode HandleToolsList(JsonNode? id)
IndexAnnotations()),
CreateToolDefinition(
"backfill_fold",
"Upgrade folded-name keys in an existing CodeIndex DB without reparsing source files. Rejects missing or blank targets instead of creating a fresh DB. Fills missing `name_folded` columns (or rewrites all keys after fold metadata drift such as version/fingerprint mismatch) and stamps FoldReady on success. / ソース再解析なしで既存の CodeIndex DB の folded-name key を更新する。欠落したDBや空のDBを新規作成せず拒否し、欠損 `name_folded` 列を埋めるか、fold metadata の drift(version / fingerprint 不一致など)時は全 key を再生成し、成功時に FoldReady を stamp する。",
"Upgrade folded-name keys in an existing CodeIndex DB without reparsing source files. Rejects missing or blank targets instead of creating a fresh DB. Fills missing `name_folded` columns (or rewrites all keys after fold metadata drift such as version/fingerprint mismatch) and stamps FoldReady on success. On transports that can carry out-of-band server messages (stdio, and HTTP clients connected to `/events`), when the tools/call request includes `_meta.progressToken`, this tool emits `notifications/progress` with that token during backfill and verification. / ソース再解析なしで既存の CodeIndex DB の folded-name key を更新する。欠落したDBや空のDBを新規作成せず拒否し、欠損 `name_folded` 列を埋めるか、fold metadata の drift(version / fingerprint 不一致など)時は全 key を再生成し、成功時に FoldReady を stamp する。out-of-band のサーバーメッセージを送れる transport(stdio、および `/events` に接続した HTTP クライアント)では、tools/call リクエストに `_meta.progressToken` が含まれる場合、backfill と検証中に同じ token の `notifications/progress` を送信する。",
new JsonObject
{
["type"] = "object",
Expand Down
11 changes: 9 additions & 2 deletions src/CodeIndex/Mcp/McpToolHandlers.cs
Original file line number Diff line number Diff line change
Expand Up @@ -2284,7 +2284,7 @@ private JsonNode ExecuteLanguages(JsonNode? id)
return CreateToolResult(id, summary, payload);
}

private JsonNode ExecuteIndex(JsonNode? id, JsonNode? args)
private JsonNode ExecuteIndex(JsonNode? id, JsonNode? args, JsonNode? progressToken = null)
{
if (!TryReadRequiredStringParameter(args, "path", out var path, out var requiredError))
return CreateToolErrorResponse(id, requiredError!);
Expand Down Expand Up @@ -2411,6 +2411,7 @@ void WriteProjectRootOnce()
// Scan and index / スキャン・インデックス
var scanResult = indexer.ScanFilesDetailed();
var files = scanResult.Files;
EmitProgressNotification(progressToken, 0, files.Count, "Index scan complete; indexing files.");
var csharpWorkspace = BuildMcpCSharpStaticInterfaceWorkspaceSymbols(writer, indexer, projectPath, files);
if (purged > 0 && hadCSharpStaticInterfaceContractsBeforePurge)
csharpWorkspace = csharpWorkspace with { HasStaticInterfaceContracts = true };
Expand Down Expand Up @@ -2496,6 +2497,7 @@ void WriteProjectRootOnce()
errors++;
}
processed++;
EmitProgressNotification(progressToken, processed, files.Count);
}

writer.OptimizeFts();
Expand All @@ -2512,6 +2514,7 @@ void WriteProjectRootOnce()
_ = priorMetadataTargetCsharp;
if (errors == 0)
{
EmitProgressNotification(progressToken, processed, files.Count, "Finalizing index metadata.");
writer.MarkBatchInProgress();
using var readinessTxn = writer.BeginTransaction();
writer.MarkGraphReady();
Expand Down Expand Up @@ -2628,6 +2631,7 @@ void WriteProjectRootOnce()
readinessTxn.Commit();
}
var (totalFiles, totalChunks, totalSymbols, totalReferences) = writer.GetCounts();
EmitProgressNotification(progressToken, files.Count, files.Count, errors == 0 ? "Indexing complete." : "Indexing completed with errors.");

var structured = new JsonObject
{
Expand Down Expand Up @@ -2673,7 +2677,7 @@ void WriteProjectRootOnce()
structured);
}

private JsonNode ExecuteBackfillFold(JsonNode? id)
private JsonNode ExecuteBackfillFold(JsonNode? id, JsonNode? progressToken = null)
{
if (!DbContext.TryValidateExistingCodeIndexDb(_dbPath, out var validationMessage, out var isNotFound))
{
Expand Down Expand Up @@ -2701,7 +2705,9 @@ private JsonNode ExecuteBackfillFold(JsonNode? id)
var storedFoldFingerprint = db.GetMetaString("fold_key_fingerprint");
var rewriteAll = storedFoldVersion != currentFoldVersion
|| storedFoldFingerprint != currentFoldFingerprint;
EmitProgressNotification(progressToken, 0, null, "Backfilling folded-name keys.");
var (symbols, symbolReferences) = writer.BackfillFoldedColumns(rewriteAll);
EmitProgressNotification(progressToken, symbols + symbolReferences, null, "Verifying folded-name keys.");
// MarkFoldReady wraps its own re-verification in BEGIN IMMEDIATE, so a concurrent
// writer cannot insert NULL-folded rows between the verify and the stamp. Issue #1535.
// MarkFoldReady は BEGIN IMMEDIATE 内で再検証するため、concurrent writer による
Expand All @@ -2711,6 +2717,7 @@ private JsonNode ExecuteBackfillFold(JsonNode? id)
return CreateToolErrorResponse(id, "Folded-name backfill verification failed: some rows still have NULL folded values. Re-run backfill_fold.");

var userVersionAfter = db.GetUserVersion();
EmitProgressNotification(progressToken, symbols + symbolReferences, symbols + symbolReferences, "Folded-name backfill complete.");

var payload = new JsonObject
{
Expand Down
Loading
Loading