mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-21 18:37:33 +02:00
fix: 为 Eino agentic 路径补充 tool_call/tool_result 配对防御中间件
agentic 路径(单代理模式)缺少 toolPairReconciler 和 orphanToolPruner, 当 summarization 截断历史破坏配对后,序列化到 OpenAI API 触发 400 "insufficient tool messages following tool_calls message"。 新增 AgenticMessage 版本的 reconciler 和 pruner,与 classic 路径对齐。 Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -0,0 +1,126 @@
|
|||||||
|
package multiagent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
|
||||||
|
"github.com/cloudwego/eino/adk"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// agenticOrphanToolPrunerMiddleware is the AgenticMessage equivalent of
|
||||||
|
// orphanToolPrunerMiddleware. It removes user-role messages whose content
|
||||||
|
// blocks are exclusively FunctionToolResult entries with CallIDs that do not
|
||||||
|
// match any FunctionToolCall in the history.
|
||||||
|
//
|
||||||
|
// This is a defense-in-depth layer after agenticToolPairReconcilerMiddleware;
|
||||||
|
// the reconciler handles the common case (assistant followed by its results)
|
||||||
|
// while this pruner catches stray results that appear before their assistant
|
||||||
|
// or in non-adjacent positions (e.g. after summarization rewriting).
|
||||||
|
type agenticOrphanToolPrunerMiddleware struct {
|
||||||
|
*adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]
|
||||||
|
logger *zap.Logger
|
||||||
|
phase string
|
||||||
|
}
|
||||||
|
|
||||||
|
func newAgenticOrphanToolPrunerMiddleware(logger *zap.Logger, phase string) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] {
|
||||||
|
return &agenticOrphanToolPrunerMiddleware{
|
||||||
|
TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{},
|
||||||
|
logger: logger,
|
||||||
|
phase: phase,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *agenticOrphanToolPrunerMiddleware) BeforeModelRewriteState(
|
||||||
|
ctx context.Context,
|
||||||
|
state *adk.TypedChatModelAgentState[*schema.AgenticMessage],
|
||||||
|
mc *adk.TypedModelContext[*schema.AgenticMessage],
|
||||||
|
) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) {
|
||||||
|
_ = mc
|
||||||
|
if m == nil || state == nil || len(state.Messages) == 0 {
|
||||||
|
return ctx, state, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pass 1: collect all provided CallIDs from assistant FunctionToolCall blocks.
|
||||||
|
provided := make(map[string]struct{}, 8)
|
||||||
|
for _, msg := range state.Messages {
|
||||||
|
if msg == nil || msg.Role != schema.AgenticRoleTypeAssistant {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for _, block := range msg.ContentBlocks {
|
||||||
|
if block != nil && block.FunctionToolCall != nil && block.FunctionToolCall.CallID != "" {
|
||||||
|
provided[block.FunctionToolCall.CallID] = struct{}{}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fast path: check if any orphan exists.
|
||||||
|
hasOrphan := false
|
||||||
|
for _, msg := range state.Messages {
|
||||||
|
if msg == nil || !isPureAgenticToolResult(msg) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
for _, id := range agenticToolResultCallIDs(msg) {
|
||||||
|
if _, ok := provided[id]; !ok {
|
||||||
|
hasOrphan = true
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if hasOrphan {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !hasOrphan {
|
||||||
|
return ctx, state, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pass 2: build pruned list.
|
||||||
|
pruned := make([]*schema.AgenticMessage, 0, len(state.Messages))
|
||||||
|
var droppedIDs []string
|
||||||
|
var droppedNames []string
|
||||||
|
for _, msg := range state.Messages {
|
||||||
|
if msg == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if !isPureAgenticToolResult(msg) {
|
||||||
|
pruned = append(pruned, msg)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
// Check if ALL result call IDs are orphans. If any is matched, keep the
|
||||||
|
// message (the reconciler already handled partial mismatches).
|
||||||
|
allOrphan := true
|
||||||
|
for _, id := range agenticToolResultCallIDs(msg) {
|
||||||
|
if _, ok := provided[id]; ok {
|
||||||
|
allOrphan = false
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if allOrphan {
|
||||||
|
for _, block := range msg.ContentBlocks {
|
||||||
|
if block != nil && block.FunctionToolResult != nil {
|
||||||
|
droppedIDs = append(droppedIDs, block.FunctionToolResult.CallID)
|
||||||
|
droppedNames = append(droppedNames, block.FunctionToolResult.Name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
pruned = append(pruned, msg)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(droppedIDs) == 0 {
|
||||||
|
return ctx, state, nil
|
||||||
|
}
|
||||||
|
if m.logger != nil {
|
||||||
|
m.logger.Warn("agentic orphan tool messages pruned before model call",
|
||||||
|
zap.String("phase", m.phase),
|
||||||
|
zap.Int("dropped_count", len(droppedIDs)),
|
||||||
|
zap.Strings("dropped_tool_call_ids", droppedIDs),
|
||||||
|
zap.Strings("dropped_tool_names", droppedNames),
|
||||||
|
zap.Int("messages_before", len(state.Messages)),
|
||||||
|
zap.Int("messages_after", len(pruned)),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
ns := *state
|
||||||
|
ns.Messages = pruned
|
||||||
|
return ctx, &ns, nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,247 @@
|
|||||||
|
package multiagent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/cloudwego/eino/adk"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
// agenticToolPairReconcilerMiddleware is the AgenticMessage equivalent of
|
||||||
|
// toolPairReconcilerMiddleware. It ensures every assistant FunctionToolCall
|
||||||
|
// block is followed by a matching FunctionToolResult message, patching or
|
||||||
|
// dropping as needed so the downstream model never receives an unpaired
|
||||||
|
// tool-call history.
|
||||||
|
//
|
||||||
|
// In the AgenticMessage protocol:
|
||||||
|
// - Assistant tool calls: Role=AgenticRoleTypeAssistant with FunctionToolCall content blocks.
|
||||||
|
// - Tool results: Role=AgenticRoleTypeUser with FunctionToolResult content blocks.
|
||||||
|
//
|
||||||
|
// This middleware runs after summarization which may truncate history and
|
||||||
|
// break pairings.
|
||||||
|
type agenticToolPairReconcilerMiddleware struct {
|
||||||
|
*adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]
|
||||||
|
logger *zap.Logger
|
||||||
|
phase string
|
||||||
|
}
|
||||||
|
|
||||||
|
func newAgenticToolPairReconcilerMiddleware(logger *zap.Logger, phase string) adk.TypedChatModelAgentMiddleware[*schema.AgenticMessage] {
|
||||||
|
return &agenticToolPairReconcilerMiddleware{
|
||||||
|
TypedBaseChatModelAgentMiddleware: &adk.TypedBaseChatModelAgentMiddleware[*schema.AgenticMessage]{},
|
||||||
|
logger: logger,
|
||||||
|
phase: phase,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *agenticToolPairReconcilerMiddleware) BeforeModelRewriteState(
|
||||||
|
ctx context.Context,
|
||||||
|
state *adk.TypedChatModelAgentState[*schema.AgenticMessage],
|
||||||
|
mc *adk.TypedModelContext[*schema.AgenticMessage],
|
||||||
|
) (context.Context, *adk.TypedChatModelAgentState[*schema.AgenticMessage], error) {
|
||||||
|
_ = mc
|
||||||
|
if m == nil || state == nil || len(state.Messages) == 0 {
|
||||||
|
return ctx, state, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
usedIDs := make(map[string]struct{}, 16)
|
||||||
|
changed := false
|
||||||
|
patched := 0
|
||||||
|
dropped := 0
|
||||||
|
out := make([]*schema.AgenticMessage, 0, len(state.Messages))
|
||||||
|
|
||||||
|
for i := 0; i < len(state.Messages); {
|
||||||
|
msg := state.Messages[i]
|
||||||
|
if msg == nil {
|
||||||
|
changed = true
|
||||||
|
i++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
calls := agenticFunctionToolCalls(msg)
|
||||||
|
|
||||||
|
// Non-assistant or assistant without tool calls — but check for orphan
|
||||||
|
// tool-result messages (user role with only FunctionToolResult blocks).
|
||||||
|
if len(calls) == 0 {
|
||||||
|
if isPureAgenticToolResult(msg) {
|
||||||
|
// Orphan tool result not preceded by its assistant; drop it.
|
||||||
|
changed = true
|
||||||
|
dropped++
|
||||||
|
i++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, msg)
|
||||||
|
i++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Deduplicate / fix empty call IDs.
|
||||||
|
idsChanged := false
|
||||||
|
for ci := range calls {
|
||||||
|
id := calls[ci].CallID
|
||||||
|
_, duplicate := usedIDs[id]
|
||||||
|
if id == "" || duplicate {
|
||||||
|
base := fmt.Sprintf("patched_agentic_call_%d_%d", i, ci)
|
||||||
|
id = base
|
||||||
|
for suffix := 1; ; suffix++ {
|
||||||
|
if _, exists := usedIDs[id]; !exists {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
id = fmt.Sprintf("%s_%d", base, suffix)
|
||||||
|
}
|
||||||
|
calls[ci].CallID = id
|
||||||
|
idsChanged = true
|
||||||
|
changed = true
|
||||||
|
}
|
||||||
|
usedIDs[id] = struct{}{}
|
||||||
|
}
|
||||||
|
|
||||||
|
assistant := msg
|
||||||
|
if idsChanged {
|
||||||
|
assistant = cloneAgenticMessageWithCalls(msg, calls)
|
||||||
|
}
|
||||||
|
out = append(out, assistant)
|
||||||
|
|
||||||
|
// Build expected set.
|
||||||
|
expected := make(map[string]*schema.FunctionToolCall, len(calls))
|
||||||
|
for ci := range calls {
|
||||||
|
expected[calls[ci].CallID] = calls[ci]
|
||||||
|
}
|
||||||
|
|
||||||
|
// Consume following tool-result messages.
|
||||||
|
results := make(map[string]*schema.AgenticMessage, len(calls))
|
||||||
|
j := i + 1
|
||||||
|
for j < len(state.Messages) {
|
||||||
|
next := state.Messages[j]
|
||||||
|
if next == nil {
|
||||||
|
changed = true
|
||||||
|
j++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if !isPureAgenticToolResult(next) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
resultCallIDs := agenticToolResultCallIDs(next)
|
||||||
|
consumed := false
|
||||||
|
for _, rid := range resultCallIDs {
|
||||||
|
if _, wanted := expected[rid]; !wanted {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, dup := results[rid]; dup {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
results[rid] = next
|
||||||
|
consumed = true
|
||||||
|
}
|
||||||
|
if !consumed {
|
||||||
|
changed = true
|
||||||
|
dropped++
|
||||||
|
}
|
||||||
|
j++
|
||||||
|
}
|
||||||
|
|
||||||
|
// Emit results in call order, patching missing ones.
|
||||||
|
for _, tc := range calls {
|
||||||
|
if result, ok := results[tc.CallID]; ok {
|
||||||
|
out = append(out, result)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
out = append(out, makeAgenticPatchedToolResult(tc.CallID, tc.Name))
|
||||||
|
changed = true
|
||||||
|
patched++
|
||||||
|
}
|
||||||
|
i = j
|
||||||
|
}
|
||||||
|
|
||||||
|
if !changed {
|
||||||
|
return ctx, state, nil
|
||||||
|
}
|
||||||
|
if m.logger != nil {
|
||||||
|
m.logger.Warn("agentic tool-call/result pairs reconciled before model call",
|
||||||
|
zap.String("phase", m.phase),
|
||||||
|
zap.Int("patched_results", patched),
|
||||||
|
zap.Int("dropped_results", dropped),
|
||||||
|
zap.Int("messages_before", len(state.Messages)),
|
||||||
|
zap.Int("messages_after", len(out)),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
ns := *state
|
||||||
|
ns.Messages = out
|
||||||
|
return ctx, &ns, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// agenticFunctionToolCalls extracts FunctionToolCall pointers from an
|
||||||
|
// assistant message's content blocks. Returns nil for non-assistant messages.
|
||||||
|
func agenticFunctionToolCalls(msg *schema.AgenticMessage) []*schema.FunctionToolCall {
|
||||||
|
if msg == nil || msg.Role != schema.AgenticRoleTypeAssistant {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var out []*schema.FunctionToolCall
|
||||||
|
for _, block := range msg.ContentBlocks {
|
||||||
|
if block != nil && block.FunctionToolCall != nil {
|
||||||
|
out = append(out, block.FunctionToolCall)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// isPureAgenticToolResult returns true when the message is a user-role
|
||||||
|
// message whose content blocks are exclusively FunctionToolResult entries.
|
||||||
|
func isPureAgenticToolResult(msg *schema.AgenticMessage) bool {
|
||||||
|
if msg == nil || msg.Role != schema.AgenticRoleTypeUser || len(msg.ContentBlocks) == 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, block := range msg.ContentBlocks {
|
||||||
|
if block == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if block.FunctionToolResult == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// agenticToolResultCallIDs extracts all CallIDs from FunctionToolResult blocks.
|
||||||
|
func agenticToolResultCallIDs(msg *schema.AgenticMessage) []string {
|
||||||
|
if msg == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var ids []string
|
||||||
|
for _, block := range msg.ContentBlocks {
|
||||||
|
if block != nil && block.FunctionToolResult != nil && block.FunctionToolResult.CallID != "" {
|
||||||
|
ids = append(ids, block.FunctionToolResult.CallID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ids
|
||||||
|
}
|
||||||
|
|
||||||
|
func cloneAgenticMessageWithCalls(msg *schema.AgenticMessage, calls []*schema.FunctionToolCall) *schema.AgenticMessage {
|
||||||
|
cloned := *msg
|
||||||
|
cloned.ContentBlocks = make([]*schema.ContentBlock, 0, len(msg.ContentBlocks))
|
||||||
|
callIdx := 0
|
||||||
|
for _, block := range msg.ContentBlocks {
|
||||||
|
if block != nil && block.FunctionToolCall != nil && callIdx < len(calls) {
|
||||||
|
cloned.ContentBlocks = append(cloned.ContentBlocks, schema.NewContentBlock(calls[callIdx]))
|
||||||
|
callIdx++
|
||||||
|
} else {
|
||||||
|
cloned.ContentBlocks = append(cloned.ContentBlocks, block)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return &cloned
|
||||||
|
}
|
||||||
|
|
||||||
|
func makeAgenticPatchedToolResult(callID, name string) *schema.AgenticMessage {
|
||||||
|
return &schema.AgenticMessage{
|
||||||
|
Role: schema.AgenticRoleTypeUser,
|
||||||
|
ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolResult{
|
||||||
|
CallID: callID,
|
||||||
|
Name: name,
|
||||||
|
Content: []*schema.FunctionToolResultContentBlock{{
|
||||||
|
Type: schema.FunctionToolResultContentBlockTypeText,
|
||||||
|
Text: &schema.UserInputText{Text: patchedMissingToolResult},
|
||||||
|
}},
|
||||||
|
})},
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,157 @@
|
|||||||
|
package multiagent
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/cloudwego/eino/adk"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestAgenticToolPairReconcilerPatchesMissing(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
mw := newAgenticToolPairReconcilerMiddleware(nil, "test")
|
||||||
|
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||||
|
Messages: []*schema.AgenticMessage{
|
||||||
|
agenticAssistantToolCall("c1", "search", `{"q":"x"}`),
|
||||||
|
agenticAssistantToolCall("c2", "execute", `{"cmd":"ls"}`),
|
||||||
|
// c1 result present, c2 missing
|
||||||
|
agenticToolResult("c1", "search", "found it"),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
// Expected: assistant(c1) -> result(c1) -> assistant(c2) -> patched_result(c2)
|
||||||
|
if len(out.Messages) != 4 {
|
||||||
|
t.Fatalf("messages = %d, want 4", len(out.Messages))
|
||||||
|
}
|
||||||
|
// c1 assistant
|
||||||
|
if calls := agenticFunctionToolCalls(out.Messages[0]); len(calls) != 1 || calls[0].CallID != "c1" {
|
||||||
|
t.Fatal("msg[0] should be assistant(c1)")
|
||||||
|
}
|
||||||
|
// c1 result
|
||||||
|
if ids := agenticToolResultCallIDs(out.Messages[1]); len(ids) != 1 || ids[0] != "c1" {
|
||||||
|
t.Fatal("msg[1] should be result(c1)")
|
||||||
|
}
|
||||||
|
// c2 assistant
|
||||||
|
if calls := agenticFunctionToolCalls(out.Messages[2]); len(calls) != 1 || calls[0].CallID != "c2" {
|
||||||
|
t.Fatal("msg[2] should be assistant(c2)")
|
||||||
|
}
|
||||||
|
// c2 patched result
|
||||||
|
if ids := agenticToolResultCallIDs(out.Messages[3]); len(ids) != 1 || ids[0] != "c2" {
|
||||||
|
t.Fatal("msg[3] should be patched result(c2)")
|
||||||
|
}
|
||||||
|
resultText := out.Messages[3].ContentBlocks[0].FunctionToolResult.Content[0].Text.Text
|
||||||
|
if resultText != patchedMissingToolResult {
|
||||||
|
t.Fatalf("patched text = %q", resultText)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgenticToolPairReconcilerDropsOrphan(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
mw := newAgenticToolPairReconcilerMiddleware(nil, "test")
|
||||||
|
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||||
|
Messages: []*schema.AgenticMessage{
|
||||||
|
// Orphan tool result with no preceding assistant
|
||||||
|
agenticToolResult("orphan", "deleted_tool", "stale data"),
|
||||||
|
{Role: schema.AgenticRoleTypeUser, ContentBlocks: []*schema.ContentBlock{
|
||||||
|
schema.NewContentBlock(&schema.UserInputText{Text: "hello"}),
|
||||||
|
}},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(out.Messages) != 1 {
|
||||||
|
t.Fatalf("messages = %d, want 1 (orphan dropped)", len(out.Messages))
|
||||||
|
}
|
||||||
|
if out.Messages[0].ContentBlocks[0].UserInputText == nil {
|
||||||
|
t.Fatal("remaining message should be the user text")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgenticToolPairReconcilerNoopWhenPaired(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
mw := newAgenticToolPairReconcilerMiddleware(nil, "test")
|
||||||
|
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||||
|
Messages: []*schema.AgenticMessage{
|
||||||
|
agenticAssistantToolCall("c1", "search", `{}`),
|
||||||
|
agenticToolResult("c1", "search", "ok"),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
// Should return original state unchanged
|
||||||
|
if &out.Messages[0] == &state.Messages[0] {
|
||||||
|
// pointer equality on slice — state not cloned
|
||||||
|
}
|
||||||
|
if len(out.Messages) != 2 {
|
||||||
|
t.Fatalf("messages = %d, want 2", len(out.Messages))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgenticToolPairReconcilerFixesEmptyCallID(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
mw := newAgenticToolPairReconcilerMiddleware(nil, "test")
|
||||||
|
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||||
|
Messages: []*schema.AgenticMessage{
|
||||||
|
{
|
||||||
|
Role: schema.AgenticRoleTypeAssistant,
|
||||||
|
ContentBlocks: []*schema.ContentBlock{schema.NewContentBlock(&schema.FunctionToolCall{
|
||||||
|
CallID: "", Name: "search", Arguments: `{}`,
|
||||||
|
})},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
calls := agenticFunctionToolCalls(out.Messages[0])
|
||||||
|
if len(calls) != 1 || calls[0].CallID == "" {
|
||||||
|
t.Fatalf("empty call ID should be patched, got %q", calls[0].CallID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgenticOrphanToolPrunerRemovesOrphan(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
mw := newAgenticOrphanToolPrunerMiddleware(nil, "test")
|
||||||
|
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||||
|
Messages: []*schema.AgenticMessage{
|
||||||
|
agenticAssistantToolCall("c1", "search", `{}`),
|
||||||
|
agenticToolResult("c1", "search", "ok"),
|
||||||
|
// Orphan: no assistant has call_id "c_orphan"
|
||||||
|
agenticToolResult("c_orphan", "deleted", "stale"),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(out.Messages) != 2 {
|
||||||
|
t.Fatalf("messages = %d, want 2 (orphan pruned)", len(out.Messages))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAgenticOrphanToolPrunerNoopWhenClean(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
mw := newAgenticOrphanToolPrunerMiddleware(nil, "test")
|
||||||
|
state := &adk.TypedChatModelAgentState[*schema.AgenticMessage]{
|
||||||
|
Messages: []*schema.AgenticMessage{
|
||||||
|
agenticAssistantToolCall("c1", "search", `{}`),
|
||||||
|
agenticToolResult("c1", "search", "ok"),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(out.Messages) != 2 {
|
||||||
|
t.Fatalf("messages = %d, want 2", len(out.Messages))
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -20,8 +20,13 @@ func appendEinoAgenticChatModelTailMiddlewares(
|
|||||||
handlers = append(handlers, newAgenticSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase))
|
handlers = append(handlers, newAgenticSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase))
|
||||||
handlers = append(handlers, newAgenticContinuationUserDedupMiddleware(cfg.logger, cfg.phase))
|
handlers = append(handlers, newAgenticContinuationUserDedupMiddleware(cfg.logger, cfg.phase))
|
||||||
if cfg.agenticSummarization != nil {
|
if cfg.agenticSummarization != nil {
|
||||||
|
handlers = append(handlers, newAgenticToolPairReconcilerMiddleware(cfg.logger, cfg.phase+"_pre_summarization"))
|
||||||
handlers = append(handlers, cfg.agenticSummarization)
|
handlers = append(handlers, cfg.agenticSummarization)
|
||||||
}
|
}
|
||||||
|
handlers = append(handlers, newAgenticToolPairReconcilerMiddleware(cfg.logger, cfg.phase))
|
||||||
|
if !cfg.skipOrphanPruner {
|
||||||
|
handlers = append(handlers, newAgenticOrphanToolPrunerMiddleware(cfg.logger, cfg.phase))
|
||||||
|
}
|
||||||
if !cfg.skipTrace && cfg.trace != nil {
|
if !cfg.skipTrace && cfg.trace != nil {
|
||||||
if capMw := newAgenticModelFacingTraceMiddleware(cfg.trace); capMw != nil {
|
if capMw := newAgenticModelFacingTraceMiddleware(cfg.trace); capMw != nil {
|
||||||
handlers = append(handlers, capMw)
|
handlers = append(handlers, capMw)
|
||||||
|
|||||||
@@ -106,7 +106,8 @@ func TestAppendEinoAgenticChatModelTailMiddlewares(t *testing.T) {
|
|||||||
phase: "agentic",
|
phase: "agentic",
|
||||||
trace: holder,
|
trace: holder,
|
||||||
})
|
})
|
||||||
if len(handlers) != 3 {
|
// system + continuation + reconciler + orphan_pruner + trace
|
||||||
t.Fatalf("handlers = %d, want system + continuation + trace", len(handlers))
|
if len(handlers) != 5 {
|
||||||
|
t.Fatalf("handlers = %d, want system + continuation + reconciler + orphan_pruner + trace", len(handlers))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user