mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-01 00:27:35 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5d1f5d2886 | ||
|
|
0f92817261 | ||
|
|
2ec5953416 | ||
|
|
b4fe2e795e | ||
|
|
904d860797 | ||
|
|
dc08199af6 | ||
|
|
84e99220ff | ||
|
|
86f1d10a8b | ||
|
|
5c643a1606 | ||
|
|
b9854192c6 | ||
|
|
5a5762e1d1 | ||
|
|
8882d70393 | ||
|
|
5747ebc612 | ||
|
|
a6b3773f00 | ||
|
|
c2b950ad53 | ||
|
|
0283fff743 | ||
|
|
018835d6b8 | ||
|
|
4b7df4e0f3 | ||
|
|
0324b41a01 | ||
|
|
4a19620137 | ||
|
|
f8110413c0 | ||
|
|
59dc7cf858 | ||
|
|
52595e07e5 | ||
|
|
e0965594bb | ||
|
|
9ef8263eaf | ||
|
|
3c54a67416 | ||
|
|
98ca395edd | ||
|
|
c616822cd6 | ||
|
|
ba1796d7ce | ||
|
|
dad14c55c1 | ||
|
|
cc0233d7b4 | ||
|
|
7b6f56e476 | ||
|
|
2522fc6ae2 | ||
|
|
99d7380450 | ||
|
|
837e41459a | ||
|
|
5cbd828cad | ||
|
|
8cb317cbd6 | ||
|
|
94a2ba0406 | ||
|
|
4ee7204509 | ||
|
|
c326adbb66 | ||
|
|
7d1e9bdac4 | ||
|
|
151b445c74 | ||
|
|
af4b25b84e | ||
|
|
f0b1955059 | ||
|
|
f5d580bbf0 | ||
|
|
44d069da2b | ||
|
|
9297e6e6ee | ||
|
|
9bbc28c14a | ||
|
|
4943b9419e | ||
|
|
446ccd3edb | ||
|
|
a00643b9c0 | ||
|
|
5c13819f66 | ||
|
|
a7190b8399 | ||
|
|
6f0eaf7667 | ||
|
|
f285d47404 | ||
|
|
273d63413b | ||
|
|
5656315b1e | ||
|
|
cee46f40fa | ||
|
|
4c55784dbc | ||
|
|
d304ff5c99 | ||
|
|
0cd4df3518 | ||
|
|
49333b856c | ||
|
|
a83490f29c | ||
|
|
dbc4ead040 | ||
|
|
4333cf1dd3 | ||
|
|
577c97aab0 | ||
|
|
93ab362b6f | ||
|
|
d9cb0b11c5 | ||
|
|
2d7f5322b3 | ||
|
|
3763da8773 | ||
|
|
3f2984b7c6 | ||
|
|
fbef2afd68 | ||
|
|
bccb324d2b | ||
|
|
b28fba3d68 | ||
|
|
8059b6d5b2 | ||
|
|
0821bb2911 | ||
|
|
e407f8203c | ||
|
|
a59253e828 | ||
|
|
4c989fbbe0 | ||
|
|
91239c9869 | ||
|
|
6892156b7c | ||
|
|
75489123f9 | ||
|
|
631ee6c447 | ||
|
|
7caf77683f | ||
|
|
fce4ffe3c4 | ||
|
|
75163f9269 | ||
|
|
06a9cea97d | ||
|
|
b62ba3b214 |
@@ -218,16 +218,23 @@ The `run.sh` script will automatically:
|
|||||||
**Networking defaults:** `run.sh` starts the server with **`--https`** and the repo **`config.yaml`** (local self-signed TLS; better for many concurrent streams). Use **`./run.sh --http`** for plain HTTP. In production, set **`server.tls_cert_path`** / **`server.tls_key_path`** in **`config.yaml`** (see comments there). For manual runs, add **`--https`** or **`CYBERSTRIKE_HTTPS=1`**; if **`-config`** is wrong, the binary prints a short usage hint on stderr.
|
**Networking defaults:** `run.sh` starts the server with **`--https`** and the repo **`config.yaml`** (local self-signed TLS; better for many concurrent streams). Use **`./run.sh --http`** for plain HTTP. In production, set **`server.tls_cert_path`** / **`server.tls_key_path`** in **`config.yaml`** (see comments there). For manual runs, add **`--https`** or **`CYBERSTRIKE_HTTPS=1`**; if **`-config`** is wrong, the binary prints a short usage hint on stderr.
|
||||||
|
|
||||||
**First-Time Configuration:**
|
**First-Time Configuration:**
|
||||||
1. **Configure OpenAI-compatible API** (required before first use)
|
1. **Configure AI channels** (required before first use)
|
||||||
- After launch, open **`https://127.0.0.1:8080/`** (or **`https://localhost:8080/`**; replace **8080** with `server.port` in `config.yaml`) and accept the self-signed certificate warning once. If you used `./run.sh --http`, use **`http://`** instead.
|
- After launch, open **`https://127.0.0.1:8080/`** (or **`https://localhost:8080/`**; replace **8080** with `server.port` in `config.yaml`) and accept the self-signed certificate warning once. If you used `./run.sh --http`, use **`http://`** instead.
|
||||||
- Go to `Settings` → Fill in your API credentials:
|
- Go to `System Settings` → `Basic Settings` → `AI Channel Configuration`, add or edit a channel, then fill in provider, Base URL, API key, model, and token limits. Click **Save changes**. The left channel list supports setting a default, copy, delete, and bulk probe.
|
||||||
```yaml
|
```yaml
|
||||||
openai:
|
ai:
|
||||||
api_key: "${OPENAI_API_KEY}"
|
default_channel: openai-main
|
||||||
base_url: "https://api.openai.com/v1" # or https://api.deepseek.com/v1
|
channels:
|
||||||
model: "gpt-4o" # or deepseek-chat, claude-3-opus, etc.
|
openai-main:
|
||||||
|
name: OpenAI Main
|
||||||
|
provider: openai_compatible
|
||||||
|
api_key: "${OPENAI_API_KEY}"
|
||||||
|
base_url: "https://api.openai.com/v1" # or https://api.deepseek.com/v1
|
||||||
|
model: "gpt-4o" # or deepseek-chat, qwen3-max, etc.
|
||||||
|
max_total_tokens: 120000
|
||||||
|
max_completion_tokens: 16384
|
||||||
```
|
```
|
||||||
- Or edit `config.yaml` directly before launching
|
- Or edit `config.yaml` directly before launching. `ai.default_channel` is used for new conversations and tasks that do not explicitly select a channel; the chat page can also select any saved channel per session.
|
||||||
2. **Login** - On first startup the console prints an auto-generated initial `admin` password; create accounts from **Platform permissions → User management**
|
2. **Login** - On first startup the console prints an auto-generated initial `admin` password; create accounts from **Platform permissions → User management**
|
||||||
3. **Install security tools (optional)** - Install tools from `tools/` as needed; missing tools are skipped or substituted at runtime. Common examples:
|
3. **Install security tools (optional)** - Install tools from `tools/` as needed; missing tools are skipped or substituted at runtime. Common examples:
|
||||||
|
|
||||||
@@ -281,19 +288,23 @@ Requirements / tips:
|
|||||||
|
|
||||||
## Configuration
|
## Configuration
|
||||||
|
|
||||||
Use [`config.example.yaml`](config.example.yaml) as the authoritative configuration template and copy only the values required for your environment. At minimum, configure the server and an OpenAI-compatible model provider:
|
Use [`config.example.yaml`](config.example.yaml) as the authoritative configuration template and copy only the values required for your environment. At minimum, configure the server and one AI channel:
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
server:
|
server:
|
||||||
host: "127.0.0.1"
|
host: "127.0.0.1"
|
||||||
port: 8080
|
port: 8080
|
||||||
openai:
|
ai:
|
||||||
api_key: "${OPENAI_API_KEY}"
|
default_channel: openai-main
|
||||||
base_url: "https://api.openai.com/v1"
|
channels:
|
||||||
model: "your-model"
|
openai-main:
|
||||||
|
provider: openai_compatible
|
||||||
|
api_key: "${OPENAI_API_KEY}"
|
||||||
|
base_url: "https://api.openai.com/v1"
|
||||||
|
model: "your-model"
|
||||||
```
|
```
|
||||||
|
|
||||||
Do not commit real credentials. Review the [configuration reference](docs/en-US/configuration.md), [recommended profiles](docs/en-US/configuration-profiles.md), and [security hardening guide](docs/en-US/security-hardening.md) before exposing the service beyond localhost.
|
`openai` is a backward-compatible runtime field; maintain new model settings in `ai.channels`. Do not commit real credentials. Review the [configuration reference](docs/en-US/configuration.md), [recommended profiles](docs/en-US/configuration-profiles.md), and [security hardening guide](docs/en-US/security-hardening.md) before exposing the service beyond localhost.
|
||||||
|
|
||||||
## Related documentation
|
## Related documentation
|
||||||
|
|
||||||
|
|||||||
+24
-13
@@ -217,16 +217,23 @@ chmod +x run.sh && ./run.sh
|
|||||||
**网络默认:** `run.sh` 会以 **`--https`** 并传入项目根 **`config.yaml`** 启动(本机自签证书,多路流式场景更稳)。只要明文 HTTP 用 **`./run.sh --http`**。生产环境在 **`config.yaml`** 的 **`server.tls_cert_path` / `server.tls_key_path`** 配正式证书(见文件内注释)。手动启动可加 **`--https`** 或环境变量 **`CYBERSTRIKE_HTTPS=1`**;`-config` 写错时程序会在终端提示正确写法。
|
**网络默认:** `run.sh` 会以 **`--https`** 并传入项目根 **`config.yaml`** 启动(本机自签证书,多路流式场景更稳)。只要明文 HTTP 用 **`./run.sh --http`**。生产环境在 **`config.yaml`** 的 **`server.tls_cert_path` / `server.tls_key_path`** 配正式证书(见文件内注释)。手动启动可加 **`--https`** 或环境变量 **`CYBERSTRIKE_HTTPS=1`**;`-config` 写错时程序会在终端提示正确写法。
|
||||||
|
|
||||||
**首次配置:**
|
**首次配置:**
|
||||||
1. **配置 AI 模型 API**(首次使用前必填)
|
1. **配置 AI 通道**(首次使用前必填)
|
||||||
- 启动后在浏览器打开 **`https://127.0.0.1:8080/`**(或 **`https://localhost:8080/`**;端口以 `config.yaml` 中 **`server.port`** 为准,默认 8080),并按提示信任自签证书。若使用 **`./run.sh --http`**,则改用 **`http://`** 访问。
|
- 启动后在浏览器打开 **`https://127.0.0.1:8080/`**(或 **`https://localhost:8080/`**;端口以 `config.yaml` 中 **`server.port`** 为准,默认 8080),并按提示信任自签证书。若使用 **`./run.sh --http`**,则改用 **`http://`** 访问。
|
||||||
- 进入 `设置` → 填写 API 配置信息:
|
- 进入 `系统设置` → `基本设置` → `AI 通道配置`,新增或编辑通道,填写 API 提供商、Base URL、API Key、模型和 Token 上限,点击 **保存更改**。左侧通道列表支持设为默认、复制、删除和批量探活。
|
||||||
```yaml
|
```yaml
|
||||||
openai:
|
ai:
|
||||||
api_key: "${OPENAI_API_KEY}"
|
default_channel: openai-main
|
||||||
base_url: "https://api.openai.com/v1" # 或 https://api.deepseek.com/v1
|
channels:
|
||||||
model: "gpt-4o" # 或 deepseek-chat, claude-3-opus 等
|
openai-main:
|
||||||
|
name: OpenAI Main
|
||||||
|
provider: openai_compatible
|
||||||
|
api_key: "${OPENAI_API_KEY}"
|
||||||
|
base_url: "https://api.openai.com/v1" # 或 https://api.deepseek.com/v1
|
||||||
|
model: "gpt-4o" # 或 deepseek-chat, qwen3-max 等
|
||||||
|
max_total_tokens: 120000
|
||||||
|
max_completion_tokens: 16384
|
||||||
```
|
```
|
||||||
- 或启动前直接编辑 `config.yaml` 文件
|
- 或启动前直接编辑 `config.yaml` 文件。`ai.default_channel` 会作为新对话和未显式选择通道任务的默认模型;对话页也可以在会话设置里选择某个已保存通道。
|
||||||
2. **登录系统** - 首次启动时控制台会显示自动生成的 `admin` 初始密码;也可在「平台权限 → 用户管理」中创建账号
|
2. **登录系统** - 首次启动时控制台会显示自动生成的 `admin` 初始密码;也可在「平台权限 → 用户管理」中创建账号
|
||||||
3. **安装安全工具(可选)** - 按需安装 `tools/` 目录中的工具;未安装的工具在执行时会自动跳过或改用替代方案。常用示例:
|
3. **安装安全工具(可选)** - 按需安装 `tools/` 目录中的工具;未安装的工具在执行时会自动跳过或改用替代方案。常用示例:
|
||||||
|
|
||||||
@@ -279,19 +286,23 @@ go build -o cyberstrike-ai cmd/server/main.go
|
|||||||
|
|
||||||
## 配置
|
## 配置
|
||||||
|
|
||||||
请以 [`config.example.yaml`](config.example.yaml) 作为权威配置模板,只复制当前环境需要的配置。最少需要配置服务监听地址和一个 OpenAI 兼容模型:
|
请以 [`config.example.yaml`](config.example.yaml) 作为权威配置模板,只复制当前环境需要的配置。最少需要配置服务监听地址和一个 AI 通道:
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
server:
|
server:
|
||||||
host: "127.0.0.1"
|
host: "127.0.0.1"
|
||||||
port: 8080
|
port: 8080
|
||||||
openai:
|
ai:
|
||||||
api_key: "${OPENAI_API_KEY}"
|
default_channel: openai-main
|
||||||
base_url: "https://api.openai.com/v1"
|
channels:
|
||||||
model: "your-model"
|
openai-main:
|
||||||
|
provider: openai_compatible
|
||||||
|
api_key: "${OPENAI_API_KEY}"
|
||||||
|
base_url: "https://api.openai.com/v1"
|
||||||
|
model: "your-model"
|
||||||
```
|
```
|
||||||
|
|
||||||
不要提交真实凭证。将服务暴露到 localhost 之外前,请阅读[配置参考](docs/zh-CN/configuration.md)、[推荐配置画像](docs/zh-CN/configuration-profiles.md)和[安全加固指南](docs/zh-CN/security-hardening.md)。
|
`openai` 是兼容旧版本的运行时字段,新配置优先维护 `ai.channels`。不要提交真实凭证。将服务暴露到 localhost 之外前,请阅读[配置参考](docs/zh-CN/configuration.md)、[推荐配置画像](docs/zh-CN/configuration-profiles.md)和[安全加固指南](docs/zh-CN/security-hardening.md)。
|
||||||
|
|
||||||
## 相关文档
|
## 相关文档
|
||||||
|
|
||||||
|
|||||||
@@ -34,6 +34,7 @@ func main() {
|
|||||||
|
|
||||||
// 注册工具
|
// 注册工具
|
||||||
executor.RegisterTools(mcpServer)
|
executor.RegisterTools(mcpServer)
|
||||||
|
mcp.RegisterExecutionControlTools(mcpServer, nil)
|
||||||
|
|
||||||
log.Logger.Info("MCP服务器(stdio模式)已启动,等待消息...")
|
log.Logger.Info("MCP服务器(stdio模式)已启动,等待消息...")
|
||||||
|
|
||||||
|
|||||||
+91
-8
@@ -4,7 +4,9 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"cyberstrike-ai/internal/app"
|
"cyberstrike-ai/internal/app"
|
||||||
"cyberstrike-ai/internal/config"
|
"cyberstrike-ai/internal/config"
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
"cyberstrike-ai/internal/logger"
|
"cyberstrike-ai/internal/logger"
|
||||||
|
"cyberstrike-ai/internal/security"
|
||||||
"cyberstrike-ai/internal/termout"
|
"cyberstrike-ai/internal/termout"
|
||||||
"flag"
|
"flag"
|
||||||
"fmt"
|
"fmt"
|
||||||
@@ -12,17 +14,21 @@ import (
|
|||||||
"os/signal"
|
"os/signal"
|
||||||
"strings"
|
"strings"
|
||||||
"syscall"
|
"syscall"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
"golang.org/x/term"
|
||||||
)
|
)
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
var configPath = flag.String("config", "config.yaml", "配置文件路径")
|
var configPath = flag.String("config", "config.yaml", "Path to the configuration file")
|
||||||
var httpsBootstrap = flag.Bool("https", false, "启用主站 HTTPS:未配置 tls_cert_path/tls_key_path 时使用内存自签证书(本地测试);与 run.sh 默认行为一致")
|
var httpsBootstrap = flag.Bool("https", false, "Enable HTTPS for the main site; uses an in-memory self-signed certificate when no cert/key is configured")
|
||||||
var httpBootstrap = flag.Bool("http", false, "强制主站使用明文 HTTP:覆盖配置文件中的 tls_enabled/tls_auto_self_sign/tls_cert_path/tls_key_path")
|
var httpBootstrap = flag.Bool("http", false, "Force plain HTTP for the main site, overriding TLS settings in the configuration file")
|
||||||
|
var resetAdminPassword = flag.Bool("reset-admin-password", false, "Interactively reset the built-in admin password and exit")
|
||||||
flag.Parse()
|
flag.Parse()
|
||||||
|
|
||||||
// 环境变量兼容(便于 systemd/docker 等不传参场景)
|
// 环境变量兼容(便于 systemd/docker 等不传参场景)
|
||||||
if *httpsBootstrap && *httpBootstrap {
|
if *httpsBootstrap && *httpBootstrap {
|
||||||
fmt.Fprintln(os.Stderr, "--http 与 --https 不能同时使用")
|
fmt.Fprintln(os.Stderr, "--http and --https cannot be used together")
|
||||||
os.Exit(2)
|
os.Exit(2)
|
||||||
}
|
}
|
||||||
if !*httpsBootstrap && !*httpBootstrap {
|
if !*httpsBootstrap && !*httpBootstrap {
|
||||||
@@ -38,24 +44,32 @@ func main() {
|
|||||||
cp = "config.yaml"
|
cp = "config.yaml"
|
||||||
}
|
}
|
||||||
if strings.HasPrefix(cp, "-") {
|
if strings.HasPrefix(cp, "-") {
|
||||||
fmt.Fprintf(os.Stderr, "无效的 -config 路径 %q。\n若同时需要 HTTPS,请写成: ./cyberstrike-ai --https -config config.yaml(-config 后必须是 yaml 文件路径)。\n", cp)
|
fmt.Fprintf(os.Stderr, "Invalid -config path %q.\nIf HTTPS is also needed, use: ./cyberstrike-ai --https -config config.yaml (-config must be followed by a yaml file path).\n", cp)
|
||||||
os.Exit(2)
|
os.Exit(2)
|
||||||
}
|
}
|
||||||
localConfig, err := config.EnsureLocalConfig(cp)
|
localConfig, err := config.EnsureLocalConfig(cp)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
fmt.Printf("加载配置失败: %v\n", err)
|
fmt.Printf("Failed to load config: %v\n", err)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
cfg, err := config.Load(cp)
|
cfg, err := config.Load(cp)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
fmt.Printf("加载配置失败: %v\n", err)
|
fmt.Printf("Failed to load config: %v\n", err)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if localConfig.Created {
|
if localConfig.Created {
|
||||||
termout.PrintConfigCreated()
|
termout.PrintConfigCreated()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if *resetAdminPassword {
|
||||||
|
if err := runResetAdminPassword(cfg); err != nil {
|
||||||
|
fmt.Fprintf(os.Stderr, "Failed to reset admin password: %v\n", err)
|
||||||
|
os.Exit(1)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
if *httpBootstrap {
|
if *httpBootstrap {
|
||||||
config.ApplyPlainHTTPBootstrap(cfg)
|
config.ApplyPlainHTTPBootstrap(cfg)
|
||||||
} else if *httpsBootstrap {
|
} else if *httpsBootstrap {
|
||||||
@@ -79,7 +93,7 @@ func main() {
|
|||||||
|
|
||||||
// MCP 启用且 auth_header_value 为空时,自动生成随机密钥并写回配置
|
// MCP 启用且 auth_header_value 为空时,自动生成随机密钥并写回配置
|
||||||
if err := config.EnsureMCPAuth(cp, cfg); err != nil {
|
if err := config.EnsureMCPAuth(cp, cfg); err != nil {
|
||||||
fmt.Printf("MCP 鉴权配置失败: %v\n", err)
|
fmt.Printf("Failed to configure MCP authentication: %v\n", err)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if cfg.MCP.Enabled {
|
if cfg.MCP.Enabled {
|
||||||
@@ -121,3 +135,72 @@ func main() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func runResetAdminPassword(cfg *config.Config) error {
|
||||||
|
dbPath := strings.TrimSpace(cfg.Database.Path)
|
||||||
|
if dbPath == "" {
|
||||||
|
dbPath = "data/conversations.db"
|
||||||
|
}
|
||||||
|
if _, err := os.Stat(dbPath); err != nil {
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
return fmt.Errorf("database does not exist: %s; start the service once to initialize it first", dbPath)
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
fmt.Println("Reset built-in admin password")
|
||||||
|
fmt.Println()
|
||||||
|
|
||||||
|
password, err := readHiddenPassword("New admin password: ")
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
password = strings.TrimSpace(password)
|
||||||
|
if len(password) < 8 {
|
||||||
|
return fmt.Errorf("new password must be at least 8 characters")
|
||||||
|
}
|
||||||
|
confirm, err := readHiddenPassword("Confirm new password: ")
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if password != strings.TrimSpace(confirm) {
|
||||||
|
return fmt.Errorf("passwords do not match")
|
||||||
|
}
|
||||||
|
|
||||||
|
hash, err := security.HashPassword(password)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
db, err := database.NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer func() { _ = db.Close() }()
|
||||||
|
|
||||||
|
admin, err := db.GetRBACUserByUsername("admin")
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("built-in admin account was not found; start the service once to initialize it first: %w", err)
|
||||||
|
}
|
||||||
|
if !admin.IsBuiltin {
|
||||||
|
return fmt.Errorf("admin account is not built in; refusing to reset it")
|
||||||
|
}
|
||||||
|
if err := db.UpdateRBACAdminPassword(hash); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
fmt.Println()
|
||||||
|
fmt.Println("Admin password has been reset.")
|
||||||
|
fmt.Println("If the service is running, existing login sessions remain valid until the service restarts or the sessions expire.")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func readHiddenPassword(prompt string) (string, error) {
|
||||||
|
fmt.Fprint(os.Stderr, prompt)
|
||||||
|
password, err := term.ReadPassword(int(os.Stdin.Fd()))
|
||||||
|
fmt.Fprintln(os.Stderr)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
return string(password), nil
|
||||||
|
}
|
||||||
|
|||||||
+27
-22
@@ -10,7 +10,7 @@
|
|||||||
# ============================================
|
# ============================================
|
||||||
|
|
||||||
# 前端显示的版本号(可选,不填则显示默认版本)
|
# 前端显示的版本号(可选,不填则显示默认版本)
|
||||||
version: "v1.7.6"
|
version: "v1.7.9"
|
||||||
# 服务器配置
|
# 服务器配置
|
||||||
server:
|
server:
|
||||||
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
host: 0.0.0.0 # 监听地址,0.0.0.0 表示监听所有网络接口
|
||||||
@@ -49,35 +49,40 @@ monitor:
|
|||||||
# 对话相关配置
|
# 对话相关配置
|
||||||
# ============================================
|
# ============================================
|
||||||
|
|
||||||
# AI 模型配置(支持 OpenAI 兼容 API)
|
# AI 通道配置(支持保存多个 OpenAI 兼容 / Claude 通道)
|
||||||
# 必填项:api_key, base_url, model 必须填写才能正常运行
|
# default_channel 指定新对话与未显式选择通道的任务使用哪个通道。
|
||||||
|
# 每个 channels.<id> 必填:api_key, base_url, model。
|
||||||
# 支持的 API 服务商:
|
# 支持的 API 服务商:
|
||||||
# - OpenAI: https://api.openai.com/v1
|
# - OpenAI: https://api.openai.com/v1
|
||||||
# - DeepSeek: https://api.deepseek.com/v1
|
# - DeepSeek: https://api.deepseek.com/v1
|
||||||
# - 其他兼容 OpenAI 协议的 API
|
# - 其他兼容 OpenAI 协议的 API
|
||||||
# 常用模型: gpt-4, gpt-3.5-turbo, deepseek-chat, claude-3-opus 等
|
# 常用模型: gpt-4, gpt-3.5-turbo, deepseek-chat, claude-3-opus 等
|
||||||
# provider: 可选值 openai(默认) | claude(自动桥接到 Anthropic Claude Messages API)
|
# provider: 可选值 openai_compatible(默认) | claude(自动桥接到 Anthropic Claude Messages API)
|
||||||
openai:
|
ai:
|
||||||
provider: openai # API 提供商: openai(默认,兼容OpenAI协议) | claude(自动桥接到Anthropic Claude Messages API)
|
default_channel: qwen-max
|
||||||
base_url: https://dashscope.aliyuncs.com/compatible-mode/v1 # API 基础 URL(必填)
|
channels:
|
||||||
api_key: sk-xxxxxxx # API 密钥(必填)
|
qwen-max:
|
||||||
model: qwen3-max # 模型名称(必填)
|
name: Qwen Max
|
||||||
max_total_tokens: 120000 # LLM 相关上下文的最大 Token 数限制(内存压缩和攻击链构建会共用此配置)
|
provider: openai_compatible
|
||||||
max_completion_tokens: 16384 # 单次生成上限(含 reasoning 与可见输出),防止依赖网关隐式默认值
|
base_url: https://dashscope.aliyuncs.com/compatible-mode/v1
|
||||||
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort;Claude 4.6+ 为 adaptive + output_config.effort(仅显式配置 effort 时下发);3.7 为 enabled+budget_tokens:10000(文档示例),effort 不映射,自定义预算用 extra_request_fields
|
api_key: sk-xxxxxxx
|
||||||
reasoning:
|
model: qwen3-max
|
||||||
mode: on # auto | on | off;off:OpenAI/Claude 不附加推理字段,DeepSeek 发送 thinking.type=disabled(其默认开启思考)
|
max_total_tokens: 120000
|
||||||
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
|
max_completion_tokens: 32768
|
||||||
allow_client_reasoning: true # false 时忽略对话请求体 reasoning,仅以下方为准
|
# Eino 路径模型推理:DeepSeek/OpenAI 为 thinking / reasoning_effort;Claude 4.6+ 为 adaptive + output_config.effort(仅显式配置 effort 时下发);3.7 为 enabled+budget_tokens:10000(文档示例),effort 不映射,自定义预算用 extra_request_fields
|
||||||
profile: openai_compat # auto | deepseek_compat | openai_compat | output_config_effort
|
reasoning:
|
||||||
# extra_request_fields: {} # 可选:管理员自定义根级 JSON 片段(高级)
|
mode: on # auto | on | off;off:OpenAI/Claude 不附加推理字段,DeepSeek 发送 thinking.type=disabled(其默认开启思考)
|
||||||
|
effort: high # low | medium | high | max | xhigh(最高档:OpenAI 常用 xhigh,部分网关用 max,原样下发);空表示不指定
|
||||||
|
allow_client_reasoning: true # false 时忽略对话请求体 reasoning,仅以下方为准
|
||||||
|
profile: openai_compat # auto | deepseek_compat | openai_compat | output_config_effort
|
||||||
|
# extra_request_fields: {} # 可选:管理员自定义根级 JSON 片段(高级)
|
||||||
# 视觉分析(analyze_image MCP 工具;图片仅在单次 VL 调用中出现,Agent 上下文只保留文字摘要)
|
# 视觉分析(analyze_image MCP 工具;图片仅在单次 VL 调用中出现,Agent 上下文只保留文字摘要)
|
||||||
vision:
|
vision:
|
||||||
enabled: false # true 且 model 非空时注册 analyze_image
|
enabled: false # true 且 model 非空时注册 analyze_image
|
||||||
model: qwen-vl # VL 模型名(enabled 时必填)
|
model: qwen-vl # VL 模型名(enabled 时必填)
|
||||||
api_key: "" # 留空则复用 openai.api_key
|
api_key: "" # 留空则复用默认 AI 通道 api_key
|
||||||
base_url: "" # 留空则复用 openai.base_url
|
base_url: "" # 留空则复用默认 AI 通道 base_url
|
||||||
provider: # 留空则复用 openai.provider(openai | claude)
|
provider: # 留空则复用默认 AI 通道 provider(openai_compatible | claude)
|
||||||
max_image_bytes: 5242880 # 原始文件上限(字节),默认 5MB
|
max_image_bytes: 5242880 # 原始文件上限(字节),默认 5MB
|
||||||
max_dimension: 2048 # 长边缩放像素
|
max_dimension: 2048 # 长边缩放像素
|
||||||
jpeg_quality: 82
|
jpeg_quality: 82
|
||||||
@@ -110,7 +115,7 @@ agent:
|
|||||||
tool_wait_timeout_seconds: 300 # 工具本轮最多等待(秒);到时返回 execution_id,worker 继续后台执行,可用 wait_tool_execution 继续等待;0=等到完成
|
tool_wait_timeout_seconds: 300 # 工具本轮最多等待(秒);到时返回 execution_id,worker 继续后台执行,可用 wait_tool_execution 继续等待;0=等到完成
|
||||||
external_mcp_max_concurrent_per_server: 5 # 单个外部 MCP server 同时运行的工具数;0=默认2;负数=不限制
|
external_mcp_max_concurrent_per_server: 5 # 单个外部 MCP server 同时运行的工具数;0=默认2;负数=不限制
|
||||||
external_mcp_max_concurrent_total: 16 # 所有外部 MCP 工具全局并发上限;0=默认16;负数=不限制
|
external_mcp_max_concurrent_total: 16 # 所有外部 MCP 工具全局并发上限;0=默认16;负数=不限制
|
||||||
external_mcp_circuit_failure_threshold: 5 # 单个外部 MCP server 连续失败多少次后熔断;0=默认3;负数=关闭熔断
|
external_mcp_circuit_failure_threshold: 15 # 单个外部 MCP server 连续失败多少次后熔断;0=默认3;负数=关闭熔断
|
||||||
external_mcp_circuit_cooldown_seconds: 60 # 熔断冷却秒数;0=默认60
|
external_mcp_circuit_cooldown_seconds: 60 # 熔断冷却秒数;0=默认60
|
||||||
shell_no_output_timeout_seconds: 1200 # execute/exec 连续无新输出则终止(秒);通用防挂死;0=默认300;-1=关闭
|
shell_no_output_timeout_seconds: 1200 # execute/exec 连续无新输出则终止(秒);通用防挂死;0=默认300;-1=关闭
|
||||||
workspace_root_dir: "" # 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离;勿用系统 /tmp
|
workspace_root_dir: "" # 会话工作目录根路径(curl/wget 下载、read_file/glob/grep 本地分析);空=tmp/workspace,其下按 projects/{id} 或 conversations/{id} 隔离;勿用系统 /tmp
|
||||||
|
|||||||
@@ -21,7 +21,7 @@ vision:
|
|||||||
timeout_seconds: 60
|
timeout_seconds: 60
|
||||||
```
|
```
|
||||||
|
|
||||||
Empty `api_key`, `base_url`, or `provider` inherits from `openai`.
|
Empty `api_key`, `base_url`, or `provider` inherits from the resolved default AI channel.
|
||||||
|
|
||||||
## Data Handling
|
## Data Handling
|
||||||
|
|
||||||
|
|||||||
@@ -45,6 +45,18 @@ Multi-agent:
|
|||||||
|
|
||||||
`orchestration` may be `deep`, `plan_execute`, or `supervisor`.
|
`orchestration` may be `deep`, `plan_execute`, or `supervisor`.
|
||||||
|
|
||||||
|
Common request body fields:
|
||||||
|
|
||||||
|
| Field | Meaning |
|
||||||
|
| --- | --- |
|
||||||
|
| `message` | User message, required. |
|
||||||
|
| `conversationId` | Continue an existing conversation; empty creates a new one. |
|
||||||
|
| `projectId` | Project for a new conversation; empty may follow `config.project.default_project_id`. |
|
||||||
|
| `role` | Use a named role. |
|
||||||
|
| `aiChannelId` | Select a channel from `ai.channels`; empty follows `ai.default_channel`. |
|
||||||
|
| `reasoning` | Per-session reasoning override, controlled by the channel's `reasoning.allow_client_reasoning`. |
|
||||||
|
| `hitl` | Per-session human-in-the-loop settings. |
|
||||||
|
|
||||||
## SSE Notes
|
## SSE Notes
|
||||||
|
|
||||||
Streaming endpoints are long-lived. Clients should:
|
Streaming endpoints are long-lived. Clients should:
|
||||||
@@ -55,6 +67,25 @@ Streaming endpoints are long-lived. Clients should:
|
|||||||
- disable proxy buffering;
|
- disable proxy buffering;
|
||||||
- pass `conversationId` when continuing a conversation.
|
- pass `conversationId` when continuing a conversation.
|
||||||
|
|
||||||
|
## File Management Sources
|
||||||
|
|
||||||
|
The file management page and `GET /api/chat-uploads` group conversation-related files by source. Directory names still use project IDs or conversation IDs for stability, while the UI prefers project names or conversation titles and keeps the full ID available in tooltips or copied paths.
|
||||||
|
|
||||||
|
| Source | `source` | Typical directory | Meaning | Mutability |
|
||||||
|
| --- | --- | --- | --- | --- |
|
||||||
|
| Workspace files | `workspace` | `tmp/workspace/projects/<projectId>/...`, `tmp/workspace/conversations/<conversationId>/...` | The Agent workspace for downloaded files, analysis scripts, intermediate results, and generated CSV/XLSX/Markdown files. If an AI-generated file is missing from the UI, check this source first. | Read-only listing; supports copy path, download, and export. |
|
||||||
|
| Conversation artifacts | `conversation_artifact` | `data/conversation_artifacts/<conversationId>/...` | Conversation-scoped deliverables or archived artifacts such as summaries, reports, or middleware-generated artifacts. | Read-only listing; supports copy path, download, and export. |
|
||||||
|
| Tool outputs | `reduction` | `tmp/reduction/projects/<projectId>/...`, `tmp/reduction/conversations/<conversationId>/...` | Persisted full tool outputs, scan raw data, or outputs saved before truncation. Useful for reviewing long command or scan results. | Read-only listing; supports copy path, download, and export. |
|
||||||
|
| Chat uploads | `upload` | `chat_uploads/<date>/<conversationId>/...` | Files manually uploaded in chat or from the file management page. Copy the server absolute path into chat when the AI should reference a file. | Supports upload, mkdir, text edit, rename, delete, copy path, download, and export. |
|
||||||
|
|
||||||
|
Related endpoints:
|
||||||
|
|
||||||
|
- `GET /api/chat-uploads`: list files filtered by source, project, conversation, or filename.
|
||||||
|
- `GET /api/chat-uploads/path`: resolve a file-management relative path or internal virtual path to a server absolute path for copy actions.
|
||||||
|
- `GET /api/chat-uploads/download`: download a file.
|
||||||
|
- `GET /api/chat-uploads/export`: export the current filtered result as a ZIP.
|
||||||
|
- `POST /api/chat-uploads`: upload into the chat uploads directory.
|
||||||
|
|
||||||
## Asset Management and Bulk Import
|
## Asset Management and Bulk Import
|
||||||
|
|
||||||
Asset endpoints:
|
Asset endpoints:
|
||||||
|
|||||||
@@ -16,11 +16,15 @@ server:
|
|||||||
# - https://trusted-integration.example
|
# - https://trusted-integration.example
|
||||||
auth:
|
auth:
|
||||||
session_duration_hours: 12
|
session_duration_hours: 12
|
||||||
openai:
|
ai:
|
||||||
provider: openai
|
default_channel: openai-main
|
||||||
base_url: https://api.openai.com/v1
|
channels:
|
||||||
api_key: sk-...
|
openai-main:
|
||||||
model: gpt-4.1
|
name: OpenAI Main
|
||||||
|
provider: openai_compatible
|
||||||
|
base_url: https://api.openai.com/v1
|
||||||
|
api_key: sk-...
|
||||||
|
model: gpt-4.1
|
||||||
agent:
|
agent:
|
||||||
max_iterations: 12000
|
max_iterations: 12000
|
||||||
tool_timeout_minutes: 60
|
tool_timeout_minutes: 60
|
||||||
@@ -30,13 +34,64 @@ Change the initial `admin` password from the Web UI after first login. Use HTTPS
|
|||||||
|
|
||||||
Valid Chromium `chrome-extension://<32-character-extension-id>` origins are recognized automatically. The extension must still obtain host permission and authenticate with a password and Bearer token. `server.cors_allowed_origins` remains available as an exact allowlist for other trusted Web integrations; wildcards are not accepted, and changing it requires a restart.
|
Valid Chromium `chrome-extension://<32-character-extension-id>` origins are recognized automatically. The extension must still obtain host permission and authenticate with a password and Bearer token. `server.cors_allowed_origins` remains available as an exact allowlist for other trusted Web integrations; wildcards are not accepted, and changing it requires a restart.
|
||||||
|
|
||||||
|
## AI Channels
|
||||||
|
|
||||||
|
`ai` is the recommended model configuration entry. In the Web UI, use **System Settings → Basic Settings → AI Channel Configuration**. Saving that form writes `ai.default_channel` and `ai.channels`. The legacy `openai` field remains as a backward-compatible runtime field; on load, CyberStrikeAI ensures a default channel exists and synchronizes the resolved `ai.default_channel` into runtime `openai`.
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
ai:
|
||||||
|
default_channel: openai-main
|
||||||
|
channels:
|
||||||
|
openai-main:
|
||||||
|
name: OpenAI Main
|
||||||
|
provider: openai_compatible
|
||||||
|
base_url: https://api.openai.com/v1
|
||||||
|
api_key: sk-...
|
||||||
|
model: gpt-4.1
|
||||||
|
max_total_tokens: 120000
|
||||||
|
max_completion_tokens: 16384
|
||||||
|
reasoning:
|
||||||
|
mode: on
|
||||||
|
effort: high
|
||||||
|
allow_client_reasoning: true
|
||||||
|
profile: openai_compat
|
||||||
|
claude-main:
|
||||||
|
name: Claude Main
|
||||||
|
provider: claude
|
||||||
|
base_url: https://api.anthropic.com/v1
|
||||||
|
api_key: sk-ant-...
|
||||||
|
model: claude-sonnet-4-5
|
||||||
|
```
|
||||||
|
|
||||||
|
| Field | Meaning |
|
||||||
|
| --- | --- |
|
||||||
|
| `ai.default_channel` | Default channel ID for new conversations and requests without an explicit channel. |
|
||||||
|
| `ai.channels.<id>` | Channel config. IDs are normalized to lowercase letters, digits, and hyphens. |
|
||||||
|
| `name` | Display name in the Web UI; falls back to the ID. |
|
||||||
|
| `provider` | `openai_compatible` or `claude`. OpenAI-compatible channels map to runtime `openai`; Claude channels bridge to Anthropic Messages API. |
|
||||||
|
| `base_url/api_key/model` | Required. Base URL usually includes a version path such as `/v1`. |
|
||||||
|
| `max_total_tokens` | Shared context budget for compression, attack-chain generation, multi-agent summaries, and similar paths. |
|
||||||
|
| `max_completion_tokens` | Per-response output cap; default is used when empty. |
|
||||||
|
| `reasoning` | Default reasoning fields for the channel. Gateway support varies; try `mode: off` first when a provider rejects requests. |
|
||||||
|
|
||||||
|
The chat page reads saved channels into the “AI Channel” selector. A non-empty request `aiChannelId` selects a channel for that run/session without sending API credentials through the prompt path. Empty `aiChannelId` follows `ai.default_channel`.
|
||||||
|
|
||||||
|
Common Web UI operations:
|
||||||
|
|
||||||
|
- Add: click `+`, fill required fields, then save.
|
||||||
|
- Set default: select a channel, click **Set as default**, then save/apply.
|
||||||
|
- Copy: duplicate the current form, useful for the same provider with a different model.
|
||||||
|
- Delete: keep at least one channel; the default channel is protected from bulk delete.
|
||||||
|
- Probe: use **Test connection** or **Bulk probe** to validate API key, Base URL, and model.
|
||||||
|
|
||||||
## Hot-Apply Boundaries
|
## Hot-Apply Boundaries
|
||||||
|
|
||||||
`POST /api/config/apply` coordinates model config, tool description mode, MCP tool registration, knowledge components, robot restarts, and C2 runtime reconciliation. It does not make every field instantly effective.
|
`POST /api/config/apply` coordinates model config, tool description mode, MCP tool registration, knowledge components, robot restarts, and C2 runtime reconciliation. It does not make every field instantly effective.
|
||||||
|
|
||||||
| Section | Usually hot-applies | Extra action |
|
| Section | Usually hot-applies | Extra action |
|
||||||
| --- | --- | --- |
|
| --- | --- | --- |
|
||||||
| `openai` | new requests use new model settings | running streams keep their current state |
|
| `ai.default_channel` / `ai.channels` | new requests use the resolved default or selected channel | running streams keep their current state; reload config for the frontend channel list |
|
||||||
|
| `openai` | compatibility field, usually synchronized from the default AI channel | prefer maintaining new config in `ai.channels` |
|
||||||
| `agent.max_iterations` | new tasks | existing tasks continue |
|
| `agent.max_iterations` | new tasks | existing tasks continue |
|
||||||
| `hitl.tool_whitelist` | new approval checks | pending approvals are not re-decided |
|
| `hitl.tool_whitelist` | new approval checks | pending approvals are not re-decided |
|
||||||
| `knowledge.enabled` | initializes/updates components | scan and index are still required |
|
| `knowledge.enabled` | initializes/updates components | scan and index are still required |
|
||||||
@@ -47,8 +102,8 @@ Valid Chromium `chrome-extension://<32-character-extension-id>` origins are reco
|
|||||||
|
|
||||||
## Fallback Relationships
|
## Fallback Relationships
|
||||||
|
|
||||||
- `vision.api_key/base_url/provider` can inherit from `openai`.
|
- `vision.api_key/base_url/provider` can inherit from the resolved default AI channel.
|
||||||
- `hitl.audit_model` can inherit from `openai`.
|
- `hitl.audit_model` can inherit from the resolved default AI channel.
|
||||||
- `knowledge.embedding.base_url/api_key` can inherit from model settings.
|
- `knowledge.embedding.base_url/api_key` can inherit from model settings.
|
||||||
- rerank config can inherit from embedding/openai.
|
- rerank config can inherit from embedding/openai.
|
||||||
- `database.knowledge_db_path` can be separate or reuse the main DB.
|
- `database.knowledge_db_path` can be separate or reuse the main DB.
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ CyberStrikeAI can run as a local testing tool, an internal team service, or a pr
|
|||||||
- Python for some MCP servers and tool scripts.
|
- Python for some MCP servers and tool scripts.
|
||||||
- SQLite files under `data/`; no external DB is required by default.
|
- SQLite files under `data/`; no external DB is required by default.
|
||||||
- Actual security tools installed in PATH. YAML files under `tools/` only describe commands.
|
- Actual security tools installed in PATH. YAML files under `tools/` only describe commands.
|
||||||
- An OpenAI-compatible model endpoint, or `openai.provider: claude` for the Claude bridge.
|
- At least one `ai.channels` entry. Use `provider: openai_compatible` for OpenAI-compatible endpoints, or `provider: claude` for the Claude bridge.
|
||||||
|
|
||||||
Important persistent paths:
|
Important persistent paths:
|
||||||
|
|
||||||
|
|||||||
@@ -23,12 +23,12 @@ hitl:
|
|||||||
provider: ""
|
provider: ""
|
||||||
base_url: ""
|
base_url: ""
|
||||||
api_key: ""
|
api_key: ""
|
||||||
model: "" # set a small model here; blank reuses openai.model
|
model: "" # set a small model here; blank reuses the default AI channel model
|
||||||
retention_days: 90
|
retention_days: 90
|
||||||
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
|
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
|
||||||
```
|
```
|
||||||
|
|
||||||
`audit_model` supports partial configuration. Empty fields inherit from the main `openai` config, so the common setup is to fill only `model` and run approvals on a cheaper small model.
|
`audit_model` supports partial configuration. Empty fields inherit from the resolved default AI channel, so the common setup is to fill only `model` and run approvals on a cheaper small model.
|
||||||
|
|
||||||
## Recommended Approval Strategy
|
## Recommended Approval Strategy
|
||||||
|
|
||||||
|
|||||||
@@ -51,6 +51,13 @@ Important: when `wait_tool_execution` reaches `timeout_seconds` and the target e
|
|||||||
| `wait_tool_execution` | Wait for a selected execution for a bounded interval |
|
| `wait_tool_execution` | Wait for a selected execution for a bounded interval |
|
||||||
| `cancel_tool_execution` | Cancel a selected execution |
|
| `cancel_tool_execution` | Cancel a selected execution |
|
||||||
|
|
||||||
|
`get_tool_execution` and `wait_tool_execution` can include a live output preview:
|
||||||
|
|
||||||
|
- `include_partial_output`: whether to return partial output, default `true`.
|
||||||
|
- `partial_output_max_bytes`: tail preview limit for this call, default `4096`, maximum `65536`.
|
||||||
|
|
||||||
|
Partial output is a bounded preview of output produced so far, not the final `result`. The canonical `result` is still written only when the tool finishes. Tools that do not support streaming output simply omit partial fields.
|
||||||
|
|
||||||
Typical flow:
|
Typical flow:
|
||||||
|
|
||||||
```text
|
```text
|
||||||
@@ -60,6 +67,8 @@ Typical flow:
|
|||||||
4. If still incomplete, continue waiting or call cancel_tool_execution
|
4. If still incomplete, continue waiting or call cancel_tool_execution
|
||||||
```
|
```
|
||||||
|
|
||||||
|
`tool_wait_timeout_seconds` applies to internal MCP tools, external MCP tools, and Eino filesystem's streaming `execute`. Eino's non-streaming filesystem tools such as `ls/read_file/write_file/edit_file/glob/grep` are recorded in execution monitoring, but they are not converted into resumable background workers.
|
||||||
|
|
||||||
## Cancellation and Session Cleanup
|
## Cancellation and Session Cleanup
|
||||||
|
|
||||||
- User stop cancels running tools for the current conversation.
|
- User stop cancels running tools for the current conversation.
|
||||||
|
|||||||
@@ -46,20 +46,29 @@ Login fails:
|
|||||||
|
|
||||||
If another administrator with `rbac:write` is available, reset the password under **Platform permissions → User management**.
|
If another administrator with `rbac:write` is available, reset the password under **Platform permissions → User management**.
|
||||||
|
|
||||||
If no administrator session is available, the built-in `admin` account can be recovered on the server. Stop CyberStrikeAI, back up the database, change to the project root, and run the command below. Enter and confirm the new password when prompted:
|
If no administrator session is available, the built-in `admin` account can be recovered on the server. Change to the project root and run:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
./run.sh --reset-admin-password
|
||||||
|
```
|
||||||
|
|
||||||
|
Enter and confirm the new password when prompted. The script hides input and stores a bcrypt hash. If the service is running, restart it afterward to invalidate existing login sessions.
|
||||||
|
|
||||||
|
If `run.sh` is not available, run the command below manually. Enter and confirm the new password when prompted:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
HASH=$(htpasswd -nBC 10 '' | cut -d: -f2 | tr -d '\n') && sqlite3 data/conversations.db "UPDATE rbac_users SET password_hash='$HASH', updated_at=CURRENT_TIMESTAMP WHERE id='admin' AND username='admin' AND is_builtin=1; SELECT changes();"
|
HASH=$(htpasswd -nBC 10 '' | cut -d: -f2 | tr -d '\n') && sqlite3 data/conversations.db "UPDATE rbac_users SET password_hash='$HASH', updated_at=CURRENT_TIMESTAMP WHERE id='admin' AND username='admin' AND is_builtin=1; SELECT changes();"
|
||||||
```
|
```
|
||||||
|
|
||||||
Output `1` means that the row was updated. The command requires `sqlite3` and `htpasswd`. If `database.path` in `config.yaml` is not the default, replace `data/conversations.db`. Password input is hidden, is not written to shell history, and is stored as a bcrypt hash. Restart the service afterward to invalidate existing login sessions.
|
Output `1` means that the row was updated. The command requires `sqlite3` and `htpasswd`. If `database.path` in `config.yaml` is not the default, replace `data/conversations.db`. Password input is hidden and is not written to shell history.
|
||||||
|
|
||||||
Model fails:
|
Model fails:
|
||||||
|
|
||||||
- wrong `base_url` path;
|
- selected AI channel does not exist; empty selection follows `ai.default_channel`;
|
||||||
|
- wrong `ai.channels.<id>.base_url` path;
|
||||||
- invalid API key;
|
- invalid API key;
|
||||||
- model unavailable;
|
- model unavailable;
|
||||||
- reasoning fields unsupported by gateway. Try `openai.reasoning.mode: off`.
|
- reasoning fields unsupported by gateway. Try `ai.channels.<id>.reasoning.mode: off`.
|
||||||
|
|
||||||
Streaming stalls:
|
Streaming stalls:
|
||||||
|
|
||||||
|
|||||||
@@ -34,7 +34,39 @@ Saved workflows can be bound to a role under **Role Management**. When `workflow
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 3. Execution model (read this before configuring)
|
## 3. Natural-language Draft Generation
|
||||||
|
|
||||||
|
The Workflows page provides a **Create from natural language** entry point. After a user describes a security operations goal, CyberStrikeAI generates an editable draft and returns a structured audit result:
|
||||||
|
|
||||||
|
- Draft generation does not save, dry-run, or execute tools automatically.
|
||||||
|
- The server endpoint is `POST /api/workflows/generate-draft`, protected by `workflow:write`.
|
||||||
|
- The response includes `graph`, `meta`, `capabilities`, `audit`, and `stats`; after applying the draft to the canvas, normal save validation still runs.
|
||||||
|
- High-risk language such as executing scripts, isolating hosts, blocking, deleting, or exploitation is marked as `high_risk` and defaults to HITL approval or `requires_human_confirmation`.
|
||||||
|
- Tool capabilities are matched against the available tool list; unmatched capabilities fall back to Agent draft nodes and are surfaced in `audit.missing_fields` / `audit.assumptions`.
|
||||||
|
- If server generation is unavailable, the frontend uses a local deterministic fallback and still returns an editable draft with risk notes.
|
||||||
|
|
||||||
|
Example request:
|
||||||
|
|
||||||
|
```http
|
||||||
|
POST /api/workflows/generate-draft
|
||||||
|
Content-Type: application/json
|
||||||
|
|
||||||
|
{
|
||||||
|
"prompt": "Scan target assets for open ports, create a vulnerability task when critical ports are found, ask the owner for approval, and output a report",
|
||||||
|
"options": {
|
||||||
|
"include_objective": true,
|
||||||
|
"allow_schedule": false,
|
||||||
|
"allow_high_risk": false
|
||||||
|
},
|
||||||
|
"available_tools": [
|
||||||
|
{ "key": "nmap", "name": "nmap", "enabled": true }
|
||||||
|
]
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 4. Execution model (read this before configuring)
|
||||||
|
|
||||||
The engine executes the workflow as a **directed graph**, starting from the **Start** node and following edges to downstream nodes.
|
The engine executes the workflow as a **directed graph**, starting from the **Start** node and following edges to downstream nodes.
|
||||||
|
|
||||||
@@ -116,7 +148,7 @@ Agent B still receives Agent A’s output even when a condition node lies betwee
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 4. Template syntax
|
## 5. Template syntax
|
||||||
|
|
||||||
### 4.1 Basic format
|
### 4.1 Basic format
|
||||||
|
|
||||||
@@ -198,7 +230,7 @@ Field bindings can read ordinary fields such as `output` or `message`, and also
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 5. Node types and configuration
|
## 6. Node types and configuration
|
||||||
|
|
||||||
### 5.1 Start
|
### 5.1 Start
|
||||||
|
|
||||||
@@ -318,7 +350,7 @@ Optional node for an end summary template (less common in role-bound flows).
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 6. Edge configuration
|
## 7. Edge configuration
|
||||||
|
|
||||||
Select an **edge** to configure its **condition** in the right panel.
|
Select an **edge** to configure its **condition** in the right panel.
|
||||||
|
|
||||||
@@ -335,7 +367,7 @@ If no edge condition is set:
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 7. Full example: passing Agent output across a condition
|
## 8. Full example: passing Agent output across a condition
|
||||||
|
|
||||||
### 7.1 Graph structure
|
### 7.1 Graph structure
|
||||||
|
|
||||||
@@ -384,7 +416,7 @@ Start → Agent (initial value) → Condition → Agent (transform) → Output
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 8. Bind to a role and run
|
## 9. Bind to a role and run
|
||||||
|
|
||||||
### 8.1 Bind in Role Management
|
### 8.1 Bind in Role Management
|
||||||
|
|
||||||
@@ -414,7 +446,7 @@ If no Output node is reached or no branch matches, `outputs` may be empty and th
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 9. Debugging, dry-run, and replay
|
## 10. Debugging, dry-run, and replay
|
||||||
|
|
||||||
### 9.1 Safe dry-run
|
### 9.1 Safe dry-run
|
||||||
|
|
||||||
@@ -487,7 +519,7 @@ Token and cost metrics depend on whether the underlying model/Agent events repor
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 10. Validation before save
|
## 11. Validation before save
|
||||||
|
|
||||||
On save, the system checks:
|
On save, the system checks:
|
||||||
|
|
||||||
@@ -510,7 +542,7 @@ On save, the system checks:
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 11. Troubleshooting
|
## 12. Troubleshooting
|
||||||
|
|
||||||
| Symptom | Likely cause | Fix |
|
| Symptom | Likely cause | Fix |
|
||||||
|---------|--------------|-----|
|
|---------|--------------|-----|
|
||||||
@@ -526,7 +558,7 @@ On save, the system checks:
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 12. Best practices
|
## 13. Best practices
|
||||||
|
|
||||||
1. **Meaningful names**: Use descriptive output variable names (`scan_result`, `parsed_targets`) instead of reusing `agent_result` everywhere.
|
1. **Meaningful names**: Use descriptive output variable names (`scan_result`, `parsed_targets`) instead of reusing `agent_result` everywhere.
|
||||||
2. **Prefer `outputs` for cross-node data**: If a condition, tool, or HITL node might sit in between, use named variables.
|
2. **Prefer `outputs` for cross-node data**: If a condition, tool, or HITL node might sit in between, use named variables.
|
||||||
@@ -539,7 +571,7 @@ On save, the system checks:
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 13. Code references (for developers)
|
## 14. Code references (for developers)
|
||||||
|
|
||||||
| Module | Path |
|
| Module | Path |
|
||||||
|--------|------|
|
|--------|------|
|
||||||
|
|||||||
@@ -12,7 +12,7 @@
|
|||||||
## 核心概念与编排
|
## 核心概念与编排
|
||||||
|
|
||||||
- [架构说明](architecture.md) · [安全模型](security-model.md) · [RBAC](rbac.md)
|
- [架构说明](architecture.md) · [安全模型](security-model.md) · [RBAC](rbac.md)
|
||||||
- [Agent 与角色](agent-and-role-guide.md) · [Skills](skills-guide.md) · [Eino 多代理](MULTI_AGENT_EINO.md)
|
- [Agent 与角色](agent-and-role-guide.md) · [Skills](skills-guide.md) · [Eino 多代理](MULTI_AGENT_EINO.md) · [Agent 最终回复治理](agent-finalization-best-practices.md)
|
||||||
- [工作流](workflow-graph.md) · [工具执行治理](tool-execution-governance.md) · [人机协同最佳实践](hitl-best-practices.md)
|
- [工作流](workflow-graph.md) · [工具执行治理](tool-execution-governance.md) · [人机协同最佳实践](hitl-best-practices.md)
|
||||||
|
|
||||||
## 功能指南
|
## 功能指南
|
||||||
|
|||||||
@@ -12,9 +12,9 @@
|
|||||||
vision:
|
vision:
|
||||||
enabled: true
|
enabled: true
|
||||||
model: qwen-vl-max # 必填
|
model: qwen-vl-max # 必填
|
||||||
api_key: # 留空 → openai.api_key
|
api_key: # 留空 → 默认 AI 通道 api_key
|
||||||
base_url: # 留空 → openai.base_url
|
base_url: # 留空 → 默认 AI 通道 base_url
|
||||||
provider: # 留空 → openai.provider
|
provider: # 留空 → 默认 AI 通道 provider
|
||||||
max_image_bytes: 5242880
|
max_image_bytes: 5242880
|
||||||
max_dimension: 2048
|
max_dimension: 2048
|
||||||
jpeg_quality: 82
|
jpeg_quality: 82
|
||||||
@@ -28,7 +28,7 @@ vision:
|
|||||||
|
|
||||||
## Web 设置
|
## Web 设置
|
||||||
|
|
||||||
**系统设置 → 基本设置 → 视觉分析(analyze_image)** 可配置启用开关、视觉模型、API Key/Base URL(留空复用 OpenAI)、预处理参数;**保存并应用** 后写入 `config.yaml` 并重新注册 MCP 工具。
|
**系统设置 → 基本设置 → 视觉分析(analyze_image)** 可配置启用开关、视觉模型、API Key/Base URL(留空复用默认 AI 通道)、预处理参数;**保存并应用** 后写入 `config.yaml` 并重新注册 MCP 工具。
|
||||||
|
|
||||||
## 路径
|
## 路径
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,333 @@
|
|||||||
|
# Agent 最终回复治理最佳实践
|
||||||
|
|
||||||
|
[返回中文文档](README.md)
|
||||||
|
|
||||||
|
调研日期:2026-07-28
|
||||||
|
|
||||||
|
本文聚焦一个具体问题:Agent 在工具调用、推理、计划或子代理协作尚未真正完成时,输出了一段“像结论”的自然语言,前端或编排层把它当作最终回复展示。结论先说清楚:成熟 Agent 系统不会用“最近一段 assistant 文本”判断任务完成,而是用运行时状态、工具状态、验证结果和显式终态事件共同决定是否 final。
|
||||||
|
|
||||||
|
## 一、核心结论
|
||||||
|
|
||||||
|
1. **最终回复是运行时事件,不是自然语言内容。**
|
||||||
|
“已拿到”“下一步”“Huge breakthrough”这类文本只能作为候选观察或进展,不能作为完成信号。
|
||||||
|
|
||||||
|
2. **过程面和交付面必须隔离。**
|
||||||
|
`thinking`、`reasoning_chain`、`planning`、`response_delta`、子代理回复、工具输出都属于过程面;只有通过 final gate 的 `response` / `final` 事件才能写入主消息气泡和 `messages.content`。
|
||||||
|
|
||||||
|
3. **复杂任务需要 verifier,而不是更长 prompt。**
|
||||||
|
Prompt 可以提醒模型谨慎,但最终完成必须由代码层判断:是否仍有待执行工具、后台 execution、未完成计划步骤、未验证证据、未记录事实/漏洞、未清理或未说明不可清理。
|
||||||
|
|
||||||
|
4. **不同 agent 模式不同,但 final 治理原则一致。**
|
||||||
|
单代理、Deep、Plan-Execute、Supervisor 都需要 final gate。区别只是 gate 的证据来源不同:单代理看工具轨迹,Deep 还要看子代理结果,Plan-Execute 要看 Replanner 的终止判断,Supervisor 要看 `exit` 与 supervisor 汇总。
|
||||||
|
|
||||||
|
## 二、成熟 Agent 的公开做法
|
||||||
|
|
||||||
|
| 系统 | 公开做法 | 对 final 治理的启发 |
|
||||||
|
|---|---|---|
|
||||||
|
| Codex | OpenAI 的 Codex prompting guide 建议不要在 prompt 中强行要求 upfront plan、preamble 或 status updates,因为这可能导致 rollout 未完成就停止。 | 不要把“模型自己说的阶段性计划/状态”当完成依据;agent harness 应负责执行循环和收尾。 |
|
||||||
|
| Claude Code | Claude Code 提供 `PreToolUse`、`PostToolUse`、`Stop` 等 hooks;`PostToolUse` 明确发生在工具成功执行之后。 | 生命周期事件比自然语言可靠。验证、审计、阻断应挂在确定的阶段边界上。 |
|
||||||
|
| Claude Code Subagents | 子代理有独立上下文、自定义系统提示、特定工具权限和独立权限;子代理适合隔离大量检索/日志/文件读取。 | 子代理输出是证据材料,不是主任务最终结论;主代理必须汇总、验收、再 final。 |
|
||||||
|
| Claude Code Plan Mode | Plan mode 先读文件并产出计划,获得批准前不编辑。 | 计划与执行是不同状态;计划完成不等于任务完成。 |
|
||||||
|
| Cursor Plan Mode | Cursor Plan Mode 会研究代码库、询问澄清问题、生成可审查计划,并等待用户确认后再构建。 | UI 层把 plan/review/build 拆开,用户不会把计划误认为最终交付。 |
|
||||||
|
| OpenCode | OpenCode 把 Build、Plan、Review、Debug、Docs 等 agent 分成不同工具权限与用途,Plan agent 只分析规划不做修改。 | 用 agent 能力边界降低误触发:能规划的 agent 不等于能执行完成。 |
|
||||||
|
| Eino ADK | Eino ADK 提供事件驱动输出、Runner 回调、中断、checkpoint,以及 Supervisor、Plan-Execute 等协作原语。Plan-Execute 由 Planner、Executor、Replanner 协作。 | 当前项目选型方向正确;需要把事件驱动能力进一步固化为 finalization contract。 |
|
||||||
|
|
||||||
|
主要参考:
|
||||||
|
|
||||||
|
- OpenAI Codex Prompting Guide: https://developers.openai.com/cookbook/examples/gpt-5/codex_prompting_guide
|
||||||
|
- Claude Code Hooks: https://docs.anthropic.com/en/docs/claude-code/hooks
|
||||||
|
- Claude Code Subagents: https://docs.anthropic.com/en/docs/claude-code/sub-agents
|
||||||
|
- Claude Code Common Workflows: https://docs.anthropic.com/en/docs/claude-code/common-workflows
|
||||||
|
- Cursor Agent Best Practices: https://cursor.com/blog/agent-best-practices
|
||||||
|
- OpenCode Agents: https://opencode.ai/docs/agents/
|
||||||
|
- CloudWeGo Eino ADK: https://www.cloudwego.io/docs/eino/core_modules/eino_adk/
|
||||||
|
- CloudWeGo Eino ADK Patterns: https://www.cloudwego.io/docs/eino/overview/eino_adk0_1/
|
||||||
|
|
||||||
|
## 三、通用最佳实践
|
||||||
|
|
||||||
|
### 1. 建立 Finalization Contract
|
||||||
|
|
||||||
|
所有执行入口统一产出一个结构化收尾对象,只有它允许触发最终回复。
|
||||||
|
|
||||||
|
```go
|
||||||
|
type FinalizationDecision struct {
|
||||||
|
Status string // in_progress | completed | blocked | failed | cancelled
|
||||||
|
Finalizable bool
|
||||||
|
CompletionReason string // verified | user_cancelled | timeout | blocked | failed
|
||||||
|
FinalText string
|
||||||
|
EvidenceVerified bool
|
||||||
|
EvidenceRefs []string
|
||||||
|
PendingToolRuns []string
|
||||||
|
PendingPlanSteps []string
|
||||||
|
PendingApprovals []string
|
||||||
|
MissingChecks []string
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
硬规则:
|
||||||
|
|
||||||
|
- `Finalizable=false` 时禁止发送 `response` 终态事件。
|
||||||
|
- `Status=in_progress` 时只能发 `progress`、`planning`、`tool_*`、`reasoning_chain` 等过程事件。
|
||||||
|
- `FinalText` 不能为空,但非空不代表可以 final。
|
||||||
|
- `PendingToolRuns`、`PendingPlanSteps`、`PendingApprovals` 任一非空时不能 `completed`。
|
||||||
|
- `EvidenceVerified=false` 时不能把候选输出写成已验证结论。
|
||||||
|
|
||||||
|
### 2. 固定 SSE 事件语义
|
||||||
|
|
||||||
|
推荐事件分层:
|
||||||
|
|
||||||
|
| 事件 | 展示位置 | 可否写 `messages.content` | 说明 |
|
||||||
|
|---|---|---:|---|
|
||||||
|
| `progress` | 任务状态/时间线 | 否 | 简短进度 |
|
||||||
|
| `planning` | 执行详情 | 否 | 主代理计划、阶段性判断 |
|
||||||
|
| `reasoning_chain` / `thinking` | 执行详情 | 否 | 推理/思考摘要 |
|
||||||
|
| `tool_call` / `tool_result` | 执行详情 | 否 | 工具事件 |
|
||||||
|
| `eino_agent_reply` | 执行详情 | 否 | 子代理返回材料 |
|
||||||
|
| `finalization_check` | 执行详情 | 否 | verifier 结果 |
|
||||||
|
| `finalization_auto_continue` | 执行详情 | 否 | verifier 触发的工程续跑,`contextInjection=false` |
|
||||||
|
| `response` | 主消息气泡 | 是 | 只能在 `data.finalized=true` 时使用 |
|
||||||
|
| `done` | 关闭流 | 否 | 仅表示流结束,不表示任务成功 |
|
||||||
|
| `error` / `cancelled` | 主消息气泡或系统提示 | 是,终态失败类 | 必须带原因 |
|
||||||
|
|
||||||
|
### 3. 把“最终候选”与“最终回复”分开
|
||||||
|
|
||||||
|
模型可以输出候选结论,但候选结论必须先进入 `final_candidate` 或 `planning`,再由 verifier 决定是否提升:
|
||||||
|
|
||||||
|
```text
|
||||||
|
assistant text
|
||||||
|
-> candidate
|
||||||
|
-> finalization gate
|
||||||
|
-> response(finalized=true)
|
||||||
|
```
|
||||||
|
|
||||||
|
不要这样做:
|
||||||
|
|
||||||
|
```text
|
||||||
|
assistant text
|
||||||
|
-> response
|
||||||
|
```
|
||||||
|
|
||||||
|
### 4. Stop-time Verification
|
||||||
|
|
||||||
|
借鉴 Claude Code hook 思路,在 agent run 停止时做一次确定性检查:
|
||||||
|
|
||||||
|
- 所有工具调用都有对应 tool result。
|
||||||
|
- 后台 execution 都处于 terminal 状态,或被明确登记为仍在运行且任务状态为 `in_progress` / `blocked`。
|
||||||
|
- Plan-Execute 没有未执行的 required step。
|
||||||
|
- Supervisor 没有未汇总的子代理结果。
|
||||||
|
- 在 evidence-required 策略下,至少存在可查询到的 completed 工具执行证据。
|
||||||
|
|
||||||
|
### 5. 子代理输出只作证据
|
||||||
|
|
||||||
|
子代理返回不能直接成为用户最终回复。主代理必须完成:
|
||||||
|
|
||||||
|
- 去重和冲突合并。
|
||||||
|
- 证据强度排序。
|
||||||
|
- 不确定性标注。
|
||||||
|
- 范围边界确认。
|
||||||
|
- 用户可读交付。
|
||||||
|
|
||||||
|
### 6. Prompt 只做软约束,代码做硬约束
|
||||||
|
|
||||||
|
Prompt 中可以写:
|
||||||
|
|
||||||
|
```text
|
||||||
|
Interim observations must be marked as progress, not final.
|
||||||
|
Do not produce a final answer until verification is complete.
|
||||||
|
```
|
||||||
|
|
||||||
|
但真正决定 final 的必须是后端字段和状态机。否则模型只要生成一段像最终结论的自然语言,UI 仍可能误判。
|
||||||
|
|
||||||
|
## 四、CyberStrikeAI 当前落地状态
|
||||||
|
|
||||||
|
当前项目已经具备一套显式 final gate:
|
||||||
|
|
||||||
|
- [internal/agentfinalizer/decision.go](../../internal/agentfinalizer/decision.go) 是唯一的最终回复决策契约。
|
||||||
|
- [internal/handler/finalization_helpers.go](../../internal/handler/finalization_helpers.go) 负责把决策结果写入 `process_details`,并且只有 `Finalizable=true` 时才调用 `UpdateAssistantMessageFinalize`。
|
||||||
|
- [internal/handler/eino_single_agent.go](../../internal/handler/eino_single_agent.go)、[internal/handler/multi_agent.go](../../internal/handler/multi_agent.go)、[internal/handler/workflow_integration.go](../../internal/handler/workflow_integration.go)、[internal/handler/batch_queue_executor.go](../../internal/handler/batch_queue_executor.go) 均已在收尾处接入 finalizer。
|
||||||
|
- [web/static/js/monitor.js](../../web/static/js/monitor.js) 只把 `data.finalized === true` 的 `response` 当最终回复;未最终化文本会显示为最终回复检查未通过。
|
||||||
|
- [web/static/js/webshell.js](../../web/static/js/webshell.js) 将流式正文标记为候选输出,只有 `response(finalized=true)` 才切换为完成态。
|
||||||
|
- [internal/agentfinalizer/decision_test.go](../../internal/agentfinalizer/decision_test.go) 覆盖 pending tool、HITL、空输出、证据策略要求但缺执行证据、失败证据不能支撑最终化、完成态证据可 final 等回归场景。
|
||||||
|
- [internal/handler/finalization_auto_continue.go](../../internal/handler/finalization_auto_continue.go) 在缺 completed 执行证据时最多自动续跑 2 段;续跑只恢复已有模型轨迹,不向 agent 注入新的 user/system 文案。
|
||||||
|
|
||||||
|
当前契约的核心规则:
|
||||||
|
|
||||||
|
1. **模型自然语言只是 candidate。**
|
||||||
|
`RunResult.Response` 不能直接升级为最终回复,必须经过 `agentfinalizer.Decide`。
|
||||||
|
|
||||||
|
2. **所有 `response` 事件必须携带终态字段。**
|
||||||
|
至少包含 `finalized`、`finalizable`、`status`、`completionReason`、`evidenceVerified`、`evidenceRefs`、`pendingExecutionIds`、`missingChecks`。
|
||||||
|
|
||||||
|
3. **未完成工具会阻断 final。**
|
||||||
|
`queued/running` 工具执行仍存在时,决策结果为 `in_progress/pending_tool_executions`。
|
||||||
|
|
||||||
|
4. **执行证据必须由结构化策略声明。**
|
||||||
|
后端不从用户自然语言、助手回复或 agent mode 名称中推断执行意图。聊天请求通过 `finalization.requireExecutionEvidence` 显式声明;WebShell、Workflow、批量、机器人等执行入口由调用点显式传入 policy。policy 要求证据时,至少需要一个可查询到的 `completed` 工具执行记录;只有 failed/cancelled 记录不能支撑最终化。
|
||||||
|
|
||||||
|
5. **缺执行证据先工程续跑,再阻断。**
|
||||||
|
Eino 单代理和 Eino 多代理主链路在 `missing_execution_evidence` 时会先通过已有 trace 自动续跑,不注入额外上下文;达到续跑上限后仍缺证据才写入 blocked。
|
||||||
|
|
||||||
|
6. **HITL 和空输出不会 final。**
|
||||||
|
workflow 等待人工确认、空 assistant 文本、Eino 空输出占位均会写入阻断文案,而不是成功总结。
|
||||||
|
|
||||||
|
## 五、贴合当前项目的推荐架构
|
||||||
|
|
||||||
|
当前采用的链路是:
|
||||||
|
|
||||||
|
```text
|
||||||
|
Agent / Eino ADK events
|
||||||
|
-> event normalizer
|
||||||
|
-> process_details
|
||||||
|
-> finalization verifier
|
||||||
|
-> response(finalized=true)
|
||||||
|
-> messages.content
|
||||||
|
```
|
||||||
|
|
||||||
|
### 1. 后端统一 Finalizer
|
||||||
|
|
||||||
|
职责:
|
||||||
|
|
||||||
|
- 接收 `RunResult` / 候选文本、`mcpExecutionIds`、会话与助手消息 ID、HITL 状态、编排模式。
|
||||||
|
- 通过数据库查询工具执行状态,识别 pending、completed、failed、cancelled 等证据状态。
|
||||||
|
- 返回 `FinalizationDecision`。
|
||||||
|
- 不调用高风险工具,只做状态和证据检查。
|
||||||
|
|
||||||
|
### 2. RunResult 终态字段
|
||||||
|
|
||||||
|
[internal/multiagent/runner.go](../../internal/multiagent/runner.go) 已扩展终态字段:
|
||||||
|
|
||||||
|
```go
|
||||||
|
type RunResult struct {
|
||||||
|
Response string
|
||||||
|
MCPExecutionIDs []string
|
||||||
|
LastAgentTraceInput string
|
||||||
|
LastAgentTraceOutput string
|
||||||
|
|
||||||
|
Finalized bool
|
||||||
|
Status string
|
||||||
|
CompletionReason string
|
||||||
|
EvidenceVerified bool
|
||||||
|
EvidenceRefs []string
|
||||||
|
PendingExecutionIDs []string
|
||||||
|
MissingChecks []string
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. 发送 `response` 的条件
|
||||||
|
|
||||||
|
在单代理、多代理、工作流、批处理收尾处统一执行:
|
||||||
|
|
||||||
|
```go
|
||||||
|
decision := h.finalizeAgentRunForDelivery(...)
|
||||||
|
if !decision.Finalizable {
|
||||||
|
sendEvent("finalization_check", "任务尚未达到最终回复条件", decision)
|
||||||
|
sendEvent("response", finalizationBlockedMessage(decision), finalizationResponsePayload(decision, extra))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
sendEvent("response", decision.FinalText, finalizationResponsePayload(decision, extra))
|
||||||
|
```
|
||||||
|
|
||||||
|
### 4. 前端只信 `finalized=true`
|
||||||
|
|
||||||
|
在 [web/static/js/monitor.js](../../web/static/js/monitor.js) 的 `case 'response'` 中执行硬判断:
|
||||||
|
|
||||||
|
```js
|
||||||
|
const responseFinalized = isFinalizedResponseData(responseData);
|
||||||
|
const bubbleText = responseFinalized
|
||||||
|
? resolvedResponseText
|
||||||
|
: (event.message || '任务尚未达到最终回复条件,暂不生成成功结论。');
|
||||||
|
markAssistantFinalizationState(assistantIdFinal, responseData);
|
||||||
|
```
|
||||||
|
|
||||||
|
WebShell 侧同理:`response_delta` 可以用于实时预览,但 UI 文案应标记为“执行中输出”,只有最终 `response(finalized=true)` 才显示为完成态。
|
||||||
|
|
||||||
|
### 5. 各模式 final gate
|
||||||
|
|
||||||
|
| 模式 | 谁可以产出最终候选 | 谁决定 final | 必须检查 |
|
||||||
|
|---|---|---|---|
|
||||||
|
| Eino 单代理 | 单代理最后助手文本 | Finalizer | 无 pending tool、证据引用完整、任务状态 terminal |
|
||||||
|
| Deep | 主代理汇总文本 | Finalizer | 子代理结果已汇总;子代理文本不能直接 final;工具状态 terminal |
|
||||||
|
| Plan-Execute | Replanner 结束后的汇总文本 | Replanner + Finalizer | Executor 单步输出不能 final;计划步骤完成或明确 blocked |
|
||||||
|
| Supervisor | Supervisor 的 `exit` / 汇总文本 | Supervisor + Finalizer | transfer 已返回;无未处理专家结果;最终由 supervisor 统一口径 |
|
||||||
|
|
||||||
|
### 6. 安全测试场景的证据 gate
|
||||||
|
|
||||||
|
安全测试、WebShell、批量验证、Workflow 和多代理执行等 evidence-required 场景,最终回复必须至少满足:
|
||||||
|
|
||||||
|
- 有明确目标和授权范围标识。
|
||||||
|
- 有可复核证据引用,例如工具 execution id、请求/响应摘要、截图路径、命令输出摘要、事实/漏洞记录 ID。
|
||||||
|
- 有身份或影响验证结果,而不是只凭 marker 文本判断。
|
||||||
|
- 已记录到项目黑板或漏洞库,或明确说明未绑定项目导致无法记录。
|
||||||
|
- 高风险动作已清理、回滚、取消,或明确说明未执行清理的原因。
|
||||||
|
- 仍在运行的扫描/命令/WebShell/C2 任务不能被隐式当作完成。
|
||||||
|
|
||||||
|
注意:这里的 gate 是治理规则,不要求最终报告暴露敏感利用细节;可以只给证据摘要和内部引用。
|
||||||
|
|
||||||
|
## 六、落地状态与后续增强
|
||||||
|
|
||||||
|
### P0:先修“误 final”(已落地)
|
||||||
|
|
||||||
|
1. 已引入 `FinalizationDecision`。
|
||||||
|
2. 主要 agent SSE `response` 事件已携带 `data.finalized/finalizable/status/completionReason` 等字段。
|
||||||
|
3. 前端 `monitor.js` 和 `webshell.js` 已按 `finalized=true` 区分候选输出和最终回复。
|
||||||
|
4. `RunResult.Response` 仍保留兼容字段名,但语义已由 finalizer 统一提升;后续可再拆成 `CandidateResponse` / `FinalResponse`,减少误用空间。
|
||||||
|
5. Plan-Execute / Deep / Supervisor / Eino Single 等模式均通过统一 handler 收尾 gate。
|
||||||
|
|
||||||
|
### P1:补证据链(部分落地)
|
||||||
|
|
||||||
|
1. 已用 `mcp_execution:<id>` 作为基础 evidence refs。
|
||||||
|
2. `finalization_check` 事件已展示 pending execution 与 missing checks。
|
||||||
|
3. 执行入口已启用显式 execution evidence policy;Eino 主链路在 policy 要求证据且缺少 completed 工具证据时先无注入续跑,达到上限后才阻断 final。
|
||||||
|
4. 后续建议:为 `record_vulnerability`、`upsert_project_fact`、项目黑板记录建立更细粒度 evidence refs。
|
||||||
|
5. 后续建议:最终报告模板固定包含“结论、证据、风险/不确定性、后续动作”。
|
||||||
|
|
||||||
|
### P2:体验和观测(后续增强)
|
||||||
|
|
||||||
|
1. 在任务卡片展示 `in_progress / verifying / finalizing / completed / blocked`。
|
||||||
|
2. 为 finalizer 加日志和指标:误拦截率、缺失证据类型、pending tool 数量。
|
||||||
|
3. 支持“继续验证”按钮,从 `FinalizationDecision.MissingChecks` 自动生成下一轮输入。
|
||||||
|
|
||||||
|
## 七、验收测试建议
|
||||||
|
|
||||||
|
至少加入这些回归测试:
|
||||||
|
|
||||||
|
1. **推理文本不 final**
|
||||||
|
模拟 `reasoning_chain` 里出现看似完成的候选结论,但本轮没有 completed 工具执行证据;预期主消息气泡不显示成功结论,只显示执行中或阻断态。
|
||||||
|
|
||||||
|
2. **主代理阶段性输出不 final**
|
||||||
|
模拟 `response_start/delta` 输出“下一步继续验证”;预期只进入 timeline `planning`。
|
||||||
|
|
||||||
|
3. **未完成后台工具不 final**
|
||||||
|
工具返回 `execution_id` 且状态 `running`;即使模型给出总结,也只能 `in_progress`。
|
||||||
|
|
||||||
|
4. **Plan-Execute Executor 输出不 final**
|
||||||
|
Executor 输出“突破成功”,但 Replanner 未结束;预期不触发 `messages.content` finalize。
|
||||||
|
|
||||||
|
5. **Supervisor 子代理输出不 final**
|
||||||
|
子代理返回确定结论,Supervisor 未 `exit`;预期只进入 `eino_agent_reply`。
|
||||||
|
|
||||||
|
6. **最终事件必须带 finalized**
|
||||||
|
前端收到旧格式 `response` 无 `finalized=true`;预期候选内容只进入详情/警告,主消息显示阻断态,不创建成功最终气泡。
|
||||||
|
|
||||||
|
7. **失败和取消可终态**
|
||||||
|
`error` / `cancelled` 仍可更新助手消息,但 `completionReason` 必须是 `failed` / `user_cancelled`,不能伪装为成功完成。
|
||||||
|
|
||||||
|
## 八、推荐默认策略
|
||||||
|
|
||||||
|
对 CyberStrikeAI,建议默认策略是:
|
||||||
|
|
||||||
|
```text
|
||||||
|
eino_single:轻量任务可用,但 final gate 必须开启
|
||||||
|
deep:复杂安全测试默认推荐
|
||||||
|
plan_execute:目标明确、需要严格“规划-执行-重规划”的任务推荐
|
||||||
|
supervisor:多专家路由任务使用,不作为默认泛化模式
|
||||||
|
```
|
||||||
|
|
||||||
|
最终治理一句话:
|
||||||
|
|
||||||
|
```text
|
||||||
|
messages.content 只能来自 FinalizationDecision.FinalText;
|
||||||
|
process_details 可以展示所有过程;
|
||||||
|
前端只能把 response(finalized=true) 当最终回复。
|
||||||
|
```
|
||||||
@@ -52,6 +52,18 @@ Content-Type: application/json
|
|||||||
- `plan_execute`
|
- `plan_execute`
|
||||||
- `supervisor`
|
- `supervisor`
|
||||||
|
|
||||||
|
常用请求体字段:
|
||||||
|
|
||||||
|
| 字段 | 说明 |
|
||||||
|
| --- | --- |
|
||||||
|
| `message` | 用户消息,必填。 |
|
||||||
|
| `conversationId` | 继续已有对话;为空时创建新对话。 |
|
||||||
|
| `projectId` | 新对话绑定项目;为空时可跟随 `config.project.default_project_id`。 |
|
||||||
|
| `role` | 使用指定角色。 |
|
||||||
|
| `aiChannelId` | 选择 `ai.channels` 中的通道 ID;为空时使用 `ai.default_channel`。 |
|
||||||
|
| `reasoning` | 会话级推理覆盖,受通道 `reasoning.allow_client_reasoning` 控制。 |
|
||||||
|
| `hitl` | 会话级人机协同配置。 |
|
||||||
|
|
||||||
对话管理:
|
对话管理:
|
||||||
|
|
||||||
- `POST /api/conversations`
|
- `POST /api/conversations`
|
||||||
@@ -62,6 +74,25 @@ Content-Type: application/json
|
|||||||
- `POST /api/conversations/:id/delete-turn`
|
- `POST /api/conversations/:id/delete-turn`
|
||||||
- `GET /api/messages/:id/process-details`
|
- `GET /api/messages/:id/process-details`
|
||||||
|
|
||||||
|
## 文件管理来源
|
||||||
|
|
||||||
|
文件管理页面和 `/api/chat-uploads` 列表接口会把对话相关文件按来源归类。底层目录仍使用项目 ID 或会话 ID 保持稳定,界面会优先显示项目名或对话标题,完整 ID 可在提示或路径中查看。
|
||||||
|
|
||||||
|
| 来源 | `source` | 典型目录 | 说明 | 可变更性 |
|
||||||
|
| --- | --- | --- | --- | --- |
|
||||||
|
| 工作目录 | `workspace` | `tmp/workspace/projects/<projectId>/...`、`tmp/workspace/conversations/<conversationId>/...` | Agent 执行任务时保存下载文件、分析脚本、中间结果和生成的 CSV/XLSX/Markdown 等。用户反馈“AI 生成的文件找不到”时,通常先看这里。 | 只读展示;支持复制路径、下载、导出。 |
|
||||||
|
| 会话产物 | `conversation_artifact` | `data/conversation_artifacts/<conversationId>/...` | 系统按会话归档的交付物或会话级产物,例如总结、报告、模型中间件生成的归档内容。 | 只读展示;支持复制路径、下载、导出。 |
|
||||||
|
| 工具输出 | `reduction` | `tmp/reduction/projects/<projectId>/...`、`tmp/reduction/conversations/<conversationId>/...` | 超长工具输出、扫描原文或被截断前落盘的结果缓存。适合回看完整工具输出。 | 只读展示;支持复制路径、下载、导出。 |
|
||||||
|
| 对话附件 | `upload` | `chat_uploads/<date>/<conversationId>/...` | 用户在对话或文件管理页手动上传的附件。需要让 AI 引用某文件时,可复制服务器绝对路径粘贴到对话中。 | 可上传、新建目录、编辑文本文件、重命名、删除、复制路径、下载、导出。 |
|
||||||
|
|
||||||
|
相关接口:
|
||||||
|
|
||||||
|
- `GET /api/chat-uploads`:按来源、项目、会话、文件名筛选文件。
|
||||||
|
- `GET /api/chat-uploads/path`:把文件管理中的相对路径或内部虚拟路径解析为服务器绝对路径,用于复制文件或目录路径。
|
||||||
|
- `GET /api/chat-uploads/download`:下载指定文件。
|
||||||
|
- `GET /api/chat-uploads/export`:导出当前筛选结果为 ZIP。
|
||||||
|
- `POST /api/chat-uploads`:上传到对话附件目录。
|
||||||
|
|
||||||
## 项目、漏洞、攻击链
|
## 项目、漏洞、攻击链
|
||||||
|
|
||||||
项目:
|
项目:
|
||||||
|
|||||||
+50
-18
@@ -29,26 +29,57 @@ log:
|
|||||||
- `auth.session_duration_hours`:登录会话有效期(小时)。登录密码由 RBAC 用户管理,首次启动时在控制台输出 `admin` 初始密码。
|
- `auth.session_duration_hours`:登录会话有效期(小时)。登录密码由 RBAC 用户管理,首次启动时在控制台输出 `admin` 初始密码。
|
||||||
- `log.output`:可以是 `stdout`、`stderr` 或文件路径。
|
- `log.output`:可以是 `stdout`、`stderr` 或文件路径。
|
||||||
|
|
||||||
## 模型配置
|
## AI 通道与模型配置
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
openai:
|
ai:
|
||||||
provider: openai
|
default_channel: openai-main
|
||||||
base_url: https://api.openai.com/v1
|
channels:
|
||||||
api_key: sk-...
|
openai-main:
|
||||||
model: gpt-4.1
|
name: OpenAI Main
|
||||||
max_total_tokens: 120000
|
provider: openai_compatible
|
||||||
reasoning:
|
base_url: https://api.openai.com/v1
|
||||||
mode: on
|
api_key: sk-...
|
||||||
effort: high
|
model: gpt-4.1
|
||||||
allow_client_reasoning: true
|
max_total_tokens: 120000
|
||||||
profile: openai_compat
|
max_completion_tokens: 16384
|
||||||
|
reasoning:
|
||||||
|
mode: on
|
||||||
|
effort: high
|
||||||
|
allow_client_reasoning: true
|
||||||
|
profile: openai_compat
|
||||||
|
claude-main:
|
||||||
|
name: Claude Main
|
||||||
|
provider: claude
|
||||||
|
base_url: https://api.anthropic.com/v1
|
||||||
|
api_key: sk-ant-...
|
||||||
|
model: claude-sonnet-4-5
|
||||||
```
|
```
|
||||||
|
|
||||||
- `provider`:`openai` 表示 OpenAI 兼容接口;`claude` 会桥接到 Anthropic Claude Messages API。
|
`ai` 是推荐的模型配置入口。系统设置页对应路径是 **系统设置 → 基本设置 → AI 通道配置**,保存后写入 `ai.default_channel` 和 `ai.channels`。旧版 `openai` 字段仍保留为兼容运行时字段;加载配置时会确保至少有一个默认通道,并把 `ai.default_channel` 解析后的配置同步到运行时 `openai`。
|
||||||
- `base_url/api_key/model`:主模型配置。
|
|
||||||
- `max_total_tokens`:上下文压缩、攻击链构建、多代理摘要等共用的总预算。
|
通道字段:
|
||||||
- `reasoning`:控制推理扩展字段。不同网关支持差异较大,异常时先尝试 `mode: off`。
|
|
||||||
|
| 字段 | 说明 |
|
||||||
|
| --- | --- |
|
||||||
|
| `ai.default_channel` | 默认通道 ID。新对话、机器人、批量任务和未显式选择通道的请求使用它。 |
|
||||||
|
| `ai.channels.<id>` | 通道配置。ID 会归一化为小写、数字和短横线,例如 `Qwen_Max` 会变成 `qwen-max`。 |
|
||||||
|
| `name` | Web UI 展示名。留空时使用通道 ID。 |
|
||||||
|
| `provider` | `openai_compatible` 或 `claude`。`openai_compatible` 会在运行时映射为 `openai`;`claude` 会桥接到 Anthropic Messages API。 |
|
||||||
|
| `base_url/api_key/model` | 必填。Base URL 通常需要包含版本路径,如 OpenAI/兼容网关的 `/v1`。 |
|
||||||
|
| `max_total_tokens` | 上下文压缩、攻击链构建、多代理摘要等共用的总预算。 |
|
||||||
|
| `max_completion_tokens` | 单次模型输出上限;未填时使用默认值。 |
|
||||||
|
| `reasoning` | 该通道的默认推理扩展字段。不同网关支持差异较大,异常时先尝试 `mode: off`。 |
|
||||||
|
|
||||||
|
对话页的“AI 通道”下拉框会读取已保存通道。请求体中的 `aiChannelId` 非空时仅对本次/本会话运行配置生效,不会把 API Key 发送给模型;为空时跟随 `ai.default_channel`。
|
||||||
|
|
||||||
|
常用操作:
|
||||||
|
|
||||||
|
- 新增:点击左侧 `+`,填写必填字段后保存。
|
||||||
|
- 设默认:选中通道后点击“设为默认”,保存并应用后新请求生效。
|
||||||
|
- 复制:以当前表单内容创建副本,适合为同一服务商配置不同模型。
|
||||||
|
- 删除:默认通道不能作为批量删除目标;删除后需保留至少一个通道。
|
||||||
|
- 探活:单通道“测试连接”或左侧“批量探活”会调用模型测试接口,适合验证 Key、Base URL 和模型名。
|
||||||
|
|
||||||
## Agent
|
## Agent
|
||||||
|
|
||||||
@@ -213,7 +244,8 @@ project:
|
|||||||
|
|
||||||
| 配置段 | 应用后通常立即生效 | 需要额外动作 |
|
| 配置段 | 应用后通常立即生效 | 需要额外动作 |
|
||||||
| --- | --- | --- |
|
| --- | --- | --- |
|
||||||
| `openai` | 新请求使用新模型配置 | 旧的流式请求不会被强制切换 |
|
| `ai.default_channel` / `ai.channels` | 新请求使用解析后的默认或选定通道 | 旧的流式请求不会被强制切换;前端通道列表需要重新读取配置 |
|
||||||
|
| `openai` | 兼容字段;通常由默认 AI 通道同步 | 新配置优先维护 `ai.channels` |
|
||||||
| `agent.max_iterations` | 新 Agent 任务生效 | 已运行任务按启动时状态继续 |
|
| `agent.max_iterations` | 新 Agent 任务生效 | 已运行任务按启动时状态继续 |
|
||||||
| `security.tool_description_mode` | 工具重新暴露时生效 | 模型已有上下文不会回滚 |
|
| `security.tool_description_mode` | 工具重新暴露时生效 | 模型已有上下文不会回滚 |
|
||||||
| `hitl.tool_whitelist` | 新工具调用审批判断生效 | 已挂起审批不自动重判 |
|
| `hitl.tool_whitelist` | 新工具调用审批判断生效 | 已挂起审批不自动重判 |
|
||||||
@@ -228,7 +260,7 @@ project:
|
|||||||
几个字段有“留空复用”的关系:
|
几个字段有“留空复用”的关系:
|
||||||
|
|
||||||
- `vision.api_key/base_url/provider` 留空时复用 `openai`。
|
- `vision.api_key/base_url/provider` 留空时复用 `openai`。
|
||||||
- `hitl.audit_model` 留空时复用 `openai`。
|
- `hitl.audit_model` 留空时复用默认 AI 通道解析后的 `openai`。
|
||||||
- `knowledge.embedding.base_url/api_key` 留空时复用主模型或 embedding 默认配置。
|
- `knowledge.embedding.base_url/api_key` 留空时复用主模型或 embedding 默认配置。
|
||||||
- `knowledge.retrieval.rerank.base_url/api_key` 留空时复用 embedding/openai。
|
- `knowledge.retrieval.rerank.base_url/api_key` 留空时复用 embedding/openai。
|
||||||
- `database.knowledge_db_path` 留空时可以复用主会话数据库,但独立文件更利于备份。
|
- `database.knowledge_db_path` 留空时可以复用主会话数据库,但独立文件更利于备份。
|
||||||
|
|||||||
@@ -10,7 +10,7 @@
|
|||||||
- Python:部分 MCP 服务或工具脚本需要 Python 运行环境。
|
- Python:部分 MCP 服务或工具脚本需要 Python 运行环境。
|
||||||
- SQLite:默认使用文件型数据库,无需单独服务。
|
- SQLite:默认使用文件型数据库,无需单独服务。
|
||||||
- 安全工具:`tools/` 中的 YAML 只是工具定义,实际命令如 `nmap`、`sqlmap`、`nuclei` 仍需安装到系统 PATH。
|
- 安全工具:`tools/` 中的 YAML 只是工具定义,实际命令如 `nmap`、`sqlmap`、`nuclei` 仍需安装到系统 PATH。
|
||||||
- 模型服务:需要 OpenAI 兼容 API,或配置 `openai.provider: claude` 走 Claude 桥接。
|
- 模型服务:至少配置一个 `ai.channels` 通道;`provider: openai_compatible` 适用于 OpenAI 兼容 API,`provider: claude` 会走 Claude 桥接。
|
||||||
|
|
||||||
建议目录:
|
建议目录:
|
||||||
|
|
||||||
|
|||||||
@@ -23,12 +23,12 @@ hitl:
|
|||||||
provider: ""
|
provider: ""
|
||||||
base_url: ""
|
base_url: ""
|
||||||
api_key: ""
|
api_key: ""
|
||||||
model: "" # 可填小模型;留空复用 openai.model
|
model: "" # 可填小模型;留空复用默认 AI 通道的模型
|
||||||
retention_days: 90
|
retention_days: 90
|
||||||
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
|
tool_whitelist: [read_file, list_dir, glob, grep, tool_search]
|
||||||
```
|
```
|
||||||
|
|
||||||
`audit_model` 的字段可以只填一部分。空字段会自动继承主 `openai` 配置,因此常见做法是只填 `model`,让审计 Agent 使用更便宜的小模型。
|
`audit_model` 的字段可以只填一部分。空字段会自动继承默认 AI 通道解析后的模型配置,因此常见做法是只填 `model`,让审计 Agent 使用更便宜的小模型。
|
||||||
|
|
||||||
## 推荐审批策略
|
## 推荐审批策略
|
||||||
|
|
||||||
|
|||||||
@@ -51,6 +51,13 @@ Agent 调用工具
|
|||||||
| `wait_tool_execution` | 等待指定 execution 一段时间 |
|
| `wait_tool_execution` | 等待指定 execution 一段时间 |
|
||||||
| `cancel_tool_execution` | 主动取消指定 execution |
|
| `cancel_tool_execution` | 主动取消指定 execution |
|
||||||
|
|
||||||
|
`get_tool_execution` 与 `wait_tool_execution` 支持返回运行中输出预览:
|
||||||
|
|
||||||
|
- `include_partial_output`:是否返回 partial output,默认 `true`。
|
||||||
|
- `partial_output_max_bytes`:本次返回的尾部预览上限,默认 `4096`,最大 `65536`。
|
||||||
|
|
||||||
|
partial output 是“已产生输出的有界预览”,不等同于最终 `result`。最终 `result` 仍只在工具结束时写入 canonical execution 记录;不支持流式输出的工具不会返回 partial 字段。
|
||||||
|
|
||||||
典型流程:
|
典型流程:
|
||||||
|
|
||||||
```text
|
```text
|
||||||
@@ -60,6 +67,8 @@ Agent 调用工具
|
|||||||
4. 仍未完成时可继续等待,或调用 cancel_tool_execution
|
4. 仍未完成时可继续等待,或调用 cancel_tool_execution
|
||||||
```
|
```
|
||||||
|
|
||||||
|
`tool_wait_timeout_seconds` 适用于内部 MCP、外部 MCP,以及 Eino filesystem 的流式 `execute`。Eino 的 `ls/read_file/write_file/edit_file/glob/grep` 等非流式 filesystem 工具会写入 execution 监控记录,但不作为后台 worker 做软等待续跑。
|
||||||
|
|
||||||
## 取消和会话清理
|
## 取消和会话清理
|
||||||
|
|
||||||
- 用户点击“停止任务”时,会取消当前会话仍在运行的工具。
|
- 用户点击“停止任务”时,会取消当前会话仍在运行的工具。
|
||||||
|
|||||||
@@ -32,29 +32,40 @@ https://127.0.0.1:8080/
|
|||||||
|
|
||||||
如果仍有其他具备 `rbac:write` 权限的管理员账号,优先在 **平台权限 → 用户管理** 中重置密码。
|
如果仍有其他具备 `rbac:write` 权限的管理员账号,优先在 **平台权限 → 用户管理** 中重置密码。
|
||||||
|
|
||||||
如果没有可用的管理员会话,可在服务器上紧急重置内置 `admin` 账号。先停止 CyberStrikeAI 服务并备份数据库,然后在项目根目录执行以下命令,按提示输入并确认新密码:
|
如果没有可用的管理员会话,可在服务器上紧急重置内置 `admin` 账号。在项目根目录执行:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
./run.sh --reset-admin-password
|
||||||
|
```
|
||||||
|
|
||||||
|
按提示输入并确认新密码。脚本会隐藏输入并写入 bcrypt 哈希。如果服务正在运行,完成后重新启动服务,使原有登录会话失效。
|
||||||
|
|
||||||
|
如果无法使用 `run.sh`,也可以手动执行以下命令,按提示输入并确认新密码:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
HASH=$(htpasswd -nBC 10 '' | cut -d: -f2 | tr -d '\n') && sqlite3 data/conversations.db "UPDATE rbac_users SET password_hash='$HASH', updated_at=CURRENT_TIMESTAMP WHERE id='admin' AND username='admin' AND is_builtin=1; SELECT changes();"
|
HASH=$(htpasswd -nBC 10 '' | cut -d: -f2 | tr -d '\n') && sqlite3 data/conversations.db "UPDATE rbac_users SET password_hash='$HASH', updated_at=CURRENT_TIMESTAMP WHERE id='admin' AND username='admin' AND is_builtin=1; SELECT changes();"
|
||||||
```
|
```
|
||||||
|
|
||||||
输出 `1` 表示修改成功。该命令需要 `sqlite3` 和 `htpasswd`;如果 `config.yaml` 中的 `database.path` 不是默认值,请替换 `data/conversations.db`。密码输入不会显示,也不会写入 Shell 历史,并以 bcrypt 哈希保存。完成后重新启动服务,使原有登录会话失效。
|
输出 `1` 表示修改成功。该命令需要 `sqlite3` 和 `htpasswd`;如果 `config.yaml` 中的 `database.path` 不是默认值,请替换 `data/conversations.db`。密码输入不会显示,也不会写入 Shell 历史。
|
||||||
|
|
||||||
## 模型无响应
|
## 模型无响应
|
||||||
|
|
||||||
检查:
|
检查:
|
||||||
|
|
||||||
- `openai.base_url` 是否包含正确路径,如 `/v1`。
|
- 当前对话选择的 AI 通道是否存在;为空时会使用 `ai.default_channel`。
|
||||||
- `openai.api_key` 是否有效。
|
- `ai.channels.<id>.base_url` 是否包含正确路径,如 `/v1`。
|
||||||
- `openai.model` 是否存在。
|
- `ai.channels.<id>.api_key` 是否有效。
|
||||||
- 服务商是否支持当前 `reasoning` 字段。
|
- `ai.channels.<id>.model` 是否存在。
|
||||||
|
- 服务商是否支持当前通道的 `reasoning` 字段。
|
||||||
|
|
||||||
可在系统设置中使用模型测试。若网关报 400,先尝试:
|
可在系统设置中使用模型测试。若网关报 400,先尝试:
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
openai:
|
ai:
|
||||||
reasoning:
|
channels:
|
||||||
mode: off
|
your-channel:
|
||||||
|
reasoning:
|
||||||
|
mode: off
|
||||||
```
|
```
|
||||||
|
|
||||||
## 流式输出中断
|
## 流式输出中断
|
||||||
|
|||||||
@@ -34,7 +34,39 @@
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 三、执行模型(先理解再配置)
|
## 三、自然语言生成草稿
|
||||||
|
|
||||||
|
工作流页面提供 **用自然语言创建** 入口。用户描述一句安全作业目标后,平台会生成一个可编辑草稿,并返回结构化审计结果:
|
||||||
|
|
||||||
|
- 只生成草稿,不会自动保存、试运行或真实执行工具。
|
||||||
|
- 服务端接口为 `POST /api/workflows/generate-draft`,权限为 `workflow:write`。
|
||||||
|
- 生成结果包含 `graph`、`meta`、`capabilities`、`audit` 和 `stats`,前端应用到画布后仍走现有保存校验。
|
||||||
|
- 高风险语义(执行脚本、隔离、封禁、删除、利用等)会标记 `high_risk`,并默认插入 HITL 审批或 `requires_human_confirmation`。
|
||||||
|
- 工具能力按已有工具列表匹配;未匹配到时降级为 Agent 草稿,并在 `audit.missing_fields` / `audit.assumptions` 中提示需要补配置。
|
||||||
|
- 如果服务端生成不可用,前端会使用本地确定性兜底生成器,继续给出可编辑草稿和风险提示。
|
||||||
|
|
||||||
|
示例请求:
|
||||||
|
|
||||||
|
```http
|
||||||
|
POST /api/workflows/generate-draft
|
||||||
|
Content-Type: application/json
|
||||||
|
|
||||||
|
{
|
||||||
|
"prompt": "对目标资产做端口扫描,如果发现高危端口就创建漏洞任务并通知负责人审批,最后输出报告",
|
||||||
|
"options": {
|
||||||
|
"include_objective": true,
|
||||||
|
"allow_schedule": false,
|
||||||
|
"allow_high_risk": false
|
||||||
|
},
|
||||||
|
"available_tools": [
|
||||||
|
{ "key": "nmap", "name": "nmap", "enabled": true }
|
||||||
|
]
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## 四、执行模型(先理解再配置)
|
||||||
|
|
||||||
工作流按 **有向图** 执行,引擎从 **开始** 节点出发,沿连线依次运行下游节点。
|
工作流按 **有向图** 执行,引擎从 **开始** 节点出发,沿连线依次运行下游节点。
|
||||||
|
|
||||||
@@ -116,7 +148,7 @@ outputs["你填的变量名"] = 节点输出内容
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 四、模板语法
|
## 五、模板语法
|
||||||
|
|
||||||
### 4.1 基本格式
|
### 4.1 基本格式
|
||||||
|
|
||||||
@@ -198,7 +230,7 @@ jq({{outputs.scan}}, ".severity") == "high"
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 五、节点类型与配置
|
## 六、节点类型与配置
|
||||||
|
|
||||||
### 5.1 开始(start)
|
### 5.1 开始(start)
|
||||||
|
|
||||||
@@ -318,7 +350,7 @@ HITL 等待信息会记录:
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 六、连线配置
|
## 七、连线配置
|
||||||
|
|
||||||
选中 **连线** 后,右侧可配置 **连线条件**。
|
选中 **连线** 后,右侧可配置 **连线条件**。
|
||||||
|
|
||||||
@@ -335,7 +367,7 @@ HITL 等待信息会记录:
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 七、完整示例:跨条件节点传递 Agent 输出
|
## 八、完整示例:跨条件节点传递 Agent 输出
|
||||||
|
|
||||||
### 7.1 流程结构
|
### 7.1 流程结构
|
||||||
|
|
||||||
@@ -384,7 +416,7 @@ HITL 等待信息会记录:
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 八、绑定角色并运行
|
## 九、绑定角色并运行
|
||||||
|
|
||||||
### 8.1 在角色管理中绑定
|
### 8.1 在角色管理中绑定
|
||||||
|
|
||||||
@@ -414,7 +446,7 @@ workflow_policy: auto
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 九、调试、试运行与复盘
|
## 十、调试、试运行与复盘
|
||||||
|
|
||||||
### 9.1 安全试运行(dry-run)
|
### 9.1 安全试运行(dry-run)
|
||||||
|
|
||||||
@@ -487,7 +519,7 @@ token 与成本是否存在取决于底层模型/Agent 事件是否上报 usage
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 十、保存前校验规则
|
## 十一、保存前校验规则
|
||||||
|
|
||||||
保存时系统会自动检查:
|
保存时系统会自动检查:
|
||||||
|
|
||||||
@@ -510,7 +542,7 @@ token 与成本是否存在取决于底层模型/Agent 事件是否上报 usage
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 十一、排错指南
|
## 十二、排错指南
|
||||||
|
|
||||||
| 现象 | 可能原因 | 处理建议 |
|
| 现象 | 可能原因 | 处理建议 |
|
||||||
|------|----------|----------|
|
|------|----------|----------|
|
||||||
@@ -526,7 +558,7 @@ token 与成本是否存在取决于底层模型/Agent 事件是否上报 usage
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 十二、最佳实践
|
## 十三、最佳实践
|
||||||
|
|
||||||
1. **命名规范**:为每个需要被引用的节点设置有意义的输出变量名,如 `scan_result`、`parsed_targets`,避免都叫 `agent_result`。
|
1. **命名规范**:为每个需要被引用的节点设置有意义的输出变量名,如 `scan_result`、`parsed_targets`,避免都叫 `agent_result`。
|
||||||
2. **跨节点传参优先用 `outputs`**:只要中间可能插入条件、工具、审批节点,就应用命名变量。
|
2. **跨节点传参优先用 `outputs`**:只要中间可能插入条件、工具、审批节点,就应用命名变量。
|
||||||
@@ -539,7 +571,7 @@ token 与成本是否存在取决于底层模型/Agent 事件是否上报 usage
|
|||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## 十三、相关代码位置(开发者参考)
|
## 十四、相关代码位置(开发者参考)
|
||||||
|
|
||||||
| 模块 | 路径 |
|
| 模块 | 路径 |
|
||||||
|------|------|
|
|------|------|
|
||||||
|
|||||||
@@ -37,6 +37,7 @@ require (
|
|||||||
go.opentelemetry.io/otel/trace v1.34.0
|
go.opentelemetry.io/otel/trace v1.34.0
|
||||||
go.uber.org/zap v1.26.0
|
go.uber.org/zap v1.26.0
|
||||||
golang.org/x/net v0.35.0
|
golang.org/x/net v0.35.0
|
||||||
|
golang.org/x/term v0.32.0
|
||||||
golang.org/x/text v0.26.0
|
golang.org/x/text v0.26.0
|
||||||
golang.org/x/time v0.14.0
|
golang.org/x/time v0.14.0
|
||||||
gopkg.in/yaml.v3 v3.0.1
|
gopkg.in/yaml.v3 v3.0.1
|
||||||
|
|||||||
Binary file not shown.
|
Before Width: | Height: | Size: 88 KiB After Width: | Height: | Size: 88 KiB |
+48
-49
@@ -24,18 +24,17 @@ import (
|
|||||||
|
|
||||||
// Agent AI代理
|
// Agent AI代理
|
||||||
type Agent struct {
|
type Agent struct {
|
||||||
openAIClient *openai.Client
|
openAIClient *openai.Client
|
||||||
config *config.OpenAIConfig
|
config *config.OpenAIConfig
|
||||||
agentConfig *config.AgentConfig
|
agentConfig *config.AgentConfig
|
||||||
mcpServer *mcp.Server
|
mcpServer *mcp.Server
|
||||||
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
|
externalMCPMgr *mcp.ExternalMCPManager // 外部MCP管理器
|
||||||
logger *zap.Logger
|
logger *zap.Logger
|
||||||
maxIterations int
|
maxIterations int
|
||||||
mu sync.RWMutex // 添加互斥锁以支持并发更新
|
mu sync.RWMutex // 添加互斥锁以支持并发更新
|
||||||
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
toolNameMapping map[string]string // 工具名称映射:OpenAI格式 -> 原始格式(用于外部MCP工具)
|
||||||
currentConversationID string // 当前对话ID(用于自动传递给工具)
|
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
||||||
promptBaseDir string // 解析 system_prompt_path 时相对路径的基准目录(通常为 config.yaml 所在目录)
|
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
||||||
toolDescriptionMode string // 工具描述模式: "short" | "full",默认 short
|
|
||||||
}
|
}
|
||||||
|
|
||||||
type agentConversationIDKey struct{}
|
type agentConversationIDKey struct{}
|
||||||
@@ -515,6 +514,14 @@ type ToolExecutionResult struct {
|
|||||||
IsError bool
|
IsError bool
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func buildToolFailureMessage(toolName, detail string, err error) string {
|
||||||
|
var b strings.Builder
|
||||||
|
fmt.Fprintf(&b, "工具调用失败\n\n")
|
||||||
|
fmt.Fprintf(&b, "工具名称: %s\n", toolName)
|
||||||
|
fmt.Fprintf(&b, "错误详情: %s", detail)
|
||||||
|
return strings.TrimRight(b.String(), "\n")
|
||||||
|
}
|
||||||
|
|
||||||
// executeToolViaMCP 通过MCP执行工具
|
// executeToolViaMCP 通过MCP执行工具
|
||||||
// 即使工具执行失败,也返回结果而不是错误,让AI能够处理错误情况
|
// 即使工具执行失败,也返回结果而不是错误,让AI能够处理错误情况
|
||||||
func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
|
func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
|
||||||
@@ -526,12 +533,6 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
|
|||||||
// 如果是record_vulnerability工具,自动添加conversation_id
|
// 如果是record_vulnerability工具,自动添加conversation_id
|
||||||
if toolName == builtin.ToolRecordVulnerability {
|
if toolName == builtin.ToolRecordVulnerability {
|
||||||
conversationID := agentConversationIDFromContext(ctx)
|
conversationID := agentConversationIDFromContext(ctx)
|
||||||
if conversationID == "" {
|
|
||||||
a.mu.RLock()
|
|
||||||
conversationID = a.currentConversationID
|
|
||||||
a.mu.RUnlock()
|
|
||||||
}
|
|
||||||
|
|
||||||
if conversationID != "" {
|
if conversationID != "" {
|
||||||
args["conversation_id"] = conversationID
|
args["conversation_id"] = conversationID
|
||||||
a.logger.Debug("自动添加conversation_id到record_vulnerability工具",
|
a.logger.Debug("自动添加conversation_id到record_vulnerability工具",
|
||||||
@@ -580,32 +581,16 @@ func (a *Agent) executeToolViaMCP(ctx context.Context, toolName string, args map
|
|||||||
// 如果调用失败(如工具不存在、超时),返回友好的错误信息而不是抛出异常
|
// 如果调用失败(如工具不存在、超时),返回友好的错误信息而不是抛出异常
|
||||||
if err != nil {
|
if err != nil {
|
||||||
detail := err.Error()
|
detail := err.Error()
|
||||||
|
timeoutMinutes := 10
|
||||||
|
if a.agentConfig != nil && a.agentConfig.ToolTimeoutMinutes > 0 {
|
||||||
|
timeoutMinutes = a.agentConfig.ToolTimeoutMinutes
|
||||||
|
}
|
||||||
if errors.Is(err, context.Canceled) {
|
if errors.Is(err, context.Canceled) {
|
||||||
detail = "工具调用已被手动终止(MCP 监控页)。智能体将携带此结果继续后续步骤,整条任务不会因此被停止。"
|
detail = "工具调用已被手动终止(MCP 监控页)。智能体将携带此结果继续后续步骤,整条任务不会因此被停止。"
|
||||||
} else if errors.Is(err, context.DeadlineExceeded) {
|
} else if errors.Is(err, context.DeadlineExceeded) {
|
||||||
min := 10
|
detail = fmt.Sprintf("工具执行超过 %d 分钟被自动终止(可在 config.yaml 的 agent.tool_timeout_minutes 中调整)", timeoutMinutes)
|
||||||
if a.agentConfig != nil && a.agentConfig.ToolTimeoutMinutes > 0 {
|
|
||||||
min = a.agentConfig.ToolTimeoutMinutes
|
|
||||||
}
|
|
||||||
detail = fmt.Sprintf("工具执行超过 %d 分钟被自动终止(可在 config.yaml 的 agent.tool_timeout_minutes 中调整)", min)
|
|
||||||
}
|
}
|
||||||
errorMsg := fmt.Sprintf(`工具调用失败
|
errorMsg := buildToolFailureMessage(toolName, detail, err)
|
||||||
|
|
||||||
工具名称: %s
|
|
||||||
错误类型: 系统错误
|
|
||||||
错误详情: %s
|
|
||||||
|
|
||||||
可能的原因:
|
|
||||||
- 工具 "%s" 不存在或未启用
|
|
||||||
- 单次执行超时(agent.tool_timeout_minutes)
|
|
||||||
- 系统配置问题
|
|
||||||
- 网络或权限问题
|
|
||||||
|
|
||||||
建议:
|
|
||||||
- 检查工具名称是否正确
|
|
||||||
- 若需更长执行时间,可适当增大 agent.tool_timeout_minutes
|
|
||||||
- 尝试使用其他替代工具
|
|
||||||
- 如果这是必需的工具,请向用户说明情况`, toolName, detail, toolName)
|
|
||||||
|
|
||||||
return &ToolExecutionResult{
|
return &ToolExecutionResult{
|
||||||
Result: errorMsg,
|
Result: errorMsg,
|
||||||
@@ -769,16 +754,8 @@ func (a *Agent) ToolsForRole(roleTools []string) []Tool {
|
|||||||
|
|
||||||
// ExecuteMCPToolForConversation 在指定会话上下文中执行 MCP 工具(行为与主 Agent 循环中的工具调用一致,如自动注入 conversation_id)。
|
// ExecuteMCPToolForConversation 在指定会话上下文中执行 MCP 工具(行为与主 Agent 循环中的工具调用一致,如自动注入 conversation_id)。
|
||||||
func (a *Agent) ExecuteMCPToolForConversation(ctx context.Context, conversationID, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
|
func (a *Agent) ExecuteMCPToolForConversation(ctx context.Context, conversationID, toolName string, args map[string]interface{}) (*ToolExecutionResult, error) {
|
||||||
a.mu.Lock()
|
|
||||||
prev := a.currentConversationID
|
|
||||||
a.currentConversationID = conversationID
|
|
||||||
a.mu.Unlock()
|
|
||||||
defer func() {
|
|
||||||
a.mu.Lock()
|
|
||||||
a.currentConversationID = prev
|
|
||||||
a.mu.Unlock()
|
|
||||||
}()
|
|
||||||
ctx = withAgentConversationID(ctx, conversationID)
|
ctx = withAgentConversationID(ctx, conversationID)
|
||||||
|
ctx = mcp.WithMCPConversationID(ctx, conversationID)
|
||||||
return a.executeToolViaMCP(ctx, toolName, args)
|
return a.executeToolViaMCP(ctx, toolName, args)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -798,6 +775,28 @@ func (a *Agent) FinishLocalToolExecution(ctx context.Context, executionID, toolN
|
|||||||
return a.mcpServer.FinishToolExecution(ctx, executionID, toolName, args, resultText, invokeErr)
|
return a.mcpServer.FinishToolExecution(ctx, executionID, toolName, args, resultText, invokeErr)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// AppendLocalToolExecutionPartialOutput records a bounded live-output preview for a running local tool.
|
||||||
|
func (a *Agent) AppendLocalToolExecutionPartialOutput(executionID, chunk string) {
|
||||||
|
if a == nil || a.mcpServer == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
a.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) RegisterLocalToolExecutionCancel(executionID string, cancel context.CancelFunc) {
|
||||||
|
if a == nil || a.mcpServer == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
a.mcpServer.RegisterToolExecutionCancel(executionID, cancel)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (a *Agent) UnregisterLocalToolExecutionCancel(executionID string) {
|
||||||
|
if a == nil || a.mcpServer == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
a.mcpServer.UnregisterToolExecutionCancel(executionID)
|
||||||
|
}
|
||||||
|
|
||||||
// RecordLocalToolExecution 将非 CallTool 路径完成的工具调用写入 MCP 监控库(与 CallTool 落库一致),返回 executionId。
|
// RecordLocalToolExecution 将非 CallTool 路径完成的工具调用写入 MCP 监控库(与 CallTool 落库一致),返回 executionId。
|
||||||
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
|
// 用于 Eino filesystem execute 等场景,使助手气泡「渗透测试详情」与常规 MCP 一致可点进监控。
|
||||||
func (a *Agent) RecordLocalToolExecution(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
func (a *Agent) RecordLocalToolExecution(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||||
|
|||||||
@@ -2,12 +2,15 @@ package agent
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"errors"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"cyberstrike-ai/internal/config"
|
"cyberstrike-ai/internal/config"
|
||||||
"cyberstrike-ai/internal/mcp"
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
"cyberstrike-ai/internal/mcp/builtin"
|
||||||
|
|
||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
)
|
)
|
||||||
@@ -69,6 +72,80 @@ func TestAgent_NewAgent_CustomConfig(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestBuildToolFailureMessageAuthorizationDenied(t *testing.T) {
|
||||||
|
msg := buildToolFailureMessage(
|
||||||
|
"list_project_facts",
|
||||||
|
"tool authorization denied: no access to project",
|
||||||
|
errors.New("tool authorization denied: no access to project"),
|
||||||
|
)
|
||||||
|
for _, want := range []string{
|
||||||
|
"工具名称: list_project_facts",
|
||||||
|
"错误详情: tool authorization denied: no access to project",
|
||||||
|
} {
|
||||||
|
if !strings.Contains(msg, want) {
|
||||||
|
t.Fatalf("message missing %q:\n%s", want, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, notWant := range []string{
|
||||||
|
"可能的原因",
|
||||||
|
"建议",
|
||||||
|
"错误类型",
|
||||||
|
"工具 \"list_project_facts\" 不存在或未启用",
|
||||||
|
"单次执行超时",
|
||||||
|
} {
|
||||||
|
if strings.Contains(msg, notWant) {
|
||||||
|
t.Fatalf("message should not include generic hint %q:\n%s", notWant, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildToolFailureMessageCanceled(t *testing.T) {
|
||||||
|
msg := buildToolFailureMessage(
|
||||||
|
"long_running_tool",
|
||||||
|
"工具调用已被手动终止(MCP 监控页)。智能体将携带此结果继续后续步骤,整条任务不会因此被停止。",
|
||||||
|
context.Canceled,
|
||||||
|
)
|
||||||
|
|
||||||
|
for _, want := range []string{
|
||||||
|
"工具名称: long_running_tool",
|
||||||
|
"错误详情: 工具调用已被手动终止",
|
||||||
|
} {
|
||||||
|
if !strings.Contains(msg, want) {
|
||||||
|
t.Fatalf("message missing %q:\n%s", want, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildToolFailureMessageDeadlineExceeded(t *testing.T) {
|
||||||
|
msg := buildToolFailureMessage(
|
||||||
|
"nmap",
|
||||||
|
"工具执行超过 15 分钟被自动终止(可在 config.yaml 的 agent.tool_timeout_minutes 中调整)",
|
||||||
|
context.DeadlineExceeded,
|
||||||
|
)
|
||||||
|
|
||||||
|
for _, want := range []string{
|
||||||
|
"工具名称: nmap",
|
||||||
|
"错误详情: 工具执行超过 15 分钟被自动终止",
|
||||||
|
} {
|
||||||
|
if !strings.Contains(msg, want) {
|
||||||
|
t.Fatalf("message missing %q:\n%s", want, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBuildToolFailureMessageUnknownKeepsGenericFallback(t *testing.T) {
|
||||||
|
msg := buildToolFailureMessage("custom_tool", "dial tcp: connection reset by peer", errors.New("dial tcp: connection reset by peer"))
|
||||||
|
|
||||||
|
for _, want := range []string{
|
||||||
|
"工具名称: custom_tool",
|
||||||
|
"错误详情: dial tcp: connection reset by peer",
|
||||||
|
} {
|
||||||
|
if !strings.Contains(msg, want) {
|
||||||
|
t.Fatalf("message missing %q:\n%s", want, msg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestAgentCancelRunningMCPToolsForConversation(t *testing.T) {
|
func TestAgentCancelRunningMCPToolsForConversation(t *testing.T) {
|
||||||
ag := setupTestAgent(t)
|
ag := setupTestAgent(t)
|
||||||
ag.mcpServer.ConfigureToolWaitTimeoutSeconds(1)
|
ag.mcpServer.ConfigureToolWaitTimeoutSeconds(1)
|
||||||
@@ -116,3 +193,93 @@ func TestAgentCancelRunningMCPToolsForConversation(t *testing.T) {
|
|||||||
}
|
}
|
||||||
t.Fatal("conv-1 execution did not become cancelled")
|
t.Fatal("conv-1 execution did not become cancelled")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestExecuteMCPToolForConversationInjectsConversationID(t *testing.T) {
|
||||||
|
ag := setupTestAgent(t)
|
||||||
|
gotArgs := make(chan map[string]interface{}, 1)
|
||||||
|
ag.mcpServer.RegisterTool(mcp.Tool{Name: builtin.ToolRecordVulnerability, InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||||
|
gotArgs <- args
|
||||||
|
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "ok"}}}, nil
|
||||||
|
})
|
||||||
|
|
||||||
|
result, err := ag.ExecuteMCPToolForConversation(context.Background(), "conv-record", builtin.ToolRecordVulnerability, map[string]interface{}{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteMCPToolForConversation: %v", err)
|
||||||
|
}
|
||||||
|
if result == nil || result.IsError {
|
||||||
|
t.Fatalf("expected successful result, got %#v", result)
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case args := <-gotArgs:
|
||||||
|
if got := args["conversation_id"]; got != "conv-record" {
|
||||||
|
t.Fatalf("conversation_id = %#v, want conv-record", got)
|
||||||
|
}
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("tool was not called")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExecuteMCPToolForConversationBindsExecutionConversation(t *testing.T) {
|
||||||
|
ag := setupTestAgent(t)
|
||||||
|
ag.mcpServer.ConfigureToolWaitTimeoutSeconds(1)
|
||||||
|
release := make(chan struct{})
|
||||||
|
ag.mcpServer.RegisterTool(mcp.Tool{Name: "slow-bind", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||||
|
select {
|
||||||
|
case <-release:
|
||||||
|
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "done"}}}, nil
|
||||||
|
case <-ctx.Done():
|
||||||
|
return nil, ctx.Err()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
result, err := ag.ExecuteMCPToolForConversation(context.Background(), "conv-bound", "slow-bind", nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteMCPToolForConversation: %v", err)
|
||||||
|
}
|
||||||
|
if result == nil || !result.IsError || result.ExecutionID == "" {
|
||||||
|
t.Fatalf("expected bounded wait result with execution id, result=%#v", result)
|
||||||
|
}
|
||||||
|
|
||||||
|
exec, ok := ag.mcpServer.GetExecution(result.ExecutionID)
|
||||||
|
if !ok || exec == nil {
|
||||||
|
t.Fatalf("missing execution %q", result.ExecutionID)
|
||||||
|
}
|
||||||
|
if exec.ConversationID != "conv-bound" {
|
||||||
|
t.Fatalf("execution conversation = %q, want conv-bound", exec.ConversationID)
|
||||||
|
}
|
||||||
|
close(release)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExecuteMCPToolForConversationConcurrentRecordIsolation(t *testing.T) {
|
||||||
|
ag := setupTestAgent(t)
|
||||||
|
seen := make(chan string, 2)
|
||||||
|
ag.mcpServer.RegisterTool(mcp.Tool{Name: builtin.ToolRecordVulnerability, InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||||
|
if conv, _ := args["conversation_id"].(string); conv != "" {
|
||||||
|
seen <- conv
|
||||||
|
}
|
||||||
|
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "ok"}}}, nil
|
||||||
|
})
|
||||||
|
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for _, conv := range []string{"conv-a", "conv-b"} {
|
||||||
|
conv := conv
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
if _, err := ag.ExecuteMCPToolForConversation(context.Background(), conv, builtin.ToolRecordVulnerability, map[string]interface{}{}); err != nil {
|
||||||
|
t.Errorf("ExecuteMCPToolForConversation %s: %v", conv, err)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
close(seen)
|
||||||
|
|
||||||
|
got := map[string]int{}
|
||||||
|
for conv := range seen {
|
||||||
|
got[conv]++
|
||||||
|
}
|
||||||
|
if got["conv-a"] != 1 || got["conv-b"] != 1 {
|
||||||
|
t.Fatalf("conversation ids = %#v, want one call for conv-a and conv-b", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,266 @@
|
|||||||
|
package agentfinalizer
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
"cyberstrike-ai/internal/multiagent"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
StatusCompleted = "completed"
|
||||||
|
StatusInProgress = "in_progress"
|
||||||
|
StatusBlocked = "blocked"
|
||||||
|
StatusFailed = "failed"
|
||||||
|
StatusCancelled = "cancelled"
|
||||||
|
StatusAwaitingHITL = "awaiting_hitl"
|
||||||
|
|
||||||
|
ReasonVerified = "verified"
|
||||||
|
ReasonPendingTools = "pending_tool_executions"
|
||||||
|
ReasonEmptyResponse = "empty_response"
|
||||||
|
ReasonAwaitingHITL = "awaiting_hitl"
|
||||||
|
ReasonFailed = "failed"
|
||||||
|
ReasonCancelled = "cancelled"
|
||||||
|
ReasonMissingEvidence = "missing_execution_evidence"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Decision is the single contract that may promote an agent run to a final
|
||||||
|
// user-facing answer. Natural-language assistant text is only a candidate until
|
||||||
|
// this object says Finalizable.
|
||||||
|
type Decision struct {
|
||||||
|
Status string `json:"status"`
|
||||||
|
Finalizable bool `json:"finalizable"`
|
||||||
|
Finalized bool `json:"finalized"`
|
||||||
|
CompletionReason string `json:"completionReason"`
|
||||||
|
FinalText string `json:"finalText,omitempty"`
|
||||||
|
EvidenceVerified bool `json:"evidenceVerified"`
|
||||||
|
EvidenceRefs []string `json:"evidenceRefs,omitempty"`
|
||||||
|
PendingExecutionIDs []string `json:"pendingExecutionIds,omitempty"`
|
||||||
|
PendingToolRuns []string `json:"pendingToolRuns,omitempty"`
|
||||||
|
MissingChecks []string `json:"missingChecks,omitempty"`
|
||||||
|
AgentMode string `json:"agentMode,omitempty"`
|
||||||
|
ConversationID string `json:"conversationId,omitempty"`
|
||||||
|
AssistantMessageID string `json:"messageId,omitempty"`
|
||||||
|
CandidateResponseLen int `json:"candidateResponseLen,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type Input struct {
|
||||||
|
Response string
|
||||||
|
MCPExecutionIDs []string
|
||||||
|
ConversationID string
|
||||||
|
AssistantMessageID string
|
||||||
|
AgentMode string
|
||||||
|
Status string
|
||||||
|
CompletionReason string
|
||||||
|
AwaitingHITL bool
|
||||||
|
RequireExecutionEvidence bool
|
||||||
|
}
|
||||||
|
|
||||||
|
func FromRunResult(db *database.DB, result *multiagent.RunResult, in Input) Decision {
|
||||||
|
if result != nil {
|
||||||
|
if strings.TrimSpace(in.Response) == "" {
|
||||||
|
in.Response = result.Response
|
||||||
|
}
|
||||||
|
if len(in.MCPExecutionIDs) == 0 {
|
||||||
|
in.MCPExecutionIDs = result.MCPExecutionIDs
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(in.Status) == "" {
|
||||||
|
in.Status = result.Status
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(in.CompletionReason) == "" {
|
||||||
|
in.CompletionReason = result.CompletionReason
|
||||||
|
}
|
||||||
|
}
|
||||||
|
d := Decide(db, in)
|
||||||
|
if result != nil {
|
||||||
|
result.Finalized = d.Finalized
|
||||||
|
result.Status = d.Status
|
||||||
|
result.CompletionReason = d.CompletionReason
|
||||||
|
result.EvidenceVerified = d.EvidenceVerified
|
||||||
|
result.EvidenceRefs = append([]string(nil), d.EvidenceRefs...)
|
||||||
|
result.PendingExecutionIDs = append([]string(nil), d.PendingExecutionIDs...)
|
||||||
|
result.MissingChecks = append([]string(nil), d.MissingChecks...)
|
||||||
|
}
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
|
||||||
|
func Decide(db *database.DB, in Input) Decision {
|
||||||
|
text := strings.TrimSpace(in.Response)
|
||||||
|
status := strings.TrimSpace(in.Status)
|
||||||
|
if status == "" {
|
||||||
|
status = StatusCompleted
|
||||||
|
}
|
||||||
|
reason := strings.TrimSpace(in.CompletionReason)
|
||||||
|
if reason == "" {
|
||||||
|
reason = ReasonVerified
|
||||||
|
}
|
||||||
|
d := Decision{
|
||||||
|
Status: status,
|
||||||
|
CompletionReason: reason,
|
||||||
|
FinalText: text,
|
||||||
|
EvidenceVerified: true,
|
||||||
|
EvidenceRefs: evidenceRefs(in.MCPExecutionIDs),
|
||||||
|
AgentMode: strings.TrimSpace(in.AgentMode),
|
||||||
|
ConversationID: strings.TrimSpace(in.ConversationID),
|
||||||
|
AssistantMessageID: strings.TrimSpace(in.AssistantMessageID),
|
||||||
|
CandidateResponseLen: len([]rune(text)),
|
||||||
|
}
|
||||||
|
|
||||||
|
if in.AwaitingHITL {
|
||||||
|
d.Status = StatusAwaitingHITL
|
||||||
|
d.CompletionReason = ReasonAwaitingHITL
|
||||||
|
d.EvidenceVerified = false
|
||||||
|
d.MissingChecks = append(d.MissingChecks, "workflow is awaiting HITL approval")
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
if isEmptyCandidate(text) {
|
||||||
|
d.Status = StatusBlocked
|
||||||
|
d.CompletionReason = ReasonEmptyResponse
|
||||||
|
d.EvidenceVerified = false
|
||||||
|
d.MissingChecks = append(d.MissingChecks, "assistant final text is empty or only an empty-response placeholder")
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
switch status {
|
||||||
|
case StatusInProgress, StatusBlocked, StatusFailed, StatusCancelled, StatusAwaitingHITL:
|
||||||
|
d.Status = status
|
||||||
|
d.EvidenceVerified = false
|
||||||
|
if d.CompletionReason == ReasonVerified {
|
||||||
|
d.CompletionReason = status
|
||||||
|
}
|
||||||
|
d.MissingChecks = append(d.MissingChecks, "agent run status is "+status)
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
|
||||||
|
pending := pendingExecutions(db, in.MCPExecutionIDs)
|
||||||
|
if len(pending) > 0 {
|
||||||
|
d.Status = StatusInProgress
|
||||||
|
d.CompletionReason = ReasonPendingTools
|
||||||
|
d.EvidenceVerified = false
|
||||||
|
d.PendingExecutionIDs = pending
|
||||||
|
d.PendingToolRuns = append([]string(nil), pending...)
|
||||||
|
d.MissingChecks = append(d.MissingChecks, "tool execution still queued or running")
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
|
||||||
|
if in.RequireExecutionEvidence && !hasCompletedEvidence(db, in.MCPExecutionIDs) {
|
||||||
|
d.Status = StatusBlocked
|
||||||
|
d.CompletionReason = ReasonMissingEvidence
|
||||||
|
d.EvidenceVerified = false
|
||||||
|
d.MissingChecks = append(d.MissingChecks, "execution evidence is required but no completed tool execution was recorded")
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
|
||||||
|
d.Finalizable = true
|
||||||
|
d.Finalized = true
|
||||||
|
d.Status = StatusCompleted
|
||||||
|
if d.CompletionReason == "" {
|
||||||
|
d.CompletionReason = ReasonVerified
|
||||||
|
}
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
|
||||||
|
func ResponsePayload(d Decision, extra map[string]interface{}) map[string]interface{} {
|
||||||
|
out := map[string]interface{}{
|
||||||
|
"finalized": d.Finalized,
|
||||||
|
"finalizable": d.Finalizable,
|
||||||
|
"status": d.Status,
|
||||||
|
"completionReason": d.CompletionReason,
|
||||||
|
"evidenceVerified": d.EvidenceVerified,
|
||||||
|
"evidenceRefs": d.EvidenceRefs,
|
||||||
|
"pendingExecutionIds": d.PendingExecutionIDs,
|
||||||
|
"pendingToolRuns": d.PendingToolRuns,
|
||||||
|
"missingChecks": d.MissingChecks,
|
||||||
|
}
|
||||||
|
if d.ConversationID != "" {
|
||||||
|
out["conversationId"] = d.ConversationID
|
||||||
|
}
|
||||||
|
if d.AssistantMessageID != "" {
|
||||||
|
out["messageId"] = d.AssistantMessageID
|
||||||
|
}
|
||||||
|
if d.AgentMode != "" {
|
||||||
|
out["agentMode"] = d.AgentMode
|
||||||
|
}
|
||||||
|
for k, v := range extra {
|
||||||
|
out[k] = v
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func isEmptyCandidate(s string) bool {
|
||||||
|
s = strings.TrimSpace(s)
|
||||||
|
if s == "" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return strings.Contains(s, "no assistant text was captured") ||
|
||||||
|
strings.Contains(s, "未捕获到助手文本输出")
|
||||||
|
}
|
||||||
|
|
||||||
|
func evidenceRefs(ids []string) []string {
|
||||||
|
out := make([]string, 0, len(ids))
|
||||||
|
seen := make(map[string]struct{}, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if id == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[id]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[id] = struct{}{}
|
||||||
|
out = append(out, "mcp_execution:"+id)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func pendingExecutions(db *database.DB, ids []string) []string {
|
||||||
|
if db == nil || len(ids) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
out := make([]string, 0)
|
||||||
|
seen := make(map[string]struct{}, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if id == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[id]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[id] = struct{}{}
|
||||||
|
exec, err := db.GetToolExecution(id)
|
||||||
|
if err != nil || exec == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
switch strings.TrimSpace(exec.Status) {
|
||||||
|
case mcp.ToolExecutionStatusQueued, mcp.ToolExecutionStatusRunning:
|
||||||
|
out = append(out, id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func hasCompletedEvidence(db *database.DB, ids []string) bool {
|
||||||
|
if db == nil || len(ids) == 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
seen := make(map[string]struct{}, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if id == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[id]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[id] = struct{}{}
|
||||||
|
exec, err := db.GetToolExecution(id)
|
||||||
|
if err != nil || exec == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(exec.Status) == mcp.ToolExecutionStatusCompleted {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
@@ -0,0 +1,132 @@
|
|||||||
|
package agentfinalizer
|
||||||
|
|
||||||
|
import (
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func newDecisionTestDB(t *testing.T) *database.DB {
|
||||||
|
t.Helper()
|
||||||
|
db, err := database.NewDB(filepath.Join(t.TempDir(), "finalizer.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() { _ = db.Close() })
|
||||||
|
return db
|
||||||
|
}
|
||||||
|
|
||||||
|
func saveDecisionTestExecution(t *testing.T, db *database.DB, id, status string) {
|
||||||
|
t.Helper()
|
||||||
|
if err := db.SaveToolExecution(&mcp.ToolExecution{
|
||||||
|
ID: id,
|
||||||
|
ToolName: "test::tool",
|
||||||
|
Arguments: map[string]interface{}{"input": id},
|
||||||
|
Status: status,
|
||||||
|
StartTime: time.Now(),
|
||||||
|
}); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution(%s): %v", id, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecideBlocksPendingToolExecutions(t *testing.T) {
|
||||||
|
db := newDecisionTestDB(t)
|
||||||
|
saveDecisionTestExecution(t, db, "run-queued", mcp.ToolExecutionStatusQueued)
|
||||||
|
saveDecisionTestExecution(t, db, "run-running", mcp.ToolExecutionStatusRunning)
|
||||||
|
saveDecisionTestExecution(t, db, "run-completed", mcp.ToolExecutionStatusCompleted)
|
||||||
|
|
||||||
|
d := Decide(db, Input{
|
||||||
|
Response: "工具还没全部结束时,这只是一段候选输出。",
|
||||||
|
MCPExecutionIDs: []string{"run-queued", "run-running", "run-completed"},
|
||||||
|
})
|
||||||
|
|
||||||
|
if d.Finalizable || d.Finalized {
|
||||||
|
t.Fatalf("pending tools should not be finalizable: %+v", d)
|
||||||
|
}
|
||||||
|
if d.Status != StatusInProgress || d.CompletionReason != ReasonPendingTools {
|
||||||
|
t.Fatalf("status/reason = %s/%s, want %s/%s", d.Status, d.CompletionReason, StatusInProgress, ReasonPendingTools)
|
||||||
|
}
|
||||||
|
if got, want := len(d.PendingExecutionIDs), 2; got != want {
|
||||||
|
t.Fatalf("pending execution count = %d, want %d (%v)", got, want, d.PendingExecutionIDs)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecideBlocksAwaitingHITLAndEmptyCandidate(t *testing.T) {
|
||||||
|
hitl := Decide(nil, Input{Response: "等待人工审批", AwaitingHITL: true})
|
||||||
|
if hitl.Finalizable || hitl.Status != StatusAwaitingHITL || hitl.CompletionReason != ReasonAwaitingHITL {
|
||||||
|
t.Fatalf("HITL decision mismatch: %+v", hitl)
|
||||||
|
}
|
||||||
|
|
||||||
|
empty := Decide(nil, Input{Response: "⚠️ Eino 执行完成,但未捕获到助手文本输出。"})
|
||||||
|
if empty.Finalizable || empty.Status != StatusBlocked || empty.CompletionReason != ReasonEmptyResponse {
|
||||||
|
t.Fatalf("empty candidate decision mismatch: %+v", empty)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecideBlocksWhenExecutionEvidenceIsRequiredButMissing(t *testing.T) {
|
||||||
|
d := Decide(nil, Input{
|
||||||
|
Response: "任务已处理完成。",
|
||||||
|
RequireExecutionEvidence: true,
|
||||||
|
})
|
||||||
|
if d.Finalizable || d.Finalized {
|
||||||
|
t.Fatalf("missing required execution evidence should not finalize: %+v", d)
|
||||||
|
}
|
||||||
|
if d.Status != StatusBlocked || d.CompletionReason != ReasonMissingEvidence {
|
||||||
|
t.Fatalf("status/reason = %s/%s, want %s/%s", d.Status, d.CompletionReason, StatusBlocked, ReasonMissingEvidence)
|
||||||
|
}
|
||||||
|
if d.EvidenceVerified {
|
||||||
|
t.Fatalf("missing required execution evidence should be marked unverified: %+v", d)
|
||||||
|
}
|
||||||
|
if len(d.MissingChecks) == 0 {
|
||||||
|
t.Fatalf("missing checks should explain the evidence gap: %+v", d)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecideBlocksWhenOnlyFailedEvidenceIsRecorded(t *testing.T) {
|
||||||
|
db := newDecisionTestDB(t)
|
||||||
|
saveDecisionTestExecution(t, db, "run-failed", mcp.ToolExecutionStatusFailed)
|
||||||
|
saveDecisionTestExecution(t, db, "run-cancelled", mcp.ToolExecutionStatusCancelled)
|
||||||
|
|
||||||
|
d := Decide(db, Input{
|
||||||
|
Response: "任务已处理完成。",
|
||||||
|
MCPExecutionIDs: []string{"run-failed", "run-cancelled"},
|
||||||
|
RequireExecutionEvidence: true,
|
||||||
|
})
|
||||||
|
|
||||||
|
if d.Finalizable || d.Finalized {
|
||||||
|
t.Fatalf("failed evidence should not satisfy required execution evidence: %+v", d)
|
||||||
|
}
|
||||||
|
if d.Status != StatusBlocked || d.CompletionReason != ReasonMissingEvidence {
|
||||||
|
t.Fatalf("status/reason = %s/%s, want %s/%s", d.Status, d.CompletionReason, StatusBlocked, ReasonMissingEvidence)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecideFinalizesCompletedEvidence(t *testing.T) {
|
||||||
|
db := newDecisionTestDB(t)
|
||||||
|
saveDecisionTestExecution(t, db, "run-ok", mcp.ToolExecutionStatusCompleted)
|
||||||
|
|
||||||
|
d := Decide(db, Input{
|
||||||
|
Response: "任务已处理完成,见工具执行记录。",
|
||||||
|
MCPExecutionIDs: []string{"run-ok"},
|
||||||
|
RequireExecutionEvidence: true,
|
||||||
|
})
|
||||||
|
|
||||||
|
if !d.Finalizable || !d.Finalized || d.Status != StatusCompleted {
|
||||||
|
t.Fatalf("completed execution should finalize: %+v", d)
|
||||||
|
}
|
||||||
|
if !d.EvidenceVerified || len(d.EvidenceRefs) != 1 {
|
||||||
|
t.Fatalf("evidence refs mismatch: %+v", d)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDecideAllowsInformationalAnswerWhenExecutionEvidenceIsNotRequired(t *testing.T) {
|
||||||
|
d := Decide(nil, Input{Response: "这是一个概念解释,不需要执行工具。"})
|
||||||
|
if !d.Finalizable || !d.Finalized || d.Status != StatusCompleted {
|
||||||
|
t.Fatalf("informational response should finalize when execution evidence is not required: %+v", d)
|
||||||
|
}
|
||||||
|
}
|
||||||
+86
-9
@@ -1313,6 +1313,7 @@ func setupRoutes(
|
|||||||
c2Routes.GET("/sessions/:id", c2Handler.GetSession)
|
c2Routes.GET("/sessions/:id", c2Handler.GetSession)
|
||||||
c2Routes.DELETE("/sessions/:id", c2Handler.DeleteSession)
|
c2Routes.DELETE("/sessions/:id", c2Handler.DeleteSession)
|
||||||
c2Routes.PUT("/sessions/:id/sleep", c2Handler.SetSessionSleep)
|
c2Routes.PUT("/sessions/:id/sleep", c2Handler.SetSessionSleep)
|
||||||
|
c2Routes.PUT("/sessions/:id/note", c2Handler.SetSessionNote)
|
||||||
c2Routes.GET("/tasks", c2Handler.ListTasks)
|
c2Routes.GET("/tasks", c2Handler.ListTasks)
|
||||||
c2Routes.DELETE("/tasks", c2Handler.DeleteTasks)
|
c2Routes.DELETE("/tasks", c2Handler.DeleteTasks)
|
||||||
c2Routes.GET("/tasks/:id", c2Handler.GetTask)
|
c2Routes.GET("/tasks/:id", c2Handler.GetTask)
|
||||||
@@ -1337,7 +1338,9 @@ func setupRoutes(
|
|||||||
|
|
||||||
// 对话附件(chat_uploads)管理
|
// 对话附件(chat_uploads)管理
|
||||||
protected.GET("/chat-uploads", chatUploadsHandler.List)
|
protected.GET("/chat-uploads", chatUploadsHandler.List)
|
||||||
|
protected.GET("/chat-uploads/export", chatUploadsHandler.Export)
|
||||||
protected.GET("/chat-uploads/download", chatUploadsHandler.Download)
|
protected.GET("/chat-uploads/download", chatUploadsHandler.Download)
|
||||||
|
protected.GET("/chat-uploads/path", chatUploadsHandler.ResolvePath)
|
||||||
protected.GET("/chat-uploads/content", chatUploadsHandler.GetContent)
|
protected.GET("/chat-uploads/content", chatUploadsHandler.GetContent)
|
||||||
protected.POST("/chat-uploads", chatUploadsHandler.Upload)
|
protected.POST("/chat-uploads", chatUploadsHandler.Upload)
|
||||||
protected.POST("/chat-uploads/mkdir", chatUploadsHandler.Mkdir)
|
protected.POST("/chat-uploads/mkdir", chatUploadsHandler.Mkdir)
|
||||||
@@ -1359,6 +1362,7 @@ func setupRoutes(
|
|||||||
protected.POST("/workflows/runs/:runId/resume", workflowHandler.ResumeRun)
|
protected.POST("/workflows/runs/:runId/resume", workflowHandler.ResumeRun)
|
||||||
protected.POST("/workflows/validate", workflowHandler.Validate)
|
protected.POST("/workflows/validate", workflowHandler.Validate)
|
||||||
protected.POST("/workflows/dry-run", workflowHandler.DryRun)
|
protected.POST("/workflows/dry-run", workflowHandler.DryRun)
|
||||||
|
protected.POST("/workflows/generate-draft", workflowHandler.GenerateDraft)
|
||||||
protected.GET("/workflows/:id/package", workflowHandler.ExportPackage)
|
protected.GET("/workflows/:id/package", workflowHandler.ExportPackage)
|
||||||
protected.POST("/workflow-package-inspections", workflowHandler.CreatePackageInspection)
|
protected.POST("/workflow-package-inspections", workflowHandler.CreatePackageInspection)
|
||||||
protected.GET("/workflow-package-inspections/:inspectionId", workflowHandler.GetPackageInspection)
|
protected.GET("/workflow-package-inspections/:inspectionId", workflowHandler.GetPackageInspection)
|
||||||
@@ -1573,22 +1577,62 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
logger.Warn("跳过 WebShell 管理工具注册:db 为空")
|
logger.Warn("跳过 WebShell 管理工具注册:db 为空")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
projectIDFromToolArgs := func(ctx context.Context, args map[string]interface{}) string {
|
||||||
|
projectID, _ := args["project_id"].(string)
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == "" {
|
||||||
|
projectID = strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx))
|
||||||
|
}
|
||||||
|
return projectID
|
||||||
|
}
|
||||||
|
explicitProjectIDFromToolArgs := func(args map[string]interface{}) string {
|
||||||
|
projectID, _ := args["project_id"].(string)
|
||||||
|
return strings.TrimSpace(projectID)
|
||||||
|
}
|
||||||
|
authorizeWebshellToolProject := func(principal authctx.Principal, permission, projectID string) *mcp.ToolResult {
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if projectID == database.ProjectFilterUnbound {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) {
|
||||||
|
return &mcp.ToolResult{
|
||||||
|
Content: []mcp.Content{{Type: "text", Text: "无权访问项目: " + projectID}},
|
||||||
|
IsError: true,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
// manage_webshell_list - 列出所有 webshell 连接
|
// manage_webshell_list - 列出所有 webshell 连接
|
||||||
listTool := mcp.Tool{
|
listTool := mcp.Tool{
|
||||||
Name: builtin.ToolManageWebshellList,
|
Name: builtin.ToolManageWebshellList,
|
||||||
Description: "列出所有已保存的 WebShell 连接,返回连接ID、URL、类型、备注等信息。",
|
Description: "列出已保存的 WebShell 连接,返回连接ID、URL、类型、所属项目、备注等信息。默认按当前对话项目边界过滤:项目对话看本项目,未绑定项目的对话看未绑定连接;显式传 project_id 时按指定项目过滤。",
|
||||||
ShortDescription: "列出所有 WebShell 连接",
|
ShortDescription: "列出所有 WebShell 连接",
|
||||||
InputSchema: map[string]interface{}{
|
InputSchema: map[string]interface{}{
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": map[string]interface{}{},
|
"properties": map[string]interface{}{
|
||||||
|
"project_id": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "项目 ID;不填时在项目会话中默认使用当前项目。",
|
||||||
|
},
|
||||||
|
},
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
listHandler := func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||||
connections := []database.WebShellConnection{}
|
connections := []database.WebShellConnection{}
|
||||||
var err error
|
var err error
|
||||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||||
connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read"))
|
projectID := explicitProjectIDFromToolArgs(args)
|
||||||
|
if projectID == "" {
|
||||||
|
projectID = mcpEffectiveProjectFilter(ctx, db)
|
||||||
|
}
|
||||||
|
if result := authorizeWebshellToolProject(principal, "webshell:read", projectID); result != nil {
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
connections, err = db.ListWebshellConnectionsForAccess(principal.UserID, principal.ScopeFor("webshell:read"), projectID)
|
||||||
} else {
|
} else {
|
||||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
||||||
}
|
}
|
||||||
@@ -1612,6 +1656,11 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
sb.WriteString(fmt.Sprintf(" 类型: %s\n", conn.Type))
|
sb.WriteString(fmt.Sprintf(" 类型: %s\n", conn.Type))
|
||||||
sb.WriteString(fmt.Sprintf(" 请求方式: %s\n", conn.Method))
|
sb.WriteString(fmt.Sprintf(" 请求方式: %s\n", conn.Method))
|
||||||
sb.WriteString(fmt.Sprintf(" 命令参数: %s\n", conn.CmdParam))
|
sb.WriteString(fmt.Sprintf(" 命令参数: %s\n", conn.CmdParam))
|
||||||
|
if conn.ProjectID != "" {
|
||||||
|
sb.WriteString(fmt.Sprintf(" 项目ID: %s\n", conn.ProjectID))
|
||||||
|
} else {
|
||||||
|
sb.WriteString(" 项目: 未绑定\n")
|
||||||
|
}
|
||||||
if conn.Remark != "" {
|
if conn.Remark != "" {
|
||||||
sb.WriteString(fmt.Sprintf(" 备注: %s\n", conn.Remark))
|
sb.WriteString(fmt.Sprintf(" 备注: %s\n", conn.Remark))
|
||||||
}
|
}
|
||||||
@@ -1686,6 +1735,14 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
cmdParam = "cmd"
|
cmdParam = "cmd"
|
||||||
}
|
}
|
||||||
remark, _ := args["remark"].(string)
|
remark, _ := args["remark"].(string)
|
||||||
|
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||||
|
if !ok {
|
||||||
|
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
||||||
|
}
|
||||||
|
projectID := projectIDFromToolArgs(ctx, args)
|
||||||
|
if result := authorizeWebshellToolProject(principal, "webshell:write", projectID); result != nil {
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
// 生成连接ID
|
// 生成连接ID
|
||||||
connID := "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12]
|
connID := "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12]
|
||||||
@@ -1697,6 +1754,7 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
Method: strings.ToLower(method),
|
Method: strings.ToLower(method),
|
||||||
CmdParam: cmdParam,
|
CmdParam: cmdParam,
|
||||||
Remark: remark,
|
Remark: remark,
|
||||||
|
ProjectID: projectID,
|
||||||
CreatedAt: time.Now(),
|
CreatedAt: time.Now(),
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1706,15 +1764,17 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
IsError: true,
|
IsError: true,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
|
||||||
_ = db.SetResourceOwner("webshell", conn.ID, principal.UserID)
|
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
|
||||||
_ = db.AssignResourceToUser(principal.UserID, "webshell", conn.ID)
|
projectLine := "项目: 未绑定"
|
||||||
|
if conn.ProjectID != "" {
|
||||||
|
projectLine = "项目ID: " + conn.ProjectID
|
||||||
}
|
}
|
||||||
|
|
||||||
return &mcp.ToolResult{
|
return &mcp.ToolResult{
|
||||||
Content: []mcp.Content{{
|
Content: []mcp.Content{{
|
||||||
Type: "text",
|
Type: "text",
|
||||||
Text: fmt.Sprintf("WebShell 连接添加成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s", conn.ID, conn.URL, conn.Type, conn.Method, conn.CmdParam),
|
Text: fmt.Sprintf("WebShell 连接添加成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n%s", conn.ID, conn.URL, conn.Type, conn.Method, conn.CmdParam, projectLine),
|
||||||
}},
|
}},
|
||||||
IsError: false,
|
IsError: false,
|
||||||
}, nil
|
}, nil
|
||||||
@@ -1759,6 +1819,10 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "新的备注",
|
"description": "新的备注",
|
||||||
},
|
},
|
||||||
|
"project_id": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "新的所属项目 ID;传空字符串可取消绑定。",
|
||||||
|
},
|
||||||
},
|
},
|
||||||
"required": []string{"connection_id"},
|
"required": []string{"connection_id"},
|
||||||
},
|
},
|
||||||
@@ -1800,6 +1864,19 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
if remark, ok := args["remark"].(string); ok {
|
if remark, ok := args["remark"].(string); ok {
|
||||||
existing.Remark = remark
|
existing.Remark = remark
|
||||||
}
|
}
|
||||||
|
if projectID, ok := args["project_id"].(string); ok {
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID != "" {
|
||||||
|
principal, ok := authctx.PrincipalFromContext(ctx)
|
||||||
|
if !ok {
|
||||||
|
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: "缺少认证身份"}}, IsError: true}, nil
|
||||||
|
}
|
||||||
|
if result := authorizeWebshellToolProject(principal, "webshell:write", projectID); result != nil {
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
existing.ProjectID = projectID
|
||||||
|
}
|
||||||
|
|
||||||
if err := db.UpdateWebshellConnection(existing); err != nil {
|
if err := db.UpdateWebshellConnection(existing); err != nil {
|
||||||
return &mcp.ToolResult{
|
return &mcp.ToolResult{
|
||||||
@@ -1811,7 +1888,7 @@ func registerWebshellManagementTools(mcpServer *mcp.Server, db *database.DB, web
|
|||||||
return &mcp.ToolResult{
|
return &mcp.ToolResult{
|
||||||
Content: []mcp.Content{{
|
Content: []mcp.Content{{
|
||||||
Type: "text",
|
Type: "text",
|
||||||
Text: fmt.Sprintf("WebShell 连接更新成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n备注: %s", existing.ID, existing.URL, existing.Type, existing.Method, existing.CmdParam, existing.Remark),
|
Text: fmt.Sprintf("WebShell 连接更新成功!\n\n连接ID: %s\nURL: %s\n类型: %s\n请求方式: %s\n命令参数: %s\n项目ID: %s\n备注: %s", existing.ID, existing.URL, existing.Type, existing.Method, existing.CmdParam, existing.ProjectID, existing.Remark),
|
||||||
}},
|
}},
|
||||||
IsError: false,
|
IsError: false,
|
||||||
}, nil
|
}, nil
|
||||||
|
|||||||
@@ -75,6 +75,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
|||||||
"bind_host": map[string]interface{}{"type": "string", "description": "绑定地址,默认 127.0.0.1;外网监听常用 0.0.0.0"},
|
"bind_host": map[string]interface{}{"type": "string", "description": "绑定地址,默认 127.0.0.1;外网监听常用 0.0.0.0"},
|
||||||
"callback_host": map[string]interface{}{"type": "string", "description": "可选:植入端/Payload 回连主机名(公网 IP 或域名)。写入 config_json;生成 oneliner/beacon 时优先于 bind_host。update 时传入空字符串可清除"},
|
"callback_host": map[string]interface{}{"type": "string", "description": "可选:植入端/Payload 回连主机名(公网 IP 或域名)。写入 config_json;生成 oneliner/beacon 时优先于 bind_host。update 时传入空字符串可清除"},
|
||||||
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535},
|
"bind_port": map[string]interface{}{"type": "integer", "description": fmt.Sprintf("绑定端口(create 必填)。须 ≠ %d(当前本服务 Web/API 端口,配置 server.port)", webListenPort), "minimum": 1, "maximum": 65535},
|
||||||
|
"project_id": map[string]interface{}{"type": "string", "description": "所属项目 ID。create 省略时默认使用当前对话绑定项目;未绑定项目的对话则创建未绑定监听器"},
|
||||||
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
|
"profile_id": map[string]interface{}{"type": "string", "description": "Malleable Profile ID"},
|
||||||
"remark": map[string]interface{}{"type": "string", "description": "备注"},
|
"remark": map[string]interface{}{"type": "string", "description": "备注"},
|
||||||
"config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false)"},
|
"config": map[string]interface{}{"type": "object", "description": "高级配置(beacon 路径/TLS/OPSEC 等),create/update 可用。tcp_reverse 可选 allow_legacy_shell:true 允许未加密经典 shell(默认 false)"},
|
||||||
@@ -87,7 +88,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
|||||||
|
|
||||||
switch action {
|
switch action {
|
||||||
case "list":
|
case "list":
|
||||||
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx))
|
listeners, err := m.DB().ListC2ListenersForAccess(c2ToolAccess(ctx), mcpEffectiveProjectFilter(ctx, m.DB()))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return makeC2Result(nil, err)
|
return makeC2Result(nil, err)
|
||||||
}
|
}
|
||||||
@@ -116,6 +117,13 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
|||||||
cfg = &c2.ListenerConfig{}
|
cfg = &c2.ListenerConfig{}
|
||||||
_ = json.Unmarshal(cfgBytes, cfg)
|
_ = json.Unmarshal(cfgBytes, cfg)
|
||||||
}
|
}
|
||||||
|
projectID := strings.TrimSpace(getString(params, "project_id"))
|
||||||
|
if projectID == "" {
|
||||||
|
projectID = mcpEffectiveProjectFilter(ctx, m.DB())
|
||||||
|
if projectID == database.ProjectFilterUnbound {
|
||||||
|
projectID = ""
|
||||||
|
}
|
||||||
|
}
|
||||||
input := c2.CreateListenerInput{
|
input := c2.CreateListenerInput{
|
||||||
Name: getString(params, "name"),
|
Name: getString(params, "name"),
|
||||||
Type: getString(params, "type"),
|
Type: getString(params, "type"),
|
||||||
@@ -123,6 +131,7 @@ tcp_reverse 默认仅接受 CSB1 加密 Beacon(AES-GCM + ImplantToken)才登
|
|||||||
BindPort: int(getFloat64(params, "bind_port")),
|
BindPort: int(getFloat64(params, "bind_port")),
|
||||||
ProfileID: getString(params, "profile_id"),
|
ProfileID: getString(params, "profile_id"),
|
||||||
Remark: getString(params, "remark"),
|
Remark: getString(params, "remark"),
|
||||||
|
ProjectID: projectID,
|
||||||
Config: cfg,
|
Config: cfg,
|
||||||
CallbackHost: getString(params, "callback_host"),
|
CallbackHost: getString(params, "callback_host"),
|
||||||
}
|
}
|
||||||
@@ -260,6 +269,7 @@ func registerC2SessionTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
|||||||
case "list":
|
case "list":
|
||||||
filter := database.ListC2SessionsFilter{
|
filter := database.ListC2SessionsFilter{
|
||||||
ListenerID: getString(params, "listener_id"),
|
ListenerID: getString(params, "listener_id"),
|
||||||
|
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||||
Status: getString(params, "status"),
|
Status: getString(params, "status"),
|
||||||
OS: getString(params, "os"),
|
OS: getString(params, "os"),
|
||||||
Search: getString(params, "search"),
|
Search: getString(params, "search"),
|
||||||
@@ -495,6 +505,7 @@ func registerC2TaskManageTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
|||||||
case "list":
|
case "list":
|
||||||
filter := database.ListC2TasksFilter{
|
filter := database.ListC2TasksFilter{
|
||||||
SessionID: getString(params, "session_id"),
|
SessionID: getString(params, "session_id"),
|
||||||
|
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||||
Status: getString(params, "status"),
|
Status: getString(params, "status"),
|
||||||
}
|
}
|
||||||
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
if limit := int(getFloat64(params, "limit")); limit > 0 {
|
||||||
@@ -645,6 +656,7 @@ func registerC2EventTool(s *mcp.Server, m *c2.Manager, l *zap.Logger) {
|
|||||||
filter := database.ListC2EventsFilter{
|
filter := database.ListC2EventsFilter{
|
||||||
Level: getString(params, "level"),
|
Level: getString(params, "level"),
|
||||||
Category: getString(params, "category"),
|
Category: getString(params, "category"),
|
||||||
|
ProjectID: mcpEffectiveProjectFilter(ctx, m.DB()),
|
||||||
SessionID: getString(params, "session_id"),
|
SessionID: getString(params, "session_id"),
|
||||||
TaskID: getString(params, "task_id"),
|
TaskID: getString(params, "task_id"),
|
||||||
Limit: int(getFloat64(params, "limit")),
|
Limit: int(getFloat64(params, "limit")),
|
||||||
|
|||||||
@@ -0,0 +1,69 @@
|
|||||||
|
package app
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/authctx"
|
||||||
|
"cyberstrike-ai/internal/c2"
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
"cyberstrike-ai/internal/mcp/builtin"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestC2ListenerCreateInheritsConversationProject(t *testing.T) {
|
||||||
|
db, err := database.NewDB(filepath.Join(t.TempDir(), "c2-tools.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
user, err := db.CreateRBACUser("c2-agent", "C2 Agent", "hash", true, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
project, err := db.CreateProject(&database.Project{Name: "engagement"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AssignResourceToUser(user.ID, "project", project.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
conversation, err := db.CreateConversation("project chat", database.ConversationCreateMeta{ProjectID: project.ID})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{
|
||||||
|
"c2:read": true, "c2:write": true,
|
||||||
|
})
|
||||||
|
ctx := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), conversation.ID), principal)
|
||||||
|
server := mcp.NewServer(zap.NewNop())
|
||||||
|
server.SetToolAuthorizer(mcpToolAuthorizer(db))
|
||||||
|
registerC2Tools(server, c2.NewManager(db, zap.NewNop(), t.TempDir()), zap.NewNop(), 8080)
|
||||||
|
|
||||||
|
result, _, err := server.CallTool(ctx, builtin.ToolC2Listener, map[string]interface{}{
|
||||||
|
"action": "create",
|
||||||
|
"name": "tcp-reverse-2222",
|
||||||
|
"type": "tcp_reverse",
|
||||||
|
"bind_host": "0.0.0.0",
|
||||||
|
"bind_port": 2222,
|
||||||
|
})
|
||||||
|
if err != nil || result == nil || result.IsError {
|
||||||
|
t.Fatalf("create listener result=%#v err=%v text=%q", result, err, toolResultText(result))
|
||||||
|
}
|
||||||
|
|
||||||
|
listeners, err := db.ListC2Listeners()
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(listeners) != 1 {
|
||||||
|
t.Fatalf("listener count=%d, want 1", len(listeners))
|
||||||
|
}
|
||||||
|
if listeners[0].ProjectID != project.ID {
|
||||||
|
t.Fatalf("listener project_id=%q, want %q", listeners[0].ProjectID, project.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -32,6 +32,9 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
|
|||||||
if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
if id == "" || db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||||
}
|
}
|
||||||
|
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, id); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
toolExecutionResource := func(permission string) error {
|
toolExecutionResource := func(permission string) error {
|
||||||
@@ -141,20 +144,26 @@ func mcpToolAuthorizer(db *database.DB) func(context.Context, string, map[string
|
|||||||
builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate:
|
builtin.ToolBatchTaskScheduleEnabled, builtin.ToolBatchTaskAdd, builtin.ToolBatchTaskUpdate:
|
||||||
return resource("tasks:write", "batch_task", "queue_id")
|
return resource("tasks:write", "batch_task", "queue_id")
|
||||||
case builtin.ToolC2Listener:
|
case builtin.ToolC2Listener:
|
||||||
return authorizeC2Action(principal, db, args, "c2_listener", "listener_id")
|
return authorizeC2Action(ctx, principal, db, args, "c2_listener", "listener_id")
|
||||||
case builtin.ToolC2Session, builtin.ToolC2Task, builtin.ToolC2File:
|
case builtin.ToolC2Session, builtin.ToolC2Task, builtin.ToolC2File:
|
||||||
if toolName == builtin.ToolC2File && mcpAuthorizationString(args, "action") == "get_result" {
|
if toolName == builtin.ToolC2File && mcpAuthorizationString(args, "action") == "get_result" {
|
||||||
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
|
return authorizeC2Action(ctx, principal, db, args, "c2_task", "task_id")
|
||||||
}
|
}
|
||||||
return authorizeC2Action(principal, db, args, "c2_session", "session_id")
|
return authorizeC2Action(ctx, principal, db, args, "c2_session", "session_id")
|
||||||
case builtin.ToolC2TaskManage:
|
case builtin.ToolC2TaskManage:
|
||||||
return authorizeC2Action(principal, db, args, "c2_task", "task_id")
|
return authorizeC2Action(ctx, principal, db, args, "c2_task", "task_id")
|
||||||
case builtin.ToolC2Payload:
|
case builtin.ToolC2Payload:
|
||||||
return resource("c2:write", "c2_listener", "listener_id")
|
return resource("c2:write", "c2_listener", "listener_id")
|
||||||
case builtin.ToolC2Event:
|
case builtin.ToolC2Event:
|
||||||
if id := mcpAuthorizationString(args, "session_id"); id != "" {
|
if id := mcpAuthorizationString(args, "session_id"); id != "" {
|
||||||
return resource("c2:read", "c2_session", "session_id")
|
return resource("c2:read", "c2_session", "session_id")
|
||||||
}
|
}
|
||||||
|
if id := mcpAuthorizationString(args, "task_id"); id != "" {
|
||||||
|
return resource("c2:read", "c2_task", "task_id")
|
||||||
|
}
|
||||||
|
if filter := mcpEffectiveProjectFilter(ctx, db); filter != "" {
|
||||||
|
return require("c2:read")
|
||||||
|
}
|
||||||
if principal.ScopeFor("c2:read") != database.RBACScopeAll {
|
if principal.ScopeFor("c2:read") != database.RBACScopeAll {
|
||||||
return fmt.Errorf("unfiltered C2 event list requires global scope")
|
return fmt.Errorf("unfiltered C2 event list requires global scope")
|
||||||
}
|
}
|
||||||
@@ -207,7 +216,7 @@ func externalMCPToolAuthorizer() func(context.Context, string, map[string]interf
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error {
|
func authorizeC2Action(ctx context.Context, principal authctx.Principal, db *database.DB, args map[string]interface{}, resourceType, argument string) error {
|
||||||
action := mcpAuthorizationString(args, "action")
|
action := mcpAuthorizationString(args, "action")
|
||||||
permission := "c2:write"
|
permission := "c2:write"
|
||||||
if action == "list" || action == "get" || action == "get_result" || action == "wait" {
|
if action == "list" || action == "get" || action == "get_result" || action == "wait" {
|
||||||
@@ -228,11 +237,27 @@ func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[st
|
|||||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) {
|
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, candidate) {
|
||||||
return fmt.Errorf("no access to %s %s", resourceType, candidate)
|
return fmt.Errorf("no access to %s %s", resourceType, candidate)
|
||||||
}
|
}
|
||||||
|
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, candidate); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
if id == "" {
|
if id == "" {
|
||||||
if action == "create" || action == "list" {
|
if action == "create" {
|
||||||
|
projectID := mcpAuthorizationString(args, "project_id")
|
||||||
|
if projectID == "" {
|
||||||
|
projectID = mcpEffectiveProjectFilter(ctx, db)
|
||||||
|
if projectID == database.ProjectFilterUnbound {
|
||||||
|
projectID = ""
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if projectID != "" && (db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID)) {
|
||||||
|
return fmt.Errorf("no access to project %s", projectID)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if action == "list" {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
return fmt.Errorf("missing resource identifier %s", argument)
|
return fmt.Errorf("missing resource identifier %s", argument)
|
||||||
@@ -240,9 +265,93 @@ func authorizeC2Action(principal authctx.Principal, db *database.DB, args map[st
|
|||||||
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
if db == nil || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), resourceType, id) {
|
||||||
return fmt.Errorf("no access to %s %s", resourceType, id)
|
return fmt.Errorf("no access to %s %s", resourceType, id)
|
||||||
}
|
}
|
||||||
|
if err := authorizeMCPProjectResourceBoundary(ctx, db, resourceType, id); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func authorizeMCPProjectResourceBoundary(ctx context.Context, db *database.DB, resourceType, resourceID string) error {
|
||||||
|
filter := mcpEffectiveProjectFilter(ctx, db)
|
||||||
|
if filter == "" || db == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
projectID, ok, err := mcpResourceProjectID(db, resourceType, resourceID)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if filter == database.ProjectFilterUnbound {
|
||||||
|
if projectID != "" {
|
||||||
|
return fmt.Errorf("resource %s %s belongs to project %s, current conversation is unbound", resourceType, resourceID, projectID)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if projectID != filter {
|
||||||
|
if projectID == "" {
|
||||||
|
return fmt.Errorf("resource %s %s is unbound, current conversation project is %s", resourceType, resourceID, filter)
|
||||||
|
}
|
||||||
|
return fmt.Errorf("resource %s %s belongs to project %s, current conversation project is %s", resourceType, resourceID, projectID, filter)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func mcpResourceProjectID(db *database.DB, resourceType, resourceID string) (string, bool, error) {
|
||||||
|
switch resourceType {
|
||||||
|
case "webshell":
|
||||||
|
conn, err := db.GetWebshellConnection(resourceID)
|
||||||
|
if err != nil {
|
||||||
|
return "", true, err
|
||||||
|
}
|
||||||
|
if conn == nil {
|
||||||
|
return "", true, fmt.Errorf("webshell not found")
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(conn.ProjectID), true, nil
|
||||||
|
case "c2_listener":
|
||||||
|
listener, err := db.GetC2Listener(resourceID)
|
||||||
|
if err != nil {
|
||||||
|
return "", true, err
|
||||||
|
}
|
||||||
|
if listener == nil {
|
||||||
|
return "", true, fmt.Errorf("listener not found")
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(listener.ProjectID), true, nil
|
||||||
|
case "c2_session":
|
||||||
|
session, err := db.GetC2Session(resourceID)
|
||||||
|
if err != nil {
|
||||||
|
return "", true, err
|
||||||
|
}
|
||||||
|
if session == nil {
|
||||||
|
return "", true, fmt.Errorf("session not found")
|
||||||
|
}
|
||||||
|
return mcpResourceProjectID(db, "c2_listener", session.ListenerID)
|
||||||
|
case "c2_task":
|
||||||
|
task, err := db.GetC2Task(resourceID)
|
||||||
|
if err != nil {
|
||||||
|
return "", true, err
|
||||||
|
}
|
||||||
|
if task == nil {
|
||||||
|
return "", true, fmt.Errorf("task not found")
|
||||||
|
}
|
||||||
|
return mcpResourceProjectIDFromC2Session(db, task.SessionID)
|
||||||
|
default:
|
||||||
|
return "", false, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func mcpResourceProjectIDFromC2Session(db *database.DB, sessionID string) (string, bool, error) {
|
||||||
|
session, err := db.GetC2Session(sessionID)
|
||||||
|
if err != nil {
|
||||||
|
return "", true, err
|
||||||
|
}
|
||||||
|
if session == nil {
|
||||||
|
return "", true, fmt.Errorf("session not found")
|
||||||
|
}
|
||||||
|
return mcpResourceProjectID(db, "c2_listener", session.ListenerID)
|
||||||
|
}
|
||||||
|
|
||||||
func mcpAuthorizationStrings(args map[string]interface{}, key string) []string {
|
func mcpAuthorizationStrings(args map[string]interface{}, key string) []string {
|
||||||
values := []string{}
|
values := []string{}
|
||||||
switch raw := args[key].(type) {
|
switch raw := args[key].(type) {
|
||||||
@@ -273,7 +382,13 @@ func authorizeProjectTool(ctx context.Context, principal authctx.Principal, db *
|
|||||||
return fmt.Errorf("no access to conversation %s", conversationID)
|
return fmt.Errorf("no access to conversation %s", conversationID)
|
||||||
}
|
}
|
||||||
projectID, err := db.GetConversationProjectID(conversationID)
|
projectID, err := db.GetConversationProjectID(conversationID)
|
||||||
if err != nil || strings.TrimSpace(projectID) == "" || !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) {
|
if err != nil {
|
||||||
|
return fmt.Errorf("no access to project: %w", err)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(projectID) == "" {
|
||||||
|
return fmt.Errorf("当前对话未绑定项目,无法使用项目黑板工具,请先在对话中选择项目或创建带项目的对话")
|
||||||
|
}
|
||||||
|
if !db.UserCanAccessResource(principal.UserID, principal.ScopeFor(permission), "project", projectID) {
|
||||||
return fmt.Errorf("no access to project %s", projectID)
|
return fmt.Errorf("no access to project %s", projectID)
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
|
|||||||
@@ -49,6 +49,91 @@ func TestMCPToolAuthorizerEnforcesPermissionAndResource(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestMCPToolAuthorizerEnforcesConversationProjectBoundary(t *testing.T) {
|
||||||
|
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-project-boundary.db"), zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
user, err := db.CreateRBACUser("boundary-user", "Boundary User", "hash", true, nil)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
project, err := db.CreateProject(&database.Project{Name: "Project 123"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
projectConv, err := db.CreateConversation("project conversation", database.ConversationCreateMeta{ProjectID: project.ID})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
unboundConv, err := db.CreateConversation("unbound conversation", database.ConversationCreateMeta{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
wsProject := database.WebShellConnection{ID: "ws_project", ProjectID: project.ID, URL: "http://127.0.0.1/project.php", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
|
wsUnbound := database.WebShellConnection{ID: "ws_unbound", URL: "http://127.0.0.1/unbound.php", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
|
if err := db.CreateWebshellConnection(&wsProject); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateWebshellConnection(&wsUnbound); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, id := range []string{wsProject.ID, wsUnbound.ID} {
|
||||||
|
if err := db.AssignResourceToUser(user.ID, "webshell", id); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
listener := &database.C2Listener{ID: "l_project", ProjectID: project.ID, Name: "project listener", Type: "tcp_reverse", BindHost: "127.0.0.1", BindPort: 5555, OwnerUserID: user.ID, CreatedAt: now}
|
||||||
|
if err := db.CreateC2Listener(listener); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AssignResourceToUser(user.ID, "c2_listener", listener.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
session := &database.C2Session{ID: "s_project", ListenerID: listener.ID, ImplantUUID: "implant-project", Status: "active", FirstSeenAt: now, LastCheckIn: now}
|
||||||
|
if err := db.UpsertC2Session(session); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
principal := authctx.NewPrincipal(user.ID, user.Username, database.RBACScopeAssigned, map[string]bool{
|
||||||
|
"webshell:read": true, "webshell:write": true,
|
||||||
|
"c2:read": true, "c2:write": true,
|
||||||
|
})
|
||||||
|
authorize := mcpToolAuthorizer(db)
|
||||||
|
unboundCtx := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), unboundConv.ID), principal)
|
||||||
|
projectCtx := authctx.WithPrincipal(mcp.WithMCPProjectID(mcp.WithMCPConversationID(context.Background(), projectConv.ID), project.ID), principal)
|
||||||
|
projectCtxFromConversationOnly := authctx.WithPrincipal(mcp.WithMCPConversationID(context.Background(), projectConv.ID), principal)
|
||||||
|
|
||||||
|
if err := authorize(unboundCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err == nil {
|
||||||
|
t.Fatal("unbound conversation was allowed to use project-bound webshell")
|
||||||
|
}
|
||||||
|
if err := authorize(unboundCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsUnbound.ID}); err != nil {
|
||||||
|
t.Fatalf("unbound webshell denied in unbound conversation: %v", err)
|
||||||
|
}
|
||||||
|
if err := authorize(projectCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err != nil {
|
||||||
|
t.Fatalf("project webshell denied in project conversation: %v", err)
|
||||||
|
}
|
||||||
|
if err := authorize(projectCtxFromConversationOnly, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsProject.ID}); err != nil {
|
||||||
|
t.Fatalf("project webshell denied when only conversation id is present: %v", err)
|
||||||
|
}
|
||||||
|
if err := authorize(projectCtx, builtin.ToolWebshellExec, map[string]interface{}{"connection_id": wsUnbound.ID}); err == nil {
|
||||||
|
t.Fatal("project conversation was allowed to use unbound webshell by id")
|
||||||
|
}
|
||||||
|
if err := authorize(unboundCtx, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err == nil {
|
||||||
|
t.Fatal("unbound conversation was allowed to use project-bound c2 session")
|
||||||
|
}
|
||||||
|
if err := authorize(projectCtx, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err != nil {
|
||||||
|
t.Fatalf("project c2 session denied in project conversation: %v", err)
|
||||||
|
}
|
||||||
|
if err := authorize(projectCtxFromConversationOnly, builtin.ToolC2Session, map[string]interface{}{"action": "get", "session_id": session.ID}); err != nil {
|
||||||
|
t.Fatalf("project c2 session denied when only conversation id is present: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
func TestEveryBuiltinMCPToolHasExplicitAuthorizationPolicy(t *testing.T) {
|
||||||
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
|
db, err := database.NewDB(filepath.Join(t.TempDir(), "mcp-policy-inventory.db"), zap.NewNop())
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@@ -0,0 +1,26 @@
|
|||||||
|
package app
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/database"
|
||||||
|
"cyberstrike-ai/internal/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
func mcpEffectiveProjectFilter(ctx context.Context, db *database.DB) string {
|
||||||
|
if projectID := strings.TrimSpace(mcp.MCPProjectIDFromContext(ctx)); projectID != "" {
|
||||||
|
return projectID
|
||||||
|
}
|
||||||
|
if conversationID := mcpAuthorizationConversationID(ctx); conversationID != "" {
|
||||||
|
if db != nil {
|
||||||
|
if projectID, err := db.GetConversationProjectID(conversationID); err == nil {
|
||||||
|
if projectID = strings.TrimSpace(projectID); projectID != "" {
|
||||||
|
return projectID
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return database.ProjectFilterUnbound
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
@@ -34,9 +34,9 @@ type Manager struct {
|
|||||||
runningListeners map[string]Listener // listener_id → 已 Start 的 listener 实例
|
runningListeners map[string]Listener // listener_id → 已 Start 的 listener 实例
|
||||||
storageDir string // 大结果(截图/下载)落盘根目录
|
storageDir string // 大结果(截图/下载)落盘根目录
|
||||||
|
|
||||||
hitlBridge HITLBridge // 危险任务在 EnqueueTask 时调它发起审批(nil 表示不接 HITL)
|
hitlBridge HITLBridge // 危险任务在 EnqueueTask 时调它发起审批(nil 表示不接 HITL)
|
||||||
hitlDangerousGate func(conversationID, mcpToolName string) bool // 与人机协同一致:为 nil 或返回 false 时不走桥
|
hitlDangerousGate func(conversationID, mcpToolName string) bool // 与人机协同一致:为 nil 或返回 false 时不走桥
|
||||||
hooks Hooks // 扩展挂钩:会话上线 / 任务完成 时通知漏洞库与攻击链
|
hooks Hooks // 扩展挂钩:会话上线 / 任务完成 时通知漏洞库与攻击链
|
||||||
}
|
}
|
||||||
|
|
||||||
// MCPToolC2Task 与 MCP builtin、c2_task 工具名一致,供 HITL 白名单与 Agent 侧对齐。
|
// MCPToolC2Task 与 MCP builtin、c2_task 工具名一致,供 HITL 白名单与 Agent 侧对齐。
|
||||||
@@ -63,7 +63,7 @@ type HITLApprovalRequest struct {
|
|||||||
|
|
||||||
// Hooks 给上层(漏洞管理 / 攻击链)注入回调
|
// Hooks 给上层(漏洞管理 / 攻击链)注入回调
|
||||||
type Hooks struct {
|
type Hooks struct {
|
||||||
OnSessionFirstSeen func(session *database.C2Session) // 新会话首次上线
|
OnSessionFirstSeen func(session *database.C2Session) // 新会话首次上线
|
||||||
OnTaskCompleted func(task *database.C2Task, sessionID string) // 任务完成(success/failed)
|
OnTaskCompleted func(task *database.C2Task, sessionID string) // 任务完成(success/failed)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -144,6 +144,7 @@ func (m *Manager) Close() {
|
|||||||
// CreateListenerInput Web/MCP 创建监听器的入参(已校验 + 已 trim)
|
// CreateListenerInput Web/MCP 创建监听器的入参(已校验 + 已 trim)
|
||||||
type CreateListenerInput struct {
|
type CreateListenerInput struct {
|
||||||
Name string
|
Name string
|
||||||
|
ProjectID string
|
||||||
Type string
|
Type string
|
||||||
BindHost string
|
BindHost string
|
||||||
BindPort int
|
BindPort int
|
||||||
@@ -195,6 +196,7 @@ func (m *Manager) CreateListener(in CreateListenerInput) (*database.C2Listener,
|
|||||||
|
|
||||||
listener := &database.C2Listener{
|
listener := &database.C2Listener{
|
||||||
ID: "l_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14],
|
ID: "l_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:14],
|
||||||
|
ProjectID: strings.TrimSpace(in.ProjectID),
|
||||||
Name: strings.TrimSpace(in.Name),
|
Name: strings.TrimSpace(in.Name),
|
||||||
Type: strings.ToLower(strings.TrimSpace(in.Type)),
|
Type: strings.ToLower(strings.TrimSpace(in.Type)),
|
||||||
BindHost: bindHost,
|
BindHost: bindHost,
|
||||||
|
|||||||
+140
-6
@@ -21,7 +21,8 @@ type Config struct {
|
|||||||
Server ServerConfig `yaml:"server"`
|
Server ServerConfig `yaml:"server"`
|
||||||
Log LogConfig `yaml:"log"`
|
Log LogConfig `yaml:"log"`
|
||||||
MCP MCPConfig `yaml:"mcp"`
|
MCP MCPConfig `yaml:"mcp"`
|
||||||
OpenAI OpenAIConfig `yaml:"openai"`
|
AI AIConfig `yaml:"ai,omitempty" json:"ai,omitempty"`
|
||||||
|
OpenAI OpenAIConfig `yaml:"openai,omitempty" json:"openai,omitempty"`
|
||||||
FOFA FofaConfig `yaml:"fofa,omitempty" json:"fofa,omitempty"`
|
FOFA FofaConfig `yaml:"fofa,omitempty" json:"fofa,omitempty"`
|
||||||
ZoomEye SpaceSearchConfig `yaml:"zoomeye,omitempty" json:"zoomeye,omitempty"`
|
ZoomEye SpaceSearchConfig `yaml:"zoomeye,omitempty" json:"zoomeye,omitempty"`
|
||||||
Quake SpaceSearchConfig `yaml:"quake,omitempty" json:"quake,omitempty"`
|
Quake SpaceSearchConfig `yaml:"quake,omitempty" json:"quake,omitempty"`
|
||||||
@@ -840,6 +841,130 @@ type OpenAIConfig struct {
|
|||||||
Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`
|
Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// AIConfig stores first-class model channels. Runtime callers resolve a channel
|
||||||
|
// into OpenAIConfig at the edge instead of moving API credentials through chat requests.
|
||||||
|
type AIConfig struct {
|
||||||
|
DefaultChannel string `yaml:"default_channel,omitempty" json:"default_channel,omitempty"`
|
||||||
|
Channels map[string]AIChannelConfig `yaml:"channels,omitempty" json:"channels,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type AIChannelConfig struct {
|
||||||
|
Name string `yaml:"name,omitempty" json:"name,omitempty"`
|
||||||
|
Provider string `yaml:"provider,omitempty" json:"provider,omitempty"`
|
||||||
|
APIKey string `yaml:"api_key" json:"api_key"`
|
||||||
|
BaseURL string `yaml:"base_url" json:"base_url"`
|
||||||
|
Model string `yaml:"model" json:"model"`
|
||||||
|
MaxTotalTokens int `yaml:"max_total_tokens,omitempty" json:"max_total_tokens,omitempty"`
|
||||||
|
MaxCompletionTokens int `yaml:"max_completion_tokens,omitempty" json:"max_completion_tokens,omitempty"`
|
||||||
|
Reasoning OpenAIReasoningConfig `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c AIChannelConfig) ToOpenAIConfig() OpenAIConfig {
|
||||||
|
provider := strings.TrimSpace(c.Provider)
|
||||||
|
if provider == "" || provider == "openai_compatible" {
|
||||||
|
provider = "openai"
|
||||||
|
}
|
||||||
|
return OpenAIConfig{
|
||||||
|
Provider: provider,
|
||||||
|
APIKey: c.APIKey,
|
||||||
|
BaseURL: c.BaseURL,
|
||||||
|
Model: c.Model,
|
||||||
|
MaxTotalTokens: c.MaxTotalTokens,
|
||||||
|
MaxCompletionTokens: c.MaxCompletionTokens,
|
||||||
|
Reasoning: c.Reasoning,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func AIChannelFromOpenAI(id, name string, oa OpenAIConfig) AIChannelConfig {
|
||||||
|
if strings.TrimSpace(name) == "" {
|
||||||
|
name = id
|
||||||
|
}
|
||||||
|
return AIChannelConfig{
|
||||||
|
Name: name,
|
||||||
|
Provider: oa.Provider,
|
||||||
|
APIKey: oa.APIKey,
|
||||||
|
BaseURL: oa.BaseURL,
|
||||||
|
Model: oa.Model,
|
||||||
|
MaxTotalTokens: oa.MaxTotalTokens,
|
||||||
|
MaxCompletionTokens: oa.MaxCompletionTokens,
|
||||||
|
Reasoning: oa.Reasoning,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func NormalizeAIChannelID(s string) string {
|
||||||
|
id := strings.ToLower(strings.TrimSpace(s))
|
||||||
|
id = strings.ReplaceAll(id, "_", "-")
|
||||||
|
var b strings.Builder
|
||||||
|
lastDash := false
|
||||||
|
for _, r := range id {
|
||||||
|
ok := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9')
|
||||||
|
if ok {
|
||||||
|
b.WriteRune(r)
|
||||||
|
lastDash = false
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if !lastDash {
|
||||||
|
b.WriteByte('-')
|
||||||
|
lastDash = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out := strings.Trim(b.String(), "-")
|
||||||
|
if out == "" {
|
||||||
|
return "default"
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *AIConfig) EnsureDefaultFromOpenAI(openAI OpenAIConfig) {
|
||||||
|
if c.Channels == nil {
|
||||||
|
c.Channels = make(map[string]AIChannelConfig)
|
||||||
|
}
|
||||||
|
def := NormalizeAIChannelID(c.DefaultChannel)
|
||||||
|
if def == "default" && strings.TrimSpace(c.DefaultChannel) == "" {
|
||||||
|
def = "default"
|
||||||
|
}
|
||||||
|
c.DefaultChannel = def
|
||||||
|
if _, ok := c.Channels[def]; !ok {
|
||||||
|
c.Channels[def] = AIChannelFromOpenAI(def, "Default", openAI)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c AIConfig) ResolveChannel(channelID string) (OpenAIConfig, string, bool) {
|
||||||
|
id := NormalizeAIChannelID(channelID)
|
||||||
|
if strings.TrimSpace(channelID) == "" {
|
||||||
|
id = NormalizeAIChannelID(c.DefaultChannel)
|
||||||
|
}
|
||||||
|
if id == "" {
|
||||||
|
id = "default"
|
||||||
|
}
|
||||||
|
if c.Channels != nil {
|
||||||
|
if ch, ok := c.Channels[id]; ok {
|
||||||
|
return ch.ToOpenAIConfig(), id, true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return OpenAIConfig{}, id, false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Config) ResolveAIChannel(channelID string) (OpenAIConfig, string, bool) {
|
||||||
|
if c == nil {
|
||||||
|
return OpenAIConfig{}, "", false
|
||||||
|
}
|
||||||
|
if oa, id, ok := c.AI.ResolveChannel(channelID); ok {
|
||||||
|
return oa, id, true
|
||||||
|
}
|
||||||
|
return c.OpenAI, NormalizeAIChannelID(channelID), strings.TrimSpace(c.OpenAI.Model) != "" || strings.TrimSpace(c.OpenAI.BaseURL) != ""
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *Config) ApplyDefaultAIChannel() {
|
||||||
|
if c == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
c.AI.EnsureDefaultFromOpenAI(c.OpenAI)
|
||||||
|
if oa, _, ok := c.AI.ResolveChannel(c.AI.DefaultChannel); ok {
|
||||||
|
c.OpenAI = oa
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (c OpenAIConfig) MaxCompletionTokensEffective() int {
|
func (c OpenAIConfig) MaxCompletionTokensEffective() int {
|
||||||
if c.MaxCompletionTokens > 0 {
|
if c.MaxCompletionTokens > 0 {
|
||||||
return c.MaxCompletionTokens
|
return c.MaxCompletionTokens
|
||||||
@@ -1224,6 +1349,7 @@ func Load(path string) (*Config, error) {
|
|||||||
if cfg.Audit.MaxDetailBytes <= 0 {
|
if cfg.Audit.MaxDetailBytes <= 0 {
|
||||||
cfg.Audit.MaxDetailBytes = 8192
|
cfg.Audit.MaxDetailBytes = 8192
|
||||||
}
|
}
|
||||||
|
cfg.ApplyDefaultAIChannel()
|
||||||
if err := validateModelOutputLimits(cfg.OpenAI, cfg.MultiAgent.EinoMiddleware); err != nil {
|
if err := validateModelOutputLimits(cfg.OpenAI, cfg.MultiAgent.EinoMiddleware); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
@@ -1720,12 +1846,20 @@ func Default() *Config {
|
|||||||
Host: "127.0.0.1",
|
Host: "127.0.0.1",
|
||||||
Port: 8081,
|
Port: 8081,
|
||||||
},
|
},
|
||||||
OpenAI: OpenAIConfig{
|
AI: AIConfig{
|
||||||
BaseURL: "https://api.openai.com/v1",
|
DefaultChannel: "default",
|
||||||
Model: "gpt-4",
|
Channels: map[string]AIChannelConfig{
|
||||||
MaxTotalTokens: 120000,
|
"default": {
|
||||||
MaxCompletionTokens: DefaultMaxCompletionTokens,
|
Name: "Default",
|
||||||
|
Provider: "openai_compatible",
|
||||||
|
BaseURL: "https://api.openai.com/v1",
|
||||||
|
Model: "gpt-4",
|
||||||
|
MaxTotalTokens: 120000,
|
||||||
|
MaxCompletionTokens: DefaultMaxCompletionTokens,
|
||||||
|
},
|
||||||
|
},
|
||||||
},
|
},
|
||||||
|
OpenAI: OpenAIConfig{},
|
||||||
Agent: AgentConfig{
|
Agent: AgentConfig{
|
||||||
MaxIterations: 30, // 默认最大迭代次数
|
MaxIterations: 30, // 默认最大迭代次数
|
||||||
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
|
ToolTimeoutMinutes: 10, // 单次工具执行默认最多 10 分钟,避免异常长时间占用
|
||||||
|
|||||||
@@ -95,6 +95,48 @@ func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestLoadUsesAIDefaultChannelAsRuntimeOpenAI(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
path := filepath.Join(dir, "config.yaml")
|
||||||
|
initial := strings.Join([]string{
|
||||||
|
"ai:",
|
||||||
|
" default_channel: deepseek",
|
||||||
|
" channels:",
|
||||||
|
" qwen:",
|
||||||
|
" name: Qwen",
|
||||||
|
" provider: openai_compatible",
|
||||||
|
" base_url: https://dashscope.example/v1",
|
||||||
|
" api_key: qwen-key",
|
||||||
|
" model: qwen-max",
|
||||||
|
" deepseek:",
|
||||||
|
" name: DeepSeek",
|
||||||
|
" provider: openai_compatible",
|
||||||
|
" base_url: https://deepseek.example/v1",
|
||||||
|
" api_key: deepseek-key",
|
||||||
|
" model: deepseek-chat",
|
||||||
|
" max_total_tokens: 64000",
|
||||||
|
"server:",
|
||||||
|
" host: 127.0.0.1",
|
||||||
|
" port: 8080",
|
||||||
|
"",
|
||||||
|
}, "\n")
|
||||||
|
if err := os.WriteFile(path, []byte(initial), 0644); err != nil {
|
||||||
|
t.Fatalf("write config: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg, err := Load(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Load: %v", err)
|
||||||
|
}
|
||||||
|
if cfg.OpenAI.Model != "deepseek-chat" || cfg.OpenAI.APIKey != "deepseek-key" || cfg.OpenAI.MaxTotalTokens != 64000 {
|
||||||
|
t.Fatalf("runtime OpenAI config did not follow ai.default_channel: %+v", cfg.OpenAI)
|
||||||
|
}
|
||||||
|
oa, id, ok := cfg.ResolveAIChannel("qwen")
|
||||||
|
if !ok || id != "qwen" || oa.Model != "qwen-max" || oa.APIKey != "qwen-key" {
|
||||||
|
t.Fatalf("ResolveAIChannel(qwen) = (%+v, %q, %v)", oa, id, ok)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) {
|
func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) {
|
||||||
var zero MultiAgentEinoMiddlewareConfig
|
var zero MultiAgentEinoMiddlewareConfig
|
||||||
if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes {
|
if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes {
|
||||||
|
|||||||
+150
-8
@@ -13,6 +13,7 @@ import (
|
|||||||
"unicode/utf8"
|
"unicode/utf8"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
|
"go.uber.org/zap"
|
||||||
"golang.org/x/net/idna"
|
"golang.org/x/net/idna"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -50,6 +51,7 @@ type Asset struct {
|
|||||||
LastScanTaskID string `json:"last_scan_task_id,omitempty"`
|
LastScanTaskID string `json:"last_scan_task_id,omitempty"`
|
||||||
VulnerabilityCount int `json:"vulnerability_count"`
|
VulnerabilityCount int `json:"vulnerability_count"`
|
||||||
RiskLevel string `json:"risk_level"`
|
RiskLevel string `json:"risk_level"`
|
||||||
|
RiskScore int `json:"-"`
|
||||||
OwnerUserID string `json:"-"`
|
OwnerUserID string `json:"-"`
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -443,15 +445,15 @@ func assetWhere(filter AssetListFilter, access RBACListAccess) (string, []interf
|
|||||||
args = append(args, *filter.Port)
|
args = append(args, *filter.Port)
|
||||||
}
|
}
|
||||||
if filter.RiskLevel != "" {
|
if filter.RiskLevel != "" {
|
||||||
query += " AND " + assetRiskLevelExpr + " = ?"
|
query += " AND " + assetRiskLevelCachedExpr + " = ?"
|
||||||
args = append(args, strings.ToLower(strings.TrimSpace(filter.RiskLevel)))
|
args = append(args, strings.ToLower(strings.TrimSpace(filter.RiskLevel)))
|
||||||
}
|
}
|
||||||
if filter.MinVulnerabilities != nil {
|
if filter.MinVulnerabilities != nil {
|
||||||
query += " AND " + assetVulnerabilityCountExpr + " >= ?"
|
query += " AND " + assetVulnerabilityCountCachedExpr + " >= ?"
|
||||||
args = append(args, *filter.MinVulnerabilities)
|
args = append(args, *filter.MinVulnerabilities)
|
||||||
}
|
}
|
||||||
if filter.MaxVulnerabilities != nil {
|
if filter.MaxVulnerabilities != nil {
|
||||||
query += " AND " + assetVulnerabilityCountExpr + " <= ?"
|
query += " AND " + assetVulnerabilityCountCachedExpr + " <= ?"
|
||||||
args = append(args, *filter.MaxVulnerabilities)
|
args = append(args, *filter.MaxVulnerabilities)
|
||||||
}
|
}
|
||||||
for _, item := range []struct {
|
for _, item := range []struct {
|
||||||
@@ -573,20 +575,24 @@ const assetVulnerabilityMatchExpr = `(
|
|||||||
|
|
||||||
const assetVulnerabilityCountExpr = `(SELECT COUNT(DISTINCT v.id) FROM vulnerabilities v WHERE ` + assetVulnerabilityMatchExpr + `)`
|
const assetVulnerabilityCountExpr = `(SELECT COUNT(DISTINCT v.id) FROM vulnerabilities v WHERE ` + assetVulnerabilityMatchExpr + `)`
|
||||||
|
|
||||||
const assetRiskScoreExpr = `COALESCE((
|
const assetRiskScoreQueryExpr = `COALESCE((
|
||||||
SELECT MAX(CASE LOWER(COALESCE(v.severity,'')) WHEN 'critical' THEN 5 WHEN 'high' THEN 4 WHEN 'medium' THEN 3 WHEN 'low' THEN 2 WHEN 'info' THEN 1 ELSE 0 END)
|
SELECT MAX(CASE LOWER(COALESCE(v.severity,'')) WHEN 'critical' THEN 5 WHEN 'high' THEN 4 WHEN 'medium' THEN 3 WHEN 'low' THEN 2 WHEN 'info' THEN 1 ELSE 0 END)
|
||||||
FROM vulnerabilities v
|
FROM vulnerabilities v
|
||||||
WHERE LOWER(COALESCE(v.status,'open')) NOT IN ('fixed','false_positive','ignored') AND ` + assetVulnerabilityMatchExpr + `
|
WHERE LOWER(COALESCE(v.status,'open')) NOT IN ('fixed','false_positive','ignored') AND ` + assetVulnerabilityMatchExpr + `
|
||||||
),0)`
|
),0)`
|
||||||
|
|
||||||
const assetRiskLevelExpr = `(CASE WHEN ` + assetEffectiveLastScanExpr + ` IS NULL THEN 'unassessed' ELSE CASE ` + assetRiskScoreExpr + `
|
const assetRiskLevelQueryExpr = `(CASE WHEN ` + assetEffectiveLastScanExpr + ` IS NULL THEN 'unassessed' ELSE CASE ` + assetRiskScoreQueryExpr + `
|
||||||
WHEN 5 THEN 'critical' WHEN 4 THEN 'high' WHEN 3 THEN 'medium' WHEN 2 THEN 'low' WHEN 1 THEN 'info' ELSE 'normal' END END)`
|
WHEN 5 THEN 'critical' WHEN 4 THEN 'high' WHEN 3 THEN 'medium' WHEN 2 THEN 'low' WHEN 1 THEN 'info' ELSE 'normal' END END)`
|
||||||
|
|
||||||
|
const assetVulnerabilityCountCachedExpr = `COALESCE(assets.vulnerability_count,0)`
|
||||||
|
const assetRiskScoreCachedExpr = `COALESCE(assets.risk_score,0)`
|
||||||
|
const assetRiskLevelCachedExpr = `COALESCE(NULLIF(assets.risk_level,''),'unassessed')`
|
||||||
|
|
||||||
const assetSelectColumns = `assets.id,COALESCE(assets.project_id,''),COALESCE(p.name,''),assets.host,assets.ip,assets.port,assets.domain,assets.protocol,assets.title,assets.server,assets.country,
|
const assetSelectColumns = `assets.id,COALESCE(assets.project_id,''),COALESCE(p.name,''),assets.host,assets.ip,assets.port,assets.domain,assets.protocol,assets.title,assets.server,assets.country,
|
||||||
assets.province,assets.city,assets.responsible_person,assets.department,assets.business_system,assets.environment,assets.criticality,
|
assets.province,assets.city,assets.responsible_person,assets.department,assets.business_system,assets.environment,assets.criticality,
|
||||||
assets.source,assets.source_query,assets.status,assets.tags_json,assets.first_seen_at,assets.last_seen_at,assets.created_at,assets.updated_at,
|
assets.source,assets.source_query,assets.status,assets.tags_json,assets.first_seen_at,assets.last_seen_at,assets.created_at,assets.updated_at,
|
||||||
` + assetEffectiveLastScanExpr + `,COALESCE(assets.last_scan_conversation_id,''),COALESCE(assets.last_scan_queue_id,''),COALESCE(assets.last_scan_task_id,''),
|
` + assetEffectiveLastScanExpr + `,COALESCE(assets.last_scan_conversation_id,''),COALESCE(assets.last_scan_queue_id,''),COALESCE(assets.last_scan_task_id,''),
|
||||||
` + assetVulnerabilityCountExpr + `,` + assetRiskLevelExpr
|
` + assetVulnerabilityCountCachedExpr + `,` + assetRiskLevelCachedExpr
|
||||||
|
|
||||||
// MarkAssetScanned links an asset to the conversation or batch subtask created from it.
|
// MarkAssetScanned links an asset to the conversation or batch subtask created from it.
|
||||||
// The link lets the asset list show the latest scan time and vulnerabilities produced by that scan.
|
// The link lets the asset list show the latest scan time and vulnerabilities produced by that scan.
|
||||||
@@ -601,6 +607,9 @@ func (db *DB) MarkAssetScanned(id, conversationID, queueID, taskID string, acces
|
|||||||
if n == 0 {
|
if n == 0 {
|
||||||
return sql.ErrNoRows
|
return sql.ErrNoRows
|
||||||
}
|
}
|
||||||
|
if err := db.RefreshAssetRiskCache(id); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -629,6 +638,9 @@ func (db *DB) CompleteAssetScan(id, conversationID string, access RBACListAccess
|
|||||||
if n == 0 {
|
if n == 0 {
|
||||||
return sql.ErrNoRows
|
return sql.ErrNoRows
|
||||||
}
|
}
|
||||||
|
if err := db.RefreshAssetRiskCache(id); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -638,6 +650,136 @@ func (db *DB) BatchTaskBelongsToQueue(taskID, queueID string) bool {
|
|||||||
return err == nil && count > 0
|
return err == nil && count > 0
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func assetRiskLevelFromScore(score int, scanned bool) string {
|
||||||
|
if !scanned {
|
||||||
|
return "unassessed"
|
||||||
|
}
|
||||||
|
switch score {
|
||||||
|
case 5:
|
||||||
|
return "critical"
|
||||||
|
case 4:
|
||||||
|
return "high"
|
||||||
|
case 3:
|
||||||
|
return "medium"
|
||||||
|
case 2:
|
||||||
|
return "low"
|
||||||
|
case 1:
|
||||||
|
return "info"
|
||||||
|
default:
|
||||||
|
return "normal"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// RefreshAssetRiskCache recalculates the denormalized fields used by the asset
|
||||||
|
// list. Keeping this in the database layer makes Web API and MCP writes share
|
||||||
|
// one consistency path.
|
||||||
|
func (db *DB) RefreshAssetRiskCache(assetID string) error {
|
||||||
|
assetID = strings.TrimSpace(assetID)
|
||||||
|
if assetID == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
var count int
|
||||||
|
if err := db.QueryRow("SELECT "+assetVulnerabilityCountExpr+" FROM assets WHERE assets.id=?", assetID).Scan(&count); err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return fmt.Errorf("刷新资产漏洞数量失败: %w", err)
|
||||||
|
}
|
||||||
|
var score int
|
||||||
|
if err := db.QueryRow("SELECT "+assetRiskScoreQueryExpr+" FROM assets WHERE assets.id=?", assetID).Scan(&score); err != nil {
|
||||||
|
return fmt.Errorf("刷新资产风险分数失败: %w", err)
|
||||||
|
}
|
||||||
|
var lastScan interface{}
|
||||||
|
if err := db.QueryRow("SELECT "+assetEffectiveLastScanExpr+" FROM assets WHERE assets.id=?", assetID).Scan(&lastScan); err != nil {
|
||||||
|
return fmt.Errorf("刷新资产扫描状态失败: %w", err)
|
||||||
|
}
|
||||||
|
level := assetRiskLevelFromScore(score, lastScan != nil)
|
||||||
|
if _, err := db.Exec(`UPDATE assets SET vulnerability_count=?, risk_score=?, risk_level=? WHERE id=?`, count, score, level, assetID); err != nil {
|
||||||
|
return fmt.Errorf("更新资产风险缓存失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) RefreshAllAssetRiskCache() error {
|
||||||
|
rows, err := db.Query(`SELECT id FROM assets`)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("查询资产列表失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
for rows.Next() {
|
||||||
|
var id string
|
||||||
|
if err := rows.Scan(&id); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if err := db.RefreshAssetRiskCache(id); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) AssetIDsForVulnerabilityConversations(conversationIDs []string) ([]string, error) {
|
||||||
|
seen := map[string]struct{}{}
|
||||||
|
cleaned := make([]string, 0, len(conversationIDs))
|
||||||
|
for _, id := range conversationIDs {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if id == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := seen[id]; ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[id] = struct{}{}
|
||||||
|
cleaned = append(cleaned, id)
|
||||||
|
}
|
||||||
|
if len(cleaned) == 0 {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
placeholders := strings.TrimRight(strings.Repeat("?,", len(cleaned)), ",")
|
||||||
|
args := make([]interface{}, 0, len(cleaned)*2)
|
||||||
|
for _, id := range cleaned {
|
||||||
|
args = append(args, id)
|
||||||
|
}
|
||||||
|
for _, id := range cleaned {
|
||||||
|
args = append(args, id)
|
||||||
|
}
|
||||||
|
rows, err := db.Query(`SELECT DISTINCT assets.id FROM assets
|
||||||
|
WHERE assets.last_scan_conversation_id IN (`+placeholders+`)
|
||||||
|
OR assets.last_scan_task_id IN (SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id IN (`+placeholders+`))`, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询受影响资产失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
assetIDs := []string{}
|
||||||
|
for rows.Next() {
|
||||||
|
var id string
|
||||||
|
if err := rows.Scan(&id); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
assetIDs = append(assetIDs, id)
|
||||||
|
}
|
||||||
|
return assetIDs, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) RefreshAssetRiskCacheForConversations(conversationIDs ...string) error {
|
||||||
|
assetIDs, err := db.AssetIDsForVulnerabilityConversations(conversationIDs)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
for _, id := range assetIDs {
|
||||||
|
if err := db.RefreshAssetRiskCache(id); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) refreshAssetRiskCacheForConversationsBestEffort(conversationIDs ...string) {
|
||||||
|
if err := db.RefreshAssetRiskCacheForConversations(conversationIDs...); err != nil && db.logger != nil {
|
||||||
|
db.logger.Warn("刷新资产风险缓存失败", zap.Error(err))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (db *DB) ListAssets(limit, offset int, filter AssetListFilter, access RBACListAccess) ([]*Asset, int, error) {
|
func (db *DB) ListAssets(limit, offset int, filter AssetListFilter, access RBACListAccess) ([]*Asset, int, error) {
|
||||||
if limit < 1 {
|
if limit < 1 {
|
||||||
limit = 20
|
limit = 20
|
||||||
@@ -726,9 +868,9 @@ func assetOrderBy(sortBy, sortOrder string) string {
|
|||||||
case "port":
|
case "port":
|
||||||
expression = "assets.port"
|
expression = "assets.port"
|
||||||
case "vulnerability_count":
|
case "vulnerability_count":
|
||||||
expression = assetVulnerabilityCountExpr
|
expression = assetVulnerabilityCountCachedExpr
|
||||||
case "risk_level":
|
case "risk_level":
|
||||||
expression = assetRiskScoreExpr
|
expression = assetRiskScoreCachedExpr
|
||||||
default:
|
default:
|
||||||
expression = "assets.last_seen_at"
|
expression = "assets.last_seen_at"
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -383,7 +383,12 @@ func TestAssetScanLinkReturnsTimeAndRelatedVulnerabilities(t *testing.T) {
|
|||||||
if linked.LastScanAt == nil || linked.LastScanConversationID != conv.ID || linked.VulnerabilityCount != 1 || linked.RiskLevel != "high" {
|
if linked.LastScanAt == nil || linked.LastScanConversationID != conv.ID || linked.VulnerabilityCount != 1 || linked.RiskLevel != "high" {
|
||||||
t.Fatalf("unexpected scan metadata: %#v", linked)
|
t.Fatalf("unexpected scan metadata: %#v", linked)
|
||||||
}
|
}
|
||||||
if _, err := db.Exec(`UPDATE vulnerabilities SET status='fixed' WHERE conversation_id=?`, conv.ID); err != nil {
|
vulns, err := db.ListVulnerabilities(10, 0, VulnerabilityListFilter{ConversationID: conv.ID})
|
||||||
|
if err != nil || len(vulns) != 1 {
|
||||||
|
t.Fatalf("list linked vulnerabilities: len=%d err=%v", len(vulns), err)
|
||||||
|
}
|
||||||
|
vulns[0].Status = "fixed"
|
||||||
|
if err := db.UpdateVulnerability(vulns[0].ID, vulns[0]); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
resolved, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll})
|
resolved, err := db.GetAsset(assets[0].ID, RBACListAccess{Scope: RBACScopeAll})
|
||||||
|
|||||||
+152
-14
@@ -46,6 +46,7 @@ func validC2TextIDForDelete(id string) bool {
|
|||||||
// C2Listener 监听器实体
|
// C2Listener 监听器实体
|
||||||
type C2Listener struct {
|
type C2Listener struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
|
Type string `json:"type"` // tcp_reverse|http_beacon|https_beacon|websocket|dns
|
||||||
BindHost string `json:"bindHost"` // 默认 127.0.0.1
|
BindHost string `json:"bindHost"` // 默认 127.0.0.1
|
||||||
@@ -165,12 +166,12 @@ func (db *DB) CreateC2Listener(l *C2Listener) error {
|
|||||||
l.ConfigJSON = "{}"
|
l.ConfigJSON = "{}"
|
||||||
}
|
}
|
||||||
query := `
|
query := `
|
||||||
INSERT INTO c2_listeners (id, name, type, bind_host, bind_port, profile_id, encryption_key,
|
INSERT INTO c2_listeners (id, project_id, name, type, bind_host, bind_port, profile_id, encryption_key,
|
||||||
implant_token, status, config_json, remark, owner_user_id, created_at, started_at, last_error)
|
implant_token, status, config_json, remark, owner_user_id, created_at, started_at, last_error)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
`
|
`
|
||||||
_, err := db.Exec(query,
|
_, err := db.Exec(query,
|
||||||
l.ID, l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
l.ID, strings.TrimSpace(l.ProjectID), l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
||||||
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.CreatedAt, l.StartedAt, l.LastError,
|
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.CreatedAt, l.StartedAt, l.LastError,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -190,12 +191,12 @@ func (db *DB) UpdateC2Listener(l *C2Listener) error {
|
|||||||
}
|
}
|
||||||
query := `
|
query := `
|
||||||
UPDATE c2_listeners SET
|
UPDATE c2_listeners SET
|
||||||
name = ?, type = ?, bind_host = ?, bind_port = ?, profile_id = ?, encryption_key = ?,
|
project_id = ?, name = ?, type = ?, bind_host = ?, bind_port = ?, profile_id = ?, encryption_key = ?,
|
||||||
implant_token = ?, status = ?, config_json = ?, remark = ?, owner_user_id = ?, started_at = ?, last_error = ?
|
implant_token = ?, status = ?, config_json = ?, remark = ?, owner_user_id = ?, started_at = ?, last_error = ?
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
`
|
`
|
||||||
res, err := db.Exec(query,
|
res, err := db.Exec(query,
|
||||||
l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
strings.TrimSpace(l.ProjectID), l.Name, l.Type, l.BindHost, l.BindPort, l.ProfileID, l.EncryptionKey,
|
||||||
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.StartedAt, l.LastError, l.ID,
|
l.ImplantToken, l.Status, l.ConfigJSON, l.Remark, l.OwnerUserID, l.StartedAt, l.LastError, l.ID,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -229,7 +230,7 @@ func (db *DB) SetC2ListenerStatus(id, status, lastError string, startedAt *time.
|
|||||||
// GetC2Listener 单条查询
|
// GetC2Listener 单条查询
|
||||||
func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
||||||
query := `
|
query := `
|
||||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||||
@@ -238,7 +239,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
|||||||
var l C2Listener
|
var l C2Listener
|
||||||
var startedAt sql.NullTime
|
var startedAt sql.NullTime
|
||||||
err := db.QueryRow(query, id).Scan(
|
err := db.QueryRow(query, id).Scan(
|
||||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
&l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||||
&l.ConfigJSON, &l.Remark,
|
&l.ConfigJSON, &l.Remark,
|
||||||
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
||||||
@@ -259,7 +260,7 @@ func (db *DB) GetC2Listener(id string) (*C2Listener, error) {
|
|||||||
// ListC2Listeners 全量列表,按创建时间倒序
|
// ListC2Listeners 全量列表,按创建时间倒序
|
||||||
func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
||||||
query := `
|
query := `
|
||||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||||
@@ -275,7 +276,7 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
|||||||
var l C2Listener
|
var l C2Listener
|
||||||
var startedAt sql.NullTime
|
var startedAt sql.NullTime
|
||||||
if err := rows.Scan(
|
if err := rows.Scan(
|
||||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
&l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||||
&l.ConfigJSON, &l.Remark,
|
&l.ConfigJSON, &l.Remark,
|
||||||
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
&l.OwnerUserID, &l.CreatedAt, &startedAt, &l.LastError,
|
||||||
@@ -293,12 +294,18 @@ func (db *DB) ListC2Listeners() ([]*C2Listener, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// ListC2ListenersForAccess lists listeners visible to the resolved RBAC scope.
|
// ListC2ListenersForAccess lists listeners visible to the resolved RBAC scope.
|
||||||
func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, error) {
|
func (db *DB) ListC2ListenersForAccess(access RBACListAccess, projectID string) ([]*C2Listener, error) {
|
||||||
conditions := []string{"1=1"}
|
conditions := []string{"1=1"}
|
||||||
args := []interface{}{}
|
args := []interface{}{}
|
||||||
|
if projectID = strings.TrimSpace(projectID); projectID == ProjectFilterUnbound {
|
||||||
|
conditions = append(conditions, "COALESCE(project_id, '') = ''")
|
||||||
|
} else if projectID != "" {
|
||||||
|
conditions = append(conditions, "COALESCE(project_id, '') = ?")
|
||||||
|
args = append(args, projectID)
|
||||||
|
}
|
||||||
appendC2ListenerAccessFilter(&conditions, &args, access)
|
appendC2ListenerAccessFilter(&conditions, &args, access)
|
||||||
query := `
|
query := `
|
||||||
SELECT id, name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
SELECT id, COALESCE(project_id, ''), name, type, bind_host, bind_port, COALESCE(profile_id, ''),
|
||||||
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
COALESCE(encryption_key, ''), COALESCE(implant_token, ''), status,
|
||||||
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
COALESCE(config_json, '{}'), COALESCE(remark, ''),
|
||||||
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
COALESCE(owner_user_id, ''), created_at, started_at, COALESCE(last_error, '')
|
||||||
@@ -316,7 +323,7 @@ func (db *DB) ListC2ListenersForAccess(access RBACListAccess) ([]*C2Listener, er
|
|||||||
var l C2Listener
|
var l C2Listener
|
||||||
var startedAt sql.NullTime
|
var startedAt sql.NullTime
|
||||||
if err := rows.Scan(
|
if err := rows.Scan(
|
||||||
&l.ID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
&l.ID, &l.ProjectID, &l.Name, &l.Type, &l.BindHost, &l.BindPort, &l.ProfileID,
|
||||||
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
&l.EncryptionKey, &l.ImplantToken, &l.Status,
|
||||||
&l.ConfigJSON, &l.Remark, &l.OwnerUserID,
|
&l.ConfigJSON, &l.Remark, &l.OwnerUserID,
|
||||||
&l.CreatedAt, &startedAt, &l.LastError,
|
&l.CreatedAt, &startedAt, &l.LastError,
|
||||||
@@ -535,6 +542,7 @@ func (db *DB) queryC2SessionWhere(whereClause string, args ...interface{}) (*C2S
|
|||||||
// ListC2SessionsFilter 列表过滤参数
|
// ListC2SessionsFilter 列表过滤参数
|
||||||
type ListC2SessionsFilter struct {
|
type ListC2SessionsFilter struct {
|
||||||
ListenerID string
|
ListenerID string
|
||||||
|
ProjectID string
|
||||||
Status string // active|sleeping|dead|killed;空表示全部
|
Status string // active|sleeping|dead|killed;空表示全部
|
||||||
OS string
|
OS string
|
||||||
Search string // 模糊匹配 hostname/username/internal_ip
|
Search string // 模糊匹配 hostname/username/internal_ip
|
||||||
@@ -550,6 +558,18 @@ func (db *DB) ListC2Sessions(filter ListC2SessionsFilter) ([]*C2Session, error)
|
|||||||
conditions = append(conditions, "listener_id = ?")
|
conditions = append(conditions, "listener_id = ?")
|
||||||
args = append(args, filter.ListenerID)
|
args = append(args, filter.ListenerID)
|
||||||
}
|
}
|
||||||
|
if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound {
|
||||||
|
conditions = append(conditions, `EXISTS (
|
||||||
|
SELECT 1 FROM c2_listeners l
|
||||||
|
WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ''
|
||||||
|
)`)
|
||||||
|
} else if strings.TrimSpace(filter.ProjectID) != "" {
|
||||||
|
conditions = append(conditions, `EXISTS (
|
||||||
|
SELECT 1 FROM c2_listeners l
|
||||||
|
WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ?
|
||||||
|
)`)
|
||||||
|
args = append(args, strings.TrimSpace(filter.ProjectID))
|
||||||
|
}
|
||||||
if filter.Status != "" {
|
if filter.Status != "" {
|
||||||
conditions = append(conditions, "status = ?")
|
conditions = append(conditions, "status = ?")
|
||||||
args = append(args, filter.Status)
|
args = append(args, filter.Status)
|
||||||
@@ -645,6 +665,18 @@ func buildC2SessionsWhere(filter ListC2SessionsFilter) ([]string, []interface{})
|
|||||||
conditions = append(conditions, "listener_id = ?")
|
conditions = append(conditions, "listener_id = ?")
|
||||||
args = append(args, filter.ListenerID)
|
args = append(args, filter.ListenerID)
|
||||||
}
|
}
|
||||||
|
if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound {
|
||||||
|
conditions = append(conditions, `EXISTS (
|
||||||
|
SELECT 1 FROM c2_listeners l
|
||||||
|
WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ''
|
||||||
|
)`)
|
||||||
|
} else if strings.TrimSpace(filter.ProjectID) != "" {
|
||||||
|
conditions = append(conditions, `EXISTS (
|
||||||
|
SELECT 1 FROM c2_listeners l
|
||||||
|
WHERE l.id = c2_sessions.listener_id AND COALESCE(l.project_id, '') = ?
|
||||||
|
)`)
|
||||||
|
args = append(args, strings.TrimSpace(filter.ProjectID))
|
||||||
|
}
|
||||||
if filter.Status != "" {
|
if filter.Status != "" {
|
||||||
conditions = append(conditions, "status = ?")
|
conditions = append(conditions, "status = ?")
|
||||||
args = append(args, filter.Status)
|
args = append(args, filter.Status)
|
||||||
@@ -947,7 +979,10 @@ func (db *DB) GetC2Task(id string) (*C2Task, error) {
|
|||||||
// ListC2TasksFilter 任务过滤
|
// ListC2TasksFilter 任务过滤
|
||||||
type ListC2TasksFilter struct {
|
type ListC2TasksFilter struct {
|
||||||
SessionID string
|
SessionID string
|
||||||
|
ProjectID string
|
||||||
Status string
|
Status string
|
||||||
|
TaskType string
|
||||||
|
Since *time.Time
|
||||||
Limit int
|
Limit int
|
||||||
Offset int
|
Offset int
|
||||||
}
|
}
|
||||||
@@ -959,10 +994,32 @@ func buildC2TasksWhere(filter ListC2TasksFilter) (where string, args []interface
|
|||||||
conditions = append(conditions, "session_id = ?")
|
conditions = append(conditions, "session_id = ?")
|
||||||
args = append(args, filter.SessionID)
|
args = append(args, filter.SessionID)
|
||||||
}
|
}
|
||||||
|
if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound {
|
||||||
|
conditions = append(conditions, `EXISTS (
|
||||||
|
SELECT 1 FROM c2_sessions s
|
||||||
|
JOIN c2_listeners l ON l.id = s.listener_id
|
||||||
|
WHERE s.id = c2_tasks.session_id AND COALESCE(l.project_id, '') = ''
|
||||||
|
)`)
|
||||||
|
} else if strings.TrimSpace(filter.ProjectID) != "" {
|
||||||
|
conditions = append(conditions, `EXISTS (
|
||||||
|
SELECT 1 FROM c2_sessions s
|
||||||
|
JOIN c2_listeners l ON l.id = s.listener_id
|
||||||
|
WHERE s.id = c2_tasks.session_id AND COALESCE(l.project_id, '') = ?
|
||||||
|
)`)
|
||||||
|
args = append(args, strings.TrimSpace(filter.ProjectID))
|
||||||
|
}
|
||||||
if filter.Status != "" {
|
if filter.Status != "" {
|
||||||
conditions = append(conditions, "status = ?")
|
conditions = append(conditions, "status = ?")
|
||||||
args = append(args, filter.Status)
|
args = append(args, filter.Status)
|
||||||
}
|
}
|
||||||
|
if strings.TrimSpace(filter.TaskType) != "" {
|
||||||
|
conditions = append(conditions, "task_type = ?")
|
||||||
|
args = append(args, strings.TrimSpace(filter.TaskType))
|
||||||
|
}
|
||||||
|
if filter.Since != nil {
|
||||||
|
conditions = append(conditions, sqliteEpochGE("created_at", ">="))
|
||||||
|
args = append(args, formatSQLiteUTC(*filter.Since))
|
||||||
|
}
|
||||||
return strings.Join(conditions, " AND "), args
|
return strings.Join(conditions, " AND "), args
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1016,6 +1073,43 @@ func (db *DB) CountC2TasksForAccess(filter ListC2TasksFilter, access RBACListAcc
|
|||||||
return n, err
|
return n, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// CountC2TasksByStatusForAccess 与 ListC2Tasks 相同过滤条件下按状态统计
|
||||||
|
func (db *DB) CountC2TasksByStatusForAccess(filter ListC2TasksFilter, access RBACListAccess) (map[string]int64, error) {
|
||||||
|
where, args := buildC2TasksWhereForAccess(filter, access)
|
||||||
|
query := `SELECT status, COUNT(*) FROM c2_tasks WHERE ` + where + ` GROUP BY status`
|
||||||
|
rows, err := db.Query(query, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
counts := map[string]int64{
|
||||||
|
"queued": 0,
|
||||||
|
"sent": 0,
|
||||||
|
"running": 0,
|
||||||
|
"success": 0,
|
||||||
|
"failed": 0,
|
||||||
|
"cancelled": 0,
|
||||||
|
"pending": 0,
|
||||||
|
}
|
||||||
|
var legacyPending int64
|
||||||
|
for rows.Next() {
|
||||||
|
var status string
|
||||||
|
var n int64
|
||||||
|
if err := rows.Scan(&status, &n); err != nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if status == "pending" {
|
||||||
|
legacyPending = n
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, ok := counts[status]; ok {
|
||||||
|
counts[status] = n
|
||||||
|
}
|
||||||
|
}
|
||||||
|
counts["pending"] = counts["queued"] + counts["sent"] + counts["running"] + legacyPending
|
||||||
|
return counts, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
// CountC2TasksQueuedOrPending 统计 queued/pending 状态任务数(仪表盘「待审任务」)
|
// CountC2TasksQueuedOrPending 统计 queued/pending 状态任务数(仪表盘「待审任务」)
|
||||||
func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
||||||
conditions := []string{"status IN ('queued', 'pending')"}
|
conditions := []string{"status IN ('queued', 'pending')"}
|
||||||
@@ -1030,8 +1124,8 @@ func (db *DB) CountC2TasksQueuedOrPending(sessionID string) (int64, error) {
|
|||||||
return n, err
|
return n, err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID string, access RBACListAccess) (int64, error) {
|
func (db *DB) CountC2TasksQueuedOrPendingForAccess(sessionID, projectID string, access RBACListAccess) (int64, error) {
|
||||||
filter := ListC2TasksFilter{SessionID: sessionID}
|
filter := ListC2TasksFilter{SessionID: sessionID, ProjectID: projectID}
|
||||||
where, args := buildC2TasksWhereForAccess(filter, access)
|
where, args := buildC2TasksWhereForAccess(filter, access)
|
||||||
query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where
|
query := `SELECT COUNT(*) FROM c2_tasks WHERE status IN ('queued', 'pending') AND ` + where
|
||||||
var n int64
|
var n int64
|
||||||
@@ -1412,6 +1506,7 @@ func (db *DB) AppendC2Event(e *C2Event) error {
|
|||||||
type ListC2EventsFilter struct {
|
type ListC2EventsFilter struct {
|
||||||
Level string
|
Level string
|
||||||
Category string
|
Category string
|
||||||
|
ProjectID string
|
||||||
SessionID string
|
SessionID string
|
||||||
TaskID string
|
TaskID string
|
||||||
Since *time.Time
|
Since *time.Time
|
||||||
@@ -1430,6 +1525,49 @@ func buildC2EventsWhere(filter ListC2EventsFilter) (where string, args []interfa
|
|||||||
conditions = append(conditions, "category = ?")
|
conditions = append(conditions, "category = ?")
|
||||||
args = append(args, filter.Category)
|
args = append(args, filter.Category)
|
||||||
}
|
}
|
||||||
|
if strings.TrimSpace(filter.ProjectID) == ProjectFilterUnbound {
|
||||||
|
conditions = append(conditions, `(
|
||||||
|
EXISTS (
|
||||||
|
SELECT 1 FROM c2_sessions s
|
||||||
|
JOIN c2_listeners l ON l.id = s.listener_id
|
||||||
|
WHERE s.id = c2_events.session_id AND COALESCE(l.project_id, '') = ''
|
||||||
|
)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM c2_tasks t
|
||||||
|
JOIN c2_sessions s ON s.id = t.session_id
|
||||||
|
JOIN c2_listeners l ON l.id = s.listener_id
|
||||||
|
WHERE t.id = c2_events.task_id AND COALESCE(l.project_id, '') = ''
|
||||||
|
)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM c2_listeners l
|
||||||
|
WHERE json_valid(c2_events.data_json)
|
||||||
|
AND l.id = json_extract(c2_events.data_json, '$.listener_id')
|
||||||
|
AND COALESCE(l.project_id, '') = ''
|
||||||
|
)
|
||||||
|
)`)
|
||||||
|
} else if strings.TrimSpace(filter.ProjectID) != "" {
|
||||||
|
conditions = append(conditions, `(
|
||||||
|
EXISTS (
|
||||||
|
SELECT 1 FROM c2_sessions s
|
||||||
|
JOIN c2_listeners l ON l.id = s.listener_id
|
||||||
|
WHERE s.id = c2_events.session_id AND COALESCE(l.project_id, '') = ?
|
||||||
|
)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM c2_tasks t
|
||||||
|
JOIN c2_sessions s ON s.id = t.session_id
|
||||||
|
JOIN c2_listeners l ON l.id = s.listener_id
|
||||||
|
WHERE t.id = c2_events.task_id AND COALESCE(l.project_id, '') = ?
|
||||||
|
)
|
||||||
|
OR EXISTS (
|
||||||
|
SELECT 1 FROM c2_listeners l
|
||||||
|
WHERE json_valid(c2_events.data_json)
|
||||||
|
AND l.id = json_extract(c2_events.data_json, '$.listener_id')
|
||||||
|
AND COALESCE(l.project_id, '') = ?
|
||||||
|
)
|
||||||
|
)`)
|
||||||
|
pid := strings.TrimSpace(filter.ProjectID)
|
||||||
|
args = append(args, pid, pid, pid)
|
||||||
|
}
|
||||||
if filter.SessionID != "" {
|
if filter.SessionID != "" {
|
||||||
conditions = append(conditions, "session_id = ?")
|
conditions = append(conditions, "session_id = ?")
|
||||||
args = append(args, filter.SessionID)
|
args = append(args, filter.SessionID)
|
||||||
|
|||||||
@@ -22,6 +22,8 @@ type Conversation struct {
|
|||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
Title string `json:"title"`
|
Title string `json:"title"`
|
||||||
ProjectID string `json:"projectId,omitempty"`
|
ProjectID string `json:"projectId,omitempty"`
|
||||||
|
RoleName string `json:"roleName,omitempty"`
|
||||||
|
AgentMode string `json:"agentMode,omitempty"`
|
||||||
Pinned bool `json:"pinned"`
|
Pinned bool `json:"pinned"`
|
||||||
CreatedAt time.Time `json:"createdAt"`
|
CreatedAt time.Time `json:"createdAt"`
|
||||||
UpdatedAt time.Time `json:"updatedAt"`
|
UpdatedAt time.Time `json:"updatedAt"`
|
||||||
@@ -57,29 +59,31 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
roleName := normalizeConversationRoleName(meta.RoleName)
|
||||||
|
agentMode := normalizeConversationAgentMode(meta.AgentMode)
|
||||||
|
|
||||||
var err error
|
var err error
|
||||||
wsID := strings.TrimSpace(webshellConnectionID)
|
wsID := strings.TrimSpace(webshellConnectionID)
|
||||||
switch {
|
switch {
|
||||||
case wsID != "" && projectID != "":
|
case wsID != "" && projectID != "":
|
||||||
_, err = db.Exec(
|
_, err = db.Exec(
|
||||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id) VALUES (?, ?, ?, ?, ?, ?)",
|
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id, role_name, agent_mode) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
|
||||||
id, title, now, now, wsID, projectID,
|
id, title, now, now, wsID, projectID, roleName, agentMode,
|
||||||
)
|
)
|
||||||
case wsID != "":
|
case wsID != "":
|
||||||
_, err = db.Exec(
|
_, err = db.Exec(
|
||||||
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)",
|
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, role_name, agent_mode) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||||
id, title, now, now, wsID,
|
id, title, now, now, wsID, roleName, agentMode,
|
||||||
)
|
)
|
||||||
case projectID != "":
|
case projectID != "":
|
||||||
_, err = db.Exec(
|
_, err = db.Exec(
|
||||||
"INSERT INTO conversations (id, title, created_at, updated_at, project_id) VALUES (?, ?, ?, ?, ?)",
|
"INSERT INTO conversations (id, title, created_at, updated_at, project_id, role_name, agent_mode) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||||
id, title, now, now, projectID,
|
id, title, now, now, projectID, roleName, agentMode,
|
||||||
)
|
)
|
||||||
default:
|
default:
|
||||||
_, err = db.Exec(
|
_, err = db.Exec(
|
||||||
"INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
|
"INSERT INTO conversations (id, title, created_at, updated_at, role_name, agent_mode) VALUES (?, ?, ?, ?, ?, ?)",
|
||||||
id, title, now, now,
|
id, title, now, now, roleName, agentMode,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -90,6 +94,8 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
|
|||||||
ID: id,
|
ID: id,
|
||||||
Title: title,
|
Title: title,
|
||||||
ProjectID: projectID,
|
ProjectID: projectID,
|
||||||
|
RoleName: roleName,
|
||||||
|
AgentMode: agentMode,
|
||||||
CreatedAt: now,
|
CreatedAt: now,
|
||||||
UpdatedAt: now,
|
UpdatedAt: now,
|
||||||
}
|
}
|
||||||
@@ -236,10 +242,12 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
|
|||||||
var pinned int
|
var pinned int
|
||||||
|
|
||||||
var projectID sql.NullString
|
var projectID sql.NullString
|
||||||
|
var roleName sql.NullString
|
||||||
|
var agentMode sql.NullString
|
||||||
err := db.QueryRow(
|
err := db.QueryRow(
|
||||||
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
|
"SELECT id, title, pinned, created_at, updated_at, project_id, role_name, agent_mode FROM conversations WHERE id = ?",
|
||||||
id,
|
id,
|
||||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
|
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName, &agentMode)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if err == sql.ErrNoRows {
|
if err == sql.ErrNoRows {
|
||||||
return nil, fmt.Errorf("对话不存在")
|
return nil, fmt.Errorf("对话不存在")
|
||||||
@@ -249,6 +257,12 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
|
|||||||
if projectID.Valid {
|
if projectID.Valid {
|
||||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||||
}
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
|
if agentMode.Valid {
|
||||||
|
conv.AgentMode = normalizeConversationAgentMode(agentMode.String)
|
||||||
|
}
|
||||||
|
|
||||||
// 尝试多种时间格式解析
|
// 尝试多种时间格式解析
|
||||||
var err1, err2 error
|
var err1, err2 error
|
||||||
@@ -322,10 +336,12 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
|||||||
var pinned int
|
var pinned int
|
||||||
|
|
||||||
var projectID sql.NullString
|
var projectID sql.NullString
|
||||||
|
var roleName sql.NullString
|
||||||
|
var agentMode sql.NullString
|
||||||
err := db.QueryRow(
|
err := db.QueryRow(
|
||||||
"SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
|
"SELECT id, title, pinned, created_at, updated_at, project_id, role_name, agent_mode FROM conversations WHERE id = ?",
|
||||||
id,
|
id,
|
||||||
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
|
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName, &agentMode)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if err == sql.ErrNoRows {
|
if err == sql.ErrNoRows {
|
||||||
return nil, fmt.Errorf("对话不存在")
|
return nil, fmt.Errorf("对话不存在")
|
||||||
@@ -335,6 +351,12 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
|||||||
if projectID.Valid {
|
if projectID.Valid {
|
||||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||||
}
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
|
if agentMode.Valid {
|
||||||
|
conv.AgentMode = normalizeConversationAgentMode(agentMode.String)
|
||||||
|
}
|
||||||
|
|
||||||
// 尝试多种时间格式解析
|
// 尝试多种时间格式解析
|
||||||
var err1, err2 error
|
var err1, err2 error
|
||||||
@@ -365,6 +387,49 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
|
|||||||
return &conv, nil
|
return &conv, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func normalizeConversationRoleName(roleName string) string {
|
||||||
|
roleName = strings.TrimSpace(roleName)
|
||||||
|
if roleName == "" {
|
||||||
|
return "默认"
|
||||||
|
}
|
||||||
|
return roleName
|
||||||
|
}
|
||||||
|
|
||||||
|
func normalizeConversationAgentMode(agentMode string) string {
|
||||||
|
agentMode = strings.ToLower(strings.TrimSpace(agentMode))
|
||||||
|
agentMode = strings.ReplaceAll(agentMode, "-", "_")
|
||||||
|
switch agentMode {
|
||||||
|
case "deep", "plan_execute", "supervisor":
|
||||||
|
return agentMode
|
||||||
|
default:
|
||||||
|
return "eino_single"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) SetConversationRoleName(id, roleName string) error {
|
||||||
|
roleName = normalizeConversationRoleName(roleName)
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE conversations SET role_name = ?, updated_at = ? WHERE id = ?",
|
||||||
|
roleName, time.Now(), id,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新对话角色失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (db *DB) SetConversationAgentMode(id, agentMode string) error {
|
||||||
|
agentMode = normalizeConversationAgentMode(agentMode)
|
||||||
|
_, err := db.Exec(
|
||||||
|
"UPDATE conversations SET agent_mode = ? WHERE id = ?",
|
||||||
|
agentMode, id,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("更新对话模式失败: %w", err)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
func conversationProjectIDColumn(alias string) string {
|
func conversationProjectIDColumn(alias string) string {
|
||||||
if alias != "" {
|
if alias != "" {
|
||||||
return alias + ".project_id"
|
return alias + ".project_id"
|
||||||
@@ -489,7 +554,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
|||||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err = db.Query(
|
rows, err = db.Query(
|
||||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
|
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name, c.agent_mode
|
||||||
FROM conversations c`+where+`
|
FROM conversations c`+where+`
|
||||||
`+orderClause+`
|
`+orderClause+`
|
||||||
LIMIT ? OFFSET ?`,
|
LIMIT ? OFFSET ?`,
|
||||||
@@ -505,7 +570,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
|||||||
}
|
}
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err = db.Query(
|
rows, err = db.Query(
|
||||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name, agent_mode FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
||||||
args...,
|
args...,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -514,45 +579,7 @@ func (db *DB) ListConversations(limit, offset int, search, sortBy, projectID str
|
|||||||
return nil, fmt.Errorf("查询对话列表失败: %w", err)
|
return nil, fmt.Errorf("查询对话列表失败: %w", err)
|
||||||
}
|
}
|
||||||
defer rows.Close()
|
defer rows.Close()
|
||||||
|
return scanConversationRows(rows)
|
||||||
var conversations []*Conversation
|
|
||||||
for rows.Next() {
|
|
||||||
var conv Conversation
|
|
||||||
var createdAt, updatedAt string
|
|
||||||
var pinned int
|
|
||||||
var projectID sql.NullString
|
|
||||||
|
|
||||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
|
|
||||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
|
||||||
}
|
|
||||||
if projectID.Valid {
|
|
||||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
|
||||||
}
|
|
||||||
|
|
||||||
// 尝试多种时间格式解析
|
|
||||||
var err1, err2 error
|
|
||||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
|
||||||
if err1 != nil {
|
|
||||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
|
||||||
}
|
|
||||||
if err1 != nil {
|
|
||||||
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
|
||||||
}
|
|
||||||
|
|
||||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
|
||||||
if err2 != nil {
|
|
||||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
|
||||||
}
|
|
||||||
if err2 != nil {
|
|
||||||
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
|
||||||
}
|
|
||||||
|
|
||||||
conv.Pinned = pinned != 0
|
|
||||||
|
|
||||||
conversations = append(conversations, &conv)
|
|
||||||
}
|
|
||||||
|
|
||||||
return conversations, nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
||||||
@@ -571,7 +598,7 @@ func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, proj
|
|||||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err = db.Query(
|
rows, err = db.Query(
|
||||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
|
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name, c.agent_mode
|
||||||
FROM conversations c`+where+`
|
FROM conversations c`+where+`
|
||||||
`+orderClause+`
|
`+orderClause+`
|
||||||
LIMIT ? OFFSET ?`, args...)
|
LIMIT ? OFFSET ?`, args...)
|
||||||
@@ -586,7 +613,7 @@ func (db *DB) ListConversationsForAccess(limit, offset int, search, sortBy, proj
|
|||||||
}
|
}
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err = db.Query(
|
rows, err = db.Query(
|
||||||
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name, agent_mode FROM conversations"+where+" "+orderClause+" LIMIT ? OFFSET ?",
|
||||||
args...)
|
args...)
|
||||||
}
|
}
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -603,12 +630,20 @@ func scanConversationRows(rows *sql.Rows) ([]*Conversation, error) {
|
|||||||
var createdAt, updatedAt string
|
var createdAt, updatedAt string
|
||||||
var pinned int
|
var pinned int
|
||||||
var projectID sql.NullString
|
var projectID sql.NullString
|
||||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
|
var roleName sql.NullString
|
||||||
|
var agentMode sql.NullString
|
||||||
|
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID, &roleName, &agentMode); err != nil {
|
||||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
||||||
}
|
}
|
||||||
if projectID.Valid {
|
if projectID.Valid {
|
||||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
conv.ProjectID = strings.TrimSpace(projectID.String)
|
||||||
}
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
|
if agentMode.Valid {
|
||||||
|
conv.AgentMode = normalizeConversationAgentMode(agentMode.String)
|
||||||
|
}
|
||||||
var err1, err2 error
|
var err1, err2 error
|
||||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
||||||
if err1 != nil {
|
if err1 != nil {
|
||||||
@@ -668,7 +703,7 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
|
|||||||
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
where, args = appendConversationProjectFilter(where, args, projectID, "c")
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err := db.Query(
|
rows, err := db.Query(
|
||||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
|
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name, c.agent_mode `+
|
||||||
where+`
|
where+`
|
||||||
`+orderClause+`
|
`+orderClause+`
|
||||||
LIMIT ? OFFSET ?`,
|
LIMIT ? OFFSET ?`,
|
||||||
@@ -678,43 +713,7 @@ func (db *DB) ListUngroupedConversations(limit, offset int, sortBy, projectID st
|
|||||||
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
|
return nil, fmt.Errorf("查询未分组对话失败: %w", err)
|
||||||
}
|
}
|
||||||
defer rows.Close()
|
defer rows.Close()
|
||||||
|
return scanConversationRows(rows)
|
||||||
var conversations []*Conversation
|
|
||||||
for rows.Next() {
|
|
||||||
var conv Conversation
|
|
||||||
var createdAt, updatedAt string
|
|
||||||
var pinned int
|
|
||||||
var projectID sql.NullString
|
|
||||||
|
|
||||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
|
|
||||||
return nil, fmt.Errorf("扫描对话失败: %w", err)
|
|
||||||
}
|
|
||||||
if projectID.Valid {
|
|
||||||
conv.ProjectID = strings.TrimSpace(projectID.String)
|
|
||||||
}
|
|
||||||
|
|
||||||
var err1, err2 error
|
|
||||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05.999999999-07:00", createdAt)
|
|
||||||
if err1 != nil {
|
|
||||||
conv.CreatedAt, err1 = time.Parse("2006-01-02 15:04:05", createdAt)
|
|
||||||
}
|
|
||||||
if err1 != nil {
|
|
||||||
conv.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
|
|
||||||
}
|
|
||||||
|
|
||||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05.999999999-07:00", updatedAt)
|
|
||||||
if err2 != nil {
|
|
||||||
conv.UpdatedAt, err2 = time.Parse("2006-01-02 15:04:05", updatedAt)
|
|
||||||
}
|
|
||||||
if err2 != nil {
|
|
||||||
conv.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
|
|
||||||
}
|
|
||||||
|
|
||||||
conv.Pinned = pinned != 0
|
|
||||||
conversations = append(conversations, &conv)
|
|
||||||
}
|
|
||||||
|
|
||||||
return conversations, rows.Err()
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, projectID, userID, scope string) ([]*Conversation, error) {
|
||||||
@@ -728,7 +727,7 @@ func (db *DB) ListUngroupedConversationsForAccess(limit, offset int, sortBy, pro
|
|||||||
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
where, args = appendConversationAccessFilter(where, args, userID, scope, "c")
|
||||||
args = append(args, limit, offset)
|
args = append(args, limit, offset)
|
||||||
rows, err := db.Query(
|
rows, err := db.Query(
|
||||||
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id `+
|
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id, c.role_name, c.agent_mode `+
|
||||||
where+`
|
where+`
|
||||||
`+orderClause+`
|
`+orderClause+`
|
||||||
LIMIT ? OFFSET ?`,
|
LIMIT ? OFFSET ?`,
|
||||||
@@ -861,6 +860,24 @@ func (db *DB) einoReductionBaseDir() string {
|
|||||||
return filepath.Join("tmp", "reduction")
|
return filepath.Join("tmp", "reduction")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// EinoReductionBaseDir returns the configured reduction cache root.
|
||||||
|
func (db *DB) EinoReductionBaseDir() string {
|
||||||
|
return db.einoReductionBaseDir()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ConversationArtifactsBaseDir returns the conversation-scoped artifacts root.
|
||||||
|
func (db *DB) ConversationArtifactsBaseDir() string {
|
||||||
|
if db == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(db.conversationArtifactsDir)
|
||||||
|
}
|
||||||
|
|
||||||
|
// EinoWorkspaceBaseDir returns the configured agent workspace root.
|
||||||
|
func (db *DB) EinoWorkspaceBaseDir() string {
|
||||||
|
return db.einoWorkspaceBaseDir()
|
||||||
|
}
|
||||||
|
|
||||||
func (db *DB) einoWorkspaceBaseDir() string {
|
func (db *DB) einoWorkspaceBaseDir() string {
|
||||||
if db == nil {
|
if db == nil {
|
||||||
return ""
|
return ""
|
||||||
@@ -1407,6 +1424,7 @@ type ProcessDetailsSummary struct {
|
|||||||
|
|
||||||
type ProcessDetailsToolExecution struct {
|
type ProcessDetailsToolExecution struct {
|
||||||
ProcessDetailID string `json:"processDetailId,omitempty"`
|
ProcessDetailID string `json:"processDetailId,omitempty"`
|
||||||
|
ResultDetailID string `json:"resultDetailId,omitempty"`
|
||||||
ToolName string `json:"toolName,omitempty"`
|
ToolName string `json:"toolName,omitempty"`
|
||||||
ToolCallID string `json:"toolCallId,omitempty"`
|
ToolCallID string `json:"toolCallId,omitempty"`
|
||||||
ExecutionID string `json:"executionId,omitempty"`
|
ExecutionID string `json:"executionId,omitempty"`
|
||||||
@@ -1445,7 +1463,8 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary
|
|||||||
seenExecIDs := make(map[string]bool)
|
seenExecIDs := make(map[string]bool)
|
||||||
// A provider may reuse a fallback toolCallId across streaming rounds. Keep a
|
// A provider may reuse a fallback toolCallId across streaming rounds. Keep a
|
||||||
// FIFO per ID instead of a single index so every persisted call gets at most
|
// FIFO per ID instead of a single index so every persisted call gets at most
|
||||||
// one result. Results without an ID fall back to the oldest unmatched call.
|
// one result. Results without a stable ID are kept separate instead of being
|
||||||
|
// guessed by order; showing no link is safer than linking to the wrong tool.
|
||||||
toolIndexesByCallID := make(map[string][]int)
|
toolIndexesByCallID := make(map[string][]int)
|
||||||
lastMatchedToolIndexByCallID := make(map[string]int)
|
lastMatchedToolIndexByCallID := make(map[string]int)
|
||||||
matchedToolIndexes := make([]bool, 0)
|
matchedToolIndexes := make([]bool, 0)
|
||||||
@@ -1520,7 +1539,7 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if idx < 0 {
|
if idx < 0 && toolCallID != "" {
|
||||||
for nextUnmatchedToolIdx < len(matchedToolIndexes) && matchedToolIndexes[nextUnmatchedToolIdx] {
|
for nextUnmatchedToolIdx < len(matchedToolIndexes) && matchedToolIndexes[nextUnmatchedToolIdx] {
|
||||||
nextUnmatchedToolIdx++
|
nextUnmatchedToolIdx++
|
||||||
}
|
}
|
||||||
@@ -1534,6 +1553,7 @@ func (db *DB) GetProcessDetailsSummary(messageID string) (*ProcessDetailsSummary
|
|||||||
if toolCallID != "" {
|
if toolCallID != "" {
|
||||||
lastMatchedToolIndexByCallID[toolCallID] = idx
|
lastMatchedToolIndexByCallID[toolCallID] = idx
|
||||||
}
|
}
|
||||||
|
summary.ToolExecutions[idx].ResultDetailID = strings.TrimSpace(detailID)
|
||||||
if summary.ToolExecutions[idx].ToolName == "" {
|
if summary.ToolExecutions[idx].ToolName == "" {
|
||||||
summary.ToolExecutions[idx].ToolName = toolName
|
summary.ToolExecutions[idx].ToolName = toolName
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5,6 +5,8 @@ type ConversationCreateMeta struct {
|
|||||||
Source string
|
Source string
|
||||||
WebShellConnectionID string
|
WebShellConnectionID string
|
||||||
ProjectID string
|
ProjectID string
|
||||||
|
RoleName string
|
||||||
|
AgentMode string
|
||||||
ClientIP string
|
ClientIP string
|
||||||
SessionHint string
|
SessionHint string
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -183,6 +183,8 @@ func (db *DB) initTables() error {
|
|||||||
title TEXT NOT NULL,
|
title TEXT NOT NULL,
|
||||||
created_at DATETIME NOT NULL,
|
created_at DATETIME NOT NULL,
|
||||||
updated_at DATETIME NOT NULL,
|
updated_at DATETIME NOT NULL,
|
||||||
|
role_name TEXT NOT NULL DEFAULT '默认',
|
||||||
|
agent_mode TEXT NOT NULL DEFAULT 'eino_single',
|
||||||
last_react_input TEXT,
|
last_react_input TEXT,
|
||||||
last_react_output TEXT
|
last_react_output TEXT
|
||||||
);`
|
);`
|
||||||
@@ -226,6 +228,10 @@ func (db *DB) initTables() error {
|
|||||||
start_time DATETIME NOT NULL,
|
start_time DATETIME NOT NULL,
|
||||||
end_time DATETIME,
|
end_time DATETIME,
|
||||||
duration_ms INTEGER,
|
duration_ms INTEGER,
|
||||||
|
partial_output TEXT,
|
||||||
|
partial_output_bytes INTEGER NOT NULL DEFAULT 0,
|
||||||
|
partial_output_truncated INTEGER NOT NULL DEFAULT 0,
|
||||||
|
partial_output_updated_at DATETIME,
|
||||||
owner_user_id TEXT,
|
owner_user_id TEXT,
|
||||||
conversation_id TEXT,
|
conversation_id TEXT,
|
||||||
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||||
@@ -415,6 +421,7 @@ func (db *DB) initTables() error {
|
|||||||
responsible_person TEXT NOT NULL DEFAULT '', department TEXT NOT NULL DEFAULT '', business_system TEXT NOT NULL DEFAULT '',
|
responsible_person TEXT NOT NULL DEFAULT '', department TEXT NOT NULL DEFAULT '', business_system TEXT NOT NULL DEFAULT '',
|
||||||
environment TEXT NOT NULL DEFAULT '', criticality TEXT NOT NULL DEFAULT '',
|
environment TEXT NOT NULL DEFAULT '', criticality TEXT NOT NULL DEFAULT '',
|
||||||
source TEXT NOT NULL DEFAULT 'manual', source_query TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'active',
|
source TEXT NOT NULL DEFAULT 'manual', source_query TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'active',
|
||||||
|
vulnerability_count INTEGER NOT NULL DEFAULT 0, risk_score INTEGER NOT NULL DEFAULT 0, risk_level TEXT NOT NULL DEFAULT 'unassessed',
|
||||||
tags_json TEXT NOT NULL DEFAULT '[]', first_seen_at DATETIME NOT NULL, last_seen_at DATETIME NOT NULL,
|
tags_json TEXT NOT NULL DEFAULT '[]', first_seen_at DATETIME NOT NULL, last_seen_at DATETIME NOT NULL,
|
||||||
created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, owner_user_id TEXT,
|
created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, owner_user_id TEXT,
|
||||||
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL
|
FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL
|
||||||
@@ -489,6 +496,7 @@ func (db *DB) initTables() error {
|
|||||||
createWebshellConnectionsTable := `
|
createWebshellConnectionsTable := `
|
||||||
CREATE TABLE IF NOT EXISTS webshell_connections (
|
CREATE TABLE IF NOT EXISTS webshell_connections (
|
||||||
id TEXT PRIMARY KEY,
|
id TEXT PRIMARY KEY,
|
||||||
|
project_id TEXT,
|
||||||
url TEXT NOT NULL,
|
url TEXT NOT NULL,
|
||||||
password TEXT NOT NULL DEFAULT '',
|
password TEXT NOT NULL DEFAULT '',
|
||||||
type TEXT NOT NULL DEFAULT 'php',
|
type TEXT NOT NULL DEFAULT 'php',
|
||||||
@@ -515,6 +523,7 @@ func (db *DB) initTables() error {
|
|||||||
createC2ListenersTable := `
|
createC2ListenersTable := `
|
||||||
CREATE TABLE IF NOT EXISTS c2_listeners (
|
CREATE TABLE IF NOT EXISTS c2_listeners (
|
||||||
id TEXT PRIMARY KEY,
|
id TEXT PRIMARY KEY,
|
||||||
|
project_id TEXT,
|
||||||
name TEXT NOT NULL,
|
name TEXT NOT NULL,
|
||||||
type TEXT NOT NULL,
|
type TEXT NOT NULL,
|
||||||
bind_host TEXT NOT NULL DEFAULT '127.0.0.1',
|
bind_host TEXT NOT NULL DEFAULT '127.0.0.1',
|
||||||
@@ -737,6 +746,9 @@ func (db *DB) initTables() error {
|
|||||||
CREATE INDEX IF NOT EXISTS idx_assets_status ON assets(status);
|
CREATE INDEX IF NOT EXISTS idx_assets_status ON assets(status);
|
||||||
CREATE INDEX IF NOT EXISTS idx_assets_owner ON assets(owner_user_id);
|
CREATE INDEX IF NOT EXISTS idx_assets_owner ON assets(owner_user_id);
|
||||||
CREATE INDEX IF NOT EXISTS idx_assets_project ON assets(project_id);
|
CREATE INDEX IF NOT EXISTS idx_assets_project ON assets(project_id);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_assets_vulnerability_count ON assets(vulnerability_count);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_assets_risk_score ON assets(risk_score);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_assets_risk_level ON assets(risk_level);
|
||||||
CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status);
|
CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status);
|
||||||
CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at);
|
CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at);
|
||||||
CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id);
|
CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id);
|
||||||
@@ -751,8 +763,10 @@ func (db *DB) initTables() error {
|
|||||||
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_created_at ON batch_task_queues(created_at);
|
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_created_at ON batch_task_queues(created_at);
|
||||||
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_title ON batch_task_queues(title);
|
CREATE INDEX IF NOT EXISTS idx_batch_task_queues_title ON batch_task_queues(title);
|
||||||
CREATE INDEX IF NOT EXISTS idx_webshell_connections_created_at ON webshell_connections(created_at);
|
CREATE INDEX IF NOT EXISTS idx_webshell_connections_created_at ON webshell_connections(created_at);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_webshell_connections_project_id ON webshell_connections(project_id);
|
||||||
CREATE INDEX IF NOT EXISTS idx_webshell_connection_states_updated_at ON webshell_connection_states(updated_at);
|
CREATE INDEX IF NOT EXISTS idx_webshell_connection_states_updated_at ON webshell_connection_states(updated_at);
|
||||||
CREATE INDEX IF NOT EXISTS idx_c2_listeners_created_at ON c2_listeners(created_at);
|
CREATE INDEX IF NOT EXISTS idx_c2_listeners_created_at ON c2_listeners(created_at);
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_c2_listeners_project_id ON c2_listeners(project_id);
|
||||||
CREATE INDEX IF NOT EXISTS idx_c2_listeners_status ON c2_listeners(status);
|
CREATE INDEX IF NOT EXISTS idx_c2_listeners_status ON c2_listeners(status);
|
||||||
CREATE INDEX IF NOT EXISTS idx_c2_sessions_listener ON c2_sessions(listener_id);
|
CREATE INDEX IF NOT EXISTS idx_c2_sessions_listener ON c2_sessions(listener_id);
|
||||||
CREATE INDEX IF NOT EXISTS idx_c2_sessions_status ON c2_sessions(status);
|
CREATE INDEX IF NOT EXISTS idx_c2_sessions_status ON c2_sessions(status);
|
||||||
@@ -951,9 +965,15 @@ func (db *DB) initTables() error {
|
|||||||
db.logger.Warn("迁移webshell_connections表失败", zap.Error(err))
|
db.logger.Warn("迁移webshell_connections表失败", zap.Error(err))
|
||||||
// 不返回错误,允许继续运行
|
// 不返回错误,允许继续运行
|
||||||
}
|
}
|
||||||
|
if err := db.migrateC2ListenersTable(); err != nil {
|
||||||
|
db.logger.Warn("迁移c2_listeners表失败", zap.Error(err))
|
||||||
|
}
|
||||||
if err := db.migrateWorkflowRunsTable(); err != nil {
|
if err := db.migrateWorkflowRunsTable(); err != nil {
|
||||||
db.logger.Warn("迁移workflow_runs表失败", zap.Error(err))
|
db.logger.Warn("迁移workflow_runs表失败", zap.Error(err))
|
||||||
}
|
}
|
||||||
|
if err := db.migrateToolExecutionsPartialOutputColumns(); err != nil {
|
||||||
|
db.logger.Warn("迁移tool_executions partial output字段失败", zap.Error(err))
|
||||||
|
}
|
||||||
if err := db.migrateRBACOwnershipColumns(); err != nil {
|
if err := db.migrateRBACOwnershipColumns(); err != nil {
|
||||||
db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err))
|
db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err))
|
||||||
}
|
}
|
||||||
@@ -961,7 +981,6 @@ func (db *DB) initTables() error {
|
|||||||
if _, err := db.Exec(createIndexes); err != nil {
|
if _, err := db.Exec(createIndexes); err != nil {
|
||||||
return fmt.Errorf("创建索引失败: %w", err)
|
return fmt.Errorf("创建索引失败: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
db.logger.Debug("数据库表初始化完成")
|
db.logger.Debug("数据库表初始化完成")
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -978,6 +997,23 @@ func (db *DB) migrateRobotUserSessionsTable() error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (db *DB) migrateToolExecutionsPartialOutputColumns() error {
|
||||||
|
for _, col := range []struct {
|
||||||
|
name string
|
||||||
|
stmt string
|
||||||
|
}{
|
||||||
|
{"partial_output", "ALTER TABLE tool_executions ADD COLUMN partial_output TEXT"},
|
||||||
|
{"partial_output_bytes", "ALTER TABLE tool_executions ADD COLUMN partial_output_bytes INTEGER NOT NULL DEFAULT 0"},
|
||||||
|
{"partial_output_truncated", "ALTER TABLE tool_executions ADD COLUMN partial_output_truncated INTEGER NOT NULL DEFAULT 0"},
|
||||||
|
{"partial_output_updated_at", "ALTER TABLE tool_executions ADD COLUMN partial_output_updated_at DATETIME"},
|
||||||
|
} {
|
||||||
|
if err := db.addColumnIfMissing("tool_executions", col.name, col.stmt); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
// migrateAssetsTable keeps databases created by the first asset-management release compatible.
|
// migrateAssetsTable keeps databases created by the first asset-management release compatible.
|
||||||
func (db *DB) migrateAssetsTable() error {
|
func (db *DB) migrateAssetsTable() error {
|
||||||
columns := []struct {
|
columns := []struct {
|
||||||
@@ -994,6 +1030,9 @@ func (db *DB) migrateAssetsTable() error {
|
|||||||
{"business_system", "ALTER TABLE assets ADD COLUMN business_system TEXT NOT NULL DEFAULT ''"},
|
{"business_system", "ALTER TABLE assets ADD COLUMN business_system TEXT NOT NULL DEFAULT ''"},
|
||||||
{"environment", "ALTER TABLE assets ADD COLUMN environment TEXT NOT NULL DEFAULT ''"},
|
{"environment", "ALTER TABLE assets ADD COLUMN environment TEXT NOT NULL DEFAULT ''"},
|
||||||
{"criticality", "ALTER TABLE assets ADD COLUMN criticality TEXT NOT NULL DEFAULT ''"},
|
{"criticality", "ALTER TABLE assets ADD COLUMN criticality TEXT NOT NULL DEFAULT ''"},
|
||||||
|
{"vulnerability_count", "ALTER TABLE assets ADD COLUMN vulnerability_count INTEGER NOT NULL DEFAULT 0"},
|
||||||
|
{"risk_score", "ALTER TABLE assets ADD COLUMN risk_score INTEGER NOT NULL DEFAULT 0"},
|
||||||
|
{"risk_level", "ALTER TABLE assets ADD COLUMN risk_level TEXT NOT NULL DEFAULT 'unassessed'"},
|
||||||
}
|
}
|
||||||
for _, column := range columns {
|
for _, column := range columns {
|
||||||
var count int
|
var count int
|
||||||
@@ -1127,6 +1166,36 @@ func (db *DB) migrateConversationsTable() error {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 检查 role_name 字段是否存在(对话绑定的业务角色,用于历史任务切换时恢复角色上下文)
|
||||||
|
err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='role_name'").Scan(&count)
|
||||||
|
if err != nil {
|
||||||
|
if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN role_name TEXT NOT NULL DEFAULT '默认'"); addErr != nil {
|
||||||
|
errMsg := strings.ToLower(addErr.Error())
|
||||||
|
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||||
|
db.logger.Warn("添加role_name字段失败", zap.Error(addErr))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if count == 0 {
|
||||||
|
if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN role_name TEXT NOT NULL DEFAULT '默认'"); err != nil {
|
||||||
|
db.logger.Warn("添加role_name字段失败", zap.Error(err))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查 agent_mode 字段是否存在(对话绑定的执行模式,用于历史任务切换时恢复对话模式)
|
||||||
|
err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='agent_mode'").Scan(&count)
|
||||||
|
if err != nil {
|
||||||
|
if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'"); addErr != nil {
|
||||||
|
errMsg := strings.ToLower(addErr.Error())
|
||||||
|
if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") {
|
||||||
|
db.logger.Warn("添加agent_mode字段失败", zap.Error(addErr))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else if count == 0 {
|
||||||
|
if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'"); err != nil {
|
||||||
|
db.logger.Warn("添加agent_mode字段失败", zap.Error(err))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1570,6 +1639,7 @@ func (db *DB) migrateWebshellConnectionsTable() error {
|
|||||||
name string
|
name string
|
||||||
stmt string
|
stmt string
|
||||||
}{
|
}{
|
||||||
|
{name: "project_id", stmt: "ALTER TABLE webshell_connections ADD COLUMN project_id TEXT"},
|
||||||
{name: "encoding", stmt: "ALTER TABLE webshell_connections ADD COLUMN encoding TEXT NOT NULL DEFAULT ''"},
|
{name: "encoding", stmt: "ALTER TABLE webshell_connections ADD COLUMN encoding TEXT NOT NULL DEFAULT ''"},
|
||||||
{name: "os", stmt: "ALTER TABLE webshell_connections ADD COLUMN os TEXT NOT NULL DEFAULT ''"},
|
{name: "os", stmt: "ALTER TABLE webshell_connections ADD COLUMN os TEXT NOT NULL DEFAULT ''"},
|
||||||
}
|
}
|
||||||
@@ -1595,6 +1665,10 @@ func (db *DB) migrateWebshellConnectionsTable() error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (db *DB) migrateC2ListenersTable() error {
|
||||||
|
return db.addColumnIfMissing("c2_listeners", "project_id", "ALTER TABLE c2_listeners ADD COLUMN project_id TEXT")
|
||||||
|
}
|
||||||
|
|
||||||
// NewKnowledgeDB 创建知识库数据库连接(只包含知识库相关的表)
|
// NewKnowledgeDB 创建知识库数据库连接(只包含知识库相关的表)
|
||||||
func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) {
|
||||||
sqlDB, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL")
|
sqlDB, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL")
|
||||||
|
|||||||
@@ -43,11 +43,19 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
|
|||||||
if exec.Duration > 0 {
|
if exec.Duration > 0 {
|
||||||
durationMs = sql.NullInt64{Int64: exec.Duration.Milliseconds(), Valid: true}
|
durationMs = sql.NullInt64{Int64: exec.Duration.Milliseconds(), Valid: true}
|
||||||
}
|
}
|
||||||
|
var partialUpdatedAt sql.NullTime
|
||||||
|
if exec.PartialOutputUpdatedAt != nil {
|
||||||
|
partialUpdatedAt = sql.NullTime{Time: *exec.PartialOutputUpdatedAt, Valid: true}
|
||||||
|
}
|
||||||
|
partialTruncated := 0
|
||||||
|
if exec.PartialOutputTruncated {
|
||||||
|
partialTruncated = 1
|
||||||
|
}
|
||||||
|
|
||||||
query := `
|
query := `
|
||||||
INSERT OR REPLACE INTO tool_executions
|
INSERT OR REPLACE INTO tool_executions
|
||||||
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, owner_user_id, conversation_id, created_at)
|
(id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, partial_output, partial_output_bytes, partial_output_truncated, partial_output_updated_at, owner_user_id, conversation_id, created_at)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
`
|
`
|
||||||
|
|
||||||
_, err = db.Exec(query,
|
_, err = db.Exec(query,
|
||||||
@@ -60,6 +68,10 @@ func (db *DB) SaveToolExecution(exec *mcp.ToolExecution) error {
|
|||||||
exec.StartTime,
|
exec.StartTime,
|
||||||
endTime,
|
endTime,
|
||||||
durationMs,
|
durationMs,
|
||||||
|
sqlNullString(exec.PartialOutput),
|
||||||
|
exec.PartialOutputBytes,
|
||||||
|
partialTruncated,
|
||||||
|
partialUpdatedAt,
|
||||||
strings.TrimSpace(exec.OwnerUserID),
|
strings.TrimSpace(exec.OwnerUserID),
|
||||||
strings.TrimSpace(exec.ConversationID),
|
strings.TrimSpace(exec.ConversationID),
|
||||||
time.Now(),
|
time.Now(),
|
||||||
@@ -90,6 +102,13 @@ func (db *DB) UpdateToolExecutionResult(id string, result *mcp.ToolResult) error
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func sqlNullString(s string) sql.NullString {
|
||||||
|
if s == "" {
|
||||||
|
return sql.NullString{}
|
||||||
|
}
|
||||||
|
return sql.NullString{String: s, Valid: true}
|
||||||
|
}
|
||||||
|
|
||||||
// CountToolExecutions 统计工具执行记录总数
|
// CountToolExecutions 统计工具执行记录总数
|
||||||
func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
|
func (db *DB) CountToolExecutions(status, toolName string) (int, error) {
|
||||||
return db.CountToolExecutionsForAccess(status, toolName, RBACListAccess{Scope: RBACScopeAll})
|
return db.CountToolExecutionsForAccess(status, toolName, RBACListAccess{Scope: RBACScopeAll})
|
||||||
@@ -267,7 +286,8 @@ type ToolStatsSummaryResult struct {
|
|||||||
TopTools []*mcp.ToolStats
|
TopTools []*mcp.ToolStats
|
||||||
}
|
}
|
||||||
|
|
||||||
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)
|
// LoadToolStatsSummary 聚合统计信息,仅返回汇总与 Top N 工具(避免全量 map 传输)。
|
||||||
|
// 监控页的失败口径只包含真实失败/异常终止;用户主动取消的 cancelled 保留在总调用中,不计入失败。
|
||||||
func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
||||||
if topN <= 0 {
|
if topN <= 0 {
|
||||||
topN = 6
|
topN = 6
|
||||||
@@ -282,19 +302,19 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
|||||||
|
|
||||||
summaryQuery := `
|
summaryQuery := `
|
||||||
SELECT COUNT(*),
|
SELECT COUNT(*),
|
||||||
COALESCE(SUM(total_calls), 0),
|
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
||||||
COALESCE(SUM(success_calls), 0),
|
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
|
||||||
COALESCE(SUM(failed_calls), 0),
|
MAX(start_time),
|
||||||
MAX(last_call_time)
|
COUNT(DISTINCT tool_name)
|
||||||
FROM tool_stats
|
FROM tool_executions
|
||||||
`
|
`
|
||||||
var lastCallRaw sql.NullString
|
var lastCallRaw sql.NullString
|
||||||
err := db.QueryRow(summaryQuery).Scan(
|
err := db.QueryRow(summaryQuery).Scan(
|
||||||
&result.Summary.ToolCount,
|
|
||||||
&result.Summary.TotalCalls,
|
&result.Summary.TotalCalls,
|
||||||
&result.Summary.SuccessCalls,
|
&result.Summary.SuccessCalls,
|
||||||
&result.Summary.FailedCalls,
|
&result.Summary.FailedCalls,
|
||||||
&lastCallRaw,
|
&lastCallRaw,
|
||||||
|
&result.Summary.ToolCount,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
@@ -310,9 +330,13 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
topQuery := `
|
topQuery := `
|
||||||
SELECT tool_name, total_calls, success_calls, failed_calls, last_call_time
|
SELECT tool_name,
|
||||||
FROM tool_stats
|
COUNT(*) AS total_calls,
|
||||||
WHERE total_calls > 0
|
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) AS success_calls,
|
||||||
|
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed_calls,
|
||||||
|
MAX(start_time) AS last_call_time
|
||||||
|
FROM tool_executions
|
||||||
|
GROUP BY tool_name
|
||||||
ORDER BY total_calls DESC, tool_name ASC
|
ORDER BY total_calls DESC, tool_name ASC
|
||||||
LIMIT ?
|
LIMIT ?
|
||||||
`
|
`
|
||||||
@@ -324,7 +348,7 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
|||||||
|
|
||||||
for rows.Next() {
|
for rows.Next() {
|
||||||
var stat mcp.ToolStats
|
var stat mcp.ToolStats
|
||||||
var lastCallTime sql.NullTime
|
var lastCallTime sql.NullString
|
||||||
if err := rows.Scan(
|
if err := rows.Scan(
|
||||||
&stat.ToolName,
|
&stat.ToolName,
|
||||||
&stat.TotalCalls,
|
&stat.TotalCalls,
|
||||||
@@ -336,7 +360,8 @@ func (db *DB) LoadToolStatsSummary(topN int) (*ToolStatsSummaryResult, error) {
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
if lastCallTime.Valid {
|
if lastCallTime.Valid {
|
||||||
stat.LastCallTime = &lastCallTime.Time
|
parsed := parseDBTime(lastCallTime.String)
|
||||||
|
stat.LastCallTime = &parsed
|
||||||
}
|
}
|
||||||
result.TopTools = append(result.TopTools, &stat)
|
result.TopTools = append(result.TopTools, &stat)
|
||||||
}
|
}
|
||||||
@@ -359,7 +384,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
|
|||||||
var lastCall sql.NullString
|
var lastCall sql.NullString
|
||||||
err := db.QueryRow(`SELECT COUNT(*),
|
err := db.QueryRow(`SELECT COUNT(*),
|
||||||
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
||||||
COALESCE(SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
|
COALESCE(SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), 0),
|
||||||
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
|
MAX(start_time), COUNT(DISTINCT tool_name)`+fromSQL, args...).Scan(
|
||||||
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
|
&result.Summary.TotalCalls, &result.Summary.SuccessCalls, &result.Summary.FailedCalls,
|
||||||
&lastCall, &result.Summary.ToolCount,
|
&lastCall, &result.Summary.ToolCount,
|
||||||
@@ -373,7 +398,7 @@ func (db *DB) LoadToolStatsSummaryForAccess(topN int, access RBACListAccess) (*T
|
|||||||
}
|
}
|
||||||
rows, err := db.Query(`SELECT tool_name, COUNT(*),
|
rows, err := db.Query(`SELECT tool_name, COUNT(*),
|
||||||
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
|
SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
|
||||||
SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+
|
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END), MAX(start_time)`+
|
||||||
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
|
fromSQL+` GROUP BY tool_name ORDER BY COUNT(*) DESC, tool_name ASC LIMIT ?`, append(args, topN)...)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
@@ -482,7 +507,9 @@ func appendToolExecutionAccessSQL(query string, args []interface{}, access RBACL
|
|||||||
// GetToolExecution 根据ID获取单条工具执行记录
|
// GetToolExecution 根据ID获取单条工具执行记录
|
||||||
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
||||||
query := `
|
query := `
|
||||||
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms, COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
|
SELECT id, tool_name, arguments, status, result, error, start_time, end_time, duration_ms,
|
||||||
|
COALESCE(partial_output, ''), COALESCE(partial_output_bytes, 0), COALESCE(partial_output_truncated, 0), partial_output_updated_at,
|
||||||
|
COALESCE(owner_user_id, ''), COALESCE(conversation_id, '')
|
||||||
FROM tool_executions
|
FROM tool_executions
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
`
|
`
|
||||||
@@ -495,6 +522,8 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
|||||||
var errorText sql.NullString
|
var errorText sql.NullString
|
||||||
var endTime sql.NullTime
|
var endTime sql.NullTime
|
||||||
var durationMs sql.NullInt64
|
var durationMs sql.NullInt64
|
||||||
|
var partialTruncated int
|
||||||
|
var partialUpdatedAt sql.NullTime
|
||||||
|
|
||||||
err := row.Scan(
|
err := row.Scan(
|
||||||
&exec.ID,
|
&exec.ID,
|
||||||
@@ -506,6 +535,10 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
|||||||
&exec.StartTime,
|
&exec.StartTime,
|
||||||
&endTime,
|
&endTime,
|
||||||
&durationMs,
|
&durationMs,
|
||||||
|
&exec.PartialOutput,
|
||||||
|
&exec.PartialOutputBytes,
|
||||||
|
&partialTruncated,
|
||||||
|
&partialUpdatedAt,
|
||||||
&exec.OwnerUserID,
|
&exec.OwnerUserID,
|
||||||
&exec.ConversationID,
|
&exec.ConversationID,
|
||||||
)
|
)
|
||||||
@@ -538,6 +571,10 @@ func (db *DB) GetToolExecution(id string) (*mcp.ToolExecution, error) {
|
|||||||
if durationMs.Valid {
|
if durationMs.Valid {
|
||||||
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
|
exec.Duration = time.Duration(durationMs.Int64) * time.Millisecond
|
||||||
}
|
}
|
||||||
|
exec.PartialOutputTruncated = partialTruncated != 0
|
||||||
|
if partialUpdatedAt.Valid {
|
||||||
|
exec.PartialOutputUpdatedAt = &partialUpdatedAt.Time
|
||||||
|
}
|
||||||
|
|
||||||
return &exec, nil
|
return &exec, nil
|
||||||
}
|
}
|
||||||
@@ -815,7 +852,7 @@ func (db *DB) PurgeToolExecutionsBefore(cutoff time.Time) (int64, error) {
|
|||||||
}
|
}
|
||||||
delta.totalCalls += count
|
delta.totalCalls += count
|
||||||
switch status {
|
switch status {
|
||||||
case "failed", "cancelled", "hard_timeout", "orphaned":
|
case "failed", "hard_timeout", "orphaned":
|
||||||
delta.failedCalls += count
|
delta.failedCalls += count
|
||||||
case "completed":
|
case "completed":
|
||||||
delta.successCalls += count
|
delta.successCalls += count
|
||||||
@@ -971,7 +1008,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
|
|||||||
query = `
|
query = `
|
||||||
SELECT date(start_time, 'localtime') AS bucket,
|
SELECT date(start_time, 'localtime') AS bucket,
|
||||||
COUNT(*) AS total,
|
COUNT(*) AS total,
|
||||||
SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
|
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
|
||||||
FROM tool_executions
|
FROM tool_executions
|
||||||
WHERE start_time >= ?
|
WHERE start_time >= ?
|
||||||
GROUP BY bucket
|
GROUP BY bucket
|
||||||
@@ -981,7 +1018,7 @@ func (db *DB) LoadCallsTimeline(since time.Time, dailyBuckets bool) ([]CallsTime
|
|||||||
query = `
|
query = `
|
||||||
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
|
SELECT strftime('%Y-%m-%d %H:00:00', start_time, 'localtime') AS bucket,
|
||||||
COUNT(*) AS total,
|
COUNT(*) AS total,
|
||||||
SUM(CASE WHEN status IN ('failed', 'cancelled', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
|
SUM(CASE WHEN status IN ('failed', 'hard_timeout', 'orphaned') THEN 1 ELSE 0 END) AS failed
|
||||||
FROM tool_executions
|
FROM tool_executions
|
||||||
WHERE start_time >= ?
|
WHERE start_time >= ?
|
||||||
GROUP BY bucket
|
GROUP BY bucket
|
||||||
|
|||||||
@@ -84,3 +84,49 @@ func TestLoadToolStatsSummaryAndListPage(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestLoadToolStatsSummaryDoesNotCountCancelledAsFailed(t *testing.T) {
|
||||||
|
dbPath := filepath.Join(t.TempDir(), "monitor-cancelled-summary.db")
|
||||||
|
db, err := NewDB(dbPath, zap.NewNop())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("NewDB: %v", err)
|
||||||
|
}
|
||||||
|
defer db.Close()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
for i, status := range []string{"completed", "cancelled", "failed"} {
|
||||||
|
exec := &mcp.ToolExecution{
|
||||||
|
ID: fmt.Sprintf("exec-%d", i),
|
||||||
|
ToolName: "exec",
|
||||||
|
Arguments: map[string]interface{}{},
|
||||||
|
Status: status,
|
||||||
|
StartTime: now.Add(time.Duration(i) * time.Second),
|
||||||
|
}
|
||||||
|
end := exec.StartTime.Add(time.Second)
|
||||||
|
exec.EndTime = &end
|
||||||
|
exec.Duration = time.Second
|
||||||
|
if err := db.SaveToolExecution(exec); err != nil {
|
||||||
|
t.Fatalf("SaveToolExecution(%s): %v", status, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, err := db.LoadToolStatsSummary(1)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("LoadToolStatsSummary: %v", err)
|
||||||
|
}
|
||||||
|
if summary.Summary.TotalCalls != 3 {
|
||||||
|
t.Fatalf("totalCalls = %d, want 3", summary.Summary.TotalCalls)
|
||||||
|
}
|
||||||
|
if summary.Summary.SuccessCalls != 1 {
|
||||||
|
t.Fatalf("successCalls = %d, want 1", summary.Summary.SuccessCalls)
|
||||||
|
}
|
||||||
|
if summary.Summary.FailedCalls != 1 {
|
||||||
|
t.Fatalf("failedCalls = %d, want 1", summary.Summary.FailedCalls)
|
||||||
|
}
|
||||||
|
if len(summary.TopTools) != 1 {
|
||||||
|
t.Fatalf("top tools = %d, want 1", len(summary.TopTools))
|
||||||
|
}
|
||||||
|
if summary.TopTools[0].FailedCalls != 1 {
|
||||||
|
t.Fatalf("top tool failedCalls = %d, want 1", summary.TopTools[0].FailedCalls)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ import (
|
|||||||
"go.uber.org/zap"
|
"go.uber.org/zap"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestProcessDetailsSummaryPairsMixedIdentifiedAndIDLessResults(t *testing.T) {
|
func TestProcessDetailsSummaryDoesNotGuessIDLessResultsByOrder(t *testing.T) {
|
||||||
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
db, conversationID, messageID := setupProcessDetailsSummaryTest(t)
|
||||||
for _, id := range []string{"call-1", "call-2", "call-3", "call-4"} {
|
for _, id := range []string{"call-1", "call-2", "call-3", "call-4"} {
|
||||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
if err := db.AddProcessDetail(messageID, conversationID, "tool_call", "call", map[string]interface{}{
|
||||||
@@ -22,23 +22,39 @@ func TestProcessDetailsSummaryPairsMixedIdentifiedAndIDLessResults(t *testing.T)
|
|||||||
{"toolName": "http-framework-test", "success": true},
|
{"toolName": "http-framework-test", "success": true},
|
||||||
{"toolName": "http-framework-test", "success": true},
|
{"toolName": "http-framework-test", "success": true},
|
||||||
}
|
}
|
||||||
|
var resultIDs []string
|
||||||
for _, result := range results {
|
for _, result := range results {
|
||||||
if err := db.AddProcessDetail(messageID, conversationID, "tool_result", "result", result); err != nil {
|
resultID, err := db.AddProcessDetailWithID(messageID, conversationID, "tool_result", "result", result)
|
||||||
|
if err != nil {
|
||||||
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
t.Fatalf("AddProcessDetail(tool_result): %v", err)
|
||||||
}
|
}
|
||||||
|
resultIDs = append(resultIDs, resultID)
|
||||||
}
|
}
|
||||||
|
|
||||||
summary, err := db.GetProcessDetailsSummary(messageID)
|
summary, err := db.GetProcessDetailsSummary(messageID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
t.Fatalf("GetProcessDetailsSummary: %v", err)
|
||||||
}
|
}
|
||||||
if len(summary.ToolExecutions) != 4 {
|
if len(summary.ToolExecutions) != 6 {
|
||||||
t.Fatalf("tool executions = %d, want 4", len(summary.ToolExecutions))
|
t.Fatalf("tool executions = %d, want 6", len(summary.ToolExecutions))
|
||||||
}
|
}
|
||||||
for i, execution := range summary.ToolExecutions {
|
for i, execution := range summary.ToolExecutions[:2] {
|
||||||
if execution.Status != "completed" {
|
if execution.Status != "completed" {
|
||||||
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
|
t.Fatalf("execution %d status = %q, want completed", i, execution.Status)
|
||||||
}
|
}
|
||||||
|
if execution.ResultDetailID != resultIDs[i] {
|
||||||
|
t.Fatalf("execution %d result detail id = %q, want %q", i, execution.ResultDetailID, resultIDs[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for i, execution := range summary.ToolExecutions[2:4] {
|
||||||
|
if execution.Status != "result_missing" {
|
||||||
|
t.Fatalf("unmatched call %d status = %q, want result_missing", i, execution.Status)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for i, execution := range summary.ToolExecutions[4:] {
|
||||||
|
if execution.Status != "completed" || execution.ToolCallID != "" {
|
||||||
|
t.Fatalf("idless result %d = %#v, want separate completed result without toolCallId", i, execution)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ import (
|
|||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
)
|
)
|
||||||
|
|
||||||
var factKeyPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9._/-]*$`)
|
var factKeyPattern = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9._/-]*$`)
|
||||||
|
|
||||||
// ValidateFactKey 校验事实 key(项目内唯一标识)。
|
// ValidateFactKey 校验事实 key(项目内唯一标识)。
|
||||||
func ValidateFactKey(key string) error {
|
func ValidateFactKey(key string) error {
|
||||||
@@ -22,7 +22,7 @@ func ValidateFactKey(key string) error {
|
|||||||
return fmt.Errorf("fact_key 过长(最多 128 字符)")
|
return fmt.Errorf("fact_key 过长(最多 128 字符)")
|
||||||
}
|
}
|
||||||
if !factKeyPattern.MatchString(key) {
|
if !factKeyPattern.MatchString(key) {
|
||||||
return fmt.Errorf("fact_key 格式无效,仅允许小写字母、数字及 . _ / -,且须以小写字母或数字开头")
|
return fmt.Errorf("fact_key 格式无效,仅允许字母、数字及 . _ / -,且须以字母或数字开头(支持驼峰命名)")
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -111,6 +111,19 @@ func (db *DB) GetProject(id string) (*Project, error) {
|
|||||||
return &p, nil
|
return &p, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// GetProjectName returns a project display name without loading the full record.
|
||||||
|
func (db *DB) GetProjectName(id string) (string, error) {
|
||||||
|
var name string
|
||||||
|
err := db.QueryRow(`SELECT name FROM projects WHERE id = ?`, id).Scan(&name)
|
||||||
|
if err != nil {
|
||||||
|
if err == sql.ErrNoRows {
|
||||||
|
return "", fmt.Errorf("项目不存在")
|
||||||
|
}
|
||||||
|
return "", fmt.Errorf("获取项目名称失败: %w", err)
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(name), nil
|
||||||
|
}
|
||||||
|
|
||||||
func projectListSearchPattern(q string) string {
|
func projectListSearchPattern(q string) string {
|
||||||
q = strings.TrimSpace(q)
|
q = strings.TrimSpace(q)
|
||||||
if q == "" {
|
if q == "" {
|
||||||
@@ -263,7 +276,7 @@ func (db *DB) UpdateProject(p *Project) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理;漏洞 project_id 置空)。
|
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理;其他资源 project_id 置空)。
|
||||||
func (db *DB) DeleteProject(id string) error {
|
func (db *DB) DeleteProject(id string) error {
|
||||||
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
if _, err := db.Exec(`UPDATE vulnerabilities SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
|
return fmt.Errorf("解除漏洞项目关联失败: %w", err)
|
||||||
@@ -271,6 +284,12 @@ func (db *DB) DeleteProject(id string) error {
|
|||||||
if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
if _, err := db.Exec(`UPDATE assets SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
return fmt.Errorf("解除资产项目关联失败: %w", err)
|
return fmt.Errorf("解除资产项目关联失败: %w", err)
|
||||||
}
|
}
|
||||||
|
if _, err := db.Exec(`UPDATE webshell_connections SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
|
return fmt.Errorf("解除 WebShell 项目关联失败: %w", err)
|
||||||
|
}
|
||||||
|
if _, err := db.Exec(`UPDATE c2_listeners SET project_id = NULL WHERE project_id = ?`, id); err != nil {
|
||||||
|
return fmt.Errorf("解除 C2 监听器项目关联失败: %w", err)
|
||||||
|
}
|
||||||
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("删除项目失败: %w", err)
|
return fmt.Errorf("删除项目失败: %w", err)
|
||||||
|
|||||||
@@ -84,7 +84,7 @@ func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int)
|
|||||||
limit = 100
|
limit = 100
|
||||||
}
|
}
|
||||||
rows, err := db.Query(
|
rows, err := db.Query(
|
||||||
`SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id
|
`SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id, role_name
|
||||||
FROM conversations WHERE project_id = ? ORDER BY updated_at DESC LIMIT ? OFFSET ?`,
|
FROM conversations WHERE project_id = ? ORDER BY updated_at DESC LIMIT ? OFFSET ?`,
|
||||||
projectID, limit, offset,
|
projectID, limit, offset,
|
||||||
)
|
)
|
||||||
@@ -99,12 +99,16 @@ func (db *DB) ListConversationsByProjectID(projectID string, limit, offset int)
|
|||||||
var createdAt, updatedAt string
|
var createdAt, updatedAt string
|
||||||
var pinned int
|
var pinned int
|
||||||
var pid sql.NullString
|
var pid sql.NullString
|
||||||
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &pid); err != nil {
|
var roleName sql.NullString
|
||||||
|
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &pid, &roleName); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if pid.Valid {
|
if pid.Valid {
|
||||||
conv.ProjectID = strings.TrimSpace(pid.String)
|
conv.ProjectID = strings.TrimSpace(pid.String)
|
||||||
}
|
}
|
||||||
|
if roleName.Valid {
|
||||||
|
conv.RoleName = normalizeConversationRoleName(roleName.String)
|
||||||
|
}
|
||||||
conv.CreatedAt = parseDBTime(createdAt)
|
conv.CreatedAt = parseDBTime(createdAt)
|
||||||
conv.UpdatedAt = parseDBTime(updatedAt)
|
conv.UpdatedAt = parseDBTime(updatedAt)
|
||||||
conv.Pinned = pinned != 0
|
conv.Pinned = pinned != 0
|
||||||
|
|||||||
@@ -373,22 +373,46 @@ func TestRBACBatchResourceAssignmentValidationAndAtomicity(t *testing.T) {
|
|||||||
|
|
||||||
func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
||||||
db := newRBACTestDB(t)
|
db := newRBACTestDB(t)
|
||||||
ws1 := WebShellConnection{ID: "ws_visible", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()}
|
ws1 := WebShellConnection{ID: "ws_visible", ProjectID: "p1", URL: "http://a", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
ws2 := WebShellConnection{ID: "ws_hidden", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()}
|
ws2 := WebShellConnection{ID: "ws_hidden", ProjectID: "p2", URL: "http://b", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
|
ws3 := WebShellConnection{ID: "ws_other_project", ProjectID: "p2", URL: "http://c", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
|
ws4 := WebShellConnection{ID: "ws_unbound", URL: "http://d", Type: "php", Method: "post", CreatedAt: time.Now()}
|
||||||
if err := db.CreateWebshellConnection(&ws1); err != nil {
|
if err := db.CreateWebshellConnection(&ws1); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if err := db.CreateWebshellConnection(&ws2); err != nil {
|
if err := db.CreateWebshellConnection(&ws2); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
if err := db.CreateWebshellConnection(&ws3); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateWebshellConnection(&ws4); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
_ = db.SetResourceOwner("webshell", ws1.ID, "u1")
|
_ = db.SetResourceOwner("webshell", ws1.ID, "u1")
|
||||||
_ = db.SetResourceOwner("webshell", ws2.ID, "u2")
|
_ = db.SetResourceOwner("webshell", ws2.ID, "u2")
|
||||||
webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn)
|
_ = db.SetResourceOwner("webshell", ws3.ID, "u1")
|
||||||
|
_ = db.SetResourceOwner("webshell", ws4.ID, "u1")
|
||||||
|
webshells, err := db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(webshells) != 3 {
|
||||||
|
t.Fatalf("webshells = %#v, want 3 owned webshells including unbound", webshells)
|
||||||
|
}
|
||||||
|
webshells, err = db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, "p1")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if len(webshells) != 1 || webshells[0].ID != ws1.ID {
|
if len(webshells) != 1 || webshells[0].ID != ws1.ID {
|
||||||
t.Fatalf("webshells = %#v, want only %s", webshells, ws1.ID)
|
t.Fatalf("webshells scoped to p1 = %#v, want only %s", webshells, ws1.ID)
|
||||||
|
}
|
||||||
|
webshells, err = db.ListWebshellConnectionsForAccess("u1", RBACScopeOwn, ProjectFilterUnbound)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(webshells) != 1 || webshells[0].ID != ws4.ID {
|
||||||
|
t.Fatalf("unbound webshells = %#v, want only %s", webshells, ws4.ID)
|
||||||
}
|
}
|
||||||
|
|
||||||
if err := db.CreateBatchQueue("q_visible", "visible", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t1", "message": "a"}}); err != nil {
|
if err := db.CreateBatchQueue("q_visible", "visible", "", "eino_single", "manual", "", nil, "", 1, []map[string]interface{}{{"id": "t1", "message": "a"}}); err != nil {
|
||||||
@@ -411,61 +435,143 @@ func TestRBACWebshellAndBatchListAccess(t *testing.T) {
|
|||||||
func TestRBACC2AccessInheritsListener(t *testing.T) {
|
func TestRBACC2AccessInheritsListener(t *testing.T) {
|
||||||
db := newRBACTestDB(t)
|
db := newRBACTestDB(t)
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
l1 := &C2Listener{ID: "l_visible", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now}
|
l1 := &C2Listener{ID: "l_visible", ProjectID: "p1", Name: "visible", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9001, OwnerUserID: "u1", CreatedAt: now}
|
||||||
l2 := &C2Listener{ID: "l_hidden", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now}
|
l2 := &C2Listener{ID: "l_hidden", ProjectID: "p2", Name: "hidden", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9002, OwnerUserID: "u2", CreatedAt: now}
|
||||||
|
l3 := &C2Listener{ID: "l_other_project", ProjectID: "p2", Name: "other project", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9003, OwnerUserID: "u1", CreatedAt: now}
|
||||||
|
l4 := &C2Listener{ID: "l_unbound", Name: "unbound", Type: "http_beacon", BindHost: "127.0.0.1", BindPort: 9004, OwnerUserID: "u1", CreatedAt: now}
|
||||||
if err := db.CreateC2Listener(l1); err != nil {
|
if err := db.CreateC2Listener(l1); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if err := db.CreateC2Listener(l2); err != nil {
|
if err := db.CreateC2Listener(l2); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
if err := db.CreateC2Listener(l3); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Listener(l4); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
if err := db.UpsertC2Session(&C2Session{ID: "s_visible", ListenerID: l1.ID, ImplantUUID: "implant-visible", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
if err := db.UpsertC2Session(&C2Session{ID: "s_visible", ListenerID: l1.ID, ImplantUUID: "implant-visible", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if err := db.UpsertC2Session(&C2Session{ID: "s_hidden", ListenerID: l2.ID, ImplantUUID: "implant-hidden", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
if err := db.UpsertC2Session(&C2Session{ID: "s_hidden", ListenerID: l2.ID, ImplantUUID: "implant-hidden", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
if err := db.UpsertC2Session(&C2Session{ID: "s_other_project", ListenerID: l3.ID, ImplantUUID: "implant-other-project", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.UpsertC2Session(&C2Session{ID: "s_unbound", ListenerID: l4.ID, ImplantUUID: "implant-unbound", Status: "active", FirstSeenAt: now, LastCheckIn: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
if err := db.CreateC2Task(&C2Task{ID: "t_visible", SessionID: "s_visible", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
if err := db.CreateC2Task(&C2Task{ID: "t_visible", SessionID: "s_visible", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
if err := db.CreateC2Task(&C2Task{ID: "t_hidden", SessionID: "s_hidden", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
if err := db.CreateC2Task(&C2Task{ID: "t_other_project", SessionID: "s_other_project", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.CreateC2Task(&C2Task{ID: "t_unbound", SessionID: "s_unbound", TaskType: "shell", Status: "queued", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
if err := db.AppendC2Event(&C2Event{ID: "e_visible", Level: "info", Category: "task", SessionID: "s_visible", TaskID: "t_visible", Message: "visible", CreatedAt: now}); err != nil {
|
if err := db.AppendC2Event(&C2Event{ID: "e_visible", Level: "info", Category: "task", SessionID: "s_visible", TaskID: "t_visible", Message: "visible", CreatedAt: now}); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if err := db.AppendC2Event(&C2Event{ID: "e_hidden", Level: "info", Category: "task", SessionID: "s_hidden", TaskID: "t_hidden", Message: "hidden", CreatedAt: now}); err != nil {
|
if err := db.AppendC2Event(&C2Event{ID: "e_hidden", Level: "info", Category: "task", SessionID: "s_hidden", TaskID: "t_hidden", Message: "hidden", CreatedAt: now}); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
if err := db.AppendC2Event(&C2Event{ID: "e_other_project", Level: "info", Category: "task", SessionID: "s_other_project", TaskID: "t_other_project", Message: "other project", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := db.AppendC2Event(&C2Event{ID: "e_unbound", Level: "info", Category: "task", SessionID: "s_unbound", TaskID: "t_unbound", Message: "unbound", CreatedAt: now}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
access := RBACListAccess{UserID: "u1", Scope: RBACScopeOwn}
|
access := RBACListAccess{UserID: "u1", Scope: RBACScopeOwn}
|
||||||
listeners, err := db.ListC2ListenersForAccess(access)
|
listeners, err := db.ListC2ListenersForAccess(access, "")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(listeners) != 3 {
|
||||||
|
t.Fatalf("listeners = %#v, want 3 owned listeners including unbound", listeners)
|
||||||
|
}
|
||||||
|
listeners, err = db.ListC2ListenersForAccess(access, "p1")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if len(listeners) != 1 || listeners[0].ID != l1.ID {
|
if len(listeners) != 1 || listeners[0].ID != l1.ID {
|
||||||
t.Fatalf("listeners = %#v, want only %s", listeners, l1.ID)
|
t.Fatalf("listeners scoped to p1 = %#v, want only %s", listeners, l1.ID)
|
||||||
|
}
|
||||||
|
listeners, err = db.ListC2ListenersForAccess(access, ProjectFilterUnbound)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(listeners) != 1 || listeners[0].ID != l4.ID {
|
||||||
|
t.Fatalf("unbound listeners = %#v, want only %s", listeners, l4.ID)
|
||||||
}
|
}
|
||||||
sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, access)
|
sessions, err := db.ListC2SessionsForAccess(ListC2SessionsFilter{}, access)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
if len(sessions) != 3 {
|
||||||
|
t.Fatalf("sessions = %#v, want 3 owned sessions including unbound", sessions)
|
||||||
|
}
|
||||||
|
sessions, err = db.ListC2SessionsForAccess(ListC2SessionsFilter{ProjectID: "p1"}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
if len(sessions) != 1 || sessions[0].ID != "s_visible" {
|
if len(sessions) != 1 || sessions[0].ID != "s_visible" {
|
||||||
t.Fatalf("sessions = %#v, want only s_visible", sessions)
|
t.Fatalf("sessions scoped to p1 = %#v, want only s_visible", sessions)
|
||||||
|
}
|
||||||
|
sessions, err = db.ListC2SessionsForAccess(ListC2SessionsFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(sessions) != 1 || sessions[0].ID != "s_unbound" {
|
||||||
|
t.Fatalf("unbound sessions = %#v, want only s_unbound", sessions)
|
||||||
}
|
}
|
||||||
tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, access)
|
tasks, err := db.ListC2TasksForAccess(ListC2TasksFilter{}, access)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
if len(tasks) != 3 {
|
||||||
|
t.Fatalf("tasks = %#v, want 3 owned tasks including unbound", tasks)
|
||||||
|
}
|
||||||
|
tasks, err = db.ListC2TasksForAccess(ListC2TasksFilter{ProjectID: "p1"}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
if len(tasks) != 1 || tasks[0].ID != "t_visible" {
|
if len(tasks) != 1 || tasks[0].ID != "t_visible" {
|
||||||
t.Fatalf("tasks = %#v, want only t_visible", tasks)
|
t.Fatalf("tasks scoped to p1 = %#v, want only t_visible", tasks)
|
||||||
|
}
|
||||||
|
tasks, err = db.ListC2TasksForAccess(ListC2TasksFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(tasks) != 1 || tasks[0].ID != "t_unbound" {
|
||||||
|
t.Fatalf("unbound tasks = %#v, want only t_unbound", tasks)
|
||||||
}
|
}
|
||||||
events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, access)
|
events, err := db.ListC2EventsForAccess(ListC2EventsFilter{}, access)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
if len(events) != 3 {
|
||||||
|
t.Fatalf("events = %#v, want 3 owned events including unbound", events)
|
||||||
|
}
|
||||||
|
events, err = db.ListC2EventsForAccess(ListC2EventsFilter{ProjectID: "p1"}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
if len(events) != 1 || events[0].ID != "e_visible" {
|
if len(events) != 1 || events[0].ID != "e_visible" {
|
||||||
t.Fatalf("events = %#v, want only e_visible", events)
|
t.Fatalf("events scoped to p1 = %#v, want only e_visible", events)
|
||||||
|
}
|
||||||
|
events, err = db.ListC2EventsForAccess(ListC2EventsFilter{ProjectID: ProjectFilterUnbound}, access)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if len(events) != 1 || events[0].ID != "e_unbound" {
|
||||||
|
t.Fatalf("unbound events = %#v, want only e_unbound", events)
|
||||||
}
|
}
|
||||||
if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") {
|
if !db.UserCanAccessResource("u1", RBACScopeOwn, "c2_task", "t_visible") {
|
||||||
t.Fatalf("expected listener ownership to allow task detail")
|
t.Fatalf("expected listener ownership to allow task detail")
|
||||||
|
|||||||
@@ -191,6 +191,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("创建漏洞失败: %w", err)
|
return nil, fmt.Errorf("创建漏洞失败: %w", err)
|
||||||
}
|
}
|
||||||
|
db.refreshAssetRiskCacheForConversationsBestEffort(vuln.ConversationID)
|
||||||
return vuln, nil
|
return vuln, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -299,6 +300,8 @@ func (db *DB) CountVulnerabilitiesForAccess(filter VulnerabilityListFilter, acce
|
|||||||
// UpdateVulnerability 更新漏洞
|
// UpdateVulnerability 更新漏洞
|
||||||
func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
||||||
vuln.UpdatedAt = time.Now()
|
vuln.UpdatedAt = time.Now()
|
||||||
|
var oldConversationID string
|
||||||
|
_ = db.QueryRow(`SELECT COALESCE(conversation_id,'') FROM vulnerabilities WHERE id = ?`, id).Scan(&oldConversationID)
|
||||||
|
|
||||||
query := `
|
query := `
|
||||||
UPDATE vulnerabilities
|
UPDATE vulnerabilities
|
||||||
@@ -318,6 +321,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
|
|||||||
return fmt.Errorf("更新漏洞失败: %w", err)
|
return fmt.Errorf("更新漏洞失败: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
db.refreshAssetRiskCacheForConversationsBestEffort(oldConversationID, vuln.ConversationID)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -337,6 +341,10 @@ func (db *DB) DeleteVulnerabilitiesByFilterForAccess(filter VulnerabilityListFil
|
|||||||
args := []interface{}{}
|
args := []interface{}{}
|
||||||
where, args = filter.appendWhere(where, args)
|
where, args = filter.appendWhere(where, args)
|
||||||
where, args = appendVulnerabilityAccessFilter(where, args, access)
|
where, args = appendVulnerabilityAccessFilter(where, args, access)
|
||||||
|
affectedConversations, err := collectVulnerabilityConversationIDs(tx, where, args)
|
||||||
|
if err != nil {
|
||||||
|
return 0, err
|
||||||
|
}
|
||||||
|
|
||||||
clearQuery := `UPDATE project_facts SET related_vulnerability_id = NULL
|
clearQuery := `UPDATE project_facts SET related_vulnerability_id = NULL
|
||||||
WHERE related_vulnerability_id IN (SELECT id FROM vulnerabilities ` + where + `)`
|
WHERE related_vulnerability_id IN (SELECT id FROM vulnerabilities ` + where + `)`
|
||||||
@@ -356,6 +364,7 @@ func (db *DB) DeleteVulnerabilitiesByFilterForAccess(filter VulnerabilityListFil
|
|||||||
if err := tx.Commit(); err != nil {
|
if err := tx.Commit(); err != nil {
|
||||||
return 0, fmt.Errorf("提交事务失败: %w", err)
|
return 0, fmt.Errorf("提交事务失败: %w", err)
|
||||||
}
|
}
|
||||||
|
db.refreshAssetRiskCacheForConversationsBestEffort(affectedConversations...)
|
||||||
return deleted, nil
|
return deleted, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -366,6 +375,8 @@ func (db *DB) DeleteVulnerability(id string) error {
|
|||||||
return fmt.Errorf("开启事务失败: %w", err)
|
return fmt.Errorf("开启事务失败: %w", err)
|
||||||
}
|
}
|
||||||
defer func() { _ = tx.Rollback() }()
|
defer func() { _ = tx.Rollback() }()
|
||||||
|
var conversationID string
|
||||||
|
_ = tx.QueryRow(`SELECT COALESCE(conversation_id,'') FROM vulnerabilities WHERE id = ?`, id).Scan(&conversationID)
|
||||||
|
|
||||||
// 删除漏洞前先解除项目事实中的关联,避免前端继续显示已删除漏洞的短 ID。
|
// 删除漏洞前先解除项目事实中的关联,避免前端继续显示已删除漏洞的短 ID。
|
||||||
if _, err := tx.Exec("UPDATE project_facts SET related_vulnerability_id = NULL WHERE related_vulnerability_id = ?", id); err != nil {
|
if _, err := tx.Exec("UPDATE project_facts SET related_vulnerability_id = NULL WHERE related_vulnerability_id = ?", id); err != nil {
|
||||||
@@ -377,9 +388,29 @@ func (db *DB) DeleteVulnerability(id string) error {
|
|||||||
if err := tx.Commit(); err != nil {
|
if err := tx.Commit(); err != nil {
|
||||||
return fmt.Errorf("提交事务失败: %w", err)
|
return fmt.Errorf("提交事务失败: %w", err)
|
||||||
}
|
}
|
||||||
|
db.refreshAssetRiskCacheForConversationsBestEffort(conversationID)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func collectVulnerabilityConversationIDs(tx *sql.Tx, where string, args []interface{}) ([]string, error) {
|
||||||
|
rows, err := tx.Query(`SELECT DISTINCT COALESCE(conversation_id,'') FROM vulnerabilities `+where, args...)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("查询受影响漏洞会话失败: %w", err)
|
||||||
|
}
|
||||||
|
defer rows.Close()
|
||||||
|
ids := []string{}
|
||||||
|
for rows.Next() {
|
||||||
|
var id string
|
||||||
|
if err := rows.Scan(&id); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(id) != "" {
|
||||||
|
ids = append(ids, id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ids, rows.Err()
|
||||||
|
}
|
||||||
|
|
||||||
// GetVulnerabilityStats 获取漏洞统计(筛选条件与 ListVulnerabilities / CountVulnerabilities 一致)
|
// GetVulnerabilityStats 获取漏洞统计(筛选条件与 ListVulnerabilities / CountVulnerabilities 一致)
|
||||||
func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) {
|
func (db *DB) GetVulnerabilityStats(filter VulnerabilityListFilter) (map[string]interface{}, error) {
|
||||||
return db.GetVulnerabilityStatsForAccess(filter, RBACListAccess{})
|
return db.GetVulnerabilityStatsForAccess(filter, RBACListAccess{})
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ import (
|
|||||||
// WebShellConnection WebShell 连接配置
|
// WebShellConnection WebShell 连接配置
|
||||||
type WebShellConnection struct {
|
type WebShellConnection struct {
|
||||||
ID string `json:"id"`
|
ID string `json:"id"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
URL string `json:"url"`
|
URL string `json:"url"`
|
||||||
Password string `json:"password"`
|
Password string `json:"password"`
|
||||||
Type string `json:"type"`
|
Type string `json:"type"`
|
||||||
@@ -60,17 +61,24 @@ func (db *DB) UpsertWebshellConnectionState(connectionID, stateJSON string) erro
|
|||||||
|
|
||||||
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
|
// ListWebshellConnections 列出所有 WebShell 连接,按创建时间倒序
|
||||||
func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
func (db *DB) ListWebshellConnections() ([]WebShellConnection, error) {
|
||||||
return db.ListWebshellConnectionsForAccess("", "")
|
return db.ListWebshellConnectionsForAccess("", "", "")
|
||||||
}
|
}
|
||||||
|
|
||||||
func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShellConnection, error) {
|
func (db *DB) ListWebshellConnectionsForAccess(userID, scope, projectID string) ([]WebShellConnection, error) {
|
||||||
query := `
|
query := `
|
||||||
SELECT id, url, password, type, method, cmd_param, remark,
|
SELECT id, COALESCE(project_id, '') AS project_id, url, password, type, method, cmd_param, remark,
|
||||||
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||||
FROM webshell_connections
|
FROM webshell_connections
|
||||||
WHERE 1=1
|
WHERE 1=1
|
||||||
`
|
`
|
||||||
args := []interface{}{}
|
args := []interface{}{}
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == ProjectFilterUnbound {
|
||||||
|
query += ` AND COALESCE(project_id, '') = ''`
|
||||||
|
} else if projectID != "" {
|
||||||
|
query += ` AND COALESCE(project_id, '') = ?`
|
||||||
|
args = append(args, projectID)
|
||||||
|
}
|
||||||
userID = strings.TrimSpace(userID)
|
userID = strings.TrimSpace(userID)
|
||||||
if userID != "" && scope != RBACScopeAll {
|
if userID != "" && scope != RBACScopeAll {
|
||||||
query += ` AND (
|
query += ` AND (
|
||||||
@@ -93,7 +101,7 @@ func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShell
|
|||||||
var list []WebShellConnection
|
var list []WebShellConnection
|
||||||
for rows.Next() {
|
for rows.Next() {
|
||||||
var c WebShellConnection
|
var c WebShellConnection
|
||||||
err := rows.Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
err := rows.Scan(&c.ID, &c.ProjectID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
db.logger.Warn("扫描 WebShell 连接行失败", zap.Error(err))
|
db.logger.Warn("扫描 WebShell 连接行失败", zap.Error(err))
|
||||||
continue
|
continue
|
||||||
@@ -106,12 +114,12 @@ func (db *DB) ListWebshellConnectionsForAccess(userID, scope string) ([]WebShell
|
|||||||
// GetWebshellConnection 根据 ID 获取一条连接
|
// GetWebshellConnection 根据 ID 获取一条连接
|
||||||
func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
||||||
query := `
|
query := `
|
||||||
SELECT id, url, password, type, method, cmd_param, remark,
|
SELECT id, COALESCE(project_id, '') AS project_id, url, password, type, method, cmd_param, remark,
|
||||||
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
COALESCE(encoding, '') AS encoding, COALESCE(os, '') AS os, created_at
|
||||||
FROM webshell_connections WHERE id = ?
|
FROM webshell_connections WHERE id = ?
|
||||||
`
|
`
|
||||||
var c WebShellConnection
|
var c WebShellConnection
|
||||||
err := db.QueryRow(query, id).Scan(&c.ID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
err := db.QueryRow(query, id).Scan(&c.ID, &c.ProjectID, &c.URL, &c.Password, &c.Type, &c.Method, &c.CmdParam, &c.Remark, &c.Encoding, &c.OS, &c.CreatedAt)
|
||||||
if err == sql.ErrNoRows {
|
if err == sql.ErrNoRows {
|
||||||
return nil, nil
|
return nil, nil
|
||||||
}
|
}
|
||||||
@@ -125,10 +133,10 @@ func (db *DB) GetWebshellConnection(id string) (*WebShellConnection, error) {
|
|||||||
// CreateWebshellConnection 创建 WebShell 连接
|
// CreateWebshellConnection 创建 WebShell 连接
|
||||||
func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
||||||
query := `
|
query := `
|
||||||
INSERT INTO webshell_connections (id, url, password, type, method, cmd_param, remark, encoding, os, created_at)
|
INSERT INTO webshell_connections (id, project_id, url, password, type, method, cmd_param, remark, encoding, os, created_at)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
`
|
`
|
||||||
_, err := db.Exec(query, c.ID, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.CreatedAt)
|
_, err := db.Exec(query, c.ID, strings.TrimSpace(c.ProjectID), c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.CreatedAt)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
db.logger.Error("创建 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
db.logger.Error("创建 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||||
return err
|
return err
|
||||||
@@ -140,10 +148,10 @@ func (db *DB) CreateWebshellConnection(c *WebShellConnection) error {
|
|||||||
func (db *DB) UpdateWebshellConnection(c *WebShellConnection) error {
|
func (db *DB) UpdateWebshellConnection(c *WebShellConnection) error {
|
||||||
query := `
|
query := `
|
||||||
UPDATE webshell_connections
|
UPDATE webshell_connections
|
||||||
SET url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?, encoding = ?, os = ?
|
SET project_id = ?, url = ?, password = ?, type = ?, method = ?, cmd_param = ?, remark = ?, encoding = ?, os = ?
|
||||||
WHERE id = ?
|
WHERE id = ?
|
||||||
`
|
`
|
||||||
result, err := db.Exec(query, c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.ID)
|
result, err := db.Exec(query, strings.TrimSpace(c.ProjectID), c.URL, c.Password, c.Type, c.Method, c.CmdParam, c.Remark, c.Encoding, c.OS, c.ID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
db.logger.Error("更新 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
db.logger.Error("更新 WebShell 连接失败", zap.Error(err), zap.String("id", c.ID))
|
||||||
return err
|
return err
|
||||||
|
|||||||
+52
-19
@@ -333,20 +333,41 @@ type ChatReasoningRequest struct {
|
|||||||
Effort string `json:"effort,omitempty"`
|
Effort string `json:"effort,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ChatFinalizationRequest is a caller-provided delivery policy. The server does
|
||||||
|
// not infer execution intent from natural-language user text.
|
||||||
|
type ChatFinalizationRequest struct {
|
||||||
|
RequireExecutionEvidence *bool `json:"requireExecutionEvidence,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
// ChatRequest 聊天请求
|
// ChatRequest 聊天请求
|
||||||
type ChatRequest struct {
|
type ChatRequest struct {
|
||||||
Message string `json:"message" binding:"required"`
|
Message string `json:"message" binding:"required"`
|
||||||
ConversationID string `json:"conversationId,omitempty"`
|
ConversationID string `json:"conversationId,omitempty"`
|
||||||
ProjectID string `json:"projectId,omitempty"` // 新对话绑定的项目(可选;未指定时可用 config.project.default_project_id)
|
ProjectID string `json:"projectId,omitempty"` // 新对话绑定的项目(可选;未指定时可用 config.project.default_project_id)
|
||||||
Role string `json:"role,omitempty"` // 角色名称
|
Role string `json:"role,omitempty"` // 角色名称
|
||||||
Attachments []ChatAttachment `json:"attachments,omitempty"`
|
Attachments []ChatAttachment `json:"attachments,omitempty"`
|
||||||
WebShellConnectionID string `json:"webshellConnectionId,omitempty"` // WebShell 管理 - AI 助手:当前选中的连接 ID,仅使用 webshell_* 工具
|
WebShellConnectionID string `json:"webshellConnectionId,omitempty"` // WebShell 管理 - AI 助手:当前选中的连接 ID,仅使用 webshell_* 工具
|
||||||
Hitl *HITLRequest `json:"hitl,omitempty"`
|
AIChannelID string `json:"aiChannelId,omitempty"` // 会话级 AI 通道;空则使用 ai.default_channel
|
||||||
Reasoning *ChatReasoningRequest `json:"reasoning,omitempty"`
|
Hitl *HITLRequest `json:"hitl,omitempty"`
|
||||||
|
Reasoning *ChatReasoningRequest `json:"reasoning,omitempty"`
|
||||||
|
Finalization ChatFinalizationRequest `json:"finalization,omitempty"`
|
||||||
// Orchestration 仅对 /api/multi-agent、/api/multi-agent/stream:deep | plan_execute | supervisor;空则等同 deep。机器人/批量等无请求体时由服务端默认 deep。/api/eino-agent* 不使用此字段。
|
// Orchestration 仅对 /api/multi-agent、/api/multi-agent/stream:deep | plan_execute | supervisor;空则等同 deep。机器人/批量等无请求体时由服务端默认 deep。/api/eino-agent* 不使用此字段。
|
||||||
Orchestration string `json:"orchestration,omitempty"`
|
Orchestration string `json:"orchestration,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (h *AgentHandler) configForAIChannel(channelID string) (*config.Config, string, error) {
|
||||||
|
if h == nil || h.config == nil {
|
||||||
|
return nil, "", fmt.Errorf("服务器配置未加载")
|
||||||
|
}
|
||||||
|
oa, resolvedID, ok := h.config.ResolveAIChannel(channelID)
|
||||||
|
if !ok {
|
||||||
|
return nil, resolvedID, fmt.Errorf("AI 通道不存在: %s", resolvedID)
|
||||||
|
}
|
||||||
|
cfgCopy := *h.config
|
||||||
|
cfgCopy.OpenAI = oa
|
||||||
|
return &cfgCopy, resolvedID, nil
|
||||||
|
}
|
||||||
|
|
||||||
func chatReasoningToClientIntent(r *ChatReasoningRequest) *reasoning.ClientIntent {
|
func chatReasoningToClientIntent(r *ChatReasoningRequest) *reasoning.ClientIntent {
|
||||||
if r == nil {
|
if r == nil {
|
||||||
return nil
|
return nil
|
||||||
@@ -654,10 +675,18 @@ func (h *AgentHandler) mergeAssistantMessagePartialOnCancel(messageID, partial s
|
|||||||
|
|
||||||
// ChatResponse 聊天响应
|
// ChatResponse 聊天响应
|
||||||
type ChatResponse struct {
|
type ChatResponse struct {
|
||||||
Response string `json:"response"`
|
Response string `json:"response"`
|
||||||
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"` // 本次对话中执行的MCP调用ID列表
|
MCPExecutionIDs []string `json:"mcpExecutionIds,omitempty"` // 本次对话中执行的MCP调用ID列表
|
||||||
ConversationID string `json:"conversationId"` // 对话ID
|
ConversationID string `json:"conversationId"` // 对话ID
|
||||||
Time time.Time `json:"time"`
|
Time time.Time `json:"time"`
|
||||||
|
Finalizable bool `json:"finalizable"`
|
||||||
|
Finalized bool `json:"finalized"`
|
||||||
|
Status string `json:"status,omitempty"`
|
||||||
|
CompletionReason string `json:"completionReason,omitempty"`
|
||||||
|
EvidenceVerified bool `json:"evidenceVerified"`
|
||||||
|
EvidenceRefs []string `json:"evidenceRefs,omitempty"`
|
||||||
|
PendingExecutionIDs []string `json:"pendingExecutionIds,omitempty"`
|
||||||
|
MissingChecks []string `json:"missingChecks,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMessageID, conversationID string, resultMA *multiagent.RunResult, errMA error) (string, string, error) {
|
func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMessageID, conversationID string, resultMA *multiagent.RunResult, errMA error) (string, string, error) {
|
||||||
@@ -673,19 +702,20 @@ func (h *AgentHandler) finalizeRobotAgentError(ctx context.Context, assistantMes
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (h *AgentHandler) finalizeRobotAgentSuccess(assistantMessageID, conversationID string, resultMA *multiagent.RunResult) (string, string, error) {
|
func (h *AgentHandler) finalizeRobotAgentSuccess(assistantMessageID, conversationID string, resultMA *multiagent.RunResult) (string, string, error) {
|
||||||
if assistantMessageID != "" {
|
decision := h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "robot", resultMA, resultMA.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput), true)
|
||||||
if errU := h.db.UpdateAssistantMessageFinalize(assistantMessageID, resultMA.Response, resultMA.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(resultMA.LastAgentTraceInput)); errU != nil {
|
responseText := decision.FinalText
|
||||||
h.logger.Warn("机器人:更新助手消息失败", zap.Error(errU))
|
if !decision.Finalizable {
|
||||||
}
|
responseText = finalizationBlockedMessage(decision)
|
||||||
} else {
|
}
|
||||||
if _, err := h.db.AddMessage(conversationID, "assistant", resultMA.Response, resultMA.MCPExecutionIDs); err != nil {
|
if assistantMessageID == "" {
|
||||||
|
if _, err := h.db.AddMessage(conversationID, "assistant", responseText, resultMA.MCPExecutionIDs); err != nil {
|
||||||
h.logger.Warn("机器人:保存助手消息失败", zap.Error(err))
|
h.logger.Warn("机器人:保存助手消息失败", zap.Error(err))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "" {
|
if resultMA.LastAgentTraceInput != "" || resultMA.LastAgentTraceOutput != "" {
|
||||||
_ = h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput)
|
_ = h.db.SaveAgentTrace(conversationID, resultMA.LastAgentTraceInput, resultMA.LastAgentTraceOutput)
|
||||||
}
|
}
|
||||||
return resultMA.Response, conversationID, nil
|
return responseText, conversationID, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (h *AgentHandler) runRobotEinoSingleWithRetry(
|
func (h *AgentHandler) runRobotEinoSingleWithRetry(
|
||||||
@@ -816,6 +846,9 @@ func (h *AgentHandler) ProcessMessageForRobot(ctx context.Context, platform stri
|
|||||||
progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, nil)
|
progressCallback := h.createProgressCallback(taskCtx, cancelWithCause, conversationID, assistantMessageID, nil)
|
||||||
|
|
||||||
robotMode := config.NormalizeAgentMode(agentMode)
|
robotMode := config.NormalizeAgentMode(agentMode)
|
||||||
|
if err := h.db.SetConversationAgentMode(conversationID, robotMode); err != nil {
|
||||||
|
h.logger.Warn("机器人:更新对话模式失败", zap.String("conversationId", conversationID), zap.String("agentMode", robotMode), zap.Error(err))
|
||||||
|
}
|
||||||
switch robotMode {
|
switch robotMode {
|
||||||
case "eino_single":
|
case "eino_single":
|
||||||
return h.runRobotEinoSingleWithRetry(taskCtx, conversationID, finalMessage, agentHistoryMessages, roleTools, progressCallback, assistantMessageID, &taskStatus)
|
return h.runRobotEinoSingleWithRetry(taskCtx, conversationID, finalMessage, agentHistoryMessages, roleTools, progressCallback, assistantMessageID, &taskStatus)
|
||||||
|
|||||||
@@ -238,6 +238,11 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
|||||||
useBatchMulti = true
|
useBatchMulti = true
|
||||||
batchOrch = "deep"
|
batchOrch = "deep"
|
||||||
}
|
}
|
||||||
|
if useBatchMulti {
|
||||||
|
_ = h.db.SetConversationAgentMode(conversationID, batchOrch)
|
||||||
|
} else {
|
||||||
|
_ = h.db.SetConversationAgentMode(conversationID, "eino_single")
|
||||||
|
}
|
||||||
|
|
||||||
var resultMA *multiagent.RunResult
|
var resultMA *multiagent.RunResult
|
||||||
var runErr error
|
var runErr error
|
||||||
@@ -268,19 +273,38 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
|||||||
|
|
||||||
h.logger.Info("批量任务执行成功", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
|
h.logger.Info("批量任务执行成功", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID))
|
||||||
|
|
||||||
resText := resultMA.Response
|
|
||||||
mcpIDs := resultMA.MCPExecutionIDs
|
mcpIDs := resultMA.MCPExecutionIDs
|
||||||
lastIn := resultMA.LastAgentTraceInput
|
lastIn := resultMA.LastAgentTraceInput
|
||||||
lastOut := resultMA.LastAgentTraceOutput
|
lastOut := resultMA.LastAgentTraceOutput
|
||||||
|
reasoningContent := multiagent.AggregatedReasoningFromTraceJSON(lastIn)
|
||||||
|
agentMode := "batch_eino_single"
|
||||||
|
if useBatchMulti {
|
||||||
|
agentMode = "batch_eino_" + batchOrch
|
||||||
|
}
|
||||||
|
decision := h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, resultMA, mcpIDs, reasoningContent, true)
|
||||||
|
resText := decision.FinalText
|
||||||
|
if !decision.Finalizable {
|
||||||
|
resText = finalizationBlockedMessage(decision)
|
||||||
|
finishStatus = decision.Status
|
||||||
|
sendEvent("finalization_check", resText, decision)
|
||||||
|
}
|
||||||
|
sendEvent("response", resText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||||
|
"conversationId": conversationID,
|
||||||
|
"messageId": assistantMessageID,
|
||||||
|
"agentMode": agentMode,
|
||||||
|
"mcpExecutionIds": mcpIDs,
|
||||||
|
"batchQueueId": queueID,
|
||||||
|
"batchTaskId": task.ID,
|
||||||
|
"batchTaskStatus": map[bool]string{true: string(BatchTaskStatusCompleted), false: string(BatchTaskStatusFailed)}[decision.Finalizable],
|
||||||
|
"candidatePreview": safeTruncateString(resultMA.Response, 500),
|
||||||
|
}))
|
||||||
|
|
||||||
if assistantMessageID != "" {
|
if assistantMessageID == "" {
|
||||||
if updateErr := h.db.UpdateAssistantMessageFinalize(assistantMessageID, resText, mcpIDs, multiagent.AggregatedReasoningFromTraceJSON(lastIn)); updateErr != nil {
|
_, err = h.db.AddMessage(conversationID, "assistant", resText, mcpIDs)
|
||||||
h.logger.Warn("更新助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.Error(updateErr))
|
} else if !decision.Finalizable {
|
||||||
if _, err = h.db.AddMessage(conversationID, "assistant", resText, mcpIDs); err != nil {
|
err = nil
|
||||||
h.logger.Error("保存助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
|
}
|
||||||
}
|
if err != nil {
|
||||||
}
|
|
||||||
} else if _, err = h.db.AddMessage(conversationID, "assistant", resText, mcpIDs); err != nil {
|
|
||||||
h.logger.Error("保存助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
|
h.logger.Error("保存助手消息失败", zap.String("queueId", queueID), zap.String("taskId", task.ID), zap.String("conversationId", conversationID), zap.Error(err))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -290,6 +314,10 @@ func (h *AgentHandler) executeOneBatchSubTask(queueID string, queue *BatchTaskQu
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if !decision.Finalizable {
|
||||||
|
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusFailed, resText, finalizationCheckMessage(decision), conversationID)
|
||||||
|
return
|
||||||
|
}
|
||||||
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusCompleted, resText, "", conversationID)
|
h.batchTaskManager.UpdateTaskStatusWithConversationID(queueID, task.ID, BatchTaskStatusCompleted, resText, "", conversationID)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+85
-2
@@ -60,7 +60,7 @@ func (h *C2Handler) SetManager(m *c2.Manager) {
|
|||||||
|
|
||||||
// ListListeners 获取监听器列表
|
// ListListeners 获取监听器列表
|
||||||
func (h *C2Handler) ListListeners(c *gin.Context) {
|
func (h *C2Handler) ListListeners(c *gin.Context) {
|
||||||
listeners, err := h.mgr().DB().ListC2ListenersForAccess(c2AccessFromContext(c))
|
listeners, err := h.mgr().DB().ListC2ListenersForAccess(c2AccessFromContext(c), c.Query("project_id"))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
return
|
return
|
||||||
@@ -77,6 +77,7 @@ func (h *C2Handler) ListListeners(c *gin.Context) {
|
|||||||
func (h *C2Handler) CreateListener(c *gin.Context) {
|
func (h *C2Handler) CreateListener(c *gin.Context) {
|
||||||
var req struct {
|
var req struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
|
ProjectID string `json:"project_id,omitempty"`
|
||||||
Type string `json:"type"`
|
Type string `json:"type"`
|
||||||
BindHost string `json:"bind_host"`
|
BindHost string `json:"bind_host"`
|
||||||
BindPort int `json:"bind_port"`
|
BindPort int `json:"bind_port"`
|
||||||
@@ -92,6 +93,7 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
|
|||||||
|
|
||||||
input := c2.CreateListenerInput{
|
input := c2.CreateListenerInput{
|
||||||
Name: req.Name,
|
Name: req.Name,
|
||||||
|
ProjectID: req.ProjectID,
|
||||||
Type: req.Type,
|
Type: req.Type,
|
||||||
BindHost: req.BindHost,
|
BindHost: req.BindHost,
|
||||||
BindPort: req.BindPort,
|
BindPort: req.BindPort,
|
||||||
@@ -100,6 +102,10 @@ func (h *C2Handler) CreateListener(c *gin.Context) {
|
|||||||
Config: req.Config,
|
Config: req.Config,
|
||||||
CallbackHost: strings.TrimSpace(req.CallbackHost),
|
CallbackHost: strings.TrimSpace(req.CallbackHost),
|
||||||
}
|
}
|
||||||
|
if !h.canAccessProject(c, input.ProjectID) {
|
||||||
|
c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
listener, err := h.mgr().CreateListener(input)
|
listener, err := h.mgr().CreateListener(input)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -158,6 +164,7 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
|||||||
|
|
||||||
var req struct {
|
var req struct {
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
|
ProjectID string `json:"project_id"`
|
||||||
BindHost string `json:"bind_host"`
|
BindHost string `json:"bind_host"`
|
||||||
BindPort int `json:"bind_port"`
|
BindPort int `json:"bind_port"`
|
||||||
ProfileID string `json:"profile_id"`
|
ProfileID string `json:"profile_id"`
|
||||||
@@ -179,6 +186,7 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
listener.Name = req.Name
|
listener.Name = req.Name
|
||||||
|
listener.ProjectID = strings.TrimSpace(req.ProjectID)
|
||||||
listener.BindHost = req.BindHost
|
listener.BindHost = req.BindHost
|
||||||
listener.BindPort = req.BindPort
|
listener.BindPort = req.BindPort
|
||||||
listener.ProfileID = req.ProfileID
|
listener.ProfileID = req.ProfileID
|
||||||
@@ -187,6 +195,10 @@ func (h *C2Handler) UpdateListener(c *gin.Context) {
|
|||||||
cfgJSON, _ := json.Marshal(req.Config)
|
cfgJSON, _ := json.Marshal(req.Config)
|
||||||
listener.ConfigJSON = string(cfgJSON)
|
listener.ConfigJSON = string(cfgJSON)
|
||||||
}
|
}
|
||||||
|
if !h.canAccessProject(c, listener.ProjectID) {
|
||||||
|
c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"})
|
||||||
|
return
|
||||||
|
}
|
||||||
if req.CallbackHost != nil {
|
if req.CallbackHost != nil {
|
||||||
cfg := &c2.ListenerConfig{}
|
cfg := &c2.ListenerConfig{}
|
||||||
raw := strings.TrimSpace(listener.ConfigJSON)
|
raw := strings.TrimSpace(listener.ConfigJSON)
|
||||||
@@ -275,6 +287,7 @@ func (h *C2Handler) StopListener(c *gin.Context) {
|
|||||||
func (h *C2Handler) ListSessions(c *gin.Context) {
|
func (h *C2Handler) ListSessions(c *gin.Context) {
|
||||||
filter := database.ListC2SessionsFilter{
|
filter := database.ListC2SessionsFilter{
|
||||||
ListenerID: c.Query("listener_id"),
|
ListenerID: c.Query("listener_id"),
|
||||||
|
ProjectID: c.Query("project_id"),
|
||||||
Status: c.Query("status"),
|
Status: c.Query("status"),
|
||||||
OS: c.Query("os"),
|
OS: c.Query("os"),
|
||||||
Search: c.Query("search"),
|
Search: c.Query("search"),
|
||||||
@@ -404,6 +417,47 @@ func (h *C2Handler) SetSessionSleep(c *gin.Context) {
|
|||||||
c.JSON(http.StatusOK, out)
|
c.JSON(http.StatusOK, out)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetSessionNote 更新会话备注(仅服务端元数据,不下发植入体)
|
||||||
|
func (h *C2Handler) SetSessionNote(c *gin.Context) {
|
||||||
|
id := c.Param("id")
|
||||||
|
var req struct {
|
||||||
|
Note string `json:"note"`
|
||||||
|
}
|
||||||
|
if err := c.ShouldBindJSON(&req); err != nil {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
note := strings.TrimSpace(req.Note)
|
||||||
|
if len(note) > 2000 {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": "note too long (max 2000 characters)"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
session, err := h.mgr().DB().GetC2Session(id)
|
||||||
|
if err != nil {
|
||||||
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if session == nil {
|
||||||
|
c.JSON(http.StatusNotFound, gin.H{"error": "session not found"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := h.mgr().DB().SetC2SessionNote(id, note); err != nil {
|
||||||
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if h.audit != nil {
|
||||||
|
h.audit.RecordOK(c, "c2", "session_note", "更新 C2 会话备注", "c2_session", id, map[string]interface{}{
|
||||||
|
"note_len": len(note),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
c.JSON(http.StatusOK, gin.H{
|
||||||
|
"updated": true,
|
||||||
|
"note": note,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
// 任务 API
|
// 任务 API
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
@@ -412,7 +466,14 @@ func (h *C2Handler) SetSessionSleep(c *gin.Context) {
|
|||||||
func (h *C2Handler) ListTasks(c *gin.Context) {
|
func (h *C2Handler) ListTasks(c *gin.Context) {
|
||||||
filter := database.ListC2TasksFilter{
|
filter := database.ListC2TasksFilter{
|
||||||
SessionID: c.Query("session_id"),
|
SessionID: c.Query("session_id"),
|
||||||
|
ProjectID: c.Query("project_id"),
|
||||||
Status: c.Query("status"),
|
Status: c.Query("status"),
|
||||||
|
TaskType: c.Query("task_type"),
|
||||||
|
}
|
||||||
|
if since := c.Query("since"); since != "" {
|
||||||
|
if t, err := database.ParseRFC3339Time(since); err == nil {
|
||||||
|
filter.Since = &t
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
paginated := false
|
paginated := false
|
||||||
@@ -447,7 +508,7 @@ func (h *C2Handler) ListTasks(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// 仪表盘「待审任务」为全局 queued/pending 数量,与列表 session 过滤无关
|
// 仪表盘「待审任务」为全局 queued/pending 数量,与列表 session 过滤无关
|
||||||
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", access)
|
pendingN, _ := h.mgr().DB().CountC2TasksQueuedOrPendingForAccess("", filter.ProjectID, access)
|
||||||
|
|
||||||
if !paginated {
|
if !paginated {
|
||||||
c.JSON(http.StatusOK, gin.H{
|
c.JSON(http.StatusOK, gin.H{
|
||||||
@@ -462,9 +523,15 @@ func (h *C2Handler) ListTasks(c *gin.Context) {
|
|||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
statusCounts, err := h.mgr().DB().CountC2TasksByStatusForAccess(filter, access)
|
||||||
|
if err != nil {
|
||||||
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
c.JSON(http.StatusOK, gin.H{
|
c.JSON(http.StatusOK, gin.H{
|
||||||
"tasks": tasks,
|
"tasks": tasks,
|
||||||
"total": total,
|
"total": total,
|
||||||
|
"status_counts": statusCounts,
|
||||||
"page": page,
|
"page": page,
|
||||||
"page_size": pageSize,
|
"page_size": pageSize,
|
||||||
"pending_queued_count": pendingN,
|
"pending_queued_count": pendingN,
|
||||||
@@ -784,6 +851,7 @@ func (h *C2Handler) ListEvents(c *gin.Context) {
|
|||||||
filter := database.ListC2EventsFilter{
|
filter := database.ListC2EventsFilter{
|
||||||
Level: c.Query("level"),
|
Level: c.Query("level"),
|
||||||
Category: c.Query("category"),
|
Category: c.Query("category"),
|
||||||
|
ProjectID: c.Query("project_id"),
|
||||||
SessionID: c.Query("session_id"),
|
SessionID: c.Query("session_id"),
|
||||||
TaskID: c.Query("task_id"),
|
TaskID: c.Query("task_id"),
|
||||||
}
|
}
|
||||||
@@ -1121,6 +1189,21 @@ func c2AccessFromContext(c *gin.Context) database.RBACListAccess {
|
|||||||
return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
|
return database.RBACListAccess{UserID: session.UserID, Scope: session.Scope}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (h *C2Handler) canAccessProject(c *gin.Context, projectID string) bool {
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == "" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
session, ok := security.CurrentSession(c)
|
||||||
|
if !ok {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if session.Scope == database.RBACScopeAll {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return h.mgr().DB().UserCanAccessResource(session.UserID, session.Scope, "project", projectID)
|
||||||
|
}
|
||||||
|
|
||||||
func (h *C2Handler) c2ResourceAllowed(c *gin.Context, resourceType, resourceID string) bool {
|
func (h *C2Handler) c2ResourceAllowed(c *gin.Context, resourceType, resourceID string) bool {
|
||||||
session, ok := security.CurrentSession(c)
|
session, ok := security.CurrentSession(c)
|
||||||
if !ok {
|
if !ok {
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -258,6 +258,7 @@ func (h *ConfigHandler) ApplyWechatRobotBinding(wc config.RobotWechatConfig) err
|
|||||||
|
|
||||||
// GetConfigResponse 获取配置响应
|
// GetConfigResponse 获取配置响应
|
||||||
type GetConfigResponse struct {
|
type GetConfigResponse struct {
|
||||||
|
AI config.AIConfig `json:"ai"`
|
||||||
OpenAI config.OpenAIConfig `json:"openai"`
|
OpenAI config.OpenAIConfig `json:"openai"`
|
||||||
Vision config.VisionConfig `json:"vision"`
|
Vision config.VisionConfig `json:"vision"`
|
||||||
FOFA config.FofaConfig `json:"fofa"`
|
FOFA config.FofaConfig `json:"fofa"`
|
||||||
@@ -363,6 +364,7 @@ func (h *ConfigHandler) GetConfig(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
c.JSON(http.StatusOK, GetConfigResponse{
|
c.JSON(http.StatusOK, GetConfigResponse{
|
||||||
|
AI: h.config.AI,
|
||||||
OpenAI: h.config.OpenAI,
|
OpenAI: h.config.OpenAI,
|
||||||
Vision: h.config.Vision,
|
Vision: h.config.Vision,
|
||||||
FOFA: h.config.FOFA,
|
FOFA: h.config.FOFA,
|
||||||
@@ -706,6 +708,7 @@ func (h *ConfigHandler) GetTools(c *gin.Context) {
|
|||||||
|
|
||||||
// UpdateConfigRequest 更新配置请求
|
// UpdateConfigRequest 更新配置请求
|
||||||
type UpdateConfigRequest struct {
|
type UpdateConfigRequest struct {
|
||||||
|
AI *config.AIConfig `json:"ai,omitempty"`
|
||||||
OpenAI *config.OpenAIConfig `json:"openai,omitempty"`
|
OpenAI *config.OpenAIConfig `json:"openai,omitempty"`
|
||||||
Vision *config.VisionConfig `json:"vision,omitempty"`
|
Vision *config.VisionConfig `json:"vision,omitempty"`
|
||||||
FOFA *config.FofaConfig `json:"fofa,omitempty"`
|
FOFA *config.FofaConfig `json:"fofa,omitempty"`
|
||||||
@@ -785,8 +788,20 @@ func (h *ConfigHandler) UpdateConfig(c *gin.Context) {
|
|||||||
defer h.mu.Unlock()
|
defer h.mu.Unlock()
|
||||||
|
|
||||||
// 更新OpenAI配置
|
// 更新OpenAI配置
|
||||||
|
if req.AI != nil {
|
||||||
|
h.config.AI = *req.AI
|
||||||
|
h.config.ApplyDefaultAIChannel()
|
||||||
|
h.logger.Info("更新 AI 通道配置",
|
||||||
|
zap.String("default_channel", h.config.AI.DefaultChannel),
|
||||||
|
zap.Int("channels", len(h.config.AI.Channels)),
|
||||||
|
)
|
||||||
|
}
|
||||||
if req.OpenAI != nil {
|
if req.OpenAI != nil {
|
||||||
h.config.OpenAI = *req.OpenAI
|
h.config.OpenAI = *req.OpenAI
|
||||||
|
h.config.AI.EnsureDefaultFromOpenAI(h.config.OpenAI)
|
||||||
|
if def := config.NormalizeAIChannelID(h.config.AI.DefaultChannel); def != "" {
|
||||||
|
h.config.AI.Channels[def] = config.AIChannelFromOpenAI(def, "Default", h.config.OpenAI)
|
||||||
|
}
|
||||||
h.logger.Info("更新OpenAI配置",
|
h.logger.Info("更新OpenAI配置",
|
||||||
zap.String("base_url", h.config.OpenAI.BaseURL),
|
zap.String("base_url", h.config.OpenAI.BaseURL),
|
||||||
zap.String("model", h.config.OpenAI.Model),
|
zap.String("model", h.config.OpenAI.Model),
|
||||||
@@ -1517,6 +1532,7 @@ func (h *ConfigHandler) ApplyConfig(c *gin.Context) {
|
|||||||
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
h.executor.SetToolOutputMaxBytes(h.config.MultiAgent.EinoMiddleware.ReductionMaxLengthForTruncEffective())
|
||||||
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
h.executor.SetToolOutputSpillRoot(h.config.MultiAgent.EinoMiddleware.ReductionRootDir)
|
||||||
h.executor.RegisterTools(h.mcpServer)
|
h.executor.RegisterTools(h.mcpServer)
|
||||||
|
mcp.RegisterExecutionControlTools(h.mcpServer, h.externalMCPMgr)
|
||||||
|
|
||||||
// 重新注册漏洞记录工具(内置工具,必须注册)
|
// 重新注册漏洞记录工具(内置工具,必须注册)
|
||||||
if h.vulnerabilityToolRegistrar != nil {
|
if h.vulnerabilityToolRegistrar != nil {
|
||||||
@@ -1682,7 +1698,8 @@ func (h *ConfigHandler) saveConfig() error {
|
|||||||
|
|
||||||
updateAgentConfig(root, h.config.Agent)
|
updateAgentConfig(root, h.config.Agent)
|
||||||
updateMCPConfig(root, h.config.MCP)
|
updateMCPConfig(root, h.config.MCP)
|
||||||
updateOpenAIConfig(root, h.config.OpenAI)
|
updateAIConfig(root, h.config.AI)
|
||||||
|
removeKeyFromMap(root.Content[0], "openai")
|
||||||
updateVisionConfig(root, h.config.Vision)
|
updateVisionConfig(root, h.config.Vision)
|
||||||
updateFOFAConfig(root, h.config.FOFA)
|
updateFOFAConfig(root, h.config.FOFA)
|
||||||
updateSpaceSearchConfig(root, "zoomeye", h.config.ZoomEye)
|
updateSpaceSearchConfig(root, "zoomeye", h.config.ZoomEye)
|
||||||
@@ -1876,6 +1893,69 @@ func updateOpenAIConfig(doc *yaml.Node, cfg config.OpenAIConfig) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func updateAIConfig(doc *yaml.Node, cfg config.AIConfig) {
|
||||||
|
root := doc.Content[0]
|
||||||
|
aiNode := ensureMap(root, "ai")
|
||||||
|
if strings.TrimSpace(cfg.DefaultChannel) != "" {
|
||||||
|
setStringInMap(aiNode, "default_channel", config.NormalizeAIChannelID(cfg.DefaultChannel))
|
||||||
|
}
|
||||||
|
channelsNode := ensureMap(aiNode, "channels")
|
||||||
|
channelsNode.Content = nil
|
||||||
|
normalized := make(map[string]config.AIChannelConfig, len(cfg.Channels))
|
||||||
|
ids := make([]string, 0, len(cfg.Channels))
|
||||||
|
for id, ch := range cfg.Channels {
|
||||||
|
nid := config.NormalizeAIChannelID(id)
|
||||||
|
if nid == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if _, exists := normalized[nid]; !exists {
|
||||||
|
ids = append(ids, nid)
|
||||||
|
}
|
||||||
|
normalized[nid] = ch
|
||||||
|
}
|
||||||
|
sort.Strings(ids)
|
||||||
|
seen := make(map[string]bool, len(ids))
|
||||||
|
for _, id := range ids {
|
||||||
|
if seen[id] {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
seen[id] = true
|
||||||
|
ch := normalized[id]
|
||||||
|
keyNode := &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: id}
|
||||||
|
channelNode := &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"}
|
||||||
|
channelsNode.Content = append(channelsNode.Content, keyNode, channelNode)
|
||||||
|
setStringInMap(channelNode, "name", ch.Name)
|
||||||
|
if strings.TrimSpace(ch.Provider) != "" {
|
||||||
|
setStringInMap(channelNode, "provider", ch.Provider)
|
||||||
|
}
|
||||||
|
setStringInMap(channelNode, "api_key", ch.APIKey)
|
||||||
|
setStringInMap(channelNode, "base_url", ch.BaseURL)
|
||||||
|
setStringInMap(channelNode, "model", ch.Model)
|
||||||
|
if ch.MaxTotalTokens > 0 {
|
||||||
|
setIntInMap(channelNode, "max_total_tokens", ch.MaxTotalTokens)
|
||||||
|
}
|
||||||
|
if ch.MaxCompletionTokens > 0 {
|
||||||
|
setIntInMap(channelNode, "max_completion_tokens", ch.MaxCompletionTokens)
|
||||||
|
}
|
||||||
|
rn := ensureMap(channelNode, "reasoning")
|
||||||
|
if strings.TrimSpace(ch.Reasoning.Mode) != "" {
|
||||||
|
setStringInMap(rn, "mode", ch.Reasoning.Mode)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(ch.Reasoning.Effort) != "" {
|
||||||
|
setStringInMap(rn, "effort", ch.Reasoning.Effort)
|
||||||
|
}
|
||||||
|
if ch.Reasoning.AllowClientReasoning != nil {
|
||||||
|
setBoolInMap(rn, "allow_client_reasoning", *ch.Reasoning.AllowClientReasoning)
|
||||||
|
}
|
||||||
|
if strings.TrimSpace(ch.Reasoning.Profile) != "" {
|
||||||
|
setStringInMap(rn, "profile", ch.Reasoning.Profile)
|
||||||
|
}
|
||||||
|
if len(rn.Content) == 0 {
|
||||||
|
removeKeyFromMap(channelNode, "reasoning")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) {
|
func updateFOFAConfig(doc *yaml.Node, cfg config.FofaConfig) {
|
||||||
root := doc.Content[0]
|
root := doc.Content[0]
|
||||||
fofaNode := ensureMap(root, "fofa")
|
fofaNode := ensureMap(root, "fofa")
|
||||||
|
|||||||
@@ -68,15 +68,15 @@ func (h *AgentHandler) tryContinueOnEinoEmptyResponse(
|
|||||||
case <-time.After(backoff):
|
case <-time.After(backoff):
|
||||||
}
|
}
|
||||||
|
|
||||||
inject := multiagent.FormatEmptyResponseContinueUserMessage()
|
h.applyEinoTraceResumeSegment(conversationID, result, curHistory, curFinalMessage, "")
|
||||||
h.applyEinoTraceResumeSegment(conversationID, result, curHistory, curFinalMessage, inject)
|
|
||||||
if progressCallback != nil {
|
if progressCallback != nil {
|
||||||
progressCallback("eino_empty_response_continue", "已恢复上下文,正在续跑…", map[string]interface{}{
|
progressCallback("eino_empty_response_continue", "已恢复上下文,正在续跑…", map[string]interface{}{
|
||||||
"conversationId": conversationID,
|
"conversationId": conversationID,
|
||||||
"source": "eino",
|
"source": "eino",
|
||||||
"attempt": *attempt,
|
"attempt": *attempt,
|
||||||
"maxAttempts": maxAttempts,
|
"maxAttempts": maxAttempts,
|
||||||
"contextSource": "empty_response_continue",
|
"contextSource": "empty_response_continue",
|
||||||
|
"contextInjection": false,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
return true
|
return true
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ import (
|
|||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/agentfinalizer"
|
||||||
"cyberstrike-ai/internal/mcp"
|
"cyberstrike-ai/internal/mcp"
|
||||||
"cyberstrike-ai/internal/multiagent"
|
"cyberstrike-ai/internal/multiagent"
|
||||||
|
|
||||||
@@ -149,6 +150,14 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
runCfg, resolvedAIChannelID, err := h.configForAIChannel(req.AIChannelID)
|
||||||
|
if err != nil {
|
||||||
|
taskStatus = "failed"
|
||||||
|
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||||
|
sendEvent("error", err.Error(), nil)
|
||||||
|
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
var result *multiagent.RunResult
|
var result *multiagent.RunResult
|
||||||
var runErr error
|
var runErr error
|
||||||
@@ -181,6 +190,8 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
|
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
|
||||||
var mainIterationOffset int
|
var mainIterationOffset int
|
||||||
var emptyResponseContinueAttempt int
|
var emptyResponseContinueAttempt int
|
||||||
|
var finalizationAutoContinueAttempt int
|
||||||
|
var decision agentfinalizer.Decision
|
||||||
|
|
||||||
for {
|
for {
|
||||||
segmentMainIterationMax := 0
|
segmentMainIterationMax := 0
|
||||||
@@ -222,8 +233,8 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
|
|
||||||
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
||||||
taskCtxLoop,
|
taskCtxLoop,
|
||||||
h.config,
|
runCfg,
|
||||||
&h.config.MultiAgent,
|
&runCfg.MultiAgent,
|
||||||
h.agent,
|
h.agent,
|
||||||
h.db,
|
h.db,
|
||||||
h.logger,
|
h.logger,
|
||||||
@@ -236,6 +247,7 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
chatReasoningToClientIntent(req.Reasoning),
|
chatReasoningToClientIntent(req.Reasoning),
|
||||||
h.agentSessionContextBlock(conversationID),
|
h.agentSessionContextBlock(conversationID),
|
||||||
)
|
)
|
||||||
|
_ = resolvedAIChannelID
|
||||||
|
|
||||||
if result != nil && len(result.MCPExecutionIDs) > 0 {
|
if result != nil && len(result.MCPExecutionIDs) > 0 {
|
||||||
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
|
cumulativeMCPExecutionIDs = mergeMCPExecutionIDLists(cumulativeMCPExecutionIDs, result.MCPExecutionIDs)
|
||||||
@@ -249,6 +261,13 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel)
|
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||||
|
if h.tryAutoContinueAfterFinalization(taskCtx, conversationID, result, decision, &finalizationAutoContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
|
||||||
|
mainIterationOffset += segmentMainIterationMax
|
||||||
|
timeoutCancel()
|
||||||
|
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel)
|
||||||
|
continue
|
||||||
|
}
|
||||||
timeoutCancel()
|
timeoutCancel()
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
@@ -349,9 +368,10 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
|
|
||||||
timeoutCancel()
|
timeoutCancel()
|
||||||
|
|
||||||
if assistantMessageID != "" {
|
if decision.CompletionReason == "" {
|
||||||
_ = h.db.UpdateAssistantMessageFinalize(assistantMessageID, result.Response, cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput))
|
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, "eino_single", result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||||
}
|
}
|
||||||
|
h.persistFinalizationDecision(conversationID, assistantMessageID, "eino_single", cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision)
|
||||||
|
|
||||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||||
@@ -359,12 +379,19 @@ func (h *AgentHandler) EinoSingleAgentLoopStream(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
sendEvent("response", result.Response, map[string]interface{}{
|
responseText := decision.FinalText
|
||||||
|
if !decision.Finalizable {
|
||||||
|
responseText = finalizationBlockedMessage(decision)
|
||||||
|
sendEvent("finalization_check", responseText, decision)
|
||||||
|
taskStatus = decision.Status
|
||||||
|
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||||
|
}
|
||||||
|
sendEvent("response", responseText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||||
"conversationId": conversationID,
|
"conversationId": conversationID,
|
||||||
"messageId": assistantMessageID,
|
"messageId": assistantMessageID,
|
||||||
"agentMode": "eino_single",
|
"agentMode": "eino_single",
|
||||||
})
|
}))
|
||||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -410,16 +437,24 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
|||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器配置未加载"})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器配置未加载"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
runCfg, _, err := h.configForAIChannel(req.AIChannelID)
|
||||||
|
if err != nil {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
curHist := prep.History
|
curHist := prep.History
|
||||||
curMsg := prep.FinalMessage
|
curMsg := prep.FinalMessage
|
||||||
var result *multiagent.RunResult
|
var result *multiagent.RunResult
|
||||||
var runErr error
|
var runErr error
|
||||||
|
var emptyResponseContinueAttempt int
|
||||||
|
var finalizationAutoContinueAttempt int
|
||||||
|
var decision agentfinalizer.Decision
|
||||||
for {
|
for {
|
||||||
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
result, runErr = multiagent.RunEinoSingleChatModelAgent(
|
||||||
taskCtx,
|
taskCtx,
|
||||||
h.config,
|
runCfg,
|
||||||
&h.config.MultiAgent,
|
&runCfg.MultiAgent,
|
||||||
h.agent,
|
h.agent,
|
||||||
h.db,
|
h.db,
|
||||||
h.logger,
|
h.logger,
|
||||||
@@ -432,28 +467,46 @@ func (h *AgentHandler) EinoSingleAgentLoop(c *gin.Context) {
|
|||||||
chatReasoningToClientIntent(req.Reasoning),
|
chatReasoningToClientIntent(req.Reasoning),
|
||||||
h.agentSessionContextBlock(prep.ConversationID),
|
h.agentSessionContextBlock(prep.ConversationID),
|
||||||
)
|
)
|
||||||
if runErr == nil {
|
if runErr != nil {
|
||||||
break
|
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
||||||
|
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
||||||
|
}
|
||||||
|
c.JSON(http.StatusInternalServerError, gin.H{"error": runErr.Error()})
|
||||||
|
return
|
||||||
}
|
}
|
||||||
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
mw := &h.config.MultiAgent.EinoMiddleware
|
||||||
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, prep.ConversationID, result, &emptyResponseContinueAttempt, &curHist, &curMsg, progressCallback) {
|
||||||
|
continue
|
||||||
}
|
}
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": runErr.Error()})
|
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, "eino_single", result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||||
return
|
if h.tryAutoContinueAfterFinalization(taskCtx, prep.ConversationID, result, decision, &finalizationAutoContinueAttempt, &curHist, &curMsg, progressCallback) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
break
|
||||||
}
|
}
|
||||||
|
|
||||||
if prep.AssistantMessageID != "" {
|
h.persistFinalizationDecision(prep.ConversationID, prep.AssistantMessageID, "eino_single", result.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision)
|
||||||
_ = h.db.UpdateAssistantMessageFinalize(prep.AssistantMessageID, result.Response, result.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput))
|
|
||||||
}
|
|
||||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||||
_ = h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput)
|
_ = h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
responseText := decision.FinalText
|
||||||
|
if !decision.Finalizable {
|
||||||
|
responseText = finalizationBlockedMessage(decision)
|
||||||
|
}
|
||||||
c.JSON(http.StatusOK, gin.H{
|
c.JSON(http.StatusOK, gin.H{
|
||||||
"response": result.Response,
|
"response": responseText,
|
||||||
"conversationId": prep.ConversationID,
|
"conversationId": prep.ConversationID,
|
||||||
"mcpExecutionIds": result.MCPExecutionIDs,
|
"mcpExecutionIds": result.MCPExecutionIDs,
|
||||||
"assistantMessageId": prep.AssistantMessageID,
|
"assistantMessageId": prep.AssistantMessageID,
|
||||||
"agentMode": "eino_single",
|
"agentMode": "eino_single",
|
||||||
|
"finalized": decision.Finalized,
|
||||||
|
"finalizable": decision.Finalizable,
|
||||||
|
"status": decision.Status,
|
||||||
|
"completionReason": decision.CompletionReason,
|
||||||
|
"evidenceVerified": decision.EvidenceVerified,
|
||||||
|
"evidenceRefs": decision.EvidenceRefs,
|
||||||
|
"pendingExecutionIds": decision.PendingExecutionIDs,
|
||||||
|
"missingChecks": decision.MissingChecks,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,77 @@
|
|||||||
|
package handler
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/agent"
|
||||||
|
"cyberstrike-ai/internal/agentfinalizer"
|
||||||
|
"cyberstrike-ai/internal/multiagent"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
const finalizationAutoContinueMaxAttempts = 2
|
||||||
|
|
||||||
|
func shouldAutoContinueAfterFinalization(d agentfinalizer.Decision, attempt int) bool {
|
||||||
|
if d.Finalizable || d.Finalized {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if attempt >= finalizationAutoContinueMaxAttempts {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return d.CompletionReason == agentfinalizer.ReasonMissingEvidence
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *AgentHandler) tryAutoContinueAfterFinalization(
|
||||||
|
taskCtx context.Context,
|
||||||
|
conversationID string,
|
||||||
|
result *multiagent.RunResult,
|
||||||
|
decision agentfinalizer.Decision,
|
||||||
|
attempt *int,
|
||||||
|
curHistory *[]agent.ChatMessage,
|
||||||
|
curFinalMessage *string,
|
||||||
|
progressCallback func(eventType, message string, data interface{}),
|
||||||
|
) bool {
|
||||||
|
if !shouldAutoContinueAfterFinalization(decision, *attempt) || result == nil || !multiagent.HasEinoResumeTrace(result) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
*attempt++
|
||||||
|
h.persistEinoAgentTraceForResume(conversationID, result)
|
||||||
|
if hist, err := h.loadHistoryFromAgentTrace(conversationID); err == nil && len(hist) > 0 {
|
||||||
|
*curHistory = hist
|
||||||
|
} else if h.logger != nil {
|
||||||
|
h.logger.Warn("finalization auto-continue could not restore trace",
|
||||||
|
zap.String("conversationId", conversationID),
|
||||||
|
zap.Error(err))
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
// Agent 无感续跑:不追加新的 user/system 文案,只使用上一段模型可见轨迹继续 Runner。
|
||||||
|
*curFinalMessage = ""
|
||||||
|
if progressCallback != nil {
|
||||||
|
progressCallback("finalization_auto_continue", "最终回复检查尚未收敛,正在基于已有轨迹继续执行…", map[string]interface{}{
|
||||||
|
"conversationId": conversationID,
|
||||||
|
"source": "finalizer",
|
||||||
|
"attempt": *attempt,
|
||||||
|
"maxAttempts": finalizationAutoContinueMaxAttempts,
|
||||||
|
"status": decision.Status,
|
||||||
|
"completionReason": decision.CompletionReason,
|
||||||
|
"missingChecks": decision.MissingChecks,
|
||||||
|
"pendingExecutionIds": decision.PendingExecutionIDs,
|
||||||
|
"contextInjection": false,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-taskCtx.Done():
|
||||||
|
return false
|
||||||
|
case <-time.After(finalizationAutoContinueBackoff(*attempt)):
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func finalizationAutoContinueBackoff(attempt int) time.Duration {
|
||||||
|
if attempt <= 1 {
|
||||||
|
return 500 * time.Millisecond
|
||||||
|
}
|
||||||
|
return time.Duration(attempt) * time.Second
|
||||||
|
}
|
||||||
@@ -0,0 +1,59 @@
|
|||||||
|
package handler
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/agentfinalizer"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestShouldAutoContinueAfterFinalization(t *testing.T) {
|
||||||
|
missingEvidence := agentfinalizer.Decision{
|
||||||
|
Status: agentfinalizer.StatusBlocked,
|
||||||
|
CompletionReason: agentfinalizer.ReasonMissingEvidence,
|
||||||
|
}
|
||||||
|
if !shouldAutoContinueAfterFinalization(missingEvidence, 0) {
|
||||||
|
t.Fatal("missing execution evidence should trigger auto-continue")
|
||||||
|
}
|
||||||
|
if shouldAutoContinueAfterFinalization(missingEvidence, finalizationAutoContinueMaxAttempts) {
|
||||||
|
t.Fatal("auto-continue should stop at max attempts")
|
||||||
|
}
|
||||||
|
|
||||||
|
finalized := agentfinalizer.Decision{
|
||||||
|
Status: agentfinalizer.StatusCompleted,
|
||||||
|
CompletionReason: agentfinalizer.ReasonVerified,
|
||||||
|
Finalizable: true,
|
||||||
|
Finalized: true,
|
||||||
|
}
|
||||||
|
if shouldAutoContinueAfterFinalization(finalized, 0) {
|
||||||
|
t.Fatal("finalized decision should not auto-continue")
|
||||||
|
}
|
||||||
|
|
||||||
|
awaitingHITL := agentfinalizer.Decision{
|
||||||
|
Status: agentfinalizer.StatusAwaitingHITL,
|
||||||
|
CompletionReason: agentfinalizer.ReasonAwaitingHITL,
|
||||||
|
}
|
||||||
|
if shouldAutoContinueAfterFinalization(awaitingHITL, 0) {
|
||||||
|
t.Fatal("awaiting HITL should not auto-continue without approval")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRequestRequiresExecutionEvidenceUsesExplicitPolicyOnly(t *testing.T) {
|
||||||
|
if requestRequiresExecutionEvidence(nil) {
|
||||||
|
t.Fatal("nil request should not require execution evidence")
|
||||||
|
}
|
||||||
|
if requestRequiresExecutionEvidence(&ChatRequest{}) {
|
||||||
|
t.Fatal("missing finalization policy should not require execution evidence")
|
||||||
|
}
|
||||||
|
require := true
|
||||||
|
if !requestRequiresExecutionEvidence(&ChatRequest{
|
||||||
|
Finalization: ChatFinalizationRequest{RequireExecutionEvidence: &require},
|
||||||
|
}) {
|
||||||
|
t.Fatal("explicit true policy should require execution evidence")
|
||||||
|
}
|
||||||
|
require = false
|
||||||
|
if requestRequiresExecutionEvidence(&ChatRequest{
|
||||||
|
Finalization: ChatFinalizationRequest{RequireExecutionEvidence: &require},
|
||||||
|
}) {
|
||||||
|
t.Fatal("explicit false policy should not require execution evidence")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,171 @@
|
|||||||
|
package handler
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/agentfinalizer"
|
||||||
|
"cyberstrike-ai/internal/multiagent"
|
||||||
|
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
func (h *AgentHandler) finalizeAgentRunForDelivery(
|
||||||
|
conversationID string,
|
||||||
|
assistantMessageID string,
|
||||||
|
agentMode string,
|
||||||
|
result *multiagent.RunResult,
|
||||||
|
mcpExecutionIDs []string,
|
||||||
|
reasoningContent string,
|
||||||
|
) agentfinalizer.Decision {
|
||||||
|
return h.finalizeAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, mcpExecutionIDs, reasoningContent, false)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *AgentHandler) finalizeAgentRunForDeliveryWithPolicy(
|
||||||
|
conversationID string,
|
||||||
|
assistantMessageID string,
|
||||||
|
agentMode string,
|
||||||
|
result *multiagent.RunResult,
|
||||||
|
mcpExecutionIDs []string,
|
||||||
|
reasoningContent string,
|
||||||
|
requireExecutionEvidence bool,
|
||||||
|
) agentfinalizer.Decision {
|
||||||
|
decision := agentfinalizer.FromRunResult(h.db, result, agentfinalizer.Input{
|
||||||
|
ConversationID: conversationID,
|
||||||
|
AssistantMessageID: assistantMessageID,
|
||||||
|
AgentMode: agentMode,
|
||||||
|
MCPExecutionIDs: mcpExecutionIDs,
|
||||||
|
RequireExecutionEvidence: requireExecutionEvidence,
|
||||||
|
})
|
||||||
|
h.persistFinalizationDecision(conversationID, assistantMessageID, agentMode, mcpExecutionIDs, reasoningContent, decision)
|
||||||
|
return decision
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *AgentHandler) decideAgentRunForDeliveryWithPolicy(
|
||||||
|
conversationID string,
|
||||||
|
assistantMessageID string,
|
||||||
|
agentMode string,
|
||||||
|
result *multiagent.RunResult,
|
||||||
|
mcpExecutionIDs []string,
|
||||||
|
requireExecutionEvidence bool,
|
||||||
|
) agentfinalizer.Decision {
|
||||||
|
return agentfinalizer.FromRunResult(h.db, result, agentfinalizer.Input{
|
||||||
|
ConversationID: conversationID,
|
||||||
|
AssistantMessageID: assistantMessageID,
|
||||||
|
AgentMode: agentMode,
|
||||||
|
MCPExecutionIDs: mcpExecutionIDs,
|
||||||
|
RequireExecutionEvidence: requireExecutionEvidence,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *AgentHandler) decideAgentRunForDelivery(
|
||||||
|
conversationID string,
|
||||||
|
assistantMessageID string,
|
||||||
|
agentMode string,
|
||||||
|
result *multiagent.RunResult,
|
||||||
|
mcpExecutionIDs []string,
|
||||||
|
) agentfinalizer.Decision {
|
||||||
|
return agentfinalizer.FromRunResult(h.db, result, agentfinalizer.Input{
|
||||||
|
ConversationID: conversationID,
|
||||||
|
AssistantMessageID: assistantMessageID,
|
||||||
|
AgentMode: agentMode,
|
||||||
|
MCPExecutionIDs: mcpExecutionIDs,
|
||||||
|
RequireExecutionEvidence: false,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *AgentHandler) persistFinalizationDecision(
|
||||||
|
conversationID string,
|
||||||
|
assistantMessageID string,
|
||||||
|
agentMode string,
|
||||||
|
mcpExecutionIDs []string,
|
||||||
|
reasoningContent string,
|
||||||
|
decision agentfinalizer.Decision,
|
||||||
|
) {
|
||||||
|
if assistantMessageID == "" || h.db == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "finalization_check", finalizationCheckMessage(decision), decision)
|
||||||
|
if decision.Finalizable {
|
||||||
|
if err := h.db.UpdateAssistantMessageFinalize(assistantMessageID, decision.FinalText, mcpExecutionIDs, reasoningContent); err != nil && h.logger != nil {
|
||||||
|
h.logger.Warn("更新最终助手消息失败", zap.Error(err), zap.String("conversationId", conversationID), zap.String("agentMode", agentMode))
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", finalizationBlockedMessage(decision), time.Now(), assistantMessageID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *AgentHandler) finalizeCandidateForDelivery(
|
||||||
|
conversationID string,
|
||||||
|
assistantMessageID string,
|
||||||
|
agentMode string,
|
||||||
|
response string,
|
||||||
|
mcpExecutionIDs []string,
|
||||||
|
awaitingHITL bool,
|
||||||
|
reasoningContent string,
|
||||||
|
) agentfinalizer.Decision {
|
||||||
|
return h.finalizeCandidateForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, response, mcpExecutionIDs, awaitingHITL, reasoningContent, false)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *AgentHandler) finalizeCandidateForDeliveryWithPolicy(
|
||||||
|
conversationID string,
|
||||||
|
assistantMessageID string,
|
||||||
|
agentMode string,
|
||||||
|
response string,
|
||||||
|
mcpExecutionIDs []string,
|
||||||
|
awaitingHITL bool,
|
||||||
|
reasoningContent string,
|
||||||
|
requireExecutionEvidence bool,
|
||||||
|
) agentfinalizer.Decision {
|
||||||
|
decision := agentfinalizer.Decide(h.db, agentfinalizer.Input{
|
||||||
|
Response: response,
|
||||||
|
ConversationID: conversationID,
|
||||||
|
AssistantMessageID: assistantMessageID,
|
||||||
|
AgentMode: agentMode,
|
||||||
|
MCPExecutionIDs: mcpExecutionIDs,
|
||||||
|
AwaitingHITL: awaitingHITL,
|
||||||
|
RequireExecutionEvidence: requireExecutionEvidence,
|
||||||
|
})
|
||||||
|
if assistantMessageID == "" || h.db == nil {
|
||||||
|
return decision
|
||||||
|
}
|
||||||
|
_ = h.db.AddProcessDetail(assistantMessageID, conversationID, "finalization_check", finalizationCheckMessage(decision), decision)
|
||||||
|
if decision.Finalizable {
|
||||||
|
if err := h.db.UpdateAssistantMessageFinalize(assistantMessageID, decision.FinalText, mcpExecutionIDs, reasoningContent); err != nil && h.logger != nil {
|
||||||
|
h.logger.Warn("更新最终助手消息失败", zap.Error(err), zap.String("conversationId", conversationID), zap.String("agentMode", agentMode))
|
||||||
|
}
|
||||||
|
return decision
|
||||||
|
}
|
||||||
|
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", finalizationBlockedMessage(decision), time.Now(), assistantMessageID)
|
||||||
|
return decision
|
||||||
|
}
|
||||||
|
|
||||||
|
func finalizationCheckMessage(d agentfinalizer.Decision) string {
|
||||||
|
if d.Finalizable {
|
||||||
|
return "最终回复检查通过。"
|
||||||
|
}
|
||||||
|
return finalizationBlockedMessage(d)
|
||||||
|
}
|
||||||
|
|
||||||
|
func finalizationBlockedMessage(d agentfinalizer.Decision) string {
|
||||||
|
parts := []string{"任务尚未达到最终回复条件,暂不生成成功结论。"}
|
||||||
|
if d.CompletionReason != "" {
|
||||||
|
parts = append(parts, "原因: "+d.CompletionReason)
|
||||||
|
}
|
||||||
|
if len(d.PendingExecutionIDs) > 0 {
|
||||||
|
parts = append(parts, fmt.Sprintf("仍有 %d 个工具执行未结束: %s", len(d.PendingExecutionIDs), strings.Join(d.PendingExecutionIDs, ", ")))
|
||||||
|
}
|
||||||
|
if len(d.MissingChecks) > 0 {
|
||||||
|
parts = append(parts, "缺失检查: "+strings.Join(d.MissingChecks, "; "))
|
||||||
|
}
|
||||||
|
return strings.Join(parts, "\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
func finalizationResponsePayload(d agentfinalizer.Decision, extra map[string]interface{}) map[string]interface{} {
|
||||||
|
return agentfinalizer.ResponsePayload(d, extra)
|
||||||
|
}
|
||||||
|
|
||||||
|
func requestRequiresExecutionEvidence(req *ChatRequest) bool {
|
||||||
|
return req != nil && req.Finalization.RequireExecutionEvidence != nil && *req.Finalization.RequireExecutionEvidence
|
||||||
|
}
|
||||||
@@ -69,7 +69,7 @@ func (h *MonitorHandler) SetAgentHandler(ah *AgentHandler) {
|
|||||||
h.agentHandler = ah
|
h.agentHandler = ah
|
||||||
}
|
}
|
||||||
|
|
||||||
const monitorPageTopTools = 3
|
const monitorPageTopTools = 6
|
||||||
|
|
||||||
// MonitorStatsSummary 工具调用汇总
|
// MonitorStatsSummary 工具调用汇总
|
||||||
type MonitorStatsSummary struct {
|
type MonitorStatsSummary struct {
|
||||||
@@ -182,7 +182,7 @@ func summarizeAccessibleExecutionPage(executions []*mcp.ToolExecution, topN int)
|
|||||||
|
|
||||||
func monitorStatusCountsAsFailed(status string) bool {
|
func monitorStatusCountsAsFailed(status string) bool {
|
||||||
switch strings.TrimSpace(strings.ToLower(status)) {
|
switch strings.TrimSpace(strings.ToLower(status)) {
|
||||||
case "failed", "cancelled", "hard_timeout", "orphaned":
|
case "failed", "hard_timeout", "orphaned":
|
||||||
return true
|
return true
|
||||||
default:
|
default:
|
||||||
return false
|
return false
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ import (
|
|||||||
"sync"
|
"sync"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/agentfinalizer"
|
||||||
"cyberstrike-ai/internal/config"
|
"cyberstrike-ai/internal/config"
|
||||||
"cyberstrike-ai/internal/mcp"
|
"cyberstrike-ai/internal/mcp"
|
||||||
"cyberstrike-ai/internal/multiagent"
|
"cyberstrike-ai/internal/multiagent"
|
||||||
@@ -158,6 +159,12 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
|||||||
|
|
||||||
stopKeepalive := runSSEKeepalive(c, &sseWriteMu)
|
stopKeepalive := runSSEKeepalive(c, &sseWriteMu)
|
||||||
defer stopKeepalive()
|
defer stopKeepalive()
|
||||||
|
runCfg, _, err := h.configForAIChannel(req.AIChannelID)
|
||||||
|
if err != nil {
|
||||||
|
sendEvent("error", err.Error(), nil)
|
||||||
|
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
var result *multiagent.RunResult
|
var result *multiagent.RunResult
|
||||||
var runErr error
|
var runErr error
|
||||||
@@ -191,6 +198,13 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
|||||||
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
|
// 同一请求内分段续跑时,主代理 iteration 事件按偏移累计,避免 UI 出现「第3轮 → 第1轮」回跳。
|
||||||
var mainIterationOffset int
|
var mainIterationOffset int
|
||||||
var emptyResponseContinueAttempt int
|
var emptyResponseContinueAttempt int
|
||||||
|
var finalizationAutoContinueAttempt int
|
||||||
|
effectiveOrch := config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration)
|
||||||
|
if o := strings.TrimSpace(req.Orchestration); o != "" {
|
||||||
|
effectiveOrch = config.NormalizeMultiAgentOrchestration(o)
|
||||||
|
}
|
||||||
|
agentMode := "eino_" + effectiveOrch
|
||||||
|
var decision agentfinalizer.Decision
|
||||||
|
|
||||||
for {
|
for {
|
||||||
segmentMainIterationMax := 0
|
segmentMainIterationMax := 0
|
||||||
@@ -232,8 +246,8 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
|||||||
|
|
||||||
result, runErr = multiagent.RunDeepAgent(
|
result, runErr = multiagent.RunDeepAgent(
|
||||||
taskCtxLoop,
|
taskCtxLoop,
|
||||||
h.config,
|
runCfg,
|
||||||
&h.config.MultiAgent,
|
&runCfg.MultiAgent,
|
||||||
h.agent,
|
h.agent,
|
||||||
h.db,
|
h.db,
|
||||||
h.logger,
|
h.logger,
|
||||||
@@ -261,6 +275,13 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
|||||||
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel)
|
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel)
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||||
|
if h.tryAutoContinueAfterFinalization(taskCtx, conversationID, result, decision, &finalizationAutoContinueAttempt, &curHistory, &curFinalMessage, progressCallback) {
|
||||||
|
mainIterationOffset += segmentMainIterationMax
|
||||||
|
timeoutCancel()
|
||||||
|
baseCtx, cancelWithCause, taskCtx, timeoutCancel = h.rebindEinoRunningTask(taskCtx, conversationID, timeoutCancel)
|
||||||
|
continue
|
||||||
|
}
|
||||||
timeoutCancel()
|
timeoutCancel()
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
@@ -361,9 +382,10 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
|||||||
|
|
||||||
timeoutCancel()
|
timeoutCancel()
|
||||||
|
|
||||||
if assistantMessageID != "" {
|
if decision.CompletionReason == "" {
|
||||||
_ = h.db.UpdateAssistantMessageFinalize(assistantMessageID, result.Response, cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput))
|
decision = h.decideAgentRunForDeliveryWithPolicy(conversationID, assistantMessageID, agentMode, result, cumulativeMCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||||
}
|
}
|
||||||
|
h.persistFinalizationDecision(conversationID, assistantMessageID, agentMode, cumulativeMCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision)
|
||||||
|
|
||||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||||
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
if err := h.db.SaveAgentTrace(conversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||||
@@ -371,16 +393,19 @@ func (h *AgentHandler) MultiAgentLoopStream(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
effectiveOrch := config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration)
|
responseText := decision.FinalText
|
||||||
if o := strings.TrimSpace(req.Orchestration); o != "" {
|
if !decision.Finalizable {
|
||||||
effectiveOrch = config.NormalizeMultiAgentOrchestration(o)
|
responseText = finalizationBlockedMessage(decision)
|
||||||
|
sendEvent("finalization_check", responseText, decision)
|
||||||
|
taskStatus = decision.Status
|
||||||
|
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||||
}
|
}
|
||||||
sendEvent("response", result.Response, map[string]interface{}{
|
sendEvent("response", responseText, finalizationResponsePayload(decision, map[string]interface{}{
|
||||||
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
"mcpExecutionIds": cumulativeMCPExecutionIDs,
|
||||||
"conversationId": conversationID,
|
"conversationId": conversationID,
|
||||||
"messageId": assistantMessageID,
|
"messageId": assistantMessageID,
|
||||||
"agentMode": "eino_" + effectiveOrch,
|
"agentMode": agentMode,
|
||||||
})
|
}))
|
||||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -421,16 +446,29 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
|||||||
taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) {
|
taskCtx = multiagent.WithHITLToolInterceptor(taskCtx, func(ctx context.Context, toolName, arguments string) (string, error) {
|
||||||
return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments)
|
return h.interceptHITLForEinoTool(ctx, cancelWithCause, prep.ConversationID, prep.AssistantMessageID, nil, toolName, arguments)
|
||||||
})
|
})
|
||||||
|
runCfg, _, err := h.configForAIChannel(req.AIChannelID)
|
||||||
|
if err != nil {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
curHist := prep.History
|
curHist := prep.History
|
||||||
curMsg := prep.FinalMessage
|
curMsg := prep.FinalMessage
|
||||||
var result *multiagent.RunResult
|
var result *multiagent.RunResult
|
||||||
var runErr error
|
var runErr error
|
||||||
|
var emptyResponseContinueAttempt int
|
||||||
|
var finalizationAutoContinueAttempt int
|
||||||
|
effectiveOrch := config.NormalizeMultiAgentOrchestration(h.config.MultiAgent.Orchestration)
|
||||||
|
if o := strings.TrimSpace(req.Orchestration); o != "" {
|
||||||
|
effectiveOrch = config.NormalizeMultiAgentOrchestration(o)
|
||||||
|
}
|
||||||
|
agentMode := "eino_" + effectiveOrch
|
||||||
|
var decision agentfinalizer.Decision
|
||||||
for {
|
for {
|
||||||
result, runErr = multiagent.RunDeepAgent(
|
result, runErr = multiagent.RunDeepAgent(
|
||||||
taskCtx,
|
taskCtx,
|
||||||
h.config,
|
runCfg,
|
||||||
&h.config.MultiAgent,
|
&runCfg.MultiAgent,
|
||||||
h.agent,
|
h.agent,
|
||||||
h.db,
|
h.db,
|
||||||
h.logger,
|
h.logger,
|
||||||
@@ -445,24 +483,30 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
|||||||
chatReasoningToClientIntent(req.Reasoning),
|
chatReasoningToClientIntent(req.Reasoning),
|
||||||
h.agentSessionContextBlock(prep.ConversationID),
|
h.agentSessionContextBlock(prep.ConversationID),
|
||||||
)
|
)
|
||||||
if runErr == nil {
|
if runErr != nil {
|
||||||
break
|
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
||||||
|
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
||||||
|
}
|
||||||
|
h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr))
|
||||||
|
errMsg := "执行失败: " + runErr.Error()
|
||||||
|
if prep.AssistantMessageID != "" {
|
||||||
|
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), prep.AssistantMessageID)
|
||||||
|
}
|
||||||
|
c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg})
|
||||||
|
return
|
||||||
}
|
}
|
||||||
if shouldPersistEinoAgentTraceAfterRunError(baseCtx) {
|
mw := &h.config.MultiAgent.EinoMiddleware
|
||||||
h.persistEinoAgentTraceForResume(prep.ConversationID, result)
|
if h.tryContinueOnEinoEmptyResponse(taskCtx, mw, prep.ConversationID, result, &emptyResponseContinueAttempt, &curHist, &curMsg, progressCallback) {
|
||||||
|
continue
|
||||||
}
|
}
|
||||||
h.logger.Error("Eino DeepAgent 执行失败", zap.Error(runErr))
|
decision = h.decideAgentRunForDeliveryWithPolicy(prep.ConversationID, prep.AssistantMessageID, agentMode, result, result.MCPExecutionIDs, requestRequiresExecutionEvidence(&req))
|
||||||
errMsg := "执行失败: " + runErr.Error()
|
if h.tryAutoContinueAfterFinalization(taskCtx, prep.ConversationID, result, decision, &finalizationAutoContinueAttempt, &curHist, &curMsg, progressCallback) {
|
||||||
if prep.AssistantMessageID != "" {
|
continue
|
||||||
_, _ = h.db.Exec("UPDATE messages SET content = ?, updated_at = ? WHERE id = ?", errMsg, time.Now(), prep.AssistantMessageID)
|
|
||||||
}
|
}
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg})
|
break
|
||||||
return
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if prep.AssistantMessageID != "" {
|
h.persistFinalizationDecision(prep.ConversationID, prep.AssistantMessageID, agentMode, result.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput), decision)
|
||||||
_ = h.db.UpdateAssistantMessageFinalize(prep.AssistantMessageID, result.Response, result.MCPExecutionIDs, multiagent.AggregatedReasoningFromTraceJSON(result.LastAgentTraceInput))
|
|
||||||
}
|
|
||||||
|
|
||||||
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
if result.LastAgentTraceInput != "" || result.LastAgentTraceOutput != "" {
|
||||||
if err := h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
if err := h.db.SaveAgentTrace(prep.ConversationID, result.LastAgentTraceInput, result.LastAgentTraceOutput); err != nil {
|
||||||
@@ -470,11 +514,23 @@ func (h *AgentHandler) MultiAgentLoop(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
responseText := decision.FinalText
|
||||||
|
if !decision.Finalizable {
|
||||||
|
responseText = finalizationBlockedMessage(decision)
|
||||||
|
}
|
||||||
c.JSON(http.StatusOK, ChatResponse{
|
c.JSON(http.StatusOK, ChatResponse{
|
||||||
Response: result.Response,
|
Response: responseText,
|
||||||
MCPExecutionIDs: result.MCPExecutionIDs,
|
MCPExecutionIDs: result.MCPExecutionIDs,
|
||||||
ConversationID: prep.ConversationID,
|
ConversationID: prep.ConversationID,
|
||||||
Time: time.Now(),
|
Time: time.Now(),
|
||||||
|
Finalizable: decision.Finalizable,
|
||||||
|
Finalized: decision.Finalized,
|
||||||
|
Status: decision.Status,
|
||||||
|
CompletionReason: decision.CompletionReason,
|
||||||
|
EvidenceVerified: decision.EvidenceVerified,
|
||||||
|
EvidenceRefs: decision.EvidenceRefs,
|
||||||
|
PendingExecutionIDs: decision.PendingExecutionIDs,
|
||||||
|
MissingChecks: decision.MissingChecks,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import (
|
|||||||
|
|
||||||
"cyberstrike-ai/internal/agent"
|
"cyberstrike-ai/internal/agent"
|
||||||
"cyberstrike-ai/internal/audit"
|
"cyberstrike-ai/internal/audit"
|
||||||
|
"cyberstrike-ai/internal/config"
|
||||||
"cyberstrike-ai/internal/database"
|
"cyberstrike-ai/internal/database"
|
||||||
"cyberstrike-ai/internal/mcp/builtin"
|
"cyberstrike-ai/internal/mcp/builtin"
|
||||||
"cyberstrike-ai/internal/security"
|
"cyberstrike-ai/internal/security"
|
||||||
@@ -25,6 +26,13 @@ type multiAgentPrepared struct {
|
|||||||
UserMessageID string
|
UserMessageID string
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func chatRequestAgentMode(req *ChatRequest, source string) string {
|
||||||
|
if strings.HasPrefix(strings.TrimSpace(source), "multi_agent") {
|
||||||
|
return config.NormalizeMultiAgentOrchestration(req.Orchestration)
|
||||||
|
}
|
||||||
|
return "eino_single"
|
||||||
|
}
|
||||||
|
|
||||||
func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context, source string) (*multiAgentPrepared, error) {
|
func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context, source string) (*multiAgentPrepared, error) {
|
||||||
if len(req.Attachments) > maxAttachments {
|
if len(req.Attachments) > maxAttachments {
|
||||||
return nil, fmt.Errorf("附件最多 %d 个", maxAttachments)
|
return nil, fmt.Errorf("附件最多 %d 个", maxAttachments)
|
||||||
@@ -56,6 +64,8 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
|||||||
var err error
|
var err error
|
||||||
meta := audit.ConversationCreateMetaFromGin(c, source)
|
meta := audit.ConversationCreateMetaFromGin(c, source)
|
||||||
meta.ProjectID = projectID
|
meta.ProjectID = projectID
|
||||||
|
meta.RoleName = req.Role
|
||||||
|
meta.AgentMode = chatRequestAgentMode(req, source)
|
||||||
if webshellID != "" {
|
if webshellID != "" {
|
||||||
meta.Source = source + "_webshell"
|
meta.Source = source + "_webshell"
|
||||||
meta.WebShellConnectionID = webshellID
|
meta.WebShellConnectionID = webshellID
|
||||||
@@ -80,6 +90,12 @@ func (h *AgentHandler) prepareMultiAgentSession(req *ChatRequest, c *gin.Context
|
|||||||
return nil, fmt.Errorf("无权访问该对话")
|
return nil, fmt.Errorf("无权访问该对话")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if err := h.db.SetConversationRoleName(conversationID, req.Role); err != nil {
|
||||||
|
h.logger.Warn("更新对话角色失败", zap.String("conversationId", conversationID), zap.String("role", req.Role), zap.Error(err))
|
||||||
|
}
|
||||||
|
if err := h.db.SetConversationAgentMode(conversationID, chatRequestAgentMode(req, source)); err != nil {
|
||||||
|
h.logger.Warn("更新对话模式失败", zap.String("conversationId", conversationID), zap.String("source", source), zap.String("orchestration", req.Orchestration), zap.Error(err))
|
||||||
|
}
|
||||||
|
|
||||||
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
agentHistoryMessages, err := h.loadHistoryFromAgentTrace(conversationID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
+176
-15
@@ -35,6 +35,17 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
scheme = "https"
|
scheme = "https"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
finalizationRequestSchema := map[string]interface{}{
|
||||||
|
"type": "object",
|
||||||
|
"description": "最终回复交付策略。后端不会从自然语言内容推断执行意图;执行入口应显式声明是否要求 completed 工具证据。",
|
||||||
|
"properties": map[string]interface{}{
|
||||||
|
"requireExecutionEvidence": map[string]interface{}{
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "为 true 时,缺少 completed 工具执行记录会触发无注入续跑或最终阻断;普通聊天可省略或设为 false。",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
spec := map[string]interface{}{
|
spec := map[string]interface{}{
|
||||||
"openapi": "3.0.0",
|
"openapi": "3.0.0",
|
||||||
"info": map[string]interface{}{
|
"info": map[string]interface{}{
|
||||||
@@ -85,6 +96,70 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
},
|
},
|
||||||
"required": []string{"projectId"},
|
"required": []string{"projectId"},
|
||||||
},
|
},
|
||||||
|
"AgentChatResponse": map[string]interface{}{
|
||||||
|
"type": "object",
|
||||||
|
"description": "Agent 非流式响应。response 只是交付文本;是否为成功最终回复必须以 finalized/finalizable/status 为准。",
|
||||||
|
"properties": map[string]interface{}{
|
||||||
|
"response": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "交付给用户的文本。finalized=false 时为阻断/未完成说明,不是成功结论。",
|
||||||
|
},
|
||||||
|
"conversationId": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "对话 ID",
|
||||||
|
},
|
||||||
|
"assistantMessageId": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "助手消息 ID(部分接口返回)",
|
||||||
|
},
|
||||||
|
"mcpExecutionIds": map[string]interface{}{
|
||||||
|
"type": "array",
|
||||||
|
"description": "本轮关联的 MCP 工具执行 ID",
|
||||||
|
"items": map[string]interface{}{"type": "string"},
|
||||||
|
},
|
||||||
|
"agentMode": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "agent 模式,例如 eino_single、eino_deep、workflow",
|
||||||
|
},
|
||||||
|
"finalized": map[string]interface{}{
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "是否已经通过最终回复检查。只有 true 才能当成功最终回复。",
|
||||||
|
},
|
||||||
|
"finalizable": map[string]interface{}{
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "候选输出是否可提升为最终回复。",
|
||||||
|
},
|
||||||
|
"status": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "最终化状态",
|
||||||
|
"enum": []string{"completed", "in_progress", "blocked", "failed", "cancelled", "awaiting_hitl"},
|
||||||
|
},
|
||||||
|
"completionReason": map[string]interface{}{
|
||||||
|
"type": "string",
|
||||||
|
"description": "最终化或阻断原因,例如 verified、pending_tool_executions、missing_execution_evidence",
|
||||||
|
},
|
||||||
|
"evidenceVerified": map[string]interface{}{
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "证据是否满足最终化要求",
|
||||||
|
},
|
||||||
|
"evidenceRefs": map[string]interface{}{
|
||||||
|
"type": "array",
|
||||||
|
"description": "证据引用,例如 mcp_execution:<id>",
|
||||||
|
"items": map[string]interface{}{"type": "string"},
|
||||||
|
},
|
||||||
|
"pendingExecutionIds": map[string]interface{}{
|
||||||
|
"type": "array",
|
||||||
|
"description": "仍处于 queued/running 的工具执行 ID",
|
||||||
|
"items": map[string]interface{}{"type": "string"},
|
||||||
|
},
|
||||||
|
"missingChecks": map[string]interface{}{
|
||||||
|
"type": "array",
|
||||||
|
"description": "未通过最终化检查的原因列表",
|
||||||
|
"items": map[string]interface{}{"type": "string"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": []string{"response", "conversationId", "finalized", "finalizable", "status", "evidenceVerified"},
|
||||||
|
},
|
||||||
"Conversation": map[string]interface{}{
|
"Conversation": map[string]interface{}{
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": map[string]interface{}{
|
"properties": map[string]interface{}{
|
||||||
@@ -1581,6 +1656,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"conversationId": map[string]interface{}{"type": "string"},
|
"conversationId": map[string]interface{}{"type": "string"},
|
||||||
"role": map[string]interface{}{"type": "string"},
|
"role": map[string]interface{}{"type": "string"},
|
||||||
"webshellConnectionId": map[string]interface{}{"type": "string"},
|
"webshellConnectionId": map[string]interface{}{"type": "string"},
|
||||||
|
"finalization": finalizationRequestSchema,
|
||||||
},
|
},
|
||||||
"required": []string{"message"},
|
"required": []string{"message"},
|
||||||
},
|
},
|
||||||
@@ -1588,7 +1664,14 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
"responses": map[string]interface{}{
|
"responses": map[string]interface{}{
|
||||||
"200": map[string]interface{}{"description": "成功,响应格式同 /api/eino-agent"},
|
"200": map[string]interface{}{
|
||||||
|
"description": "成功。只有 finalized=true 表示成功最终回复;finalized=false 时 response 为未完成/阻断说明。",
|
||||||
|
"content": map[string]interface{}{
|
||||||
|
"application/json": map[string]interface{}{
|
||||||
|
"schema": map[string]interface{}{"$ref": "#/components/schemas/AgentChatResponse"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
"400": map[string]interface{}{"description": "参数错误"},
|
"400": map[string]interface{}{"description": "参数错误"},
|
||||||
"401": map[string]interface{}{"description": "未授权"},
|
"401": map[string]interface{}{"description": "未授权"},
|
||||||
"500": map[string]interface{}{"description": "执行失败"},
|
"500": map[string]interface{}{"description": "执行失败"},
|
||||||
@@ -1599,7 +1682,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"post": map[string]interface{}{
|
"post": map[string]interface{}{
|
||||||
"tags": []string{"对话交互"},
|
"tags": []string{"对话交互"},
|
||||||
"summary": "发送消息并获取 AI 回复(Eino ADK 单代理,SSE)",
|
"summary": "发送消息并获取 AI 回复(Eino ADK 单代理,SSE)",
|
||||||
"description": "向 AI 发送消息并获取流式回复(SSE)。由 Eino **单代理** ADK 执行;事件类型与多代理流式一致(含 `tool_call` / `response_delta` / `thinking` 等)。**不依赖** `multi_agent.enabled`。",
|
"description": "向 AI 发送消息并获取流式回复(SSE)。由 Eino **单代理** ADK 执行;事件类型与多代理流式一致(含 `tool_call` / `response_delta` / `thinking` 等)。`response_start` / `response_delta` 仅为候选/过程输出;只有 `type: response` 且 `data.finalized=true` 才表示成功最终回复。缺 completed 执行证据时可能先发送 `finalization_auto_continue`,表示服务端基于已有 trace 无注入续跑。`data.finalized=false` 时 message 为未完成/阻断说明。**不依赖** `multi_agent.enabled`。",
|
||||||
"operationId": "sendMessageEinoSingleAgentStream",
|
"operationId": "sendMessageEinoSingleAgentStream",
|
||||||
"requestBody": map[string]interface{}{
|
"requestBody": map[string]interface{}{
|
||||||
"required": true,
|
"required": true,
|
||||||
@@ -1612,6 +1695,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"conversationId": map[string]interface{}{"type": "string"},
|
"conversationId": map[string]interface{}{"type": "string"},
|
||||||
"role": map[string]interface{}{"type": "string"},
|
"role": map[string]interface{}{"type": "string"},
|
||||||
"webshellConnectionId": map[string]interface{}{"type": "string"},
|
"webshellConnectionId": map[string]interface{}{"type": "string"},
|
||||||
|
"finalization": finalizationRequestSchema,
|
||||||
},
|
},
|
||||||
"required": []string{"message"},
|
"required": []string{"message"},
|
||||||
},
|
},
|
||||||
@@ -1625,7 +1709,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"text/event-stream": map[string]interface{}{
|
"text/event-stream": map[string]interface{}{
|
||||||
"schema": map[string]interface{}{
|
"schema": map[string]interface{}{
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "SSE 流",
|
"description": "SSE 流。终态 response 事件 data 包含 finalized、finalizable、status、completionReason、evidenceVerified、evidenceRefs、pendingExecutionIds、missingChecks;过程事件可能包含 finalization_auto_continue。",
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -1663,6 +1747,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "WebShell 连接 ID(可选,与 Eino 单/多代理流式行为一致)",
|
"description": "WebShell 连接 ID(可选,与 Eino 单/多代理流式行为一致)",
|
||||||
},
|
},
|
||||||
|
"finalization": finalizationRequestSchema,
|
||||||
"orchestration": map[string]interface{}{
|
"orchestration": map[string]interface{}{
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "Eino 预置编排:deep | plan_execute | supervisor;缺省 deep",
|
"description": "Eino 预置编排:deep | plan_execute | supervisor;缺省 deep",
|
||||||
@@ -1676,7 +1761,12 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
},
|
},
|
||||||
"responses": map[string]interface{}{
|
"responses": map[string]interface{}{
|
||||||
"200": map[string]interface{}{
|
"200": map[string]interface{}{
|
||||||
"description": "成功,响应格式同 /api/eino-agent",
|
"description": "成功。只有 finalized=true 表示成功最终回复;finalized=false 时 response 为未完成/阻断说明。",
|
||||||
|
"content": map[string]interface{}{
|
||||||
|
"application/json": map[string]interface{}{
|
||||||
|
"schema": map[string]interface{}{"$ref": "#/components/schemas/AgentChatResponse"},
|
||||||
|
},
|
||||||
|
},
|
||||||
},
|
},
|
||||||
"400": map[string]interface{}{"description": "参数错误"},
|
"400": map[string]interface{}{"description": "参数错误"},
|
||||||
"401": map[string]interface{}{"description": "未授权"},
|
"401": map[string]interface{}{"description": "未授权"},
|
||||||
@@ -1689,7 +1779,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"post": map[string]interface{}{
|
"post": map[string]interface{}{
|
||||||
"tags": []string{"对话交互"},
|
"tags": []string{"对话交互"},
|
||||||
"summary": "发送消息并获取 AI 回复(Eino 多代理,SSE)",
|
"summary": "发送消息并获取 AI 回复(Eino 多代理,SSE)",
|
||||||
"description": "与 `POST /api/eino-agent/stream` 类似;由 Eino 多代理执行。`orchestration` 指定 deep / plan_execute / supervisor,缺省 deep。**前提**:`multi_agent.enabled: true`;未启用时 SSE 内首条为 `type: error` 后接 `done`。支持 `webshellConnectionId`。",
|
"description": "与 `POST /api/eino-agent/stream` 类似;由 Eino 多代理执行。`orchestration` 指定 deep / plan_execute / supervisor,缺省 deep。`response_start` / `response_delta` 仅为候选/过程输出;只有 `type: response` 且 `data.finalized=true` 才表示成功最终回复。缺 completed 执行证据时可能先发送 `finalization_auto_continue`,表示服务端基于已有 trace 无注入续跑。**前提**:`multi_agent.enabled: true`;未启用时 SSE 内首条为 `type: error` 后接 `done`。支持 `webshellConnectionId`。",
|
||||||
"operationId": "sendMessageMultiAgentStream",
|
"operationId": "sendMessageMultiAgentStream",
|
||||||
"requestBody": map[string]interface{}{
|
"requestBody": map[string]interface{}{
|
||||||
"required": true,
|
"required": true,
|
||||||
@@ -1702,6 +1792,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"conversationId": map[string]interface{}{"type": "string"},
|
"conversationId": map[string]interface{}{"type": "string"},
|
||||||
"role": map[string]interface{}{"type": "string"},
|
"role": map[string]interface{}{"type": "string"},
|
||||||
"webshellConnectionId": map[string]interface{}{"type": "string"},
|
"webshellConnectionId": map[string]interface{}{"type": "string"},
|
||||||
|
"finalization": finalizationRequestSchema,
|
||||||
"orchestration": map[string]interface{}{
|
"orchestration": map[string]interface{}{
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "deep | plan_execute | supervisor;缺省 deep",
|
"description": "deep | plan_execute | supervisor;缺省 deep",
|
||||||
@@ -1720,7 +1811,7 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"text/event-stream": map[string]interface{}{
|
"text/event-stream": map[string]interface{}{
|
||||||
"schema": map[string]interface{}{
|
"schema": map[string]interface{}{
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "SSE 流",
|
"description": "SSE 流。终态 response 事件 data 包含 finalized、finalizable、status、completionReason、evidenceVerified、evidenceRefs、pendingExecutionIds、missingChecks;过程事件可能包含 finalization_auto_continue。",
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -5798,10 +5889,15 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"get": map[string]interface{}{
|
"get": map[string]interface{}{
|
||||||
"tags": []string{"对话附件"},
|
"tags": []string{"对话附件"},
|
||||||
"summary": "列出附件",
|
"summary": "列出附件",
|
||||||
"description": "获取对话附件文件列表,可按对话ID过滤。",
|
"description": "获取对话文件列表,包含手动上传附件、工具输出和会话产物,可按会话、项目、来源、文件名搜索和分页过滤。",
|
||||||
"operationId": "listChatUploads",
|
"operationId": "listChatUploads",
|
||||||
"parameters": []map[string]interface{}{
|
"parameters": []map[string]interface{}{
|
||||||
{"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
{"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
||||||
|
{"name": "project", "in": "query", "required": false, "description": "按项目ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
||||||
|
{"name": "source", "in": "query", "required": false, "description": "按来源过滤:upload/reduction/workspace/conversation_artifact/all", "schema": map[string]interface{}{"type": "string", "enum": []string{"all", "upload", "reduction", "workspace", "conversation_artifact"}}},
|
||||||
|
{"name": "search", "in": "query", "required": false, "description": "按文件名或子路径搜索", "schema": map[string]interface{}{"type": "string"}},
|
||||||
|
{"name": "page", "in": "query", "required": false, "description": "页码,从1开始", "schema": map[string]interface{}{"type": "integer", "default": 1}},
|
||||||
|
{"name": "pageSize", "in": "query", "required": false, "description": "每页数量,传 all 返回全部", "schema": map[string]interface{}{"oneOf": []map[string]interface{}{{"type": "integer"}, {"type": "string", "enum": []string{"all"}}}}},
|
||||||
},
|
},
|
||||||
"responses": map[string]interface{}{
|
"responses": map[string]interface{}{
|
||||||
"200": map[string]interface{}{
|
"200": map[string]interface{}{
|
||||||
@@ -5816,18 +5912,27 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
"items": map[string]interface{}{
|
"items": map[string]interface{}{
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": map[string]interface{}{
|
"properties": map[string]interface{}{
|
||||||
"relativePath": map[string]interface{}{"type": "string"},
|
"relativePath": map[string]interface{}{"type": "string"},
|
||||||
"absolutePath": map[string]interface{}{"type": "string"},
|
"absolutePath": map[string]interface{}{"type": "string"},
|
||||||
"name": map[string]interface{}{"type": "string"},
|
"name": map[string]interface{}{"type": "string"},
|
||||||
"size": map[string]interface{}{"type": "integer"},
|
"size": map[string]interface{}{"type": "integer"},
|
||||||
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
"modifiedUnix": map[string]interface{}{"type": "integer"},
|
||||||
"date": map[string]interface{}{"type": "string"},
|
"date": map[string]interface{}{"type": "string"},
|
||||||
"conversationId": map[string]interface{}{"type": "string"},
|
"conversationId": map[string]interface{}{"type": "string"},
|
||||||
"subPath": map[string]interface{}{"type": "string"},
|
"conversationTitle": map[string]interface{}{"type": "string"},
|
||||||
|
"projectId": map[string]interface{}{"type": "string"},
|
||||||
|
"projectName": map[string]interface{}{"type": "string"},
|
||||||
|
"subPath": map[string]interface{}{"type": "string"},
|
||||||
|
"source": map[string]interface{}{"type": "string", "description": "upload/reduction/workspace/conversation_artifact"},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
"folders": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}},
|
"folders": map[string]interface{}{"type": "array", "items": map[string]interface{}{"type": "string"}},
|
||||||
|
"total": map[string]interface{}{"type": "integer"},
|
||||||
|
"page": map[string]interface{}{"type": "integer"},
|
||||||
|
"pageSize": map[string]interface{}{
|
||||||
|
"type": "integer",
|
||||||
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -5902,6 +6007,31 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
"/api/chat-uploads/export": map[string]interface{}{
|
||||||
|
"get": map[string]interface{}{
|
||||||
|
"tags": []string{"对话附件"},
|
||||||
|
"summary": "导出附件",
|
||||||
|
"description": "按当前过滤条件导出对话文件 ZIP,包含 manifest.json。",
|
||||||
|
"operationId": "exportChatUploads",
|
||||||
|
"parameters": []map[string]interface{}{
|
||||||
|
{"name": "conversation", "in": "query", "required": false, "description": "按对话ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
||||||
|
{"name": "project", "in": "query", "required": false, "description": "按项目ID过滤", "schema": map[string]interface{}{"type": "string"}},
|
||||||
|
{"name": "source", "in": "query", "required": false, "description": "按来源过滤:upload/reduction/workspace/conversation_artifact/all", "schema": map[string]interface{}{"type": "string", "enum": []string{"all", "upload", "reduction", "workspace", "conversation_artifact"}}},
|
||||||
|
{"name": "search", "in": "query", "required": false, "description": "按文件名或子路径搜索", "schema": map[string]interface{}{"type": "string"}},
|
||||||
|
},
|
||||||
|
"responses": map[string]interface{}{
|
||||||
|
"200": map[string]interface{}{
|
||||||
|
"description": "ZIP文件下载",
|
||||||
|
"content": map[string]interface{}{
|
||||||
|
"application/zip": map[string]interface{}{
|
||||||
|
"schema": map[string]interface{}{"type": "string", "format": "binary"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"401": map[string]interface{}{"description": "未授权"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
"/api/chat-uploads/download": map[string]interface{}{
|
"/api/chat-uploads/download": map[string]interface{}{
|
||||||
"get": map[string]interface{}{
|
"get": map[string]interface{}{
|
||||||
"tags": []string{"对话附件"},
|
"tags": []string{"对话附件"},
|
||||||
@@ -5925,6 +6055,37 @@ func (h *OpenAPIHandler) GetOpenAPISpec(c *gin.Context) {
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
"/api/chat-uploads/path": map[string]interface{}{
|
||||||
|
"get": map[string]interface{}{
|
||||||
|
"tags": []string{"对话附件"},
|
||||||
|
"summary": "解析附件路径",
|
||||||
|
"description": "将文件管理中的相对路径或内部虚拟路径解析为服务器绝对路径,用于复制文件/目录路径。",
|
||||||
|
"operationId": "resolveChatUploadPath",
|
||||||
|
"parameters": []map[string]interface{}{
|
||||||
|
{"name": "path", "in": "query", "required": true, "description": "相对路径或虚拟路径(如 __workspace__/projects/<id>/csv)", "schema": map[string]interface{}{"type": "string"}},
|
||||||
|
{"name": "kind", "in": "query", "required": false, "description": "路径类型:file/directory,默认 file", "schema": map[string]interface{}{"type": "string", "enum": []string{"file", "directory"}}},
|
||||||
|
},
|
||||||
|
"responses": map[string]interface{}{
|
||||||
|
"200": map[string]interface{}{
|
||||||
|
"description": "解析成功",
|
||||||
|
"content": map[string]interface{}{
|
||||||
|
"application/json": map[string]interface{}{
|
||||||
|
"schema": map[string]interface{}{
|
||||||
|
"type": "object",
|
||||||
|
"properties": map[string]interface{}{
|
||||||
|
"absolutePath": map[string]interface{}{"type": "string"},
|
||||||
|
"isDir": map[string]interface{}{"type": "boolean"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"401": map[string]interface{}{"description": "未授权"},
|
||||||
|
"403": map[string]interface{}{"description": "无权访问"},
|
||||||
|
"404": map[string]interface{}{"description": "路径不存在"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
"/api/chat-uploads/content": map[string]interface{}{
|
"/api/chat-uploads/content": map[string]interface{}{
|
||||||
"get": map[string]interface{}{
|
"get": map[string]interface{}{
|
||||||
"tags": []string{"对话附件"},
|
"tags": []string{"对话附件"},
|
||||||
|
|||||||
@@ -61,7 +61,7 @@ var apiDocI18nSummaryToKey = map[string]string{
|
|||||||
"获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState",
|
"获取连接状态": "getWebshellConnectionState", "保存连接状态": "saveWebshellConnectionState",
|
||||||
"获取AI对话历史": "getWebshellAIHistory", "列出AI对话": "listWebshellAIConversations",
|
"获取AI对话历史": "getWebshellAIHistory", "列出AI对话": "listWebshellAIConversations",
|
||||||
"执行WebShell命令": "webshellExec", "WebShell文件操作": "webshellFileOp",
|
"执行WebShell命令": "webshellExec", "WebShell文件操作": "webshellFileOp",
|
||||||
"列出附件": "listChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload",
|
"列出附件": "listChatUploads", "导出附件": "exportChatUploads", "上传附件": "uploadChatFile", "删除附件": "deleteChatUpload",
|
||||||
"下载附件": "downloadChatUpload", "获取附件文本内容": "getChatUploadContent",
|
"下载附件": "downloadChatUpload", "获取附件文本内容": "getChatUploadContent",
|
||||||
"写入附件文本内容": "putChatUploadContent", "创建附件目录": "mkdirChatUpload", "重命名附件": "renameChatUpload",
|
"写入附件文本内容": "putChatUploadContent", "创建附件目录": "mkdirChatUpload", "重命名附件": "renameChatUpload",
|
||||||
"企业微信回调验证": "wecomCallbackVerify", "企业微信消息回调": "wecomCallbackMessage",
|
"企业微信回调验证": "wecomCallbackVerify", "企业微信消息回调": "wecomCallbackMessage",
|
||||||
|
|||||||
@@ -6,6 +6,8 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
|
"net/url"
|
||||||
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
@@ -146,6 +148,131 @@ func TestChatUploadPathAuthorizationFollowsConversationAccess(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestChatUploadsListIncludesAuthorizedProjectWorkspaceFiles(t *testing.T) {
|
||||||
|
db, user := setupConversationRBACTest(t)
|
||||||
|
fsBase := t.TempDir()
|
||||||
|
workspaceBase := filepath.Join(fsBase, "workspace")
|
||||||
|
reductionBase := filepath.Join(fsBase, "reduction")
|
||||||
|
db.SetEinoConversationDirs("", "", reductionBase, workspaceBase)
|
||||||
|
allowedProject, _ := db.CreateProject(&database.Project{Name: "allowed"})
|
||||||
|
hiddenProject, _ := db.CreateProject(&database.Project{Name: "hidden"})
|
||||||
|
conversation, _ := db.CreateConversation("project conversation", database.ConversationCreateMeta{ProjectID: allowedProject.ID})
|
||||||
|
if err := db.AssignResourceToUser(user.ID, "project", allowedProject.ID); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
allowedFile := filepath.Join(workspaceBase, "projects", allowedProject.ID, "csv", "assets.csv")
|
||||||
|
hiddenFile := filepath.Join(workspaceBase, "projects", hiddenProject.ID, "csv", "secret.csv")
|
||||||
|
for _, path := range []string{allowedFile, hiddenFile} {
|
||||||
|
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(path, []byte("name\nexample\n"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
h := NewChatUploadsHandler(zap.NewNop(), db)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
c, _ := gin.CreateTestContext(w)
|
||||||
|
c.Request = httptest.NewRequest(http.MethodGet, "/api/chat-uploads?source=workspace&pageSize=all&conversation="+conversation.ID, nil)
|
||||||
|
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned})
|
||||||
|
h.List(c)
|
||||||
|
if w.Code != http.StatusOK {
|
||||||
|
t.Fatalf("status = %d, want 200: %s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
var response struct {
|
||||||
|
Files []ChatUploadFileItem `json:"files"`
|
||||||
|
Total int `json:"total"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if response.Total != 1 || len(response.Files) != 1 {
|
||||||
|
t.Fatalf("files = %#v, total = %d, want only authorized workspace file", response.Files, response.Total)
|
||||||
|
}
|
||||||
|
got := response.Files[0]
|
||||||
|
if got.Source != chatUploadSourceWorkspace || got.Name != "assets.csv" || got.ProjectID != allowedProject.ID {
|
||||||
|
t.Fatalf("workspace file = %#v", got)
|
||||||
|
}
|
||||||
|
if got.ProjectName != allowedProject.Name {
|
||||||
|
t.Fatalf("projectName = %q, want %q", got.ProjectName, allowedProject.Name)
|
||||||
|
}
|
||||||
|
if got.ConversationID != conversation.ID {
|
||||||
|
t.Fatalf("conversationId = %q, want %q", got.ConversationID, conversation.ID)
|
||||||
|
}
|
||||||
|
if got.ConversationTitle != conversation.Title {
|
||||||
|
t.Fatalf("conversationTitle = %q, want %q", got.ConversationTitle, conversation.Title)
|
||||||
|
}
|
||||||
|
if got.AbsolutePath != allowedFile {
|
||||||
|
t.Fatalf("absolutePath = %q, want %q", got.AbsolutePath, allowedFile)
|
||||||
|
}
|
||||||
|
|
||||||
|
w = httptest.NewRecorder()
|
||||||
|
c, _ = gin.CreateTestContext(w)
|
||||||
|
resolveURL := "/api/chat-uploads/path?kind=directory&path=__workspace__%2Fprojects%2F" + allowedProject.ID
|
||||||
|
c.Request = httptest.NewRequest(http.MethodGet, resolveURL, nil)
|
||||||
|
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned})
|
||||||
|
h.ResolvePath(c)
|
||||||
|
if w.Code != http.StatusOK {
|
||||||
|
t.Fatalf("resolve status = %d, want 200: %s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
var resolved struct {
|
||||||
|
AbsolutePath string `json:"absolutePath"`
|
||||||
|
IsDir bool `json:"isDir"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &resolved); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
wantDir := filepath.Join(workspaceBase, "projects", allowedProject.ID)
|
||||||
|
if !resolved.IsDir || resolved.AbsolutePath != wantDir {
|
||||||
|
t.Fatalf("resolved = %#v, want dir %q", resolved, wantDir)
|
||||||
|
}
|
||||||
|
|
||||||
|
w = httptest.NewRecorder()
|
||||||
|
c, _ = gin.CreateTestContext(w)
|
||||||
|
c.Request = httptest.NewRequest(http.MethodGet, "/api/chat-uploads/path?kind=directory&path=__workspace__%2Fprojects", nil)
|
||||||
|
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned})
|
||||||
|
h.ResolvePath(c)
|
||||||
|
if w.Code != http.StatusOK {
|
||||||
|
t.Fatalf("resolve projects container status = %d, want 200: %s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &resolved); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
wantContainer := filepath.Join(workspaceBase, "projects")
|
||||||
|
if !resolved.IsDir || resolved.AbsolutePath != wantContainer {
|
||||||
|
t.Fatalf("resolved container = %#v, want dir %q", resolved, wantContainer)
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range []struct {
|
||||||
|
path string
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{"__workspace__/", workspaceBase},
|
||||||
|
{"__reduction__/", reductionBase},
|
||||||
|
{"__conversation_artifact__/", db.ConversationArtifactsBaseDir()},
|
||||||
|
} {
|
||||||
|
if err := os.MkdirAll(tc.want, 0o755); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
w = httptest.NewRecorder()
|
||||||
|
c, _ = gin.CreateTestContext(w)
|
||||||
|
c.Request = httptest.NewRequest(http.MethodGet, "/api/chat-uploads/path?kind=directory&path="+url.QueryEscape(tc.path), nil)
|
||||||
|
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: database.RBACScopeAssigned})
|
||||||
|
h.ResolvePath(c)
|
||||||
|
if w.Code != http.StatusOK {
|
||||||
|
t.Fatalf("resolve root %q status = %d, want 200: %s", tc.path, w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &resolved); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
wantAbs, _ := filepath.Abs(tc.want)
|
||||||
|
if !resolved.IsDir || resolved.AbsolutePath != wantAbs {
|
||||||
|
t.Fatalf("resolved root %q = %#v, want dir %q", tc.path, resolved, wantAbs)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestPrepareMultiAgentSessionRejectsForeignConversation(t *testing.T) {
|
func TestPrepareMultiAgentSessionRejectsForeignConversation(t *testing.T) {
|
||||||
db, user := setupConversationRBACTest(t)
|
db, user := setupConversationRBACTest(t)
|
||||||
hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{})
|
hidden, _ := db.CreateConversation("hidden", database.ConversationCreateMeta{})
|
||||||
|
|||||||
@@ -352,26 +352,28 @@ func NewWebShellHandler(logger *zap.Logger, db *database.DB) *WebShellHandler {
|
|||||||
|
|
||||||
// CreateConnectionRequest 创建连接请求
|
// CreateConnectionRequest 创建连接请求
|
||||||
type CreateConnectionRequest struct {
|
type CreateConnectionRequest struct {
|
||||||
URL string `json:"url" binding:"required"`
|
ProjectID string `json:"project_id"`
|
||||||
Password string `json:"password"`
|
URL string `json:"url" binding:"required"`
|
||||||
Type string `json:"type"`
|
Password string `json:"password"`
|
||||||
Method string `json:"method"`
|
Type string `json:"type"`
|
||||||
CmdParam string `json:"cmd_param"`
|
Method string `json:"method"`
|
||||||
Remark string `json:"remark"`
|
CmdParam string `json:"cmd_param"`
|
||||||
Encoding string `json:"encoding"`
|
Remark string `json:"remark"`
|
||||||
OS string `json:"os"`
|
Encoding string `json:"encoding"`
|
||||||
|
OS string `json:"os"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// UpdateConnectionRequest 更新连接请求
|
// UpdateConnectionRequest 更新连接请求
|
||||||
type UpdateConnectionRequest struct {
|
type UpdateConnectionRequest struct {
|
||||||
URL string `json:"url" binding:"required"`
|
ProjectID string `json:"project_id"`
|
||||||
Password string `json:"password"`
|
URL string `json:"url" binding:"required"`
|
||||||
Type string `json:"type"`
|
Password string `json:"password"`
|
||||||
Method string `json:"method"`
|
Type string `json:"type"`
|
||||||
CmdParam string `json:"cmd_param"`
|
Method string `json:"method"`
|
||||||
Remark string `json:"remark"`
|
CmdParam string `json:"cmd_param"`
|
||||||
Encoding string `json:"encoding"`
|
Remark string `json:"remark"`
|
||||||
OS string `json:"os"`
|
Encoding string `json:"encoding"`
|
||||||
|
OS string `json:"os"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// ListConnections 列出所有 WebShell 连接(GET /api/webshell/connections)
|
// ListConnections 列出所有 WebShell 连接(GET /api/webshell/connections)
|
||||||
@@ -381,7 +383,7 @@ func (h *WebShellHandler) ListConnections(c *gin.Context) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
session, _ := security.CurrentSession(c)
|
session, _ := security.CurrentSession(c)
|
||||||
list, err := h.db.ListWebshellConnectionsForAccess(session.UserID, session.Scope)
|
list, err := h.db.ListWebshellConnectionsForAccess(session.UserID, session.Scope, c.Query("project_id"))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||||
return
|
return
|
||||||
@@ -412,6 +414,11 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
|||||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
projectID := strings.TrimSpace(req.ProjectID)
|
||||||
|
if !h.canAccessProject(c, projectID) {
|
||||||
|
c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"})
|
||||||
|
return
|
||||||
|
}
|
||||||
method := strings.ToLower(strings.TrimSpace(req.Method))
|
method := strings.ToLower(strings.TrimSpace(req.Method))
|
||||||
if method != "get" && method != "post" {
|
if method != "get" && method != "post" {
|
||||||
method = "post"
|
method = "post"
|
||||||
@@ -422,6 +429,7 @@ func (h *WebShellHandler) CreateConnection(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
conn := &database.WebShellConnection{
|
conn := &database.WebShellConnection{
|
||||||
ID: "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12],
|
ID: "ws_" + strings.ReplaceAll(uuid.New().String(), "-", "")[:12],
|
||||||
|
ProjectID: projectID,
|
||||||
URL: req.URL,
|
URL: req.URL,
|
||||||
Password: strings.TrimSpace(req.Password),
|
Password: strings.TrimSpace(req.Password),
|
||||||
Type: shellType,
|
Type: shellType,
|
||||||
@@ -477,6 +485,11 @@ func (h *WebShellHandler) UpdateConnection(c *gin.Context) {
|
|||||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid url"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
projectID := strings.TrimSpace(req.ProjectID)
|
||||||
|
if !h.canAccessProject(c, projectID) {
|
||||||
|
c.JSON(http.StatusForbidden, gin.H{"error": "project access denied"})
|
||||||
|
return
|
||||||
|
}
|
||||||
method := strings.ToLower(strings.TrimSpace(req.Method))
|
method := strings.ToLower(strings.TrimSpace(req.Method))
|
||||||
if method != "get" && method != "post" {
|
if method != "get" && method != "post" {
|
||||||
method = "post"
|
method = "post"
|
||||||
@@ -486,15 +499,16 @@ func (h *WebShellHandler) UpdateConnection(c *gin.Context) {
|
|||||||
shellType = "php"
|
shellType = "php"
|
||||||
}
|
}
|
||||||
conn := &database.WebShellConnection{
|
conn := &database.WebShellConnection{
|
||||||
ID: id,
|
ID: id,
|
||||||
URL: req.URL,
|
ProjectID: projectID,
|
||||||
Password: strings.TrimSpace(req.Password),
|
URL: req.URL,
|
||||||
Type: shellType,
|
Password: strings.TrimSpace(req.Password),
|
||||||
Method: method,
|
Type: shellType,
|
||||||
CmdParam: strings.TrimSpace(req.CmdParam),
|
Method: method,
|
||||||
Remark: strings.TrimSpace(req.Remark),
|
CmdParam: strings.TrimSpace(req.CmdParam),
|
||||||
Encoding: normalizeWebshellEncoding(req.Encoding),
|
Remark: strings.TrimSpace(req.Remark),
|
||||||
OS: normalizeWebshellOS(req.OS),
|
Encoding: normalizeWebshellEncoding(req.Encoding),
|
||||||
|
OS: normalizeWebshellOS(req.OS),
|
||||||
}
|
}
|
||||||
if err := h.db.UpdateWebshellConnection(conn); err != nil {
|
if err := h.db.UpdateWebshellConnection(conn); err != nil {
|
||||||
if err == sql.ErrNoRows {
|
if err == sql.ErrNoRows {
|
||||||
@@ -721,15 +735,22 @@ func (h *WebShellHandler) Exec(c *gin.Context) {
|
|||||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"})
|
c.JSON(http.StatusBadRequest, gin.H{"error": "url and command are required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
|
// Pre-save connectivity tests send form credentials without connection_id.
|
||||||
if !allowed {
|
// Saved connections must go through resource ACL; DB credentials are authoritative.
|
||||||
|
if cid := strings.TrimSpace(req.ConnectionID); cid != "" {
|
||||||
|
conn, allowed := h.authorizedWebshellConnection(c, cid, req.URL)
|
||||||
|
if !allowed {
|
||||||
|
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// Never let a caller pair an authorized ID with attacker-controlled
|
||||||
|
// transport credentials or a URL.
|
||||||
|
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
||||||
|
req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding
|
||||||
|
} else if !security.SessionHasPermission(c, "webshell:write") {
|
||||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
// The database record is authoritative. Never let a caller pair an
|
|
||||||
// authorized ID with attacker-controlled transport credentials or a URL.
|
|
||||||
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
|
||||||
req.Method, req.CmdParam, req.Encoding = conn.Method, conn.CmdParam, conn.Encoding
|
|
||||||
|
|
||||||
parsed, err := url.Parse(req.URL)
|
parsed, err := url.Parse(req.URL)
|
||||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||||
@@ -823,13 +844,18 @@ func (h *WebShellHandler) FileOp(c *gin.Context) {
|
|||||||
c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"})
|
c.JSON(http.StatusBadRequest, gin.H{"error": "url and action are required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
conn, allowed := h.authorizedWebshellConnection(c, req.ConnectionID, req.URL)
|
if cid := strings.TrimSpace(req.ConnectionID); cid != "" {
|
||||||
if !allowed {
|
conn, allowed := h.authorizedWebshellConnection(c, cid, req.URL)
|
||||||
|
if !allowed {
|
||||||
|
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
||||||
|
req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS
|
||||||
|
} else if !security.SessionHasPermission(c, "webshell:write") {
|
||||||
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
c.JSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
req.URL, req.Password, req.Type = conn.URL, conn.Password, conn.Type
|
|
||||||
req.Method, req.CmdParam, req.Encoding, req.OS = conn.Method, conn.CmdParam, conn.Encoding, conn.OS
|
|
||||||
|
|
||||||
parsed, err := url.Parse(req.URL)
|
parsed, err := url.Parse(req.URL)
|
||||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||||
@@ -928,6 +954,21 @@ func (h *WebShellHandler) authorizedWebshellConnection(c *gin.Context, connectio
|
|||||||
return conn, true
|
return conn, true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (h *WebShellHandler) canAccessProject(c *gin.Context, projectID string) bool {
|
||||||
|
projectID = strings.TrimSpace(projectID)
|
||||||
|
if projectID == "" || h.db == nil {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
session, ok := security.CurrentSession(c)
|
||||||
|
if !ok {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if session.Scope == database.RBACScopeAll {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return h.db.UserCanAccessResource(session.UserID, session.Scope, "project", projectID)
|
||||||
|
}
|
||||||
|
|
||||||
// ExecWithConnection 在指定 WebShell 连接上执行命令(供 MCP/Agent 等非 HTTP 调用)
|
// ExecWithConnection 在指定 WebShell 连接上执行命令(供 MCP/Agent 等非 HTTP 调用)
|
||||||
func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) {
|
func (h *WebShellHandler) ExecWithConnection(conn *database.WebShellConnection, command string) (output string, ok bool, errMsg string) {
|
||||||
if conn == nil {
|
if conn == nil {
|
||||||
|
|||||||
@@ -40,15 +40,47 @@ func TestWebshellExecRequiresConnectionAccessWhenConnectionIDProvided(t *testing
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestWebshellExecRejectsAdHocURLWithoutConnectionID(t *testing.T) {
|
func TestWebshellExecAllowsAdHocURLWithoutConnectionID(t *testing.T) {
|
||||||
gin.SetMode(gin.TestMode)
|
gin.SetMode(gin.TestMode)
|
||||||
_, user, _, _ := setupWebshellRBACTest(t)
|
_, user, _, _ := setupWebshellRBACTest(t)
|
||||||
handler := NewWebShellHandler(zap.NewNop(), nil)
|
handler := NewWebShellHandler(zap.NewNop(), nil)
|
||||||
|
// Ad-hoc probe (connectivity test before save) must not be rejected as "无权访问".
|
||||||
|
// The target URL will fail to connect; we only assert auth allows the request through.
|
||||||
w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{
|
w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{
|
||||||
"url": "http://127.0.0.1/admin", "command": "id",
|
"url": "http://127.0.0.1:1/admin", "command": "id",
|
||||||
}, handler.Exec)
|
}, handler.Exec)
|
||||||
|
if w.Code == http.StatusForbidden {
|
||||||
|
t.Fatalf("ad-hoc URL status = %d, want non-403: %s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
var resp ExecResponse
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||||
|
t.Fatalf("decode response: %v body=%s", err, w.Body.String())
|
||||||
|
}
|
||||||
|
if resp.OK {
|
||||||
|
t.Fatalf("expected connection failure for closed port, got ok=true")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestWebshellExecRejectsAdHocWithoutWritePermission(t *testing.T) {
|
||||||
|
gin.SetMode(gin.TestMode)
|
||||||
|
user := &database.RBACUser{ID: "u_ro", Username: "readonly"}
|
||||||
|
handler := NewWebShellHandler(zap.NewNop(), nil)
|
||||||
|
payload, _ := json.Marshal(map[string]interface{}{
|
||||||
|
"url": "http://127.0.0.1/admin", "command": "id",
|
||||||
|
})
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
c, _ := gin.CreateTestContext(w)
|
||||||
|
c.Request = httptest.NewRequest(http.MethodPost, "/api/webshell/exec", bytes.NewReader(payload))
|
||||||
|
c.Request.Header.Set("Content-Type", "application/json")
|
||||||
|
c.Set(security.ContextSessionKey, security.Session{
|
||||||
|
UserID: user.ID,
|
||||||
|
Username: user.Username,
|
||||||
|
Permissions: map[string]bool{"webshell:read": true},
|
||||||
|
Scope: database.RBACScopeAssigned,
|
||||||
|
})
|
||||||
|
handler.Exec(c)
|
||||||
if w.Code != http.StatusForbidden {
|
if w.Code != http.StatusForbidden {
|
||||||
t.Fatalf("ad-hoc URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ package handler
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
@@ -47,6 +48,12 @@ type workflowDryRunRequest struct {
|
|||||||
Inputs map[string]interface{} `json:"inputs,omitempty"`
|
Inputs map[string]interface{} `json:"inputs,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type workflowGenerateDraftRequest struct {
|
||||||
|
Prompt string `json:"prompt"`
|
||||||
|
Options workflowrunner.DraftOptions `json:"options"`
|
||||||
|
AvailableTools []workflowrunner.DraftTool `json:"available_tools,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
func (h *WorkflowHandler) List(c *gin.Context) {
|
func (h *WorkflowHandler) List(c *gin.Context) {
|
||||||
includeDisabled := strings.EqualFold(c.Query("includeDisabled"), "true") || c.Query("include_disabled") == "1"
|
includeDisabled := strings.EqualFold(c.Query("includeDisabled"), "true") || c.Query("include_disabled") == "1"
|
||||||
items, err := h.db.ListWorkflowDefinitions(includeDisabled)
|
items, err := h.db.ListWorkflowDefinitions(includeDisabled)
|
||||||
@@ -126,6 +133,44 @@ func (h *WorkflowHandler) DryRun(c *gin.Context) {
|
|||||||
c.JSON(http.StatusOK, gin.H{"result": result})
|
c.JSON(http.StatusOK, gin.H{"result": result})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (h *WorkflowHandler) GenerateDraft(c *gin.Context) {
|
||||||
|
var req workflowGenerateDraftRequest
|
||||||
|
if err := c.ShouldBindJSON(&req); err != nil {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": "无效的请求参数: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
draftReq := workflowrunner.DraftRequest{
|
||||||
|
Prompt: req.Prompt,
|
||||||
|
Options: req.Options,
|
||||||
|
AvailableTools: req.AvailableTools,
|
||||||
|
}
|
||||||
|
var result *workflowrunner.DraftResult
|
||||||
|
var llmErr error
|
||||||
|
if h.cfg != nil {
|
||||||
|
if llmCfg, _, ok := h.cfg.ResolveAIChannel(""); ok && strings.TrimSpace(llmCfg.APIKey) != "" && strings.TrimSpace(llmCfg.Model) != "" {
|
||||||
|
result, llmErr = workflowrunner.GenerateDraftFromLLM(c.Request.Context(), draftReq, llmCfg, h.logger)
|
||||||
|
} else {
|
||||||
|
llmErr = errors.New("AI 通道未配置 api_key 或 model")
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
llmErr = errors.New("工作流生成器未加载平台 AI 配置")
|
||||||
|
}
|
||||||
|
if llmErr != nil {
|
||||||
|
c.JSON(http.StatusBadGateway, gin.H{"error": "大模型生成失败: " + llmErr.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if h.audit != nil {
|
||||||
|
h.audit.RecordOK(c, "workflow", "generate_draft", "自然语言生成工作流草稿", "", "", map[string]interface{}{
|
||||||
|
"generator": result.Generator,
|
||||||
|
"nodes": result.Stats["nodes"],
|
||||||
|
"edges": result.Stats["edges"],
|
||||||
|
"high_risk": result.Audit.HighRisk,
|
||||||
|
"savable": result.Audit.Savable,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
c.JSON(http.StatusOK, gin.H{"result": result})
|
||||||
|
}
|
||||||
|
|
||||||
func (h *WorkflowHandler) Update(c *gin.Context) {
|
func (h *WorkflowHandler) Update(c *gin.Context) {
|
||||||
h.save(c, c.Param("id"))
|
h.save(c, c.Param("id"))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -152,20 +152,37 @@ func (h *AgentHandler) runRoleWorkflowStreamIfBound(
|
|||||||
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
sendEvent("done", "", map[string]interface{}{"conversationId": conversationID})
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
if prep.AssistantMessageID != "" {
|
decision := h.finalizeCandidateForDeliveryWithPolicy(
|
||||||
_ = h.db.UpdateAssistantMessageFinalize(prep.AssistantMessageID, result.Response, nil, "")
|
prep.ConversationID,
|
||||||
|
prep.AssistantMessageID,
|
||||||
|
"workflow",
|
||||||
|
result.Response,
|
||||||
|
nil,
|
||||||
|
result.AwaitingHITL,
|
||||||
|
"",
|
||||||
|
true,
|
||||||
|
)
|
||||||
|
responseText := decision.FinalText
|
||||||
|
if !decision.Finalizable {
|
||||||
|
responseText = finalizationBlockedMessage(decision)
|
||||||
|
taskStatus = decision.Status
|
||||||
|
h.tasks.UpdateTaskStatus(conversationID, taskStatus)
|
||||||
|
sendEvent("finalization_check", responseText, decision)
|
||||||
}
|
}
|
||||||
payload := map[string]interface{}{
|
payload := finalizationResponsePayload(decision, map[string]interface{}{
|
||||||
"conversationId": prep.ConversationID,
|
"conversationId": prep.ConversationID,
|
||||||
"messageId": prep.AssistantMessageID,
|
"messageId": prep.AssistantMessageID,
|
||||||
"agentMode": "workflow",
|
"agentMode": "workflow",
|
||||||
"workflowRunId": result.RunID,
|
"workflowRunId": result.RunID,
|
||||||
}
|
})
|
||||||
if result.AwaitingHITL {
|
if result.AwaitingHITL {
|
||||||
payload["workflowStatus"] = "awaiting_hitl"
|
payload["workflowStatus"] = "awaiting_hitl"
|
||||||
payload["awaitingHitl"] = true
|
payload["awaitingHitl"] = true
|
||||||
|
} else {
|
||||||
|
payload["workflowStatus"] = result.Status
|
||||||
|
payload["awaitingHitl"] = false
|
||||||
}
|
}
|
||||||
sendEvent("response", result.Response, payload)
|
sendEvent("response", responseText, payload)
|
||||||
sendEvent("done", "", map[string]interface{}{"conversationId": prep.ConversationID})
|
sendEvent("done", "", map[string]interface{}{"conversationId": prep.ConversationID})
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
@@ -251,17 +268,37 @@ func (h *AgentHandler) runRoleWorkflowJSONIfBound(c *gin.Context, req *ChatReque
|
|||||||
c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg, "conversationId": conversationID})
|
c.JSON(http.StatusInternalServerError, gin.H{"error": errMsg, "conversationId": conversationID})
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
if prep.AssistantMessageID != "" {
|
decision := h.finalizeCandidateForDeliveryWithPolicy(
|
||||||
_ = h.db.UpdateAssistantMessageFinalize(prep.AssistantMessageID, result.Response, nil, "")
|
prep.ConversationID,
|
||||||
|
prep.AssistantMessageID,
|
||||||
|
"workflow",
|
||||||
|
result.Response,
|
||||||
|
nil,
|
||||||
|
result.AwaitingHITL,
|
||||||
|
"",
|
||||||
|
true,
|
||||||
|
)
|
||||||
|
responseText := decision.FinalText
|
||||||
|
if !decision.Finalizable {
|
||||||
|
responseText = finalizationBlockedMessage(decision)
|
||||||
|
taskStatus = decision.Status
|
||||||
}
|
}
|
||||||
c.JSON(http.StatusOK, gin.H{
|
c.JSON(http.StatusOK, gin.H{
|
||||||
"response": result.Response,
|
"response": responseText,
|
||||||
"conversationId": prep.ConversationID,
|
"conversationId": prep.ConversationID,
|
||||||
"assistantMessageId": prep.AssistantMessageID,
|
"assistantMessageId": prep.AssistantMessageID,
|
||||||
"agentMode": "workflow",
|
"agentMode": "workflow",
|
||||||
"workflowRunId": result.RunID,
|
"workflowRunId": result.RunID,
|
||||||
"workflowStatus": result.Status,
|
"workflowStatus": result.Status,
|
||||||
"awaitingHitl": result.AwaitingHITL,
|
"awaitingHitl": result.AwaitingHITL,
|
||||||
|
"finalized": decision.Finalized,
|
||||||
|
"finalizable": decision.Finalizable,
|
||||||
|
"status": decision.Status,
|
||||||
|
"completionReason": decision.CompletionReason,
|
||||||
|
"evidenceVerified": decision.EvidenceVerified,
|
||||||
|
"evidenceRefs": decision.EvidenceRefs,
|
||||||
|
"pendingExecutionIds": decision.PendingExecutionIDs,
|
||||||
|
"missingChecks": decision.MissingChecks,
|
||||||
})
|
})
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -76,3 +77,26 @@ func TestWorkflowPackageHandlerInspectionAndCreateImport(t *testing.T) {
|
|||||||
t.Fatalf("saved=%#v", saved)
|
t.Fatalf("saved=%#v", saved)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestWorkflowHandlerGenerateDraft(t *testing.T) {
|
||||||
|
gin.SetMode(gin.TestMode)
|
||||||
|
h := NewWorkflowHandler(nil, zap.NewNop())
|
||||||
|
body := bytes.NewBufferString(`{"prompt":"对目标资产做端口扫描,如果发现高危端口就执行加固脚本,最后输出报告","options":{"include_objective":true},"available_tools":[{"key":"nmap","name":"nmap","enabled":true}]}`)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
c, _ := gin.CreateTestContext(w)
|
||||||
|
c.Request = httptest.NewRequest(http.MethodPost, "/api/workflows/generate-draft", body)
|
||||||
|
c.Request.Header.Set("Content-Type", "application/json")
|
||||||
|
h.GenerateDraft(c)
|
||||||
|
if w.Code != http.StatusBadGateway {
|
||||||
|
t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
var resp struct {
|
||||||
|
Error string `json:"error"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if !strings.Contains(resp.Error, "大模型生成失败") {
|
||||||
|
t.Fatalf("unexpected error: %#v", resp.Error)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ import (
|
|||||||
const (
|
const (
|
||||||
defaultExecutionWaitTimeout = 60 * time.Second
|
defaultExecutionWaitTimeout = 60 * time.Second
|
||||||
maxExecutionWaitTimeout = 10 * time.Minute
|
maxExecutionWaitTimeout = 10 * time.Minute
|
||||||
|
defaultPartialPreviewBytes = 4096
|
||||||
|
maxPartialPreviewBytes = 64 * 1024
|
||||||
)
|
)
|
||||||
|
|
||||||
// RegisterExecutionControlTools exposes execution handle operations to Eino as
|
// RegisterExecutionControlTools exposes execution handle operations to Eino as
|
||||||
@@ -32,7 +34,9 @@ func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager)
|
|||||||
InputSchema: map[string]interface{}{
|
InputSchema: map[string]interface{}{
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": map[string]interface{}{
|
"properties": map[string]interface{}{
|
||||||
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
||||||
|
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览,默认 true"},
|
||||||
|
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数,默认 4096,最大 65536"},
|
||||||
},
|
},
|
||||||
"required": []string{"execution_id"},
|
"required": []string{"execution_id"},
|
||||||
},
|
},
|
||||||
@@ -45,7 +49,7 @@ func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager)
|
|||||||
if exec == nil {
|
if exec == nil {
|
||||||
return textToolResult("未找到该 execution_id: "+id, true), nil
|
return textToolResult("未找到该 execution_id: "+id, true), nil
|
||||||
}
|
}
|
||||||
return textToolResult(formatExecutionForModel(exec), false), nil
|
return textToolResult(formatExecutionForModel(exec, executionFormatOptionsFromArgs(args)), false), nil
|
||||||
})
|
})
|
||||||
|
|
||||||
server.RegisterTool(Tool{
|
server.RegisterTool(Tool{
|
||||||
@@ -55,8 +59,10 @@ func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager)
|
|||||||
InputSchema: map[string]interface{}{
|
InputSchema: map[string]interface{}{
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": map[string]interface{}{
|
"properties": map[string]interface{}{
|
||||||
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
||||||
"timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数,默认 60,最大 600"},
|
"timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数,默认 60,最大 600"},
|
||||||
|
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览,默认 true"},
|
||||||
|
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数,默认 4096,最大 65536"},
|
||||||
},
|
},
|
||||||
"required": []string{"execution_id"},
|
"required": []string{"execution_id"},
|
||||||
},
|
},
|
||||||
@@ -73,7 +79,7 @@ func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager)
|
|||||||
if snap == nil || snap.Execution == nil {
|
if snap == nil || snap.Execution == nil {
|
||||||
return textToolResult("未找到该 execution_id: "+id, true), nil
|
return textToolResult("未找到该 execution_id: "+id, true), nil
|
||||||
}
|
}
|
||||||
body := formatExecutionForModel(snap.Execution)
|
body := formatExecutionForModel(snap.Execution, executionFormatOptionsFromArgs(args))
|
||||||
if errors.Is(err, ErrExecutionWaitTimeout) {
|
if errors.Is(err, ErrExecutionWaitTimeout) {
|
||||||
body += "\n\n本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。"
|
body += "\n\n本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。"
|
||||||
}
|
}
|
||||||
@@ -144,7 +150,25 @@ func lookupToolExecution(server *Server, external *ExternalMCPManager, id string
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func formatExecutionForModel(exec *ToolExecution) string {
|
type executionFormatOptions struct {
|
||||||
|
includePartialOutput bool
|
||||||
|
partialMaxBytes int
|
||||||
|
}
|
||||||
|
|
||||||
|
func executionFormatOptionsFromArgs(args map[string]interface{}) executionFormatOptions {
|
||||||
|
includePartial := true
|
||||||
|
if raw, ok := args["include_partial_output"]; ok {
|
||||||
|
if b, ok := raw.(bool); ok {
|
||||||
|
includePartial = b
|
||||||
|
} else if s := strings.TrimSpace(fmt.Sprint(raw)); s != "" {
|
||||||
|
includePartial = strings.EqualFold(s, "true") || s == "1" || strings.EqualFold(s, "yes")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
maxBytes := intArg(args, "partial_output_max_bytes", defaultPartialPreviewBytes, maxPartialPreviewBytes)
|
||||||
|
return executionFormatOptions{includePartialOutput: includePartial, partialMaxBytes: maxBytes}
|
||||||
|
}
|
||||||
|
|
||||||
|
func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) string {
|
||||||
if exec == nil {
|
if exec == nil {
|
||||||
return "execution: null"
|
return "execution: null"
|
||||||
}
|
}
|
||||||
@@ -167,6 +191,15 @@ func formatExecutionForModel(exec *ToolExecution) string {
|
|||||||
payload["result"] = ToolResultPlainText(exec.Result)
|
payload["result"] = ToolResultPlainText(exec.Result)
|
||||||
payload["is_error"] = exec.Result.IsError
|
payload["is_error"] = exec.Result.IsError
|
||||||
}
|
}
|
||||||
|
if opts.includePartialOutput && exec.PartialOutput != "" {
|
||||||
|
partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes)
|
||||||
|
payload["partial_output"] = partial
|
||||||
|
payload["partial_output_bytes"] = exec.PartialOutputBytes
|
||||||
|
payload["partial_output_truncated"] = exec.PartialOutputTruncated || len([]byte(partial)) < len([]byte(exec.PartialOutput))
|
||||||
|
if exec.PartialOutputUpdatedAt != nil {
|
||||||
|
payload["partial_output_updated_at"] = exec.PartialOutputUpdatedAt.Format(time.RFC3339)
|
||||||
|
}
|
||||||
|
}
|
||||||
b, err := json.MarshalIndent(payload, "", " ")
|
b, err := json.MarshalIndent(payload, "", " ")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error)
|
return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error)
|
||||||
@@ -174,6 +207,17 @@ func formatExecutionForModel(exec *ToolExecution) string {
|
|||||||
return string(b)
|
return string(b)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func tailStringBytes(s string, maxBytes int) string {
|
||||||
|
if maxBytes <= 0 {
|
||||||
|
maxBytes = defaultPartialPreviewBytes
|
||||||
|
}
|
||||||
|
b := []byte(s)
|
||||||
|
if len(b) <= maxBytes {
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
return string(b[len(b)-maxBytes:])
|
||||||
|
}
|
||||||
|
|
||||||
func textToolResult(text string, isErr bool) *ToolResult {
|
func textToolResult(text string, isErr bool) *ToolResult {
|
||||||
return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr}
|
return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr}
|
||||||
}
|
}
|
||||||
@@ -222,3 +266,31 @@ func durationSecondsArg(args map[string]interface{}, key string, def, max time.D
|
|||||||
}
|
}
|
||||||
return d
|
return d
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func intArg(args map[string]interface{}, key string, def, max int) int {
|
||||||
|
if args == nil {
|
||||||
|
return def
|
||||||
|
}
|
||||||
|
var n int
|
||||||
|
switch v := args[key].(type) {
|
||||||
|
case int:
|
||||||
|
n = v
|
||||||
|
case int64:
|
||||||
|
n = int(v)
|
||||||
|
case float64:
|
||||||
|
n = int(v)
|
||||||
|
case json.Number:
|
||||||
|
i, _ := v.Int64()
|
||||||
|
n = int(i)
|
||||||
|
case string:
|
||||||
|
i, _ := strconv.Atoi(strings.TrimSpace(v))
|
||||||
|
n = i
|
||||||
|
}
|
||||||
|
if n <= 0 {
|
||||||
|
return def
|
||||||
|
}
|
||||||
|
if max > 0 && n > max {
|
||||||
|
return max
|
||||||
|
}
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
|||||||
@@ -348,6 +348,22 @@ func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) {
|
|||||||
return s.getPersistedSnapshot(executionID)
|
return s.getPersistedSnapshot(executionID)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (s *ExecutionService) AppendPartialOutput(executionID, chunk string) bool {
|
||||||
|
id := strings.TrimSpace(executionID)
|
||||||
|
if s == nil || id == "" || chunk == "" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
entry := s.entries[id]
|
||||||
|
if entry == nil || entry.exec == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
appendPartialOutput(entry.exec, chunk, defaultPartialOutputMaxBytes, now)
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
func (s *ExecutionService) Cancel(executionID, note string) bool {
|
func (s *ExecutionService) Cancel(executionID, note string) bool {
|
||||||
id := strings.TrimSpace(executionID)
|
id := strings.TrimSpace(executionID)
|
||||||
if id == "" || s == nil {
|
if id == "" || s == nil {
|
||||||
@@ -582,5 +598,28 @@ func cloneToolExecution(in *ToolExecution) *ToolExecution {
|
|||||||
t := *in.EndTime
|
t := *in.EndTime
|
||||||
out.EndTime = &t
|
out.EndTime = &t
|
||||||
}
|
}
|
||||||
|
if in.PartialOutputUpdatedAt != nil {
|
||||||
|
t := *in.PartialOutputUpdatedAt
|
||||||
|
out.PartialOutputUpdatedAt = &t
|
||||||
|
}
|
||||||
return &out
|
return &out
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func appendPartialOutput(exec *ToolExecution, chunk string, maxBytes int, updatedAt time.Time) {
|
||||||
|
if exec == nil || chunk == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if maxBytes <= 0 {
|
||||||
|
maxBytes = defaultPartialOutputMaxBytes
|
||||||
|
}
|
||||||
|
exec.PartialOutputBytes += int64(len([]byte(chunk)))
|
||||||
|
combined := exec.PartialOutput + chunk
|
||||||
|
if len([]byte(combined)) > maxBytes {
|
||||||
|
b := []byte(combined)
|
||||||
|
combined = string(b[len(b)-maxBytes:])
|
||||||
|
exec.PartialOutputTruncated = true
|
||||||
|
}
|
||||||
|
exec.PartialOutput = combined
|
||||||
|
t := updatedAt
|
||||||
|
exec.PartialOutputUpdatedAt = &t
|
||||||
|
}
|
||||||
|
|||||||
@@ -674,48 +674,6 @@ func (m *ExternalMCPManager) updateToolCache(name string, tools []Tool) {
|
|||||||
|
|
||||||
// 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) {
|
||||||
_, authenticated := authctx.PrincipalFromContext(ctx)
|
|
||||||
m.mu.RLock()
|
|
||||||
authorizer := m.toolAuthorizer
|
|
||||||
m.mu.RUnlock()
|
|
||||||
if authorizer != nil {
|
|
||||||
if err := authorizer(ctx, toolName, args); err != nil {
|
|
||||||
return nil, "", fmt.Errorf("external tool authorization denied: %w", err)
|
|
||||||
}
|
|
||||||
} else if authenticated {
|
|
||||||
return nil, "", fmt.Errorf("external tool authorization policy is not configured")
|
|
||||||
}
|
|
||||||
// 解析工具名称:name::toolName
|
|
||||||
var mcpName, actualToolName string
|
|
||||||
if idx := findSubstring(toolName, "::"); idx > 0 {
|
|
||||||
mcpName = toolName[:idx]
|
|
||||||
actualToolName = toolName[idx+2:]
|
|
||||||
} else {
|
|
||||||
return nil, "", fmt.Errorf("无效的工具名称格式: %s", toolName)
|
|
||||||
}
|
|
||||||
|
|
||||||
client, exists := m.GetClient(mcpName)
|
|
||||||
if !exists {
|
|
||||||
return nil, "", fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
|
|
||||||
}
|
|
||||||
if err := m.checkExternalMCPCircuit(mcpName); err != nil {
|
|
||||||
return nil, "", err
|
|
||||||
}
|
|
||||||
|
|
||||||
// 检查连接状态,如果未连接或状态为error,不允许调用
|
|
||||||
if !client.IsConnected() {
|
|
||||||
status := client.GetStatus()
|
|
||||||
if status == "error" {
|
|
||||||
// 获取错误信息(如果有)
|
|
||||||
errorMsg := m.GetError(mcpName)
|
|
||||||
if errorMsg != "" {
|
|
||||||
return nil, "", fmt.Errorf("外部MCP连接失败: %s (错误: %s)", mcpName, errorMsg)
|
|
||||||
}
|
|
||||||
return nil, "", fmt.Errorf("外部MCP连接失败: %s", mcpName)
|
|
||||||
}
|
|
||||||
return nil, "", fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status)
|
|
||||||
}
|
|
||||||
|
|
||||||
if m.executionService == nil {
|
if m.executionService == nil {
|
||||||
m.executionService = NewExecutionService(m.storage, m.logger)
|
m.executionService = NewExecutionService(m.storage, m.logger)
|
||||||
m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes)
|
m.executionService.ConfigureToolResultMaxBytes(m.toolResultMaxBytes)
|
||||||
@@ -725,12 +683,57 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
|||||||
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
|
||||||
ownerUserID = principal.UserID
|
ownerUserID = principal.UserID
|
||||||
}
|
}
|
||||||
|
var mcpName, actualToolName string
|
||||||
|
var client ExternalMCPClient
|
||||||
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
|
handle, err := m.executionService.Submit(ctx, ExecutionRequest{
|
||||||
ToolName: toolName,
|
ToolName: toolName,
|
||||||
Arguments: args,
|
Arguments: args,
|
||||||
ConversationID: MCPConversationIDFromContext(ctx),
|
ConversationID: MCPConversationIDFromContext(ctx),
|
||||||
OwnerUserID: ownerUserID,
|
OwnerUserID: ownerUserID,
|
||||||
PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) {
|
PreRun: func(runCtx context.Context, exec *ToolExecution) (func(), error) {
|
||||||
|
_, authenticated := authctx.PrincipalFromContext(runCtx)
|
||||||
|
m.mu.RLock()
|
||||||
|
authorizer := m.toolAuthorizer
|
||||||
|
m.mu.RUnlock()
|
||||||
|
if authorizer != nil {
|
||||||
|
if err := authorizer(runCtx, toolName, args); err != nil {
|
||||||
|
return nil, fmt.Errorf("external tool authorization denied: %w", err)
|
||||||
|
}
|
||||||
|
} else if authenticated {
|
||||||
|
return nil, fmt.Errorf("external tool authorization policy is not configured")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 解析工具名称:name::toolName
|
||||||
|
if idx := findSubstring(toolName, "::"); idx > 0 {
|
||||||
|
mcpName = toolName[:idx]
|
||||||
|
actualToolName = toolName[idx+2:]
|
||||||
|
} else {
|
||||||
|
return nil, fmt.Errorf("无效的工具名称格式: %s", toolName)
|
||||||
|
}
|
||||||
|
|
||||||
|
var exists bool
|
||||||
|
client, exists = m.GetClient(mcpName)
|
||||||
|
if !exists {
|
||||||
|
return nil, fmt.Errorf("外部MCP客户端不存在: %s", mcpName)
|
||||||
|
}
|
||||||
|
if err := m.checkExternalMCPCircuit(mcpName); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查连接状态,如果未连接或状态为error,不允许调用
|
||||||
|
if !client.IsConnected() {
|
||||||
|
status := client.GetStatus()
|
||||||
|
if status == "error" {
|
||||||
|
// 获取错误信息(如果有)
|
||||||
|
errorMsg := m.GetError(mcpName)
|
||||||
|
if errorMsg != "" {
|
||||||
|
return nil, fmt.Errorf("外部MCP连接失败: %s (错误: %s)", mcpName, errorMsg)
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("外部MCP连接失败: %s", mcpName)
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("外部MCP客户端未连接: %s (状态: %s)", mcpName, status)
|
||||||
|
}
|
||||||
|
|
||||||
release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName)
|
release, acquireErr := m.acquireExternalMCPCallSlot(runCtx, mcpName)
|
||||||
if acquireErr != nil {
|
if acquireErr != nil {
|
||||||
return nil, acquireErr
|
return nil, acquireErr
|
||||||
@@ -745,8 +748,10 @@ func (m *ExternalMCPManager) CallTool(ctx context.Context, toolName string, args
|
|||||||
return result, callErr
|
return result, callErr
|
||||||
},
|
},
|
||||||
OnDone: func(exec *ToolExecution) {
|
OnDone: func(exec *ToolExecution) {
|
||||||
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted
|
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
|
||||||
m.recordExternalMCPResult(mcpName, failed)
|
if mcpName != "" {
|
||||||
|
m.recordExternalMCPResult(mcpName, failed)
|
||||||
|
}
|
||||||
m.updateStats(toolName, failed)
|
m.updateStats(toolName, failed)
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -20,10 +20,20 @@ func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) {
|
|||||||
return errors.New("denied by policy")
|
return errors.New("denied by policy")
|
||||||
})
|
})
|
||||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true}))
|
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true}))
|
||||||
_, _, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{})
|
_, executionID, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{})
|
||||||
if err == nil || !strings.Contains(err.Error(), "authorization denied") {
|
if err == nil || !strings.Contains(err.Error(), "authorization denied") {
|
||||||
t.Fatalf("external call bypassed authorizer: %v", err)
|
t.Fatalf("external call bypassed authorizer: %v", err)
|
||||||
}
|
}
|
||||||
|
if executionID == "" {
|
||||||
|
t.Fatal("denied external call should still return an execution id")
|
||||||
|
}
|
||||||
|
execution, ok := manager.GetExecution(executionID)
|
||||||
|
if !ok || execution == nil {
|
||||||
|
t.Fatalf("missing denied external execution %q", executionID)
|
||||||
|
}
|
||||||
|
if execution.Status != ToolExecutionStatusFailed || !strings.Contains(execution.Error, "denied by policy") {
|
||||||
|
t.Fatalf("denied external execution = %#v, want failed with policy error", execution)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) {
|
func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) {
|
||||||
|
|||||||
+57
-22
@@ -59,6 +59,8 @@ type Server struct {
|
|||||||
spillRootDir string
|
spillRootDir string
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const defaultPartialOutputMaxBytes = 64 * 1024
|
||||||
|
|
||||||
// SetToolAuthorizer installs the common policy decision point for every
|
// SetToolAuthorizer installs the common policy decision point for every
|
||||||
// user-attributed tool call, whether it originates from HTTP or an Agent.
|
// user-attributed tool call, whether it originates from HTTP or an Agent.
|
||||||
func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) {
|
func (s *Server) SetToolAuthorizer(authorizer func(context.Context, string, map[string]interface{}) error) {
|
||||||
@@ -602,13 +604,13 @@ func (s *Server) handleCallTool(requestCtx context.Context, msg *Message) *Messa
|
|||||||
st, msg := executionStatusAndMessage(err)
|
st, msg := executionStatusAndMessage(err)
|
||||||
execution.Status = st
|
execution.Status = st
|
||||||
execution.Error = msg
|
execution.Error = msg
|
||||||
failed = true
|
failed = st != "cancelled"
|
||||||
} else if result != nil && result.IsError {
|
} else if result != nil && result.IsError {
|
||||||
if cancelledWithUserNote {
|
if cancelledWithUserNote {
|
||||||
execution.Status = "cancelled"
|
execution.Status = "cancelled"
|
||||||
execution.Error = ""
|
execution.Error = ""
|
||||||
execution.Result = result
|
execution.Result = result
|
||||||
failed = true
|
failed = false
|
||||||
} else {
|
} else {
|
||||||
execution.Status = "failed"
|
execution.Status = "failed"
|
||||||
if len(result.Content) > 0 {
|
if len(result.Content) > 0 {
|
||||||
@@ -893,25 +895,6 @@ func (s *Server) GetAllTools() []Tool {
|
|||||||
|
|
||||||
// CallTool 直接调用工具(用于内部调用)
|
// CallTool 直接调用工具(用于内部调用)
|
||||||
func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
|
func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]interface{}) (*ToolResult, string, error) {
|
||||||
_, authenticated := authctx.PrincipalFromContext(ctx)
|
|
||||||
s.mu.RLock()
|
|
||||||
authorizer := s.toolAuthorizer
|
|
||||||
s.mu.RUnlock()
|
|
||||||
if authorizer != nil {
|
|
||||||
if err := authorizer(ctx, toolName, args); err != nil {
|
|
||||||
return nil, "", fmt.Errorf("tool authorization denied: %w", err)
|
|
||||||
}
|
|
||||||
} else if authenticated {
|
|
||||||
return nil, "", errors.New("tool authorization policy is not configured")
|
|
||||||
}
|
|
||||||
s.mu.RLock()
|
|
||||||
handler, exists := s.tools[toolName]
|
|
||||||
s.mu.RUnlock()
|
|
||||||
|
|
||||||
if !exists {
|
|
||||||
return nil, "", fmt.Errorf("工具 %s 未找到", toolName)
|
|
||||||
}
|
|
||||||
|
|
||||||
if s.executionService == nil {
|
if s.executionService == nil {
|
||||||
s.executionService = NewExecutionService(s.storage, s.logger)
|
s.executionService = NewExecutionService(s.storage, s.logger)
|
||||||
s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes)
|
s.executionService.ConfigureToolResultMaxBytes(s.toolResultMaxBytes)
|
||||||
@@ -927,10 +910,25 @@ func (s *Server) CallTool(ctx context.Context, toolName string, args map[string]
|
|||||||
ConversationID: MCPConversationIDFromContext(ctx),
|
ConversationID: MCPConversationIDFromContext(ctx),
|
||||||
OwnerUserID: ownerUserID,
|
OwnerUserID: ownerUserID,
|
||||||
Run: func(runCtx context.Context) (*ToolResult, error) {
|
Run: func(runCtx context.Context) (*ToolResult, error) {
|
||||||
|
_, authenticated := authctx.PrincipalFromContext(runCtx)
|
||||||
|
s.mu.RLock()
|
||||||
|
authorizer := s.toolAuthorizer
|
||||||
|
handler, exists := s.tools[toolName]
|
||||||
|
s.mu.RUnlock()
|
||||||
|
if authorizer != nil {
|
||||||
|
if err := authorizer(runCtx, toolName, args); err != nil {
|
||||||
|
return nil, fmt.Errorf("tool authorization denied: %w", err)
|
||||||
|
}
|
||||||
|
} else if authenticated {
|
||||||
|
return nil, errors.New("tool authorization policy is not configured")
|
||||||
|
}
|
||||||
|
if !exists {
|
||||||
|
return nil, fmt.Errorf("工具 %s 未找到", toolName)
|
||||||
|
}
|
||||||
return handler(runCtx, args)
|
return handler(runCtx, args)
|
||||||
},
|
},
|
||||||
OnDone: func(exec *ToolExecution) {
|
OnDone: func(exec *ToolExecution) {
|
||||||
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted
|
failed := exec != nil && exec.Status != ToolExecutionStatusCompleted && exec.Status != ToolExecutionStatusCancelled
|
||||||
s.updateStats(toolName, failed)
|
s.updateStats(toolName, failed)
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
@@ -1123,6 +1121,25 @@ func (s *Server) FinishToolExecution(ctx context.Context, executionID, toolName
|
|||||||
return id
|
return id
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// AppendToolExecutionPartialOutput records a bounded tail preview for a running local execution.
|
||||||
|
// The final Result remains authoritative and is written only when the tool finishes.
|
||||||
|
func (s *Server) AppendToolExecutionPartialOutput(executionID, chunk string) {
|
||||||
|
if s == nil || strings.TrimSpace(executionID) == "" || chunk == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
id := strings.TrimSpace(executionID)
|
||||||
|
if s.executionService != nil && s.executionService.AppendPartialOutput(id, chunk) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
s.mu.Lock()
|
||||||
|
exec := s.executions[id]
|
||||||
|
if exec != nil {
|
||||||
|
appendPartialOutput(exec, chunk, defaultPartialOutputMaxBytes, now)
|
||||||
|
}
|
||||||
|
s.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
|
// RecordCompletedToolInvocation 将已在其它路径完成的工具调用写入监控存储(格式与 CallTool 结束后一致),
|
||||||
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
|
// 用于 Eino ADK filesystem execute 等未经过 CallTool 的场景;返回 executionId 供助手消息 mcpExecutionIds 关联。
|
||||||
func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
func (s *Server) RecordCompletedToolInvocation(ctx context.Context, toolName string, args map[string]interface{}, resultText string, invokeErr error) string {
|
||||||
@@ -1206,6 +1223,24 @@ func (s *Server) unregisterRunningCancel(id string) {
|
|||||||
s.runningCancelsMu.Unlock()
|
s.runningCancelsMu.Unlock()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// RegisterToolExecutionCancel lets non-ExecutionService tool paths, such as Eino
|
||||||
|
// filesystem execute, participate in cancel_tool_execution by execution_id.
|
||||||
|
func (s *Server) RegisterToolExecutionCancel(id string, cancel context.CancelFunc) {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if s == nil || id == "" || cancel == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.registerRunningCancel(id, cancel)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Server) UnregisterToolExecutionCancel(id string) {
|
||||||
|
id = strings.TrimSpace(id)
|
||||||
|
if s == nil || id == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.unregisterRunningCancel(id)
|
||||||
|
}
|
||||||
|
|
||||||
func (s *Server) readAbortUserNote(id string) string {
|
func (s *Server) readAbortUserNote(id string) string {
|
||||||
s.runningCancelsMu.Lock()
|
s.runningCancelsMu.Lock()
|
||||||
defer s.runningCancelsMu.Unlock()
|
defer s.runningCancelsMu.Unlock()
|
||||||
|
|||||||
@@ -23,9 +23,20 @@ func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
|
|||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
})
|
})
|
||||||
if _, _, err := server.CallTool(context.Background(), "echo", nil); err == nil {
|
_, deniedExecutionID, err := server.CallTool(context.Background(), "echo", nil)
|
||||||
|
if err == nil {
|
||||||
t.Fatal("tool call without principal was allowed")
|
t.Fatal("tool call without principal was allowed")
|
||||||
}
|
}
|
||||||
|
if deniedExecutionID == "" {
|
||||||
|
t.Fatal("denied tool call should still return an execution id")
|
||||||
|
}
|
||||||
|
deniedExecution, ok := server.GetExecution(deniedExecutionID)
|
||||||
|
if !ok || deniedExecution == nil {
|
||||||
|
t.Fatalf("missing denied execution %q", deniedExecutionID)
|
||||||
|
}
|
||||||
|
if deniedExecution.Status != ToolExecutionStatusFailed || !strings.Contains(deniedExecution.Error, "principal required") {
|
||||||
|
t.Fatalf("denied execution = %#v, want failed with authorization error", deniedExecution)
|
||||||
|
}
|
||||||
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true}))
|
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true}))
|
||||||
_, executionID, err := server.CallTool(ctx, "echo", nil)
|
_, executionID, err := server.CallTool(ctx, "echo", nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -180,3 +191,41 @@ func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) {
|
|||||||
}
|
}
|
||||||
close(release)
|
close(release)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestGetToolExecutionIncludesBoundedPartialOutput(t *testing.T) {
|
||||||
|
server := NewServer(zap.NewNop())
|
||||||
|
RegisterExecutionControlTools(server, nil)
|
||||||
|
|
||||||
|
executionID := server.BeginToolExecution(context.Background(), "execute", map[string]interface{}{"command": "demo"})
|
||||||
|
if executionID == "" {
|
||||||
|
t.Fatal("missing execution id")
|
||||||
|
}
|
||||||
|
server.AppendToolExecutionPartialOutput(executionID, "first\n")
|
||||||
|
server.AppendToolExecutionPartialOutput(executionID, strings.Repeat("x", 32))
|
||||||
|
|
||||||
|
result, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
|
||||||
|
"execution_id": executionID,
|
||||||
|
"partial_output_max_bytes": 8,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get_tool_execution: %v", err)
|
||||||
|
}
|
||||||
|
body := ToolResultPlainText(result)
|
||||||
|
if !strings.Contains(body, `"partial_output": "xxxxxxxx"`) {
|
||||||
|
t.Fatalf("missing bounded partial output: %s", body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(body, `"partial_output_bytes": 38`) {
|
||||||
|
t.Fatalf("missing partial byte count: %s", body)
|
||||||
|
}
|
||||||
|
|
||||||
|
result, _, err = server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
|
||||||
|
"execution_id": executionID,
|
||||||
|
"include_partial_output": false,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get_tool_execution without partial: %v", err)
|
||||||
|
}
|
||||||
|
if body := ToolResultPlainText(result); strings.Contains(body, "partial_output") {
|
||||||
|
t.Fatalf("partial output should be omitted: %s", body)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -199,6 +199,12 @@ type ToolExecution struct {
|
|||||||
StartTime time.Time `json:"startTime"`
|
StartTime time.Time `json:"startTime"`
|
||||||
EndTime *time.Time `json:"endTime,omitempty"`
|
EndTime *time.Time `json:"endTime,omitempty"`
|
||||||
Duration time.Duration `json:"duration,omitempty"`
|
Duration time.Duration `json:"duration,omitempty"`
|
||||||
|
// PartialOutput is a bounded tail preview of output produced by a running tool.
|
||||||
|
// It is intentionally separate from Result, which remains the final canonical tool result.
|
||||||
|
PartialOutput string `json:"partialOutput,omitempty"`
|
||||||
|
PartialOutputBytes int64 `json:"partialOutputBytes,omitempty"`
|
||||||
|
PartialOutputTruncated bool `json:"partialOutputTruncated,omitempty"`
|
||||||
|
PartialOutputUpdatedAt *time.Time `json:"partialOutputUpdatedAt,omitempty"`
|
||||||
// ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
|
// ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
|
||||||
ConversationID string `json:"conversationId,omitempty"`
|
ConversationID string `json:"conversationId,omitempty"`
|
||||||
OwnerUserID string `json:"-"`
|
OwnerUserID string `json:"-"`
|
||||||
|
|||||||
@@ -12,6 +12,9 @@ import (
|
|||||||
// 与 CallTool 路径一致,使监控页能展示「执行中」状态。
|
// 与 CallTool 路径一致,使监控页能展示「执行中」状态。
|
||||||
func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
|
func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, recorder einomcp.ExecutionRecorder) (
|
||||||
begin func(toolCallID, command string) string,
|
begin func(toolCallID, command string) string,
|
||||||
|
appendPartial func(executionID, toolCallID, chunk string),
|
||||||
|
registerCancel func(executionID string, cancel context.CancelFunc),
|
||||||
|
unregisterCancel func(executionID string),
|
||||||
finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
finish func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||||
) {
|
) {
|
||||||
begin = func(toolCallID, command string) string {
|
begin = func(toolCallID, command string) string {
|
||||||
@@ -25,6 +28,24 @@ func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, record
|
|||||||
}
|
}
|
||||||
return id
|
return id
|
||||||
}
|
}
|
||||||
|
appendPartial = func(executionID, toolCallID, chunk string) {
|
||||||
|
if ag == nil || executionID == "" || chunk == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ag.AppendLocalToolExecutionPartialOutput(executionID, chunk)
|
||||||
|
}
|
||||||
|
registerCancel = func(executionID string, cancel context.CancelFunc) {
|
||||||
|
if ag == nil || executionID == "" || cancel == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ag.RegisterLocalToolExecutionCancel(executionID, cancel)
|
||||||
|
}
|
||||||
|
unregisterCancel = func(executionID string) {
|
||||||
|
if ag == nil || executionID == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ag.UnregisterLocalToolExecutionCancel(executionID)
|
||||||
|
}
|
||||||
finish = func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) {
|
finish = func(executionID, toolCallID, command, stdout string, success bool, invokeErr error) {
|
||||||
if ag == nil {
|
if ag == nil {
|
||||||
return
|
return
|
||||||
@@ -43,5 +64,5 @@ func newEinoExecuteMonitorCallbacks(ctx context.Context, ag *agent.Agent, record
|
|||||||
recorder(id, toolCallID)
|
recorder(id, toolCallID)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return begin, finish
|
return begin, appendPartial, registerCancel, unregisterCancel, finish
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -63,11 +63,16 @@ type einoStreamingShellWrap struct {
|
|||||||
outputChunk func(toolName, toolCallID, chunk string)
|
outputChunk func(toolName, toolCallID, chunk string)
|
||||||
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
|
// toolTimeoutMinutes 与 agent.tool_timeout_minutes 对齐;>0 时对单次 execute 套用 context 超时(与 MCP 工具经 executeToolViaMCP 行为一致)。0 表示仅依赖上层 ctx(如整任务 10h 上限)。
|
||||||
toolTimeoutMinutes int
|
toolTimeoutMinutes int
|
||||||
|
// toolWaitTimeoutSeconds 与 agent.tool_wait_timeout_seconds 对齐;>0 时本轮等待到期后返回 execution_id,shell 继续后台运行。
|
||||||
|
toolWaitTimeoutSeconds int
|
||||||
// shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭。
|
// shellNoOutputTimeoutSec:无任何输出时的空闲秒数;0=关闭。
|
||||||
shellNoOutputTimeoutSec int
|
shellNoOutputTimeoutSec int
|
||||||
// beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。
|
// beginMonitor 在 execute 开始时写入 running 状态;finishMonitor 在流结束后更新为 completed/failed。
|
||||||
beginMonitor func(toolCallID, command string) string
|
beginMonitor func(toolCallID, command string) string
|
||||||
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
|
appendPartialMonitor func(executionID, toolCallID, chunk string)
|
||||||
|
registerCancelMonitor func(executionID string, cancel context.CancelFunc)
|
||||||
|
unregisterCancelMonitor func(executionID string)
|
||||||
|
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
|
||||||
@@ -104,6 +109,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
if w.toolTimeoutMinutes > 0 {
|
if w.toolTimeoutMinutes > 0 {
|
||||||
execCtx, timeoutCancel = context.WithTimeout(execCtx, time.Duration(w.toolTimeoutMinutes)*time.Minute)
|
execCtx, timeoutCancel = context.WithTimeout(execCtx, time.Duration(w.toolTimeoutMinutes)*time.Minute)
|
||||||
}
|
}
|
||||||
|
if monitorExecID != "" && w.registerCancelMonitor != nil {
|
||||||
|
w.registerCancelMonitor(monitorExecID, execCancel)
|
||||||
|
}
|
||||||
if execReg != nil && convID != "" {
|
if execReg != nil && convID != "" {
|
||||||
execReg.RegisterActiveEinoExecute(convID, execCancel)
|
execReg.RegisterActiveEinoExecute(convID, execCancel)
|
||||||
}
|
}
|
||||||
@@ -116,6 +124,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
if execCancel != nil {
|
if execCancel != nil {
|
||||||
execCancel()
|
execCancel()
|
||||||
}
|
}
|
||||||
|
if monitorExecID != "" && w.unregisterCancelMonitor != nil {
|
||||||
|
w.unregisterCancelMonitor(monitorExecID)
|
||||||
|
}
|
||||||
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
|
if einoExecuteRecvErrIsToolTimeout(err, execCtx) {
|
||||||
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
||||||
if w.finishMonitor != nil {
|
if w.finishMonitor != nil {
|
||||||
@@ -146,7 +157,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
|
|
||||||
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32)
|
outR, outW := schema.Pipe[*filesystem.ExecuteResponse](32)
|
||||||
|
|
||||||
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int) {
|
go func(inner *schema.StreamReader[*filesystem.ExecuteResponse], command string, cancel context.CancelFunc, timeoutCleanup context.CancelFunc, tctx context.Context, conversationID string, reg mcp.EinoExecuteRunRegistry, toolReg mcp.ToolRunRegistry, execID string, toolCallID string, noOutputSec int, waitTimeoutSec int) {
|
||||||
var innerCloseOnce sync.Once
|
var innerCloseOnce sync.Once
|
||||||
closeInner := func() {
|
closeInner := func() {
|
||||||
innerCloseOnce.Do(func() { inner.Close() })
|
innerCloseOnce.Do(func() { inner.Close() })
|
||||||
@@ -164,6 +175,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
if toolReg != nil && conversationID != "" && execID != "" {
|
if toolReg != nil && conversationID != "" && execID != "" {
|
||||||
defer toolReg.UnregisterRunningTool(conversationID, execID)
|
defer toolReg.UnregisterRunningTool(conversationID, execID)
|
||||||
}
|
}
|
||||||
|
if w.unregisterCancelMonitor != nil && execID != "" {
|
||||||
|
defer w.unregisterCancelMonitor(execID)
|
||||||
|
}
|
||||||
|
|
||||||
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
|
// ctx 取消时关闭内层流,避免 amass 等长时间无换行输出时 Recv 永久阻塞。
|
||||||
stopWatch := make(chan struct{})
|
stopWatch := make(chan struct{})
|
||||||
@@ -181,11 +195,30 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
var invokeErr error
|
var invokeErr error
|
||||||
exitCode := 0
|
exitCode := 0
|
||||||
hasExitCode := false
|
hasExitCode := false
|
||||||
|
softReturned := false
|
||||||
|
var outCloseOnce sync.Once
|
||||||
|
closeOut := func() {
|
||||||
|
outCloseOnce.Do(func() { outW.Close() })
|
||||||
|
}
|
||||||
|
defer closeOut()
|
||||||
|
sendOut := func(resp *filesystem.ExecuteResponse, err error) bool {
|
||||||
|
if softReturned {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return outW.Send(resp, err)
|
||||||
|
}
|
||||||
|
|
||||||
idleWatch := security.NewShellInactivityWatch(noOutputSec)
|
idleWatch := security.NewShellInactivityWatch(noOutputSec)
|
||||||
if idleWatch != nil {
|
if idleWatch != nil {
|
||||||
defer idleWatch.Stop()
|
defer idleWatch.Stop()
|
||||||
}
|
}
|
||||||
|
var waitTimeoutCh <-chan time.Time
|
||||||
|
var waitTimer *time.Timer
|
||||||
|
if waitTimeoutSec > 0 {
|
||||||
|
waitTimer = time.NewTimer(time.Duration(waitTimeoutSec) * time.Second)
|
||||||
|
waitTimeoutCh = waitTimer.C
|
||||||
|
defer waitTimer.Stop()
|
||||||
|
}
|
||||||
|
|
||||||
type execRecvMsg struct {
|
type execRecvMsg struct {
|
||||||
resp *filesystem.ExecuteResponse
|
resp *filesystem.ExecuteResponse
|
||||||
@@ -206,8 +239,11 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
success = false
|
success = false
|
||||||
invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
|
invokeErr = fmt.Errorf("shell inactivity timeout (%ds)", idleWatch.Sec)
|
||||||
msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
msg := security.ShellNoOutputTimeoutMessage(idleWatch.Sec)
|
||||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
|
_ = sendOut(&filesystem.ExecuteResponse{Output: msg}, nil)
|
||||||
sb.WriteString(msg)
|
sb.WriteString(msg)
|
||||||
|
if w.appendPartialMonitor != nil && execID != "" {
|
||||||
|
w.appendPartialMonitor(execID, toolCallID, msg)
|
||||||
|
}
|
||||||
if w.outputChunk != nil && toolCallID != "" {
|
if w.outputChunk != nil && toolCallID != "" {
|
||||||
w.outputChunk("execute", toolCallID, msg)
|
w.outputChunk("execute", toolCallID, msg)
|
||||||
}
|
}
|
||||||
@@ -227,6 +263,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
case <-idleCh:
|
case <-idleCh:
|
||||||
fireInactivityTimeout()
|
fireInactivityTimeout()
|
||||||
break recvLoop
|
break recvLoop
|
||||||
|
case <-waitTimeoutCh:
|
||||||
|
if execID != "" && !softReturned {
|
||||||
|
msg := einoExecuteSoftWaitTimeoutResult(execID, waitTimeoutSec)
|
||||||
|
_ = outW.Send(&filesystem.ExecuteResponse{Output: msg}, nil)
|
||||||
|
softReturned = true
|
||||||
|
closeOut()
|
||||||
|
}
|
||||||
|
waitTimeoutCh = nil
|
||||||
case msg := <-recvCh:
|
case msg := <-recvCh:
|
||||||
rerr := msg.err
|
rerr := msg.err
|
||||||
resp := msg.resp
|
resp := msg.resp
|
||||||
@@ -244,7 +288,7 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
invokeErr = context.Canceled
|
invokeErr = context.Canceled
|
||||||
break recvLoop
|
break recvLoop
|
||||||
}
|
}
|
||||||
_ = outW.Send(nil, rerr)
|
_ = sendOut(nil, rerr)
|
||||||
break recvLoop
|
break recvLoop
|
||||||
}
|
}
|
||||||
if resp != nil {
|
if resp != nil {
|
||||||
@@ -263,11 +307,14 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
}
|
}
|
||||||
sb.WriteString(resp.Output)
|
sb.WriteString(resp.Output)
|
||||||
appended = resp.Output
|
appended = resp.Output
|
||||||
|
if w.appendPartialMonitor != nil && execID != "" {
|
||||||
|
w.appendPartialMonitor(execID, toolCallID, appended)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
|
if w.outputChunk != nil && strings.TrimSpace(appended) != "" {
|
||||||
w.outputChunk("execute", toolCallID, appended)
|
w.outputChunk("execute", toolCallID, appended)
|
||||||
}
|
}
|
||||||
if outW.Send(resp, nil) {
|
if sendOut(resp, nil) {
|
||||||
success = false
|
success = false
|
||||||
invokeErr = fmt.Errorf("execute stream closed by consumer")
|
invokeErr = fmt.Errorf("execute stream closed by consumer")
|
||||||
break recvLoop
|
break recvLoop
|
||||||
@@ -304,7 +351,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
// ADK 从本 Pipe 拼出 tool 消息正文;仅 Notify 尾标不会进入模型上下文。超时句写入流,与 UI 一致。
|
// ADK 从本 Pipe 拼出 tool 消息正文;仅 Notify 尾标不会进入模型上下文。超时句写入流,与 UI 一致。
|
||||||
if invokeErr != nil && errors.Is(invokeErr, context.DeadlineExceeded) {
|
if invokeErr != nil && errors.Is(invokeErr, context.DeadlineExceeded) {
|
||||||
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
hint := "\n\n" + einoExecuteTimeoutUserHint() + "\n"
|
||||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: hint}, nil)
|
_ = sendOut(&filesystem.ExecuteResponse{Output: hint}, nil)
|
||||||
|
if w.appendPartialMonitor != nil && execID != "" {
|
||||||
|
w.appendPartialMonitor(execID, toolCallID, hint)
|
||||||
|
}
|
||||||
if w.outputChunk != nil && tid != "" {
|
if w.outputChunk != nil && tid != "" {
|
||||||
w.outputChunk("execute", tid, hint)
|
w.outputChunk("execute", tid, hint)
|
||||||
}
|
}
|
||||||
@@ -313,9 +363,9 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
// 中断时循环内已逐行写入 stdout;此处只追加 USER INTERRUPT NOTE,避免整段输出重复。
|
// 中断时循环内已逐行写入 stdout;此处只追加 USER INTERRUPT NOTE,避免整段输出重复。
|
||||||
if invokeErr != nil && errors.Is(invokeErr, context.Canceled) && abortNote != "" {
|
if invokeErr != nil && errors.Is(invokeErr, context.Canceled) && abortNote != "" {
|
||||||
if partialStreamed != "" {
|
if partialStreamed != "" {
|
||||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: "\n\n" + mcp.AbortNoteBannerForModel + "\n" + abortNote}, nil)
|
_ = sendOut(&filesystem.ExecuteResponse{Output: "\n\n" + mcp.AbortNoteBannerForModel + "\n" + abortNote}, nil)
|
||||||
} else if text := strings.TrimSpace(sb.String()); text != "" {
|
} else if text := strings.TrimSpace(sb.String()); text != "" {
|
||||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
|
_ = sendOut(&filesystem.ExecuteResponse{Output: text + "\n"}, nil)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
rawOutput := sb.String()
|
rawOutput := sb.String()
|
||||||
@@ -323,7 +373,10 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
if !success && hasExitCode && exitCode != 0 {
|
if !success && hasExitCode && exitCode != 0 {
|
||||||
statusLine := security.ExecuteFailureStatusLine(exitCode)
|
statusLine := security.ExecuteFailureStatusLine(exitCode)
|
||||||
if !strings.Contains(rawOutput, "命令执行失败:") {
|
if !strings.Contains(rawOutput, "命令执行失败:") {
|
||||||
_ = outW.Send(&filesystem.ExecuteResponse{Output: statusLine}, nil)
|
_ = sendOut(&filesystem.ExecuteResponse{Output: statusLine}, nil)
|
||||||
|
if w.appendPartialMonitor != nil && execID != "" {
|
||||||
|
w.appendPartialMonitor(execID, toolCallID, statusLine)
|
||||||
|
}
|
||||||
sb.WriteString(statusLine)
|
sb.WriteString(statusLine)
|
||||||
}
|
}
|
||||||
fireBody = einomcp.ToolErrorPrefix + security.FormatCommandFailureResult(exitCode, rawOutput)
|
fireBody = einomcp.ToolErrorPrefix + security.FormatCommandFailureResult(exitCode, rawOutput)
|
||||||
@@ -332,10 +385,25 @@ func (w *einoStreamingShellWrap) ExecuteStreaming(ctx context.Context, input *fi
|
|||||||
w.finishMonitor(execID, toolCallID, command, sb.String(), success, invokeErr)
|
w.finishMonitor(execID, toolCallID, command, sb.String(), success, invokeErr)
|
||||||
}
|
}
|
||||||
if w.invokeNotify != nil {
|
if w.invokeNotify != nil {
|
||||||
w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, fireBody, invokeErr)
|
if !softReturned {
|
||||||
|
w.invokeNotify.Fire(toolCallID, "execute", agentTag, success, fireBody, invokeErr)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
outW.Close()
|
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec, w.toolWaitTimeoutSeconds)
|
||||||
}(sr, userCmd, execCancel, timeoutCancel, execCtx, convID, execReg, toolRunReg, monitorExecID, tid, w.shellNoOutputTimeoutSec)
|
|
||||||
|
|
||||||
return outR, nil
|
return outR, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func einoExecuteSoftWaitTimeoutResult(executionID string, waitTimeoutSec int) string {
|
||||||
|
waitText := "configured wait timeout"
|
||||||
|
if waitTimeoutSec > 0 {
|
||||||
|
waitText = fmt.Sprintf("%ds", waitTimeoutSec)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf(`工具已提交到后台执行,当前仍在运行。
|
||||||
|
|
||||||
|
execution_id: %s
|
||||||
|
status: running
|
||||||
|
wait_timeout: %s
|
||||||
|
|
||||||
|
你可以继续推理、改用其他工具,或调用 get_tool_execution / wait_tool_execution 读取 partial_output 并继续等待;也可以调用 cancel_tool_execution 取消。`, executionID, waitText)
|
||||||
|
}
|
||||||
|
|||||||
@@ -153,6 +153,81 @@ func TestEinoStreamingShellWrap_InactivityAfterPartialOutput(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestEinoStreamingShellWrap_SoftWaitTimeoutReturnsExecutionIDAndKeepsRunning(t *testing.T) {
|
||||||
|
inner := &mockStreamingShellPartialThenHang{}
|
||||||
|
partialCh := make(chan string, 4)
|
||||||
|
cancelCh := make(chan context.CancelFunc, 1)
|
||||||
|
unregistered := make(chan string, 1)
|
||||||
|
wrap := &einoStreamingShellWrap{
|
||||||
|
inner: inner,
|
||||||
|
toolWaitTimeoutSeconds: 1,
|
||||||
|
beginMonitor: func(toolCallID, command string) string {
|
||||||
|
return "exec-soft-wait"
|
||||||
|
},
|
||||||
|
appendPartialMonitor: func(executionID, toolCallID, chunk string) {
|
||||||
|
partialCh <- chunk
|
||||||
|
},
|
||||||
|
registerCancelMonitor: func(executionID string, cancel context.CancelFunc) {
|
||||||
|
if executionID == "exec-soft-wait" {
|
||||||
|
cancelCh <- cancel
|
||||||
|
}
|
||||||
|
},
|
||||||
|
unregisterCancelMonitor: func(executionID string) {
|
||||||
|
unregistered <- executionID
|
||||||
|
},
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
sr, err := wrap.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: "sudo whoami"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("ExecuteStreaming: %v", err)
|
||||||
|
}
|
||||||
|
defer sr.Close()
|
||||||
|
|
||||||
|
var got strings.Builder
|
||||||
|
for {
|
||||||
|
resp, rerr := sr.Recv()
|
||||||
|
if errors.Is(rerr, io.EOF) {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if rerr != nil {
|
||||||
|
t.Fatalf("recv: %v", rerr)
|
||||||
|
}
|
||||||
|
if resp != nil {
|
||||||
|
got.WriteString(resp.Output)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
body := got.String()
|
||||||
|
if !strings.Contains(body, "execution_id: exec-soft-wait") || !strings.Contains(body, "status: running") {
|
||||||
|
t.Fatalf("expected background execution marker, got: %q", body)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case chunk := <-partialCh:
|
||||||
|
if !strings.Contains(chunk, "[sudo] password") {
|
||||||
|
t.Fatalf("unexpected partial chunk: %q", chunk)
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
t.Fatal("expected streamed partial output before soft wait return")
|
||||||
|
}
|
||||||
|
if !inner.called {
|
||||||
|
t.Fatal("inner shell did not run")
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case registeredCancel := <-cancelCh:
|
||||||
|
registeredCancel()
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("expected execution cancel registration")
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case id := <-unregistered:
|
||||||
|
if id != "exec-soft-wait" {
|
||||||
|
t.Fatalf("unexpected unregistered id: %q", id)
|
||||||
|
}
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("expected execution cancel unregister")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
type mockStreamingShellHanging struct {
|
type mockStreamingShellHanging struct {
|
||||||
called bool
|
called bool
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -82,7 +82,7 @@ func RunEinoSingleChatModelAgent(
|
|||||||
}
|
}
|
||||||
|
|
||||||
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
|
toolInvokeNotify := einomcp.NewToolInvokeNotifyHolder()
|
||||||
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||||
mainDefs := ag.ToolsForRole(roleTools)
|
mainDefs := ag.ToolsForRole(roleTools)
|
||||||
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
|
mainTools, err := einomcp.ToolsFromDefinitions(ag, holder, mainDefs, recorder, nil, toolInvokeNotify, einoSingleAgentName)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -139,7 +139,7 @@ func RunEinoSingleChatModelAgent(
|
|||||||
}
|
}
|
||||||
if einoSkillMW != nil {
|
if einoSkillMW != nil {
|
||||||
if einoFSTools && einoLoc != nil {
|
if einoFSTools && einoLoc != nil {
|
||||||
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
fsMw, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, einoSingleAgentName, einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||||
if fsErr != nil {
|
if fsErr != nil {
|
||||||
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
|
return nil, fmt.Errorf("eino single filesystem 中间件: %w", fsErr)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -107,8 +107,12 @@ func subAgentFilesystemMiddleware(
|
|||||||
invokeNotify *einomcp.ToolInvokeNotifyHolder,
|
invokeNotify *einomcp.ToolInvokeNotifyHolder,
|
||||||
einoAgentName string,
|
einoAgentName string,
|
||||||
beginMonitor func(toolCallID, command string) string,
|
beginMonitor func(toolCallID, command string) string,
|
||||||
|
appendPartialMonitor func(executionID, toolCallID, chunk string),
|
||||||
|
registerCancelMonitor func(executionID string, cancel context.CancelFunc),
|
||||||
|
unregisterCancelMonitor func(executionID string),
|
||||||
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
finishMonitor func(executionID, toolCallID, command, stdout string, success bool, invokeErr error),
|
||||||
toolTimeoutMinutes int,
|
toolTimeoutMinutes int,
|
||||||
|
toolWaitTimeoutSeconds int,
|
||||||
shellNoOutputTimeoutSec int,
|
shellNoOutputTimeoutSec int,
|
||||||
outputChunk func(toolName, toolCallID, chunk string),
|
outputChunk func(toolName, toolCallID, chunk string),
|
||||||
) (adk.ChatModelAgentMiddleware, error) {
|
) (adk.ChatModelAgentMiddleware, error) {
|
||||||
@@ -123,8 +127,12 @@ func subAgentFilesystemMiddleware(
|
|||||||
einoAgentName: strings.TrimSpace(einoAgentName),
|
einoAgentName: strings.TrimSpace(einoAgentName),
|
||||||
outputChunk: outputChunk,
|
outputChunk: outputChunk,
|
||||||
beginMonitor: beginMonitor,
|
beginMonitor: beginMonitor,
|
||||||
|
appendPartialMonitor: appendPartialMonitor,
|
||||||
|
registerCancelMonitor: registerCancelMonitor,
|
||||||
|
unregisterCancelMonitor: unregisterCancelMonitor,
|
||||||
finishMonitor: finishMonitor,
|
finishMonitor: finishMonitor,
|
||||||
toolTimeoutMinutes: toolTimeoutMinutes,
|
toolTimeoutMinutes: toolTimeoutMinutes,
|
||||||
|
toolWaitTimeoutSeconds: toolWaitTimeoutSeconds,
|
||||||
shellNoOutputTimeoutSec: shellNoOutputTimeoutSec,
|
shellNoOutputTimeoutSec: shellNoOutputTimeoutSec,
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
@@ -138,6 +146,13 @@ func agentToolTimeoutMinutes(cfg *config.Config) int {
|
|||||||
return cfg.Agent.ToolTimeoutMinutes
|
return cfg.Agent.ToolTimeoutMinutes
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func agentToolWaitTimeoutSeconds(cfg *config.Config) int {
|
||||||
|
if cfg == nil {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return cfg.Agent.ToolWaitTimeoutSeconds
|
||||||
|
}
|
||||||
|
|
||||||
// agentShellNoOutputTimeoutSeconds:0=默认 300s(5 分钟);-1=关闭;>0=自定义秒数。
|
// agentShellNoOutputTimeoutSeconds:0=默认 300s(5 分钟);-1=关闭;>0=自定义秒数。
|
||||||
func agentShellNoOutputTimeoutSeconds(cfg *config.Config) int {
|
func agentShellNoOutputTimeoutSeconds(cfg *config.Config) int {
|
||||||
if cfg == nil {
|
if cfg == nil {
|
||||||
|
|||||||
@@ -60,6 +60,7 @@ func isEinoTransientRunError(err error) bool {
|
|||||||
"bad gateway",
|
"bad gateway",
|
||||||
"gateway timeout",
|
"gateway timeout",
|
||||||
"internal server error",
|
"internal server error",
|
||||||
|
"unexpected internal error",
|
||||||
"connection reset",
|
"connection reset",
|
||||||
"connection refused",
|
"connection refused",
|
||||||
"connection closed",
|
"connection closed",
|
||||||
@@ -72,6 +73,7 @@ func isEinoTransientRunError(err error) bool {
|
|||||||
"dial tcp",
|
"dial tcp",
|
||||||
"tls handshake timeout",
|
"tls handshake timeout",
|
||||||
"stream error",
|
"stream error",
|
||||||
|
"failed to receive stream chunk",
|
||||||
"goaway", // http2: server sent GOAWAY and closed the connection
|
"goaway", // http2: server sent GOAWAY and closed the connection
|
||||||
"unexpected eof",
|
"unexpected eof",
|
||||||
`": eof`, // net/http: Post "url": EOF (often wraps io.EOF)
|
`": eof`, // net/http: Post "url": EOF (often wraps io.EOF)
|
||||||
@@ -136,7 +138,8 @@ func einoTransientRunErrorUserDetail(err error) (kind, summary string) {
|
|||||||
case strings.Contains(lower, "overloaded") ||
|
case strings.Contains(lower, "overloaded") ||
|
||||||
strings.Contains(lower, "capacity") ||
|
strings.Contains(lower, "capacity") ||
|
||||||
strings.Contains(lower, "temporarily unavailable") ||
|
strings.Contains(lower, "temporarily unavailable") ||
|
||||||
strings.Contains(lower, "service unavailable"):
|
strings.Contains(lower, "service unavailable") ||
|
||||||
|
strings.Contains(lower, "unexpected internal error"):
|
||||||
kind = "upstream_busy"
|
kind = "upstream_busy"
|
||||||
case strings.Contains(lower, "connection reset") ||
|
case strings.Contains(lower, "connection reset") ||
|
||||||
strings.Contains(lower, "connection refused") ||
|
strings.Contains(lower, "connection refused") ||
|
||||||
@@ -153,6 +156,7 @@ func einoTransientRunErrorUserDetail(err error) (kind, summary string) {
|
|||||||
strings.Contains(lower, "unexpected eof"):
|
strings.Contains(lower, "unexpected eof"):
|
||||||
kind = "network"
|
kind = "network"
|
||||||
case strings.Contains(lower, "stream error") ||
|
case strings.Contains(lower, "stream error") ||
|
||||||
|
strings.Contains(lower, "failed to receive stream chunk") ||
|
||||||
strings.Contains(lower, "unexpected end of json"):
|
strings.Contains(lower, "unexpected end of json"):
|
||||||
kind = "stream"
|
kind = "stream"
|
||||||
default:
|
default:
|
||||||
|
|||||||
@@ -34,6 +34,7 @@ func TestIsEinoTransientRunError(t *testing.T) {
|
|||||||
{"rate limit", errors.New(`{"error":"rate limit exceeded"}`), true},
|
{"rate limit", errors.New(`{"error":"rate limit exceeded"}`), true},
|
||||||
{"connection reset", errors.New("read tcp: connection reset by peer"), true},
|
{"connection reset", errors.New("read tcp: connection reset by peer"), true},
|
||||||
{"http2 goaway", errors.New("failed to receive stream chunk: error, http2: server sent GOAWAY and closed the connection; LastStreamID=791, ErrCode=NO_ERROR"), true},
|
{"http2 goaway", errors.New("failed to receive stream chunk: error, http2: server sent GOAWAY and closed the connection; LastStreamID=791, ErrCode=NO_ERROR"), true},
|
||||||
|
{"unexpected internal stream chunk", errors.New("failed to receive stream chunk: error, The service encountered an unexpected internal error. Request id: 0217851391106464f01ec66621d0980a42fd45436ed75957a6a0a"), true},
|
||||||
{"unexpected eof", errors.New("unexpected EOF"), true},
|
{"unexpected eof", errors.New("unexpected EOF"), true},
|
||||||
{"503", errors.New("upstream returned 503"), true},
|
{"503", errors.New("upstream returned 503"), true},
|
||||||
{"iteration limit", errors.New("max iteration reached"), false},
|
{"iteration limit", errors.New("max iteration reached"), false},
|
||||||
@@ -74,6 +75,7 @@ func TestEinoTransientRunErrorUserDetail(t *testing.T) {
|
|||||||
{"upstream", errors.New("upstream returned 503"), "upstream_server"},
|
{"upstream", errors.New("upstream returned 503"), "upstream_server"},
|
||||||
{"network", errors.New("read tcp: connection reset by peer"), "network"},
|
{"network", errors.New("read tcp: connection reset by peer"), "network"},
|
||||||
{"stream", errors.New("unexpected end of JSON"), "stream"},
|
{"stream", errors.New("unexpected end of JSON"), "stream"},
|
||||||
|
{"stream chunk", errors.New("failed to receive stream chunk: error, The service encountered an unexpected internal error. Request id: abc"), "upstream_busy"},
|
||||||
}
|
}
|
||||||
for _, tc := range cases {
|
for _, tc := range cases {
|
||||||
tc := tc
|
tc := tc
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package multiagent
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
|
"io"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/cloudwego/eino/adk"
|
"github.com/cloudwego/eino/adk"
|
||||||
@@ -65,11 +66,51 @@ func hitlClearReturnDirectlyIfTransfer(ctx context.Context, toolName string) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func hitlEditedArgumentsNotice(original, edited string) string {
|
||||||
|
original = strings.TrimSpace(original)
|
||||||
|
edited = strings.TrimSpace(edited)
|
||||||
|
if edited == "" || edited == original {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return "[HITL] Human reviewer approved this tool call with edited arguments.\n" +
|
||||||
|
"Original arguments: " + original + "\n" +
|
||||||
|
"Executed arguments: " + edited + "\n\n"
|
||||||
|
}
|
||||||
|
|
||||||
|
func hitlPrependEditedArgumentsNotice(result, original, edited string) string {
|
||||||
|
notice := hitlEditedArgumentsNotice(original, edited)
|
||||||
|
if notice == "" {
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
return notice + result
|
||||||
|
}
|
||||||
|
|
||||||
|
func hitlCollectStringStream(sr *schema.StreamReader[string]) (string, error) {
|
||||||
|
if sr == nil {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
defer sr.Close()
|
||||||
|
var b strings.Builder
|
||||||
|
for {
|
||||||
|
chunk, err := sr.Recv()
|
||||||
|
if errors.Is(err, io.EOF) {
|
||||||
|
return b.String(), nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return b.String(), err
|
||||||
|
}
|
||||||
|
b.WriteString(chunk)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
||||||
return func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
return func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint {
|
||||||
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) {
|
||||||
|
originalArgs := ""
|
||||||
|
editedArgs := ""
|
||||||
if input != nil {
|
if input != nil {
|
||||||
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
||||||
|
originalArgs = input.Arguments
|
||||||
edited, err := fn(ctx, input.Name, input.Arguments)
|
edited, err := fn(ctx, input.Name, input.Arguments)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if IsHumanRejectError(err) {
|
if IsHumanRejectError(err) {
|
||||||
@@ -85,11 +126,17 @@ func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if edited != "" {
|
if edited != "" {
|
||||||
|
editedArgs = edited
|
||||||
input.Arguments = edited
|
input.Arguments = edited
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return next(ctx, input)
|
out, err := next(ctx, input)
|
||||||
|
if err != nil || out == nil {
|
||||||
|
return out, err
|
||||||
|
}
|
||||||
|
out.Result = hitlPrependEditedArgumentsNotice(out.Result, originalArgs, editedArgs)
|
||||||
|
return out, nil
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -97,8 +144,11 @@ func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware {
|
|||||||
func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware {
|
func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware {
|
||||||
return func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
|
return func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint {
|
||||||
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
|
return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) {
|
||||||
|
originalArgs := ""
|
||||||
|
editedArgs := ""
|
||||||
if input != nil {
|
if input != nil {
|
||||||
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil {
|
||||||
|
originalArgs = input.Arguments
|
||||||
edited, err := fn(ctx, input.Name, input.Arguments)
|
edited, err := fn(ctx, input.Name, input.Arguments)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if IsHumanRejectError(err) {
|
if IsHumanRejectError(err) {
|
||||||
@@ -111,11 +161,26 @@ func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware {
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
if edited != "" {
|
if edited != "" {
|
||||||
|
editedArgs = edited
|
||||||
input.Arguments = edited
|
input.Arguments = edited
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return next(ctx, input)
|
out, err := next(ctx, input)
|
||||||
|
if err != nil || out == nil {
|
||||||
|
return out, err
|
||||||
|
}
|
||||||
|
if hitlEditedArgumentsNotice(originalArgs, editedArgs) == "" {
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
result, collectErr := hitlCollectStringStream(out.Result)
|
||||||
|
if collectErr != nil {
|
||||||
|
return nil, collectErr
|
||||||
|
}
|
||||||
|
out.Result = schema.StreamReaderFromArray([]string{
|
||||||
|
hitlPrependEditedArgumentsNotice(result, originalArgs, editedArgs),
|
||||||
|
})
|
||||||
|
return out, nil
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -96,7 +96,17 @@ func (m *modelOutputGuardMiddleware) AfterModelRewriteState(
|
|||||||
badIndex := -1
|
badIndex := -1
|
||||||
argumentBytes := 0
|
argumentBytes := 0
|
||||||
if strings.EqualFold(strings.TrimSpace(finishReason), "length") {
|
if strings.EqualFold(strings.TrimSpace(finishReason), "length") {
|
||||||
reason = "output_limit"
|
if len(last.ToolCalls) == 0 {
|
||||||
|
reason = "output_limit"
|
||||||
|
} else {
|
||||||
|
for i, tc := range last.ToolCalls {
|
||||||
|
r, n := validateGeneratedToolCall(tc, m.cfg)
|
||||||
|
if r != "" {
|
||||||
|
reason, badIndex, argumentBytes = "output_limit", i, n
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
} else {
|
} else {
|
||||||
for i, tc := range last.ToolCalls {
|
for i, tc := range last.ToolCalls {
|
||||||
r, n := validateGeneratedToolCall(tc, m.cfg)
|
r, n := validateGeneratedToolCall(tc, m.cfg)
|
||||||
|
|||||||
@@ -50,6 +50,18 @@ func TestModelOutputGuardRejectsTruncatedToolCallBeforeExecution(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestModelOutputGuardAllowsValidToolCallDespiteLengthFinish(t *testing.T) {
|
||||||
|
original := `{"command":"echo ok"}`
|
||||||
|
state, err := runModelOutputGuard(t, []adk.Message{schema.UserMessage("run"), guardedAssistant(original, "length")}, config.MultiAgentEinoMiddlewareConfig{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
got := state.Messages[len(state.Messages)-1].ToolCalls[0].Function.Arguments
|
||||||
|
if got != original {
|
||||||
|
t.Fatalf("valid arguments should pass unchanged: %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestModelOutputGuardRejectsInvalidJSONShapes(t *testing.T) {
|
func TestModelOutputGuardRejectsInvalidJSONShapes(t *testing.T) {
|
||||||
for _, arguments := range []string{"", `[]`, `{"command":`} {
|
for _, arguments := range []string{"", `[]`, `{"command":`} {
|
||||||
t.Run(arguments, func(t *testing.T) {
|
t.Run(arguments, func(t *testing.T) {
|
||||||
|
|||||||
@@ -40,6 +40,13 @@ type RunResult struct {
|
|||||||
MCPExecutionIDs []string
|
MCPExecutionIDs []string
|
||||||
LastAgentTraceInput string // 已序列化的消息带(JSON):原生循环或 Eino 均写入,供续跑/攻击链等恢复上下文
|
LastAgentTraceInput string // 已序列化的消息带(JSON):原生循环或 Eino 均写入,供续跑/攻击链等恢复上下文
|
||||||
LastAgentTraceOutput string // 本轮助手侧对外展示文本(摘要或最终回复)
|
LastAgentTraceOutput string // 本轮助手侧对外展示文本(摘要或最终回复)
|
||||||
|
Finalized bool
|
||||||
|
Status string
|
||||||
|
CompletionReason string
|
||||||
|
EvidenceVerified bool
|
||||||
|
EvidenceRefs []string
|
||||||
|
PendingExecutionIDs []string
|
||||||
|
MissingChecks []string
|
||||||
}
|
}
|
||||||
|
|
||||||
// toolCallPendingInfo tracks a tool_call emitted to the UI so we can later
|
// toolCallPendingInfo tracks a tool_call emitted to the UI so we can later
|
||||||
@@ -135,7 +142,7 @@ func RunDeepAgent(
|
|||||||
mcpIDs = append(mcpIDs, id)
|
mcpIDs = append(mcpIDs, id)
|
||||||
mcpIDsMu.Unlock()
|
mcpIDsMu.Unlock()
|
||||||
}
|
}
|
||||||
einoExecBegin, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish := newEinoExecuteMonitorCallbacks(ctx, ag, recorder)
|
||||||
|
|
||||||
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
|
// 与单代理流式一致:在 response_start / response_delta 的 data 中带当前 mcpExecutionIds,供主聊天绑定复制与展示。
|
||||||
snapshotMCPIDs := func() []string {
|
snapshotMCPIDs := func() []string {
|
||||||
@@ -240,7 +247,7 @@ func RunDeepAgent(
|
|||||||
}
|
}
|
||||||
if einoSkillMW != nil {
|
if einoSkillMW != nil {
|
||||||
if einoFSTools && einoLoc != nil {
|
if einoFSTools && einoLoc != nil {
|
||||||
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
subFs, fsErr := subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, id, einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||||
if fsErr != nil {
|
if fsErr != nil {
|
||||||
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
|
return nil, fmt.Errorf("子代理 %q filesystem 中间件: %w", id, fsErr)
|
||||||
}
|
}
|
||||||
@@ -387,8 +394,12 @@ func RunDeepAgent(
|
|||||||
einoAgentName: orchestratorName,
|
einoAgentName: orchestratorName,
|
||||||
outputChunk: nil,
|
outputChunk: nil,
|
||||||
beginMonitor: einoExecBegin,
|
beginMonitor: einoExecBegin,
|
||||||
|
appendPartialMonitor: einoExecAppendPartial,
|
||||||
|
registerCancelMonitor: einoExecRegisterCancel,
|
||||||
|
unregisterCancelMonitor: einoExecUnregisterCancel,
|
||||||
finishMonitor: einoExecFinish,
|
finishMonitor: einoExecFinish,
|
||||||
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
|
toolTimeoutMinutes: agentToolTimeoutMinutes(appCfg),
|
||||||
|
toolWaitTimeoutSeconds: agentToolWaitTimeoutSeconds(appCfg),
|
||||||
shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg),
|
shellNoOutputTimeoutSec: agentShellNoOutputTimeoutSeconds(appCfg),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -486,7 +497,7 @@ func RunDeepAgent(
|
|||||||
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
|
// 构建 filesystem 中间件(与 Deep sub-agent 一致)
|
||||||
var peFsMw adk.ChatModelAgentMiddleware
|
var peFsMw adk.ChatModelAgentMiddleware
|
||||||
if einoSkillMW != nil && einoFSTools && einoLoc != nil {
|
if einoSkillMW != nil && einoFSTools && einoLoc != nil {
|
||||||
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecBegin, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
peFsMw, err = subAgentFilesystemMiddleware(ctx, einoLoc, toolInvokeNotify, "executor", einoExecBegin, einoExecAppendPartial, einoExecRegisterCancel, einoExecUnregisterCancel, einoExecFinish, agentToolTimeoutMinutes(appCfg), agentToolWaitTimeoutSeconds(appCfg), agentShellNoOutputTimeoutSeconds(appCfg), nil)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("plan_execute filesystem 中间件: %w", err)
|
return nil, fmt.Errorf("plan_execute filesystem 中间件: %w", err)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -893,7 +893,7 @@ func NewEinoHTTPClient(cfg *config.OpenAIConfig, base *http.Client) *http.Client
|
|||||||
if transport == nil {
|
if transport == nil {
|
||||||
transport = http.DefaultTransport
|
transport = http.DefaultTransport
|
||||||
}
|
}
|
||||||
transport = &reasoningToolChoiceCompatRoundTripper{base: transport}
|
transport = &reasoningToolChoiceCompatRoundTripper{base: transport, cfg: cfg}
|
||||||
if isClaudeProvider(cfg) {
|
if isClaudeProvider(cfg) {
|
||||||
transport = &claudeRoundTripper{
|
transport = &claudeRoundTripper{
|
||||||
base: transport,
|
base: transport,
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
package openai
|
package openai
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"strings"
|
||||||
|
|
||||||
"github.com/bytedance/sonic"
|
"github.com/bytedance/sonic"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -52,6 +54,28 @@ func StripReasoningIfForcedToolChoice(rawBody []byte) ([]byte, error) {
|
|||||||
return out, nil
|
return out, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// StripToolChoiceForThinkingMode removes tool_choice while preserving tools and
|
||||||
|
// thinking fields. DeepSeek thinking mode can use tools, but rejects the
|
||||||
|
// tool_choice parameter itself on some agent requests.
|
||||||
|
func StripToolChoiceForThinkingMode(rawBody []byte) ([]byte, error) {
|
||||||
|
var payload map[string]any
|
||||||
|
if err := sonic.Unmarshal(rawBody, &payload); err != nil {
|
||||||
|
return rawBody, nil
|
||||||
|
}
|
||||||
|
if !thinkingModeEnabledByPayload(payload) {
|
||||||
|
return rawBody, nil
|
||||||
|
}
|
||||||
|
if _, ok := payload["tool_choice"]; !ok {
|
||||||
|
return rawBody, nil
|
||||||
|
}
|
||||||
|
delete(payload, "tool_choice")
|
||||||
|
out, err := sonic.Marshal(payload)
|
||||||
|
if err != nil {
|
||||||
|
return rawBody, err
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
func stripReasoningFields(payload map[string]any) bool {
|
func stripReasoningFields(payload map[string]any) bool {
|
||||||
changed := false
|
changed := false
|
||||||
for _, key := range reasoningPayloadKeys {
|
for _, key := range reasoningPayloadKeys {
|
||||||
@@ -77,3 +101,17 @@ func forcedToolChoiceIncompatibleWithThinking(payload map[string]any) bool {
|
|||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func thinkingModeEnabledByPayload(payload map[string]any) bool {
|
||||||
|
thinking, ok := payload["thinking"]
|
||||||
|
if !ok || thinking == nil {
|
||||||
|
// DeepSeek enables thinking by default unless explicitly disabled.
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if m, ok := thinking.(map[string]any); ok {
|
||||||
|
if typ, ok := m["type"].(string); ok && strings.EqualFold(strings.TrimSpace(typ), "disabled") {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|||||||
@@ -5,6 +5,8 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/config"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestStripReasoningFromChatCompletionBody(t *testing.T) {
|
func TestStripReasoningFromChatCompletionBody(t *testing.T) {
|
||||||
@@ -82,6 +84,58 @@ func TestStripReasoningIfForcedToolChoice(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestStripToolChoiceForThinkingMode(t *testing.T) {
|
||||||
|
cases := []struct {
|
||||||
|
name string
|
||||||
|
in string
|
||||||
|
wantToolChoice bool
|
||||||
|
wantThinking bool
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "enabled thinking removes tool_choice",
|
||||||
|
in: `{"model":"deepseek-v4","messages":[],"thinking":{"type":"enabled"},"tool_choice":"required","tools":[{"type":"function","function":{"name":"scan"}}]}`,
|
||||||
|
wantToolChoice: false,
|
||||||
|
wantThinking: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "default thinking removes tool_choice",
|
||||||
|
in: `{"model":"deepseek-v4","messages":[],"tool_choice":"auto","tools":[]}`,
|
||||||
|
wantToolChoice: false,
|
||||||
|
wantThinking: false,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "disabled thinking keeps tool_choice",
|
||||||
|
in: `{"model":"deepseek-v4","messages":[],"thinking":{"type":"disabled"},"tool_choice":"required","tools":[]}`,
|
||||||
|
wantToolChoice: true,
|
||||||
|
wantThinking: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "no tool_choice unchanged",
|
||||||
|
in: `{"model":"deepseek-v4","messages":[],"thinking":{"type":"enabled"},"tools":[]}`,
|
||||||
|
wantToolChoice: false,
|
||||||
|
wantThinking: true,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
for _, tc := range cases {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
out, err := StripToolChoiceForThinkingMode([]byte(tc.in))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
s := string(out)
|
||||||
|
if strings.Contains(s, "tool_choice") != tc.wantToolChoice {
|
||||||
|
t.Fatalf("tool_choice presence mismatch, got %s", s)
|
||||||
|
}
|
||||||
|
if strings.Contains(s, "thinking") != tc.wantThinking {
|
||||||
|
t.Fatalf("thinking presence mismatch, got %s", s)
|
||||||
|
}
|
||||||
|
if !strings.Contains(s, "tools") {
|
||||||
|
t.Fatalf("expected tools preserved, got %s", s)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestReasoningToolChoiceCompatRoundTripper(t *testing.T) {
|
func TestReasoningToolChoiceCompatRoundTripper(t *testing.T) {
|
||||||
var gotBody string
|
var gotBody string
|
||||||
rt := &reasoningToolChoiceCompatRoundTripper{
|
rt := &reasoningToolChoiceCompatRoundTripper{
|
||||||
@@ -113,6 +167,44 @@ func TestReasoningToolChoiceCompatRoundTripper(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestReasoningToolChoiceCompatRoundTripperDeepSeek(t *testing.T) {
|
||||||
|
var gotBody string
|
||||||
|
rt := &reasoningToolChoiceCompatRoundTripper{
|
||||||
|
cfg: &config.OpenAIConfig{
|
||||||
|
BaseURL: "https://api.deepseek.com/v1",
|
||||||
|
Model: "deepseek-v4",
|
||||||
|
},
|
||||||
|
base: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
|
||||||
|
b, _ := io.ReadAll(req.Body)
|
||||||
|
gotBody = string(b)
|
||||||
|
return &http.Response{
|
||||||
|
StatusCode: 200,
|
||||||
|
Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`)),
|
||||||
|
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||||
|
}, nil
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
req, err := http.NewRequest(http.MethodPost, "https://api.deepseek.com/v1/chat/completions", strings.NewReader(
|
||||||
|
`{"model":"deepseek-v4","thinking":{"type":"enabled"},"tool_choice":"required","tools":[],"messages":[]}`,
|
||||||
|
))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
_, err = rt.RoundTrip(req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if strings.Contains(gotBody, "tool_choice") {
|
||||||
|
t.Fatalf("expected DeepSeek tool_choice stripped in transit, got %s", gotBody)
|
||||||
|
}
|
||||||
|
if !strings.Contains(gotBody, "thinking") {
|
||||||
|
t.Fatalf("expected thinking preserved for DeepSeek, got %s", gotBody)
|
||||||
|
}
|
||||||
|
if !strings.Contains(gotBody, "tools") {
|
||||||
|
t.Fatalf("expected tools preserved for DeepSeek, got %s", gotBody)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
type roundTripperFunc func(*http.Request) (*http.Response, error)
|
type roundTripperFunc func(*http.Request) (*http.Response, error)
|
||||||
|
|
||||||
func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
|
|||||||
@@ -6,6 +6,8 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
|
"cyberstrike-ai/internal/config"
|
||||||
)
|
)
|
||||||
|
|
||||||
// reasoningToolChoiceCompatRoundTripper strips thinking/reasoning fields from
|
// reasoningToolChoiceCompatRoundTripper strips thinking/reasoning fields from
|
||||||
@@ -13,6 +15,7 @@ import (
|
|||||||
// when thinking mode is enabled on the same request.
|
// when thinking mode is enabled on the same request.
|
||||||
type reasoningToolChoiceCompatRoundTripper struct {
|
type reasoningToolChoiceCompatRoundTripper struct {
|
||||||
base http.RoundTripper
|
base http.RoundTripper
|
||||||
|
cfg *config.OpenAIConfig
|
||||||
}
|
}
|
||||||
|
|
||||||
func (rt *reasoningToolChoiceCompatRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
func (rt *reasoningToolChoiceCompatRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
@@ -32,7 +35,13 @@ func (rt *reasoningToolChoiceCompatRoundTripper) RoundTrip(req *http.Request) (*
|
|||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
patched, perr := StripReasoningIfForcedToolChoice(body)
|
patched := body
|
||||||
|
var perr error
|
||||||
|
if isDeepSeekToolChoiceCompatProfile(rt.cfg) {
|
||||||
|
patched, perr = StripToolChoiceForThinkingMode(body)
|
||||||
|
} else {
|
||||||
|
patched, perr = StripReasoningIfForcedToolChoice(body)
|
||||||
|
}
|
||||||
if perr != nil {
|
if perr != nil {
|
||||||
patched = body
|
patched = body
|
||||||
}
|
}
|
||||||
@@ -41,3 +50,19 @@ func (rt *reasoningToolChoiceCompatRoundTripper) RoundTrip(req *http.Request) (*
|
|||||||
req.Header.Set("Content-Length", strconv.Itoa(len(patched)))
|
req.Header.Set("Content-Length", strconv.Itoa(len(patched)))
|
||||||
return rt.base.RoundTrip(req)
|
return rt.base.RoundTrip(req)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func isDeepSeekToolChoiceCompatProfile(cfg *config.OpenAIConfig) bool {
|
||||||
|
if cfg == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
profile := strings.ToLower(strings.TrimSpace(cfg.Reasoning.ProfileEffective()))
|
||||||
|
if profile == "deepseek" || profile == "deepseek_compat" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if profile != "" && profile != "auto" {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
baseURL := strings.ToLower(cfg.BaseURL)
|
||||||
|
model := strings.ToLower(cfg.Model)
|
||||||
|
return strings.Contains(baseURL, "deepseek") || strings.Contains(model, "deepseek")
|
||||||
|
}
|
||||||
|
|||||||
@@ -9,8 +9,8 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
bodyDepFactLine = regexp.MustCompile(`(?im)^[\s\-*]*依赖事实\s*[::]\s*([a-z0-9][a-z0-9._/-]*)`)
|
bodyDepFactLine = regexp.MustCompile(`(?im)^[\s\-*]*依赖事实\s*[::]\s*([a-zA-Z0-9][a-zA-Z0-9._/-]*)`)
|
||||||
bodyRelFactLine = regexp.MustCompile(`(?im)^[\s\-*]*相关\s*fact_key\s*[::]\s*([a-z0-9][a-z0-9._/-]*)`)
|
bodyRelFactLine = regexp.MustCompile(`(?im)^[\s\-*]*相关\s*fact_key\s*[::]\s*([a-zA-Z0-9][a-zA-Z0-9._/-]*)`)
|
||||||
bodyAssocSection = regexp.MustCompile(`(?im)^##\s*关联\s*$`)
|
bodyAssocSection = regexp.MustCompile(`(?im)^##\s*关联\s*$`)
|
||||||
bodySyncLinksHead = "结构化关系边(自动同步)"
|
bodySyncLinksHead = "结构化关系边(自动同步)"
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -74,6 +74,21 @@ func (e *Executor) SetToolOutputSpillRoot(rootDir string) {
|
|||||||
e.spillRootDir = strings.TrimSpace(rootDir)
|
e.spillRootDir = strings.TrimSpace(rootDir)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (e *Executor) wrapToolOutputCallback(ctx context.Context, cb ToolOutputCallback) ToolOutputCallback {
|
||||||
|
executionID := mcp.MCPExecutionIDFromContext(ctx)
|
||||||
|
if e == nil || e.mcpServer == nil || strings.TrimSpace(executionID) == "" {
|
||||||
|
return cb
|
||||||
|
}
|
||||||
|
return func(chunk string) {
|
||||||
|
if chunk != "" {
|
||||||
|
e.mcpServer.AppendToolExecutionPartialOutput(executionID, chunk)
|
||||||
|
}
|
||||||
|
if cb != nil {
|
||||||
|
cb(chunk)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (e *Executor) spillOptsFromContext(ctx context.Context) tooloutput.SpillOpts {
|
func (e *Executor) spillOptsFromContext(ctx context.Context) tooloutput.SpillOpts {
|
||||||
root := ""
|
root := ""
|
||||||
if e != nil {
|
if e != nil {
|
||||||
@@ -184,8 +199,9 @@ func (e *Executor) ExecuteTool(ctx context.Context, toolName string, args map[st
|
|||||||
var output string
|
var output string
|
||||||
var err error
|
var err error
|
||||||
spill := e.spillOptsFromContext(ctx)
|
spill := e.spillOptsFromContext(ctx)
|
||||||
// 如果上层提供了 stdout/stderr 增量回调,则边执行边读取并回调。
|
// 如果上层提供了 stdout/stderr 增量回调,或当前处于 MCP execution 中,则边执行边读取并回调。
|
||||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" {
|
||||||
|
cb = e.wrapToolOutputCallback(ctx, cb)
|
||||||
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
|
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
|
||||||
if err != nil && shouldRetryWithPTY(output) {
|
if err != nil && shouldRetryWithPTY(output) {
|
||||||
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
|
e.logger.Info("检测到工具需要 TTY,使用 PTY 重试",
|
||||||
@@ -948,8 +964,9 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
|||||||
var output string
|
var output string
|
||||||
var err error
|
var err error
|
||||||
spill := e.spillOptsFromContext(ctx)
|
spill := e.spillOptsFromContext(ctx)
|
||||||
// 若上层提供工具输出增量回调,则边执行边流式读取。
|
// 若上层提供工具输出增量回调,或当前处于 MCP execution 中,则边执行边流式读取。
|
||||||
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); ok && cb != nil {
|
if cb, ok := ctx.Value(ToolOutputCallbackCtxKey).(ToolOutputCallback); (ok && cb != nil) || mcp.MCPExecutionIDFromContext(ctx) != "" {
|
||||||
|
cb = e.wrapToolOutputCallback(ctx, cb)
|
||||||
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
|
output, err = streamCommandOutput(ctx, cmd, cb, ResolveShellNoOutputTimeoutSeconds(e.shellNoOutputTimeoutSec), e.toolOutputMaxBytes, spill)
|
||||||
if err != nil && shouldRetryWithPTY(output) {
|
if err != nil && shouldRetryWithPTY(output) {
|
||||||
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
|
e.logger.Info("检测到系统命令需要 TTY,使用 PTY 重试")
|
||||||
|
|||||||
@@ -73,6 +73,43 @@ func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T)
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestExecToolSoftWaitExposesPartialOutput(t *testing.T) {
|
||||||
|
executor, server := setupTestExecutor(t)
|
||||||
|
server.ConfigureToolWaitTimeoutSeconds(1)
|
||||||
|
mcp.RegisterExecutionControlTools(server, nil)
|
||||||
|
server.RegisterTool(mcp.Tool{Name: "exec", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
|
||||||
|
return executor.ExecuteTool(ctx, "exec", args)
|
||||||
|
})
|
||||||
|
|
||||||
|
result, executionID, err := server.CallTool(context.Background(), "exec", map[string]interface{}{
|
||||||
|
"command": "for i in 1 2 3 4; do echo partial-$i; sleep 0.3; done; sleep 5",
|
||||||
|
"shell": "sh",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CallTool exec: %v", err)
|
||||||
|
}
|
||||||
|
if executionID == "" || result == nil || !result.IsError {
|
||||||
|
t.Fatalf("expected soft wait timeout, id=%q result=%#v", executionID, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
status, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
|
||||||
|
"execution_id": executionID,
|
||||||
|
"include_partial_output": true,
|
||||||
|
"partial_output_max_bytes": 4096,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("get_tool_execution: %v", err)
|
||||||
|
}
|
||||||
|
body := mcp.ToolResultPlainText(status)
|
||||||
|
if !strings.Contains(body, `"status": "running"`) {
|
||||||
|
t.Fatalf("expected running execution, got: %s", body)
|
||||||
|
}
|
||||||
|
if !strings.Contains(body, "partial-") || !strings.Contains(body, "partial_output") {
|
||||||
|
t.Fatalf("expected partial output in execution status, got: %s", body)
|
||||||
|
}
|
||||||
|
server.CancelToolExecution(executionID)
|
||||||
|
}
|
||||||
|
|
||||||
func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
|
func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
|
||||||
executor, _ := setupTestExecutor(t)
|
executor, _ := setupTestExecutor(t)
|
||||||
res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{
|
res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{
|
||||||
|
|||||||
@@ -172,6 +172,8 @@ func permissionForRequest(method, fullPath string) string {
|
|||||||
return "workflow:read"
|
return "workflow:read"
|
||||||
case strings.HasPrefix(path, "/workflow-package-inspections"), strings.HasPrefix(path, "/workflow-package-imports"):
|
case strings.HasPrefix(path, "/workflow-package-inspections"), strings.HasPrefix(path, "/workflow-package-imports"):
|
||||||
return "workflow:write"
|
return "workflow:write"
|
||||||
|
case path == "/workflows/generate-draft":
|
||||||
|
return "workflow:write"
|
||||||
case strings.HasPrefix(path, "/workflows"):
|
case strings.HasPrefix(path, "/workflows"):
|
||||||
if path == "/workflows/validate" || path == "/workflows/dry-run" || strings.HasSuffix(path, "/resume") {
|
if path == "/workflows/validate" || path == "/workflows/dry-run" || strings.HasSuffix(path, "/resume") {
|
||||||
return "workflow:execute"
|
return "workflow:execute"
|
||||||
@@ -265,7 +267,7 @@ func isProcessGlobalMutationPath(path string) bool {
|
|||||||
}
|
}
|
||||||
if strings.HasPrefix(path, "/workflows") {
|
if strings.HasPrefix(path, "/workflows") {
|
||||||
// Workflow runs inherit conversation access; definitions are global.
|
// Workflow runs inherit conversation access; definitions are global.
|
||||||
return !strings.HasPrefix(path, "/workflows/runs/") && path != "/workflows/validate" && path != "/workflows/dry-run"
|
return !strings.HasPrefix(path, "/workflows/runs/") && path != "/workflows/validate" && path != "/workflows/dry-run" && path != "/workflows/generate-draft"
|
||||||
}
|
}
|
||||||
if strings.HasPrefix(path, "/workflow-package-inspections") || strings.HasPrefix(path, "/workflow-package-imports") {
|
if strings.HasPrefix(path, "/workflow-package-inspections") || strings.HasPrefix(path, "/workflow-package-imports") {
|
||||||
return true
|
return true
|
||||||
|
|||||||
@@ -157,9 +157,15 @@ func TestWorkflowRunPermissionIsSeparateFromDefinitionManagement(t *testing.T) {
|
|||||||
if got := permissionForRequest(http.MethodPost, "/api/workflows/runs/run-1/resume"); got != "workflow:execute" {
|
if got := permissionForRequest(http.MethodPost, "/api/workflows/runs/run-1/resume"); got != "workflow:execute" {
|
||||||
t.Fatalf("resume permission = %q, want workflow:execute", got)
|
t.Fatalf("resume permission = %q, want workflow:execute", got)
|
||||||
}
|
}
|
||||||
|
if got := permissionForRequest(http.MethodPost, "/api/workflows/generate-draft"); got != "workflow:write" {
|
||||||
|
t.Fatalf("generate draft permission = %q, want workflow:write", got)
|
||||||
|
}
|
||||||
if got := permissionForRequest(http.MethodPut, "/api/workflows/workflow-1"); got != "workflow:write" {
|
if got := permissionForRequest(http.MethodPut, "/api/workflows/workflow-1"); got != "workflow:write" {
|
||||||
t.Fatalf("definition permission = %q, want workflow:write", got)
|
t.Fatalf("definition permission = %q, want workflow:write", got)
|
||||||
}
|
}
|
||||||
|
if isProcessGlobalMutationPath("/workflows/generate-draft") {
|
||||||
|
t.Fatalf("generate draft should not be treated as a process-global mutation")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestRBACDenyHookReceivesDeniedDecision(t *testing.T) {
|
func TestRBACDenyHookReceivesDeniedDecision(t *testing.T) {
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user