Add files via upload

2026-06-29 17:30:14 +02:00 · 2026-06-29 14:35:52 +08:00
parent dd3621e89d
commit 39f1c72755
3 changed files with 195 additions and 24 deletions
@@ -828,33 +828,13 @@ func Load(path string) (*Config, error) {

 	// 如果配置了工具目录，从目录加载工具配置
 	if cfg.Security.ToolsDir != "" {
-		configDir := filepath.Dir(path)
-		toolsDir := cfg.Security.ToolsDir
-
-		// 如果是相对路径，相对于配置文件所在目录
-		if !filepath.IsAbs(toolsDir) {
-			toolsDir = filepath.Join(configDir, toolsDir)
-		}
-
-		tools, err := LoadToolsFromDir(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)
 		}
-
-		// 合并工具配置：目录中的工具优先，主配置中的工具作为补充
-		existingTools := make(map[string]bool)
-		for _, tool := range tools {
-			existingTools[tool.Name] = true
-		}
-
-		// 添加主配置中不存在于目录中的工具（向后兼容）
-		for _, tool := range cfg.Security.Tools {
-			if !existingTools[tool.Name] {
-				tools = append(tools, tool)
-			}
-		}
-
-		cfg.Security.Tools = tools
+		cfg.Security.Tools = merged
 	}

 	// 外部 MCP：迁移 + 环境变量展开
@@ -1126,6 +1106,75 @@ func PrintMCPConfigJSON(mcp MCPConfig) {
 	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
@@ -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)
+	}
+}
@@ -1333,6 +1333,17 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
 		h.logger.Info("已更新嵌入模型配置记录")
 	}

+	// 从 tools 目录重新加载工具配置（新增/修改/删除 yaml 后无需重启）
+	if err := config.ReloadSecurityToolsFromDir(h.config, h.configPath); err != nil {
+		h.logger.Error("重新加载工具配置失败", zap.Error(err))
+		if h.audit != nil {
+			h.audit.RecordFail(c, "config", "apply", "应用配置失败：重新加载工具", map[string]interface{}{"error": err.Error()})
+		}
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "重新加载工具配置失败: " + err.Error()})
+		return
+	}
+	h.logger.Info("已从 tools 目录重新加载工具配置", zap.Int("tools_count", len(h.config.Security.Tools)))
+
 	// 重新注册工具（根据新的启用状态）
 	h.logger.Info("重新注册工具")