Handle multi-block sampling responses in SampleStructuredAsync

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<string>(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<T>(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<T>(text, _jsonOptions);
-            DebugLogger.Log("Sampling", $"Parsed structured response: {typeof(T).Name}");
-            return parsed;
+    /// <summary>
+    /// Attempt to deserialize a single candidate response into <typeparamref name="T"/>,
+    /// stripping code fences and escaping stray control characters first.
+    /// Returns false (with a diagnostic message) instead of throwing on invalid JSON.
+    /// </summary>
+    private static bool TryDeserialize<T>(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<T>(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;
         }
     }
 
     /// <summary>
-    /// 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.
     /// </summary>
-    private static string? ExtractText(CreateMessageResult result)
+    private static async Task<CreateMessageResult> 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);
+    }
+
+    /// <summary>
+    /// Extract the non-empty text content blocks from a sampling result, preserving order.
+    /// </summary>
+    private static List<string> ExtractTextBlocks(CreateMessageResult result)
     {
-        var texts = result.Content
+        return result.Content
             .OfType<TextContentBlock>()
             .Select(b => b.Text)
+            .Where(t => !string.IsNullOrWhiteSpace(t))
             .ToList();
+    }
+
+    /// <summary>
+    /// Extract the text content from a sampling result.
+    /// Concatenates all TextContentBlocks in case the response spans multiple blocks.
+    /// </summary>
+    private static string? ExtractText(CreateMessageResult result)
+    {
+        var texts = ExtractTextBlocks(result);
         return texts.Count > 0 ? string.Join("\n", texts) : null;
     }
 

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];
+    }
+
+    /// <summary>
+    /// Simulate a client that returns the answer across multiple content blocks
+    /// (e.g. a truncated false-start block followed by the complete answer).
+    /// </summary>
+    public FakeSamplingMcpServer(params string[] responseBlocks)
+    {
+        _responseBlocks = responseBlocks.Length > 0 ? responseBlocks : ["sampling response"];
     }
 
     public override ClientCapabilities? ClientCapabilities => new()
@@ -43,7 +52,7 @@ public override Task<JsonRpcResponse> 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"
             };
 

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<CommentExplanationTestModel>(
+            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<CommentExplanationTestModel>(
+            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<CommentExplanationTestModel>(
+            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; }