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) } } }