From c841809f9e4f1bdfbfe58aee773ed22bb9de4b9c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=85=AC=E6=98=8E?=
 <83812544+Ed1s0nZ@users.noreply.github.com>
Date: Wed, 22 Apr 2026 10:03:46 +0800
Subject: [PATCH] Add files via upload

---
 cmd/server/main.go            | 32 ++++++++++++++++++++++++++++----
 cmd/test-config/main.go       | 16 +++++-----------
 cmd/test-external-mcp/main.go | 18 ++++++------------
 3 files changed, 39 insertions(+), 27 deletions(-)

diff --git a/cmd/server/main.go b/cmd/server/main.go
index 9a962a3e..ce21de8c 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -1,11 +1,15 @@
 package main
 
 import (
+	"context"
 	"cyberstrike-ai/internal/app"
 	"cyberstrike-ai/internal/config"
 	"cyberstrike-ai/internal/logger"
 	"flag"
 	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
 )
 
 func main() {
@@ -31,15 +35,35 @@ func main() {
 	// 初始化日志
 	log := logger.New(cfg.Log.Level, cfg.Log.Output)
 
+	// 创建可取消的根 context，用于优雅关闭
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// 监听系统信号
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+
 	// 创建应用
 	application, err := app.New(cfg, log)
 	if err != nil {
 		log.Fatal("应用初始化失败", "error", err)
 	}
 
-	// 启动服务器
-	if err := application.Run(); err != nil {
-		log.Fatal("服务器启动失败", "error", err)
+	// 在后台监听信号
+	go func() {
+		sig := <-sigCh
+		log.Info("收到系统信号，开始优雅关闭: " + sig.String())
+		application.Shutdown()
+		cancel()
+	}()
+
+	// 启动服务器（传入 context 以支持优雅关闭）
+	if err := application.RunWithContext(ctx); err != nil {
+		// context 取消导致的关闭不视为错误
+		if ctx.Err() != nil {
+			log.Info("服务器已优雅关闭")
+		} else {
+			log.Fatal("服务器启动失败", "error", err)
+		}
 	}
 }
-
diff --git a/cmd/test-config/main.go b/cmd/test-config/main.go
index ab69be78..ef884dbc 100644
--- a/cmd/test-config/main.go
+++ b/cmd/test-config/main.go
@@ -37,21 +37,15 @@ func main() {
 		fmt.Printf("  URL: %s\n", srv.URL)
 		fmt.Printf("  Description: %s\n", srv.Description)
 		fmt.Printf("  Timeout: %d seconds\n", srv.Timeout)
-		fmt.Printf("  Enabled: %v\n", srv.Enabled)
-		fmt.Printf("  Disabled: %v\n", srv.Disabled)
+		fmt.Printf("  ExternalMCPEnable: %v\n", srv.ExternalMCPEnable)
 		fmt.Println()
 	}
 }
 
 func getTransport(srv config.ExternalMCPServerConfig) string {
-	if srv.Transport != "" {
-		return srv.Transport
+	t := srv.GetTransportType()
+	if t == "" {
+		return "unknown"
 	}
-	if srv.Command != "" {
-		return "stdio"
-	}
-	if srv.URL != "" {
-		return "http"
-	}
-	return "unknown"
+	return t
 }
diff --git a/cmd/test-external-mcp/main.go b/cmd/test-external-mcp/main.go
index 4b74bcbf..2b0078d2 100644
--- a/cmd/test-external-mcp/main.go
+++ b/cmd/test-external-mcp/main.go
@@ -52,8 +52,7 @@ func main() {
 		}
 		fmt.Printf("  Description: %s\n", srv.Description)
 		fmt.Printf("  Timeout: %d seconds\n", srv.Timeout)
-		fmt.Printf("  Enabled: %v\n", srv.Enabled)
-		fmt.Printf("  Disabled: %v\n", srv.Disabled)
+		fmt.Printf("  ExternalMCPEnable: %v\n", srv.ExternalMCPEnable)
 	}
 
 	// 获取统计信息
@@ -67,7 +66,7 @@ func main() {
 	// 测试启动（仅测试启用的）
 	fmt.Println("\n=== 测试启动 ===")
 	for name, srv := range cfg.ExternalMCP.Servers {
-		if srv.Enabled && !srv.Disabled {
+		if srv.ExternalMCPEnable {
 			fmt.Printf("\n尝试启动 %s...\n", name)
 			// 注意：实际启动可能会失败，因为需要真实的MCP服务器
 			err := manager.StartClient(name)
@@ -131,15 +130,10 @@ func main() {
 }
 
 func getTransport(srv config.ExternalMCPServerConfig) string {
-	if srv.Transport != "" {
-		return srv.Transport
+	t := srv.GetTransportType()
+	if t == "" {
+		return "unknown"
 	}
-	if srv.Command != "" {
-		return "stdio"
-	}
-	if srv.URL != "" {
-		return "http"
-	}
-	return "unknown"
+	return t
 }