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 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 := parseToolsTurn(raw, 200)
	if err != nil || len(turn.ToolCalls) != 1 || turn.ToolCalls[0].Function.Name != "opsxdb_list_tables" {
		t.Fatalf("parse: %+v %v", turn, err)
	}
}