CalvinBackup/CyberStrikeAI
1
0
Fork 0
mirror of https://github.com/Ed1s0nZ/CyberStrikeAI.git synced 2026-08-01 08:37:41 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
公明
2026-07-21 20:57:07 +08:00
committed by GitHub
parent f30e8e7ca6
commit 07229f0fde
7 changed files with 141 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/handler/agent.go
+8 -7
View File
@@ -212,7 +212,7 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
if h == nil || conversationID == "" || h.tasks == nil {
return
}
h.cancelActiveMCPToolForConversation(conversationID)
h.cancelRunningMCPToolsForConversation(conversationID)
h.tasks.AbortActiveEinoExecute(conversationID, "")
if ok, err := h.tasks.CancelTask(conversationID, ErrTaskCancelled); ok {
h.logger.Info("已取消会话运行中任务", zap.String("conversationId", conversationID))
@@ -221,12 +221,13 @@ func (h *AgentHandler) CancelRunningTaskForConversation(conversationID string) {
}
}
func (h *AgentHandler) cancelActiveMCPToolForConversation(conversationID string) {
if h == nil || h.tasks == nil || h.agent == nil {
func (h *AgentHandler) cancelRunningMCPToolsForConversation(conversationID string) {
if h == nil || h.agent == nil {
return
}
if execID := h.tasks.ActiveMCPExecutionID(conversationID); execID != "" {
h.agent.CancelMCPToolExecutionWithNote(execID, "")
n := h.agent.CancelRunningMCPToolsForConversation(conversationID, "会话已结束，自动终止仍在运行的工具")
if n > 0 && h.logger != nil {
h.logger.Info("已终止会话仍在运行的 MCP 工具", zap.String("conversationId", conversationID), zap.Int("count", n))
}
}
@@ -266,7 +267,7 @@ func NewAgentHandler(agent *agent.Agent, db *database.DB, cfg *config.Config, lo
batchCronParser: cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow | cron.Descriptor),
auditLLM: openai.NewClient(llmCfg, llmHTTP, logger),
}
tm.SetToolCanceler(handler.cancelActiveMCPToolForConversation)
tm.SetToolCanceler(handler.cancelRunningMCPToolsForConversation)
if err := handler.hitlManager.EnsureSchema(); err != nil {
logger.Warn("初始化 HITL 表失败", zap.Error(err))
}
@@ -1553,7 +1554,7 @@ func (h *AgentHandler) CancelAgentLoop(c *gin.Context) {
var cause error = ErrTaskCancelled
msg := "已提交取消请求，任务将在当前步骤完成后停止。"
h.cancelActiveMCPToolForConversation(req.ConversationID)
h.cancelRunningMCPToolsForConversation(req.ConversationID)
h.tasks.AbortActiveEinoExecute(req.ConversationID, "")
ok, err := h.tasks.CancelTask(req.ConversationID, cause)
if err != nil {
internal/handler/config.go
+64 -3
View File
@@ -725,9 +725,14 @@ type UpdateConfigRequest struct {
// AgentConfigUpdate 用于 PATCH /api/config 的 agent 段：仅 JSON 中出现的字段（指针非 nil）覆盖内存配置。
// 避免旧版「整包替换 *AgentConfig」时，未传的整型字段被反序列化为 0 误覆盖（例如 tool_timeout_minutes 变成 0）。
type AgentConfigUpdate struct {
MaxIterations *int `json:"max_iterations,omitempty"`
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
SystemPromptPath *string `json:"system_prompt_path,omitempty"`
MaxIterations *int `json:"max_iterations,omitempty"`
ToolTimeoutMinutes *int `json:"tool_timeout_minutes,omitempty"`
ToolWaitTimeoutSeconds *int `json:"tool_wait_timeout_seconds,omitempty"`
ExternalMCPMaxConcurrentPerServer *int `json:"external_mcp_max_concurrent_per_server,omitempty"`
ExternalMCPMaxConcurrentTotal *int `json:"external_mcp_max_concurrent_total,omitempty"`
ExternalMCPCircuitFailureThreshold *int `json:"external_mcp_circuit_failure_threshold,omitempty"`
ExternalMCPCircuitCooldownSeconds *int `json:"external_mcp_circuit_cooldown_seconds,omitempty"`
SystemPromptPath *string `json:"system_prompt_path,omitempty"`
}
func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
@@ -740,6 +745,21 @@ func applyAgentConfigUpdate(dst *config.AgentConfig, src *AgentConfigUpdate) {
if src.ToolTimeoutMinutes != nil {
dst.ToolTimeoutMinutes = *src.ToolTimeoutMinutes
}
if src.ToolWaitTimeoutSeconds != nil {
dst.ToolWaitTimeoutSeconds = *src.ToolWaitTimeoutSeconds
}
if src.ExternalMCPMaxConcurrentPerServer != nil {
dst.ExternalMCPMaxConcurrentPerServer = *src.ExternalMCPMaxConcurrentPerServer
}
if src.ExternalMCPMaxConcurrentTotal != nil {
dst.ExternalMCPMaxConcurrentTotal = *src.ExternalMCPMaxConcurrentTotal
}
if src.ExternalMCPCircuitFailureThreshold != nil {
dst.ExternalMCPCircuitFailureThreshold = *src.ExternalMCPCircuitFailureThreshold
}
if src.ExternalMCPCircuitCooldownSeconds != nil {
dst.ExternalMCPCircuitCooldownSeconds = *src.ExternalMCPCircuitCooldownSeconds
}
if src.SystemPromptPath != nil {
dst.SystemPromptPath = *src.SystemPromptPath
}
@@ -815,12 +835,32 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
h.logger.Info("更新Agent配置",
zap.Int("max_iterations", h.config.Agent.MaxIterations),
zap.Int("tool_timeout_minutes", h.config.Agent.ToolTimeoutMinutes),
zap.Int("tool_wait_timeout_seconds", h.config.Agent.ToolWaitTimeoutSeconds),
zap.Int("external_mcp_max_concurrent_per_server", h.config.Agent.ExternalMCPMaxConcurrentPerServer),
zap.Int("external_mcp_max_concurrent_total", h.config.Agent.ExternalMCPMaxConcurrentTotal),
zap.Int("external_mcp_circuit_failure_threshold", h.config.Agent.ExternalMCPCircuitFailureThreshold),
zap.Int("external_mcp_circuit_cooldown_seconds", h.config.Agent.ExternalMCPCircuitCooldownSeconds),
)
if h.agent != nil && req.Agent.MaxIterations != nil {
h.agent.UpdateMaxIterations(h.config.Agent.MaxIterations)
}
if h.executor != nil {
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
}
if h.mcpServer != nil {
h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes)
h.mcpServer.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
h.mcpServer.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
}
if h.externalMCPMgr != nil {
h.externalMCPMgr.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
h.externalMCPMgr.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
MaxConcurrentPerServer: h.config.Agent.ExternalMCPMaxConcurrentPerServer,
MaxConcurrentTotal: h.config.Agent.ExternalMCPMaxConcurrentTotal,
CircuitFailureThreshold: h.config.Agent.ExternalMCPCircuitFailureThreshold,
CircuitCooldown: time.Duration(h.config.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
})
}
}
@@ -1471,6 +1511,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
h.mcpServer.ClearTools()
// 重新注册安全工具
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.executor.RegisterTools(h.mcpServer)
// 重新注册漏洞记录工具（内置工具，必须注册）
@@ -1542,6 +1583,21 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
}
if h.mcpServer != nil {
h.mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(h.config.Agent.ToolTimeoutMinutes)
h.mcpServer.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
h.mcpServer.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
}
if h.executor != nil {
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
}
if h.externalMCPMgr != nil {
h.externalMCPMgr.ConfigureToolWaitTimeoutSeconds(h.config.Agent.ToolWaitTimeoutSeconds)
h.externalMCPMgr.ConfigureToolResultMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
h.externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{
MaxConcurrentPerServer: h.config.Agent.ExternalMCPMaxConcurrentPerServer,
MaxConcurrentTotal: h.config.Agent.ExternalMCPMaxConcurrentTotal,
CircuitFailureThreshold: h.config.Agent.ExternalMCPCircuitFailureThreshold,
CircuitCooldown: time.Duration(h.config.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second,
})
}
// 更新AttackChainHandler的OpenAI配置
@@ -1731,6 +1787,11 @@ func updateAgentConfig(doc *yaml.Node, agent config.AgentConfig) {
agentNode := ensureMap(root, "agent")
setIntInMap(agentNode, "max_iterations", agent.MaxIterations)
setIntInMap(agentNode, "tool_timeout_minutes", agent.ToolTimeoutMinutes)
setIntInMap(agentNode, "tool_wait_timeout_seconds", agent.ToolWaitTimeoutSeconds)
setIntInMap(agentNode, "external_mcp_max_concurrent_per_server", agent.ExternalMCPMaxConcurrentPerServer)
setIntInMap(agentNode, "external_mcp_max_concurrent_total", agent.ExternalMCPMaxConcurrentTotal)
setIntInMap(agentNode, "external_mcp_circuit_failure_threshold", agent.ExternalMCPCircuitFailureThreshold)
setIntInMap(agentNode, "external_mcp_circuit_cooldown_seconds", agent.ExternalMCPCircuitCooldownSeconds)
setStringInMap(agentNode, "system_prompt_path", agent.SystemPromptPath)
}
internal/handler/conversation.go
+1
View File
@@ -373,6 +373,7 @@ func summarizeProcessDetailData(eventType string, data interface{}) interface{}
"success": true, "isError": true, "executionId": true,
"einoAgent": true, "einoRole": true, "einoScope": true, "orchestration": true,
"agentFacing": true,
"status": true, "modelFacingIsError": true, "resultPreview": true,
}
out := make(map[string]interface{}, len(allow)+1)
for k, v := range m {
internal/handler/monitor.go
+13 -4
View File
@@ -167,7 +167,7 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
stats[exec.ToolName] = stat
}
stat.TotalCalls++
if exec.Status == "failed" || exec.Status == "cancelled" {
if monitorStatusCountsAsFailed(exec.Status) {
stat.FailedCalls++
} else if exec.Status == "completed" {
stat.SuccessCalls++
@@ -180,6 +180,15 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
return summarizeToolStats(stats, topN)
}
func monitorStatusCountsAsFailed(status string) bool {
switch strings.TrimSpace(strings.ToLower(status)) {
case "failed", "cancelled", "hard_timeout", "orphaned":
return true
default:
return false
}
}
func (h *MonitorHandler) monitorRetentionDays() int {
if h.monitorRetention != nil {
return h.monitorRetention.RetentionDays()
@@ -813,7 +822,7 @@ func (h *MonitorHandler) loadCallsTimeline(cfg callsTimelineConfig) []CallsTimel
key := truncateToBucket(exec.StartTime, cfg.bucketSize, cfg.dailyBuckets)
entry := bucketMap[key]
entry.total++
if exec.Status == "failed" || exec.Status == "cancelled" {
if monitorStatusCountsAsFailed(exec.Status) {
entry.failed++
}
bucketMap[key] = entry
@@ -876,7 +885,7 @@ func (h *MonitorHandler) DeleteExecution(c *gin.Context) {
totalCalls := 1
successCalls := 0
failedCalls := 0
if exec.Status == "failed" || exec.Status == "cancelled" {
if monitorStatusCountsAsFailed(exec.Status) {
failedCalls = 1
} else if exec.Status == "completed" {
successCalls = 1
@@ -951,7 +960,7 @@ func (h *MonitorHandler) DeleteExecutions(c *gin.Context) {
stats := toolStats[exec.ToolName]
stats.totalCalls++
if exec.Status == "failed" || exec.Status == "cancelled" {
if monitorStatusCountsAsFailed(exec.Status) {
stats.failedCalls++
} else if exec.Status == "completed" {
stats.successCalls++
internal/handler/openapi.go
+25 -3
View File
@@ -307,7 +307,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"status": map[string]interface{}{
"type": "string",
"description": "执行状态",
"enum": []string{"success", "failed", "running"},
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
},
"result": map[string]interface{}{
"type": "string",
@@ -778,7 +778,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"status": map[string]interface{}{
"type": "string",
"description": "执行状态",
"enum": []string{"success", "failed", "running"},
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
},
"createdAt": map[string]interface{}{
"type": "string",
@@ -839,6 +839,9 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"type": "object",
"description": "配置信息（含 openai、vision、multi_agent 等）",
"properties": map[string]interface{}{
"agent": map[string]interface{}{
"$ref": "#/components/schemas/AgentConfig",
},
"vision": map[string]interface{}{
"$ref": "#/components/schemas/VisionConfig",
},
@@ -848,11 +851,30 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"type": "object",
"description": "更新配置请求",
"properties": map[string]interface{}{
"agent": map[string]interface{}{
"$ref": "#/components/schemas/AgentConfig",
},
"vision": map[string]interface{}{
"$ref": "#/components/schemas/VisionConfig",
},
},
},
"AgentConfig": map[string]interface{}{
"type": "object",
"description": "Agent 运行与外部 MCP 防卡死保护配置",
"properties": map[string]interface{}{
"max_iterations": map[string]interface{}{"type": "integer", "description": "最大迭代次数"},
"tool_timeout_minutes": map[string]interface{}{"type": "integer", "description": "单次工具执行硬超时（分钟）"},
"tool_wait_timeout_seconds": map[string]interface{}{"type": "integer", "description": "工具单轮等待秒数；到时返回 execution_id，worker 继续后台执行"},
"external_mcp_max_concurrent_per_server": map[string]interface{}{"type": "integer", "description": "单个外部 MCP server 并发上限；0=默认2；负数=不限制"},
"external_mcp_max_concurrent_total": map[string]interface{}{"type": "integer", "description": "外部 MCP 全局并发上限；0=默认16；负数=不限制"},
"external_mcp_circuit_failure_threshold": map[string]interface{}{"type": "integer", "description": "连续失败熔断阈值；0=默认3；负数=关闭熔断"},
"external_mcp_circuit_cooldown_seconds": map[string]interface{}{"type": "integer", "description": "熔断冷却秒数；0=默认60"},
"shell_no_output_timeout_seconds": map[string]interface{}{"type": "integer", "description": "execute/exec 连续无输出终止秒数"},
"workspace_root_dir": map[string]interface{}{"type": "string", "description": "会话工作目录根路径"},
"system_prompt_path": map[string]interface{}{"type": "string", "description": "单代理系统提示文件路径"},
},
},
"VisionConfig": map[string]interface{}{
"type": "object",
"description": "视觉分析（analyze_image MCP 工具）；enabled 且 model 非空时注册工具",
@@ -3483,7 +3505,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
"description": "状态筛选",
"schema": map[string]interface{}{
"type": "string",
"enum": []string{"success", "failed", "running"},
"enum": []string{"queued", "running", "completed", "failed", "cancelled", "hard_timeout", "orphaned"},
},
},
{
internal/handler/task_manager.go
+10 -2
View File
@@ -252,7 +252,7 @@ type AgentTaskManager struct {
maxHistorySize int // 最大历史记录数
historyRetention time.Duration // 历史记录保留时间
eventBus *TaskEventBus // 可选：任务结束时关闭镜像 SSE 订阅
// toolCanceler 在用户整轮停止任务时终止当前 MCP 工具（非「中断并继续」）。
// toolCanceler 在用户整轮停止任务或会话结束时终止该会话仍在运行的 MCP 工具（非「中断并继续」）。
toolCanceler func(conversationID string)
}
@@ -284,7 +284,7 @@ func (m *AgentTaskManager) SetTaskEventBus(b *TaskEventBus) {
m.eventBus = b
}
// SetToolCanceler 设置整轮停止任务时终止当前 MCP 工具的回调（由 AgentHandler 注入）。
// SetToolCanceler 设置整轮停止任务/会话结束时终止仍在运行 MCP 工具的回调（由 AgentHandler 注入）。
func (m *AgentTaskManager) SetToolCanceler(fn func(conversationID string)) {
m.mu.Lock()
defer m.mu.Unlock()
@@ -444,6 +444,8 @@ func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string)
if finalStatus != "" {
task.Status = finalStatus
}
toolCanceler := m.toolCanceler
activeEinoExecuteCancel := task.activeEinoExecuteCancel
// 保存到历史记录
completedTask := &CompletedTask{
@@ -464,6 +466,12 @@ func (m *AgentTaskManager) FinishTask(conversationID string, finalStatus string)
delete(m.tasks, conversationID)
bus := m.eventBus
m.mu.Unlock()
if toolCanceler != nil {
toolCanceler(conversationID)
}
if activeEinoExecuteCancel != nil {
activeEinoExecuteCancel()
}
if bus != nil {
bus.CloseConversation(conversationID)
}
internal/handler/task_manager_tool_cancel_test.go
+20
View File
@@ -78,3 +78,23 @@ func TestCancelTaskDefaultCauseIsTaskCancelled(t *testing.T) {
t.Fatalf("expected tool canceler path for default cancel cause")
}
}
func TestFinishTaskInvokesToolCancelerOnSessionEnd(t *testing.T) {
tm := NewAgentTaskManager()
calls := 0
tm.SetToolCanceler(func(conversationID string) {
if conversationID == "conv-3" {
calls++
}
})
_, cancel := context.WithCancelCause(context.Background())
if _, err := tm.StartTask("conv-3", "hello", cancel); err != nil {
t.Fatalf("StartTask: %v", err)
}
tm.FinishTask("conv-3", "completed")
if calls != 1 {
t.Fatalf("expected one tool cleanup on FinishTask, got %d", calls)
}
}