mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-01 00:27:35 +02:00
297 lines
10 KiB
Go
297 lines
10 KiB
Go
package mcp
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"cyberstrike-ai/internal/mcp/builtin"
|
|
)
|
|
|
|
const (
|
|
defaultExecutionWaitTimeout = 60 * time.Second
|
|
maxExecutionWaitTimeout = 10 * time.Minute
|
|
defaultPartialPreviewBytes = 4096
|
|
maxPartialPreviewBytes = 64 * 1024
|
|
)
|
|
|
|
// RegisterExecutionControlTools exposes execution handle operations to Eino as
|
|
// ordinary MCP tools. This keeps the agent loop native: the model calls a tool,
|
|
// receives a bounded result, and may call wait_tool_execution again if needed.
|
|
func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) {
|
|
if server == nil {
|
|
return
|
|
}
|
|
|
|
server.RegisterTool(Tool{
|
|
Name: builtin.ToolGetToolExecution,
|
|
Description: "查询后台工具 execution 的当前状态、结果和错误。用于外部 MCP 工具等待超时后，凭 execution_id 继续查看进度。",
|
|
ShortDescription: "查询后台工具执行状态",
|
|
InputSchema: map[string]interface{}{
|
|
"type": "object",
|
|
"properties": map[string]interface{}{
|
|
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
|
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览，默认 true"},
|
|
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数，默认 4096，最大 65536"},
|
|
},
|
|
"required": []string{"execution_id"},
|
|
},
|
|
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
|
id := stringArg(args, "execution_id")
|
|
if id == "" {
|
|
return textToolResult("execution_id 必填", true), nil
|
|
}
|
|
exec := lookupToolExecution(server, external, id)
|
|
if exec == nil {
|
|
return textToolResult("未找到该 execution_id: "+id, true), nil
|
|
}
|
|
return textToolResult(formatExecutionForModel(exec, executionFormatOptionsFromArgs(args)), false), nil
|
|
})
|
|
|
|
server.RegisterTool(Tool{
|
|
Name: builtin.ToolWaitToolExecution,
|
|
Description: "继续等待一个后台工具 execution 完成。每次等待都有 timeout_seconds 上限；若仍未完成，会返回当前状态，模型可稍后再次调用。",
|
|
ShortDescription: "有界等待后台工具执行",
|
|
InputSchema: map[string]interface{}{
|
|
"type": "object",
|
|
"properties": map[string]interface{}{
|
|
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
|
"timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数，默认 60，最大 600"},
|
|
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览，默认 true"},
|
|
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数，默认 4096，最大 65536"},
|
|
},
|
|
"required": []string{"execution_id"},
|
|
},
|
|
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
|
id := stringArg(args, "execution_id")
|
|
if id == "" {
|
|
return textToolResult("execution_id 必填", true), nil
|
|
}
|
|
wait := durationSecondsArg(args, "timeout_seconds", defaultExecutionWaitTimeout, maxExecutionWaitTimeout)
|
|
snap, err := waitToolExecutionSnapshot(ctx, server, external, id, wait)
|
|
if err != nil && !errors.Is(err, ErrExecutionWaitTimeout) {
|
|
return textToolResult("等待 execution 失败: "+err.Error(), true), nil
|
|
}
|
|
if snap == nil || snap.Execution == nil {
|
|
return textToolResult("未找到该 execution_id: "+id, true), nil
|
|
}
|
|
body := formatExecutionForModel(snap.Execution, executionFormatOptionsFromArgs(args))
|
|
if errors.Is(err, ErrExecutionWaitTimeout) {
|
|
body += "\n\n本次等待已到达 timeout_seconds，上述 execution 仍未完成。可继续等待、取消，或采用其他步骤。"
|
|
}
|
|
return textToolResult(body, false), nil
|
|
})
|
|
|
|
server.RegisterTool(Tool{
|
|
Name: builtin.ToolCancelToolExecution,
|
|
Description: "取消一个后台工具 execution。用于外部 MCP 工具长时间运行、误调用或用户要求停止时。",
|
|
ShortDescription: "取消后台工具执行",
|
|
InputSchema: map[string]interface{}{
|
|
"type": "object",
|
|
"properties": map[string]interface{}{
|
|
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
|
|
"reason": map[string]interface{}{"type": "string", "description": "取消原因，可选，会写入终止说明"},
|
|
},
|
|
"required": []string{"execution_id"},
|
|
},
|
|
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
|
|
id := stringArg(args, "execution_id")
|
|
if id == "" {
|
|
return textToolResult("execution_id 必填", true), nil
|
|
}
|
|
reason := stringArg(args, "reason")
|
|
if server.CancelToolExecutionWithNote(id, reason) {
|
|
return textToolResult("已请求取消内部工具 execution: "+id, false), nil
|
|
}
|
|
if external != nil && external.CancelToolExecutionWithNote(id, reason) {
|
|
return textToolResult("已请求取消外部 MCP execution: "+id, false), nil
|
|
}
|
|
return textToolResult("未找到进行中的 execution，或该 execution 已结束: "+id, true), nil
|
|
})
|
|
}
|
|
|
|
func waitToolExecutionSnapshot(ctx context.Context, server *Server, external *ExternalMCPManager, id string, wait time.Duration) (*ExecutionSnapshot, error) {
|
|
if server != nil && server.executionService != nil && server.executionService.getEntry(id) != nil {
|
|
return server.executionService.Wait(ctx, id, wait)
|
|
}
|
|
if external != nil && external.executionService != nil && external.executionService.getEntry(id) != nil {
|
|
return external.executionService.Wait(ctx, id, wait)
|
|
}
|
|
if server != nil && server.executionService != nil {
|
|
if snap, err := server.executionService.Get(id); err == nil {
|
|
return snap, nil
|
|
}
|
|
}
|
|
if external != nil && external.executionService != nil {
|
|
return external.executionService.Get(id)
|
|
}
|
|
exec := lookupToolExecution(server, external, id)
|
|
if exec == nil {
|
|
return nil, fmt.Errorf("execution not found: %s", id)
|
|
}
|
|
return &ExecutionSnapshot{Execution: exec}, nil
|
|
}
|
|
|
|
func lookupToolExecution(server *Server, external *ExternalMCPManager, id string) *ToolExecution {
|
|
if server != nil {
|
|
if exec, ok := server.GetExecution(id); ok && exec != nil {
|
|
return exec
|
|
}
|
|
}
|
|
if external != nil {
|
|
if exec, ok := external.GetExecution(id); ok && exec != nil {
|
|
return exec
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type executionFormatOptions struct {
|
|
includePartialOutput bool
|
|
partialMaxBytes int
|
|
}
|
|
|
|
func executionFormatOptionsFromArgs(args map[string]interface{}) executionFormatOptions {
|
|
includePartial := true
|
|
if raw, ok := args["include_partial_output"]; ok {
|
|
if b, ok := raw.(bool); ok {
|
|
includePartial = b
|
|
} else if s := strings.TrimSpace(fmt.Sprint(raw)); s != "" {
|
|
includePartial = strings.EqualFold(s, "true") || s == "1" || strings.EqualFold(s, "yes")
|
|
}
|
|
}
|
|
maxBytes := intArg(args, "partial_output_max_bytes", defaultPartialPreviewBytes, maxPartialPreviewBytes)
|
|
return executionFormatOptions{includePartialOutput: includePartial, partialMaxBytes: maxBytes}
|
|
}
|
|
|
|
func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) string {
|
|
if exec == nil {
|
|
return "execution: null"
|
|
}
|
|
payload := map[string]interface{}{
|
|
"execution_id": exec.ID,
|
|
"tool": exec.ToolName,
|
|
"status": exec.Status,
|
|
"started_at": exec.StartTime.Format(time.RFC3339),
|
|
}
|
|
if exec.EndTime != nil {
|
|
payload["ended_at"] = exec.EndTime.Format(time.RFC3339)
|
|
}
|
|
if exec.Duration > 0 {
|
|
payload["duration"] = exec.Duration.String()
|
|
}
|
|
if exec.Error != "" {
|
|
payload["error"] = exec.Error
|
|
}
|
|
if exec.Result != nil {
|
|
payload["result"] = ToolResultPlainText(exec.Result)
|
|
payload["is_error"] = exec.Result.IsError
|
|
}
|
|
if opts.includePartialOutput && exec.PartialOutput != "" {
|
|
partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes)
|
|
payload["partial_output"] = partial
|
|
payload["partial_output_bytes"] = exec.PartialOutputBytes
|
|
payload["partial_output_truncated"] = exec.PartialOutputTruncated || len([]byte(partial)) < len([]byte(exec.PartialOutput))
|
|
if exec.PartialOutputUpdatedAt != nil {
|
|
payload["partial_output_updated_at"] = exec.PartialOutputUpdatedAt.Format(time.RFC3339)
|
|
}
|
|
}
|
|
b, err := json.MarshalIndent(payload, "", " ")
|
|
if err != nil {
|
|
return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error)
|
|
}
|
|
return string(b)
|
|
}
|
|
|
|
func tailStringBytes(s string, maxBytes int) string {
|
|
if maxBytes <= 0 {
|
|
maxBytes = defaultPartialPreviewBytes
|
|
}
|
|
b := []byte(s)
|
|
if len(b) <= maxBytes {
|
|
return s
|
|
}
|
|
return string(b[len(b)-maxBytes:])
|
|
}
|
|
|
|
func textToolResult(text string, isErr bool) *ToolResult {
|
|
return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr}
|
|
}
|
|
|
|
func stringArg(args map[string]interface{}, key string) string {
|
|
if args == nil {
|
|
return ""
|
|
}
|
|
raw, ok := args[key]
|
|
if !ok || raw == nil {
|
|
return ""
|
|
}
|
|
switch v := raw.(type) {
|
|
case string:
|
|
return strings.TrimSpace(v)
|
|
default:
|
|
return strings.TrimSpace(fmt.Sprint(v))
|
|
}
|
|
}
|
|
|
|
func durationSecondsArg(args map[string]interface{}, key string, def, max time.Duration) time.Duration {
|
|
if args == nil {
|
|
return def
|
|
}
|
|
var seconds float64
|
|
switch v := args[key].(type) {
|
|
case int:
|
|
seconds = float64(v)
|
|
case int64:
|
|
seconds = float64(v)
|
|
case float64:
|
|
seconds = v
|
|
case json.Number:
|
|
f, _ := v.Float64()
|
|
seconds = f
|
|
case string:
|
|
f, _ := strconv.ParseFloat(strings.TrimSpace(v), 64)
|
|
seconds = f
|
|
}
|
|
if seconds <= 0 {
|
|
return def
|
|
}
|
|
d := time.Duration(seconds * float64(time.Second))
|
|
if max > 0 && d > max {
|
|
return max
|
|
}
|
|
return d
|
|
}
|
|
|
|
func intArg(args map[string]interface{}, key string, def, max int) int {
|
|
if args == nil {
|
|
return def
|
|
}
|
|
var n int
|
|
switch v := args[key].(type) {
|
|
case int:
|
|
n = v
|
|
case int64:
|
|
n = int(v)
|
|
case float64:
|
|
n = int(v)
|
|
case json.Number:
|
|
i, _ := v.Int64()
|
|
n = int(i)
|
|
case string:
|
|
i, _ := strconv.Atoi(strings.TrimSpace(v))
|
|
n = i
|
|
}
|
|
if n <= 0 {
|
|
return def
|
|
}
|
|
if max > 0 && n > max {
|
|
return max
|
|
}
|
|
return n
|
|
}
|