Add files via upload

2026-07-16 00:47:29 +02:00 · 2026-07-12 13:20:08 +08:00
parent 0f1a6ad25a
commit fe3c845ff8
3 changed files with 136 additions and 2 deletions
@@ -18,8 +18,9 @@ import (
 //  5. soft model-input budget (warn/compact only, never fail locally)
 //  6. final tool-call/result reconciliation
 //  7. orphan tool prune (defense in depth)
-//  8. telemetry
-//  9. model-facing trace snapshot
+//  8. malformed tool_search history repair
+//  9. telemetry
+//  10. model-facing trace snapshot
 type einoChatModelTailConfig struct {
 	logger           *zap.Logger
 	phase            string
@@ -48,6 +49,7 @@ func appendEinoChatModelTailMiddlewares(handlers []adk.ChatModelAgentMiddleware,
 	if !cfg.skipOrphanPruner {
 		handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase))
 	}
+	handlers = append(handlers, newToolSearchResultSanitizerMiddleware(cfg.logger, cfg.phase))
 	if !cfg.skipTelemetry {
 		if teleMw := newEinoModelInputTelemetryMiddleware(cfg.logger, cfg.modelName, cfg.conversationID, cfg.phase); teleMw != nil {
 			handlers = append(handlers, teleMw)
@@ -0,0 +1,76 @@
+package multiagent
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+
+	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/schema"
+	"go.uber.org/zap"
+)
+
+// toolSearchResultSanitizerMiddleware prevents malformed historical tool_search
+// results (for example an HTML gateway error page) from crashing Eino's dynamic
+// tool loader on every retry. Eino expects every tool_search result to be a JSON
+// object containing selectedTools.
+type toolSearchResultSanitizerMiddleware struct {
+	adk.BaseChatModelAgentMiddleware
+	logger *zap.Logger
+	phase  string
+}
+
+func newToolSearchResultSanitizerMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware {
+	return &toolSearchResultSanitizerMiddleware{logger: logger, phase: phase}
+}
+
+type toolSearchResultEnvelope struct {
+	SelectedTools []string `json:"selectedTools"`
+}
+
+func validToolSearchResult(content string) bool {
+	var result toolSearchResultEnvelope
+	if err := json.Unmarshal([]byte(content), &result); err != nil {
+		return false
+	}
+	// Reject JSON values such as null. They unmarshal without an error but do not
+	// satisfy the object-shaped contract used by the toolsearch middleware.
+	return strings.HasPrefix(strings.TrimSpace(content), "{")
+}
+
+func (m *toolSearchResultSanitizerMiddleware) BeforeModelRewriteState(
+	ctx context.Context,
+	state *adk.ChatModelAgentState,
+	_ *adk.ModelContext,
+) (context.Context, *adk.ChatModelAgentState, error) {
+	if m == nil || state == nil || len(state.Messages) == 0 {
+		return ctx, state, nil
+	}
+
+	var rewritten []adk.Message
+	repaired := 0
+	for i, msg := range state.Messages {
+		if msg == nil || msg.Role != schema.Tool || !IsToolSearchTool(msg.ToolName) || validToolSearchResult(msg.Content) {
+			continue
+		}
+		if rewritten == nil {
+			rewritten = append([]adk.Message(nil), state.Messages...)
+		}
+		clone := *msg
+		clone.Content = `{"selectedTools":[],"_recovered":true,"reason":"invalid historical tool_search result"}`
+		rewritten[i] = &clone
+		repaired++
+	}
+
+	if repaired == 0 {
+		return ctx, state, nil
+	}
+	if m.logger != nil {
+		m.logger.Warn("invalid historical tool_search results repaired before model call",
+			zap.String("phase", m.phase),
+			zap.Int("repaired_count", repaired))
+	}
+	ns := *state
+	ns.Messages = rewritten
+	return ctx, &ns, nil
+}
@@ -0,0 +1,56 @@
+package multiagent
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cloudwego/eino/adk"
+	"github.com/cloudwego/eino/schema"
+)
+
+func TestToolSearchResultSanitizerRepairsMalformedHistory(t *testing.T) {
+	good := &schema.Message{Role: schema.Tool, ToolName: "tool_search", Content: `{"selectedTools":["grep"]}`}
+	bad := &schema.Message{Role: schema.Tool, ToolName: "tool_search", Content: "<html>502 Bad Gateway</html>"}
+	other := &schema.Message{Role: schema.Tool, ToolName: "grep", Content: "plain text is valid for other tools"}
+	state := &adk.ChatModelAgentState{Messages: []adk.Message{good, bad, other}}
+
+	mw := newToolSearchResultSanitizerMiddleware(nil, "test")
+	_, got, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
+	if err != nil {
+		t.Fatalf("BeforeModelRewriteState: %v", err)
+	}
+	if got.Messages[0] != good || got.Messages[0].Content != good.Content {
+		t.Fatal("valid tool_search result was unexpectedly changed")
+	}
+	if got.Messages[1] == bad || !validToolSearchResult(got.Messages[1].Content) {
+		t.Fatalf("malformed result was not safely replaced: %q", got.Messages[1].Content)
+	}
+	if got.Messages[2] != other {
+		t.Fatal("non-tool_search result was unexpectedly changed")
+	}
+	if bad.Content != "<html>502 Bad Gateway</html>" {
+		t.Fatal("middleware mutated the original message")
+	}
+}
+
+func TestToolSearchResultSanitizerFastPath(t *testing.T) {
+	msg := &schema.Message{Role: schema.Tool, ToolName: "tool_search", Content: `{"selectedTools":[]}`}
+	state := &adk.ChatModelAgentState{Messages: []adk.Message{msg}}
+	mw := newToolSearchResultSanitizerMiddleware(nil, "test")
+
+	_, got, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
+	if err != nil {
+		t.Fatalf("BeforeModelRewriteState: %v", err)
+	}
+	if got != state {
+		t.Fatal("valid history should use the allocation-free fast path")
+	}
+}
+
+func TestValidToolSearchResultRejectsNonObjectJSON(t *testing.T) {
+	for _, content := range []string{"null", `[]`, `"text"`, `{"selectedTools":"grep"}`} {
+		if validToolSearchResult(content) {
+			t.Fatalf("expected invalid tool_search result: %s", content)
+		}
+	}
+}