diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 00000000..941f2430 --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,2108 @@ +package config + +import ( + "crypto/rand" + "encoding/hex" + "encoding/json" + "fmt" + "io/fs" + "os" + "path/filepath" + "strconv" + "strings" + + "cyberstrike-ai/internal/termout" + + "gopkg.in/yaml.v3" +) + +type Config struct { + Version string `yaml:"version,omitempty" json:"version,omitempty"` // 前端显示的版本号，如 v1.3.3 + Server ServerConfig `yaml:"server"` + Log LogConfig `yaml:"log"` + MCP MCPConfig `yaml:"mcp"` + AI AIConfig `yaml:"ai,omitempty" json:"ai,omitempty"` + OpenAI OpenAIConfig `yaml:"openai,omitempty" json:"openai,omitempty"` + FOFA FofaConfig `yaml:"fofa,omitempty" json:"fofa,omitempty"` + ZoomEye SpaceSearchConfig `yaml:"zoomeye,omitempty" json:"zoomeye,omitempty"` + Quake SpaceSearchConfig `yaml:"quake,omitempty" json:"quake,omitempty"` + Shodan SpaceSearchConfig `yaml:"shodan,omitempty" json:"shodan,omitempty"` + Agent AgentConfig `yaml:"agent"` + Hitl HitlConfig `yaml:"hitl,omitempty" json:"hitl,omitempty"` + Security SecurityConfig `yaml:"security"` + Database DatabaseConfig `yaml:"database"` + Auth AuthConfig `yaml:"auth"` + Audit AuditConfig `yaml:"audit,omitempty" json:"audit,omitempty"` + Monitor MonitorConfig `yaml:"monitor,omitempty" json:"monitor,omitempty"` + ExternalMCP ExternalMCPConfig `yaml:"external_mcp,omitempty"` + Knowledge KnowledgeConfig `yaml:"knowledge,omitempty"` + C2 C2Config `yaml:"c2,omitempty" json:"c2,omitempty"` // 内置 C2 总开关；未配置时默认启用 + Robots RobotsConfig `yaml:"robots,omitempty" json:"robots,omitempty"` // 企业微信/钉钉/飞书等机器人配置 + RolesDir string `yaml:"roles_dir,omitempty" json:"roles_dir,omitempty"` // 角色配置文件目录（新方式） + Roles map[string]RoleConfig `yaml:"roles,omitempty" json:"roles,omitempty"` // 向后兼容：支持在主配置文件中定义角色 + SkillsDir string `yaml:"skills_dir,omitempty" json:"skills_dir,omitempty"` // Skills配置文件目录 + AgentsDir string `yaml:"agents_dir,omitempty" json:"agents_dir,omitempty"` // 多代理子 Agent Markdown 定义目录（*.md，YAML front matter） + MultiAgent MultiAgentConfig `yaml:"multi_agent,omitempty" json:"multi_agent,omitempty"` + Project ProjectConfig `yaml:"project,omitempty" json:"project,omitempty"` + Vision VisionConfig `yaml:"vision,omitempty" json:"vision,omitempty"` +} + +type EnsureLocalConfigResult struct { + Created bool + ExamplePath string +} + +const ( + DefaultMaxCompletionTokens = 16384 + DefaultSummarizationUserIntentLedgerMaxRunes = 96000 + DefaultSummarizationUserIntentLedgerEntryMaxRunes = 16000 + DefaultLatestUserMessageMaxRunes = 48000 + DefaultLatestUserMessageHeadRunes = 24000 + DefaultLatestUserMessageTailRunes = 24000 + DefaultSummarizationOutputReserveTokens = 8192 +) + +// ProjectConfig 项目黑板（跨对话共享事实）配置。 +type ProjectConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目 + FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"` + FactIndexPathMaxRunes int `yaml:"fact_index_path_max_runes,omitempty" json:"fact_index_path_max_runes,omitempty"` + FactSummaryMaxRunes int `yaml:"fact_summary_max_runes,omitempty" json:"fact_summary_max_runes,omitempty"` + DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"` +} + +// FactIndexMaxRunesEffective 自动注入黑板索引的最大 rune 数。 +func (c ProjectConfig) FactIndexMaxRunesEffective() int { + if c.FactIndexMaxRunes <= 0 { + return 3500 + } + return c.FactIndexMaxRunes +} + +// FactIndexPathMaxRunesEffective 攻击路径速览段的最大 rune 数（从 fact_index_max_runes 预算中预留）。 +func (c ProjectConfig) FactIndexPathMaxRunesEffective() int { + if c.FactIndexPathMaxRunes <= 0 { + return 1000 + } + return c.FactIndexPathMaxRunes +} + +// FactSummaryMaxRunesEffective upsert 时 summary 最大 rune 数（索引一行，宜含验证要点）。 +func (c ProjectConfig) FactSummaryMaxRunesEffective() int { + if c.FactSummaryMaxRunes <= 0 { + return 200 + } + return c.FactSummaryMaxRunes +} + +// MultiAgentConfig 基于 CloudWeGo Eino adk/prebuilt 的多代理编排（deep | plan_execute | supervisor）。 +type MultiAgentConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + RobotDefaultAgentMode string `yaml:"robot_default_agent_mode,omitempty" json:"robot_default_agent_mode,omitempty"` // eino_single | deep | plan_execute | supervisor + BatchUseMultiAgent bool `yaml:"batch_use_multi_agent" json:"batch_use_multi_agent"` // 为 true 时批量任务队列中每子任务走 Eino 多代理 + // Orchestration 已弃用：保留仅兼容旧版 config.yaml；编排由聊天/WebShell 请求体 orchestration 决定，未传时按 deep。 + Orchestration string `yaml:"orchestration,omitempty" json:"orchestration,omitempty"` + // MaxIteration 已废弃：统一使用 agent.max_iterations（YAML 中保留字段仅为兼容旧配置，运行时不读取）。 + MaxIteration int `yaml:"max_iteration,omitempty" json:"max_iteration,omitempty"` + // PlanExecuteLoopMaxIterations plan_execute 模式下 execute↔replan 外层循环上限；0 表示用 Eino 默认 10。 + PlanExecuteLoopMaxIterations int `yaml:"plan_execute_loop_max_iterations,omitempty" json:"plan_execute_loop_max_iterations,omitempty"` + // SubAgentMaxIterations 已废弃：子代理与主代理均使用 agent.max_iterations（Markdown max_iterations>0 可覆盖）。 + SubAgentMaxIterations int `yaml:"sub_agent_max_iterations,omitempty" json:"sub_agent_max_iterations,omitempty"` + WithoutGeneralSubAgent bool `yaml:"without_general_sub_agent" json:"without_general_sub_agent"` + WithoutWriteTodos bool `yaml:"without_write_todos" json:"without_write_todos"` + OrchestratorInstruction string `yaml:"orchestrator_instruction" json:"orchestrator_instruction"` + // OrchestratorInstructionPlanExecute plan_execute 主代理（规划侧）系统提示；非空且 agents/orchestrator-plan-execute.md 正文为空或未存在时生效。不与 Deep 的 orchestrator_instruction 混用。 + OrchestratorInstructionPlanExecute string `yaml:"orchestrator_instruction_plan_execute,omitempty" json:"orchestrator_instruction_plan_execute,omitempty"` + // OrchestratorInstructionSupervisor supervisor 主代理系统提示（transfer/exit 说明仍由运行追加）；非空且 agents/orchestrator-supervisor.md 正文为空或未存在时生效。 + OrchestratorInstructionSupervisor string `yaml:"orchestrator_instruction_supervisor,omitempty" json:"orchestrator_instruction_supervisor,omitempty"` + SubAgents []MultiAgentSubConfig `yaml:"sub_agents" json:"sub_agents"` + // SubAgentUserContextMaxRunes caps user-context supplement for sub-agent task descriptions. + // 0 (default) preserves all user turns verbatim; >0 caps total runes; negative disables injection. + SubAgentUserContextMaxRunes int `yaml:"sub_agent_user_context_max_runes,omitempty" json:"sub_agent_user_context_max_runes,omitempty"` + // EinoSkills configures CloudWeGo Eino ADK skill middleware + optional local filesystem/execute on DeepAgent. + EinoSkills MultiAgentEinoSkillsConfig `yaml:"eino_skills,omitempty" json:"eino_skills,omitempty"` + // EinoMiddleware wires optional ADK middleware (patchtoolcalls, toolsearch, plantask, reduction) and Deep extras. + EinoMiddleware MultiAgentEinoMiddlewareConfig `yaml:"eino_middleware,omitempty" json:"eino_middleware,omitempty"` + // EinoCallbacks attaches CloudWeGo eino callbacks.InitCallbacks on ADK Runner context (structured logs + optional SSE trace). + EinoCallbacks MultiAgentEinoCallbacksConfig `yaml:"eino_callbacks,omitempty" json:"eino_callbacks,omitempty"` +} + +// SubAgentUserContextMaxRunesEffective returns max runes for sub-agent task supplement; 0 = unlimited; negative = disabled. +func (c MultiAgentConfig) SubAgentUserContextMaxRunesEffective() int { + return c.SubAgentUserContextMaxRunes +} + +// MultiAgentEinoCallbacksConfig enables Eino unified callbacks on each ADK agent run (deep / plan_execute / supervisor / eino_single). +// Modes: log_only (zap + optional OTel; no SSE to browser), sse (adds client SSE eino_trace_* when sse_trace_to_client), full (sse rules + stream callback copies closed). +type MultiAgentEinoCallbacksConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` // log_only | sse | full; empty with enabled=true defaults to log_only + // SseTraceToClient when true emits eino_trace_* SSE for UI (use only for admin/debug; nil/false recommended in production). + SseTraceToClient *bool `yaml:"sse_trace_to_client,omitempty" json:"sse_trace_to_client,omitempty"` + // Otel configures OpenTelemetry trace export (independent of mode; exporter none disables export even if enabled). + Otel MultiAgentEinoCallbacksOtelConfig `yaml:"otel,omitempty" json:"otel,omitempty"` + // MaxInputSummaryRunes / MaxOutputSummaryRunes cap text placed in SSE payloads and debug logs (not full payloads). + MaxInputSummaryRunes int `yaml:"max_input_summary_runes,omitempty" json:"max_input_summary_runes,omitempty"` + MaxOutputSummaryRunes int `yaml:"max_output_summary_runes,omitempty" json:"max_output_summary_runes,omitempty"` + // ZapVerbose when true logs input/output summaries at zap.Debug on start/end; false uses Info with short fields only. + ZapVerbose bool `yaml:"zap_verbose,omitempty" json:"zap_verbose,omitempty"` +} + +// MultiAgentEinoCallbacksOtelConfig OpenTelemetry for Eino callback spans (W3C trace in collector / stdout). +type MultiAgentEinoCallbacksOtelConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + ServiceName string `yaml:"service_name,omitempty" json:"service_name,omitempty"` + Exporter string `yaml:"exporter,omitempty" json:"exporter,omitempty"` // none | stdout | otlphttp + OTLPEndpoint string `yaml:"otlp_endpoint,omitempty" json:"otlp_endpoint,omitempty"` // host:port, e.g. localhost:4318 (path /v1/traces) + SampleRatio float64 `yaml:"sample_ratio,omitempty" json:"sample_ratio,omitempty"` // 0–1, default 1.0 +} + +// EinoCallbacksModeEffective returns off | log_only | sse | full. +func (c MultiAgentEinoCallbacksConfig) EinoCallbacksModeEffective() string { + if !c.Enabled { + return "off" + } + m := strings.TrimSpace(strings.ToLower(c.Mode)) + switch m { + case "log_only": + return "log_only" + case "sse": + return "sse" + case "full": + return "full" + case "": + return "log_only" + default: + return "log_only" + } +} + +// SseTraceToClientEffective is false unless explicitly set true (best practice: do not expose framework traces to end users by default). +func (c MultiAgentEinoCallbacksConfig) SseTraceToClientEffective() bool { + if c.SseTraceToClient == nil { + return false + } + return *c.SseTraceToClient +} + +// ShouldEmitEinoTraceSSE is true when client-visible trace events should be sent over progress/SSE. +func (c MultiAgentEinoCallbacksConfig) ShouldEmitEinoTraceSSE(mode string) bool { + if !c.SseTraceToClientEffective() { + return false + } + return mode == "sse" || mode == "full" +} + +// OtelExporterEffective returns none | stdout | otlphttp. +func (c MultiAgentEinoCallbacksOtelConfig) OtelExporterEffective() string { + e := strings.TrimSpace(strings.ToLower(c.Exporter)) + switch e { + case "none", "stdout", "otlphttp": + return e + case "": + if c.Enabled { + return "stdout" + } + return "none" + default: + return "none" + } +} + +// OtelTracingActive is true when spans should be started (enabled + non-none exporter). +func (c MultiAgentEinoCallbacksConfig) OtelTracingActive() bool { + if !c.Otel.Enabled { + return false + } + return c.Otel.OtelExporterEffective() != "none" +} + +func (c MultiAgentEinoCallbacksOtelConfig) ServiceNameEffective() string { + s := strings.TrimSpace(c.ServiceName) + if s != "" { + return s + } + return "cyberstrike-ai" +} + +func (c MultiAgentEinoCallbacksOtelConfig) SampleRatioEffective() float64 { + r := c.SampleRatio + if r <= 0 { + return 1.0 + } + if r > 1 { + return 1.0 + } + return r +} + +func (c MultiAgentEinoCallbacksConfig) EinoCallbacksMaxInputSummaryRunes() int { + if c.MaxInputSummaryRunes > 0 { + return c.MaxInputSummaryRunes + } + return 400 +} + +func (c MultiAgentEinoCallbacksConfig) EinoCallbacksMaxOutputSummaryRunes() int { + if c.MaxOutputSummaryRunes > 0 { + return c.MaxOutputSummaryRunes + } + return 400 +} + +// MultiAgentEinoMiddlewareConfig optional Eino ADK middleware and Deep / supervisor tuning. +type MultiAgentEinoMiddlewareConfig struct { + // PatchToolCalls inserts placeholder tool results for dangling assistant tool_calls (nil = enabled). + PatchToolCalls *bool `yaml:"patch_tool_calls,omitempty" json:"patch_tool_calls,omitempty"` + // ToolSearch enables dynamictool/toolsearch: hide tail tools until model calls tool_search (reduces prompt tools). + 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). + PlantaskRelDir string `yaml:"plantask_rel_dir,omitempty" json:"plantask_rel_dir,omitempty"` + // 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"` // 非空：落盘根目录（默认 tmp/reduction）；其下按 projects/{id} 或 conversations/{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"` + // SummarizationOutputReserveTokens reserves completion headroom for the summarization model call (default 8192). + SummarizationOutputReserveTokens int `yaml:"summarization_output_reserve_tokens,omitempty" json:"summarization_output_reserve_tokens,omitempty"` + // SummarizationEmitInternalEvents controls middleware internal event emission (default true). + SummarizationEmitInternalEvents *bool `yaml:"summarization_emit_internal_events,omitempty" json:"summarization_emit_internal_events,omitempty"` + // SummarizationUserIntentLedgerMaxRunes caps the DB-backed immutable user input ledger injected into model context. + SummarizationUserIntentLedgerMaxRunes int `yaml:"summarization_user_intent_ledger_max_runes,omitempty" json:"summarization_user_intent_ledger_max_runes,omitempty"` + // SummarizationUserIntentLedgerEntryMaxRunes caps each user message entry inside the immutable user input ledger. + SummarizationUserIntentLedgerEntryMaxRunes int `yaml:"summarization_user_intent_ledger_entry_max_runes,omitempty" json:"summarization_user_intent_ledger_entry_max_runes,omitempty"` + // LatestUserMessageMaxRunes caps the current user turn inserted into model context; full text is persisted as an artifact when capped. + LatestUserMessageMaxRunes int `yaml:"latest_user_message_max_runes,omitempty" json:"latest_user_message_max_runes,omitempty"` + // LatestUserMessageHeadRunes keeps the head preview for an oversized current user turn. + LatestUserMessageHeadRunes int `yaml:"latest_user_message_head_runes,omitempty" json:"latest_user_message_head_runes,omitempty"` + // LatestUserMessageTailRunes keeps the tail preview for an oversized current user turn. + LatestUserMessageTailRunes int `yaml:"latest_user_message_tail_runes,omitempty" json:"latest_user_message_tail_runes,omitempty"` + // SummarizationRetryMaxAttempts 已废弃：summarization 与 Eino 原生 ModelRetry 共用 model_retry_max_retries 及 isEinoTransientRunError。 + SummarizationRetryMaxAttempts int `yaml:"summarization_retry_max_attempts,omitempty" json:"summarization_retry_max_attempts,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. + DeepOutputKey string `yaml:"deep_output_key,omitempty" json:"deep_output_key,omitempty"` + // DeepModelRetryMaxRetries 已废弃：请用 model_retry_max_retries；保留字段仅为兼容旧配置。 + DeepModelRetryMaxRetries int `yaml:"deep_model_retry_max_retries,omitempty" json:"deep_model_retry_max_retries,omitempty"` + // ModelRetryMaxRetries configures Eino ADK native ChatModel retry attempts; 0=default 4. + ModelRetryMaxRetries int `yaml:"model_retry_max_retries,omitempty" json:"model_retry_max_retries,omitempty"` + // ModelRetryMaxBackoffSec caps native model retry backoff seconds; 0=default 30. + ModelRetryMaxBackoffSec int `yaml:"model_retry_max_backoff_sec,omitempty" json:"model_retry_max_backoff_sec,omitempty"` + // ModelFailoverChannels lists ai.channels IDs to try after native model retry is exhausted. + ModelFailoverChannels []string `yaml:"model_failover_channels,omitempty" json:"model_failover_channels,omitempty"` + // ModelFailoverMaxRetries caps distinct failover channel attempts; 0=all configured failover channels. + ModelFailoverMaxRetries int `yaml:"model_failover_max_retries,omitempty" json:"model_failover_max_retries,omitempty"` + // RunRetryMaxAttempts 已废弃：模型临时错误由 Eino 原生 ModelRetry 处理；仅保留给非模型层 run loop 兜底与 summarization 旧字段。 + RunRetryMaxAttempts int `yaml:"run_retry_max_attempts,omitempty" json:"run_retry_max_attempts,omitempty"` + // RunRetryMaxBackoffSec 已废弃：请用 model_retry_max_backoff_sec；仅保留给非模型层 run loop 兜底与 summarization 旧字段。 + RunRetryMaxBackoffSec int `yaml:"run_retry_max_backoff_sec,omitempty" json:"run_retry_max_backoff_sec,omitempty"` + // EmptyResponseContinueMaxAttempts Run 成功但未捕获助手正文时 Handler 层退避续跑次数；0=默认 5。 + EmptyResponseContinueMaxAttempts int `yaml:"empty_response_continue_max_attempts,omitempty" json:"empty_response_continue_max_attempts,omitempty"` + // TaskToolDescriptionPrefix when non-empty sets deep.Config TaskToolDescriptionGenerator (sub-agent names appended). + 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) SummarizationOutputReserveTokensEffective() int { + if c.SummarizationOutputReserveTokens > 0 { + return c.SummarizationOutputReserveTokens + } + return DefaultSummarizationOutputReserveTokens +} + +func (c MultiAgentEinoMiddlewareConfig) SummarizationEmitInternalEventsEffective() bool { + if c.SummarizationEmitInternalEvents != nil { + return *c.SummarizationEmitInternalEvents + } + return true +} + +func (c MultiAgentEinoMiddlewareConfig) SummarizationUserIntentLedgerMaxRunesEffective() int { + if c.SummarizationUserIntentLedgerMaxRunes > 0 { + return c.SummarizationUserIntentLedgerMaxRunes + } + return DefaultSummarizationUserIntentLedgerMaxRunes +} + +func (c MultiAgentEinoMiddlewareConfig) SummarizationUserIntentLedgerEntryMaxRunesEffective() int { + if c.SummarizationUserIntentLedgerEntryMaxRunes > 0 { + return c.SummarizationUserIntentLedgerEntryMaxRunes + } + return DefaultSummarizationUserIntentLedgerEntryMaxRunes +} + +func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageMaxRunesEffective() int { + if c.LatestUserMessageMaxRunes > 0 { + return c.LatestUserMessageMaxRunes + } + return DefaultLatestUserMessageMaxRunes +} + +func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageHeadRunesEffective() int { + if c.LatestUserMessageHeadRunes > 0 { + return c.LatestUserMessageHeadRunes + } + return DefaultLatestUserMessageHeadRunes +} + +func (c MultiAgentEinoMiddlewareConfig) LatestUserMessageTailRunesEffective() int { + if c.LatestUserMessageTailRunes > 0 { + return c.LatestUserMessageTailRunes + } + return DefaultLatestUserMessageTailRunes +} + +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). + Disable bool `yaml:"disable" json:"disable"` + // FilesystemTools registers read_file/glob/grep/write/edit/execute (eino-ext local backend). Nil/omitted = true. + FilesystemTools *bool `yaml:"filesystem_tools,omitempty" json:"filesystem_tools,omitempty"` + // SkillToolName overrides the default Eino tool name "skill". + SkillToolName string `yaml:"skill_tool_name,omitempty" json:"skill_tool_name,omitempty"` +} + +// EinoSkillFilesystemToolsEffective returns whether Deep/sub-agents should attach local filesystem + streaming shell. +func (c MultiAgentEinoSkillsConfig) EinoSkillFilesystemToolsEffective() bool { + if c.FilesystemTools != nil { + return *c.FilesystemTools + } + return true +} + +// PatchToolCallsEffective returns whether patchtoolcalls middleware should run (default true). +func (c MultiAgentEinoMiddlewareConfig) PatchToolCallsEffective() bool { + if c.PatchToolCalls != nil { + return *c.PatchToolCalls + } + return true +} + +// MultiAgentSubConfig 子代理（Eino ChatModelAgent）：deep 下由 task 调度；supervisor 下由 transfer 委派；plan_execute 不使用子代理列表。 +type MultiAgentSubConfig struct { + ID string `yaml:"id" json:"id"` + Name string `yaml:"name" json:"name"` + Description string `yaml:"description" json:"description"` + Instruction string `yaml:"instruction" json:"instruction"` + BindRole string `yaml:"bind_role,omitempty" json:"bind_role,omitempty"` // 可选：关联主配置 roles 中的角色名；未配 role_tools 时沿用该角色的 tools + RoleTools []string `yaml:"role_tools" json:"role_tools"` // 与单 Agent 角色工具相同 key；空表示全部工具（bind_role 可补全 tools） + MaxIterations int `yaml:"max_iterations" json:"max_iterations"` + Kind string `yaml:"kind,omitempty" json:"kind,omitempty"` // 仅 Markdown：kind=orchestrator 表示 Deep 主代理（与 orchestrator.md 二选一约定） +} + +// MultiAgentPublic 返回给前端的精简信息（不含子代理指令全文）。 +type MultiAgentPublic struct { + Enabled bool `json:"enabled"` + RobotDefaultAgentMode string `json:"robot_default_agent_mode,omitempty"` + BatchUseMultiAgent bool `json:"batch_use_multi_agent"` + SubAgentCount int `json:"sub_agent_count"` + Orchestration string `json:"orchestration,omitempty"` + PlanExecuteLoopMaxIterations int `json:"plan_execute_loop_max_iterations"` + SummarizationUserIntentLedgerMaxRunes int `json:"summarization_user_intent_ledger_max_runes"` + SummarizationUserIntentLedgerEntryMaxRunes int `json:"summarization_user_intent_ledger_entry_max_runes"` + LatestUserMessageMaxRunes int `json:"latest_user_message_max_runes"` + LatestUserMessageHeadRunes int `json:"latest_user_message_head_runes"` + LatestUserMessageTailRunes int `json:"latest_user_message_tail_runes"` + ModelRetryMaxRetries int `json:"model_retry_max_retries"` + ModelRetryMaxBackoffSec int `json:"model_retry_max_backoff_sec"` + ModelFailoverChannels []string `json:"model_failover_channels,omitempty"` + ModelFailoverMaxRetries int `json:"model_failover_max_retries"` + ToolSearchAlwaysVisibleTools []string `json:"tool_search_always_visible_tools,omitempty"` + ToolSearchAlwaysVisibleEffectiveTools []string `json:"tool_search_always_visible_effective_tools,omitempty"` +} + +// NormalizeAgentMode 解析代理模式（eino_single | deep | plan_execute | supervisor）；空值默认 eino_single。 +func NormalizeAgentMode(mode string) string { + s := strings.TrimSpace(strings.ToLower(mode)) + switch s { + case "", "eino_single": + return "eino_single" + case "deep": + return "deep" + case "plan_execute", "plan-execute", "planexecute", "pe": + return "plan_execute" + case "supervisor", "super", "sv": + return "supervisor" + default: + return "eino_single" + } +} + +// NormalizeRobotAgentMode 解析机器人默认对话模式。 +func NormalizeRobotAgentMode(ma MultiAgentConfig) string { + return NormalizeAgentMode(ma.RobotDefaultAgentMode) +} + +// NormalizeMultiAgentOrchestration 返回 deep、plan_execute 或 supervisor。 +func NormalizeMultiAgentOrchestration(s string) string { + v := strings.TrimSpace(strings.ToLower(s)) + switch v { + case "plan_execute", "plan-execute", "planexecute", "pe": + return "plan_execute" + case "supervisor", "super", "sv": + return "supervisor" + default: + return "deep" + } +} + +// MultiAgentAPIUpdate 设置页/API 仅更新多代理标量字段；写入 YAML 时不覆盖 sub_agents 等块。 +type MultiAgentAPIUpdate struct { + Enabled bool `json:"enabled"` + RobotDefaultAgentMode string `json:"robot_default_agent_mode,omitempty"` + BatchUseMultiAgent bool `json:"batch_use_multi_agent"` + PlanExecuteLoopMaxIterations *int `json:"plan_execute_loop_max_iterations,omitempty"` + SummarizationUserIntentLedgerMaxRunes *int `json:"summarization_user_intent_ledger_max_runes,omitempty"` + SummarizationUserIntentLedgerEntryMaxRunes *int `json:"summarization_user_intent_ledger_entry_max_runes,omitempty"` + LatestUserMessageMaxRunes *int `json:"latest_user_message_max_runes,omitempty"` + LatestUserMessageHeadRunes *int `json:"latest_user_message_head_runes,omitempty"` + LatestUserMessageTailRunes *int `json:"latest_user_message_tail_runes,omitempty"` + ModelRetryMaxRetries *int `json:"model_retry_max_retries,omitempty"` + ModelRetryMaxBackoffSec *int `json:"model_retry_max_backoff_sec,omitempty"` + ModelFailoverChannels *[]string `json:"model_failover_channels,omitempty"` + ModelFailoverMaxRetries *int `json:"model_failover_max_retries,omitempty"` + // 指针区分「JSON 未传该字段」与「传空数组要清空」；省略时不应覆盖 YAML 中的常驻工具白名单。 + ToolSearchAlwaysVisibleTools *[]string `json:"tool_search_always_visible_tools,omitempty"` +} + +// RobotsConfig 机器人配置（企业微信、钉钉、飞书、微信 iLink、Telegram、Slack、Discord、QQ 等） +type RobotsConfig struct { + Session RobotSessionConfig `yaml:"session,omitempty" json:"session,omitempty"` // 机器人会话隔离策略 + Wechat RobotWechatConfig `yaml:"wechat,omitempty" json:"wechat,omitempty"` // 微信（iLink 扫码绑定） + Wecom RobotWecomConfig `yaml:"wecom,omitempty" json:"wecom,omitempty"` // 企业微信 + Dingtalk RobotDingtalkConfig `yaml:"dingtalk,omitempty" json:"dingtalk,omitempty"` // 钉钉 + Lark RobotLarkConfig `yaml:"lark,omitempty" json:"lark,omitempty"` // 飞书 + Telegram RobotTelegramConfig `yaml:"telegram,omitempty" json:"telegram,omitempty"` // Telegram + Slack RobotSlackConfig `yaml:"slack,omitempty" json:"slack,omitempty"` // Slack + Discord RobotDiscordConfig `yaml:"discord,omitempty" json:"discord,omitempty"` // Discord + QQ RobotQQConfig `yaml:"qq,omitempty" json:"qq,omitempty"` // QQ 机器人 +} + +// RobotWechatConfig 微信 iLink 机器人配置（个人微信 ClawBot / iLink 协议） +type RobotWechatConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + BotToken string `yaml:"bot_token,omitempty" json:"bot_token,omitempty"` + ILinkBotID string `yaml:"ilink_bot_id,omitempty" json:"ilink_bot_id,omitempty"` + ILinkUserID string `yaml:"ilink_user_id,omitempty" json:"ilink_user_id,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` // 默认 https://ilinkai.weixin.qq.com + BotType string `yaml:"bot_type,omitempty" json:"bot_type,omitempty"` // get_bot_qrcode 参数，默认 3 + BotAgent string `yaml:"bot_agent,omitempty" json:"bot_agent,omitempty"` // base_info.bot_agent + GetUpdatesBuf string `yaml:"get_updates_buf,omitempty" json:"get_updates_buf,omitempty"` // 长轮询游标（运行时） + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +const ( + RobotAuthModeUserBinding = "user_binding" + RobotAuthModeServiceAccount = "service_account" +) + +// RobotAuthorizationConfig controls how a verified platform sender becomes +// an RBAC principal. service_account is intentionally fail-closed unless an +// explicit non-admin service user and sender allowlist are both configured. +type RobotAuthorizationConfig struct { + Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` + ServiceUserID string `yaml:"service_user_id,omitempty" json:"service_user_id,omitempty"` + AllowedExternalUsers []string `yaml:"allowed_external_users,omitempty" json:"allowed_external_users,omitempty"` +} + +func (c RobotAuthorizationConfig) EffectiveMode() string { + mode := strings.ToLower(strings.TrimSpace(c.Mode)) + if mode == "" { + return RobotAuthModeUserBinding + } + return mode +} + +func (c RobotAuthorizationConfig) ExternalUserAllowed(externalUserID string) bool { + externalUserID = strings.TrimSpace(externalUserID) + if externalUserID == "" { + return false + } + for _, allowed := range c.AllowedExternalUsers { + if strings.TrimSpace(allowed) == externalUserID { + return true + } + } + return false +} + +// RobotSessionConfig 机器人会话隔离策略 +type RobotSessionConfig struct { + StrictUserIdentity *bool `yaml:"strict_user_identity,omitempty" json:"strict_user_identity,omitempty"` // true 时只允许真实用户标识，不允许会话/群 ID 兜底 +} + +// StrictUserIdentityEnabled 返回是否启用严格用户身份模式；未配置时默认 true。 +func (c RobotSessionConfig) StrictUserIdentityEnabled() bool { + if c.StrictUserIdentity == nil { + return true + } + return *c.StrictUserIdentity +} + +// RobotWecomConfig 企业微信机器人配置 +type RobotWecomConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + Token string `yaml:"token" json:"token"` // 回调 URL 校验 Token + EncodingAESKey string `yaml:"encoding_aes_key" json:"encoding_aes_key"` // EncodingAESKey + CorpID string `yaml:"corp_id" json:"corp_id"` // 企业 ID + Secret string `yaml:"secret" json:"secret"` // 应用 Secret + AgentID int64 `yaml:"agent_id" json:"agent_id"` // 应用 AgentId + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// ValidateWecomConfig 校验企业微信机器人配置；启用时必须配置 token，否则回调无法防伪造。 +func ValidateWecomConfig(w RobotWecomConfig) error { + if !w.Enabled { + return nil + } + if strings.TrimSpace(w.Token) == "" { + return fmt.Errorf("robots.wecom.enabled 为 true 时必须配置 robots.wecom.token") + } + return nil +} + +// RobotDingtalkConfig 钉钉机器人配置 +type RobotDingtalkConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + ClientID string `yaml:"client_id" json:"client_id"` // 应用 Key (AppKey) + ClientSecret string `yaml:"client_secret" json:"client_secret"` // 应用 Secret + AllowConversationIDFallback bool `yaml:"allow_conversation_id_fallback" json:"allow_conversation_id_fallback"` // sender_id 缺失时是否允许回退到会话 ID + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotLarkConfig 飞书机器人配置 +type RobotLarkConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + AppID string `yaml:"app_id" json:"app_id"` // 应用 App ID + AppSecret string `yaml:"app_secret" json:"app_secret"` // 应用 App Secret + VerifyToken string `yaml:"verify_token" json:"verify_token"` // 事件订阅 Verification Token（可选） + AllowChatIDFallback bool `yaml:"allow_chat_id_fallback" json:"allow_chat_id_fallback"` // 用户 ID 缺失时是否允许回退到 chat_id + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotTelegramConfig Telegram 机器人配置（Bot API 长轮询） +type RobotTelegramConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + BotToken string `yaml:"bot_token" json:"bot_token"` + BotUsername string `yaml:"bot_username,omitempty" json:"bot_username,omitempty"` // 可选，用于群聊 @ 识别；留空则启动时 getMe + AllowGroupMessages bool `yaml:"allow_group_messages" json:"allow_group_messages"` // 群聊中仅响应 @ 机器人 + UpdateOffset int64 `yaml:"update_offset,omitempty" json:"update_offset,omitempty"` + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotSlackConfig Slack 机器人配置（Socket Mode，无需公网回调） +type RobotSlackConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + BotToken string `yaml:"bot_token" json:"bot_token"` // xoxb- + AppToken string `yaml:"app_token" json:"app_token"` // xapp-（connections:write） + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotDiscordConfig Discord 机器人配置（Gateway WebSocket） +type RobotDiscordConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + BotToken string `yaml:"bot_token" json:"bot_token"` + AllowGuildMessages bool `yaml:"allow_guild_messages" json:"allow_guild_messages"` // 服务器频道中仅响应 @ 机器人 + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +// RobotQQConfig QQ 机器人配置（QQ 开放平台 WebSocket） +type RobotQQConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + AppID string `yaml:"app_id" json:"app_id"` + ClientSecret string `yaml:"client_secret" json:"client_secret"` + Sandbox bool `yaml:"sandbox" json:"sandbox"` // 沙箱环境（上线前测试） + Auth RobotAuthorizationConfig `yaml:"auth,omitempty" json:"auth,omitempty"` +} + +func (c RobotsConfig) AuthorizationFor(platform string) RobotAuthorizationConfig { + switch strings.ToLower(strings.TrimSpace(platform)) { + case "wechat": + return c.Wechat.Auth + case "wecom": + return c.Wecom.Auth + case "dingtalk": + return c.Dingtalk.Auth + case "lark": + return c.Lark.Auth + case "telegram": + return c.Telegram.Auth + case "slack": + return c.Slack.Auth + case "discord": + return c.Discord.Auth + case "qq": + return c.QQ.Auth + default: + return RobotAuthorizationConfig{} + } +} + +func ValidateRobotAuthorization(c RobotAuthorizationConfig, path string) error { + switch c.EffectiveMode() { + case RobotAuthModeUserBinding: + return nil + case RobotAuthModeServiceAccount: + serviceUserID := strings.TrimSpace(c.ServiceUserID) + if serviceUserID == "" { + return fmt.Errorf("%s.auth.service_user_id 不能为空", path) + } + if len(c.AllowedExternalUsers) == 0 { + return fmt.Errorf("%s.auth.allowed_external_users 至少配置一个真实发送者", path) + } + seen := map[string]bool{} + for _, userID := range c.AllowedExternalUsers { + userID = strings.TrimSpace(userID) + if userID == "" || userID == "*" { + return fmt.Errorf("%s.auth.allowed_external_users 不允许空值或通配符", path) + } + if seen[userID] { + return fmt.Errorf("%s.auth.allowed_external_users 包含重复用户", path) + } + seen[userID] = true + } + return nil + default: + return fmt.Errorf("%s.auth.mode 仅支持 user_binding 或 service_account", path) + } +} + +func ValidateRobotsAuthorization(c RobotsConfig) error { + items := []struct { + path string + auth RobotAuthorizationConfig + }{ + {"robots.wechat", c.Wechat.Auth}, {"robots.wecom", c.Wecom.Auth}, + {"robots.dingtalk", c.Dingtalk.Auth}, {"robots.lark", c.Lark.Auth}, + {"robots.telegram", c.Telegram.Auth}, {"robots.slack", c.Slack.Auth}, + {"robots.discord", c.Discord.Auth}, {"robots.qq", c.QQ.Auth}, + } + for _, item := range items { + if err := ValidateRobotAuthorization(item.auth, item.path); err != nil { + return err + } + } + return nil +} + +func (c RobotsConfig) ServiceAccountUserIDs() map[string]string { + out := map[string]string{} + for _, platform := range []string{"wechat", "wecom", "dingtalk", "lark", "telegram", "slack", "discord", "qq"} { + auth := c.AuthorizationFor(platform) + if auth.EffectiveMode() == RobotAuthModeServiceAccount { + out[platform] = strings.TrimSpace(auth.ServiceUserID) + } + } + return out +} + +type ServerConfig struct { + Host string `yaml:"host" json:"host"` + Port int `yaml:"port" json:"port"` + // CORSAllowedOrigins contains additional, exact origins that may call the API. + // Same-origin browser requests are always allowed. Wildcards are intentionally unsupported. + CORSAllowedOrigins []string `yaml:"cors_allowed_origins,omitempty" json:"cors_allowed_origins,omitempty"` + // TLSEnabled 为 true 时主 Web UI 使用 HTTPS；现代浏览器在同源下会协商 HTTP/2，缓解 HTTP/1.1 每源并发连接数限制。 + TLSEnabled bool `yaml:"tls_enabled,omitempty" json:"tls_enabled,omitempty"` + // TLSCertPath / TLSKeyPath 非空时从 PEM 文件加载证书（生产环境推荐）。 + TLSCertPath string `yaml:"tls_cert_path,omitempty" json:"tls_cert_path,omitempty"` + TLSKeyPath string `yaml:"tls_key_path,omitempty" json:"tls_key_path,omitempty"` + // TLSAutoSelfSign 为 true 且未配置有效证书路径时，启动时生成内存自签证书（仅本地/测试；浏览器会提示不受信任）。 + TLSAutoSelfSign bool `yaml:"tls_auto_self_sign,omitempty" json:"tls_auto_self_sign,omitempty"` + // TLSHTTPRedirect 为 false 时禁用 HTTP→HTTPS 跳转；省略或为 true 且已启用 HTTPS 时，明文 HTTP 访问将 308 跳转到 HTTPS（同端口嗅探分流）。 + TLSHTTPRedirect *bool `yaml:"tls_http_redirect,omitempty" json:"tls_http_redirect,omitempty"` +} + +type LogConfig struct { + Level string `yaml:"level"` + Output string `yaml:"output"` +} + +type MCPConfig struct { + Enabled bool `yaml:"enabled"` + Host string `yaml:"host"` + Port int `yaml:"port"` + AuthHeader string `yaml:"auth_header,omitempty"` // 可选的全局服务凭证 header；普通调用优先使用用户 Bearer Token + AuthHeaderValue string `yaml:"auth_header_value,omitempty"` // 全局服务凭证，仅 allow_global_access=true 时接受 + AllowGlobalAccess bool `yaml:"allow_global_access,omitempty"` // 静态服务密钥是否映射为全局服务身份（默认关闭） +} + +type OpenAIConfig struct { + Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` // API 提供商: "openai"(默认) 或 "claude"，claude 时自动桥接为 Anthropic Messages API + APIKey string `yaml:"api_key" json:"api_key"` + BaseURL string `yaml:"base_url" json:"base_url"` + Model string `yaml:"model" json:"model"` + MaxTotalTokens int `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"` + MaxCompletionTokens int `yaml:"max_completion_tokens,omitempty" json:"max_completion_tokens,omitempty"` + // Reasoning 控制 Eino ChatModel 的 thinking / reasoning_effort / output_config 等（Eino 单/多代理路径生效）。 + Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"` +} + +// AIConfig stores first-class model channels. Runtime callers resolve a channel +// into OpenAIConfig at the edge instead of moving API credentials through chat requests. +type AIConfig struct { + DefaultChannel string `yaml:"default_channel,omitempty" json:"default_channel,omitempty"` + Channels map[string]AIChannelConfig `yaml:"channels,omitempty" json:"channels,omitempty"` +} + +type AIChannelConfig struct { + Name string `yaml:"name,omitempty" json:"name,omitempty"` + Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` + APIKey string `yaml:"api_key" json:"api_key"` + BaseURL string `yaml:"base_url" json:"base_url"` + Model string `yaml:"model" json:"model"` + MaxTotalTokens int `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"` + MaxCompletionTokens int `yaml:"max_completion_tokens,omitempty" json:"max_completion_tokens,omitempty"` + Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"` +} + +func (c AIChannelConfig) ToOpenAIConfig() OpenAIConfig { + provider := strings.TrimSpace(c.Provider) + if provider == "" || provider == "openai_compatible" { + provider = "openai" + } + return OpenAIConfig{ + Provider: provider, + APIKey: c.APIKey, + BaseURL: c.BaseURL, + Model: c.Model, + MaxTotalTokens: c.MaxTotalTokens, + MaxCompletionTokens: c.MaxCompletionTokens, + Reasoning: c.Reasoning, + } +} + +func AIChannelFromOpenAI(id, name string, oa OpenAIConfig) AIChannelConfig { + if strings.TrimSpace(name) == "" { + name = id + } + return AIChannelConfig{ + Name: name, + Provider: oa.Provider, + APIKey: oa.APIKey, + BaseURL: oa.BaseURL, + Model: oa.Model, + MaxTotalTokens: oa.MaxTotalTokens, + MaxCompletionTokens: oa.MaxCompletionTokens, + Reasoning: oa.Reasoning, + } +} + +func NormalizeAIChannelID(s string) string { + id := strings.ToLower(strings.TrimSpace(s)) + id = strings.ReplaceAll(id, "_", "-") + var b strings.Builder + lastDash := false + for _, r := range id { + ok := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') + if ok { + b.WriteRune(r) + lastDash = false + continue + } + if !lastDash { + b.WriteByte('-') + lastDash = true + } + } + out := strings.Trim(b.String(), "-") + if out == "" { + return "default" + } + return out +} + +func (c *AIConfig) EnsureDefaultFromOpenAI(openAI OpenAIConfig) { + if c.Channels == nil { + c.Channels = make(map[string]AIChannelConfig) + } + def := NormalizeAIChannelID(c.DefaultChannel) + if def == "default" && strings.TrimSpace(c.DefaultChannel) == "" { + def = "default" + } + c.DefaultChannel = def + if _, ok := c.Channels[def]; !ok { + c.Channels[def] = AIChannelFromOpenAI(def, "Default", openAI) + } +} + +func (c AIConfig) ResolveChannel(channelID string) (OpenAIConfig, string, bool) { + id := NormalizeAIChannelID(channelID) + if strings.TrimSpace(channelID) == "" { + id = NormalizeAIChannelID(c.DefaultChannel) + } + if id == "" { + id = "default" + } + if c.Channels != nil { + if ch, ok := c.Channels[id]; ok { + return ch.ToOpenAIConfig(), id, true + } + } + return OpenAIConfig{}, id, false +} + +func (c *Config) ResolveAIChannel(channelID string) (OpenAIConfig, string, bool) { + if c == nil { + return OpenAIConfig{}, "", false + } + if oa, id, ok := c.AI.ResolveChannel(channelID); ok { + return oa, id, true + } + return c.OpenAI, NormalizeAIChannelID(channelID), strings.TrimSpace(c.OpenAI.Model) != "" || strings.TrimSpace(c.OpenAI.BaseURL) != "" +} + +func (c *Config) ApplyDefaultAIChannel() { + if c == nil { + return + } + c.AI.EnsureDefaultFromOpenAI(c.OpenAI) + if oa, _, ok := c.AI.ResolveChannel(c.AI.DefaultChannel); ok { + c.OpenAI = oa + } +} + +func (c OpenAIConfig) MaxCompletionTokensEffective() int { + if c.MaxCompletionTokens > 0 { + return c.MaxCompletionTokens + } + return DefaultMaxCompletionTokens +} + +// IsDeepSeekEndpointOrModel reports whether the channel targets DeepSeek's +// official-compatible API or a DeepSeek model family. This is separate from the +// reasoning profile: profile controls field mapping, while DeepSeek has provider +// constraints such as default thinking mode and no tool_choice in thinking mode. +func (c OpenAIConfig) IsDeepSeekEndpointOrModel() bool { + baseURL := strings.ToLower(strings.TrimSpace(c.BaseURL)) + model := strings.ToLower(strings.TrimSpace(c.Model)) + return strings.Contains(baseURL, "deepseek") || strings.Contains(model, "deepseek") +} + +// OpenAIReasoningConfig 全局默认与网关 profile（对话页可通过 ChatRequest.reasoning 覆盖，受 AllowClientReasoning 约束）。 +type OpenAIReasoningConfig struct { + // Mode: auto（默认）| on | off | default（与 auto 相同）。 + // off 在 OpenAI/Claude profile 下省略推理字段；DeepSeek profile 下发送 thinking.type=disabled（其默认开启思考）。 + Mode string `yaml:"mode,omitempty" json:"mode,omitempty"` + // Effort: low | medium | high | max | xhigh；max/xhigh 为不同网关最高档命名，原样下发、不互转。空表示不单独指定强度。 + Effort string `yaml:"effort,omitempty" json:"effort,omitempty"` + // AllowClientReasoning 为 false 时忽略请求体 reasoning；nil 或未设置等同于 true。 + AllowClientReasoning *bool `yaml:"allow_client_reasoning,omitempty" json:"allow_client_reasoning,omitempty"` + // Profile: auto | deepseek_compat | openai_compat | output_config_effort + Profile string `yaml:"profile,omitempty" json:"profile,omitempty"` + // ExtraRequestFields 合并进 Chat Completions 根 JSON（管理员用；与自动字段同名时后者覆盖）。 + // Mode=off 时会移除其中的推理控制字段，但保留其他扩展字段；DeepSeek profile 随后补充显式关闭开关。 + ExtraRequestFields map[string]interface{} `yaml:"extra_request_fields,omitempty" json:"extra_request_fields,omitempty"` +} + +// ModeEffective returns auto when empty or default. +func (c OpenAIReasoningConfig) ModeEffective() string { + m := strings.ToLower(strings.TrimSpace(c.Mode)) + if m == "" || m == "default" { + return "auto" + } + return m +} + +// ProfileEffective returns auto when empty. +func (c OpenAIReasoningConfig) ProfileEffective() string { + p := strings.ToLower(strings.TrimSpace(c.Profile)) + if p == "" { + return "auto" + } + return p +} + +// AllowClientReasoningEffective true when client may send ChatRequest.reasoning. +func (c OpenAIReasoningConfig) AllowClientReasoningEffective() bool { + if c.AllowClientReasoning == nil { + return true + } + return *c.AllowClientReasoning +} + +type FofaConfig struct { + // APIKey 为 FOFA API Key（建议使用只读权限的 Key） + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` // 默认 https://fofa.info/api/v1/search/all +} + +type SpaceSearchConfig struct { + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` +} + +type SecurityConfig struct { + Tools []ToolConfig `yaml:"tools,omitempty"` // 向后兼容：支持在主配置文件中定义工具 + ToolsDir string `yaml:"tools_dir,omitempty"` // 工具配置文件目录（新方式） + ToolDescriptionMode string `yaml:"tool_description_mode,omitempty"` // 工具描述模式: "short" | "full"，默认 short +} + +type DatabaseConfig struct { + Path string `yaml:"path"` // 会话数据库路径 + KnowledgeDBPath string `yaml:"knowledge_db_path,omitempty"` // 知识库数据库路径（可选，为空则使用会话数据库） +} + +type AgentConfig struct { + MaxIterations int `yaml:"max_iterations" json:"max_iterations"` + ToolTimeoutMinutes int `yaml:"tool_timeout_minutes" json:"tool_timeout_minutes"` // 单次工具执行最大时长（分钟），超时自动终止，防止长时间挂起；0 表示不限制（不推荐） + ToolWaitTimeoutSeconds int `yaml:"tool_wait_timeout_seconds" json:"tool_wait_timeout_seconds"` // 工具本轮等待秒数；到时返回 execution_id，worker 继续后台执行；0 表示等到完成 + ExternalMCPMaxConcurrentPerServer int `yaml:"external_mcp_max_concurrent_per_server" json:"external_mcp_max_concurrent_per_server"` // 单个外部 MCP server 同时运行的工具数；0 表示默认 2 + ExternalMCPMaxConcurrentTotal int `yaml:"external_mcp_max_concurrent_total" json:"external_mcp_max_concurrent_total"` // 所有外部 MCP 工具全局并发；0 表示默认 16 + ExternalMCPCircuitFailureThreshold int `yaml:"external_mcp_circuit_failure_threshold" json:"external_mcp_circuit_failure_threshold"` // 单个 MCP server 连续失败多少次后打开熔断；0 表示默认 3；负数关闭 + ExternalMCPCircuitCooldownSeconds int `yaml:"external_mcp_circuit_cooldown_seconds" json:"external_mcp_circuit_cooldown_seconds"` // 熔断后冷却秒数；0 表示默认 60 + // ShellNoOutputTimeoutSeconds execute/exec 无任何 stdout/stderr 时的空闲终止秒数（通用防挂死，不维护命令黑名单）；0=默认 300（5 分钟）；-1=关闭。 + ShellNoOutputTimeoutSeconds int `yaml:"shell_no_output_timeout_seconds" json:"shell_no_output_timeout_seconds"` + // WorkspaceRootDir 会话工作目录根路径（curl/wget 下载、read_file/glob/grep 本地分析）；空=tmp/workspace，其下按 projects/{id} 或 conversations/{id} 隔离。 + WorkspaceRootDir string `yaml:"workspace_root_dir,omitempty" json:"workspace_root_dir,omitempty"` + // SystemPromptPath 单代理系统提示 Markdown/文本文件路径（相对 config.yaml 所在目录，或可写绝对路径）。非空且可读时替换内置单代理提示；留空用内置。 + SystemPromptPath string `yaml:"system_prompt_path,omitempty" json:"system_prompt_path,omitempty"` +} + +// HitlConfig 人机协同全局选项；与会话侧栏/API 中的白名单合并为并集后参与判定。 +// tool_whitelist 可在侧栏「应用」时合并写入 config.yaml 并立即生效。 +// audit_agent_prompt / audit_agent_prompt_review_edit 可在人机协同页编辑并立即生效；空则使用内置默认。 +type HitlConfig struct { + // AuditModel 审计 Agent 专用模型；字段留空时继承 OpenAI 主配置，便于用小模型做审批。 + AuditModel OpenAIConfig `yaml:"audit_model,omitempty" json:"audit_model,omitempty"` + // ToolWhitelist 全局免审批工具名（与白名单内工具不触发 HITL 审批）。 + ToolWhitelist []string `yaml:"tool_whitelist,omitempty" json:"tool_whitelist,omitempty"` + // AuditAgentPrompt 审批模式（approval）下审计 Agent 系统提示词。 + AuditAgentPrompt string `yaml:"audit_agent_prompt,omitempty" json:"audit_agent_prompt,omitempty"` + // AuditAgentPromptReviewEdit 审查编辑模式（review_edit）下审计 Agent 系统提示词。 + AuditAgentPromptReviewEdit string `yaml:"audit_agent_prompt_review_edit,omitempty" json:"audit_agent_prompt_review_edit,omitempty"` + // RetentionDays 已决策审计日志（hitl_interrupts 非 pending）保留天数；省略时默认 90；0 表示不自动清理。 + RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"` + // DefaultReviewer 全局默认审批方（human | audit_agent）；未选会话时切换会写入 config.yaml；新建会话无独立配置时沿用。 + DefaultReviewer string `yaml:"default_reviewer,omitempty" json:"default_reviewer,omitempty"` +} + +// EffectiveDefaultReviewer returns human or audit_agent; omitted or unknown values default to human. +func (h HitlConfig) EffectiveDefaultReviewer() string { + switch strings.ToLower(strings.TrimSpace(h.DefaultReviewer)) { + case "audit_agent", "agent", "ai": + return "audit_agent" + default: + return "human" + } +} + +// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90. +func (h HitlConfig) RetentionDaysEffective() int { + if h.RetentionDays == nil { + return 90 + } + if *h.RetentionDays < 0 { + return 0 + } + return *h.RetentionDays +} + +// AuditModelEffective returns the audit-agent model config with empty fields inherited from the main model config. +func (h HitlConfig) AuditModelEffective(main OpenAIConfig) OpenAIConfig { + out := main + am := h.AuditModel + if strings.TrimSpace(am.Provider) != "" { + out.Provider = strings.TrimSpace(am.Provider) + } + if strings.TrimSpace(am.BaseURL) != "" { + out.BaseURL = strings.TrimSpace(am.BaseURL) + } + if strings.TrimSpace(am.APIKey) != "" { + out.APIKey = strings.TrimSpace(am.APIKey) + } + if strings.TrimSpace(am.Model) != "" { + out.Model = strings.TrimSpace(am.Model) + } + if am.MaxTotalTokens > 0 { + out.MaxTotalTokens = am.MaxTotalTokens + } + if am.MaxCompletionTokens > 0 { + out.MaxCompletionTokens = am.MaxCompletionTokens + } + return out +} + +const hitlAuditAgentPromptBase = `你是 CyberStrikeAI 人机协同审计 Agent。审查 Agent 即将执行的工具调用是否会对系统造成实质性损害。 + +你会收到 JSON，包含 hitlMode、toolName、arguments/argumentsObj、userMessage、thinking、reasoningChain、planning 等字段。 + +规则优先级： +- 必须 reject 的规则 +- 必须 approve 的规则 +- 默认放行规则 + +规则匹配要求： +- 如果同时命中 reject 和 approve，必须 reject。 +- 每次裁决必须给出一个规则编号，不得输出“命中规则：无”。 +- 未命中明确 reject 或 approve 规则，但未发现明确破坏行为时，命中 D1：默认放行。 +- 不得仅凭“本地执行”“设置超时”“文件名包含 test”判断脚本无害；必须检查实际 command、代码或请求参数。 +- 如果脚本具体内容不可见，应表述为“当前参数未显示明确高危操作”，不能断言“无持久化影响”。 + +裁决基调（默认放行）： +- 绝大多数渗透测试操作（包括各类攻击 payload）都应放行 +- 在裁决之前需要判断出：1、此请求的是干什么（理解请求内容）？2、此请求造成什么后果？ +- 仅在「明确会直接破坏目标业务系统正常运行，或造成真实用户数据丢失/不可逆篡改」时 → reject +- 信息不足、参数略宽泛、仅为探测意图时 → approve +- 拿不准时 → approve（宁可放过可恢复的测试，不要误拦合法攻击） + +必须 reject 的高危情形（示例，非穷举）： +- 修改/重置任意用户或管理员密码、强制下线真实用户 +- 修改服务器或服务配置（系统配置、Web/数据库参数、防火墙、启动项、计划任务） +- 修改/创建/删除用户、角色、权限（账户体系篡改） +- 删除或大规模破坏数据/文件（DROP/TRUNCATE/DELETE 全表、rm 重要文件、清空数据库） +- 停止、禁用、重启业务服务（systemctl stop、kill 关键进程、shutdown、reboot） +- 大规模压测或高并发造成 DoS + +以下情形必须 approve（即使看起来“攻击性强”或包含写操作）： +- 各类漏洞探测与利用 payload（SQL 注入含读写型、命令注入、反序列化、文件包含、路径遍历、SSRF、SSTI、XXE 等） +- 文件上传测试，包括上传 webshell、写入 webshell 文件 +- 命令执行探测、反弹 Shell、dnslog / 数据外带探测 +- 数据库命令执行、写文件、写计划任务、内存马注入等 getshell 利用 +- 常规速率的口令爆破（非 DoS 级） +- 信息收集、端口/目录扫描、枚举、抓取 +- 读取、查询、探测类操作（含敏感文件） +- 攻陷后无害侦察命令（whoami/id/uname/ls/cat 等） +- 参数宽泛但无明确破坏意图的请求 + +判定原则： +- “写”本身不是拦截理由——渗透测试大量依赖写操作（注入、上传、爆破、getshell） +- 关键看是否破坏「业务可用性 / 关键数据完整性」 +- 只有明确命中上方 reject 清单才 reject` + +const hitlAuditAgentPromptApprovalOutput = ` +仅输出一行 JSON，不要 markdown 代码块： +{"decision":"approve"|"reject","comment":"实际操作：...；成功后的后果：...；命中规则：..."}` + +const hitlAuditAgentPromptReviewEditOutput = ` +仅输出一行 JSON，不要 markdown 代码块： +{"decision":"approve"|"reject","comment":"实际操作：...；成功后的后果：...；命中规则：...","editedArguments":{...}} + +editedArguments 规则（仅 approve 且需要改参时填写，否则省略该字段）： +- 提供完整替换后的工具参数对象，键名与 argumentsObj 一致 +- 只做最小必要修改以收窄范围、消除风险（如限制 path、去掉危险 flag） +- 禁止扩大攻击面：不得扩大目标范围、提升权限或引入破坏性参数 +- 无法安全改参时应 reject，不要勉强 approve` + +// DefaultHitlAuditAgentPrompt 内置审批模式审计 Agent 提示词。 +func DefaultHitlAuditAgentPrompt() string { + return hitlAuditAgentPromptBase + hitlAuditAgentPromptApprovalOutput +} + +// DefaultHitlAuditAgentPromptReviewEdit 内置审查编辑模式审计 Agent 提示词。 +func DefaultHitlAuditAgentPromptReviewEdit() string { + return hitlAuditAgentPromptBase + hitlAuditAgentPromptReviewEditOutput +} + +// EffectiveAuditAgentPrompt 返回审批模式生效的审计 Agent 提示词。 +func (c HitlConfig) EffectiveAuditAgentPrompt() string { + return c.EffectiveAuditAgentPromptForMode("approval") +} + +// EffectiveAuditAgentPromptForMode 按 HITL 模式返回生效的审计 Agent 提示词。 +func (c HitlConfig) EffectiveAuditAgentPromptForMode(mode string) string { + if normalizeHitlModeForPrompt(mode) == "review_edit" { + if s := strings.TrimSpace(c.AuditAgentPromptReviewEdit); s != "" { + return s + } + return DefaultHitlAuditAgentPromptReviewEdit() + } + if s := strings.TrimSpace(c.AuditAgentPrompt); s != "" { + return s + } + return DefaultHitlAuditAgentPrompt() +} + +func normalizeHitlModeForPrompt(mode string) string { + switch strings.ToLower(strings.TrimSpace(mode)) { + case "review_edit": + return "review_edit" + default: + return "approval" + } +} + +type AuthConfig struct { + SessionDurationHours int `yaml:"session_duration_hours" json:"session_duration_hours"` +} + +// MonitorConfig MCP 状态监控（tool_executions）保留策略。 +type MonitorConfig struct { + // RetentionDays 执行记录保留天数；省略时默认 90；0 表示不自动清理。 + RetentionDays *int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"` +} + +// RetentionDaysEffective returns retention; 0 means keep forever; omitted defaults to 90. +func (m MonitorConfig) RetentionDaysEffective() int { + if m.RetentionDays == nil { + return 90 + } + if *m.RetentionDays < 0 { + return 0 + } + return *m.RetentionDays +} + +// AuditConfig platform operation audit log settings (not chat/tool execution bodies). +type AuditConfig struct { + // Enabled nil or true enables persistence; explicit false disables. + Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"` + RetentionDays int `yaml:"retention_days,omitempty" json:"retention_days,omitempty"` + MaxDetailBytes int `yaml:"max_detail_bytes,omitempty" json:"max_detail_bytes,omitempty"` + // AuthFailureCooldownSeconds: per-IP cooldown for auth login/change_password failure audit rows; -1 disables; 0 uses default 60. + AuthFailureCooldownSeconds int `yaml:"auth_failure_cooldown_seconds,omitempty" json:"auth_failure_cooldown_seconds,omitempty"` +} + +// EnabledEffective returns true unless audit.enabled is explicitly false. +func (a AuditConfig) EnabledEffective() bool { + if a.Enabled == nil { + return true + } + return *a.Enabled +} + +// RetentionDaysEffective returns retention; 0 means keep forever. +func (a AuditConfig) RetentionDaysEffective() int { + if a.RetentionDays < 0 { + return 0 + } + return a.RetentionDays +} + +// MaxDetailBytesEffective caps serialized detail JSON size. +func (a AuditConfig) MaxDetailBytesEffective() int { + if a.MaxDetailBytes <= 0 { + return 8192 + } + return a.MaxDetailBytes +} + +// AuthFailureCooldownEffective returns seconds between duplicate auth-failure audit rows per IP (default 60; -1 disables). +func (a AuditConfig) AuthFailureCooldownEffective() int { + if a.AuthFailureCooldownSeconds < 0 { + return 0 + } + if a.AuthFailureCooldownSeconds == 0 { + return 60 + } + return a.AuthFailureCooldownSeconds +} + +// ExternalMCPConfig 外部MCP配置 +type ExternalMCPConfig struct { + Servers map[string]ExternalMCPServerConfig `yaml:"servers,omitempty" json:"servers,omitempty"` +} + +// ExternalMCPServerConfig 外部MCP服务器配置（遵循官方 MCP 配置格式，兼容 Claude Desktop / Cursor / VS Code）。 +// 所有字符串字段均支持 ${VAR} 和 ${VAR:-default} 环境变量展开语法。 +type ExternalMCPServerConfig struct { + // 传输类型: "stdio" | "sse" | "http"（Streamable HTTP）。 + // stdio 模式可省略，有 command 字段时自动推断。 + Type string `yaml:"type,omitempty" json:"type,omitempty"` + + // stdio 模式配置 + Command string `yaml:"command,omitempty" json:"command,omitempty"` + Args []string `yaml:"args,omitempty" json:"args,omitempty"` + Env map[string]string `yaml:"env,omitempty" json:"env,omitempty"` + + // HTTP/SSE 模式配置 + URL string `yaml:"url,omitempty" json:"url,omitempty"` + Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` + + // 官方标准字段 + Disabled bool `yaml:"disabled,omitempty" json:"disabled,omitempty"` // 禁用服务器（官方字段） + AutoApprove []string `yaml:"autoApprove,omitempty" json:"autoApprove,omitempty"` // 自动批准的工具列表（官方字段） + + // SDK 高级配置（对应 MCP Go SDK 传输层参数） + MaxRetries int `yaml:"max_retries,omitempty" json:"max_retries,omitempty"` // Streamable HTTP 断线重连次数（默认 5） + TerminateDuration int `yaml:"terminate_duration,omitempty" json:"terminate_duration,omitempty"` // stdio 进程优雅关闭等待秒数（默认 5） + KeepAlive int `yaml:"keep_alive,omitempty" json:"keep_alive,omitempty"` // 客户端心跳间隔秒数（0 = 禁用） + + // 通用配置 + Description string `yaml:"description,omitempty" json:"description,omitempty"` + Timeout int `yaml:"timeout,omitempty" json:"timeout,omitempty"` // 连接超时（秒） + ExternalMCPEnable bool `yaml:"external_mcp_enable,omitempty" json:"external_mcp_enable,omitempty"` // 是否启用 + ToolEnabled map[string]bool `yaml:"tool_enabled,omitempty" json:"tool_enabled,omitempty"` // 每个工具的启用状态 +} + +// GetTransportType 返回实际传输类型。优先读 Type，否则根据 Command/URL 自动推断。 +func (c ExternalMCPServerConfig) GetTransportType() string { + if c.Type != "" { + return c.Type + } + if c.Command != "" { + return "stdio" + } + if c.URL != "" { + return "http" + } + return "" +} + +type ToolConfig struct { + Name string `yaml:"name"` + Command string `yaml:"command"` + Args []string `yaml:"args,omitempty"` // 固定参数（可选） + ShortDescription string `yaml:"short_description,omitempty"` // 简短描述（用于工具列表，减少token消耗） + Description string `yaml:"description"` // 详细描述（用于工具文档） + Enabled bool `yaml:"enabled"` + Parameters []ParameterConfig `yaml:"parameters,omitempty"` // 参数定义（可选） + ArgMapping string `yaml:"arg_mapping,omitempty"` // 参数映射方式: "auto", "manual", "template"（可选） + AllowedExitCodes []int `yaml:"allowed_exit_codes,omitempty"` // 允许的退出码列表（某些工具在成功时也返回非零退出码） +} + +// ParameterConfig 参数配置 +type ParameterConfig struct { + Name string `yaml:"name"` // 参数名称 + Type string `yaml:"type"` // 参数类型: string, int, bool, array + Description string `yaml:"description"` // 参数描述 + Required bool `yaml:"required,omitempty"` // 是否必需 + Default interface{} `yaml:"default,omitempty"` // 默认值 + ItemType string `yaml:"item_type,omitempty"` // 当 type 为 array 时，数组元素类型，如 string, number, object + Flag string `yaml:"flag,omitempty"` // 命令行标志，如 "-u", "--url", "-p" + Position *int `yaml:"position,omitempty"` // 位置参数的位置（从0开始） + Format string `yaml:"format,omitempty"` // 参数格式: "flag", "positional", "combined" (flag=value), "template" + Template string `yaml:"template,omitempty"` // 模板字符串，如 "{flag} {value}" 或 "{value}" + Options []string `yaml:"options,omitempty"` // 可选值列表（用于枚举） +} + +func Load(path string) (*Config, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("读取配置文件失败: %w", err) + } + + var cfg Config + if err := yaml.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("解析配置文件失败: %w", err) + } + + if cfg.Auth.SessionDurationHours <= 0 { + cfg.Auth.SessionDurationHours = 12 + } + if cfg.Audit.MaxDetailBytes <= 0 { + cfg.Audit.MaxDetailBytes = 8192 + } + cfg.ApplyDefaultAIChannel() + if err := validateOpenAIOutputLimits(cfg.OpenAI); err != nil { + return nil, err + } + // 如果配置了工具目录，从目录加载工具配置 + if cfg.Security.ToolsDir != "" { + inlineTools := append([]ToolConfig(nil), cfg.Security.Tools...) + toolsDir := ResolveToolsDir(cfg.Security.ToolsDir, path) + merged, err := MergeToolsFromDir(toolsDir, inlineTools) + if err != nil { + return nil, fmt.Errorf("从工具目录加载工具配置失败: %w", err) + } + cfg.Security.Tools = merged + } + + // 外部 MCP：迁移 + 环境变量展开 + if cfg.ExternalMCP.Servers != nil { + for name, serverCfg := range cfg.ExternalMCP.Servers { + // 官方 disabled 字段 → ExternalMCPEnable + if serverCfg.Disabled { + serverCfg.ExternalMCPEnable = false + } else if !serverCfg.ExternalMCPEnable { + // 默认启用 + serverCfg.ExternalMCPEnable = true + } + + // 展开所有 ${VAR} / ${VAR:-default} 环境变量引用 + ExpandConfigEnv(&serverCfg) + + cfg.ExternalMCP.Servers[name] = serverCfg + } + } + + // 从角色目录加载角色配置 + if cfg.RolesDir != "" { + configDir := filepath.Dir(path) + rolesDir := cfg.RolesDir + + // 如果是相对路径，相对于配置文件所在目录 + if !filepath.IsAbs(rolesDir) { + rolesDir = filepath.Join(configDir, rolesDir) + } + + roles, err := LoadRolesFromDir(rolesDir) + if err != nil { + return nil, fmt.Errorf("从角色目录加载角色配置失败: %w", err) + } + + cfg.Roles = roles + } else { + // 如果未配置 roles_dir，初始化为空 map + if cfg.Roles == nil { + cfg.Roles = make(map[string]RoleConfig) + } + } + + if err := ValidateWecomConfig(cfg.Robots.Wecom); err != nil { + return nil, err + } + if err := ValidateRobotsAuthorization(cfg.Robots); err != nil { + return nil, err + } + + return &cfg, nil +} + +func validateOpenAIOutputLimits(openAI OpenAIConfig) error { + if openAI.MaxCompletionTokens < 0 { + return fmt.Errorf("openai.max_completion_tokens 必须为正数") + } + return nil +} + +func EnsureLocalConfig(path string) (EnsureLocalConfigResult, error) { + path = strings.TrimSpace(path) + if path == "" { + path = "config.yaml" + } + + if _, err := os.Stat(path); err == nil { + return EnsureLocalConfigResult{}, nil + } else if !os.IsNotExist(err) { + return EnsureLocalConfigResult{}, fmt.Errorf("检查配置文件失败: %w", err) + } + + examplePath := filepath.Join(filepath.Dir(path), "config.example.yaml") + if _, err := os.Stat(examplePath); err != nil { + if os.IsNotExist(err) { + if alt := "config.example.yaml"; examplePath != alt { + if _, altErr := os.Stat(alt); altErr == nil { + examplePath = alt + } else { + return EnsureLocalConfigResult{}, fmt.Errorf("配置文件 %s 不存在，且未找到模板 %s", path, examplePath) + } + } else { + return EnsureLocalConfigResult{}, fmt.Errorf("配置文件 %s 不存在，且未找到模板 %s", path, examplePath) + } + } else { + return EnsureLocalConfigResult{}, fmt.Errorf("检查配置模板失败: %w", err) + } + } + + data, err := os.ReadFile(examplePath) + if err != nil { + return EnsureLocalConfigResult{}, fmt.Errorf("读取配置模板失败: %w", err) + } + + if dir := filepath.Dir(path); dir != "." && dir != "" { + if err := os.MkdirAll(dir, 0700); err != nil { + return EnsureLocalConfigResult{}, fmt.Errorf("创建配置目录失败: %w", err) + } + } + if err := os.WriteFile(path, data, fs.FileMode(0600)); err != nil { + return EnsureLocalConfigResult{}, fmt.Errorf("创建配置文件失败: %w", err) + } + + return EnsureLocalConfigResult{ + Created: true, + ExamplePath: examplePath, + }, nil +} + +func PrintBootstrapAdminPassword(password string) { + termout.PrintBootstrapAdminCredentials(password) +} + +// generateRandomToken 生成用于 MCP 鉴权的随机字符串（64 位十六进制） +func generateRandomToken() (string, error) { + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return "", err + } + return hex.EncodeToString(b), nil +} + +// persistMCPAuth 将 MCP 的 auth_header / auth_header_value 写回配置文件 +func persistMCPAuth(path string, mcp *MCPConfig) error { + data, err := os.ReadFile(path) + if err != nil { + return err + } + lines := strings.Split(string(data), "\n") + inMcpBlock := false + mcpIndent := -1 + + for i, line := range lines { + trimmed := strings.TrimSpace(line) + if !inMcpBlock { + if strings.HasPrefix(trimmed, "mcp:") { + inMcpBlock = true + mcpIndent = len(line) - len(strings.TrimLeft(line, " ")) + } + continue + } + if trimmed == "" || strings.HasPrefix(trimmed, "#") { + continue + } + leadingSpaces := len(line) - len(strings.TrimLeft(line, " ")) + if leadingSpaces <= mcpIndent { + inMcpBlock = false + mcpIndent = -1 + if strings.HasPrefix(trimmed, "mcp:") { + inMcpBlock = true + mcpIndent = leadingSpaces + } + continue + } + + prefix := line[:leadingSpaces] + rest := strings.TrimSpace(line[leadingSpaces:]) + comment := "" + if idx := strings.Index(line, "#"); idx >= 0 { + comment = strings.TrimRight(line[idx:], " ") + } + withComment := "" + if comment != "" { + if !strings.HasPrefix(comment, " ") { + withComment = " " + } + withComment += comment + } + + if strings.HasPrefix(rest, "auth_header_value:") { + lines[i] = fmt.Sprintf("%sauth_header_value: %q%s", prefix, mcp.AuthHeaderValue, withComment) + } else if strings.HasPrefix(rest, "auth_header:") { + lines[i] = fmt.Sprintf("%sauth_header: %q%s", prefix, mcp.AuthHeader, withComment) + } + } + + return os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0644) +} + +// EnsureMCPAuth only provisions the privileged static service credential when +// global service access was explicitly enabled. +func EnsureMCPAuth(path string, cfg *Config) error { + if !cfg.MCP.Enabled || !cfg.MCP.AllowGlobalAccess || strings.TrimSpace(cfg.MCP.AuthHeaderValue) != "" { + return nil + } + token, err := generateRandomToken() + if err != nil { + return fmt.Errorf("生成 MCP 鉴权密钥失败: %w", err) + } + cfg.MCP.AuthHeaderValue = token + if strings.TrimSpace(cfg.MCP.AuthHeader) == "" { + cfg.MCP.AuthHeader = "X-MCP-Token" + } + return persistMCPAuth(path, &cfg.MCP) +} + +// PrintMCPConfigJSON 向终端输出 MCP 配置的 JSON，可直接复制到 Cursor / Claude Code 的 mcp 配置中使用 +func PrintMCPConfigJSON(mcp MCPConfig) { + if !mcp.Enabled { + return + } + hostForURL := strings.TrimSpace(mcp.Host) + if hostForURL == "" || hostForURL == "0.0.0.0" { + hostForURL = "localhost" + } + url := fmt.Sprintf("http://%s:%d/mcp", hostForURL, mcp.Port) + headers := map[string]string{"Authorization": "Bearer "} + if mcp.AllowGlobalAccess && mcp.AuthHeader != "" { + delete(headers, "Authorization") + headers[mcp.AuthHeader] = mcp.AuthHeaderValue + } + serverEntry := map[string]interface{}{ + "url": url, + } + if len(headers) > 0 { + serverEntry["headers"] = headers + } + // Claude Code 需要 type: "http" + serverEntry["type"] = "http" + out := map[string]interface{}{ + "mcpServers": map[string]interface{}{ + "cyberstrike-ai": serverEntry, + }, + } + b, _ := json.MarshalIndent(out, "", " ") + fmt.Println("[CyberStrikeAI] MCP 配置（可复制到 Cursor / Claude Code 使用）：") + fmt.Println(" Cursor: 放入 ~/.cursor/mcp.json 的 mcpServers，或项目 .cursor/mcp.json") + fmt.Println(" Claude Code: 放入 .mcp.json 或 ~/.claude.json 的 mcpServers") + fmt.Println("----------------------------------------------------------------") + fmt.Println(string(b)) + fmt.Println("----------------------------------------------------------------") +} + +// ResolveToolsDir 将 tools_dir 解析为绝对路径（相对路径相对于 configPath 所在目录）。 +func ResolveToolsDir(toolsDir, configPath string) string { + toolsDir = strings.TrimSpace(toolsDir) + if toolsDir == "" { + return "" + } + if filepath.IsAbs(toolsDir) { + return toolsDir + } + return filepath.Join(filepath.Dir(configPath), toolsDir) +} + +// MergeToolsFromDir 从目录加载工具并与 inline 列表合并：目录中的工具优先，主配置中的工具作为补充。 +func MergeToolsFromDir(toolsDir string, inlineTools []ToolConfig) ([]ToolConfig, error) { + dirTools, err := LoadToolsFromDir(toolsDir) + if err != nil { + return nil, err + } + existing := make(map[string]bool, len(dirTools)) + for _, tool := range dirTools { + existing[tool.Name] = true + } + merged := append([]ToolConfig(nil), dirTools...) + for _, tool := range inlineTools { + if !existing[tool.Name] { + merged = append(merged, tool) + } + } + return merged, nil +} + +// loadInlineSecurityToolsFromYAML 读取 config.yaml 中 security.tools（不含 tools_dir 扫描结果）。 +func loadInlineSecurityToolsFromYAML(configPath string) ([]ToolConfig, error) { + data, err := os.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("读取配置文件失败: %w", err) + } + var partial struct { + Security struct { + Tools []ToolConfig `yaml:"tools"` + } `yaml:"security"` + } + if err := yaml.Unmarshal(data, &partial); err != nil { + return nil, fmt.Errorf("解析配置文件失败: %w", err) + } + if partial.Security.Tools == nil { + return []ToolConfig{}, nil + } + return partial.Security.Tools, nil +} + +// ReloadSecurityToolsFromDir 从 tools_dir 重新加载工具并更新 cfg.Security.Tools（ApplyConfig 热重载用）。 +func ReloadSecurityToolsFromDir(cfg *Config, configPath string) error { + if cfg == nil || strings.TrimSpace(cfg.Security.ToolsDir) == "" { + return nil + } + inlineTools, err := loadInlineSecurityToolsFromYAML(configPath) + if err != nil { + return err + } + toolsDir := ResolveToolsDir(cfg.Security.ToolsDir, configPath) + merged, err := MergeToolsFromDir(toolsDir, inlineTools) + if err != nil { + return fmt.Errorf("从工具目录加载工具配置失败: %w", err) + } + cfg.Security.Tools = merged + return nil +} + +// LoadToolsFromDir 从目录加载所有工具配置文件 +func LoadToolsFromDir(dir string) ([]ToolConfig, error) { + var tools []ToolConfig + + // 检查目录是否存在 + if _, err := os.Stat(dir); os.IsNotExist(err) { + return tools, nil // 目录不存在时返回空列表，不报错 + } + + // 读取目录中的所有 .yaml 和 .yml 文件 + entries, err := os.ReadDir(dir) + if err != nil { + return nil, fmt.Errorf("读取工具目录失败: %w", err) + } + + for _, entry := range entries { + if entry.IsDir() { + continue + } + + name := entry.Name() + if !strings.HasSuffix(name, ".yaml") && !strings.HasSuffix(name, ".yml") { + continue + } + + filePath := filepath.Join(dir, name) + tool, err := LoadToolFromFile(filePath) + if err != nil { + // 记录错误但继续加载其他文件 + fmt.Printf("警告: 加载工具配置文件 %s 失败: %v\n", filePath, err) + continue + } + + tools = append(tools, *tool) + } + + return tools, nil +} + +// LoadToolFromFile 从单个文件加载工具配置 +func LoadToolFromFile(path string) (*ToolConfig, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("读取文件失败: %w", err) + } + + var tool ToolConfig + if err := yaml.Unmarshal(data, &tool); err != nil { + return nil, fmt.Errorf("解析工具配置失败: %w", err) + } + + // 验证必需字段 + if tool.Name == "" { + return nil, fmt.Errorf("工具名称不能为空") + } + if tool.Command == "" { + return nil, fmt.Errorf("工具命令不能为空") + } + + return &tool, nil +} + +// LoadRolesFromDir 从目录加载所有角色配置文件 +func LoadRolesFromDir(dir string) (map[string]RoleConfig, error) { + roles := make(map[string]RoleConfig) + + // 检查目录是否存在 + if _, err := os.Stat(dir); os.IsNotExist(err) { + return roles, nil // 目录不存在时返回空map，不报错 + } + + // 读取目录中的所有 .yaml 和 .yml 文件 + entries, err := os.ReadDir(dir) + if err != nil { + return nil, fmt.Errorf("读取角色目录失败: %w", err) + } + + for _, entry := range entries { + if entry.IsDir() { + continue + } + + name := entry.Name() + if !strings.HasSuffix(name, ".yaml") && !strings.HasSuffix(name, ".yml") { + continue + } + + filePath := filepath.Join(dir, name) + role, err := LoadRoleFromFile(filePath) + if err != nil { + // 记录错误但继续加载其他文件 + fmt.Printf("警告: 加载角色配置文件 %s 失败: %v\n", filePath, err) + continue + } + + // 使用角色名称作为key + roleName := role.Name + if roleName == "" { + // 如果角色名称为空，使用文件名（去掉扩展名）作为名称 + roleName = strings.TrimSuffix(strings.TrimSuffix(name, ".yaml"), ".yml") + role.Name = roleName + } + + roles[roleName] = *role + } + + return roles, nil +} + +// LoadRoleFromFile 从单个文件加载角色配置 +func LoadRoleFromFile(path string) (*RoleConfig, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("读取文件失败: %w", err) + } + + var role RoleConfig + if err := yaml.Unmarshal(data, &role); err != nil { + return nil, fmt.Errorf("解析角色配置失败: %w", err) + } + + // 处理 icon 字段：如果包含 Unicode 转义格式（\U0001F3C6），转换为实际的 Unicode 字符 + // Go 的 yaml 库可能不会自动解析 \U 转义序列，需要手动转换 + if role.Icon != "" { + icon := role.Icon + // 去除可能的引号 + icon = strings.Trim(icon, `"`) + + // 检查是否是 Unicode 转义格式 \U0001F3C6（8位十六进制）或 \uXXXX（4位十六进制） + if len(icon) >= 3 && icon[0] == '\\' { + if icon[1] == 'U' && len(icon) >= 10 { + // \U0001F3C6 格式（8位十六进制） + if codePoint, err := strconv.ParseInt(icon[2:10], 16, 32); err == nil { + role.Icon = string(rune(codePoint)) + } + } else if icon[1] == 'u' && len(icon) >= 6 { + // \uXXXX 格式（4位十六进制） + if codePoint, err := strconv.ParseInt(icon[2:6], 16, 32); err == nil { + role.Icon = string(rune(codePoint)) + } + } + } + } + + // 验证必需字段 + if role.Name == "" { + // 如果名称为空，尝试从文件名获取 + baseName := filepath.Base(path) + role.Name = strings.TrimSuffix(strings.TrimSuffix(baseName, ".yaml"), ".yml") + } + + return &role, nil +} + +func Default() *Config { + strictRobotIdentity := true + return &Config{ + Server: ServerConfig{ + Host: "0.0.0.0", + Port: 8080, + }, + Log: LogConfig{ + Level: "info", + Output: "stdout", + }, + MCP: MCPConfig{ + Enabled: false, + Host: "127.0.0.1", + Port: 8081, + }, + AI: AIConfig{ + DefaultChannel: "default", + Channels: map[string]AIChannelConfig{ + "default": { + Name: "Default", + Provider: "openai_compatible", + BaseURL: "https://api.openai.com/v1", + Model: "gpt-4", + MaxTotalTokens: 120000, + MaxCompletionTokens: DefaultMaxCompletionTokens, + }, + }, + }, + OpenAI: OpenAIConfig{}, + Agent: AgentConfig{ + MaxIterations: 30, // 默认最大迭代次数 + ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟，避免异常长时间占用 + ToolWaitTimeoutSeconds: 60, // 外部 MCP 工具单轮最多等待 60 秒，超时后返回 execution_id 可继续等待 + ExternalMCPMaxConcurrentPerServer: 2, // 单个外部 MCP server 默认最多 2 个工具同时执行 + ExternalMCPMaxConcurrentTotal: 16, // 外部 MCP 工具全局默认最多 16 个同时执行 + ExternalMCPCircuitFailureThreshold: 3, // 单个 server 连续 3 次失败后临时熔断 + ExternalMCPCircuitCooldownSeconds: 60, // 熔断默认冷却 60 秒 + ShellNoOutputTimeoutSeconds: 300, // execute/exec 无新输出空闲终止（秒）；-1 关闭 + }, + Security: SecurityConfig{ + Tools: []ToolConfig{}, // 工具配置应该从 config.yaml 或 tools/ 目录加载 + ToolsDir: "tools", // 默认工具目录 + }, + Database: DatabaseConfig{ + Path: "data/conversations.db", + KnowledgeDBPath: "data/knowledge.db", // 默认知识库数据库路径 + }, + Auth: AuthConfig{ + SessionDurationHours: 12, + }, + Audit: func() AuditConfig { + on := true + return AuditConfig{ + RetentionDays: 90, + MaxDetailBytes: 8192, + Enabled: &on, + } + }(), + Monitor: func() MonitorConfig { + days := 90 + return MonitorConfig{RetentionDays: &days} + }(), + Robots: RobotsConfig{ + Session: RobotSessionConfig{ + StrictUserIdentity: &strictRobotIdentity, + }, + }, + Knowledge: KnowledgeConfig{ + Enabled: true, + BasePath: "knowledge_base", + Embedding: EmbeddingConfig{ + Provider: "openai", + Model: "text-embedding-3-small", + BaseURL: "https://api.openai.com/v1", + }, + Retrieval: RetrievalConfig{ + TopK: 5, + SimilarityThreshold: 0.65, + MultiQuery: MultiQueryConfig{MaxQueries: 4}, + Rerank: RerankConfig{}, + PostRetrieve: PostRetrieveConfig{ + PrefetchTopK: 20, + }, + }, + Indexing: IndexingConfig{ + ChunkStrategy: "markdown_then_recursive", + RequestTimeoutSeconds: 120, + ChunkSize: 768, // 增加到 768，更好的上下文保持 + ChunkOverlap: 50, + MaxChunksPerItem: 20, // 限制单个知识项最多 20 个块，避免消耗过多配额 + BatchSize: 64, + PreferSourceFile: false, + MaxRPM: 100, // 默认 100 RPM，避免 429 错误 + RateLimitDelayMs: 600, // 600ms 间隔，对应 100 RPM + MaxRetries: 3, + RetryDelayMs: 1000, + SubIndexes: nil, + }, + }, + } +} + +// C2Config 内置 C2 模块开关（与知识库 enabled 语义一致：关闭后不初始化监听器、不注册 C2 MCP 工具）。 +type C2Config struct { + // Enabled 为 nil 表示未写配置，按 true 处理（兼容旧 config.yaml） + Enabled *bool `yaml:"enabled,omitempty" json:"enabled,omitempty"` +} + +// EnabledEffective 返回是否启用 C2；未显式配置时默认启用。 +func (c C2Config) EnabledEffective() bool { + if c.Enabled == nil { + return true + } + return *c.Enabled +} + +// C2Public 返回给前端的 C2 状态（仅标量）。 +type C2Public struct { + Enabled bool `json:"enabled"` +} + +// Public 将内部配置转为 API 响应。 +func (c C2Config) Public() C2Public { + return C2Public{Enabled: c.EnabledEffective()} +} + +// C2APIUpdate 设置页/API 更新 C2 开关。 +type C2APIUpdate struct { + Enabled bool `json:"enabled"` +} + +// KnowledgeConfig 知识库配置 +type KnowledgeConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` // 是否启用知识检索 + BasePath string `yaml:"base_path" json:"base_path"` // 知识库路径 + Embedding EmbeddingConfig `yaml:"embedding" json:"embedding"` + Retrieval RetrievalConfig `yaml:"retrieval" json:"retrieval"` + Indexing IndexingConfig `yaml:"indexing,omitempty" json:"indexing,omitempty"` // 索引构建配置 +} + +// IndexingConfig 索引构建配置（用于控制知识库索引构建时的行为） +type IndexingConfig struct { + // ChunkStrategy: "markdown_then_recursive"（默认，Eino Markdown 标题切分后再递归切）或 "recursive"（仅递归切分） + ChunkStrategy string `yaml:"chunk_strategy,omitempty" json:"chunk_strategy,omitempty"` + // RequestTimeoutSeconds 嵌入 HTTP 客户端超时（秒），0 表示使用默认 120 + RequestTimeoutSeconds int `yaml:"request_timeout_seconds,omitempty" json:"request_timeout_seconds,omitempty"` + // 分块配置 + ChunkSize int `yaml:"chunk_size,omitempty" json:"chunk_size,omitempty"` // 每个块的最大 token 数（估算），默认 512 + ChunkOverlap int `yaml:"chunk_overlap,omitempty" json:"chunk_overlap,omitempty"` // 块之间的重叠 token 数，默认 50 + MaxChunksPerItem int `yaml:"max_chunks_per_item,omitempty" json:"max_chunks_per_item,omitempty"` // 单个知识项的最大块数量，0 表示不限制 + + // PreferSourceFile 为 true 时优先用 Eino FileLoader 从 file_path 读原文再索引（与库内 content 不一致时以磁盘为准） + PreferSourceFile bool `yaml:"prefer_source_file,omitempty" json:"prefer_source_file,omitempty"` + + // 速率限制配置（用于避免 API 速率限制） + RateLimitDelayMs int `yaml:"rate_limit_delay_ms,omitempty" json:"rate_limit_delay_ms,omitempty"` // 请求间隔时间（毫秒），0 表示不使用固定延迟 + MaxRPM int `yaml:"max_rpm,omitempty" json:"max_rpm,omitempty"` // 每分钟最大请求数，0 表示不限制 + + // 重试配置（用于处理临时错误） + MaxRetries int `yaml:"max_retries,omitempty" json:"max_retries,omitempty"` // 最大重试次数，默认 3 + RetryDelayMs int `yaml:"retry_delay_ms,omitempty" json:"retry_delay_ms,omitempty"` // 重试间隔（毫秒），默认 1000 + + // BatchSize 嵌入批大小（SQLite 索引写入），0 表示默认 64 + BatchSize int `yaml:"batch_size,omitempty" json:"batch_size,omitempty"` + // SubIndexes 传入 Eino indexer.WithSubIndexes（逻辑分区标记，随 Document 元数据传递） + SubIndexes []string `yaml:"sub_indexes,omitempty" json:"sub_indexes,omitempty"` +} + +// EmbeddingConfig 嵌入配置 +type EmbeddingConfig struct { + Provider string `yaml:"provider" json:"provider"` // 嵌入模型提供商 + Model string `yaml:"model" json:"model"` // 模型名称 + BaseURL string `yaml:"base_url" json:"base_url"` // API Base URL + APIKey string `yaml:"api_key" json:"api_key"` // API Key（从OpenAI配置继承） +} + +// PostRetrieveConfig 检索后处理：固定对正文做规范化去重（最佳实践）、上下文预算截断；PrefetchTopK 用于多取候选再收敛到 top_k。 +type PostRetrieveConfig struct { + // PrefetchTopK 向量检索阶段每条 MultiQuery 变体最多保留的候选数；0 表示使用内置默认 max(top_k*4, 20)。 + PrefetchTopK int `yaml:"prefetch_top_k,omitempty" json:"prefetch_top_k,omitempty"` + // MaxContextChars 返回文档内容总 Unicode 字符数上限（整段 chunk，不截断半段）；0 表示不限制。 + MaxContextChars int `yaml:"max_context_chars,omitempty" json:"max_context_chars,omitempty"` + // MaxContextTokens 返回文档内容总 token 上限（tiktoken，按嵌入模型名映射，失败则 cl100k_base）；0 表示不限制。 + MaxContextTokens int `yaml:"max_context_tokens,omitempty" json:"max_context_tokens,omitempty"` +} + +// MultiQueryConfig Eino MultiQuery 查询改写（始终启用，无关闭开关）。 +type MultiQueryConfig struct { + // MaxQueries LLM 生成的检索变体上限（含原问语义覆盖）；0 表示默认 4。 + MaxQueries int `yaml:"max_queries,omitempty" json:"max_queries,omitempty"` +} + +func (c MultiQueryConfig) MaxQueriesEffective() int { + if c.MaxQueries <= 0 { + return 4 + } + if c.MaxQueries > 8 { + return 8 + } + return c.MaxQueries +} + +// RerankConfig 检索精排（始终启用）；支持 dashscope 与 Cohere 兼容 HTTP API。 +type RerankConfig struct { + // Provider: dashscope | cohere；空则按 base_url 自动推断。 + Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` + Model string `yaml:"model,omitempty" json:"model,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` +} + +func (c RerankConfig) ProviderEffective(baseURL string) string { + p := strings.TrimSpace(strings.ToLower(c.Provider)) + if p != "" { + return p + } + u := strings.ToLower(baseURL) + if strings.Contains(u, "dashscope") { + return "dashscope" + } + return "cohere" +} + +func (c RerankConfig) ModelEffective(provider string) string { + if m := strings.TrimSpace(c.Model); m != "" { + return m + } + if provider == "dashscope" { + return "gte-rerank" + } + return "rerank-multilingual-v3.0" +} + +// RetrievalConfig 检索配置 +type RetrievalConfig struct { + TopK int `yaml:"top_k" json:"top_k"` // 检索Top-K + SimilarityThreshold float64 `yaml:"similarity_threshold" json:"similarity_threshold"` // 余弦相似度阈值 + // SubIndexFilter 非空时仅保留 sub_indexes 含该标签（逗号分隔之一）的行；sub_indexes 为空的旧行仍返回。 + SubIndexFilter string `yaml:"sub_index_filter,omitempty" json:"sub_index_filter,omitempty"` + MultiQuery MultiQueryConfig `yaml:"multi_query" json:"multi_query"` + Rerank RerankConfig `yaml:"rerank" json:"rerank"` + // PostRetrieve 检索后处理（去重、预算截断）；精排在 MultiQuery 融合后执行。 + PostRetrieve PostRetrieveConfig `yaml:"post_retrieve,omitempty" json:"post_retrieve,omitempty"` +} + +// RolesConfig 角色配置（已废弃，使用 map[string]RoleConfig 替代） +// 保留此类型以兼容旧代码，但建议直接使用 map[string]RoleConfig +type RolesConfig struct { + Roles map[string]RoleConfig `yaml:"roles,omitempty" json:"roles,omitempty"` +} + +// RoleConfig 单个角色配置 +type RoleConfig struct { + Name string `yaml:"name" json:"name"` // 角色名称 + Description string `yaml:"description" json:"description"` // 角色描述 + UserPrompt string `yaml:"user_prompt" json:"user_prompt"` // 用户提示词(追加到用户消息前) + Icon string `yaml:"icon,omitempty" json:"icon,omitempty"` // 角色图标（可选） + Tools []string `yaml:"tools,omitempty" json:"tools,omitempty"` // 关联的工具列表（toolKey格式，如 "toolName" 或 "mcpName::toolName"） + MCPs []string `yaml:"mcps,omitempty" json:"mcps,omitempty"` // 向后兼容：关联的MCP服务器列表（已废弃，使用tools替代） + WorkflowID string `yaml:"workflow_id,omitempty" json:"workflow_id,omitempty"` // 可选：绑定工作流 ID + WorkflowVersion string `yaml:"workflow_version,omitempty" json:"workflow_version,omitempty"` // latest 或具体版本号；空等同 latest + WorkflowPolicy string `yaml:"workflow_policy,omitempty" json:"workflow_policy,omitempty"` // auto | off；空且 workflow_id 非空时按 auto + Enabled bool `yaml:"enabled" json:"enabled"` // 是否启用 +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go new file mode 100644 index 00000000..a774f8c6 --- /dev/null +++ b/internal/config/config_test.go @@ -0,0 +1,207 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestEnsureLocalConfigCreatesFromExample(t *testing.T) { + dir := t.TempDir() + examplePath := filepath.Join(dir, "config.example.yaml") + configPath := filepath.Join(dir, "config.yaml") + + example := []byte(`auth: + session_duration_hours: 12 +server: + host: 127.0.0.1 + port: 8080 +`) + if err := os.WriteFile(examplePath, example, 0644); err != nil { + t.Fatalf("write example: %v", err) + } + + result, err := EnsureLocalConfig(configPath) + if err != nil { + t.Fatalf("EnsureLocalConfig: %v", err) + } + if !result.Created { + t.Fatal("Created = false, want true") + } + if result.ExamplePath != examplePath { + t.Fatalf("ExamplePath = %q, want %q", result.ExamplePath, examplePath) + } + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load generated config: %v", err) + } + if cfg.Auth.SessionDurationHours != 12 { + t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours) + } + + second, err := EnsureLocalConfig(configPath) + if err != nil { + t.Fatalf("EnsureLocalConfig existing: %v", err) + } + if second.Created { + t.Fatal("Created = true for existing config, want false") + } +} + +func TestLoadIgnoresLegacyAuthPasswordField(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + initial := strings.Join([]string{ + "auth:", + ` password: "legacy-password"`, + " session_duration_hours: 12", + "server:", + " host: 127.0.0.1", + " port: 8080", + "", + }, "\n") + if err := os.WriteFile(path, []byte(initial), 0644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.Auth.SessionDurationHours != 12 { + t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours) + } +} + +func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) { + main := OpenAIConfig{ + Provider: "openai", + BaseURL: "https://api.example.com/v1", + APIKey: "main-key", + Model: "large-model", + } + + got := (HitlConfig{ + AuditModel: OpenAIConfig{Model: "small-reviewer"}, + }).AuditModelEffective(main) + + if got.Provider != main.Provider || got.BaseURL != main.BaseURL || got.APIKey != main.APIKey { + t.Fatalf("expected provider/base_url/api_key to inherit main config, got %+v", got) + } + if got.Model != "small-reviewer" { + t.Fatalf("expected audit model override, got %q", got.Model) + } +} + +func TestLoadUsesAIDefaultChannelAsRuntimeOpenAI(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "config.yaml") + initial := strings.Join([]string{ + "ai:", + " default_channel: deepseek", + " channels:", + " qwen:", + " name: Qwen", + " provider: openai_compatible", + " base_url: https://dashscope.example/v1", + " api_key: qwen-key", + " model: qwen-max", + " deepseek:", + " name: DeepSeek", + " provider: openai_compatible", + " base_url: https://deepseek.example/v1", + " api_key: deepseek-key", + " model: deepseek-chat", + " max_total_tokens: 64000", + "server:", + " host: 127.0.0.1", + " port: 8080", + "", + }, "\n") + if err := os.WriteFile(path, []byte(initial), 0644); err != nil { + t.Fatalf("write config: %v", err) + } + + cfg, err := Load(path) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.OpenAI.Model != "deepseek-chat" || cfg.OpenAI.APIKey != "deepseek-key" || cfg.OpenAI.MaxTotalTokens != 64000 { + t.Fatalf("runtime OpenAI config did not follow ai.default_channel: %+v", cfg.OpenAI) + } + oa, id, ok := cfg.ResolveAIChannel("qwen") + if !ok || id != "qwen" || oa.Model != "qwen-max" || oa.APIKey != "qwen-key" { + t.Fatalf("ResolveAIChannel(qwen) = (%+v, %q, %v)", oa, id, ok) + } +} + +func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) { + var zero MultiAgentEinoMiddlewareConfig + if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes { + t.Fatalf("default ledger max runes = %d, want %d", got, DefaultSummarizationUserIntentLedgerMaxRunes) + } + if got := zero.SummarizationUserIntentLedgerEntryMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerEntryMaxRunes { + t.Fatalf("default ledger entry max runes = %d, want %d", got, DefaultSummarizationUserIntentLedgerEntryMaxRunes) + } + + custom := MultiAgentEinoMiddlewareConfig{ + SummarizationUserIntentLedgerMaxRunes: 12345, + SummarizationUserIntentLedgerEntryMaxRunes: 2345, + } + if got := custom.SummarizationUserIntentLedgerMaxRunesEffective(); got != 12345 { + t.Fatalf("custom ledger max runes = %d", got) + } + if got := custom.SummarizationUserIntentLedgerEntryMaxRunesEffective(); got != 2345 { + t.Fatalf("custom ledger entry max runes = %d", got) + } +} + +func TestSummarizationOutputReserveTokensEffective(t *testing.T) { + var zero MultiAgentEinoMiddlewareConfig + if got := zero.SummarizationOutputReserveTokensEffective(); got != DefaultSummarizationOutputReserveTokens { + t.Fatalf("default output reserve = %d, want %d", got, DefaultSummarizationOutputReserveTokens) + } + custom := MultiAgentEinoMiddlewareConfig{SummarizationOutputReserveTokens: 4096} + if got := custom.SummarizationOutputReserveTokensEffective(); got != 4096 { + t.Fatalf("custom output reserve = %d", got) + } +} + +func TestOpenAIOutputLimitValidation(t *testing.T) { + if got := (OpenAIConfig{}).MaxCompletionTokensEffective(); got != DefaultMaxCompletionTokens { + t.Fatalf("max completion default=%d", got) + } + if err := validateOpenAIOutputLimits(OpenAIConfig{MaxCompletionTokens: -1}); err == nil { + t.Fatal("negative completion limit must fail") + } +} + +func TestLatestUserMessageRunesEffective(t *testing.T) { + var zero MultiAgentEinoMiddlewareConfig + if got := zero.LatestUserMessageMaxRunesEffective(); got != DefaultLatestUserMessageMaxRunes { + t.Fatalf("default latest user max runes = %d, want %d", got, DefaultLatestUserMessageMaxRunes) + } + if got := zero.LatestUserMessageHeadRunesEffective(); got != DefaultLatestUserMessageHeadRunes { + t.Fatalf("default latest user head runes = %d, want %d", got, DefaultLatestUserMessageHeadRunes) + } + if got := zero.LatestUserMessageTailRunesEffective(); got != DefaultLatestUserMessageTailRunes { + t.Fatalf("default latest user tail runes = %d, want %d", got, DefaultLatestUserMessageTailRunes) + } + + custom := MultiAgentEinoMiddlewareConfig{ + LatestUserMessageMaxRunes: 100, + LatestUserMessageHeadRunes: 40, + LatestUserMessageTailRunes: 60, + } + if got := custom.LatestUserMessageMaxRunesEffective(); got != 100 { + t.Fatalf("custom latest user max runes = %d", got) + } + if got := custom.LatestUserMessageHeadRunesEffective(); got != 40 { + t.Fatalf("custom latest user head runes = %d", got) + } + if got := custom.LatestUserMessageTailRunesEffective(); got != 60 { + t.Fatalf("custom latest user tail runes = %d", got) + } +} diff --git a/internal/config/envexpand.go b/internal/config/envexpand.go new file mode 100644 index 00000000..0ffc1784 --- /dev/null +++ b/internal/config/envexpand.go @@ -0,0 +1,66 @@ +package config + +import ( + "os" + "strings" +) + +// expandEnvVar 展开字符串中的 ${VAR} 和 ${VAR:-default} 环境变量引用。 +// 与官方 MCP 配置格式一致（Claude Desktop / Cursor / VS Code 均支持此语法）。 +func expandEnvVar(s string) string { + var b strings.Builder + i := 0 + for i < len(s) { + // 查找 ${ + idx := strings.Index(s[i:], "${") + if idx < 0 { + b.WriteString(s[i:]) + break + } + b.WriteString(s[i : i+idx]) + i += idx + 2 // skip ${ + + // 查找对应的 } + end := strings.IndexByte(s[i:], '}') + if end < 0 { + // 没有 }，原样保留 + b.WriteString("${") + continue + } + expr := s[i : i+end] + i += end + 1 // skip } + + // 解析 VAR:-default + varName := expr + defaultVal := "" + hasDefault := false + if colonIdx := strings.Index(expr, ":-"); colonIdx >= 0 { + varName = expr[:colonIdx] + defaultVal = expr[colonIdx+2:] + hasDefault = true + } + + val := os.Getenv(varName) + if val == "" && hasDefault { + val = defaultVal + } + b.WriteString(val) + } + return b.String() +} + +// ExpandConfigEnv 展开 ExternalMCPServerConfig 中所有支持环境变量的字段。 +// 展开范围：Command、Args、Env values、URL、Headers values。 +func ExpandConfigEnv(cfg *ExternalMCPServerConfig) { + cfg.Command = expandEnvVar(cfg.Command) + for i, arg := range cfg.Args { + cfg.Args[i] = expandEnvVar(arg) + } + for k, v := range cfg.Env { + cfg.Env[k] = expandEnvVar(v) + } + cfg.URL = expandEnvVar(cfg.URL) + for k, v := range cfg.Headers { + cfg.Headers[k] = expandEnvVar(v) + } +} diff --git a/internal/config/envexpand_test.go b/internal/config/envexpand_test.go new file mode 100644 index 00000000..a17c4514 --- /dev/null +++ b/internal/config/envexpand_test.go @@ -0,0 +1,81 @@ +package config + +import ( + "os" + "testing" +) + +func TestExpandEnvVar(t *testing.T) { + os.Setenv("TEST_MCP_VAR", "hello") + os.Setenv("TEST_MCP_PATH", "/usr/local/bin") + defer os.Unsetenv("TEST_MCP_VAR") + defer os.Unsetenv("TEST_MCP_PATH") + + tests := []struct { + name string + input string + expect string + }{ + {"plain string", "no vars here", "no vars here"}, + {"empty string", "", ""}, + {"simple var", "${TEST_MCP_VAR}", "hello"}, + {"var in middle", "prefix-${TEST_MCP_VAR}-suffix", "prefix-hello-suffix"}, + {"multiple vars", "${TEST_MCP_PATH}/${TEST_MCP_VAR}", "/usr/local/bin/hello"}, + {"missing var empty", "${NONEXISTENT_MCP_VAR_XYZ}", ""}, + {"default value used", "${NONEXISTENT_MCP_VAR_XYZ:-fallback}", "fallback"}, + {"default not used", "${TEST_MCP_VAR:-unused}", "hello"}, + {"default with path", "${NONEXISTENT_MCP_VAR_XYZ:-/tmp/default}", "/tmp/default"}, + {"unclosed brace", "${UNCLOSED", "${UNCLOSED"}, + {"dollar without brace", "$PLAIN", "$PLAIN"}, + {"empty var name", "${}", ""}, + {"default empty var", "${:-default}", "default"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := expandEnvVar(tt.input) + if got != tt.expect { + t.Errorf("expandEnvVar(%q) = %q, want %q", tt.input, got, tt.expect) + } + }) + } +} + +func TestExpandConfigEnv(t *testing.T) { + os.Setenv("TEST_MCP_CMD", "python3") + os.Setenv("TEST_MCP_TOKEN", "secret123") + defer os.Unsetenv("TEST_MCP_CMD") + defer os.Unsetenv("TEST_MCP_TOKEN") + + cfg := &ExternalMCPServerConfig{ + Command: "${TEST_MCP_CMD}", + Args: []string{"--token", "${TEST_MCP_TOKEN}", "${MISSING:-default_arg}"}, + Env: map[string]string{"API_KEY": "${TEST_MCP_TOKEN}", "LEVEL": "${MISSING:-INFO}"}, + URL: "https://${MISSING:-example.com}/mcp", + Headers: map[string]string{"Authorization": "Bearer ${TEST_MCP_TOKEN}"}, + } + + ExpandConfigEnv(cfg) + + if cfg.Command != "python3" { + t.Errorf("Command = %q, want %q", cfg.Command, "python3") + } + if cfg.Args[1] != "secret123" { + t.Errorf("Args[1] = %q, want %q", cfg.Args[1], "secret123") + } + if cfg.Args[2] != "default_arg" { + t.Errorf("Args[2] = %q, want %q", cfg.Args[2], "default_arg") + } + if cfg.Env["API_KEY"] != "secret123" { + t.Errorf("Env[API_KEY] = %q, want %q", cfg.Env["API_KEY"], "secret123") + } + if cfg.Env["LEVEL"] != "INFO" { + t.Errorf("Env[LEVEL] = %q, want %q", cfg.Env["LEVEL"], "INFO") + } + if cfg.URL != "https://example.com/mcp" { + t.Errorf("URL = %q, want %q", cfg.URL, "https://example.com/mcp") + } + if cfg.Headers["Authorization"] != "Bearer secret123" { + t.Errorf("Headers[Authorization] = %q, want %q", cfg.Headers["Authorization"], "Bearer secret123") + } +} diff --git a/internal/config/hitl_prompt_test.go b/internal/config/hitl_prompt_test.go new file mode 100644 index 00000000..d2fbdcb8 --- /dev/null +++ b/internal/config/hitl_prompt_test.go @@ -0,0 +1,31 @@ +package config + +import ( + "strings" + "testing" +) + +func TestDefaultHitlAuditAgentPromptIncludesPrioritizedRules(t *testing.T) { + prompt := DefaultHitlAuditAgentPrompt() + for _, want := range []string{ + "如果同时命中 reject 和 approve，必须 reject", + "修改/重置任意用户或管理员密码", + "修改/创建/删除用户、角色、权限", + "停止、禁用、重启业务服务", + "命中规则：...", + } { + if !strings.Contains(prompt, want) { + t.Fatalf("default approval prompt missing %q", want) + } + } +} + +func TestDefaultHitlAuditAgentPromptReviewEditKeepsEditedArguments(t *testing.T) { + prompt := DefaultHitlAuditAgentPromptReviewEdit() + if !strings.Contains(prompt, `"editedArguments":{...}`) { + t.Fatal("review-edit prompt must preserve editedArguments output") + } + if !strings.Contains(prompt, "命中规则：...") { + t.Fatal("review-edit prompt must require a matched rule") + } +} diff --git a/internal/config/robots_validate_test.go b/internal/config/robots_validate_test.go new file mode 100644 index 00000000..f012ce54 --- /dev/null +++ b/internal/config/robots_validate_test.go @@ -0,0 +1,69 @@ +package config + +import "testing" + +func TestValidateWecomConfig(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + cfg RobotWecomConfig + wantErr bool + }{ + { + name: "disabled without token", + cfg: RobotWecomConfig{Enabled: false, Token: ""}, + wantErr: false, + }, + { + name: "enabled with token", + cfg: RobotWecomConfig{Enabled: true, Token: "secret"}, + wantErr: false, + }, + { + name: "enabled without token", + cfg: RobotWecomConfig{Enabled: true, Token: ""}, + wantErr: true, + }, + { + name: "enabled with whitespace token", + cfg: RobotWecomConfig{Enabled: true, Token: " "}, + wantErr: true, + }, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + err := ValidateWecomConfig(tt.cfg) + if (err != nil) != tt.wantErr { + t.Fatalf("ValidateWecomConfig() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func TestValidateRobotAuthorization(t *testing.T) { + tests := []struct { + name string + cfg RobotAuthorizationConfig + wantErr bool + }{ + {name: "default user binding", cfg: RobotAuthorizationConfig{}, wantErr: false}, + {name: "explicit user binding", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeUserBinding}, wantErr: false}, + {name: "service account", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1", AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: false}, + {name: "missing service user", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: true}, + {name: "admin allowed with exact sender", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "admin", AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: false}, + {name: "allowlist required", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1"}, wantErr: true}, + {name: "wildcard forbidden", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1", AllowedExternalUsers: []string{"*"}}, wantErr: true}, + {name: "unknown mode", cfg: RobotAuthorizationConfig{Mode: "open"}, wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if err := ValidateRobotAuthorization(tt.cfg, "robots.lark"); (err != nil) != tt.wantErr { + t.Fatalf("ValidateRobotAuthorization() error=%v wantErr=%v", err, tt.wantErr) + } + }) + } +} diff --git a/internal/config/server_https_bootstrap.go b/internal/config/server_https_bootstrap.go new file mode 100644 index 00000000..6f67a601 --- /dev/null +++ b/internal/config/server_https_bootstrap.go @@ -0,0 +1,60 @@ +package config + +import "strings" + +// MainWebUIUsesHTTPS 判断主 Web UI 是否以 HTTPS 监听（与 internal/app.prepareMainServerTLS 前置条件一致）。 +func MainWebUIUsesHTTPS(s *ServerConfig) bool { + if s == nil { + return false + } + if s.TLSEnabled { + return true + } + if s.TLSAutoSelfSign { + return true + } + cert := strings.TrimSpace(s.TLSCertPath) + key := strings.TrimSpace(s.TLSKeyPath) + return cert != "" && key != "" +} + +// ServerHTTPRedirectEnabled 是否在主站启用 HTTPS 时把明文 HTTP 请求重定向到 HTTPS（默认开启）。 +func ServerHTTPRedirectEnabled(s *ServerConfig) bool { + if s == nil || !MainWebUIUsesHTTPS(s) { + return false + } + if s.TLSHTTPRedirect == nil { + return true + } + return *s.TLSHTTPRedirect +} + +// ApplyDevHTTPSBootstrap 供 --https / 一键脚本使用：强制开启主站 TLS。 +// 若已配置 tls_cert_path 与 tls_key_path 则仅用 PEM，不开启自签；否则启用 tls_auto_self_sign（内存证书，仅本地测试）。 +func ApplyDevHTTPSBootstrap(cfg *Config) { + if cfg == nil { + return + } + cfg.Server.TLSEnabled = true + cert := strings.TrimSpace(cfg.Server.TLSCertPath) + key := strings.TrimSpace(cfg.Server.TLSKeyPath) + if cert != "" && key != "" { + cfg.Server.TLSAutoSelfSign = false + return + } + cfg.Server.TLSAutoSelfSign = true +} + +// ApplyPlainHTTPBootstrap 供 --http / 一键脚本使用：强制主站使用明文 HTTP。 +// 它会覆盖配置文件中的 TLS 开关、自签证书以及证书路径，避免 --http 仍被配置中的 HTTPS 选项重新启用。 +func ApplyPlainHTTPBootstrap(cfg *Config) { + if cfg == nil { + return + } + cfg.Server.TLSEnabled = false + cfg.Server.TLSAutoSelfSign = false + cfg.Server.TLSCertPath = "" + cfg.Server.TLSKeyPath = "" + disabled := false + cfg.Server.TLSHTTPRedirect = &disabled +} diff --git a/internal/config/server_https_bootstrap_test.go b/internal/config/server_https_bootstrap_test.go new file mode 100644 index 00000000..0a7836ca --- /dev/null +++ b/internal/config/server_https_bootstrap_test.go @@ -0,0 +1,31 @@ +package config + +import "testing" + +func TestApplyPlainHTTPBootstrapDisablesConfiguredTLS(t *testing.T) { + enabled := true + cfg := &Config{ + Server: ServerConfig{ + TLSEnabled: true, + TLSAutoSelfSign: true, + TLSCertPath: "/tmp/server.crt", + TLSKeyPath: "/tmp/server.key", + TLSHTTPRedirect: &enabled, + }, + } + + ApplyPlainHTTPBootstrap(cfg) + + if MainWebUIUsesHTTPS(&cfg.Server) { + t.Fatal("expected --http bootstrap to disable main web UI HTTPS") + } + if ServerHTTPRedirectEnabled(&cfg.Server) { + t.Fatal("expected --http bootstrap to disable HTTP to HTTPS redirect") + } + if cfg.Server.TLSCertPath != "" || cfg.Server.TLSKeyPath != "" { + t.Fatalf("expected TLS cert paths to be cleared, got cert=%q key=%q", cfg.Server.TLSCertPath, cfg.Server.TLSKeyPath) + } + if cfg.Server.TLSHTTPRedirect == nil || *cfg.Server.TLSHTTPRedirect { + t.Fatal("expected TLSHTTPRedirect to be explicitly disabled") + } +} diff --git a/internal/config/tools_reload_test.go b/internal/config/tools_reload_test.go new file mode 100644 index 00000000..0fbbd074 --- /dev/null +++ b/internal/config/tools_reload_test.go @@ -0,0 +1,111 @@ +package config + +import ( + "os" + "path/filepath" + "testing" +) + +func TestReloadSecurityToolsFromDir(t *testing.T) { + root := t.TempDir() + toolsDir := filepath.Join(root, "tools") + if err := os.MkdirAll(toolsDir, 0755); err != nil { + t.Fatal(err) + } + + configPath := filepath.Join(root, "config.yaml") + if err := os.WriteFile(configPath, []byte(`security: + tools_dir: tools + tools: + - name: inline-only + command: inline-cmd + enabled: true + description: inline tool +`), 0644); err != nil { + t.Fatal(err) + } + + writeTool := func(name, command string) { + t.Helper() + content := "name: " + name + "\ncommand: " + command + "\nenabled: true\ndescription: test\n" + if err := os.WriteFile(filepath.Join(toolsDir, name+".yaml"), []byte(content), 0644); err != nil { + t.Fatal(err) + } + } + + writeTool("alpha", "alpha-cmd") + + cfg := &Config{ + Security: SecurityConfig{ + ToolsDir: "tools", + Tools: []ToolConfig{ + {Name: "stale", Command: "stale-cmd", Enabled: true, Description: "should be removed"}, + }, + }, + } + + if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil { + t.Fatalf("reload: %v", err) + } + if len(cfg.Security.Tools) != 2 { + t.Fatalf("expected 2 tools, got %d", len(cfg.Security.Tools)) + } + + names := map[string]string{} + for _, tool := range cfg.Security.Tools { + names[tool.Name] = tool.Command + } + if names["alpha"] != "alpha-cmd" { + t.Fatalf("alpha tool missing or wrong command: %#v", names) + } + if names["inline-only"] != "inline-cmd" { + t.Fatalf("inline-only tool missing: %#v", names) + } + if _, ok := names["stale"]; ok { + t.Fatal("stale in-memory tool should not survive reload") + } + + writeTool("beta", "beta-cmd") + if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil { + t.Fatalf("second reload: %v", err) + } + if len(cfg.Security.Tools) != 3 { + t.Fatalf("expected 3 tools after add, got %d", len(cfg.Security.Tools)) + } + foundBeta := false + for _, tool := range cfg.Security.Tools { + if tool.Name == "beta" { + foundBeta = true + break + } + } + if !foundBeta { + t.Fatal("beta tool not found after second reload") + } +} + +func TestMergeToolsFromDir_DirOverridesInline(t *testing.T) { + root := t.TempDir() + toolsDir := filepath.Join(root, "tools") + if err := os.MkdirAll(toolsDir, 0755); err != nil { + t.Fatal(err) + } + content := "name: shared\ncommand: dir-cmd\nenabled: true\ndescription: from dir\n" + if err := os.WriteFile(filepath.Join(toolsDir, "shared.yaml"), []byte(content), 0644); err != nil { + t.Fatal(err) + } + + inline := []ToolConfig{ + {Name: "shared", Command: "inline-cmd", Enabled: true, Description: "from inline"}, + } + merged, err := MergeToolsFromDir(toolsDir, inline) + if err != nil { + t.Fatal(err) + } + if len(merged) != 1 { + t.Fatalf("expected 1 tool, got %d", len(merged)) + } + if merged[0].Command != "dir-cmd" { + t.Fatalf("dir tool should win, got command %q", merged[0].Command) + } +} diff --git a/internal/config/vision.go b/internal/config/vision.go new file mode 100644 index 00000000..1052d3b9 --- /dev/null +++ b/internal/config/vision.go @@ -0,0 +1,97 @@ +package config + +import "strings" + +// VisionConfig 独立视觉模型与 analyze_image 工具参数；enabled 时注册 MCP 工具 analyze_image。 +type VisionConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` + APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"` + BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"` + Model string `yaml:"model,omitempty" json:"model,omitempty"` + Provider string `yaml:"provider,omitempty" json:"provider,omitempty"` + TimeoutSeconds int `yaml:"timeout_seconds,omitempty" json:"timeout_seconds,omitempty"` + MaxImageBytes int64 `yaml:"max_image_bytes,omitempty" json:"max_image_bytes,omitempty"` + MaxDimension int `yaml:"max_dimension,omitempty" json:"max_dimension,omitempty"` + JPEGQuality int `yaml:"jpeg_quality,omitempty" json:"jpeg_quality,omitempty"` + MaxPayloadBytes int64 `yaml:"max_payload_bytes,omitempty" json:"max_payload_bytes,omitempty"` + SkipPreprocessBelowBytes int64 `yaml:"skip_preprocess_below_bytes,omitempty" json:"skip_preprocess_below_bytes,omitempty"` // 0=始终压缩；默认 2MB 且长边已<=max_dimension 时原图直传 + Detail string `yaml:"detail,omitempty" json:"detail,omitempty"` // low | high | auto +} + +func (v VisionConfig) TimeoutSecondsEffective() int { + if v.TimeoutSeconds <= 0 { + return 60 + } + return v.TimeoutSeconds +} + +func (v VisionConfig) MaxImageBytesEffective() int64 { + if v.MaxImageBytes <= 0 { + return 5 * 1024 * 1024 + } + return v.MaxImageBytes +} + +func (v VisionConfig) MaxDimensionEffective() int { + if v.MaxDimension <= 0 { + return 2048 + } + return v.MaxDimension +} + +func (v VisionConfig) JPEGQualityEffective() int { + if v.JPEGQuality <= 0 || v.JPEGQuality > 100 { + return 82 + } + return v.JPEGQuality +} + +func (v VisionConfig) MaxPayloadBytesEffective() int64 { + if v.MaxPayloadBytes <= 0 { + return 512 * 1024 + } + return v.MaxPayloadBytes +} + +// SkipPreprocessBelowBytesEffective 低于该字节数且长边<=max_dimension、且<=max_payload 时可原图直传；0 表示始终压缩。 +func (v VisionConfig) SkipPreprocessBelowBytesEffective() int64 { + if v.SkipPreprocessBelowBytes < 0 { + return 0 + } + return v.SkipPreprocessBelowBytes +} + +func (v VisionConfig) DetailEffective() string { + d := strings.ToLower(strings.TrimSpace(v.Detail)) + switch d { + case "high", "low", "auto": + return d + default: + return "low" + } +} + +// OpenAICfgEffective 合并主 openai 配置与 vision 覆盖项，供 VL ChatModel 使用。 +// vision.api_key / base_url / provider 留空或省略时，沿用 main（openai）对应字段；vision.model 必填（由 Ready 校验）。 +func (v VisionConfig) OpenAICfgEffective(main OpenAIConfig) OpenAIConfig { + out := main + if k := strings.TrimSpace(v.APIKey); k != "" { + out.APIKey = k + } + if u := strings.TrimSpace(v.BaseURL); u != "" { + out.BaseURL = u + } + if m := strings.TrimSpace(v.Model); m != "" { + out.Model = m + } + if p := strings.TrimSpace(v.Provider); p != "" { + out.Provider = p + } + out.Reasoning.Mode = "off" + return out +} + +// Ready 表示已启用且模型名非空。 +func (v VisionConfig) Ready() bool { + return v.Enabled && strings.TrimSpace(v.Model) != "" +} diff --git a/internal/config/vision_test.go b/internal/config/vision_test.go new file mode 100644 index 00000000..0620a181 --- /dev/null +++ b/internal/config/vision_test.go @@ -0,0 +1,55 @@ +package config + +import "testing" + +func TestVisionConfig_OpenAICfgEffective_fallbackToMain(t *testing.T) { + main := OpenAIConfig{ + APIKey: "main-key", + BaseURL: "https://main.example/v1", + Model: "main-model", + Provider: "openai", + } + v := VisionConfig{Model: "qwen-vl-max"} + out := v.OpenAICfgEffective(main) + if out.APIKey != main.APIKey || out.BaseURL != main.BaseURL || out.Provider != main.Provider { + t.Fatalf("expected openai fallback, got key=%q url=%q provider=%q", out.APIKey, out.BaseURL, out.Provider) + } + if out.Model != "qwen-vl-max" { + t.Fatalf("model: %s", out.Model) + } +} + +func TestVisionConfig_OpenAICfgEffective(t *testing.T) { + main := OpenAIConfig{ + APIKey: "main-key", + BaseURL: "https://main.example/v1", + Model: "main-model", + Provider: "openai", + Reasoning: OpenAIReasoningConfig{Mode: "on"}, + } + v := VisionConfig{ + Model: "vl-model", + APIKey: "vl-key", + BaseURL: "https://vl.example/v1", + Provider: "claude", + } + out := v.OpenAICfgEffective(main) + if out.APIKey != "vl-key" || out.BaseURL != "https://vl.example/v1" || out.Model != "vl-model" { + t.Fatalf("unexpected merge: %+v", out) + } + if out.Provider != "claude" { + t.Fatalf("provider: %s", out.Provider) + } + if out.Reasoning.Mode != "off" { + t.Fatalf("reasoning should be off for vision, got %s", out.Reasoning.Mode) + } +} + +func TestVisionConfig_Ready(t *testing.T) { + if (VisionConfig{Enabled: true, Model: "x"}).Ready() != true { + t.Fatal("expected ready") + } + if (VisionConfig{Enabled: true}).Ready() != false { + t.Fatal("expected not ready without model") + } +} diff --git a/internal/mcp/builtin/constants.go b/internal/mcp/builtin/constants.go new file mode 100644 index 00000000..d38bcaae --- /dev/null +++ b/internal/mcp/builtin/constants.go @@ -0,0 +1,195 @@ +package builtin + +// 内置工具名称常量 +// 所有代码中使用内置工具名称的地方都应该使用这些常量，而不是硬编码字符串 +const ( + // 漏洞管理工具 + ToolRecordVulnerability = "record_vulnerability" + ToolListVulnerabilities = "list_vulnerabilities" + ToolGetVulnerability = "get_vulnerability" + + // 资产管理工具 + ToolCreateAsset = "create_asset" + ToolGetAsset = "get_asset" + ToolQueryAssets = "query_assets" + ToolUpdateAsset = "update_asset" + ToolDeleteAsset = "delete_asset" + ToolCompleteAssetScan = "complete_asset_scan" + + // 项目黑板（事实）工具 + ToolUpsertProjectFact = "upsert_project_fact" + ToolGetProjectFact = "get_project_fact" + ToolListProjectFacts = "list_project_facts" + ToolSearchProjectFacts = "search_project_facts" + ToolDeprecateProjectFact = "deprecate_project_fact" + ToolRestoreProjectFact = "restore_project_fact" + + // 知识库工具 + ToolListKnowledgeRiskTypes = "list_knowledge_risk_types" + ToolSearchKnowledgeBase = "search_knowledge_base" + + // 视觉分析（本地图片 → VL 模型 → 文本摘要） + ToolAnalyzeImage = "analyze_image" + + // 长耗时工具执行控制（后台 execution 查询/等待/取消） + ToolGetToolExecution = "get_tool_execution" + ToolWaitToolExecution = "wait_tool_execution" + ToolCancelToolExecution = "cancel_tool_execution" + + // WebShell 助手工具（AI 在 WebShell 管理 - AI 助手 中使用） + ToolWebshellExec = "webshell_exec" + ToolWebshellFileList = "webshell_file_list" + ToolWebshellFileRead = "webshell_file_read" + ToolWebshellFileWrite = "webshell_file_write" + + // WebShell 连接管理工具（用于通过 MCP 管理 webshell 连接） + ToolManageWebshellList = "manage_webshell_list" + ToolManageWebshellAdd = "manage_webshell_add" + ToolManageWebshellUpdate = "manage_webshell_update" + ToolManageWebshellDelete = "manage_webshell_delete" + ToolManageWebshellTest = "manage_webshell_test" + + // 批量任务队列（与 Web 端批量任务一致，供模型创建/启停/查询队列） + ToolBatchTaskList = "batch_task_list" + ToolBatchTaskGet = "batch_task_get" + ToolBatchTaskCreate = "batch_task_create" + ToolBatchTaskStart = "batch_task_start" + ToolBatchTaskRerun = "batch_task_rerun" + ToolBatchTaskPause = "batch_task_pause" + ToolBatchTaskDelete = "batch_task_delete" + ToolBatchTaskUpdateMetadata = "batch_task_update_metadata" + ToolBatchTaskUpdateSchedule = "batch_task_update_schedule" + ToolBatchTaskScheduleEnabled = "batch_task_schedule_enabled" + ToolBatchTaskAdd = "batch_task_add_task" + ToolBatchTaskUpdate = "batch_task_update_task" + ToolBatchTaskRemove = "batch_task_remove_task" + + // C2 工具集（合并同类项，8 个统一工具） + ToolC2Listener = "c2_listener" // 监听器管理（create/start/stop/list/get/update/delete） + ToolC2Session = "c2_session" // 会话管理（list/get/set_sleep/kill/delete） + ToolC2Task = "c2_task" // 任务下发（统一 task_type 参数） + ToolC2TaskManage = "c2_task_manage" // 任务管理（get_result/wait/list/cancel） + ToolC2Payload = "c2_payload" // Payload 生成（oneliner/build） + ToolC2Event = "c2_event" // 事件查询 + ToolC2Profile = "c2_profile" // Malleable Profile 管理（list/get/create/update/delete） + ToolC2File = "c2_file" // 文件管理（list/get_result） +) + +// IsBuiltinTool 检查工具名称是否是内置工具 +func IsBuiltinTool(toolName string) bool { + switch toolName { + case ToolRecordVulnerability, + ToolListVulnerabilities, + ToolGetVulnerability, + ToolCreateAsset, + ToolGetAsset, + ToolQueryAssets, + ToolUpdateAsset, + ToolDeleteAsset, + ToolCompleteAssetScan, + ToolUpsertProjectFact, + ToolGetProjectFact, + ToolListProjectFacts, + ToolSearchProjectFacts, + ToolDeprecateProjectFact, + ToolRestoreProjectFact, + ToolListKnowledgeRiskTypes, + ToolSearchKnowledgeBase, + ToolAnalyzeImage, + ToolGetToolExecution, + ToolWaitToolExecution, + ToolCancelToolExecution, + ToolWebshellExec, + ToolWebshellFileList, + ToolWebshellFileRead, + ToolWebshellFileWrite, + ToolManageWebshellList, + ToolManageWebshellAdd, + ToolManageWebshellUpdate, + ToolManageWebshellDelete, + ToolManageWebshellTest, + ToolBatchTaskList, + ToolBatchTaskGet, + ToolBatchTaskCreate, + ToolBatchTaskStart, + ToolBatchTaskRerun, + ToolBatchTaskPause, + ToolBatchTaskDelete, + ToolBatchTaskUpdateMetadata, + ToolBatchTaskUpdateSchedule, + ToolBatchTaskScheduleEnabled, + ToolBatchTaskAdd, + ToolBatchTaskUpdate, + ToolBatchTaskRemove, + // C2 工具 + ToolC2Listener, + ToolC2Session, + ToolC2Task, + ToolC2TaskManage, + ToolC2Payload, + ToolC2Event, + ToolC2Profile, + ToolC2File: + return true + default: + return false + } +} + +// GetAllBuiltinTools 返回所有内置工具名称列表 +func GetAllBuiltinTools() []string { + return []string{ + ToolRecordVulnerability, + ToolListVulnerabilities, + ToolGetVulnerability, + ToolCreateAsset, + ToolGetAsset, + ToolQueryAssets, + ToolUpdateAsset, + ToolDeleteAsset, + ToolCompleteAssetScan, + ToolUpsertProjectFact, + ToolGetProjectFact, + ToolListProjectFacts, + ToolSearchProjectFacts, + ToolDeprecateProjectFact, + ToolRestoreProjectFact, + ToolListKnowledgeRiskTypes, + ToolSearchKnowledgeBase, + ToolAnalyzeImage, + ToolGetToolExecution, + ToolWaitToolExecution, + ToolCancelToolExecution, + ToolWebshellExec, + ToolWebshellFileList, + ToolWebshellFileRead, + ToolWebshellFileWrite, + ToolManageWebshellList, + ToolManageWebshellAdd, + ToolManageWebshellUpdate, + ToolManageWebshellDelete, + ToolManageWebshellTest, + ToolBatchTaskList, + ToolBatchTaskGet, + ToolBatchTaskCreate, + ToolBatchTaskStart, + ToolBatchTaskRerun, + ToolBatchTaskPause, + ToolBatchTaskDelete, + ToolBatchTaskUpdateMetadata, + ToolBatchTaskUpdateSchedule, + ToolBatchTaskScheduleEnabled, + ToolBatchTaskAdd, + ToolBatchTaskUpdate, + ToolBatchTaskRemove, + // C2 工具 + ToolC2Listener, + ToolC2Session, + ToolC2Task, + ToolC2TaskManage, + ToolC2Payload, + ToolC2Event, + ToolC2Profile, + ToolC2File, + } +} diff --git a/internal/mcp/client_sdk.go b/internal/mcp/client_sdk.go new file mode 100644 index 00000000..0d7ebfb3 --- /dev/null +++ b/internal/mcp/client_sdk.go @@ -0,0 +1,475 @@ +// Package mcp 外部 MCP 客户端 - 基于官方 go-sdk 实现，保证协议兼容性 +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "os/exec" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "go.uber.org/zap" +) + +const ( + clientName = "CyberStrikeAI" + clientVersion = "1.0.0" +) + +// sdkClient 基于官方 MCP Go SDK 的外部 MCP 客户端，实现 ExternalMCPClient 接口 +type sdkClient struct { + session *mcp.ClientSession + client *mcp.Client + logger *zap.Logger + mu sync.RWMutex + status string // "disconnected", "connecting", "connected", "error" +} + +// newSDKClientFromSession 用已连接成功的 session 构造（供 createSDKClient 内部使用） +func newSDKClientFromSession(session *mcp.ClientSession, client *mcp.Client, logger *zap.Logger) *sdkClient { + return &sdkClient{ + session: session, + client: client, + logger: logger, + status: "connected", + } +} + +// lazySDKClient 延迟连接：Initialize() 时才调用官方 SDK 建立连接，对外实现 ExternalMCPClient +type lazySDKClient struct { + serverCfg config.ExternalMCPServerConfig + logger *zap.Logger + sessionCancel context.CancelFunc + inner ExternalMCPClient // connected SDK client + mu sync.RWMutex + status string +} + +func newLazySDKClient(serverCfg config.ExternalMCPServerConfig, logger *zap.Logger) *lazySDKClient { + return &lazySDKClient{ + serverCfg: serverCfg, + logger: logger, + status: "connecting", + } +} + +func (c *lazySDKClient) setStatus(s string) { + c.mu.Lock() + defer c.mu.Unlock() + c.status = s +} + +func (c *lazySDKClient) GetStatus() string { + c.mu.RLock() + defer c.mu.RUnlock() + if c.inner != nil { + return c.inner.GetStatus() + } + return c.status +} + +func (c *lazySDKClient) IsConnected() bool { + c.mu.RLock() + inner := c.inner + c.mu.RUnlock() + if inner != nil { + return inner.IsConnected() + } + return false +} + +func (c *lazySDKClient) Initialize(ctx context.Context) error { + c.mu.Lock() + if c.inner != nil { + c.mu.Unlock() + return nil + } + c.mu.Unlock() + + sessionCtx, sessionCancel := context.WithCancel(context.Background()) + type connectResult struct { + inner ExternalMCPClient + err error + } + resultCh := make(chan connectResult) + abandoned := make(chan struct{}) + go func() { + inner, err := createSDKClient(sessionCtx, c.serverCfg, c.logger) + select { + case resultCh <- connectResult{inner: inner, err: err}: + case <-abandoned: + if inner != nil { + _ = inner.Close() + } + sessionCancel() + } + }() + + var result connectResult + select { + case result = <-resultCh: + case <-ctx.Done(): + close(abandoned) + sessionCancel() + c.setStatus("error") + return ctx.Err() + } + + if err := ctx.Err(); err != nil { + sessionCancel() + if result.inner != nil { + _ = result.inner.Close() + } + c.setStatus("error") + return err + } + + if result.err != nil { + sessionCancel() + c.setStatus("error") + return result.err + } + + c.mu.Lock() + if c.inner != nil { + c.mu.Unlock() + sessionCancel() + if result.inner != nil { + _ = result.inner.Close() + } + return nil + } + c.inner = result.inner + c.sessionCancel = sessionCancel + c.mu.Unlock() + c.setStatus("connected") + return nil +} + +func (c *lazySDKClient) ListTools(ctx context.Context) ([]Tool, error) { + c.mu.RLock() + inner := c.inner + c.mu.RUnlock() + if inner == nil { + return nil, fmt.Errorf("未连接") + } + return inner.ListTools(ctx) +} + +func (c *lazySDKClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) { + c.mu.RLock() + inner := c.inner + c.mu.RUnlock() + if inner == nil { + return nil, fmt.Errorf("未连接") + } + return inner.CallTool(ctx, name, args) +} + +func (c *lazySDKClient) Close() error { + c.mu.Lock() + inner := c.inner + sessionCancel := c.sessionCancel + c.inner = nil + c.sessionCancel = nil + c.mu.Unlock() + c.setStatus("disconnected") + if sessionCancel != nil { + sessionCancel() + } + if inner != nil { + return inner.Close() + } + return nil +} + +// markDisconnected 在检测到传输层断连时关闭底层 session，避免 IsConnected 仍返回 true。 +func (c *lazySDKClient) markDisconnected() { + c.mu.Lock() + inner := c.inner + sessionCancel := c.sessionCancel + c.inner = nil + c.sessionCancel = nil + c.mu.Unlock() + if sessionCancel != nil { + sessionCancel() + } + if inner != nil { + _ = inner.Close() + } + c.setStatus("disconnected") +} + +func (c *sdkClient) setStatus(s string) { + c.mu.Lock() + defer c.mu.Unlock() + c.status = s +} + +func (c *sdkClient) GetStatus() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.status +} + +func (c *sdkClient) IsConnected() bool { + return c.GetStatus() == "connected" +} + +func (c *sdkClient) Initialize(ctx context.Context) error { + // sdkClient 由 createSDKClient 在 Connect 成功后才创建，因此 Initialize 时已经连接 + // 此方法仅用于满足 ExternalMCPClient 接口，实际连接在 createSDKClient 中完成 + return nil +} + +func (c *sdkClient) ListTools(ctx context.Context) ([]Tool, error) { + if c.session == nil { + return nil, fmt.Errorf("未连接") + } + res, err := c.session.ListTools(ctx, nil) + if err != nil { + return nil, err + } + if res == nil { + return nil, nil + } + return sdkToolsToOur(res.Tools), nil +} + +func (c *sdkClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) { + if c.session == nil { + return nil, fmt.Errorf("未连接") + } + params := &mcp.CallToolParams{ + Name: name, + Arguments: args, + } + res, err := c.session.CallTool(ctx, params) + if err != nil { + return nil, err + } + return sdkCallToolResultToOurs(res), nil +} + +func (c *sdkClient) Close() error { + c.setStatus("disconnected") + if c.session != nil { + err := c.session.Close() + c.session = nil + return err + } + return nil +} + +// sdkToolsToOur 将 SDK 的 []*mcp.Tool 转为我们的 []Tool +func sdkToolsToOur(tools []*mcp.Tool) []Tool { + if len(tools) == 0 { + return nil + } + out := make([]Tool, 0, len(tools)) + for _, t := range tools { + if t == nil { + continue + } + schema := make(map[string]interface{}) + if t.InputSchema != nil { + // SDK InputSchema 可能为 *jsonschema.Schema 或 map，统一转为 map + if m, ok := t.InputSchema.(map[string]interface{}); ok { + schema = m + } else { + _ = json.Unmarshal(mustJSON(t.InputSchema), &schema) + } + } + desc := t.Description + shortDesc := desc + if t.Annotations != nil && t.Annotations.Title != "" { + shortDesc = t.Annotations.Title + } + out = append(out, Tool{ + Name: t.Name, + Description: desc, + ShortDescription: shortDesc, + InputSchema: schema, + }) + } + return out +} + +// sdkCallToolResultToOurs 将 SDK 的 *mcp.CallToolResult 转为我们的 *ToolResult +func sdkCallToolResultToOurs(res *mcp.CallToolResult) *ToolResult { + if res == nil { + return &ToolResult{Content: []Content{}} + } + content := sdkContentToOurs(res.Content) + return &ToolResult{ + Content: content, + IsError: res.IsError, + } +} + +func sdkContentToOurs(list []mcp.Content) []Content { + if len(list) == 0 { + return nil + } + out := make([]Content, 0, len(list)) + for _, c := range list { + switch v := c.(type) { + case *mcp.TextContent: + out = append(out, Content{Type: "text", Text: v.Text}) + default: + out = append(out, Content{Type: "text", Text: fmt.Sprintf("%v", c)}) + } + } + return out +} + +func mustJSON(v interface{}) []byte { + b, _ := json.Marshal(v) + return b +} + +// createSDKClient 根据配置创建并连接外部 MCP 客户端（使用官方 SDK），返回实现 ExternalMCPClient 的 *sdkClient +// 若连接失败返回 (nil, error)。ctx 用于连接超时与取消。 +func createSDKClient(ctx context.Context, serverCfg config.ExternalMCPServerConfig, logger *zap.Logger) (ExternalMCPClient, error) { + timeout := time.Duration(serverCfg.Timeout) * time.Second + if timeout <= 0 { + timeout = 30 * time.Second + } + + transport := serverCfg.GetTransportType() + if transport == "" { + return nil, fmt.Errorf("配置缺少 command 或 url，且未指定 type/transport") + } + + // 构造 ClientOptions：KeepAlive 心跳 + var clientOpts *mcp.ClientOptions + if serverCfg.KeepAlive > 0 { + clientOpts = &mcp.ClientOptions{ + KeepAlive: time.Duration(serverCfg.KeepAlive) * time.Second, + } + } + + client := mcp.NewClient(&mcp.Implementation{ + Name: clientName, + Version: clientVersion, + }, clientOpts) + + var t mcp.Transport + switch transport { + case "stdio": + if serverCfg.Command == "" { + return nil, fmt.Errorf("stdio 模式需要配置 command") + } + // 必须用 exec.Command 而非 CommandContext：doConnect 返回后 ctx 会被 cancel， + // 若用 CommandContext(ctx) 会立刻杀掉子进程，导致 ListTools 等后续请求失败、显示 0 工具 + cmd := exec.Command(serverCfg.Command, serverCfg.Args...) + if len(serverCfg.Env) > 0 { + cmd.Env = append(cmd.Env, envMapToSlice(serverCfg.Env)...) + } + ct := &mcp.CommandTransport{Command: cmd} + if serverCfg.TerminateDuration > 0 { + ct.TerminateDuration = time.Duration(serverCfg.TerminateDuration) * time.Second + } + t = ct + case "sse": + if serverCfg.URL == "" { + return nil, fmt.Errorf("sse 模式需要配置 url") + } + // SSE 是长连接（GET 流持续打开），不能设置 http.Client.Timeout（会在超时后杀掉整个连接导致 EOF）。 + // 超时由每次 ListTools/CallTool 的 context 单独控制。 + httpClient := httpClientForLongLived(serverCfg.Headers) + t = &mcp.SSEClientTransport{ + Endpoint: serverCfg.URL, + HTTPClient: httpClient, + } + case "http": + if serverCfg.URL == "" { + return nil, fmt.Errorf("http 模式需要配置 url") + } + httpClient := httpClientWithTimeoutAndHeaders(timeout, serverCfg.Headers) + st := &mcp.StreamableClientTransport{ + Endpoint: serverCfg.URL, + HTTPClient: httpClient, + } + if serverCfg.MaxRetries > 0 { + st.MaxRetries = serverCfg.MaxRetries + } + t = st + default: + return nil, fmt.Errorf("不支持的传输模式: %s（支持: stdio, sse, http）", transport) + } + + session, err := client.Connect(ctx, t, nil) + if err != nil { + return nil, fmt.Errorf("连接失败: %w", err) + } + + return newSDKClientFromSession(session, client, logger), nil +} + +func envMapToSlice(env map[string]string) []string { + m := make(map[string]string) + for _, s := range os.Environ() { + if i := strings.IndexByte(s, '='); i > 0 { + m[s[:i]] = s[i+1:] + } + } + for k, v := range env { + m[k] = v + } + out := make([]string, 0, len(m)) + for k, v := range m { + out = append(out, k+"="+v) + } + return out +} + +func httpClientWithTimeoutAndHeaders(timeout time.Duration, headers map[string]string) *http.Client { + transport := http.DefaultTransport + if len(headers) > 0 { + transport = &headerRoundTripper{ + headers: headers, + base: http.DefaultTransport, + } + } + return &http.Client{ + Timeout: timeout, + Transport: transport, + } +} + +// httpClientForLongLived 创建不设超时的 HTTP 客户端，用于 SSE 等长连接传输。 +// SSE 的 GET 流会持续打开，http.Client.Timeout 会在超时后强制关闭连接导致 EOF。 +// 超时由调用方通过 context 控制。 +func httpClientForLongLived(headers map[string]string) *http.Client { + transport := http.DefaultTransport + if len(headers) > 0 { + transport = &headerRoundTripper{ + headers: headers, + base: http.DefaultTransport, + } + } + return &http.Client{ + Transport: transport, + // 不设 Timeout，SSE 长连接的超时由 per-request context 控制 + } +} + +type headerRoundTripper struct { + headers map[string]string + base http.RoundTripper +} + +func (h *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + for k, v := range h.headers { + req.Header.Set(k, v) + } + return h.base.RoundTrip(req) +} diff --git a/internal/mcp/connection_recovery.go b/internal/mcp/connection_recovery.go new file mode 100644 index 00000000..a2ed9bfb --- /dev/null +++ b/internal/mcp/connection_recovery.go @@ -0,0 +1,192 @@ +package mcp + +import ( + "context" + "errors" + "io" + "strings" + "time" + + "go.uber.org/zap" +) + +const ( + // externalReconnectMinInterval 两次自动重连之间的最短间隔 + externalReconnectMinInterval = 30 * time.Second + // externalReconnectMaxBackoff 指数退避上限 + externalReconnectMaxBackoff = 5 * time.Minute +) + +// isConnectionDeadError 判断错误是否表示底层传输已断开（而非调用方主动取消或超时）。 +func isConnectionDeadError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return false + } + if errors.Is(err, io.EOF) { + return true + } + s := strings.ToLower(err.Error()) + return strings.Contains(s, "eof") || + strings.Contains(s, "client is closing") || + strings.Contains(s, "connection closed") || + strings.Contains(s, "connection reset") || + strings.Contains(s, "broken pipe") +} + +// handleConnectionDead 在 ListTools/CallTool 等操作失败且判定为断连时，标记客户端并调度重连。 +func (m *ExternalMCPManager) handleConnectionDead(name string, client ExternalMCPClient, err error) { + if !isConnectionDeadError(err) { + return + } + m.logger.Warn("检测到外部MCP连接已断开，将尝试自动重连", + zap.String("name", name), + zap.Error(err), + ) + m.markClientDisconnected(name, client, err) + m.scheduleReconnect(name) +} + +func (m *ExternalMCPManager) markClientDisconnected(name string, client ExternalMCPClient, err error) { + if lazy, ok := client.(*lazySDKClient); ok { + lazy.markDisconnected() + } + m.mu.Lock() + if err != nil { + m.errors[name] = "连接已断开: " + err.Error() + } + m.mu.Unlock() + m.toolCountsMu.Lock() + m.toolCounts[name] = 0 + m.toolCountsMu.Unlock() +} + +func (m *ExternalMCPManager) onClientConnected(name string) { + m.clearReconnectState(name) +} + +func (m *ExternalMCPManager) clearReconnectState(name string) { + m.reconnectMu.Lock() + delete(m.reconnectAttempts, name) + delete(m.reconnectLastTry, name) + delete(m.reconnecting, name) + m.reconnectMu.Unlock() +} + +func (m *ExternalMCPManager) reconnectBackoff(attempts int) time.Duration { + if attempts <= 0 { + return 0 + } + d := externalReconnectMinInterval + for i := 1; i < attempts && d < externalReconnectMaxBackoff; i++ { + d *= 2 + } + if d > externalReconnectMaxBackoff { + return externalReconnectMaxBackoff + } + return d +} + +func (m *ExternalMCPManager) scheduleReconnect(name string) { + m.mu.RLock() + cfg, exists := m.configs[name] + enabled := exists && m.isEnabled(cfg) + m.mu.RUnlock() + if !enabled { + return + } + go m.tryReconnect(name) +} + +func (m *ExternalMCPManager) tryReconnect(name string) { + m.reconnectMu.Lock() + if m.reconnecting[name] { + m.reconnectMu.Unlock() + return + } + attempts := m.reconnectAttempts[name] + if wait := m.reconnectBackoff(attempts); wait > 0 { + if last, ok := m.reconnectLastTry[name]; ok { + if elapsed := time.Since(last); elapsed < wait { + remaining := wait - elapsed + m.reconnectMu.Unlock() + m.scheduleReconnectAfter(name, remaining) + return + } + } + } + m.reconnecting[name] = true + m.reconnectMu.Unlock() + + defer func() { + m.reconnectMu.Lock() + delete(m.reconnecting, name) + m.reconnectMu.Unlock() + }() + + m.mu.RLock() + cfg, exists := m.configs[name] + enabled := exists && m.isEnabled(cfg) + client, hasClient := m.clients[name] + connecting := hasClient && client.GetStatus() == "connecting" + m.mu.RUnlock() + + if !enabled { + m.logger.Debug("跳过自动重连（外部MCP已停用）", zap.String("name", name)) + return + } + if connecting { + m.logger.Debug("跳过自动重连（连接正在进行中）", zap.String("name", name)) + return + } + + m.reconnectMu.Lock() + m.reconnectLastTry[name] = time.Now() + m.reconnectAttempts[name] = attempts + 1 + attemptNum := m.reconnectAttempts[name] + m.reconnectMu.Unlock() + + m.logger.Info("正在自动重连外部MCP", + zap.String("name", name), + zap.Int("attempt", attemptNum), + ) + + if err := m.startClient(name, true); err != nil { + m.logger.Warn("自动重连外部MCP失败", + zap.String("name", name), + zap.Error(err), + ) + } +} + +// scheduleReconnectAfterFailure 在自动重连失败后，按当前退避间隔预约下一次重试。 +func (m *ExternalMCPManager) scheduleReconnectAfterFailure(name string) { + m.mu.RLock() + cfg, exists := m.configs[name] + enabled := exists && m.isEnabled(cfg) + m.mu.RUnlock() + if !enabled { + return + } + m.reconnectMu.Lock() + wait := m.reconnectBackoff(m.reconnectAttempts[name]) + m.reconnectMu.Unlock() + m.logger.Info("自动重连失败，将按退避间隔再次尝试", + zap.String("name", name), + zap.Duration("after", wait), + ) + m.scheduleReconnectAfter(name, wait) +} + +// scheduleReconnectAfter 在 delay 后触发 tryReconnect（delay<=0 时立即执行）。 +func (m *ExternalMCPManager) scheduleReconnectAfter(name string, delay time.Duration) { + if delay <= 0 { + go m.tryReconnect(name) + return + } + time.AfterFunc(delay, func() { + m.tryReconnect(name) + }) +} diff --git a/internal/mcp/connection_recovery_test.go b/internal/mcp/connection_recovery_test.go new file mode 100644 index 00000000..f04e4622 --- /dev/null +++ b/internal/mcp/connection_recovery_test.go @@ -0,0 +1,215 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "io" + "testing" + "time" + + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +func TestIsConnectionDeadError(t *testing.T) { + t.Parallel() + cases := []struct { + name string + err error + want bool + }{ + {"nil", nil, false}, + {"eof", io.EOF, true}, + {"wrapped eof", fmt.Errorf("connection closed: %w", io.EOF), true}, + {"client closing", errors.New(`calling "tools/list": client is closing: EOF`), true}, + {"connection reset", errors.New("read tcp: connection reset by peer"), true}, + {"canceled", context.Canceled, false}, + {"deadline", context.DeadlineExceeded, false}, + {"other", errors.New("invalid params"), false}, + } + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + if got := isConnectionDeadError(tc.err); got != tc.want { + t.Fatalf("isConnectionDeadError(%v) = %v, want %v", tc.err, got, tc.want) + } + }) + } +} + +func TestLazySDKClient_MarkDisconnected(t *testing.T) { + c := &lazySDKClient{status: "connected"} + c.inner = &sdkClient{status: "connected"} + c.markDisconnected() + if c.IsConnected() { + t.Fatal("expected disconnected after markDisconnected") + } + if c.GetStatus() != "disconnected" { + t.Fatalf("expected status disconnected, got %s", c.GetStatus()) + } +} + +func TestHandleConnectionDead_MarksLazyClientDisconnected(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + + name := "dead-mcp" + cfg := config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://example.com/mcp", + ExternalMCPEnable: true, + } + m.mu.Lock() + m.configs[name] = cfg + client := newLazySDKClient(cfg, logger) + client.inner = &sdkClient{status: "connected"} + client.status = "connected" + m.clients[name] = client + m.mu.Unlock() + + deadErr := errors.New(`connection closed: calling "tools/list": client is closing: EOF`) + m.handleConnectionDead(name, client, deadErr) + + if client.IsConnected() { + t.Fatal("expected disconnected after handleConnectionDead") + } + if m.GetError(name) == "" { + t.Fatal("expected error message to be recorded") + } + counts := m.GetToolCounts() + if counts[name] != 0 { + t.Fatalf("expected tool count 0 after disconnect, got %d", counts[name]) + } +} + +func TestReconnectBackoff(t *testing.T) { + t.Parallel() + if d := (&ExternalMCPManager{}).reconnectBackoff(0); d != 0 { + t.Fatalf("attempt 0: got %v", d) + } + if d := (&ExternalMCPManager{}).reconnectBackoff(1); d != externalReconnectMinInterval { + t.Fatalf("attempt 1: got %v", d) + } + if d := (&ExternalMCPManager{}).reconnectBackoff(10); d != externalReconnectMaxBackoff { + t.Fatalf("attempt 10: got %v, want cap %v", d, externalReconnectMaxBackoff) + } +} + +func TestTryReconnect_RateLimited(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + + name := "rate-limited" + m.reconnectMu.Lock() + m.reconnectLastTry[name] = time.Now() + m.reconnectAttempts[name] = 2 + m.reconnectMu.Unlock() + + m.tryReconnect(name) + + m.reconnectMu.Lock() + attempts := m.reconnectAttempts[name] + m.reconnectMu.Unlock() + if attempts != 2 { + t.Fatalf("rate limited reconnect should not increment attempts, got %d", attempts) + } +} + +func TestTryReconnect_SkipsWhenDisabled(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + + name := "disabled-mcp" + m.mu.Lock() + m.configs[name] = config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://example.com/mcp", + ExternalMCPEnable: false, + } + m.mu.Unlock() + + m.tryReconnect(name) + + m.reconnectMu.Lock() + attempts := m.reconnectAttempts[name] + m.reconnectMu.Unlock() + if attempts != 0 { + t.Fatalf("disabled MCP should not increment reconnect attempts, got %d", attempts) + } +} + +func TestTryReconnect_SkipsWhenConnecting(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + + name := "connecting-mcp" + cfg := config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://example.com/mcp", + ExternalMCPEnable: true, + } + client := newLazySDKClient(cfg, logger) + client.setStatus("connecting") + + m.mu.Lock() + m.configs[name] = cfg + m.clients[name] = client + m.mu.Unlock() + + m.tryReconnect(name) + + m.reconnectMu.Lock() + attempts := m.reconnectAttempts[name] + m.reconnectMu.Unlock() + if attempts != 0 { + t.Fatalf("connecting MCP should not increment reconnect attempts, got %d", attempts) + } +} + +func TestStartClientAutoReconnect_SkipsWhenDisabled(t *testing.T) { + logger := zap.NewNop() + m := NewExternalMCPManager(logger) + m.stopRefresh = make(chan struct{}) + + name := "stopped" + m.mu.Lock() + m.configs[name] = config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://example.com/mcp", + ExternalMCPEnable: false, + } + m.mu.Unlock() + + if err := m.startClient(name, true); err != nil { + t.Fatalf("startClient: %v", err) + } + + m.mu.RLock() + cfg := m.configs[name] + _, hasClient := m.clients[name] + m.mu.RUnlock() + if cfg.ExternalMCPEnable { + t.Fatal("auto reconnect should not enable stopped MCP") + } + if hasClient { + t.Fatal("auto reconnect should not create client when disabled") + } +} + +func TestOnClientConnected_ClearsReconnectState(t *testing.T) { + m := &ExternalMCPManager{ + reconnectAttempts: map[string]int{"x": 3}, + reconnectLastTry: map[string]time.Time{"x": time.Now()}, + reconnecting: map[string]bool{"x": true}, + } + m.onClientConnected("x") + + m.reconnectMu.Lock() + defer m.reconnectMu.Unlock() + if len(m.reconnectAttempts) != 0 || len(m.reconnectLastTry) != 0 || len(m.reconnecting) != 0 { + t.Fatal("expected reconnect state cleared") + } +} diff --git a/internal/mcp/execution_control_tools.go b/internal/mcp/execution_control_tools.go new file mode 100644 index 00000000..7d74af25 --- /dev/null +++ b/internal/mcp/execution_control_tools.go @@ -0,0 +1,296 @@ +package mcp + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strconv" + "strings" + "time" + + "cyberstrike-ai/internal/mcp/builtin" +) + +const ( + defaultExecutionWaitTimeout = 60 * time.Second + maxExecutionWaitTimeout = 10 * time.Minute + defaultPartialPreviewBytes = 4096 + maxPartialPreviewBytes = 64 * 1024 +) + +// RegisterExecutionControlTools exposes execution handle operations to Eino as +// ordinary MCP tools. This keeps the agent loop native: the model calls a tool, +// receives a bounded result, and may call wait_tool_execution again if needed. +func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) { + if server == nil { + return + } + + server.RegisterTool(Tool{ + Name: builtin.ToolGetToolExecution, + Description: "查询后台工具 execution 的当前状态、结果和错误。用于外部 MCP 工具等待超时后，凭 execution_id 继续查看进度。", + ShortDescription: "查询后台工具执行状态", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览，默认 true"}, + "partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数，默认 4096，最大 65536"}, + }, + "required": []string{"execution_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + id := stringArg(args, "execution_id") + if id == "" { + return textToolResult("execution_id 必填", true), nil + } + exec := lookupToolExecution(server, external, id) + if exec == nil { + return textToolResult("未找到该 execution_id: "+id, true), nil + } + return textToolResult(formatExecutionForModel(exec, executionFormatOptionsFromArgs(args)), false), nil + }) + + server.RegisterTool(Tool{ + Name: builtin.ToolWaitToolExecution, + Description: "继续等待一个后台工具 execution 完成。每次等待都有 timeout_seconds 上限；若仍未完成，会返回当前状态，模型可稍后再次调用。", + ShortDescription: "有界等待后台工具执行", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数，默认 60，最大 600"}, + "include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览，默认 true"}, + "partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数，默认 4096，最大 65536"}, + }, + "required": []string{"execution_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + id := stringArg(args, "execution_id") + if id == "" { + return textToolResult("execution_id 必填", true), nil + } + wait := durationSecondsArg(args, "timeout_seconds", defaultExecutionWaitTimeout, maxExecutionWaitTimeout) + snap, err := waitToolExecutionSnapshot(ctx, server, external, id, wait) + if err != nil && !errors.Is(err, ErrExecutionWaitTimeout) { + return textToolResult("等待 execution 失败: "+err.Error(), true), nil + } + if snap == nil || snap.Execution == nil { + return textToolResult("未找到该 execution_id: "+id, true), nil + } + body := formatExecutionForModel(snap.Execution, executionFormatOptionsFromArgs(args)) + if errors.Is(err, ErrExecutionWaitTimeout) { + body += "\n\n本次等待已到达 timeout_seconds，上述 execution 仍未完成。可继续等待、取消，或采用其他步骤。" + } + return textToolResult(body, false), nil + }) + + server.RegisterTool(Tool{ + Name: builtin.ToolCancelToolExecution, + Description: "取消一个后台工具 execution。用于外部 MCP 工具长时间运行、误调用或用户要求停止时。", + ShortDescription: "取消后台工具执行", + InputSchema: map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"}, + "reason": map[string]interface{}{"type": "string", "description": "取消原因，可选，会写入终止说明"}, + }, + "required": []string{"execution_id"}, + }, + }, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + id := stringArg(args, "execution_id") + if id == "" { + return textToolResult("execution_id 必填", true), nil + } + reason := stringArg(args, "reason") + if server.CancelToolExecutionWithNote(id, reason) { + return textToolResult("已请求取消内部工具 execution: "+id, false), nil + } + if external != nil && external.CancelToolExecutionWithNote(id, reason) { + return textToolResult("已请求取消外部 MCP execution: "+id, false), nil + } + return textToolResult("未找到进行中的 execution，或该 execution 已结束: "+id, true), nil + }) +} + +func waitToolExecutionSnapshot(ctx context.Context, server *Server, external *ExternalMCPManager, id string, wait time.Duration) (*ExecutionSnapshot, error) { + if server != nil && server.executionService != nil && server.executionService.getEntry(id) != nil { + return server.executionService.Wait(ctx, id, wait) + } + if external != nil && external.executionService != nil && external.executionService.getEntry(id) != nil { + return external.executionService.Wait(ctx, id, wait) + } + if server != nil && server.executionService != nil { + if snap, err := server.executionService.Get(id); err == nil { + return snap, nil + } + } + if external != nil && external.executionService != nil { + return external.executionService.Get(id) + } + exec := lookupToolExecution(server, external, id) + if exec == nil { + return nil, fmt.Errorf("execution not found: %s", id) + } + return &ExecutionSnapshot{Execution: exec}, nil +} + +func lookupToolExecution(server *Server, external *ExternalMCPManager, id string) *ToolExecution { + if server != nil { + if exec, ok := server.GetExecution(id); ok && exec != nil { + return exec + } + } + if external != nil { + if exec, ok := external.GetExecution(id); ok && exec != nil { + return exec + } + } + return nil +} + +type executionFormatOptions struct { + includePartialOutput bool + partialMaxBytes int +} + +func executionFormatOptionsFromArgs(args map[string]interface{}) executionFormatOptions { + includePartial := true + if raw, ok := args["include_partial_output"]; ok { + if b, ok := raw.(bool); ok { + includePartial = b + } else if s := strings.TrimSpace(fmt.Sprint(raw)); s != "" { + includePartial = strings.EqualFold(s, "true") || s == "1" || strings.EqualFold(s, "yes") + } + } + maxBytes := intArg(args, "partial_output_max_bytes", defaultPartialPreviewBytes, maxPartialPreviewBytes) + return executionFormatOptions{includePartialOutput: includePartial, partialMaxBytes: maxBytes} +} + +func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) string { + if exec == nil { + return "execution: null" + } + payload := map[string]interface{}{ + "execution_id": exec.ID, + "tool": exec.ToolName, + "status": exec.Status, + "started_at": exec.StartTime.Format(time.RFC3339), + } + if exec.EndTime != nil { + payload["ended_at"] = exec.EndTime.Format(time.RFC3339) + } + if exec.Duration > 0 { + payload["duration"] = exec.Duration.String() + } + if exec.Error != "" { + payload["error"] = exec.Error + } + if exec.Result != nil { + payload["result"] = ToolResultPlainText(exec.Result) + payload["is_error"] = exec.Result.IsError + } + if opts.includePartialOutput && exec.PartialOutput != "" { + partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes) + payload["partial_output"] = partial + payload["partial_output_bytes"] = exec.PartialOutputBytes + payload["partial_output_truncated"] = exec.PartialOutputTruncated || len([]byte(partial)) < len([]byte(exec.PartialOutput)) + if exec.PartialOutputUpdatedAt != nil { + payload["partial_output_updated_at"] = exec.PartialOutputUpdatedAt.Format(time.RFC3339) + } + } + b, err := json.MarshalIndent(payload, "", " ") + if err != nil { + return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error) + } + return string(b) +} + +func tailStringBytes(s string, maxBytes int) string { + if maxBytes <= 0 { + maxBytes = defaultPartialPreviewBytes + } + b := []byte(s) + if len(b) <= maxBytes { + return s + } + return string(b[len(b)-maxBytes:]) +} + +func textToolResult(text string, isErr bool) *ToolResult { + return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr} +} + +func stringArg(args map[string]interface{}, key string) string { + if args == nil { + return "" + } + raw, ok := args[key] + if !ok || raw == nil { + return "" + } + switch v := raw.(type) { + case string: + return strings.TrimSpace(v) + default: + return strings.TrimSpace(fmt.Sprint(v)) + } +} + +func durationSecondsArg(args map[string]interface{}, key string, def, max time.Duration) time.Duration { + if args == nil { + return def + } + var seconds float64 + switch v := args[key].(type) { + case int: + seconds = float64(v) + case int64: + seconds = float64(v) + case float64: + seconds = v + case json.Number: + f, _ := v.Float64() + seconds = f + case string: + f, _ := strconv.ParseFloat(strings.TrimSpace(v), 64) + seconds = f + } + if seconds <= 0 { + return def + } + d := time.Duration(seconds * float64(time.Second)) + if max > 0 && d > max { + return max + } + return d +} + +func intArg(args map[string]interface{}, key string, def, max int) int { + if args == nil { + return def + } + var n int + switch v := args[key].(type) { + case int: + n = v + case int64: + n = int(v) + case float64: + n = int(v) + case json.Number: + i, _ := v.Int64() + n = int(i) + case string: + i, _ := strconv.Atoi(strings.TrimSpace(v)) + n = i + } + if n <= 0 { + return def + } + if max > 0 && n > max { + return max + } + return n +} diff --git a/internal/mcp/execution_service.go b/internal/mcp/execution_service.go new file mode 100644 index 00000000..1e2ad6d5 --- /dev/null +++ b/internal/mcp/execution_service.go @@ -0,0 +1,625 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/authctx" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +const ( + ToolExecutionStatusQueued = "queued" + ToolExecutionStatusRunning = "running" + ToolExecutionStatusCompleted = "completed" + ToolExecutionStatusFailed = "failed" + ToolExecutionStatusCancelled = "cancelled" + ToolExecutionStatusHardTimeout = "hard_timeout" + ToolExecutionStatusOrphaned = "orphaned" +) + +var ErrExecutionWaitTimeout = errors.New("tool execution wait timeout") + +// ExecutionRunFunc is the blocking operation owned by a worker. +type ExecutionRunFunc func(context.Context) (*ToolResult, error) + +type ExecutionPreRunFunc func(context.Context, *ToolExecution) (func(), error) + +// ExecutionDoneFunc observes the final persisted state. It is invoked once, +// including for late completions after an agent has stopped waiting. +type ExecutionDoneFunc func(*ToolExecution) + +type ExecutionRequest struct { + ID string + ToolName string + Arguments map[string]interface{} + ConversationID string + OwnerUserID string + HardTimeout time.Duration + PreRun ExecutionPreRunFunc + Run ExecutionRunFunc + OnDone ExecutionDoneFunc +} + +type ExecutionHandle struct { + ID string +} + +type ExecutionSnapshot struct { + Execution *ToolExecution +} + +type executionEntry struct { + exec *ToolExecution + cancel context.CancelFunc + done chan struct{} + preRun ExecutionPreRunFunc + run ExecutionRunFunc + result *ToolResult + err error +} + +// ExecutionService keeps Eino-facing tool calls synchronous while moving the +// untrusted blocking work into cancellable workers with explicit execution IDs. +type ExecutionService struct { + storage MonitorStorage + logger *zap.Logger + + mu sync.Mutex + entries map[string]*executionEntry + abortUserNotes map[string]string + maxInMemory int + resultMaxBytes int + spillRootDir string +} + +func NewExecutionService(storage MonitorStorage, logger *zap.Logger) *ExecutionService { + if logger == nil { + logger = zap.NewNop() + } + return &ExecutionService{ + storage: storage, + logger: logger, + entries: make(map[string]*executionEntry), + abortUserNotes: make(map[string]string), + maxInMemory: 1000, + resultMaxBytes: DefaultToolResultMaxBytes, + } +} + +func (s *ExecutionService) ConfigureToolResultMaxBytes(maxBytes int) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.resultMaxBytes = maxBytes +} + +// ConfigureToolResultSpillRoot sets the reduction-compatible root used when +// oversized tool results are spilled to local files (empty → tmp/reduction). +func (s *ExecutionService) ConfigureToolResultSpillRoot(rootDir string) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.spillRootDir = strings.TrimSpace(rootDir) +} + +func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*ExecutionHandle, error) { + if s == nil { + return nil, fmt.Errorf("execution service is nil") + } + if req.Run == nil { + return nil, fmt.Errorf("execution run func is nil") + } + id := strings.TrimSpace(req.ID) + if id == "" { + id = uuid.New().String() + } + start := time.Now() + exec := &ToolExecution{ + ID: id, + ToolName: strings.TrimSpace(req.ToolName), + Arguments: cloneArgsMap(req.Arguments), + Status: ToolExecutionStatusQueued, + StartTime: start, + ConversationID: strings.TrimSpace(req.ConversationID), + OwnerUserID: strings.TrimSpace(req.OwnerUserID), + } + if exec.ConversationID == "" { + exec.ConversationID = MCPConversationIDFromContext(ctx) + } + if exec.OwnerUserID == "" { + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + exec.OwnerUserID = principal.UserID + } + } + + runCtx := detachedExecutionContext(ctx) + var cancel context.CancelFunc + if req.HardTimeout > 0 { + runCtx, cancel = context.WithTimeout(runCtx, req.HardTimeout) + } else { + runCtx, cancel = context.WithCancel(runCtx) + } + entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run} + + s.mu.Lock() + if _, exists := s.entries[id]; exists { + s.mu.Unlock() + cancel() + return nil, fmt.Errorf("execution already exists: %s", id) + } + s.entries[id] = entry + s.cleanupOldEntriesLocked() + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(exec); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", id)) + } + } + notifyToolRunBegin(ctx, id) + + go s.runWorker(runCtx, entry, req.OnDone) + return &ExecutionHandle{ID: id}, nil +} + +func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, onDone ExecutionDoneFunc) { + id := entry.exec.ID + ctx = WithMCPExecutionID(ctx, id) + if conv := strings.TrimSpace(entry.exec.ConversationID); conv != "" { + ctx = WithMCPConversationID(ctx, conv) + } + var release func() + defer func() { + if release != nil { + release() + } + entry.cancel() + notifyToolRunEnd(ctx, id) + close(entry.done) + }() + + if entry.preRun != nil { + var preErr error + release, preErr = entry.preRun(ctx, cloneToolExecution(entry.exec)) + if preErr != nil { + s.finishEntry(ctx, entry, nil, preErr, onDone) + return + } + } + s.markEntryRunning(entry) + + result, err := entryResultRecover(ctx, entry.exec.ToolName, s.logger, func() (*ToolResult, error) { + return nilSafeRun(ctx, entry) + }) + s.finishEntry(ctx, entry, result, err, onDone) +} + +func (s *ExecutionService) markEntryRunning(entry *executionEntry) { + if s == nil || entry == nil || entry.exec == nil { + return + } + s.mu.Lock() + if !isExecutionTerminal(entry.exec.Status) { + entry.exec.Status = ToolExecutionStatusRunning + } + runningExec := cloneToolExecution(entry.exec) + s.mu.Unlock() + if s.storage != nil { + if err := s.storage.SaveToolExecution(runningExec); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", runningExec.ID)) + } + } +} + +func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) { + id := entry.exec.ID + cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err) + + now := time.Now() + s.mu.Lock() + spill := ToolResultSpillConfig{ + RootDir: s.spillRootDir, + ConversationID: entry.exec.ConversationID, + ExecutionID: id, + } + if ctx != nil { + if pid := MCPProjectIDFromContext(ctx); pid != "" { + spill.ProjectID = pid + } + if conv := MCPConversationIDFromContext(ctx); conv != "" { + spill.ConversationID = conv + } + } + result = NormalizeToolResultForStorageWithSpill(result, s.resultMaxBytes, spill) + entry.result = result + entry.err = err + entry.exec.EndTime = &now + entry.exec.Duration = now.Sub(entry.exec.StartTime) + if err != nil { + switch { + case errors.Is(err, context.DeadlineExceeded): + entry.exec.Status = ToolExecutionStatusHardTimeout + entry.exec.Error = "工具执行超过硬超时限制" + case errors.Is(err, context.Canceled): + entry.exec.Status = ToolExecutionStatusCancelled + entry.exec.Error = "已手动终止或任务已取消" + default: + entry.exec.Status = ToolExecutionStatusFailed + entry.exec.Error = err.Error() + } + } else if result != nil && result.IsError { + if cancelledWithUserNote { + entry.exec.Status = ToolExecutionStatusCancelled + entry.exec.Error = "" + } else if isBackgroundWaitToolResult(result) { + entry.exec.Status = ToolExecutionStatusCompleted + entry.exec.Error = "" + } else { + entry.exec.Status = ToolExecutionStatusFailed + entry.exec.Error = firstToolResultText(result, "工具执行返回错误结果") + } + entry.exec.Result = result + } else { + entry.exec.Status = ToolExecutionStatusCompleted + if result == nil { + result = &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回结果"}}} + entry.result = result + } + entry.exec.Result = result + } + finalExec := cloneToolExecution(entry.exec) + s.mu.Unlock() + + if s.storage != nil { + if saveErr := s.storage.SaveToolExecution(finalExec); saveErr != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(saveErr), zap.String("executionId", id)) + } + } + if onDone != nil { + onDone(finalExec) + } +} + +func nilSafeRun(ctx context.Context, entry *executionEntry) (*ToolResult, error) { + if entry == nil { + return nil, fmt.Errorf("execution entry is nil") + } + if entry.run == nil { + return nil, fmt.Errorf("execution run func not wired") + } + return entry.run(ctx) +} + +func entryResultRecover(ctx context.Context, toolName string, logger *zap.Logger, fn func() (*ToolResult, error)) (res *ToolResult, err error) { + defer func() { + if r := recover(); r != nil { + if logger != nil { + logger.Error("tool execution worker panic recovered", zap.Any("recover", r), zap.String("toolName", toolName), zap.Stack("stack")) + } + err = fmt.Errorf("tool execution panic: %v", r) + } + }() + return fn() +} + +func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout time.Duration) (*ExecutionSnapshot, error) { + entry := s.getEntry(executionID) + if entry == nil { + return s.getPersistedSnapshot(executionID) + } + if isExecutionTerminal(entry.exec.Status) { + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil + } + + var timeoutCh <-chan time.Time + var timer *time.Timer + if timeout > 0 { + timer = time.NewTimer(timeout) + timeoutCh = timer.C + defer timer.Stop() + } + + select { + case <-entry.done: + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil + case <-timeoutCh: + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ErrExecutionWaitTimeout + case <-ctxDone(ctx): + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ctx.Err() + } +} + +func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) { + entry := s.getEntry(executionID) + if entry != nil { + return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil + } + return s.getPersistedSnapshot(executionID) +} + +func (s *ExecutionService) AppendPartialOutput(executionID, chunk string) bool { + id := strings.TrimSpace(executionID) + if s == nil || id == "" || chunk == "" { + return false + } + now := time.Now() + s.mu.Lock() + defer s.mu.Unlock() + entry := s.entries[id] + if entry == nil || entry.exec == nil { + return false + } + appendPartialOutput(entry.exec, chunk, defaultPartialOutputMaxBytes, now) + return true +} + +func (s *ExecutionService) Cancel(executionID, note string) bool { + id := strings.TrimSpace(executionID) + if id == "" || s == nil { + return false + } + s.mu.Lock() + entry := s.entries[id] + if entry == nil || isExecutionTerminal(entry.exec.Status) { + s.mu.Unlock() + return false + } + if strings.TrimSpace(note) != "" { + s.abortUserNotes[id] = strings.TrimSpace(note) + } + cancel := entry.cancel + s.mu.Unlock() + if cancel != nil { + cancel() + } + return true +} + +func (s *ExecutionService) ActiveRunningExecutionIDs() map[string]struct{} { + if s == nil { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + out := make(map[string]struct{}) + for id, entry := range s.entries { + if entry != nil && entry.exec != nil && !isExecutionTerminal(entry.exec.Status) { + out[id] = struct{}{} + } + } + if len(out) == 0 { + return nil + } + return out +} + +func (s *ExecutionService) CancelAll(note string) { + if s == nil { + return + } + s.mu.Lock() + cancels := make([]context.CancelFunc, 0, len(s.entries)) + for id, entry := range s.entries { + if entry == nil || isExecutionTerminal(entry.exec.Status) { + continue + } + if strings.TrimSpace(note) != "" { + s.abortUserNotes[id] = strings.TrimSpace(note) + } + if entry.cancel != nil { + cancels = append(cancels, entry.cancel) + } + } + s.mu.Unlock() + for _, cancel := range cancels { + cancel() + } +} + +func (s *ExecutionService) getEntry(executionID string) *executionEntry { + if s == nil { + return nil + } + id := strings.TrimSpace(executionID) + if id == "" { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + return s.entries[id] +} + +func (s *ExecutionService) getPersistedSnapshot(executionID string) (*ExecutionSnapshot, error) { + id := strings.TrimSpace(executionID) + if id == "" { + return nil, fmt.Errorf("execution_id is required") + } + if s != nil && s.storage != nil { + exec, err := s.storage.GetToolExecution(id) + if err == nil && exec != nil { + return &ExecutionSnapshot{Execution: exec}, nil + } + if err != nil { + return nil, err + } + } + return nil, fmt.Errorf("execution not found: %s", id) +} + +func (s *ExecutionService) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) { + note := strings.TrimSpace(s.takeAbortUserNote(executionID)) + if note == "" { + return false + } + hasErr := err != nil && *err != nil + hasRes := result != nil && *result != nil + if !hasErr && !hasRes { + return false + } + partial := "" + if hasRes { + partial = ToolResultPlainText(*result) + } + if partial == "" && hasErr { + partial = (*err).Error() + } + merged := MergePartialToolOutputAndAbortNote(partial, note) + if err != nil { + *err = nil + } + if result != nil { + *result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true} + } + return true +} + +func (s *ExecutionService) takeAbortUserNote(id string) string { + s.mu.Lock() + defer s.mu.Unlock() + note := s.abortUserNotes[id] + delete(s.abortUserNotes, id) + return note +} + +func (s *ExecutionService) cleanupOldEntriesLocked() { + if s.maxInMemory <= 0 || len(s.entries) <= s.maxInMemory { + return + } + type oldEntry struct { + id string + startTime time.Time + } + var terminal []oldEntry + for id, entry := range s.entries { + if entry != nil && entry.exec != nil && isExecutionTerminal(entry.exec.Status) { + terminal = append(terminal, oldEntry{id: id, startTime: entry.exec.StartTime}) + } + } + for len(s.entries) > s.maxInMemory && len(terminal) > 0 { + oldest := 0 + for i := 1; i < len(terminal); i++ { + if terminal[i].startTime.Before(terminal[oldest].startTime) { + oldest = i + } + } + delete(s.entries, terminal[oldest].id) + terminal = append(terminal[:oldest], terminal[oldest+1:]...) + } +} + +func firstToolResultText(result *ToolResult, fallback string) string { + if result != nil { + for _, c := range result.Content { + if strings.TrimSpace(c.Text) != "" { + return c.Text + } + } + } + return fallback +} + +func isBackgroundWaitToolResult(result *ToolResult) bool { + text := strings.ToLower(strings.TrimSpace(ToolResultPlainText(result))) + if text == "" { + return false + } + hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`) + hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") || + strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) || + strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`) + hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") || + strings.Contains(text, "本次等待已到达") || + strings.Contains(text, "wait_timeout:") || + strings.Contains(text, "background execution") || + strings.Contains(text, "still running") || + strings.Contains(text, "仍未完成") + return hasExecutionID && hasRunningStatus && hasSoftWaitSignal +} + +func isExecutionTerminal(status string) bool { + switch strings.TrimSpace(strings.ToLower(status)) { + case ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned: + return true + default: + return false + } +} + +func ctxDone(ctx context.Context) <-chan struct{} { + if ctx == nil { + return nil + } + return ctx.Done() +} + +func detachedExecutionContext(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + return context.WithoutCancel(ctx) +} + +func cloneArgsMap(in map[string]interface{}) map[string]interface{} { + if in == nil { + return map[string]interface{}{} + } + out := make(map[string]interface{}, len(in)) + for k, v := range in { + out[k] = v + } + return out +} + +func cloneToolExecution(in *ToolExecution) *ToolExecution { + if in == nil { + return nil + } + out := *in + out.Arguments = cloneArgsMap(in.Arguments) + if in.Result != nil { + res := *in.Result + if in.Result.Content != nil { + res.Content = append([]Content(nil), in.Result.Content...) + } + out.Result = &res + } + if in.EndTime != nil { + t := *in.EndTime + out.EndTime = &t + } + if in.PartialOutputUpdatedAt != nil { + t := *in.PartialOutputUpdatedAt + out.PartialOutputUpdatedAt = &t + } + return &out +} + +func appendPartialOutput(exec *ToolExecution, chunk string, maxBytes int, updatedAt time.Time) { + if exec == nil || chunk == "" { + return + } + if maxBytes <= 0 { + maxBytes = defaultPartialOutputMaxBytes + } + exec.PartialOutputBytes += int64(len([]byte(chunk))) + combined := exec.PartialOutput + chunk + if len([]byte(combined)) > maxBytes { + b := []byte(combined) + combined = string(b[len(b)-maxBytes:]) + exec.PartialOutputTruncated = true + } + exec.PartialOutput = combined + t := updatedAt + exec.PartialOutputUpdatedAt = &t +} diff --git a/internal/mcp/execution_service_test.go b/internal/mcp/execution_service_test.go new file mode 100644 index 00000000..29f981c5 --- /dev/null +++ b/internal/mcp/execution_service_test.go @@ -0,0 +1,41 @@ +package mcp + +import ( + "context" + "testing" +) + +func TestExecutionServiceBackgroundWaitResultCompletesWaitTool(t *testing.T) { + service := NewExecutionService(nil, nil) + handle, err := service.Submit(context.Background(), ExecutionRequest{ + ToolName: "wait_tool_execution", + Run: func(context.Context) (*ToolResult, error) { + return &ToolResult{ + Content: []Content{{Type: "text", Text: `{ + "execution_id": "3eaaa391-050b-4be1-a870-48a855923cb7", + "tool": "exec", + "status": "running" +} + +本次等待已到达 timeout_seconds，上述 execution 仍未完成。可继续等待、取消，或采用其他步骤。`}}, + IsError: true, + }, nil + }, + }) + if err != nil { + t.Fatalf("Submit: %v", err) + } + snap, err := service.Wait(context.Background(), handle.ID, 0) + if err != nil { + t.Fatalf("Wait: %v", err) + } + if snap == nil || snap.Execution == nil { + t.Fatal("missing execution snapshot") + } + if snap.Execution.Status != ToolExecutionStatusCompleted { + t.Fatalf("status = %q, want %q", snap.Execution.Status, ToolExecutionStatusCompleted) + } + if snap.Execution.Result == nil || !snap.Execution.Result.IsError { + t.Fatal("model-facing result should remain IsError") + } +} diff --git a/internal/mcp/external_manager.go b/internal/mcp/external_manager.go new file mode 100644 index 00000000..f51b0457 --- /dev/null +++ b/internal/mcp/external_manager.go @@ -0,0 +1,1615 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +const ( + // externalToolListCacheTTL 已连接外部 MCP 的工具列表缓存有效期，避免每次 API 请求都打远程 ListTools。 + externalToolListCacheTTL = 60 * time.Second + // externalToolCountRefreshInterval 后台刷新工具数量的间隔（仅刷新缓存过期或缺失的客户端）。 + externalToolCountRefreshInterval = 60 * time.Second +) + +// toolListCacheEntry 外部 MCP 工具列表缓存条目 +type toolListCacheEntry struct { + tools []Tool + updatedAt time.Time +} + +// listToolsInflight 合并同一 MCP 上并发的 ListTools 请求 +type listToolsInflight struct { + done chan struct{} + tools []Tool + err error +} + +type ExternalMCPResilienceConfig struct { + MaxConcurrentPerServer int + MaxConcurrentTotal int + CircuitFailureThreshold int + CircuitCooldown time.Duration +} + +type externalMCPServerRuntime struct { + semaphore chan struct{} + consecutiveFailures int + circuitOpenUntil time.Time +} + +// ExternalMCPManager 外部MCP管理器 +type ExternalMCPManager struct { + clients map[string]ExternalMCPClient + configs map[string]config.ExternalMCPServerConfig + logger *zap.Logger + storage MonitorStorage // 可选的持久化存储 + executions map[string]*ToolExecution // 执行记录 + stats map[string]*ToolStats // 工具统计信息 + errors map[string]string // 错误信息 + toolCounts map[string]int // 工具数量缓存 + toolCountsMu sync.RWMutex // 工具数量缓存的锁 + toolCache map[string]toolListCacheEntry // 工具列表缓存：MCP名称 -> 工具列表 + toolCacheMu sync.RWMutex // 工具列表缓存的锁 + listToolsMu sync.Mutex + listToolsInflight map[string]*listToolsInflight + stopRefresh chan struct{} // 停止后台刷新的信号 + refreshWg sync.WaitGroup // 等待后台刷新goroutine完成 + refreshing atomic.Bool // 防止 refreshToolCounts 并发堆积 + mu sync.RWMutex + runningCancels map[string]context.CancelFunc + abortUserNotes map[string]string + reconnectMu sync.Mutex + reconnecting map[string]bool + reconnectLastTry map[string]time.Time + reconnectAttempts map[string]int + toolAuthorizer func(context.Context, string, map[string]interface{}) error + executionService *ExecutionService + toolWaitTimeout time.Duration + toolResultMaxBytes int + spillRootDir string + resilience ExternalMCPResilienceConfig + serverRuntimes map[string]*externalMCPServerRuntime + globalSemaphore chan struct{} +} + +// NewExternalMCPManager 创建外部MCP管理器 +func NewExternalMCPManager(logger *zap.Logger) *ExternalMCPManager { + return NewExternalMCPManagerWithStorage(logger, nil) +} + +// SetToolAuthorizer installs the policy decision point for all external MCP +// invocations. App wiring configures this before any Agent can call a tool. +func (m *ExternalMCPManager) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) { + m.mu.Lock() + m.toolAuthorizer = authorizer + m.mu.Unlock() +} + +// NewExternalMCPManagerWithStorage 创建外部MCP管理器（带持久化存储） +func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage) *ExternalMCPManager { + manager := &ExternalMCPManager{ + clients: make(map[string]ExternalMCPClient), + configs: make(map[string]config.ExternalMCPServerConfig), + logger: logger, + storage: storage, + executions: make(map[string]*ToolExecution), + stats: make(map[string]*ToolStats), + errors: make(map[string]string), + toolCounts: make(map[string]int), + toolCache: make(map[string]toolListCacheEntry), + listToolsInflight: make(map[string]*listToolsInflight), + stopRefresh: make(chan struct{}), + runningCancels: make(map[string]context.CancelFunc), + abortUserNotes: make(map[string]string), + reconnecting: make(map[string]bool), + reconnectLastTry: make(map[string]time.Time), + reconnectAttempts: make(map[string]int), + toolWaitTimeout: 60 * time.Second, + toolResultMaxBytes: DefaultToolResultMaxBytes, + resilience: ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: 2, + MaxConcurrentTotal: 16, + CircuitFailureThreshold: 3, + CircuitCooldown: 60 * time.Second, + }, + serverRuntimes: make(map[string]*externalMCPServerRuntime), + globalSemaphore: make(chan struct{}, 16), + } + manager.executionService = NewExecutionService(storage, logger) + // 启动后台刷新工具数量的goroutine + manager.startToolCountRefresh() + return manager +} + +func (m *ExternalMCPManager) ConfigureToolResultMaxBytes(maxBytes int) { + if m == nil { + return + } + m.mu.Lock() + m.toolResultMaxBytes = maxBytes + m.mu.Unlock() + if m.executionService != nil { + m.executionService.ConfigureToolResultMaxBytes(maxBytes) + } +} + +// ConfigureToolResultSpillRoot sets the local directory root used when oversized +// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction). +func (m *ExternalMCPManager) ConfigureToolResultSpillRoot(rootDir string) { + if m == nil { + return + } + m.mu.Lock() + m.spillRootDir = strings.TrimSpace(rootDir) + m.mu.Unlock() + if m.executionService != nil { + m.executionService.ConfigureToolResultSpillRoot(rootDir) + } +} + +// ConfigureToolWaitTimeoutSeconds controls how long an agent-facing tool call +// waits for an external MCP execution before returning an execution_id that can +// be polled with wait_tool_execution. seconds<=0 waits until completion. +func (m *ExternalMCPManager) ConfigureToolWaitTimeoutSeconds(seconds int) { + if m == nil { + return + } + m.mu.Lock() + defer m.mu.Unlock() + if seconds <= 0 { + m.toolWaitTimeout = 0 + return + } + m.toolWaitTimeout = time.Duration(seconds) * time.Second +} + +func (m *ExternalMCPManager) ConfigureResilience(cfg ExternalMCPResilienceConfig) { + if m == nil { + return + } + normalized := normalizeExternalMCPResilienceConfig(cfg) + m.mu.Lock() + defer m.mu.Unlock() + m.resilience = normalized + m.serverRuntimes = make(map[string]*externalMCPServerRuntime) + if normalized.MaxConcurrentTotal > 0 { + m.globalSemaphore = make(chan struct{}, normalized.MaxConcurrentTotal) + } else { + m.globalSemaphore = nil + } +} + +func normalizeExternalMCPResilienceConfig(cfg ExternalMCPResilienceConfig) ExternalMCPResilienceConfig { + if cfg.MaxConcurrentPerServer == 0 { + cfg.MaxConcurrentPerServer = 2 + } + if cfg.MaxConcurrentTotal == 0 { + cfg.MaxConcurrentTotal = 16 + } + if cfg.CircuitFailureThreshold == 0 { + cfg.CircuitFailureThreshold = 3 + } + if cfg.CircuitCooldown <= 0 { + cfg.CircuitCooldown = 60 * time.Second + } + if cfg.MaxConcurrentPerServer < 0 { + cfg.MaxConcurrentPerServer = 0 + } + if cfg.MaxConcurrentTotal < 0 { + cfg.MaxConcurrentTotal = 0 + } + return cfg +} + +// LoadConfigs 加载配置 +func (m *ExternalMCPManager) LoadConfigs(cfg *config.ExternalMCPConfig) { + m.mu.Lock() + defer m.mu.Unlock() + + if cfg == nil || cfg.Servers == nil { + return + } + + m.configs = make(map[string]config.ExternalMCPServerConfig) + for name, serverCfg := range cfg.Servers { + m.configs[name] = serverCfg + } +} + +// GetConfigs 获取所有配置 +func (m *ExternalMCPManager) GetConfigs() map[string]config.ExternalMCPServerConfig { + m.mu.RLock() + defer m.mu.RUnlock() + + result := make(map[string]config.ExternalMCPServerConfig) + for k, v := range m.configs { + result[k] = v + } + return result +} + +// AddOrUpdateConfig 添加或更新配置 +func (m *ExternalMCPManager) AddOrUpdateConfig(name string, serverCfg config.ExternalMCPServerConfig) error { + m.mu.Lock() + defer m.mu.Unlock() + + // 如果已存在客户端，先关闭 + if client, exists := m.clients[name]; exists { + client.Close() + delete(m.clients, name) + } + + m.configs[name] = serverCfg + + // 如果启用，自动连接 + if m.isEnabled(serverCfg) { + go m.connectClient(name, serverCfg) + } + + return nil +} + +// RemoveConfig 移除配置 +func (m *ExternalMCPManager) RemoveConfig(name string) error { + m.mu.Lock() + defer m.mu.Unlock() + + // 关闭客户端 + if client, exists := m.clients[name]; exists { + client.Close() + delete(m.clients, name) + } + + delete(m.configs, name) + m.clearReconnectState(name) + + // 清理工具数量缓存 + m.toolCountsMu.Lock() + delete(m.toolCounts, name) + m.toolCountsMu.Unlock() + + // 清理工具列表缓存 + m.toolCacheMu.Lock() + delete(m.toolCache, name) + m.toolCacheMu.Unlock() + + return nil +} + +// StartClient 启动客户端（用户手动启动；连接失败不自动重试） +func (m *ExternalMCPManager) StartClient(name string) error { + return m.startClient(name, false) +} + +// startClient 启动客户端。autoReconnect 为 true 时用于断连自愈：尊重停用状态，失败后按退避继续重试。 +func (m *ExternalMCPManager) startClient(name string, autoReconnect bool) error { + m.mu.Lock() + serverCfg, exists := m.configs[name] + m.mu.Unlock() + + if !exists { + return fmt.Errorf("配置不存在: %s", name) + } + + if autoReconnect && !m.isEnabled(serverCfg) { + return nil + } + + // 检查是否已经有连接的客户端 + m.mu.RLock() + existingClient, hasClient := m.clients[name] + m.mu.RUnlock() + + if hasClient { + // 检查客户端是否已连接 + if existingClient.IsConnected() { + // 客户端已连接，直接返回成功（目标状态已达成） + if !autoReconnect { + m.mu.Lock() + serverCfg.ExternalMCPEnable = true + m.configs[name] = serverCfg + m.mu.Unlock() + } + return nil + } + // 如果有客户端但未连接，先关闭 + existingClient.Close() + m.mu.Lock() + delete(m.clients, name) + m.mu.Unlock() + } + + if autoReconnect { + m.mu.RLock() + serverCfg, exists = m.configs[name] + enabled := exists && m.isEnabled(serverCfg) + m.mu.RUnlock() + if !enabled { + return nil + } + } + + // 更新配置为启用 + m.mu.Lock() + serverCfg.ExternalMCPEnable = true + m.configs[name] = serverCfg + // 清除之前的错误信息（重新启动时） + delete(m.errors, name) + m.mu.Unlock() + + // 立即创建客户端并设置为"connecting"状态，这样前端可以立即看到状态 + client := m.createClient(serverCfg) + if client == nil { + return fmt.Errorf("无法创建客户端：不支持的传输模式") + } + + // 设置状态为connecting + m.setClientStatus(client, "connecting") + + // 立即保存客户端，这样前端查询时就能看到"connecting"状态 + m.mu.Lock() + m.clients[name] = client + m.mu.Unlock() + + // 在后台异步进行实际连接 + go func(reconnect bool) { + if err := m.doConnect(name, serverCfg, client); err != nil { + m.logger.Error("连接外部MCP客户端失败", + zap.String("name", name), + zap.Bool("auto_reconnect", reconnect), + zap.Error(err), + ) + // 连接失败，设置状态为error并保存错误信息 + m.setClientStatus(client, "error") + m.mu.Lock() + m.errors[name] = err.Error() + m.mu.Unlock() + // 触发工具数量刷新（连接失败，工具数量应为0） + m.triggerToolCountRefresh() + if reconnect { + m.scheduleReconnectAfterFailure(name) + } + } else { + // 连接成功，清除错误信息 + m.mu.Lock() + delete(m.errors, name) + m.mu.Unlock() + m.onClientConnected(name) + // 异步拉取工具列表（singleflight 去重，结果同时写入 toolCache 与 toolCounts） + go m.refreshToolCache(name, client) + } + }(autoReconnect) + + return nil +} + +// StopClient 停止客户端 +func (m *ExternalMCPManager) StopClient(name string) error { + m.mu.Lock() + defer m.mu.Unlock() + + serverCfg, exists := m.configs[name] + if !exists { + return fmt.Errorf("配置不存在: %s", name) + } + + // 关闭客户端 + if client, exists := m.clients[name]; exists { + client.Close() + delete(m.clients, name) + } + + // 清除错误信息 + delete(m.errors, name) + + // 更新工具数量缓存（停止后工具数量为0） + m.toolCountsMu.Lock() + m.toolCounts[name] = 0 + m.toolCountsMu.Unlock() + + m.toolCacheMu.Lock() + delete(m.toolCache, name) + m.toolCacheMu.Unlock() + + // 更新配置为禁用 + serverCfg.ExternalMCPEnable = false + m.configs[name] = serverCfg + + m.clearReconnectState(name) + + return nil +} + +// GetClient 获取客户端 +func (m *ExternalMCPManager) GetClient(name string) (ExternalMCPClient, bool) { + m.mu.RLock() + defer m.mu.RUnlock() + + client, exists := m.clients[name] + return client, exists +} + +// GetError 获取错误信息 +func (m *ExternalMCPManager) GetError(name string) string { + m.mu.RLock() + defer m.mu.RUnlock() + + return m.errors[name] +} + +// GetAllTools 获取所有外部MCP的工具 +// 优先从已连接的客户端获取，如果连接断开则返回缓存的工具列表 +// 策略： +// - error 状态：不使用缓存，直接跳过（配置错误或服务不可用） +// - disconnected/connecting 状态：使用缓存（临时断开） +// - connected 状态：正常获取，失败时降级使用缓存 +func (m *ExternalMCPManager) GetAllTools(ctx context.Context) ([]Tool, error) { + m.mu.RLock() + clients := make(map[string]ExternalMCPClient) + for k, v := range m.clients { + clients[k] = v + } + m.mu.RUnlock() + + var allTools []Tool + var hasError bool + var lastError error + + // 使用较短的超时时间进行快速检查（3秒），避免阻塞 + quickCtx, quickCancel := context.WithTimeout(ctx, 3*time.Second) + defer quickCancel() + + for name, client := range clients { + tools, err := m.getToolsForClient(name, client, quickCtx) + if err != nil { + // 记录错误，但继续处理其他客户端 + hasError = true + if lastError == nil { + lastError = err + } + continue + } + + // 为工具添加前缀，避免冲突 + for _, tool := range tools { + tool.Name = fmt.Sprintf("%s::%s", name, tool.Name) + allTools = append(allTools, tool) + } + } + + // 如果有错误但至少返回了一些工具，不返回错误（部分成功） + if hasError && len(allTools) == 0 { + return nil, fmt.Errorf("获取外部MCP工具失败: %w", lastError) + } + + return allTools, nil +} + +// getToolsForClient 获取指定客户端的工具列表 +// 返回工具列表和错误（如果完全无法获取） +func (m *ExternalMCPManager) getToolsForClient(name string, client ExternalMCPClient, ctx context.Context) ([]Tool, error) { + status := client.GetStatus() + + // error 状态：不使用缓存，直接返回错误 + if status == "error" { + m.logger.Debug("跳过连接失败的外部MCP（不使用缓存）", + zap.String("name", name), + zap.String("status", status), + ) + return nil, fmt.Errorf("外部MCP连接失败: %s", name) + } + + // 已连接：缓存优先，仅在缺失或过期时打远程 ListTools + if client.IsConnected() { + if tools, ok := m.getFreshCachedTools(name); ok { + return tools, nil + } + if tools, ok := m.getAnyCachedTools(name); ok { + m.triggerToolListRefresh(name, client) + return tools, nil + } + tools, err := m.listToolsDeduped(ctx, name, client) + if err != nil { + return m.getCachedTools(name, "连接正常但获取失败", err) + } + return tools, nil + } + + // 未连接：根据状态决定是否使用缓存 + if status == "disconnected" || status == "connecting" { + return m.getCachedTools(name, fmt.Sprintf("客户端临时断开（状态: %s）", status), nil) + } + + // 其他未知状态，不使用缓存 + m.logger.Debug("跳过外部MCP（未知状态）", + zap.String("name", name), + zap.String("status", status), + ) + return nil, fmt.Errorf("外部MCP状态未知: %s (状态: %s)", name, status) +} + +// getCachedTools 获取缓存的工具列表（含空列表缓存） +func (m *ExternalMCPManager) getCachedTools(name, reason string, originalErr error) ([]Tool, error) { + if tools, ok := m.getAnyCachedTools(name); ok { + m.logger.Debug("使用缓存的工具列表", + zap.String("name", name), + zap.String("reason", reason), + zap.Int("count", len(tools)), + zap.Error(originalErr), + ) + return tools, nil + } + + if originalErr != nil { + return nil, fmt.Errorf("获取外部MCP工具失败且无缓存: %w", originalErr) + } + return nil, fmt.Errorf("外部MCP无缓存工具: %s", name) +} + +func (m *ExternalMCPManager) isToolCacheFresh(updatedAt time.Time) bool { + return !updatedAt.IsZero() && time.Since(updatedAt) < externalToolListCacheTTL +} + +func cloneTools(tools []Tool) []Tool { + if len(tools) == 0 { + return nil + } + out := make([]Tool, len(tools)) + copy(out, tools) + return out +} + +func (m *ExternalMCPManager) getFreshCachedTools(name string) ([]Tool, bool) { + m.toolCacheMu.RLock() + entry, ok := m.toolCache[name] + m.toolCacheMu.RUnlock() + if !ok || !m.isToolCacheFresh(entry.updatedAt) { + return nil, false + } + return cloneTools(entry.tools), true +} + +func (m *ExternalMCPManager) getAnyCachedTools(name string) ([]Tool, bool) { + m.toolCacheMu.RLock() + entry, ok := m.toolCache[name] + m.toolCacheMu.RUnlock() + if !ok { + return nil, false + } + return cloneTools(entry.tools), true +} + +// listToolsDeduped 对同一 MCP 合并并发 ListTools，并更新 toolCache / toolCounts。 +func (m *ExternalMCPManager) listToolsDeduped(ctx context.Context, name string, client ExternalMCPClient) ([]Tool, error) { + m.listToolsMu.Lock() + if inflight, exists := m.listToolsInflight[name]; exists { + m.listToolsMu.Unlock() + select { + case <-inflight.done: + if inflight.err != nil { + return nil, inflight.err + } + return cloneTools(inflight.tools), nil + case <-ctx.Done(): + return nil, ctx.Err() + } + } + inflight := &listToolsInflight{done: make(chan struct{})} + m.listToolsInflight[name] = inflight + m.listToolsMu.Unlock() + + inflight.tools, inflight.err = client.ListTools(ctx) + if inflight.err == nil { + m.updateToolCache(name, inflight.tools) + } + + m.listToolsMu.Lock() + delete(m.listToolsInflight, name) + close(inflight.done) + m.listToolsMu.Unlock() + + if inflight.err != nil { + m.handleConnectionDead(name, client, inflight.err) + return nil, inflight.err + } + return cloneTools(inflight.tools), nil +} + +// InvalidateToolCache 清除指定外部 MCP 的工具列表缓存（手动刷新时使用） +func (m *ExternalMCPManager) InvalidateToolCache(name string) { + m.toolCacheMu.Lock() + delete(m.toolCache, name) + m.toolCacheMu.Unlock() +} + +// InvalidateAllToolCaches 清除所有外部 MCP 工具列表缓存 +func (m *ExternalMCPManager) InvalidateAllToolCaches() { + m.toolCacheMu.Lock() + m.toolCache = make(map[string]toolListCacheEntry) + m.toolCacheMu.Unlock() +} + +func (m *ExternalMCPManager) triggerToolListRefresh(name string, client ExternalMCPClient) { + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + _, _ = m.listToolsDeduped(ctx, name, client) + }() +} + +// updateToolCache 更新工具列表缓存与工具数量 +func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) { + stored := cloneTools(tools) + m.toolCacheMu.Lock() + m.toolCache[name] = toolListCacheEntry{tools: stored, updatedAt: time.Now()} + m.toolCacheMu.Unlock() + + m.toolCountsMu.Lock() + m.toolCounts[name] = len(stored) + m.toolCountsMu.Unlock() + + if len(stored) == 0 { + m.logger.Warn("外部MCP返回空工具列表", + zap.String("name", name), + zap.String("hint", "服务可能暂时不可用，工具列表为空"), + ) + } else { + m.logger.Debug("工具列表缓存已更新", + zap.String("name", name), + zap.Int("count", len(stored)), + ) + } +} + +// CallTool 调用外部MCP工具（返回执行ID） +func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) { + if m.executionService == nil { + m.executionService = NewExecutionService(m.storage, m.logger) + m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes) + m.executionService.ConfigureToolResultSpillRoot(m.spillRootDir) + } + var ownerUserID string + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + ownerUserID = principal.UserID + } + var mcpName, actualToolName string + var client ExternalMCPClient + handle, err := m.executionService.Submit(ctx, ExecutionRequest{ + ToolName: toolName, + Arguments: args, + ConversationID: MCPConversationIDFromContext(ctx), + OwnerUserID: ownerUserID, + PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) { + _, authenticated := authctx.PrincipalFromContext(runCtx) + m.mu.RLock() + authorizer := m.toolAuthorizer + m.mu.RUnlock() + if authorizer != nil { + if err := authorizer(runCtx, toolName, args); err != nil { + return nil, fmt.Errorf("external tool authorization denied: %w", err) + } + } else if authenticated { + return nil, fmt.Errorf("external tool authorization policy is not configured") + } + + // 解析工具名称：name::toolName + if idx := findSubstring(toolName, "::"); idx > 0 { + mcpName = toolName[:idx] + actualToolName = toolName[idx+2:] + } else { + return nil, fmt.Errorf("无效的工具名称格式: %s", toolName) + } + + var exists bool + client, exists = m.GetClient(mcpName) + if !exists { + return nil, fmt.Errorf("外部MCP客户端不存在: %s", mcpName) + } + if err := m.checkExternalMCPCircuit(mcpName); err != nil { + return nil, err + } + + // 检查连接状态，如果未连接或状态为error，不允许调用 + if !client.IsConnected() { + status := client.GetStatus() + if status == "error" { + // 获取错误信息（如果有） + errorMsg := m.GetError(mcpName) + if errorMsg != "" { + return nil, fmt.Errorf("外部MCP连接失败: %s (错误: %s)", mcpName, errorMsg) + } + return nil, fmt.Errorf("外部MCP连接失败: %s", mcpName) + } + return nil, fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status) + } + + release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName) + if acquireErr != nil { + return nil, acquireErr + } + return release, nil + }, + Run: func(runCtx context.Context) (*ToolResult, error) { + result, callErr := client.CallTool(runCtx, actualToolName, args) + if callErr != nil { + m.handleConnectionDead(mcpName, client, callErr) + } + return result, callErr + }, + OnDone: func(exec *ToolExecution) { + failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled + if mcpName != "" { + m.recordExternalMCPResult(mcpName, failed) + } + m.updateStats(toolName, failed) + }, + }) + if err != nil { + return nil, "", err + } + + m.mu.RLock() + waitTimeout := m.toolWaitTimeout + m.mu.RUnlock() + snapshot, waitErr := m.executionService.Wait(ctx, handle.ID, waitTimeout) + if errors.Is(waitErr, ErrExecutionWaitTimeout) { + return externalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil + } + if waitErr != nil { + return nil, handle.ID, waitErr + } + if snapshot == nil || snapshot.Execution == nil { + return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回执行快照"}}, IsError: true}, handle.ID, nil + } + if snapshot.Execution.Result != nil { + return snapshot.Execution.Result, handle.ID, nil + } + if snapshot.Execution.Error != "" { + return nil, handle.ID, errors.New(snapshot.Execution.Error) + } + return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回结果"}}, IsError: false}, handle.ID, nil +} + +func externalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult { + execID := "" + status := ToolExecutionStatusRunning + toolName := "" + elapsed := time.Duration(0) + if snapshot != nil && snapshot.Execution != nil { + execID = snapshot.Execution.ID + status = snapshot.Execution.Status + toolName = snapshot.Execution.ToolName + elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second) + } + waitText := "unbounded" + if waitTimeout > 0 { + waitText = waitTimeout.Round(time.Second).String() + } + msg := fmt.Sprintf(`工具已提交到后台执行，但本次等待已到达上限。 + +execution_id: %s +tool: %s +status: %s +wait_timeout: %s +elapsed: %s + +你可以继续推理、改用其他工具，或调用 wait_tool_execution 继续等待该 execution_id；也可以调用 cancel_tool_execution 取消。`, execID, toolName, status, waitText, elapsed) + return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true} +} + +func (m *ExternalMCPManager) checkExternalMCPCircuit(mcpName string) error { + if m == nil { + return nil + } + name := strings.TrimSpace(mcpName) + if name == "" { + return nil + } + m.mu.Lock() + defer m.mu.Unlock() + if m.resilience.CircuitFailureThreshold < 0 { + return nil + } + rt := m.externalMCPRuntimeLocked(name) + if rt == nil || rt.circuitOpenUntil.IsZero() { + return nil + } + now := time.Now() + if now.Before(rt.circuitOpenUntil) { + return fmt.Errorf("外部MCP服务 %s 已临时熔断，预计 %s 后重试", name, time.Until(rt.circuitOpenUntil).Round(time.Second)) + } + rt.circuitOpenUntil = time.Time{} + return nil +} + +func (m *ExternalMCPManager) acquireExternalMCPCallSlot(ctx context.Context, mcpName string) (func(), error) { + if m == nil { + return func() {}, nil + } + name := strings.TrimSpace(mcpName) + m.mu.Lock() + rt := m.externalMCPRuntimeLocked(name) + serverSem := chan struct{}(nil) + if rt != nil { + serverSem = rt.semaphore + } + globalSem := m.globalSemaphore + m.mu.Unlock() + + releaseGlobal := false + if globalSem != nil { + select { + case globalSem <- struct{}{}: + releaseGlobal = true + case <-ctxDone(ctx): + return func() {}, contextErr(ctx) + } + } + releaseServer := false + if serverSem != nil { + select { + case serverSem <- struct{}{}: + releaseServer = true + case <-ctxDone(ctx): + if releaseGlobal { + <-globalSem + } + return func() {}, contextErr(ctx) + } + } + return func() { + if releaseServer { + <-serverSem + } + if releaseGlobal { + <-globalSem + } + }, nil +} + +func contextErr(ctx context.Context) error { + if ctx == nil || ctx.Err() == nil { + return context.Canceled + } + return ctx.Err() +} + +func (m *ExternalMCPManager) recordExternalMCPResult(mcpName string, failed bool) { + if m == nil || strings.TrimSpace(mcpName) == "" { + return + } + m.mu.Lock() + defer m.mu.Unlock() + rt := m.externalMCPRuntimeLocked(mcpName) + if rt == nil { + return + } + if !failed { + rt.consecutiveFailures = 0 + rt.circuitOpenUntil = time.Time{} + return + } + if m.resilience.CircuitFailureThreshold < 0 { + return + } + rt.consecutiveFailures++ + if rt.consecutiveFailures >= m.resilience.CircuitFailureThreshold { + rt.circuitOpenUntil = time.Now().Add(m.resilience.CircuitCooldown) + m.logger.Warn("外部MCP服务触发熔断", + zap.String("name", mcpName), + zap.Int("consecutiveFailures", rt.consecutiveFailures), + zap.Duration("cooldown", m.resilience.CircuitCooldown), + ) + } +} + +func (m *ExternalMCPManager) externalMCPRuntimeLocked(mcpName string) *externalMCPServerRuntime { + if m.serverRuntimes == nil { + m.serverRuntimes = make(map[string]*externalMCPServerRuntime) + } + name := strings.TrimSpace(mcpName) + if name == "" { + return nil + } + if rt := m.serverRuntimes[name]; rt != nil { + return rt + } + var sem chan struct{} + if m.resilience.MaxConcurrentPerServer > 0 { + sem = make(chan struct{}, m.resilience.MaxConcurrentPerServer) + } + rt := &externalMCPServerRuntime{semaphore: sem} + m.serverRuntimes[name] = rt + return rt +} + +func (m *ExternalMCPManager) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) { + note := strings.TrimSpace(m.readAbortUserNote(executionID)) + if note == "" { + return false + } + hasErr := err != nil && *err != nil + hasRes := result != nil && *result != nil + if !hasErr && !hasRes { + return false + } + _ = m.takeAbortUserNote(executionID) + partial := "" + if hasRes { + partial = ToolResultPlainText(*result) + } + if partial == "" && hasErr { + partial = (*err).Error() + } + merged := MergePartialToolOutputAndAbortNote(partial, note) + *err = nil + *result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true} + return true +} + +func (m *ExternalMCPManager) readAbortUserNote(id string) string { + m.mu.Lock() + defer m.mu.Unlock() + if m.abortUserNotes == nil { + return "" + } + return m.abortUserNotes[id] +} + +func (m *ExternalMCPManager) takeAbortUserNote(id string) string { + m.mu.Lock() + defer m.mu.Unlock() + if m.abortUserNotes == nil { + return "" + } + n := m.abortUserNotes[id] + delete(m.abortUserNotes, id) + return n +} + +// cleanupOldExecutions 清理旧的执行记录（保持内存中的记录数量在限制内） +func (m *ExternalMCPManager) cleanupOldExecutions() { + const maxExecutionsInMemory = 1000 + if len(m.executions) <= maxExecutionsInMemory { + return + } + + // 按开始时间排序，删除最旧的记录 + type execTime struct { + id string + startTime time.Time + } + var execs []execTime + for id, exec := range m.executions { + execs = append(execs, execTime{id: id, startTime: exec.StartTime}) + } + + // 按时间排序 + for i := 0; i < len(execs)-1; i++ { + for j := i + 1; j < len(execs); j++ { + if execs[i].startTime.After(execs[j].startTime) { + execs[i], execs[j] = execs[j], execs[i] + } + } + } + + // 删除最旧的记录 + toDelete := len(m.executions) - maxExecutionsInMemory + for i := 0; i < toDelete && i < len(execs); i++ { + delete(m.executions, execs[i].id) + } +} + +// GetExecution 获取执行记录（先从内存查找，再从数据库查找） +func (m *ExternalMCPManager) GetExecution(id string) (*ToolExecution, bool) { + if m.executionService != nil { + if snap, err := m.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil { + return snap.Execution, true + } + } + m.mu.RLock() + exec, exists := m.executions[id] + m.mu.RUnlock() + + if exists { + return exec, true + } + + if m.storage != nil { + exec, err := m.storage.GetToolExecution(id) + if err == nil { + return exec, true + } + } + + return nil, false +} + +func (m *ExternalMCPManager) registerRunningCancel(id string, cancel context.CancelFunc) { + m.mu.Lock() + m.runningCancels[id] = cancel + m.mu.Unlock() +} + +func (m *ExternalMCPManager) unregisterRunningCancel(id string) { + m.mu.Lock() + delete(m.runningCancels, id) + m.mu.Unlock() +} + +// CancelToolExecutionWithNote 取消外部 MCP 工具；note 非空时与已返回输出合并后交给模型。 +func (m *ExternalMCPManager) CancelToolExecutionWithNote(id string, note string) bool { + if m.executionService != nil && m.executionService.Cancel(id, note) { + return true + } + m.mu.Lock() + cancel, ok := m.runningCancels[id] + if !ok || cancel == nil { + m.mu.Unlock() + return false + } + if strings.TrimSpace(note) != "" { + if m.abortUserNotes == nil { + m.abortUserNotes = make(map[string]string) + } + m.abortUserNotes[id] = strings.TrimSpace(note) + } + m.mu.Unlock() + cancel() + return true +} + +// CancelToolExecution 取消正在执行的外部 MCP 工具（无用户说明）。 +func (m *ExternalMCPManager) CancelToolExecution(id string) bool { + return m.CancelToolExecutionWithNote(id, "") +} + +// ActiveRunningExecutionIDs 返回当前进程内仍登记 cancel 的外部 MCP executionId 快照。 +func (m *ExternalMCPManager) ActiveRunningExecutionIDs() map[string]struct{} { + if m == nil { + return nil + } + if m.executionService != nil { + if ids := m.executionService.ActiveRunningExecutionIDs(); len(ids) > 0 { + return ids + } + } + m.mu.Lock() + defer m.mu.Unlock() + if len(m.runningCancels) == 0 { + return nil + } + out := make(map[string]struct{}, len(m.runningCancels)) + for id := range m.runningCancels { + out[id] = struct{}{} + } + return out +} + +// updateStats 更新统计信息 +func (m *ExternalMCPManager) updateStats(toolName string, failed bool) { + now := time.Now() + if m.storage != nil { + totalCalls := 1 + successCalls := 0 + failedCalls := 0 + if failed { + failedCalls = 1 + } else { + successCalls = 1 + } + if err := m.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil { + m.logger.Warn("保存统计信息到数据库失败", zap.Error(err)) + } + return + } + + m.mu.Lock() + defer m.mu.Unlock() + + if m.stats[toolName] == nil { + m.stats[toolName] = &ToolStats{ + ToolName: toolName, + } + } + + stats := m.stats[toolName] + stats.TotalCalls++ + stats.LastCallTime = &now + + if failed { + stats.FailedCalls++ + } else { + stats.SuccessCalls++ + } +} + +// GetStats 获取MCP服务器统计信息 +func (m *ExternalMCPManager) GetStats() map[string]interface{} { + m.mu.RLock() + defer m.mu.RUnlock() + + total := len(m.configs) + enabled := 0 + disabled := 0 + connected := 0 + + for name, cfg := range m.configs { + if m.isEnabled(cfg) { + enabled++ + if client, exists := m.clients[name]; exists && client.IsConnected() { + connected++ + } + } else { + disabled++ + } + } + + return map[string]interface{}{ + "total": total, + "enabled": enabled, + "disabled": disabled, + "connected": connected, + } +} + +// GetToolStats 获取工具统计信息（合并内存和数据库） +// 只返回外部MCP工具的统计信息（工具名称包含 "::"） +func (m *ExternalMCPManager) GetToolStats() map[string]*ToolStats { + result := make(map[string]*ToolStats) + + // 从数据库加载统计信息（如果使用数据库存储） + if m.storage != nil { + dbStats, err := m.storage.LoadToolStats() + if err == nil { + // 只保留外部MCP工具的统计信息（工具名称包含 "::"） + for k, v := range dbStats { + if findSubstring(k, "::") > 0 { + result[k] = v + } + } + } else { + m.logger.Warn("从数据库加载统计信息失败", zap.Error(err)) + } + } + + // 合并内存中的统计信息 + m.mu.RLock() + for k, v := range m.stats { + // 如果数据库中已有该工具的统计信息，合并它们 + if existing, exists := result[k]; exists { + // 创建新的统计信息对象，避免修改共享对象 + merged := &ToolStats{ + ToolName: k, + TotalCalls: existing.TotalCalls + v.TotalCalls, + SuccessCalls: existing.SuccessCalls + v.SuccessCalls, + FailedCalls: existing.FailedCalls + v.FailedCalls, + } + // 使用最新的调用时间 + if v.LastCallTime != nil && (existing.LastCallTime == nil || v.LastCallTime.After(*existing.LastCallTime)) { + merged.LastCallTime = v.LastCallTime + } else if existing.LastCallTime != nil { + timeCopy := *existing.LastCallTime + merged.LastCallTime = &timeCopy + } + result[k] = merged + } else { + // 如果数据库中没有，直接使用内存中的统计信息 + statCopy := *v + result[k] = &statCopy + } + } + m.mu.RUnlock() + + return result +} + +// GetToolCount 获取指定外部MCP的工具数量（从缓存读取，不阻塞） +func (m *ExternalMCPManager) GetToolCount(name string) (int, error) { + // 先从缓存读取 + m.toolCountsMu.RLock() + if count, exists := m.toolCounts[name]; exists { + m.toolCountsMu.RUnlock() + return count, nil + } + m.toolCountsMu.RUnlock() + + // 如果缓存中没有，检查客户端状态 + client, exists := m.GetClient(name) + if !exists { + return 0, fmt.Errorf("客户端不存在: %s", name) + } + + if !client.IsConnected() { + // 未连接，缓存为0 + m.toolCountsMu.Lock() + m.toolCounts[name] = 0 + m.toolCountsMu.Unlock() + return 0, nil + } + + // 如果已连接但缓存中没有，触发异步刷新并返回0（避免阻塞） + m.triggerToolCountRefresh() + return 0, nil +} + +// GetToolCounts 获取所有外部MCP的工具数量（从缓存读取，不阻塞） +func (m *ExternalMCPManager) GetToolCounts() map[string]int { + m.toolCountsMu.RLock() + defer m.toolCountsMu.RUnlock() + + // 返回缓存的副本，避免外部修改 + result := make(map[string]int) + for k, v := range m.toolCounts { + result[k] = v + } + return result +} + +// refreshToolCounts 刷新工具数量缓存（后台异步执行） +// 使用 atomic flag 防止并发堆积：如果上一次刷新尚未完成，本次触发直接跳过。 +func (m *ExternalMCPManager) refreshToolCounts() { + if !m.refreshing.CompareAndSwap(false, true) { + return // 上一次刷新尚未完成，跳过 + } + defer m.refreshing.Store(false) + + m.mu.RLock() + clients := make(map[string]ExternalMCPClient) + for k, v := range m.clients { + clients[k] = v + } + m.mu.RUnlock() + + newCounts := make(map[string]int) + + // 使用goroutine并发获取每个客户端的工具数量，避免串行阻塞 + type countResult struct { + name string + count int + } + resultChan := make(chan countResult, len(clients)) + + for name, client := range clients { + go func(n string, c ExternalMCPClient) { + if !c.IsConnected() { + resultChan <- countResult{name: n, count: 0} + return + } + + // 缓存仍新鲜时直接复用，避免与 GetAllTools 重复打远程 + if _, fresh := m.getFreshCachedTools(n); fresh { + m.toolCountsMu.RLock() + count := m.toolCounts[n] + m.toolCountsMu.RUnlock() + resultChan <- countResult{name: n, count: count} + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + tools, err := m.listToolsDeduped(ctx, n, c) + cancel() + + if err != nil { + if !isConnectionDeadError(err) { + m.logger.Warn("获取外部MCP工具数量失败，请检查连接或服务端 tools/list", + zap.String("name", n), + zap.Error(err), + ) + } + resultChan <- countResult{name: n, count: -1} + return + } + + resultChan <- countResult{name: n, count: len(tools)} + }(name, client) + } + + // 收集结果 + m.toolCountsMu.RLock() + oldCounts := make(map[string]int) + for k, v := range m.toolCounts { + oldCounts[k] = v + } + m.toolCountsMu.RUnlock() + + for i := 0; i < len(clients); i++ { + result := <-resultChan + if result.count >= 0 { + newCounts[result.name] = result.count + } else { + // 获取失败，保留旧值 + if oldCount, exists := oldCounts[result.name]; exists { + newCounts[result.name] = oldCount + } else { + newCounts[result.name] = 0 + } + } + } + + // 更新缓存 + m.toolCountsMu.Lock() + // 更新所有获取到的值 + for name, count := range newCounts { + m.toolCounts[name] = count + } + // 对于未连接的客户端，设置为0 + for name, client := range clients { + if !client.IsConnected() { + m.toolCounts[name] = 0 + } + } + m.toolCountsMu.Unlock() +} + +// refreshToolCache 刷新指定MCP的工具列表缓存 +func (m *ExternalMCPManager) refreshToolCache(name string, client ExternalMCPClient) { + if !client.IsConnected() { + return + } + if client.GetStatus() == "error" { + m.logger.Debug("跳过刷新工具列表缓存（连接失败）", + zap.String("name", name), + ) + return + } + + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if _, err := m.listToolsDeduped(ctx, name, client); err != nil { + m.logger.Debug("刷新工具列表缓存失败", + zap.String("name", name), + zap.Error(err), + ) + } +} + +// startToolCountRefresh 启动后台刷新工具数量的goroutine +func (m *ExternalMCPManager) startToolCountRefresh() { + m.refreshWg.Add(1) + go func() { + defer m.refreshWg.Done() + ticker := time.NewTicker(externalToolCountRefreshInterval) + defer ticker.Stop() + + // 立即执行一次刷新 + m.refreshToolCounts() + + for { + select { + case <-ticker.C: + m.refreshToolCounts() + case <-m.stopRefresh: + return + } + } + }() +} + +// triggerToolCountRefresh 触发立即刷新工具数量（异步） +func (m *ExternalMCPManager) triggerToolCountRefresh() { + go m.refreshToolCounts() +} + +// createClient 创建客户端（不连接）。统一使用官方 MCP Go SDK 的 lazy 客户端，连接在 Initialize 时完成。 +func (m *ExternalMCPManager) createClient(serverCfg config.ExternalMCPServerConfig) ExternalMCPClient { + transport := serverCfg.GetTransportType() + + switch transport { + case "http": + if serverCfg.URL == "" { + return nil + } + return newLazySDKClient(serverCfg, m.logger) + case "stdio": + if serverCfg.Command == "" { + return nil + } + return newLazySDKClient(serverCfg, m.logger) + case "sse": + if serverCfg.URL == "" { + return nil + } + return newLazySDKClient(serverCfg, m.logger) + default: + if transport == "" { + return nil + } + // 未知传输类型也尝试使用 lazy client + return newLazySDKClient(serverCfg, m.logger) + } +} + +// doConnect 执行实际连接 +func (m *ExternalMCPManager) doConnect(name string, serverCfg config.ExternalMCPServerConfig, client ExternalMCPClient) error { + timeout := time.Duration(serverCfg.Timeout) * time.Second + if timeout <= 0 { + timeout = 30 * time.Second + } + + // 初始化连接 + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + if err := client.Initialize(ctx); err != nil { + return err + } + + m.logger.Info("外部MCP客户端已连接", + zap.String("name", name), + ) + + return nil +} + +// setClientStatus 设置客户端状态（通过类型断言） +func (m *ExternalMCPManager) setClientStatus(client ExternalMCPClient, status string) { + if c, ok := client.(*lazySDKClient); ok { + c.setStatus(status) + } +} + +// connectClient 连接客户端（异步）- 保留用于向后兼容 +func (m *ExternalMCPManager) connectClient(name string, serverCfg config.ExternalMCPServerConfig) error { + client := m.createClient(serverCfg) + if client == nil { + return fmt.Errorf("无法创建客户端：不支持的传输模式") + } + + // 设置状态为connecting + m.setClientStatus(client, "connecting") + + // 初始化连接 + timeout := time.Duration(serverCfg.Timeout) * time.Second + if timeout <= 0 { + timeout = 30 * time.Second + } + + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + if err := client.Initialize(ctx); err != nil { + m.logger.Error("初始化外部MCP客户端失败", + zap.String("name", name), + zap.Error(err), + ) + return err + } + + // 保存客户端 + m.mu.Lock() + m.clients[name] = client + m.mu.Unlock() + + m.logger.Info("外部MCP客户端已连接", + zap.String("name", name), + ) + + m.onClientConnected(name) + + // 连接成功，触发工具数量刷新和工具列表缓存刷新 + m.triggerToolCountRefresh() + m.mu.RLock() + if client, exists := m.clients[name]; exists { + m.refreshToolCache(name, client) + } + m.mu.RUnlock() + + return nil +} + +// isEnabled 检查是否启用 +func (m *ExternalMCPManager) isEnabled(cfg config.ExternalMCPServerConfig) bool { + return cfg.ExternalMCPEnable +} + +// findSubstring 查找子字符串（简单实现） +func findSubstring(s, substr string) int { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return i + } + } + return -1 +} + +// StartAllEnabled 启动所有启用的客户端 +func (m *ExternalMCPManager) StartAllEnabled() { + m.mu.RLock() + configs := make(map[string]config.ExternalMCPServerConfig) + for k, v := range m.configs { + configs[k] = v + } + m.mu.RUnlock() + + for name, cfg := range configs { + if m.isEnabled(cfg) { + go func(n string, c config.ExternalMCPServerConfig) { + if err := m.connectClient(n, c); err != nil { + // 检查是否是连接被拒绝的错误（服务可能还没启动） + errStr := strings.ToLower(err.Error()) + isConnectionRefused := strings.Contains(errStr, "connection refused") || + strings.Contains(errStr, "dial tcp") || + strings.Contains(errStr, "connect: connection refused") + + if isConnectionRefused { + // 连接被拒绝，说明目标服务可能还没启动，这是正常的 + // 使用 Warn 级别，提示用户这是正常的，可以通过手动启动或等待服务启动后自动连接 + fields := []zap.Field{ + zap.String("name", n), + zap.String("message", "目标服务可能尚未启动，这是正常的。服务启动后可通过界面手动连接，或等待自动重试"), + zap.Error(err), + } + + transport := c.GetTransportType() + + if transport == "http" && c.URL != "" { + fields = append(fields, zap.String("url", c.URL)) + } else if transport == "stdio" && c.Command != "" { + fields = append(fields, zap.String("command", c.Command)) + } + + m.logger.Warn("外部MCP服务暂未就绪", fields...) + } else { + // 其他错误，使用 Error 级别 + m.logger.Error("启动外部MCP客户端失败", + zap.String("name", n), + zap.Error(err), + ) + } + } + }(name, cfg) + } + } +} + +// StopAll 停止所有客户端 +func (m *ExternalMCPManager) StopAll() { + if m.executionService != nil { + m.executionService.CancelAll("外部 MCP 管理器正在停止") + } + clients := make(map[string]ExternalMCPClient) + m.mu.Lock() + for name, client := range m.clients { + clients[name] = client + delete(m.clients, name) + } + m.mu.Unlock() + + for name, client := range clients { + if client != nil { + _ = client.Close() + } + m.clearReconnectState(name) + } + + // 清理所有工具数量缓存 + m.toolCountsMu.Lock() + m.toolCounts = make(map[string]int) + m.toolCountsMu.Unlock() + + // 清理所有工具列表缓存 + m.toolCacheMu.Lock() + m.toolCache = make(map[string]toolListCacheEntry) + m.toolCacheMu.Unlock() + + // 停止后台刷新（使用 select 避免重复关闭 channel） + select { + case <-m.stopRefresh: + // 已经关闭，不需要再次关闭 + default: + close(m.stopRefresh) + } + m.refreshWg.Wait() +} diff --git a/internal/mcp/external_manager_async_test.go b/internal/mcp/external_manager_async_test.go new file mode 100644 index 00000000..6c280ed9 --- /dev/null +++ b/internal/mcp/external_manager_async_test.go @@ -0,0 +1,230 @@ +package mcp + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + "time" + + "go.uber.org/zap" +) + +type blockingExternalMCPClient struct { + started chan struct{} + calls chan string + release chan struct{} + result *ToolResult + count atomic.Int32 +} + +func newBlockingExternalMCPClient(resultText string) *blockingExternalMCPClient { + return &blockingExternalMCPClient{ + started: make(chan struct{}), + calls: make(chan string, 8), + release: make(chan struct{}), + result: &ToolResult{Content: []Content{{Type: "text", Text: resultText}}}, + } +} + +func (c *blockingExternalMCPClient) Initialize(ctx context.Context) error { return nil } +func (c *blockingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) { + return []Tool{{Name: "slow_tool"}}, nil +} +func (c *blockingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) { + c.count.Add(1) + select { + case c.calls <- name: + default: + } + select { + case <-c.started: + default: + close(c.started) + } + select { + case <-c.release: + return c.result, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} +func (c *blockingExternalMCPClient) Close() error { return nil } +func (c *blockingExternalMCPClient) IsConnected() bool { return true } +func (c *blockingExternalMCPClient) GetStatus() string { return "connected" } + +type failingExternalMCPClient struct{} + +func (c *failingExternalMCPClient) Initialize(ctx context.Context) error { return nil } +func (c *failingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) { + return []Tool{{Name: "fail_tool"}}, nil +} +func (c *failingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) { + return nil, errors.New("boom") +} +func (c *failingExternalMCPClient) Close() error { return nil } +func (c *failingExternalMCPClient) IsConnected() bool { return true } +func (c *failingExternalMCPClient) GetStatus() string { return "connected" } + +func TestExternalMCPManager_CallToolBoundedWaitThenContinue(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + manager.ConfigureToolWaitTimeoutSeconds(1) + manager.toolWaitTimeout = 10 * time.Millisecond + client := newBlockingExternalMCPClient("slow result ready") + manager.clients["lab"] = client + + callCtx, callCancel := context.WithCancel(context.Background()) + result, executionID, err := manager.CallTool(callCtx, "lab::slow_tool", map[string]interface{}{"target": "example"}) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if executionID == "" { + t.Fatal("expected execution id") + } + if result == nil || !result.IsError { + t.Fatalf("expected soft timeout tool result, got %#v", result) + } + text := ToolResultPlainText(result) + if !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") { + t.Fatalf("timeout result should include execution id and wait guidance, got %q", text) + } + + select { + case <-client.started: + default: + t.Fatal("worker did not start") + } + callCancel() + close(client.release) + + snapshot, err := manager.executionService.Wait(context.Background(), executionID, time.Second) + if err != nil { + t.Fatalf("Wait returned error: %v", err) + } + if snapshot == nil || snapshot.Execution == nil { + t.Fatal("expected execution snapshot") + } + if snapshot.Execution.Status != ToolExecutionStatusCompleted { + t.Fatalf("status = %q, want completed", snapshot.Execution.Status) + } + if got := ToolResultPlainText(snapshot.Execution.Result); got != "slow result ready" { + t.Fatalf("result = %q, want slow result ready", got) + } +} + +func TestExecutionControlWaitToolReturnsCompletedResult(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + manager.toolWaitTimeout = 10 * time.Millisecond + client := newBlockingExternalMCPClient("control wait result") + manager.clients["lab"] = client + + result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if result == nil || !result.IsError || executionID == "" { + t.Fatalf("expected soft timeout and execution id, got result=%#v id=%q", result, executionID) + } + + server := NewServer(zap.NewNop()) + RegisterExecutionControlTools(server, manager) + close(client.release) + + waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "timeout_seconds": 1, + }) + if err != nil { + t.Fatalf("wait_tool_execution returned error: %v", err) + } + if waitResult == nil || waitResult.IsError { + t.Fatalf("expected successful wait result, got %#v", waitResult) + } + body := ToolResultPlainText(waitResult) + if !strings.Contains(body, `"status": "completed"`) || !strings.Contains(body, "control wait result") { + t.Fatalf("wait result body missing completed status/result: %s", body) + } +} + +func TestExternalMCPManager_PerServerConcurrencyLimitsWorkers(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + manager.toolWaitTimeout = 10 * time.Millisecond + manager.ConfigureResilience(ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: 1, + MaxConcurrentTotal: 4, + CircuitFailureThreshold: -1, + CircuitCooldown: time.Second, + }) + client := newBlockingExternalMCPClient("ok") + manager.clients["lab"] = client + + done1 := make(chan struct{}) + go func() { + _, _, _ = manager.CallTool(context.Background(), "lab::slow_tool", nil) + close(done1) + }() + select { + case <-client.calls: + case <-time.After(time.Second): + t.Fatal("first worker did not enter client") + } + + type callOutcome struct { + executionID string + err error + } + done2 := make(chan callOutcome, 1) + go func() { + _, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil) + done2 <- callOutcome{executionID: executionID, err: err} + }() + select { + case <-client.calls: + t.Fatal("second worker entered client before per-server slot was released") + case <-time.After(50 * time.Millisecond): + } + var second callOutcome + select { + case second = <-done2: + case <-time.After(time.Second): + t.Fatal("second call did not return after bounded wait") + } + if second.err != nil || second.executionID == "" { + t.Fatalf("second call should return queued execution id after bounded wait, id=%q err=%v", second.executionID, second.err) + } + snapshot, err := manager.executionService.Get(second.executionID) + if err != nil { + t.Fatalf("Get queued execution: %v", err) + } + if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusQueued { + t.Fatalf("second execution status = %#v, want queued", snapshot) + } + close(client.release) + select { + case <-client.calls: + case <-time.After(time.Second): + t.Fatal("second worker did not enter client after slot release") + } + <-done1 +} + +func TestExternalMCPManager_CircuitBreakerOpensAfterFailures(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + manager.ConfigureResilience(ExternalMCPResilienceConfig{ + MaxConcurrentPerServer: 2, + MaxConcurrentTotal: 4, + CircuitFailureThreshold: 1, + CircuitCooldown: time.Minute, + }) + manager.clients["lab"] = &failingExternalMCPClient{} + + _, _, err := manager.CallTool(context.Background(), "lab::fail_tool", nil) + if err == nil || !strings.Contains(err.Error(), "boom") { + t.Fatalf("expected first call to fail with client error, got %v", err) + } + _, _, err = manager.CallTool(context.Background(), "lab::fail_tool", nil) + if err == nil || !strings.Contains(err.Error(), "熔断") { + t.Fatalf("expected circuit breaker rejection, got %v", err) + } +} diff --git a/internal/mcp/external_manager_test.go b/internal/mcp/external_manager_test.go new file mode 100644 index 00000000..3baff567 --- /dev/null +++ b/internal/mcp/external_manager_test.go @@ -0,0 +1,261 @@ +package mcp + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/config" + + "go.uber.org/zap" +) + +func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) { + manager := NewExternalMCPManager(zap.NewNop()) + t.Cleanup(manager.StopAll) + manager.SetToolAuthorizer(func(context.Context, string, map[string]interface{}) error { + return errors.New("denied by policy") + }) + ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true})) + _, executionID, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{}) + if err == nil || !strings.Contains(err.Error(), "authorization denied") { + t.Fatalf("external call bypassed authorizer: %v", err) + } + if executionID == "" { + t.Fatal("denied external call should still return an execution id") + } + execution, ok := manager.GetExecution(executionID) + if !ok || execution == nil { + t.Fatalf("missing denied external execution %q", executionID) + } + if execution.Status != ToolExecutionStatusFailed || !strings.Contains(execution.Error, "denied by policy") { + t.Fatalf("denied external execution = %#v, want failed with policy error", execution) + } +} + +func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + // 测试添加stdio配置 + stdioCfg := config.ExternalMCPServerConfig{ + Command: "python3", + Args: []string{"/path/to/script.py"}, + Description: "Test stdio MCP", + Timeout: 30, + ExternalMCPEnable: true, + } + + err := manager.AddOrUpdateConfig("test-stdio", stdioCfg) + if err != nil { + t.Fatalf("添加stdio配置失败: %v", err) + } + + // 测试添加HTTP配置 + httpCfg := config.ExternalMCPServerConfig{ + Type: "http", + URL: "http://127.0.0.1:8081/mcp", + Description: "Test HTTP MCP", + Timeout: 30, + ExternalMCPEnable: false, + } + + err = manager.AddOrUpdateConfig("test-http", httpCfg) + if err != nil { + t.Fatalf("添加HTTP配置失败: %v", err) + } + + // 验证配置已保存 + configs := manager.GetConfigs() + if len(configs) != 2 { + t.Fatalf("期望2个配置，实际%d个", len(configs)) + } + + if configs["test-stdio"].Command != stdioCfg.Command { + t.Errorf("stdio配置命令不匹配") + } + + if configs["test-http"].URL != httpCfg.URL { + t.Errorf("HTTP配置URL不匹配") + } +} + +func TestExternalMCPManager_RemoveConfig(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + cfg := config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: false, + } + + manager.AddOrUpdateConfig("test-remove", cfg) + + // 移除配置 + err := manager.RemoveConfig("test-remove") + if err != nil { + t.Fatalf("移除配置失败: %v", err) + } + + configs := manager.GetConfigs() + if _, exists := configs["test-remove"]; exists { + t.Error("配置应该已被移除") + } +} + +func TestExternalMCPManager_GetStats(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + // 添加多个配置 + manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: true, + }) + + manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{ + URL: "http://127.0.0.1:8081/mcp", + ExternalMCPEnable: true, + }) + + manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: false, + }) + + stats := manager.GetStats() + + if stats["total"].(int) != 3 { + t.Errorf("期望总数3，实际%d", stats["total"]) + } + + if stats["enabled"].(int) != 2 { + t.Errorf("期望启用数2，实际%d", stats["enabled"]) + } + + if stats["disabled"].(int) != 1 { + t.Errorf("期望停用数1，实际%d", stats["disabled"]) + } +} + +func TestExternalMCPManager_LoadConfigs(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + externalMCPConfig := config.ExternalMCPConfig{ + Servers: map[string]config.ExternalMCPServerConfig{ + "loaded1": { + Command: "python3", + ExternalMCPEnable: true, + }, + "loaded2": { + URL: "http://127.0.0.1:8081/mcp", + ExternalMCPEnable: false, + }, + }, + } + + manager.LoadConfigs(&externalMCPConfig) + + configs := manager.GetConfigs() + if len(configs) != 2 { + t.Fatalf("期望2个配置，实际%d个", len(configs)) + } + + if configs["loaded1"].Command != "python3" { + t.Error("配置1加载失败") + } + + if configs["loaded2"].URL != "http://127.0.0.1:8081/mcp" { + t.Error("配置2加载失败") + } +} + +// TestLazySDKClient_InitializeFails 验证无效配置时 SDK 客户端 Initialize 失败并设置 error 状态 +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, + } + c := newLazySDKClient(cfg, logger) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := c.Initialize(ctx) + if err == nil { + t.Fatal("expected error when connecting to invalid server") + } + if c.GetStatus() != "error" { + t.Errorf("expected status error, got %s", c.GetStatus()) + } + c.Close() +} + +func TestExternalMCPManager_StartStopClient(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + // 添加一个禁用的配置 + cfg := config.ExternalMCPServerConfig{ + Command: "python3", + ExternalMCPEnable: false, + } + + manager.AddOrUpdateConfig("test-start-stop", cfg) + + // 尝试启动（可能会失败，因为没有真实的服务器） + err := manager.StartClient("test-start-stop") + if err != nil { + t.Logf("启动失败（可能是没有服务器）: %v", err) + } + + // 停止 + err = manager.StopClient("test-start-stop") + if err != nil { + t.Fatalf("停止失败: %v", err) + } + + // 验证配置已更新为禁用 + configs := manager.GetConfigs() + if configs["test-start-stop"].ExternalMCPEnable { + t.Error("配置应该已被禁用") + } +} + +func TestExternalMCPManager_CallTool(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + // 测试调用不存在的工具 + _, _, err := manager.CallTool(context.Background(), "nonexistent::tool", map[string]interface{}{}) + if err == nil { + t.Error("应该返回错误") + } + + // 测试无效的工具名称格式 + _, _, err = manager.CallTool(context.Background(), "invalid-tool-name", map[string]interface{}{}) + if err == nil { + t.Error("应该返回错误（无效格式）") + } +} + +func TestExternalMCPManager_GetAllTools(t *testing.T) { + logger := zap.NewNop() + manager := NewExternalMCPManager(logger) + + ctx := context.Background() + tools, err := manager.GetAllTools(ctx) + if err != nil { + t.Fatalf("获取工具列表失败: %v", err) + } + + // 如果没有连接的客户端，应该返回空列表 + if len(tools) != 0 { + t.Logf("获取到%d个工具", len(tools)) + } +} diff --git a/internal/mcp/run_context.go b/internal/mcp/run_context.go new file mode 100644 index 00000000..7612032e --- /dev/null +++ b/internal/mcp/run_context.go @@ -0,0 +1,147 @@ +package mcp + +import ( + "context" + "strings" +) + +// ToolRunRegistry 在工具开始/结束时登记当前 executionId，供对话页「仅终止当前工具」与监控页共用取消逻辑。 +type ToolRunRegistry interface { + RegisterRunningTool(conversationID, executionID string) + UnregisterRunningTool(conversationID, executionID string) +} + +// EinoExecuteRunRegistry 登记进行中的 Eino filesystem execute，供「中断并继续」终止 amass 等长命令。 +type EinoExecuteRunRegistry interface { + RegisterActiveEinoExecute(conversationID string, cancel context.CancelFunc) + UnregisterActiveEinoExecute(conversationID string) + AbortActiveEinoExecute(conversationID, note string) bool + TakeEinoExecuteAbortNote(conversationID string) string +} + +type toolRunRegistryCtxKey struct{} +type einoExecuteRunRegistryCtxKey struct{} +type mcpConversationIDCtxKey struct{} +type mcpExecutionIDCtxKey struct{} +type mcpProjectIDCtxKey struct{} + +// WithToolRunRegistry 将登记器注入 ctx（Eino / 原生 Agent 任务 ctx）。 +func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context { + if ctx == nil || reg == nil { + return ctx + } + return context.WithValue(ctx, toolRunRegistryCtxKey{}, reg) +} + +// ToolRunRegistryFromContext 取出登记器（无则 nil）。 +func ToolRunRegistryFromContext(ctx context.Context) ToolRunRegistry { + if ctx == nil { + return nil + } + v, _ := ctx.Value(toolRunRegistryCtxKey{}).(ToolRunRegistry) + return v +} + +// WithEinoExecuteRunRegistry 将 Eino execute 取消登记器注入 ctx。 +func WithEinoExecuteRunRegistry(ctx context.Context, reg EinoExecuteRunRegistry) context.Context { + if ctx == nil || reg == nil { + return ctx + } + return context.WithValue(ctx, einoExecuteRunRegistryCtxKey{}, reg) +} + +// EinoExecuteRunRegistryFromContext 取出 Eino execute 登记器（无则 nil）。 +func EinoExecuteRunRegistryFromContext(ctx context.Context) EinoExecuteRunRegistry { + if ctx == nil { + return nil + } + v, _ := ctx.Value(einoExecuteRunRegistryCtxKey{}).(EinoExecuteRunRegistry) + return v +} + +// WithMCPConversationID 将对话 ID 注入 ctx，供 CallTool 内与 executionId 关联。 +func WithMCPConversationID(ctx context.Context, conversationID string) context.Context { + if ctx == nil { + return nil + } + id := strings.TrimSpace(conversationID) + if id == "" { + return ctx + } + return context.WithValue(ctx, mcpConversationIDCtxKey{}, id) +} + +// MCPConversationIDFromContext 读取对话 ID。 +func MCPConversationIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + v, _ := ctx.Value(mcpConversationIDCtxKey{}).(string) + return v +} + +// WithMCPExecutionID 将当前工具 executionId 注入 ctx，供超长输出落盘文件名对齐。 +func WithMCPExecutionID(ctx context.Context, executionID string) context.Context { + if ctx == nil { + return nil + } + id := strings.TrimSpace(executionID) + if id == "" { + return ctx + } + return context.WithValue(ctx, mcpExecutionIDCtxKey{}, id) +} + +// MCPExecutionIDFromContext 读取当前工具 executionId。 +func MCPExecutionIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + v, _ := ctx.Value(mcpExecutionIDCtxKey{}).(string) + return v +} + +// WithMCPProjectID 将项目 ID 注入 ctx，供 reduction/trunc 落盘路径与项目隔离对齐。 +func WithMCPProjectID(ctx context.Context, projectID string) context.Context { + if ctx == nil { + return nil + } + id := strings.TrimSpace(projectID) + if id == "" { + return ctx + } + return context.WithValue(ctx, mcpProjectIDCtxKey{}, id) +} + +// MCPProjectIDFromContext 读取项目 ID。 +func MCPProjectIDFromContext(ctx context.Context) string { + if ctx == nil { + return "" + } + v, _ := ctx.Value(mcpProjectIDCtxKey{}).(string) + return v +} + +func notifyToolRunBegin(ctx context.Context, executionID string) { + reg := ToolRunRegistryFromContext(ctx) + if reg == nil { + return + } + conv := MCPConversationIDFromContext(ctx) + if conv == "" || strings.TrimSpace(executionID) == "" { + return + } + reg.RegisterRunningTool(conv, executionID) +} + +func notifyToolRunEnd(ctx context.Context, executionID string) { + reg := ToolRunRegistryFromContext(ctx) + if reg == nil { + return + } + conv := MCPConversationIDFromContext(ctx) + if conv == "" || strings.TrimSpace(executionID) == "" { + return + } + reg.UnregisterRunningTool(conv, executionID) +} diff --git a/internal/mcp/server.go b/internal/mcp/server.go new file mode 100644 index 00000000..3760183b --- /dev/null +++ b/internal/mcp/server.go @@ -0,0 +1,1704 @@ +package mcp + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "sort" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/mcp/builtin" + + "github.com/google/uuid" + "go.uber.org/zap" +) + +// MonitorStorage 监控数据存储接口 +type MonitorStorage interface { + SaveToolExecution(exec *ToolExecution) error + UpdateToolExecutionResult(id string, result *ToolResult) error + LoadToolExecutions() ([]*ToolExecution, error) + GetToolExecution(id string) (*ToolExecution, error) + SaveToolStats(toolName string, stats *ToolStats) error + LoadToolStats() (map[string]*ToolStats, error) + UpdateToolStats(toolName string, totalCalls, successCalls, failedCalls int, lastCallTime *time.Time) error +} + +// Server MCP服务器 +type Server struct { + tools map[string]ToolHandler + toolDefs map[string]Tool // 工具定义 + executions map[string]*ToolExecution + stats map[string]*ToolStats + prompts map[string]*Prompt // 提示词模板 + resources map[string]*Resource // 资源 + storage MonitorStorage // 可选的持久化存储 + mu sync.RWMutex + logger *zap.Logger + maxExecutionsInMemory int // 内存中最大执行记录数 + sseClients map[string]*sseClient + runningCancels map[string]context.CancelFunc + runningCancelsMu sync.Mutex + abortUserNotes map[string]string // 监控页终止时附带的用户说明，与 executionID 对应 + // httpToolTimeoutMinutes 同步 agent.tool_timeout_minutes，用于 POST /api/mcp 的 tools/call（不经 Agent 包装的路径）。 + // nil 表示未配置，沿用默认 30 分钟；指向 0 表示不限制；>0 为分钟数。 + httpToolTimeoutMinutes *int + httpToolTimeoutMu sync.RWMutex + toolAuthorizer func(context.Context, string, map[string]interface{}) error + executionService *ExecutionService + toolWaitTimeout time.Duration + toolResultMaxBytes int + spillRootDir string +} + +const defaultPartialOutputMaxBytes = 64 * 1024 + +// SetToolAuthorizer installs the common policy decision point for every +// user-attributed tool call, whether it originates from HTTP or an Agent. +func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) { + if s == nil { + return + } + s.mu.Lock() + s.toolAuthorizer = authorizer + s.mu.Unlock() +} + +type sseClient struct { + id string + send chan []byte +} + +// ToolHandler 工具处理函数 +type ToolHandler func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) + +func executionStatusAndMessage(err error) (status string, errMsg string) { + if errors.Is(err, context.Canceled) { + return "cancelled", "已手动终止（MCP 监控）" + } + return "failed", err.Error() +} + +// NewServer 创建新的MCP服务器 +func NewServer(logger *zap.Logger) *Server { + return NewServerWithStorage(logger, nil) +} + +// NewServerWithStorage 创建新的MCP服务器（带持久化存储） +func NewServerWithStorage(logger *zap.Logger, storage MonitorStorage) *Server { + s := &Server{ + tools: make(map[string]ToolHandler), + toolDefs: make(map[string]Tool), + executions: make(map[string]*ToolExecution), + stats: make(map[string]*ToolStats), + prompts: make(map[string]*Prompt), + resources: make(map[string]*Resource), + storage: storage, + logger: logger, + maxExecutionsInMemory: 1000, // 默认最多在内存中保留1000条执行记录 + sseClients: make(map[string]*sseClient), + runningCancels: make(map[string]context.CancelFunc), + abortUserNotes: make(map[string]string), + toolWaitTimeout: 60 * time.Second, + toolResultMaxBytes: DefaultToolResultMaxBytes, + } + s.executionService = NewExecutionService(storage, logger) + + // 初始化默认提示词和资源 + s.initDefaultPrompts() + s.initDefaultResources() + + return s +} + +func (s *Server) ConfigureToolResultMaxBytes(maxBytes int) { + if s == nil { + return + } + s.mu.Lock() + s.toolResultMaxBytes = maxBytes + s.mu.Unlock() + if s.executionService != nil { + s.executionService.ConfigureToolResultMaxBytes(maxBytes) + } +} + +// ConfigureToolResultSpillRoot sets the local directory root used when oversized +// tool results are spilled (aligned with reduction_root_dir; empty → tmp/reduction). +func (s *Server) ConfigureToolResultSpillRoot(rootDir string) { + if s == nil { + return + } + s.mu.Lock() + s.spillRootDir = strings.TrimSpace(rootDir) + s.mu.Unlock() + if s.executionService != nil { + s.executionService.ConfigureToolResultSpillRoot(rootDir) + } +} + +// ConfigureHTTPToolCallTimeoutFromAgentMinutes 将 agent.tool_timeout_minutes 同步到经 HTTP POST /api/mcp 触发的 tools/call。 +// minutes<=0 表示不设置硬性截止时间（与配置「0 不限制」一致）；minutes>0 为该次调用的最长等待时间。 +// 未调用前对 tools/call 使用默认 30 分钟（与历史硬编码一致）。 +func (s *Server) ConfigureHTTPToolCallTimeoutFromAgentMinutes(minutes int) { + if s == nil { + return + } + v := minutes + if v < 0 { + v = 0 + } + s.httpToolTimeoutMu.Lock() + defer s.httpToolTimeoutMu.Unlock() + s.httpToolTimeoutMinutes = &v +} + +func (s *Server) ConfigureToolWaitTimeoutSeconds(seconds int) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if seconds <= 0 { + s.toolWaitTimeout = 0 + return + } + s.toolWaitTimeout = time.Duration(seconds) * time.Second +} + +func (s *Server) effectiveHTTPToolCallDeadline(parent context.Context) (context.Context, context.CancelFunc) { + const defaultDur = 30 * time.Minute + if parent == nil { + parent = context.Background() + } + if s == nil { + return context.WithTimeout(parent, defaultDur) + } + s.httpToolTimeoutMu.RLock() + mPtr := s.httpToolTimeoutMinutes + s.httpToolTimeoutMu.RUnlock() + if mPtr == nil { + return context.WithTimeout(parent, defaultDur) + } + if *mPtr <= 0 { + return context.WithCancel(parent) + } + return context.WithTimeout(parent, time.Duration(*mPtr)*time.Minute) +} + +// RegisterTool 注册工具 +func (s *Server) RegisterTool(tool Tool, handler ToolHandler) { + s.mu.Lock() + defer s.mu.Unlock() + s.tools[tool.Name] = handler + s.toolDefs[tool.Name] = tool + + // 自动为工具创建资源文档 + resourceURI := fmt.Sprintf("tool://%s", tool.Name) + s.resources[resourceURI] = &Resource{ + URI: resourceURI, + Name: fmt.Sprintf("%s工具文档", tool.Name), + Description: tool.Description, + MimeType: "text/plain", + } +} + +// ClearTools 清空所有工具（用于重新加载配置） +func (s *Server) ClearTools() { + s.mu.Lock() + defer s.mu.Unlock() + + // 清空工具和工具定义 + s.tools = make(map[string]ToolHandler) + s.toolDefs = make(map[string]Tool) + + // 清空工具相关的资源（保留其他资源） + newResources := make(map[string]*Resource) + for uri, resource := range s.resources { + // 保留非工具资源 + if !strings.HasPrefix(uri, "tool://") { + newResources[uri] = resource + } + } + s.resources = newResources +} + +// HandleHTTP 处理HTTP请求 +func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodGet && strings.Contains(r.Header.Get("Accept"), "text/event-stream") { + s.handleSSE(w, r) + return + } + + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // 官方 MCP SSE 规范：带 sessionid 的 POST 表示消息发往该 SSE 会话，响应通过 SSE 流返回 + if sessionID := r.URL.Query().Get("sessionid"); sessionID != "" { + s.serveSSESessionMessage(w, r, sessionID) + return + } + + // 简单 POST：请求体为 JSON-RPC，响应在 body 中返回 + body, err := io.ReadAll(r.Body) + if err != nil { + s.sendError(w, nil, -32700, "Parse error", err.Error()) + return + } + + var msg Message + if err := json.Unmarshal(body, &msg); err != nil { + s.sendError(w, nil, -32700, "Parse error", err.Error()) + return + } + + response := s.handleMessage(r.Context(), &msg) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(response) +} + +// serveSSESessionMessage 处理发往 SSE 会话的 POST：读取 JSON-RPC 请求，处理后将响应通过该会话的 SSE 流推送 +func (s *Server) serveSSESessionMessage(w http.ResponseWriter, r *http.Request, sessionID string) { + s.mu.RLock() + client, exists := s.sseClients[sessionID] + s.mu.RUnlock() + if !exists || client == nil { + http.Error(w, "session not found", http.StatusNotFound) + return + } + + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, "failed to read body", http.StatusBadRequest) + return + } + + var msg Message + if err := json.Unmarshal(body, &msg); err != nil { + http.Error(w, "failed to parse body", http.StatusBadRequest) + return + } + + response := s.handleMessage(r.Context(), &msg) + if response == nil { + w.WriteHeader(http.StatusAccepted) + return + } + + respBytes, err := json.Marshal(response) + if err != nil { + http.Error(w, "failed to encode response", http.StatusInternalServerError) + return + } + + select { + case client.send <- respBytes: + w.WriteHeader(http.StatusAccepted) + default: + http.Error(w, "session send buffer full", http.StatusServiceUnavailable) + } +} + +// handleSSE 处理 SSE 连接，兼容官方 MCP 2024-11-05 SSE 规范： +// 1. 首个事件必须为 event: endpoint，data 为客户端 POST 消息的 URL（含 sessionid） +// 2. 后续事件为 event: message，data 为 JSON-RPC 响应 +func (s *Server) handleSSE(w http.ResponseWriter, r *http.Request) { + flusher, ok := w.(http.Flusher) + if !ok { + http.Error(w, "Streaming unsupported", http.StatusInternalServerError) + return + } + + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + w.Header().Set("X-Accel-Buffering", "no") + + sessionID := uuid.New().String() + client := &sseClient{ + id: sessionID, + send: make(chan []byte, 32), + } + + s.addSSEClient(client) + defer s.removeSSEClient(client.id) + + // 官方规范：首个事件为 endpoint，data 为消息端点 URL（客户端将向该 URL POST 请求） + scheme := "http" + if r.TLS != nil { + scheme = "https" + } + if r.URL.Scheme != "" { + scheme = r.URL.Scheme + } + endpointURL := fmt.Sprintf("%s://%s%s?sessionid=%s", scheme, r.Host, r.URL.Path, sessionID) + fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", endpointURL) + flusher.Flush() + + ticker := time.NewTicker(15 * time.Second) + defer ticker.Stop() + + for { + select { + case <-r.Context().Done(): + return + case msg, ok := <-client.send: + if !ok { + return + } + fmt.Fprintf(w, "event: message\ndata: %s\n\n", msg) + flusher.Flush() + case <-ticker.C: + fmt.Fprintf(w, ": ping\n\n") + flusher.Flush() + } + } +} + +// addSSEClient 注册SSE客户端 +func (s *Server) addSSEClient(client *sseClient) { + s.mu.Lock() + defer s.mu.Unlock() + s.sseClients[client.id] = client +} + +// removeSSEClient 移除SSE客户端 +func (s *Server) removeSSEClient(id string) { + s.mu.Lock() + defer s.mu.Unlock() + if client, exists := s.sseClients[id]; exists { + close(client.send) + delete(s.sseClients, id) + } +} + +// handleMessage 处理MCP消息 +func (s *Server) handleMessage(ctx context.Context, msg *Message) *Message { + // 检查是否是通知（notification）- 通知没有id字段，不需要响应 + isNotification := msg.ID.Value() == nil || msg.ID.String() == "" + + // 如果不是通知且ID为空，生成新的UUID + if !isNotification && msg.ID.String() == "" { + msg.ID = MessageID{value: uuid.New().String()} + } + + switch msg.Method { + case "initialize": + return s.handleInitialize(msg) + case "tools/list": + return s.handleListTools(msg) + case "tools/call": + return s.handleCallTool(ctx, msg) + case "prompts/list": + return s.handleListPrompts(msg) + case "prompts/get": + return s.handleGetPrompt(msg) + case "resources/list": + return s.handleListResources(msg) + case "resources/read": + return s.handleReadResource(msg) + case "sampling/request": + return s.handleSamplingRequest(msg) + case "notifications/initialized": + // 通知类型，不需要响应 + s.logger.Debug("收到 initialized 通知") + return nil + case "": + // 空方法名，可能是通知，不返回错误 + if isNotification { + s.logger.Debug("收到无方法名的通知消息") + return nil + } + fallthrough + default: + // 如果是通知，不返回错误响应 + if isNotification { + s.logger.Debug("收到未知通知", zap.String("method", msg.Method)) + return nil + } + // 对于请求，返回方法未找到错误 + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32601, Message: "Method not found"}, + } + } +} + +// handleInitialize 处理初始化请求 +func (s *Server) handleInitialize(msg *Message) *Message { + var req InitializeRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + + response := InitializeResponse{ + ProtocolVersion: ProtocolVersion, + Capabilities: ServerCapabilities{ + Tools: map[string]interface{}{ + "listChanged": true, + }, + Prompts: map[string]interface{}{ + "listChanged": true, + }, + Resources: map[string]interface{}{ + "subscribe": true, + "listChanged": true, + }, + Sampling: map[string]interface{}{}, + }, + ServerInfo: ServerInfo{ + Name: "CyberStrikeAI", + Version: "1.0.0", + }, + } + + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// handleListTools 处理列出工具请求 +func (s *Server) handleListTools(msg *Message) *Message { + s.mu.RLock() + tools := make([]Tool, 0, len(s.toolDefs)) + for _, tool := range s.toolDefs { + tools = append(tools, tool) + } + s.mu.RUnlock() + s.logger.Debug("tools/list 请求", zap.Int("返回工具数", len(tools))) + + response := ListToolsResponse{Tools: tools} + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// handleCallTool 处理工具调用请求 +func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Message { + var req CallToolRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + _, authenticated := authctx.PrincipalFromContext(requestCtx) + s.mu.RLock() + authorizer := s.toolAuthorizer + s.mu.RUnlock() + if authorizer != nil { + if err := authorizer(requestCtx, req.Name, req.Arguments); err != nil { + return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Forbidden", Data: err.Error()}} + } + } else if authenticated { + return &Message{ID: msg.ID, Type: MessageTypeError, Version: "2.0", Error: &Error{Code: -32003, Message: "Tool authorization policy is not configured"}} + } + + executionID := uuid.New().String() + execution := &ToolExecution{ + ID: executionID, + ToolName: req.Name, + Arguments: req.Arguments, + Status: "running", + StartTime: time.Now(), + } + if principal, ok := authctx.PrincipalFromContext(requestCtx); ok { + execution.OwnerUserID = principal.UserID + } + execution.ConversationID = MCPConversationIDFromContext(requestCtx) + + s.mu.Lock() + s.executions[executionID] = execution + // 如果内存中的执行记录超过限制，清理最旧的记录 + s.cleanupOldExecutions() + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(execution); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + } + + s.mu.RLock() + handler, exists := s.tools[req.Name] + s.mu.RUnlock() + + if !exists { + execution.Status = "failed" + execution.Error = "Tool not found" + now := time.Now() + execution.EndTime = &now + execution.Duration = now.Sub(execution.StartTime) + + if s.storage != nil { + if err := s.storage.SaveToolExecution(execution); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + s.mu.Lock() + delete(s.executions, executionID) + s.mu.Unlock() + } + + s.updateStats(req.Name, true) + + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32601, Message: "Tool not found"}, + } + } + + baseCtx, timeoutCancel := s.effectiveHTTPToolCallDeadline(requestCtx) + defer timeoutCancel() + execCtx, runCancel := context.WithCancel(baseCtx) + s.registerRunningCancel(executionID, runCancel) + defer func() { + runCancel() + s.unregisterRunningCancel(executionID) + }() + + s.logger.Info("开始执行工具", + zap.String("toolName", req.Name), + zap.Any("arguments", req.Arguments), + ) + + result, err := handler(execCtx, req.Arguments) + cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(executionID, &result, &err) + now := time.Now() + var failed bool + var finalResult *ToolResult + + s.mu.Lock() + execution.EndTime = &now + execution.Duration = now.Sub(execution.StartTime) + + if err != nil { + st, msg := executionStatusAndMessage(err) + execution.Status = st + execution.Error = msg + failed = st != "cancelled" + } else if result != nil && result.IsError { + if cancelledWithUserNote { + execution.Status = "cancelled" + execution.Error = "" + execution.Result = result + failed = false + } else { + execution.Status = "failed" + if len(result.Content) > 0 { + execution.Error = result.Content[0].Text + } else { + execution.Error = "工具执行返回错误结果" + } + execution.Result = result + failed = true + } + } else { + execution.Status = "completed" + if result == nil { + result = &ToolResult{ + Content: []Content{ + {Type: "text", Text: "工具执行完成，但未返回结果"}, + }, + } + } + execution.Result = result + failed = false + } + + finalResult = execution.Result + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(execution); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + } + + s.updateStats(req.Name, failed) + + if s.storage != nil { + s.mu.Lock() + delete(s.executions, executionID) + s.mu.Unlock() + } + + if err != nil { + s.logger.Error("工具执行失败", + zap.String("toolName", req.Name), + zap.Error(err), + ) + + errText := fmt.Sprintf("工具执行失败: %v", err) + if errors.Is(err, context.Canceled) { + errText = "工具执行已手动终止（MCP 监控）。后续编排步骤可继续。" + } + errorResult, _ := json.Marshal(CallToolResponse{ + Content: []Content{ + {Type: "text", Text: errText}, + }, + IsError: true, + }) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: errorResult, + } + } + + if finalResult != nil && finalResult.IsError { + s.logger.Warn("工具执行返回错误结果", + zap.String("toolName", req.Name), + ) + + errorResult, _ := json.Marshal(CallToolResponse{ + Content: finalResult.Content, + IsError: true, + }) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: errorResult, + } + } + + if finalResult == nil { + finalResult = &ToolResult{ + Content: []Content{ + {Type: "text", Text: "工具执行完成，但未返回结果"}, + }, + } + } + + resultJSON, _ := json.Marshal(CallToolResponse{ + Content: finalResult.Content, + IsError: false, + }) + + s.logger.Info("工具执行完成", + zap.String("toolName", req.Name), + zap.Bool("isError", finalResult.IsError), + ) + + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: resultJSON, + } +} + +// updateStats 更新统计信息 +func (s *Server) updateStats(toolName string, failed bool) { + now := time.Now() + if s.storage != nil { + totalCalls := 1 + successCalls := 0 + failedCalls := 0 + if failed { + failedCalls = 1 + } else { + successCalls = 1 + } + if err := s.storage.UpdateToolStats(toolName, totalCalls, successCalls, failedCalls, &now); err != nil { + s.logger.Warn("保存统计信息到数据库失败", zap.Error(err)) + } + return + } + + s.mu.Lock() + defer s.mu.Unlock() + + if s.stats[toolName] == nil { + s.stats[toolName] = &ToolStats{ + ToolName: toolName, + } + } + + stats := s.stats[toolName] + stats.TotalCalls++ + stats.LastCallTime = &now + + if failed { + stats.FailedCalls++ + } else { + stats.SuccessCalls++ + } +} + +// GetExecution 获取执行记录（先从内存查找，再从数据库查找） +func (s *Server) GetExecution(id string) (*ToolExecution, bool) { + if s.executionService != nil { + if snap, err := s.executionService.Get(id); err == nil && snap != nil && snap.Execution != nil { + return snap.Execution, true + } + } + s.mu.RLock() + exec, exists := s.executions[id] + s.mu.RUnlock() + + if exists { + return exec, true + } + + if s.storage != nil { + exec, err := s.storage.GetToolExecution(id) + if err == nil { + return exec, true + } + } + + return nil, false +} + +// loadHistoricalData 从数据库加载历史数据 +func (s *Server) loadHistoricalData() { + if s.storage == nil { + return + } + + // 加载历史执行记录（最近1000条） + executions, err := s.storage.LoadToolExecutions() + if err != nil { + s.logger.Warn("加载历史执行记录失败", zap.Error(err)) + } else { + s.mu.Lock() + for _, exec := range executions { + // 只加载最近 maxExecutionsInMemory 条，避免内存占用过大 + if len(s.executions) < s.maxExecutionsInMemory { + s.executions[exec.ID] = exec + } else { + break + } + } + s.mu.Unlock() + s.logger.Info("加载历史执行记录", zap.Int("count", len(executions))) + } + + // 加载历史统计信息 + stats, err := s.storage.LoadToolStats() + if err != nil { + s.logger.Warn("加载历史统计信息失败", zap.Error(err)) + } else { + s.mu.Lock() + for k, v := range stats { + s.stats[k] = v + } + s.mu.Unlock() + s.logger.Info("加载历史统计信息", zap.Int("count", len(stats))) + } +} + +// GetAllExecutions 获取所有执行记录（合并内存和数据库） +func (s *Server) GetAllExecutions() []*ToolExecution { + if s.storage != nil { + dbExecutions, err := s.storage.LoadToolExecutions() + if err == nil { + execMap := make(map[string]*ToolExecution) + for _, exec := range dbExecutions { + if _, exists := execMap[exec.ID]; !exists { + execMap[exec.ID] = exec + } + } + + s.mu.RLock() + for id, exec := range s.executions { + if _, exists := execMap[id]; !exists { + execMap[id] = exec + } + } + s.mu.RUnlock() + + result := make([]*ToolExecution, 0, len(execMap)) + for _, exec := range execMap { + result = append(result, exec) + } + return result + } else { + s.logger.Warn("从数据库加载执行记录失败", zap.Error(err)) + } + } + + s.mu.RLock() + defer s.mu.RUnlock() + + memExecutions := make([]*ToolExecution, 0, len(s.executions)) + for _, exec := range s.executions { + memExecutions = append(memExecutions, exec) + } + return memExecutions +} + +// GetStats 获取统计信息（合并内存和数据库） +func (s *Server) GetStats() map[string]*ToolStats { + if s.storage != nil { + dbStats, err := s.storage.LoadToolStats() + if err == nil { + return dbStats + } + s.logger.Warn("从数据库加载统计信息失败", zap.Error(err)) + } + + s.mu.RLock() + defer s.mu.RUnlock() + + memStats := make(map[string]*ToolStats) + for k, v := range s.stats { + statCopy := *v + memStats[k] = &statCopy + } + + return memStats +} + +// GetAllTools 获取所有已注册的工具（用于Agent动态获取工具列表） +func (s *Server) GetAllTools() []Tool { + s.mu.RLock() + defer s.mu.RUnlock() + + tools := make([]Tool, 0, len(s.toolDefs)) + for _, tool := range s.toolDefs { + tools = append(tools, tool) + } + return tools +} + +// CallTool 直接调用工具（用于内部调用） +func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) { + if s.executionService == nil { + s.executionService = NewExecutionService(s.storage, s.logger) + s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes) + s.executionService.ConfigureToolResultSpillRoot(s.spillRootDir) + } + var ownerUserID string + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + ownerUserID = principal.UserID + } + handle, err := s.executionService.Submit(ctx, ExecutionRequest{ + ToolName: toolName, + Arguments: args, + ConversationID: MCPConversationIDFromContext(ctx), + OwnerUserID: ownerUserID, + Run: func(runCtx context.Context) (*ToolResult, error) { + _, authenticated := authctx.PrincipalFromContext(runCtx) + s.mu.RLock() + authorizer := s.toolAuthorizer + handler, exists := s.tools[toolName] + s.mu.RUnlock() + if authorizer != nil { + if err := authorizer(runCtx, toolName, args); err != nil { + return nil, fmt.Errorf("tool authorization denied: %w", err) + } + } else if authenticated { + return nil, errors.New("tool authorization policy is not configured") + } + if !exists { + return nil, fmt.Errorf("工具 %s 未找到", toolName) + } + return handler(runCtx, args) + }, + OnDone: func(exec *ToolExecution) { + failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled + s.updateStats(toolName, failed) + }, + }) + if err != nil { + return nil, "", err + } + + s.mu.RLock() + waitTimeout := s.toolWaitTimeout + s.mu.RUnlock() + if isExecutionControlTool(toolName) { + waitTimeout = 0 + } + snapshot, waitErr := s.executionService.Wait(ctx, handle.ID, waitTimeout) + if errors.Is(waitErr, ErrExecutionWaitTimeout) { + return internalMCPWaitTimeoutResult(snapshot, waitTimeout), handle.ID, nil + } + if waitErr != nil { + return nil, handle.ID, waitErr + } + if snapshot == nil || snapshot.Execution == nil { + return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回执行快照"}}, IsError: true}, handle.ID, nil + } + if snapshot.Execution.Result != nil { + return snapshot.Execution.Result, handle.ID, nil + } + if snapshot.Execution.Error != "" { + return nil, handle.ID, errors.New(snapshot.Execution.Error) + } + return &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成，但未返回结果"}}, IsError: false}, handle.ID, nil +} + +func internalMCPWaitTimeoutResult(snapshot *ExecutionSnapshot, waitTimeout time.Duration) *ToolResult { + execID := "" + status := ToolExecutionStatusRunning + toolName := "" + elapsed := time.Duration(0) + if snapshot != nil && snapshot.Execution != nil { + execID = snapshot.Execution.ID + status = snapshot.Execution.Status + toolName = snapshot.Execution.ToolName + elapsed = time.Since(snapshot.Execution.StartTime).Round(time.Second) + } + waitText := "unbounded" + if waitTimeout > 0 { + waitText = waitTimeout.Round(time.Second).String() + } + msg := fmt.Sprintf(`工具已提交到后台执行，但本次等待已到达上限。 + +execution_id: %s +tool: %s +status: %s +wait_timeout: %s +elapsed: %s + +你可以继续推理、改用其他工具，或调用 wait_tool_execution 继续等待该 execution_id；也可以调用 cancel_tool_execution 取消。`, execID, toolName, status, waitText, elapsed) + return &ToolResult{Content: []Content{{Type: "text", Text: msg}}, IsError: true} +} + +func isExecutionControlTool(toolName string) bool { + switch strings.TrimSpace(toolName) { + case builtin.ToolGetToolExecution, builtin.ToolWaitToolExecution, builtin.ToolCancelToolExecution: + return true + default: + return false + } +} + +// BeginToolExecution 创建 running 状态的执行记录，供 Eino 等非 CallTool 路径在工具开始时落库。 +func (s *Server) BeginToolExecution(ctx context.Context, toolName string, args map[string]interface{}) string { + if s == nil { + return "" + } + if args == nil { + args = map[string]interface{}{} + } + executionID := uuid.New().String() + execution := &ToolExecution{ + ID: executionID, + ToolName: toolName, + Arguments: args, + Status: "running", + StartTime: time.Now(), + } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + execution.OwnerUserID = principal.UserID + } + execution.ConversationID = MCPConversationIDFromContext(ctx) + + s.mu.Lock() + s.executions[executionID] = execution + s.cleanupOldExecutions() + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(execution); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + } + return executionID +} + +// FinishToolExecution 完成先前 BeginToolExecution 创建的记录；executionID 为空时等同 RecordCompletedToolInvocation。 +func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName string, args map[string]interface{}, resultText string, invokeErr error) string { + if s == nil { + return "" + } + if args == nil { + args = map[string]interface{}{} + } + id := strings.TrimSpace(executionID) + if id == "" { + id = uuid.New().String() + } + + now := time.Now() + failed := invokeErr != nil + var finalResult *ToolResult + + s.mu.Lock() + maxBytes := s.toolResultMaxBytes + spillRoot := s.spillRootDir + exec, inMem := s.executions[id] + if !inMem || exec == nil { + exec = &ToolExecution{ + ID: id, + ToolName: toolName, + Arguments: args, + StartTime: now, + } + s.executions[id] = exec + } else if toolName != "" { + exec.ToolName = toolName + } + if len(args) > 0 { + exec.Arguments = args + } + if principal, ok := authctx.PrincipalFromContext(ctx); ok { + exec.OwnerUserID = principal.UserID + } + if conversationID := MCPConversationIDFromContext(ctx); conversationID != "" { + exec.ConversationID = conversationID + } + exec.EndTime = &now + if exec.StartTime.IsZero() { + exec.StartTime = now + } + exec.Duration = now.Sub(exec.StartTime) + + spill := ToolResultSpillConfig{ + RootDir: spillRoot, + ProjectID: MCPProjectIDFromContext(ctx), + ConversationID: exec.ConversationID, + ExecutionID: id, + } + if failed { + st, msg := executionStatusAndMessage(invokeErr) + exec.Status = st + exec.Error = msg + if strings.TrimSpace(resultText) != "" { + finalResult = &ToolResult{Content: []Content{{Type: "text", Text: resultText}}} + finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill) + exec.Result = finalResult + } + } else { + exec.Status = "completed" + text := resultText + if strings.TrimSpace(text) == "" { + text = "（无输出）" + } + finalResult = &ToolResult{Content: []Content{{Type: "text", Text: text}}} + finalResult = NormalizeToolResultForStorageWithSpill(finalResult, maxBytes, spill) + exec.Result = finalResult + } + s.mu.Unlock() + + if s.storage != nil { + if err := s.storage.SaveToolExecution(exec); err != nil { + s.logger.Warn("保存执行记录到数据库失败", zap.Error(err)) + } + } + + s.updateStats(exec.ToolName, failed) + + if s.storage != nil { + s.mu.Lock() + delete(s.executions, id) + s.mu.Unlock() + } + return id +} + +// AppendToolExecutionPartialOutput records a bounded tail preview for a running local execution. +// The final Result remains authoritative and is written only when the tool finishes. +func (s *Server) AppendToolExecutionPartialOutput(executionID, chunk string) { + if s == nil || strings.TrimSpace(executionID) == "" || chunk == "" { + return + } + id := strings.TrimSpace(executionID) + if s.executionService != nil && s.executionService.AppendPartialOutput(id, chunk) { + return + } + now := time.Now() + s.mu.Lock() + exec := s.executions[id] + if exec != nil { + appendPartialOutput(exec, chunk, defaultPartialOutputMaxBytes, now) + } + s.mu.Unlock() +} + +// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储（格式与 CallTool 结束后一致）， +// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景；返回 executionId 供助手消息 mcpExecutionIds 关联。 +func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string { + return s.FinishToolExecution(ctx, "", toolName, args, resultText, invokeErr) +} + +// UpdateToolExecutionResult 将监控库中的工具结果更新为送入模型的展示正文（如 reduction 后的 persisted-output）。 +func (s *Server) UpdateToolExecutionResult(executionID string, result *ToolResult) error { + if s == nil { + return nil + } + executionID = strings.TrimSpace(executionID) + if executionID == "" || result == nil { + return nil + } + s.mu.Lock() + spill := ToolResultSpillConfig{ + RootDir: s.spillRootDir, + ExecutionID: executionID, + } + if exec, ok := s.executions[executionID]; ok && exec != nil { + spill.ConversationID = exec.ConversationID + result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill) + exec.Result = result + } else { + result = NormalizeToolResultForStorageWithSpill(result, s.toolResultMaxBytes, spill) + } + s.mu.Unlock() + if s.storage != nil { + return s.storage.UpdateToolExecutionResult(executionID, result) + } + return nil +} + +// cleanupOldExecutions 清理旧的执行记录，防止内存无限增长 +func (s *Server) cleanupOldExecutions() { + if len(s.executions) <= s.maxExecutionsInMemory { + return + } + + // 按开始时间排序，找出最旧的记录 + type execWithTime struct { + id string + startTime time.Time + } + execs := make([]execWithTime, 0, len(s.executions)) + for id, exec := range s.executions { + execs = append(execs, execWithTime{ + id: id, + startTime: exec.StartTime, + }) + } + + // 使用 sort 包进行高效排序（最旧的在前） + sort.Slice(execs, func(i, j int) bool { + return execs[i].startTime.Before(execs[j].startTime) + }) + + // 删除最旧的记录，保留 maxExecutionsInMemory 条 + toDelete := len(s.executions) - s.maxExecutionsInMemory + for i := 0; i < toDelete; i++ { + delete(s.executions, execs[i].id) + } + + s.logger.Debug("清理旧的执行记录", + zap.Int("before", len(execs)), + zap.Int("after", len(s.executions)), + zap.Int("deleted", toDelete), + ) +} + +func (s *Server) registerRunningCancel(id string, cancel context.CancelFunc) { + s.runningCancelsMu.Lock() + s.runningCancels[id] = cancel + s.runningCancelsMu.Unlock() +} + +func (s *Server) unregisterRunningCancel(id string) { + s.runningCancelsMu.Lock() + delete(s.runningCancels, id) + s.runningCancelsMu.Unlock() +} + +// RegisterToolExecutionCancel lets non-ExecutionService tool paths, such as Eino +// filesystem execute, participate in cancel_tool_execution by execution_id. +func (s *Server) RegisterToolExecutionCancel(id string, cancel context.CancelFunc) { + id = strings.TrimSpace(id) + if s == nil || id == "" || cancel == nil { + return + } + s.registerRunningCancel(id, cancel) +} + +func (s *Server) UnregisterToolExecutionCancel(id string) { + id = strings.TrimSpace(id) + if s == nil || id == "" { + return + } + s.unregisterRunningCancel(id) +} + +func (s *Server) readAbortUserNote(id string) string { + s.runningCancelsMu.Lock() + defer s.runningCancelsMu.Unlock() + if s.abortUserNotes == nil { + return "" + } + return s.abortUserNotes[id] +} + +func (s *Server) takeAbortUserNote(id string) string { + s.runningCancelsMu.Lock() + defer s.runningCancelsMu.Unlock() + if s.abortUserNotes == nil { + return "" + } + n := s.abortUserNotes[id] + delete(s.abortUserNotes, id) + return n +} + +// applyAbortUserNoteToCancelledToolResult 监控页「终止并填写说明」时合并「工具已输出 + 用户说明」交给模型。 +// exec 等工具会把失败写在 *ToolResult 里并返回 err==nil，若仅在 err!=nil 时合并会漏掉说明，甚至误 clear 掉 note。 +func (s *Server) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) { + note := strings.TrimSpace(s.readAbortUserNote(executionID)) + if note == "" { + return false + } + hasErr := err != nil && *err != nil + hasRes := result != nil && *result != nil + if !hasErr && !hasRes { + return false + } + _ = s.takeAbortUserNote(executionID) + partial := "" + if hasRes { + partial = ToolResultPlainText(*result) + } + if partial == "" && hasErr { + partial = (*err).Error() + } + merged := MergePartialToolOutputAndAbortNote(partial, note) + *err = nil + *result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true} + return true +} + +// CancelToolExecutionWithNote 取消内部工具；note 非空时与工具已返回文本合并后交给上层模型。 +func (s *Server) CancelToolExecutionWithNote(id string, note string) bool { + if s.executionService != nil && s.executionService.Cancel(id, note) { + return true + } + s.runningCancelsMu.Lock() + cancel, ok := s.runningCancels[id] + if !ok || cancel == nil { + s.runningCancelsMu.Unlock() + return false + } + if strings.TrimSpace(note) != "" { + if s.abortUserNotes == nil { + s.abortUserNotes = make(map[string]string) + } + s.abortUserNotes[id] = strings.TrimSpace(note) + } + s.runningCancelsMu.Unlock() + cancel() + return true +} + +// CancelToolExecution 取消正在执行的内部工具调用（无用户说明）。 +func (s *Server) CancelToolExecution(id string) bool { + return s.CancelToolExecutionWithNote(id, "") +} + +// ActiveRunningExecutionIDs 返回当前进程内仍登记 cancel 的 executionId 快照。 +func (s *Server) ActiveRunningExecutionIDs() map[string]struct{} { + if s == nil { + return nil + } + out := make(map[string]struct{}) + if s.executionService != nil { + for id := range s.executionService.ActiveRunningExecutionIDs() { + out[id] = struct{}{} + } + } + s.runningCancelsMu.Lock() + defer s.runningCancelsMu.Unlock() + if len(s.runningCancels) == 0 && len(out) == 0 { + return nil + } + for id := range s.runningCancels { + out[id] = struct{}{} + } + return out +} + +// initDefaultPrompts 初始化默认提示词模板 +func (s *Server) initDefaultPrompts() { + s.mu.Lock() + defer s.mu.Unlock() + + // 网络安全测试提示词 + s.prompts["security_scan"] = &Prompt{ + Name: "security_scan", + Description: "生成网络安全扫描任务的提示词", + Arguments: []PromptArgument{ + {Name: "target", Description: "扫描目标（IP地址或域名）", Required: true}, + {Name: "scan_type", Description: "扫描类型（port, vuln, web等）", Required: false}, + }, + } + + // 渗透测试提示词 + s.prompts["penetration_test"] = &Prompt{ + Name: "penetration_test", + Description: "生成渗透测试任务的提示词", + Arguments: []PromptArgument{ + {Name: "target", Description: "测试目标", Required: true}, + {Name: "scope", Description: "测试范围", Required: false}, + }, + } +} + +// initDefaultResources 初始化默认资源 +// 注意：工具资源现在在 RegisterTool 时自动创建，此函数保留用于其他非工具资源 +func (s *Server) initDefaultResources() { + // 工具资源已改为在 RegisterTool 时自动创建，无需在此硬编码 +} + +// handleListPrompts 处理列出提示词请求 +func (s *Server) handleListPrompts(msg *Message) *Message { + s.mu.RLock() + prompts := make([]Prompt, 0, len(s.prompts)) + for _, prompt := range s.prompts { + prompts = append(prompts, *prompt) + } + s.mu.RUnlock() + + response := ListPromptsResponse{ + Prompts: prompts, + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// handleGetPrompt 处理获取提示词请求 +func (s *Server) handleGetPrompt(msg *Message) *Message { + var req GetPromptRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + + s.mu.RLock() + prompt, exists := s.prompts[req.Name] + s.mu.RUnlock() + + if !exists { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32601, Message: "Prompt not found"}, + } + } + + // 根据提示词名称生成消息 + messages := s.generatePromptMessages(prompt, req.Arguments) + + response := GetPromptResponse{ + Messages: messages, + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// generatePromptMessages 生成提示词消息 +func (s *Server) generatePromptMessages(prompt *Prompt, args map[string]interface{}) []PromptMessage { + messages := []PromptMessage{} + + switch prompt.Name { + case "security_scan": + target, _ := args["target"].(string) + scanType, _ := args["scan_type"].(string) + if scanType == "" { + scanType = "comprehensive" + } + + content := fmt.Sprintf(`请对目标 %s 执行%s安全扫描。包括： +1. 端口扫描和服务识别 +2. 漏洞检测 +3. Web应用安全测试 +4. 生成详细的安全报告`, target, scanType) + + messages = append(messages, PromptMessage{ + Role: "user", + Content: content, + }) + + case "penetration_test": + target, _ := args["target"].(string) + scope, _ := args["scope"].(string) + + content := fmt.Sprintf(`请对目标 %s 执行渗透测试。`, target) + if scope != "" { + content += fmt.Sprintf("测试范围：%s", scope) + } + content += "\n请按照OWASP Top 10进行全面的安全测试。" + + messages = append(messages, PromptMessage{ + Role: "user", + Content: content, + }) + + default: + messages = append(messages, PromptMessage{ + Role: "user", + Content: "请执行安全测试任务", + }) + } + + return messages +} + +// handleListResources 处理列出资源请求 +func (s *Server) handleListResources(msg *Message) *Message { + s.mu.RLock() + resources := make([]Resource, 0, len(s.resources)) + for _, resource := range s.resources { + resources = append(resources, *resource) + } + s.mu.RUnlock() + + response := ListResourcesResponse{ + Resources: resources, + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// handleReadResource 处理读取资源请求 +func (s *Server) handleReadResource(msg *Message) *Message { + var req ReadResourceRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + + s.mu.RLock() + resource, exists := s.resources[req.URI] + s.mu.RUnlock() + + if !exists { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32601, Message: "Resource not found"}, + } + } + + // 生成资源内容 + content := s.generateResourceContent(resource) + + response := ReadResourceResponse{ + Contents: []ResourceContent{content}, + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// generateResourceContent 生成资源内容 +func (s *Server) generateResourceContent(resource *Resource) ResourceContent { + content := ResourceContent{ + URI: resource.URI, + MimeType: resource.MimeType, + } + + // 如果是工具资源，生成详细文档 + if strings.HasPrefix(resource.URI, "tool://") { + toolName := strings.TrimPrefix(resource.URI, "tool://") + content.Text = s.generateToolDocumentation(toolName, resource) + } else { + // 其他资源使用描述或默认内容 + content.Text = resource.Description + } + + return content +} + +// generateToolDocumentation 生成工具文档 +// 注意：硬编码的工具文档已移除，现在只使用工具定义中的信息 +func (s *Server) generateToolDocumentation(toolName string, resource *Resource) string { + // 获取工具定义以获取更详细的信息 + s.mu.RLock() + tool, hasTool := s.toolDefs[toolName] + s.mu.RUnlock() + + // 使用工具定义中的描述信息 + if hasTool { + doc := fmt.Sprintf("%s\n\n", resource.Description) + if tool.InputSchema != nil { + if props, ok := tool.InputSchema["properties"].(map[string]interface{}); ok { + doc += "参数说明：\n" + for paramName, paramInfo := range props { + if paramMap, ok := paramInfo.(map[string]interface{}); ok { + if desc, ok := paramMap["description"].(string); ok { + doc += fmt.Sprintf("- %s: %s\n", paramName, desc) + } + } + } + } + } + return doc + } + return resource.Description +} + +// handleSamplingRequest 处理采样请求 +func (s *Server) handleSamplingRequest(msg *Message) *Message { + var req SamplingRequest + if err := json.Unmarshal(msg.Params, &req); err != nil { + return &Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32602, Message: "Invalid params"}, + } + } + + // 注意：采样功能通常需要连接到实际的LLM服务 + // 这里返回一个占位符响应，实际实现需要集成LLM API + s.logger.Warn("Sampling request received but not fully implemented", + zap.Any("request", req), + ) + + response := SamplingResponse{ + Content: []SamplingContent{ + { + Type: "text", + Text: "采样功能需要配置LLM服务。请使用Agent Loop API进行AI对话。", + }, + }, + StopReason: "length", + } + result, _ := json.Marshal(response) + return &Message{ + ID: msg.ID, + Type: MessageTypeResponse, + Version: "2.0", + Result: result, + } +} + +// RegisterPrompt 注册提示词模板 +func (s *Server) RegisterPrompt(prompt *Prompt) { + s.mu.Lock() + defer s.mu.Unlock() + s.prompts[prompt.Name] = prompt +} + +// RegisterResource 注册资源 +func (s *Server) RegisterResource(resource *Resource) { + s.mu.Lock() + defer s.mu.Unlock() + s.resources[resource.URI] = resource +} + +// HandleStdio 处理标准输入输出（用于 stdio 传输模式） +// MCP 协议使用换行分隔的 JSON-RPC 消息；管道下需每次写入后 Flush，否则客户端会读不到响应 +func (s *Server) HandleStdio() error { + decoder := json.NewDecoder(os.Stdin) + stdout := bufio.NewWriter(os.Stdout) + encoder := json.NewEncoder(stdout) + // 注意：不设置缩进，MCP 协议期望紧凑的 JSON 格式 + + for { + var msg Message + if err := decoder.Decode(&msg); err != nil { + if err == io.EOF { + break + } + // 日志输出到 stderr，避免干扰 stdout 的 JSON-RPC 通信 + s.logger.Error("读取消息失败", zap.Error(err)) + // 发送错误响应 + errorMsg := Message{ + ID: msg.ID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: -32700, Message: "Parse error", Data: err.Error()}, + } + if err := encoder.Encode(errorMsg); err != nil { + return fmt.Errorf("发送错误响应失败: %w", err) + } + if err := stdout.Flush(); err != nil { + return fmt.Errorf("刷新 stdout 失败: %w", err) + } + continue + } + + // 处理消息 + response := s.handleMessage(context.Background(), &msg) + + // 如果是通知（response 为 nil），不需要发送响应 + if response == nil { + continue + } + + // 发送响应 + if err := encoder.Encode(response); err != nil { + return fmt.Errorf("发送响应失败: %w", err) + } + if err := stdout.Flush(); err != nil { + return fmt.Errorf("刷新 stdout 失败: %w", err) + } + } + + return nil +} + +// sendError 发送错误响应 +func (s *Server) sendError(w http.ResponseWriter, id interface{}, code int, message, data string) { + var msgID MessageID + if id != nil { + msgID = MessageID{value: id} + } + response := Message{ + ID: msgID, + Type: MessageTypeError, + Version: "2.0", + Error: &Error{Code: code, Message: message, Data: data}, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(response) +} diff --git a/internal/mcp/server_authorization_test.go b/internal/mcp/server_authorization_test.go new file mode 100644 index 00000000..bea489a3 --- /dev/null +++ b/internal/mcp/server_authorization_test.go @@ -0,0 +1,231 @@ +package mcp + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/authctx" + + "go.uber.org/zap" +) + +func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) { + server := NewServer(zap.NewNop()) + server.RegisterTool(Tool{Name: "echo", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil + }) + server.SetToolAuthorizer(func(ctx context.Context, toolName string, args map[string]interface{}) error { + if _, ok := authctx.PrincipalFromContext(ctx); !ok { + return errors.New("principal required") + } + return nil + }) + _, deniedExecutionID, err := server.CallTool(context.Background(), "echo", nil) + if err == nil { + t.Fatal("tool call without principal was allowed") + } + if deniedExecutionID == "" { + t.Fatal("denied tool call should still return an execution id") + } + deniedExecution, ok := server.GetExecution(deniedExecutionID) + if !ok || deniedExecution == nil { + t.Fatalf("missing denied execution %q", deniedExecutionID) + } + if deniedExecution.Status != ToolExecutionStatusFailed || !strings.Contains(deniedExecution.Error, "principal required") { + t.Fatalf("denied execution = %#v, want failed with authorization error", deniedExecution) + } + ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true})) + _, executionID, err := server.CallTool(ctx, "echo", nil) + if err != nil { + t.Fatal(err) + } + execution, ok := server.GetExecution(executionID) + if !ok || execution.OwnerUserID != "u1" { + t.Fatalf("execution owner = %#v, want u1", execution) + } +} + +func TestServerCallToolBoundedWaitForInternalTool(t *testing.T) { + server := NewServer(zap.NewNop()) + server.toolWaitTimeout = 10 * time.Millisecond + release := make(chan struct{}) + started := make(chan struct{}) + server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + close(started) + select { + case <-release: + return &ToolResult{Content: []Content{{Type: "text", Text: "internal done"}}}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }) + + callCtx, callCancel := context.WithCancel(context.Background()) + result, executionID, err := server.CallTool(callCtx, "slow", nil) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if executionID == "" || result == nil || !result.IsError { + t.Fatalf("expected soft timeout with execution id, result=%#v id=%q", result, executionID) + } + if text := ToolResultPlainText(result); !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") { + t.Fatalf("timeout result missing execution guidance: %q", text) + } + select { + case <-started: + default: + t.Fatal("internal worker did not start") + } + callCancel() + close(release) + + snapshot, err := server.executionService.Wait(context.Background(), executionID, time.Second) + if err != nil { + t.Fatalf("wait internal execution: %v", err) + } + if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusCompleted { + t.Fatalf("snapshot = %#v, want completed", snapshot) + } + if got := ToolResultPlainText(snapshot.Execution.Result); got != "internal done" { + t.Fatalf("result = %q, want internal done", got) + } +} + +func TestWaitToolExecutionWaitsForInternalActiveExecution(t *testing.T) { + server := NewServer(zap.NewNop()) + server.toolWaitTimeout = 10 * time.Millisecond + release := make(chan struct{}) + server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + select { + case <-release: + return &ToolResult{Content: []Content{{Type: "text", Text: "wait saw completion"}}}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }) + RegisterExecutionControlTools(server, nil) + + result, executionID, err := server.CallTool(context.Background(), "slow", nil) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if result == nil || !result.IsError || executionID == "" { + t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID) + } + + done := make(chan *ToolResult, 1) + errCh := make(chan error, 1) + go func() { + waitResult, _, waitErr := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "timeout_seconds": 1, + }) + if waitErr != nil { + errCh <- waitErr + return + } + done <- waitResult + }() + + select { + case <-done: + t.Fatal("wait_tool_execution returned before target execution completed") + case err := <-errCh: + t.Fatalf("wait_tool_execution errored before release: %v", err) + case <-time.After(50 * time.Millisecond): + } + close(release) + + select { + case err := <-errCh: + t.Fatalf("wait_tool_execution returned error: %v", err) + case waitResult := <-done: + if waitResult == nil || waitResult.IsError { + t.Fatalf("expected successful wait result, got %#v", waitResult) + } + if body := ToolResultPlainText(waitResult); !strings.Contains(body, "wait saw completion") || !strings.Contains(body, `"status": "completed"`) { + t.Fatalf("wait result missing completed target: %s", body) + } + case <-time.After(time.Second): + t.Fatal("wait_tool_execution did not return after target completion") + } +} + +func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) { + server := NewServer(zap.NewNop()) + server.toolWaitTimeout = 10 * time.Millisecond + release := make(chan struct{}) + server.RegisterTool(Tool{Name: "slow_observed", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + <-release + return &ToolResult{Content: []Content{{Type: "text", Text: "done"}}}, nil + }) + RegisterExecutionControlTools(server, nil) + + result, executionID, err := server.CallTool(context.Background(), "slow_observed", nil) + if err != nil { + t.Fatalf("CallTool returned error: %v", err) + } + if result == nil || !result.IsError || executionID == "" { + t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID) + } + + waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "timeout_seconds": 0.01, + }) + if err != nil { + t.Fatalf("wait_tool_execution returned error: %v", err) + } + if waitResult == nil { + t.Fatal("missing wait result") + } + if waitResult.IsError { + t.Fatalf("wait timeout should be a successful observation, got %#v", waitResult) + } + body := ToolResultPlainText(waitResult) + if !strings.Contains(body, `"status": "running"`) || !strings.Contains(body, "本次等待已到达") { + t.Fatalf("wait timeout body missing running status/guidance: %s", body) + } + close(release) +} + +func TestGetToolExecutionIncludesBoundedPartialOutput(t *testing.T) { + server := NewServer(zap.NewNop()) + RegisterExecutionControlTools(server, nil) + + executionID := server.BeginToolExecution(context.Background(), "execute", map[string]interface{}{"command": "demo"}) + if executionID == "" { + t.Fatal("missing execution id") + } + server.AppendToolExecutionPartialOutput(executionID, "first\n") + server.AppendToolExecutionPartialOutput(executionID, strings.Repeat("x", 32)) + + result, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "partial_output_max_bytes": 8, + }) + if err != nil { + t.Fatalf("get_tool_execution: %v", err) + } + body := ToolResultPlainText(result) + if !strings.Contains(body, `"partial_output": "xxxxxxxx"`) { + t.Fatalf("missing bounded partial output: %s", body) + } + if !strings.Contains(body, `"partial_output_bytes": 38`) { + t.Fatalf("missing partial byte count: %s", body) + } + + result, _, err = server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "include_partial_output": false, + }) + if err != nil { + t.Fatalf("get_tool_execution without partial: %v", err) + } + if body := ToolResultPlainText(result); strings.Contains(body, "partial_output") { + t.Fatalf("partial output should be omitted: %s", body) + } +} diff --git a/internal/mcp/tool_result_guard.go b/internal/mcp/tool_result_guard.go new file mode 100644 index 00000000..851db4df --- /dev/null +++ b/internal/mcp/tool_result_guard.go @@ -0,0 +1,65 @@ +package mcp + +import "cyberstrike-ai/internal/tooloutput" + +const DefaultToolResultMaxBytes = 12000 + +// ToolResultSpillConfig controls where oversized tool results are written on disk +// before the in-memory/DB/agent-facing payload is truncated. +type ToolResultSpillConfig struct { + RootDir string + ProjectID string + ConversationID string + ExecutionID string +} + +// NormalizeToolResultForStorage returns the canonical result used by both the +// agent-facing response and monitor persistence. When maxBytes is exceeded the +// full text is spilled under the reduction cache tree and replaced with a +// notice that includes the file path. +func NormalizeToolResultForStorage(result *ToolResult, maxBytes int) *ToolResult { + return NormalizeToolResultForStorageWithSpill(result, maxBytes, ToolResultSpillConfig{}) +} + +// NormalizeToolResultForStorageWithSpill is NormalizeToolResultForStorage with +// an explicit spill location (conversation/execution scoped). +func NormalizeToolResultForStorageWithSpill(result *ToolResult, maxBytes int, spill ToolResultSpillConfig) *ToolResult { + if result == nil { + return nil + } + out := cloneToolResult(result) + if maxBytes <= 0 { + return out + } + + total := 0 + for _, c := range out.Content { + if c.Type == "text" { + total += len(c.Text) + } + } + if total <= maxBytes { + return out + } + + full := ToolResultPlainText(out) + bound := tooloutput.BoundWithSpill(full, maxBytes, tooloutput.SpillOpts{ + RootDir: spill.RootDir, + ProjectID: spill.ProjectID, + ConversationID: spill.ConversationID, + ExecutionID: spill.ExecutionID, + }) + out.Content = []Content{{Type: "text", Text: bound}} + return out +} + +func cloneToolResult(in *ToolResult) *ToolResult { + if in == nil { + return nil + } + out := *in + if in.Content != nil { + out.Content = append([]Content(nil), in.Content...) + } + return &out +} diff --git a/internal/mcp/tool_result_guard_test.go b/internal/mcp/tool_result_guard_test.go new file mode 100644 index 00000000..fd7c8d55 --- /dev/null +++ b/internal/mcp/tool_result_guard_test.go @@ -0,0 +1,158 @@ +package mcp + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "go.uber.org/zap" +) + +type inMemoryMonitorStorage struct { + executions map[string]*ToolExecution +} + +func newInMemoryMonitorStorage() *inMemoryMonitorStorage { + return &inMemoryMonitorStorage{executions: map[string]*ToolExecution{}} +} + +func (s *inMemoryMonitorStorage) SaveToolExecution(exec *ToolExecution) error { + if exec != nil { + s.executions[exec.ID] = cloneToolExecution(exec) + } + return nil +} + +func (s *inMemoryMonitorStorage) UpdateToolExecutionResult(id string, result *ToolResult) error { + exec := s.executions[id] + if exec == nil { + exec = &ToolExecution{ID: id} + s.executions[id] = exec + } + exec.Result = cloneToolResult(result) + return nil +} + +func (s *inMemoryMonitorStorage) LoadToolExecutions() ([]*ToolExecution, error) { + out := make([]*ToolExecution, 0, len(s.executions)) + for _, exec := range s.executions { + out = append(out, cloneToolExecution(exec)) + } + return out, nil +} + +func (s *inMemoryMonitorStorage) GetToolExecution(id string) (*ToolExecution, error) { + if exec := s.executions[id]; exec != nil { + return cloneToolExecution(exec), nil + } + return nil, nil +} + +func (s *inMemoryMonitorStorage) SaveToolStats(string, *ToolStats) error { return nil } + +func (s *inMemoryMonitorStorage) LoadToolStats() (map[string]*ToolStats, error) { + return map[string]*ToolStats{}, nil +} + +func (s *inMemoryMonitorStorage) UpdateToolStats(string, int, int, int, *time.Time) error { + return nil +} + +func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) { + storage := newInMemoryMonitorStorage() + server := NewServerWithStorage(zap.NewNop(), storage) + server.ConfigureToolWaitTimeoutSeconds(0) + server.ConfigureToolResultMaxBytes(400) + spillRoot := t.TempDir() + server.ConfigureToolResultSpillRoot(spillRoot) + server.RegisterTool(Tool{Name: "big", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) { + return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("x", 800)}}}, nil + }) + + ctx := WithMCPConversationID(context.Background(), "conv-spill") + result, executionID, err := server.CallTool(ctx, "big", nil) + if err != nil { + t.Fatalf("CallTool: %v", err) + } + if executionID == "" { + t.Fatal("missing execution id") + } + returned := ToolResultPlainText(result) + if !strings.Contains(returned, "") || !strings.Contains(returned, "Full output saved to:") { + t.Fatalf("returned result was not spilled: %q", returned) + } + if len(returned) > 400 { + t.Fatalf("returned result exceeded hard limit: len=%d text=%q", len(returned), returned) + } + + spillPath := filepath.Join(spillRoot, "conversations", "conv-spill", "trunc", executionID) + abs, err := filepath.Abs(spillPath) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(returned, abs) { + t.Fatalf("missing spill path %q in %q", abs, returned) + } + body, err := os.ReadFile(abs) + if err != nil { + t.Fatalf("read spill file: %v", err) + } + if string(body) != strings.Repeat("x", 800) { + t.Fatalf("spill body mismatch: len=%d", len(body)) + } + + inMem, ok := server.GetExecution(executionID) + if !ok || inMem == nil || inMem.Result == nil { + t.Fatalf("missing in-memory execution: %#v", inMem) + } + stored := storage.executions[executionID] + if stored == nil || stored.Result == nil { + t.Fatalf("missing stored execution: %#v", stored) + } + if ToolResultPlainText(inMem.Result) != returned { + t.Fatalf("in-memory result != returned\nmem=%q\nret=%q", ToolResultPlainText(inMem.Result), returned) + } + if ToolResultPlainText(stored.Result) != returned { + t.Fatalf("stored result != returned\nstored=%q\nret=%q", ToolResultPlainText(stored.Result), returned) + } +} + +func TestExecutionServiceStoresGuardedResult(t *testing.T) { + service := NewExecutionService(nil, zap.NewNop()) + service.ConfigureToolResultMaxBytes(400) + spillRoot := t.TempDir() + service.ConfigureToolResultSpillRoot(spillRoot) + handle, err := service.Submit(context.Background(), ExecutionRequest{ + ToolName: "big", + ConversationID: "svc-conv", + Run: func(context.Context) (*ToolResult, error) { + return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("a", 800)}}}, nil + }, + }) + if err != nil { + t.Fatalf("Submit: %v", err) + } + snap, err := service.Wait(context.Background(), handle.ID, time.Second) + if err != nil { + t.Fatalf("Wait: %v", err) + } + got := ToolResultPlainText(snap.Execution.Result) + if !strings.Contains(got, "") { + t.Fatalf("service result was not spilled: %q", got) + } + if len(got) > 400 { + t.Fatalf("service result exceeded hard limit: len=%d text=%q", len(got), got) + } + path := filepath.Join(spillRoot, "conversations", "svc-conv", "trunc", handle.ID) + abs, _ := filepath.Abs(path) + body, err := os.ReadFile(abs) + if err != nil { + t.Fatalf("read spill: %v", err) + } + if string(body) != strings.Repeat("a", 800) { + t.Fatalf("unexpected spill body len=%d", len(body)) + } +} diff --git a/internal/mcp/types.go b/internal/mcp/types.go new file mode 100644 index 00000000..6922e047 --- /dev/null +++ b/internal/mcp/types.go @@ -0,0 +1,338 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" +) + +// ExternalMCPClient 外部 MCP 客户端接口（由 client_sdk.go 基于官方 SDK 实现） +type ExternalMCPClient interface { + Initialize(ctx context.Context) error + ListTools(ctx context.Context) ([]Tool, error) + CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) + Close() error + IsConnected() bool + GetStatus() string +} + +// MCP消息类型 +const ( + MessageTypeRequest = "request" + MessageTypeResponse = "response" + MessageTypeError = "error" + MessageTypeNotify = "notify" +) + +// MCP协议版本 +const ProtocolVersion = "2024-11-05" + +// MessageID 表示JSON-RPC 2.0的id字段，可以是字符串、数字或null +type MessageID struct { + value interface{} +} + +// UnmarshalJSON 自定义反序列化，支持字符串、数字和null +func (m *MessageID) UnmarshalJSON(data []byte) error { + // 尝试解析为null + if string(data) == "null" { + m.value = nil + return nil + } + + // 尝试解析为字符串 + var str string + if err := json.Unmarshal(data, &str); err == nil { + m.value = str + return nil + } + + // 尝试解析为数字 + var num json.Number + if err := json.Unmarshal(data, &num); err == nil { + m.value = num + return nil + } + + return fmt.Errorf("invalid id type") +} + +// MarshalJSON 自定义序列化 +func (m MessageID) MarshalJSON() ([]byte, error) { + if m.value == nil { + return []byte("null"), nil + } + return json.Marshal(m.value) +} + +// String 返回字符串表示 +func (m MessageID) String() string { + if m.value == nil { + return "" + } + return fmt.Sprintf("%v", m.value) +} + +// Value 返回原始值 +func (m MessageID) Value() interface{} { + return m.value +} + +// Message 表示MCP消息（符合JSON-RPC 2.0规范） +type Message struct { + ID MessageID `json:"id,omitempty"` + Type string `json:"-"` // 内部使用，不序列化到JSON + Method string `json:"method,omitempty"` + Params json.RawMessage `json:"params,omitempty"` + Result json.RawMessage `json:"result,omitempty"` + Error *Error `json:"error,omitempty"` + Version string `json:"jsonrpc,omitempty"` // JSON-RPC 2.0 版本标识 +} + +// Error 表示MCP错误 +type Error struct { + Code int `json:"code"` + Message string `json:"message"` + Data interface{} `json:"data,omitempty"` +} + +// Tool 表示MCP工具定义 +type Tool struct { + Name string `json:"name"` + Description string `json:"description"` // 详细描述 + ShortDescription string `json:"shortDescription,omitempty"` // 简短描述（用于工具列表，减少token消耗） + InputSchema map[string]interface{} `json:"inputSchema"` +} + +// ToolCall 表示工具调用 +type ToolCall struct { + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments"` +} + +// ToolResult 表示工具执行结果 +type ToolResult struct { + Content []Content `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +// Content 表示内容 +type Content struct { + Type string `json:"type"` + Text string `json:"text"` +} + +// InitializeRequest 初始化请求 +type InitializeRequest struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities map[string]interface{} `json:"capabilities"` + ClientInfo ClientInfo `json:"clientInfo"` +} + +// ClientInfo 客户端信息 +type ClientInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// InitializeResponse 初始化响应 +type InitializeResponse struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities ServerCapabilities `json:"capabilities"` + ServerInfo ServerInfo `json:"serverInfo"` +} + +// ServerCapabilities 服务器能力 +type ServerCapabilities struct { + Tools map[string]interface{} `json:"tools,omitempty"` + Prompts map[string]interface{} `json:"prompts,omitempty"` + Resources map[string]interface{} `json:"resources,omitempty"` + Sampling map[string]interface{} `json:"sampling,omitempty"` +} + +// ServerInfo 服务器信息 +type ServerInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// ListToolsRequest 列出工具请求 +type ListToolsRequest struct{} + +// ListToolsResponse 列出工具响应 +type ListToolsResponse struct { + Tools []Tool `json:"tools"` +} + +// ListPromptsResponse 列出提示词响应 +type ListPromptsResponse struct { + Prompts []Prompt `json:"prompts"` +} + +// ListResourcesResponse 列出资源响应 +type ListResourcesResponse struct { + Resources []Resource `json:"resources"` +} + +// CallToolRequest 调用工具请求 +type CallToolRequest struct { + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments"` +} + +// CallToolResponse 调用工具响应 +type CallToolResponse struct { + Content []Content `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +// ToolExecution 工具执行记录 +type ToolExecution struct { + ID string `json:"id"` + ToolName string `json:"toolName"` + Arguments map[string]interface{} `json:"arguments"` + Status string `json:"status"` // pending, running, completed, failed, cancelled + Result *ToolResult `json:"result,omitempty"` + Error string `json:"error,omitempty"` + StartTime time.Time `json:"startTime"` + EndTime *time.Time `json:"endTime,omitempty"` + Duration time.Duration `json:"duration,omitempty"` + // PartialOutput is a bounded tail preview of output produced by a running tool. + // It is intentionally separate from Result, which remains the final canonical tool result. + PartialOutput string `json:"partialOutput,omitempty"` + PartialOutputBytes int64 `json:"partialOutputBytes,omitempty"` + PartialOutputTruncated bool `json:"partialOutputTruncated,omitempty"` + PartialOutputUpdatedAt *time.Time `json:"partialOutputUpdatedAt,omitempty"` + // ConversationID 仅 API 展示用（进行中的 Agent 任务），不写入 tool_executions 表。 + ConversationID string `json:"conversationId,omitempty"` + OwnerUserID string `json:"-"` +} + +// ToolStats 工具统计信息 +type ToolStats struct { + ToolName string `json:"toolName"` + TotalCalls int `json:"totalCalls"` + SuccessCalls int `json:"successCalls"` + FailedCalls int `json:"failedCalls"` + LastCallTime *time.Time `json:"lastCallTime,omitempty"` +} + +// Prompt 提示词模板 +type Prompt struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Arguments []PromptArgument `json:"arguments,omitempty"` +} + +// PromptArgument 提示词参数 +type PromptArgument struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Required bool `json:"required,omitempty"` +} + +// GetPromptRequest 获取提示词请求 +type GetPromptRequest struct { + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments,omitempty"` +} + +// GetPromptResponse 获取提示词响应 +type GetPromptResponse struct { + Messages []PromptMessage `json:"messages"` +} + +// PromptMessage 提示词消息 +type PromptMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +// Resource 资源 +type Resource struct { + URI string `json:"uri"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + MimeType string `json:"mimeType,omitempty"` +} + +// ReadResourceRequest 读取资源请求 +type ReadResourceRequest struct { + URI string `json:"uri"` +} + +// ReadResourceResponse 读取资源响应 +type ReadResourceResponse struct { + Contents []ResourceContent `json:"contents"` +} + +// ResourceContent 资源内容 +type ResourceContent struct { + URI string `json:"uri"` + MimeType string `json:"mimeType,omitempty"` + Text string `json:"text,omitempty"` + Blob string `json:"blob,omitempty"` +} + +// SamplingRequest 采样请求 +type SamplingRequest struct { + Messages []SamplingMessage `json:"messages"` + Model string `json:"model,omitempty"` + MaxTokens int `json:"maxTokens,omitempty"` + Temperature float64 `json:"temperature,omitempty"` + TopP float64 `json:"topP,omitempty"` +} + +// SamplingMessage 采样消息 +type SamplingMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +// SamplingResponse 采样响应 +type SamplingResponse struct { + Content []SamplingContent `json:"content"` + Model string `json:"model,omitempty"` + StopReason string `json:"stopReason,omitempty"` +} + +// SamplingContent 采样内容 +type SamplingContent struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` +} + +// ToolResultPlainText 拼接工具结果中的文本（手动终止时作为「工具原始输出」）。 +func ToolResultPlainText(r *ToolResult) string { + if r == nil || len(r.Content) == 0 { + return "" + } + var b strings.Builder + for _, c := range r.Content { + b.WriteString(c.Text) + } + return strings.TrimSpace(b.String()) +} + +// AbortNoteBannerForModel 标出后续文本来自「用户手动终止工具时在弹窗中填写」，避免与 stdout/stderr 混淆。 +const AbortNoteBannerForModel = "---\n" + + "【用户终止说明｜USER INTERRUPT NOTE】\n" + + "（以下由操作者填写，用于指示模型如何继续；不是工具原始输出。）\n" + + "（Written by the operator when stopping this tool; not raw tool output.）\n" + + "---" + +// MergePartialToolOutputAndAbortNote 格式：工具原始输出 + 醒目标题 + 用户终止说明（无说明则原样返回 partial）。 +func MergePartialToolOutputAndAbortNote(partial, userNote string) string { + partial = strings.TrimSpace(partial) + userNote = strings.TrimSpace(userNote) + if userNote == "" { + return partial + } + section := AbortNoteBannerForModel + "\n" + userNote + if partial == "" { + return section + } + return partial + "\n\n" + section +} diff --git a/internal/security/auth_manager.go b/internal/security/auth_manager.go new file mode 100644 index 00000000..74e3b79a --- /dev/null +++ b/internal/security/auth_manager.go @@ -0,0 +1,266 @@ +package security + +import ( + "database/sql" + "errors" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/database" + + "github.com/google/uuid" +) + +// Predefined errors for authentication operations. +var ( + ErrInvalidPassword = errors.New("invalid password") +) + +// Session represents an authenticated user session. +type Session struct { + Token string + ExpiresAt time.Time + UserID string + Username string + DisplayName string + Roles []string + Permissions map[string]bool + PermissionScopes map[string]string + Scope string +} + +// AuthManager manages password-based authentication and session lifecycle. +type AuthManager struct { + sessionDuration time.Duration + db *database.DB + + mu sync.RWMutex + sessions map[string]Session +} + +// NewAuthManager creates a new AuthManager instance. +func NewAuthManager(sessionDurationHours int) *AuthManager { + if sessionDurationHours <= 0 { + sessionDurationHours = 12 + } + + return &AuthManager{ + sessionDuration: time.Duration(sessionDurationHours) * time.Hour, + sessions: make(map[string]Session), + } +} + +// AttachRBACStore enables multi-user RBAC authentication. When no users exist yet, +// it bootstraps the built-in admin account and returns the generated initial password. +func (a *AuthManager) AttachRBACStore(db *database.DB) (generatedAdminPassword string, err error) { + if db == nil { + return "", errors.New("database is required for authentication") + } + + needsAdminPassword, err := db.RBACNeedsAdminPassword() + if err != nil { + return "", err + } + + adminPasswordHash := "" + if needsAdminPassword { + generatedAdminPassword, err = GenerateStrongPassword(24) + if err != nil { + return "", err + } + adminPasswordHash, err = HashPassword(generatedAdminPassword) + if err != nil { + return "", err + } + } + + if err := db.BootstrapRBAC(adminPasswordHash, PermissionCatalog); err != nil { + return "", err + } + + a.mu.Lock() + a.db = db + a.mu.Unlock() + return generatedAdminPassword, nil +} + +// Authenticate validates the password and creates a new session. +func (a *AuthManager) Authenticate(username, password string) (string, time.Time, error) { + session, err := a.authenticateSession(username, password) + if err != nil { + return "", time.Time{}, err + } + a.mu.Lock() + a.sessions[session.Token] = session + a.mu.Unlock() + return session.Token, session.ExpiresAt, nil +} + +func (a *AuthManager) authenticateSession(username, password string) (Session, error) { + token := uuid.NewString() + expiresAt := time.Now().Add(a.sessionDuration) + + a.mu.RLock() + db := a.db + a.mu.RUnlock() + if db == nil { + return Session{}, errors.New("authentication store is not configured") + } + + username = strings.TrimSpace(strings.ToLower(username)) + if username == "" { + username = "admin" + } + user, err := db.GetRBACUserByUsername(username) + if err != nil { + if err == sql.ErrNoRows { + return Session{}, ErrInvalidPassword + } + return Session{}, err + } + if !user.Enabled || !VerifyPasswordHash(password, user.PasswordHash) { + return Session{}, ErrInvalidPassword + } + access, err := db.ResolveRBACAccess(user.ID) + if err != nil { + return Session{}, err + } + roleIDs := make([]string, 0, len(access.Roles)) + for _, role := range access.Roles { + roleIDs = append(roleIDs, role.ID) + } + return Session{ + Token: token, + ExpiresAt: expiresAt, + UserID: user.ID, + Username: user.Username, + DisplayName: user.DisplayName, + Roles: roleIDs, + Permissions: access.Permissions, + PermissionScopes: access.PermissionScopes, + Scope: access.Scope, + }, nil +} + +func (s Session) ScopeFor(permission string) string { + if scope := strings.TrimSpace(s.PermissionScopes[strings.TrimSpace(permission)]); scope != "" { + return scope + } + return strings.TrimSpace(s.Scope) +} + +// ValidateToken checks whether the provided token is still valid. +func (a *AuthManager) ValidateToken(token string) (Session, bool) { + if strings.TrimSpace(token) == "" { + return Session{}, false + } + + a.mu.RLock() + session, ok := a.sessions[token] + a.mu.RUnlock() + if !ok { + return Session{}, false + } + + if time.Now().After(session.ExpiresAt) { + a.mu.Lock() + delete(a.sessions, token) + a.mu.Unlock() + return Session{}, false + } + + return session, true +} + +// CheckPassword verifies whether the provided password matches the current password. +func (a *AuthManager) CheckPassword(password string) bool { + return a.CheckUserPassword("admin", password) +} + +// CheckUserPassword verifies whether the provided password matches a user. +func (a *AuthManager) CheckUserPassword(username, password string) bool { + a.mu.RLock() + db := a.db + a.mu.RUnlock() + if db == nil { + return false + } + user, err := db.GetRBACUserByUsername(username) + if err != nil { + return false + } + return VerifyPasswordHash(password, user.PasswordHash) +} + +func (a *AuthManager) UpdateUserPassword(userID, password string) error { + password = strings.TrimSpace(password) + if password == "" { + return errors.New("auth password must be configured") + } + hash, err := HashPassword(password) + if err != nil { + return err + } + a.mu.RLock() + db := a.db + a.mu.RUnlock() + if db == nil { + return errors.New("authentication store is not configured") + } + if err := db.UpdateRBACUserPassword(userID, hash); err != nil { + return err + } + a.mu.Lock() + for token, session := range a.sessions { + if session.UserID == userID { + delete(a.sessions, token) + } + } + a.mu.Unlock() + return nil +} + +// RevokeToken invalidates the specified token. +func (a *AuthManager) RevokeToken(token string) { + if strings.TrimSpace(token) == "" { + return + } + + a.mu.Lock() + delete(a.sessions, token) + a.mu.Unlock() +} + +func (a *AuthManager) RevokeUserSessions(userID string) { + userID = strings.TrimSpace(userID) + if userID == "" { + return + } + a.mu.Lock() + for token, session := range a.sessions { + if session.UserID == userID { + delete(a.sessions, token) + } + } + a.mu.Unlock() +} + +func (a *AuthManager) RevokeAllSessions() { + a.mu.Lock() + a.sessions = make(map[string]Session) + a.mu.Unlock() +} + +// SessionDurationHours returns the configured session duration in hours. +func (a *AuthManager) SessionDurationHours() int { + return int(a.sessionDuration / time.Hour) +} + +func allPermissions() map[string]bool { + out := make(map[string]bool, len(PermissionCatalog)) + for key := range PermissionCatalog { + out[key] = true + } + return out +} diff --git a/internal/security/auth_manager_bootstrap_test.go b/internal/security/auth_manager_bootstrap_test.go new file mode 100644 index 00000000..fff311b3 --- /dev/null +++ b/internal/security/auth_manager_bootstrap_test.go @@ -0,0 +1,38 @@ +package security + +import ( + "path/filepath" + "testing" + + "cyberstrike-ai/internal/database" + + "go.uber.org/zap" +) + +func TestAttachRBACStoreBootstrapsAdminPassword(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "auth-bootstrap.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + manager := NewAuthManager(12) + generated, err := manager.AttachRBACStore(db) + if err != nil { + t.Fatalf("AttachRBACStore: %v", err) + } + if generated == "" { + t.Fatal("expected generated admin password on first bootstrap") + } + if !manager.CheckUserPassword("admin", generated) { + t.Fatal("generated password should authenticate admin") + } + + second, err := manager.AttachRBACStore(db) + if err != nil { + t.Fatalf("AttachRBACStore second call: %v", err) + } + if second != "" { + t.Fatalf("expected no password on second bootstrap, got %q", second) + } +} diff --git a/internal/security/auth_manager_test.go b/internal/security/auth_manager_test.go new file mode 100644 index 00000000..25e1b591 --- /dev/null +++ b/internal/security/auth_manager_test.go @@ -0,0 +1,94 @@ +package security + +import ( + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +func TestAuthManagerAuthenticatesCreatedRBACUser(t *testing.T) { + db, err := database.NewDB(filepath.Join(t.TempDir(), "auth-rbac.db"), zap.NewNop()) + if err != nil { + t.Fatalf("NewDB: %v", err) + } + t.Cleanup(func() { _ = db.Close() }) + + manager := NewAuthManager(12) + if _, err := manager.AttachRBACStore(db); err != nil { + t.Fatalf("AttachRBACStore: %v", err) + } + hash, err := HashPassword("operator-secret") + if err != nil { + t.Fatalf("HashPassword: %v", err) + } + user, err := db.CreateRBACUser("operator1", "Operator One", hash, true, []string{database.RBACSystemRoleViewer}) + if err != nil { + t.Fatalf("CreateRBACUser: %v", err) + } + + token, _, err := manager.Authenticate("operator1", "operator-secret") + if err != nil { + t.Fatalf("Authenticate created user: %v", err) + } + session, ok := manager.ValidateToken(token) + if !ok { + t.Fatalf("expected created user session to validate") + } + if session.UserID != user.ID || session.Username != "operator1" { + t.Fatalf("session user = %s/%s, want %s/operator1", session.UserID, session.Username, user.ID) + } + if !session.Permissions["auth:self"] || !session.Permissions["chat:read"] { + t.Fatalf("expected viewer permissions in session, got %#v", session.Permissions) + } + + if _, _, err := manager.Authenticate("", "operator-secret"); err == nil { + t.Fatalf("empty username must not authenticate non-admin user") + } + + router := gin.New() + router.Use(AuthMiddleware(manager)) + router.GET("/principal", func(c *gin.Context) { + principal, ok := authctx.PrincipalFromContext(c.Request.Context()) + if !ok || principal.UserID != user.ID || !principal.HasPermission("chat:read") || principal.ScopeFor("chat:read") != database.RBACScopeAssigned { + c.Status(http.StatusInternalServerError) + return + } + c.Status(http.StatusNoContent) + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/principal", nil) + req.Header.Set("Authorization", "Bearer "+token) + router.ServeHTTP(w, req) + if w.Code != http.StatusNoContent { + t.Fatalf("principal propagation status = %d", w.Code) + } +} + +func TestQueryTokenOnlyAllowedForSSEAndWebSocketGET(t *testing.T) { + requestToken := func(method, accept, upgrade string) string { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(method, "/api/test?token=secret", nil) + c.Request.Header.Set("Accept", accept) + c.Request.Header.Set("Upgrade", upgrade) + return extractTokenFromRequest(c) + } + if got := requestToken(http.MethodGet, "application/json", ""); got != "" { + t.Fatalf("ordinary GET accepted query token %q", got) + } + if got := requestToken(http.MethodPost, "text/event-stream", ""); got != "" { + t.Fatalf("POST accepted query token %q", got) + } + if got := requestToken(http.MethodGet, "text/event-stream", ""); got != "secret" { + t.Fatalf("SSE token = %q", got) + } + if got := requestToken(http.MethodGet, "", "websocket"); got != "secret" { + t.Fatalf("WebSocket token = %q", got) + } +} diff --git a/internal/security/auth_middleware.go b/internal/security/auth_middleware.go new file mode 100644 index 00000000..8b4d5e24 --- /dev/null +++ b/internal/security/auth_middleware.go @@ -0,0 +1,151 @@ +package security + +import ( + "net/http" + "strings" + + "cyberstrike-ai/internal/authctx" + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" +) + +const ( + ContextAuthTokenKey = "authToken" + ContextSessionExpiry = "authSessionExpiry" + ContextUserIDKey = "authUserID" + ContextUsernameKey = "authUsername" + ContextUserScopeKey = "authUserScope" + ContextSessionKey = "authSession" +) + +// AuthMiddleware enforces authentication on protected routes. +func AuthMiddleware(manager *AuthManager) gin.HandlerFunc { + return func(c *gin.Context) { + token := extractTokenFromRequest(c) + session, ok := manager.ValidateToken(token) + if !ok { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{ + "error": "未授权访问，请先登录", + }) + return + } + + c.Set(ContextAuthTokenKey, session.Token) + c.Set(ContextSessionExpiry, session.ExpiresAt) + c.Set(ContextUserIDKey, session.UserID) + c.Set(ContextUsernameKey, session.Username) + c.Set(ContextUserScopeKey, session.Scope) + c.Set(ContextSessionKey, session) + // Gin context values do not survive into Agent/MCP/background contexts. + // Attach an immutable principal to the request context as the canonical + // identity for every downstream execution layer. + principal := authctx.NewPrincipalWithScopes(session.UserID, session.Username, session.Scope, session.Permissions, session.PermissionScopes) + c.Request = c.Request.WithContext(authctx.WithPrincipal(c.Request.Context(), principal)) + c.Next() + } +} + +func RequirePermission(permission string) gin.HandlerFunc { + permission = strings.TrimSpace(permission) + return func(c *gin.Context) { + if permission == "" || SessionHasPermission(c, permission) { + c.Next() + return + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "权限不足", + "permission": permission, + }) + } +} + +func RequireAnyPermission(permissions ...string) gin.HandlerFunc { + return func(c *gin.Context) { + for _, permission := range permissions { + if SessionHasPermission(c, permission) { + c.Next() + return + } + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "权限不足", + "permissions": permissions, + }) + } +} + +func RequireResourcePermission(db *database.DB, permission, resourceType, paramName string) gin.HandlerFunc { + return func(c *gin.Context) { + if !SessionHasPermission(c, permission) { + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "权限不足", + "permission": permission, + }) + return + } + if db == nil { + c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"error": "资源鉴权服务不可用"}) + return + } + resourceID := strings.TrimSpace(c.Param(paramName)) + if resourceID == "" { + c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "资源 ID 不能为空"}) + return + } + session, ok := CurrentSession(c) + if !ok || !db.UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID) { + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "无权访问该资源", + "resource_type": resourceType, + "resource_id": resourceID, + }) + return + } + c.Next() + } +} + +func CurrentSession(c *gin.Context) (Session, bool) { + if c == nil { + return Session{}, false + } + v, ok := c.Get(ContextSessionKey) + if !ok { + return Session{}, false + } + session, ok := v.(Session) + return session, ok +} + +func SessionHasPermission(c *gin.Context, permission string) bool { + session, ok := CurrentSession(c) + if !ok { + return false + } + return session.Permissions[permission] +} + +func extractTokenFromRequest(c *gin.Context) string { + authHeader := c.GetHeader("Authorization") + if authHeader != "" { + if len(authHeader) > 7 && strings.EqualFold(authHeader[0:7], "Bearer ") { + return strings.TrimSpace(authHeader[7:]) + } + return strings.TrimSpace(authHeader) + } + + if token := c.Query("token"); token != "" && c.Request.Method == http.MethodGet { + acceptsSSE := strings.Contains(strings.ToLower(c.GetHeader("Accept")), "text/event-stream") + upgradesWebSocket := strings.EqualFold(strings.TrimSpace(c.GetHeader("Upgrade")), "websocket") + if acceptsSSE || upgradesWebSocket { + return strings.TrimSpace(token) + } + } + + if cookie, err := c.Cookie("auth_token"); err == nil { + return strings.TrimSpace(cookie) + } + + return "" +} diff --git a/internal/security/command_failure_format.go b/internal/security/command_failure_format.go new file mode 100644 index 00000000..dc5af2c5 --- /dev/null +++ b/internal/security/command_failure_format.go @@ -0,0 +1,56 @@ +package security + +import ( + "errors" + "fmt" + "os/exec" + "strings" +) + +// FormatCommandFailureResult 与 exec 工具 ToolResult 文案一致（不含 ToolErrorPrefix）。 +func FormatCommandFailureResult(exitCode int, output string) string { + output = strings.TrimSpace(output) + errMsg := fmt.Sprintf("exit status %d", exitCode) + if output == "" { + return fmt.Sprintf("命令执行失败: %s", errMsg) + } + if strings.HasPrefix(output, "命令执行失败:") { + return output + } + return fmt.Sprintf("命令执行失败: %s\n输出: %s", errMsg, output) +} + +// FormatCommandFailureFromErr 根据 exec/execute 返回的 error 生成统一失败文案（IsError 正文）。 +func FormatCommandFailureFromErr(err error, output string) string { + if err == nil { + return strings.TrimSpace(output) + } + var exitError *exec.ExitError + if errors.As(err, &exitError) { + return FormatCommandFailureResult(exitError.ExitCode(), output) + } + output = strings.TrimSpace(output) + if output == "" { + return fmt.Sprintf("命令执行失败: %v", err) + } + if strings.HasPrefix(output, "命令执行失败:") { + return output + } + return fmt.Sprintf("命令执行失败: %v\n输出: %s", err, output) +} + +// ExecuteFailureStatusLine 流式 execute 结束时追加的单行状态（输出正文已在流中推送过）。 +func ExecuteFailureStatusLine(exitCode int) string { + return fmt.Sprintf("\n命令执行失败: exit status %d", exitCode) +} + +// IsCommandFailureResult 判断工具结果正文是否表示命令非零退出（用于 execute / exec 对齐 isError）。 +func IsCommandFailureResult(content string) bool { + return strings.Contains(content, "命令执行失败:") +} + +// IsLegacyShellExitNoise 过滤旧版 shell 流中冗余的 exit code 行。 +func IsLegacyShellExitNoise(s string) bool { + trimmed := strings.TrimSpace(s) + return strings.HasPrefix(trimmed, "command exited with non-zero code ") +} diff --git a/internal/security/command_failure_format_test.go b/internal/security/command_failure_format_test.go new file mode 100644 index 00000000..d7ca53a2 --- /dev/null +++ b/internal/security/command_failure_format_test.go @@ -0,0 +1,54 @@ +package security + +import ( + "errors" + "os/exec" + "strings" + "testing" +) + +func TestFormatCommandFailureResult(t *testing.T) { + got := FormatCommandFailureResult(1, "sudo: password required") + want := "命令执行失败: exit status 1\n输出: sudo: password required" + if got != want { + t.Fatalf("got %q want %q", got, want) + } + if FormatCommandFailureResult(2, "") != "命令执行失败: exit status 2" { + t.Fatal("empty output format") + } + if FormatCommandFailureResult(1, "命令执行失败: exit status 1") != "命令执行失败: exit status 1" { + t.Fatal("should not double-wrap") + } +} + +func TestIsCommandFailureResult(t *testing.T) { + if !IsCommandFailureResult("sudo: err\n命令执行失败: exit status 1") { + t.Fatal("expected true") + } + if IsCommandFailureResult("sudo: err only") { + t.Fatal("expected false") + } +} + +func TestFormatCommandFailureFromErr(t *testing.T) { + cmd := exec.Command("sh", "-c", "exit 42") + err := cmd.Run() + got := FormatCommandFailureFromErr(err, "oops") + if got != "命令执行失败: exit status 42\n输出: oops" { + t.Fatalf("got %q", got) + } + timeoutErr := errors.New("shell inactivity timeout (300s)") + got2 := FormatCommandFailureFromErr(timeoutErr, "already timed out") + if !strings.Contains(got2, "shell inactivity timeout") || !strings.Contains(got2, "already timed out") { + t.Fatalf("got %q", got2) + } +} + +func TestIsLegacyShellExitNoise(t *testing.T) { + if !IsLegacyShellExitNoise("command exited with non-zero code 1\n") { + t.Fatal("expected legacy noise") + } + if IsLegacyShellExitNoise("sudo: failed") { + t.Fatal("unexpected noise") + } +} diff --git a/internal/security/executor.go b/internal/security/executor.go new file mode 100644 index 00000000..a1f6ae97 --- /dev/null +++ b/internal/security/executor.go @@ -0,0 +1,1626 @@ +package security + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "runtime" + "strconv" + "strings" + "sync" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + "cyberstrike-ai/internal/tooloutput" + + "github.com/creack/pty" + "github.com/google/uuid" + "go.uber.org/zap" +) + +// ToolOutputCallback 用于在工具执行过程中把 stdout/stderr 增量推给上层（SSE）。 +// 通过 context 传递，避免修改 MCP ToolHandler 签名导致的“写死工具”问题。 +type ToolOutputCallback func(chunk string) + +type toolOutputCallbackCtxKey struct{} + +// ToolOutputCallbackCtxKey 是 context 中的 key，供 Agent 写入回调，Executor 读取并流式回调。 +var ToolOutputCallbackCtxKey = toolOutputCallbackCtxKey{} + +// Executor 安全工具执行器 +type Executor struct { + config *config.SecurityConfig + toolIndex map[string]*config.ToolConfig // 工具索引，用于 O(1) 查找 + mcpServer *mcp.Server + logger *zap.Logger + shellNoOutputTimeoutSec int // execute/exec 无新输出空闲秒数；0=默认 300；-1=关闭（见 SetShellNoOutputTimeoutSeconds） + toolOutputMaxBytes int + spillRootDir string +} + +// NewExecutor 创建新的执行器 +func NewExecutor(cfg *config.SecurityConfig, mcpServer *mcp.Server, logger *zap.Logger) *Executor { + executor := &Executor{ + config: cfg, + toolIndex: make(map[string]*config.ToolConfig), + mcpServer: mcpServer, + logger: logger, + } + // 构建工具索引 + executor.buildToolIndex() + return executor +} + +// SetShellNoOutputTimeoutSeconds 配置 exec 工具无输出空闲终止（与 agent.shell_no_output_timeout_seconds 一致）。 +func (e *Executor) SetShellNoOutputTimeoutSeconds(sec int) { + e.shellNoOutputTimeoutSec = sec +} + +// SetToolOutputMaxBytes limits stdout/stderr retained and streamed by exec-like +// tools. It should stay aligned with MCP result normalization so every channel +// sees the same bounded payload. Oversized full output is spilled to disk first. +func (e *Executor) SetToolOutputMaxBytes(maxBytes int) { + e.toolOutputMaxBytes = maxBytes +} + +// SetToolOutputSpillRoot sets the reduction-compatible root for spilling full +// exec stdout/stderr when the in-memory bound is exceeded (empty → tmp/reduction). +func (e *Executor) SetToolOutputSpillRoot(rootDir string) { + e.spillRootDir = strings.TrimSpace(rootDir) +} + +func (e *Executor) wrapToolOutputCallback(ctx context.Context, cb ToolOutputCallback) ToolOutputCallback { + executionID := mcp.MCPExecutionIDFromContext(ctx) + if e == nil || e.mcpServer == nil || strings.TrimSpace(executionID) == "" { + return cb + } + return func(chunk string) { + if chunk != "" { + e.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk) + } + if cb != nil { + cb(chunk) + } + } +} + +func (e *Executor) spillOptsFromContext(ctx context.Context) tooloutput.SpillOpts { + root := "" + if e != nil { + root = e.spillRootDir + } + opts := tooloutput.SpillOpts{RootDir: root} + if ctx != nil { + opts.ConversationID = mcp.MCPConversationIDFromContext(ctx) + opts.ProjectID = mcp.MCPProjectIDFromContext(ctx) + opts.ExecutionID = mcp.MCPExecutionIDFromContext(ctx) + } + if opts.ExecutionID == "" { + opts.ExecutionID = uuid.NewString() + } + return opts +} + +// buildToolIndex 构建工具索引，将 O(n) 查找优化为 O(1) +func (e *Executor) buildToolIndex() { + e.toolIndex = make(map[string]*config.ToolConfig) + for i := range e.config.Tools { + if e.config.Tools[i].Enabled { + e.toolIndex[e.config.Tools[i].Name] = &e.config.Tools[i] + } + } + e.logger.Debug("工具索引构建完成", + zap.Int("totalTools", len(e.config.Tools)), + zap.Int("enabledTools", len(e.toolIndex)), + ) +} + +// ExecuteTool 执行安全工具 +func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[string]interface{}) (*mcp.ToolResult, error) { + e.logger.Debug("ExecuteTool被调用", + zap.String("toolName", toolName), + zap.Any("args", args), + ) + + // 特殊处理：exec工具直接执行系统命令 + if toolName == "exec" { + e.logger.Debug("执行exec工具") + return e.executeSystemCommand(ctx, args) + } + + // 使用索引查找工具配置（O(1) 查找） + toolConfig, exists := e.toolIndex[toolName] + if !exists { + e.logger.Error("工具未找到或未启用", + zap.String("toolName", toolName), + zap.Int("totalTools", len(e.config.Tools)), + zap.Int("enabledTools", len(e.toolIndex)), + ) + return nil, fmt.Errorf("工具 %s 未找到或未启用", toolName) + } + + e.logger.Debug("找到工具配置", + zap.String("toolName", toolName), + zap.String("command", toolConfig.Command), + zap.Strings("args", toolConfig.Args), + ) + + // 特殊处理：内部工具（command 以 "internal:" 开头） + if strings.HasPrefix(toolConfig.Command, "internal:") { + e.logger.Debug("执行内部工具", + zap.String("toolName", toolName), + zap.String("command", toolConfig.Command), + ) + return e.executeInternalTool(ctx, toolName, toolConfig.Command, args) + } + + // 构建命令 - 根据工具类型使用不同的参数格式 + cmdArgs := e.buildCommandArgs(toolName, toolConfig, args) + + e.logger.Debug("构建命令参数完成", + zap.String("toolName", toolName), + zap.Strings("cmdArgs", cmdArgs), + zap.Int("argsCount", len(cmdArgs)), + ) + + // 验证命令参数 + if len(cmdArgs) == 0 { + e.logger.Warn("命令参数为空", + zap.String("toolName", toolName), + zap.Any("inputArgs", args), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("错误: 工具 %s 缺少必需的参数。接收到的参数: %v", toolName, args), + }, + }, + IsError: true, + }, nil + } + + // 执行命令 + cmd := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd) + attachNonInteractiveStdin(cmd) + _ = prepareShellCmdSession(cmd) + + e.logger.Debug("执行安全工具", + zap.String("tool", toolName), + zap.Strings("args", cmdArgs), + ) + + var output string + var err error + spill := e.spillOptsFromContext(ctx) + // 如果上层提供了 stdout/stderr 增量回调，或当前处于 MCP execution 中，则边执行边读取并回调。 + if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" { + cb = e.wrapToolOutputCallback(ctx, cb) + output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到工具需要 TTY，使用 PTY 重试", + zap.String("tool", toolName), + ) + cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd2) + _ = prepareShellCmdSession(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill) + } + } else { + // 非流式：内存缓冲 + ctx 取消杀进程组；行为对齐原 CombinedOutput，避免双流管道 fan-in 死锁。 + output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到工具需要 TTY，使用 PTY 重试", + zap.String("tool", toolName), + ) + cmd2 := exec.CommandContext(ctx, toolConfig.Command, cmdArgs...) + applyDefaultTerminalEnv(cmd2) + _ = prepareShellCmdSession(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill) + } + } + if err != nil { + // 检查退出码是否在允许列表中 + exitCode := getExitCode(err) + if exitCode != nil && toolConfig.AllowedExitCodes != nil { + for _, allowedCode := range toolConfig.AllowedExitCodes { + if *exitCode == allowedCode { + e.logger.Debug("工具执行完成（退出码在允许列表中）", + zap.String("tool", toolName), + zap.Int("exitCode", *exitCode), + zap.String("output", string(output)), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: string(output), + }, + }, + IsError: false, + }, nil + } + } + } + + e.logger.Error("工具执行失败", + zap.String("tool", toolName), + zap.Error(err), + zap.Int("exitCode", getExitCodeValue(err)), + zap.String("output", string(output)), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("工具执行失败: %v\n输出: %s", err, string(output)), + }, + }, + IsError: true, + }, nil + } + + e.logger.Debug("工具执行成功", + zap.String("tool", toolName), + zap.String("output", string(output)), + ) + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: string(output), + }, + }, + IsError: false, + }, nil +} + +// RegisterTools 注册工具到MCP服务器 +func (e *Executor) RegisterTools(mcpServer *mcp.Server) { + e.logger.Debug("开始注册工具", + zap.Int("totalTools", len(e.config.Tools)), + zap.Int("enabledTools", len(e.toolIndex)), + ) + + // 重新构建索引（以防配置更新） + e.buildToolIndex() + + for i, toolConfig := range e.config.Tools { + if !toolConfig.Enabled { + e.logger.Debug("跳过未启用的工具", + zap.String("tool", toolConfig.Name), + ) + continue + } + + // 创建工具配置的副本，避免闭包问题 + toolName := toolConfig.Name + toolConfigCopy := toolConfig + + // 根据配置决定暴露给 AI/API 的描述：short_description 或 description + useFullDescription := strings.TrimSpace(strings.ToLower(e.config.ToolDescriptionMode)) == "full" + shortDesc := toolConfigCopy.ShortDescription + if shortDesc == "" { + // 如果没有简短描述，从详细描述中提取第一行或前10000个字符 + desc := toolConfigCopy.Description + if len(desc) > 10000 { + if idx := strings.Index(desc, "\n"); idx > 0 && idx < 10000 { + shortDesc = strings.TrimSpace(desc[:idx]) + } else { + shortDesc = desc[:10000] + "..." + } + } else { + shortDesc = desc + } + } + if useFullDescription { + shortDesc = "" // 使用 description 时清空 ShortDescription，下游会回退到 Description + } + + tool := mcp.Tool{ + Name: toolConfigCopy.Name, + Description: toolConfigCopy.Description, + ShortDescription: shortDesc, + InputSchema: e.buildInputSchema(&toolConfigCopy), + } + + handler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + e.logger.Debug("工具handler被调用", + zap.String("toolName", toolName), + zap.Any("args", args), + ) + return e.ExecuteTool(ctx, toolName, args) + } + + mcpServer.RegisterTool(tool, handler) + e.logger.Debug("注册安全工具成功", + zap.String("tool", toolConfigCopy.Name), + zap.String("command", toolConfigCopy.Command), + zap.Int("index", i), + ) + } + + e.logger.Debug("工具注册完成", + zap.Int("registeredCount", len(e.config.Tools)), + ) +} + +// buildCommandArgs 构建命令参数 +func (e *Executor) buildCommandArgs(toolName string, toolConfig *config.ToolConfig, args map[string]interface{}) []string { + cmdArgs := make([]string, 0) + + // 如果配置中定义了参数映射，使用配置中的映射规则 + if len(toolConfig.Parameters) > 0 { + // 检查是否有 scan_type 参数，如果有则替换默认的扫描类型参数 + hasScanType := false + var scanTypeValue string + if scanType, ok := args["scan_type"].(string); ok && scanType != "" { + hasScanType = true + scanTypeValue = scanType + } + + // 添加固定参数（如果指定了 scan_type，可能需要过滤掉默认的扫描类型参数） + if hasScanType && toolName == "nmap" { + // 对于 nmap，如果指定了 scan_type，跳过默认的 -sT -sV -sC + // 这些参数会被 scan_type 参数替换 + } else { + cmdArgs = append(cmdArgs, toolConfig.Args...) + } + + // 按位置参数排序 + positionalParams := make([]config.ParameterConfig, 0) + flagParams := make([]config.ParameterConfig, 0) + + for _, param := range toolConfig.Parameters { + if param.Position != nil { + positionalParams = append(positionalParams, param) + } else { + flagParams = append(flagParams, param) + } + } + + // 对于需要子命令的工具（如 gobuster dir），position 0 必须紧跟在命令名后、所有 flag 之前 + for _, param := range positionalParams { + if param.Name == "additional_args" || param.Name == "scan_type" || param.Name == "action" { + continue + } + if param.Position != nil && *param.Position == 0 { + value := e.getParamValue(args, param) + if value == nil && param.Default != nil { + value = param.Default + } + if value != nil { + cmdArgs = append(cmdArgs, e.formatParamValue(param, value)) + } + break + } + } + + // 处理标志参数 + for _, param := range flagParams { + // 跳过特殊参数，它们会在后面单独处理 + // action 参数仅用于工具内部逻辑，不传递给命令 + if param.Name == "additional_args" || param.Name == "scan_type" || param.Name == "action" { + continue + } + + value := e.getParamValue(args, param) + if value == nil { + if param.Required { + // 必需参数缺失，返回空数组让上层处理错误 + e.logger.Warn("缺少必需的标志参数", + zap.String("tool", toolName), + zap.String("param", param.Name), + ) + return []string{} + } + continue + } + + // 布尔值特殊处理：如果为 false，跳过；如果为 true，只添加标志 + if param.Type == "bool" { + var boolVal bool + var ok bool + + // 尝试多种类型转换 + if boolVal, ok = value.(bool); ok { + // 已经是布尔值 + } else if numVal, ok := value.(float64); ok { + // JSON 数字类型（float64） + boolVal = numVal != 0 + ok = true + } else if numVal, ok := value.(int); ok { + // int 类型 + boolVal = numVal != 0 + ok = true + } else if strVal, ok := value.(string); ok { + // 字符串类型 + boolVal = strVal == "true" || strVal == "1" || strVal == "yes" + ok = true + } + + if ok { + if !boolVal { + continue // false 时不添加任何参数 + } + // true 时只添加标志，不添加值 + if param.Flag != "" { + cmdArgs = append(cmdArgs, param.Flag) + } + continue + } + } + + formattedValue := e.formatParamValue(param, value) + if strings.TrimSpace(formattedValue) == "" { + if param.Required { + e.logger.Warn("必需参数为空", + zap.String("tool", toolName), + zap.String("param", param.Name), + ) + return []string{} + } + continue + } + + format := param.Format + if format == "" { + format = "flag" // 默认格式 + } + + switch format { + case "flag": + // --flag value 或 -f value + if param.Flag != "" { + cmdArgs = append(cmdArgs, param.Flag) + } + cmdArgs = append(cmdArgs, formattedValue) + case "combined": + // --flag=value 或 -f=value + if param.Flag != "" { + cmdArgs = append(cmdArgs, fmt.Sprintf("%s=%s", param.Flag, formattedValue)) + } else { + cmdArgs = append(cmdArgs, formattedValue) + } + case "template": + // 使用模板字符串 + if param.Template != "" { + template := param.Template + template = strings.ReplaceAll(template, "{flag}", param.Flag) + template = strings.ReplaceAll(template, "{value}", formattedValue) + template = strings.ReplaceAll(template, "{name}", param.Name) + cmdArgs = append(cmdArgs, strings.Fields(template)...) + } else { + // 如果没有模板，使用默认格式 + if param.Flag != "" { + cmdArgs = append(cmdArgs, param.Flag) + } + cmdArgs = append(cmdArgs, formattedValue) + } + case "positional": + // 位置参数（已在上面处理） + cmdArgs = append(cmdArgs, formattedValue) + default: + // 默认：直接添加值 + cmdArgs = append(cmdArgs, formattedValue) + } + } + + // 然后处理位置参数（位置参数通常在标志参数之后） + // 对位置参数按位置排序 + // 首先找到最大的位置值，确定需要处理多少个位置 + maxPosition := -1 + for _, param := range positionalParams { + if param.Position != nil && *param.Position > maxPosition { + maxPosition = *param.Position + } + } + + // 按位置顺序处理参数，确保即使某些位置没有参数或使用默认值，也能正确传递 + // position 0 已在前面插入（子命令优先），此处从 1 开始 + for i := 0; i <= maxPosition; i++ { + if i == 0 { + continue + } + for _, param := range positionalParams { + // 跳过特殊参数，它们会在后面单独处理 + // action 参数仅用于工具内部逻辑，不传递给命令 + if param.Name == "additional_args" || param.Name == "scan_type" || param.Name == "action" { + continue + } + + if param.Position != nil && *param.Position == i { + value := e.getParamValue(args, param) + if value == nil { + if param.Required { + // 必需参数缺失，返回空数组让上层处理错误 + e.logger.Warn("缺少必需的位置参数", + zap.String("tool", toolName), + zap.String("param", param.Name), + zap.Int("position", *param.Position), + ) + return []string{} + } + // 对于非必需参数，如果值为 nil，尝试使用默认值 + if param.Default != nil { + value = param.Default + } else { + // 如果没有默认值，跳过这个位置，继续处理下一个位置 + break + } + } + // 只有当值不为 nil 时才添加到命令参数中 + if value != nil { + cmdArgs = append(cmdArgs, e.formatParamValue(param, value)) + } + break + } + } + // 如果某个位置没有找到对应的参数，继续处理下一个位置 + // 这样可以确保位置参数的顺序正确 + } + + // 特殊处理：additional_args 参数（需要按空格分割成多个参数） + if additionalArgs, ok := args["additional_args"].(string); ok && additionalArgs != "" { + // 按空格分割，但保留引号内的内容 + additionalArgsList := e.parseAdditionalArgs(additionalArgs) + cmdArgs = append(cmdArgs, additionalArgsList...) + } + + // 特殊处理：scan_type 参数（需要按空格分割并插入到合适位置） + if hasScanType { + scanTypeArgs := e.parseAdditionalArgs(scanTypeValue) + if len(scanTypeArgs) > 0 { + // 对于 nmap，scan_type 应该替换默认的扫描类型参数 + // 由于我们已经跳过了默认的 args，现在需要将 scan_type 插入到合适位置 + // 找到 target 参数的位置（通常是最后一个位置参数） + insertPos := len(cmdArgs) + for i := len(cmdArgs) - 1; i >= 0; i-- { + // target 通常是最后一个非标志参数 + if !strings.HasPrefix(cmdArgs[i], "-") { + insertPos = i + break + } + } + // 在 target 之前插入 scan_type 参数 + newArgs := make([]string, 0, len(cmdArgs)+len(scanTypeArgs)) + newArgs = append(newArgs, cmdArgs[:insertPos]...) + newArgs = append(newArgs, scanTypeArgs...) + newArgs = append(newArgs, cmdArgs[insertPos:]...) + cmdArgs = newArgs + } + } + + return cmdArgs + } + + // 如果没有定义参数配置，使用固定参数和通用处理 + // 添加固定参数 + cmdArgs = append(cmdArgs, toolConfig.Args...) + + // 通用处理：将参数转换为命令行参数 + for key, value := range args { + if key == "_tool_name" { + continue + } + // 使用 --key value 格式 + cmdArgs = append(cmdArgs, fmt.Sprintf("--%s", key)) + if strValue, ok := value.(string); ok { + cmdArgs = append(cmdArgs, strValue) + } else { + cmdArgs = append(cmdArgs, fmt.Sprintf("%v", value)) + } + } + + return cmdArgs +} + +// parseAdditionalArgs 解析 additional_args 字符串，按空格分割但保留引号内的内容 +func (e *Executor) parseAdditionalArgs(argsStr string) []string { + if argsStr == "" { + return []string{} + } + + result := make([]string, 0) + var current strings.Builder + inQuotes := false + var quoteChar rune + escapeNext := false + + runes := []rune(argsStr) + for i := 0; i < len(runes); i++ { + r := runes[i] + + if escapeNext { + current.WriteRune(r) + escapeNext = false + continue + } + + if r == '\\' { + // 检查下一个字符是否是引号 + if i+1 < len(runes) && (runes[i+1] == '"' || runes[i+1] == '\'') { + // 转义的引号：跳过反斜杠，将引号作为普通字符写入 + i++ + current.WriteRune(runes[i]) + } else { + // 其他转义字符：写入反斜杠，下一个字符会在下次迭代处理 + escapeNext = true + current.WriteRune(r) + } + continue + } + + if !inQuotes && (r == '"' || r == '\'') { + inQuotes = true + quoteChar = r + continue + } + + if inQuotes && r == quoteChar { + inQuotes = false + quoteChar = 0 + continue + } + + if !inQuotes && (r == ' ' || r == '\t' || r == '\n') { + if current.Len() > 0 { + result = append(result, current.String()) + current.Reset() + } + continue + } + + current.WriteRune(r) + } + + // 处理最后一个参数（如果存在） + if current.Len() > 0 { + result = append(result, current.String()) + } + + // 如果解析结果为空，使用简单的空格分割作为降级方案 + if len(result) == 0 { + result = strings.Fields(argsStr) + } + + return result +} + +// getParamValue 获取参数值，支持默认值 +func (e *Executor) getParamValue(args map[string]interface{}, param config.ParameterConfig) interface{} { + // 从参数中获取值 + if value, ok := args[param.Name]; ok && value != nil { + return value + } + + // 如果参数是必需的但没有提供，返回 nil（让上层处理错误） + if param.Required { + return nil + } + + // 返回默认值 + return param.Default +} + +// formatParamValue 格式化参数值 +func (e *Executor) formatParamValue(param config.ParameterConfig, value interface{}) string { + switch param.Type { + case "bool": + // 布尔值应该在上层处理，这里不应该被调用 + if boolVal, ok := value.(bool); ok { + return fmt.Sprintf("%v", boolVal) + } + return "false" + case "array": + // 数组：转换为逗号分隔的字符串 + if arr, ok := value.([]interface{}); ok { + strs := make([]string, 0, len(arr)) + for _, item := range arr { + strs = append(strs, fmt.Sprintf("%v", item)) + } + return strings.Join(strs, ",") + } + return fmt.Sprintf("%v", value) + case "object": + // 对象/字典：序列化为 JSON 字符串 + if jsonBytes, err := json.Marshal(value); err == nil { + return string(jsonBytes) + } + // 如果 JSON 序列化失败，回退到默认格式化 + return fmt.Sprintf("%v", value) + default: + formattedValue := fmt.Sprintf("%v", value) + // 特殊处理：对于 ports 参数（通常是 nmap 等工具的端口参数），清理空格 + // nmap 不接受端口列表中有空格，例如 "80,443, 22" 应该变成 "80,443,22" + if param.Name == "ports" { + // 移除所有空格，但保留逗号和其他字符 + formattedValue = strings.ReplaceAll(formattedValue, " ", "") + } + return formattedValue + } +} + +// IsBackgroundShellCommand 检测命令是否为完全后台命令（末尾有独立 &，且不在引号内）。 +// command1 & command2 不算完全后台（command2 仍在前台执行）。 +func IsBackgroundShellCommand(command string) bool { + command = strings.TrimSpace(command) + if command == "" { + return false + } + positions := findStandaloneAmpersandPositions(command) + if len(positions) == 0 { + return false + } + last := positions[len(positions)-1] + afterAmpersand := strings.TrimSpace(command[last+1:]) + if afterAmpersand != "" { + return false + } + beforeAmpersand := strings.TrimSpace(command[:last]) + return beforeAmpersand != "" +} + +// executeSystemCommand 执行系统命令 +func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + // 获取命令 + command, ok := args["command"].(string) + if !ok { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: "错误: 缺少command参数", + }, + }, + IsError: true, + }, nil + } + + if command == "" { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: "错误: command参数不能为空", + }, + }, + IsError: true, + }, nil + } + + // 安全检查：记录执行的命令 + e.logger.Warn("执行系统命令", + zap.String("command", command), + ) + + command = PrepareShellCommandForExecute(command) + + // 获取shell类型（可选，默认为sh） + shell := "sh" + if s, ok := args["shell"].(string); ok && s != "" { + shell = s + } + + // 获取工作目录（可选） + workDir := "" + if wd, ok := args["workdir"].(string); ok && wd != "" { + workDir = wd + } + + // 检测是否为后台命令（包含 & 符号，但不在引号内） + isBackground := IsBackgroundShellCommand(command) + + // 构建命令 + var cmd *exec.Cmd + if workDir != "" { + cmd = exec.CommandContext(ctx, shell, "-c", command) + cmd.Dir = workDir + } else { + cmd = exec.CommandContext(ctx, shell, "-c", command) + } + ConfigureShellCmdForAgentExecute(cmd) + + // 执行命令 + e.logger.Info("执行系统命令", + zap.String("command", command), + zap.String("shell", shell), + zap.String("workdir", workDir), + zap.Bool("isBackground", isBackground), + ) + + // 如果是后台命令，使用特殊处理来获取实际的后台进程PID + if isBackground { + // 移除命令末尾的 & 符号 + commandWithoutAmpersand := strings.TrimSuffix(strings.TrimSpace(command), "&") + commandWithoutAmpersand = strings.TrimSpace(commandWithoutAmpersand) + + // 构建新命令：后台作业重定向标准流后 echo $pid（与 RedirectBackgroundJobStdio 一致）。 + pidCommand := RedirectBackgroundJobStdio(commandWithoutAmpersand+" &") + " pid=$!; echo $pid" + + // 创建新命令来获取PID + var pidCmd *exec.Cmd + if workDir != "" { + pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand) + pidCmd.Dir = workDir + } else { + pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand) + } + ConfigureShellCmdForAgentExecute(pidCmd) + + // 获取stdout管道 + stdout, err := pidCmd.StdoutPipe() + if err != nil { + e.logger.Error("创建stdout管道失败", + zap.String("command", command), + zap.Error(err), + ) + // 如果创建管道失败，使用shell进程的PID作为fallback + if err := pidCmd.Start(); err != nil { + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("后台命令启动失败: %v", err), + }, + }, + IsError: true, + }, nil + } + pid := pidCmd.Process.Pid + go pidCmd.Wait() // 在后台等待，避免僵尸进程 + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程ID: %d (可能不准确，获取PID失败)\n\n注意: 后台进程将继续运行，不会等待其完成。", command, pid), + }, + }, + IsError: false, + }, nil + } + + // 启动命令 + if err := pidCmd.Start(); err != nil { + stdout.Close() + e.logger.Error("后台命令启动失败", + zap.String("command", command), + zap.Error(err), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("后台命令启动失败: %v", err), + }, + }, + IsError: true, + }, nil + } + + // 读取第一行输出（PID） + reader := bufio.NewReader(stdout) + pidLine, err := reader.ReadString('\n') + stdout.Close() + + var actualPid int + if err != nil && err != io.EOF { + e.logger.Warn("读取后台进程PID失败", + zap.String("command", command), + zap.Error(err), + ) + // 如果读取失败，使用shell进程的PID + actualPid = pidCmd.Process.Pid + } else { + // 解析PID + pidStr := strings.TrimSpace(pidLine) + if parsedPid, err := strconv.Atoi(pidStr); err == nil { + actualPid = parsedPid + } else { + e.logger.Warn("解析后台进程PID失败", + zap.String("command", command), + zap.String("pidLine", pidStr), + zap.Error(err), + ) + // 如果解析失败，使用shell进程的PID + actualPid = pidCmd.Process.Pid + } + } + + // 在goroutine中等待shell进程，避免僵尸进程 + go func() { + if err := pidCmd.Wait(); err != nil { + e.logger.Debug("后台命令shell进程执行完成", + zap.String("command", command), + zap.Error(err), + ) + } + }() + + e.logger.Info("后台命令已启动", + zap.String("command", command), + zap.Int("actualPid", actualPid), + ) + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程ID: %d\n\n注意: 后台进程将继续运行，不会等待其完成。", command, actualPid), + }, + }, + IsError: false, + }, nil + } + + // 非后台命令：等待输出 + var output string + var err error + spill := e.spillOptsFromContext(ctx) + // 若上层提供工具输出增量回调，或当前处于 MCP execution 中，则边执行边流式读取。 + if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" { + cb = e.wrapToolOutputCallback(ctx, cb) + output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到系统命令需要 TTY，使用 PTY 重试") + cmd2 := exec.CommandContext(ctx, shell, "-c", command) + if workDir != "" { + cmd2.Dir = workDir + } + ConfigureShellCmdForAgentExecute(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, cb, e.toolOutputMaxBytes, spill) + } + } else { + output, err = combinedOutputCancellableWithLimit(ctx, cmd, e.toolOutputMaxBytes, spill) + if err != nil && shouldRetryWithPTY(output) { + e.logger.Info("检测到系统命令需要 TTY，使用 PTY 重试") + cmd2 := exec.CommandContext(ctx, shell, "-c", command) + if workDir != "" { + cmd2.Dir = workDir + } + ConfigureShellCmdForAgentExecute(cmd2) + output, err = runCommandWithPTY(ctx, cmd2, nil, e.toolOutputMaxBytes, spill) + } + } + if err != nil { + e.logger.Error("系统命令执行失败", + zap.String("command", command), + zap.Error(err), + zap.String("output", string(output)), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: FormatCommandFailureFromErr(err, output), + }, + }, + IsError: true, + }, nil + } + + e.logger.Info("系统命令执行成功", + zap.String("command", command), + zap.String("output_length", fmt.Sprintf("%d", len(output))), + ) + + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: string(output), + }, + }, + IsError: false, + }, nil +} + +// combinedOutputCancellable 行为对齐 cmd.CombinedOutput（stdout/stderr 写入内存缓冲）， +// 但在 ctx 取消时 terminateCmdTree 终止整棵进程树。 +// 非流式路径不使用双流管道 fan-in，避免 stderr 撑满管道缓冲区时与 stdout 互相阻塞导致死锁。 +// 无输出空闲检测由上层 agent.tool_timeout_minutes 兜底，不改变原 CombinedOutput 语义。 +func combinedOutputCancellable(ctx context.Context, cmd *exec.Cmd) (string, error) { + return combinedOutputCancellableWithLimit(ctx, cmd, 0, tooloutput.SpillOpts{}) +} + +func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxBytes int, spill tooloutput.SpillOpts) (string, error) { + var tee *tooloutput.Tee + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + stdoutBuf := newBoundedOutputCollector(maxBytes, tee) + stderrBuf := newBoundedOutputCollector(maxBytes, tee) + cmd.Stdout = stdoutBuf + cmd.Stderr = stderrBuf + + session, err := StartShellSession(cmd) + if err != nil { + return "", err + } + + done := make(chan error, 1) + go func() { + done <- session.Wait() + }() + + stopWatch := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + TerminateShellCmdSession(session) + case <-stopWatch: + } + }() + defer close(stopWatch) + + var waitErr error + select { + case waitErr = <-done: + case <-ctx.Done(): + waitErr = <-done + return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), ctx.Err() + } + return finalizeJoinedBoundedOutputs(stdoutBuf, stderrBuf, maxBytes, tee), waitErr +} + +func joinCommandOutput(stdout, stderr string) string { + if stderr == "" { + return stdout + } + if stdout == "" { + return stderr + } + return stdout + stderr +} + +type boundedOutputCollector struct { + builder strings.Builder + maxBytes int + seenBytes int + truncated bool + tee *tooloutput.Tee +} + +func newBoundedOutputCollector(maxBytes int, tee *tooloutput.Tee) *boundedOutputCollector { + return &boundedOutputCollector{maxBytes: maxBytes, tee: tee} +} + +func (b *boundedOutputCollector) Write(p []byte) (int, error) { + b.WriteStringLimited(string(p)) + return len(p), nil +} + +func (b *boundedOutputCollector) WriteStringLimited(s string) string { + if b == nil { + return "" + } + if b.tee != nil { + _, _ = b.tee.Write([]byte(s)) + } + if b.maxBytes <= 0 { + b.seenBytes += len(s) + b.builder.WriteString(s) + return s + } + b.seenBytes += len(s) + if b.builder.Len() >= b.maxBytes { + b.truncated = true + return "" + } + remaining := b.maxBytes - b.builder.Len() + if len(s) <= remaining { + b.builder.WriteString(s) + return s + } + kept := truncateStringBytes(s, remaining) + b.builder.WriteString(kept) + b.truncated = true + return kept +} + +func (b *boundedOutputCollector) String() string { + if b == nil { + return "" + } + return b.builder.String() +} + +func finalizeJoinedBoundedOutputs(stdout, stderr *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string { + if tee != nil { + _ = tee.Close() + } + truncated := (stdout != nil && stdout.truncated) || (stderr != nil && stderr.truncated) + seen := 0 + if stdout != nil { + seen += stdout.seenBytes + } + if stderr != nil { + seen += stderr.seenBytes + } + joined := joinCommandOutput( + func() string { + if stdout == nil { + return "" + } + return stdout.String() + }(), + func() string { + if stderr == nil { + return "" + } + return stderr.String() + }(), + ) + if maxBytes > 0 && !truncated && len(joined) > maxBytes { + truncated = true + seen = len(joined) + } + path := "" + if tee != nil { + path = tee.Path() + } + if truncated && maxBytes > 0 { + if path != "" { + return tooloutput.FormatPersistedFromFile(path, seen, maxBytes) + } + if len(joined) > maxBytes { + return truncateStringBytes(joined, maxBytes) + } + return joined + } + if path != "" { + _ = os.Remove(path) + } + if maxBytes > 0 && len(joined) > maxBytes { + return truncateStringBytes(joined, maxBytes) + } + return joined +} + +func finalizeBoundedOutput(collector *boundedOutputCollector, maxBytes int, tee *tooloutput.Tee) string { + if tee != nil { + _ = tee.Close() + } + if collector == nil { + return "" + } + path := "" + if tee != nil { + path = tee.Path() + } + if collector.truncated && maxBytes > 0 { + if path != "" { + return tooloutput.FormatPersistedFromFile(path, collector.seenBytes, maxBytes) + } + return truncateStringBytes(collector.String(), maxBytes) + } + if path != "" { + _ = os.Remove(path) + } + out := collector.String() + if maxBytes > 0 && len(out) > maxBytes { + return tooloutput.BoundWithSpill(out, maxBytes, tooloutput.SpillOpts{}) + } + return out +} + +func limitOutputString(s string, maxBytes int, spill tooloutput.SpillOpts) string { + if maxBytes <= 0 || len(s) <= maxBytes { + return s + } + return tooloutput.BoundWithSpill(s, maxBytes, spill) +} + +func truncateStringBytes(s string, maxBytes int) string { + if maxBytes <= 0 { + return "" + } + if len(s) <= maxBytes { + return s + } + cut := maxBytes + for cut > 0 && (s[cut]&0xC0) == 0x80 { + cut-- + } + if cut <= 0 { + return "" + } + return s[:cut] +} + +// streamCommandOutput 以“边读边回调”的方式读取命令 stdout/stderr。 +// 使用定长块读取，避免按行读取在无换行输出时永久阻塞；ctx 取消时终止进程树。 +func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, noOutputSec int, maxBytes int, spill tooloutput.SpillOpts) (string, error) { + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + return "", err + } + stderrPipe, err := cmd.StderrPipe() + if err != nil { + _ = stdoutPipe.Close() + return "", err + } + session, err := StartShellSession(cmd) + if err != nil { + _ = stdoutPipe.Close() + _ = stderrPipe.Close() + return "", err + } + + stopWatch := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + TerminateShellCmdSession(session) + case <-stopWatch: + } + }() + defer close(stopWatch) + + chunks := make(chan string, 64) + var wg sync.WaitGroup + readFn := func(r io.Reader) { + defer wg.Done() + buf := make([]byte, 8192) + for { + n, readErr := r.Read(buf) + if n > 0 { + chunks <- string(buf[:n]) + } + if readErr != nil { + return + } + } + } + + wg.Add(2) + go readFn(stdoutPipe) + go readFn(stderrPipe) + + go func() { + wg.Wait() + close(chunks) + }() + + tee := (*tooloutput.Tee)(nil) + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + outBuilder := newBoundedOutputCollector(maxBytes, tee) + var deltaBuilder strings.Builder + lastFlush := time.Now() + + flush := func() { + if deltaBuilder.Len() == 0 { + return + } + if cb != nil { + cb(deltaBuilder.String()) + } + deltaBuilder.Reset() + lastFlush = time.Now() + } + + idleWatch := NewShellInactivityWatch(noOutputSec) + if idleWatch != nil { + defer idleWatch.Stop() + } + + fireInactivity := func() { + TerminateShellCmdSession(session) + msg := ShellNoOutputTimeoutMessage(idleWatch.Sec) + msg = outBuilder.WriteStringLimited(msg) + if cb != nil { + cb(msg) + } + _ = session.Wait() + } + +chunksLoop: + for { + var idleCh <-chan struct{} + if idleWatch != nil { + idleCh = idleWatch.Expired + } + select { + case <-ctx.Done(): + TerminateShellCmdSession(session) + flush() + _ = session.Wait() + return outBuilder.String(), ctx.Err() + case <-idleCh: + fireInactivity() + return finalizeBoundedOutput(outBuilder, maxBytes, tee), fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec) + case chunk, ok := <-chunks: + if !ok { + break chunksLoop + } + if chunk != "" && idleWatch != nil { + idleWatch.Bump() + } + keptChunk := outBuilder.WriteStringLimited(chunk) + deltaBuilder.WriteString(keptChunk) + if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond { + flush() + } + } + } + flush() + + // 等待命令结束，返回最终退出状态 + waitErr := session.Wait() + return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr +} + +// applyDefaultTerminalEnv 为外部工具补齐常见的终端环境变量。 +// 注意：这不会创建 TTY，只是减少某些工具在非交互环境下的“奇怪排版/检测失败”。 +func applyDefaultTerminalEnv(cmd *exec.Cmd) { + if cmd == nil { + return + } + // 仅在未显式设置 Env 时，继承当前进程环境 + if cmd.Env == nil { + cmd.Env = os.Environ() + } + cmd.Env = ApplyNonInteractivePagerEnv(cmd.Env) + // 如果用户已设置 TERM/COLUMNS/LINES，则不覆盖 + has := func(k string) bool { + prefix := k + "=" + for _, e := range cmd.Env { + if strings.HasPrefix(e, prefix) { + return true + } + } + return false + } + if !has("TERM") { + cmd.Env = append(cmd.Env, "TERM=xterm-256color") + } + if !has("COLUMNS") { + cmd.Env = append(cmd.Env, "COLUMNS=256") + } + if !has("LINES") { + cmd.Env = append(cmd.Env, "LINES=40") + } +} + +func shouldRetryWithPTY(output string) bool { + o := strings.ToLower(output) + // autorecon / python termios 常见报错 + if strings.Contains(o, "inappropriate ioctl for device") { + return true + } + if strings.Contains(o, "termios.error") { + return true + } + // 兜底：stdin 不是 tty + if strings.Contains(o, "not a tty") { + return true + } + return false +} + +// runCommandWithPTY 为子进程分配 PTY，适配需要交互式终端的工具（如 autorecon）。 +// 若 cb != nil，将持续回调增量输出（用于 SSE）。 +func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback, maxBytes int, spill tooloutput.SpillOpts) (string, error) { + if runtime.GOOS == "windows" { + // PTY 方案为类 Unix；Windows 走原逻辑 + if cb != nil { + return streamCommandOutput(ctx, cmd, cb, 0, maxBytes, spill) + } + _ = prepareShellCmdSession(cmd) + return combinedOutputCancellableWithLimit(ctx, cmd, maxBytes, spill) + } + + _ = prepareShellCmdSession(cmd) + ptmx, err := pty.Start(cmd) + if err != nil { + return "", err + } + defer func() { _ = ptmx.Close() }() + + rootPID := 0 + if cmd.Process != nil { + rootPID = cmd.Process.Pid + } + + // ctx 取消时尽快终止子进程 + done := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + _ = ptmx.Close() // 触发读退出 + terminateProcessGroup(rootPID, cmd) + case <-done: + } + }() + defer close(done) + + tee := (*tooloutput.Tee)(nil) + if maxBytes > 0 { + tee = tooloutput.NewTee(spill) + defer func() { _ = tee.Close() }() + } + outBuilder := newBoundedOutputCollector(maxBytes, tee) + var deltaBuilder strings.Builder + lastFlush := time.Now() + flush := func() { + if cb == nil || deltaBuilder.Len() == 0 { + deltaBuilder.Reset() + lastFlush = time.Now() + return + } + cb(deltaBuilder.String()) + deltaBuilder.Reset() + lastFlush = time.Now() + } + + buf := make([]byte, 4096) + for { + n, readErr := ptmx.Read(buf) + if n > 0 { + chunk := string(buf[:n]) + // 统一换行为 \n，避免前端错位 + chunk = strings.ReplaceAll(chunk, "\r\n", "\n") + chunk = strings.ReplaceAll(chunk, "\r", "\n") + keptChunk := outBuilder.WriteStringLimited(chunk) + deltaBuilder.WriteString(keptChunk) + if deltaBuilder.Len() >= 2048 || time.Since(lastFlush) >= 200*time.Millisecond { + flush() + } + } + if readErr != nil { + break + } + } + flush() + + waitErr := cmd.Wait() + return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr +} + +// executeInternalTool 执行内部工具（不执行外部命令） +func (e *Executor) executeInternalTool(ctx context.Context, toolName string, command string, args map[string]interface{}) (*mcp.ToolResult, error) { + internalToolType := strings.TrimPrefix(command, "internal:") + e.logger.Warn("未知的内部工具", + zap.String("toolName", toolName), + zap.String("internalToolType", internalToolType), + ) + return &mcp.ToolResult{ + Content: []mcp.Content{ + { + Type: "text", + Text: fmt.Sprintf("错误: 未知的内部工具类型: %s", internalToolType), + }, + }, + IsError: true, + }, nil +} + +// buildInputSchema 构建输入模式 +func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]interface{} { + schema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + "required": []string{}, + } + + // 如果配置中定义了参数，优先使用配置中的参数定义 + if len(toolConfig.Parameters) > 0 { + properties := make(map[string]interface{}) + required := []string{} + + for _, param := range toolConfig.Parameters { + // 跳过 name 为空的参数（避免 YAML 中 name: null 或空导致非法 schema） + if strings.TrimSpace(param.Name) == "" { + e.logger.Debug("跳过无名称的参数", + zap.String("tool", toolConfig.Name), + zap.String("type", param.Type), + ) + continue + } + // 转换类型为OpenAI/JSON Schema标准类型（空类型默认为 string） + openAIType := e.convertToOpenAIType(param.Type) + + prop := map[string]interface{}{ + "type": openAIType, + "description": param.Description, + } + + // JSON Schema/OpenAI 要求 array 类型必须包含 items，否则 API 报 invalid_function_parameters + if openAIType == "array" { + itemType := strings.TrimSpace(param.ItemType) + if itemType == "" { + itemType = "string" + } + prop["items"] = map[string]interface{}{ + "type": e.convertToOpenAIType(itemType), + } + } + + // 添加默认值 + if param.Default != nil { + prop["default"] = param.Default + } + + // 添加枚举选项 + if len(param.Options) > 0 { + prop["enum"] = param.Options + } + + properties[param.Name] = prop + + // 添加到必需参数列表 + if param.Required { + required = append(required, param.Name) + } + } + + schema["properties"] = properties + schema["required"] = required + return schema + } + + // 如果没有定义参数配置，返回空schema + // 这种情况下工具可能只使用固定参数（args字段） + // 或者需要通过YAML配置文件定义参数 + e.logger.Warn("工具未定义参数配置，返回空schema", + zap.String("tool", toolConfig.Name), + ) + return schema +} + +// convertToOpenAIType 将配置中的类型转换为OpenAI/JSON Schema标准类型 +func (e *Executor) convertToOpenAIType(configType string) string { + // 空或 null 类型统一视为 string，避免非法 schema 导致工具调用失败 + if strings.TrimSpace(configType) == "" { + return "string" + } + switch configType { + case "bool": + return "boolean" + case "int", "integer": + return "number" + case "float", "double": + return "number" + case "string", "array", "object": + return configType + default: + // 默认返回原类型，但记录警告 + e.logger.Warn("未知的参数类型，使用原类型", + zap.String("type", configType), + ) + return configType + } +} + +// getExitCode 从错误中提取退出码，如果不是ExitError则返回nil +func getExitCode(err error) *int { + if err == nil { + return nil + } + if exitError, ok := err.(*exec.ExitError); ok { + if exitError.ProcessState != nil { + exitCode := exitError.ExitCode() + return &exitCode + } + } + return nil +} + +// getExitCodeValue 从错误中提取退出码值，如果不是ExitError则返回-1 +func getExitCodeValue(err error) int { + if code := getExitCode(err); code != nil { + return *code + } + return -1 +} diff --git a/internal/security/executor_test.go b/internal/security/executor_test.go new file mode 100644 index 00000000..4b62889a --- /dev/null +++ b/internal/security/executor_test.go @@ -0,0 +1,282 @@ +package security + +import ( + "context" + "os/exec" + "runtime" + "strings" + "testing" + "time" + + "cyberstrike-ai/internal/config" + "cyberstrike-ai/internal/mcp" + + "go.uber.org/zap" +) + +// setupTestExecutor 创建测试用的执行器 +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 +} + +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("错误消息应该包含'未知的内部工具类型'") + } +} + +func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T) { + executor, _ := setupTestExecutor(t) + // 子进程先向 stdout 写无换行字符再长时间 sleep；若与 echo $pid 共享管道且未重定向子进程 stdout， + // ReadString('\n') 会阻塞到子进程退出。后台包装须将子进程标准流与 PID 行分离。 + ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second) + defer cancel() + args := map[string]interface{}{ + "command": `(sh -c 'printf x; sleep 120') &`, + "shell": "sh", + } + res, err := executor.executeSystemCommand(ctx, args) + if err != nil { + t.Fatalf("executeSystemCommand: %v", err) + } + if res == nil || res.IsError { + t.Fatalf("expected success, got %+v", res) + } + txt := res.Content[0].Text + if !strings.Contains(txt, "后台命令已启动") { + t.Fatalf("unexpected body: %q", txt) + } +} + +func TestExecToolSoftWaitExposesPartialOutput(t *testing.T) { + executor, server := setupTestExecutor(t) + server.ConfigureToolWaitTimeoutSeconds(1) + mcp.RegisterExecutionControlTools(server, nil) + server.RegisterTool(mcp.Tool{Name: "exec", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) { + return executor.ExecuteTool(ctx, "exec", args) + }) + + result, executionID, err := server.CallTool(context.Background(), "exec", map[string]interface{}{ + "command": "for i in 1 2 3 4; do echo partial-$i; sleep 0.3; done; sleep 5", + "shell": "sh", + }) + if err != nil { + t.Fatalf("CallTool exec: %v", err) + } + if executionID == "" || result == nil || !result.IsError { + t.Fatalf("expected soft wait timeout, id=%q result=%#v", executionID, result) + } + + status, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{ + "execution_id": executionID, + "include_partial_output": true, + "partial_output_max_bytes": 4096, + }) + if err != nil { + t.Fatalf("get_tool_execution: %v", err) + } + body := mcp.ToolResultPlainText(status) + if !strings.Contains(body, `"status": "running"`) { + t.Fatalf("expected running execution, got: %s", body) + } + if !strings.Contains(body, "partial-") || !strings.Contains(body, "partial_output") { + t.Fatalf("expected partial output in execution status, got: %s", body) + } + server.CancelToolExecution(executionID) +} + +func TestExecuteSystemCommand_FailureFormat(t *testing.T) { + executor, _ := setupTestExecutor(t) + res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{ + "command": "echo fail-msg >&2; exit 7", + "shell": "sh", + }) + if err != nil { + t.Fatalf("executeSystemCommand: %v", err) + } + if res == nil || !res.IsError { + t.Fatalf("expected IsError, got %+v", res) + } + text := res.Content[0].Text + if text != FormatCommandFailureResult(7, "fail-msg\n") && text != FormatCommandFailureResult(7, "fail-msg") { + t.Fatalf("unexpected failure text: %q", text) + } + if !strings.Contains(text, "exit status 7") || !strings.Contains(text, "fail-msg") { + t.Fatalf("unexpected failure text: %q", text) + } +} + +func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) { + executor, _ := setupTestExecutor(t) + spillRoot := t.TempDir() + executor.SetToolOutputMaxBytes(200) + executor.SetToolOutputSpillRoot(spillRoot) + ctx := mcp.WithMCPConversationID(context.Background(), "exec-spill") + res, err := executor.executeSystemCommand(ctx, map[string]interface{}{ + "command": "i=0; while [ $i -lt 2000 ]; do printf 0123456789; i=$((i+1)); done", + "shell": "sh", + }) + if err != nil { + t.Fatalf("executeSystemCommand: %v", err) + } + if res == nil || res.IsError { + t.Fatalf("expected success, got %+v", res) + } + text := res.Content[0].Text + if !strings.Contains(text, "") || !strings.Contains(text, "Full output saved to:") { + t.Fatalf("missing persisted-output notice: %q", text) + } + if len(text) > 200 { + t.Fatalf("output exceeded hard limit: len=%d text=%q", len(text), text) + } + if strings.Contains(text, strings.Repeat("0123456789", 20)) { + t.Fatalf("output kept too much data: len=%d", len(text)) + } +} + +func TestExecuteSystemCommand_StreamingOutputIsSourceLimited(t *testing.T) { + executor, _ := setupTestExecutor(t) + spillRoot := t.TempDir() + executor.SetToolOutputMaxBytes(200) + executor.SetToolOutputSpillRoot(spillRoot) + var streamed strings.Builder + ctx := context.WithValue(context.Background(), ToolOutputCallbackCtxKey, ToolOutputCallback(func(chunk string) { + streamed.WriteString(chunk) + })) + ctx = mcp.WithMCPConversationID(ctx, "exec-stream-spill") + res, err := executor.executeSystemCommand(ctx, map[string]interface{}{ + "command": "i=0; while [ $i -lt 2000 ]; do printf abcdefghij; i=$((i+1)); done", + "shell": "sh", + }) + if err != nil { + t.Fatalf("executeSystemCommand: %v", err) + } + text := res.Content[0].Text + if !strings.Contains(text, "") { + t.Fatalf("missing persisted-output notice: %q", text) + } + if len(text) > 200 { + t.Fatalf("returned output exceeded hard limit: len=%d text=%q", len(text), text) + } + // SSE only streams the bounded prefix; final agent-facing body is the spill notice. + if len(streamed.String()) > 200 { + t.Fatalf("streamed prefix exceeded hard limit: len=%d", len(streamed.String())) + } + if streamed.Len() == 0 { + t.Fatal("expected some streamed prefix before truncation") + } + if strings.Contains(text, strings.Repeat("abcdefghij", 50)) { + t.Fatalf("returned output kept too much raw data: len=%d", len(text)) + } +} + +func TestBuildCommandArgs_NmapSkipsEmptyOptionalFlags(t *testing.T) { + pos1 := 1 + executor, _ := setupTestExecutor(t) + toolConfig := &config.ToolConfig{ + Name: "nmap", + Command: "nmap", + Args: []string{"-sT", "-sV", "-sC"}, + Parameters: []config.ParameterConfig{ + {Name: "target", Type: "string", Required: true, Position: &pos1, Format: "positional"}, + {Name: "ports", Type: "string", Flag: "-p", Format: "flag"}, + {Name: "timing", Type: "string", Template: "-T{value}", Format: "template"}, + {Name: "nse_scripts", Type: "string", Flag: "--script", Format: "flag"}, + {Name: "os_detection", Type: "bool", Flag: "-O", Format: "flag", Default: false}, + {Name: "aggressive", Type: "bool", Flag: "-A", Format: "flag", Default: false}, + {Name: "scan_type", Type: "string", Format: "template", Template: "{value}"}, + {Name: "additional_args", Type: "string", Format: "positional"}, + }, + } + + args := map[string]interface{}{ + "target": "110.52.223.114", + "ports": "21, 22, 80, 443", + "timing": "4", + "nse_scripts": "", + "scan_type": "", + "os_detection": false, + "aggressive": false, + "additional_args": "-Pn", + } + + cmdArgs := executor.buildCommandArgs("nmap", toolConfig, args) + joined := strings.Join(cmdArgs, " ") + + if strings.Contains(joined, "--script") { + t.Fatalf("empty nse_scripts must not emit --script, got: %v", cmdArgs) + } + if !strings.Contains(joined, "110.52.223.114") { + t.Fatalf("target missing from args: %v", cmdArgs) + } + // target 应出现在 -Pn 之前，避免被误当作 --script 的参数 + pnIdx := indexOf(cmdArgs, "-Pn") + targetIdx := indexOf(cmdArgs, "110.52.223.114") + if pnIdx < 0 || targetIdx < 0 || targetIdx >= pnIdx { + t.Fatalf("expected target before -Pn, got: %v", cmdArgs) + } +} + +func indexOf(slice []string, s string) int { + for i, v := range slice { + if v == s { + return i + } + } + return -1 +} + +// TestCombinedOutputCancellable_ContextCancelKillsTree 验证 ctx 取消时能在数秒内结束（杀进程组，非挂死）。 +func TestCombinedOutputCancellable_ContextCancelKillsTree(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("unix process group kill") + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cmd := exec.CommandContext(ctx, "sh", "-c", "sleep 300") + ConfigureShellCmdForAgentExecute(cmd) + + done := make(chan error, 1) + go func() { + _, err := combinedOutputCancellable(ctx, cmd) + done <- err + }() + + time.Sleep(150 * time.Millisecond) + cancel() + + select { + case err := <-done: + if err == nil { + t.Fatal("expected context cancel error") + } + case <-time.After(5 * time.Second): + t.Fatal("combinedOutputCancellable did not return within 5s after context cancel") + } +} diff --git a/internal/security/password.go b/internal/security/password.go new file mode 100644 index 00000000..3eb356de --- /dev/null +++ b/internal/security/password.go @@ -0,0 +1,24 @@ +package security + +import ( + "crypto/rand" + "encoding/base64" +) + +// GenerateStrongPassword returns a URL-safe random password of the given length. +func GenerateStrongPassword(length int) (string, error) { + if length <= 0 { + length = 24 + } + + randomBytes := make([]byte, length) + if _, err := rand.Read(randomBytes); err != nil { + return "", err + } + + password := base64.RawURLEncoding.EncodeToString(randomBytes) + if len(password) > length { + password = password[:length] + } + return password, nil +} diff --git a/internal/security/procattr_unix.go b/internal/security/procattr_unix.go new file mode 100644 index 00000000..8f516ec8 --- /dev/null +++ b/internal/security/procattr_unix.go @@ -0,0 +1,41 @@ +//go:build !windows + +package security + +import ( + "os/exec" + "syscall" +) + +// prepareShellCmdSession 让 shell 子进程在独立会话中运行，便于超时/取消时整组 SIGKILL（含子进程）。 +func prepareShellCmdSession(cmd *exec.Cmd) error { + if cmd == nil { + return nil + } + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{} + } + cmd.SysProcAttr.Setsid = true + return nil +} + +// terminateProcessGroup 对 rootPID 对应进程组发 SIGKILL；rootPID 为 0 时回退到 cmd.Process.Pid。 +func terminateProcessGroup(rootPID int, cmd *exec.Cmd) { + pid := rootPID + if pid <= 0 && cmd != nil && cmd.Process != nil { + pid = cmd.Process.Pid + } + if pid <= 0 { + return + } + if err := syscall.Kill(-pid, syscall.SIGKILL); err != nil { + if cmd != nil && cmd.Process != nil { + _ = cmd.Process.Kill() + } + } +} + +// terminateCmdTree 尽力终止 cmd 及其进程组（Unix 下 Setsid 后 PGID == 首进程 PID）。 +func terminateCmdTree(cmd *exec.Cmd) { + terminateProcessGroup(0, cmd) +} diff --git a/internal/security/procattr_windows.go b/internal/security/procattr_windows.go new file mode 100644 index 00000000..af7da8c1 --- /dev/null +++ b/internal/security/procattr_windows.go @@ -0,0 +1,43 @@ +//go:build windows + +package security + +import ( + "os/exec" + "strconv" + "syscall" +) + +func prepareShellCmdSession(cmd *exec.Cmd) error { + if cmd == nil { + return nil + } + // 独立进程组，便于 taskkill /T 终止整棵子进程树。 + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{} + } + cmd.SysProcAttr.CreationFlags = syscall.CREATE_NEW_PROCESS_GROUP + return nil +} + +// terminateProcessGroup 使用 taskkill /F /T 终止进程及其子进程；rootPID 为 0 时回退到 cmd.Process.Pid。 +func terminateProcessGroup(rootPID int, cmd *exec.Cmd) { + pid := rootPID + if pid <= 0 && cmd != nil && cmd.Process != nil { + pid = cmd.Process.Pid + } + if pid <= 0 { + return + } + tk := exec.Command("taskkill", "/F", "/T", "/PID", strconv.Itoa(pid)) + if err := tk.Run(); err != nil { + if cmd != nil && cmd.Process != nil { + _ = cmd.Process.Kill() + } + } +} + +// terminateCmdTree 使用 taskkill /F /T 终止进程及其子进程（Windows 上 Process.Kill 无法保证杀掉 python 等孙进程）。 +func terminateCmdTree(cmd *exec.Cmd) { + terminateProcessGroup(0, cmd) +} diff --git a/internal/security/ratelimit.go b/internal/security/ratelimit.go new file mode 100644 index 00000000..71795710 --- /dev/null +++ b/internal/security/ratelimit.go @@ -0,0 +1,81 @@ +package security + +import ( + "net/http" + "sync" + "time" + + "github.com/gin-gonic/gin" +) + +// rateLimitEntry 记录某个 IP 的请求窗口信息 +type rateLimitEntry struct { + count int + windowAt time.Time +} + +// RateLimiter 基于 IP 的滑动窗口速率限制器 +type RateLimiter struct { + mu sync.Mutex + entries map[string]*rateLimitEntry + limit int // 窗口内允许的最大请求数 + window time.Duration // 窗口时长 +} + +// NewRateLimiter 创建速率限制器 +func NewRateLimiter(limit int, window time.Duration) *RateLimiter { + rl := &RateLimiter{ + entries: make(map[string]*rateLimitEntry), + limit: limit, + window: window, + } + // 后台定期清理过期条目，防止内存泄漏 + go rl.cleanup() + return rl +} + +// cleanup 每分钟清理一次过期条目 +func (rl *RateLimiter) cleanup() { + ticker := time.NewTicker(1 * time.Minute) + defer ticker.Stop() + for range ticker.C { + rl.mu.Lock() + now := time.Now() + for ip, entry := range rl.entries { + if now.Sub(entry.windowAt) > rl.window { + delete(rl.entries, ip) + } + } + rl.mu.Unlock() + } +} + +// allow 检查指定 IP 是否允许通过 +func (rl *RateLimiter) allow(ip string) bool { + rl.mu.Lock() + defer rl.mu.Unlock() + + now := time.Now() + entry, ok := rl.entries[ip] + if !ok || now.Sub(entry.windowAt) > rl.window { + rl.entries[ip] = &rateLimitEntry{count: 1, windowAt: now} + return true + } + + entry.count++ + return entry.count <= rl.limit +} + +// RateLimitMiddleware 返回 Gin 中间件，对超限请求返回 429 +func RateLimitMiddleware(rl *RateLimiter) gin.HandlerFunc { + return func(c *gin.Context) { + ip := c.ClientIP() + if !rl.allow(ip) { + c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{ + "error": "rate limit exceeded, please try again later", + }) + return + } + c.Next() + } +} diff --git a/internal/security/rbac.go b/internal/security/rbac.go new file mode 100644 index 00000000..1c5b766f --- /dev/null +++ b/internal/security/rbac.go @@ -0,0 +1,119 @@ +package security + +import ( + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "fmt" + "strings" + + "golang.org/x/crypto/bcrypt" +) + +// Platform permissions use module:action naming. They are intentionally +// separate from AI testing roles under roles/. +var PermissionCatalog = map[string]string{ + "auth:self": "Manage own session and password", + "dashboard:read": "View dashboard summaries", + "chat:read": "View conversations", + "chat:write": "Create and update conversations", + "chat:delete": "Delete conversations and turns", + "agent:execute": "Run AI agents and workflows", + "agent:local-execute": "Use local filesystem, shell, and configured command tools from an agent", + "hitl:read": "View HITL queues and logs", + "hitl:write": "Approve, dismiss, and configure HITL", + "tasks:read": "View task queues", + "tasks:write": "Create and run task queues", + "tasks:delete": "Delete task queues", + "project:read": "View projects and project facts", + "project:write": "Create and update projects and facts", + "project:delete": "Delete projects and facts", + "vulnerability:read": "View vulnerabilities", + "vulnerability:write": "Create and update vulnerabilities", + "vulnerability:delete": "Delete vulnerabilities", + "asset:read": "View managed assets and asset summaries", + "asset:write": "Create, import, and update assets", + "asset:delete": "Delete managed assets", + "webshell:read": "View WebShell connections", + "webshell:write": "Manage and use WebShell connections", + "webshell:delete": "Delete WebShell connections", + "c2:read": "View C2 listeners, sessions, tasks, events, and profiles", + "c2:write": "Operate C2 listeners, sessions, tasks, payloads, files, and profiles", + "c2:delete": "Delete C2 objects", + "mcp:read": "View MCP status and external MCP configuration", + "mcp:execute": "Invoke the authenticated MCP endpoint", + "mcp:external:execute": "Invoke tools exposed by configured external MCP servers", + "mcp:write": "Manage external MCP server configuration and lifecycle", + "knowledge:read": "View knowledge base and retrieval logs", + "knowledge:write": "Create, update, index, and scan knowledge base", + "knowledge:delete": "Delete knowledge items and retrieval logs", + "skills:read": "View skills and skill stats", + "skills:write": "Create and update skills", + "skills:delete": "Delete skills and stats", + "agents:read": "View markdown agents", + "agents:write": "Create and update markdown agents", + "agents:delete": "Delete markdown agents", + "roles:read": "View AI testing roles", + "roles:write": "Create and update AI testing roles", + "roles:delete": "Delete AI testing roles", + "workflow:read": "View workflow definitions and runs", + "workflow:execute": "Validate, dry-run, and resume authorized workflow runs", + "workflow:write": "Create and update workflow definitions", + "workflow:delete": "Delete workflows", + "config:read": "View system configuration", + "config:write": "Update and apply system configuration", + "terminal:execute": "Run terminal commands", + "audit:read": "View and export audit logs", + "audit:delete": "Delete audit logs", + "rbac:read": "View users, platform roles, permissions, and assignments", + "rbac:write": "Manage users, platform roles, permissions, and assignments", + "notification:read": "View notifications", + "notification:write": "Mark notifications as read", + "robot:read": "View robot binding status", + "robot:write": "Manage robot bindings and test robot callbacks", + "files:read": "View chat uploads", + "files:write": "Upload, edit, and rename chat files", + "files:delete": "Delete chat files", + "attackchain:read": "View attack chains", + "attackchain:write": "Regenerate attack chains", + "fofa:execute": "Run FOFA searches and query parsing", + "openapi:read": "Read OpenAPI aggregation results", + "group:read": "View conversation groups", + "group:write": "Create and update conversation groups", + "group:delete": "Delete conversation groups", + "monitor:read": "View execution monitor", + "monitor:write": "Cancel monitor executions", + "monitor:delete": "Delete monitor executions", +} + +func HashPassword(password string) (string, error) { + password = strings.TrimSpace(password) + if password == "" { + return "", fmt.Errorf("password is empty") + } + hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost) + if err != nil { + return "", err + } + return string(hash), nil +} + +func VerifyPasswordHash(password, encoded string) bool { + if strings.HasPrefix(encoded, "$2a$") || strings.HasPrefix(encoded, "$2b$") || strings.HasPrefix(encoded, "$2y$") { + return bcrypt.CompareHashAndPassword([]byte(encoded), []byte(strings.TrimSpace(password))) == nil + } + parts := strings.Split(encoded, "$") + if len(parts) != 3 || parts[0] != "sha256" { + return false + } + salt, err := hex.DecodeString(parts[1]) + if err != nil { + return false + } + expected, err := hex.DecodeString(parts[2]) + if err != nil { + return false + } + sum := sha256.Sum256(append(salt, []byte(strings.TrimSpace(password))...)) + return subtle.ConstantTimeCompare(sum[:], expected) == 1 +} diff --git a/internal/security/rbac_middleware.go b/internal/security/rbac_middleware.go new file mode 100644 index 00000000..6a719612 --- /dev/null +++ b/internal/security/rbac_middleware.go @@ -0,0 +1,282 @@ +package security + +import ( + "net/http" + "strings" + + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" +) + +// RBACMiddleware maps protected API routes to platform permissions. It keeps +// enforcement centralized so route declarations stay readable. +func RBACMiddleware(db *database.DB) gin.HandlerFunc { + return RBACMiddlewareWithDenyHook(db, nil) +} + +type RBACDenyHook func(c *gin.Context, reason, permission string) + +func RBACMiddlewareWithDenyHook(db *database.DB, denyHook RBACDenyHook) gin.HandlerFunc { + return func(c *gin.Context) { + permission := permissionForRequest(c.Request.Method, c.FullPath()) + if permission == "" { + if denyHook != nil { + denyHook(c, "unmapped_route", "") + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "未配置访问权限", + }) + return + } + permission, allowed := sessionHasRoutePermission(c, c.Request.Method, c.FullPath()) + if !allowed { + if denyHook != nil { + denyHook(c, "permission_denied", permission) + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{ + "error": "权限不足", + "permission": permission, + }) + return + } + // Bind the scope of the permission authorizing this request. Scope is + // permission-specific; using the user's broadest role scope here would + // let an unrelated global read role widen a write permission. + session, _ := CurrentSession(c) + session.Scope = session.ScopeFor(permission) + c.Set(ContextSessionKey, session) + c.Set(ContextUserScopeKey, session.Scope) + if db != nil && !resourceAllowed(c, db) { + if denyHook != nil { + denyHook(c, "resource_denied", permission) + } + c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"}) + return + } + c.Next() + } +} + +func sessionHasRoutePermission(c *gin.Context, method, fullPath string) (string, bool) { + path := strings.TrimPrefix(fullPath, "/api") + if alts := permissionAlternativesForRequest(method, path); len(alts) > 0 { + for _, permission := range alts { + if SessionHasPermission(c, permission) { + return permission, true + } + } + return alts[0], false + } + permission := permissionForRequest(method, fullPath) + if permission == "" { + return "", false + } + return permission, SessionHasPermission(c, permission) +} + +func permissionAlternativesForRequest(method, path string) []string { + if method != http.MethodGet && method != http.MethodHead { + return nil + } + switch { + case strings.HasPrefix(path, "/config/tools"): + // MCP 管理页只需 mcp:read；系统设置页仍可用 config:read 访问同一接口。 + return []string{"mcp:read", "config:read"} + default: + return nil + } +} + +func permissionForRequest(method, fullPath string) string { + path := strings.TrimPrefix(fullPath, "/api") + switch { + case path == "/rbac/me": + return "auth:self" + case path == "/rbac/resources": + // The picker enumerates resource names and IDs and is only needed by + // administrators who can actually create assignments. + return "rbac:write" + case strings.HasPrefix(path, "/rbac"): + if method == http.MethodGet { + return "rbac:read" + } + return "rbac:write" + case strings.HasPrefix(path, "/robot/wechat/status"): + return "robot:read" + case strings.HasPrefix(path, "/robot"): + return "robot:write" + case strings.HasPrefix(path, "/eino-agent"), strings.HasPrefix(path, "/multi-agent"): + if strings.Contains(path, "/markdown-agents") { + return crudPermission(method, "agents") + } + return "agent:execute" + case strings.HasPrefix(path, "/hitl"): + if method == http.MethodGet || method == http.MethodHead { + return "hitl:read" + } + return "hitl:write" + case strings.HasPrefix(path, "/agent-loop"), strings.HasPrefix(path, "/batch-tasks"): + return crudPermission(method, "tasks") + case strings.HasPrefix(path, "/conversations"), strings.HasPrefix(path, "/messages"), strings.HasPrefix(path, "/process-details"): + return crudPermission(method, "chat") + case strings.HasPrefix(path, "/groups"): + return crudPermission(method, "group") + case strings.HasPrefix(path, "/monitor"): + return crudPermission(method, "monitor") + case strings.HasPrefix(path, "/notifications"): + if method == http.MethodGet { + return "notification:read" + } + return "notification:write" + case strings.HasPrefix(path, "/config"): + return crudPermission(method, "config") + case strings.HasPrefix(path, "/terminal"): + return "terminal:execute" + case strings.HasPrefix(path, "/audit"): + return crudPermission(method, "audit") + case path == "/mcp": + return "mcp:execute" + case strings.HasPrefix(path, "/external-mcp"): + if method == http.MethodGet || method == http.MethodHead { + return "mcp:read" + } + return "mcp:write" + case strings.HasPrefix(path, "/attack-chain"): + return crudPermission(method, "attackchain") + case strings.HasPrefix(path, "/knowledge"): + if path == "/knowledge/search" { + return "knowledge:read" + } + return crudPermission(method, "knowledge") + case strings.HasPrefix(path, "/vulnerabilities"): + return crudPermission(method, "vulnerability") + case path == "/assets/batch-delete", path == "/assets/merge": + return "asset:delete" + case strings.HasPrefix(path, "/assets"): + return crudPermission(method, "asset") + case strings.HasPrefix(path, "/vulnerability-alerts"): + // This endpoint only changes the authenticated user's own preference. + return "vulnerability:read" + case strings.HasPrefix(path, "/projects"): + return crudPermission(method, "project") + case strings.HasPrefix(path, "/webshell"): + return crudPermission(method, "webshell") + case strings.HasPrefix(path, "/c2"): + return crudPermission(method, "c2") + case strings.HasPrefix(path, "/chat-uploads"): + return crudPermission(method, "files") + case strings.HasPrefix(path, "/roles"): + return crudPermission(method, "roles") + case path == "/workflows/:id/package": + return "workflow:read" + case strings.HasPrefix(path, "/workflow-package-inspections"), strings.HasPrefix(path, "/workflow-package-imports"): + return "workflow:write" + case path == "/workflows/generate-draft": + return "workflow:write" + case strings.HasPrefix(path, "/workflows"): + if path == "/workflows/validate" || path == "/workflows/dry-run" || strings.HasSuffix(path, "/resume") { + return "workflow:execute" + } + return crudPermission(method, "workflow") + case strings.HasPrefix(path, "/skills"): + return crudPermission(method, "skills") + case strings.HasPrefix(path, "/openapi"): + return "openapi:read" + case strings.HasPrefix(path, "/fofa"): + return "fofa:execute" + default: + return "" + } +} + +func crudPermission(method, module string) string { + switch method { + case http.MethodGet, http.MethodHead: + return module + ":read" + case http.MethodDelete: + return module + ":delete" + default: + return module + ":write" + } +} + +func resourceAllowed(c *gin.Context, db *database.DB) bool { + session, ok := CurrentSession(c) + if !ok || session.Scope == database.RBACScopeAll { + return ok + } + path := strings.TrimPrefix(c.FullPath(), "/api") + switch { + case path == "/monitor/stats", path == "/monitor/calls-timeline": + // These APIs currently operate on process-global state. Until every MCP + // invocation and persisted execution record carries an immutable owner, + // allowing an assigned/own-scoped session would be a cross-user bypass. + return session.Scope == database.RBACScopeAll + case strings.HasPrefix(path, "/c2/profiles") && c.Request.Method != http.MethodGet: + return session.Scope == database.RBACScopeAll + case (strings.HasPrefix(path, "/hitl/tool-whitelist") || strings.HasPrefix(path, "/hitl/default-reviewer") || strings.HasPrefix(path, "/hitl/audit-strategy")) && c.Request.Method != http.MethodGet: + return session.Scope == database.RBACScopeAll + case isMutationMethod(c.Request.Method) && isProcessGlobalMutationPath(path): + // These definitions/configurations are shared by every user and do not + // carry owners. A module write permission with assigned/own scope must + // not silently become a process-global administrative capability. + return session.Scope == database.RBACScopeAll + case strings.HasPrefix(path, "/projects/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "project", c.Param("id")) + case strings.HasPrefix(path, "/conversations/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "conversation", c.Param("id")) + case strings.HasPrefix(path, "/messages/:id/process-details"): + return db.UserCanAccessMessage(session.UserID, session.Scope, c.Param("id")) + case strings.HasPrefix(path, "/process-details/:id"): + return db.UserCanAccessProcessDetail(session.UserID, session.Scope, c.Param("id")) + case strings.HasPrefix(path, "/attack-chain/:conversationId"): + return db.UserCanAccessResource(session.UserID, session.Scope, "conversation", c.Param("conversationId")) + case strings.HasPrefix(path, "/webshell/connections/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "webshell", c.Param("id")) + case strings.HasPrefix(path, "/batch-tasks/:queueId"): + return db.UserCanAccessResource(session.UserID, session.Scope, "batch_task", c.Param("queueId")) + case strings.HasPrefix(path, "/vulnerabilities/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "vulnerability", c.Param("id")) + case strings.HasPrefix(path, "/assets/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "asset", c.Param("id")) + case strings.HasPrefix(path, "/c2/listeners/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "c2_listener", c.Param("id")) + case strings.HasPrefix(path, "/c2/sessions/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "c2_session", c.Param("id")) + case strings.HasPrefix(path, "/c2/tasks/:id"): + return db.UserCanAccessResource(session.UserID, session.Scope, "c2_task", c.Param("id")) + default: + return true + } +} + +func isMutationMethod(method string) bool { + switch method { + case http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete: + return true + default: + return false + } +} + +func isProcessGlobalMutationPath(path string) bool { + if strings.HasPrefix(path, "/roles") || strings.HasPrefix(path, "/skills") || + strings.HasPrefix(path, "/external-mcp") || strings.HasPrefix(path, "/robot") { + return true + } + if strings.HasPrefix(path, "/workflows") { + // Workflow runs inherit conversation access; definitions are global. + return !strings.HasPrefix(path, "/workflows/runs/") && path != "/workflows/validate" && path != "/workflows/dry-run" && path != "/workflows/generate-draft" + } + if strings.HasPrefix(path, "/workflow-package-inspections") || strings.HasPrefix(path, "/workflow-package-imports") { + return true + } + if strings.HasPrefix(path, "/knowledge") { + return path != "/knowledge/search" + } + if strings.HasPrefix(path, "/eino-agent/markdown-agents") || strings.HasPrefix(path, "/multi-agent/markdown-agents") { + return true + } + return false +} diff --git a/internal/security/rbac_middleware_test.go b/internal/security/rbac_middleware_test.go new file mode 100644 index 00000000..6d1c54a5 --- /dev/null +++ b/internal/security/rbac_middleware_test.go @@ -0,0 +1,269 @@ +package security + +import ( + "net/http" + "net/http/httptest" + "testing" + + "cyberstrike-ai/internal/database" + + "github.com/gin-gonic/gin" +) + +func TestRBACMiddlewareUsesMatchedFullPath(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "u1", + Username: "operator", + Permissions: map[string]bool{"project:read": true}, + Scope: database.RBACScopeAll, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.GET("/api/projects/:id", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/projects/p1", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func TestRBACMiddlewareRejectsMissingPermission(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "u1", + Username: "viewer", + Permissions: map[string]bool{"project:read": true}, + Scope: database.RBACScopeAll, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.POST("/api/projects", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/projects", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden) + } +} + +func TestRBACMiddlewareRejectsUnmappedProtectedRoute(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "u1", + Username: "admin", + Permissions: allPermissions(), + Scope: database.RBACScopeAll, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.GET("/api/new-module", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/new-module", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden) + } +} + +func TestRBACMiddlewareMapsOpenAPISpec(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "u1", + Username: "viewer", + Permissions: map[string]bool{"openapi:read": true}, + Scope: database.RBACScopeAll, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.GET("/api/openapi/spec", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"ok": true}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/openapi/spec", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func TestRBACResourcePickerRequiresWritePermission(t *testing.T) { + if got := permissionForRequest(http.MethodGet, "/api/rbac/resources"); got != "rbac:write" { + t.Fatalf("picker permission = %q, want rbac:write", got) + } + if got := permissionForRequest(http.MethodGet, "/api/rbac/resource-assignments"); got != "rbac:read" { + t.Fatalf("assignment list permission = %q, want rbac:read", got) + } +} + +func TestMCPInvocationPermissionIsSeparateFromMCPAdministration(t *testing.T) { + if got := permissionForRequest(http.MethodPost, "/api/mcp"); got != "mcp:execute" { + t.Fatalf("MCP invocation permission = %q, want mcp:execute", got) + } + if got := permissionForRequest(http.MethodPut, "/api/external-mcp/example"); got != "mcp:write" { + t.Fatalf("external MCP admin permission = %q, want mcp:write", got) + } +} + +func TestConfigToolsReadAllowsMCPReadWithoutConfigRead(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "viewer", + Username: "viewer", + Permissions: map[string]bool{"mcp:read": true}, + Scope: database.RBACScopeAssigned, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + router.GET("/api/config/tools", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"tools": []any{}}) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/config/tools", nil) + router.ServeHTTP(w, req) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String()) + } +} + +func TestWorkflowRunPermissionIsSeparateFromDefinitionManagement(t *testing.T) { + if got := permissionForRequest(http.MethodPost, "/api/workflows/runs/run-1/resume"); got != "workflow:execute" { + t.Fatalf("resume permission = %q, want workflow:execute", got) + } + if got := permissionForRequest(http.MethodPost, "/api/workflows/generate-draft"); got != "workflow:write" { + t.Fatalf("generate draft permission = %q, want workflow:write", got) + } + if got := permissionForRequest(http.MethodPut, "/api/workflows/workflow-1"); got != "workflow:write" { + t.Fatalf("definition permission = %q, want workflow:write", got) + } + if isProcessGlobalMutationPath("/workflows/generate-draft") { + t.Fatalf("generate draft should not be treated as a process-global mutation") + } +} + +func TestRBACDenyHookReceivesDeniedDecision(t *testing.T) { + gin.SetMode(gin.TestMode) + called := false + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{UserID: "viewer", Permissions: map[string]bool{"project:read": true}, Scope: database.RBACScopeAssigned}) + c.Next() + }) + router.Use(RBACMiddlewareWithDenyHook(nil, func(_ *gin.Context, reason, permission string) { + called = reason == "permission_denied" && permission == "project:write" + })) + router.POST("/api/projects", func(c *gin.Context) { c.Status(http.StatusNoContent) }) + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/projects", nil)) + if w.Code != http.StatusForbidden || !called { + t.Fatalf("denial = status %d, hook called %v", w.Code, called) + } +} + +func TestRBACMiddlewareBindsPermissionSpecificScope(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "mixed", Scope: database.RBACScopeAll, + Permissions: map[string]bool{"project:read": true, "project:write": true}, + PermissionScopes: map[string]string{"project:read": database.RBACScopeAll, "project:write": database.RBACScopeOwn}, + }) + c.Next() + }) + router.Use(RBACMiddleware(nil)) + handler := func(c *gin.Context) { + session, _ := CurrentSession(c) + c.String(http.StatusOK, session.Scope) + } + router.GET("/api/projects/:id", handler) + router.PUT("/api/projects/:id", handler) + + for _, tc := range []struct{ method, want string }{ + {http.MethodGet, database.RBACScopeAll}, + {http.MethodPut, database.RBACScopeOwn}, + } { + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(tc.method, "/api/projects/p1", nil)) + if w.Code != http.StatusOK || w.Body.String() != tc.want { + t.Fatalf("%s scope response = %d/%q, want 200/%q", tc.method, w.Code, w.Body.String(), tc.want) + } + } +} + +func TestRBACMiddlewareRejectsAssignedScopeForGlobalMonitorAggregates(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, tc := range []struct { + method string + path string + permission string + }{ + {method: http.MethodGet, path: "/api/monitor/stats", permission: "monitor:read"}, + } { + t.Run(tc.path, func(t *testing.T) { + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{ + UserID: "assigned-user", Permissions: map[string]bool{tc.permission: true}, Scope: database.RBACScopeAssigned, + }) + c.Next() + }) + router.Use(RBACMiddleware(&database.DB{})) + router.Handle(tc.method, tc.path, func(c *gin.Context) { c.Status(http.StatusOK) }) + + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(tc.method, tc.path, nil)) + if w.Code != http.StatusForbidden { + t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden) + } + }) + } +} + +func TestAssignedScopeCannotMutateProcessGlobalAssets(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, path := range []string{"/api/roles/demo", "/api/skills/demo", "/api/external-mcp/demo", "/api/workflows/demo", "/api/knowledge/items/demo"} { + t.Run(path, func(t *testing.T) { + permission := permissionForRequest(http.MethodPut, path) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(ContextSessionKey, Session{UserID: "operator", Scope: database.RBACScopeAssigned, Permissions: map[string]bool{permission: true}, PermissionScopes: map[string]string{permission: database.RBACScopeAssigned}}) + c.Next() + }) + router.Use(RBACMiddleware(&database.DB{})) + router.PUT(path, func(c *gin.Context) { c.Status(http.StatusNoContent) }) + w := httptest.NewRecorder() + router.ServeHTTP(w, httptest.NewRequest(http.MethodPut, path, nil)) + if w.Code != http.StatusForbidden { + t.Fatalf("global mutation status = %d, want 403", w.Code) + } + }) + } +} diff --git a/internal/security/route_inventory_test.go b/internal/security/route_inventory_test.go new file mode 100644 index 00000000..c08bfe99 --- /dev/null +++ b/internal/security/route_inventory_test.go @@ -0,0 +1,60 @@ +package security + +import ( + "go/ast" + "go/parser" + "go/token" + "net/http" + "path/filepath" + "strconv" + "testing" +) + +func TestEveryProtectedRouteHasCatalogPermission(t *testing.T) { + file, err := parser.ParseFile(token.NewFileSet(), filepath.Join("..", "app", "app.go"), nil, 0) + if err != nil { + t.Fatal(err) + } + methods := map[string]string{ + "GET": http.MethodGet, "POST": http.MethodPost, "PUT": http.MethodPut, + "PATCH": http.MethodPatch, "DELETE": http.MethodDelete, + } + prefixes := map[string]string{"protected": "", "c2Routes": "/c2", "knowledgeRoutes": "/knowledge"} + found := 0 + ast.Inspect(file, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok || len(call.Args) == 0 { + return true + } + sel, ok := call.Fun.(*ast.SelectorExpr) + if !ok { + return true + } + ident, ok := sel.X.(*ast.Ident) + if !ok { + return true + } + prefix, protected := prefixes[ident.Name] + method, routeMethod := methods[sel.Sel.Name] + literal, literalPath := call.Args[0].(*ast.BasicLit) + if !protected || !routeMethod || !literalPath || literal.Kind != token.STRING { + return true + } + path, err := strconv.Unquote(literal.Value) + if err != nil { + t.Errorf("invalid route literal %s", literal.Value) + return true + } + found++ + permission := permissionForRequest(method, "/api"+prefix+path) + if permission == "" { + t.Errorf("unmapped protected route: %s %s%s", method, prefix, path) + } else if _, ok := PermissionCatalog[permission]; !ok { + t.Errorf("route %s %s%s maps to unknown permission %q", method, prefix, path, permission) + } + return true + }) + if found < 100 { + t.Fatalf("route inventory unexpectedly small: %d", found) + } +} diff --git a/internal/security/shell_background_io.go b/internal/security/shell_background_io.go new file mode 100644 index 00000000..b54b1a0a --- /dev/null +++ b/internal/security/shell_background_io.go @@ -0,0 +1,111 @@ +package security + +import "strings" + +const backgroundJobStdioRedirect = " /dev/null 2>&1" + +// findStandaloneAmpersandPositions 返回不在引号内的独立 & 下标（排除 &&）。 +func findStandaloneAmpersandPositions(command string) []int { + command = strings.TrimSpace(command) + if command == "" { + return nil + } + + var positions []int + inSingleQuote := false + inDoubleQuote := false + escaped := false + + for i := 0; i < len(command); i++ { + r := command[i] + if escaped { + escaped = false + continue + } + if r == '\\' { + escaped = true + continue + } + if r == '\'' && !inDoubleQuote { + inSingleQuote = !inSingleQuote + continue + } + if r == '"' && !inSingleQuote { + inDoubleQuote = !inDoubleQuote + continue + } + if r != '&' || inSingleQuote || inDoubleQuote { + continue + } + if i+1 < len(command) && command[i+1] == '&' { + continue + } + if i > 0 && command[i-1] == '&' { + continue + } + + isStandalone := i == 0 + if !isStandalone { + prev := command[i-1] + isStandalone = prev == ' ' || prev == '\t' || prev == '\n' || prev == '\r' + } + if !isStandalone { + continue + } + if i == len(command)-1 { + positions = append(positions, i) + continue + } + next := command[i+1] + if next == ' ' || next == '\t' || next == '\n' || next == '\r' { + positions = append(positions, i) + } + } + return positions +} + +func segmentHasStdioRedirect(segment string) bool { + lower := strings.ToLower(strings.TrimSpace(segment)) + if lower == "" { + return false + } + if strings.Contains(lower, ">/dev/null") || strings.Contains(lower, "2>/dev/null") { + return true + } + if strings.Contains(lower, "&>") || strings.Contains(lower, "&>>") { + return true + } + if strings.Contains(lower, "2>&1") && strings.Contains(lower, "/dev/null") { + return true + } + return false +} + +// RedirectBackgroundJobStdio 为每个独立 & 前的后台段注入 /dev/null 2>&1， +// 避免后台子进程占用 execute/exec 管道导致挂死。 +func RedirectBackgroundJobStdio(command string) string { + positions := findStandaloneAmpersandPositions(command) + if len(positions) == 0 { + return command + } + + out := command + for j := len(positions) - 1; j >= 0; j-- { + i := positions[j] + before := out[:i] + after := out[i:] + trimmed := strings.TrimRight(before, " \t\r\n") + if segmentHasStdioRedirect(trimmed) { + continue + } + trailing := before[len(trimmed):] + out = trimmed + backgroundJobStdioRedirect + trailing + after + } + return out +} + +// PrepareShellCommandForExecute 组合 execute/exec 用的非交互包装与后台 IO 重定向。 +// 须先注入 exec /dev/null 2>&1 &") { + t.Fatalf("expected redirect before &: %q", out) + } + if !strings.Contains(out, "echo started") { + t.Fatalf("foreground tail preserved: %q", out) + } +} + +func TestRedirectBackgroundJobStdio_trailingOnly(t *testing.T) { + in := "sleep 120 &" + out := RedirectBackgroundJobStdio(in) + want := "sleep 120 /dev/null 2>&1 &" + if strings.TrimSpace(out) != want { + t.Fatalf("got %q want %q", out, want) + } +} + +func TestRedirectBackgroundJobStdio_skipsAlreadyRedirected(t *testing.T) { + in := "sleep 1 >/dev/null 2>&1 & echo ok" + out := RedirectBackgroundJobStdio(in) + if out != in { + t.Fatalf("should not double-redirect: %q", out) + } +} + +func TestRedirectBackgroundJobStdio_skipsAndAnd(t *testing.T) { + in := "test -f /etc/passwd && echo ok" + out := RedirectBackgroundJobStdio(in) + if out != in { + t.Fatalf("&& must not be treated as background &: %q", out) + } +} + +func TestPrepareShellCommandForExecute(t *testing.T) { + out := PrepareShellCommandForExecute("java -jar x & echo hi") + if !strings.Contains(out, "exec /dev/null 2>&1 &") { + t.Fatalf("missing background redirect: %q", out) + } +} + +func TestIsBackgroundShellCommand_usesSharedParser(t *testing.T) { + if !IsBackgroundShellCommand("sleep 1 &") { + t.Fatal("trailing & should be background") + } + if IsBackgroundShellCommand("sleep 1 & echo hi") { + t.Fatal("mixed should not be fully background") + } +} diff --git a/internal/security/shell_execute_stream.go b/internal/security/shell_execute_stream.go new file mode 100644 index 00000000..02c5cb74 --- /dev/null +++ b/internal/security/shell_execute_stream.go @@ -0,0 +1,211 @@ +package security + +import ( + "context" + "errors" + "fmt" + "io" + "os/exec" + "sync" + + "github.com/cloudwego/eino/adk/filesystem" + "github.com/cloudwego/eino/schema" +) + +// ConfigureShellCmdForAgentExecute 与 exec 工具一致：非交互 stdin、pager/TERM 环境、独立进程组。 +func ConfigureShellCmdForAgentExecute(cmd *exec.Cmd) { + if cmd == nil { + return + } + applyDefaultTerminalEnv(cmd) + attachNonInteractiveStdin(cmd) + _ = prepareShellCmdSession(cmd) +} + +// TerminateShellCmdTree 尽力终止 shell 及其子进程组（与 exec/execute 超时取消一致）。 +func TerminateShellCmdTree(cmd *exec.Cmd) { + terminateCmdTree(cmd) +} + +// TerminateShellCmdSession 使用 Start 时缓存的进程组 ID 终止（shell 已退出时仍有效）。 +func TerminateShellCmdSession(session *ShellSession) { + TerminateShellSession(session) +} + +// EinoStreamingShell 为 Eino ADK execute 工具提供流式 shell，行为与 exec 对齐： +// 并发读取 stdout/stderr（定长块，非按行），避免官方 local.ExecuteStreaming 先排空 stdout +// 导致 stderr 错误（如 sudo 密码提示）长时间不可见、UI 一直显示「执行中」。 +type EinoStreamingShell struct{} + +// NewEinoStreamingShell 创建 execute 流式 shell 实现。 +func NewEinoStreamingShell() *EinoStreamingShell { + return &EinoStreamingShell{} +} + +// ExecuteStreaming 实现 filesystem.StreamingShell。 +func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) { + if input == nil || input.Command == "" { + return nil, fmt.Errorf("command is required") + } + + sr, w := schema.Pipe[*filesystem.ExecuteResponse](100) + if input.RunInBackendGround { + go runShellInBackground(ctx, input.Command, w) + return sr, nil + } + go streamShellForeground(ctx, input.Command, w) + return sr, nil +} + +func runShellInBackground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) { + defer w.Close() + + command = PrepareShellCommandForExecute(command) + cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command) + applyDefaultTerminalEnv(cmd) + attachNonInteractiveStdin(cmd) + stdout, err := cmd.StdoutPipe() + if err != nil { + _ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err)) + return + } + stderr, err := cmd.StderrPipe() + if err != nil { + _ = stdout.Close() + _ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err)) + return + } + session, err := StartShellSession(cmd) + if err != nil { + _ = stdout.Close() + _ = stderr.Close() + _ = w.Send(nil, fmt.Errorf("failed to start command: %w", err)) + return + } + + done := make(chan struct{}) + go func() { + drainShellPipes(stdout, stderr) + _ = session.Wait() + close(done) + }() + + select { + case <-done: + case <-ctx.Done(): + TerminateShellCmdSession(session) + } + + exitCode := 0 + _ = w.Send(&filesystem.ExecuteResponse{ + Output: "command started in background\n", + ExitCode: &exitCode, + }, nil) +} + +func drainShellPipes(stdout, stderr io.Reader) { + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + _, _ = io.Copy(io.Discard, stdout) + }() + go func() { + defer wg.Done() + _, _ = io.Copy(io.Discard, stderr) + }() + wg.Wait() +} + +func streamShellForeground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) { + defer w.Close() + + command = PrepareShellCommandForExecute(command) + cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command) + applyDefaultTerminalEnv(cmd) + attachNonInteractiveStdin(cmd) + + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + _ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err)) + return + } + stderrPipe, err := cmd.StderrPipe() + if err != nil { + _ = stdoutPipe.Close() + _ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err)) + return + } + session, err := StartShellSession(cmd) + if err != nil { + _ = stdoutPipe.Close() + _ = stderrPipe.Close() + _ = w.Send(nil, fmt.Errorf("failed to start command: %w", err)) + return + } + + stopWatch := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + TerminateShellCmdSession(session) + case <-stopWatch: + } + }() + defer close(stopWatch) + + chunks := make(chan string, 64) + var wg sync.WaitGroup + readFn := func(r io.Reader) { + defer wg.Done() + buf := make([]byte, 8192) + for { + n, readErr := r.Read(buf) + if n > 0 { + chunks <- string(buf[:n]) + } + if readErr != nil { + return + } + } + } + + wg.Add(2) + go readFn(stdoutPipe) + go readFn(stderrPipe) + go func() { + wg.Wait() + close(chunks) + }() + + hadOutput := false + for chunk := range chunks { + if chunk == "" { + continue + } + hadOutput = true + if w.Send(&filesystem.ExecuteResponse{Output: chunk}, nil) { + TerminateShellCmdSession(session) + return + } + } + + waitErr := session.Wait() + if waitErr == nil { + exitCode := 0 + _ = w.Send(&filesystem.ExecuteResponse{ExitCode: &exitCode}, nil) + return + } + + var exitError *exec.ExitError + if errors.As(waitErr, &exitError) { + exitCode := exitError.ExitCode() + resp := &filesystem.ExecuteResponse{ExitCode: &exitCode} + if !hadOutput { + resp.Output = FormatCommandFailureResult(exitCode, "") + } + _ = w.Send(resp, nil) + return + } + _ = w.Send(nil, fmt.Errorf("command failed: %w", waitErr)) +} diff --git a/internal/security/shell_execute_stream_test.go b/internal/security/shell_execute_stream_test.go new file mode 100644 index 00000000..938f2994 --- /dev/null +++ b/internal/security/shell_execute_stream_test.go @@ -0,0 +1,152 @@ +package security + +import ( + "context" + "errors" + "io" + "strings" + "testing" + "time" + + "github.com/cloudwego/eino/adk/filesystem" +) + +func TestEinoStreamingShell_StreamsStderrBeforeStdoutEOF(t *testing.T) { + shell := NewEinoStreamingShell() + cmd := PrepareNonInteractiveShellCommand("echo err-only >&2; exit 1") + sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + start := time.Now() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + } + } + if time.Since(start) > 3*time.Second { + t.Fatalf("expected fast completion, took %v", time.Since(start)) + } + if !strings.Contains(got.String(), "err-only") { + t.Fatalf("expected stderr in output, got: %q", got.String()) + } +} + +func TestEinoStreamingShell_SudoFailsFast(t *testing.T) { + shell := NewEinoStreamingShell() + cmd := PrepareNonInteractiveShellCommand("sudo whoami && sudo cat /etc/os-release") + sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + start := time.Now() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp == nil { + continue + } + got.WriteString(resp.Output) + } + if time.Since(start) > 5*time.Second { + t.Fatalf("sudo should fail quickly, took %v output=%q", time.Since(start), got.String()) + } + out := got.String() + if strings.Contains(out, "command exited with non-zero code") { + t.Fatalf("legacy exit line present: %q", out) + } + if !strings.Contains(out, "sudo") && !strings.Contains(out, "password") && !strings.Contains(out, "terminal") { + t.Fatalf("expected sudo error text, got: %q", out) + } +} + +func TestEinoStreamingShell_StderrWhileStdoutBlocks(t *testing.T) { + shell := NewEinoStreamingShell() + // 模拟 sudo：stderr 先有输出，stdout 侧进程仍挂起；旧 eino local 在首包 stderr 前不会向流写任何内容。 + cmd := PrepareNonInteractiveShellCommand(`echo "password prompt" >&2; sleep 30`) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + sr, err := shell.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: cmd}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + start := time.Now() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + break + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + if strings.Contains(got.String(), "password prompt") { + break + } + } + } + if time.Since(start) > 1500*time.Millisecond { + t.Fatalf("expected stderr promptly, took %v output=%q", time.Since(start), got.String()) + } + if !strings.Contains(got.String(), "password prompt") { + t.Fatalf("expected early stderr, got: %q", got.String()) + } +} + +// TestEinoStreamingShell_BackgroundJobDoesNotHoldPipe 模拟 cmd & 后继续前台逻辑：重定向后应快速结束。 +func TestEinoStreamingShell_BackgroundJobDoesNotHoldPipe(t *testing.T) { + if testing.Short() { + t.Skip("skipping shell integration in -short") + } + shell := NewEinoStreamingShell() + cmd := `(sh -c 'printf x; sleep 120') & echo started; sleep 0` + sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd}) + if err != nil { + t.Fatalf("ExecuteStreaming: %v", err) + } + defer sr.Close() + + start := time.Now() + var got strings.Builder + for { + resp, rerr := sr.Recv() + if errors.Is(rerr, io.EOF) { + break + } + if rerr != nil { + t.Fatalf("recv: %v", rerr) + } + if resp != nil && resp.Output != "" { + got.WriteString(resp.Output) + } + } + if time.Since(start) > 3*time.Second { + t.Fatalf("expected fast completion, took %v output=%q", time.Since(start), got.String()) + } + if !strings.Contains(got.String(), "started") { + t.Fatalf("expected foreground echo, got: %q", got.String()) + } +} diff --git a/internal/security/shell_noninteractive.go b/internal/security/shell_noninteractive.go new file mode 100644 index 00000000..c6c567f0 --- /dev/null +++ b/internal/security/shell_noninteractive.go @@ -0,0 +1,163 @@ +package security + +import ( + "fmt" + "os" + "os/exec" + "strings" + "sync" + "time" +) + +// ShellNoOutputTimeoutMessage 长时间无新 stdout/stderr 时的提示（软失败，模型可见）。 +func ShellNoOutputTimeoutMessage(idleSec int) string { + return fmt.Sprintf(`命令已终止：超过 %d 秒没有新的输出，疑似在等待交互输入或已挂起。 + +长时静默任务请使用末尾 & 后台运行，或增大 agent.shell_no_output_timeout_seconds（-1=关闭此检测）。 + +Command terminated: no new output for %d seconds (possible interactive wait or hung process).`, idleSec, idleSec) +} + +// ShellInactivityWatch 在 noOutputSec 内无任何新输出时向 expired 发送信号；每次 Bump 重置计时。 +// 与「仅有首包输出就永久取消计时」不同，可兜住 sudo 打印 Password 提示后继续挂起等情况。 +type ShellInactivityWatch struct { + Sec int + mu sync.Mutex + timer *time.Timer + Expired chan struct{} +} + +func NewShellInactivityWatch(noOutputSec int) *ShellInactivityWatch { + sec := ResolveShellNoOutputTimeoutSeconds(noOutputSec) + if sec <= 0 { + return nil + } + w := &ShellInactivityWatch{ + Sec: sec, + Expired: make(chan struct{}, 1), + } + w.Bump() + return w +} + +func (w *ShellInactivityWatch) Bump() { + if w == nil || w.Sec <= 0 { + return + } + w.mu.Lock() + defer w.mu.Unlock() + if w.timer != nil { + w.timer.Stop() + } + w.timer = time.AfterFunc(time.Duration(w.Sec)*time.Second, func() { + select { + case w.Expired <- struct{}{}: + default: + } + }) +} + +func (w *ShellInactivityWatch) Stop() { + if w == nil { + return + } + w.mu.Lock() + defer w.mu.Unlock() + if w.timer != nil { + w.timer.Stop() + w.timer = nil + } +} + +// ResolveShellNoOutputTimeoutSeconds：0=默认 300（5 分钟）；-1=关闭；>0=自定义。 +func ResolveShellNoOutputTimeoutSeconds(sec int) int { + if sec < 0 { + return 0 + } + if sec == 0 { + return 300 + } + return sec +} + +// PrependNonInteractiveShellExports 为 sh -c 注入通用非交互环境（pager 等），不维护命令黑名单。 +func PrependNonInteractiveShellExports(shellCommand string) string { + if strings.TrimSpace(shellCommand) == "" { + return shellCommand + } + upper := strings.ToUpper(shellCommand) + var pairs []string + add := func(key, val string) { + if strings.Contains(upper, strings.ToUpper(key)) { + return + } + pairs = append(pairs, key+"="+val) + } + add("GIT_PAGER", "cat") + add("PAGER", "cat") + add("SYSTEMD_PAGER", "cat") + add("DEBIAN_FRONTEND", "noninteractive") + if len(pairs) == 0 { + return shellCommand + } + return "export " + strings.Join(pairs, " ") + "\n" + shellCommand +} + +// PrependNonInteractiveStdinRedirect 为 sh -c 关闭 stdin（与 attachNonInteractiveStdin 等价）， +// 使 read/input()/sudo -S 等从 stdin 读取的程序快速失败而非挂起。已含 "`)) + attachNonInteractiveStdin(cmd) + + start := time.Now() + out, err := cmd.CombinedOutput() + elapsed := time.Since(start) + if elapsed > 2*time.Second { + t.Fatalf("read with closed stdin took %v, want <2s", elapsed) + } + if err != nil { + t.Fatalf("unexpected error: %v output=%q", err, out) + } + if !strings.Contains(string(out), "x=<>") { + t.Fatalf("unexpected output: %q", out) + } +} + +// TestNonInteractiveStdinReadBlocksWithoutRedirect 对照：stdin 为永不写入的管道时 read 会挂起。 +func TestNonInteractiveStdinReadBlocksWithoutRedirect(t *testing.T) { + if testing.Short() { + t.Skip("skipping shell integration in -short") + } + r, w, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + defer r.Close() + // 保持 w 打开且不写数据，模拟「等待用户输入」 + + cmd := exec.Command("sh", "-c", `read x; echo done`) + cmd.Stdin = r + + done := make(chan error, 1) + go func() { done <- cmd.Run() }() + + select { + case err := <-done: + t.Fatalf("expected hang, but command finished: %v", err) + case <-time.After(500 * time.Millisecond): + if cmd.Process != nil { + _ = cmd.Process.Kill() + } + _ = w.Close() + <-done // 等待 goroutine 退出 + } +} diff --git a/internal/security/shell_session.go b/internal/security/shell_session.go new file mode 100644 index 00000000..72cb15e1 --- /dev/null +++ b/internal/security/shell_session.go @@ -0,0 +1,47 @@ +package security + +import "os/exec" + +// ShellSession 在 Start 时记录根 shell 的进程组 ID，取消/超时时可杀整组（即使 cmd.Process 已失效）。 +type ShellSession struct { + Cmd *exec.Cmd + rootPID int +} + +// StartShellSession 配置独立进程组并启动 shell，缓存 rootPID（Unix 下即 PGID）。 +func StartShellSession(cmd *exec.Cmd) (*ShellSession, error) { + if err := prepareShellCmdSession(cmd); err != nil { + return nil, err + } + if err := cmd.Start(); err != nil { + return nil, err + } + pid := 0 + if cmd.Process != nil { + pid = cmd.Process.Pid + } + return &ShellSession{Cmd: cmd, rootPID: pid}, nil +} + +// Wait 等待 shell 退出。 +func (s *ShellSession) Wait() error { + if s == nil || s.Cmd == nil { + return nil + } + return s.Cmd.Wait() +} + +// Terminate 终止 shell 及其进程组。 +func (s *ShellSession) Terminate() { + if s == nil { + return + } + terminateProcessGroup(s.rootPID, s.Cmd) +} + +// TerminateShellSession 终止由 StartShellSession 启动的会话。 +func TerminateShellSession(session *ShellSession) { + if session != nil { + session.Terminate() + } +} diff --git a/internal/security/shell_session_test.go b/internal/security/shell_session_test.go new file mode 100644 index 00000000..40520e3b --- /dev/null +++ b/internal/security/shell_session_test.go @@ -0,0 +1,65 @@ +package security + +import ( + "context" + "os/exec" + "runtime" + "testing" + "time" +) + +func TestShellSession_TerminateUsesCachedRootPID(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("unix process group kill") + } + + cmd := exec.Command("sh", "-c", "sleep 300") + ConfigureShellCmdForAgentExecute(cmd) + + session, err := StartShellSession(cmd) + if err != nil { + t.Fatalf("StartShellSession: %v", err) + } + + time.Sleep(100 * time.Millisecond) + session.Terminate() + + done := make(chan error, 1) + go func() { done <- session.Wait() }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("session did not finish within 5s after Terminate") + } +} + +func TestShellSession_TerminateAfterContextCancel(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("unix process group kill") + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + cmd := exec.CommandContext(ctx, "sh", "-c", "sleep 300") + ConfigureShellCmdForAgentExecute(cmd) + + session, err := StartShellSession(cmd) + if err != nil { + t.Fatalf("StartShellSession: %v", err) + } + + time.Sleep(100 * time.Millisecond) + cancel() + TerminateShellCmdSession(session) + + done := make(chan error, 1) + go func() { done <- session.Wait() }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("session did not finish within 5s after cancel+terminate") + } +} diff --git a/internal/security/workflow_package_rbac_test.go b/internal/security/workflow_package_rbac_test.go new file mode 100644 index 00000000..2f1fb1bc --- /dev/null +++ b/internal/security/workflow_package_rbac_test.go @@ -0,0 +1,20 @@ +package security + +import ( + "net/http" + "testing" +) + +func TestWorkflowPackageRoutesHaveExplicitWorkflowPermissions(t *testing.T) { + if got := permissionForRequest(http.MethodGet, "/api/workflows/:id/package"); got != "workflow:read" { + t.Fatalf("export permission=%q", got) + } + for _, path := range []string{"/api/workflow-package-inspections", "/api/workflow-package-inspections/:inspectionId", "/api/workflow-package-imports", "/api/workflow-package-imports/:importId"} { + if got := permissionForRequest(http.MethodGet, path); got != "workflow:write" { + t.Fatalf("%s permission=%q", path, got) + } + } + if !isProcessGlobalMutationPath("/workflow-package-imports") || !isProcessGlobalMutationPath("/workflow-package-inspections") { + t.Fatal("package mutations must require all-resource scope") + } +}