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

package llamacpp

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"

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

// TestClient_Stream_ThinkingBudgetCutsPureReasoning: with MaxThinkingTokens
// set, a round that is still pure reasoning past the character budget must be
// cut with FinishThinkingBudget — and nothing after the cut may be delivered.
func TestClient_Stream_ThinkingBudgetCutsPureReasoning(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		// 5 chars per delta; budget = 2 tokens * 4 chars = 8 chars, so the
		// second delta (total 10) tips it over.
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"aaaaa\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"bbbbb\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"never delivered\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"never delivered\"},\"finish_reason\":\"stop\"}]}\n"))
		w.Write([]byte("data: [DONE]\n"))
	}))
	defer server.Close()

	client, err := New(Config{BaseURL: server.URL + "/v1", MaxThinkingTokens: 2})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var chunks []llm.StreamChunk
	for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if len(chunks) != 2 {
		t.Fatalf("expected 2 chunks (reasoning + budget cut), got %d: %+v", len(chunks), chunks)
	}
	last := chunks[len(chunks)-1]
	if last.FinishReason != llm.FinishThinkingBudget {
		t.Errorf("expected finish reason %q, got %q", llm.FinishThinkingBudget, last.FinishReason)
	}
	if last.ReasoningDelta != "bbbbb" {
		t.Errorf("expected the tipping reasoning delta on the final chunk, got %q", last.ReasoningDelta)
	}
	for _, c := range chunks {
		if c.Delta != "" {
			t.Errorf("no content should have been delivered, got %q", c.Delta)
		}
	}
}

// TestClient_Stream_ThinkingBudgetSparesStartedAnswer: once the model has
// begun its actual answer, exceeding the reasoning budget must NOT cut the
// stream — the spiral risk is over and real work is in flight.
func TestClient_Stream_ThinkingBudgetSparesStartedAnswer(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"aaaaa\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hola\"},\"finish_reason\":null}]}\n"))
		// Over budget, but the answer already started.
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"bbbbbbbbbb\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" mundo\"},\"finish_reason\":\"stop\"}]}\n"))
		w.Write([]byte("data: [DONE]\n"))
	}))
	defer server.Close()

	client, err := New(Config{BaseURL: server.URL + "/v1", MaxThinkingTokens: 2})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var content string
	var finish string
	for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		content += chunk.Delta
		if chunk.FinishReason != "" {
			finish = chunk.FinishReason
		}
	}

	if content != "Hola mundo" {
		t.Errorf("expected the full answer, got %q", content)
	}
	if finish != "stop" {
		t.Errorf("expected a normal stop, got %q", finish)
	}
}

// TestClient_Stream_NoThinkingBudgetMeansUnlimited: MaxThinkingTokens 0 keeps
// today's behavior — reasoning streams without any client-side cap.
func TestClient_Stream_NoThinkingBudgetMeansUnlimited(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n"))
		w.Write([]byte("data: [DONE]\n"))
	}))
	defer server.Close()

	client, err := New(Config{BaseURL: server.URL + "/v1"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var content, finish string
	for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		content += chunk.Delta
		if chunk.FinishReason != "" {
			finish = chunk.FinishReason
		}
	}

	if content != "ok" || finish != "stop" {
		t.Errorf("expected uncut stream (content %q, finish %q), got content %q finish %q", "ok", "stop", content, finish)
	}
}