opsxdb/server/internal/aitools/runtime.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)
}

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 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 (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")
	}
	// P0 纯只读：写/DDL 一律拒绝；不做 write mode / confirm token。
	if ok, reason := IsReadOnlySQL(sqlText); !ok {
		return fail(reason)
	}

	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 = false // 永不带确认放行

	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")
	}
	cmd, err := CheckRedisReadOnly(command)
	if err != nil {
		return fail(err.Error())
	}
	_ = cmd

	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 = false

	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)
}