From f2ce891f719cd235b2476914c1cb240b3e941c40 Mon Sep 17 00:00:00 2001
From: Ryan Stan Lin <RyanStanLinPublic@gmail.com>
Date: Mon, 21 Jul 2025 19:22:28 +0800
Subject: [PATCH] Fix: Correctly handle chat history in StreamContentAsync for
 function calls

---
 src/GenerativeAI/AiModels/ChatSession.cs | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/GenerativeAI/AiModels/ChatSession.cs b/src/GenerativeAI/AiModels/ChatSession.cs
index b2be786d..04826e0b 100644
--- a/src/GenerativeAI/AiModels/ChatSession.cs
+++ b/src/GenerativeAI/AiModels/ChatSession.cs
@@ -173,20 +173,28 @@ protected override List<Content> BeforeRegeneration(GenerateContentRequest origi
     /// <inheritdoc />
     public override async IAsyncEnumerable<GenerateContentResponse> StreamContentAsync(GenerateContentRequest request, [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
+        var historyCountBeforeRequest = this.History.Count;
         var sb = new StringBuilder();
-        await foreach (var response in base.StreamContentAsync(request,cancellationToken).ConfigureAwait(false))
+        await foreach (var response in base.StreamContentAsync(request, cancellationToken).ConfigureAwait(false))
         {
             if (cancellationToken.IsCancellationRequested)
                 yield break;
-            
             sb.Append(response.Text());
             yield return response;
         }
-
-        var lastRequestContent = request.Contents.Last();
-        var lastResponseContent = RequestExtensions.FormatGenerateContentInput(sb.ToString(), Roles.Model);
-        
-        UpdateHistory(lastRequestContent, lastResponseContent);
+        var finalModelResponseContent = RequestExtensions.FormatGenerateContentInput(sb.ToString(), Roles.Model);
+        var userPromptContent = request.Contents
+            .Skip(historyCountBeforeRequest)
+            .FirstOrDefault(c => c.Role?.Equals(Roles.User, StringComparison.OrdinalIgnoreCase) == true && 
+                                 c.Parts.All(p => p.FunctionResponse == null));
+        if (userPromptContent != null)
+        {
+            UpdateHistory(userPromptContent, finalModelResponseContent);
+        }
+        else
+        {
+            this.LastResponseContent = finalModelResponseContent;
+        }
     }
 
     private void UpdateHistory(GenerateContentRequest request, GenerateContentResponse response)