Compare commits

...

19 Commits

Author SHA1 Message Date
公明 44b8d0b427 Add files via upload 2026-02-08 22:01:56 +08:00
公明 3a26d77c94 Update config.yaml 2026-02-08 21:29:08 +08:00
公明 0be6746794 Add files via upload 2026-02-08 21:28:20 +08:00
公明 06bfed508a Update requirements.txt 2026-02-08 20:56:15 +08:00
公明 0d617ebd66 Add files via upload 2026-02-08 20:46:51 +08:00
公明 9a52ec25ea Add files via upload 2026-02-08 20:40:50 +08:00
公明 594b7676e1 Add files via upload 2026-02-08 20:38:50 +08:00
公明 fd5d1dff10 Add files via upload 2026-02-08 20:20:43 +08:00
公明 b8218d9f77 Add tool description mode to security config
Add tool description mode configuration options.
2026-02-08 20:11:04 +08:00
公明 7b7c689efd Add files via upload 2026-02-08 20:09:23 +08:00
公明 27b16e0d54 Add files via upload 2026-02-07 00:04:59 +08:00
公明 2b6b678439 Add files via upload 2026-02-04 19:39:29 +08:00
公明 be104d1a05 Add files via upload 2026-02-04 19:37:46 +08:00
公明 f64bda3678 Add files via upload 2026-02-04 19:07:32 +08:00
公明 4b8dbb1bd6 Delete internal/mcp/client.go 2026-02-04 10:14:49 +08:00
公明 783d80ee37 Add files via upload 2026-02-04 02:00:52 +08:00
公明 a27c13b734 Add files via upload 2026-02-04 01:57:01 +08:00
公明 3cbf398636 Add files via upload 2026-02-04 01:39:11 +08:00
公明 84d54b1ea9 Add files via upload 2026-02-04 01:34:25 +08:00
24 changed files with 2042 additions and 1528 deletions
+7
View File
@@ -9,6 +9,9 @@
# 系统设置 # 系统设置
# ============================================ # ============================================
# 前端显示的版本号(可选,不填则显示默认版本)
version: "v1.3.4"
# 服务器配置 # 服务器配置
server: server:
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口 host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
@@ -67,6 +70,10 @@ database:
# 推荐方式:在 tools/ 目录下为每个工具创建独立的配置文件 # 推荐方式:在 tools/ 目录下为每个工具创建独立的配置文件
security: security:
tools_dir: tools # 工具配置文件目录(相对于配置文件所在目录) tools_dir: tools # 工具配置文件目录(相对于配置文件所在目录)
# 工具描述模式:加载 tools 下工具时,暴露给 AI/API 使用的描述来源
# short - 优先使用 short_description(简短描述,省 token),为空时用 description
# full - 使用 description(详细描述)
tool_description_mode: full
# ============================================ # ============================================
# MCP 相关配置 # MCP 相关配置
+7 -1
View File
@@ -1,11 +1,14 @@
module cyberstrike-ai module cyberstrike-ai
go 1.21 go 1.23.0
toolchain go1.24.4
require ( require (
github.com/gin-gonic/gin v1.9.1 github.com/gin-gonic/gin v1.9.1
github.com/google/uuid v1.5.0 github.com/google/uuid v1.5.0
github.com/mattn/go-sqlite3 v1.14.18 github.com/mattn/go-sqlite3 v1.14.18
github.com/modelcontextprotocol/go-sdk v1.2.0
github.com/pkoukk/tiktoken-go v0.1.8 github.com/pkoukk/tiktoken-go v0.1.8
go.uber.org/zap v1.26.0 go.uber.org/zap v1.26.0
gopkg.in/yaml.v3 v3.0.1 gopkg.in/yaml.v3 v3.0.1
@@ -21,6 +24,7 @@ require (
github.com/go-playground/universal-translator v0.18.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect
github.com/go-playground/validator/v10 v10.14.0 // indirect github.com/go-playground/validator/v10 v10.14.0 // indirect
github.com/goccy/go-json v0.10.2 // indirect github.com/goccy/go-json v0.10.2 // indirect
github.com/google/jsonschema-go v0.3.0 // indirect
github.com/json-iterator/go v1.1.12 // indirect github.com/json-iterator/go v1.1.12 // indirect
github.com/klauspost/cpuid/v2 v2.2.4 // indirect github.com/klauspost/cpuid/v2 v2.2.4 // indirect
github.com/leodido/go-urn v1.2.4 // indirect github.com/leodido/go-urn v1.2.4 // indirect
@@ -30,10 +34,12 @@ require (
github.com/pelletier/go-toml/v2 v2.0.8 // indirect github.com/pelletier/go-toml/v2 v2.0.8 // indirect
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
github.com/ugorji/go/codec v1.2.11 // indirect github.com/ugorji/go/codec v1.2.11 // indirect
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
go.uber.org/multierr v1.11.0 // indirect go.uber.org/multierr v1.11.0 // indirect
golang.org/x/arch v0.3.0 // indirect golang.org/x/arch v0.3.0 // indirect
golang.org/x/crypto v0.14.0 // indirect golang.org/x/crypto v0.14.0 // indirect
golang.org/x/net v0.17.0 // indirect golang.org/x/net v0.17.0 // indirect
golang.org/x/oauth2 v0.30.0 // indirect
golang.org/x/sys v0.13.0 // indirect golang.org/x/sys v0.13.0 // indirect
golang.org/x/text v0.13.0 // indirect golang.org/x/text v0.13.0 // indirect
google.golang.org/protobuf v1.30.0 // indirect google.golang.org/protobuf v1.30.0 // indirect
+14 -2
View File
@@ -25,10 +25,15 @@ github.com/go-playground/validator/v10 v10.14.0 h1:vgvQWe3XCz3gIeFDm/HnTIbj6UGmg
github.com/go-playground/validator/v10 v10.14.0/go.mod h1:9iXMNT7sEkjXb0I+enO7QXmzG6QCsPWY4zveKFVRSyU= github.com/go-playground/validator/v10 v10.14.0/go.mod h1:9iXMNT7sEkjXb0I+enO7QXmzG6QCsPWY4zveKFVRSyU=
github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8=
github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q=
github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
github.com/google/uuid v1.5.0 h1:1p67kYwdtXjb0gL0BPiP1Av9wiZPo5A8z2cWkTZ+eyU= github.com/google/uuid v1.5.0 h1:1p67kYwdtXjb0gL0BPiP1Av9wiZPo5A8z2cWkTZ+eyU=
github.com/google/uuid v1.5.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.5.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
@@ -42,6 +47,8 @@ github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APP
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-sqlite3 v1.14.18 h1:JL0eqdCOq6DJVNPSvArO/bIV9/P7fbGrV00LZHc+5aI= github.com/mattn/go-sqlite3 v1.14.18 h1:JL0eqdCOq6DJVNPSvArO/bIV9/P7fbGrV00LZHc+5aI=
github.com/mattn/go-sqlite3 v1.14.18/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= github.com/mattn/go-sqlite3 v1.14.18/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
github.com/modelcontextprotocol/go-sdk v1.2.0 h1:Y23co09300CEk8iZ/tMxIX1dVmKZkzoSBZOpJwUnc/s=
github.com/modelcontextprotocol/go-sdk v1.2.0/go.mod h1:6fM3LCm3yV7pAs8isnKLn07oKtB0MP9LHd3DfAcKw10=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -68,6 +75,8 @@ github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS
github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU= github.com/ugorji/go/codec v1.2.11 h1:BMaWp1Bb6fHwEtbplGBGJ498wD+LKlNSl25MjdZY4dU=
github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= github.com/ugorji/go/codec v1.2.11/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
go.uber.org/goleak v1.2.0 h1:xqgm/S+aQvhWFTtR0XK3Jvg7z8kGV8P4X14IzwN3Eqk= go.uber.org/goleak v1.2.0 h1:xqgm/S+aQvhWFTtR0XK3Jvg7z8kGV8P4X14IzwN3Eqk=
go.uber.org/goleak v1.2.0/go.mod h1:XJYK+MuIchqpmGmUSAzotztawfKvYLUIgg7guXrwVUo= go.uber.org/goleak v1.2.0/go.mod h1:XJYK+MuIchqpmGmUSAzotztawfKvYLUIgg7guXrwVUo=
go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0= go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
@@ -81,13 +90,16 @@ golang.org/x/crypto v0.14.0 h1:wBqGXzWJW6m1XrIKlAH0Hs1JJ7+9KBwnIO8v66Q9cHc=
golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4= golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4=
golang.org/x/net v0.17.0 h1:pVaXccu2ozPjCXewfr1S7xza/zcXTity9cCdXQYSjIM= golang.org/x/net v0.17.0 h1:pVaXccu2ozPjCXewfr1S7xza/zcXTity9cCdXQYSjIM=
golang.org/x/net v0.17.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE= golang.org/x/net v0.17.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE=
golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU=
golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE= golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k= golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=
golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= golang.org/x/tools v0.34.0 h1:qIpSLOxeCYGg9TrcJokLBG4KFA6d795g0xkBkiESGlo=
golang.org/x/tools v0.34.0/go.mod h1:pAP9OwEaY1CAW3HOmg3hLZC5Z0CCmzjAF2UQMSqNARg=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng= google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
+27 -13
View File
@@ -529,8 +529,10 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
maxIterations := a.maxIterations maxIterations := a.maxIterations
for i := 0; i < maxIterations; i++ { for i := 0; i < maxIterations; i++ {
// 每轮调用前先尝试压缩,防止历史消息持续膨胀 // 先获取本轮可用工具并统计 tools token,再压缩,以便压缩时预留 tools 占用的空间
messages = a.applyMemoryCompression(ctx, messages) tools := a.getAvailableTools(roleTools)
toolsTokens := a.countToolsTokens(tools)
messages = a.applyMemoryCompression(ctx, messages, toolsTokens)
// 检查是否是最后一次迭代 // 检查是否是最后一次迭代
isLastIteration := (i == maxIterations-1) isLastIteration := (i == maxIterations-1)
@@ -562,17 +564,17 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
default: default:
} }
// 获取可用工具 // 记录当前上下文的 Token 用量(messages + tools),展示压缩器运行状态
tools := a.getAvailableTools(roleTools)
// 记录当前上下文的Token用量,展示压缩器运行状态
if a.memoryCompressor != nil { if a.memoryCompressor != nil {
totalTokens, systemCount, regularCount := a.memoryCompressor.totalTokensFor(messages) messagesTokens, systemCount, regularCount := a.memoryCompressor.totalTokensFor(messages)
totalTokens := messagesTokens + toolsTokens
a.logger.Info("memory compressor context stats", a.logger.Info("memory compressor context stats",
zap.Int("iteration", i+1), zap.Int("iteration", i+1),
zap.Int("messagesCount", len(messages)), zap.Int("messagesCount", len(messages)),
zap.Int("systemMessages", systemCount), zap.Int("systemMessages", systemCount),
zap.Int("regularMessages", regularCount), zap.Int("regularMessages", regularCount),
zap.Int("messagesTokens", messagesTokens),
zap.Int("toolsTokens", toolsTokens),
zap.Int("totalTokens", totalTokens), zap.Int("totalTokens", totalTokens),
zap.Int("maxTotalTokens", a.memoryCompressor.maxTotalTokens), zap.Int("maxTotalTokens", a.memoryCompressor.maxTotalTokens),
) )
@@ -788,7 +790,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
Role: "user", Role: "user",
Content: "这是最后一次迭代。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试,请提供详细的下一步执行计划。请直接回复,不要调用工具。", Content: "这是最后一次迭代。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试,请提供详细的下一步执行计划。请直接回复,不要调用工具。",
}) })
messages = a.applyMemoryCompression(ctx, messages) messages = a.applyMemoryCompression(ctx, messages, 0) // 总结时不带 tools,不预留
// 立即调用OpenAI获取总结 // 立即调用OpenAI获取总结
summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具,强制AI直接回复 summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具,强制AI直接回复
if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 { if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 {
@@ -828,7 +830,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
Role: "user", Role: "user",
Content: "这是最后一次迭代。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试,请提供详细的下一步执行计划。请直接回复,不要调用工具。", Content: "这是最后一次迭代。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试,请提供详细的下一步执行计划。请直接回复,不要调用工具。",
}) })
messages = a.applyMemoryCompression(ctx, messages) messages = a.applyMemoryCompression(ctx, messages, 0) // 总结时不带 tools,不预留
// 立即调用OpenAI获取总结 // 立即调用OpenAI获取总结
summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具,强制AI直接回复 summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具,强制AI直接回复
if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 { if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 {
@@ -867,7 +869,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
Content: fmt.Sprintf("已达到最大迭代次数(%d轮)。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试,请提供详细的下一步执行计划。请直接回复,不要调用工具。", a.maxIterations), Content: fmt.Sprintf("已达到最大迭代次数(%d轮)。请总结到目前为止的所有测试结果、发现的问题和已完成的工作。如果需要继续测试,请提供详细的下一步执行计划。请直接回复,不要调用工具。", a.maxIterations),
} }
messages = append(messages, finalSummaryPrompt) messages = append(messages, finalSummaryPrompt)
messages = a.applyMemoryCompression(ctx, messages) messages = a.applyMemoryCompression(ctx, messages, 0) // 总结时不带 tools,不预留
summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具,强制AI直接回复 summaryResponse, err := a.callOpenAI(ctx, messages, []Tool{}) // 不提供工具,强制AI直接回复
if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 { if err == nil && summaryResponse != nil && len(summaryResponse.Choices) > 0 {
@@ -1425,13 +1427,13 @@ func (a *Agent) formatToolError(toolName string, args map[string]interface{}, er
return errorMsg return errorMsg
} }
// applyMemoryCompression 在调用LLM前对消息进行压缩,避免超过token限制 // applyMemoryCompression 在调用LLM前对消息进行压缩,避免超过 token 限制。reservedTokens 为预留给 tools 的 token 数,传 0 表示不预留。
func (a *Agent) applyMemoryCompression(ctx context.Context, messages []ChatMessage) []ChatMessage { func (a *Agent) applyMemoryCompression(ctx context.Context, messages []ChatMessage, reservedTokens int) []ChatMessage {
if a.memoryCompressor == nil { if a.memoryCompressor == nil {
return messages return messages
} }
compressed, changed, err := a.memoryCompressor.CompressHistory(ctx, messages) compressed, changed, err := a.memoryCompressor.CompressHistory(ctx, messages, reservedTokens)
if err != nil { if err != nil {
a.logger.Warn("上下文压缩失败,将使用原始消息继续", zap.Error(err)) a.logger.Warn("上下文压缩失败,将使用原始消息继续", zap.Error(err))
return messages return messages
@@ -1447,6 +1449,18 @@ func (a *Agent) applyMemoryCompression(ctx context.Context, messages []ChatMessa
return messages return messages
} }
// countToolsTokens 统计 tools 序列化后的 token 数,用于日志与压缩时预留空间。mc 为 nil 时返回 0。
func (a *Agent) countToolsTokens(tools []Tool) int {
if len(tools) == 0 || a.memoryCompressor == nil {
return 0
}
data, err := json.Marshal(tools)
if err != nil {
return 0
}
return a.memoryCompressor.CountTextTokens(string(data))
}
// handleMissingToolError 当LLM调用不存在的工具时,向其追加提示消息并允许继续迭代 // handleMissingToolError 当LLM调用不存在的工具时,向其追加提示消息并允许继续迭代
func (a *Agent) handleMissingToolError(errMsg string, messages *[]ChatMessage) (bool, string) { func (a *Agent) handleMissingToolError(errMsg string, messages *[]ChatMessage) (bool, string) {
lowerMsg := strings.ToLower(errMsg) lowerMsg := strings.ToLower(errMsg)
+15 -3
View File
@@ -158,8 +158,8 @@ func (mc *MemoryCompressor) UpdateConfig(cfg *config.OpenAIConfig) {
} }
} }
// CompressHistory 根据Token限制压缩历史消息。 // CompressHistory 根据 Token 限制压缩历史消息。reservedTokens 为预留给 tools 等非消息内容的 token 数,压缩时使用 (maxTotalTokens - reservedTokens) 作为消息上限。
func (mc *MemoryCompressor) CompressHistory(ctx context.Context, messages []ChatMessage) ([]ChatMessage, bool, error) { func (mc *MemoryCompressor) CompressHistory(ctx context.Context, messages []ChatMessage, reservedTokens int) ([]ChatMessage, bool, error) {
if len(messages) == 0 { if len(messages) == 0 {
return messages, false, nil return messages, false, nil
} }
@@ -171,8 +171,13 @@ func (mc *MemoryCompressor) CompressHistory(ctx context.Context, messages []Chat
return messages, false, nil return messages, false, nil
} }
effectiveMax := mc.maxTotalTokens
if reservedTokens > 0 && reservedTokens < mc.maxTotalTokens {
effectiveMax = mc.maxTotalTokens - reservedTokens
}
totalTokens := mc.countTotalTokens(systemMsgs, regularMsgs) totalTokens := mc.countTotalTokens(systemMsgs, regularMsgs)
if totalTokens <= int(float64(mc.maxTotalTokens)*0.9) { if totalTokens <= int(float64(effectiveMax)*0.9) {
return messages, false, nil return messages, false, nil
} }
@@ -184,6 +189,8 @@ func (mc *MemoryCompressor) CompressHistory(ctx context.Context, messages []Chat
mc.logger.Info("memory compression triggered", mc.logger.Info("memory compression triggered",
zap.Int("total_tokens", totalTokens), zap.Int("total_tokens", totalTokens),
zap.Int("max_total_tokens", mc.maxTotalTokens), zap.Int("max_total_tokens", mc.maxTotalTokens),
zap.Int("reserved_tokens", reservedTokens),
zap.Int("effective_max", effectiveMax),
zap.Int("system_messages", len(systemMsgs)), zap.Int("system_messages", len(systemMsgs)),
zap.Int("regular_messages", len(regularMsgs)), zap.Int("regular_messages", len(regularMsgs)),
zap.Int("old_messages", len(oldMsgs)), zap.Int("old_messages", len(oldMsgs)),
@@ -282,6 +289,11 @@ func (mc *MemoryCompressor) countTokens(text string) int {
return count return count
} }
// CountTextTokens 对外暴露的文本 Token 计数,用于统计 tools 等非消息内容的 token(如 agent 侧序列化 tools 后计数)。
func (mc *MemoryCompressor) CountTextTokens(text string) int {
return mc.countTokens(text)
}
// totalTokensFor provides token statistics without mutating the message list. // totalTokensFor provides token statistics without mutating the message list.
func (mc *MemoryCompressor) totalTokensFor(messages []ChatMessage) (totalTokens int, systemCount int, regularCount int) { func (mc *MemoryCompressor) totalTokensFor(messages []ChatMessage) (totalTokens int, systemCount int, regularCount int) {
if len(messages) == 0 { if len(messages) == 0 {
+5 -1
View File
@@ -746,7 +746,11 @@ func setupRoutes(
// 前端页面 // 前端页面
router.GET("/", func(c *gin.Context) { router.GET("/", func(c *gin.Context) {
c.HTML(http.StatusOK, "index.html", nil) version := app.config.Version
if version == "" {
version = "v1.0.0"
}
c.HTML(http.StatusOK, "index.html", gin.H{"Version": version})
}) })
} }
+4 -1
View File
@@ -13,6 +13,7 @@ import (
) )
type Config struct { type Config struct {
Version string `yaml:"version,omitempty" json:"version,omitempty"` // 前端显示的版本号,如 v1.3.3
Server ServerConfig `yaml:"server"` Server ServerConfig `yaml:"server"`
Log LogConfig `yaml:"log"` Log LogConfig `yaml:"log"`
MCP MCPConfig `yaml:"mcp"` MCP MCPConfig `yaml:"mcp"`
@@ -54,6 +55,7 @@ type OpenAIConfig struct {
type SecurityConfig struct { type SecurityConfig struct {
Tools []ToolConfig `yaml:"tools,omitempty"` // 向后兼容:支持在主配置文件中定义工具 Tools []ToolConfig `yaml:"tools,omitempty"` // 向后兼容:支持在主配置文件中定义工具
ToolsDir string `yaml:"tools_dir,omitempty"` // 工具配置文件目录(新方式) ToolsDir string `yaml:"tools_dir,omitempty"` // 工具配置文件目录(新方式)
ToolDescriptionMode string `yaml:"tool_description_mode,omitempty"` // 工具描述模式: "short" | "full",默认 short
} }
type DatabaseConfig struct { type DatabaseConfig struct {
@@ -88,8 +90,9 @@ type ExternalMCPServerConfig struct {
Env map[string]string `yaml:"env,omitempty" json:"env,omitempty"` // 环境变量(用于stdio模式) Env map[string]string `yaml:"env,omitempty" json:"env,omitempty"` // 环境变量(用于stdio模式)
// HTTP模式配置 // HTTP模式配置
Transport string `yaml:"transport,omitempty" json:"transport,omitempty"` // "http" 或 "stdio" Transport string `yaml:"transport,omitempty" json:"transport,omitempty"` // "stdio" | "sse" | "http"(Streamable) | "simple_http"(自建/简单POST端点,如本机 http://127.0.0.1:8081/mcp)
URL string `yaml:"url,omitempty" json:"url,omitempty"` URL string `yaml:"url,omitempty" json:"url,omitempty"`
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` // HTTP/SSE 请求头(如 x-api-key
// 通用配置 // 通用配置
Description string `yaml:"description,omitempty" json:"description,omitempty"` Description string `yaml:"description,omitempty" json:"description,omitempty"`
+133 -145
View File
@@ -174,18 +174,10 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
configToolMap[tool.Name] = true configToolMap[tool.Name] = true
tools = append(tools, ToolConfigInfo{ tools = append(tools, ToolConfigInfo{
Name: tool.Name, Name: tool.Name,
Description: tool.ShortDescription, Description: h.pickToolDescription(tool.ShortDescription, tool.Description),
Enabled: tool.Enabled, Enabled: tool.Enabled,
IsExternal: false, IsExternal: false,
}) })
// 如果没有简短描述,使用详细描述的前100个字符
if tools[len(tools)-1].Description == "" {
desc := tool.Description
if len(desc) > 100 {
desc = desc[:100] + "..."
}
tools[len(tools)-1].Description = desc
}
} }
// 从MCP服务器获取所有已注册的工具(包括直接注册的工具,如知识检索工具) // 从MCP服务器获取所有已注册的工具(包括直接注册的工具,如知识检索工具)
@@ -201,8 +193,8 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
if description == "" { if description == "" {
description = mcpTool.Description description = mcpTool.Description
} }
if len(description) > 100 { if len(description) > 10000 {
description = description[:100] + "..." description = description[:10000] + "..."
} }
tools = append(tools, ToolConfigInfo{ tools = append(tools, ToolConfigInfo{
Name: mcpTool.Name, Name: mcpTool.Name,
@@ -215,61 +207,10 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
// 获取外部MCP工具 // 获取外部MCP工具
if h.externalMCPMgr != nil { if h.externalMCPMgr != nil {
// 增加超时时间到30秒,因为通过代理连接远程服务器可能需要更长时间 ctx := context.Background()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) externalTools := h.getExternalMCPTools(ctx)
defer cancel() for _, toolInfo := range externalTools {
tools = append(tools, toolInfo)
externalTools, err := h.externalMCPMgr.GetAllTools(ctx)
if err == nil {
externalMCPConfigs := h.externalMCPMgr.GetConfigs()
for _, externalTool := range externalTools {
var mcpName, actualToolName string
if idx := strings.Index(externalTool.Name, "::"); idx > 0 {
mcpName = externalTool.Name[:idx]
actualToolName = externalTool.Name[idx+2:]
} else {
continue
}
enabled := false
if cfg, exists := externalMCPConfigs[mcpName]; exists {
// 首先检查外部MCP是否启用
if !cfg.ExternalMCPEnable && !(cfg.Enabled && !cfg.Disabled) {
enabled = false // MCP未启用,所有工具都禁用
} else {
// MCP已启用,检查单个工具的启用状态
// 如果ToolEnabled为空或未设置该工具,默认为启用(向后兼容)
if cfg.ToolEnabled == nil {
enabled = true // 未设置工具状态,默认为启用
} else if toolEnabled, exists := cfg.ToolEnabled[actualToolName]; exists {
enabled = toolEnabled // 使用配置的工具状态
} else {
enabled = true // 工具未在配置中,默认为启用
}
}
}
client, exists := h.externalMCPMgr.GetClient(mcpName)
if !exists || !client.IsConnected() {
enabled = false
}
description := externalTool.ShortDescription
if description == "" {
description = externalTool.Description
}
if len(description) > 100 {
description = description[:100] + "..."
}
tools = append(tools, ToolConfigInfo{
Name: actualToolName,
Description: description,
Enabled: enabled,
IsExternal: true,
ExternalMCP: mcpName,
})
}
} }
} }
@@ -342,18 +283,10 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
configToolMap[tool.Name] = true configToolMap[tool.Name] = true
toolInfo := ToolConfigInfo{ toolInfo := ToolConfigInfo{
Name: tool.Name, Name: tool.Name,
Description: tool.ShortDescription, Description: h.pickToolDescription(tool.ShortDescription, tool.Description),
Enabled: tool.Enabled, Enabled: tool.Enabled,
IsExternal: false, IsExternal: false,
} }
// 如果没有简短描述,使用详细描述的前100个字符
if toolInfo.Description == "" {
desc := tool.Description
if len(desc) > 100 {
desc = desc[:100] + "..."
}
toolInfo.Description = desc
}
// 根据角色配置标注工具状态 // 根据角色配置标注工具状态
if roleName != "" { if roleName != "" {
@@ -405,8 +338,8 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
if description == "" { if description == "" {
description = mcpTool.Description description = mcpTool.Description
} }
if len(description) > 100 { if len(description) > 10000 {
description = description[:100] + "..." description = description[:10000] + "..."
} }
toolInfo := ToolConfigInfo{ toolInfo := ToolConfigInfo{
@@ -451,88 +384,33 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
// 获取外部MCP工具 // 获取外部MCP工具
if h.externalMCPMgr != nil { if h.externalMCPMgr != nil {
// 增加超时时间到30秒,因为通过代理连接远程服务器可能需要更长时间 // 创建context用于获取外部工具
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) ctx := context.Background()
defer cancel() externalTools := h.getExternalMCPTools(ctx)
externalTools, err := h.externalMCPMgr.GetAllTools(ctx) // 应用搜索过滤和角色配置
if err != nil { for _, toolInfo := range externalTools {
h.logger.Warn("获取外部MCP工具失败", zap.Error(err)) // 搜索过滤
} else {
// 获取外部MCP配置,用于判断启用状态
externalMCPConfigs := h.externalMCPMgr.GetConfigs()
for _, externalTool := range externalTools {
// 解析工具名称:mcpName::toolName
var mcpName, actualToolName string
if idx := strings.Index(externalTool.Name, "::"); idx > 0 {
mcpName = externalTool.Name[:idx]
actualToolName = externalTool.Name[idx+2:]
} else {
continue // 跳过格式不正确的工具
}
// 获取外部工具的启用状态
enabled := false
if cfg, exists := externalMCPConfigs[mcpName]; exists {
// 首先检查外部MCP是否启用
if !cfg.ExternalMCPEnable && !(cfg.Enabled && !cfg.Disabled) {
enabled = false // MCP未启用,所有工具都禁用
} else {
// MCP已启用,检查单个工具的启用状态
// 如果ToolEnabled为空或未设置该工具,默认为启用(向后兼容)
if cfg.ToolEnabled == nil {
enabled = true // 未设置工具状态,默认为启用
} else if toolEnabled, exists := cfg.ToolEnabled[actualToolName]; exists {
enabled = toolEnabled // 使用配置的工具状态
} else {
enabled = true // 工具未在配置中,默认为启用
}
}
}
// 检查外部MCP是否已连接
client, exists := h.externalMCPMgr.GetClient(mcpName)
if !exists || !client.IsConnected() {
enabled = false // 未连接时视为禁用
}
description := externalTool.ShortDescription
if description == "" {
description = externalTool.Description
}
if len(description) > 100 {
description = description[:100] + "..."
}
// 如果有关键词,进行搜索过滤
if searchTermLower != "" { if searchTermLower != "" {
nameLower := strings.ToLower(actualToolName) nameLower := strings.ToLower(toolInfo.Name)
descLower := strings.ToLower(description) descLower := strings.ToLower(toolInfo.Description)
if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) { if !strings.Contains(nameLower, searchTermLower) && !strings.Contains(descLower, searchTermLower) {
continue // 不匹配,跳过 continue // 不匹配,跳过
} }
} }
toolInfo := ToolConfigInfo{
Name: actualToolName, // 显示实际工具名称,不带前缀
Description: description,
Enabled: enabled,
IsExternal: true,
ExternalMCP: mcpName,
}
// 根据角色配置标注工具状态 // 根据角色配置标注工具状态
if roleName != "" { if roleName != "" {
if roleUsesAllTools { if roleUsesAllTools {
// 角色使用所有工具,标注启用的工具为role_enabled=true // 角色使用所有工具,标注启用的工具为role_enabled=true
toolInfo.RoleEnabled = &enabled roleEnabled := toolInfo.Enabled
toolInfo.RoleEnabled = &roleEnabled
} else { } else {
// 角色配置了工具列表,检查工具是否在列表中 // 角色配置了工具列表,检查工具是否在列表中
// 外部工具使用 "mcpName::toolName" 格式作为key // 外部工具使用 "mcpName::toolName" 格式作为key
externalToolKey := externalTool.Name // 这是 "mcpName::toolName" 格式 externalToolKey := fmt.Sprintf("%s::%s", toolInfo.ExternalMCP, toolInfo.Name)
if roleToolsSet[externalToolKey] { if roleToolsSet[externalToolKey] {
roleEnabled := enabled // 工具必须在角色列表中且本身启用 roleEnabled := toolInfo.Enabled // 工具必须在角色列表中且本身启用
toolInfo.RoleEnabled = &roleEnabled toolInfo.RoleEnabled = &roleEnabled
} else { } else {
// 不在角色列表中,标记为false // 不在角色列表中,标记为false
@@ -545,7 +423,6 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
allTools = append(allTools, toolInfo) allTools = append(allTools, toolInfo)
} }
} }
}
// 如果角色配置了工具列表,过滤工具(只保留列表中的工具,但保留其他工具并标记为禁用) // 如果角色配置了工具列表,过滤工具(只保留列表中的工具,但保留其他工具并标记为禁用)
// 注意:这里我们不直接过滤掉工具,而是保留所有工具,但通过 role_enabled 字段标注状态 // 注意:这里我们不直接过滤掉工具,而是保留所有工具,但通过 role_enabled 字段标注状态
@@ -1259,3 +1136,114 @@ func setFloatInMap(mapNode *yaml.Node, key string, value float64) {
valueNode.Value = fmt.Sprintf("%g", value) valueNode.Value = fmt.Sprintf("%g", value)
} }
} }
// getExternalMCPTools 获取外部MCP工具列表(公共方法)
// 返回 ToolConfigInfo 列表,已处理启用状态和描述信息
func (h *ConfigHandler) getExternalMCPTools(ctx context.Context) []ToolConfigInfo {
var result []ToolConfigInfo
if h.externalMCPMgr == nil {
return result
}
// 使用较短的超时时间(5秒)进行快速失败,避免阻塞页面加载
timeoutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
externalTools, err := h.externalMCPMgr.GetAllTools(timeoutCtx)
if err != nil {
// 记录警告但不阻塞,继续返回已缓存的工具(如果有)
h.logger.Warn("获取外部MCP工具失败(可能连接断开),尝试返回缓存的工具",
zap.Error(err),
zap.String("hint", "如果外部MCP工具未显示,请检查连接状态或点击刷新按钮"),
)
}
// 如果获取到了工具(即使有错误),继续处理
if len(externalTools) == 0 {
return result
}
externalMCPConfigs := h.externalMCPMgr.GetConfigs()
for _, externalTool := range externalTools {
// 解析工具名称:mcpName::toolName
mcpName, actualToolName := h.parseExternalToolName(externalTool.Name)
if mcpName == "" || actualToolName == "" {
continue // 跳过格式不正确的工具
}
// 计算启用状态
enabled := h.calculateExternalToolEnabled(mcpName, actualToolName, externalMCPConfigs)
// 处理描述信息
description := h.pickToolDescription(externalTool.ShortDescription, externalTool.Description)
result = append(result, ToolConfigInfo{
Name: actualToolName,
Description: description,
Enabled: enabled,
IsExternal: true,
ExternalMCP: mcpName,
})
}
return result
}
// parseExternalToolName 解析外部工具名称(格式:mcpName::toolName
func (h *ConfigHandler) parseExternalToolName(fullName string) (mcpName, toolName string) {
idx := strings.Index(fullName, "::")
if idx > 0 {
return fullName[:idx], fullName[idx+2:]
}
return "", ""
}
// calculateExternalToolEnabled 计算外部工具的启用状态
func (h *ConfigHandler) calculateExternalToolEnabled(mcpName, toolName string, configs map[string]config.ExternalMCPServerConfig) bool {
cfg, exists := configs[mcpName]
if !exists {
return false
}
// 首先检查外部MCP是否启用
if !cfg.ExternalMCPEnable && !(cfg.Enabled && !cfg.Disabled) {
return false // MCP未启用,所有工具都禁用
}
// MCP已启用,检查单个工具的启用状态
// 如果ToolEnabled为空或未设置该工具,默认为启用(向后兼容)
if cfg.ToolEnabled == nil {
// 未设置工具状态,默认为启用
} else if toolEnabled, exists := cfg.ToolEnabled[toolName]; exists {
// 使用配置的工具状态
if !toolEnabled {
return false
}
}
// 工具未在配置中,默认为启用
// 最后检查外部MCP是否已连接
client, exists := h.externalMCPMgr.GetClient(mcpName)
if !exists || !client.IsConnected() {
return false // 未连接时视为禁用
}
return true
}
// pickToolDescription 根据 security.tool_description_mode 选择 short 或 full 描述并限制长度
func (h *ConfigHandler) pickToolDescription(shortDesc, fullDesc string) string {
useFull := strings.TrimSpace(strings.ToLower(h.config.Security.ToolDescriptionMode)) == "full"
description := shortDesc
if useFull {
description = fullDesc
} else if description == "" {
description = fullDesc
}
if len(description) > 10000 {
description = description[:10000] + "..."
}
return description
}
+8 -1
View File
@@ -8,6 +8,7 @@ import (
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"go.uber.org/zap" "go.uber.org/zap"
"gopkg.in/yaml.v3" "gopkg.in/yaml.v3"
@@ -459,6 +460,13 @@ func updateExternalMCPConfig(doc *yaml.Node, cfg config.ExternalMCPConfig, origi
if serverCfg.URL != "" { if serverCfg.URL != "" {
setStringInMap(serverNode, "url", serverCfg.URL) setStringInMap(serverNode, "url", serverCfg.URL)
} }
// 保存 headers 字段(HTTP/SSE 请求头)
if serverCfg.Headers != nil && len(serverCfg.Headers) > 0 {
headersNode := ensureMap(serverNode, "headers")
for k, v := range serverCfg.Headers {
setStringInMap(headersNode, k, v)
}
}
if serverCfg.Description != "" { if serverCfg.Description != "" {
setStringInMap(serverNode, "description", serverCfg.Description) setStringInMap(serverNode, "description", serverCfg.Description)
} }
@@ -532,4 +540,3 @@ type ExternalMCPResponse struct {
ToolCount int `json:"tool_count"` // 工具数量 ToolCount int `json:"tool_count"` // 工具数量
Error string `json:"error,omitempty"` // 错误信息(仅在status为error时存在) Error string `json:"error,omitempty"` // 错误信息(仅在status为error时存在)
} }
+2 -2
View File
@@ -11,6 +11,7 @@ import (
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"go.uber.org/zap" "go.uber.org/zap"
) )
@@ -187,7 +188,7 @@ func TestExternalMCPHandler_AddOrUpdateExternalMCP_InvalidConfig(t *testing.T) {
{ {
name: "缺少command和url", name: "缺少command和url",
configJSON: `{"enabled": true}`, configJSON: `{"enabled": true}`,
expectedErr: "需要指定commandstdio模式)或urlhttp模式)", expectedErr: "需要指定commandstdio模式)或urlhttp/sse模式)",
}, },
{ {
name: "stdio模式缺少command", name: "stdio模式缺少command",
@@ -515,4 +516,3 @@ func TestExternalMCPHandler_UpdateExistingConfig(t *testing.T) {
t.Errorf("期望command为空,实际%s", response.Config.Command) t.Errorf("期望command为空,实际%s", response.Config.Command)
} }
} }
File diff suppressed because it is too large Load Diff
+551
View File
@@ -0,0 +1,551 @@
// Package mcp 外部 MCP 客户端 - 基于官方 go-sdk 实现,保证协议兼容性
package mcp
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"os/exec"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/config"
"github.com/google/uuid"
"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
inner ExternalMCPClient // 连接成功后为 *sdkClient
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()
inner, err := createSDKClient(ctx, c.serverCfg, c.logger)
if err != nil {
c.setStatus("error")
return err
}
c.mu.Lock()
c.inner = inner
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
c.inner = nil
c.mu.Unlock()
c.setStatus("disconnected")
if inner != nil {
return inner.Close()
}
return nil
}
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
}
// simpleHTTPClient 简单 JSON-RPC over HTTP:每次请求一次 POST、响应在 body。实现 ExternalMCPClient。
// 用于自建 MCP(如 http://127.0.0.1:8081/mcp)或其它仅支持简单 POST 的端点。
type simpleHTTPClient struct {
url string
client *http.Client
logger *zap.Logger
mu sync.RWMutex
status string
}
func newSimpleHTTPClient(ctx context.Context, url string, timeout time.Duration, headers map[string]string, logger *zap.Logger) (ExternalMCPClient, error) {
c := &simpleHTTPClient{
url: url,
client: httpClientWithTimeoutAndHeaders(timeout, headers),
logger: logger,
status: "connecting",
}
if err := c.initialize(ctx); err != nil {
return nil, err
}
c.mu.Lock()
c.status = "connected"
c.mu.Unlock()
return c, nil
}
func (c *simpleHTTPClient) setStatus(s string) {
c.mu.Lock()
defer c.mu.Unlock()
c.status = s
}
func (c *simpleHTTPClient) GetStatus() string {
c.mu.RLock()
defer c.mu.RUnlock()
return c.status
}
func (c *simpleHTTPClient) IsConnected() bool {
return c.GetStatus() == "connected"
}
func (c *simpleHTTPClient) Initialize(context.Context) error {
return nil // 已在 newSimpleHTTPClient 中完成
}
func (c *simpleHTTPClient) initialize(ctx context.Context) error {
params := InitializeRequest{
ProtocolVersion: ProtocolVersion,
Capabilities: make(map[string]interface{}),
ClientInfo: ClientInfo{Name: clientName, Version: clientVersion},
}
paramsJSON, _ := json.Marshal(params)
req := &Message{
ID: MessageID{value: "1"},
Method: "initialize",
Version: "2.0",
Params: paramsJSON,
}
resp, err := c.sendRequest(ctx, req)
if err != nil {
return fmt.Errorf("initialize: %w", err)
}
if resp.Error != nil {
return fmt.Errorf("initialize: %s (code %d)", resp.Error.Message, resp.Error.Code)
}
// 发送 notifications/initialized(协议要求)
notify := &Message{
ID: MessageID{value: nil},
Method: "notifications/initialized",
Version: "2.0",
Params: json.RawMessage("{}"),
}
_ = c.sendNotification(notify)
return nil
}
func (c *simpleHTTPClient) sendRequest(ctx context.Context, msg *Message) (*Message, error) {
body, err := json.Marshal(msg)
if err != nil {
return nil, err
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
if err != nil {
return nil, err
}
httpReq.Header.Set("Content-Type", "application/json")
resp, err := c.client.Do(httpReq)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
b, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(b))
}
var out Message
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
return nil, err
}
return &out, nil
}
func (c *simpleHTTPClient) sendNotification(msg *Message) error {
body, _ := json.Marshal(msg)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
httpReq, _ := http.NewRequestWithContext(ctx, http.MethodPost, c.url, bytes.NewReader(body))
httpReq.Header.Set("Content-Type", "application/json")
resp, err := c.client.Do(httpReq)
if err != nil {
return err
}
resp.Body.Close()
return nil
}
func (c *simpleHTTPClient) ListTools(ctx context.Context) ([]Tool, error) {
req := &Message{
ID: MessageID{value: uuid.New().String()},
Method: "tools/list",
Version: "2.0",
Params: json.RawMessage("{}"),
}
resp, err := c.sendRequest(ctx, req)
if err != nil {
return nil, err
}
if resp.Error != nil {
return nil, fmt.Errorf("tools/list: %s (code %d)", resp.Error.Message, resp.Error.Code)
}
var listResp ListToolsResponse
if err := json.Unmarshal(resp.Result, &listResp); err != nil {
return nil, err
}
return listResp.Tools, nil
}
func (c *simpleHTTPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
params := CallToolRequest{Name: name, Arguments: args}
paramsJSON, _ := json.Marshal(params)
req := &Message{
ID: MessageID{value: uuid.New().String()},
Method: "tools/call",
Version: "2.0",
Params: paramsJSON,
}
resp, err := c.sendRequest(ctx, req)
if err != nil {
return nil, err
}
if resp.Error != nil {
return nil, fmt.Errorf("tools/call: %s (code %d)", resp.Error.Message, resp.Error.Code)
}
var callResp CallToolResponse
if err := json.Unmarshal(resp.Result, &callResp); err != nil {
return nil, err
}
return &ToolResult{Content: callResp.Content, IsError: callResp.IsError}, nil
}
func (c *simpleHTTPClient) Close() error {
c.setStatus("disconnected")
return nil
}
// 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.Transport
if transport == "" {
if serverCfg.Command != "" {
transport = "stdio"
} else if serverCfg.URL != "" {
transport = "http"
} else {
return nil, fmt.Errorf("配置缺少 command 或 url")
}
}
client := mcp.NewClient(&mcp.Implementation{
Name: clientName,
Version: clientVersion,
}, nil)
var t mcp.Transport
switch transport {
case "stdio":
if serverCfg.Command == "" {
return nil, fmt.Errorf("stdio 模式需要配置 command")
}
// 必须用 exec.Command 而非 CommandContextdoConnect 返回后 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)...)
}
t = &mcp.CommandTransport{Command: cmd}
case "sse":
if serverCfg.URL == "" {
return nil, fmt.Errorf("sse 模式需要配置 url")
}
httpClient := httpClientWithTimeoutAndHeaders(timeout, 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)
t = &mcp.StreamableClientTransport{
Endpoint: serverCfg.URL,
HTTPClient: httpClient,
}
case "simple_http":
// 简单 JSON-RPC HTTP:每次请求一次 POST、响应在 body。用于自建 MCP 或兼容旧端点(如 http://127.0.0.1:8081/mcp
if serverCfg.URL == "" {
return nil, fmt.Errorf("simple_http 模式需要配置 url")
}
return newSimpleHTTPClient(ctx, serverCfg.URL, timeout, serverCfg.Headers, logger)
default:
return nil, fmt.Errorf("不支持的传输模式: %s", 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,
}
}
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)
}
+200 -32
View File
@@ -25,6 +25,8 @@ type ExternalMCPManager struct {
errors map[string]string // 错误信息 errors map[string]string // 错误信息
toolCounts map[string]int // 工具数量缓存 toolCounts map[string]int // 工具数量缓存
toolCountsMu sync.RWMutex // 工具数量缓存的锁 toolCountsMu sync.RWMutex // 工具数量缓存的锁
toolCache map[string][]Tool // 工具列表缓存:MCP名称 -> 工具列表
toolCacheMu sync.RWMutex // 工具列表缓存的锁
stopRefresh chan struct{} // 停止后台刷新的信号 stopRefresh chan struct{} // 停止后台刷新的信号
refreshWg sync.WaitGroup // 等待后台刷新goroutine完成 refreshWg sync.WaitGroup // 等待后台刷新goroutine完成
mu sync.RWMutex mu sync.RWMutex
@@ -46,6 +48,7 @@ func NewExternalMCPManagerWithStorage(logger *zap.Logger, storage MonitorStorage
stats: make(map[string]*ToolStats), stats: make(map[string]*ToolStats),
errors: make(map[string]string), errors: make(map[string]string),
toolCounts: make(map[string]int), toolCounts: make(map[string]int),
toolCache: make(map[string][]Tool),
stopRefresh: make(chan struct{}), stopRefresh: make(chan struct{}),
} }
// 启动后台刷新工具数量的goroutine // 启动后台刷新工具数量的goroutine
@@ -119,6 +122,11 @@ func (m *ExternalMCPManager) RemoveConfig(name string) error {
delete(m.toolCounts, name) delete(m.toolCounts, name)
m.toolCountsMu.Unlock() m.toolCountsMu.Unlock()
// 清理工具列表缓存
m.toolCacheMu.Lock()
delete(m.toolCache, name)
m.toolCacheMu.Unlock()
return nil return nil
} }
@@ -196,8 +204,15 @@ func (m *ExternalMCPManager) StartClient(name string) error {
m.mu.Lock() m.mu.Lock()
delete(m.errors, name) delete(m.errors, name)
m.mu.Unlock() m.mu.Unlock()
// 连接成功,立即刷新工具数量 // 立即刷新工具数量和工具列表缓存
m.triggerToolCountRefresh() m.triggerToolCountRefresh()
m.refreshToolCache(name, client)
// 2 秒后再刷新一次,覆盖 SSE/Streamable 等需稍等就绪的远端
go func() {
time.Sleep(2 * time.Second)
m.triggerToolCountRefresh()
m.refreshToolCache(name, client)
}()
} }
}() }()
@@ -253,6 +268,11 @@ func (m *ExternalMCPManager) GetError(name string) string {
} }
// GetAllTools 获取所有外部MCP的工具 // GetAllTools 获取所有外部MCP的工具
// 优先从已连接的客户端获取,如果连接断开则返回缓存的工具列表
// 策略:
// - error 状态:不使用缓存,直接跳过(配置错误或服务不可用)
// - disconnected/connecting 状态:使用缓存(临时断开)
// - connected 状态:正常获取,失败时降级使用缓存
func (m *ExternalMCPManager) GetAllTools(ctx context.Context) ([]Tool, error) { func (m *ExternalMCPManager) GetAllTools(ctx context.Context) ([]Tool, error) {
m.mu.RLock() m.mu.RLock()
clients := make(map[string]ExternalMCPClient) clients := make(map[string]ExternalMCPClient)
@@ -262,17 +282,21 @@ func (m *ExternalMCPManager) GetAllTools(ctx context.Context) ([]Tool, error) {
m.mu.RUnlock() m.mu.RUnlock()
var allTools []Tool var allTools []Tool
for name, client := range clients { var hasError bool
if !client.IsConnected() { var lastError error
continue
}
tools, err := client.ListTools(ctx) // 使用较短的超时时间进行快速检查(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 { if err != nil {
m.logger.Warn("获取外部MCP工具列表失败", // 记录错误,但继续处理其他客户端
zap.String("name", name), hasError = true
zap.Error(err), if lastError == nil {
) lastError = err
}
continue continue
} }
@@ -283,9 +307,97 @@ func (m *ExternalMCPManager) GetAllTools(ctx context.Context) ([]Tool, error) {
} }
} }
// 如果有错误但至少返回了一些工具,不返回错误(部分成功)
if hasError && len(allTools) == 0 {
return nil, fmt.Errorf("获取外部MCP工具失败: %w", lastError)
}
return allTools, nil 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)
}
// 已连接:尝试获取最新工具列表
if client.IsConnected() {
tools, err := client.ListTools(ctx)
if err != nil {
// 获取失败,尝试使用缓存
return m.getCachedTools(name, "连接正常但获取失败", err)
}
// 获取成功,更新缓存
m.updateToolCache(name, tools)
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) {
m.toolCacheMu.RLock()
cachedTools, hasCache := m.toolCache[name]
m.toolCacheMu.RUnlock()
if hasCache && len(cachedTools) > 0 {
m.logger.Debug("使用缓存的工具列表",
zap.String("name", name),
zap.String("reason", reason),
zap.Int("count", len(cachedTools)),
zap.Error(originalErr),
)
return cachedTools, nil
}
// 无缓存,返回错误
if originalErr != nil {
return nil, fmt.Errorf("获取外部MCP工具失败且无缓存: %w", originalErr)
}
return nil, fmt.Errorf("外部MCP无缓存工具: %s", name)
}
// updateToolCache 更新工具列表缓存
func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
m.toolCacheMu.Lock()
m.toolCache[name] = tools
m.toolCacheMu.Unlock()
// 如果返回空列表,记录警告
if len(tools) == 0 {
m.logger.Warn("外部MCP返回空工具列表",
zap.String("name", name),
zap.String("hint", "服务可能暂时不可用,工具列表为空"),
)
} else {
m.logger.Debug("工具列表缓存已更新",
zap.String("name", name),
zap.Int("count", len(tools)),
)
}
}
// CallTool 调用外部MCP工具(返回执行ID) // CallTool 调用外部MCP工具(返回执行ID)
func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) { func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
// 解析工具名称:name::toolName // 解析工具名称:name::toolName
@@ -302,8 +414,18 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName) return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
} }
// 检查连接状态,如果未连接或状态为error,不允许调用
if !client.IsConnected() { if !client.IsConnected() {
return nil, "", fmt.Errorf("外部MCP客户端未连接: %s", mcpName) 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)
} }
// 创建执行记录 // 创建执行记录
@@ -630,11 +752,20 @@ func (m *ExternalMCPManager) refreshToolCounts() {
cancel() cancel()
if err != nil { if err != nil {
m.logger.Debug("获取外部MCP工具数量失败", errStr := err.Error()
// SSE 连接 EOF:远端可能关闭了流或未按规范在流上推送响应,仅首次用 Warn 提示
if strings.Contains(errStr, "EOF") || strings.Contains(errStr, "client is closing") {
m.logger.Warn("获取外部MCP工具数量失败(SSE 流已关闭或服务端未在流上返回 tools/list 响应)",
zap.String("name", n),
zap.String("hint", "若为 SSE 连接,请确认服务端保持 GET 流打开并按 MCP 规范以 event: message 推送 JSON-RPC 响应"),
zap.Error(err),
)
} else {
m.logger.Warn("获取外部MCP工具数量失败,请检查连接或服务端 tools/list",
zap.String("name", n), zap.String("name", n),
zap.Error(err), zap.Error(err),
) )
// 如果获取失败,保留旧值(在更新时处理) }
resultChan <- countResult{name: n, count: -1} // -1 表示使用旧值 resultChan <- countResult{name: n, count: -1} // -1 表示使用旧值
return return
} }
@@ -680,6 +811,40 @@ func (m *ExternalMCPManager) refreshToolCounts() {
m.toolCountsMu.Unlock() m.toolCountsMu.Unlock()
} }
// refreshToolCache 刷新指定MCP的工具列表缓存
func (m *ExternalMCPManager) refreshToolCache(name string, client ExternalMCPClient) {
if !client.IsConnected() {
return
}
// 检查状态,如果是error状态,不更新缓存
status := client.GetStatus()
if status == "error" {
m.logger.Debug("跳过刷新工具列表缓存(连接失败)",
zap.String("name", name),
zap.String("status", status),
)
return
}
// 使用较短的超时时间(5秒)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
tools, err := client.ListTools(ctx)
if err != nil {
m.logger.Debug("刷新工具列表缓存失败",
zap.String("name", name),
zap.Error(err),
)
// 刷新失败时不更新缓存,保留旧缓存(如果有)
return
}
// 使用统一的缓存更新方法
m.updateToolCache(name, tools)
}
// startToolCountRefresh 启动后台刷新工具数量的goroutine // startToolCountRefresh 启动后台刷新工具数量的goroutine
func (m *ExternalMCPManager) startToolCountRefresh() { func (m *ExternalMCPManager) startToolCountRefresh() {
m.refreshWg.Add(1) m.refreshWg.Add(1)
@@ -707,21 +872,13 @@ func (m *ExternalMCPManager) triggerToolCountRefresh() {
go m.refreshToolCounts() go m.refreshToolCounts()
} }
// createClient 创建客户端(不连接) // createClient 创建客户端(不连接)。统一使用官方 MCP Go SDK 的 lazy 客户端,连接在 Initialize 时完成。
func (m *ExternalMCPManager) createClient(serverCfg config.ExternalMCPServerConfig) ExternalMCPClient { func (m *ExternalMCPManager) createClient(serverCfg config.ExternalMCPServerConfig) ExternalMCPClient {
timeout := time.Duration(serverCfg.Timeout) * time.Second
if timeout <= 0 {
timeout = 30 * time.Second
}
// 根据传输模式创建客户端
transport := serverCfg.Transport transport := serverCfg.Transport
if transport == "" { if transport == "" {
// 如果没有指定transport,根据是否有command或url判断
if serverCfg.Command != "" { if serverCfg.Command != "" {
transport = "stdio" transport = "stdio"
} else if serverCfg.URL != "" { } else if serverCfg.URL != "" {
// 默认使用http,但可以通过transport字段指定sse
transport = "http" transport = "http"
} else { } else {
return nil return nil
@@ -733,17 +890,23 @@ func (m *ExternalMCPManager) createClient(serverCfg config.ExternalMCPServerConf
if serverCfg.URL == "" { if serverCfg.URL == "" {
return nil return nil
} }
return NewHTTPMCPClient(serverCfg.URL, timeout, m.logger) return newLazySDKClient(serverCfg, m.logger)
case "simple_http":
// 简单 HTTP(一次 POST 一次响应),用于自建 MCP 等
if serverCfg.URL == "" {
return nil
}
return newLazySDKClient(serverCfg, m.logger)
case "stdio": case "stdio":
if serverCfg.Command == "" { if serverCfg.Command == "" {
return nil return nil
} }
return NewStdioMCPClient(serverCfg.Command, serverCfg.Args, serverCfg.Env, timeout, m.logger) return newLazySDKClient(serverCfg, m.logger)
case "sse": case "sse":
if serverCfg.URL == "" { if serverCfg.URL == "" {
return nil return nil
} }
return NewSSEMCPClient(serverCfg.URL, timeout, m.logger) return newLazySDKClient(serverCfg, m.logger)
default: default:
return nil return nil
} }
@@ -773,12 +936,7 @@ func (m *ExternalMCPManager) doConnect(name string, serverCfg config.ExternalMCP
// setClientStatus 设置客户端状态(通过类型断言) // setClientStatus 设置客户端状态(通过类型断言)
func (m *ExternalMCPManager) setClientStatus(client ExternalMCPClient, status string) { func (m *ExternalMCPManager) setClientStatus(client ExternalMCPClient, status string) {
switch c := client.(type) { if c, ok := client.(*lazySDKClient); ok {
case *HTTPMCPClient:
c.setStatus(status)
case *StdioMCPClient:
c.setStatus(status)
case *SSEMCPClient:
c.setStatus(status) c.setStatus(status)
} }
} }
@@ -819,8 +977,13 @@ func (m *ExternalMCPManager) connectClient(name string, serverCfg config.Externa
zap.String("name", name), zap.String("name", name),
) )
// 连接成功,触发工具数量刷新 // 连接成功,触发工具数量刷新和工具列表缓存刷新
m.triggerToolCountRefresh() m.triggerToolCountRefresh()
m.mu.RLock()
if client, exists := m.clients[name]; exists {
m.refreshToolCache(name, client)
}
m.mu.RUnlock()
return nil return nil
} }
@@ -926,6 +1089,11 @@ func (m *ExternalMCPManager) StopAll() {
m.toolCounts = make(map[string]int) m.toolCounts = make(map[string]int)
m.toolCountsMu.Unlock() m.toolCountsMu.Unlock()
// 清理所有工具列表缓存
m.toolCacheMu.Lock()
m.toolCache = make(map[string][]Tool)
m.toolCacheMu.Unlock()
// 停止后台刷新(使用 select 避免重复关闭 channel // 停止后台刷新(使用 select 避免重复关闭 channel
select { select {
case <-m.stopRefresh: case <-m.stopRefresh:
+15 -37
View File
@@ -151,48 +151,26 @@ func TestExternalMCPManager_LoadConfigs(t *testing.T) {
} }
} }
func TestHTTPMCPClient_Initialize(t *testing.T) { // TestLazySDKClient_InitializeFails 验证无效配置时 SDK 客户端 Initialize 失败并设置 error 状态
// 注意:这个测试需要一个真实的HTTP MCP服务器 func TestLazySDKClient_InitializeFails(t *testing.T) {
// 如果没有服务器,这个测试会失败
// 在实际测试中,可以使用mock服务器
logger := zap.NewNop() logger := zap.NewNop()
client := NewHTTPMCPClient("http://127.0.0.1:8081/mcp", 5*time.Second, logger) // 使用不存在的 HTTP 地址,Initialize 应失败
cfg := config.ExternalMCPServerConfig{
Transport: "http",
URL: "http://127.0.0.1:19999/nonexistent",
Timeout: 2,
}
c := newLazySDKClient(cfg, logger)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel() defer cancel()
err := c.Initialize(ctx)
// 这个测试可能会失败,如果没有真实的服务器 if err == nil {
// 在实际环境中,应该使用mock服务器 t.Fatal("expected error when connecting to invalid server")
err := client.Initialize(ctx)
if err != nil {
t.Logf("初始化失败(可能是没有服务器): %v", err)
} }
if c.GetStatus() != "error" {
status := client.GetStatus() t.Errorf("expected status error, got %s", c.GetStatus())
if status == "" {
t.Error("状态不应该为空")
} }
c.Close()
client.Close()
}
func TestStdioMCPClient_Initialize(t *testing.T) {
// 注意:这个测试需要一个真实的stdio MCP服务器
// 如果没有服务器,这个测试会失败
logger := zap.NewNop()
client := NewStdioMCPClient("echo", []string{"test"}, nil, 5*time.Second, logger)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
// 这个测试可能会失败,因为echo不是MCP服务器
// 在实际环境中,应该使用真实的MCP服务器或mock
err := client.Initialize(ctx)
if err != nil {
t.Logf("初始化失败(echo不是MCP服务器): %v", err)
}
client.Close()
} }
func TestExternalMCPManager_StartStopClient(t *testing.T) { func TestExternalMCPManager_StartStopClient(t *testing.T) {
+76 -10
View File
@@ -1,6 +1,7 @@
package mcp package mcp
import ( import (
"bufio"
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
@@ -125,6 +126,13 @@ func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) {
return 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) body, err := io.ReadAll(r.Body)
if err != nil { if err != nil {
s.sendError(w, nil, -32700, "Parse error", err.Error()) s.sendError(w, nil, -32700, "Parse error", err.Error())
@@ -137,14 +145,56 @@ func (s *Server) HandleHTTP(w http.ResponseWriter, r *http.Request) {
return return
} }
// 处理消息
response := s.handleMessage(&msg) response := s.handleMessage(&msg)
w.Header().Set("Content-Type", "application/json") w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(response) json.NewEncoder(w).Encode(response)
} }
// handleSSE 处理SSE连接(用于MCP HTTP传输的事件通道) // 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(&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: endpointdata 为客户端 POST 消息的 URL(含 sessionid
// 2. 后续事件为 event: messagedata 为 JSON-RPC 响应
func (s *Server) handleSSE(w http.ResponseWriter, r *http.Request) { func (s *Server) handleSSE(w http.ResponseWriter, r *http.Request) {
flusher, ok := w.(http.Flusher) flusher, ok := w.(http.Flusher)
if !ok { if !ok {
@@ -157,16 +207,25 @@ func (s *Server) handleSSE(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Connection", "keep-alive") w.Header().Set("Connection", "keep-alive")
w.Header().Set("X-Accel-Buffering", "no") w.Header().Set("X-Accel-Buffering", "no")
sessionID := uuid.New().String()
client := &sseClient{ client := &sseClient{
id: uuid.New().String(), id: sessionID,
send: make(chan []byte, 8), send: make(chan []byte, 32),
} }
s.addSSEClient(client) s.addSSEClient(client)
defer s.removeSSEClient(client.id) defer s.removeSSEClient(client.id)
// 发送初始ready事件,告知客户端连接成功 // 官方规范:首个事件为 endpoint,data 为消息端点 URL(客户端将向该 URL POST 请求)
fmt.Fprintf(w, "event: message\ndata: {\"type\":\"ready\",\"status\":\"ok\"}\n\n") 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() flusher.Flush()
ticker := time.NewTicker(15 * time.Second) ticker := time.NewTicker(15 * time.Second)
@@ -183,7 +242,6 @@ func (s *Server) handleSSE(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "event: message\ndata: %s\n\n", msg) fmt.Fprintf(w, "event: message\ndata: %s\n\n", msg)
flusher.Flush() flusher.Flush()
case <-ticker.C: case <-ticker.C:
// 心跳保持连接
fmt.Fprintf(w, ": ping\n\n") fmt.Fprintf(w, ": ping\n\n")
flusher.Flush() flusher.Flush()
} }
@@ -311,6 +369,7 @@ func (s *Server) handleListTools(msg *Message) *Message {
tools = append(tools, tool) tools = append(tools, tool)
} }
s.mu.RUnlock() s.mu.RUnlock()
s.logger.Debug("tools/list 请求", zap.Int("返回工具数", len(tools)))
response := ListToolsResponse{Tools: tools} response := ListToolsResponse{Tools: tools}
result, _ := json.Marshal(response) result, _ := json.Marshal(response)
@@ -1110,10 +1169,11 @@ func (s *Server) RegisterResource(resource *Resource) {
} }
// HandleStdio 处理标准输入输出(用于 stdio 传输模式) // HandleStdio 处理标准输入输出(用于 stdio 传输模式)
// MCP 协议使用换行分隔的 JSON-RPC 消息 // MCP 协议使用换行分隔的 JSON-RPC 消息;管道下需每次写入后 Flush,否则客户端会读不到响应
func (s *Server) HandleStdio() error { func (s *Server) HandleStdio() error {
decoder := json.NewDecoder(os.Stdin) decoder := json.NewDecoder(os.Stdin)
encoder := json.NewEncoder(os.Stdout) stdout := bufio.NewWriter(os.Stdout)
encoder := json.NewEncoder(stdout)
// 注意:不设置缩进,MCP 协议期望紧凑的 JSON 格式 // 注意:不设置缩进,MCP 协议期望紧凑的 JSON 格式
for { for {
@@ -1134,6 +1194,9 @@ func (s *Server) HandleStdio() error {
if err := encoder.Encode(errorMsg); err != nil { if err := encoder.Encode(errorMsg); err != nil {
return fmt.Errorf("发送错误响应失败: %w", err) return fmt.Errorf("发送错误响应失败: %w", err)
} }
if err := stdout.Flush(); err != nil {
return fmt.Errorf("刷新 stdout 失败: %w", err)
}
continue continue
} }
@@ -1149,6 +1212,9 @@ func (s *Server) HandleStdio() error {
if err := encoder.Encode(response); err != nil { if err := encoder.Encode(response); err != nil {
return fmt.Errorf("发送响应失败: %w", err) return fmt.Errorf("发送响应失败: %w", err)
} }
if err := stdout.Flush(); err != nil {
return fmt.Errorf("刷新 stdout 失败: %w", err)
}
} }
return nil return nil
+11 -1
View File
@@ -1,11 +1,22 @@
package mcp package mcp
import ( import (
"context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"time" "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消息类型 // MCP消息类型
const ( const (
MessageTypeRequest = "request" MessageTypeRequest = "request"
@@ -282,4 +293,3 @@ type SamplingContent struct {
Type string `json:"type"` Type string `json:"type"`
Text string `json:"text,omitempty"` Text string `json:"text,omitempty"`
} }
+43 -8
View File
@@ -224,22 +224,25 @@ func (e *Executor) RegisterTools(mcpServer *mcp.Server) {
toolName := toolConfig.Name toolName := toolConfig.Name
toolConfigCopy := toolConfig toolConfigCopy := toolConfig
// 使用简短描述(如果存在),否则使用详细描述的前100个字符 // 根据配置决定暴露给 AI/API 的描述:short_description 或 description
useFullDescription := strings.TrimSpace(strings.ToLower(e.config.ToolDescriptionMode)) == "full"
shortDesc := toolConfigCopy.ShortDescription shortDesc := toolConfigCopy.ShortDescription
if shortDesc == "" { if shortDesc == "" {
// 如果没有简短描述,从详细描述中提取第一行或前100个字符 // 如果没有简短描述,从详细描述中提取第一行或前10000个字符
desc := toolConfigCopy.Description desc := toolConfigCopy.Description
if len(desc) > 100 { if len(desc) > 10000 {
// 尝试找到第一个换行符 if idx := strings.Index(desc, "\n"); idx > 0 && idx < 10000 {
if idx := strings.Index(desc, "\n"); idx > 0 && idx < 100 {
shortDesc = strings.TrimSpace(desc[:idx]) shortDesc = strings.TrimSpace(desc[:idx])
} else { } else {
shortDesc = desc[:100] + "..." shortDesc = desc[:10000] + "..."
} }
} else { } else {
shortDesc = desc shortDesc = desc
} }
} }
if useFullDescription {
shortDesc = "" // 使用 description 时清空 ShortDescription,下游会回退到 Description
}
tool := mcp.Tool{ tool := mcp.Tool{
Name: toolConfigCopy.Name, Name: toolConfigCopy.Name,
@@ -303,7 +306,23 @@ func (e *Executor) buildCommandArgs(toolName string, toolConfig *config.ToolConf
} }
} }
// 先处理标志参数(对于大多数命令,标志应该在位置参数之前 // 对于需要子命令的工具(如 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 { for _, param := range flagParams {
// 跳过特殊参数,它们会在后面单独处理 // 跳过特殊参数,它们会在后面单独处理
@@ -416,7 +435,11 @@ func (e *Executor) buildCommandArgs(toolName string, toolConfig *config.ToolConf
} }
// 按位置顺序处理参数,确保即使某些位置没有参数或使用默认值,也能正确传递 // 按位置顺序处理参数,确保即使某些位置没有参数或使用默认值,也能正确传递
// position 0 已在前面插入(子命令优先),此处从 1 开始
for i := 0; i <= maxPosition; i++ { for i := 0; i <= maxPosition; i++ {
if i == 0 {
continue
}
for _, param := range positionalParams { for _, param := range positionalParams {
// 跳过特殊参数,它们会在后面单独处理 // 跳过特殊参数,它们会在后面单独处理
// action 参数仅用于工具内部逻辑,不传递给命令 // action 参数仅用于工具内部逻辑,不传递给命令
@@ -1190,7 +1213,15 @@ func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]in
required := []string{} required := []string{}
for _, param := range toolConfig.Parameters { for _, param := range toolConfig.Parameters {
// 转换类型为OpenAI/JSON Schema标准类型 // 跳过 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) openAIType := e.convertToOpenAIType(param.Type)
prop := map[string]interface{}{ prop := map[string]interface{}{
@@ -1232,6 +1263,10 @@ func (e *Executor) buildInputSchema(toolConfig *config.ToolConfig) map[string]in
// convertToOpenAIType 将配置中的类型转换为OpenAI/JSON Schema标准类型 // convertToOpenAIType 将配置中的类型转换为OpenAI/JSON Schema标准类型
func (e *Executor) convertToOpenAIType(configType string) string { func (e *Executor) convertToOpenAIType(configType string) string {
// 空或 null 类型统一视为 string,避免非法 schema 导致工具调用失败
if strings.TrimSpace(configType) == "" {
return "string"
}
switch configType { switch configType {
case "bool": case "bool":
return "boolean" return "boolean"
+3 -3
View File
@@ -1,8 +1,8 @@
# Python HTTP helpers leveraged by tools like api-fuzzer, dnslog, http-intruder, http-framework-test # Python HTTP helpers leveraged by tools like api-fuzzer, dnslog, http-intruder, http-framework-test
requests>=2.32.3 requests>=2.32.3
httpx>=0.27.0
charset-normalizer>=3.3.2 # dirsearch:用 python3 -m dirsearch 时由本依赖提供(含 defusedxml 等)
chardet>=5.2.0 dirsearch>=0.4.3
# Python exploitation / analysis frameworks referenced by tool recipes # Python exploitation / analysis frameworks referenced by tool recipes
# angr>=9.2.96 # angr>=9.2.96
+1 -1
View File
@@ -407,7 +407,7 @@ go run cmd/test-config/main.go
### Q: 参数顺序如何控制? ### Q: 参数顺序如何控制?
A: 使用 `position` 字段控制位置参数的顺序。标志参数会按照在 `parameters` 列表中的顺序添加`additional_args` 会追加到命令末尾。 A: 使用 `position` 字段控制位置参数的顺序。**位置 0 的参数(如 gobuster 的 `dir` 子命令)会紧跟在命令名后、所有标志参数之前**,以便兼容需要“子命令 + 选项”形式的 CLI。其余标志参数按在 `parameters` 列表中的顺序添加,再按 position 1、2… 添加其余位置参数`additional_args` 会追加到命令末尾。
## 工具配置模板 ## 工具配置模板
+370 -2
View File
@@ -488,8 +488,26 @@ header {
font-size: 1.5rem; font-size: 1.5rem;
font-weight: 600; font-weight: 600;
letter-spacing: -0.5px; letter-spacing: -0.5px;
margin: 0; }
color: var(--text-primary);
.header-logo-link {
cursor: pointer;
transition: opacity 0.2s ease;
}
.header-logo-link:hover {
opacity: 0.85;
}
.version-badge {
display: inline-block;
margin-left: 6px;
font-size: 0.6875rem;
font-weight: 400;
color: var(--text-muted);
letter-spacing: 0.02em;
vertical-align: 0.35em;
user-select: none;
} }
.header-right { .header-right {
@@ -7920,6 +7938,356 @@ header {
color: var(--text-primary); color: var(--text-primary);
} }
/* 仪表盘页面样式 */
.dashboard-page {
height: 100%;
display: flex;
flex-direction: column;
overflow: auto;
background: linear-gradient(180deg, #f0f4ff 0%, #f8fafc 24%, var(--bg-secondary) 100%);
}
.dashboard-page .page-header {
flex-shrink: 0;
}
.dashboard-content {
flex: 1;
padding: 20px 24px 32px;
overflow: auto;
}
.dashboard-cards {
display: grid;
grid-template-columns: repeat(3, 1fr);
gap: 20px;
margin-bottom: 28px;
}
@media (max-width: 900px) {
.dashboard-cards { grid-template-columns: 1fr; }
}
.dashboard-card {
position: relative;
background: var(--bg-primary);
border-radius: 16px;
padding: 24px;
display: flex;
align-items: center;
gap: 20px;
cursor: pointer;
transition: transform 0.25s ease, box-shadow 0.25s ease;
box-shadow: 0 4px 14px rgba(0,0,0,0.06), 0 1px 3px rgba(0,0,0,0.04);
border: 1px solid rgba(0,0,0,0.06);
overflow: hidden;
}
.dashboard-card:hover {
transform: translateY(-4px);
box-shadow: 0 12px 28px rgba(0,0,0,0.1), 0 4px 12px rgba(0,0,0,0.06);
}
.dashboard-card-glow {
position: absolute;
top: -40%;
right: -20%;
width: 60%;
height: 100%;
border-radius: 50%;
opacity: 0.08;
pointer-events: none;
}
.dashboard-card-tasks .dashboard-card-glow { background: radial-gradient(circle, #0066ff 0%, transparent 70%); }
.dashboard-card-vulns .dashboard-card-glow { background: radial-gradient(circle, #dc3545 0%, transparent 70%); }
.dashboard-card-chat .dashboard-card-glow { background: radial-gradient(circle, #0d9488 0%, transparent 70%); }
.dashboard-card-icon {
width: 56px;
height: 56px;
border-radius: 14px;
display: flex;
align-items: center;
justify-content: center;
flex-shrink: 0;
transition: transform 0.25s ease;
}
.dashboard-card:hover .dashboard-card-icon {
transform: scale(1.08);
}
.dashboard-card-tasks .dashboard-card-icon {
background: linear-gradient(145deg, #e0edff 0%, #cce0ff 100%);
color: #0066ff;
}
.dashboard-card-vulns .dashboard-card-icon {
background: linear-gradient(145deg, #ffebee 0%, #ffcdd2 100%);
color: #c62828;
}
.dashboard-card-chat .dashboard-card-icon {
background: linear-gradient(145deg, #e0f2f1 0%, #b2dfdb 100%);
color: #0d9488;
}
.dashboard-card-body {
min-width: 0;
}
.dashboard-card-value {
font-size: 2rem;
font-weight: 800;
color: var(--text-primary);
line-height: 1.1;
letter-spacing: -0.03em;
}
.dashboard-card-label {
font-size: 0.9rem;
color: var(--text-secondary);
margin-top: 4px;
font-weight: 500;
}
.dashboard-card-desc {
font-size: 0.8rem;
color: var(--text-muted);
margin-top: 2px;
}
.dashboard-card-cta {
font-size: 1.15rem;
font-weight: 700;
color: var(--text-primary);
}
.dashboard-card-hint {
font-size: 0.75rem;
color: var(--text-muted);
margin-top: 6px;
opacity: 0.9;
}
.dashboard-section {
margin-bottom: 28px;
}
.dashboard-section-chart {
background: var(--bg-primary);
border-radius: 16px;
padding: 22px 24px;
box-shadow: 0 2px 12px rgba(0,0,0,0.04);
border: 1px solid rgba(0,0,0,0.05);
}
.dashboard-section-overview {
background: var(--bg-primary);
border-radius: 16px;
padding: 22px 24px;
box-shadow: 0 2px 12px rgba(0,0,0,0.04);
border: 1px solid rgba(0,0,0,0.05);
}
.dashboard-overview-grid {
display: grid;
grid-template-columns: repeat(auto-fit, minmax(240px, 1fr));
gap: 16px;
}
.dashboard-overview-item {
padding: 14px 18px;
background: var(--bg-secondary);
border-radius: 12px;
border: 1px solid var(--border-color);
cursor: pointer;
transition: border-color 0.2s, background 0.2s;
}
.dashboard-overview-item:hover {
border-color: var(--accent-color);
background: rgba(0, 102, 255, 0.04);
}
.dashboard-overview-label {
display: block;
font-size: 0.8rem;
font-weight: 600;
color: var(--text-secondary);
margin-bottom: 6px;
}
.dashboard-overview-value {
font-size: 0.9rem;
color: var(--text-primary);
}
.dashboard-section-title {
display: flex;
align-items: center;
gap: 10px;
font-size: 1rem;
font-weight: 600;
color: var(--text-primary);
margin: 0 0 16px 0;
}
.dashboard-section-dot {
width: 6px;
height: 6px;
border-radius: 50%;
background: var(--accent-color);
flex-shrink: 0;
}
.dashboard-chart-wrap {
display: flex;
flex-direction: column;
gap: 16px;
}
.dashboard-stacked-bar {
display: flex;
width: 100%;
height: 28px;
border-radius: 14px;
overflow: hidden;
background: var(--bg-tertiary);
box-shadow: inset 0 1px 2px rgba(0,0,0,0.06);
}
.dashboard-bar-seg {
height: 100%;
min-width: 2px;
transition: width 0.4s ease;
}
.dashboard-bar-seg.seg-critical { background: linear-gradient(90deg, #b91c1c, #dc2626); }
.dashboard-bar-seg.seg-high { background: linear-gradient(90deg, #ea580c, #f97316); }
.dashboard-bar-seg.seg-medium { background: linear-gradient(90deg, #ca8a04, #eab308); }
.dashboard-bar-seg.seg-low { background: linear-gradient(90deg, #059669, #10b981); }
.dashboard-bar-seg.seg-info { background: linear-gradient(90deg, #4b5563, #6b7280); }
.dashboard-legend {
display: flex;
flex-wrap: wrap;
gap: 20px 28px;
}
.dashboard-legend-item {
display: inline-flex;
align-items: center;
gap: 8px;
font-size: 0.875rem;
}
.dashboard-legend-dot {
width: 10px;
height: 10px;
border-radius: 50%;
flex-shrink: 0;
}
.dashboard-legend-dot.critical { background: #dc2626; }
.dashboard-legend-dot.high { background: #f97316; }
.dashboard-legend-dot.medium { background: #eab308; }
.dashboard-legend-dot.low { background: #10b981; }
.dashboard-legend-dot.info { background: #6b7280; }
.dashboard-legend-label {
color: var(--text-secondary);
}
.dashboard-legend-value {
font-weight: 700;
color: var(--text-primary);
min-width: 1.5em;
}
.dashboard-quick-links {
display: flex;
flex-wrap: wrap;
gap: 12px;
}
.dashboard-quick-link {
display: inline-flex;
align-items: center;
gap: 10px;
padding: 12px 20px;
background: var(--bg-primary);
border: 1px solid var(--border-color);
border-radius: 12px;
color: var(--text-primary);
font-size: 0.9rem;
font-weight: 500;
text-decoration: none;
cursor: pointer;
transition: all 0.2s ease;
box-shadow: 0 1px 4px rgba(0,0,0,0.04);
}
.dashboard-quick-link:hover {
border-color: var(--accent-color);
background: rgba(0, 102, 255, 0.06);
color: var(--accent-color);
box-shadow: 0 4px 12px rgba(0, 102, 255, 0.12);
}
.dashboard-quick-icon {
display: flex;
align-items: center;
justify-content: center;
color: var(--text-secondary);
}
.dashboard-quick-link:hover .dashboard-quick-icon {
color: var(--accent-color);
}
.dashboard-cta-block {
margin-top: 8px;
padding: 24px;
background: linear-gradient(135deg, #f0f9ff 0%, #e0f2fe 100%);
border: 1px solid rgba(14, 165, 233, 0.2);
border-radius: 16px;
}
.dashboard-cta-inner {
display: flex;
flex-wrap: wrap;
align-items: center;
justify-content: space-between;
gap: 16px;
}
.dashboard-cta-text {
font-size: 1rem;
font-weight: 600;
color: var(--text-primary);
margin: 0;
}
.dashboard-cta-btn {
padding: 12px 24px;
background: linear-gradient(135deg, #0284c7 0%, #0369a1 100%);
border: none;
border-radius: 12px;
color: #fff;
font-size: 0.95rem;
font-weight: 600;
cursor: pointer;
transition: transform 0.15s, box-shadow 0.2s;
box-shadow: 0 4px 14px rgba(2, 132, 199, 0.35);
}
.dashboard-cta-btn:hover {
transform: translateY(-2px);
box-shadow: 0 6px 20px rgba(2, 132, 199, 0.4);
}
.vulnerability-controls { .vulnerability-controls {
margin-bottom: 24px; margin-bottom: 24px;
} }
+118
View File
@@ -0,0 +1,118 @@
// 仪表盘页面:拉取运行中任务、漏洞统计、批量任务、工具与 Skills 统计并渲染
async function refreshDashboard() {
const runningEl = document.getElementById('dashboard-running-tasks');
const vulnTotalEl = document.getElementById('dashboard-vuln-total');
const severityIds = ['critical', 'high', 'medium', 'low', 'info'];
if (runningEl) runningEl.textContent = '…';
if (vulnTotalEl) vulnTotalEl.textContent = '…';
severityIds.forEach(s => {
const el = document.getElementById('dashboard-severity-' + s);
if (el) el.textContent = '0';
const barEl = document.getElementById('dashboard-bar-' + s);
if (barEl) barEl.style.width = '0%';
});
setDashboardOverviewPlaceholder('…');
if (typeof apiFetch === 'undefined') {
if (runningEl) runningEl.textContent = '-';
if (vulnTotalEl) vulnTotalEl.textContent = '-';
setDashboardOverviewPlaceholder('-');
return;
}
try {
const [tasksRes, vulnRes, batchRes, monitorRes, skillsRes] = await Promise.all([
apiFetch('/api/agent-loop/tasks').then(r => r.ok ? r.json() : null).catch(() => null),
apiFetch('/api/vulnerabilities/stats').then(r => r.ok ? r.json() : null).catch(() => null),
apiFetch('/api/batch-tasks?limit=500&page=1').then(r => r.ok ? r.json() : null).catch(() => null),
apiFetch('/api/monitor/stats').then(r => r.ok ? r.json() : null).catch(() => null),
apiFetch('/api/skills/stats').then(r => r.ok ? r.json() : null).catch(() => null)
]);
if (tasksRes && Array.isArray(tasksRes.tasks)) {
if (runningEl) runningEl.textContent = String(tasksRes.tasks.length);
} else {
if (runningEl) runningEl.textContent = '-';
}
if (vulnRes && typeof vulnRes.total === 'number') {
if (vulnTotalEl) vulnTotalEl.textContent = String(vulnRes.total);
const bySeverity = vulnRes.by_severity || {};
const total = vulnRes.total || 0;
severityIds.forEach(sev => {
const count = bySeverity[sev] || 0;
const el = document.getElementById('dashboard-severity-' + sev);
if (el) el.textContent = String(count);
const barEl = document.getElementById('dashboard-bar-' + sev);
if (barEl) barEl.style.width = total > 0 ? (count / total * 100) + '%' : '0%';
});
} else {
if (vulnTotalEl) vulnTotalEl.textContent = '-';
severityIds.forEach(sev => {
const barEl = document.getElementById('dashboard-bar-' + sev);
if (barEl) barEl.style.width = '0%';
});
}
// 批量任务队列:按状态统计
if (batchRes && Array.isArray(batchRes.queues)) {
const queues = batchRes.queues;
let pending = 0, running = 0, done = 0;
queues.forEach(q => {
const s = (q.status || '').toLowerCase();
if (s === 'pending' || s === 'paused') pending++;
else if (s === 'running') running++;
else if (s === 'completed' || s === 'cancelled') done++;
});
setEl('dashboard-batch-pending', String(pending));
setEl('dashboard-batch-running', String(running));
setEl('dashboard-batch-done', String(done));
} else {
setEl('dashboard-batch-pending', '-');
setEl('dashboard-batch-running', '-');
setEl('dashboard-batch-done', '-');
}
// 工具调用:monitor/stats 为 { toolName: { TotalCalls, ... } }
if (monitorRes && typeof monitorRes === 'object') {
const names = Object.keys(monitorRes);
let totalCalls = 0;
names.forEach(k => {
const v = monitorRes[k];
const n = v && (v.totalCalls ?? v.TotalCalls);
if (typeof n === 'number') totalCalls += n;
});
setEl('dashboard-tools-count', String(names.length));
setEl('dashboard-tools-calls', String(totalCalls));
} else {
setEl('dashboard-tools-count', '-');
setEl('dashboard-tools-calls', '-');
}
// Skills{ total_skills, total_calls, ... }
if (skillsRes && typeof skillsRes === 'object') {
setEl('dashboard-skills-count', String(skillsRes.total_skills ?? '-'));
setEl('dashboard-skills-calls', String(skillsRes.total_calls ?? '-'));
} else {
setEl('dashboard-skills-count', '-');
setEl('dashboard-skills-calls', '-');
}
} catch (e) {
console.warn('仪表盘拉取统计失败', e);
if (runningEl) runningEl.textContent = '-';
if (vulnTotalEl) vulnTotalEl.textContent = '-';
setDashboardOverviewPlaceholder('-');
}
}
function setEl(id, text) {
const el = document.getElementById(id);
if (el) el.textContent = text;
}
function setDashboardOverviewPlaceholder(t) {
['dashboard-batch-pending', 'dashboard-batch-running', 'dashboard-batch-done',
'dashboard-tools-count', 'dashboard-tools-calls', 'dashboard-skills-count', 'dashboard-skills-calls'].forEach(id => setEl(id, t));
}
+17 -3
View File
@@ -8,7 +8,7 @@ function initRouter() {
if (hash) { if (hash) {
const hashParts = hash.split('?'); const hashParts = hash.split('?');
const pageId = hashParts[0]; const pageId = hashParts[0];
if (pageId && ['chat', 'vulnerabilities', 'mcp-monitor', 'mcp-management', 'knowledge-management', 'knowledge-retrieval-logs', 'roles-management', 'skills-monitor', 'skills-management', 'settings', 'tasks'].includes(pageId)) { if (pageId && ['dashboard', 'chat', 'vulnerabilities', 'mcp-monitor', 'mcp-management', 'knowledge-management', 'knowledge-retrieval-logs', 'roles-management', 'skills-monitor', 'skills-management', 'settings', 'tasks'].includes(pageId)) {
switchPage(pageId); switchPage(pageId);
// 如果是chat页面且带有conversation参数,加载对应对话 // 如果是chat页面且带有conversation参数,加载对应对话
@@ -237,6 +237,11 @@ function showSubmenuPopup(navItem, menuId) {
// 初始化页面 // 初始化页面
function initPage(pageId) { function initPage(pageId) {
switch(pageId) { switch(pageId) {
case 'dashboard':
if (typeof refreshDashboard === 'function') {
refreshDashboard();
}
break;
case 'chat': case 'chat':
// 对话页面已由chat.js初始化 // 对话页面已由chat.js初始化
break; break;
@@ -254,16 +259,25 @@ function initPage(pageId) {
break; break;
case 'mcp-management': case 'mcp-management':
// 初始化MCP管理 // 初始化MCP管理
// 先加载外部MCP列表(快速),然后加载工具列表
if (typeof loadExternalMCPs === 'function') { if (typeof loadExternalMCPs === 'function') {
loadExternalMCPs(); loadExternalMCPs().catch(err => {
console.warn('加载外部MCP列表失败:', err);
});
} }
// 加载工具列表(MCP工具配置已移到MCP管理页面) // 加载工具列表(MCP工具配置已移到MCP管理页面)
// 使用异步加载,避免阻塞页面渲染
if (typeof loadToolsList === 'function') { if (typeof loadToolsList === 'function') {
// 确保工具分页设置已初始化 // 确保工具分页设置已初始化
if (typeof getToolsPageSize === 'function' && typeof toolsPagination !== 'undefined') { if (typeof getToolsPageSize === 'function' && typeof toolsPagination !== 'undefined') {
toolsPagination.pageSize = getToolsPageSize(); toolsPagination.pageSize = getToolsPageSize();
} }
loadToolsList(1, ''); // 延迟加载,让页面先渲染
setTimeout(() => {
loadToolsList(1, '').catch(err => {
console.error('加载工具列表失败:', err);
});
}, 100);
} }
break; break;
case 'vulnerabilities': case 'vulnerabilities':
+74 -17
View File
@@ -183,6 +183,14 @@ let toolsSearchKeyword = '';
// 加载工具列表(分页) // 加载工具列表(分页)
async function loadToolsList(page = 1, searchKeyword = '') { async function loadToolsList(page = 1, searchKeyword = '') {
const toolsList = document.getElementById('tools-list');
// 显示加载状态
if (toolsList) {
// 清空整个容器,包括可能存在的分页控件
toolsList.innerHTML = '<div class="tools-list-items"><div class="loading" style="padding: 20px; text-align: center; color: var(--text-muted);">⏳ 正在加载工具列表...</div></div>';
}
try { try {
// 在加载新页面之前,先保存当前页的状态到全局映射 // 在加载新页面之前,先保存当前页的状态到全局映射
saveCurrentPageToolStates(); saveCurrentPageToolStates();
@@ -193,7 +201,15 @@ async function loadToolsList(page = 1, searchKeyword = '') {
url += `&search=${encodeURIComponent(searchKeyword)}`; url += `&search=${encodeURIComponent(searchKeyword)}`;
} }
const response = await apiFetch(url); // 使用较短的超时时间(10秒),避免长时间等待
const controller = new AbortController();
const timeoutId = setTimeout(() => controller.abort(), 10000);
const response = await apiFetch(url, {
signal: controller.signal
});
clearTimeout(timeoutId);
if (!response.ok) { if (!response.ok) {
throw new Error('获取工具列表失败'); throw new Error('获取工具列表失败');
} }
@@ -224,9 +240,12 @@ async function loadToolsList(page = 1, searchKeyword = '') {
renderToolsPagination(); renderToolsPagination();
} catch (error) { } catch (error) {
console.error('加载工具列表失败:', error); console.error('加载工具列表失败:', error);
const toolsList = document.getElementById('tools-list');
if (toolsList) { if (toolsList) {
toolsList.innerHTML = `<div class="error">加载工具列表失败: ${escapeHtml(error.message)}</div>`; const isTimeout = error.name === 'AbortError' || error.message.includes('timeout');
const errorMsg = isTimeout
? '加载工具列表超时,可能是外部MCP连接较慢。请点击"刷新"按钮重试,或检查外部MCP连接状态。'
: `加载工具列表失败: ${escapeHtml(error.message)}`;
toolsList.innerHTML = `<div class="error" style="padding: 20px; text-align: center;">${errorMsg}</div>`;
} }
} }
} }
@@ -281,9 +300,21 @@ function renderToolsList() {
const toolsList = document.getElementById('tools-list'); const toolsList = document.getElementById('tools-list');
if (!toolsList) return; if (!toolsList) return;
// 只渲染列表部分,分页控件单独渲染 // 移除可能存在的分页控件(会在 renderToolsPagination 中重新添加)
const listContainer = toolsList.querySelector('.tools-list-items') || document.createElement('div'); const oldPagination = toolsList.querySelector('.tools-pagination');
if (oldPagination) {
oldPagination.remove();
}
// 获取或创建列表容器
let listContainer = toolsList.querySelector('.tools-list-items');
if (!listContainer) {
listContainer = document.createElement('div');
listContainer.className = 'tools-list-items'; listContainer.className = 'tools-list-items';
toolsList.appendChild(listContainer);
}
// 清空列表容器内容(移除加载提示)
listContainer.innerHTML = ''; listContainer.innerHTML = '';
if (allTools.length === 0) { if (allTools.length === 0) {
@@ -974,15 +1005,17 @@ async function changePassword() {
let currentEditingMCPName = null; let currentEditingMCPName = null;
// 加载外部MCP列表 // 拉取外部MCP列表数据(供轮询使用,返回 { servers, stats }
async function fetchExternalMCPs() {
const response = await apiFetch('/api/external-mcp');
if (!response.ok) throw new Error('获取外部MCP列表失败');
return response.json();
}
// 加载外部MCP列表并渲染
async function loadExternalMCPs() { async function loadExternalMCPs() {
try { try {
const response = await apiFetch('/api/external-mcp'); const data = await fetchExternalMCPs();
if (!response.ok) {
throw new Error('获取外部MCP列表失败');
}
const data = await response.json();
renderExternalMCPList(data.servers || {}); renderExternalMCPList(data.servers || {});
renderExternalMCPStats(data.stats || {}); renderExternalMCPStats(data.stats || {});
} catch (error) { } catch (error) {
@@ -994,6 +1027,29 @@ async function loadExternalMCPs() {
} }
} }
// 轮询列表直到指定 MCP 的工具数量已更新(每秒拉一次,拿到即停,无固定延迟)
// name 为 null 时仅按 maxAttempts 次数轮询,不判断 tool_count
async function pollExternalMCPToolCount(name, maxAttempts = 10) {
const pollIntervalMs = 1000;
for (let attempt = 0; attempt < maxAttempts; attempt++) {
await new Promise(r => setTimeout(r, pollIntervalMs));
try {
const data = await fetchExternalMCPs();
renderExternalMCPList(data.servers || {});
renderExternalMCPStats(data.stats || {});
if (name != null) {
const server = data.servers && data.servers[name];
if (server && server.tool_count > 0) break;
}
} catch (e) {
console.warn('轮询工具数量失败:', e);
}
}
if (typeof window !== 'undefined' && typeof window.refreshMentionTools === 'function') {
window.refreshMentionTools();
}
}
// 渲染外部MCP列表 // 渲染外部MCP列表
function renderExternalMCPList(servers) { function renderExternalMCPList(servers) {
const list = document.getElementById('external-mcp-list'); const list = document.getElementById('external-mcp-list');
@@ -1350,10 +1406,11 @@ async function saveExternalMCP() {
closeExternalMCPModal(); closeExternalMCPModal();
await loadExternalMCPs(); await loadExternalMCPs();
// 刷新对话界面的工具列表,使新添加的MCP工具立即可用
if (typeof window !== 'undefined' && typeof window.refreshMentionTools === 'function') { if (typeof window !== 'undefined' && typeof window.refreshMentionTools === 'function') {
window.refreshMentionTools(); window.refreshMentionTools();
} }
// 轮询几次以拉取后端异步更新的工具数量(无固定延迟,拿到即停)
pollExternalMCPToolCount(null, 5);
alert('保存成功'); alert('保存成功');
} catch (error) { } catch (error) {
console.error('保存外部MCP失败:', error); console.error('保存外部MCP失败:', error);
@@ -1427,12 +1484,12 @@ async function toggleExternalMCP(name, currentStatus) {
const status = statusData.status || 'disconnected'; const status = statusData.status || 'disconnected';
if (status === 'connected') { if (status === 'connected') {
// 已经连接,立即刷新
await loadExternalMCPs(); await loadExternalMCPs();
// 刷新对话界面的工具列表
if (typeof window !== 'undefined' && typeof window.refreshMentionTools === 'function') { if (typeof window !== 'undefined' && typeof window.refreshMentionTools === 'function') {
window.refreshMentionTools(); window.refreshMentionTools();
} }
// 轮询直到该 MCP 工具数量已更新(每秒拉一次,无固定延迟)
pollExternalMCPToolCount(name, 10);
return; return;
} }
} }
@@ -1490,12 +1547,12 @@ async function pollExternalMCPStatus(name, maxAttempts = 30) {
const button = document.getElementById(buttonId); const button = document.getElementById(buttonId);
if (status === 'connected') { if (status === 'connected') {
// 连接成功,刷新列表
await loadExternalMCPs(); await loadExternalMCPs();
// 刷新对话界面的工具列表
if (typeof window !== 'undefined' && typeof window.refreshMentionTools === 'function') { if (typeof window !== 'undefined' && typeof window.refreshMentionTools === 'function') {
window.refreshMentionTools(); window.refreshMentionTools();
} }
// 轮询直到该 MCP 工具数量已更新(每秒拉一次,无固定延迟)
pollExternalMCPToolCount(name, 10);
return; return;
} else if (status === 'error' || status === 'disconnected') { } else if (status === 'error' || status === 'disconnected') {
// 连接失败,刷新列表并显示错误 // 连接失败,刷新列表并显示错误
+130 -2
View File
@@ -29,12 +29,12 @@
<div class="container"> <div class="container">
<header> <header>
<div class="header-content"> <div class="header-content">
<div class="logo"> <div class="logo header-logo-link" onclick="switchPage('dashboard')" role="button" title="返回仪表盘">
<img src="/static/logo.png" alt="CyberStrikeAI Logo" style="width: 32px; height: 32px; margin-right: 8px;"> <img src="/static/logo.png" alt="CyberStrikeAI Logo" style="width: 32px; height: 32px; margin-right: 8px;">
<h1>CyberStrikeAI</h1> <h1>CyberStrikeAI</h1>
<span class="version-badge" title="当前版本">{{.Version}}</span>
</div> </div>
<div class="header-right"> <div class="header-right">
<p class="header-subtitle">AI 驱动的自动化安全测试平台</p>
<div class="header-actions"> <div class="header-actions">
<button class="openapi-doc-btn" onclick="window.open('/api-docs', '_blank')" title="OpenAPI 文档"> <button class="openapi-doc-btn" onclick="window.open('/api-docs', '_blank')" title="OpenAPI 文档">
<span>API 文档</span> <span>API 文档</span>
@@ -71,6 +71,17 @@
</svg> </svg>
</div> </div>
<nav class="main-sidebar-nav"> <nav class="main-sidebar-nav">
<div class="nav-item" data-page="dashboard">
<div class="nav-item-content" data-title="仪表盘" onclick="switchPage('dashboard')">
<svg width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
<rect x="3" y="3" width="7" height="9"></rect>
<rect x="14" y="3" width="7" height="5"></rect>
<rect x="14" y="12" width="7" height="9"></rect>
<rect x="3" y="16" width="7" height="5"></rect>
</svg>
<span>仪表盘</span>
</div>
</div>
<div class="nav-item" data-page="chat"> <div class="nav-item" data-page="chat">
<div class="nav-item-content" data-title="对话" onclick="switchPage('chat')"> <div class="nav-item-content" data-title="对话" onclick="switchPage('chat')">
<svg width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"> <svg width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
@@ -192,6 +203,122 @@
<!-- 内容区域 --> <!-- 内容区域 -->
<div class="content-area"> <div class="content-area">
<!-- 仪表盘页面 -->
<div id="page-dashboard" class="page">
<div class="dashboard-page">
<div class="page-header">
<h2>仪表盘</h2>
<div class="page-header-actions">
<button class="btn-secondary" onclick="refreshDashboard()" title="刷新数据">刷新</button>
</div>
</div>
<div class="dashboard-content">
<div class="dashboard-cards" id="dashboard-cards">
<div class="dashboard-card dashboard-card-tasks" onclick="switchPage('tasks')">
<div class="dashboard-card-glow"></div>
<div class="dashboard-card-icon">
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M9 11l3 3L22 4"></path><path d="M21 12v7a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2V5a2 2 0 0 1 2-2h11"></path></svg>
</div>
<div class="dashboard-card-body">
<div class="dashboard-card-value" id="dashboard-running-tasks">-</div>
<div class="dashboard-card-label">运行中任务</div>
<div class="dashboard-card-hint">点击查看任务管理</div>
</div>
</div>
<div class="dashboard-card dashboard-card-vulns" onclick="switchPage('vulnerabilities')">
<div class="dashboard-card-glow"></div>
<div class="dashboard-card-icon">
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M12 22s8-4 8-10V5l-8-3-8 3v7c0 6 8 10 8 10z"></path></svg>
</div>
<div class="dashboard-card-body">
<div class="dashboard-card-value" id="dashboard-vuln-total">-</div>
<div class="dashboard-card-label">漏洞总数</div>
<div class="dashboard-card-hint">点击查看漏洞管理</div>
</div>
</div>
<div class="dashboard-card dashboard-card-chat" onclick="switchPage('chat')">
<div class="dashboard-card-glow"></div>
<div class="dashboard-card-icon">
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 15a2 2 0 0 1-2 2H7l-4 4V5a2 2 0 0 1 2-2h14a2 2 0 0 1 2 2z"></path></svg>
</div>
<div class="dashboard-card-body">
<div class="dashboard-card-cta">开始对话</div>
<div class="dashboard-card-desc">与 AI 进行安全测试</div>
<div class="dashboard-card-hint">立即开始</div>
</div>
</div>
</div>
<div class="dashboard-section dashboard-section-chart">
<h3 class="dashboard-section-title"><span class="dashboard-section-dot"></span>漏洞严重程度分布</h3>
<div class="dashboard-chart-wrap">
<div class="dashboard-stacked-bar" id="dashboard-stacked-bar">
<span class="dashboard-bar-seg seg-critical" id="dashboard-bar-critical" style="width: 0%"></span>
<span class="dashboard-bar-seg seg-high" id="dashboard-bar-high" style="width: 0%"></span>
<span class="dashboard-bar-seg seg-medium" id="dashboard-bar-medium" style="width: 0%"></span>
<span class="dashboard-bar-seg seg-low" id="dashboard-bar-low" style="width: 0%"></span>
<span class="dashboard-bar-seg seg-info" id="dashboard-bar-info" style="width: 0%"></span>
</div>
<div class="dashboard-legend" id="dashboard-vuln-bars">
<div class="dashboard-legend-item"><span class="dashboard-legend-dot critical"></span><span class="dashboard-legend-label">严重</span><span class="dashboard-legend-value" id="dashboard-severity-critical">0</span></div>
<div class="dashboard-legend-item"><span class="dashboard-legend-dot high"></span><span class="dashboard-legend-label">高危</span><span class="dashboard-legend-value" id="dashboard-severity-high">0</span></div>
<div class="dashboard-legend-item"><span class="dashboard-legend-dot medium"></span><span class="dashboard-legend-label">中危</span><span class="dashboard-legend-value" id="dashboard-severity-medium">0</span></div>
<div class="dashboard-legend-item"><span class="dashboard-legend-dot low"></span><span class="dashboard-legend-label">低危</span><span class="dashboard-legend-value" id="dashboard-severity-low">0</span></div>
<div class="dashboard-legend-item"><span class="dashboard-legend-dot info"></span><span class="dashboard-legend-label">信息</span><span class="dashboard-legend-value" id="dashboard-severity-info">0</span></div>
</div>
</div>
</div>
<div class="dashboard-section dashboard-section-overview">
<h3 class="dashboard-section-title"><span class="dashboard-section-dot"></span>运行概览</h3>
<div class="dashboard-overview-grid">
<div class="dashboard-overview-item" onclick="switchPage('tasks')">
<span class="dashboard-overview-label">批量任务队列</span>
<span class="dashboard-overview-value"><span id="dashboard-batch-pending">-</span> 待执行 / <span id="dashboard-batch-running">-</span> 执行中 / <span id="dashboard-batch-done">-</span> 已完成</span>
</div>
<div class="dashboard-overview-item" onclick="switchPage('mcp-monitor')">
<span class="dashboard-overview-label">工具调用</span>
<span class="dashboard-overview-value"><span id="dashboard-tools-count">-</span> 个工具,共 <span id="dashboard-tools-calls">-</span> 次调用</span>
</div>
<div class="dashboard-overview-item" onclick="switchPage('skills-monitor')">
<span class="dashboard-overview-label">Skills</span>
<span class="dashboard-overview-value"><span id="dashboard-skills-count">-</span> 个 Skill,共 <span id="dashboard-skills-calls">-</span> 次调用</span>
</div>
</div>
</div>
<div class="dashboard-section">
<h3 class="dashboard-section-title"><span class="dashboard-section-dot"></span>快捷入口</h3>
<div class="dashboard-quick-links">
<a class="dashboard-quick-link" onclick="switchPage('tasks')">
<span class="dashboard-quick-icon"><svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M9 11l3 3L22 4"></path><path d="M21 12v7a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2V5a2 2 0 0 1 2-2h11"></path></svg></span>
<span>任务管理</span>
</a>
<a class="dashboard-quick-link" onclick="switchPage('vulnerabilities')">
<span class="dashboard-quick-icon"><svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M12 22s8-4 8-10V5l-8-3-8 3v7c0 6 8 10 8 10z"></path></svg></span>
<span>漏洞管理</span>
</a>
<a class="dashboard-quick-link" onclick="switchPage('chat')">
<span class="dashboard-quick-icon"><svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M21 15a2 2 0 0 1-2 2H7l-4 4V5a2 2 0 0 1 2-2h14a2 2 0 0 1 2 2z"></path></svg></span>
<span>对话</span>
</a>
<a class="dashboard-quick-link" onclick="switchPage('mcp-monitor')">
<span class="dashboard-quick-icon"><svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M13 2L3 14h9l-1 8 10-12h-9l1-8z"></path></svg></span>
<span>MCP 监控</span>
</a>
<a class="dashboard-quick-link" onclick="switchPage('skills-monitor')">
<span class="dashboard-quick-icon"><svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M14.5 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V7.5L14.5 2z"></path><polyline points="14 2 14 8 20 8"></polyline></svg></span>
<span>Skills 监控</span>
</a>
</div>
</div>
<div class="dashboard-cta-block">
<div class="dashboard-cta-inner">
<p class="dashboard-cta-text">准备好开始安全测试?</p>
<button class="dashboard-cta-btn" onclick="switchPage('chat')">前往对话 →</button>
</div>
</div>
</div>
</div>
</div>
<!-- 对话页面 --> <!-- 对话页面 -->
<div id="page-chat" class="page active"> <div id="page-chat" class="page active">
<div class="chat-page-layout"> <div class="chat-page-layout">
@@ -1701,6 +1828,7 @@ version: 1.0.0<br>
<script src="/static/js/builtin-tools.js"></script> <script src="/static/js/builtin-tools.js"></script>
<script src="/static/js/auth.js"></script> <script src="/static/js/auth.js"></script>
<script src="/static/js/router.js"></script> <script src="/static/js/router.js"></script>
<script src="/static/js/dashboard.js"></script>
<script src="/static/js/monitor.js"></script> <script src="/static/js/monitor.js"></script>
<script src="/static/js/chat.js"></script> <script src="/static/js/chat.js"></script>
<script src="/static/js/settings.js"></script> <script src="/static/js/settings.js"></script>