logan
2336c9433c
- Agent 支持 SSE（步骤边推、正文分块），JSON 非流式 fallback - tool calling 覆盖 openai-chat / anthropic / responses（单测解析回灌） - P1：mcp.json mode=read_write 才可写；需 confirm_danger；Agent 剥离模型自带令牌 - GET /api/ai/mcp-status、POST /api/ai/agent/confirm；审计 source=mcp|agent
534 lines
14 KiB
Go
534 lines
14 KiB
Go
package aitools
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"fmt"
|
||
"strings"
|
||
|
||
"opsxdb/server/internal/db"
|
||
)
|
||
|
||
const (
|
||
ToolListConnections = "opsxdb_list_connections"
|
||
ToolListDatabases = "opsxdb_list_databases"
|
||
ToolListTables = "opsxdb_list_tables"
|
||
ToolDescribeTable = "opsxdb_describe_table"
|
||
ToolGetSchemaContext = "opsxdb_get_schema_context"
|
||
ToolExecuteQuery = "opsxdb_execute_query"
|
||
ToolExecuteRedisCommand = "opsxdb_execute_redis_command"
|
||
|
||
schemaContextCap = 4000
|
||
)
|
||
|
||
// ToolNames P0 注册表（顺序稳定，便于 tools/list 断言）。
|
||
func ToolNames() []string {
|
||
return []string{
|
||
ToolListConnections,
|
||
ToolListDatabases,
|
||
ToolListTables,
|
||
ToolDescribeTable,
|
||
ToolGetSchemaContext,
|
||
ToolExecuteQuery,
|
||
ToolExecuteRedisCommand,
|
||
}
|
||
}
|
||
|
||
// Runtime 可被 MCP stdio 与将来侧栏 Agent 进程内共用。 // Logan.Li
|
||
type Runtime struct {
|
||
Policy Policy
|
||
// LoadConns / FindConn 可注入，便于单测；nil 则读磁盘。
|
||
LoadConns func() ([]SavedConnection, error)
|
||
FindConn func(id string) (*SavedConnection, error)
|
||
// ExecSQL 可注入；nil 则 db.ExecSQL。
|
||
ExecSQL func(ctx context.Context, c db.ConnReq) db.ExecResult
|
||
ListDB func(ctx context.Context, c db.ConnReq) ([]string, error)
|
||
ListTbl func(ctx context.Context, c db.ConnReq) ([]string, error)
|
||
ListCol func(ctx context.Context, c db.ConnReq) ([]db.ColumnInfo, error)
|
||
Audit func(AuditEntry)
|
||
// RejectCallerConfirm 为 true 时（Agent 路径），忽略调用方 confirm_danger，
|
||
// 写操作只返回 needConfirm，由用户在 UI 显式确认后再走 /api/ai/agent/confirm。
|
||
RejectCallerConfirm bool
|
||
// AllowConfirmedWrite 为 true 时，已校验的确认令牌可放行写（仅 confirm 端点设置）。
|
||
AllowConfirmedWrite bool
|
||
// Source 写入审计：mcp | agent。
|
||
Source string
|
||
}
|
||
|
||
func NewRuntime(policy Policy) *Runtime {
|
||
return &Runtime{Policy: policy}
|
||
}
|
||
|
||
func DefaultRuntime() (*Runtime, error) {
|
||
p, err := LoadPolicy()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return NewRuntime(p), nil
|
||
}
|
||
|
||
func (r *Runtime) audit(e AuditEntry) {
|
||
if e.Mode == "" {
|
||
e.Mode = r.Policy.Mode
|
||
if e.Mode == "" {
|
||
e.Mode = ModeReadOnly
|
||
}
|
||
}
|
||
if e.Source == "" {
|
||
e.Source = r.Source
|
||
}
|
||
if r.Audit != nil {
|
||
r.Audit(e)
|
||
return
|
||
}
|
||
AppendAudit(e)
|
||
}
|
||
|
||
func (r *Runtime) loadConns() ([]SavedConnection, error) {
|
||
if r.LoadConns != nil {
|
||
return r.LoadConns()
|
||
}
|
||
return LoadConnections()
|
||
}
|
||
|
||
func (r *Runtime) findConn(id string) (*SavedConnection, error) {
|
||
if r.FindConn != nil {
|
||
return r.FindConn(id)
|
||
}
|
||
return FindConnection(id)
|
||
}
|
||
|
||
func (r *Runtime) resolveAllowed(connectionID string) (*SavedConnection, error) {
|
||
if err := r.Policy.AllowConnection(connectionID); err != nil {
|
||
return nil, err
|
||
}
|
||
c, err := r.findConn(connectionID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return c, nil
|
||
}
|
||
|
||
func (r *Runtime) execSQL(ctx context.Context, c db.ConnReq) db.ExecResult {
|
||
if r.ExecSQL != nil {
|
||
return r.ExecSQL(ctx, c)
|
||
}
|
||
return db.ExecSQL(ctx, c)
|
||
}
|
||
|
||
func (r *Runtime) listDB(ctx context.Context, c db.ConnReq) ([]string, error) {
|
||
if r.ListDB != nil {
|
||
return r.ListDB(ctx, c)
|
||
}
|
||
return db.ListDatabases(ctx, c)
|
||
}
|
||
|
||
func (r *Runtime) listTbl(ctx context.Context, c db.ConnReq) ([]string, error) {
|
||
if r.ListTbl != nil {
|
||
return r.ListTbl(ctx, c)
|
||
}
|
||
return db.ListTables(ctx, c)
|
||
}
|
||
|
||
func (r *Runtime) listCol(ctx context.Context, c db.ConnReq) ([]db.ColumnInfo, error) {
|
||
if r.ListCol != nil {
|
||
return r.ListCol(ctx, c)
|
||
}
|
||
return db.ListColumns(ctx, c)
|
||
}
|
||
|
||
// Call 统一入口：MCP 与 Agent 都走这里。args 为 JSON object。
|
||
func (r *Runtime) Call(ctx context.Context, tool string, args map[string]any) (any, error) {
|
||
switch tool {
|
||
case ToolListConnections:
|
||
return r.listConnections(ctx)
|
||
case ToolListDatabases:
|
||
return r.listDatabases(ctx, args)
|
||
case ToolListTables:
|
||
return r.listTables(ctx, args)
|
||
case ToolDescribeTable:
|
||
return r.describeTable(ctx, args)
|
||
case ToolGetSchemaContext:
|
||
return r.getSchemaContext(ctx, args)
|
||
case ToolExecuteQuery:
|
||
return r.executeQuery(ctx, args)
|
||
case ToolExecuteRedisCommand:
|
||
return r.executeRedis(ctx, args)
|
||
default:
|
||
return nil, fmt.Errorf("unknown tool: %s", tool)
|
||
}
|
||
}
|
||
|
||
func (r *Runtime) listConnections(ctx context.Context) (any, error) {
|
||
_ = ctx
|
||
list, err := r.loadConns()
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolListConnections, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
out := make([]ConnectionSummary, 0, len(list))
|
||
for _, c := range list {
|
||
if err := r.Policy.AllowConnection(c.ID); err != nil {
|
||
continue
|
||
}
|
||
out = append(out, c.Summary())
|
||
}
|
||
r.audit(AuditEntry{Tool: ToolListConnections, OK: true, Summary: fmt.Sprintf("%d connections", len(out))})
|
||
return map[string]any{"connections": out}, nil
|
||
}
|
||
|
||
func argString(args map[string]any, key string) string {
|
||
if args == nil {
|
||
return ""
|
||
}
|
||
v, ok := args[key]
|
||
if !ok || v == nil {
|
||
return ""
|
||
}
|
||
switch t := v.(type) {
|
||
case string:
|
||
return strings.TrimSpace(t)
|
||
default:
|
||
return strings.TrimSpace(fmt.Sprint(t))
|
||
}
|
||
}
|
||
|
||
func argInt(args map[string]any, key string, def int) int {
|
||
if args == nil {
|
||
return def
|
||
}
|
||
v, ok := args[key]
|
||
if !ok || v == nil {
|
||
return def
|
||
}
|
||
switch t := v.(type) {
|
||
case float64:
|
||
return int(t)
|
||
case int:
|
||
return t
|
||
case int64:
|
||
return int(t)
|
||
case json.Number:
|
||
n, _ := t.Int64()
|
||
return int(n)
|
||
case string:
|
||
var n int
|
||
fmt.Sscanf(t, "%d", &n)
|
||
if n == 0 {
|
||
return def
|
||
}
|
||
return n
|
||
default:
|
||
return def
|
||
}
|
||
}
|
||
|
||
func argBool(args map[string]any, key string) bool {
|
||
if args == nil {
|
||
return false
|
||
}
|
||
v, ok := args[key]
|
||
if !ok || v == nil {
|
||
return false
|
||
}
|
||
switch t := v.(type) {
|
||
case bool:
|
||
return t
|
||
case string:
|
||
return strings.EqualFold(t, "true") || t == "1"
|
||
case float64:
|
||
return t != 0
|
||
case int:
|
||
return t != 0
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
|
||
func (r *Runtime) listDatabases(ctx context.Context, args map[string]any) (any, error) {
|
||
id := argString(args, "connection_id")
|
||
c, err := r.resolveAllowed(id)
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolListDatabases, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
req := c.ToConnReq()
|
||
dbs, err := r.listDB(ctx, req)
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolListDatabases, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
r.audit(AuditEntry{Tool: ToolListDatabases, ConnectionID: id, OK: true, Summary: fmt.Sprintf("%d databases", len(dbs))})
|
||
return map[string]any{"databases": dbs}, nil
|
||
}
|
||
|
||
func (r *Runtime) listTables(ctx context.Context, args map[string]any) (any, error) {
|
||
id := argString(args, "connection_id")
|
||
database := argString(args, "database")
|
||
c, err := r.resolveAllowed(id)
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolListTables, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
req := c.ToConnReq()
|
||
if database != "" {
|
||
req.Database = database
|
||
}
|
||
tables, err := r.listTbl(ctx, req)
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolListTables, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
r.audit(AuditEntry{Tool: ToolListTables, ConnectionID: id, OK: true, Summary: fmt.Sprintf("%d tables", len(tables))})
|
||
return map[string]any{"tables": tables, "database": req.Database}, nil
|
||
}
|
||
|
||
func (r *Runtime) describeTable(ctx context.Context, args map[string]any) (any, error) {
|
||
id := argString(args, "connection_id")
|
||
database := argString(args, "database")
|
||
table := argString(args, "table")
|
||
if table == "" {
|
||
err := fmt.Errorf("table required")
|
||
r.audit(AuditEntry{Tool: ToolDescribeTable, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
c, err := r.resolveAllowed(id)
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolDescribeTable, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
req := c.ToConnReq()
|
||
if database != "" {
|
||
req.Database = database
|
||
}
|
||
req.Table = table
|
||
cols, err := r.listCol(ctx, req)
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolDescribeTable, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
r.audit(AuditEntry{Tool: ToolDescribeTable, ConnectionID: id, OK: true, Summary: table + fmt.Sprintf(" (%d cols)", len(cols))})
|
||
return map[string]any{"table": table, "database": req.Database, "columns": cols}, nil
|
||
}
|
||
|
||
func (r *Runtime) getSchemaContext(ctx context.Context, args map[string]any) (any, error) {
|
||
id := argString(args, "connection_id")
|
||
database := argString(args, "database")
|
||
maxTables := argInt(args, "max_tables", 40)
|
||
if maxTables <= 0 {
|
||
maxTables = 40
|
||
}
|
||
if maxTables > 200 {
|
||
maxTables = 200
|
||
}
|
||
c, err := r.resolveAllowed(id)
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolGetSchemaContext, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
req := c.ToConnReq()
|
||
if database != "" {
|
||
req.Database = database
|
||
}
|
||
tables, err := r.listTbl(ctx, req)
|
||
if err != nil {
|
||
r.audit(AuditEntry{Tool: ToolGetSchemaContext, ConnectionID: id, OK: false, Error: err.Error()})
|
||
return nil, err
|
||
}
|
||
var b strings.Builder
|
||
n := 0
|
||
for _, t := range tables {
|
||
if n >= maxTables {
|
||
break
|
||
}
|
||
colReq := req
|
||
colReq.Table = t
|
||
cols, err := r.listCol(ctx, colReq)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
parts := make([]string, 0, len(cols))
|
||
for _, col := range cols {
|
||
parts = append(parts, strings.TrimSpace(col.Name+" "+col.Type))
|
||
}
|
||
dbName := req.Database
|
||
if dbName == "" {
|
||
dbName = "default"
|
||
}
|
||
line := fmt.Sprintf("%s.%s(%s)\n", dbName, t, strings.Join(parts, ", "))
|
||
if b.Len()+len(line) > schemaContextCap {
|
||
break
|
||
}
|
||
b.WriteString(line)
|
||
n++
|
||
}
|
||
text := b.String()
|
||
if len(text) > schemaContextCap {
|
||
text = text[:schemaContextCap]
|
||
}
|
||
r.audit(AuditEntry{Tool: ToolGetSchemaContext, ConnectionID: id, OK: true, Summary: fmt.Sprintf("%d tables, %d chars", n, len(text))})
|
||
return map[string]any{"schema": text, "truncated": len(tables) > n || len(text) >= schemaContextCap}, nil
|
||
}
|
||
|
||
func (r *Runtime) executeQuery(ctx context.Context, args map[string]any) (any, error) {
|
||
id := argString(args, "connection_id")
|
||
sqlText := argString(args, "sql")
|
||
database := argString(args, "database")
|
||
limit := r.Policy.ClampLimit(argInt(args, "limit", 0))
|
||
|
||
fail := func(msg string) (any, error) {
|
||
r.audit(AuditEntry{Tool: ToolExecuteQuery, ConnectionID: id, OK: false, Error: msg, Summary: truncate(sqlText, 80)})
|
||
return nil, fmt.Errorf("%s", msg)
|
||
}
|
||
|
||
if sqlText == "" {
|
||
return fail("sql required")
|
||
}
|
||
callerConfirm := argBool(args, "confirm_danger") || argBool(args, "confirmDanger")
|
||
if r.RejectCallerConfirm {
|
||
callerConfirm = false
|
||
}
|
||
confirmed := r.AllowConfirmedWrite || callerConfirm
|
||
|
||
readOK, readReason := IsReadOnlySQL(sqlText)
|
||
if !readOK {
|
||
if !r.Policy.WritesEnabled() {
|
||
return fail(readReason)
|
||
}
|
||
// P1：可写已开启，但无确认令牌则返回 needConfirm（永不静默写）。
|
||
if !confirmed {
|
||
r.audit(AuditEntry{Tool: ToolExecuteQuery, ConnectionID: id, OK: false, Error: "needConfirm", Summary: truncate(sqlText, 80)})
|
||
return map[string]any{
|
||
"needConfirm": true,
|
||
"sql": sqlText,
|
||
"message": "write requires explicit user confirmation",
|
||
"reason": readReason,
|
||
}, nil
|
||
}
|
||
}
|
||
|
||
c, err := r.resolveAllowed(id)
|
||
if err != nil {
|
||
return fail(err.Error())
|
||
}
|
||
if strings.EqualFold(c.Driver, "redis") {
|
||
return fail("redis connections do not accept SQL; use opsxdb_execute_redis_command")
|
||
}
|
||
|
||
req := c.ToConnReq()
|
||
if database != "" {
|
||
req.Database = database
|
||
}
|
||
req.SQL = sqlText
|
||
req.Limit = limit
|
||
req.ConfirmDanger = confirmed && !readOK
|
||
|
||
res := r.execSQL(ctx, req)
|
||
if res.NeedConfirm {
|
||
// 只读层已拦；若底层仍标 needConfirm，视为拒绝
|
||
return fail("dangerous sql blocked (needConfirm): " + res.Error)
|
||
}
|
||
if res.Error != "" {
|
||
return fail(res.Error)
|
||
}
|
||
r.audit(AuditEntry{
|
||
Tool: ToolExecuteQuery, ConnectionID: id, OK: true,
|
||
Summary: fmt.Sprintf("rows=%d truncated=%v sql=%s", res.RowCount, res.Truncated, truncate(sqlText, 60)),
|
||
})
|
||
return map[string]any{
|
||
"columns": res.Columns,
|
||
"rows": res.Rows,
|
||
"rowCount": res.RowCount,
|
||
"timeMs": res.TimeMs,
|
||
"truncated": res.Truncated,
|
||
"limit": limit,
|
||
}, nil
|
||
}
|
||
|
||
func (r *Runtime) executeRedis(ctx context.Context, args map[string]any) (any, error) {
|
||
id := argString(args, "connection_id")
|
||
command := argString(args, "command")
|
||
database := argString(args, "database")
|
||
|
||
fail := func(msg string) (any, error) {
|
||
r.audit(AuditEntry{Tool: ToolExecuteRedisCommand, ConnectionID: id, OK: false, Error: msg, Summary: truncate(command, 80)})
|
||
return nil, fmt.Errorf("%s", msg)
|
||
}
|
||
if command == "" {
|
||
return fail("command required")
|
||
}
|
||
callerConfirm := argBool(args, "confirm_danger") || argBool(args, "confirmDanger")
|
||
if r.RejectCallerConfirm {
|
||
callerConfirm = false
|
||
}
|
||
confirmed := r.AllowConfirmedWrite || callerConfirm
|
||
|
||
cmd, rerr := CheckRedisReadOnly(command)
|
||
if rerr != nil {
|
||
if redisDeniedAlways[cmd] || !r.Policy.WritesEnabled() {
|
||
return fail(rerr.Error())
|
||
}
|
||
if !confirmed {
|
||
r.audit(AuditEntry{Tool: ToolExecuteRedisCommand, ConnectionID: id, OK: false, Error: "needConfirm", Summary: truncate(command, 80)})
|
||
return map[string]any{
|
||
"needConfirm": true,
|
||
"command": command,
|
||
"message": "redis write requires explicit user confirmation",
|
||
"reason": rerr.Error(),
|
||
}, nil
|
||
}
|
||
}
|
||
|
||
c, err := r.resolveAllowed(id)
|
||
if err != nil {
|
||
return fail(err.Error())
|
||
}
|
||
if !strings.EqualFold(c.Driver, "redis") {
|
||
return fail("connection is not redis")
|
||
}
|
||
|
||
req := c.ToConnReq()
|
||
if database != "" {
|
||
req.Database = database
|
||
}
|
||
req.SQL = command
|
||
req.ConfirmDanger = confirmed && rerr != nil
|
||
|
||
res := r.execSQL(ctx, req)
|
||
if res.NeedConfirm {
|
||
return fail("dangerous redis command blocked: " + res.Error)
|
||
}
|
||
if res.Error != "" {
|
||
return fail(res.Error)
|
||
}
|
||
r.audit(AuditEntry{
|
||
Tool: ToolExecuteRedisCommand, ConnectionID: id, OK: true,
|
||
Summary: fmt.Sprintf("cmd=%s rows=%d", cmd, res.RowCount),
|
||
})
|
||
return map[string]any{
|
||
"columns": res.Columns,
|
||
"rows": res.Rows,
|
||
"rowCount": res.RowCount,
|
||
"timeMs": res.TimeMs,
|
||
"truncated": res.Truncated,
|
||
}, nil
|
||
}
|
||
|
||
func truncate(s string, n int) string {
|
||
s = strings.Join(strings.Fields(s), " ")
|
||
if len(s) <= n {
|
||
return s
|
||
}
|
||
return s[:n] + "…"
|
||
}
|
||
|
||
// ResultJSON 将 Call 结果序列化为 MCP 文本内容。
|
||
func ResultJSON(v any) string {
|
||
b, err := json.MarshalIndent(v, "", " ")
|
||
if err != nil {
|
||
return fmt.Sprintf("%v", v)
|
||
}
|
||
return string(b)
|
||
}
|