CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-07-16 08:57:30 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-12 13:20:08 +08:00
committed by GitHub
parent 0f1a6ad25a
commit fe3c845ff8
3 changed files with 136 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/multiagent/eino_chat_model_tail_middleware.go
+4 -2
View File
@@ -18,8 +18,9 @@ import (
// 5. soft model-input budget (warn/compact only, never fail locally) // 5. soft model-input budget (warn/compact only, never fail locally)
// 6. final tool-call/result reconciliation // 6. final tool-call/result reconciliation
// 7. orphan tool prune (defense in depth) // 7. orphan tool prune (defense in depth)
// 8. telemetry // 8. malformed tool_search history repair
// 9. model-facing trace snapshot // 9. telemetry
// 10. model-facing trace snapshot
type einoChatModelTailConfig struct { type einoChatModelTailConfig struct {
logger *zap.Logger logger *zap.Logger
phase string phase string
@@ -48,6 +49,7 @@ func appendEinoChatModelTailMiddlewares(handlers []adk.ChatModelAgentMiddleware,
if !cfg.skipOrphanPruner { if !cfg.skipOrphanPruner {
handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase)) handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase))
} }
handlers = append(handlers, newToolSearchResultSanitizerMiddleware(cfg.logger, cfg.phase))
if !cfg.skipTelemetry { if !cfg.skipTelemetry {
if teleMw := newEinoModelInputTelemetryMiddleware(cfg.logger, cfg.modelName, cfg.conversationID, cfg.phase); teleMw != nil { if teleMw := newEinoModelInputTelemetryMiddleware(cfg.logger, cfg.modelName, cfg.conversationID, cfg.phase); teleMw != nil {
handlers = append(handlers, teleMw) handlers = append(handlers, teleMw)
internal/multiagent/tool_search_result_sanitizer_middleware.go
+76
View File
@@ -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
}
internal/multiagent/tool_search_result_sanitizer_middleware_test.go
+56
View File
@@ -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)
}
}
}