From 3dfb3b4e8240876481023deadad9cc6775f8658a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=85=AC=E6=98=8E?=
 <83812544+Ed1s0nZ@users.noreply.github.com>
Date: Fri, 24 Apr 2026 19:59:15 +0800
Subject: [PATCH] Add files via upload

---
 internal/handler/agent.go             | 11 ++++++++---
 internal/handler/eino_single_agent.go |  8 +++++---
 internal/handler/multi_agent.go       | 11 ++++++++---
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/internal/handler/agent.go b/internal/handler/agent.go
index 5de90488..a5342c87 100644
--- a/internal/handler/agent.go
+++ b/internal/handler/agent.go
@@ -603,11 +603,13 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
 
 	baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context())
 	defer cancelWithCause(nil)
-	progressCallback := h.createProgressCallback(baseCtx, cancelWithCause, conversationID, "", nil)
-	baseCtx = h.injectReactHITLInterceptor(baseCtx, cancelWithCause, conversationID, "", nil)
+	taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
+	defer timeoutCancel()
+	progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, "", nil)
+	taskCtx = h.injectReactHITLInterceptor(taskCtx, cancelWithCause, conversationID, "", nil)
 
 	// 执行Agent Loop，传入历史消息和对话ID（使用包含角色提示词的finalMessage和角色工具列表）
-	result, err := h.agent.AgentLoopWithProgress(baseCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools)
+	result, err := h.agent.AgentLoopWithProgress(taskCtx, finalMessage, agentHistoryMessages, conversationID, progressCallback, roleTools)
 	if err != nil {
 		h.logger.Error("Agent Loop执行失败", zap.Error(err))
 
@@ -1209,6 +1211,9 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
 		}
 		eventJSON, _ := json.Marshal(event)
 		fmt.Fprintf(c.Writer, "data: %s\n\n", eventJSON)
+		done := StreamEvent{Type: "done", Message: ""}
+		doneJSON, _ := json.Marshal(done)
+		fmt.Fprintf(c.Writer, "data: %s\n\n", doneJSON)
 		c.Writer.Flush()
 		return
 	}
diff --git a/internal/handler/eino_single_agent.go b/internal/handler/eino_single_agent.go
index 7f18bf8c..36fc3c1b 100644
--- a/internal/handler/eino_single_agent.go
+++ b/internal/handler/eino_single_agent.go
@@ -281,8 +281,10 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
 	}
 	baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context())
 	defer cancelWithCause(nil)
-	progressCallback := h.createProgressCallback(baseCtx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, progressCallbackRaw)
-	baseCtx = multiagent.WithHITLToolInterceptor(baseCtx, func(ctx context.Context, toolName, arguments string) (string, error) {
+	taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
+	defer timeoutCancel()
+	progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, progressCallbackRaw)
+	taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) {
 		return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments)
 	})
 
@@ -292,7 +294,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
 	}
 
 	result, runErr := multiagent.RunEinoSingleChatModelAgent(
-		baseCtx,
+		taskCtx,
 		h.config,
 		&h.config.MultiAgent,
 		h.agent,
diff --git a/internal/handler/multi_agent.go b/internal/handler/multi_agent.go
index 5114a222..0684ad87 100644
--- a/internal/handler/multi_agent.go
+++ b/internal/handler/multi_agent.go
@@ -40,6 +40,9 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
 		event := StreamEvent{Type: "error", Message: "请求参数错误: " + err.Error()}
 		b, _ := json.Marshal(event)
 		fmt.Fprintf(c.Writer, "data: %s\n\n", b)
+		done := StreamEvent{Type: "done", Message: ""}
+		db, _ := json.Marshal(done)
+		fmt.Fprintf(c.Writer, "data: %s\n\n", db)
 		c.Writer.Flush()
 		return
 	}
@@ -293,13 +296,15 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
 
 	baseCtx, cancelWithCause := context.WithCancelCause(c.Request.Context())
 	defer cancelWithCause(nil)
-	progressCallback := h.createProgressCallback(baseCtx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil)
-	baseCtx = multiagent.WithHITLToolInterceptor(baseCtx, func(ctx context.Context, toolName, arguments string) (string, error) {
+	taskCtx, timeoutCancel := context.WithTimeout(baseCtx, 600*time.Minute)
+	defer timeoutCancel()
+	progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil)
+	taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) {
 		return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments)
 	})
 
 	result, runErr := multiagent.RunDeepAgent(
-		baseCtx,
+		taskCtx,
 		h.config,
 		&h.config.MultiAgent,
 		h.agent,