Add files via upload

internal/multiagent/eino_summarize.go

@@ -130,6 +130,14 @@ func newEinoSummarizationMiddleware(
 }
 
 // summarizeFinalizeWithRecentAssistantToolTrail 在摘要消息后保留最近 assistant/tool 轨迹，避免压缩后执行链断裂。
+//
+// 关键不变量：tool_call ↔ tool_result 的 pair 必须整体保留或整体丢弃。
+// 把消息切成 round（回合）为原子单位：
+//   - user(...) 单条为一个 round；
+//   - assistant(tool_calls=[...]) 及其后连续的 role=tool 消息合成一个 round；
+//   - 其它 assistant(reply, 无 tool_calls) 单条为一个 round。
+//
+// 倒序挑 round（预算不够即放弃该 round），保证 tool 消息不会跨 round 被孤立。
 func summarizeFinalizeWithRecentAssistantToolTrail(
 	ctx context.Context,
 	originalMessages []adk.Message,
@@ -157,80 +165,136 @@ func summarizeFinalizeWithRecentAssistantToolTrail(
 		return out, nil
 	}
 
-	selectedReverse := make([]adk.Message, 0, 8)
-	seen := make(map[adk.Message]struct{})
-	totalTokens := 0
-	assistantToolKept := 0
-	const minAssistantToolTrail = 4
+	rounds := splitMessagesIntoRounds(nonSystem)
+	if len(rounds) == 0 {
+		out := make([]adk.Message, 0, len(systemMsgs)+1)
+		out = append(out, systemMsgs...)
+		out = append(out, summary)
+		return out, nil
+	}
 
-	tryKeep := func(msg adk.Message) (bool, error) {
-		if msg == nil {
-			return false, nil
+	// 目标：至少保留 minRounds 个 round 的执行轨迹；在预算允许时尽量多保留。
+	// 优先确保最后一个 round（通常是最新的 tool 往返或 assistant 回复）存在。
+	const minRounds = 2
+
+	selectedRoundsReverse := make([]messageRound, 0, 8)
+	selectedCount := 0
+	totalTokens := 0
+
+	tokensOfRound := func(r messageRound) (int, error) {
+		if len(r.messages) == 0 {
+			return 0, nil
 		}
-		if _, ok := seen[msg]; ok {
-			return false, nil
-		}
-		n, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: []adk.Message{msg}})
+		n, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: r.messages})
 		if err != nil {
-			return false, err
+			return 0, err
 		}
 		if n <= 0 {
-			n = 1
+			n = len(r.messages)
 		}
+		return n, nil
+	}
+
+	for i := len(rounds) - 1; i >= 0; i-- {
+		r := rounds[i]
+		n, err := tokensOfRound(r)
+		if err != nil {
+			return nil, err
+		}
+		// 预算不够：已经保留了足够 round 则停，否则跳过该 round 继续往前找
+		// （避免一个超大 round 挤占全部预算，至少保证有轨迹）。
 		if totalTokens+n > recentTrailTokenBudget {
-			return false, nil
+			if selectedCount >= minRounds {
+				break
+			}
+			continue
 		}
 		totalTokens += n
-		selectedReverse = append(selectedReverse, msg)
-		seen[msg] = struct{}{}
-		return true, nil
+		selectedRoundsReverse = append(selectedRoundsReverse, r)
+		selectedCount++
 	}
 
-	// 优先保留最近 assistant/tool，确保执行轨迹可续跑。
-	for i := len(nonSystem) - 1; i >= 0; i-- {
-		msg := nonSystem[i]
-		if msg.Role != schema.Assistant && msg.Role != schema.Tool {
-			continue
-		}
-		ok, err := tryKeep(msg)
-		if err != nil {
-			return nil, err
-		}
-		if ok {
-			assistantToolKept++
-		}
-		if assistantToolKept >= minAssistantToolTrail {
-			break
-		}
+	// 还原时间顺序
+	selectedMsgs := make([]adk.Message, 0, 8)
+	for i := len(selectedRoundsReverse) - 1; i >= 0; i-- {
+		selectedMsgs = append(selectedMsgs, selectedRoundsReverse[i].messages...)
 	}
 
-	// 在预算内回填更多最近消息，保持短链路上下文。
-	for i := len(nonSystem) - 1; i >= 0; i-- {
-		_, exists := seen[nonSystem[i]]
-		if exists {
-			continue
-		}
-		ok, err := tryKeep(nonSystem[i])
-		if err != nil {
-			return nil, err
-		}
-		if !ok {
-			break
-		}
-	}
-
-	selected := make([]adk.Message, 0, len(selectedReverse))
-	for i := len(selectedReverse) - 1; i >= 0; i-- {
-		selected = append(selected, selectedReverse[i])
-	}
-
-	out := make([]adk.Message, 0, len(systemMsgs)+1+len(selected))
+	out := make([]adk.Message, 0, len(systemMsgs)+1+len(selectedMsgs))
 	out = append(out, systemMsgs...)
 	out = append(out, summary)
-	out = append(out, selected...)
+	out = append(out, selectedMsgs...)
 	return out, nil
 }
 
+// messageRound 表示一个"不可分割"的消息回合。
+//   - 对 assistant(tool_calls) + 随后若干 tool 消息的组合，round 内全部 call_id 成对完整；
+//   - 对独立的 user / assistant(reply) 消息，round 仅包含该条消息。
+type messageRound struct {
+	messages []adk.Message
+}
+
+// splitMessagesIntoRounds 将非 system 消息切分为若干 round，保证：
+//   - 每个 assistant(tool_calls) 与其对应的 role=tool 响应消息在同一个 round；
+//   - 孤立（无对应 assistant(tool_calls)）的 role=tool 消息不会单独成为 round，
+//     而是被丢弃（这些消息在 pair 完整性层面已属孤儿，保留反而会触发 LLM 400）。
+func splitMessagesIntoRounds(msgs []adk.Message) []messageRound {
+	if len(msgs) == 0 {
+		return nil
+	}
+	rounds := make([]messageRound, 0, len(msgs))
+	i := 0
+	for i < len(msgs) {
+		msg := msgs[i]
+		if msg == nil {
+			i++
+			continue
+		}
+		switch {
+		case msg.Role == schema.Assistant && len(msg.ToolCalls) > 0:
+			// 收集该 assistant 提供的 call_id 集合。
+			provided := make(map[string]struct{}, len(msg.ToolCalls))
+			for _, tc := range msg.ToolCalls {
+				if tc.ID != "" {
+					provided[tc.ID] = struct{}{}
+				}
+			}
+			round := messageRound{messages: []adk.Message{msg}}
+			j := i + 1
+			for j < len(msgs) {
+				next := msgs[j]
+				if next == nil {
+					j++
+					continue
+				}
+				if next.Role != schema.Tool {
+					break
+				}
+				if next.ToolCallID != "" {
+					if _, ok := provided[next.ToolCallID]; !ok {
+						// 下一条 tool 不属于当前 assistant，认为当前 round 结束。
+						break
+					}
+				}
+				round.messages = append(round.messages, next)
+				j++
+			}
+			rounds = append(rounds, round)
+			i = j
+		case msg.Role == schema.Tool:
+			// 孤儿 tool 消息：既不跟随在一个 assistant(tool_calls) 后，
+			// 说明它对应的 assistant 已被上游裁剪；直接丢弃，下一步到 orphan pruner
+			// 兜底也不会出错，但在 round 切分这里就剔除更干净。
+			i++
+		default:
+			// user / assistant(reply) / 其它：单条成 round。
+			rounds = append(rounds, messageRound{messages: []adk.Message{msg}})
+			i++
+		}
+	}
+	return rounds
+}
+
 func einoSummarizationTokenCounter(openAIModel string) summarization.TokenCounterFunc {
 	tc := agent.NewTikTokenCounter()
 	return func(ctx context.Context, input *summarization.TokenCounterInput) (int, error) {