opsxdb/server/internal/ai/harness_test.go

package ai

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

	"opsxdb/server/internal/aitools"
)

func TestTrimMessagesKeepsSystemAndDropsOrphanTool(t *testing.T) {
	h := DefaultHarness()
	h.MaxContextTokens = 40
	msgs := []Message{
		{Role: "system", Content: "sys"},
		{Role: "user", Content: strings.Repeat("u", 400)},
		{Role: "assistant", Content: "a", ToolCalls: []ToolCall{{ID: "c1", Type: "function", Function: ToolCallFunc{Name: "opsxdb_list_tables", Arguments: "{}"}}}},
		{Role: "tool", ToolCallID: "c1", Name: "opsxdb_list_tables", Content: `{"tables":["users"]}`},
		{Role: "user", Content: "continue please"},
	}
	out := TrimMessages(msgs, h)
	hasSys := false
	for _, m := range out {
		if m.Role == "system" {
			hasSys = true
		}
	}
	if !hasSys {
		t.Fatal("system dropped")
	}
	orphan := TrimMessages([]Message{
		{Role: "system", Content: "s"},
		{Role: "tool", ToolCallID: "x", Content: "nope"},
		{Role: "user", Content: "hi"},
	}, h)
	for _, m := range orphan {
		if m.Role == "tool" {
			t.Fatalf("expected orphan tool dropped: %+v", orphan)
		}
	}
}

func TestDetectRepeatedToolCall(t *testing.T) {
	fp := ToolCallFingerprint("opsxdb_list_tables", `{"connection_id":"c1"}`)
	hist := []string{fp, fp, fp}
	if !DetectRepeatedToolCall(hist, fp, 3) {
		t.Fatal("want reject at 3")
	}
	if DetectRepeatedToolCall(hist[:2], fp, 3) {
		t.Fatal("should allow under limit")
	}
}

func TestRedactSensitive(t *testing.T) {
	out := RedactSensitive("Bearer tokensecret password=secret mysql://user:pass@host/db")
	if strings.Contains(out, "tokensecret") || strings.Contains(out, "password=secret") || strings.Contains(out, ":pass@") {
		t.Fatalf("not redacted: %s", out)
	}
	if !strings.Contains(RedactSensitive("sk-abcdefghijklmnop"), "[REDACTED]") {
		t.Fatal("sk key")
	}
}

func TestRunAgentMaxRoundsIncomplete(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
	}
	sawRunID := false
	rt.Audit = func(e aitools.AuditEntry) {
		if e.RunID != "" {
			sawRunID = true
		}
	}
	var turns []ToolsTurn
	for i := 0; i < 3; i++ {
		turns = append(turns, ToolsTurn{ToolCalls: []ToolCall{{
			ID: "c", Type: "function",
			Function: ToolCallFunc{Name: aitools.ToolListConnections, Arguments: "{}"},
		}}})
	}
	llm := &scriptedLLM{turns: turns}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "loop", MaxRounds: 2})
	if err != nil {
		t.Fatal(err)
	}
	if res.Outcome != OutcomeIncomplete {
		t.Fatalf("outcome=%s want incomplete", res.Outcome)
	}
	if res.IncompleteReason != ReasonMaxSteps {
		t.Fatalf("reason=%s", res.IncompleteReason)
	}
	if !res.Resumable || res.RunID == "" {
		t.Fatalf("resumable/runId: %+v", res)
	}
	if !sawRunID {
		t.Fatal("expected audit run_id")
	}
}

func TestRunAgentCompletedHasOutcome(t *testing.T) {
	rt := aitools.NewRuntime(aitools.Policy{Mode: aitools.ModeReadOnly, AllowedConnectionIDs: []string{"c1"}, DefaultLimit: 100, MaxLimit: 500})
	rt.Audit = func(aitools.AuditEntry) {}
	llm := &scriptedLLM{turns: []ToolsTurn{{Content: "hello"}}}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "hi"})
	if err != nil {
		t.Fatal(err)
	}
	if res.Outcome != OutcomeCompleted {
		t.Fatalf("outcome=%s", res.Outcome)
	}
	if res.RunID == "" || res.StepsBudget != defaultAgentMaxRounds {
		t.Fatalf("%+v", res)
	}
}

func TestRunAgentRepeatCallRejected(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) {}
	args := `{"connection_id":"c1"}`
	var turns []ToolsTurn
	// 同一轮多个相同 tool_call，以及跨轮：SameCallLimit=3 指 history 已有 3 次则拒第 4 次
	// 每轮一个相同调用，第 4 轮应被拒但仍继续到最终答案
	for i := 0; i < 4; i++ {
		turns = append(turns, ToolsTurn{ToolCalls: []ToolCall{{
			ID: "c" + string(rune('1'+i)), Type: "function",
			Function: ToolCallFunc{Name: aitools.ToolListConnections, Arguments: args},
		}}})
	}
	turns = append(turns, ToolsTurn{Content: "done"})
	llm := &scriptedLLM{turns: turns}
	res, err := RunAgent(context.Background(), llm, rt, AgentReq{Prompt: "x", MaxRounds: 8})
	if err != nil {
		t.Fatal(err)
	}
	rejected := 0
	for _, s := range res.Steps {
		if strings.Contains(s.Summary, "repeated tool call") {
			rejected++
		}
	}
	if rejected < 1 {
		t.Fatalf("expected at least one repeated rejection, steps=%+v", res.Steps)
	}
	_ = json.Marshal
}