214 lines
7.3 KiB
Go
214 lines
7.3 KiB
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)
|
|
}
|
|
}
|
|
}
|