opsxdb/server/internal/ai/agent_test.go

package ai

import (
	"context"
	"encoding/json"
	"strings"
	"testing"

	"opsxdb/server/internal/aitools"
	"opsxdb/server/internal/db"
)

// scriptedLLM 按调用次序返回预设 ToolsTurn，用于证明 Agent tool loop 不依赖真 LLM。
type scriptedLLM struct {
	turns []ToolsTurn
	i     int
	seen  [][]Message
}

func (s *scriptedLLM) ChatWithTools(_ context.Context, _ string, messages []Message, tools []aitools.OpenAITool) (ToolsTurn, error) {
	cp := append([]Message(nil), messages...)
	s.seen = append(s.seen, cp)
	if len(tools) < 7 {
		return ToolsTurn{}, nil // 测试里仍继续；下面断言
	}
	if s.i >= len(s.turns) {
		return ToolsTurn{Content: "fallback"}, nil
	}
	t := s.turns[s.i]
	s.i++
	return t, nil
}

func TestOpenAIToolsMatchRuntime(t *testing.T) {
	tools := aitools.OpenAITools()
	names := aitools.ToolNames()
	if len(tools) != len(names) {
		t.Fatalf("tools %d vs names %d", len(tools), len(names))
	}
	for i, n := range names {
		if tools[i].Function.Name != n {
			t.Fatalf("tool[%d]=%s want %s", i, tools[i].Function.Name, n)
		}
		if tools[i].Type != "function" {
			t.Fatalf("type want function")
		}
	}
}

func TestRunAgentToolLoop(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{
		Mode:                 aitools.ModeReadOnly,
		AllowedConnectionIDs: []string{"c1"},
		DefaultLimit:         100,
		MaxLimit:             500,
	})
	rt.FindConn = func(id string) (*aitools.SavedConnection, error) {
		return &aitools.SavedConnection{ID: id, Name: "demo", Driver: "mysql", Host: "127.0.0.1", Port: 3306, Database: "demo"}, nil
	}
	rt.ListTbl = func(_ context.Context, c db.ConnReq) ([]string, error) {
		if c.Database != "demo" && c.Database != "" {
			return nil, nil
		}
		return []string{"users", "orders"}, nil
	}
	rt.ExecSQL = func(_ context.Context, c db.ConnReq) db.ExecResult {
		if !strings.Contains(strings.ToUpper(c.SQL), "SELECT") {
			return db.ExecResult{Error: "unexpected"}
		}
		return db.ExecResult{
			Columns:  []string{"id", "name"},
			Rows:     []map[string]interface{}{{"id": 1, "name": "alice"}, {"id": 2, "name": "bob"}},
			RowCount: 2,
			TimeMs:   1,
		}
	}
	rt.Audit = func(aitools.AuditEntry) {}

	argsList, _ := json.Marshal(map[string]any{"connection_id": "c1", "database": "demo"})
	argsQuery, _ := json.Marshal(map[string]any{
		"connection_id": "c1",
		"database":      "demo",
		"sql":           "SELECT id, name FROM users LIMIT 5",
	})

	llm := &scriptedLLM{turns: []ToolsTurn{
		{
			ToolCalls: []ToolCall{{
				ID: "call_1", Type: "function",
				Function: ToolCallFunc{Name: aitools.ToolListTables, Arguments: string(argsList)},
			}},
		},
		{
			ToolCalls: []ToolCall{{
				ID: "call_2", Type: "function",
				Function: ToolCallFunc{Name: aitools.ToolExecuteQuery, Arguments: string(argsQuery)},
			}},
		},
		{Content: "demo 库有 users、orders。users 样例两行。\n\n```sql\nSELECT id, name FROM users LIMIT 5;\n```"},
	}}

	res, err := RunAgent(context.Background(), llm, rt, AgentReq{
		Prompt:       "列出 demo 库有哪些表，并查 users 几行",
		ConnectionID: "c1",
		Driver:       "mysql",
		Database:     "demo",
	})
	if err != nil {
		t.Fatalf("RunAgent: %v", err)
	}
	if len(res.Steps) != 2 {
		t.Fatalf("want 2 steps, got %d: %+v", len(res.Steps), res.Steps)
	}
	if !res.Steps[0].OK || res.Steps[0].Tool != aitools.ToolListTables {
		t.Fatalf("step0: %+v", res.Steps[0])
	}
	if !strings.Contains(res.Steps[0].Summary, "2 tables") {
		t.Fatalf("list_tables summary: %s", res.Steps[0].Summary)
	}
	if !res.Steps[1].OK || res.Steps[1].Tool != aitools.ToolExecuteQuery {
		t.Fatalf("step1: %+v", res.Steps[1])
	}
	if !strings.Contains(res.Steps[1].Summary, "rows=2") {
		t.Fatalf("execute_query summary: %s", res.Steps[1].Summary)
	}
	if !strings.Contains(res.Content, "users") {
		t.Fatalf("final content: %s", res.Content)
	}
	if res.SQL == "" {
		t.Fatal("expected extracted SQL")
	}
	if res.Rounds != 3 {
		t.Fatalf("rounds=%d", res.Rounds)
	}
	// 第二轮 messages 应已含 tool 角色回灌
	if len(llm.seen) < 2 {
		t.Fatal("expected multiple LLM calls")
	}
	foundTool := false
	for _, m := range llm.seen[1] {
		if m.Role == "tool" && m.ToolCallID == "call_1" {
			foundTool = true
			if !strings.Contains(m.Content, "users") {
				t.Fatalf("tool result missing tables: %s", m.Content)
			}
		}
	}
	if !foundTool {
		t.Fatal("second LLM call should include tool result for call_1")
	}
}

func TestRunAgentRejectsWriteViaRuntime(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{
		Mode:                 aitools.ModeReadOnly,
		AllowedConnectionIDs: []string{"c1"},
		DefaultLimit:         100,
		MaxLimit:             500,
	})
	rt.FindConn = func(id string) (*aitools.SavedConnection, error) {
		return &aitools.SavedConnection{ID: id, Name: "demo", Driver: "mysql"}, nil
	}
	rt.Audit = func(aitools.AuditEntry) {}

	args, _ := json.Marshal(map[string]any{"connection_id": "c1", "sql": "DROP TABLE users"})
	llm := &scriptedLLM{turns: []ToolsTurn{
		{ToolCalls: []ToolCall{{
			ID: "x", Type: "function",
			Function: ToolCallFunc{Name: aitools.ToolExecuteQuery, Arguments: string(args)},
		}}},
		{Content: "无法删除表：只读策略拒绝。"},
	}}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "删掉 users"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Steps[0].OK {
		t.Fatal("DROP should fail")
	}
	errMsg := strings.ToLower(res.Steps[0].Error + " " + res.Steps[0].Summary)
	if !strings.Contains(errMsg, "read-only") && !strings.Contains(errMsg, "not a read-only") {
		t.Fatalf("expected read-only deny, got %+v", res.Steps[0])
	}
}

func TestAgentCannotLeaveCurrentMySQLConnectionOrDatabase(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadOnly, AllowedConnectionIDs: []string{"current", "other"}, DefaultLimit: 10, MaxLimit: 20})
	rt.FindConn = func(id string) (*aitools.SavedConnection, error) {
		return &aitools.SavedConnection{ID: id, Driver: "mysql", Database: "demo"}, nil
	}
	var queries int
	rt.ExecSQL = func(_ context.Context, req db.ConnReq) db.ExecResult {
		queries++
		if req.Database != "demo" {
			t.Fatalf("wrong database: %s", req.Database)
		}
		return db.ExecResult{Columns: []string{"id"}, Rows: []map[string]interface{}{{"id": 1}}, RowCount: 1}
	}
	rt.Audit = func(aitools.AuditEntry) {}
	call := func(id, conn, database string) ToolCall {
		args, _ := json.Marshal(map[string]any{"connection_id": conn, "database": database, "sql": "SELECT id FROM users LIMIT 1"})
		return ToolCall{ID: id, Type: "function", Function: ToolCallFunc{Name: aitools.ToolExecuteQuery, Arguments: string(args)}}
	}
	llm := &scriptedLLM{turns: []ToolsTurn{
		{ToolCalls: []ToolCall{call("other", "other", "demo"), call("db", "current", "secret"), call("ok", "current", "demo")}},
		{Content: "只查询了当前库。"},
	}}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "查询当前库", ConnectionID: "current", Driver: "mysql", Database: "demo"})
	if err != nil {
		t.Fatal(err)
	}
	if len(res.Steps) != 3 || res.Steps[0].OK || res.Steps[1].OK || !res.Steps[2].OK || queries != 1 {
		t.Fatalf("Agent escaped scope or rejected valid query: steps=%+v queries=%d", res.Steps, queries)
	}
}

func TestAgentRedisKeepOneKeyNeedsConfirmation(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadWrite, AllowedConnectionIDs: []string{"redis"}})
	rt.RejectCallerConfirm = true
	rt.FindConn = func(id string) (*aitools.SavedConnection, error) {
		return &aitools.SavedConnection{ID: id, Driver: "redis", Database: "0"}, nil
	}
	var executed []string
	rt.ExecSQL = func(_ context.Context, req db.ConnReq) db.ExecResult {
		executed = append(executed, req.SQL)
		return db.ExecResult{Columns: []string{"value"}, Rows: []map[string]interface{}{{"value": 3}}, RowCount: 1}
	}
	rt.Audit = func(aitools.AuditEntry) {}
	call := func(id, command string) ToolCall {
		args, _ := json.Marshal(map[string]any{"connection_id": "redis", "command": command, "confirm_danger": true})
		return ToolCall{ID: id, Type: "function", Function: ToolCallFunc{Name: aitools.ToolExecuteRedisCommand, Arguments: string(args)}}
	}
	llm := &scriptedLLM{turns: []ToolsTurn{
		{ToolCalls: []ToolCall{call("read", "EXISTS key1 key2 key3")}},
		{ToolCalls: []ToolCall{call("delete", "DEL key2 key3")}},
	}}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "key1 key2 key3 只保留 key1", ConnectionID: "redis", Driver: "redis", Database: "0"})
	if err != nil {
		t.Fatal(err)
	}
	if len(executed) != 1 || executed[0] != "EXISTS key1 key2 key3" || len(res.Steps) != 2 || !res.Steps[1].NeedConfirm {
		t.Fatalf("删除必须停在用户确认前: executed=%v steps=%+v", executed, res.Steps)
	}
	if err := aitools.VerifyAgentBind("redis", "0", "DEL key2 key3", res.Steps[1].ConfirmBind); err != nil {
		t.Fatal(err)
	}
}

func TestAgentContinueRefreshesWritePolicy(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadOnly})
	llm := &scriptedLLM{turns: []ToolsTurn{{Content: "当前只读"}}}
	_, err := RunAgent(context.Background(), llm, rt, AgentReq{
		Prompt: "继续", Continue: true,
		Prior: []Message{{Role: "system", Content: "旧策略：允许写入"}, {Role: "user", Content: "只保留 key1"}},
	})
	if err != nil {
		t.Fatal(err)
	}
	if len(llm.seen) != 1 || !strings.Contains(llm.seen[0][0].Content, "当前策略为只读") || strings.Contains(llm.seen[0][0].Content, "旧策略") {
		t.Fatalf("续聊必须使用当前权限: %+v", llm.seen)
	}
}

func TestAgentRedisReadThenWriteNeedsCurrentDBConfirmation(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadWrite, AllowedConnectionIDs: []string{"redis"}, DefaultLimit: 10, MaxLimit: 20})
	rt.FindConn = func(id string) (*aitools.SavedConnection, error) {
		return &aitools.SavedConnection{ID: id, Driver: "redis", Database: "0"}, nil
	}
	var commands []string
	rt.ExecSQL = func(_ context.Context, req db.ConnReq) db.ExecResult {
		if req.Database != "1" {
			t.Fatalf("Redis DB drifted: %s", req.Database)
		}
		commands = append(commands, req.SQL)
		return db.ExecResult{Columns: []string{"value"}, Rows: []map[string]interface{}{{"value": "old"}}, RowCount: 1}
	}
	rt.Audit = func(aitools.AuditEntry) {}
	call := func(id, cmd string) ToolCall {
		args, _ := json.Marshal(map[string]any{"connection_id": "redis", "command": cmd})
		return ToolCall{ID: id, Type: "function", Function: ToolCallFunc{Name: aitools.ToolExecuteRedisCommand, Arguments: string(args)}}
	}
	llm := &scriptedLLM{turns: []ToolsTurn{
		{ToolCalls: []ToolCall{call("read", "GET smoke:key")}},
		{ToolCalls: []ToolCall{call("write", "SET smoke:key new")}},
	}}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "读取然后修改键", ConnectionID: "redis", Driver: "redis", Database: "1"})
	if err != nil {
		t.Fatal(err)
	}
	if len(commands) != 1 || commands[0] != "GET smoke:key" || len(res.Steps) != 2 || !res.Steps[0].OK || !res.Steps[1].NeedConfirm {
		t.Fatalf("unexpected Redis execution/confirmation: commands=%v steps=%+v", commands, res.Steps)
	}
	if err := aitools.VerifyAgentBind("redis", "1", "SET smoke:key new", res.Steps[1].ConfirmBind); err != nil {
		t.Fatal(err)
	}
	if err := aitools.VerifyAgentBind("redis", "0", "SET smoke:key new", res.Steps[1].ConfirmBind); err == nil {
		t.Fatal("other DB must not reuse the confirmation")
	}
}

func TestParseToolsTurn(t *testing.T) {
	content := "hi"
	raw, _ := json.Marshal(map[string]any{
		"choices": []map[string]any{{
			"finish_reason": "tool_calls",
			"message": map[string]any{
				"role":    "assistant",
				"content": content,
				"tool_calls": []map[string]any{{
					"id":   "c1",
					"type": "function",
					"function": map[string]string{
						"name":      "opsxdb_list_tables",
						"arguments": `{"connection_id":"x"}`,
					},
				}},
			},
		}},
	})
	turn, err := parseToolsTurnChat(raw, 200)
	if err != nil || len(turn.ToolCalls) != 1 || turn.ToolCalls[0].Function.Name != "opsxdb_list_tables" {
		t.Fatalf("parse: %+v %v", turn, err)
	}
}