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

package llamacpp

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

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

func TestClient_Name(t *testing.T) {
	client := &Client{}
	if client.Name() != "llama.cpp" {
		t.Errorf("expected 'llama.cpp', got %q", client.Name())
	}
}

func TestClient_Capabilities(t *testing.T) {
	client := &Client{}
	caps := client.Capabilities()
	if !caps.SupportsTools {
		t.Error("expected SupportsTools to be true")
	}
	if !caps.SupportsVision {
		t.Error("expected SupportsVision to be true")
	}
	if !caps.SupportsVideo {
		t.Error("expected SupportsVideo to be true")
	}
}

func TestClient_Generate(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(llamaChatResponse{
			ID:    "llama-123",
			Model: "llama3",
			Choices: []llamaChoice{
				{
					Index:        0,
					FinishReason: "stop",
					Message: llamaMessageResult{
						Role:    "assistant",
						Content: "Hello from llama.cpp!",
					},
				},
			},
			Usage: llamaUsage{
				PromptTokens:     10,
				CompletionTokens: 5,
				TotalTokens:      15,
			},
		})
	}))
	defer server.Close()

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

	resp, err := client.Generate(context.Background(), llm.CompletionRequest{
		Model: "llama3",
		Messages: []llm.Message{
			{Role: llm.RoleUser, Content: "Hi"},
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Content != "Hello from llama.cpp!" {
		t.Errorf("expected 'Hello from llama.cpp!', got %q", resp.Content)
	}
	if resp.ID != "llama-123" {
		t.Errorf("expected 'llama-123', got %q", resp.ID)
	}
	if resp.Usage.InputTokens != 10 {
		t.Errorf("expected 10 input tokens, got %d", resp.Usage.InputTokens)
	}
}

// TestClient_BuildRequest_SendsToolCallHistory is the regression test for a
// bug where an assistant message's ToolCalls and a tool message's
// ToolCallID were silently dropped when building the wire request: the
// model would see a "tool" message with nothing tying it to a prior
// assistant turn, lose track of what it had already tried, and re-attempt
// the same thing over and over (reported in production as Rony repeatedly
// re-greeting and re-searching for a file instead of ever finishing).
func TestClient_BuildRequest_SendsToolCallHistory(t *testing.T) {
	var gotBody string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		gotBody = string(body)
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(llamaChatResponse{
			Choices: []llamaChoice{{Message: llamaMessageResult{Content: "ok"}}},
		})
	}))
	defer server.Close()

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

	_, err = client.Generate(context.Background(), llm.CompletionRequest{
		Messages: []llm.Message{
			{Role: llm.RoleUser, Content: "busca el archivo"},
			{
				Role:    llm.RoleAssistant,
				Content: "voy a buscar",
				ToolCalls: []llm.ToolCall{
					{ID: "call-1", Name: "glob", Arguments: json.RawMessage(`{"pattern":"*.md"}`)},
				},
			},
			{Role: llm.RoleTool, ToolCallID: "call-1", Content: "No matches found."},
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var sent struct {
		Messages []struct {
			Role       string `json:"role"`
			Content    string `json:"content"`
			ToolCallID string `json:"tool_call_id"`
			ToolCalls  []struct {
				ID       string `json:"id"`
				Function struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				} `json:"function"`
			} `json:"tool_calls"`
		} `json:"messages"`
	}
	if err := json.Unmarshal([]byte(gotBody), &sent); err != nil {
		t.Fatalf("failed to parse sent body: %v\nbody: %s", err, gotBody)
	}

	if len(sent.Messages) != 3 {
		t.Fatalf("expected 3 messages sent, got %d: %s", len(sent.Messages), gotBody)
	}

	assistantMsg := sent.Messages[1]
	if assistantMsg.Role != "assistant" {
		t.Fatalf("expected message 1 to be the assistant turn, got role %q", assistantMsg.Role)
	}
	if len(assistantMsg.ToolCalls) != 1 || assistantMsg.ToolCalls[0].ID != "call-1" {
		t.Fatalf("expected the assistant message to carry its tool_calls with id 'call-1', got %+v", assistantMsg.ToolCalls)
	}
	if assistantMsg.ToolCalls[0].Function.Name != "glob" {
		t.Errorf("expected function name 'glob', got %q", assistantMsg.ToolCalls[0].Function.Name)
	}

	toolMsg := sent.Messages[2]
	if toolMsg.Role != "tool" {
		t.Fatalf("expected message 2 to be the tool result, got role %q", toolMsg.Role)
	}
	if toolMsg.ToolCallID != "call-1" {
		t.Errorf("expected tool_call_id 'call-1' on the tool message, got %q (body: %s)", toolMsg.ToolCallID, gotBody)
	}
}

func TestClient_Generate_Error(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("server error"))
	}))
	defer server.Close()

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

	_, err = client.Generate(context.Background(), llm.CompletionRequest{})
	if err == nil {
		t.Fatal("expected error, got nil")
	}
}

func TestClient_Generate_ToolCall(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(llamaChatResponse{
			ID:    "llama-tool-1",
			Model: "llama3",
			Choices: []llamaChoice{
				{
					Index:        0,
					FinishReason: "tool_calls",
					Message: llamaMessageResult{
						Role:    "assistant",
						Content: "",
						ToolCalls: []llamaToolCall{
							{
								ID:   "call-1",
								Type: "function",
								Function: llamaFunction{
									Name:      "calculate",
									Arguments: `{"a":1,"b":2}`,
								},
							},
						},
					},
				},
			},
			Usage: llamaUsage{
				PromptTokens:     20,
				CompletionTokens: 10,
				TotalTokens:      30,
			},
		})
	}))
	defer server.Close()

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

	resp, err := client.Generate(context.Background(), llm.CompletionRequest{
		Model: "llama3",
		Messages: []llm.Message{
			{Role: llm.RoleUser, Content: "What is 1 + 2?"},
		},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(resp.ToolCalls) != 1 {
		t.Fatalf("expected 1 tool call, got %d", len(resp.ToolCalls))
	}
	if resp.ToolCalls[0].Name != "calculate" {
		t.Errorf("expected 'calculate', got %q", resp.ToolCalls[0].Name)
	}
}

func TestClient_Stream(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\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n"))
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"!\"},\"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 chunks []llm.StreamChunk
	stream := client.Stream(context.Background(), llm.CompletionRequest{})
	for chunk, err := range stream {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if len(chunks) != 3 {
		t.Errorf("expected 3 chunks, got %d", len(chunks))
	}
	if chunks[0].Delta != "Hello" {
		t.Errorf("expected 'Hello', got %q", chunks[0].Delta)
	}
	if chunks[1].Delta != " world" {
		t.Errorf("expected ' world', got %q", chunks[1].Delta)
	}
	if chunks[2].Delta != "!" {
		t.Errorf("expected '!', got %q", chunks[2].Delta)
	}
}

func TestClient_Stream_FinishReason(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\":{\"content\":\"done\"},\"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 chunks []llm.StreamChunk
	stream := client.Stream(context.Background(), llm.CompletionRequest{})
	for chunk, err := range stream {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if len(chunks) != 1 {
		t.Fatalf("expected 1 chunk, got %d", len(chunks))
	}
	if chunks[0].FinishReason != "stop" {
		t.Errorf("expected 'stop' finish reason, got %q", chunks[0].FinishReason)
	}
}

// TestClient_Stream_ToolCallSingleEvent covers the simplest case: a server
// that sends the whole tool call (id, name, complete arguments) in one delta
// followed immediately by finish_reason "tool_calls".
func TestClient_Stream_ToolCallSingleEvent(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":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"write","arguments":"{\"path\":\"a.txt\",\"content\":\"hi\"}"}}]},"finish_reason":"tool_calls"}]}` + "\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 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) != 1 {
		t.Fatalf("expected 1 chunk, got %d: %+v", len(chunks), chunks)
	}
	if len(chunks[0].ToolCalls) != 1 {
		t.Fatalf("expected 1 tool call in the chunk, got %d", len(chunks[0].ToolCalls))
	}
	call := chunks[0].ToolCalls[0]
	if call.ID != "call-1" || call.Name != "write" {
		t.Errorf("expected call-1/write, got %+v", call)
	}
	if string(call.Arguments) != `{"path":"a.txt","content":"hi"}` {
		t.Errorf("unexpected arguments: %s", call.Arguments)
	}
}

// TestClient_Stream_ToolCallFragmentsAssembled is the regression test for
// the actual bug reported in production: llama.cpp (like any OpenAI-style
// server) streams a tool call's arguments in many small deltas keyed by
// index, with the name/id only present in the first fragment. The old
// Stream() implementation never even read choice.Delta.ToolCalls, so every
// fragment was silently dropped and the agent loop never saw a tool call at
// all - the model would narrate "I'll write the file" and nothing would
// happen. This verifies the fragments are buffered and only surfaced, fully
// assembled, once finish_reason arrives.
func TestClient_Stream_ToolCallFragmentsAssembled(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":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"write","arguments":""}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":"}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a.txt\",\"content\""}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":":\"hi\"}"}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}` + "\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 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)
	}

	// The four fragment-only events must not surface as separate empty
	// chunks; only the finish_reason event, carrying the fully assembled
	// call, should be yielded.
	if len(chunks) != 1 {
		t.Fatalf("expected 1 chunk (fragments buffered, only the assembled call yielded), got %d: %+v", len(chunks), chunks)
	}
	if chunks[0].FinishReason != "tool_calls" {
		t.Errorf("expected finish_reason 'tool_calls', got %q", chunks[0].FinishReason)
	}
	if len(chunks[0].ToolCalls) != 1 {
		t.Fatalf("expected 1 assembled tool call, got %d", len(chunks[0].ToolCalls))
	}
	call := chunks[0].ToolCalls[0]
	if call.ID != "call-1" || call.Name != "write" {
		t.Errorf("expected call-1/write, got %+v", call)
	}
	if string(call.Arguments) != `{"path":"a.txt","content":"hi"}` {
		t.Errorf("expected assembled arguments %q, got %q", `{"path":"a.txt","content":"hi"}`, call.Arguments)
	}
}

// TestClient_Stream_ToolCallWithPrecedingContent verifies that reasoning or
// content deltas that arrive before a tool call (e.g. a model "thinking"
// before deciding to call a tool) are still streamed normally, and don't get
// mixed up with the buffered tool-call fragments.
func TestClient_Stream_ToolCallWithPrecedingContent(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":{"content":"Voy a escribir el archivo."}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call-9","type":"function","function":{"name":"write","arguments":"{}"}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}` + "\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 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 (content, then the assembled tool call), got %d: %+v", len(chunks), chunks)
	}
	if chunks[0].Delta != "Voy a escribir el archivo." {
		t.Errorf("expected the content delta first, got %q", chunks[0].Delta)
	}
	if len(chunks[0].ToolCalls) != 0 {
		t.Errorf("expected the content chunk to carry no tool calls, got %+v", chunks[0].ToolCalls)
	}
	if len(chunks[1].ToolCalls) != 1 || chunks[1].ToolCalls[0].Name != "write" {
		t.Errorf("expected the second chunk to carry the assembled write call, got %+v", chunks[1].ToolCalls)
	}
}

// TestClient_Stream_ParallelToolCalls verifies two tool calls streamed in
// parallel (interleaved by index) are assembled independently and returned
// in call order.
func TestClient_Stream_ParallelToolCalls(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":{"tool_calls":[{"index":0,"id":"call-a","type":"function","function":{"name":"read","arguments":"{\"path\":"}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":1,"id":"call-b","type":"function","function":{"name":"glob","arguments":"{\"pattern\":"}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"a.txt\"}"}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"*.go\"}"}}]}}]}` + "\n"))
		w.Write([]byte(`data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}` + "\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 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) != 1 {
		t.Fatalf("expected 1 chunk, got %d: %+v", len(chunks), chunks)
	}
	if len(chunks[0].ToolCalls) != 2 {
		t.Fatalf("expected 2 assembled tool calls, got %d", len(chunks[0].ToolCalls))
	}
	if chunks[0].ToolCalls[0].Name != "read" || string(chunks[0].ToolCalls[0].Arguments) != `{"path":"a.txt"}` {
		t.Errorf("unexpected first call: %+v", chunks[0].ToolCalls[0])
	}
	if chunks[0].ToolCalls[1].Name != "glob" || string(chunks[0].ToolCalls[1].Arguments) != `{"pattern":"*.go"}` {
		t.Errorf("unexpected second call: %+v", chunks[0].ToolCalls[1])
	}
}

func TestClient_Stream_RequestsAndParsesUsage(t *testing.T) {
	var gotBody string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		gotBody = string(body)
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}]}\n"))
		w.Write([]byte("data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2,\"total_tokens\":12}}\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 chunks []llm.StreamChunk
	stream := client.Stream(context.Background(), llm.CompletionRequest{})
	for chunk, err := range stream {
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		chunks = append(chunks, chunk)
	}

	if !strings.Contains(gotBody, `"stream_options":{"include_usage":true}`) {
		t.Errorf("expected the request to ask for usage via stream_options, got body: %s", gotBody)
	}

	if len(chunks) != 2 {
		t.Fatalf("expected 2 chunks (content + usage-only), got %d", len(chunks))
	}
	usage := chunks[len(chunks)-1].Usage
	if usage.InputTokens != 10 || usage.OutputTokens != 2 || usage.TotalTokens != 12 {
		t.Errorf("expected usage to be parsed from the final event, got %+v", usage)
	}
}