iasysadmin/internal/ctxmgr/manager_test.go

package ctxmgr

import (
	"context"
	"errors"
	"strings"
	"testing"

	"sysagent/internal/llm"
)

func TestManagerMessagesAndCalibration(t *testing.T) {
	m := New("system prompt", 1000)
	m.Add(llm.Message{Role: llm.RoleUser, Content: "hola"})
	msgs := m.Messages()
	if len(msgs) != 2 || msgs[0].Role != llm.RoleSystem || msgs[1].Content != "hola" {
		t.Fatalf("messages = %+v", msgs)
	}

	raw := m.RawEstimate()
	tok, calibrated := m.Usage()
	if calibrated || tok != raw {
		t.Errorf("before server report: tok=%d calibrated=%v", tok, calibrated)
	}

	// El servidor dice que la petición costó la mitad de lo estimado.
	m.RecordUsage(llm.Usage{PromptTokens: raw / 2}, raw)
	tok, calibrated = m.Usage()
	if !calibrated || tok != raw/2 {
		t.Errorf("after calibration: tok=%d calibrated=%v raw=%d", tok, calibrated, raw)
	}
	// Factores absurdos se acotan.
	m.RecordUsage(llm.Usage{PromptTokens: 1}, 1000000)
	if got := m.Calibrate(1000); got != 250 {
		t.Errorf("factor must be clamped to 0.25, got %d", got)
	}

	m.SetOverhead(100)
	if m.RawEstimate() != raw+100 {
		t.Error("overhead not included")
	}

	m.Clear()
	if len(m.History()) != 0 || len(m.Messages()) != 1 {
		t.Error("Clear must keep only the system prompt")
	}
}

func TestEstimateTokensIsConservative(t *testing.T) {
	if EstimateTokens("") != 0 {
		t.Error("empty must be 0")
	}
	// 12 bytes -> 4 tokens con 3 bytes/token (un tokenizer real daría ~3).
	if got := EstimateTokens("abcdefghijkl"); got != 4 {
		t.Errorf("got %d", got)
	}
}

func TestElideToolOutputs(t *testing.T) {
	m := New("s", 0)
	m.Add(
		llm.Message{Role: llm.RoleUser, Content: "tarea"},
		llm.Message{Role: llm.RoleTool, Content: "salida 1"},
		llm.Message{Role: llm.RoleUser, Content: ToolResultPrefix + " salida 2"},
		llm.Message{Role: llm.RoleTool, Content: "salida 3"},
	)
	if n := m.ElideToolOutputs(1); n != 2 {
		t.Fatalf("elided %d, want 2", n)
	}
	h := m.History()
	if h[0].Content != "tarea" || h[1].Content != elided || h[2].Content != elided || h[3].Content != "salida 3" {
		t.Errorf("history = %+v", h)
	}
	if n := m.ElideToolOutputs(1); n != 0 {
		t.Errorf("second pass elided %d", n)
	}
	m.Truncate(1)
	if m.Len() != 1 {
		t.Errorf("Len after Truncate = %d", m.Len())
	}
}

func tc(name string) []llm.ToolCall {
	return []llm.ToolCall{{ID: "x", Type: "function", Function: llm.FunctionCall{Name: name, Arguments: "{}"}}}
}

func TestSplitNeverSeparatesToolCallsFromResults(t *testing.T) {
	m := New("s", 0)
	big := strings.Repeat("x", 3000)
	m.Add(
		llm.Message{Role: llm.RoleUser, Content: "tarea"},                // 0 (no es corte: índice 0)
		llm.Message{Role: llm.RoleAssistant, ToolCalls: tc("exec")},      // 1 corte
		llm.Message{Role: llm.RoleTool, Content: big},                    // 2
		llm.Message{Role: llm.RoleAssistant, ToolCalls: tc("exec")},      // 3 corte
		llm.Message{Role: llm.RoleTool, Content: big},                    // 4
		llm.Message{Role: llm.RoleUser, Content: ToolResultPrefix + big}, // 5 (no es corte)
		llm.Message{Role: llm.RoleAssistant, Content: "final"},           // 6 corte
	)
	old, kept, ok := m.Split(2200)
	if !ok {
		t.Fatal("expected split")
	}
	if len(old)+len(kept) != 7 {
		t.Fatalf("lost messages: %d + %d", len(old), len(kept))
	}
	if kept[0].Role == llm.RoleTool || IsToolResult(kept[0]) {
		t.Errorf("kept starts with a tool result: %+v", kept[0])
	}
	// Con 2200 tokens cabe desde el índice 3 (~2000 tokens), no desde el 1.
	if len(kept) != 4 || kept[0].Role != llm.RoleAssistant {
		t.Errorf("kept = %d messages starting with %s", len(kept), kept[0].Role)
	}

	// Presupuesto mínimo: se conserva solo el último tramo seguro.
	_, kept, _ = m.Split(1)
	if len(kept) != 1 || kept[0].Content != "final" {
		t.Errorf("minimal kept = %+v", kept)
	}

	single := New("s", 0)
	single.Add(llm.Message{Role: llm.RoleUser, Content: "solo"})
	if _, _, ok := single.Split(10); ok {
		t.Error("nothing to split with a single message")
	}
}

// fakeClient registra peticiones y devuelve respuestas programadas.
type fakeClient struct {
	reqs    []llm.ChatRequest
	replies []string
	err     error
}

func (f *fakeClient) Chat(_ context.Context, req llm.ChatRequest) (llm.ChatResponse, error) {
	f.reqs = append(f.reqs, req)
	if f.err != nil {
		return llm.ChatResponse{}, f.err
	}
	r := f.replies[min(len(f.reqs)-1, len(f.replies)-1)]
	return llm.ChatResponse{Message: llm.Message{Role: llm.RoleAssistant, Content: r}, FinishReason: "stop"}, nil
}

func fullSummary(extra string) string {
	var b strings.Builder
	for _, s := range RequiredSections {
		b.WriteString("## " + s + "\n- none\n")
	}
	return b.String() + extra
}

func TestCompressorSingleAndChunked(t *testing.T) {
	fc := &fakeClient{replies: []string{fullSummary("- db01 10.10.10.15:5432 PostgreSQL 17")}}
	c := &Compressor{Client: fc, ChunkTokens: 1000}
	out, err := c.Compress(context.Background(), "", "2026 [tool] systemctl restart nginx → exit=0", "[user]\nhola\n")
	if err != nil || !strings.Contains(out, "10.10.10.15:5432") {
		t.Fatalf("out=%q err=%v", out, err)
	}
	if len(fc.reqs) != 1 {
		t.Fatalf("requests = %d", len(fc.reqs))
	}
	sys := fc.reqs[0].Messages[0].Content
	for _, want := range []string{"Preserve exactly", "IP addresses", "Do not invent", "## CREDENTIAL_REFERENCES"} {
		if !strings.Contains(sys, want) {
			t.Errorf("compression prompt missing %q", want)
		}
	}
	if !strings.Contains(fc.reqs[0].Messages[1].Content, "systemctl restart nginx") {
		t.Error("journal not sent")
	}

	// Transcripción grande: varias peticiones encadenadas; el diario va en la última.
	fc2 := &fakeClient{replies: []string{fullSummary("parcial"), fullSummary("final")}}
	c2 := &Compressor{Client: fc2, ChunkTokens: 1000}
	transcript := strings.Repeat("línea de transcripción bastante larga\n", 300)
	out, err = c2.Compress(context.Background(), "previo", "JOURNAL-DATA", transcript)
	if err != nil || !strings.Contains(out, "final") {
		t.Fatalf("chunked: %q %v", out, err)
	}
	if len(fc2.reqs) < 2 {
		t.Fatalf("expected chunking, got %d requests", len(fc2.reqs))
	}
	if strings.Contains(fc2.reqs[0].Messages[1].Content, "JOURNAL-DATA") ||
		!strings.Contains(fc2.reqs[len(fc2.reqs)-1].Messages[1].Content, "JOURNAL-DATA") {
		t.Error("journal must be sent only with the last chunk")
	}
	if !strings.Contains(fc2.reqs[1].Messages[1].Content, "parcial") {
		t.Error("partial summary must be chained into the next chunk")
	}
}

func TestCompressorRejectsBadOutput(t *testing.T) {
	c := &Compressor{Client: &fakeClient{replies: []string{"Resumen: todo bien."}}, ChunkTokens: 1000}
	if _, err := c.Compress(context.Background(), "", "", "[user]\nx\n"); err == nil {
		t.Error("summary without sections must be rejected")
	}
	c = &Compressor{Client: &fakeClient{err: errors.New("down")}, ChunkTokens: 1000}
	if _, err := c.Compress(context.Background(), "", "", "x"); err == nil {
		t.Error("client error must propagate")
	}
}

func TestRenderTranscript(t *testing.T) {
	out := RenderTranscript([]llm.Message{
		{Role: llm.RoleUser, Content: "revisa nginx"},
		{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{{Function: llm.FunctionCall{Name: "exec", Arguments: `{"command": "nginx", "args": ["-t"]}`}}}},
		{Role: llm.RoleTool, Name: "exec", Content: strings.Repeat("y", 5000)},
	}, 100)
	for _, want := range []string{"[user]\nrevisa nginx", `calls tool exec] {"args":["-t"],"command":"nginx"}`, "[tool result exec]", "bytes omitted"} {
		if !strings.Contains(out, want) {
			t.Errorf("missing %q in:\n%s", want, out)
		}
	}
}