From 542d1d24119fb20d41a5ced703d3fd1b7e240ce9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Fri, 10 Jul 2026 16:48:49 +0800 Subject: [PATCH] Add files via upload --- internal/multiagent/eino_adk_run_loop.go | 45 +++-- .../eino_chat_model_tail_middleware.go | 17 +- .../eino_model_facing_persistence_test.go | 34 ++++ internal/multiagent/eino_orchestration.go | 1 + internal/multiagent/eino_single_runner.go | 1 + internal/multiagent/eino_skills.go | 38 ++++- .../multiagent/eino_skills_reduction_test.go | 27 +++ internal/multiagent/eino_summarize.go | 141 ++++++++++------ internal/multiagent/eino_summarize_test.go | 66 ++++---- internal/multiagent/eino_transient_retry.go | 5 +- .../multiagent/eino_transient_retry_test.go | 1 + internal/multiagent/latest_user_message.go | 27 +++ .../model_input_budget_middleware.go | 106 ++++++++++++ .../model_input_budget_middleware_test.go | 50 ++++++ internal/multiagent/runner.go | 82 ++++++++- .../runner_reasoning_history_test.go | 61 +++++++ .../multiagent/runner_tool_call_id_test.go | 37 +++++ .../tool_pair_reconciler_middleware.go | 156 ++++++++++++++++++ .../tool_pair_reconciler_middleware_test.go | 141 ++++++++++++++++ 19 files changed, 922 insertions(+), 114 deletions(-) create mode 100644 internal/multiagent/eino_model_facing_persistence_test.go create mode 100644 internal/multiagent/eino_skills_reduction_test.go create mode 100644 internal/multiagent/model_input_budget_middleware.go create mode 100644 internal/multiagent/model_input_budget_middleware_test.go create mode 100644 internal/multiagent/runner_tool_call_id_test.go create mode 100644 internal/multiagent/tool_pair_reconciler_middleware.go create mode 100644 internal/multiagent/tool_pair_reconciler_middleware_test.go diff --git a/internal/multiagent/eino_adk_run_loop.go b/internal/multiagent/eino_adk_run_loop.go index d55e7cf0..9848bc05 100644 --- a/internal/multiagent/eino_adk_run_loop.go +++ b/internal/multiagent/eino_adk_run_loop.go @@ -79,7 +79,7 @@ type einoADKRunLoopArgs struct { EinoRoleTag func(agent string) string CheckpointDir string // RunRetryMaxAttempts / RunRetryMaxBackoffSec：429、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 轨迹中回填用户可见回复。 // 典型场景：监督者仅调用 exit（final_result 落在 Tool 消息中），或工具结果已写入历史但 lastAssistant 未更新。 // diff --git a/internal/multiagent/eino_chat_model_tail_middleware.go b/internal/multiagent/eino_chat_model_tail_middleware.go index c1d7d4f4..5d7e4689 100644 --- a/internal/multiagent/eino_chat_model_tail_middleware.go +++ b/internal/multiagent/eino_chat_model_tail_middleware.go @@ -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)) } diff --git a/internal/multiagent/eino_model_facing_persistence_test.go b/internal/multiagent/eino_model_facing_persistence_test.go new file mode 100644 index 00000000..eddf1bed --- /dev/null +++ b/internal/multiagent/eino_model_facing_persistence_test.go @@ -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") + } +} diff --git a/internal/multiagent/eino_orchestration.go b/internal/multiagent/eino_orchestration.go index 434429dd..be8018ff 100644 --- a/internal/multiagent/eino_orchestration.go +++ b/internal/multiagent/eino_orchestration.go @@ -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, }) diff --git a/internal/multiagent/eino_single_runner.go b/internal/multiagent/eino_single_runner.go index 1eaff65c..6489dd86 100644 --- a/internal/multiagent/eino_single_runner.go +++ b/internal/multiagent/eino_single_runner.go @@ -150,6 +150,7 @@ func RunEinoSingleChatModelAgent( phase: "eino_single", summarization: mainSumMw, modelName: appCfg.OpenAI.Model, + maxTotalTokens: appCfg.OpenAI.MaxTotalTokens, conversationID: conversationID, trace: modelFacingTrace, }) diff --git a/internal/multiagent/eino_skills.go b/internal/multiagent/eino_skills.go index 1af6dbda..e5fc5c75 100644 --- a/internal/multiagent/eino_skills.go +++ b/internal/multiagent/eino_skills.go @@ -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{ diff --git a/internal/multiagent/eino_skills_reduction_test.go b/internal/multiagent/eino_skills_reduction_test.go new file mode 100644 index 00000000..7f9fe8b5 --- /dev/null +++ b/internal/multiagent/eino_skills_reduction_test.go @@ -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) + } +} diff --git a/internal/multiagent/eino_summarize.go b/internal/multiagent/eino_summarize.go index eee3589c..331f81f6 100644 --- a/internal/multiagent/eino_summarize.go +++ b/internal/multiagent/eino_summarize.go @@ -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 中已有的项目黑板索引段。 diff --git a/internal/multiagent/eino_summarize_test.go b/internal/multiagent/eino_summarize_test.go index 13479d0d..f9c83539 100644 --- a/internal/multiagent/eino_summarize_test.go +++ b/internal/multiagent/eino_summarize_test.go @@ -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) } } diff --git a/internal/multiagent/eino_transient_retry.go b/internal/multiagent/eino_transient_retry.go index bb1d4fae..2736f6cd 100644 --- a/internal/multiagent/eino_transient_retry.go +++ b/internal/multiagent/eino_transient_retry.go @@ -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 } diff --git a/internal/multiagent/eino_transient_retry_test.go b/internal/multiagent/eino_transient_retry_test.go index ef8f0608..d52aa961 100644 --- a/internal/multiagent/eino_transient_retry_test.go +++ b/internal/multiagent/eino_transient_retry_test.go @@ -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 { diff --git a/internal/multiagent/latest_user_message.go b/internal/multiagent/latest_user_message.go index de7a33c9..e126dbf1 100644 --- a/internal/multiagent/latest_user_message.go +++ b/internal/multiagent/latest_user_message.go @@ -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 diff --git a/internal/multiagent/model_input_budget_middleware.go b/internal/multiagent/model_input_budget_middleware.go new file mode 100644 index 00000000..85621fcf --- /dev/null +++ b/internal/multiagent/model_input_budget_middleware.go @@ -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, + ) +} diff --git a/internal/multiagent/model_input_budget_middleware_test.go b/internal/multiagent/model_input_budget_middleware_test.go new file mode 100644 index 00000000..4a20373b --- /dev/null +++ b/internal/multiagent/model_input_budget_middleware_test.go @@ -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) + } +} diff --git a/internal/multiagent/runner.go b/internal/multiagent/runner.go index dea22bd1..f29c7f80 100644 --- a/internal/multiagent/runner.go +++ b/internal/multiagent/runner.go @@ -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. diff --git a/internal/multiagent/runner_reasoning_history_test.go b/internal/multiagent/runner_reasoning_history_test.go index 8027c486..792582ef 100644 --- a/internal/multiagent/runner_reasoning_history_test.go +++ b/internal/multiagent/runner_reasoning_history_test.go @@ -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) + } +} diff --git a/internal/multiagent/runner_tool_call_id_test.go b/internal/multiagent/runner_tool_call_id_test.go new file mode 100644 index 00000000..e3e89356 --- /dev/null +++ b/internal/multiagent/runner_tool_call_id_test.go @@ -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]) + } +} diff --git a/internal/multiagent/tool_pair_reconciler_middleware.go b/internal/multiagent/tool_pair_reconciler_middleware.go new file mode 100644 index 00000000..fcd157b8 --- /dev/null +++ b/internal/multiagent/tool_pair_reconciler_middleware.go @@ -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 +} diff --git a/internal/multiagent/tool_pair_reconciler_middleware_test.go b/internal/multiagent/tool_pair_reconciler_middleware_test.go new file mode 100644 index 00000000..fac59ab0 --- /dev/null +++ b/internal/multiagent/tool_pair_reconciler_middleware_test.go @@ -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") + } +}