From dd651888f827f1011829bd9995bf162270d3d9f6 Mon Sep 17 00:00:00 2001 From: Michael Nash Date: Thu, 2 Jul 2026 10:42:34 -0700 Subject: [PATCH] Handle multi-block sampling responses in SampleStructuredAsync The client can split a sampling answer across multiple content blocks, sometimes a truncated false-start followed by the complete answer. Joining them blindly produced invalid JSON. Try each block individually (last first), then fall back to the joined text. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/PrCopilot/Tools/SamplingHelper.cs | 131 ++++++++++++------ .../PrCopilot.Tests/FakeSamplingMcpServer.cs | 15 +- .../PrCopilot.Tests/SamplingHelperTests.cs | 51 +++++++ 3 files changed, 155 insertions(+), 42 deletions(-) diff --git a/PrCopilot/src/PrCopilot/Tools/SamplingHelper.cs b/PrCopilot/src/PrCopilot/Tools/SamplingHelper.cs index f879e96..b5c056c 100644 --- a/PrCopilot/src/PrCopilot/Tools/SamplingHelper.cs +++ b/PrCopilot/src/PrCopilot/Tools/SamplingHelper.cs @@ -38,22 +38,7 @@ internal static class SamplingHelper { DebugLogger.Log("Sampling", $"Requesting text sample (maxTokens={maxTokens}, temp={temperature?.ToString() ?? "default"})"); - var request = new CreateMessageRequestParams - { - MaxTokens = maxTokens, - SystemPrompt = systemPrompt, - Temperature = temperature, - Messages = - [ - new SamplingMessage - { - Role = Role.User, - Content = [new TextContentBlock { Text = userMessage }] - } - ] - }; - - var result = await server.SampleAsync(request, cancellationToken); + var result = await SampleRawAsync(server, systemPrompt, userMessage, maxTokens, temperature, cancellationToken); var text = ExtractText(result); DebugLogger.Log("Sampling", $"Received response: model={result.Model}, stopReason={result.StopReason ?? "null"}, length={text?.Length ?? 0}"); @@ -75,52 +60,120 @@ internal static class SamplingHelper float? temperature = null, CancellationToken cancellationToken = default) where T : class { - var text = await SampleTextAsync( - server, systemPrompt, userMessage, - maxTokens, temperature, cancellationToken); + DebugLogger.Log("Sampling", $"Requesting structured sample (maxTokens={maxTokens}, temp={temperature?.ToString() ?? "default"})"); + + var result = await SampleRawAsync(server, systemPrompt, userMessage, maxTokens, temperature, cancellationToken); + var blocks = ExtractTextBlocks(result); + DebugLogger.Log("Sampling", $"Received response: model={result.Model}, stopReason={result.StopReason ?? "null"}, blocks={blocks.Count}"); - if (string.IsNullOrWhiteSpace(text)) + if (blocks.Count == 0) return null; - try + // The client may split the answer across multiple content blocks — sometimes an + // earlier block is a truncated false-start and the final block is the complete + // answer. Joining them blindly produces invalid JSON, so try each block on its + // own (last first, since the final block is usually the complete response), then + // fall back to the joined text for providers that split one object across blocks. + var candidates = new List(blocks.Count + 1); + for (int i = blocks.Count - 1; i >= 0; i--) + candidates.Add(blocks[i]); + if (blocks.Count > 1) + candidates.Add(string.Join("\n", blocks)); + + string? lastError = null; + foreach (var candidate in candidates) { - // Strip markdown code fences if present (LLMs often wrap JSON in ```json ... ```) - text = StripCodeFences(text); + if (TryDeserialize(candidate, out var parsed, out lastError)) + { + DebugLogger.Log("Sampling", $"Parsed structured response: {typeof(T).Name}"); + return parsed; + } + } - // LLMs frequently emit literal control characters (newlines, tabs, etc.) - // inside JSON string values, which is invalid per RFC 8259. - text = SanitizeJsonControlChars(text); + DebugLogger.Log("Sampling", $"Failed to parse JSON response as {typeof(T).Name} from {blocks.Count} block(s): {lastError}"); + DebugLogger.Log("Sampling", $"Raw blocks: {string.Join(" || ", blocks).Truncate(500)}"); + return null; + } - var parsed = JsonSerializer.Deserialize(text, _jsonOptions); - DebugLogger.Log("Sampling", $"Parsed structured response: {typeof(T).Name}"); - return parsed; + /// + /// Attempt to deserialize a single candidate response into , + /// stripping code fences and escaping stray control characters first. + /// Returns false (with a diagnostic message) instead of throwing on invalid JSON. + /// + private static bool TryDeserialize(string text, out T? parsed, out string? error) where T : class + { + parsed = null; + error = null; + try + { + // Strip markdown code fences if present (LLMs often wrap JSON in ```json ... ```) + // and escape literal control chars inside strings (invalid per RFC 8259). + var cleaned = SanitizeJsonControlChars(StripCodeFences(text)); + parsed = JsonSerializer.Deserialize(cleaned, _jsonOptions); + return parsed != null; } catch (JsonException ex) { - DebugLogger.Log("Sampling", $"Failed to parse JSON response as {typeof(T).Name}: {ex.Message}"); - DebugLogger.Log("Sampling", $"Raw response ({text?.Length ?? 0} chars): {text.Truncate(500)}"); - return null; + error = ex.Message; + return false; } catch (FileNotFoundException ex) { // .NET single-file publishing can throw FileNotFoundException instead of JsonException // when satellite assemblies for JSON error messages are missing. - DebugLogger.Log("Sampling", $"Failed to parse JSON response as {typeof(T).Name}: [{ex.GetType().Name}] {ex.Message}"); - DebugLogger.Log("Sampling", $"Raw response ({text?.Length ?? 0} chars): {text.Truncate(500)}"); - return null; + error = $"[{ex.GetType().Name}] {ex.Message}"; + return false; } } /// - /// Extract the text content from a sampling result. - /// Concatenates all TextContentBlocks in case the response spans multiple blocks. + /// Issue a sampling request and return the raw result. /// - private static string? ExtractText(CreateMessageResult result) + private static async Task SampleRawAsync( + McpServer server, + string systemPrompt, + string userMessage, + int maxTokens, + float? temperature, + CancellationToken cancellationToken) + { + var request = new CreateMessageRequestParams + { + MaxTokens = maxTokens, + SystemPrompt = systemPrompt, + Temperature = temperature, + Messages = + [ + new SamplingMessage + { + Role = Role.User, + Content = [new TextContentBlock { Text = userMessage }] + } + ] + }; + + return await server.SampleAsync(request, cancellationToken); + } + + /// + /// Extract the non-empty text content blocks from a sampling result, preserving order. + /// + private static List ExtractTextBlocks(CreateMessageResult result) { - var texts = result.Content + return result.Content .OfType() .Select(b => b.Text) + .Where(t => !string.IsNullOrWhiteSpace(t)) .ToList(); + } + + /// + /// Extract the text content from a sampling result. + /// Concatenates all TextContentBlocks in case the response spans multiple blocks. + /// + private static string? ExtractText(CreateMessageResult result) + { + var texts = ExtractTextBlocks(result); return texts.Count > 0 ? string.Join("\n", texts) : null; } diff --git a/PrCopilot/tests/PrCopilot.Tests/FakeSamplingMcpServer.cs b/PrCopilot/tests/PrCopilot.Tests/FakeSamplingMcpServer.cs index 84c2c04..679a81a 100644 --- a/PrCopilot/tests/PrCopilot.Tests/FakeSamplingMcpServer.cs +++ b/PrCopilot/tests/PrCopilot.Tests/FakeSamplingMcpServer.cs @@ -13,13 +13,22 @@ namespace PrCopilot.Tests; internal class FakeSamplingMcpServer : FakeMcpServer #pragma warning restore MCPEXP002 { - private readonly string _responseText; + private readonly string[] _responseBlocks; public CreateMessageRequestParams? LastRequest { get; private set; } public FakeSamplingMcpServer(string responseText = "sampling response") { - _responseText = responseText; + _responseBlocks = [responseText]; + } + + /// + /// Simulate a client that returns the answer across multiple content blocks + /// (e.g. a truncated false-start block followed by the complete answer). + /// + public FakeSamplingMcpServer(params string[] responseBlocks) + { + _responseBlocks = responseBlocks.Length > 0 ? responseBlocks : ["sampling response"]; } public override ClientCapabilities? ClientCapabilities => new() @@ -43,7 +52,7 @@ public override Task SendRequestAsync(JsonRpcRequest request, C { Model = "test-model", Role = Role.Assistant, - Content = [new TextContentBlock { Text = _responseText }], + Content = [.. _responseBlocks.Select(b => new TextContentBlock { Text = b })], StopReason = "endTurn" }; diff --git a/PrCopilot/tests/PrCopilot.Tests/SamplingHelperTests.cs b/PrCopilot/tests/PrCopilot.Tests/SamplingHelperTests.cs index 627ef90..4e3b6e5 100644 --- a/PrCopilot/tests/PrCopilot.Tests/SamplingHelperTests.cs +++ b/PrCopilot/tests/PrCopilot.Tests/SamplingHelperTests.cs @@ -348,6 +348,57 @@ public async Task SampleStructuredAsync_ParsesJsonWithLiteralNewlines() Assert.Equal("implement", result.RecommendationType); } + [Fact] + public async Task SampleStructuredAsync_MultipleBlocks_UsesCompleteBlock() + { + // The exact new failure: the client returns a truncated false-start block + // followed by the complete answer block. Joining them yields invalid JSON. + var truncated = "{\"explanation\": \"The reviewer notes that a descriptor message uses the wrong ind"; + var complete = "{\"explanation\": \"The reviewer flags a grammar issue.\", \"recommendation\": \"Change 'a Azure' to 'an Azure'.\", \"recommendationType\": \"implement\"}"; + var server = new FakeSamplingMcpServer(truncated, complete); + + var result = await SamplingHelper.SampleStructuredAsync( + server, "system", "user", maxTokens: 100); + + Assert.NotNull(result); + Assert.Equal("The reviewer flags a grammar issue.", result!.Explanation); + Assert.Contains("an Azure", result.Recommendation); + Assert.Equal("implement", result.RecommendationType); + } + + [Fact] + public async Task SampleStructuredAsync_MultipleBlocks_CompleteBlockFirst() + { + // Same protection regardless of ordering: complete block first, truncated after. + var complete = "{\"explanation\": \"done\", \"recommendation\": \"fix it\", \"recommendationType\": \"implement\"}"; + var truncated = "{\"explanation\": \"partial follow-up that never fin"; + var server = new FakeSamplingMcpServer(complete, truncated); + + var result = await SamplingHelper.SampleStructuredAsync( + server, "system", "user", maxTokens: 100); + + Assert.NotNull(result); + Assert.Equal("done", result!.Explanation); + Assert.Equal("fix it", result.Recommendation); + } + + [Fact] + public async Task SampleStructuredAsync_MultipleBlocks_OneObjectSplitAcrossBlocks_FallsBackToJoin() + { + // Back-compat: a single JSON object split across blocks still parses via the join fallback. + var part1 = "{\"explanation\": \"split\", \"recommendation\": "; + var part2 = "\"joined value\", \"recommendationType\": \"clarify\"}"; + var server = new FakeSamplingMcpServer(part1, part2); + + var result = await SamplingHelper.SampleStructuredAsync( + server, "system", "user", maxTokens: 100); + + Assert.NotNull(result); + Assert.Equal("split", result!.Explanation); + Assert.Equal("joined value", result.Recommendation); + Assert.Equal("clarify", result.RecommendationType); + } + private class TestResponse { public string? Name { get; set; }