瀏覽代碼

fix:解决kiro-cli的ACP协议框架上的一些有风险的问题

Jerry Liang 2 月之前
父節點
當前提交
4a64212163
共有 3 個文件被更改,包括 602 次插入 和 300 次删除
  1. 29 4
      service/controllers/aiChatControllers.go
  2. 41 10
      service/services/aiChatServices.go
  3. 532 286
      service/utils/acp/acp.go

+ 29 - 4
service/controllers/aiChatControllers.go

@@ -89,27 +89,52 @@ func AiChat(c *fiber.Ctx) error {
 	c.Set("X-Accel-Buffering", "no")
 
 	c.Context().SetBodyStreamWriter(func(w *bufio.Writer) {
+		// 连接生命周期上下文:写入失败(客户端断开)时取消,
+		// 令 service/acp 层尽快停止本轮并解除阻塞。
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		// flush 封装:返回 false 表示连接已断开
+		flush := func() bool {
+			if err := w.Flush(); err != nil {
+				cancel()
+				return false
+			}
+			return true
+		}
+
 		// 发送 session_info 事件,包含实际生效的 agent 和 model
 		infoContent, _ := json.Marshal(fiber.Map{"agentId": sess.AgentId, "modelId": sess.ModelId})
 		infoEvent, _ := json.Marshal(services.AIChatEvent{Type: "session_info", Content: string(infoContent)})
 		fmt.Fprintf(w, "event: session_info\ndata: %s\n\n", string(infoEvent))
-		w.Flush()
+		if !flush() {
+			return
+		}
 
 		// 如果有欢迎消息且是新会话,先发送 welcome 事件
 		if sess.WelcomeMessage != "" {
 			welcomeEvent, _ := json.Marshal(services.AIChatEvent{Type: "welcome", Content: sess.WelcomeMessage})
 			fmt.Fprintf(w, "event: welcome\ndata: %s\n\n", string(welcomeEvent))
-			w.Flush()
+			if !flush() {
+				return
+			}
 		}
 
 		eventCh := make(chan services.AIChatEvent, 50)
 
-		go acpManager.Prompt(sess, req.Message, eventCh)
+		go acpManager.Prompt(ctx, sess, req.Message, eventCh)
 
 		for event := range eventCh {
 			data, _ := json.Marshal(event)
 			fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.Type, string(data))
-			w.Flush()
+			if !flush() {
+				// 客户端断开:取消上下文,让 service 层停止转发并排空。
+				// 继续 range 直到 eventCh 被关闭,避免 goroutine 泄漏。
+				cancel()
+				for range eventCh {
+				}
+				return
+			}
 		}
 	})
 	return nil

+ 41 - 10
service/services/aiChatServices.go

@@ -1,6 +1,7 @@
 package services
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"fwh-toolbox-server/repositories"
@@ -330,8 +331,11 @@ func (m *ACPManager) ResolvePermission(userId string, optionId string) error {
 
 // ============ 消息发送 ============
 
-// Prompt 发送消息并流式读取响应
-func (m *ACPManager) Prompt(sess *ACPSession, message string, eventCh chan<- AIChatEvent) {
+// Prompt 发送消息并流式读取响应。
+// ctx 由上层(HTTP 处理器)在连接断开时取消:一旦取消,立即通知 Agent 取消本轮
+// (Client.Cancel)、唤醒 acp 层阻塞的事件发送(Client.AbortEvents),并停止向
+// eventCh 转发,避免 goroutine 因客户端断开而永久阻塞。
+func (m *ACPManager) Prompt(ctx context.Context, sess *ACPSession, message string, eventCh chan<- AIChatEvent) {
 	sess.Mutex.Lock()
 	defer sess.Mutex.Unlock()
 
@@ -353,19 +357,37 @@ func (m *ACPManager) Prompt(sess *ACPSession, message string, eventCh chan<- AIC
 	acpEvents := make(chan acp.Event, 50)
 	go sess.Client.Prompt(sess.SessionId, contextMessage, acpEvents)
 
+	// 监听上层取消:断开时通知 Agent 取消并解除 acp 层发送阻塞
+	go func() {
+		<-ctx.Done()
+		sess.Client.Cancel(sess.SessionId)
+		sess.Client.AbortEvents()
+	}()
+
+	// 转发单条事件,支持在上层取消时中断,避免阻塞在满通道上
+	send := func(ev AIChatEvent) bool {
+		select {
+		case eventCh <- ev:
+			return true
+		case <-ctx.Done():
+			return false
+		}
+	}
+
 	// 转换 acp.Event → AIChatEvent
 	for event := range acpEvents {
+		var ok bool
 		switch event.Type {
 		case "text":
-			eventCh <- AIChatEvent{Type: "text", Content: event.Content}
+			ok = send(AIChatEvent{Type: "text", Content: event.Content})
 		case "tool_call":
-			eventCh <- AIChatEvent{
+			ok = send(AIChatEvent{
 				Type:        "tool_call",
 				ToolCallId:  event.ToolCallId,
 				ToolName:    event.ToolName,
 				Description: acp.DescribeTool(event.ToolName),
 				Status:      acp.NormalizeStatus(event.ToolStatus),
-			}
+			})
 		case "tool_update":
 			// update 一般只变更状态;若本次携带了新 title 才重新映射简述
 			ev := AIChatEvent{
@@ -377,15 +399,24 @@ func (m *ACPManager) Prompt(sess *ACPSession, message string, eventCh chan<- AIC
 				ev.ToolName = event.ToolName
 				ev.Description = acp.DescribeTool(event.ToolName)
 			}
-			eventCh <- ev
+			ok = send(ev)
 		case "permission":
-			eventCh <- AIChatEvent{Type: "permission", Content: toJSON(event.Permission)}
+			ok = send(AIChatEvent{Type: "permission", Content: toJSON(event.Permission)})
 		case "metadata":
-			eventCh <- AIChatEvent{Type: "metadata", Content: event.Content}
+			ok = send(AIChatEvent{Type: "metadata", Content: event.Content})
 		case "done":
-			eventCh <- AIChatEvent{Type: "done", SessionId: event.SessionId, StopReason: event.StopReason}
+			ok = send(AIChatEvent{Type: "done", SessionId: event.SessionId, StopReason: event.StopReason})
 		case "error":
-			eventCh <- AIChatEvent{Type: "error", Content: event.Content}
+			ok = send(AIChatEvent{Type: "error", Content: event.Content})
+		default:
+			ok = true
+		}
+		if !ok {
+			// 上层已取消:继续排空 acpEvents 直至 acp 层因 abort/取消关闭它,
+			// 但不再向 eventCh 转发,防止阻塞。
+			for range acpEvents {
+			}
+			break
 		}
 	}
 

File diff suppressed because it is too large
+ 532 - 286
service/utils/acp/acp.go


Some files were not shown because too many files changed in this diff