Revert "refactor: reimplement tool patching via reverse traversal (#1050)"

This reverts commit 7491283.

Author: tangchaojun-bytedance

adk/middlewares/patchtoolcalls/patchtoolcalls.go

@@ -69,8 +69,8 @@ type typedMiddleware[M adk.MessageType] struct {
 }
 
 func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state *adk.TypedChatModelAgentState[M],
-	mc *adk.TypedModelContext[M],
-) (context.Context, *adk.TypedChatModelAgentState[M], error) {
+	mc *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) {
+
 	if len(state.Messages) == 0 {
 		return ctx, state, nil
 	}
@@ -89,82 +89,112 @@ func (m *typedMiddleware[M]) BeforeModelRewriteState(ctx context.Context, state
 func patchToolCallsForMessage[M adk.MessageType](ctx context.Context,
 	gen func(ctx context.Context, toolName, toolCallID string) (string, error),
 	state *adk.TypedChatModelAgentState[*schema.Message],
-	_ *adk.TypedModelContext[M],
-) (context.Context, *adk.TypedChatModelAgentState[M], error) {
-	// seenIDs stores unique tool call IDs collected by reverse traversal
-	seenIDs := make(map[string]struct{})
+	_ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) {
+
 	patched := make([]*schema.Message, 0, len(state.Messages))
 
-	// Iterate messages in reverse order to track existing tool call IDs
-	for i := len(state.Messages) - 1; i >= 0; i-- {
-		msg := state.Messages[i]
+	for i, msg := range state.Messages {
+		patched = append(patched, msg)
 
-		if msg.Role == schema.Tool {
-			seenIDs[msg.ToolCallID] = struct{}{}
+		if msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 {
+			continue
 		}
 
-		if msg.Role == schema.Assistant && len(msg.ToolCalls) > 0 {
-			for _, tc := range msg.ToolCalls {
-				if _, exists := seenIDs[tc.ID]; !exists {
-					toolMsg, err := createPatchedToolMessage(ctx, gen, tc)
-					if err != nil {
-						return ctx, nil, err
-					}
-					patched = append(patched, toolMsg)
-				}
+		for _, tc := range msg.ToolCalls {
+			if hasCorrespondingToolMessage(state.Messages[i+1:], tc.ID) {
+				continue
 			}
-		}
 
-		patched = append(patched, msg)
+			toolMsg, err := createPatchedToolMessage(ctx, gen, tc)
+			if err != nil {
+				return ctx, nil, err
+			}
+			patched = append(patched, toolMsg)
+		}
 	}
 
 	nState := *state
-	nState.Messages = reverse(patched)
+	nState.Messages = patched
 	return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil
 }
 
 func patchToolCallsForAgenticMessage[M adk.MessageType](ctx context.Context,
 	gen func(ctx context.Context, toolName, toolCallID string) (string, error),
 	state *adk.TypedChatModelAgentState[*schema.AgenticMessage],
-	_ *adk.TypedModelContext[M],
-) (context.Context, *adk.TypedChatModelAgentState[M], error) {
-	// seenIDs stores unique tool call IDs collected by reverse traversal
-	seenIDs := make(map[string]struct{})
+	_ *adk.TypedModelContext[M]) (context.Context, *adk.TypedChatModelAgentState[M], error) {
+
 	patched := make([]*schema.AgenticMessage, 0, len(state.Messages))
 
-	// Iterate messages in reverse order to track existing tool call IDs
-	for i := len(state.Messages) - 1; i >= 0; i-- {
-		msg := state.Messages[i]
+	for i, msg := range state.Messages {
+		patched = append(patched, msg)
+
+		if msg.Role != schema.AgenticRoleTypeAssistant {
+			continue
+		}
 
+		// Collect tool call IDs from this assistant message.
+		var toolCalls []struct {
+			callID string
+			name   string
+		}
 		for _, block := range msg.ContentBlocks {
-			if block == nil {
-				continue
+			if block != nil && block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil {
+				toolCalls = append(toolCalls, struct {
+					callID string
+					name   string
+				}{callID: block.FunctionToolCall.CallID, name: block.FunctionToolCall.Name})
 			}
-			if block.Type == schema.ContentBlockTypeFunctionToolResult && block.FunctionToolResult != nil {
-				seenIDs[block.FunctionToolResult.CallID] = struct{}{}
-			}
-			if block.Type == schema.ContentBlockTypeToolSearchResult && block.ToolSearchFunctionToolResult != nil {
-				seenIDs[block.ToolSearchFunctionToolResult.CallID] = struct{}{}
+		}
+		if len(toolCalls) == 0 {
+			continue
+		}
+
+		for _, tc := range toolCalls {
+			if hasCorrespondingAgenticToolResult(state.Messages[i+1:], tc.callID) {
+				continue
 			}
-			if block.Type == schema.ContentBlockTypeFunctionToolCall && block.FunctionToolCall != nil {
-				if _, exists := seenIDs[block.FunctionToolCall.CallID]; !exists {
-					toolMsg, err := createPatchedAgenticToolMessage(ctx, gen, block.FunctionToolCall.Name, block.FunctionToolCall.CallID)
-					if err != nil {
-						return ctx, nil, err
-					}
-					patched = append(patched, toolMsg)
-				}
+
+			toolMsg, err := createPatchedAgenticToolMessage(ctx, gen, tc.name, tc.callID)
+			if err != nil {
+				return ctx, nil, err
 			}
+			patched = append(patched, toolMsg)
 		}
-
-		patched = append(patched, msg)
 	}
 
 	nState := *state
-	nState.Messages = reverse(patched)
+	nState.Messages = patched
 	return ctx, any(&nState).(*adk.TypedChatModelAgentState[M]), nil
 }
 
+func hasCorrespondingToolMessage(messages []*schema.Message, toolCallID string) bool {
+	for _, msg := range messages {
+		if msg.Role == schema.Tool && msg.ToolCallID == toolCallID {
+			return true
+		}
+	}
+	return false
+}
+
+func hasCorrespondingAgenticToolResult(messages []*schema.AgenticMessage, toolCallID string) bool {
+	for _, msg := range messages {
+		for _, block := range msg.ContentBlocks {
+			if block == nil {
+				continue
+			}
+			if block.Type == schema.ContentBlockTypeFunctionToolResult &&
+				block.FunctionToolResult != nil && block.FunctionToolResult.CallID == toolCallID {
+				return true
+			}
+			if block.Type == schema.ContentBlockTypeToolSearchResult &&
+				block.ToolSearchFunctionToolResult != nil && block.ToolSearchFunctionToolResult.CallID == toolCallID {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 func createPatchedToolMessage(ctx context.Context, gen func(ctx context.Context, toolName, toolCallID string) (string, error), tc schema.ToolCall) (*schema.Message, error) {
 	if gen != nil {
 		content, err := gen(ctx, tc.Function.Name, tc.ID)
@@ -211,13 +241,6 @@ func createPatchedAgenticToolMessage(ctx context.Context, gen func(ctx context.C
 	}, nil
 }
 
-func reverse[M adk.MessageType](s []M) []M {
-	for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 {
-		s[i], s[j] = s[j], s[i]
-	}
-	return s
-}
-
 const (
 	defaultPatchedToolMessageTemplate        = "Tool call %s with id %s was canceled - another message came in before it could be completed."
 	defaultPatchedToolMessageTemplateChinese = "工具调用 %s（ID 为 %s）已被取消——在其完成之前收到了另一条消息。"