From f469568d85e95c8db6964e05013c341198693473 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Tue, 4 Mar 2025 12:03:19 +0800
Subject: [PATCH] Improve request context handling in chat streaming

- Use request context for client connection tracking
- Replace CloseNotify with context cancellation
- Enhance streaming method to handle request context termination
---
 neo/assistant/api.go | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index f8b0bef8..69bda80c 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -319,7 +319,7 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
 	var result interface{} = nil
 	var err error = nil
 
-	// Chat with AI in background
+	requestCtx := c.Request.Context()
 	go func() {
 		var res interface{} = nil
 		res, err = ast.streamChat(c, ctx, messages, options, clientBreak, contents, callback...)
@@ -338,7 +338,8 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
 			return nil, err
 		}
 		return result, nil
-	case <-c.Writer.CloseNotify():
+
+	case <-requestCtx.Done():
 		clientBreak <- true
 		return nil, nil
 	}