mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-07-13 23:47:24 +02:00
Add files via upload
This commit is contained in:
@@ -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 未更新。
|
||||
//
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 中已有的项目黑板索引段。
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user