mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-05-18 14:04:52 +02:00
Compare commits
12 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bcfb7b8da1 | |||
| f420ae0265 | |||
| e3f59b29ab | |||
| 87cba37203 | |||
| 4773b9e963 | |||
| eda5f9bba1 | |||
| 1318607813 | |||
| 5100924abe | |||
| 44079674dd | |||
| d959390e27 | |||
| 62a0d8cb71 | |||
| b53cae3a02 |
+15
-5
@@ -10,7 +10,7 @@
|
||||
# ============================================
|
||||
|
||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||
version: "v1.5.10"
|
||||
version: "v1.5.13"
|
||||
# 服务器配置
|
||||
server:
|
||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||
@@ -70,7 +70,7 @@ multi_agent:
|
||||
robot_use_multi_agent: true # true 时企业微信/钉钉/飞书机器人也走 Eino 多代理(成本更高)
|
||||
batch_use_multi_agent: false # true 时「批量任务」队列中每个子任务也走 Eino 多代理(成本更高)
|
||||
max_iteration: 0 # 主代理 / plan_execute 执行器最大轮次,0 表示沿用 agent.max_iterations
|
||||
# plan_execute 专用:execute↔replan 外层循环上限,0 表示 Eino 默认 10。Executor 未暴露 Handlers:patch/reduction/plantask 不作用于 PE,但 tool_search 工具列表拆分仍通过共享 ToolsConfig 作用于执行器。
|
||||
# plan_execute 专用:execute↔replan 外层循环上限,0 表示 Eino 默认 10。当前实现下 Executor 会挂载 patch/reduction/tool_search 等前置中间件。
|
||||
plan_execute_loop_max_iterations: 0
|
||||
sub_agent_max_iterations: 120
|
||||
sub_agent_user_context_max_runes: 0 # 子代理 task 描述中自动注入用户原始请求的字符上限;0=默认2000,负数=禁用
|
||||
@@ -87,15 +87,25 @@ multi_agent:
|
||||
# Eino ADK 中间件与 Deep/Supervisor 调参(结构体见 internal/config/config.go → MultiAgentEinoMiddlewareConfig)
|
||||
eino_middleware:
|
||||
patch_tool_calls: true # true:修补历史中无 tool_result 的悬空 tool_call(流式中断/重试后更稳);false:关闭;字段省略时默认等同 true
|
||||
tool_search_enable: false # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
|
||||
tool_search_enable: true # true:工具数 ≥ min 时启用 tool_search,仅前 N 个工具常驻,其余按正则按需解锁,省 token、减误选;false:全量工具进上下文
|
||||
tool_search_min_tools: 20 # 达到该数量才启用 tool_search(避免工具很少时多此一举);与 always_visible 配合使用
|
||||
tool_search_always_visible: 12 # 始终直接暴露给模型的工具个数(顺序与角色工具列表一致);其余工具进入动态池,需 tool_search 解锁
|
||||
tool_search_always_visible_tools: [read_file, glob, grep, write_file, edit_file, execute, task, transfer_to_agent, exit, write_todos, skill, tool_search, TaskCreate, TaskGet, TaskUpdate, TaskList, record_vulnerability, list_knowledge_risk_types, search_knowledge_base, webshell_exec, webshell_file_list, webshell_file_read, webshell_file_write, manage_webshell_list, manage_webshell_add, manage_webshell_update, manage_webshell_delete, manage_webshell_test, batch_task_list, batch_task_get, batch_task_start, batch_task_rerun, batch_task_pause, batch_task_update_metadata, batch_task_update_schedule, batch_task_schedule_enabled, batch_task_update_task, batch_task_remove_task, batch_task_delete, batch_task_create, batch_task_add_task, http-framework-test] # 后端内置常驻工具白名单(优先于 always_visible 数量策略)
|
||||
plantask_enable: false # true:主代理(Deep / Supervisor 主)挂载 TaskCreate/Get/Update/List;需 eino_skills 可用且 skills_dir 存在,否则仅打日志并跳过
|
||||
plantask_rel_dir: .eino/plantask # 结构化任务文件相对 skills_dir 的子目录,其下再按会话 ID 分子目录存放
|
||||
reduction_enable: false # true:大工具输出截断/落盘以控上下文;依赖与 plantask 相同的 eino local 写盘后端,无后端时不挂载
|
||||
reduction_enable: true # true:大工具输出截断/落盘以控上下文;依赖与 plantask 相同的 eino local 写盘后端,无后端时不挂载
|
||||
reduction_max_length_for_trunc: 50000 # 单条工具结果超过该字符数(bytes)时截断并落盘(由 reduction 中间件处理)
|
||||
reduction_max_tokens_for_clear: 160000 # 历史工具结果清理阈值(tokens),超阈值时在模型调用前清理旧结果
|
||||
reduction_root_dir: "" # 非空:截断/清理内容落盘根路径;空:使用系统临时目录下按会话隔离的默认路径
|
||||
reduction_clear_exclude: [] # 不参与「清理阶段」的工具名额外列表(会与 task/transfer/exit 等内置排除项合并);需要时用 YAML 列表填写
|
||||
reduction_sub_agents: false # true:子代理也挂 reduction;false:仅编排主代理使用 reduction
|
||||
reduction_sub_agents: true # true:子代理也挂 reduction;false:仅编排主代理使用 reduction
|
||||
summarization_trigger_ratio: 0.8 # summarization 触发比例(max_total_tokens * ratio),建议 0.75~0.85
|
||||
summarization_emit_internal_events: true # true:发出 summarization 内部事件(便于诊断)
|
||||
history_input_budget_ratio: 0.35 # 历史入队预算比例(max_total_tokens * ratio)
|
||||
plan_execute_user_input_budget_ratio: 0.35 # plan_execute 中 userInput 预算比例(planner/replanner/executor 共用)
|
||||
plan_execute_executed_steps_budget_ratio: 0.2 # plan_execute 中 executed_steps 预算比例
|
||||
plan_execute_max_step_result_runes: 4000 # plan_execute 每步结果最大字符数(超出截断)
|
||||
plan_execute_keep_last_steps: 8 # plan_execute 仅保留最近 N 步正文,早期步骤折叠为标题
|
||||
checkpoint_dir: "" # 非空:为 adk.NewRunner 启用按会话子目录的文件型 CheckPointStore,便于中断恢复持久化;Resume 的 HTTP/前端流程需另行对接
|
||||
deep_output_key: "" # 非空:将最终助手输出写入 adk session 的键名(Deep 与 Supervisor 主代理);空表示不写入
|
||||
deep_model_retry_max_retries: 0 # >0:ChatModel 调用失败时的框架级最大重试次数(Deep 与 Supervisor 主);0:不重试
|
||||
|
||||
+51
-32
@@ -39,6 +39,7 @@ type Agent struct {
|
||||
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
||||
currentConversationID string // 当前对话ID(用于自动传递给工具)
|
||||
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
||||
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
||||
}
|
||||
|
||||
// ResultStorage 结果存储接口(直接使用 storage 包的类型)
|
||||
@@ -162,6 +163,7 @@ func NewAgent(cfg *config.OpenAIConfig, agentCfg *config.AgentConfig, mcpServer
|
||||
resultStorage: resultStorage,
|
||||
largeResultThreshold: largeResultThreshold,
|
||||
toolNameMapping: make(map[string]string), // 初始化工具名称映射
|
||||
toolDescriptionMode: "short",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -336,10 +338,10 @@ func (fc *FunctionCall) UnmarshalJSON(data []byte) error {
|
||||
|
||||
// AgentLoopResult Agent Loop执行结果
|
||||
type AgentLoopResult struct {
|
||||
Response string
|
||||
MCPExecutionIDs []string
|
||||
LastReActInput string // 最后一轮ReAct的输入(压缩后的messages,JSON格式)
|
||||
LastReActOutput string // 最终大模型的输出
|
||||
Response string
|
||||
MCPExecutionIDs []string
|
||||
LastAgentTraceInput string // 最后一轮代理消息轨迹(压缩后的 messages,JSON;与 multiagent.RunResult 字段对齐)
|
||||
LastAgentTraceOutput string // 最终助手输出文本
|
||||
}
|
||||
|
||||
// ProgressCallback 进度回调函数类型
|
||||
@@ -471,7 +473,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
}
|
||||
|
||||
// 用于保存当前的messages,以便在异常情况下也能保存ReAct输入
|
||||
var currentReActInput string
|
||||
var currentAgentTraceInput string
|
||||
|
||||
maxIterations := a.maxIterations
|
||||
thinkingStreamSeq := 0
|
||||
@@ -490,9 +492,9 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
if err != nil {
|
||||
a.logger.Warn("序列化ReAct输入失败", zap.Error(err))
|
||||
} else {
|
||||
currentReActInput = string(messagesJSON)
|
||||
currentAgentTraceInput = string(messagesJSON)
|
||||
// 更新result中的值,确保始终保存最新的ReAct输入(压缩后的)
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
}
|
||||
|
||||
// 检查上下文是否已取消
|
||||
@@ -500,13 +502,13 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
case <-ctx.Done():
|
||||
// 上下文被取消(可能是用户主动暂停或其他原因)
|
||||
a.logger.Info("检测到上下文取消,保存当前ReAct数据", zap.Error(ctx.Err()))
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
if ctx.Err() == context.Canceled {
|
||||
result.Response = "任务已被取消。"
|
||||
} else {
|
||||
result.Response = fmt.Sprintf("任务执行中断: %v", ctx.Err())
|
||||
}
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
return result, ctx.Err()
|
||||
default:
|
||||
}
|
||||
@@ -600,10 +602,10 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
})
|
||||
if err != nil {
|
||||
// API调用失败,保存当前的ReAct输入和错误信息作为输出
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
errorMsg := fmt.Sprintf("调用OpenAI失败: %v", err)
|
||||
result.Response = errorMsg
|
||||
result.LastReActOutput = errorMsg
|
||||
result.LastAgentTraceOutput = errorMsg
|
||||
a.logger.Warn("OpenAI调用失败,已保存ReAct数据", zap.Error(err))
|
||||
return result, fmt.Errorf("调用OpenAI失败: %w", err)
|
||||
}
|
||||
@@ -629,19 +631,19 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
continue
|
||||
}
|
||||
// OpenAI返回错误,保存当前的ReAct输入和错误信息作为输出
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
errorMsg := fmt.Sprintf("OpenAI错误: %s", response.Error.Message)
|
||||
result.Response = errorMsg
|
||||
result.LastReActOutput = errorMsg
|
||||
result.LastAgentTraceOutput = errorMsg
|
||||
return result, fmt.Errorf("OpenAI错误: %s", response.Error.Message)
|
||||
}
|
||||
|
||||
if len(response.Choices) == 0 {
|
||||
// 没有收到响应,保存当前的ReAct输入和错误信息作为输出
|
||||
result.LastReActInput = currentReActInput
|
||||
result.LastAgentTraceInput = currentAgentTraceInput
|
||||
errorMsg := "没有收到响应"
|
||||
result.Response = errorMsg
|
||||
result.LastReActOutput = errorMsg
|
||||
result.LastAgentTraceOutput = errorMsg
|
||||
return result, fmt.Errorf("没有收到响应")
|
||||
}
|
||||
|
||||
@@ -816,7 +818,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
result.Response = streamText
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
sendProgress("progress", "总结生成完成", nil)
|
||||
return result, nil
|
||||
}
|
||||
@@ -863,14 +865,14 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
result.Response = streamText
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
sendProgress("progress", "总结生成完成", nil)
|
||||
return result, nil
|
||||
}
|
||||
// 如果获取总结失败,使用当前回复作为结果
|
||||
if choice.Message.Content != "" {
|
||||
result.Response = choice.Message.Content
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
return result, nil
|
||||
}
|
||||
// 如果都没有内容,跳出循环,让后续逻辑处理
|
||||
@@ -881,7 +883,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
if choice.FinishReason == "stop" {
|
||||
sendProgress("progress", "正在生成最终回复...", nil)
|
||||
result.Response = choice.Message.Content
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
return result, nil
|
||||
}
|
||||
}
|
||||
@@ -910,19 +912,19 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
|
||||
})
|
||||
if strings.TrimSpace(streamText) != "" {
|
||||
result.Response = streamText
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
sendProgress("progress", "总结生成完成", nil)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// 如果无法生成总结,返回友好的提示
|
||||
result.Response = fmt.Sprintf("已达到最大迭代次数(%d轮)。系统已执行了多轮测试,但由于达到迭代上限,无法继续自动执行。建议您查看已执行的工具结果,或提出新的测试请求以继续测试。", a.maxIterations)
|
||||
result.LastReActOutput = result.Response
|
||||
result.LastAgentTraceOutput = result.Response
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// getAvailableTools 获取可用工具
|
||||
// 从MCP服务器动态获取工具列表,使用简短描述以减少token消耗
|
||||
// 从MCP服务器动态获取工具列表,描述模式由 tool_description_mode 控制
|
||||
// roleTools: 角色配置的工具列表(toolKey格式),如果为空或nil,则使用所有工具(默认角色)
|
||||
func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||
// 构建角色工具集合(用于快速查找)
|
||||
@@ -946,11 +948,7 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||
continue // 不在角色工具列表中,跳过
|
||||
}
|
||||
}
|
||||
// 使用简短描述(如果存在),否则使用详细描述
|
||||
description := mcpTool.ShortDescription
|
||||
if description == "" {
|
||||
description = mcpTool.Description
|
||||
}
|
||||
description := a.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||
|
||||
// 转换schema中的类型为OpenAI标准类型
|
||||
convertedSchema := a.convertSchemaTypes(mcpTool.InputSchema)
|
||||
@@ -1024,11 +1022,7 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||
continue
|
||||
}
|
||||
|
||||
// 使用简短描述(如果存在),否则使用详细描述
|
||||
description := externalTool.ShortDescription
|
||||
if description == "" {
|
||||
description = externalTool.Description
|
||||
}
|
||||
description := a.pickToolDescription(externalTool.ShortDescription, externalTool.Description)
|
||||
|
||||
// 转换schema中的类型为OpenAI标准类型
|
||||
convertedSchema := a.convertSchemaTypes(externalTool.InputSchema)
|
||||
@@ -1063,6 +1057,19 @@ func (a *Agent) getAvailableTools(roleTools []string) []Tool {
|
||||
return tools
|
||||
}
|
||||
|
||||
func (a *Agent) pickToolDescription(shortDesc, fullDesc string) string {
|
||||
a.mu.RLock()
|
||||
mode := strings.TrimSpace(strings.ToLower(a.toolDescriptionMode))
|
||||
a.mu.RUnlock()
|
||||
if mode == "full" {
|
||||
return fullDesc
|
||||
}
|
||||
if shortDesc != "" {
|
||||
return shortDesc
|
||||
}
|
||||
return fullDesc
|
||||
}
|
||||
|
||||
// convertSchemaTypes 递归转换schema中的类型为OpenAI标准类型
|
||||
func (a *Agent) convertSchemaTypes(schema map[string]interface{}) map[string]interface{} {
|
||||
if schema == nil {
|
||||
@@ -1665,6 +1672,18 @@ func (a *Agent) UpdateMaxIterations(maxIterations int) {
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateToolDescriptionMode 更新工具描述模式(short/full)
|
||||
func (a *Agent) UpdateToolDescriptionMode(mode string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
mode = strings.TrimSpace(strings.ToLower(mode))
|
||||
if mode != "full" {
|
||||
mode = "short"
|
||||
}
|
||||
a.toolDescriptionMode = mode
|
||||
a.logger.Info("Agent工具描述模式已更新", zap.String("tool_description_mode", mode))
|
||||
}
|
||||
|
||||
// formatToolError 格式化工具错误信息,提供更友好的错误描述
|
||||
func (a *Agent) formatToolError(toolName string, args map[string]interface{}, err error) string {
|
||||
errorMsg := fmt.Sprintf(`工具执行失败
|
||||
|
||||
@@ -18,62 +18,62 @@ import (
|
||||
func setupTestAgent(t *testing.T) (*Agent, *storage.FileResultStorage) {
|
||||
logger := zap.NewNop()
|
||||
mcpServer := mcp.NewServer(logger)
|
||||
|
||||
|
||||
openAICfg := &config.OpenAIConfig{
|
||||
APIKey: "test-key",
|
||||
BaseURL: "https://api.test.com/v1",
|
||||
Model: "test-model",
|
||||
}
|
||||
|
||||
|
||||
agentCfg := &config.AgentConfig{
|
||||
MaxIterations: 10,
|
||||
LargeResultThreshold: 100, // 设置较小的阈值便于测试
|
||||
ResultStorageDir: "",
|
||||
}
|
||||
|
||||
|
||||
agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 10)
|
||||
|
||||
|
||||
// 创建测试存储
|
||||
tmpDir := filepath.Join(os.TempDir(), "test_agent_storage_"+time.Now().Format("20060102_150405"))
|
||||
testStorage, err := storage.NewFileResultStorage(tmpDir, logger)
|
||||
if err != nil {
|
||||
t.Fatalf("创建测试存储失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
agent.SetResultStorage(testStorage)
|
||||
|
||||
|
||||
return agent, testStorage
|
||||
}
|
||||
|
||||
func TestAgent_FormatMinimalNotification(t *testing.T) {
|
||||
agent, testStorage := setupTestAgent(t)
|
||||
_ = testStorage // 避免未使用变量警告
|
||||
|
||||
|
||||
executionID := "test_exec_001"
|
||||
toolName := "nmap_scan"
|
||||
size := 50000
|
||||
lineCount := 1000
|
||||
filePath := "tmp/test_exec_001.txt"
|
||||
|
||||
|
||||
notification := agent.formatMinimalNotification(executionID, toolName, size, lineCount, filePath)
|
||||
|
||||
|
||||
// 验证通知包含必要信息
|
||||
if !strings.Contains(notification, executionID) {
|
||||
t.Errorf("通知中应该包含执行ID: %s", executionID)
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(notification, toolName) {
|
||||
t.Errorf("通知中应该包含工具名称: %s", toolName)
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(notification, "50000") {
|
||||
t.Errorf("通知中应该包含大小信息")
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(notification, "1000") {
|
||||
t.Errorf("通知中应该包含行数信息")
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(notification, "query_execution_result") {
|
||||
t.Errorf("通知中应该包含查询工具的使用说明")
|
||||
}
|
||||
@@ -81,7 +81,7 @@ func TestAgent_FormatMinimalNotification(t *testing.T) {
|
||||
|
||||
func TestAgent_ExecuteToolViaMCP_LargeResult(t *testing.T) {
|
||||
agent, _ := setupTestAgent(t)
|
||||
|
||||
|
||||
// 创建模拟的MCP工具结果(大结果)
|
||||
largeResult := &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
@@ -92,59 +92,59 @@ func TestAgent_ExecuteToolViaMCP_LargeResult(t *testing.T) {
|
||||
},
|
||||
IsError: false,
|
||||
}
|
||||
|
||||
|
||||
// 模拟MCP服务器返回大结果
|
||||
// 由于我们需要模拟CallTool的行为,这里需要创建一个mock或者使用实际的MCP服务器
|
||||
// 为了简化测试,我们直接测试结果处理逻辑
|
||||
|
||||
|
||||
// 设置阈值
|
||||
agent.mu.Lock()
|
||||
agent.largeResultThreshold = 1000 // 设置较小的阈值
|
||||
agent.mu.Unlock()
|
||||
|
||||
|
||||
// 创建执行ID
|
||||
executionID := "test_exec_large_001"
|
||||
toolName := "test_tool"
|
||||
|
||||
|
||||
// 格式化结果
|
||||
var resultText strings.Builder
|
||||
for _, content := range largeResult.Content {
|
||||
resultText.WriteString(content.Text)
|
||||
resultText.WriteString("\n")
|
||||
}
|
||||
|
||||
|
||||
resultStr := resultText.String()
|
||||
resultSize := len(resultStr)
|
||||
|
||||
|
||||
// 检测大结果并保存
|
||||
agent.mu.RLock()
|
||||
threshold := agent.largeResultThreshold
|
||||
storage := agent.resultStorage
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if resultSize > threshold && storage != nil {
|
||||
// 保存大结果
|
||||
err := storage.SaveResult(executionID, toolName, resultStr)
|
||||
if err != nil {
|
||||
t.Fatalf("保存大结果失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
// 生成通知
|
||||
lines := strings.Split(resultStr, "\n")
|
||||
filePath := storage.GetResultPath(executionID)
|
||||
notification := agent.formatMinimalNotification(executionID, toolName, resultSize, len(lines), filePath)
|
||||
|
||||
|
||||
// 验证通知格式
|
||||
if !strings.Contains(notification, executionID) {
|
||||
t.Errorf("通知中应该包含执行ID")
|
||||
}
|
||||
|
||||
|
||||
// 验证结果已保存
|
||||
savedResult, err := storage.GetResult(executionID)
|
||||
if err != nil {
|
||||
t.Fatalf("获取保存的结果失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if savedResult != resultStr {
|
||||
t.Errorf("保存的结果与原始结果不匹配")
|
||||
}
|
||||
@@ -155,7 +155,7 @@ func TestAgent_ExecuteToolViaMCP_LargeResult(t *testing.T) {
|
||||
|
||||
func TestAgent_ExecuteToolViaMCP_SmallResult(t *testing.T) {
|
||||
agent, _ := setupTestAgent(t)
|
||||
|
||||
|
||||
// 创建小结果
|
||||
smallResult := &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
@@ -166,32 +166,32 @@ func TestAgent_ExecuteToolViaMCP_SmallResult(t *testing.T) {
|
||||
},
|
||||
IsError: false,
|
||||
}
|
||||
|
||||
|
||||
// 设置较大的阈值
|
||||
agent.mu.Lock()
|
||||
agent.largeResultThreshold = 100000 // 100KB
|
||||
agent.mu.Unlock()
|
||||
|
||||
|
||||
// 格式化结果
|
||||
var resultText strings.Builder
|
||||
for _, content := range smallResult.Content {
|
||||
resultText.WriteString(content.Text)
|
||||
resultText.WriteString("\n")
|
||||
}
|
||||
|
||||
|
||||
resultStr := resultText.String()
|
||||
resultSize := len(resultStr)
|
||||
|
||||
|
||||
// 检测大结果
|
||||
agent.mu.RLock()
|
||||
threshold := agent.largeResultThreshold
|
||||
storage := agent.resultStorage
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if resultSize > threshold && storage != nil {
|
||||
t.Fatal("小结果不应该被保存")
|
||||
}
|
||||
|
||||
|
||||
// 小结果应该直接返回
|
||||
if resultSize <= threshold {
|
||||
// 这是预期的行为
|
||||
@@ -203,26 +203,26 @@ func TestAgent_ExecuteToolViaMCP_SmallResult(t *testing.T) {
|
||||
|
||||
func TestAgent_SetResultStorage(t *testing.T) {
|
||||
agent, _ := setupTestAgent(t)
|
||||
|
||||
|
||||
// 创建新的存储
|
||||
tmpDir := filepath.Join(os.TempDir(), "test_new_storage_"+time.Now().Format("20060102_150405"))
|
||||
newStorage, err := storage.NewFileResultStorage(tmpDir, zap.NewNop())
|
||||
if err != nil {
|
||||
t.Fatalf("创建新存储失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
// 设置新存储
|
||||
agent.SetResultStorage(newStorage)
|
||||
|
||||
|
||||
// 验证存储已更新
|
||||
agent.mu.RLock()
|
||||
currentStorage := agent.resultStorage
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if currentStorage != newStorage {
|
||||
t.Fatal("存储未正确更新")
|
||||
}
|
||||
|
||||
|
||||
// 清理
|
||||
os.RemoveAll(tmpDir)
|
||||
}
|
||||
@@ -230,24 +230,24 @@ func TestAgent_SetResultStorage(t *testing.T) {
|
||||
func TestAgent_NewAgent_DefaultValues(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
mcpServer := mcp.NewServer(logger)
|
||||
|
||||
|
||||
openAICfg := &config.OpenAIConfig{
|
||||
APIKey: "test-key",
|
||||
BaseURL: "https://api.test.com/v1",
|
||||
Model: "test-model",
|
||||
}
|
||||
|
||||
|
||||
// 测试默认配置
|
||||
agent := NewAgent(openAICfg, nil, mcpServer, nil, logger, 0)
|
||||
|
||||
|
||||
if agent.maxIterations != 30 {
|
||||
t.Errorf("默认迭代次数不匹配。期望: 30, 实际: %d", agent.maxIterations)
|
||||
}
|
||||
|
||||
|
||||
agent.mu.RLock()
|
||||
threshold := agent.largeResultThreshold
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if threshold != 50*1024 {
|
||||
t.Errorf("默认阈值不匹配。期望: %d, 实际: %d", 50*1024, threshold)
|
||||
}
|
||||
@@ -256,31 +256,30 @@ func TestAgent_NewAgent_DefaultValues(t *testing.T) {
|
||||
func TestAgent_NewAgent_CustomConfig(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
mcpServer := mcp.NewServer(logger)
|
||||
|
||||
|
||||
openAICfg := &config.OpenAIConfig{
|
||||
APIKey: "test-key",
|
||||
BaseURL: "https://api.test.com/v1",
|
||||
Model: "test-model",
|
||||
}
|
||||
|
||||
|
||||
agentCfg := &config.AgentConfig{
|
||||
MaxIterations: 20,
|
||||
LargeResultThreshold: 100 * 1024, // 100KB
|
||||
ResultStorageDir: "custom_tmp",
|
||||
}
|
||||
|
||||
|
||||
agent := NewAgent(openAICfg, agentCfg, mcpServer, nil, logger, 15)
|
||||
|
||||
|
||||
if agent.maxIterations != 15 {
|
||||
t.Errorf("迭代次数不匹配。期望: 15, 实际: %d", agent.maxIterations)
|
||||
}
|
||||
|
||||
|
||||
agent.mu.RLock()
|
||||
threshold := agent.largeResultThreshold
|
||||
agent.mu.RUnlock()
|
||||
|
||||
|
||||
if threshold != 100*1024 {
|
||||
t.Errorf("阈值不匹配。期望: %d, 实际: %d", 100*1024, threshold)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -256,11 +256,11 @@ func orchestratorConfigFromOrchestrator(o *OrchestratorMarkdown) config.MultiAge
|
||||
return config.MultiAgentSubConfig{}
|
||||
}
|
||||
return config.MultiAgentSubConfig{
|
||||
ID: o.EinoName,
|
||||
Name: o.DisplayName,
|
||||
Description: o.Description,
|
||||
Instruction: o.Instruction,
|
||||
Kind: "orchestrator",
|
||||
ID: o.EinoName,
|
||||
Name: o.DisplayName,
|
||||
Description: o.Description,
|
||||
Instruction: o.Instruction,
|
||||
Kind: "orchestrator",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -133,6 +133,7 @@ func New(cfg *config.Config, log *logger.Logger) (*App, error) {
|
||||
maxIterations = 30 // 默认值
|
||||
}
|
||||
agent := agent.NewAgent(&cfg.OpenAI, &cfg.Agent, mcpServer, externalMCPMgr, log.Logger, maxIterations)
|
||||
agent.UpdateToolDescriptionMode(cfg.Security.ToolDescriptionMode)
|
||||
|
||||
// 设置结果存储到Agent
|
||||
agent.SetResultStorage(resultStorage)
|
||||
|
||||
@@ -145,7 +145,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
}
|
||||
|
||||
// 1. 优先尝试从数据库获取保存的最后一轮ReAct输入和输出
|
||||
reactInputJSON, modelOutput, err := b.db.GetReActData(conversationID)
|
||||
reactInputJSON, modelOutput, err := b.db.GetAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
b.logger.Warn("获取保存的ReAct数据失败,将使用消息历史构建", zap.Error(err))
|
||||
// 继续使用原来的逻辑
|
||||
@@ -170,7 +170,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
messageCount = len(tempMessages)
|
||||
}
|
||||
|
||||
dataSource = "database_last_react_input"
|
||||
dataSource = "database_last_agent_trace"
|
||||
b.logger.Info("使用保存的ReAct数据构建攻击链",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("dataSource", dataSource),
|
||||
@@ -183,7 +183,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
// userInput = b.extractUserInputFromReActInput(reactInputJSON)
|
||||
|
||||
// 将JSON格式的messages转换为可读格式
|
||||
reactInputFinal = b.formatReActInputFromJSON(reactInputJSON)
|
||||
reactInputFinal = b.formatAgentTraceInputFromJSON(reactInputJSON)
|
||||
} else {
|
||||
// 2. 如果没有保存的ReAct数据,从对话消息构建
|
||||
dataSource = "messages_table"
|
||||
@@ -201,7 +201,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
}
|
||||
|
||||
// 提取最后一轮ReAct的输入(历史消息+当前用户输入)
|
||||
reactInputFinal = b.buildReActInput(messages)
|
||||
reactInputFinal = b.buildAgentTraceInput(messages)
|
||||
|
||||
// 提取大模型最后的输出(最后一条assistant消息)
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
@@ -212,7 +212,7 @@ func (b *Builder) BuildChainFromConversation(ctx context.Context, conversationID
|
||||
}
|
||||
}
|
||||
|
||||
// 多代理:保存的 last_react_input 可能仅为首轮用户消息,不含工具轨迹;补充最后一轮助手的过程详情(与单代理「最后一轮 ReAct」对齐)
|
||||
// 多代理:保存的轨迹列可能仅为首轮用户消息,不含工具轨迹;补充最后一轮助手的过程详情(与单代理完整轨迹对齐)
|
||||
hasMCPOnAssistant := false
|
||||
var lastAssistantID string
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
@@ -366,8 +366,8 @@ func (b *Builder) formatProcessDetailsForAttackChain(details []database.ProcessD
|
||||
return strings.TrimSpace(sb.String())
|
||||
}
|
||||
|
||||
// buildReActInput 构建最后一轮ReAct的输入(历史消息+当前用户输入)
|
||||
func (b *Builder) buildReActInput(messages []database.Message) string {
|
||||
// buildAgentTraceInput 构建最后一轮ReAct的输入(历史消息+当前用户输入)
|
||||
func (b *Builder) buildAgentTraceInput(messages []database.Message) string {
|
||||
var builder strings.Builder
|
||||
for _, msg := range messages {
|
||||
builder.WriteString(fmt.Sprintf("[%s]: %s\n\n", msg.Role, msg.Content))
|
||||
@@ -396,8 +396,8 @@ func (b *Builder) buildReActInput(messages []database.Message) string {
|
||||
// return ""
|
||||
// }
|
||||
|
||||
// formatReActInputFromJSON 将JSON格式的messages数组转换为可读的字符串格式
|
||||
func (b *Builder) formatReActInputFromJSON(reactInputJSON string) string {
|
||||
// formatAgentTraceInputFromJSON 将JSON格式的messages数组转换为可读的字符串格式
|
||||
func (b *Builder) formatAgentTraceInputFromJSON(reactInputJSON string) string {
|
||||
var messages []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(reactInputJSON), &messages); err != nil {
|
||||
b.logger.Warn("解析ReAct输入JSON失败", zap.Error(err))
|
||||
|
||||
@@ -72,6 +72,8 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
ToolSearchEnable bool `yaml:"tool_search_enable,omitempty" json:"tool_search_enable,omitempty"`
|
||||
ToolSearchMinTools int `yaml:"tool_search_min_tools,omitempty" json:"tool_search_min_tools,omitempty"` // default 20; applies when len(tools) >= this
|
||||
ToolSearchAlwaysVisible int `yaml:"tool_search_always_visible,omitempty" json:"tool_search_always_visible,omitempty"` // default 12; first N tools stay always visible
|
||||
// ToolSearchAlwaysVisibleTools keeps specified tool names always visible (never hidden by tool_search).
|
||||
ToolSearchAlwaysVisibleTools []string `yaml:"tool_search_always_visible_tools,omitempty" json:"tool_search_always_visible_tools,omitempty"`
|
||||
// Plantask adds TaskCreate/Get/Update/List (file-backed under skills dir); requires eino_skills + local backend.
|
||||
PlantaskEnable bool `yaml:"plantask_enable,omitempty" json:"plantask_enable,omitempty"`
|
||||
// PlantaskRelDir relative to skills_dir for per-conversation task boards (default .eino/plantask).
|
||||
@@ -79,8 +81,24 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
// Reduction truncates/offloads large tool outputs (requires eino local backend for Write).
|
||||
ReductionEnable bool `yaml:"reduction_enable,omitempty" json:"reduction_enable,omitempty"`
|
||||
ReductionRootDir string `yaml:"reduction_root_dir,omitempty" json:"reduction_root_dir,omitempty"` // default: os temp + conversation id
|
||||
ReductionMaxLengthForTrunc int `yaml:"reduction_max_length_for_trunc,omitempty" json:"reduction_max_length_for_trunc,omitempty"` // default 12000
|
||||
ReductionMaxTokensForClear int `yaml:"reduction_max_tokens_for_clear,omitempty" json:"reduction_max_tokens_for_clear,omitempty"` // default 50000
|
||||
ReductionClearExclude []string `yaml:"reduction_clear_exclude,omitempty" json:"reduction_clear_exclude,omitempty"`
|
||||
ReductionSubAgents bool `yaml:"reduction_sub_agents,omitempty" json:"reduction_sub_agents,omitempty"` // also attach to sub-agents
|
||||
// SummarizationTriggerRatio controls summarization trigger threshold as max_total_tokens * ratio (default 0.8).
|
||||
SummarizationTriggerRatio float64 `yaml:"summarization_trigger_ratio,omitempty" json:"summarization_trigger_ratio,omitempty"`
|
||||
// SummarizationEmitInternalEvents controls middleware internal event emission (default true).
|
||||
SummarizationEmitInternalEvents *bool `yaml:"summarization_emit_internal_events,omitempty" json:"summarization_emit_internal_events,omitempty"`
|
||||
// HistoryInputBudgetRatio caps pre-agent history tokens as max_total_tokens * ratio (default 0.35).
|
||||
HistoryInputBudgetRatio float64 `yaml:"history_input_budget_ratio,omitempty" json:"history_input_budget_ratio,omitempty"`
|
||||
// PlanExecuteUserInputBudgetRatio caps planner/replanner/executor userInput prompt budget ratio (default 0.35).
|
||||
PlanExecuteUserInputBudgetRatio float64 `yaml:"plan_execute_user_input_budget_ratio,omitempty" json:"plan_execute_user_input_budget_ratio,omitempty"`
|
||||
// PlanExecuteExecutedStepsBudgetRatio caps executed_steps prompt budget ratio (default 0.2).
|
||||
PlanExecuteExecutedStepsBudgetRatio float64 `yaml:"plan_execute_executed_steps_budget_ratio,omitempty" json:"plan_execute_executed_steps_budget_ratio,omitempty"`
|
||||
// PlanExecuteMaxStepResultRunes caps each executed step result length for prompt view (default 4000).
|
||||
PlanExecuteMaxStepResultRunes int `yaml:"plan_execute_max_step_result_runes,omitempty" json:"plan_execute_max_step_result_runes,omitempty"`
|
||||
// PlanExecuteKeepLastSteps keeps only the tail steps in prompt view (default 8).
|
||||
PlanExecuteKeepLastSteps int `yaml:"plan_execute_keep_last_steps,omitempty" json:"plan_execute_keep_last_steps,omitempty"`
|
||||
// CheckpointDir when non-empty enables adk.Runner CheckPointStore (file-backed) for interrupt/resume persistence.
|
||||
CheckpointDir string `yaml:"checkpoint_dir,omitempty" json:"checkpoint_dir,omitempty"`
|
||||
// DeepOutputKey passed to deep.Config OutputKey (session final text); empty = off.
|
||||
@@ -91,6 +109,97 @@ type MultiAgentEinoMiddlewareConfig struct {
|
||||
TaskToolDescriptionPrefix string `yaml:"task_tool_description_prefix,omitempty" json:"task_tool_description_prefix,omitempty"`
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) SummarizationTriggerRatioEffective() float64 {
|
||||
v := c.SummarizationTriggerRatio
|
||||
if v <= 0 {
|
||||
return 0.8
|
||||
}
|
||||
if v < 0.5 {
|
||||
return 0.5
|
||||
}
|
||||
if v > 0.95 {
|
||||
return 0.95
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) SummarizationEmitInternalEventsEffective() bool {
|
||||
if c.SummarizationEmitInternalEvents != nil {
|
||||
return *c.SummarizationEmitInternalEvents
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) HistoryInputBudgetRatioEffective() float64 {
|
||||
v := c.HistoryInputBudgetRatio
|
||||
if v <= 0 {
|
||||
return 0.35
|
||||
}
|
||||
if v < 0.15 {
|
||||
return 0.15
|
||||
}
|
||||
if v > 0.6 {
|
||||
return 0.6
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteUserInputBudgetRatioEffective() float64 {
|
||||
v := c.PlanExecuteUserInputBudgetRatio
|
||||
if v <= 0 {
|
||||
return 0.35
|
||||
}
|
||||
if v < 0.1 {
|
||||
return 0.1
|
||||
}
|
||||
if v > 0.6 {
|
||||
return 0.6
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteExecutedStepsBudgetRatioEffective() float64 {
|
||||
v := c.PlanExecuteExecutedStepsBudgetRatio
|
||||
if v <= 0 {
|
||||
return 0.2
|
||||
}
|
||||
if v < 0.08 {
|
||||
return 0.08
|
||||
}
|
||||
if v > 0.5 {
|
||||
return 0.5
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteMaxStepResultRunesEffective() int {
|
||||
if c.PlanExecuteMaxStepResultRunes > 0 {
|
||||
return c.PlanExecuteMaxStepResultRunes
|
||||
}
|
||||
return 4000
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) PlanExecuteKeepLastStepsEffective() int {
|
||||
if c.PlanExecuteKeepLastSteps > 0 {
|
||||
return c.PlanExecuteKeepLastSteps
|
||||
}
|
||||
return 8
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) ReductionMaxLengthForTruncEffective() int {
|
||||
if c.ReductionMaxLengthForTrunc > 0 {
|
||||
return c.ReductionMaxLengthForTrunc
|
||||
}
|
||||
return 12000
|
||||
}
|
||||
|
||||
func (c MultiAgentEinoMiddlewareConfig) ReductionMaxTokensForClearEffective() int {
|
||||
if c.ReductionMaxTokensForClear > 0 {
|
||||
return c.ReductionMaxTokensForClear
|
||||
}
|
||||
return 50000
|
||||
}
|
||||
|
||||
// MultiAgentEinoSkillsConfig toggles Eino official skill progressive disclosure and host filesystem tools.
|
||||
type MultiAgentEinoSkillsConfig struct {
|
||||
// Disable skips skill middleware (and does not attach local FS tools for Deep).
|
||||
@@ -137,6 +246,8 @@ type MultiAgentPublic struct {
|
||||
SubAgentCount int `json:"sub_agent_count"`
|
||||
Orchestration string `json:"orchestration,omitempty"`
|
||||
PlanExecuteLoopMaxIterations int `json:"plan_execute_loop_max_iterations"`
|
||||
ToolSearchAlwaysVisibleTools []string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
ToolSearchAlwaysVisibleEffectiveTools []string `json:"tool_search_always_visible_effective_tools,omitempty"`
|
||||
}
|
||||
|
||||
// NormalizeMultiAgentOrchestration 返回 deep、plan_execute 或 supervisor。
|
||||
@@ -158,6 +269,7 @@ type MultiAgentAPIUpdate struct {
|
||||
RobotUseMultiAgent bool `json:"robot_use_multi_agent"`
|
||||
BatchUseMultiAgent bool `json:"batch_use_multi_agent"`
|
||||
PlanExecuteLoopMaxIterations *int `json:"plan_execute_loop_max_iterations,omitempty"`
|
||||
ToolSearchAlwaysVisibleTools []string `json:"tool_search_always_visible_tools,omitempty"`
|
||||
}
|
||||
|
||||
// RobotsConfig 机器人配置(企业微信、钉钉、飞书等)
|
||||
|
||||
@@ -165,4 +165,3 @@ func (db *DB) DeleteAttackChain(conversationID string) error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -4,6 +4,8 @@ import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -308,7 +310,7 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
||||
func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversation, error) {
|
||||
var rows *sql.Rows
|
||||
var err error
|
||||
|
||||
|
||||
if search != "" {
|
||||
// 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积
|
||||
searchPattern := "%" + search + "%"
|
||||
@@ -327,7 +329,7 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
|
||||
limit, offset,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("查询对话列表失败: %w", err)
|
||||
}
|
||||
@@ -416,25 +418,34 @@ func (db *DB) DeleteConversation(id string) error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("删除对话失败: %w", err)
|
||||
}
|
||||
// Best-effort cleanup for conversation-scoped filesystem artifacts
|
||||
// (e.g., summarization transcript, reduction/checkpoint files under conversation_artifacts/<id>).
|
||||
if base := strings.TrimSpace(db.conversationArtifactsDir); base != "" {
|
||||
artDir := filepath.Join(base, id)
|
||||
if rmErr := os.RemoveAll(artDir); rmErr != nil {
|
||||
db.logger.Warn("删除会话 artifacts 目录失败", zap.String("conversationId", id), zap.String("dir", artDir), zap.Error(rmErr))
|
||||
}
|
||||
}
|
||||
|
||||
db.logger.Info("对话及其所有相关数据已删除", zap.String("conversationId", id))
|
||||
return nil
|
||||
}
|
||||
|
||||
// SaveReActData 保存最后一轮ReAct的输入和输出
|
||||
func (db *DB) SaveReActData(conversationID, reactInput, reactOutput string) error {
|
||||
// SaveAgentTrace 保存最后一轮代理消息轨迹与助手输出摘要。
|
||||
// SQLite 列名仍为 last_react_input / last_react_output,与历史库表兼容;语义上为「全模式代理轨迹」,非仅 ReAct。
|
||||
func (db *DB) SaveAgentTrace(conversationID, traceInputJSON, assistantOutput string) error {
|
||||
_, err := db.Exec(
|
||||
"UPDATE conversations SET last_react_input = ?, last_react_output = ?, updated_at = ? WHERE id = ?",
|
||||
reactInput, reactOutput, time.Now(), conversationID,
|
||||
traceInputJSON, assistantOutput, time.Now(), conversationID,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("保存ReAct数据失败: %w", err)
|
||||
return fmt.Errorf("保存代理轨迹失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetReActData 获取最后一轮ReAct的输入和输出
|
||||
func (db *DB) GetReActData(conversationID string) (reactInput, reactOutput string, err error) {
|
||||
// GetAgentTrace 读取 conversations 中保存的代理轨迹(列名 last_react_*)。
|
||||
func (db *DB) GetAgentTrace(conversationID string) (traceInputJSON, assistantOutput string, err error) {
|
||||
var input, output sql.NullString
|
||||
err = db.QueryRow(
|
||||
"SELECT last_react_input, last_react_output FROM conversations WHERE id = ?",
|
||||
@@ -444,17 +455,17 @@ func (db *DB) GetReActData(conversationID string) (reactInput, reactOutput strin
|
||||
if err == sql.ErrNoRows {
|
||||
return "", "", fmt.Errorf("对话不存在")
|
||||
}
|
||||
return "", "", fmt.Errorf("获取ReAct数据失败: %w", err)
|
||||
return "", "", fmt.Errorf("获取代理轨迹失败: %w", err)
|
||||
}
|
||||
|
||||
if input.Valid {
|
||||
reactInput = input.String
|
||||
traceInputJSON = input.String
|
||||
}
|
||||
if output.Valid {
|
||||
reactOutput = output.String
|
||||
assistantOutput = output.String
|
||||
}
|
||||
|
||||
return reactInput, reactOutput, nil
|
||||
return traceInputJSON, assistantOutput, nil
|
||||
}
|
||||
|
||||
// ConversationHasToolProcessDetails 对话是否存在已落库的工具调用/结果(用于多代理等场景下 MCP execution id 未汇总时的攻击链判定)。
|
||||
|
||||
@@ -3,6 +3,8 @@ package database
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -21,7 +23,8 @@ func configureDBPool(db *sql.DB) {
|
||||
// DB 数据库连接
|
||||
type DB struct {
|
||||
*sql.DB
|
||||
logger *zap.Logger
|
||||
logger *zap.Logger
|
||||
conversationArtifactsDir string
|
||||
}
|
||||
|
||||
// NewDB 创建数据库连接
|
||||
@@ -41,6 +44,13 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
DB: db,
|
||||
logger: logger,
|
||||
}
|
||||
// Keep conversation-scoped artifacts near database files, so cleanup can follow conversation lifecycle.
|
||||
baseDir := filepath.Join(filepath.Dir(dbPath), "conversation_artifacts")
|
||||
if mkErr := os.MkdirAll(baseDir, 0o755); mkErr == nil {
|
||||
database.conversationArtifactsDir = baseDir
|
||||
} else if logger != nil {
|
||||
logger.Warn("创建 conversation artifacts 目录失败", zap.String("dir", baseDir), zap.Error(mkErr))
|
||||
}
|
||||
|
||||
// 初始化表
|
||||
if err := database.initTables(); err != nil {
|
||||
@@ -52,7 +62,7 @@ func NewDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||
|
||||
// initTables 初始化数据库表
|
||||
func (db *DB) initTables() error {
|
||||
// 创建对话表
|
||||
// 创建对话表(last_react_input / last_react_output 存「代理消息轨迹」JSON 与助手摘要,列名保留以兼容已有库)
|
||||
createConversationsTable := `
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
id TEXT PRIMARY KEY,
|
||||
|
||||
@@ -12,7 +12,7 @@ import (
|
||||
// Vulnerability 漏洞
|
||||
type Vulnerability struct {
|
||||
ID string `json:"id"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
ConversationTag string `json:"conversation_tag,omitempty"`
|
||||
TaskTag string `json:"task_tag,omitempty"`
|
||||
TaskID string `json:"task_id,omitempty"`
|
||||
@@ -367,4 +367,3 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
|
||||
"task_tags": taskTags,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
+84
-92
@@ -497,10 +497,10 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
defer h.hitlManager.DeactivateConversation(conversationID)
|
||||
}
|
||||
|
||||
// 优先尝试从保存的ReAct数据恢复历史上下文
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
// 优先尝试从保存的代理轨迹恢复历史上下文
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
h.logger.Warn("从ReAct数据加载历史消息失败,使用消息表", zap.Error(err))
|
||||
h.logger.Warn("从代理轨迹加载历史消息失败,使用消息表", zap.Error(err))
|
||||
// 回退到使用数据库消息表
|
||||
historyMessages, err := h.db.GetMessages(conversationID)
|
||||
if err != nil {
|
||||
@@ -518,7 +518,7 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
h.logger.Info("从消息表加载历史消息", zap.Int("count", len(agentHistoryMessages)))
|
||||
}
|
||||
} else {
|
||||
h.logger.Info("从ReAct数据恢复历史上下文", zap.Int("count", len(agentHistoryMessages)))
|
||||
h.logger.Info("从代理轨迹恢复历史上下文", zap.Int("count", len(agentHistoryMessages)))
|
||||
}
|
||||
|
||||
// 校验附件数量(非流式)
|
||||
@@ -613,12 +613,12 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
if err != nil {
|
||||
h.logger.Error("Agent Loop执行失败", zap.Error(err))
|
||||
|
||||
// 即使执行失败,也尝试保存ReAct数据(如果result中有)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if saveErr := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); saveErr != nil {
|
||||
h.logger.Warn("保存失败任务的ReAct数据失败", zap.Error(saveErr))
|
||||
// 即使执行失败,也尝试保存代理轨迹(如果 result 中有)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if saveErr := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); saveErr != nil {
|
||||
h.logger.Warn("保存失败任务的代理轨迹失败", zap.Error(saveErr))
|
||||
} else {
|
||||
h.logger.Info("已保存失败任务的ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存失败任务的代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -634,12 +634,12 @@ func (h *AgentHandler) AgentLoop(c *gin.Context) {
|
||||
// 因为AI已经生成了回复,用户应该能看到
|
||||
}
|
||||
|
||||
// 保存最后一轮ReAct的输入和输出
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存ReAct数据失败", zap.Error(err))
|
||||
// 保存最后一轮代理轨迹与助手输出
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -666,7 +666,7 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
}
|
||||
}
|
||||
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
historyMessages, getErr := h.db.GetMessages(conversationID)
|
||||
if getErr != nil {
|
||||
@@ -722,6 +722,7 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
"deep",
|
||||
)
|
||||
if errMA != nil {
|
||||
h.persistEinoAgentTraceForResume(conversationID, resultMA)
|
||||
errMsg := "执行失败: " + errMA.Error()
|
||||
if assistantMessageID != "" {
|
||||
_, _ = h.db.Exec("UPDATE messages SET content = ? WHERE id = ?", errMsg, assistantMessageID)
|
||||
@@ -747,8 +748,8 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
h.logger.Warn("机器人:保存助手消息失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
if resultMA.LastReActInput != "" || resultMA.LastReActOutput != "" {
|
||||
_ = h.db.SaveReActData(conversationID, resultMA.LastReActInput, resultMA.LastReActOutput)
|
||||
if resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "" {
|
||||
_ = h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput)
|
||||
}
|
||||
return resultMA.Response, conversationID, nil
|
||||
}
|
||||
@@ -782,8 +783,8 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, conversationI
|
||||
h.logger.Warn("机器人:保存助手消息失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
_ = h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput)
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
_ = h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput)
|
||||
}
|
||||
return result.Response, conversationID, nil
|
||||
}
|
||||
@@ -1359,10 +1360,10 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
}
|
||||
ssePublishConversationID = conversationID
|
||||
|
||||
// 优先尝试从保存的ReAct数据恢复历史上下文
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
// 优先尝试从保存的代理轨迹恢复历史上下文
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
h.logger.Warn("从ReAct数据加载历史消息失败,使用消息表", zap.Error(err))
|
||||
h.logger.Warn("从代理轨迹加载历史消息失败,使用消息表", zap.Error(err))
|
||||
// 回退到使用数据库消息表
|
||||
historyMessages, err := h.db.GetMessages(conversationID)
|
||||
if err != nil {
|
||||
@@ -1380,7 +1381,7 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
h.logger.Info("从消息表加载历史消息", zap.Int("count", len(agentHistoryMessages)))
|
||||
}
|
||||
} else {
|
||||
h.logger.Info("从ReAct数据恢复历史上下文", zap.Int("count", len(agentHistoryMessages)))
|
||||
h.logger.Info("从代理轨迹恢复历史上下文", zap.Int("count", len(agentHistoryMessages)))
|
||||
}
|
||||
|
||||
// 校验附件数量
|
||||
@@ -1579,12 +1580,12 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
h.db.AddProcessDetail(assistantMessageID, conversationID, "cancelled", cancelMsg, nil)
|
||||
}
|
||||
|
||||
// 即使任务被取消,也尝试保存ReAct数据(如果result中有)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的ReAct数据失败", zap.Error(err))
|
||||
// 即使任务被取消,也尝试保存代理轨迹(如果 result 中有)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存取消任务的ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存取消任务的代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1614,12 +1615,12 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
h.db.AddProcessDetail(assistantMessageID, conversationID, "timeout", timeoutMsg, nil)
|
||||
}
|
||||
|
||||
// 即使任务超时,也尝试保存ReAct数据(如果result中有)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存超时任务的ReAct数据失败", zap.Error(err))
|
||||
// 即使任务超时,也尝试保存代理轨迹(如果 result 中有)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存超时任务的代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存超时任务的ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存超时任务的代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1649,12 +1650,12 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
h.db.AddProcessDetail(assistantMessageID, conversationID, "error", errorMsg, nil)
|
||||
}
|
||||
|
||||
// 即使任务失败,也尝试保存ReAct数据(如果result中有)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存失败任务的ReAct数据失败", zap.Error(err))
|
||||
// 即使任务失败,也尝试保存代理轨迹(如果 result 中有)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存失败任务的代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存失败任务的ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存失败任务的代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1694,12 +1695,12 @@ func (h *AgentHandler) AgentLoopStream(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
// 保存最后一轮ReAct的输入和输出
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存ReAct数据失败", zap.Error(err))
|
||||
// 保存最后一轮代理轨迹与助手输出
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存ReAct数据", zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存代理轨迹", zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2499,6 +2500,9 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
cancel()
|
||||
|
||||
if runErr != nil {
|
||||
if useRunResult {
|
||||
h.persistEinoAgentTraceForResume(conversationID, resultMA)
|
||||
}
|
||||
// 检查是否是取消错误
|
||||
// 1. 直接检查是否是 context.Canceled(包括包装后的错误)
|
||||
// 2. 检查错误消息中是否包含"context canceled"或"cancelled"关键字
|
||||
@@ -2542,14 +2546,14 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
h.logger.Warn("保存取消消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(errMsg))
|
||||
}
|
||||
}
|
||||
// 保存ReAct数据(如果存在)
|
||||
if result != nil && (result.LastReActInput != "" || result.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的ReAct数据失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
// 保存代理轨迹(如果存在)
|
||||
if result != nil && (result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
}
|
||||
} else if useRunResult && resultMA != nil && (resultMA.LastReActInput != "" || resultMA.LastReActOutput != "") {
|
||||
if err := h.db.SaveReActData(conversationID, resultMA.LastReActInput, resultMA.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的ReAct数据失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
} else if useRunResult && resultMA != nil && (resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "") {
|
||||
if err := h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存取消任务的代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, "cancelled", cancelMsg, "", conversationID)
|
||||
@@ -2581,13 +2585,13 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
if useRunResult {
|
||||
resText = resultMA.Response
|
||||
mcpIDs = resultMA.MCPExecutionIDs
|
||||
lastIn = resultMA.LastReActInput
|
||||
lastOut = resultMA.LastReActOutput
|
||||
lastIn = resultMA.LastAgentTraceInput
|
||||
lastOut = resultMA.LastAgentTraceOutput
|
||||
} else {
|
||||
resText = result.Response
|
||||
mcpIDs = result.MCPExecutionIDs
|
||||
lastIn = result.LastReActInput
|
||||
lastOut = result.LastReActOutput
|
||||
lastIn = result.LastAgentTraceInput
|
||||
lastOut = result.LastAgentTraceOutput
|
||||
}
|
||||
|
||||
// 更新助手消息内容
|
||||
@@ -2618,12 +2622,12 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
}
|
||||
}
|
||||
|
||||
// 保存ReAct数据
|
||||
// 保存代理轨迹
|
||||
if lastIn != "" || lastOut != "" {
|
||||
if err := h.db.SaveReActData(conversationID, lastIn, lastOut); err != nil {
|
||||
h.logger.Warn("保存ReAct数据失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
if err := h.db.SaveAgentTrace(conversationID, lastIn, lastOut); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(err))
|
||||
} else {
|
||||
h.logger.Info("已保存ReAct数据", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
|
||||
h.logger.Info("已保存代理轨迹", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2642,36 +2646,33 @@ func (h *AgentHandler) executeBatchQueue(queueID string) {
|
||||
}
|
||||
}
|
||||
|
||||
// loadHistoryFromReActData 从保存的ReAct数据恢复历史消息上下文
|
||||
// 采用与攻击链生成类似的拼接逻辑:优先使用保存的last_react_input和last_react_output,若不存在则回退到消息表
|
||||
func (h *AgentHandler) loadHistoryFromReActData(conversationID string) ([]agent.ChatMessage, error) {
|
||||
// 获取保存的ReAct输入和输出
|
||||
reactInputJSON, reactOutput, err := h.db.GetReActData(conversationID)
|
||||
// loadHistoryFromAgentTrace 从库中保存的代理消息轨迹恢复历史(列 last_react_*;含单代理与 Eino)。
|
||||
// 逻辑与攻击链一致:优先用已保存的 JSON 消息带 + 最后一轮助手摘要,否则回退消息表。
|
||||
func (h *AgentHandler) loadHistoryFromAgentTrace(conversationID string) ([]agent.ChatMessage, error) {
|
||||
traceInputJSON, assistantOut, err := h.db.GetAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("获取ReAct数据失败: %w", err)
|
||||
return nil, fmt.Errorf("获取代理轨迹失败: %w", err)
|
||||
}
|
||||
|
||||
// 如果last_react_input为空,回退到使用消息表(与攻击链生成逻辑一致)
|
||||
if reactInputJSON == "" {
|
||||
return nil, fmt.Errorf("ReAct数据为空,将使用消息表")
|
||||
if traceInputJSON == "" {
|
||||
return nil, fmt.Errorf("代理轨迹为空,将使用消息表")
|
||||
}
|
||||
|
||||
dataSource := "database_last_react_input"
|
||||
dataSource := "database_last_agent_trace"
|
||||
|
||||
// 解析JSON格式的messages数组
|
||||
var messagesArray []map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(reactInputJSON), &messagesArray); err != nil {
|
||||
return nil, fmt.Errorf("解析ReAct输入JSON失败: %w", err)
|
||||
if err := json.Unmarshal([]byte(traceInputJSON), &messagesArray); err != nil {
|
||||
return nil, fmt.Errorf("解析代理轨迹 JSON 失败: %w", err)
|
||||
}
|
||||
|
||||
messageCount := len(messagesArray)
|
||||
|
||||
h.logger.Info("使用保存的ReAct数据恢复历史上下文",
|
||||
h.logger.Info("使用保存的代理轨迹恢复历史上下文",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("dataSource", dataSource),
|
||||
zap.Int("reactInputSize", len(reactInputJSON)),
|
||||
zap.Int("traceInputSize", len(traceInputJSON)),
|
||||
zap.Int("messageCount", messageCount),
|
||||
zap.Int("reactOutputSize", len(reactOutput)),
|
||||
zap.Int("assistantOutSize", len(assistantOut)),
|
||||
)
|
||||
// fmt.Println("messagesArray:", messagesArray)//debug
|
||||
|
||||
@@ -2755,53 +2756,44 @@ func (h *AgentHandler) loadHistoryFromReActData(conversationID string) ([]agent.
|
||||
agentMessages = append(agentMessages, msg)
|
||||
}
|
||||
|
||||
// 如果存在last_react_output,需要将其作为最后一条assistant消息
|
||||
// 因为last_react_input是在迭代开始前保存的,不包含最后一轮的最终输出
|
||||
if reactOutput != "" {
|
||||
// 检查最后一条消息是否是assistant消息且没有tool_calls
|
||||
// 如果有tool_calls,说明后面应该还有tool消息和最终的assistant回复
|
||||
// 若存在 last_react_output(助手摘要),合并为最后一条 assistant(与保存格式一致)
|
||||
if assistantOut != "" {
|
||||
if len(agentMessages) > 0 {
|
||||
lastMsg := &agentMessages[len(agentMessages)-1]
|
||||
if strings.EqualFold(lastMsg.Role, "assistant") && len(lastMsg.ToolCalls) == 0 {
|
||||
// 最后一条是assistant消息且没有tool_calls,用最终输出更新其content
|
||||
lastMsg.Content = reactOutput
|
||||
lastMsg.Content = assistantOut
|
||||
} else {
|
||||
// 最后一条不是assistant消息,或者有tool_calls,添加最终输出作为新的assistant消息
|
||||
agentMessages = append(agentMessages, agent.ChatMessage{
|
||||
Role: "assistant",
|
||||
Content: reactOutput,
|
||||
Content: assistantOut,
|
||||
})
|
||||
}
|
||||
} else {
|
||||
// 如果没有消息,直接添加最终输出
|
||||
agentMessages = append(agentMessages, agent.ChatMessage{
|
||||
Role: "assistant",
|
||||
Content: reactOutput,
|
||||
Content: assistantOut,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if len(agentMessages) == 0 {
|
||||
return nil, fmt.Errorf("从ReAct数据解析的消息为空")
|
||||
return nil, fmt.Errorf("从代理轨迹解析的消息为空")
|
||||
}
|
||||
|
||||
// 修复可能存在的失配tool消息,避免OpenAI报错
|
||||
// 这可以防止出现"messages with role 'tool' must be a response to a preceeding message with 'tool_calls'"错误
|
||||
if h.agent != nil {
|
||||
if fixed := h.agent.RepairOrphanToolMessages(&agentMessages); fixed {
|
||||
h.logger.Info("修复了从ReAct数据恢复的历史消息中的失配tool消息",
|
||||
h.logger.Info("修复了从代理轨迹恢复的历史消息中的失配 tool 消息",
|
||||
zap.String("conversationId", conversationID),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
h.logger.Info("从ReAct数据恢复历史消息完成",
|
||||
h.logger.Info("从代理轨迹恢复历史消息完成",
|
||||
zap.String("conversationId", conversationID),
|
||||
zap.String("dataSource", dataSource),
|
||||
zap.Int("originalMessageCount", messageCount),
|
||||
zap.Int("finalMessageCount", len(agentMessages)),
|
||||
zap.Bool("hasReactOutput", reactOutput != ""),
|
||||
zap.Bool("hasAssistantOut", assistantOut != ""),
|
||||
)
|
||||
fmt.Println("agentMessages:", agentMessages) //debug
|
||||
return agentMessages, nil
|
||||
}
|
||||
|
||||
@@ -83,7 +83,7 @@ func (h *AttackChainHandler) GetAttackChain(c *gin.Context) {
|
||||
// 使用锁机制防止同一对话的并发生成
|
||||
lockInterface, _ := h.generatingLocks.LoadOrStore(conversationID, &sync.Mutex{})
|
||||
lock := lockInterface.(*sync.Mutex)
|
||||
|
||||
|
||||
// 尝试获取锁,如果正在生成则返回错误
|
||||
acquired := lock.TryLock()
|
||||
if !acquired {
|
||||
@@ -144,7 +144,7 @@ func (h *AttackChainHandler) RegenerateAttackChain(c *gin.Context) {
|
||||
// 使用锁机制防止并发生成
|
||||
lockInterface, _ := h.generatingLocks.LoadOrStore(conversationID, &sync.Mutex{})
|
||||
lock := lockInterface.(*sync.Mutex)
|
||||
|
||||
|
||||
acquired := lock.TryLock()
|
||||
if !acquired {
|
||||
h.logger.Info("攻击链正在生成中,请稍后再试", zap.String("conversationId", conversationID))
|
||||
@@ -170,4 +170,3 @@ func (h *AttackChainHandler) RegenerateAttackChain(c *gin.Context) {
|
||||
|
||||
c.JSON(http.StatusOK, chain)
|
||||
}
|
||||
|
||||
|
||||
+39
-14
@@ -17,6 +17,7 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/knowledge"
|
||||
"cyberstrike-ai/internal/mcp"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
"cyberstrike-ai/internal/openai"
|
||||
"cyberstrike-ai/internal/security"
|
||||
|
||||
@@ -90,6 +91,7 @@ type AttackChainUpdater interface {
|
||||
type AgentUpdater interface {
|
||||
UpdateConfig(cfg *config.OpenAIConfig)
|
||||
UpdateMaxIterations(maxIterations int)
|
||||
UpdateToolDescriptionMode(mode string)
|
||||
}
|
||||
|
||||
// NewConfigHandler 创建新的配置处理器
|
||||
@@ -232,13 +234,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
if configToolMap[mcpTool.Name] {
|
||||
continue
|
||||
}
|
||||
description := mcpTool.ShortDescription
|
||||
if description == "" {
|
||||
description = mcpTool.Description
|
||||
}
|
||||
if len(description) > 10000 {
|
||||
description = description[:10000] + "..."
|
||||
}
|
||||
description := h.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||
tools = append(tools, ToolConfigInfo{
|
||||
Name: mcpTool.Name,
|
||||
Description: description,
|
||||
@@ -275,6 +271,11 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
||||
SubAgentCount: subAgentCount,
|
||||
Orchestration: config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration),
|
||||
PlanExecuteLoopMaxIterations: h.config.MultiAgent.PlanExecuteLoopMaxIterations,
|
||||
ToolSearchAlwaysVisibleTools: append([]string(nil), h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools...),
|
||||
ToolSearchAlwaysVisibleEffectiveTools: mergeToolNameLists(
|
||||
h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools,
|
||||
builtin.GetAllBuiltinTools(),
|
||||
),
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, GetConfigResponse{
|
||||
@@ -430,13 +431,7 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
|
||||
continue
|
||||
}
|
||||
|
||||
description := mcpTool.ShortDescription
|
||||
if description == "" {
|
||||
description = mcpTool.Description
|
||||
}
|
||||
if len(description) > 10000 {
|
||||
description = description[:10000] + "..."
|
||||
}
|
||||
description := h.pickToolDescription(mcpTool.ShortDescription, mcpTool.Description)
|
||||
|
||||
toolInfo := ToolConfigInfo{
|
||||
Name: mcpTool.Name,
|
||||
@@ -689,11 +684,13 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
||||
if req.MultiAgent.PlanExecuteLoopMaxIterations != nil {
|
||||
h.config.MultiAgent.PlanExecuteLoopMaxIterations = *req.MultiAgent.PlanExecuteLoopMaxIterations
|
||||
}
|
||||
h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools = dedupeToolNameList(req.MultiAgent.ToolSearchAlwaysVisibleTools)
|
||||
h.logger.Info("更新多代理配置",
|
||||
zap.Bool("enabled", h.config.MultiAgent.Enabled),
|
||||
zap.Bool("robot_use_multi_agent", h.config.MultiAgent.RobotUseMultiAgent),
|
||||
zap.Bool("batch_use_multi_agent", h.config.MultiAgent.BatchUseMultiAgent),
|
||||
zap.Int("plan_execute_loop_max_iterations", h.config.MultiAgent.PlanExecuteLoopMaxIterations),
|
||||
zap.Int("tool_search_always_visible_tools", len(h.config.MultiAgent.EinoMiddleware.ToolSearchAlwaysVisibleTools)),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -1061,6 +1058,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
||||
if h.agent != nil {
|
||||
h.agent.UpdateConfig(&h.config.OpenAI)
|
||||
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
|
||||
h.agent.UpdateToolDescriptionMode(h.config.Security.ToolDescriptionMode)
|
||||
h.logger.Info("Agent配置已更新")
|
||||
}
|
||||
|
||||
@@ -1383,6 +1381,33 @@ func updateMultiAgentConfig(doc *yaml.Node, cfg config.MultiAgentConfig) {
|
||||
setBoolInMap(maNode, "robot_use_multi_agent", cfg.RobotUseMultiAgent)
|
||||
setBoolInMap(maNode, "batch_use_multi_agent", cfg.BatchUseMultiAgent)
|
||||
setIntInMap(maNode, "plan_execute_loop_max_iterations", cfg.PlanExecuteLoopMaxIterations)
|
||||
mwNode := ensureMap(maNode, "eino_middleware")
|
||||
setFlowStringSliceInMap(mwNode, "tool_search_always_visible_tools", dedupeToolNameList(cfg.EinoMiddleware.ToolSearchAlwaysVisibleTools))
|
||||
}
|
||||
|
||||
func dedupeToolNameList(in []string) []string {
|
||||
if len(in) == 0 {
|
||||
return []string{}
|
||||
}
|
||||
seen := make(map[string]struct{}, len(in))
|
||||
out := make([]string, 0, len(in))
|
||||
for _, name := range in {
|
||||
n := strings.TrimSpace(name)
|
||||
if n == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(n)
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
out = append(out, n)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func mergeToolNameLists(a, b []string) []string {
|
||||
return dedupeToolNameList(append(append([]string{}, a...), b...))
|
||||
}
|
||||
|
||||
func ensureMap(parent *yaml.Node, path ...string) *yaml.Node {
|
||||
|
||||
@@ -230,4 +230,3 @@ func (h *ConversationHandler) DeleteConversationTurn(c *gin.Context) {
|
||||
"message": "ok",
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -175,6 +175,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
)
|
||||
|
||||
if runErr != nil {
|
||||
h.persistEinoAgentTraceForResume(conversationID, result)
|
||||
cause := context.Cause(baseCtx)
|
||||
if errors.Is(cause, ErrTaskCancelled) {
|
||||
taskStatus = "cancelled"
|
||||
@@ -239,9 +240,9 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
||||
)
|
||||
}
|
||||
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存 ReAct 数据失败", zap.Error(err))
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -306,6 +307,7 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
progressCallback,
|
||||
)
|
||||
if runErr != nil {
|
||||
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": runErr.Error()})
|
||||
return
|
||||
}
|
||||
@@ -323,8 +325,8 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
||||
prep.AssistantMessageID,
|
||||
)
|
||||
}
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
_ = h.db.SaveReActData(prep.ConversationID, result.LastReActInput, result.LastReActOutput)
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
_ = h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
|
||||
@@ -247,7 +247,7 @@ func TestExternalMCPHandler_DeleteExternalMCP(t *testing.T) {
|
||||
|
||||
// 先添加一个配置
|
||||
configObj := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
}
|
||||
handler.manager.AddOrUpdateConfig("test-delete", configObj)
|
||||
@@ -276,11 +276,11 @@ func TestExternalMCPHandler_GetExternalMCPs(t *testing.T) {
|
||||
|
||||
// 添加多个配置
|
||||
handler.manager.AddOrUpdateConfig("test1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("test2", config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: false,
|
||||
})
|
||||
|
||||
@@ -319,15 +319,15 @@ func TestExternalMCPHandler_GetExternalMCPStats(t *testing.T) {
|
||||
|
||||
// 添加配置
|
||||
handler.manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
handler.manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
})
|
||||
|
||||
req := httptest.NewRequest("GET", "/api/external-mcp/stats", nil)
|
||||
@@ -360,7 +360,7 @@ func TestExternalMCPHandler_StartStopExternalMCP(t *testing.T) {
|
||||
|
||||
// 添加一个禁用的配置
|
||||
handler.manager.AddOrUpdateConfig("test-start-stop", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
})
|
||||
|
||||
// 测试启动(可能会失败,因为没有真实的服务器)
|
||||
@@ -416,7 +416,7 @@ func TestExternalMCPHandler_AddOrUpdateExternalMCP_EmptyName(t *testing.T) {
|
||||
router, _, _ := setupTestRouter()
|
||||
|
||||
configObj := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
}
|
||||
|
||||
@@ -459,14 +459,14 @@ func TestExternalMCPHandler_UpdateExistingConfig(t *testing.T) {
|
||||
|
||||
// 先添加配置
|
||||
config1 := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
}
|
||||
handler.manager.AddOrUpdateConfig("test-update", config1)
|
||||
|
||||
// 更新配置
|
||||
config2 := config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: true,
|
||||
}
|
||||
|
||||
|
||||
@@ -131,16 +131,16 @@ func (h *MarkdownAgentsHandler) GetMarkdownAgent(c *gin.Context) {
|
||||
}
|
||||
|
||||
type markdownAgentBody struct {
|
||||
Filename string `json:"filename"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Tools []string `json:"tools"`
|
||||
Instruction string `json:"instruction"`
|
||||
BindRole string `json:"bind_role"`
|
||||
MaxIterations int `json:"max_iterations"`
|
||||
Kind string `json:"kind"`
|
||||
Raw string `json:"raw"`
|
||||
Filename string `json:"filename"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Tools []string `json:"tools"`
|
||||
Instruction string `json:"instruction"`
|
||||
BindRole string `json:"bind_role"`
|
||||
MaxIterations int `json:"max_iterations"`
|
||||
Kind string `json:"kind"`
|
||||
Raw string `json:"raw"`
|
||||
}
|
||||
|
||||
// CreateMarkdownAgent POST /api/multi-agent/markdown-agents
|
||||
|
||||
@@ -42,11 +42,11 @@ func (h *MonitorHandler) SetExternalMCPManager(mgr *mcp.ExternalMCPManager) {
|
||||
type MonitorResponse struct {
|
||||
Executions []*mcp.ToolExecution `json:"executions"`
|
||||
Stats map[string]*mcp.ToolStats `json:"stats"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Total int `json:"total,omitempty"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PageSize int `json:"page_size,omitempty"`
|
||||
TotalPages int `json:"total_pages,omitempty"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Total int `json:"total,omitempty"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PageSize int `json:"page_size,omitempty"`
|
||||
TotalPages int `json:"total_pages,omitempty"`
|
||||
}
|
||||
|
||||
// Monitor 获取监控信息
|
||||
@@ -213,7 +213,6 @@ func (h *MonitorHandler) loadStats() map[string]*mcp.ToolStats {
|
||||
return stats
|
||||
}
|
||||
|
||||
|
||||
// GetExecution 获取特定执行记录
|
||||
func (h *MonitorHandler) GetExecution(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
@@ -416,5 +415,3 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
|
||||
h.logger.Info("尝试批量删除内存中的执行记录", zap.Int("count", len(request.IDs)))
|
||||
c.JSON(http.StatusOK, gin.H{"message": "执行记录已删除(如果存在)"})
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -185,6 +185,7 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
)
|
||||
|
||||
if runErr != nil {
|
||||
h.persistEinoAgentTraceForResume(conversationID, result)
|
||||
cause := context.Cause(baseCtx)
|
||||
if errors.Is(cause, ErrTaskCancelled) {
|
||||
taskStatus = "cancelled"
|
||||
@@ -249,9 +250,9 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
||||
)
|
||||
}
|
||||
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(conversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存 ReAct 数据失败", zap.Error(err))
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -318,6 +319,7 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
strings.TrimSpace(req.Orchestration),
|
||||
)
|
||||
if runErr != nil {
|
||||
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
||||
h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr))
|
||||
errMsg := "执行失败: " + runErr.Error()
|
||||
if prep.AssistantMessageID != "" {
|
||||
@@ -341,9 +343,9 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
)
|
||||
}
|
||||
|
||||
if result.LastReActInput != "" || result.LastReActOutput != "" {
|
||||
if err := h.db.SaveReActData(prep.ConversationID, result.LastReActInput, result.LastReActOutput); err != nil {
|
||||
h.logger.Warn("保存 ReAct 数据失败", zap.Error(err))
|
||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||
if err := h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存代理轨迹失败", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -355,6 +357,19 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
// persistEinoAgentTraceForResume 在 Eino 运行异常结束时写入代理轨迹(库列 last_react_*),供下一请求 loadHistoryFromAgentTrace 软续跑。
|
||||
func (h *AgentHandler) persistEinoAgentTraceForResume(conversationID string, result *multiagent.RunResult) {
|
||||
if h == nil || result == nil {
|
||||
return
|
||||
}
|
||||
if result.LastAgentTraceInput == "" && result.LastAgentTraceOutput == "" {
|
||||
return
|
||||
}
|
||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||
h.logger.Warn("保存 Eino 续跑上下文失败", zap.String("conversationId", conversationID), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
func multiAgentHTTPErrorStatus(err error) (int, string) {
|
||||
msg := err.Error()
|
||||
switch {
|
||||
|
||||
@@ -49,7 +49,7 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest) (*multiAgentPr
|
||||
}
|
||||
}
|
||||
|
||||
agentHistoryMessages, err := h.loadHistoryFromReActData(conversationID)
|
||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||
if err != nil {
|
||||
historyMessages, getErr := h.db.GetMessages(conversationID)
|
||||
if getErr != nil {
|
||||
|
||||
+26
-26
@@ -4445,7 +4445,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"messageId"},
|
||||
"properties": map[string]interface{}{
|
||||
"messageId": map[string]interface{}{
|
||||
@@ -4689,7 +4689,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"scheduleEnabled"},
|
||||
"properties": map[string]interface{}{
|
||||
"scheduleEnabled": map[string]interface{}{"type": "boolean", "description": "是否启用自动调度"},
|
||||
@@ -4761,7 +4761,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"query"},
|
||||
"properties": map[string]interface{}{
|
||||
"query": map[string]interface{}{"type": "string", "description": "FOFA查询语法", "example": "domain=\"example.com\""},
|
||||
@@ -4810,7 +4810,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"text"},
|
||||
"properties": map[string]interface{}{
|
||||
"text": map[string]interface{}{"type": "string", "description": "自然语言描述", "example": "查找使用WordPress的网站"},
|
||||
@@ -4853,7 +4853,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"api_key", "model"},
|
||||
"properties": map[string]interface{}{
|
||||
"provider": map[string]interface{}{"type": "string", "description": "LLM提供商(openai/claude)", "example": "openai"},
|
||||
@@ -4900,7 +4900,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"command"},
|
||||
"properties": map[string]interface{}{
|
||||
"command": map[string]interface{}{"type": "string", "description": "要执行的命令"},
|
||||
@@ -4943,7 +4943,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"command"},
|
||||
"properties": map[string]interface{}{
|
||||
"command": map[string]interface{}{"type": "string", "description": "要执行的命令"},
|
||||
@@ -5027,7 +5027,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"url"},
|
||||
"properties": map[string]interface{}{
|
||||
"url": map[string]interface{}{"type": "string", "description": "WebShell URL"},
|
||||
@@ -5231,7 +5231,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"url", "command"},
|
||||
"properties": map[string]interface{}{
|
||||
"url": map[string]interface{}{"type": "string", "description": "WebShell URL"},
|
||||
@@ -5277,7 +5277,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"url", "action", "path"},
|
||||
"properties": map[string]interface{}{
|
||||
"url": map[string]interface{}{"type": "string", "description": "WebShell URL"},
|
||||
@@ -5339,14 +5339,14 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"items": map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"relativePath": map[string]interface{}{"type": "string"},
|
||||
"absolutePath": map[string]interface{}{"type": "string"},
|
||||
"name": map[string]interface{}{"type": "string"},
|
||||
"size": map[string]interface{}{"type": "integer"},
|
||||
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
||||
"date": map[string]interface{}{"type": "string"},
|
||||
"relativePath": map[string]interface{}{"type": "string"},
|
||||
"absolutePath": map[string]interface{}{"type": "string"},
|
||||
"name": map[string]interface{}{"type": "string"},
|
||||
"size": map[string]interface{}{"type": "integer"},
|
||||
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
||||
"date": map[string]interface{}{"type": "string"},
|
||||
"conversationId": map[string]interface{}{"type": "string"},
|
||||
"subPath": map[string]interface{}{"type": "string"},
|
||||
"subPath": map[string]interface{}{"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -5369,7 +5369,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"multipart/form-data": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"file"},
|
||||
"properties": map[string]interface{}{
|
||||
"file": map[string]interface{}{"type": "string", "format": "binary", "description": "上传的文件"},
|
||||
@@ -5410,7 +5410,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"path"},
|
||||
"properties": map[string]interface{}{
|
||||
"path": map[string]interface{}{"type": "string", "description": "文件相对路径"},
|
||||
@@ -5485,7 +5485,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"path", "content"},
|
||||
"properties": map[string]interface{}{
|
||||
"path": map[string]interface{}{"type": "string", "description": "文件相对路径"},
|
||||
@@ -5512,7 +5512,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"name"},
|
||||
"properties": map[string]interface{}{
|
||||
"parent": map[string]interface{}{"type": "string", "description": "父目录相对路径"},
|
||||
@@ -5552,7 +5552,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"path", "newName"},
|
||||
"properties": map[string]interface{}{
|
||||
"path": map[string]interface{}{"type": "string", "description": "当前文件相对路径"},
|
||||
@@ -5646,7 +5646,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"platform", "text"},
|
||||
"properties": map[string]interface{}{
|
||||
"platform": map[string]interface{}{"type": "string", "description": "平台类型", "enum": []string{"dingtalk", "lark", "wecom"}},
|
||||
@@ -5712,7 +5712,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"name"},
|
||||
"properties": map[string]interface{}{
|
||||
"filename": map[string]interface{}{"type": "string", "description": "文件名(可选,自动生成)"},
|
||||
@@ -5932,7 +5932,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"path"},
|
||||
"properties": map[string]interface{}{
|
||||
"path": map[string]interface{}{"type": "string", "description": "文件相对路径"},
|
||||
@@ -5974,7 +5974,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
||||
"content": map[string]interface{}{
|
||||
"application/json": map[string]interface{}{
|
||||
"schema": map[string]interface{}{
|
||||
"type": "object",
|
||||
"type": "object",
|
||||
"required": []string{"ids"},
|
||||
"properties": map[string]interface{}{
|
||||
"ids": map[string]interface{}{
|
||||
|
||||
@@ -26,7 +26,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"创建分组": "createGroup", "列出分组": "listGroups", "获取分组": "getGroup", "更新分组": "updateGroup",
|
||||
"删除分组": "deleteGroup", "获取分组中的对话": "getGroupConversations", "添加对话到分组": "addConversationToGroup",
|
||||
"从分组移除对话": "removeConversationFromGroup",
|
||||
"列出漏洞": "listVulnerabilities", "创建漏洞": "createVulnerability", "获取漏洞统计": "getVulnerabilityStats",
|
||||
"列出漏洞": "listVulnerabilities", "创建漏洞": "createVulnerability", "获取漏洞统计": "getVulnerabilityStats",
|
||||
"获取漏洞": "getVulnerability", "更新漏洞": "updateVulnerability", "删除漏洞": "deleteVulnerability",
|
||||
"列出角色": "listRoles", "创建角色": "createRole", "获取角色": "getRole", "更新角色": "updateRole", "删除角色": "deleteRole",
|
||||
"获取可用Skills列表": "getAvailableSkills", "列出Skills": "listSkills", "创建Skill": "createSkill",
|
||||
@@ -52,9 +52,9 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"重跑批量任务队列": "rerunBatchQueue", "修改队列元数据": "updateBatchQueueMetadata",
|
||||
"修改队列调度配置": "updateBatchQueueSchedule", "开关Cron自动调度": "setBatchQueueScheduleEnabled",
|
||||
"获取所有分组映射": "getAllGroupMappings",
|
||||
"FOFA搜索": "fofaSearch", "自然语言解析为FOFA语法": "fofaParse",
|
||||
"FOFA搜索": "fofaSearch", "自然语言解析为FOFA语法": "fofaParse",
|
||||
"测试OpenAI API连接": "testOpenAI",
|
||||
"执行终端命令": "terminalRun", "流式执行终端命令": "terminalRunStream", "WebSocket终端": "terminalWS",
|
||||
"执行终端命令": "terminalRun", "流式执行终端命令": "terminalRunStream", "WebSocket终端": "terminalWS",
|
||||
"列出WebShell连接": "listWebshellConnections", "创建WebShell连接": "createWebshellConnection",
|
||||
"更新WebShell连接": "updateWebshellConnection", "删除WebShell连接": "deleteWebshellConnection",
|
||||
"获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState",
|
||||
@@ -69,7 +69,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
||||
"获取Markdown代理详情": "getMarkdownAgent", "更新Markdown代理": "updateMarkdownAgent", "删除Markdown代理": "deleteMarkdownAgent",
|
||||
"列出技能包文件": "listSkillPackageFiles", "获取技能包文件内容": "getSkillPackageFile", "写入技能包文件": "putSkillPackageFile",
|
||||
"批量获取工具名称": "batchGetToolNames",
|
||||
"获取知识库统计": "getKnowledgeStats",
|
||||
"获取知识库统计": "getKnowledgeStats",
|
||||
}
|
||||
|
||||
var apiDocI18nResponseDescToKey = map[string]string{
|
||||
@@ -78,7 +78,7 @@ var apiDocI18nResponseDescToKey = map[string]string{
|
||||
"对话不存在或结果不存在": "conversationOrResultNotFound", "请求参数错误(如task为空)": "badRequestTaskEmpty",
|
||||
"请求参数错误或分组名称已存在": "badRequestGroupNameExists", "分组不存在": "groupNotFound",
|
||||
"请求参数错误(如配置格式不正确、缺少必需字段等)": "badRequestConfig",
|
||||
"请求参数错误(如query为空)": "badRequestQueryEmpty", "方法不允许(仅支持POST请求)": "methodNotAllowed",
|
||||
"请求参数错误(如query为空)": "badRequestQueryEmpty", "方法不允许(仅支持POST请求)": "methodNotAllowed",
|
||||
"登录成功": "loginSuccess", "密码错误": "invalidPassword", "登出成功": "logoutSuccess",
|
||||
"密码修改成功": "passwordChanged", "Token有效": "tokenValid", "Token无效或已过期": "tokenInvalid",
|
||||
"对话创建成功": "conversationCreated", "服务器内部错误": "internalError", "更新成功": "updateSuccess",
|
||||
@@ -89,7 +89,7 @@ var apiDocI18nResponseDescToKey = map[string]string{
|
||||
"消息发送成功,返回AI回复": "messageSent", "流式响应(Server-Sent Events)": "streamResponse",
|
||||
// 新增缺失端点响应
|
||||
"参数错误或删除失败": "badRequestOrDeleteFailed",
|
||||
"参数错误": "paramError", "仅已完成或已取消的队列可以重跑": "onlyCompletedOrCancelledCanRerun",
|
||||
"参数错误": "paramError", "仅已完成或已取消的队列可以重跑": "onlyCompletedOrCancelledCanRerun",
|
||||
"参数错误或队列正在运行中": "badRequestOrQueueRunning", "设置成功": "setSuccess",
|
||||
"搜索成功": "searchSuccess", "解析成功": "parseSuccess", "测试结果": "testResult",
|
||||
"执行完成": "executionDone", "SSE事件流": "sseEventStream", "WebSocket连接已建立": "wsEstablished",
|
||||
|
||||
+14
-14
@@ -28,20 +28,20 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
robotCmdHelp = "帮助"
|
||||
robotCmdList = "列表"
|
||||
robotCmdListAlt = "对话列表"
|
||||
robotCmdSwitch = "切换"
|
||||
robotCmdContinue = "继续"
|
||||
robotCmdNew = "新对话"
|
||||
robotCmdClear = "清空"
|
||||
robotCmdCurrent = "当前"
|
||||
robotCmdStop = "停止"
|
||||
robotCmdRoles = "角色"
|
||||
robotCmdRolesList = "角色列表"
|
||||
robotCmdSwitchRole = "切换角色"
|
||||
robotCmdDelete = "删除"
|
||||
robotCmdVersion = "版本"
|
||||
robotCmdHelp = "帮助"
|
||||
robotCmdList = "列表"
|
||||
robotCmdListAlt = "对话列表"
|
||||
robotCmdSwitch = "切换"
|
||||
robotCmdContinue = "继续"
|
||||
robotCmdNew = "新对话"
|
||||
robotCmdClear = "清空"
|
||||
robotCmdCurrent = "当前"
|
||||
robotCmdStop = "停止"
|
||||
robotCmdRoles = "角色"
|
||||
robotCmdRolesList = "角色列表"
|
||||
robotCmdSwitchRole = "切换角色"
|
||||
robotCmdDelete = "删除"
|
||||
robotCmdVersion = "版本"
|
||||
)
|
||||
|
||||
// RobotHandler 企业微信/钉钉/飞书等机器人回调处理
|
||||
|
||||
+13
-13
@@ -65,19 +65,19 @@ func (h *SkillsHandler) GetSkills(c *gin.Context) {
|
||||
allSkillsInfo := make([]map[string]interface{}, 0, len(allSummaries))
|
||||
for _, s := range allSummaries {
|
||||
skillInfo := map[string]interface{}{
|
||||
"id": s.ID,
|
||||
"name": s.Name,
|
||||
"dir_name": s.DirName,
|
||||
"description": s.Description,
|
||||
"version": s.Version,
|
||||
"path": s.Path,
|
||||
"tags": s.Tags,
|
||||
"triggers": s.Triggers,
|
||||
"script_count": s.ScriptCount,
|
||||
"file_count": s.FileCount,
|
||||
"progressive": s.Progressive,
|
||||
"file_size": s.FileSize,
|
||||
"mod_time": s.ModTime,
|
||||
"id": s.ID,
|
||||
"name": s.Name,
|
||||
"dir_name": s.DirName,
|
||||
"description": s.Description,
|
||||
"version": s.Version,
|
||||
"path": s.Path,
|
||||
"tags": s.Tags,
|
||||
"triggers": s.Triggers,
|
||||
"script_count": s.ScriptCount,
|
||||
"file_count": s.FileCount,
|
||||
"progressive": s.Progressive,
|
||||
"file_size": s.FileSize,
|
||||
"mod_time": s.ModTime,
|
||||
}
|
||||
allSkillsInfo = append(allSkillsInfo, skillInfo)
|
||||
}
|
||||
|
||||
@@ -109,4 +109,3 @@ func (h *TerminalHandler) RunCommandWS(c *gin.Context) {
|
||||
|
||||
<-doneChan
|
||||
}
|
||||
|
||||
|
||||
@@ -28,7 +28,7 @@ func NewVulnerabilityHandler(db *database.DB, logger *zap.Logger) *Vulnerability
|
||||
|
||||
// CreateVulnerabilityRequest 创建漏洞请求
|
||||
type CreateVulnerabilityRequest struct {
|
||||
ConversationID string `json:"conversation_id" binding:"required"`
|
||||
ConversationID string `json:"conversation_id" binding:"required"`
|
||||
ConversationTag string `json:"conversation_tag"`
|
||||
TaskTag string `json:"task_tag"`
|
||||
Title string `json:"title" binding:"required"`
|
||||
@@ -51,18 +51,18 @@ func (h *VulnerabilityHandler) CreateVulnerability(c *gin.Context) {
|
||||
}
|
||||
|
||||
vuln := &database.Vulnerability{
|
||||
ConversationID: req.ConversationID,
|
||||
ConversationID: req.ConversationID,
|
||||
ConversationTag: req.ConversationTag,
|
||||
TaskTag: req.TaskTag,
|
||||
Title: req.Title,
|
||||
Description: req.Description,
|
||||
Severity: req.Severity,
|
||||
Status: req.Status,
|
||||
Type: req.Type,
|
||||
Target: req.Target,
|
||||
Proof: req.Proof,
|
||||
Impact: req.Impact,
|
||||
Recommendation: req.Recommendation,
|
||||
Title: req.Title,
|
||||
Description: req.Description,
|
||||
Severity: req.Severity,
|
||||
Status: req.Status,
|
||||
Type: req.Type,
|
||||
Target: req.Target,
|
||||
Proof: req.Proof,
|
||||
Impact: req.Impact,
|
||||
Recommendation: req.Recommendation,
|
||||
}
|
||||
|
||||
created, err := h.db.CreateVulnerability(vuln)
|
||||
@@ -172,15 +172,15 @@ func (h *VulnerabilityHandler) ListVulnerabilities(c *gin.Context) {
|
||||
type UpdateVulnerabilityRequest struct {
|
||||
ConversationTag string `json:"conversation_tag"`
|
||||
TaskTag string `json:"task_tag"`
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
Target string `json:"target"`
|
||||
Proof string `json:"proof"`
|
||||
Impact string `json:"impact"`
|
||||
Recommendation string `json:"recommendation"`
|
||||
Title string `json:"title"`
|
||||
Description string `json:"description"`
|
||||
Severity string `json:"severity"`
|
||||
Status string `json:"status"`
|
||||
Type string `json:"type"`
|
||||
Target string `json:"target"`
|
||||
Proof string `json:"proof"`
|
||||
Impact string `json:"impact"`
|
||||
Recommendation string `json:"recommendation"`
|
||||
}
|
||||
|
||||
// UpdateVulnerability 更新漏洞
|
||||
@@ -460,4 +460,3 @@ func sanitizeExportName(raw string) string {
|
||||
replacer := strings.NewReplacer("/", "-", "\\", "-", ":", "-", "*", "-", "?", "-", "\"", "-", "<", "-", ">", "-", "|", "-")
|
||||
return replacer.Replace(name)
|
||||
}
|
||||
|
||||
|
||||
@@ -16,9 +16,9 @@ const (
|
||||
|
||||
// DSL keys for [VectorEinoRetriever.Retrieve] via [retriever.WithDSLInfo].
|
||||
const (
|
||||
DSLRiskType = "risk_type"
|
||||
DSLSimilarityThreshold = "similarity_threshold"
|
||||
DSLSubIndexFilter = "sub_index_filter"
|
||||
DSLRiskType = "risk_type"
|
||||
DSLSimilarityThreshold = "similarity_threshold"
|
||||
DSLSubIndexFilter = "sub_index_filter"
|
||||
)
|
||||
|
||||
// FormatEmbeddingInput matches the historical indexing format so existing embeddings
|
||||
|
||||
@@ -8,8 +8,8 @@ import (
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/components/document"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
)
|
||||
|
||||
|
||||
@@ -11,9 +11,9 @@ import (
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
fileloader "github.com/cloudwego/eino-ext/components/document/loader/file"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/components/document"
|
||||
"github.com/cloudwego/eino/components/indexer"
|
||||
"github.com/cloudwego/eino/compose"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -35,14 +35,14 @@ type Indexer struct {
|
||||
lastErrorTime time.Time
|
||||
errorCount int
|
||||
|
||||
rebuildMu sync.RWMutex
|
||||
isRebuilding bool
|
||||
rebuildTotalItems int
|
||||
rebuildCurrent int
|
||||
rebuildFailed int
|
||||
rebuildStartTime time.Time
|
||||
rebuildLastItemID string
|
||||
rebuildLastChunks int
|
||||
rebuildMu sync.RWMutex
|
||||
isRebuilding bool
|
||||
rebuildTotalItems int
|
||||
rebuildCurrent int
|
||||
rebuildFailed int
|
||||
rebuildStartTime time.Time
|
||||
rebuildLastItemID string
|
||||
rebuildLastChunks int
|
||||
}
|
||||
|
||||
// NewIndexer 创建索引器并编译 Eino 索引链;kcfg 为完整知识库配置(含 indexing 与路径相关行为)。
|
||||
|
||||
@@ -108,9 +108,9 @@ func (r *RetrievalLog) MarshalJSON() ([]byte, error) {
|
||||
|
||||
// CategoryWithItems 分类及其下的知识项(用于按分类分页)
|
||||
type CategoryWithItems struct {
|
||||
Category string `json:"category"` // 分类名称
|
||||
ItemCount int `json:"itemCount"` // 该分类下的知识项总数
|
||||
Items []*KnowledgeItemSummary `json:"items"` // 该分类下的知识项列表
|
||||
Category string `json:"category"` // 分类名称
|
||||
ItemCount int `json:"itemCount"` // 该分类下的知识项总数
|
||||
Items []*KnowledgeItemSummary `json:"items"` // 该分类下的知识项列表
|
||||
}
|
||||
|
||||
// SearchRequest 搜索请求
|
||||
|
||||
+10
-10
@@ -55,14 +55,14 @@ func New(level, output string) *Logger {
|
||||
}
|
||||
|
||||
func (l *Logger) Fatal(msg string, fields ...interface{}) {
|
||||
zapFields := make([]zap.Field, 0, len(fields))
|
||||
for _, f := range fields {
|
||||
switch v := f.(type) {
|
||||
case error:
|
||||
zapFields = append(zapFields, zap.Error(v))
|
||||
default:
|
||||
zapFields = append(zapFields, zap.Any("field", v))
|
||||
}
|
||||
}
|
||||
l.Logger.Fatal(msg, zapFields...)
|
||||
zapFields := make([]zap.Field, 0, len(fields))
|
||||
for _, f := range fields {
|
||||
switch v := f.(type) {
|
||||
case error:
|
||||
zapFields = append(zapFields, zap.Error(v))
|
||||
default:
|
||||
zapFields = append(zapFields, zap.Any("field", v))
|
||||
}
|
||||
}
|
||||
l.Logger.Fatal(msg, zapFields...)
|
||||
}
|
||||
|
||||
@@ -62,7 +62,7 @@ func TestExternalMCPManager_RemoveConfig(t *testing.T) {
|
||||
manager := NewExternalMCPManager(logger)
|
||||
|
||||
cfg := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: false,
|
||||
}
|
||||
|
||||
@@ -86,17 +86,17 @@ func TestExternalMCPManager_GetStats(t *testing.T) {
|
||||
|
||||
// 添加多个配置
|
||||
manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
|
||||
manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: true,
|
||||
})
|
||||
|
||||
manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: false,
|
||||
})
|
||||
|
||||
@@ -122,11 +122,11 @@ func TestExternalMCPManager_LoadConfigs(t *testing.T) {
|
||||
externalMCPConfig := config.ExternalMCPConfig{
|
||||
Servers: map[string]config.ExternalMCPServerConfig{
|
||||
"loaded1": {
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: true,
|
||||
},
|
||||
"loaded2": {
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
URL: "http://127.0.0.1:8081/mcp",
|
||||
ExternalMCPEnable: false,
|
||||
},
|
||||
},
|
||||
@@ -153,9 +153,9 @@ func TestLazySDKClient_InitializeFails(t *testing.T) {
|
||||
logger := zap.NewNop()
|
||||
// 使用不存在的 HTTP 地址,Initialize 应失败
|
||||
cfg := config.ExternalMCPServerConfig{
|
||||
Type: "http",
|
||||
URL: "http://127.0.0.1:19999/nonexistent",
|
||||
Timeout: 2,
|
||||
Type: "http",
|
||||
URL: "http://127.0.0.1:19999/nonexistent",
|
||||
Timeout: 2,
|
||||
}
|
||||
c := newLazySDKClient(cfg, logger)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
@@ -176,7 +176,7 @@ func TestExternalMCPManager_StartStopClient(t *testing.T) {
|
||||
|
||||
// 添加一个禁用的配置
|
||||
cfg := config.ExternalMCPServerConfig{
|
||||
Command: "python3",
|
||||
Command: "python3",
|
||||
ExternalMCPEnable: false,
|
||||
}
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -109,11 +110,11 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
}
|
||||
}()
|
||||
|
||||
var lastRunMsgs []adk.Message
|
||||
var lastAssistant string
|
||||
var lastPlanExecuteExecutor string
|
||||
msgs := append([]adk.Message(nil), baseMsgs...)
|
||||
runAccumulatedMsgs := append([]adk.Message(nil), msgs...)
|
||||
baseAccumulatedCount := len(runAccumulatedMsgs)
|
||||
|
||||
emptyHint := strings.TrimSpace(args.EmptyResponseMessage)
|
||||
if emptyHint == "" {
|
||||
@@ -132,125 +133,150 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
pendingByID := make(map[string]toolCallPendingInfo)
|
||||
pendingQueueByAgent := make(map[string][]string)
|
||||
markPending := func(tc toolCallPendingInfo) {
|
||||
if tc.ToolCallID == "" {
|
||||
return
|
||||
}
|
||||
pendingByID[tc.ToolCallID] = tc
|
||||
pendingQueueByAgent[tc.EinoAgent] = append(pendingQueueByAgent[tc.EinoAgent], tc.ToolCallID)
|
||||
if tc.ToolCallID == "" {
|
||||
return
|
||||
}
|
||||
popNextPendingForAgent := func(agentName string) (toolCallPendingInfo, bool) {
|
||||
q := pendingQueueByAgent[agentName]
|
||||
for len(q) > 0 {
|
||||
id := q[0]
|
||||
q = q[1:]
|
||||
pendingQueueByAgent[agentName] = q
|
||||
if tc, ok := pendingByID[id]; ok {
|
||||
delete(pendingByID, id)
|
||||
return tc, true
|
||||
}
|
||||
pendingByID[tc.ToolCallID] = tc
|
||||
pendingQueueByAgent[tc.EinoAgent] = append(pendingQueueByAgent[tc.EinoAgent], tc.ToolCallID)
|
||||
}
|
||||
popNextPendingForAgent := func(agentName string) (toolCallPendingInfo, bool) {
|
||||
q := pendingQueueByAgent[agentName]
|
||||
for len(q) > 0 {
|
||||
id := q[0]
|
||||
q = q[1:]
|
||||
pendingQueueByAgent[agentName] = q
|
||||
if tc, ok := pendingByID[id]; ok {
|
||||
delete(pendingByID, id)
|
||||
return tc, true
|
||||
}
|
||||
return toolCallPendingInfo{}, false
|
||||
}
|
||||
removePendingByID := func(toolCallID string) {
|
||||
if toolCallID == "" {
|
||||
return
|
||||
}
|
||||
delete(pendingByID, toolCallID)
|
||||
return toolCallPendingInfo{}, false
|
||||
}
|
||||
removePendingByID := func(toolCallID string) {
|
||||
if toolCallID == "" {
|
||||
return
|
||||
}
|
||||
flushAllPendingAsFailed := func(err error) {
|
||||
if progress == nil {
|
||||
pendingByID = make(map[string]toolCallPendingInfo)
|
||||
pendingQueueByAgent = make(map[string][]string)
|
||||
return
|
||||
}
|
||||
msg := ""
|
||||
if err != nil {
|
||||
msg = err.Error()
|
||||
}
|
||||
for _, tc := range pendingByID {
|
||||
toolName := tc.ToolName
|
||||
if strings.TrimSpace(toolName) == "" {
|
||||
toolName = "unknown"
|
||||
}
|
||||
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), map[string]interface{}{
|
||||
"toolName": toolName,
|
||||
"success": false,
|
||||
"isError": true,
|
||||
"result": msg,
|
||||
"resultPreview": msg,
|
||||
"toolCallId": tc.ToolCallID,
|
||||
"conversationId": conversationID,
|
||||
"einoAgent": tc.EinoAgent,
|
||||
"einoRole": tc.EinoRole,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
delete(pendingByID, toolCallID)
|
||||
}
|
||||
flushAllPendingAsFailed := func(err error) {
|
||||
if progress == nil {
|
||||
pendingByID = make(map[string]toolCallPendingInfo)
|
||||
pendingQueueByAgent = make(map[string][]string)
|
||||
return
|
||||
}
|
||||
msg := ""
|
||||
if err != nil {
|
||||
msg = err.Error()
|
||||
}
|
||||
for _, tc := range pendingByID {
|
||||
toolName := tc.ToolName
|
||||
if strings.TrimSpace(toolName) == "" {
|
||||
toolName = "unknown"
|
||||
}
|
||||
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), map[string]interface{}{
|
||||
"toolName": toolName,
|
||||
"success": false,
|
||||
"isError": true,
|
||||
"result": msg,
|
||||
"resultPreview": msg,
|
||||
"toolCallId": tc.ToolCallID,
|
||||
"conversationId": conversationID,
|
||||
"einoAgent": tc.EinoAgent,
|
||||
"einoRole": tc.EinoRole,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
pendingByID = make(map[string]toolCallPendingInfo)
|
||||
pendingQueueByAgent = make(map[string][]string)
|
||||
}
|
||||
|
||||
runnerCfg := adk.RunnerConfig{
|
||||
Agent: da,
|
||||
EnableStreaming: true,
|
||||
}
|
||||
if cp := strings.TrimSpace(args.CheckpointDir); cp != "" {
|
||||
cpDir := filepath.Join(cp, sanitizeEinoPathSegment(conversationID))
|
||||
st, stErr := newFileCheckPointStore(cpDir)
|
||||
if stErr != nil {
|
||||
if logger != nil {
|
||||
logger.Warn("eino checkpoint store disabled", zap.String("dir", cpDir), zap.Error(stErr))
|
||||
}
|
||||
} else {
|
||||
runnerCfg.CheckPointStore = st
|
||||
if logger != nil {
|
||||
logger.Info("eino runner: checkpoint store enabled", zap.String("dir", cpDir))
|
||||
}
|
||||
runnerCfg := adk.RunnerConfig{
|
||||
Agent: da,
|
||||
EnableStreaming: true,
|
||||
}
|
||||
var cpStore *fileCheckPointStore
|
||||
var checkPointID string
|
||||
if cp := strings.TrimSpace(args.CheckpointDir); cp != "" {
|
||||
cpDir := filepath.Join(cp, sanitizeEinoPathSegment(conversationID))
|
||||
st, stErr := newFileCheckPointStore(cpDir)
|
||||
if stErr != nil {
|
||||
if logger != nil {
|
||||
logger.Warn("eino checkpoint store disabled", zap.String("dir", cpDir), zap.Error(stErr))
|
||||
}
|
||||
} else {
|
||||
cpStore = st
|
||||
checkPointID = buildEinoCheckpointID(orchMode)
|
||||
runnerCfg.CheckPointStore = st
|
||||
if logger != nil {
|
||||
logger.Info("eino runner: checkpoint store enabled",
|
||||
zap.String("dir", cpDir),
|
||||
zap.String("checkPointID", checkPointID))
|
||||
}
|
||||
}
|
||||
}
|
||||
runner := adk.NewRunner(ctx, runnerCfg)
|
||||
iter := runner.Run(ctx, msgs)
|
||||
handleRunErr := func(runErr error) error {
|
||||
if runErr == nil {
|
||||
return nil
|
||||
var iter *adk.AsyncIterator[*adk.AgentEvent]
|
||||
if cpStore != nil && checkPointID != "" {
|
||||
if _, existed, getErr := cpStore.Get(ctx, checkPointID); getErr != nil {
|
||||
if logger != nil {
|
||||
logger.Warn("eino checkpoint preflight get failed", zap.String("checkPointID", checkPointID), zap.Error(getErr))
|
||||
}
|
||||
if errors.Is(runErr, context.DeadlineExceeded) {
|
||||
flushAllPendingAsFailed(runErr)
|
||||
if progress != nil {
|
||||
progress("error", runErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"errorKind": "timeout",
|
||||
})
|
||||
} else if existed {
|
||||
if progress != nil {
|
||||
progress("progress", "检测到断点,正在从中断节点恢复执行...", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"orchestration": orchMode,
|
||||
"checkPointID": checkPointID,
|
||||
})
|
||||
}
|
||||
if logger != nil {
|
||||
logger.Info("eino runner: resume from checkpoint", zap.String("checkPointID", checkPointID))
|
||||
}
|
||||
resumeIter, resumeErr := runner.Resume(ctx, checkPointID)
|
||||
if resumeErr == nil {
|
||||
iter = resumeIter
|
||||
} else {
|
||||
if logger != nil {
|
||||
logger.Warn("eino runner: resume failed, fallback to fresh run",
|
||||
zap.String("checkPointID", checkPointID),
|
||||
zap.Error(resumeErr))
|
||||
}
|
||||
return runErr
|
||||
}
|
||||
// context.Canceled 是唯一应当直接终止编排的错误(用户关闭页面、主动停止等)。
|
||||
if errors.Is(runErr, context.Canceled) {
|
||||
flushAllPendingAsFailed(runErr)
|
||||
if progress != nil {
|
||||
progress("error", runErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
return runErr
|
||||
}
|
||||
if isEinoIterationLimitError(runErr) {
|
||||
flushAllPendingAsFailed(runErr)
|
||||
if progress != nil {
|
||||
progress("iteration_limit_reached", runErr.Error(), map[string]interface{}{
|
||||
progress("progress", "断点恢复失败,已回退为全新执行。", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
progress("error", runErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"errorKind": "iteration_limit",
|
||||
"checkPointID": checkPointID,
|
||||
})
|
||||
}
|
||||
return runErr
|
||||
}
|
||||
}
|
||||
}
|
||||
if iter == nil {
|
||||
if checkPointID != "" {
|
||||
iter = runner.Run(ctx, msgs, adk.WithCheckPointID(checkPointID))
|
||||
} else {
|
||||
iter = runner.Run(ctx, msgs)
|
||||
}
|
||||
}
|
||||
handleRunErr := func(runErr error) error {
|
||||
if runErr == nil {
|
||||
return nil
|
||||
}
|
||||
if errors.Is(runErr, context.DeadlineExceeded) {
|
||||
flushAllPendingAsFailed(runErr)
|
||||
if progress != nil {
|
||||
progress("error", runErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"errorKind": "timeout",
|
||||
})
|
||||
}
|
||||
return runErr
|
||||
}
|
||||
// context.Canceled 是唯一应当直接终止编排的错误(用户关闭页面、主动停止等)。
|
||||
if errors.Is(runErr, context.Canceled) {
|
||||
flushAllPendingAsFailed(runErr)
|
||||
if progress != nil {
|
||||
progress("error", runErr.Error(), map[string]interface{}{
|
||||
@@ -260,249 +286,190 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
}
|
||||
return runErr
|
||||
}
|
||||
|
||||
for {
|
||||
// 检测 context 取消(用户关闭浏览器、请求超时等),flush pending 工具状态避免 UI 卡在 "执行中"。
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
flushAllPendingAsFailed(ctx.Err())
|
||||
if progress != nil {
|
||||
progress("error", "Request cancelled / 请求已取消", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
ev, ok := iter.Next()
|
||||
if !ok {
|
||||
if len(pendingByID) > 0 {
|
||||
orphanCount := len(pendingByID)
|
||||
flushAllPendingAsFailed(errors.New("pending tool call missing result before run completion"))
|
||||
if progress != nil {
|
||||
progress("eino_pending_orphaned", "pending tool calls were force-closed at run end", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"orchestration": orchMode,
|
||||
"pendingCount": orphanCount,
|
||||
})
|
||||
}
|
||||
}
|
||||
lastRunMsgs = runAccumulatedMsgs
|
||||
break
|
||||
}
|
||||
if ev == nil {
|
||||
continue
|
||||
}
|
||||
if ev.Err != nil {
|
||||
if retErr := handleRunErr(ev.Err); retErr != nil {
|
||||
return nil, retErr
|
||||
}
|
||||
}
|
||||
if ev.AgentName != "" && progress != nil {
|
||||
iterEinoAgent := orchestratorName
|
||||
if orchMode == "plan_execute" {
|
||||
if a := strings.TrimSpace(ev.AgentName); a != "" {
|
||||
iterEinoAgent = a
|
||||
}
|
||||
}
|
||||
if streamsMainAssistant(ev.AgentName) {
|
||||
if einoMainRound == 0 {
|
||||
einoMainRound = 1
|
||||
progress("iteration", "", map[string]interface{}{
|
||||
"iteration": 1,
|
||||
"einoScope": "main",
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": iterEinoAgent,
|
||||
"orchestration": orchMode,
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
} else if einoLastAgent != "" && !streamsMainAssistant(einoLastAgent) {
|
||||
einoMainRound++
|
||||
progress("iteration", "", map[string]interface{}{
|
||||
"iteration": einoMainRound,
|
||||
"einoScope": "main",
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": iterEinoAgent,
|
||||
"orchestration": orchMode,
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
}
|
||||
einoLastAgent = ev.AgentName
|
||||
progress("progress", fmt.Sprintf("[Eino] %s", ev.AgentName), map[string]interface{}{
|
||||
if isEinoIterationLimitError(runErr) {
|
||||
flushAllPendingAsFailed(runErr)
|
||||
if progress != nil {
|
||||
progress("iteration_limit_reached", runErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
"source": "eino",
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
progress("error", runErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"errorKind": "iteration_limit",
|
||||
})
|
||||
}
|
||||
if ev.Output == nil || ev.Output.MessageOutput == nil {
|
||||
continue
|
||||
}
|
||||
mv := ev.Output.MessageOutput
|
||||
return runErr
|
||||
}
|
||||
flushAllPendingAsFailed(runErr)
|
||||
if progress != nil {
|
||||
progress("error", runErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
return runErr
|
||||
}
|
||||
|
||||
if mv.IsStreaming && mv.MessageStream != nil {
|
||||
streamHeaderSent := false
|
||||
var reasoningStreamID string
|
||||
var toolStreamFragments []schema.ToolCall
|
||||
var subAssistantBuf strings.Builder
|
||||
var subReplyStreamID string
|
||||
var mainAssistantBuf strings.Builder
|
||||
var streamRecvErr error
|
||||
for {
|
||||
chunk, rerr := mv.MessageStream.Recv()
|
||||
if rerr != nil {
|
||||
if errors.Is(rerr, io.EOF) {
|
||||
break
|
||||
}
|
||||
if logger != nil {
|
||||
logger.Warn("eino stream recv error, flushing incomplete stream",
|
||||
zap.Error(rerr),
|
||||
zap.String("agent", ev.AgentName),
|
||||
zap.Int("toolFragments", len(toolStreamFragments)))
|
||||
}
|
||||
streamRecvErr = rerr
|
||||
break
|
||||
}
|
||||
if chunk == nil {
|
||||
continue
|
||||
}
|
||||
if progress != nil && strings.TrimSpace(chunk.ReasoningContent) != "" {
|
||||
if reasoningStreamID == "" {
|
||||
reasoningStreamID = fmt.Sprintf("eino-reasoning-%s-%d", conversationID, atomic.AddInt64(&reasoningStreamSeq, 1))
|
||||
progress("thinking_stream_start", " ", map[string]interface{}{
|
||||
"streamId": reasoningStreamID,
|
||||
"source": "eino",
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}
|
||||
progress("thinking_stream_delta", chunk.ReasoningContent, map[string]interface{}{
|
||||
"streamId": reasoningStreamID,
|
||||
})
|
||||
}
|
||||
if chunk.Content != "" {
|
||||
if progress != nil && streamsMainAssistant(ev.AgentName) {
|
||||
if !streamHeaderSent {
|
||||
progress("response_start", "", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"messageGeneratedBy": "eino:" + ev.AgentName,
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
streamHeaderSent = true
|
||||
}
|
||||
progress("response_delta", chunk.Content, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
mainAssistantBuf.WriteString(chunk.Content)
|
||||
} else if !streamsMainAssistant(ev.AgentName) {
|
||||
if progress != nil {
|
||||
if subReplyStreamID == "" {
|
||||
subReplyStreamID = fmt.Sprintf("eino-sub-reply-%s-%d", conversationID, atomic.AddInt64(&einoSubReplyStreamSeq, 1))
|
||||
progress("eino_agent_reply_stream_start", "", map[string]interface{}{
|
||||
"streamId": subReplyStreamID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": "sub",
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
progress("eino_agent_reply_stream_delta", chunk.Content, map[string]interface{}{
|
||||
"streamId": subReplyStreamID,
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
}
|
||||
subAssistantBuf.WriteString(chunk.Content)
|
||||
}
|
||||
}
|
||||
if len(chunk.ToolCalls) > 0 {
|
||||
toolStreamFragments = append(toolStreamFragments, chunk.ToolCalls...)
|
||||
}
|
||||
}
|
||||
if streamsMainAssistant(ev.AgentName) {
|
||||
if s := strings.TrimSpace(mainAssistantBuf.String()); s != "" {
|
||||
lastAssistant = s
|
||||
runAccumulatedMsgs = append(runAccumulatedMsgs, schema.AssistantMessage(s, nil))
|
||||
if orchMode == "plan_execute" && strings.EqualFold(strings.TrimSpace(ev.AgentName), "executor") {
|
||||
lastPlanExecuteExecutor = UnwrapPlanExecuteUserText(s)
|
||||
}
|
||||
}
|
||||
}
|
||||
if subAssistantBuf.Len() > 0 && progress != nil {
|
||||
if s := strings.TrimSpace(subAssistantBuf.String()); s != "" {
|
||||
if subReplyStreamID != "" {
|
||||
progress("eino_agent_reply_stream_end", s, map[string]interface{}{
|
||||
"streamId": subReplyStreamID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": "sub",
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
} else {
|
||||
progress("eino_agent_reply", s, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": "sub",
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
var lastToolChunk *schema.Message
|
||||
if merged := mergeStreamingToolCallFragments(toolStreamFragments); len(merged) > 0 {
|
||||
lastToolChunk = &schema.Message{ToolCalls: merged}
|
||||
}
|
||||
tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, progress, toolEmitSeen, subAgentToolStep, markPending)
|
||||
if streamRecvErr != nil {
|
||||
if progress != nil {
|
||||
progress("eino_stream_error", streamRecvErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
})
|
||||
}
|
||||
if retErr := handleRunErr(streamRecvErr); retErr != nil {
|
||||
return nil, retErr
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
takePartial := func(runErr error) (*RunResult, error) {
|
||||
if len(runAccumulatedMsgs) <= baseAccumulatedCount {
|
||||
return nil, runErr
|
||||
}
|
||||
ids := snapshotMCPIDs()
|
||||
return buildEinoRunResultFromAccumulated(
|
||||
orchMode, runAccumulatedMsgs, lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, true,
|
||||
), runErr
|
||||
}
|
||||
|
||||
msg, gerr := mv.GetMessage()
|
||||
if gerr != nil || msg == nil {
|
||||
continue
|
||||
for {
|
||||
// 检测 context 取消(用户关闭浏览器、请求超时等),flush pending 工具状态避免 UI 卡在 "执行中"。
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
flushAllPendingAsFailed(ctx.Err())
|
||||
if progress != nil {
|
||||
progress("error", "Request cancelled / 请求已取消", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
runAccumulatedMsgs = append(runAccumulatedMsgs, msg)
|
||||
tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, progress, toolEmitSeen, subAgentToolStep, markPending)
|
||||
return takePartial(ctx.Err())
|
||||
default:
|
||||
}
|
||||
|
||||
if mv.Role == schema.Assistant {
|
||||
if progress != nil && strings.TrimSpace(msg.ReasoningContent) != "" {
|
||||
progress("thinking", strings.TrimSpace(msg.ReasoningContent), map[string]interface{}{
|
||||
ev, ok := iter.Next()
|
||||
if !ok {
|
||||
// iter 结束并不总是“正常完成”:
|
||||
// 当取消/超时发生在 iter.Next() 阻塞期间时,可能直接返回 !ok。
|
||||
// 此时必须保留 checkpoint,避免后续恢复时被误判为“无断点”而全量重跑。
|
||||
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||
flushAllPendingAsFailed(ctxErr)
|
||||
if progress != nil {
|
||||
progress("error", ctxErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}
|
||||
body := strings.TrimSpace(msg.Content)
|
||||
if body != "" {
|
||||
if streamsMainAssistant(ev.AgentName) {
|
||||
if progress != nil {
|
||||
return takePartial(ctxErr)
|
||||
}
|
||||
if len(pendingByID) > 0 {
|
||||
orphanCount := len(pendingByID)
|
||||
flushAllPendingAsFailed(errors.New("pending tool call missing result before run completion"))
|
||||
if progress != nil {
|
||||
progress("eino_pending_orphaned", "pending tool calls were force-closed at run end", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"orchestration": orchMode,
|
||||
"pendingCount": orphanCount,
|
||||
})
|
||||
}
|
||||
}
|
||||
if cpStore != nil && checkPointID != "" {
|
||||
if p, pErr := cpStore.path(checkPointID); pErr == nil {
|
||||
if rmErr := os.Remove(p); rmErr != nil && !os.IsNotExist(rmErr) && logger != nil {
|
||||
logger.Warn("eino checkpoint cleanup failed", zap.String("path", p), zap.Error(rmErr))
|
||||
}
|
||||
}
|
||||
}
|
||||
break
|
||||
}
|
||||
if ev == nil {
|
||||
continue
|
||||
}
|
||||
if ev.Err != nil {
|
||||
if retErr := handleRunErr(ev.Err); retErr != nil {
|
||||
return takePartial(retErr)
|
||||
}
|
||||
}
|
||||
if ev.AgentName != "" && progress != nil {
|
||||
iterEinoAgent := orchestratorName
|
||||
if orchMode == "plan_execute" {
|
||||
if a := strings.TrimSpace(ev.AgentName); a != "" {
|
||||
iterEinoAgent = a
|
||||
}
|
||||
}
|
||||
if streamsMainAssistant(ev.AgentName) {
|
||||
if einoMainRound == 0 {
|
||||
einoMainRound = 1
|
||||
progress("iteration", "", map[string]interface{}{
|
||||
"iteration": 1,
|
||||
"einoScope": "main",
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": iterEinoAgent,
|
||||
"orchestration": orchMode,
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
} else if einoLastAgent != "" && !streamsMainAssistant(einoLastAgent) {
|
||||
einoMainRound++
|
||||
progress("iteration", "", map[string]interface{}{
|
||||
"iteration": einoMainRound,
|
||||
"einoScope": "main",
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": iterEinoAgent,
|
||||
"orchestration": orchMode,
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
}
|
||||
einoLastAgent = ev.AgentName
|
||||
progress("progress", fmt.Sprintf("[Eino] %s", ev.AgentName), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}
|
||||
if ev.Output == nil || ev.Output.MessageOutput == nil {
|
||||
continue
|
||||
}
|
||||
mv := ev.Output.MessageOutput
|
||||
|
||||
if mv.IsStreaming && mv.MessageStream != nil {
|
||||
streamHeaderSent := false
|
||||
var reasoningStreamID string
|
||||
var toolStreamFragments []schema.ToolCall
|
||||
var subAssistantBuf strings.Builder
|
||||
var subReplyStreamID string
|
||||
var mainAssistantBuf strings.Builder
|
||||
var streamRecvErr error
|
||||
for {
|
||||
chunk, rerr := mv.MessageStream.Recv()
|
||||
if rerr != nil {
|
||||
if errors.Is(rerr, io.EOF) {
|
||||
break
|
||||
}
|
||||
if logger != nil {
|
||||
logger.Warn("eino stream recv error, flushing incomplete stream",
|
||||
zap.Error(rerr),
|
||||
zap.String("agent", ev.AgentName),
|
||||
zap.Int("toolFragments", len(toolStreamFragments)))
|
||||
}
|
||||
streamRecvErr = rerr
|
||||
break
|
||||
}
|
||||
if chunk == nil {
|
||||
continue
|
||||
}
|
||||
if progress != nil && strings.TrimSpace(chunk.ReasoningContent) != "" {
|
||||
if reasoningStreamID == "" {
|
||||
reasoningStreamID = fmt.Sprintf("eino-reasoning-%s-%d", conversationID, atomic.AddInt64(&reasoningStreamSeq, 1))
|
||||
progress("thinking_stream_start", " ", map[string]interface{}{
|
||||
"streamId": reasoningStreamID,
|
||||
"source": "eino",
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}
|
||||
progress("thinking_stream_delta", chunk.ReasoningContent, map[string]interface{}{
|
||||
"streamId": reasoningStreamID,
|
||||
})
|
||||
}
|
||||
if chunk.Content != "" {
|
||||
if progress != nil && streamsMainAssistant(ev.AgentName) {
|
||||
if !streamHeaderSent {
|
||||
progress("response_start", "", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
@@ -511,20 +478,61 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
progress("response_delta", body, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
streamHeaderSent = true
|
||||
}
|
||||
progress("response_delta", chunk.Content, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
mainAssistantBuf.WriteString(chunk.Content)
|
||||
} else if !streamsMainAssistant(ev.AgentName) {
|
||||
if progress != nil {
|
||||
if subReplyStreamID == "" {
|
||||
subReplyStreamID = fmt.Sprintf("eino-sub-reply-%s-%d", conversationID, atomic.AddInt64(&einoSubReplyStreamSeq, 1))
|
||||
progress("eino_agent_reply_stream_start", "", map[string]interface{}{
|
||||
"streamId": subReplyStreamID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": "sub",
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
progress("eino_agent_reply_stream_delta", chunk.Content, map[string]interface{}{
|
||||
"streamId": subReplyStreamID,
|
||||
"conversationId": conversationID,
|
||||
})
|
||||
}
|
||||
lastAssistant = body
|
||||
if orchMode == "plan_execute" && strings.EqualFold(strings.TrimSpace(ev.AgentName), "executor") {
|
||||
lastPlanExecuteExecutor = UnwrapPlanExecuteUserText(body)
|
||||
}
|
||||
} else if progress != nil {
|
||||
progress("eino_agent_reply", body, map[string]interface{}{
|
||||
subAssistantBuf.WriteString(chunk.Content)
|
||||
}
|
||||
}
|
||||
if len(chunk.ToolCalls) > 0 {
|
||||
toolStreamFragments = append(toolStreamFragments, chunk.ToolCalls...)
|
||||
}
|
||||
}
|
||||
if streamsMainAssistant(ev.AgentName) {
|
||||
if s := strings.TrimSpace(mainAssistantBuf.String()); s != "" {
|
||||
lastAssistant = s
|
||||
runAccumulatedMsgs = append(runAccumulatedMsgs, schema.AssistantMessage(s, nil))
|
||||
if orchMode == "plan_execute" && strings.EqualFold(strings.TrimSpace(ev.AgentName), "executor") {
|
||||
lastPlanExecuteExecutor = UnwrapPlanExecuteUserText(s)
|
||||
}
|
||||
}
|
||||
}
|
||||
if subAssistantBuf.Len() > 0 && progress != nil {
|
||||
if s := strings.TrimSpace(subAssistantBuf.String()); s != "" {
|
||||
if subReplyStreamID != "" {
|
||||
progress("eino_agent_reply_stream_end", s, map[string]interface{}{
|
||||
"streamId": subReplyStreamID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": "sub",
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
})
|
||||
} else {
|
||||
progress("eino_agent_reply", s, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": "sub",
|
||||
@@ -533,65 +541,157 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if mv.Role == schema.Tool && progress != nil {
|
||||
toolName := msg.ToolName
|
||||
if toolName == "" {
|
||||
toolName = mv.ToolName
|
||||
var lastToolChunk *schema.Message
|
||||
if merged := mergeStreamingToolCallFragments(toolStreamFragments); len(merged) > 0 {
|
||||
lastToolChunk = &schema.Message{ToolCalls: merged}
|
||||
}
|
||||
tryEmitToolCallsOnce(lastToolChunk, ev.AgentName, orchestratorName, conversationID, progress, toolEmitSeen, subAgentToolStep, markPending)
|
||||
if streamRecvErr != nil {
|
||||
if progress != nil {
|
||||
progress("eino_stream_error", streamRecvErr.Error(), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
})
|
||||
}
|
||||
|
||||
content := msg.Content
|
||||
isErr := false
|
||||
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
|
||||
isErr = true
|
||||
content = strings.TrimPrefix(content, einomcp.ToolErrorPrefix)
|
||||
if retErr := handleRunErr(streamRecvErr); retErr != nil {
|
||||
return takePartial(retErr)
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
preview := content
|
||||
if len(preview) > 200 {
|
||||
preview = preview[:200] + "..."
|
||||
}
|
||||
data := map[string]interface{}{
|
||||
"toolName": toolName,
|
||||
"success": !isErr,
|
||||
"isError": isErr,
|
||||
"result": content,
|
||||
"resultPreview": preview,
|
||||
msg, gerr := mv.GetMessage()
|
||||
if gerr != nil || msg == nil {
|
||||
continue
|
||||
}
|
||||
runAccumulatedMsgs = append(runAccumulatedMsgs, msg)
|
||||
tryEmitToolCallsOnce(mergeMessageToolCalls(msg), ev.AgentName, orchestratorName, conversationID, progress, toolEmitSeen, subAgentToolStep, markPending)
|
||||
|
||||
if mv.Role == schema.Assistant {
|
||||
if progress != nil && strings.TrimSpace(msg.ReasoningContent) != "" {
|
||||
progress("thinking", strings.TrimSpace(msg.ReasoningContent), map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"source": "eino",
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
"source": "eino",
|
||||
}
|
||||
toolCallID := strings.TrimSpace(msg.ToolCallID)
|
||||
if toolCallID == "" {
|
||||
if inferred, ok := popNextPendingForAgent(ev.AgentName); ok {
|
||||
toolCallID = inferred.ToolCallID
|
||||
} else if inferred, ok := popNextPendingForAgent(orchestratorName); ok {
|
||||
toolCallID = inferred.ToolCallID
|
||||
} else if inferred, ok := popNextPendingForAgent(""); ok {
|
||||
toolCallID = inferred.ToolCallID
|
||||
} else {
|
||||
for id := range pendingByID {
|
||||
toolCallID = id
|
||||
delete(pendingByID, id)
|
||||
break
|
||||
}
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}
|
||||
body := strings.TrimSpace(msg.Content)
|
||||
if body != "" {
|
||||
if streamsMainAssistant(ev.AgentName) {
|
||||
if progress != nil {
|
||||
progress("response_start", "", map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"messageGeneratedBy": "eino:" + ev.AgentName,
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
progress("response_delta", body, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"mcpExecutionIds": snapshotMCPIDs(),
|
||||
"einoRole": "orchestrator",
|
||||
"einoAgent": ev.AgentName,
|
||||
"orchestration": orchMode,
|
||||
})
|
||||
}
|
||||
} else {
|
||||
removePendingByID(toolCallID)
|
||||
lastAssistant = body
|
||||
if orchMode == "plan_execute" && strings.EqualFold(strings.TrimSpace(ev.AgentName), "executor") {
|
||||
lastPlanExecuteExecutor = UnwrapPlanExecuteUserText(body)
|
||||
}
|
||||
} else if progress != nil {
|
||||
progress("eino_agent_reply", body, map[string]interface{}{
|
||||
"conversationId": conversationID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": "sub",
|
||||
"source": "eino",
|
||||
})
|
||||
}
|
||||
if toolCallID != "" {
|
||||
data["toolCallId"] = toolCallID
|
||||
}
|
||||
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
|
||||
}
|
||||
}
|
||||
|
||||
if mv.Role == schema.Tool && progress != nil {
|
||||
toolName := msg.ToolName
|
||||
if toolName == "" {
|
||||
toolName = mv.ToolName
|
||||
}
|
||||
|
||||
content := msg.Content
|
||||
isErr := false
|
||||
if strings.HasPrefix(content, einomcp.ToolErrorPrefix) {
|
||||
isErr = true
|
||||
content = strings.TrimPrefix(content, einomcp.ToolErrorPrefix)
|
||||
}
|
||||
|
||||
preview := content
|
||||
if len(preview) > 200 {
|
||||
preview = preview[:200] + "..."
|
||||
}
|
||||
data := map[string]interface{}{
|
||||
"toolName": toolName,
|
||||
"success": !isErr,
|
||||
"isError": isErr,
|
||||
"result": content,
|
||||
"resultPreview": preview,
|
||||
"conversationId": conversationID,
|
||||
"einoAgent": ev.AgentName,
|
||||
"einoRole": einoRoleTag(ev.AgentName),
|
||||
"source": "eino",
|
||||
}
|
||||
toolCallID := strings.TrimSpace(msg.ToolCallID)
|
||||
if toolCallID == "" {
|
||||
if inferred, ok := popNextPendingForAgent(ev.AgentName); ok {
|
||||
toolCallID = inferred.ToolCallID
|
||||
} else if inferred, ok := popNextPendingForAgent(orchestratorName); ok {
|
||||
toolCallID = inferred.ToolCallID
|
||||
} else if inferred, ok := popNextPendingForAgent(""); ok {
|
||||
toolCallID = inferred.ToolCallID
|
||||
} else {
|
||||
for id := range pendingByID {
|
||||
toolCallID = id
|
||||
delete(pendingByID, id)
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
removePendingByID(toolCallID)
|
||||
}
|
||||
if toolCallID != "" {
|
||||
data["toolCallId"] = toolCallID
|
||||
}
|
||||
progress("tool_result", fmt.Sprintf("工具结果 (%s)", toolName), data)
|
||||
}
|
||||
}
|
||||
|
||||
mcpIDsMu.Lock()
|
||||
ids := append([]string(nil), *mcpIDs...)
|
||||
mcpIDsMu.Unlock()
|
||||
|
||||
histJSON, _ := json.Marshal(lastRunMsgs)
|
||||
out := buildEinoRunResultFromAccumulated(
|
||||
orchMode, runAccumulatedMsgs, lastAssistant, lastPlanExecuteExecutor, emptyHint, ids, false,
|
||||
)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func einoPartialRunLastOutputHint() string {
|
||||
return "[执行未正常结束(用户停止、超时或异常)。续跑时请基于上文已产生的工具与结果继续,勿重复已完成步骤。]\n" +
|
||||
"[Run ended abnormally; continue from the trace above without repeating completed steps.]"
|
||||
}
|
||||
|
||||
func buildEinoRunResultFromAccumulated(
|
||||
orchMode string,
|
||||
runAccumulatedMsgs []adk.Message,
|
||||
lastAssistant string,
|
||||
lastPlanExecuteExecutor string,
|
||||
emptyHint string,
|
||||
mcpIDs []string,
|
||||
partial bool,
|
||||
) *RunResult {
|
||||
histJSON, _ := json.Marshal(runAccumulatedMsgs)
|
||||
cleaned := strings.TrimSpace(lastAssistant)
|
||||
if orchMode == "plan_execute" {
|
||||
if e := strings.TrimSpace(lastPlanExecuteExecutor); e != "" {
|
||||
@@ -607,15 +707,29 @@ func runEinoADKAgentLoop(ctx context.Context, args *einoADKRunLoopArgs, baseMsgs
|
||||
if rs := []rune(cleaned); len(rs) > maxResponseRunes {
|
||||
cleaned = string(rs[:maxResponseRunes]) + "\n\n... (response truncated / 响应已截断)"
|
||||
}
|
||||
lastOut := cleaned
|
||||
resp := cleaned
|
||||
if partial && cleaned == "" {
|
||||
lastOut = einoPartialRunLastOutputHint()
|
||||
resp = emptyHint
|
||||
}
|
||||
out := &RunResult{
|
||||
Response: cleaned,
|
||||
MCPExecutionIDs: ids,
|
||||
LastReActInput: string(histJSON),
|
||||
LastReActOutput: cleaned,
|
||||
Response: resp,
|
||||
MCPExecutionIDs: mcpIDs,
|
||||
LastAgentTraceInput: string(histJSON),
|
||||
LastAgentTraceOutput: lastOut,
|
||||
}
|
||||
if out.Response == "" {
|
||||
if !partial && out.Response == "" {
|
||||
out.Response = emptyHint
|
||||
out.LastReActOutput = out.Response
|
||||
out.LastAgentTraceOutput = out.Response
|
||||
}
|
||||
return out, nil
|
||||
return out
|
||||
}
|
||||
|
||||
func buildEinoCheckpointID(orchMode string) string {
|
||||
mode := sanitizeEinoPathSegment(strings.TrimSpace(orchMode))
|
||||
if mode == "" {
|
||||
mode = "default"
|
||||
}
|
||||
return "runner-" + mode
|
||||
}
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
|
||||
"github.com/bytedance/sonic"
|
||||
"github.com/cloudwego/eino/adk"
|
||||
"github.com/cloudwego/eino/schema"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
type einoModelInputTelemetryMiddleware struct {
|
||||
adk.BaseChatModelAgentMiddleware
|
||||
logger *zap.Logger
|
||||
modelName string
|
||||
conversationID string
|
||||
phase string
|
||||
}
|
||||
|
||||
func newEinoModelInputTelemetryMiddleware(
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
phase string,
|
||||
) adk.ChatModelAgentMiddleware {
|
||||
if logger == nil {
|
||||
return nil
|
||||
}
|
||||
return &einoModelInputTelemetryMiddleware{
|
||||
logger: logger,
|
||||
modelName: strings.TrimSpace(modelName),
|
||||
conversationID: strings.TrimSpace(conversationID),
|
||||
phase: strings.TrimSpace(phase),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *einoModelInputTelemetryMiddleware) BeforeModelRewriteState(
|
||||
ctx context.Context,
|
||||
state *adk.ChatModelAgentState,
|
||||
mc *adk.ModelContext,
|
||||
) (context.Context, *adk.ChatModelAgentState, error) {
|
||||
if m == nil || m.logger == nil || state == nil {
|
||||
return ctx, state, nil
|
||||
}
|
||||
tokens := estimateTokensForMessagesAndTools(ctx, m.modelName, state.Messages, mcTools(mc))
|
||||
m.logger.Info("eino model input estimated",
|
||||
zap.String("phase", m.phase),
|
||||
zap.String("conversation_id", m.conversationID),
|
||||
zap.Int("messages", len(state.Messages)),
|
||||
zap.Int("tools", len(mcTools(mc))),
|
||||
zap.Int("input_tokens_estimated", tokens),
|
||||
)
|
||||
return ctx, state, nil
|
||||
}
|
||||
|
||||
func mcTools(mc *adk.ModelContext) []*schema.ToolInfo {
|
||||
if mc == nil || len(mc.Tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
return mc.Tools
|
||||
}
|
||||
|
||||
func estimateTokensForMessagesAndTools(
|
||||
_ context.Context,
|
||||
modelName string,
|
||||
messages []adk.Message,
|
||||
tools []*schema.ToolInfo,
|
||||
) int {
|
||||
var sb strings.Builder
|
||||
for _, msg := range messages {
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
sb.WriteString(string(msg.Role))
|
||||
sb.WriteByte('\n')
|
||||
sb.WriteString(msg.Content)
|
||||
sb.WriteByte('\n')
|
||||
if msg.ReasoningContent != "" {
|
||||
sb.WriteString(msg.ReasoningContent)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
if len(msg.ToolCalls) > 0 {
|
||||
if b, err := sonic.Marshal(msg.ToolCalls); err == nil {
|
||||
sb.Write(b)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, tl := range tools {
|
||||
if tl == nil {
|
||||
continue
|
||||
}
|
||||
cp := *tl
|
||||
cp.Extra = nil
|
||||
if text, err := sonic.MarshalString(cp); err == nil {
|
||||
sb.WriteString(text)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
}
|
||||
text := sb.String()
|
||||
if text == "" {
|
||||
return 0
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
if n, err := tc.Count(modelName, text); err == nil {
|
||||
return n
|
||||
}
|
||||
return (len(text) + 3) / 4
|
||||
}
|
||||
|
||||
func logPlanExecuteModelInputEstimate(
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
phase string,
|
||||
msgs []adk.Message,
|
||||
) {
|
||||
if logger == nil {
|
||||
return
|
||||
}
|
||||
tokens := estimateTokensForMessagesAndTools(context.Background(), modelName, msgs, nil)
|
||||
logger.Info("eino model input estimated",
|
||||
zap.String("phase", phase),
|
||||
zap.String("conversation_id", strings.TrimSpace(conversationID)),
|
||||
zap.Int("messages", len(msgs)),
|
||||
zap.Int("tools", 0),
|
||||
zap.Int("input_tokens_estimated", tokens),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
"cyberstrike-ai/internal/mcp/builtin"
|
||||
|
||||
localbk "github.com/cloudwego/eino-ext/adk/backend/local"
|
||||
"github.com/cloudwego/eino/adk"
|
||||
@@ -65,6 +66,66 @@ func splitToolsForToolSearch(all []tool.BaseTool, alwaysVisible int) (static []t
|
||||
return append([]tool.BaseTool(nil), all[:alwaysVisible]...), append([]tool.BaseTool(nil), all[alwaysVisible:]...), true
|
||||
}
|
||||
|
||||
func splitToolsForToolSearchByNames(all []tool.BaseTool, names []string, fallbackAlwaysVisible int) (static []tool.BaseTool, dynamic []tool.BaseTool, ok bool) {
|
||||
nameSet := make(map[string]struct{}, len(names))
|
||||
for _, n := range names {
|
||||
n = strings.TrimSpace(strings.ToLower(n))
|
||||
if n == "" {
|
||||
continue
|
||||
}
|
||||
nameSet[n] = struct{}{}
|
||||
}
|
||||
if len(nameSet) == 0 {
|
||||
return splitToolsForToolSearch(all, fallbackAlwaysVisible)
|
||||
}
|
||||
static = make([]tool.BaseTool, 0, len(all))
|
||||
dynamic = make([]tool.BaseTool, 0, len(all))
|
||||
for _, t := range all {
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
info, err := t.Info(context.Background())
|
||||
name := ""
|
||||
if err == nil && info != nil {
|
||||
name = strings.TrimSpace(strings.ToLower(info.Name))
|
||||
}
|
||||
if _, keep := nameSet[name]; keep {
|
||||
static = append(static, t)
|
||||
continue
|
||||
}
|
||||
dynamic = append(dynamic, t)
|
||||
}
|
||||
if len(static) == 0 || len(dynamic) == 0 {
|
||||
// fallback: preserve previous behavior when whitelist misses all or includes all.
|
||||
return splitToolsForToolSearch(all, fallbackAlwaysVisible)
|
||||
}
|
||||
return static, dynamic, true
|
||||
}
|
||||
|
||||
func mergeAlwaysVisibleToolNames(configured []string) []string {
|
||||
merged := make([]string, 0, len(configured)+32)
|
||||
seen := make(map[string]struct{}, len(configured)+32)
|
||||
add := func(name string) {
|
||||
n := strings.TrimSpace(strings.ToLower(name))
|
||||
if n == "" {
|
||||
return
|
||||
}
|
||||
if _, ok := seen[n]; ok {
|
||||
return
|
||||
}
|
||||
seen[n] = struct{}{}
|
||||
merged = append(merged, n)
|
||||
}
|
||||
for _, n := range configured {
|
||||
add(n)
|
||||
}
|
||||
// Always include hardcoded backend builtin MCP tools from constants.
|
||||
for _, n := range builtin.GetAllBuiltinTools() {
|
||||
add(n)
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
func buildReductionMiddleware(ctx context.Context, mw config.MultiAgentEinoMiddlewareConfig, convID string, loc *localbk.Local, logger *zap.Logger) (adk.ChatModelAgentMiddleware, error) {
|
||||
if loc == nil {
|
||||
return nil, fmt.Errorf("reduction: local backend nil")
|
||||
@@ -87,6 +148,8 @@ func buildReductionMiddleware(ctx context.Context, mw config.MultiAgentEinoMiddl
|
||||
RootDir: root,
|
||||
ReadFileToolName: "read_file",
|
||||
ClearExcludeTools: excl,
|
||||
MaxLengthForTrunc: mw.ReductionMaxLengthForTruncEffective(),
|
||||
MaxTokensForClear: int64(mw.ReductionMaxTokensForClearEffective()),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -142,7 +205,7 @@ func prependEinoMiddlewares(
|
||||
alwaysVis = 12
|
||||
}
|
||||
if mw.ToolSearchEnable && len(tools) >= minTools {
|
||||
static, dynamic, split := splitToolsForToolSearch(tools, alwaysVis)
|
||||
static, dynamic, split := splitToolsForToolSearchByNames(tools, mergeAlwaysVisibleToolNames(mw.ToolSearchAlwaysVisibleTools), alwaysVis)
|
||||
if split && len(dynamic) > 0 {
|
||||
ts, terr := toolsearch.New(ctx, &toolsearch.Config{DynamicTools: dynamic})
|
||||
if terr != nil {
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/cloudwego/eino/adk"
|
||||
)
|
||||
|
||||
func applyBeforeModelRewriteHandlers(
|
||||
ctx context.Context,
|
||||
msgs []adk.Message,
|
||||
handlers []adk.ChatModelAgentMiddleware,
|
||||
) ([]adk.Message, error) {
|
||||
if len(msgs) == 0 || len(handlers) == 0 {
|
||||
return msgs, nil
|
||||
}
|
||||
state := &adk.ChatModelAgentState{Messages: msgs}
|
||||
modelCtx := &adk.ModelContext{}
|
||||
curCtx := ctx
|
||||
for _, h := range handlers {
|
||||
if h == nil {
|
||||
continue
|
||||
}
|
||||
nextCtx, nextState, err := h.BeforeModelRewriteState(curCtx, state, modelCtx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("before model rewrite: %w", err)
|
||||
}
|
||||
if nextCtx != nil {
|
||||
curCtx = nextCtx
|
||||
}
|
||||
if nextState != nil {
|
||||
state = nextState
|
||||
}
|
||||
}
|
||||
return state.Messages, nil
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino-ext/components/model/openai"
|
||||
@@ -25,7 +26,12 @@ type PlanExecuteRootArgs struct {
|
||||
LoopMaxIter int
|
||||
// AppCfg / Logger 非空时为 Executor 挂载与 Deep/Supervisor 一致的 Eino summarization 中间件。
|
||||
AppCfg *config.Config
|
||||
MwCfg *config.MultiAgentEinoMiddlewareConfig
|
||||
// ConversationID is used for transcript/isolation paths in middleware.
|
||||
ConversationID string
|
||||
Logger *zap.Logger
|
||||
// ModelName is used for model input token estimation logs.
|
||||
ModelName string
|
||||
// ExecPreMiddlewares 是由 prependEinoMiddlewares 构建的前置中间件(patchtoolcalls, reduction, toolsearch, plantask),
|
||||
// 与 Deep/Supervisor 主代理的 mainOrchestratorPre 一致。
|
||||
ExecPreMiddlewares []adk.ChatModelAgentMiddleware
|
||||
@@ -33,6 +39,8 @@ type PlanExecuteRootArgs struct {
|
||||
SkillMiddleware adk.ChatModelAgentMiddleware
|
||||
// FilesystemMiddleware 是 Eino filesystem 中间件,当 eino_skills.filesystem_tools 启用时提供本机文件读写与 Shell 能力(可选)。
|
||||
FilesystemMiddleware adk.ChatModelAgentMiddleware
|
||||
// PlannerReplannerRewriteHandlers applies BeforeModelRewriteState pipeline for planner/replanner input.
|
||||
PlannerReplannerRewriteHandlers []adk.ChatModelAgentMiddleware
|
||||
}
|
||||
|
||||
// NewPlanExecuteRoot 返回 plan → execute → replan 预置编排根节点(与 Deep / Supervisor 并列)。
|
||||
@@ -50,7 +58,7 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
|
||||
plannerCfg := &planexecute.PlannerConfig{
|
||||
ToolCallingChatModel: tcm,
|
||||
}
|
||||
if fn := planExecutePlannerGenInput(a.OrchInstruction); fn != nil {
|
||||
if fn := planExecutePlannerGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID, a.PlannerReplannerRewriteHandlers); fn != nil {
|
||||
plannerCfg.GenInputFn = fn
|
||||
}
|
||||
planner, err := planexecute.NewPlanner(ctx, plannerCfg)
|
||||
@@ -59,7 +67,7 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
|
||||
}
|
||||
replanner, err := planexecute.NewReplanner(ctx, &planexecute.ReplannerConfig{
|
||||
ChatModel: tcm,
|
||||
GenInputFn: planExecuteReplannerGenInput(a.OrchInstruction),
|
||||
GenInputFn: planExecuteReplannerGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID, a.PlannerReplannerRewriteHandlers),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plan_execute replanner: %w", err)
|
||||
@@ -81,17 +89,20 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
|
||||
}
|
||||
// 4. summarization(最后,与 Deep/Supervisor 一致)
|
||||
if a.AppCfg != nil {
|
||||
sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.Logger)
|
||||
sumMw, sumErr := newEinoSummarizationMiddleware(ctx, a.ExecModel, a.AppCfg, a.MwCfg, a.ConversationID, a.Logger)
|
||||
if sumErr != nil {
|
||||
return nil, fmt.Errorf("plan_execute executor summarization: %w", sumErr)
|
||||
}
|
||||
execHandlers = append(execHandlers, sumMw)
|
||||
}
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(a.Logger, a.ModelName, a.ConversationID, "plan_execute_executor"); teleMw != nil {
|
||||
execHandlers = append(execHandlers, teleMw)
|
||||
}
|
||||
executor, err := newPlanExecuteExecutor(ctx, &planexecute.ExecutorConfig{
|
||||
Model: a.ExecModel,
|
||||
ToolsConfig: a.ToolsCfg,
|
||||
MaxIterations: a.ExecMaxIter,
|
||||
GenInputFn: planExecuteExecutorGenInput(a.OrchInstruction),
|
||||
GenInputFn: planExecuteExecutorGenInput(a.OrchInstruction, a.AppCfg, a.MwCfg, a.Logger, a.ModelName, a.ConversationID),
|
||||
}, execHandlers)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("plan_execute executor: %w", err)
|
||||
@@ -110,20 +121,42 @@ func NewPlanExecuteRoot(ctx context.Context, a *PlanExecuteRootArgs) (adk.Resuma
|
||||
|
||||
// planExecutePlannerGenInput 将 orchestrator instruction 作为 SystemMessage 注入 planner 输入。
|
||||
// 返回 nil 时 Eino 使用内置默认 planner prompt。
|
||||
func planExecutePlannerGenInput(orchInstruction string) planexecute.GenPlannerModelInputFn {
|
||||
func planExecutePlannerGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
rewriteHandlers []adk.ChatModelAgentMiddleware,
|
||||
) planexecute.GenPlannerModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
if oi == "" {
|
||||
if oi == "" && appCfg == nil {
|
||||
return nil
|
||||
}
|
||||
return func(ctx context.Context, userInput []adk.Message) ([]adk.Message, error) {
|
||||
userInput = capPlanExecuteUserInputMessages(userInput, appCfg, mwCfg)
|
||||
msgs := make([]adk.Message, 0, 1+len(userInput))
|
||||
msgs = append(msgs, schema.SystemMessage(oi))
|
||||
if oi != "" {
|
||||
msgs = append(msgs, schema.SystemMessage(oi))
|
||||
}
|
||||
msgs = append(msgs, userInput...)
|
||||
if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 {
|
||||
msgs = rewritten
|
||||
}
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_planner", msgs)
|
||||
return msgs, nil
|
||||
}
|
||||
}
|
||||
|
||||
func planExecuteExecutorGenInput(orchInstruction string) planexecute.GenModelInputFn {
|
||||
func planExecuteExecutorGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
) planexecute.GenModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
return func(ctx context.Context, in *planexecute.ExecutionContext) ([]adk.Message, error) {
|
||||
planContent, err := in.Plan.MarshalJSON()
|
||||
@@ -131,9 +164,9 @@ func planExecuteExecutorGenInput(orchInstruction string) planexecute.GenModelInp
|
||||
return nil, err
|
||||
}
|
||||
userMsgs, err := planexecute.ExecutorPrompt.Format(ctx, map[string]any{
|
||||
"input": planExecuteFormatInput(in.UserInput),
|
||||
"input": planExecuteFormatInput(capPlanExecuteUserInputMessages(in.UserInput, appCfg, mwCfg)),
|
||||
"plan": string(planContent),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, appCfg, mwCfg),
|
||||
"step": in.Plan.FirstStep(),
|
||||
})
|
||||
if err != nil {
|
||||
@@ -142,6 +175,7 @@ func planExecuteExecutorGenInput(orchInstruction string) planexecute.GenModelInp
|
||||
if oi != "" {
|
||||
userMsgs = append([]adk.Message{schema.SystemMessage(oi)}, userMsgs...)
|
||||
}
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_executor_gen_input", userMsgs)
|
||||
return userMsgs, nil
|
||||
}
|
||||
}
|
||||
@@ -155,18 +189,22 @@ func planExecuteFormatInput(input []adk.Message) string {
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func planExecuteFormatExecutedSteps(results []planexecute.ExecutedStep) string {
|
||||
capped := capPlanExecuteExecutedSteps(results)
|
||||
var sb strings.Builder
|
||||
for _, result := range capped {
|
||||
sb.WriteString(fmt.Sprintf("Step: %s\nResult: %s\n\n", result.Step, result.Result))
|
||||
}
|
||||
return sb.String()
|
||||
func planExecuteFormatExecutedSteps(results []planexecute.ExecutedStep, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) string {
|
||||
capped := capPlanExecuteExecutedStepsWithConfig(results, mwCfg)
|
||||
return renderPlanExecuteStepsByBudget(capped, appCfg, mwCfg)
|
||||
}
|
||||
|
||||
// planExecuteReplannerGenInput 与 Eino 默认 Replanner 输入一致,但 executed_steps 经 cap 后再写入 prompt,
|
||||
// 且在 orchInstruction 非空时 prepend SystemMessage 使 replanner 也能接收全局指令。
|
||||
func planExecuteReplannerGenInput(orchInstruction string) planexecute.GenModelInputFn {
|
||||
func planExecuteReplannerGenInput(
|
||||
orchInstruction string,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
logger *zap.Logger,
|
||||
modelName string,
|
||||
conversationID string,
|
||||
rewriteHandlers []adk.ChatModelAgentMiddleware,
|
||||
) planexecute.GenModelInputFn {
|
||||
oi := strings.TrimSpace(orchInstruction)
|
||||
return func(ctx context.Context, in *planexecute.ExecutionContext) ([]adk.Message, error) {
|
||||
planContent, err := in.Plan.MarshalJSON()
|
||||
@@ -175,8 +213,8 @@ func planExecuteReplannerGenInput(orchInstruction string) planexecute.GenModelIn
|
||||
}
|
||||
msgs, err := planexecute.ReplannerPrompt.Format(ctx, map[string]any{
|
||||
"plan": string(planContent),
|
||||
"input": planExecuteFormatInput(in.UserInput),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps),
|
||||
"input": planExecuteFormatInput(capPlanExecuteUserInputMessages(in.UserInput, appCfg, mwCfg)),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, appCfg, mwCfg),
|
||||
"plan_tool": planexecute.PlanToolInfo.Name,
|
||||
"respond_tool": planexecute.RespondToolInfo.Name,
|
||||
})
|
||||
@@ -186,10 +224,120 @@ func planExecuteReplannerGenInput(orchInstruction string) planexecute.GenModelIn
|
||||
if oi != "" {
|
||||
msgs = append([]adk.Message{schema.SystemMessage(oi)}, msgs...)
|
||||
}
|
||||
if rewritten, rerr := applyBeforeModelRewriteHandlers(ctx, msgs, rewriteHandlers); rerr == nil && len(rewritten) > 0 {
|
||||
msgs = rewritten
|
||||
}
|
||||
logPlanExecuteModelInputEstimate(logger, modelName, conversationID, "plan_execute_replanner", msgs)
|
||||
return msgs, nil
|
||||
}
|
||||
}
|
||||
|
||||
func capPlanExecuteUserInputMessages(input []adk.Message, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) []adk.Message {
|
||||
if len(input) == 0 {
|
||||
return input
|
||||
}
|
||||
maxTotal := 120000
|
||||
modelName := "gpt-4o"
|
||||
if appCfg != nil {
|
||||
if appCfg.OpenAI.MaxTotalTokens > 0 {
|
||||
maxTotal = appCfg.OpenAI.MaxTotalTokens
|
||||
}
|
||||
if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" {
|
||||
modelName = m
|
||||
}
|
||||
}
|
||||
// Reserve most tokens for planner/replanner prompt and tool schema.
|
||||
ratio := 0.35
|
||||
if mwCfg != nil {
|
||||
ratio = mwCfg.PlanExecuteUserInputBudgetRatioEffective()
|
||||
}
|
||||
budget := int(float64(maxTotal) * ratio)
|
||||
if budget < 4096 {
|
||||
budget = 4096
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
out := make([]adk.Message, 0, len(input))
|
||||
used := 0
|
||||
for i := len(input) - 1; i >= 0; i-- {
|
||||
msg := input[i]
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
n, err := tc.Count(modelName, string(msg.Role)+"\n"+msg.Content)
|
||||
if err != nil {
|
||||
n = (len(msg.Content) + 3) / 4
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
if used+n > budget {
|
||||
break
|
||||
}
|
||||
used += n
|
||||
out = append(out, msg)
|
||||
}
|
||||
for i, j := 0, len(out)-1; i < j; i, j = i+1, j-1 {
|
||||
out[i], out[j] = out[j], out[i]
|
||||
}
|
||||
if len(out) == 0 {
|
||||
// Keep the latest user message at least.
|
||||
return []adk.Message{input[len(input)-1]}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func renderPlanExecuteStepsByBudget(steps []planexecute.ExecutedStep, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) string {
|
||||
if len(steps) == 0 {
|
||||
return ""
|
||||
}
|
||||
maxTotal := 120000
|
||||
modelName := "gpt-4o"
|
||||
if appCfg != nil {
|
||||
if appCfg.OpenAI.MaxTotalTokens > 0 {
|
||||
maxTotal = appCfg.OpenAI.MaxTotalTokens
|
||||
}
|
||||
if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" {
|
||||
modelName = m
|
||||
}
|
||||
}
|
||||
ratio := 0.2
|
||||
if mwCfg != nil {
|
||||
ratio = mwCfg.PlanExecuteExecutedStepsBudgetRatioEffective()
|
||||
}
|
||||
budget := int(float64(maxTotal) * ratio)
|
||||
if budget < 3072 {
|
||||
budget = 3072
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
var kept []string
|
||||
used := 0
|
||||
skipped := 0
|
||||
for i := len(steps) - 1; i >= 0; i-- {
|
||||
block := fmt.Sprintf("Step: %s\nResult: %s\n\n", steps[i].Step, steps[i].Result)
|
||||
n, err := tc.Count(modelName, block)
|
||||
if err != nil {
|
||||
n = (len(block) + 3) / 4
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
if used+n > budget {
|
||||
skipped = i + 1
|
||||
break
|
||||
}
|
||||
used += n
|
||||
kept = append(kept, block)
|
||||
}
|
||||
var sb strings.Builder
|
||||
if skipped > 0 {
|
||||
sb.WriteString(fmt.Sprintf("Earlier executed steps omitted due to context budget: %d steps.\n\n", skipped))
|
||||
}
|
||||
for i := len(kept) - 1; i >= 0; i-- {
|
||||
sb.WriteString(kept[i])
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// planExecuteStreamsMainAssistant 将规划/执行/重规划各阶段助手流式输出映射到主对话区。
|
||||
func planExecuteStreamsMainAssistant(agent string) bool {
|
||||
if agent == "" {
|
||||
|
||||
@@ -125,7 +125,7 @@ func RunEinoSingleChatModelAgent(
|
||||
return nil, fmt.Errorf("eino single 模型: %w", err)
|
||||
}
|
||||
|
||||
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, logger)
|
||||
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("eino single summarization: %w", err)
|
||||
}
|
||||
@@ -145,6 +145,9 @@ func RunEinoSingleChatModelAgent(
|
||||
handlers = append(handlers, einoSkillMW)
|
||||
}
|
||||
handlers = append(handlers, mainSumMw)
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "eino_single"); teleMw != nil {
|
||||
handlers = append(handlers, teleMw)
|
||||
}
|
||||
|
||||
maxIter := ma.MaxIteration
|
||||
if maxIter <= 0 {
|
||||
@@ -165,11 +168,29 @@ func RunEinoSingleChatModelAgent(
|
||||
},
|
||||
EmitInternalEvents: true,
|
||||
}
|
||||
ins := injectToolNamesOnlyInstruction(ctx, ag.EinoSingleAgentSystemInstruction(), mainTools)
|
||||
if logger != nil {
|
||||
names := collectToolNames(ctx, mainTools)
|
||||
mountedNames := collectToolNames(ctx, mainToolsForCfg)
|
||||
hasToolSearch := false
|
||||
for _, n := range names {
|
||||
if strings.EqualFold(strings.TrimSpace(n), "tool_search") {
|
||||
hasToolSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Info("eino tool-name injection",
|
||||
zap.String("scope", "eino_single"),
|
||||
zap.Int("tool_names", len(names)),
|
||||
zap.Int("mounted_tool_names", len(mountedNames)),
|
||||
zap.Bool("has_tool_search", hasToolSearch),
|
||||
)
|
||||
}
|
||||
|
||||
chatCfg := &adk.ChatModelAgentConfig{
|
||||
Name: einoSingleAgentName,
|
||||
Description: "Eino ADK ChatModelAgent with MCP tools for authorized security testing.",
|
||||
Instruction: ag.EinoSingleAgentSystemInstruction(),
|
||||
Instruction: ins,
|
||||
Model: mainModel,
|
||||
ToolsConfig: mainToolsCfg,
|
||||
MaxIterations: maxIter,
|
||||
@@ -188,7 +209,7 @@ func RunEinoSingleChatModelAgent(
|
||||
return nil, fmt.Errorf("eino single NewChatModelAgent: %w", err)
|
||||
}
|
||||
|
||||
baseMsgs := historyToMessages(history)
|
||||
baseMsgs := historyToMessages(history, appCfg, &ma.EinoMiddleware)
|
||||
baseMsgs = append(baseMsgs, schema.UserMessage(userMessage))
|
||||
|
||||
streamsMainAssistant := func(agent string) bool {
|
||||
|
||||
@@ -3,6 +3,8 @@ package multiagent
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"cyberstrike-ai/internal/agent"
|
||||
@@ -32,6 +34,8 @@ func newEinoSummarizationMiddleware(
|
||||
ctx context.Context,
|
||||
summaryModel model.BaseChatModel,
|
||||
appCfg *config.Config,
|
||||
mwCfg *config.MultiAgentEinoMiddlewareConfig,
|
||||
conversationID string,
|
||||
logger *zap.Logger,
|
||||
) (adk.ChatModelAgentMiddleware, error) {
|
||||
if summaryModel == nil || appCfg == nil {
|
||||
@@ -41,7 +45,14 @@ func newEinoSummarizationMiddleware(
|
||||
if maxTotal <= 0 {
|
||||
maxTotal = 120000
|
||||
}
|
||||
trigger := int(float64(maxTotal) * 0.9)
|
||||
triggerRatio := 0.8
|
||||
emitInternalEvents := true
|
||||
if mwCfg != nil {
|
||||
triggerRatio = mwCfg.SummarizationTriggerRatioEffective()
|
||||
emitInternalEvents = mwCfg.SummarizationEmitInternalEventsEffective()
|
||||
}
|
||||
// Keep enough safety margin for tokenizer/model-side accounting mismatch.
|
||||
trigger := int(float64(maxTotal) * triggerRatio)
|
||||
if trigger < 4096 {
|
||||
trigger = maxTotal
|
||||
if trigger < 4096 {
|
||||
@@ -65,6 +76,18 @@ func newEinoSummarizationMiddleware(
|
||||
if recentTrailMax > trigger/2 {
|
||||
recentTrailMax = trigger / 2
|
||||
}
|
||||
transcriptPath := ""
|
||||
if conv := strings.TrimSpace(conversationID); conv != "" {
|
||||
baseRoot := filepath.Join(os.TempDir(), "cyberstrike-summarization")
|
||||
if dbPath := strings.TrimSpace(appCfg.Database.Path); dbPath != "" {
|
||||
// Persist with the same lifecycle as local conversation storage.
|
||||
baseRoot = filepath.Join(filepath.Dir(dbPath), "conversation_artifacts", sanitizeEinoPathSegment(conv), "summarization")
|
||||
}
|
||||
base := baseRoot
|
||||
if mkErr := os.MkdirAll(base, 0o755); mkErr == nil {
|
||||
transcriptPath = filepath.Join(base, "transcript.txt")
|
||||
}
|
||||
}
|
||||
|
||||
mw, err := summarization.New(ctx, &summarization.Config{
|
||||
Model: summaryModel,
|
||||
@@ -73,7 +96,8 @@ func newEinoSummarizationMiddleware(
|
||||
},
|
||||
TokenCounter: tokenCounter,
|
||||
UserInstruction: einoSummarizeUserInstruction,
|
||||
EmitInternalEvents: false,
|
||||
EmitInternalEvents: emitInternalEvents,
|
||||
TranscriptFilePath: transcriptPath,
|
||||
PreserveUserMessages: &summarization.PreserveUserMessages{
|
||||
Enabled: true,
|
||||
MaxTokens: preserveMax,
|
||||
@@ -85,11 +109,16 @@ func newEinoSummarizationMiddleware(
|
||||
if logger == nil {
|
||||
return nil
|
||||
}
|
||||
beforeTokens, _ := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: before.Messages})
|
||||
afterTokens, _ := tokenCounter(ctx, &summarization.TokenCounterInput{Messages: after.Messages})
|
||||
logger.Info("eino summarization 已压缩上下文",
|
||||
zap.Int("messages_before", len(before.Messages)),
|
||||
zap.Int("messages_after", len(after.Messages)),
|
||||
zap.Int("tokens_before_estimated", beforeTokens),
|
||||
zap.Int("tokens_after_estimated", afterTokens),
|
||||
zap.Int("max_total_tokens", maxTotal),
|
||||
zap.Int("trigger_context_tokens", trigger),
|
||||
zap.String("transcript_file", transcriptPath),
|
||||
)
|
||||
return nil
|
||||
},
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
package multiagent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/cloudwego/eino/components/tool"
|
||||
)
|
||||
|
||||
// injectToolNamesOnlyInstruction prepends a compact tool-name-only section into
|
||||
// the system instruction so the model can reference current callable names.
|
||||
func injectToolNamesOnlyInstruction(ctx context.Context, instruction string, tools []tool.BaseTool) string {
|
||||
names := collectToolNames(ctx, tools)
|
||||
if len(names) == 0 {
|
||||
return strings.TrimSpace(instruction)
|
||||
}
|
||||
hasToolSearch := false
|
||||
for _, n := range names {
|
||||
if strings.EqualFold(strings.TrimSpace(n), "tool_search") {
|
||||
hasToolSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString("以下是当前会话中可调用的工具名称列表(仅名称,无参数定义):\n")
|
||||
for _, name := range names {
|
||||
sb.WriteString("- ")
|
||||
sb.WriteString(name)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
sb.WriteString("\n使用规则:\n")
|
||||
sb.WriteString("1) 上述仅为名称列表,不包含参数定义。\n")
|
||||
if hasToolSearch {
|
||||
sb.WriteString("2) 在调用具体工具前,应先使用 tool_search 查看工具详情与参数要求,再发起调用。\n")
|
||||
} else {
|
||||
sb.WriteString("2) 调用具体工具前,请先确认该工具的参数要求;不确定时先澄清再调用。\n")
|
||||
}
|
||||
sb.WriteString("3) 不要臆造不存在的工具名。\n\n")
|
||||
if s := strings.TrimSpace(instruction); s != "" {
|
||||
sb.WriteString(s)
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func collectToolNames(ctx context.Context, tools []tool.BaseTool) []string {
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
seen := make(map[string]struct{}, len(tools))
|
||||
out := make([]string, 0, len(tools))
|
||||
for _, t := range tools {
|
||||
if t == nil {
|
||||
continue
|
||||
}
|
||||
info, err := t.Info(ctx)
|
||||
if err != nil || info == nil {
|
||||
continue
|
||||
}
|
||||
name := strings.TrimSpace(info.Name)
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(name)
|
||||
if _, ok := seen[key]; ok {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
out = append(out, name)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -59,4 +59,3 @@ func (m *noNestedTaskMiddleware) WrapInvokableToolCall(
|
||||
return endpoint(ctx2, argumentsInJSON, opts...)
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -71,7 +71,7 @@ func planExecuteDefaultGenExecutorInput(ctx context.Context, in *planexecute.Exe
|
||||
return planexecute.ExecutorPrompt.Format(ctx, map[string]any{
|
||||
"input": planExecuteFormatInput(in.UserInput),
|
||||
"plan": string(planContent),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps),
|
||||
"executed_steps": planExecuteFormatExecutedSteps(in.ExecutedSteps, nil, nil),
|
||||
"step": in.Plan.FirstStep(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"cyberstrike-ai/internal/config"
|
||||
|
||||
"github.com/cloudwego/eino/adk/prebuilt/planexecute"
|
||||
)
|
||||
|
||||
@@ -12,8 +14,11 @@ import (
|
||||
// 此处仅约束「写入模型 prompt 的视图」,不修改 Eino session 中的原始 ExecutedSteps。
|
||||
|
||||
const (
|
||||
planExecuteMaxStepResultRunes = 12000
|
||||
planExecuteKeepLastSteps = 16
|
||||
defaultPlanExecuteMaxStepResultRunes = 4000
|
||||
defaultPlanExecuteKeepLastSteps = 8
|
||||
// Backward-compatible aliases for tests and existing references.
|
||||
planExecuteMaxStepResultRunes = defaultPlanExecuteMaxStepResultRunes
|
||||
planExecuteKeepLastSteps = defaultPlanExecuteKeepLastSteps
|
||||
)
|
||||
|
||||
func truncateRunesWithSuffix(s string, maxRunes int, suffix string) string {
|
||||
@@ -29,16 +34,26 @@ func truncateRunesWithSuffix(s string, maxRunes int, suffix string) string {
|
||||
|
||||
// capPlanExecuteExecutedSteps 折叠较早步骤、截断单步过长结果,供 prompt 使用。
|
||||
func capPlanExecuteExecutedSteps(steps []planexecute.ExecutedStep) []planexecute.ExecutedStep {
|
||||
return capPlanExecuteExecutedStepsWithConfig(steps, nil)
|
||||
}
|
||||
|
||||
func capPlanExecuteExecutedStepsWithConfig(steps []planexecute.ExecutedStep, mwCfg *config.MultiAgentEinoMiddlewareConfig) []planexecute.ExecutedStep {
|
||||
if len(steps) == 0 {
|
||||
return steps
|
||||
}
|
||||
maxStepResultRunes := defaultPlanExecuteMaxStepResultRunes
|
||||
keepLastSteps := defaultPlanExecuteKeepLastSteps
|
||||
if mwCfg != nil {
|
||||
maxStepResultRunes = mwCfg.PlanExecuteMaxStepResultRunesEffective()
|
||||
keepLastSteps = mwCfg.PlanExecuteKeepLastStepsEffective()
|
||||
}
|
||||
out := make([]planexecute.ExecutedStep, 0, len(steps)+1)
|
||||
start := 0
|
||||
if len(steps) > planExecuteKeepLastSteps {
|
||||
start = len(steps) - planExecuteKeepLastSteps
|
||||
if len(steps) > keepLastSteps {
|
||||
start = len(steps) - keepLastSteps
|
||||
var b strings.Builder
|
||||
b.WriteString(fmt.Sprintf("(上文已完成 %d 步;此处仅保留步骤标题以节省上下文,完整输出已省略。后续 %d 步仍保留正文。)\n",
|
||||
start, planExecuteKeepLastSteps))
|
||||
start, keepLastSteps))
|
||||
for i := 0; i < start; i++ {
|
||||
b.WriteString(fmt.Sprintf("- %s\n", steps[i].Step))
|
||||
}
|
||||
@@ -50,8 +65,8 @@ func capPlanExecuteExecutedSteps(steps []planexecute.ExecutedStep) []planexecute
|
||||
suffix := "\n…[step result truncated]"
|
||||
for i := start; i < len(steps); i++ {
|
||||
e := steps[i]
|
||||
if utf8.RuneCountInString(e.Result) > planExecuteMaxStepResultRunes {
|
||||
e.Result = truncateRunesWithSuffix(e.Result, planExecuteMaxStepResultRunes, suffix)
|
||||
if utf8.RuneCountInString(e.Result) > maxStepResultRunes {
|
||||
e.Result = truncateRunesWithSuffix(e.Result, maxStepResultRunes, suffix)
|
||||
}
|
||||
out = append(out, e)
|
||||
}
|
||||
|
||||
+111
-14
@@ -30,10 +30,10 @@ import (
|
||||
|
||||
// RunResult 与单 Agent 循环结果字段对齐,便于复用存储与 SSE 收尾逻辑。
|
||||
type RunResult struct {
|
||||
Response string
|
||||
MCPExecutionIDs []string
|
||||
LastReActInput string
|
||||
LastReActOutput string
|
||||
Response string
|
||||
MCPExecutionIDs []string
|
||||
LastAgentTraceInput string // 已序列化的消息带(JSON):原生循环或 Eino 均写入,供续跑/攻击链等恢复上下文
|
||||
LastAgentTraceOutput string // 本轮助手侧对外展示文本(摘要或最终回复)
|
||||
}
|
||||
|
||||
// toolCallPendingInfo tracks a tool_call emitted to the UI so we can later
|
||||
@@ -237,7 +237,7 @@ func RunDeepAgent(
|
||||
subMax = subDefaultIter
|
||||
}
|
||||
|
||||
subSumMw, err := newEinoSummarizationMiddleware(ctx, subModel, appCfg, logger)
|
||||
subSumMw, err := newEinoSummarizationMiddleware(ctx, subModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("子代理 %q summarization 中间件: %w", id, err)
|
||||
}
|
||||
@@ -257,11 +257,33 @@ func RunDeepAgent(
|
||||
subHandlers = append(subHandlers, einoSkillMW)
|
||||
}
|
||||
subHandlers = append(subHandlers, subSumMw)
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "sub_agent"); teleMw != nil {
|
||||
subHandlers = append(subHandlers, teleMw)
|
||||
}
|
||||
|
||||
subInstrFinal := injectToolNamesOnlyInstruction(ctx, instr, subTools)
|
||||
if logger != nil {
|
||||
subNames := collectToolNames(ctx, subTools)
|
||||
mountedNames := collectToolNames(ctx, subToolsForCfg)
|
||||
hasToolSearch := false
|
||||
for _, n := range subNames {
|
||||
if strings.EqualFold(strings.TrimSpace(n), "tool_search") {
|
||||
hasToolSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Info("eino tool-name injection",
|
||||
zap.String("scope", "sub_agent"),
|
||||
zap.String("agent", id),
|
||||
zap.Int("tool_names", len(subNames)),
|
||||
zap.Int("mounted_tool_names", len(mountedNames)),
|
||||
zap.Bool("has_tool_search", hasToolSearch),
|
||||
)
|
||||
}
|
||||
sa, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
||||
Name: id,
|
||||
Description: desc,
|
||||
Instruction: instr,
|
||||
Instruction: subInstrFinal,
|
||||
Model: subModel,
|
||||
ToolsConfig: adk.ToolsConfig{
|
||||
ToolsNodeConfig: compose.ToolsNodeConfig{
|
||||
@@ -289,7 +311,7 @@ func RunDeepAgent(
|
||||
return nil, fmt.Errorf("多代理主模型: %w", err)
|
||||
}
|
||||
|
||||
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, logger)
|
||||
mainSumMw, err := newEinoSummarizationMiddleware(ctx, mainModel, appCfg, &ma.EinoMiddleware, conversationID, logger)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("多代理主 summarization 中间件: %w", err)
|
||||
}
|
||||
@@ -313,6 +335,25 @@ func RunDeepAgent(
|
||||
orchDescription = d
|
||||
}
|
||||
}
|
||||
orchInstruction = injectToolNamesOnlyInstruction(ctx, orchInstruction, mainTools)
|
||||
if logger != nil {
|
||||
mainNames := collectToolNames(ctx, mainTools)
|
||||
mountedNames := collectToolNames(ctx, mainToolsForCfg)
|
||||
hasToolSearch := false
|
||||
for _, n := range mainNames {
|
||||
if strings.EqualFold(strings.TrimSpace(n), "tool_search") {
|
||||
hasToolSearch = true
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Info("eino tool-name injection",
|
||||
zap.String("scope", "orchestrator"),
|
||||
zap.String("orchestration", orchMode),
|
||||
zap.Int("tool_names", len(mainNames)),
|
||||
zap.Int("mounted_tool_names", len(mountedNames)),
|
||||
zap.Bool("has_tool_search", hasToolSearch),
|
||||
)
|
||||
}
|
||||
|
||||
supInstr := strings.TrimSpace(orchInstruction)
|
||||
if orchMode == "supervisor" {
|
||||
@@ -352,6 +393,9 @@ func RunDeepAgent(
|
||||
deepHandlers = append(deepHandlers, einoSkillMW)
|
||||
}
|
||||
deepHandlers = append(deepHandlers, mainSumMw)
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "deep_orchestrator"); teleMw != nil {
|
||||
deepHandlers = append(deepHandlers, teleMw)
|
||||
}
|
||||
|
||||
supHandlers := []adk.ChatModelAgentMiddleware{}
|
||||
if len(mainOrchestratorPre) > 0 {
|
||||
@@ -361,6 +405,9 @@ func RunDeepAgent(
|
||||
supHandlers = append(supHandlers, einoSkillMW)
|
||||
}
|
||||
supHandlers = append(supHandlers, mainSumMw)
|
||||
if teleMw := newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "supervisor_orchestrator"); teleMw != nil {
|
||||
supHandlers = append(supHandlers, teleMw)
|
||||
}
|
||||
|
||||
mainToolsCfg := adk.ToolsConfig{
|
||||
ToolsNodeConfig: compose.ToolsNodeConfig{
|
||||
@@ -399,10 +446,17 @@ func RunDeepAgent(
|
||||
ExecMaxIter: deepMaxIter,
|
||||
LoopMaxIter: ma.PlanExecuteLoopMaxIterations,
|
||||
AppCfg: appCfg,
|
||||
MwCfg: &ma.EinoMiddleware,
|
||||
ConversationID: conversationID,
|
||||
Logger: logger,
|
||||
ModelName: appCfg.OpenAI.Model,
|
||||
ExecPreMiddlewares: mainOrchestratorPre,
|
||||
SkillMiddleware: einoSkillMW,
|
||||
FilesystemMiddleware: peFsMw,
|
||||
PlannerReplannerRewriteHandlers: []adk.ChatModelAgentMiddleware{
|
||||
mainSumMw,
|
||||
newEinoModelInputTelemetryMiddleware(logger, appCfg.OpenAI.Model, conversationID, "plan_execute_planner_replanner_rewrite"),
|
||||
},
|
||||
})
|
||||
if perr != nil {
|
||||
return nil, perr
|
||||
@@ -468,7 +522,7 @@ func RunDeepAgent(
|
||||
da = dDeep
|
||||
}
|
||||
|
||||
baseMsgs := historyToMessages(history)
|
||||
baseMsgs := historyToMessages(history, appCfg, &ma.EinoMiddleware)
|
||||
baseMsgs = append(baseMsgs, schema.UserMessage(userMessage))
|
||||
|
||||
streamsMainAssistant := func(agent string) bool {
|
||||
@@ -505,34 +559,77 @@ func RunDeepAgent(
|
||||
}, baseMsgs)
|
||||
}
|
||||
|
||||
func historyToMessages(history []agent.ChatMessage) []adk.Message {
|
||||
func historyToMessages(history []agent.ChatMessage, appCfg *config.Config, mwCfg *config.MultiAgentEinoMiddlewareConfig) []adk.Message {
|
||||
if len(history) == 0 {
|
||||
return nil
|
||||
}
|
||||
// 放宽条数上限:跨轮历史交给 Eino Summarization(阈值对齐 openai.max_total_tokens)在调用模型前压缩,避免在入队前硬截断为 40 条。
|
||||
const maxHistoryMessages = 300
|
||||
// Keep a bounded tail first; then enforce a token budget.
|
||||
const maxHistoryMessages = 200
|
||||
start := 0
|
||||
if len(history) > maxHistoryMessages {
|
||||
start = len(history) - maxHistoryMessages
|
||||
}
|
||||
out := make([]adk.Message, 0, len(history[start:]))
|
||||
raw := make([]adk.Message, 0, len(history[start:]))
|
||||
for _, h := range history[start:] {
|
||||
switch h.Role {
|
||||
case "user":
|
||||
if strings.TrimSpace(h.Content) != "" {
|
||||
out = append(out, schema.UserMessage(h.Content))
|
||||
raw = append(raw, schema.UserMessage(h.Content))
|
||||
}
|
||||
case "assistant":
|
||||
if strings.TrimSpace(h.Content) == "" && len(h.ToolCalls) > 0 {
|
||||
continue
|
||||
}
|
||||
if strings.TrimSpace(h.Content) != "" {
|
||||
out = append(out, schema.AssistantMessage(h.Content, nil))
|
||||
raw = append(raw, schema.AssistantMessage(h.Content, nil))
|
||||
}
|
||||
default:
|
||||
continue
|
||||
}
|
||||
}
|
||||
if len(raw) == 0 {
|
||||
return raw
|
||||
}
|
||||
maxTotal := 120000
|
||||
modelName := "gpt-4o"
|
||||
if appCfg != nil {
|
||||
if appCfg.OpenAI.MaxTotalTokens > 0 {
|
||||
maxTotal = appCfg.OpenAI.MaxTotalTokens
|
||||
}
|
||||
if m := strings.TrimSpace(appCfg.OpenAI.Model); m != "" {
|
||||
modelName = m
|
||||
}
|
||||
}
|
||||
ratio := 0.35
|
||||
if mwCfg != nil {
|
||||
ratio = mwCfg.HistoryInputBudgetRatioEffective()
|
||||
}
|
||||
budget := int(float64(maxTotal) * ratio)
|
||||
if budget < 4096 {
|
||||
budget = 4096
|
||||
}
|
||||
tc := agent.NewTikTokenCounter()
|
||||
outRev := make([]adk.Message, 0, len(raw))
|
||||
used := 0
|
||||
for i := len(raw) - 1; i >= 0; i-- {
|
||||
msg := raw[i]
|
||||
n, err := tc.Count(modelName, string(msg.Role)+"\n"+msg.Content)
|
||||
if err != nil {
|
||||
n = (len(msg.Content) + 3) / 4
|
||||
}
|
||||
if n <= 0 {
|
||||
n = 1
|
||||
}
|
||||
if used+n > budget {
|
||||
break
|
||||
}
|
||||
used += n
|
||||
outRev = append(outRev, msg)
|
||||
}
|
||||
out := make([]adk.Message, 0, len(outRev))
|
||||
for i := len(outRev) - 1; i >= 0; i-- {
|
||||
out = append(out, outRev[i])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
|
||||
@@ -192,13 +192,13 @@ func convertOpenAIToClaude(payload interface{}) (*claudeRequest, error) {
|
||||
fnName, _ := fn["name"].(string)
|
||||
fnArgs, _ := fn["arguments"]
|
||||
|
||||
// 防御:缺少 name 或 id 的 tool_call 会被 Claude 拒绝
|
||||
if strings.TrimSpace(fnName) == "" {
|
||||
fnName = "unknown_function"
|
||||
}
|
||||
if strings.TrimSpace(tcID) == "" {
|
||||
tcID = fmt.Sprintf("call_%d", time.Now().UnixNano())
|
||||
}
|
||||
// 防御:缺少 name 或 id 的 tool_call 会被 Claude 拒绝
|
||||
if strings.TrimSpace(fnName) == "" {
|
||||
fnName = "unknown_function"
|
||||
}
|
||||
if strings.TrimSpace(tcID) == "" {
|
||||
tcID = fmt.Sprintf("call_%d", time.Now().UnixNano())
|
||||
}
|
||||
|
||||
var inputRaw json.RawMessage
|
||||
switch v := fnArgs.(type) {
|
||||
|
||||
+14
-14
@@ -281,9 +281,9 @@ func (c *Client) ChatCompletionStream(ctx context.Context, payload interface{},
|
||||
|
||||
// StreamToolCall 流式工具调用的累积结果(arguments 以字符串形式拼接,留给上层再解析为 JSON)。
|
||||
type StreamToolCall struct {
|
||||
Index int
|
||||
ID string
|
||||
Type string
|
||||
Index int
|
||||
ID string
|
||||
Type string
|
||||
FunctionName string
|
||||
FunctionArgsStr string
|
||||
}
|
||||
@@ -348,10 +348,10 @@ func (c *Client) ChatCompletionStreamWithToolCalls(
|
||||
Arguments string `json:"arguments,omitempty"`
|
||||
}
|
||||
type toolCallDelta struct {
|
||||
Index int `json:"index,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function toolCallFunctionDelta `json:"function,omitempty"`
|
||||
Index int `json:"index,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Type string `json:"type,omitempty"`
|
||||
Function toolCallFunctionDelta `json:"function,omitempty"`
|
||||
}
|
||||
type streamDelta2 struct {
|
||||
Content string `json:"content,omitempty"`
|
||||
@@ -371,10 +371,10 @@ func (c *Client) ChatCompletionStreamWithToolCalls(
|
||||
}
|
||||
|
||||
type toolCallAccum struct {
|
||||
id string
|
||||
typ string
|
||||
name string
|
||||
args strings.Builder
|
||||
id string
|
||||
typ string
|
||||
name string
|
||||
args strings.Builder
|
||||
}
|
||||
toolCallAccums := make(map[int]*toolCallAccum)
|
||||
|
||||
@@ -475,9 +475,9 @@ func (c *Client) ChatCompletionStreamWithToolCalls(
|
||||
for _, idx := range indices {
|
||||
acc := toolCallAccums[idx]
|
||||
tc := StreamToolCall{
|
||||
Index: idx,
|
||||
ID: acc.id,
|
||||
Type: acc.typ,
|
||||
Index: idx,
|
||||
ID: acc.id,
|
||||
Type: acc.typ,
|
||||
FunctionName: acc.name,
|
||||
FunctionArgsStr: acc.args.String(),
|
||||
}
|
||||
|
||||
@@ -19,11 +19,11 @@ import (
|
||||
func setupTestExecutor(t *testing.T) (*Executor, *mcp.Server) {
|
||||
logger := zap.NewNop()
|
||||
mcpServer := mcp.NewServer(logger)
|
||||
|
||||
|
||||
cfg := &config.SecurityConfig{
|
||||
Tools: []config.ToolConfig{},
|
||||
}
|
||||
|
||||
|
||||
executor := NewExecutor(cfg, mcpServer, logger)
|
||||
return executor, mcpServer
|
||||
}
|
||||
@@ -32,12 +32,12 @@ func setupTestExecutor(t *testing.T) (*Executor, *mcp.Server) {
|
||||
func setupTestStorage(t *testing.T) *storage.FileResultStorage {
|
||||
tmpDir := filepath.Join(os.TempDir(), "test_executor_storage_"+time.Now().Format("20060102_150405"))
|
||||
logger := zap.NewNop()
|
||||
|
||||
|
||||
storage, err := storage.NewFileResultStorage(tmpDir, logger)
|
||||
if err != nil {
|
||||
t.Fatalf("创建测试存储失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
return storage
|
||||
}
|
||||
|
||||
@@ -45,46 +45,46 @@ func TestExecutor_ExecuteInternalTool_QueryExecutionResult(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
testStorage := setupTestStorage(t)
|
||||
executor.SetResultStorage(testStorage)
|
||||
|
||||
|
||||
// 准备测试数据
|
||||
executionID := "test_exec_001"
|
||||
toolName := "nmap_scan"
|
||||
result := "Line 1: Port 22 open\nLine 2: Port 80 open\nLine 3: Port 443 open\nLine 4: error occurred"
|
||||
|
||||
|
||||
// 保存测试结果
|
||||
err := testStorage.SaveResult(executionID, toolName, result)
|
||||
if err != nil {
|
||||
t.Fatalf("保存测试结果失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
|
||||
// 测试1: 基本查询(第一页)
|
||||
args := map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
"page": float64(1),
|
||||
"limit": float64(2),
|
||||
}
|
||||
|
||||
|
||||
toolResult, err := executor.executeQueryExecutionResult(ctx, args)
|
||||
if err != nil {
|
||||
t.Fatalf("执行查询失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if toolResult.IsError {
|
||||
t.Fatalf("查询应该成功,但返回了错误: %s", toolResult.Content[0].Text)
|
||||
}
|
||||
|
||||
|
||||
// 验证结果包含预期内容
|
||||
resultText := toolResult.Content[0].Text
|
||||
if !strings.Contains(resultText, executionID) {
|
||||
t.Errorf("结果中应该包含执行ID: %s", executionID)
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(resultText, "第 1/") {
|
||||
t.Errorf("结果中应该包含分页信息")
|
||||
}
|
||||
|
||||
|
||||
// 测试2: 搜索功能
|
||||
args2 := map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
@@ -92,21 +92,21 @@ func TestExecutor_ExecuteInternalTool_QueryExecutionResult(t *testing.T) {
|
||||
"page": float64(1),
|
||||
"limit": float64(10),
|
||||
}
|
||||
|
||||
|
||||
toolResult2, err := executor.executeQueryExecutionResult(ctx, args2)
|
||||
if err != nil {
|
||||
t.Fatalf("执行搜索失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if toolResult2.IsError {
|
||||
t.Fatalf("搜索应该成功,但返回了错误: %s", toolResult2.Content[0].Text)
|
||||
}
|
||||
|
||||
|
||||
resultText2 := toolResult2.Content[0].Text
|
||||
if !strings.Contains(resultText2, "error") {
|
||||
t.Errorf("搜索结果中应该包含关键词: error")
|
||||
}
|
||||
|
||||
|
||||
// 测试3: 过滤功能
|
||||
args3 := map[string]interface{}{
|
||||
"execution_id": executionID,
|
||||
@@ -114,46 +114,46 @@ func TestExecutor_ExecuteInternalTool_QueryExecutionResult(t *testing.T) {
|
||||
"page": float64(1),
|
||||
"limit": float64(10),
|
||||
}
|
||||
|
||||
|
||||
toolResult3, err := executor.executeQueryExecutionResult(ctx, args3)
|
||||
if err != nil {
|
||||
t.Fatalf("执行过滤失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if toolResult3.IsError {
|
||||
t.Fatalf("过滤应该成功,但返回了错误: %s", toolResult3.Content[0].Text)
|
||||
}
|
||||
|
||||
|
||||
resultText3 := toolResult3.Content[0].Text
|
||||
if !strings.Contains(resultText3, "Port") {
|
||||
t.Errorf("过滤结果中应该包含关键词: Port")
|
||||
}
|
||||
|
||||
|
||||
// 测试4: 缺少必需参数
|
||||
args4 := map[string]interface{}{
|
||||
"page": float64(1),
|
||||
}
|
||||
|
||||
|
||||
toolResult4, err := executor.executeQueryExecutionResult(ctx, args4)
|
||||
if err != nil {
|
||||
t.Fatalf("执行查询失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if !toolResult4.IsError {
|
||||
t.Fatal("缺少execution_id应该返回错误")
|
||||
}
|
||||
|
||||
|
||||
// 测试5: 不存在的执行ID
|
||||
args5 := map[string]interface{}{
|
||||
"execution_id": "nonexistent_id",
|
||||
"page": float64(1),
|
||||
}
|
||||
|
||||
|
||||
toolResult5, err := executor.executeQueryExecutionResult(ctx, args5)
|
||||
if err != nil {
|
||||
t.Fatalf("执行查询失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if !toolResult5.IsError {
|
||||
t.Fatal("不存在的执行ID应该返回错误")
|
||||
}
|
||||
@@ -161,22 +161,22 @@ func TestExecutor_ExecuteInternalTool_QueryExecutionResult(t *testing.T) {
|
||||
|
||||
func TestExecutor_ExecuteInternalTool_UnknownTool(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"test": "value",
|
||||
}
|
||||
|
||||
|
||||
// 测试未知的内部工具类型
|
||||
toolResult, err := executor.executeInternalTool(ctx, "unknown_tool", "internal:unknown_tool", args)
|
||||
if err != nil {
|
||||
t.Fatalf("执行内部工具失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if !toolResult.IsError {
|
||||
t.Fatal("未知的工具类型应该返回错误")
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(toolResult.Content[0].Text, "未知的内部工具类型") {
|
||||
t.Errorf("错误消息应该包含'未知的内部工具类型'")
|
||||
}
|
||||
@@ -185,21 +185,21 @@ func TestExecutor_ExecuteInternalTool_UnknownTool(t *testing.T) {
|
||||
func TestExecutor_ExecuteInternalTool_NoStorage(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
// 不设置存储,测试未初始化的情况
|
||||
|
||||
|
||||
ctx := context.Background()
|
||||
args := map[string]interface{}{
|
||||
"execution_id": "test_id",
|
||||
}
|
||||
|
||||
|
||||
toolResult, err := executor.executeQueryExecutionResult(ctx, args)
|
||||
if err != nil {
|
||||
t.Fatalf("执行查询失败: %v", err)
|
||||
}
|
||||
|
||||
|
||||
if !toolResult.IsError {
|
||||
t.Fatal("未初始化的存储应该返回错误")
|
||||
}
|
||||
|
||||
|
||||
if !strings.Contains(toolResult.Content[0].Text, "结果存储未初始化") {
|
||||
t.Errorf("错误消息应该包含'结果存储未初始化'")
|
||||
}
|
||||
@@ -207,7 +207,7 @@ func TestExecutor_ExecuteInternalTool_NoStorage(t *testing.T) {
|
||||
|
||||
func TestPaginateLines(t *testing.T) {
|
||||
lines := []string{"Line 1", "Line 2", "Line 3", "Line 4", "Line 5"}
|
||||
|
||||
|
||||
// 测试第一页
|
||||
page := paginateLines(lines, 1, 2)
|
||||
if page.Page != 1 {
|
||||
@@ -225,7 +225,7 @@ func TestPaginateLines(t *testing.T) {
|
||||
if len(page.Lines) != 2 {
|
||||
t.Errorf("第一页行数不匹配。期望: 2, 实际: %d", len(page.Lines))
|
||||
}
|
||||
|
||||
|
||||
// 测试第二页
|
||||
page2 := paginateLines(lines, 2, 2)
|
||||
if len(page2.Lines) != 2 {
|
||||
@@ -234,13 +234,13 @@ func TestPaginateLines(t *testing.T) {
|
||||
if page2.Lines[0] != "Line 3" {
|
||||
t.Errorf("第二页第一行不匹配。期望: Line 3, 实际: %s", page2.Lines[0])
|
||||
}
|
||||
|
||||
|
||||
// 测试最后一页
|
||||
page3 := paginateLines(lines, 3, 2)
|
||||
if len(page3.Lines) != 1 {
|
||||
t.Errorf("第三页行数不匹配。期望: 1, 实际: %d", len(page3.Lines))
|
||||
}
|
||||
|
||||
|
||||
// 测试超出范围的页码(应该返回最后一页)
|
||||
page4 := paginateLines(lines, 4, 2)
|
||||
if page4.Page != 3 {
|
||||
@@ -249,13 +249,13 @@ func TestPaginateLines(t *testing.T) {
|
||||
if len(page4.Lines) != 1 {
|
||||
t.Errorf("最后一页应该只有1行。实际: %d行", len(page4.Lines))
|
||||
}
|
||||
|
||||
|
||||
// 测试无效页码(小于1)
|
||||
page0 := paginateLines(lines, 0, 2)
|
||||
if page0.Page != 1 {
|
||||
t.Errorf("无效页码应该被修正为1。实际: %d", page0.Page)
|
||||
}
|
||||
|
||||
|
||||
// 测试空列表
|
||||
emptyPage := paginateLines([]string{}, 1, 10)
|
||||
if emptyPage.TotalLines != 0 {
|
||||
@@ -265,4 +265,3 @@ func TestPaginateLines(t *testing.T) {
|
||||
t.Errorf("空列表应该返回空结果。实际: %d行", len(emptyPage.Lines))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,10 +16,10 @@ type rateLimitEntry struct {
|
||||
|
||||
// RateLimiter 基于 IP 的滑动窗口速率限制器
|
||||
type RateLimiter struct {
|
||||
mu sync.Mutex
|
||||
entries map[string]*rateLimitEntry
|
||||
limit int // 窗口内允许的最大请求数
|
||||
window time.Duration // 窗口时长
|
||||
mu sync.Mutex
|
||||
entries map[string]*rateLimitEntry
|
||||
limit int // 窗口内允许的最大请求数
|
||||
window time.Duration // 窗口时长
|
||||
}
|
||||
|
||||
// NewRateLimiter 创建速率限制器
|
||||
|
||||
@@ -162,4 +162,3 @@ func truncateRunes(s string, max int) string {
|
||||
}
|
||||
return string(r[:max]) + "…"
|
||||
}
|
||||
|
||||
|
||||
@@ -49,12 +49,12 @@ func ParseSkillMD(raw []byte) (*SkillManifest, string, error) {
|
||||
}
|
||||
|
||||
type skillFrontMatterExport struct {
|
||||
Name string `yaml:"name"`
|
||||
Description string `yaml:"description"`
|
||||
License string `yaml:"license,omitempty"`
|
||||
Compatibility string `yaml:"compatibility,omitempty"`
|
||||
Metadata map[string]any `yaml:"metadata,omitempty"`
|
||||
AllowedTools string `yaml:"allowed-tools,omitempty"`
|
||||
Name string `yaml:"name"`
|
||||
Description string `yaml:"description"`
|
||||
License string `yaml:"license,omitempty"`
|
||||
Compatibility string `yaml:"compatibility,omitempty"`
|
||||
Metadata map[string]any `yaml:"metadata,omitempty"`
|
||||
AllowedTools string `yaml:"allowed-tools,omitempty"`
|
||||
}
|
||||
|
||||
// BuildSkillMD serializes SKILL.md per agentskills.io.
|
||||
|
||||
@@ -9,10 +9,10 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
maxPackageFiles = 4000
|
||||
maxPackageDepth = 24
|
||||
maxScriptsDepth = 24
|
||||
defaultMaxRead = 10 << 20
|
||||
maxPackageFiles = 4000
|
||||
maxPackageDepth = 24
|
||||
maxScriptsDepth = 24
|
||||
defaultMaxRead = 10 << 20
|
||||
)
|
||||
|
||||
// SafeRelPath resolves rel inside root (no ..).
|
||||
|
||||
@@ -4472,6 +4472,25 @@ header {
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
|
||||
.tool-resident-toggle {
|
||||
display: inline-flex;
|
||||
align-items: center;
|
||||
gap: 4px;
|
||||
font-size: 0.75rem;
|
||||
color: var(--text-secondary);
|
||||
border: 1px solid var(--border-color);
|
||||
border-radius: 12px;
|
||||
padding: 2px 8px;
|
||||
background: var(--bg-secondary);
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
.tool-resident-toggle input[type="checkbox"] {
|
||||
width: 14px;
|
||||
height: 14px;
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
.external-tool-badge {
|
||||
display: inline-flex;
|
||||
align-items: center;
|
||||
|
||||
@@ -576,6 +576,10 @@
|
||||
"addExternal": "Add external MCP",
|
||||
"toolConfig": "MCP tool config",
|
||||
"saveToolConfig": "Save tool config",
|
||||
"alwaysVisibleLabel": "Pinned",
|
||||
"alwaysVisibleHint": "Always keep visible in Tool Search results",
|
||||
"alwaysVisibleBuiltinLabel": "Builtin default",
|
||||
"alwaysVisibleBuiltinHint": "Backend builtin tool is pinned by default and cannot be disabled",
|
||||
"externalConfig": "External MCP config",
|
||||
"loadingTools": "Loading tools...",
|
||||
"loadToolsTimeout": "Tools load timeout. External MCP may be slow. Click Refresh to retry or check connection.",
|
||||
|
||||
@@ -576,6 +576,10 @@
|
||||
"addExternal": "添加外部MCP",
|
||||
"toolConfig": "MCP 工具配置",
|
||||
"saveToolConfig": "保存工具配置",
|
||||
"alwaysVisibleLabel": "常驻",
|
||||
"alwaysVisibleHint": "始终常驻在 Tool Search 可见列表(不被 tool_search 隐藏)",
|
||||
"alwaysVisibleBuiltinLabel": "内置默认",
|
||||
"alwaysVisibleBuiltinHint": "后端内置工具默认常驻,不可关闭",
|
||||
"externalConfig": "外部 MCP 配置",
|
||||
"loadingTools": "正在加载工具列表...",
|
||||
"loadToolsTimeout": "加载工具列表超时,可能是外部MCP连接较慢。请点击\"刷新\"按钮重试,或检查外部MCP连接状态。",
|
||||
|
||||
+26
-14
@@ -301,26 +301,38 @@ async function initPage(pageId) {
|
||||
break;
|
||||
case 'mcp-management':
|
||||
// 初始化MCP管理
|
||||
const startLoadMcpTools = () => {
|
||||
// 加载工具列表(MCP工具配置已移到MCP管理页面)
|
||||
// 使用异步加载,避免阻塞页面渲染
|
||||
if (typeof loadToolsList === 'function') {
|
||||
// 确保工具分页设置已初始化
|
||||
if (typeof getToolsPageSize === 'function' && typeof toolsPagination !== 'undefined') {
|
||||
toolsPagination.pageSize = getToolsPageSize();
|
||||
}
|
||||
// 延迟加载,让页面先渲染
|
||||
setTimeout(() => {
|
||||
loadToolsList(1, '').catch(err => {
|
||||
console.error('加载工具列表失败:', err);
|
||||
});
|
||||
}, 100);
|
||||
}
|
||||
};
|
||||
// 先拉取全局配置,确保 tool_search 常驻状态按后端生效集合展示
|
||||
if (typeof loadConfig === 'function') {
|
||||
loadConfig(false)
|
||||
.catch(err => {
|
||||
console.warn('加载配置失败(将继续加载工具列表):', err);
|
||||
})
|
||||
.finally(startLoadMcpTools);
|
||||
} else {
|
||||
startLoadMcpTools();
|
||||
}
|
||||
// 先加载外部MCP列表(快速),然后加载工具列表
|
||||
if (typeof loadExternalMCPs === 'function') {
|
||||
loadExternalMCPs().catch(err => {
|
||||
console.warn('加载外部MCP列表失败:', err);
|
||||
});
|
||||
}
|
||||
// 加载工具列表(MCP工具配置已移到MCP管理页面)
|
||||
// 使用异步加载,避免阻塞页面渲染
|
||||
if (typeof loadToolsList === 'function') {
|
||||
// 确保工具分页设置已初始化
|
||||
if (typeof getToolsPageSize === 'function' && typeof toolsPagination !== 'undefined') {
|
||||
toolsPagination.pageSize = getToolsPageSize();
|
||||
}
|
||||
// 延迟加载,让页面先渲染
|
||||
setTimeout(() => {
|
||||
loadToolsList(1, '').catch(err => {
|
||||
console.error('加载工具列表失败:', err);
|
||||
});
|
||||
}, 100);
|
||||
}
|
||||
break;
|
||||
case 'vulnerabilities':
|
||||
// 初始化漏洞管理页面
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
// 设置相关功能
|
||||
let currentConfig = null;
|
||||
let allTools = [];
|
||||
let alwaysVisibleToolNames = new Set();
|
||||
let alwaysVisibleBuiltinToolNames = new Set();
|
||||
// 全局工具状态映射,用于保存用户在所有页面的修改
|
||||
// key: 唯一工具标识符(toolKey),value: { enabled: boolean, is_external: boolean, external_mcp: string }
|
||||
let toolStateMap = new Map();
|
||||
@@ -100,6 +102,14 @@ async function loadConfig(loadTools = true) {
|
||||
}
|
||||
|
||||
currentConfig = await response.json();
|
||||
const alwaysVisibleList = currentConfig?.multi_agent?.tool_search_always_visible_effective_tools;
|
||||
const alwaysVisibleConfigured = currentConfig?.multi_agent?.tool_search_always_visible_tools;
|
||||
alwaysVisibleToolNames = new Set(Array.isArray(alwaysVisibleList) ? alwaysVisibleList.filter(Boolean) : []);
|
||||
alwaysVisibleBuiltinToolNames = new Set(
|
||||
alwaysVisibleToolNames.size > 0 && Array.isArray(alwaysVisibleConfigured)
|
||||
? Array.from(alwaysVisibleToolNames).filter(name => !alwaysVisibleConfigured.includes(name))
|
||||
: []
|
||||
);
|
||||
|
||||
// 填充OpenAI配置
|
||||
const providerEl = document.getElementById('openai-provider');
|
||||
@@ -498,6 +508,8 @@ function renderToolsList() {
|
||||
is_external: tool.is_external || false,
|
||||
external_mcp: tool.external_mcp || ''
|
||||
};
|
||||
const alwaysVisibleChecked = alwaysVisibleToolNames.has(tool.name);
|
||||
const alwaysVisibleLocked = alwaysVisibleBuiltinToolNames.has(tool.name);
|
||||
|
||||
// 外部工具标签,显示来源信息(可点击跳转到对应 MCP 卡片)
|
||||
let externalBadge = '';
|
||||
@@ -521,6 +533,11 @@ function renderToolsList() {
|
||||
<div class="tool-item-name">
|
||||
${escapeHtml(tool.name)}
|
||||
${externalBadge}
|
||||
<label class="tool-resident-toggle" title="${typeof window.t === 'function' ? window.t('mcp.alwaysVisibleHint') : '始终常驻在 Tool Search 可见列表'}" onclick="event.stopPropagation()">
|
||||
<input type="checkbox" ${alwaysVisibleChecked ? 'checked' : ''} ${alwaysVisibleLocked ? 'disabled' : ''} onchange="handleToolAlwaysVisibleChange('${escapeHtml(tool.name)}', this.checked)" />
|
||||
<span>${typeof window.t === 'function' ? window.t('mcp.alwaysVisibleLabel') : '常驻'}</span>
|
||||
</label>
|
||||
${alwaysVisibleLocked ? `<span class="external-tool-badge" title="${typeof window.t === 'function' ? window.t('mcp.alwaysVisibleBuiltinHint') : '后端内置工具默认常驻,不可关闭'}">${typeof window.t === 'function' ? window.t('mcp.alwaysVisibleBuiltinLabel') : '内置默认'}</span>` : ''}
|
||||
<span class="tool-expand-icon">▶</span>
|
||||
</div>
|
||||
<div class="tool-item-desc">${escapeHtml(tool.description || (typeof window.t === 'function' ? window.t('mcp.noDescription') : '无描述'))}</div>
|
||||
@@ -716,6 +733,16 @@ function handleToolCheckboxChange(toolKey, enabled) {
|
||||
updateToolsStats();
|
||||
}
|
||||
|
||||
function handleToolAlwaysVisibleChange(toolName, alwaysVisible) {
|
||||
const name = (toolName || '').trim();
|
||||
if (!name) return;
|
||||
if (alwaysVisible) {
|
||||
alwaysVisibleToolNames.add(name);
|
||||
} else {
|
||||
alwaysVisibleToolNames.delete(name);
|
||||
}
|
||||
}
|
||||
|
||||
// 全选工具
|
||||
function selectAllTools() {
|
||||
document.querySelectorAll('#tools-list input[type="checkbox"]').forEach(checkbox => {
|
||||
@@ -886,9 +913,11 @@ async function updateToolsStats() {
|
||||
}
|
||||
|
||||
const tStats = typeof window.t === 'function' ? window.t : (k) => k;
|
||||
const pinnedCount = alwaysVisibleToolNames.size;
|
||||
statsEl.innerHTML = `
|
||||
<span title="${tStats('mcp.currentPageEnabled')}">✅ ${tStats('mcp.currentPageEnabled')}: <strong>${currentPageEnabled}</strong> / ${currentPageTotal}</span>
|
||||
<span title="${tStats('mcp.totalEnabled')}">📊 ${tStats('mcp.totalEnabled')}: <strong>${totalEnabled}</strong> / ${totalTools}</span>
|
||||
<span title="${tStats('mcp.alwaysVisibleHint')}">📌 ${tStats('mcp.alwaysVisibleLabel')}: <strong>${pinnedCount}</strong></span>
|
||||
`;
|
||||
}
|
||||
|
||||
@@ -1230,6 +1259,13 @@ async function saveToolsConfig() {
|
||||
const config = {
|
||||
openai: currentConfig.openai || {},
|
||||
agent: currentConfig.agent || {},
|
||||
multi_agent: {
|
||||
enabled: currentConfig?.multi_agent?.enabled === true,
|
||||
robot_use_multi_agent: currentConfig?.multi_agent?.robot_use_multi_agent === true,
|
||||
batch_use_multi_agent: currentConfig?.multi_agent?.batch_use_multi_agent === true,
|
||||
plan_execute_loop_max_iterations: Number(currentConfig?.multi_agent?.plan_execute_loop_max_iterations || 0),
|
||||
tool_search_always_visible_tools: Array.from(alwaysVisibleToolNames).filter(name => !alwaysVisibleBuiltinToolNames.has(name))
|
||||
},
|
||||
tools: []
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user