From c925c2e74badb186afb7e234f8089d654c3514dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=85=AC=E6=98=8E?= <83812544+Ed1s0nZ@users.noreply.github.com> Date: Tue, 21 Jul 2026 20:49:10 +0800 Subject: [PATCH] Add files via upload --- internal/app/app.go | 12 ++++++ internal/app/mcp_authorization.go | 14 +++++++ internal/app/mcp_authorization_test.go | 42 ++++++++++++++++++++- internal/database/monitor.go | 18 ++++----- internal/database/monitor_reconcile_test.go | 8 ++-- 5 files changed, 80 insertions(+), 14 deletions(-) diff --git a/internal/app/app.go b/internal/app/app.go index ad4764bc..805c1b73 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -148,10 +148,13 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error mcpServer := mcp.NewServerWithStorage(log.Logger, db) mcpServer.SetToolAuthorizer(mcpToolAuthorizer(db)) mcpServer.ConfigureHTTPToolCallTimeoutFromAgentMinutes(cfg.Agent.ToolTimeoutMinutes) + mcpServer.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds) + mcpServer.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) // 创建安全工具执行器 executor := security.NewExecutor(&cfg.Security, mcpServer, log.Logger) executor.SetShellNoOutputTimeoutSeconds(cfg.Agent.ShellNoOutputTimeoutSeconds) + executor.SetToolOutputMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) // 注册工具 executor.RegisterTools(mcpServer) @@ -165,6 +168,15 @@ func New(cfg *config.Config, log *logger.Logger, configPath string) (*App, error // 创建外部MCP管理器（使用与内部MCP服务器相同的存储） externalMCPMgr := mcp.NewExternalMCPManagerWithStorage(log.Logger, db) externalMCPMgr.SetToolAuthorizer(externalMCPToolAuthorizer()) + externalMCPMgr.ConfigureToolWaitTimeoutSeconds(cfg.Agent.ToolWaitTimeoutSeconds) + externalMCPMgr.ConfigureToolResultMaxBytes(cfg.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective()) + externalMCPMgr.ConfigureResilience(mcp.ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: cfg.Agent.ExternalMCPMaxConcurrentPerServer, + MaxConcurrentTotal: cfg.Agent.ExternalMCPMaxConcurrentTotal, + CircuitFailureThreshold: cfg.Agent.ExternalMCPCircuitFailureThreshold, + CircuitCooldown: time.Duration(cfg.Agent.ExternalMCPCircuitCooldownSeconds) * time.Second, + }) + mcp.RegisterExecutionControlTools(mcpServer, externalMCPMgr) if cfg.ExternalMCP.Servers != nil { externalMCPMgr.LoadConfigs(&cfg.ExternalMCP) // 启动所有启用的外部MCP客户端 diff --git a/internal/app/mcp_authorization.go b/internal/app/mcp_authorization.go index a92b6728..1bfeb0a0 100644 --- a/internal/app/mcp_authorization.go +++ b/internal/app/mcp_authorization.go @@ -34,6 +34,16 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string } return nil } + toolExecutionResource := func(permission string) error { + if err := require(permission); err != nil { + return err + } + id := mcpAuthorizationString(args, "execution_id") + if id == "" || db == nil || !db.UserCanAccessToolExecution(principal.UserID, principal.ScopeFor(permission), id) { + return fmt.Errorf("no access to tool execution %s", id) + } + return nil + } switch toolName { case builtin.ToolWebshellExec, builtin.ToolWebshellFileWrite: @@ -108,6 +118,10 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string return require("knowledge:read") case builtin.ToolAnalyzeImage: return require("agent:execute") + case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution: + return toolExecutionResource("monitor:read") + case builtin.ToolCancelToolExecution: + return toolExecutionResource("monitor:write") case builtin.ToolBatchTaskList: return require("tasks:read") case builtin.ToolBatchTaskGet: diff --git a/internal/app/mcp_authorization_test.go b/internal/app/mcp_authorization_test.go index 6b185b1a..e8b29d6e 100644 --- a/internal/app/mcp_authorization_test.go +++ b/internal/app/mcp_authorization_test.go @@ -9,6 +9,7 @@ import ( "cyberstrike-ai/internal/authctx" "cyberstrike-ai/internal/database" + "cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp/builtin" "cyberstrike-ai/internal/security" @@ -62,7 +63,7 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) { authorize := mcpToolAuthorizer(db) args := map[string]interface{}{ "action": "get", "connection_id": "x", "queue_id": "x", "listener_id": "x", - "session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x", + "session_id": "x", "task_id": "x", "id": "x", "conversation_id": "x", "execution_id": "x", } for _, toolName := range builtin.GetAllBuiltinTools() { err := authorize(ctx, toolName, args) @@ -72,6 +73,45 @@ func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) { } } +func TestMCPExecutionControlAuthorizationUsesExecutionScope(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-exec-authz.db"), zap.NewNop()) + if err != nil { + t.Fatal(err) + } + defer db.Close() + user, err := db.CreateRBACUser("exec-user", "Exec User", "hash", true, nil) + if err != nil { + t.Fatal(err) + } + if err := db.SaveToolExecution(&mcp.ToolExecution{ + ID: "exec-owned", + ToolName: "lab::slow", + Status: "running", + StartTime: time.Now(), + OwnerUserID: user.ID, + }); err != nil { + t.Fatal(err) + } + if err := db.SaveToolExecution(&mcp.ToolExecution{ + ID: "exec-hidden", + ToolName: "lab::slow", + Status: "running", + StartTime: time.Now(), + }); err != nil { + t.Fatal(err) + } + + principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{"monitor:read": true, "monitor:write": true}) + ctx := authctx.WithPrincipal(context.Background(), principal) + authorize := mcpToolAuthorizer(db) + if err := authorize(ctx, builtin.ToolWaitToolExecution, map[string]interface{}{"execution_id": "exec-owned"}); err != nil { + t.Fatalf("owned execution denied: %v", err) + } + if err := authorize(ctx, builtin.ToolCancelToolExecution, map[string]interface{}{"execution_id": "exec-hidden"}); err == nil { + t.Fatal("foreign execution was allowed") + } +} + func TestMCPAssetToolAuthorizationUsesAssetPermissionsAndScope(t *testing.T) { db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-asset-authz.db"), zap.NewNop()) if err != nil { diff --git a/internal/database/monitor.go b/internal/database/monitor.go index 77f3006c..c9f27f08 100644 --- a/internal/database/monitor.go +++ b/internal/database/monitor.go @@ -359,7 +359,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T var lastCall sql.NullString err := db.QueryRow(`SELECT COUNT(*), COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0), - COALESCE(SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), 0), + COALESCE(SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0), MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan( &result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls, &lastCall, &result.Summary.ToolCount, @@ -373,7 +373,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T } rows, err := db.Query(`SELECT tool_name, COUNT(*), SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), - SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END), MAX(start_time)`+ + SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+ fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...) if err != nil { return nil, err @@ -565,7 +565,7 @@ func (db *DB) UserCanAccessToolExecution(userID, scope, executionID string) bool return conversation != "" && db.UserCanAccessResource(userID, scope, "conversation", conversation) } -// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 cancelled（如进程重启后无对应执行协程）。 +// CancelOrphanedRunningToolExecutions 将仍为 running 的记录批量标记为 orphaned（如进程重启后无对应执行协程）。 func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg string) (int64, error) { errMsg = strings.TrimSpace(errMsg) if errMsg == "" { @@ -573,7 +573,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri } query := ` UPDATE tool_executions - SET status = 'cancelled', + SET status = 'orphaned', error = ?, end_time = ?, duration_ms = MAX(0, CAST((julianday(?) - julianday(start_time)) * 86400000 AS INTEGER)) @@ -586,7 +586,7 @@ func (db *DB) CancelOrphanedRunningToolExecutions(endTime time.Time, errMsg stri return res.RowsAffected() } -// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 cancelled。 +// FinalizeStaleRunningToolExecutions 将「非活跃且超过 minAge」的 running 记录标记为 orphaned。 // activeIDs 为当前进程内仍登记 cancel 的 executionId；不在集合内且已超时的视为孤儿记录。 func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time.Duration, activeIDs map[string]struct{}, errMsg string) (int64, error) { errMsg = strings.TrimSpace(errMsg) @@ -639,7 +639,7 @@ func (db *DB) FinalizeStaleRunningToolExecutions(endTime time.Time, minAge time. } res, err := db.Exec(` UPDATE tool_executions - SET status = 'cancelled', error = ?, end_time = ?, duration_ms = ? + SET status = 'orphaned', error = ?, end_time = ?, duration_ms = ? WHERE id = ? AND status = 'running' `, errMsg, endTime, durationMs, row.id) if err != nil { @@ -815,7 +815,7 @@ func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) { } delta.totalCalls += count switch status { - case "failed", "cancelled": + case "failed", "cancelled", "hard_timeout", "orphaned": delta.failedCalls += count case "completed": delta.successCalls += count @@ -971,7 +971,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime query = ` SELECT date(start_time, 'localtime') AS bucket, COUNT(*) AS total, - SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed + SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed FROM tool_executions WHERE start_time >= ? GROUP BY bucket @@ -981,7 +981,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime query = ` SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket, COUNT(*) AS total, - SUM(CASE WHEN status IN ('failed', 'cancelled') THEN 1 ELSE 0 END) AS failed + SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed FROM tool_executions WHERE start_time >= ? GROUP BY bucket diff --git a/internal/database/monitor_reconcile_test.go b/internal/database/monitor_reconcile_test.go index c2f96a3c..72e60d6c 100644 --- a/internal/database/monitor_reconcile_test.go +++ b/internal/database/monitor_reconcile_test.go @@ -43,8 +43,8 @@ func TestCancelOrphanedRunningToolExecutions(t *testing.T) { if err != nil { t.Fatalf("GetToolExecution: %v", err) } - if got.Status != "cancelled" { - t.Fatalf("expected cancelled, got %s", got.Status) + if got.Status != "orphaned" { + t.Fatalf("expected orphaned, got %s", got.Status) } if got.EndTime == nil { t.Fatal("expected end_time to be set") @@ -88,8 +88,8 @@ func TestFinalizeStaleRunningToolExecutions_skipsActive(t *testing.T) { if err != nil { t.Fatalf("GetToolExecution stale: %v", err) } - if stale.Status != "cancelled" { - t.Fatalf("stale expected cancelled, got %s", stale.Status) + if stale.Status != "orphaned" { + t.Fatalf("stale expected orphaned, got %s", stale.Status) } activeExec, err := db.GetToolExecution("active")