rony-llm-agent/pkg/agent/thinking_budget_test.go

package agent

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

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
)

// TestRunStream_RecoversFromThinkingBudgetCut covers the reasoning-spiral
// failure seen live with Qwen3.6 + llama.cpp: the model thinks for tens of
// thousands of tokens without ever acting, the provider cuts the round with
// FinishThinkingBudget, and the loop must re-prompt for direct action instead
// of silently ending the turn with nothing.
func TestRunStream_RecoversFromThinkingBudgetCut(t *testing.T) {
	executed := 0
	stub := &scriptedLLM{responses: []llm.CompletionResponse{
		// Round 1: pure reasoning, cut by the provider's budget enforcement.
		{Reasoning: "hmm let me think about this again and again", StopReason: llm.FinishThinkingBudget},
		// Round 2 (after the nudge): a real tool call.
		{ToolCalls: []llm.ToolCall{{ID: "1", Name: "edit", Arguments: json.RawMessage(`{}`)}}},
		// Round 3: final answer.
		{Content: "Listo."},
	}}

	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})

	var final strings.Builder
	for chunk, err := range loop.RunStream(context.Background(), "arregla x.py") {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		final.WriteString(chunk.Delta)
	}

	if executed != 1 {
		t.Fatalf("expected the post-nudge tool call to execute once, got %d", executed)
	}
	if !strings.Contains(final.String(), "Listo.") {
		t.Fatalf("expected the turn to continue to a final answer, got %q", final.String())
	}
	foundNudge := false
	for _, m := range stub.lastMessages {
		if m.Role == llm.RoleUser && strings.Contains(m.Content, "exceeded the thinking budget") {
			foundNudge = true
		}
	}
	if !foundNudge {
		t.Fatal("expected the thinking-budget nudge in the follow-up request messages")
	}
}

// TestRunStream_ThinkingBudgetNudgeGivesUpAfterLimit keeps a model that
// spirals every single round from ping-ponging forever: after
// maxThinkingBudgetNudges the turn ends.
func TestRunStream_ThinkingBudgetNudgeGivesUpAfterLimit(t *testing.T) {
	executed := 0
	spiral := llm.CompletionResponse{Reasoning: "thinking forever", StopReason: llm.FinishThinkingBudget}
	stub := &scriptedLLM{responses: []llm.CompletionResponse{spiral, spiral, spiral, spiral}}

	loop := New(Config{LLM: stub, Tools: editTestRegistry(t, &executed), MaxIters: 10})

	for _, err := range loop.RunStream(context.Background(), "haz algo") {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
	}

	if stub.calls != maxThinkingBudgetNudges+1 {
		t.Fatalf("expected %d rounds (original + nudges), got %d", maxThinkingBudgetNudges+1, stub.calls)
	}
	if executed != 0 {
		t.Fatalf("no tool should have executed, got %d", executed)
	}
}