From fe3c845ff847569626c5ec3caeee8c37a9cdee05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Sun, 12 Jul 2026 13:20:08 +0800 Subject: [PATCH] Add files via upload --- .../eino_chat_model_tail_middleware.go | 6 +- ...tool_search_result_sanitizer_middleware.go | 76 +++++++++++++++++++ ...search_result_sanitizer_middleware_test.go | 56 ++++++++++++++ 3 files changed, 136 insertions(+), 2 deletions(-) create mode 100644 internal/multiagent/tool_search_result_sanitizer_middleware.go create mode 100644 internal/multiagent/tool_search_result_sanitizer_middleware_test.go diff --git a/internal/multiagent/eino_chat_model_tail_middleware.go b/internal/multiagent/eino_chat_model_tail_middleware.go index ad49b9b5..307e1644 100644 --- a/internal/multiagent/eino_chat_model_tail_middleware.go +++ b/internal/multiagent/eino_chat_model_tail_middleware.go @@ -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) diff --git a/internal/multiagent/tool_search_result_sanitizer_middleware.go b/internal/multiagent/tool_search_result_sanitizer_middleware.go new file mode 100644 index 00000000..df9bb247 --- /dev/null +++ b/internal/multiagent/tool_search_result_sanitizer_middleware.go @@ -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 +} diff --git a/internal/multiagent/tool_search_result_sanitizer_middleware_test.go b/internal/multiagent/tool_search_result_sanitizer_middleware_test.go new file mode 100644 index 00000000..c5b026cc --- /dev/null +++ b/internal/multiagent/tool_search_result_sanitizer_middleware_test.go @@ -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: "502 Bad Gateway"} + 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 != "502 Bad Gateway" { + 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) + } + } +}