Add files via upload

This commit is contained in:
公明
2026-07-10 16:48:49 +08:00
committed by GitHub
parent 0ff8c58fbd
commit 542d1d2411
19 changed files with 922 additions and 114 deletions
+31 -14
View File
@@ -79,7 +79,7 @@ type einoADKRunLoopArgs struct {
EinoRoleTag func(agent string) string
CheckpointDir string
// RunRetryMaxAttempts / RunRetryMaxBackoffSec429、5xx、网络抖动时的指数退避续跑(0=默认 10 次 / 30s 上限)。
RunRetryMaxAttempts int
RunRetryMaxAttempts int
RunRetryMaxBackoffSec int
McpIDsMu *sync.Mutex
@@ -301,9 +301,6 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
tryEmitToolResultProgress := func(toolName, content, toolCallID string, isErr bool, agentName string) {
// 仅由 ADK schema.Tool 事件调用;MCP/execute 桥在 reduction 前的 ToolInvokeNotify 不得推送 tool_result
// 否则全量输出会先占位并触发 toolResultSent 去重,导致 UI/监控展示与 agent 实际收到的截断正文不一致。
if progress == nil {
return
}
toolName = strings.TrimSpace(toolName)
if toolName == "" {
toolName = "unknown"
@@ -351,7 +348,9 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
args.FilesystemMonitorAgent.UpdateMCPExecutionDisplayResult(execID, content)
}
}
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
if progress != nil {
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
}
}
if args.EinoCallbacks != nil {
@@ -560,7 +559,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
}
ids := snapshotMCPIDs()
return buildEinoRunResultFromAccumulated(
orchMode, runAccumulatedMsgs, persistTraceSource(args, runAccumulatedMsgs),
orchMode, runAccumulatedMsgs, modelFacingTraceSnapshot(args),
lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, true,
), runErr
}
@@ -840,7 +839,7 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
"einoRole": "orchestrator",
"einoAgent": ev.AgentName,
"orchestration": orchMode,
"iteration": einoMainRound,
"iteration": einoMainRound,
"streamId": mainStreamID,
}, mainAssistantBuf))
mainAssistWireAccum, _ = normalizeStreamingDelta(mainAssistWireAccum, contentDelta)
@@ -1089,19 +1088,22 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
mcpIDsMu.Unlock()
out := buildEinoRunResultFromAccumulated(
orchMode, runAccumulatedMsgs, persistTraceSource(args, runAccumulatedMsgs),
orchMode, runAccumulatedMsgs, modelFacingTraceSnapshot(args),
lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, false,
)
return out, nil
}
func persistTraceSource(args *einoADKRunLoopArgs, fallback []adk.Message) []adk.Message {
// modelFacingTraceSnapshot returns only the state that actually reached the model boundary.
// Never fall back to event-stream accumulation here: it can contain pre-reduction tool output
// that the model never received (for example when summarization failed before the first call).
func modelFacingTraceSnapshot(args *einoADKRunLoopArgs) []adk.Message {
if args != nil && args.ModelFacingTrace != nil {
if snap := args.ModelFacingTrace.Snapshot(); len(snap) > 0 {
return snap
}
}
return fallback
return nil
}
func einoPartialRunLastOutputHint() string {
@@ -1233,10 +1235,13 @@ func buildEinoRunResultFromAccumulated(
partial bool,
) *RunResult {
traceForJSON := persistMsgs
if len(traceForJSON) == 0 {
traceForJSON = runAccumulatedMsgs
traceJSON := ""
if len(traceForJSON) > 0 {
traceForJSON = markModelFacingTraceForPersistence(traceForJSON)
if histJSON, err := json.Marshal(traceForJSON); err == nil {
traceJSON = string(histJSON)
}
}
histJSON, _ := json.Marshal(traceForJSON)
cleaned := strings.TrimSpace(lastAssistant)
if orchMode == "plan_execute" {
if e := strings.TrimSpace(lastPlanExecuteExecutor); e != "" {
@@ -1266,7 +1271,7 @@ func buildEinoRunResultFromAccumulated(
out := &RunResult{
Response: resp,
MCPExecutionIDs: mcpIDs,
LastAgentTraceInput: string(histJSON),
LastAgentTraceInput: traceJSON,
LastAgentTraceOutput: lastOut,
}
if !partial && out.Response == "" {
@@ -1276,6 +1281,18 @@ func buildEinoRunResultFromAccumulated(
return out
}
func markModelFacingTraceForPersistence(msgs []adk.Message) []adk.Message {
out := cloneADKMessagesForTrace(msgs)
if len(out) == 0 || out[0] == nil {
return out
}
if out[0].Extra == nil {
out[0].Extra = make(map[string]any, 1)
}
out[0].Extra[agent.ModelFacingTraceVersionKey] = 1
return out
}
// einoExtractFallbackAssistantFromMsgs 在「主通道未产出助手正文」时,从 Eino ADK 轨迹中回填用户可见回复。
// 典型场景:监督者仅调用 exitfinal_result 落在 Tool 消息中),或工具结果已写入历史但 lastAssistant 未更新。
//
@@ -11,15 +11,19 @@ import (
// Order (best practice):
// 1. system merge — accurate token count for summarization
// 2. continuation user dedup — drop stale session-resume injections
// 3. summarization
// 4. orphan tool prune
// 5. telemetry
// 6. model-facing trace snapshot
// 3. pre-summarization tool-call/result reconciliation
// 4. summarization
// 5. total model-input hard budget
// 6. final tool-call/result reconciliation
// 7. orphan tool prune (defense in depth)
// 8. telemetry
// 9. model-facing trace snapshot
type einoChatModelTailConfig struct {
logger *zap.Logger
phase string
summarization adk.ChatModelAgentMiddleware
modelName string
maxTotalTokens int
conversationID string
trace *modelFacingTraceHolder
skipOrphanPruner bool
@@ -31,8 +35,13 @@ func appendEinoChatModelTailMiddlewares(handlers []adk.ChatModelAgentMiddleware,
handlers = append(handlers, newSystemMessageNormalizerMiddleware(cfg.logger, cfg.phase))
handlers = append(handlers, newContinuationUserDedupMiddleware(cfg.logger, cfg.phase))
if cfg.summarization != nil {
// Summarization invokes the model internally, so its input needs the same
// structural guarantee as the agent's final model call.
handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase+"_pre_summarization"))
handlers = append(handlers, cfg.summarization)
}
handlers = append(handlers, newModelInputBudgetMiddleware(cfg.maxTotalTokens, cfg.modelName, cfg.logger, cfg.phase))
handlers = append(handlers, newToolPairReconcilerMiddleware(cfg.logger, cfg.phase))
if !cfg.skipOrphanPruner {
handlers = append(handlers, newOrphanToolPrunerMiddleware(cfg.logger, cfg.phase))
}
@@ -0,0 +1,34 @@
package multiagent
import (
"strings"
"testing"
"cyberstrike-ai/internal/agent"
"github.com/cloudwego/eino/schema"
)
func TestBuildEinoRunResultNeverPersistsRawAccumulationWithoutModelFacingTrace(t *testing.T) {
raw := []schema.Message{*schema.ToolMessage(strings.Repeat("raw-tool-output", 1000), "call-1")}
rawMsgs := make([]*schema.Message, len(raw))
for i := range raw {
rawMsgs[i] = &raw[i]
}
result := buildEinoRunResultFromAccumulated("deep", rawMsgs, nil, "", "", "empty", nil, true)
if result.LastAgentTraceInput != "" {
t.Fatalf("pre-model raw accumulation must not be persisted: %d bytes", len(result.LastAgentTraceInput))
}
modelFacing := []*schema.Message{schema.UserMessage("bounded-model-view")}
result = buildEinoRunResultFromAccumulated("deep", rawMsgs, modelFacing, "ok", "", "empty", nil, false)
if !strings.Contains(result.LastAgentTraceInput, "bounded-model-view") {
t.Fatalf("model-facing trace missing: %s", result.LastAgentTraceInput)
}
if strings.Contains(result.LastAgentTraceInput, "raw-tool-output") {
t.Fatal("raw accumulation leaked into persisted model-facing trace")
}
if !agent.IsModelFacingTraceJSON(result.LastAgentTraceInput) {
t.Fatal("persisted model-facing trace is missing its version marker")
}
}
@@ -131,6 +131,7 @@ func buildPlanExecuteExecutorHandlers(ctx context.Context, a *PlanExecuteRootArg
phase: "plan_execute_executor",
summarization: sumMw,
modelName: a.ModelName,
maxTotalTokens: a.AppCfg.OpenAI.MaxTotalTokens,
conversationID: a.ConversationID,
trace: a.ModelFacingTrace,
})
@@ -150,6 +150,7 @@ func RunEinoSingleChatModelAgent(
phase: "eino_single",
summarization: mainSumMw,
modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
conversationID: conversationID,
trace: modelFacingTrace,
})
+31 -7
View File
@@ -18,8 +18,9 @@ import (
"go.uber.org/zap"
)
// prepareEinoSkills builds Eino official skill backend + middleware, and a shared local disk backend
// for skill discovery and (optionally) filesystem/execute tools. Returns nils when disabled or dir missing.
// prepareEinoSkills builds Eino official skill backend + middleware, and a shared local disk backend.
// The local backend is also required by reduction, so reduction must not silently disappear merely
// because Skills are disabled or skills_dir is unavailable.
// skillsRoot is the absolute skills directory (empty when skills are not active).
func prepareEinoSkills(
ctx context.Context,
@@ -27,15 +28,34 @@ func prepareEinoSkills(
ma *config.MultiAgentConfig,
logger *zap.Logger,
) (loc *localbk.Local, skillMW adk.ChatModelAgentMiddleware, fsTools bool, skillsRoot string, err error) {
if ma == nil || ma.EinoSkills.Disable {
if ma == nil {
return nil, nil, false, "", nil
}
needLocalBackend := ma.EinoMiddleware.ReductionEnable
newLocalBackend := func() (*localbk.Local, error) {
backend, backendErr := localbk.NewBackend(ctx, &localbk.Config{})
if backendErr != nil {
return nil, fmt.Errorf("eino local backend: %w", backendErr)
}
return backend, nil
}
if ma.EinoSkills.Disable {
if !needLocalBackend {
return nil, nil, false, "", nil
}
loc, err = newLocalBackend()
return loc, nil, false, "", err
}
root := strings.TrimSpace(skillsDir)
if root == "" {
if logger != nil {
logger.Warn("eino skills: skills_dir empty, skip")
}
return nil, nil, false, "", nil
if !needLocalBackend {
return nil, nil, false, "", nil
}
loc, err = newLocalBackend()
return loc, nil, false, "", err
}
abs, err := filepath.Abs(root)
if err != nil {
@@ -45,12 +65,16 @@ func prepareEinoSkills(
if logger != nil {
logger.Warn("eino skills: directory missing, skip", zap.String("dir", abs), zap.Error(err))
}
return nil, nil, false, "", nil
if !needLocalBackend {
return nil, nil, false, "", nil
}
loc, err = newLocalBackend()
return loc, nil, false, "", err
}
loc, err = localbk.NewBackend(ctx, &localbk.Config{})
loc, err = newLocalBackend()
if err != nil {
return nil, nil, false, "", fmt.Errorf("eino local backend: %w", err)
return nil, nil, false, "", err
}
skillBE, err := skill.NewBackendFromFilesystem(ctx, &skill.BackendFromFilesystemConfig{
@@ -0,0 +1,27 @@
package multiagent
import (
"context"
"testing"
"cyberstrike-ai/internal/config"
)
func TestPrepareEinoSkillsStillCreatesReductionBackendWhenSkillsDisabled(t *testing.T) {
ma := &config.MultiAgentConfig{
EinoSkills: config.MultiAgentEinoSkillsConfig{Disable: true},
EinoMiddleware: config.MultiAgentEinoMiddlewareConfig{
ReductionEnable: true,
},
}
loc, skillMW, fsTools, skillsRoot, err := prepareEinoSkills(context.Background(), "", ma, nil)
if err != nil {
t.Fatal(err)
}
if loc == nil {
t.Fatal("reduction backend must exist even when Skills are disabled")
}
if skillMW != nil || fsTools || skillsRoot != "" {
t.Fatalf("Skills unexpectedly enabled: mw=%v fs=%v root=%q", skillMW, fsTools, skillsRoot)
}
}
+94 -47
View File
@@ -6,7 +6,6 @@ import (
"os"
"path/filepath"
"strings"
"time"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
@@ -107,6 +106,12 @@ func newEinoSummarizationMiddleware(
userLedgerMaxRunes = mwCfg.SummarizationUserIntentLedgerMaxRunesEffective()
userLedgerEntryMaxRunes = mwCfg.SummarizationUserIntentLedgerEntryMaxRunesEffective()
}
// The ledger is merged into the leading system message and cannot be removed as
// an ordinary conversation round. Bound it relative to the configured window so
// it cannot crowd out the summary/latest turn or trip the final fail-closed guard.
ledgerWindowCap := modelFacingRuneBudget(maxTotal, 0.20)
userLedgerMaxRunes = minPositiveInt(userLedgerMaxRunes, ledgerWindowCap)
userLedgerEntryMaxRunes = minPositiveInt(userLedgerEntryMaxRunes, userLedgerMaxRunes)
// Keep enough safety margin for tokenizer/model-side accounting mismatch.
trigger := int(float64(maxTotal) * triggerRatio)
if trigger < 4096 {
@@ -132,6 +137,14 @@ func newEinoSummarizationMiddleware(
if recentTrailMax > trigger/2 {
recentTrailMax = trigger / 2
}
// The summarization request itself needs output headroom. A trigger is not a hard
// request limit: one large turn can jump far beyond it. Bound the actual summary
// model input to 60% of the configured context window and keep complete recent
// rounds so tool_call/tool_result pairs are never split.
summaryInputMax := int(float64(maxTotal) * 0.6)
if summaryInputMax < 4096 {
summaryInputMax = 4096
}
transcriptPath := ""
if conv := strings.TrimSpace(conversationID); conv != "" {
baseRoot := filepath.Join(os.TempDir(), "cyberstrike-summarization")
@@ -169,6 +182,30 @@ func newEinoSummarizationMiddleware(
mw, err := summarization.New(ctx, &summarization.Config{
Model: summaryModel,
ModelOptions: summaryModelOpts,
GenModelInput: func(ctx context.Context, sysInstruction, userInstruction adk.Message, originalMsgs []adk.Message) ([]adk.Message, error) {
if transcriptPath != "" && len(originalMsgs) > 0 {
if werr := writeSummarizationTranscript(transcriptPath, originalMsgs); werr != nil && logger != nil {
logger.Warn("eino summarization transcript preflight 写入失败",
zap.String("path", transcriptPath), zap.Error(werr))
}
}
input, dropped, berr := buildBudgetedSummarizationModelInput(
ctx, sysInstruction, userInstruction, originalMsgs, tokenCounter, summaryInputMax,
)
if logger != nil && (berr != nil || dropped > 0) {
fields := []zap.Field{
zap.Int("max_input_tokens", summaryInputMax),
zap.Int("dropped_rounds", dropped),
}
if berr != nil {
fields = append(fields, zap.Error(berr))
logger.Warn("eino summarization input budget failed", fields...)
} else {
logger.Info("eino summarization input bounded", fields...)
}
}
return input, berr
},
Trigger: &summarization.TriggerCondition{
ContextTokens: trigger,
},
@@ -195,7 +232,7 @@ func newEinoSummarizationMiddleware(
},
Finalize: func(ctx context.Context, originalMessages []adk.Message, summary adk.Message) ([]adk.Message, error) {
summary = stripAnalysisFromSummarizationMessage(summary)
userLedger := buildOriginalUserIntentLedgerMessageFromStore(db, conversationID, originalMessages, userLedgerMaxRunes, userLedgerEntryMaxRunes, logger)
userLedger := buildOriginalUserIntentLedgerMessage(originalMessages, userLedgerMaxRunes, userLedgerEntryMaxRunes)
compactionMessages := stripOriginalUserIntentLedgerFromMessages(originalMessages)
out, ferr := summarizeFinalizeWithRecentAssistantToolTrail(ctx, compactionMessages, summary, tokenCounter, recentTrailMax)
if ferr != nil {
@@ -238,59 +275,69 @@ func newEinoSummarizationMiddleware(
return mw, nil
}
func buildOriginalUserIntentLedgerMessageFromStore(
db *database.DB,
conversationID string,
fallbackMessages []adk.Message,
maxRunes int,
entryMaxRunes int,
logger *zap.Logger,
) adk.Message {
conversationID = strings.TrimSpace(conversationID)
if db == nil || conversationID == "" {
return buildOriginalUserIntentLedgerMessage(fallbackMessages, maxRunes, entryMaxRunes)
}
msgs, err := db.GetMessages(conversationID)
// buildBudgetedSummarizationModelInput builds the exact payload sent to the summary model.
// It retains the newest complete conversation rounds within budget and emits an explicit
// marker when older rounds are omitted. The full pre-compaction transcript is persisted
// separately; omitted raw messages never become model-facing history again.
func buildBudgetedSummarizationModelInput(
ctx context.Context,
sysInstruction adk.Message,
userInstruction adk.Message,
originalMsgs []adk.Message,
tokenCounter summarization.TokenCounterFunc,
maxTokens int,
) ([]adk.Message, int, error) {
base := []adk.Message{sysInstruction, userInstruction}
baseTokens, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: base})
if err != nil {
if logger != nil {
logger.Warn("summarization: 从数据库读取原始用户消息失败,回退到当前上下文",
zap.String("conversationId", conversationID),
zap.Error(err),
)
}
return buildOriginalUserIntentLedgerMessage(fallbackMessages, maxRunes, entryMaxRunes)
return nil, 0, err
}
userMsgs := make([]adk.Message, 0, len(msgs))
for _, msg := range msgs {
if strings.ToLower(strings.TrimSpace(msg.Role)) != "user" {
continue
}
content := strings.TrimSpace(msg.Content)
if content == "" {
continue
}
userMsgs = append(userMsgs, schema.UserMessage(formatStoredUserLedgerText(msg)))
remaining := maxTokens - baseTokens
markerTemplate := schema.UserMessage("[Context budget guard omitted older conversation rounds; summarize the retained recent rounds and preserve the omission marker.]")
markerTokens, err := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: []adk.Message{markerTemplate}})
if err != nil {
return nil, 0, err
}
if len(userMsgs) == 0 {
return buildOriginalUserIntentLedgerMessage(fallbackMessages, maxRunes, entryMaxRunes)
remaining -= markerTokens
if remaining < 0 {
remaining = 0
}
return buildOriginalUserIntentLedgerMessage(userMsgs, maxRunes, entryMaxRunes)
}
func formatStoredUserLedgerText(msg database.Message) string {
var parts []string
if id := strings.TrimSpace(msg.ID); id != "" {
parts = append(parts, "message_id="+id)
contextMsgs := make([]adk.Message, 0, len(originalMsgs))
for _, msg := range originalMsgs {
if msg != nil && msg.Role != schema.System {
contextMsgs = append(contextMsgs, msg)
}
}
if !msg.CreatedAt.IsZero() {
parts = append(parts, "created_at="+msg.CreatedAt.Format(time.RFC3339Nano))
rounds := splitMessagesIntoRounds(contextMsgs)
selectedReverse := make([]messageRound, 0, len(rounds))
used := 0
for i := len(rounds) - 1; i >= 0; i-- {
n, countErr := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: rounds[i].messages})
if countErr != nil {
return nil, 0, countErr
}
if used+n > remaining {
break
}
used += n
selectedReverse = append(selectedReverse, rounds[i])
}
meta := strings.Join(parts, " ")
content := strings.TrimSpace(msg.Content)
if meta == "" {
return content
dropped := len(rounds) - len(selectedReverse)
input := make([]adk.Message, 0, 3+len(contextMsgs))
input = append(input, sysInstruction)
if dropped > 0 {
input = append(input, schema.UserMessage(fmt.Sprintf(
"[Context budget guard omitted %d older conversation round(s); summarize the retained recent rounds and preserve the omission marker.]",
dropped,
)))
}
return meta + "\n" + content
for i := len(selectedReverse) - 1; i >= 0; i-- {
input = append(input, selectedReverse[i].messages...)
}
input = append(input, userInstruction)
return input, dropped, nil
}
// refreshFactIndexInMessages 在 summarization 压缩后,用 DB 最新索引替换 system 中已有的项目黑板索引段。
+32 -34
View File
@@ -57,6 +57,33 @@ func variableTokenCounter() summarization.TokenCounterFunc {
}
}
func TestBuildBudgetedSummarizationModelInputKeepsRecentCompleteRounds(t *testing.T) {
msgs := []adk.Message{
schema.UserMessage("old-user"),
schema.AssistantMessage("old-answer", nil),
schema.UserMessage("latest-user"),
assistantToolCallsMsg("", "call-latest"),
schema.ToolMessage("latest-tool-result", "call-latest"),
}
input, dropped, err := buildBudgetedSummarizationModelInput(
context.Background(), schema.SystemMessage("summary-system"), schema.UserMessage("summary-instruction"),
msgs, fixedTokenCounter(2), 7,
)
if err != nil {
t.Fatal(err)
}
if dropped == 0 {
t.Fatal("expected older rounds to be omitted")
}
joined := formatSummarizationTranscript(input)
if strings.Contains(joined, "old-user") || strings.Contains(joined, "old-answer") {
t.Fatalf("old rounds leaked into bounded input: %s", joined)
}
if !strings.Contains(joined, "latest-user") || !strings.Contains(joined, "latest-tool-result") {
t.Fatalf("latest complete rounds missing: %s", joined)
}
}
func TestSplitMessagesIntoRounds_Complex(t *testing.T) {
msgs := []adk.Message{
schema.UserMessage("q1"),
@@ -499,46 +526,17 @@ func TestRefreshFactIndexInMessages(t *testing.T) {
}
}
func TestBuildOriginalUserIntentLedgerMessageFromStore_UsesDatabaseAsCanonicalSource(t *testing.T) {
t.Parallel()
dbPath := filepath.Join(t.TempDir(), "summarize-user-ledger.db")
db, err := database.NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
defer db.Close()
conv, err := db.CreateConversation("ledger", database.ConversationCreateMeta{})
if err != nil {
t.Fatal(err)
}
stored, err := db.AddMessage(conv.ID, "user", "原始用户输入:只测 staging,不要碰 prod。", nil)
if err != nil {
t.Fatal(err)
}
if _, err := db.AddMessage(conv.ID, "assistant", "处理中...", nil); err != nil {
t.Fatal(err)
}
ledger := buildOriginalUserIntentLedgerMessageFromStore(
db,
conv.ID,
[]adk.Message{schema.UserMessage("内存里的派生/续跑消息,不应作为权威源。")},
func TestBuildOriginalUserIntentLedgerUsesOnlyModelFacingMessages(t *testing.T) {
ledger := buildOriginalUserIntentLedgerMessage(
[]adk.Message{schema.UserMessage("模型实际看到的裁剪预览")},
config.DefaultSummarizationUserIntentLedgerMaxRunes,
config.DefaultSummarizationUserIntentLedgerEntryMaxRunes,
zap.NewNop(),
)
if ledger == nil {
t.Fatal("expected ledger message")
}
body := ledger.Content
if !strings.Contains(body, stored.ID) {
t.Fatalf("ledger should include canonical DB message id %q: %q", stored.ID, body)
}
if !strings.Contains(body, "只测 staging,不要碰 prod") {
t.Fatalf("ledger should include stored user content: %q", body)
}
if strings.Contains(body, "内存里的派生") {
t.Fatalf("fallback in-memory user message should not be used when DB has user messages: %q", body)
if !strings.Contains(body, "模型实际看到的裁剪预览") {
t.Fatalf("ledger should preserve the model-facing user message: %q", body)
}
}
+4 -1
View File
@@ -35,6 +35,9 @@ func isEinoTransientRunError(err error) bool {
if msg == "" {
return false
}
if strings.Contains(msg, "model input exceeds configured hard budget") {
return false
}
transientMarkers := []string{
"406",
"429",
@@ -180,7 +183,7 @@ const (
// einoMessagesForRunRestart 在退避后重新 Run 时选用最完整的上下文:
// 1) ModelFacingTrace(与模型实际入参一致) 2) 事件流累积的 runAccumulatedMsgs 3) 初始 msgs。
func einoMessagesForRunRestart(args *einoADKRunLoopArgs, baseMsgs, accumulated []adk.Message, baseCount int) ([]adk.Message, einoRunRestartContextSource) {
if trace := persistTraceSource(args, nil); len(trace) > 0 {
if trace := modelFacingTraceSnapshot(args); len(trace) > 0 {
// modelFacingTrace includes prior Instruction system message(s); genModelInput will prepend again.
return stripADKSystemMessages(trace), einoRestartContextModelTrace
}
@@ -33,6 +33,7 @@ func TestIsEinoTransientRunError(t *testing.T) {
{"iteration limit", errors.New("max iteration reached"), false},
{"canceled", context.Canceled, false},
{"deadline", context.DeadlineExceeded, false},
{"model hard budget containing 500", errors.New("model input exceeds configured hard budget after preserving the latest round: tokens=20500 max=19500 phase=test"), false},
{"auth", errors.New("invalid api key"), false},
}
for _, tc := range cases {
@@ -21,6 +21,9 @@ func prepareLatestUserMessageForModel(userMessage string, appCfg *config.Config,
mwCfg = &zero
}
maxRunes := mwCfg.LatestUserMessageMaxRunesEffective()
if appCfg != nil {
maxRunes = minPositiveInt(maxRunes, modelFacingRuneBudget(appCfg.OpenAI.MaxTotalTokens, 0.20))
}
if maxRunes <= 0 || utf8RuneLen(userMessage) <= maxRunes {
return userMessage
}
@@ -65,6 +68,30 @@ func prepareLatestUserMessageForModel(userMessage string, appCfg *config.Config,
return strings.TrimSpace(sb.String())
}
func modelFacingRuneBudget(maxTotalTokens int, ratio float64) int {
if maxTotalTokens <= 0 {
maxTotalTokens = 120000
}
if ratio <= 0 || ratio >= 1 {
ratio = 0.20
}
budget := int(float64(maxTotalTokens) * ratio)
if budget < 1024 {
budget = 1024
}
return budget
}
func minPositiveInt(a, b int) int {
if a <= 0 {
return b
}
if b <= 0 || a < b {
return a
}
return b
}
func normalizeLatestUserPreviewBudget(maxRunes, headRunes, tailRunes int) (int, int) {
if maxRunes <= 0 {
return headRunes, tailRunes
@@ -0,0 +1,106 @@
package multiagent
import (
"context"
"fmt"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/adk/middlewares/summarization"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
// modelInputBudgetMiddleware is the final deterministic guard before a normal model call.
// Summarization remains the primary compaction strategy; this middleware only removes the
// oldest complete rounds if the finalized state still exceeds 65% of the configured
// window, leaving serialization/tokenizer headroom for the outbound HTTP guard.
type modelInputBudgetMiddleware struct {
adk.BaseChatModelAgentMiddleware
maxTokens int
counter summarization.TokenCounterFunc
logger *zap.Logger
phase string
}
func newModelInputBudgetMiddleware(maxTotalTokens int, modelName string, logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware {
if maxTotalTokens <= 0 {
maxTotalTokens = 120000
}
limit := int(float64(maxTotalTokens) * 0.65)
if limit < 4096 {
limit = 4096
}
return &modelInputBudgetMiddleware{
maxTokens: limit,
counter: einoSummarizationTokenCounter(modelName),
logger: logger,
phase: phase,
}
}
func (m *modelInputBudgetMiddleware) BeforeModelRewriteState(
ctx context.Context,
state *adk.ChatModelAgentState,
mc *adk.ModelContext,
) (context.Context, *adk.ChatModelAgentState, error) {
if m == nil || state == nil || len(state.Messages) == 0 {
return ctx, state, nil
}
count := func(msgs []adk.Message) (int, error) {
input := &summarization.TokenCounterInput{Messages: msgs}
if mc != nil {
input.Tools = mc.Tools
}
return m.counter(ctx, input)
}
before, err := count(state.Messages)
if err != nil {
return ctx, state, err
}
if before <= m.maxTokens {
return ctx, state, nil
}
systems := make([]adk.Message, 0, 1)
contextMsgs := make([]adk.Message, 0, len(state.Messages))
for _, msg := range state.Messages {
if msg != nil && msg.Role == schema.System && len(contextMsgs) == 0 {
systems = append(systems, msg)
continue
}
if msg != nil {
contextMsgs = append(contextMsgs, msg)
}
}
rounds := splitMessagesIntoRounds(contextMsgs)
dropped := 0
var candidate []adk.Message
for len(rounds) > 1 {
rounds = rounds[1:]
dropped++
candidate = append(candidate[:0], systems...)
for _, round := range rounds {
candidate = append(candidate, round.messages...)
}
after, countErr := count(candidate)
if countErr != nil {
return ctx, state, countErr
}
if after <= m.maxTokens {
out := *state
out.Messages = append([]adk.Message(nil), candidate...)
if m.logger != nil {
m.logger.Warn("eino model input hard budget applied",
zap.String("phase", m.phase), zap.Int("tokens_before", before),
zap.Int("tokens_after", after), zap.Int("max_tokens", m.maxTokens),
zap.Int("dropped_rounds", dropped))
}
return ctx, &out, nil
}
}
return ctx, state, fmt.Errorf(
"model input exceeds configured hard budget after preserving the latest round: tokens=%d max=%d phase=%s",
before, m.maxTokens, m.phase,
)
}
@@ -0,0 +1,50 @@
package multiagent
import (
"context"
"strings"
"testing"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func TestModelInputBudgetDropsOldestCompleteRoundsAndPersistsLatest(t *testing.T) {
mw := &modelInputBudgetMiddleware{
maxTokens: 7,
counter: fixedTokenCounter(4),
phase: "test",
}
state := &adk.ChatModelAgentState{Messages: []adk.Message{
schema.SystemMessage("system"),
schema.UserMessage("old-user"),
schema.AssistantMessage("old-answer", nil),
schema.UserMessage("latest-user"),
assistantToolCallsMsg("", "latest-call"),
schema.ToolMessage("latest-result", "latest-call"),
}}
_, out, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
if err != nil {
t.Fatal(err)
}
joined := formatSummarizationTranscript(out.Messages)
if strings.Contains(joined, "old-user") || strings.Contains(joined, "old-answer") {
t.Fatalf("old rounds retained after hard budget: %s", joined)
}
if !strings.Contains(joined, "latest-user") || !strings.Contains(joined, "latest-result") {
t.Fatalf("latest rounds lost after hard budget: %s", joined)
}
}
func TestModelInputBudgetFailsLocallyWhenLatestRoundAloneCannotFit(t *testing.T) {
mw := &modelInputBudgetMiddleware{maxTokens: 2, counter: fixedTokenCounter(4), phase: "test"}
state := &adk.ChatModelAgentState{Messages: []adk.Message{
schema.SystemMessage("system"),
assistantToolCallsMsg("", "latest-call"),
schema.ToolMessage("latest-result", "latest-call"),
}}
_, _, err := mw.BeforeModelRewriteState(context.Background(), state, nil)
if err == nil || !strings.Contains(err.Error(), "hard budget") {
t.Fatalf("expected local hard-budget error, got %v", err)
}
}
+75 -7
View File
@@ -10,7 +10,9 @@ import (
"sort"
"strings"
"sync"
"sync/atomic"
"time"
"unicode/utf8"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/agents"
@@ -50,6 +52,8 @@ type toolCallPendingInfo struct {
EinoRole string
}
var fallbackToolCallSequence atomic.Uint64
// RunDeepAgent 使用 Eino 多代理预置编排执行一轮对话(deep / plan_execute / supervisor;流式事件通过 progress 回调输出)。
// orchestrationOverride 非空时优先(如聊天/WebShell 请求体);否则用 multi_agent.orchestration(遗留 yaml);皆空则按 deep。
// reasoningClient 来自 ChatRequest.reasoning;可为 nil(机器人/批量等走全局 openai.reasoning)。
@@ -247,6 +251,7 @@ func RunDeepAgent(
phase: "sub_agent:" + id,
summarization: subSumMw,
modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
conversationID: conversationID,
})
@@ -406,6 +411,7 @@ func RunDeepAgent(
phase: "deep_orchestrator",
summarization: mainSumMw,
modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
conversationID: conversationID,
trace: modelFacingTrace,
})
@@ -422,6 +428,7 @@ func RunDeepAgent(
phase: "supervisor_orchestrator",
summarization: mainSumMw,
modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
conversationID: conversationID,
trace: modelFacingTrace,
})
@@ -495,6 +502,7 @@ func RunDeepAgent(
phase: "plan_execute_planner_replanner",
summarization: mainSumMw,
modelName: appCfg.OpenAI.Model,
maxTotalTokens: appCfg.OpenAI.MaxTotalTokens,
conversationID: conversationID,
skipTrace: true,
}),
@@ -642,11 +650,19 @@ func chatToolCallsToSchema(tcs []agent.ToolCall) []schema.ToolCall {
return out
}
// historyToMessages 将轨迹恢复的 ChatMessage 转为 Eino ADK 消息:**不裁剪条数、不按 token 预算截断**,
// 并保留 user / assistant(含仅 tool_calls/ tool,与库中 last_react 轨迹一致。
// historyToMessages 将已保存的 model-facing 轨迹转为 Eino ADK 消息。
// 新轨迹应已是模型实际看到的内容;对旧版本遗留的超大 tool 正文再做一次上限规范化,
// 防止原始工具输出通过 last_react/checkpoint 绕过 reduction。
func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) []adk.Message {
_ = appCfg
_ = mwCfg
toolContentMax := config.MultiAgentEinoMiddlewareConfig{}.ReductionMaxLengthForTruncEffective()
userContentMaxRunes := config.MultiAgentEinoMiddlewareConfig{}.LatestUserMessageMaxRunesEffective()
if mwCfg != nil {
toolContentMax = mwCfg.ReductionMaxLengthForTruncEffective()
userContentMaxRunes = mwCfg.LatestUserMessageMaxRunesEffective()
}
if appCfg != nil {
userContentMaxRunes = minPositiveInt(userContentMaxRunes, modelFacingRuneBudget(appCfg.OpenAI.MaxTotalTokens, 0.20))
}
if len(history) == 0 {
return nil
}
@@ -656,7 +672,11 @@ func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg
switch role {
case "user":
if strings.TrimSpace(h.Content) != "" {
raw = append(raw, schema.UserMessage(h.Content))
content := h.Content
if !h.ModelFacingTrace {
content = normalizeRestoredUserContent(content, userContentMaxRunes)
}
raw = append(raw, schema.UserMessage(content))
}
case "assistant":
toolSchema := chatToolCallsToSchema(h.ToolCalls)
@@ -676,7 +696,11 @@ func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg
if tn := strings.TrimSpace(h.ToolName); tn != "" {
opts = append(opts, schema.WithToolName(tn))
}
raw = append(raw, schema.ToolMessage(h.Content, h.ToolCallID, opts...))
content := h.Content
if !h.ModelFacingTrace {
content = normalizeRestoredToolContent(content, toolContentMax)
}
raw = append(raw, schema.ToolMessage(content, h.ToolCallID, opts...))
default:
continue
}
@@ -684,6 +708,47 @@ func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg
return raw
}
func normalizeRestoredUserContent(content string, maxRunes int) string {
if maxRunes <= 0 || utf8.RuneCountInString(content) <= maxRunes {
return content
}
runes := []rune(content)
const marker = "\n\n...[historical user input normalized to the model-facing budget]...\n\n"
markerRunes := []rune(marker)
budget := maxRunes - len(markerRunes)
if budget <= 0 {
end := maxRunes
if end > len(markerRunes) {
end = len(markerRunes)
}
return string(markerRunes[:end])
}
head := budget / 2
tail := budget - head
return string(runes[:head]) + marker + string(runes[len(runes)-tail:])
}
func normalizeRestoredToolContent(content string, maxBytes int) string {
if maxBytes <= 0 || len(content) <= maxBytes {
return content
}
const marker = "\n\n...[legacy tool output discarded during model-facing history migration]...\n\n"
budget := maxBytes - len(marker)
if budget <= 0 {
return marker
}
head := budget / 2
tail := budget - head
for head > 0 && !utf8.RuneStart(content[head]) {
head--
}
tailStart := len(content) - tail
for tailStart < len(content) && !utf8.RuneStart(content[tailStart]) {
tailStart++
}
return content[:head] + marker + content[tailStart:]
}
// mergeStreamingToolCallFragments 将流式多帧的 ToolCall 按 index 合并 arguments(与 schema.concatToolCalls 行为一致)。
func mergeStreamingToolCallFragments(fragments []schema.ToolCall) []schema.ToolCall {
if len(fragments) == 0 {
@@ -870,7 +935,10 @@ func emitToolCallsFromMessage(
display := toolCallDisplayName(tc)
toolCallID := tc.ID
if toolCallID == "" && tc.Index != nil {
toolCallID = fmt.Sprintf("eino-stream-%d", *tc.Index)
// Stream indexes restart from zero for every model turn. Include a
// process-wide sequence so pending/result de-duplication cannot collide
// with an earlier batch in the same agent run.
toolCallID = fmt.Sprintf("eino-stream-%d-%d", fallbackToolCallSequence.Add(1), *tc.Index)
}
// Record pending tool calls for later tool_result correlation / recovery flushing.
// We intentionally record even for unknown tools to avoid "running" badge getting stuck.
@@ -1,9 +1,12 @@
package multiagent
import (
"strings"
"testing"
"unicode/utf8"
"cyberstrike-ai/internal/agent"
"cyberstrike-ai/internal/config"
)
func TestHistoryToMessagesPreservesReasoningContent(t *testing.T) {
@@ -20,3 +23,61 @@ func TestHistoryToMessagesPreservesReasoningContent(t *testing.T) {
t.Fatalf("got reasoning=%q content=%q", am.ReasoningContent, am.Content)
}
}
func TestHistoryToMessagesNormalizesLegacyRawToolOutput(t *testing.T) {
mw := &config.MultiAgentEinoMiddlewareConfig{ReductionMaxLengthForTrunc: 128}
h := []agent.ChatMessage{
{Role: "assistant", ToolCalls: []agent.ToolCall{{ID: "t1", Type: "function", Function: agent.FunctionCall{Name: "http-framework-test"}}}},
{Role: "tool", ToolCallID: "t1", ToolName: "http-framework-test", Content: strings.Repeat("响应正文", 1000)},
}
msgs := historyToMessages(h, nil, mw)
if len(msgs) != 2 {
t.Fatalf("len=%d", len(msgs))
}
if len(msgs[1].Content) > 128 {
t.Fatalf("normalized tool bytes=%d, want <=128", len(msgs[1].Content))
}
if !strings.Contains(msgs[1].Content, "legacy tool output discarded") {
t.Fatalf("missing migration marker: %q", msgs[1].Content)
}
}
func TestHistoryToMessagesRestoresModelFacingTraceByteForByte(t *testing.T) {
mw := &config.MultiAgentEinoMiddlewareConfig{
ReductionMaxLengthForTrunc: 128,
LatestUserMessageMaxRunes: 64,
}
userContent := strings.Repeat("model-facing-user-", 100)
toolContent := strings.Repeat("model-facing-tool-", 100)
h := []agent.ChatMessage{
{Role: "user", Content: userContent, ModelFacingTrace: true},
{Role: "assistant", ModelFacingTrace: true, ToolCalls: []agent.ToolCall{{ID: "t1", Type: "function", Function: agent.FunctionCall{Name: "http-framework-test"}}}},
{Role: "tool", ToolCallID: "t1", ToolName: "http-framework-test", Content: toolContent, ModelFacingTrace: true},
}
msgs := historyToMessages(h, nil, mw)
if len(msgs) != 3 {
t.Fatalf("len=%d", len(msgs))
}
if msgs[0].Content != userContent {
t.Fatal("model-facing user content changed during restore")
}
if msgs[2].Content != toolContent {
t.Fatal("model-facing tool content changed during restore")
}
}
func TestHistoryToMessagesNeverReinjectsRawOversizedUserFallback(t *testing.T) {
appCfg := &config.Config{OpenAI: config.OpenAIConfig{MaxTotalTokens: 10000}}
mw := &config.MultiAgentEinoMiddlewareConfig{LatestUserMessageMaxRunes: 8000}
h := []agent.ChatMessage{{Role: "user", Content: strings.Repeat("原始用户输入", 2000)}}
msgs := historyToMessages(h, appCfg, mw)
if len(msgs) != 1 {
t.Fatalf("len=%d", len(msgs))
}
if got := utf8.RuneCountInString(msgs[0].Content); got > 2000 {
t.Fatalf("restored user runes=%d, want <=2000", got)
}
if !strings.Contains(msgs[0].Content, "historical user input normalized") {
t.Fatalf("missing normalization marker: %q", msgs[0].Content)
}
}
@@ -0,0 +1,37 @@
package multiagent
import (
"testing"
"github.com/cloudwego/eino/schema"
)
func TestEmitToolCallsUsesUniqueFallbackIDsAcrossBatches(t *testing.T) {
index := 0
msg := &schema.Message{ToolCalls: []schema.ToolCall{{
Index: &index,
Function: schema.FunctionCall{
Name: "http-framework-test",
Arguments: `{}`,
},
}}}
var ids []string
progress := func(eventType, _ string, raw interface{}) {
if eventType != "tool_call" {
return
}
data, _ := raw.(map[string]interface{})
if id, _ := data["toolCallId"].(string); id != "" {
ids = append(ids, id)
}
}
for i := 0; i < 2; i++ {
emitToolCallsFromMessage(msg, "agent", "agent", "conversation", "deep", progress, nil, make(map[string]int), nil)
}
if len(ids) != 2 {
t.Fatalf("fallback IDs = %v, want two IDs", ids)
}
if ids[0] == ids[1] {
t.Fatalf("fallback ID was reused across batches: %q", ids[0])
}
}
@@ -0,0 +1,156 @@
package multiagent
import (
"context"
"fmt"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
"go.uber.org/zap"
)
const patchedMissingToolResult = "[Tool execution result was lost or interrupted; continue without relying on this call.]"
// toolPairReconcilerMiddleware is the final structural guard before a model call.
// It makes every assistant tool-call batch immediately followed by exactly one tool
// result per call ID, and drops tool messages that cannot belong to that batch.
//
// This intentionally runs after summarization/reduction/budget middleware: those
// middlewares rewrite history and can otherwise re-introduce a partial tool round
// after the ordinary patchtoolcalls middleware has already run.
type toolPairReconcilerMiddleware struct {
adk.BaseChatModelAgentMiddleware
logger *zap.Logger
phase string
}
func newToolPairReconcilerMiddleware(logger *zap.Logger, phase string) adk.ChatModelAgentMiddleware {
return &toolPairReconcilerMiddleware{logger: logger, phase: phase}
}
func (m *toolPairReconcilerMiddleware) BeforeModelRewriteState(
ctx context.Context,
state *adk.ChatModelAgentState,
mc *adk.ModelContext,
) (context.Context, *adk.ChatModelAgentState, 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([]adk.Message, 0, len(state.Messages))
for i := 0; i < len(state.Messages); {
msg := state.Messages[i]
if msg == nil {
changed = true
i++
continue
}
if msg.Role == schema.Tool {
// Valid tool results are consumed with their immediately preceding assistant.
changed = true
dropped++
i++
continue
}
if msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 {
out = append(out, msg)
i++
continue
}
assistant := msg
calls := append([]schema.ToolCall(nil), msg.ToolCalls...)
expected := make(map[string]schema.ToolCall, len(calls))
idsChanged := false
for callIndex := range calls {
id := calls[callIndex].ID
_, duplicate := usedIDs[id]
if id == "" || duplicate {
base := fmt.Sprintf("patched_tool_call_%d_%d", i, callIndex)
id = base
for suffix := 1; ; suffix++ {
if _, exists := usedIDs[id]; !exists {
break
}
id = fmt.Sprintf("%s_%d", base, suffix)
}
calls[callIndex].ID = id
idsChanged = true
changed = true
}
usedIDs[id] = struct{}{}
expected[id] = calls[callIndex]
}
if idsChanged {
cloned := *msg
cloned.ToolCalls = calls
assistant = &cloned
}
out = append(out, assistant)
results := make(map[string]adk.Message, len(calls))
j := i + 1
for j < len(state.Messages) {
toolMsg := state.Messages[j]
if toolMsg == nil {
changed = true
j++
continue
}
if toolMsg.Role != schema.Tool {
break
}
id := toolMsg.ToolCallID
if _, wanted := expected[id]; !wanted {
changed = true
dropped++
j++
continue
}
if _, duplicate := results[id]; duplicate {
changed = true
dropped++
j++
continue
}
results[id] = toolMsg
j++
}
for _, tc := range calls {
if result, ok := results[tc.ID]; ok {
out = append(out, result)
continue
}
out = append(out, schema.ToolMessage(
patchedMissingToolResult,
tc.ID,
schema.WithToolName(tc.Function.Name),
))
changed = true
patched++
}
i = j
}
if !changed {
return ctx, state, nil
}
if m.logger != nil {
m.logger.Warn("eino 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
}
@@ -0,0 +1,141 @@
package multiagent
import (
"context"
"testing"
"github.com/cloudwego/eino/adk"
"github.com/cloudwego/eino/schema"
)
func runToolPairReconciler(t *testing.T, msgs []adk.Message) []adk.Message {
t.Helper()
mw := newToolPairReconcilerMiddleware(nil, "test").(*toolPairReconcilerMiddleware)
_, out, err := mw.BeforeModelRewriteState(
context.Background(),
&adk.ChatModelAgentState{Messages: msgs},
&adk.ModelContext{},
)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
return out.Messages
}
func assertCompleteImmediateToolPairs(t *testing.T, msgs []adk.Message) {
t.Helper()
for i, msg := range msgs {
if msg == nil || msg.Role != schema.Assistant || len(msg.ToolCalls) == 0 {
continue
}
want := make(map[string]struct{}, len(msg.ToolCalls))
for _, tc := range msg.ToolCalls {
if tc.ID == "" {
t.Fatal("empty tool call id remained")
}
if _, duplicate := want[tc.ID]; duplicate {
t.Fatalf("duplicate tool call id remained: %s", tc.ID)
}
want[tc.ID] = struct{}{}
}
seen := make(map[string]struct{}, len(want))
for j := i + 1; j < len(msgs) && msgs[j] != nil && msgs[j].Role == schema.Tool; j++ {
id := msgs[j].ToolCallID
if _, ok := want[id]; !ok {
t.Fatalf("unexpected tool result %q after assistant %d", id, i)
}
if _, duplicate := seen[id]; duplicate {
t.Fatalf("duplicate tool result %q", id)
}
seen[id] = struct{}{}
}
if len(seen) != len(want) {
t.Fatalf("assistant %d: want %d results, got %d", i, len(want), len(seen))
}
}
}
func TestToolPairReconcilerPatchesPartialMultiToolBatch(t *testing.T) {
msgs := []adk.Message{
schema.UserMessage("start"),
assistantToolCallsMsg("", "c1", "c2"),
schema.ToolMessage("r1", "c1"),
schema.UserMessage("continue"),
}
out := runToolPairReconciler(t, msgs)
assertCompleteImmediateToolPairs(t, out)
if len(out) != 5 || out[3].Role != schema.Tool || out[3].ToolCallID != "c2" {
t.Fatalf("missing c2 result was not inserted in place: %+v", out)
}
if out[3].Content != patchedMissingToolResult {
t.Fatalf("unexpected patched content: %q", out[3].Content)
}
}
func TestToolPairReconcilerDropsMisplacedDuplicateAndOrphanResults(t *testing.T) {
msgs := []adk.Message{
schema.ToolMessage("old", "orphan"),
assistantToolCallsMsg("", "c1"),
schema.ToolMessage("first", "c1"),
schema.ToolMessage("duplicate", "c1"),
schema.ToolMessage("wrong", "other"),
schema.UserMessage("next"),
schema.ToolMessage("late", "c1"),
}
out := runToolPairReconciler(t, msgs)
assertCompleteImmediateToolPairs(t, out)
toolCount := 0
for _, msg := range out {
if msg.Role == schema.Tool {
toolCount++
if msg.Content != "first" || msg.ToolCallID != "c1" {
t.Fatalf("unexpected retained tool result: %+v", msg)
}
}
}
if toolCount != 1 {
t.Fatalf("want one retained tool result, got %d", toolCount)
}
}
func TestToolPairReconcilerRepairsEmptyAndRepeatedCallIDs(t *testing.T) {
msgs := []adk.Message{
assistantToolCallsMsg("", "", "same"),
schema.ToolMessage("same-1", "same"),
assistantToolCallsMsg("", "same"),
schema.ToolMessage("same-2", "same"),
}
out := runToolPairReconciler(t, msgs)
assertCompleteImmediateToolPairs(t, out)
all := make(map[string]struct{})
for _, msg := range out {
if msg.Role != schema.Assistant {
continue
}
for _, tc := range msg.ToolCalls {
if _, duplicate := all[tc.ID]; duplicate {
t.Fatalf("global duplicate tool call id remained: %s", tc.ID)
}
all[tc.ID] = struct{}{}
}
}
}
func TestToolPairReconcilerNoOpForValidHistory(t *testing.T) {
msgs := []adk.Message{
schema.UserMessage("start"),
assistantToolCallsMsg("", "c1", "c2"),
schema.ToolMessage("r1", "c1"),
schema.ToolMessage("r2", "c2"),
schema.AssistantMessage("done", nil),
}
mw := newToolPairReconcilerMiddleware(nil, "test").(*toolPairReconcilerMiddleware)
in := &adk.ChatModelAgentState{Messages: msgs}
_, out, err := mw.BeforeModelRewriteState(context.Background(), in, &adk.ModelContext{})
if err != nil {
t.Fatal(err)
}
if out != in {
t.Fatal("valid history should use the no-op fast path")
}
}