Fix Post-Session Processing Service (#64)

Author: MikeAlhayek

src/Primitives/CrestApps.Core.AI.Chat/Services/PostSessionProcessingService.cs

@@ -5,6 +5,7 @@
 using CrestApps.Core.AI.Orchestration;
 using CrestApps.Core.AI.Services;
 using CrestApps.Core.Support.Json;
+using CrestApps.Core.Templates.Parsing;
 using CrestApps.Core.Templates.Services;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.Logging;
@@ -22,6 +23,7 @@ public sealed class PostSessionProcessingService
     private readonly IAIDeploymentManager _deploymentManager;
     private readonly IAIToolsService _toolsService;
     private readonly ITemplateService _aiTemplateService;
+    private readonly ITemplateParser _markdownTemplateParser;
     private readonly IServiceProvider _serviceProvider;
     private readonly TimeProvider _timeProvider;
     private readonly DefaultAIOptions _defaultOptions;
@@ -35,6 +37,7 @@ public sealed class PostSessionProcessingService
     /// <param name="clientFactory">The client factory.</param>
     /// <param name="toolsService">The tools service.</param>
     /// <param name="aiTemplateService">The ai template service.</param>
+    /// <param name="templateParsers">The registered template parsers.</param>
     /// <param name="defaultOptions">The default options.</param>
     /// <param name="serviceProvider">The service provider.</param>
     /// <param name="timeProvider">The time provider.</param>
@@ -44,6 +47,7 @@ public PostSessionProcessingService(
         IAIClientFactory clientFactory,
         IAIToolsService toolsService,
         ITemplateService aiTemplateService,
+        IEnumerable<ITemplateParser> templateParsers,
         DefaultAIOptions defaultOptions,
         IServiceProvider serviceProvider,
         TimeProvider timeProvider,
@@ -54,6 +58,7 @@ public PostSessionProcessingService(
         _deploymentManager = deploymentManager;
         _toolsService = toolsService;
         _aiTemplateService = aiTemplateService;
+        _markdownTemplateParser = ResolveMarkdownTemplateParser(templateParsers);
         _serviceProvider = serviceProvider;
         _timeProvider = timeProvider;
         _defaultOptions = defaultOptions;
@@ -479,79 +484,96 @@ private async Task<Dictionary<string, PostSessionResult>> ProcessWithToolsAsync(
     private PostSessionProcessingResponse TryParsePostSessionResponse(string sessionId, string responseText)
     {
         // Strategy 1: Direct JSON deserialization.
-        try
+        if (TryDeserializePostSessionResponse(responseText, out var directResult))
         {
-            var result = JsonSerializer.Deserialize<PostSessionProcessingResponse>(
-                responseText, JSOptions.CaseInsensitive);
+            return directResult;
+        }
 
-            if (result?.Tasks != null && result.Tasks.Count > 0)
+        if (_logger.IsEnabled(LogLevel.Debug))
+        {
+            _logger.LogDebug(
+                "Post-session response for session '{SessionId}' is not valid JSON. Trying fallback extraction.",
+                sessionId);
+        }
+
+        // Strategy 2: Extract JSON from markdown code fences.
+        var jsonBlock = JsonExtractor.ExtractFromCodeFence(responseText);
+
+        if (TryDeserializePostSessionResponse(jsonBlock, out var fencedResult))
+        {
+            if (_logger.IsEnabled(LogLevel.Debug))
             {
-                return result;
+                _logger.LogDebug(
+                    "Post-session response for session '{SessionId}' parsed successfully from code fence.",
+                    sessionId);
             }
+
+            return fencedResult;
         }
-        catch (JsonException)
+
+        // Strategy 3: Extract the first JSON object from surrounding text.
+        var jsonObject = JsonExtractor.ExtractJsonObject(responseText);
+
+        if (jsonObject != null &&
+            jsonObject != responseText &&
+            TryDeserializePostSessionResponse(jsonObject, out var objectResult))
         {
             if (_logger.IsEnabled(LogLevel.Debug))
             {
                 _logger.LogDebug(
-                    "Post-session response for session '{SessionId}' is not valid JSON. Trying fallback extraction.",
+                    "Post-session response for session '{SessionId}' parsed successfully from embedded JSON object.",
                     sessionId);
             }
+
+            return objectResult;
         }
 
-        // Strategy 2: Extract JSON from markdown code fences.
-        var jsonBlock = JsonExtractor.ExtractFromCodeFence(responseText);
+        // Strategy 4: Normalize markdown/front matter, then retry the same lenient extraction.
+        var normalizedBody = _markdownTemplateParser.Parse(responseText).Body?.Trim();
 
-        if (jsonBlock != null)
+        if (!string.IsNullOrWhiteSpace(normalizedBody) &&
+            !string.Equals(normalizedBody, responseText.Trim(), StringComparison.Ordinal))
         {
-            try
+            if (TryDeserializePostSessionResponse(normalizedBody, out var normalizedResult))
             {
-                var result = JsonSerializer.Deserialize<PostSessionProcessingResponse>(
-                    jsonBlock, JSOptions.CaseInsensitive);
-
-                if (result?.Tasks != null && result.Tasks.Count > 0)
+                if (_logger.IsEnabled(LogLevel.Debug))
                 {
-                    if (_logger.IsEnabled(LogLevel.Debug))
-                    {
-                        _logger.LogDebug(
-                            "Post-session response for session '{SessionId}' parsed successfully from code fence.",
-                            sessionId);
-                    }
-
-                    return result;
+                    _logger.LogDebug(
+                        "Post-session response for session '{SessionId}' parsed successfully from normalized markdown body.",
+                        sessionId);
                 }
+
+                return normalizedResult;
             }
-            catch (JsonException)
-            {
-                // Code fence content wasn't valid JSON either, continue to next strategy.
-            }
-        }
 
-        // Strategy 3: Extract the first JSON object from surrounding text.
-        var jsonObject = JsonExtractor.ExtractJsonObject(responseText);
+            var normalizedJsonBlock = JsonExtractor.ExtractFromCodeFence(normalizedBody);
 
-        if (jsonObject != null && jsonObject != responseText)
-        {
-            try
+            if (TryDeserializePostSessionResponse(normalizedJsonBlock, out var normalizedFencedResult))
             {
-                var result = JsonSerializer.Deserialize<PostSessionProcessingResponse>(
-                    jsonObject, JSOptions.CaseInsensitive);
-
-                if (result?.Tasks != null && result.Tasks.Count > 0)
+                if (_logger.IsEnabled(LogLevel.Debug))
                 {
-                    if (_logger.IsEnabled(LogLevel.Debug))
-                    {
-                        _logger.LogDebug(
-                            "Post-session response for session '{SessionId}' parsed successfully from embedded JSON object.",
-                            sessionId);
-                    }
-
-                    return result;
+                    _logger.LogDebug(
+                        "Post-session response for session '{SessionId}' parsed successfully from normalized markdown code fence.",
+                        sessionId);
                 }
+
+                return normalizedFencedResult;
             }
-            catch (JsonException)
+
+            var normalizedJsonObject = JsonExtractor.ExtractJsonObject(normalizedBody);
+
+            if (normalizedJsonObject != null &&
+                normalizedJsonObject != normalizedBody &&
+                TryDeserializePostSessionResponse(normalizedJsonObject, out var normalizedObjectResult))
             {
-                // Extracted text wasn't valid JSON either.
+                if (_logger.IsEnabled(LogLevel.Debug))
+                {
+                    _logger.LogDebug(
+                        "Post-session response for session '{SessionId}' parsed successfully from normalized markdown embedded JSON object.",
+                        sessionId);
+                }
+
+                return normalizedObjectResult;
             }
         }
 
@@ -617,7 +639,7 @@ private async Task<Dictionary<string, PostSessionResult>> TryRecoverStructuredTo
                 CreateResponseLogPreview(recoveryResponseText));
         }
 
-        PostSessionProcessingResponse result;
+        PostSessionProcessingResponse result = null;
 
         try
         {
@@ -631,8 +653,6 @@ private async Task<Dictionary<string, PostSessionResult>> TryRecoverStructuredTo
                     "Structured recovery for post-session tool response on session '{SessionId}' did not return JSON content.",
                     sessionId);
             }
-
-            return null;
         }
         catch (JsonException)
         {
@@ -642,11 +662,24 @@ private async Task<Dictionary<string, PostSessionResult>> TryRecoverStructuredTo
                     "Structured recovery for post-session tool response on session '{SessionId}' returned invalid JSON content.",
                     sessionId);
             }
+        }
 
-            return null;
+        if (result?.Tasks is { Count: > 0 })
+        {
+            if (_logger.IsEnabled(LogLevel.Debug))
+            {
+                _logger.LogDebug(
+                    "Structured recovery for post-session tool response on session '{SessionId}' succeeded with {TaskCount} task result(s).",
+                    sessionId,
+                    result.Tasks.Count);
+            }
+
+            return ApplyResults(tasks, result.Tasks);
         }
 
-        if (result?.Tasks is null || result.Tasks.Count == 0)
+        var recoveredFromText = TryParsePostSessionResponse(sessionId, recoveryResponseText);
+
+        if (recoveredFromText?.Tasks is null || recoveredFromText.Tasks.Count == 0)
         {
             if (_logger.IsEnabled(LogLevel.Debug))
             {
@@ -661,12 +694,46 @@ private async Task<Dictionary<string, PostSessionResult>> TryRecoverStructuredTo
         if (_logger.IsEnabled(LogLevel.Debug))
         {
             _logger.LogDebug(
-                "Structured recovery for post-session tool response on session '{SessionId}' succeeded with {TaskCount} task result(s).",
+                "Structured recovery for post-session tool response on session '{SessionId}' succeeded by parsing assistant text with {TaskCount} task result(s).",
                 sessionId,
-                result.Tasks.Count);
+                recoveredFromText.Tasks.Count);
+        }
+
+        return ApplyResults(tasks, recoveredFromText.Tasks);
+    }
+
+    private static bool TryDeserializePostSessionResponse(
+        string responseText,
+        out PostSessionProcessingResponse response)
+    {
+        if (string.IsNullOrWhiteSpace(responseText))
+        {
+            response = null;
+            return false;
+        }
+
+        try
+        {
+            response = JsonSerializer.Deserialize<PostSessionProcessingResponse>(
+                responseText,
+                JSOptions.CaseInsensitive);
+
+            return response?.Tasks is { Count: > 0 };
+        }
+        catch (JsonException)
+        {
+            response = null;
+            return false;
         }
+    }
+
+    private static ITemplateParser ResolveMarkdownTemplateParser(IEnumerable<ITemplateParser> templateParsers)
+    {
+        ArgumentNullException.ThrowIfNull(templateParsers);
 
-        return ApplyResults(tasks, result.Tasks);
+        return templateParsers.FirstOrDefault(parser =>
+            parser.SupportedExtensions.Any(extension => string.Equals(extension, ".md", StringComparison.OrdinalIgnoreCase)))
+            ?? throw new InvalidOperationException("No markdown template parser is registered for post-session response recovery.");
     }
 
     private Dictionary<string, PostSessionResult> CreateFailedResults(

tests/CrestApps.Core.Tests/Core/Services/PostSession/PostSessionProcessingServiceTests.cs

@@ -3,6 +3,7 @@
 using CrestApps.Core.AI.Clients;
 using CrestApps.Core.AI.Deployments;
 using CrestApps.Core.AI.Models;
+using CrestApps.Core.Templates.Parsing;
 using CrestApps.Core.Templates.Services;
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.Logging.Abstractions;
@@ -518,7 +519,16 @@ private static PostSessionProcessingService CreateService(IChatClient chatClient
             MaximumIterationsPerRequest = 10,
         };
 
-        return new PostSessionProcessingService(mockClientFactory.Object, toolsService ?? mockToolsService.Object, templateService ?? mockTemplateService.Object, defaultOptions, new Mock<IServiceProvider>().Object, TimeProvider.System, NullLoggerFactory.Instance, mockDeploymentManager.Object);
+        return new PostSessionProcessingService(
+            mockClientFactory.Object,
+            toolsService ?? mockToolsService.Object,
+            templateService ?? mockTemplateService.Object,
+            [new DefaultMarkdownTemplateParser()],
+            defaultOptions,
+            new Mock<IServiceProvider>().Object,
+            TimeProvider.System,
+            NullLoggerFactory.Instance,
+            mockDeploymentManager.Object);
     }
 
     /// <summary>
@@ -675,6 +685,63 @@ public async Task ProcessAsync_WithTools_WhenResponseIsTruncatedJson_ShouldRecov
         mockChatClient.Verify(c => c.GetResponseAsync(It.IsAny<IEnumerable<ChatMessage>>(), It.IsAny<ChatOptions>(), It.IsAny<CancellationToken>()), Times.Exactly(2));
     }
 
+    [Fact]
+    public async Task ProcessAsync_WithTools_WhenStructuredRetryReturnsMarkdownWrappedJson_ShouldRecoverFromAssistantText()
+    {
+        // Arrange: the structured retry still returns markdown-wrapped text even though the JSON is recoverable.
+        var profile = CreateProfile();
+        profile.AlterSettings<AIProfilePostSessionSettings>(s =>
+        {
+            s.EnablePostSessionProcessing = true;
+            s.PostSessionTasks = [new PostSessionTask
+            {
+                Name = "summary",
+                Type = PostSessionTaskType.Semantic,
+                Instructions = "Summarize the conversation.",
+            }, ];
+            s.ToolNames = ["sendEmail"];
+        });
+        var session = CreateSession();
+        var prompts = CreatePrompts();
+        var mockTool = new TestAIFunction("sendEmail");
+        var mockToolsService = new Mock<IAIToolsService>();
+        mockToolsService.Setup(t => t.GetByNameAsync("sendEmail")).ReturnsAsync(mockTool);
+        var mockChatClient = new Mock<IChatClient>();
+        mockChatClient.SetupSequence(c => c
+            .GetResponseAsync(It.IsAny<IEnumerable<ChatMessage>>(), It.IsAny<ChatOptions>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, "{")))
+            .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, """
+                ---
+                Title: Post-session result
+                ---
+                ```json
+                {
+                  "tasks": [
+                    {
+                      "name": "summary",
+                      "value": "Customer requested a quote and the assistant collected the zip code."
+                    }
+                  ]
+                }
+                ```
+                """)));
+        var mockTemplateService = new Mock<ITemplateService>();
+        mockTemplateService.Setup(t => t
+            .RenderAsync(It.IsAny<string>(), It.IsAny<IDictionary<string, object>>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync("Rendered prompt");
+        var service = CreateService(chatClient: mockChatClient.Object, toolsService: mockToolsService.Object, templateService: mockTemplateService.Object);
+
+        // Act
+        var result = await service.ProcessAsync(profile, session, prompts, TestContext.Current.CancellationToken);
+
+        // Assert: the recovery pass should fall back to the assistant text and extract the structured result.
+        Assert.NotNull(result);
+        Assert.True(result.ContainsKey("summary"));
+        Assert.Equal("Customer requested a quote and the assistant collected the zip code.", result["summary"].Value);
+        Assert.Equal(PostSessionTaskResultStatus.Succeeded, result["summary"].Status);
+        mockChatClient.Verify(c => c.GetResponseAsync(It.IsAny<IEnumerable<ChatMessage>>(), It.IsAny<ChatOptions>(), It.IsAny<CancellationToken>()), Times.Exactly(2));
+    }
+
     [Fact]
     public async Task ProcessAsync_WithTools_WhenSingleSemanticTaskAndNonJsonResponse_ShouldReturnFailedResult()
     {