Add files via upload

internal/multiagent/runner.go

@@ -4,16 +4,12 @@ package multiagent
 import (
 	"context"
 	"encoding/json"
-	"errors"
 	"fmt"
-	"io"
 	"net"
-	"path/filepath"
 	"net/http"
 	"sort"
 	"strings"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"cyberstrike-ai/internal/agent"
@@ -484,485 +480,21 @@ func RunDeepAgent(
 		return "sub"
 	}
 
-	var lastRunMsgs []adk.Message
-	var lastAssistant string
-	// plan_execute：最后一轮 assistant 常被 replanner 的 JSON 覆盖，单独保留 executor 对用户文本。
-	var lastPlanExecuteExecutor string
-
-	// retryHints tracks the corrective hint to append for each retry attempt.
-	// Index i corresponds to the hint that will be appended on attempt i+1.
-	var retryHints []adk.Message
-
-attemptLoop:
-	for attempt := 0; attempt < maxToolCallRecoveryAttempts; attempt++ {
-		msgs := make([]adk.Message, 0, len(baseMsgs)+len(retryHints))
-		msgs = append(msgs, baseMsgs...)
-		msgs = append(msgs, retryHints...)
-
-		if attempt > 0 {
-			mcpIDsMu.Lock()
-			mcpIDs = mcpIDs[:0]
-			mcpIDsMu.Unlock()
-		}
-
-		// 仅保留主代理最后一次 assistant 输出；每轮重试重置，避免拼接失败轮次的片段。
-		lastAssistant = ""
-		lastPlanExecuteExecutor = ""
-		var reasoningStreamSeq int64
-		var einoSubReplyStreamSeq int64
-		toolEmitSeen := make(map[string]struct{})
-		var einoMainRound int
-		var einoLastAgent string
-		subAgentToolStep := make(map[string]int)
-		// Track tool calls emitted in this attempt so we can:
-		// - attach toolCallId to tool_result when framework omits it
-		// - flush running tool calls as failed when a recoverable tool execution error happens
-		pendingByID := make(map[string]toolCallPendingInfo)
-		pendingQueueByAgent := make(map[string][]string)
-		markPending := func(tc toolCallPendingInfo) {
-			if tc.ToolCallID == "" {
-				return
-			}
-			pendingByID[tc.ToolCallID] = tc
-			pendingQueueByAgent[tc.EinoAgent] = append(pendingQueueByAgent[tc.EinoAgent], tc.ToolCallID)
-		}
-		popNextPendingForAgent := func(agentName string) (toolCallPendingInfo, bool) {
-			q := pendingQueueByAgent[agentName]
-			for len(q) > 0 {
-				id := q[0]
-				q = q[1:]
-				pendingQueueByAgent[agentName] = q
-				if tc, ok := pendingByID[id]; ok {
-					delete(pendingByID, id)
-					return tc, true
-				}
-			}
-			return toolCallPendingInfo{}, false
-		}
-		removePendingByID := func(toolCallID string) {
-			if toolCallID == "" {
-				return
-			}
-			delete(pendingByID, toolCallID)
-			// queue cleanup is lazy in popNextPendingForAgent
-		}
-		flushAllPendingAsFailed := func(err error) {
-			if progress == nil {
-				pendingByID = make(map[string]toolCallPendingInfo)
-				pendingQueueByAgent = make(map[string][]string)
-				return
-			}
-			msg := ""
-			if err != nil {
-				msg = err.Error()
-			}
-			for _, tc := range pendingByID {
-				toolName := tc.ToolName
-				if strings.TrimSpace(toolName) == "" {
-					toolName = "unknown"
-				}
-				progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), map[string]interface{}{
-					"toolName":       toolName,
-					"success":        false,
-					"isError":        true,
-					"result":         msg,
-					"resultPreview":  msg,
-					"toolCallId":     tc.ToolCallID,
-					"conversationId": conversationID,
-					"einoAgent":      tc.EinoAgent,
-					"einoRole":       tc.EinoRole,
-					"source":         "eino",
-				})
-			}
-			pendingByID = make(map[string]toolCallPendingInfo)
-			pendingQueueByAgent = make(map[string][]string)
-		}
-
-		runnerCfg := adk.RunnerConfig{
-			Agent:           da,
-			EnableStreaming: true,
-		}
-		if cp := strings.TrimSpace(ma.EinoMiddleware.CheckpointDir); cp != "" {
-			cpDir := filepath.Join(cp, sanitizeEinoPathSegment(conversationID))
-			st, stErr := newFileCheckPointStore(cpDir)
-			if stErr != nil {
-				if logger != nil {
-					logger.Warn("eino checkpoint store disabled", zap.String("dir", cpDir), zap.Error(stErr))
-				}
-			} else {
-				runnerCfg.CheckPointStore = st
-				if logger != nil {
-					logger.Info("eino runner: checkpoint store enabled", zap.String("dir", cpDir))
-				}
-			}
-		}
-		runner := adk.NewRunner(ctx, runnerCfg)
-		iter := runner.Run(ctx, msgs)
-
-		for {
-			ev, ok := iter.Next()
-			if !ok {
-				lastRunMsgs = msgs
-				break attemptLoop
-			}
-			if ev == nil {
-				continue
-			}
-			if ev.Err != nil {
-				canRetry := attempt+1 < maxToolCallRecoveryAttempts
-
-				// Recoverable: API-level JSON argument validation error.
-				if canRetry && isRecoverableToolCallArgumentsJSONError(ev.Err) {
-					if logger != nil {
-						logger.Warn("eino: recoverable tool-call JSON error from model/API", zap.Error(ev.Err), zap.Int("attempt", attempt))
-					}
-					retryHints = append(retryHints, toolCallArgumentsJSONRetryHint())
-					if progress != nil {
-						progress("eino_recovery", toolCallArgumentsJSONRecoveryTimelineMessage(attempt), map[string]interface{}{
-							"conversationId": conversationID,
-							"source":         "eino",
-							"einoRetry":      attempt,
-							"runIndex":       attempt + 1,
-							"maxRuns":        maxToolCallRecoveryAttempts,
-							"reason":         "invalid_tool_arguments_json",
-						})
-					}
-					continue attemptLoop
-				}
-
-				// Recoverable: tool execution error (unknown sub-agent, tool not found, bad JSON in args, etc.).
-				if canRetry && isRecoverableToolExecutionError(ev.Err) {
-					if logger != nil {
-						logger.Warn("eino: recoverable tool execution error, will retry with corrective hint",
-							zap.Error(ev.Err), zap.Int("attempt", attempt))
-					}
-					// Ensure UI/tool timeline doesn't get stuck at "running" for tool calls that
-					// will never receive a proper tool_result due to the recoverable error.
-					flushAllPendingAsFailed(ev.Err)
-					retryHints = append(retryHints, toolExecutionRetryHint())
-					if progress != nil {
-						progress("eino_recovery", toolExecutionRecoveryTimelineMessage(attempt), map[string]interface{}{
-							"conversationId": conversationID,
-							"source":         "eino",
-							"einoRetry":      attempt,
-							"runIndex":       attempt + 1,
-							"maxRuns":        maxToolCallRecoveryAttempts,
-							"reason":         "tool_execution_error",
-						})
-					}
-					continue attemptLoop
-				}
-
-				// Non-recoverable error.
-				flushAllPendingAsFailed(ev.Err)
-				if progress != nil {
-					progress("error", ev.Err.Error(), map[string]interface{}{
-						"conversationId": conversationID,
-						"source":         "eino",
-					})
-				}
-				return nil, ev.Err
-			}
-			if ev.AgentName != "" && progress != nil {
-				iterEinoAgent := orchestratorName
-				if orchMode == "plan_execute" {
-					if a := strings.TrimSpace(ev.AgentName); a != "" {
-						iterEinoAgent = a
-					}
-				}
-				if streamsMainAssistant(ev.AgentName) {
-					if einoMainRound == 0 {
-						einoMainRound = 1
-						progress("iteration", "", map[string]interface{}{
-							"iteration":      1,
-							"einoScope":      "main",
-							"einoRole":       "orchestrator",
-							"einoAgent":      iterEinoAgent,
-							"orchestration":  orchMode,
-							"conversationId": conversationID,
-							"source":         "eino",
-						})
-					} else if einoLastAgent != "" && !streamsMainAssistant(einoLastAgent) {
-						einoMainRound++
-						progress("iteration", "", map[string]interface{}{
-							"iteration":      einoMainRound,
-							"einoScope":      "main",
-							"einoRole":       "orchestrator",
-							"einoAgent":      iterEinoAgent,
-							"orchestration":  orchMode,
-							"conversationId": conversationID,
-							"source":         "eino",
-						})
-					}
-				}
-				einoLastAgent = ev.AgentName
-				progress("progress", fmt.Sprintf("[Eino] %s", ev.AgentName), map[string]interface{}{
-					"conversationId": conversationID,
-					"einoAgent":      ev.AgentName,
-					"einoRole":       einoRoleTag(ev.AgentName),
-					"orchestration":  orchMode,
-				})
-			}
-			if ev.Output == nil || ev.Output.MessageOutput == nil {
-				continue
-			}
-			mv := ev.Output.MessageOutput
-
-			if mv.IsStreaming && mv.MessageStream != nil {
-				streamHeaderSent := false
-				var reasoningStreamID string
-				var toolStreamFragments []schema.ToolCall
-				var subAssistantBuf strings.Builder
-				var subReplyStreamID string
-				var mainAssistantBuf strings.Builder
-				for {
-					chunk, rerr := mv.MessageStream.Recv()
-					if rerr != nil {
-						if errors.Is(rerr, io.EOF) {
-							break
-						}
-						if logger != nil {
-							logger.Warn("eino stream recv", zap.Error(rerr))
-						}
-						break
-					}
-					if chunk == nil {
-						continue
-					}
-					if progress != nil && strings.TrimSpace(chunk.ReasoningContent) != "" {
-						if reasoningStreamID == "" {
-							reasoningStreamID = fmt.Sprintf("eino-reasoning-%s-%d", conversationID, atomic.AddInt64(&reasoningStreamSeq, 1))
-							progress("thinking_stream_start", " ", map[string]interface{}{
-								"streamId":      reasoningStreamID,
-								"source":        "eino",
-								"einoAgent":     ev.AgentName,
-								"einoRole":      einoRoleTag(ev.AgentName),
-								"orchestration": orchMode,
-							})
-						}
-						progress("thinking_stream_delta", chunk.ReasoningContent, map[string]interface{}{
-							"streamId": reasoningStreamID,
-						})
-					}
-					if chunk.Content != "" {
-						if progress != nil && streamsMainAssistant(ev.AgentName) {
-							if !streamHeaderSent {
-								progress("response_start", "", map[string]interface{}{
-									"conversationId":     conversationID,
-									"mcpExecutionIds":    snapshotMCPIDs(),
-									"messageGeneratedBy": "eino:" + ev.AgentName,
-									"einoRole":           "orchestrator",
-									"einoAgent":          ev.AgentName,
-									"orchestration":      orchMode,
-								})
-								streamHeaderSent = true
-							}
-							progress("response_delta", chunk.Content, map[string]interface{}{
-								"conversationId":  conversationID,
-								"mcpExecutionIds": snapshotMCPIDs(),
-								"einoRole":        "orchestrator",
-								"einoAgent":       ev.AgentName,
-								"orchestration":   orchMode,
-							})
-							mainAssistantBuf.WriteString(chunk.Content)
-						} else if !streamsMainAssistant(ev.AgentName) {
-							if progress != nil {
-								if subReplyStreamID == "" {
-									subReplyStreamID = fmt.Sprintf("eino-sub-reply-%s-%d", conversationID, atomic.AddInt64(&einoSubReplyStreamSeq, 1))
-									progress("eino_agent_reply_stream_start", "", map[string]interface{}{
-										"streamId":       subReplyStreamID,
-										"einoAgent":      ev.AgentName,
-										"einoRole":       "sub",
-										"conversationId": conversationID,
-										"source":         "eino",
-									})
-								}
-								progress("eino_agent_reply_stream_delta", chunk.Content, map[string]interface{}{
-									"streamId":       subReplyStreamID,
-									"conversationId": conversationID,
-								})
-							}
-							subAssistantBuf.WriteString(chunk.Content)
-						}
-					}
-					// 收集流式 tool_calls 全部分片；arguments 在最后一帧常为 ""，需按 index/id 合并后才能展示 subagent_type/description。
-					if len(chunk.ToolCalls) > 0 {
-						toolStreamFragments = append(toolStreamFragments, chunk.ToolCalls...)
-					}
-				}
-				if streamsMainAssistant(ev.AgentName) {
-					if s := strings.TrimSpace(mainAssistantBuf.String()); s != "" {
-						lastAssistant = s
-						if orchMode == "plan_execute" && strings.EqualFold(strings.TrimSpace(ev.AgentName), "executor") {
-							lastPlanExecuteExecutor = UnwrapPlanExecuteUserText(s)
-						}
-					}
-				}
-				if subAssistantBuf.Len() > 0 && progress != nil {
-					if s := strings.TrimSpace(subAssistantBuf.String()); s != "" {
-						if subReplyStreamID != "" {
-							progress("eino_agent_reply_stream_end", s, map[string]interface{}{
-								"streamId":       subReplyStreamID,
-								"einoAgent":      ev.AgentName,
-								"einoRole":       "sub",
-								"conversationId": conversationID,
-								"source":         "eino",
-							})
-						} else {
-							progress("eino_agent_reply", s, map[string]interface{}{
-								"conversationId": conversationID,
-								"einoAgent":      ev.AgentName,
-								"einoRole":       "sub",
-								"source":         "eino",
-							})
-						}
-					}
-				}
-				var lastToolChunk *schema.Message
-				if merged := mergeStreamingToolCallFragments(toolStreamFragments); len(merged) > 0 {
-					lastToolChunk = &schema.Message{ToolCalls: merged}
-				}
-				tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, progress, toolEmitSeen, subAgentToolStep, markPending)
-				continue
-			}
-
-			msg, gerr := mv.GetMessage()
-			if gerr != nil || msg == nil {
-				continue
-			}
-			tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, progress, toolEmitSeen, subAgentToolStep, markPending)
-
-			if mv.Role == schema.Assistant {
-				if progress != nil && strings.TrimSpace(msg.ReasoningContent) != "" {
-					progress("thinking", strings.TrimSpace(msg.ReasoningContent), map[string]interface{}{
-						"conversationId": conversationID,
-						"source":         "eino",
-						"einoAgent":      ev.AgentName,
-						"einoRole":       einoRoleTag(ev.AgentName),
-						"orchestration":  orchMode,
-					})
-				}
-				body := strings.TrimSpace(msg.Content)
-				if body != "" {
-					if streamsMainAssistant(ev.AgentName) {
-						if progress != nil {
-							progress("response_start", "", map[string]interface{}{
-								"conversationId":     conversationID,
-								"mcpExecutionIds":    snapshotMCPIDs(),
-								"messageGeneratedBy": "eino:" + ev.AgentName,
-								"einoRole":           "orchestrator",
-								"einoAgent":          ev.AgentName,
-								"orchestration":      orchMode,
-							})
-							progress("response_delta", body, map[string]interface{}{
-								"conversationId":  conversationID,
-								"mcpExecutionIds": snapshotMCPIDs(),
-								"einoRole":        "orchestrator",
-								"einoAgent":       ev.AgentName,
-								"orchestration":   orchMode,
-							})
-						}
-						lastAssistant = body
-						if orchMode == "plan_execute" && strings.EqualFold(strings.TrimSpace(ev.AgentName), "executor") {
-							lastPlanExecuteExecutor = UnwrapPlanExecuteUserText(body)
-						}
-					} else if progress != nil {
-						progress("eino_agent_reply", body, map[string]interface{}{
-							"conversationId": conversationID,
-							"einoAgent":      ev.AgentName,
-							"einoRole":       "sub",
-							"source":         "eino",
-						})
-					}
-				}
-			}
-
-			if mv.Role == schema.Tool && progress != nil {
-				toolName := msg.ToolName
-				if toolName == "" {
-					toolName = mv.ToolName
-				}
-
-				// bridge 工具在 res.IsError=true 时会返回带前缀的内容；这里解析为 success/isError，避免前端误判为成功。
-				content := msg.Content
-				isErr := false
-				if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
-					isErr = true
-					content = strings.TrimPrefix(content, einomcp.ToolErrorPrefix)
-				}
-
-				preview := content
-				if len(preview) > 200 {
-					preview = preview[:200] + "..."
-				}
-				data := map[string]interface{}{
-					"toolName":       toolName,
-					"success":        !isErr,
-					"isError":        isErr,
-					"result":         content,
-					"resultPreview":  preview,
-					"conversationId": conversationID,
-					"einoAgent":      ev.AgentName,
-					"einoRole":       einoRoleTag(ev.AgentName),
-					"source":         "eino",
-				}
-				toolCallID := strings.TrimSpace(msg.ToolCallID)
-				// Some framework paths (e.g. UnknownToolsHandler) may omit ToolCallID on tool messages.
-				// Infer from the tool_call emission order for this agent to keep UI state consistent.
-				if toolCallID == "" {
-					// In some internal tool execution paths, ev.AgentName may be empty for tool-role
-					// messages. Try several fallbacks to avoid leaving UI tool_call status stuck.
-					if inferred, ok := popNextPendingForAgent(ev.AgentName); ok {
-						toolCallID = inferred.ToolCallID
-					} else if inferred, ok := popNextPendingForAgent(orchestratorName); ok {
-						toolCallID = inferred.ToolCallID
-					} else if inferred, ok := popNextPendingForAgent(""); ok {
-						toolCallID = inferred.ToolCallID
-					} else {
-						// last resort: pick any pending toolCallID
-						for id := range pendingByID {
-							toolCallID = id
-							delete(pendingByID, id)
-							break
-						}
-					}
-				} else {
-					removePendingByID(toolCallID)
-				}
-				if toolCallID != "" {
-					data["toolCallId"] = toolCallID
-				}
-				progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
-			}
-		}
-	}
-
-	mcpIDsMu.Lock()
-	ids := append([]string(nil), mcpIDs...)
-	mcpIDsMu.Unlock()
-
-	histJSON, _ := json.Marshal(lastRunMsgs)
-	cleaned := strings.TrimSpace(lastAssistant)
-	if orchMode == "plan_execute" {
-		if e := strings.TrimSpace(lastPlanExecuteExecutor); e != "" {
-			cleaned = e
-		} else {
-			cleaned = UnwrapPlanExecuteUserText(cleaned)
-		}
-	}
-	cleaned = dedupeRepeatedParagraphs(cleaned, 80)
-	cleaned = dedupeParagraphsByLineFingerprint(cleaned, 100)
-	out := &RunResult{
-		Response:        cleaned,
-		MCPExecutionIDs: ids,
-		LastReActInput:  string(histJSON),
-		LastReActOutput: cleaned,
-	}
-	if out.Response == "" {
-		out.Response = "（Eino 多代理编排已完成，但未捕获到助手文本输出。请查看过程详情或日志。）"
-		out.LastReActOutput = out.Response
-	}
-	return out, nil
+	return runEinoADKAgentLoop(ctx, &einoADKRunLoopArgs{
+		OrchMode:             orchMode,
+		OrchestratorName:     orchestratorName,
+		ConversationID:       conversationID,
+		Progress:             progress,
+		Logger:               logger,
+		SnapshotMCPIDs:       snapshotMCPIDs,
+		StreamsMainAssistant: streamsMainAssistant,
+		EinoRoleTag:          einoRoleTag,
+		CheckpointDir:        ma.EinoMiddleware.CheckpointDir,
+		McpIDsMu:             &mcpIDsMu,
+		McpIDs:               &mcpIDs,
+		DA:                   da,
+		EmptyResponseMessage: "（Eino 多代理编排已完成，但未捕获到助手文本输出。请查看过程详情或日志。）",
+	}, baseMsgs)
 }
 
 func historyToMessages(history []agent.ChatMessage) []adk.Message {