optimize SSE

src/Infrastructure/BotSharp.Abstraction/Infrastructures/Enums/StateConst.cs

@@ -14,6 +14,11 @@ public static class StateConst
 
     public const string LANGUAGE = "language";
 
+    /// <summary>
+    /// Set from the incoming request: whether the caller wants the reply streamed back as it is generated.
+    /// </summary>
+    public const string USE_STREAM_MESSAGE = "use_stream_message";
+
     public const string SUB_CONVERSATION_ID = "sub_conversation_id";
     public const string ORIGIN_CONVERSATION_ID = "origin_conversation_id";
     public const string WEB_DRIVER_TASK_ID = "web_driver_task_id";

src/Infrastructure/BotSharp.Core/Routing/Reasoning/InstructExecutor.cs

@@ -66,7 +66,7 @@ await HookEmitter.Emit<IRoutingHook>(_services, async hook => await hook.OnRouti
         else
         {
             var state = _services.GetRequiredService<IConversationStateService>();
-            var useStreamMsg = state.GetState("use_stream_message");
+            var useStreamMsg = state.GetState(StateConst.USE_STREAM_MESSAGE);
             var options = new InvokeAgentOptions()
             {
                 From = InvokeSource.Routing,

src/Infrastructure/BotSharp.Core/Routing/RoutingService.cs

@@ -1,3 +1,4 @@
+using BotSharp.Abstraction.Infrastructures.Enums;
 using BotSharp.Abstraction.Routing.Models;
 using BotSharp.Abstraction.Routing.Settings;
 
@@ -53,7 +54,7 @@ public async Task<RoleDialogModel> InstructDirect(Agent agent, RoleDialogModel m
         else
         {
             var state = _services.GetRequiredService<IConversationStateService>();
-            var useStreamMsg = state.GetState("use_stream_message");
+            var useStreamMsg = state.GetState(StateConst.USE_STREAM_MESSAGE);
             var options = new InvokeAgentOptions()
             {
                 From = InvokeSource.Routing,

src/Infrastructure/BotSharp.OpenAPI/Controllers/Conversation/ConversationController.cs

@@ -1,5 +1,6 @@
-using BotSharp.Abstraction.Files.Constants;
+using BotSharp.Abstraction.Files.Constants;
 using BotSharp.Abstraction.Files.Enums;
+using BotSharp.Abstraction.Infrastructures.Enums;
 using BotSharp.Abstraction.MessageHub.Models;
 using BotSharp.Abstraction.MessageHub.Services;
 using BotSharp.Abstraction.Options;
@@ -19,6 +20,8 @@ public partial class ConversationController : ControllerBase
     private readonly JsonSerializerOptions _jsonOptions;
 
     private const string StreamingFlag = "streaming";
+    private const string DoneFlag = "done";
+    private const string IndicatingFlag = "indicating";
 
     public ConversationController(
         IServiceProvider services,
@@ -490,34 +493,53 @@ public async Task SendMessageSse([FromRoute] string agentId, [FromRoute] string
             MessageId = inputMsg.MessageId,
         };
 
+        IConversationCancellationService? convCancellation = null;
+        if (input.IsStreamingMessage)
+        {
+            convCancellation = _services.GetRequiredService<IConversationCancellationService>();
+            convCancellation.RegisterConversation(conversationId);
+        }
+
         Response.StatusCode = 200;
         Response.Headers.Append(Microsoft.Net.Http.Headers.HeaderNames.ContentType, "text/event-stream");
         Response.Headers.Append(Microsoft.Net.Http.Headers.HeaderNames.CacheControl, "no-cache");
         Response.Headers.Append(Microsoft.Net.Http.Headers.HeaderNames.Connection, "keep-alive");
 
-        await conv.SendMessage(agentId, inputMsg,
-                replyMessage: input.Postback,
-                // responsed generated
-                async msg =>
-                {
-                    response.Text = !string.IsNullOrEmpty(msg.SecondaryContent) ? msg.SecondaryContent : msg.Content;
-                    response.MessageLabel = msg.MessageLabel;
-                    response.Function = msg.FunctionName;
-                    response.RichContent = msg.SecondaryRichContent ?? msg.RichContent;
-                    response.Instruction = msg.Instruction;
-                    response.Data = msg.Data;
-                    response.Thought = msg.Thought;
-                    response.MetaData = msg.MetaData;
-                    response.States = state.GetStates();
-
-                    await OnChunkReceived(Response, response);
-                });
-
-        response.States = state.GetStates();
-        response.MessageId = inputMsg.MessageId;
-        response.ConversationId = conversationId;
+        var cancelled = false;
+        try
+        {
+            await conv.SendMessage(agentId, inputMsg,
+                    replyMessage: input.Postback,
+                    // responsed generated
+                    async msg =>
+                    {
+                        response.Text = !string.IsNullOrEmpty(msg.SecondaryContent) ? msg.SecondaryContent : msg.Content;
+                        response.MessageLabel = msg.MessageLabel;
+                        response.Function = msg.FunctionName;
+                        response.RichContent = msg.SecondaryRichContent ?? msg.RichContent;
+                        response.Instruction = msg.Instruction;
+                        response.Data = msg.Data;
+                        response.Thought = msg.Thought;
+                        response.MetaData = msg.MetaData;
+                        response.States = state.GetStates();
+
+                        await OnChunkReceived(Response, response);
+                    });
+        }
+        catch (OperationCanceledException) when (input.IsStreamingMessage)
+        {
+            // The 200 and part of the stream are already on the wire, so this cannot surface as an error
+            // response. The done frame below reports the cancellation instead.
+            cancelled = true;
+        }
+        finally
+        {
+            convCancellation?.UnregisterConversation(conversationId);
+        }
 
-        // await OnEventCompleted(Response);
+        // Nothing else in the stream marks the end of a response: without this frame a client cannot tell
+        // a finished reply from a dropped connection or a proxy timeout.
+        await OnEventCompleted(Response, conversationId, cancelled);
     }
 
     [HttpPost("/conversation/{conversationId}/stop-streaming")]
@@ -554,7 +576,7 @@ private void SetStates(IConversationService conv, NewMessageModel input)
             conv.States.SetState("sampling_factor", input.SamplingFactor, source: StateSource.External);
         }
 
-        conv.States.SetState("use_stream_message", input.IsStreamingMessage, source: StateSource.Application);
+        conv.States.SetState(StateConst.USE_STREAM_MESSAGE, input.IsStreamingMessage, source: StateSource.Application);
     }
 
     private FileContentResult BuildFileResult(string file)
@@ -575,12 +597,17 @@ private async Task OnChunkReceived(HttpResponse response, object message)
         await response.Body.WriteAsync(buffer, 0, buffer.Length);
     }
 
-    private async Task OnEventCompleted(HttpResponse response)
+    private async Task OnEventCompleted(HttpResponse response, string conversationId, bool cancelled)
     {
-        var buffer = Encoding.UTF8.GetBytes("data:[DONE]\n");
-        await response.Body.WriteAsync(buffer, 0, buffer.Length);
+        var completion = new StreamingCompletion
+        {
+            Function = DoneFlag,
+            ConversationId = conversationId,
+            Cancelled = cancelled
+        };
 
-        buffer = Encoding.UTF8.GetBytes("\n");
+        var json = JsonSerializer.Serialize(completion, _jsonOptions);
+        var buffer = Encoding.UTF8.GetBytes($"data:{json}\n\n");
         await response.Body.WriteAsync(buffer, 0, buffer.Length);
     }
 
@@ -611,7 +638,7 @@ private async Task OnReceiveToolCallIndication(string conversationId, RoleDialog
             ConversationId = conversationId,
             MessageId = msg.MessageId,
             Text = msg.Indication,
-            Function = "indicating",
+            Function = IndicatingFlag,
             Instruction = msg.Instruction,
             States = []
         };
@@ -630,26 +657,5 @@ private async Task OnReceiveStreamingDelta(string conversationId, RoleDialogMode
         };
         await OnChunkReceived(Response, delta);
     }
-
-    /// <summary>
-    /// Serializing a full ChatResponseModel per token sent 602 bytes for 5 characters of text, 323 of them
-    /// an empty Sender repeated every token. message_id has to stay: consumers reject a non-indicating
-    /// frame without one.
-    /// </summary>
-    private sealed class StreamingDelta
-    {
-        [System.Text.Json.Serialization.JsonPropertyName("conversation_id")]
-        public string ConversationId { get; set; } = string.Empty;
-
-        [System.Text.Json.Serialization.JsonPropertyName("message_id")]
-        public string MessageId { get; set; } = string.Empty;
-
-        public string? Function { get; set; }
-
-        public string Text { get; set; } = string.Empty;
-
-        [System.Text.Json.Serialization.JsonIgnore(Condition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingNull)]
-        public Dictionary<string, string?>? Thought { get; set; }
-    }
     #endregion
 }

src/Infrastructure/BotSharp.OpenAPI/ViewModels/Conversations/Response/StreamingCompletion.cs

@@ -0,0 +1,20 @@
+using System.Text.Json.Serialization;
+
+namespace BotSharp.OpenAPI.ViewModels.Conversations;
+
+/// <summary>
+/// Payload of the terminating frame. Flagged in the body rather than with an SSE event name because
+/// consumers read this stream line by line and drop anything that is not a data: line.
+///
+/// Deliberately carries no message_id: consumers take any non-indicating frame that has one for a real
+/// agent reply, and would render this one as an empty message.
+/// </summary>
+public sealed class StreamingCompletion
+{
+    public string Function { get; set; } = string.Empty;
+
+    [JsonPropertyName("conversation_id")]
+    public string ConversationId { get; set; } = string.Empty;
+
+    public bool Cancelled { get; set; }
+}

src/Infrastructure/BotSharp.OpenAPI/ViewModels/Conversations/Response/StreamingDelta.cs

@@ -0,0 +1,24 @@
+using System.Text.Json.Serialization;
+
+namespace BotSharp.OpenAPI.ViewModels.Conversations;
+
+/// <summary>
+/// One token of a streamed reply. Serializing a full ChatResponseModel per token sent 602 bytes for 5
+/// characters of text, 323 of them an empty Sender repeated every token. message_id has to stay:
+/// consumers reject a non-indicating frame without one.
+/// </summary>
+public sealed class StreamingDelta
+{
+    [JsonPropertyName("conversation_id")]
+    public string ConversationId { get; set; } = string.Empty;
+
+    [JsonPropertyName("message_id")]
+    public string MessageId { get; set; } = string.Empty;
+
+    public string? Function { get; set; }
+
+    public string Text { get; set; } = string.Empty;
+
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    public Dictionary<string, string?>? Thought { get; set; }
+}