opsxdb/server/internal/ai/session_test.go

package ai

import (
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

func TestSessionPersistAndLoadTrajectory(t *testing.T) {
	dir := t.TempDir()
	t.Setenv("HOME", dir)
	id := NewSessionID()
	traj := []Message{
		{Role: "system", Content: "sys"},
		{Role: "user", Content: "list tables"},
		{Role: "assistant", Content: "", ToolCalls: []ToolCall{{ID: "c1", Type: "function", Function: ToolCallFunc{Name: "opsxdb_list_tables", Arguments: "{}"}}}},
		{Role: "tool", ToolCallID: "c1", Name: "opsxdb_list_tables", Content: `{"tables":["users"]}`},
	}
	res := AgentRes{RunID: "run_x", Outcome: OutcomeIncomplete, IncompleteReason: ReasonMaxSteps, Resumable: true, Rounds: 8, StepsBudget: 8, SessionID: id}
	s, err := PersistAgentRun(id, "list tables", traj, res)
	if err != nil {
		t.Fatal(err)
	}
	if s.Status != "incomplete" {
		t.Fatalf("status=%s", s.Status)
	}
	loaded, err := LoadSession(id)
	if err != nil {
		t.Fatal(err)
	}
	if len(loaded.Messages) != 4 {
		t.Fatalf("msgs=%d", len(loaded.Messages))
	}
	if loaded.Messages[3].ToolCallID != "c1" {
		t.Fatal("tool trajectory lost")
	}
	list, err := ListSessions(10)
	if err != nil || len(list) < 1 {
		t.Fatalf("list: %v %#v", err, list)
	}
	path := filepath.Join(dir, ".opsxdb", "ai-sessions", id+".json")
	if _, err := os.Stat(path); err != nil {
		t.Fatal(err)
	}
}

func TestRunAgentContinueInjectsPrior(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadOnly, AllowedConnectionIDs: []string{"c1"}, DefaultLimit: 100, MaxLimit: 500})
	rt.LoadConns = func() ([]aitools.SavedConnection, error) {
		return []aitools.SavedConnection{{ID: "c1", Name: "d", Driver: "mysql"}}, nil
	}
	rt.Audit = func(aitools.AuditEntry) {}
	prior := []Message{
		{Role: "system", Content: "sys"},
		{Role: "user", Content: "prev"},
		{Role: "assistant", Content: "partial", ToolCalls: []ToolCall{{ID: "c1", Type: "function", Function: ToolCallFunc{Name: aitools.ToolListConnections, Arguments: "{}"}}}},
		{Role: "tool", ToolCallID: "c1", Name: aitools.ToolListConnections, Content: `{"connections":[]}`},
	}
	llm := &scriptedLLM{turns: []ToolsTurn{{Content: "continued answer"}}}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{
		Prompt: "加预算继续", Continue: true, Prior: prior, MaxRounds: 4, SessionID: "sess_test",
	})
	if err != nil {
		t.Fatal(err)
	}
	if res.Outcome != OutcomeCompleted {
		t.Fatalf("outcome=%s", res.Outcome)
	}
	if len(llm.seen) < 1 {
		t.Fatal("no llm call")
	}
	sawTool := false
	for _, m := range llm.seen[0] {
		if m.Role == "tool" && m.ToolCallID == "c1" {
			sawTool = true
		}
	}
	if !sawTool {
		t.Fatalf("prior tool not injected: %+v", llm.seen[0])
	}
	if res.SessionID == "" || len(res.Trajectory) < 4 {
		t.Fatalf("session/traj: sid=%s traj=%d", res.SessionID, len(res.Trajectory))
	}
}

func TestRunAgentAwaitingConfirmOutcome(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{
		Mode:                 aitools.ModeReadWrite,
		AllowedConnectionIDs: []string{"c1"},
		DefaultLimit:         100,
		MaxLimit:             500,
	})
	rt.FindConn = func(id string) (*aitools.SavedConnection, error) {
		return &aitools.SavedConnection{ID: id, Name: "d", Driver: "mysql", Host: "127.0.0.1", Port: 3306, Database: "demo"}, nil
	}
	rt.Audit = func(aitools.AuditEntry) {}
	execCalls := 0
	rt.ExecSQL = func(_ context.Context, _ db.ConnReq) db.ExecResult {
		execCalls++
		return db.ExecResult{RowCount: 1}
	}
	args := `{"connection_id":"c1","database":"demo","sql":"INSERT INTO users(name) VALUES('x')"}`
	llm := &scriptedLLM{turns: []ToolsTurn{{
		ToolCalls: []ToolCall{{ID: "w1", Type: "function", Function: ToolCallFunc{Name: aitools.ToolExecuteQuery, Arguments: args}}},
	}}}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "insert one", MaxRounds: 4})
	if err != nil {
		t.Fatal(err)
	}
	if res.Outcome != OutcomeAwaitingConfirm {
		t.Fatalf("outcome=%s want awaiting_confirm", res.Outcome)
	}
	if !res.Resumable {
		t.Fatal("expected resumable")
	}
	if execCalls != 0 {
		t.Fatal("must not execute write without confirm")
	}
	found := false
	for _, s := range res.Steps {
		if s.NeedConfirm {
			found = true
		}
	}
	if !found {
		t.Fatalf("steps missing needConfirm: %+v", res.Steps)
	}
}

func TestCheckpointHint(t *testing.T) {
	if CheckpointHint(4) == "" || CheckpointHint(6) == "" {
		t.Fatal("expected hints")
	}
	if CheckpointHint(1) != "" {
		t.Fatal("no hint round 1")
	}
}

func TestRunAgentInjectsCheckpoint(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadOnly, AllowedConnectionIDs: []string{"c1"}, DefaultLimit: 100, MaxLimit: 500})
	rt.LoadConns = func() ([]aitools.SavedConnection, error) {
		return []aitools.SavedConnection{{ID: "c1", Name: "d", Driver: "mysql"}}, nil
	}
	rt.Audit = func(aitools.AuditEntry) {}
	var turns []ToolsTurn
	for i := 0; i < 4; i++ {
		turns = append(turns, ToolsTurn{ToolCalls: []ToolCall{{
			ID: "c" + string(rune('a'+i)), Type: "function",
			Function: ToolCallFunc{Name: aitools.ToolListConnections, Arguments: `{"n":` + string(rune('0'+i)) + `}`},
		}}})
	}
	turns = append(turns, ToolsTurn{Content: "ok"})
	llm := &scriptedLLM{turns: turns}
	_, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "x", MaxRounds: 8})
	if err != nil {
		t.Fatal(err)
	}
	if len(llm.seen) < 4 {
		t.Fatalf("seen=%d", len(llm.seen))
	}
	found := false
	for _, m := range llm.seen[3] {
		if m.Role == "system" && strings.Contains(m.Content, "检查点") {
			found = true
		}
	}
	if !found {
		t.Fatalf("checkpoint not in round4 msgs")
	}
}