hyperguild/internal/exec/litellm_test.go

package exec_test

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	iexec "git.d-ma.be/mathias/hyperguild/internal/exec"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func chatResponse(t *testing.T, content string) []byte {
	t.Helper()
	resp := map[string]any{
		"choices": []map[string]any{
			{"message": map[string]any{"role": "assistant", "content": content}},
		},
	}
	data, err := json.Marshal(resp)
	require.NoError(t, err)
	return data
}

func TestLiteLLMReturnsText(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "/v1/chat/completions", r.URL.Path)
		assert.Equal(t, "application/json", r.Header.Get("Content-Type"))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(chatResponse(t, "here is my analysis"))
	}))
	defer srv.Close()

	ex := iexec.NewLiteLLM(srv.URL, "", 5*time.Second)
	text, dur, err := ex.Complete(context.Background(), "ollama/devstral", "system prompt", "user prompt")
	require.NoError(t, err)
	assert.Equal(t, "here is my analysis", text)
	assert.GreaterOrEqual(t, dur, int64(0))
}

func TestLiteLLMSendsAuthHeader(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, "Bearer secret", r.Header.Get("Authorization"))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(chatResponse(t, "ok"))
	}))
	defer srv.Close()

	ex := iexec.NewLiteLLM(srv.URL, "secret", 5*time.Second)
	_, _, err := ex.Complete(context.Background(), "model", "sys", "user")
	require.NoError(t, err)
}

func TestLiteLLMErrorOnNonOKStatus(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
	}))
	defer srv.Close()

	ex := iexec.NewLiteLLM(srv.URL, "", 5*time.Second)
	_, _, err := ex.Complete(context.Background(), "model", "sys", "user")
	assert.ErrorContains(t, err, "503")
}

func TestLiteLLMErrorOnEmptyChoices(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"choices":[]}`))
	}))
	defer srv.Close()

	ex := iexec.NewLiteLLM(srv.URL, "", 5*time.Second)
	_, _, err := ex.Complete(context.Background(), "model", "sys", "user")
	assert.ErrorContains(t, err, "no choices")
}

func TestLiteLLMStripsTrailingResultJSON(t *testing.T) {
	content := "## Hypotheses\n\n**H1 (high):** nil map access.\n\n```json\n{\n  \"status\": \"pass\",\n  \"phase\": \"debug\",\n  \"skill\": \"debug\"\n}\n```"
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(chatResponse(t, content))
	}))
	defer srv.Close()

	ex := iexec.NewLiteLLM(srv.URL, "", 5*time.Second)
	text, _, err := ex.Complete(context.Background(), "model", "sys", "user")
	require.NoError(t, err)
	assert.Contains(t, text, "nil map access")
	assert.NotContains(t, text, `"status"`)
	assert.NotContains(t, text, "```json")
}

func TestLiteLLMKeepsNonResultJSONFence(t *testing.T) {
	// A json block that is part of the actual answer (no status/phase) should be kept.
	content := "Use this config:\n\n```json\n{\"model\": \"koala/phi4\"}\n```"
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(chatResponse(t, content))
	}))
	defer srv.Close()

	ex := iexec.NewLiteLLM(srv.URL, "", 5*time.Second)
	text, _, err := ex.Complete(context.Background(), "model", "sys", "user")
	require.NoError(t, err)
	assert.Contains(t, text, `"model"`)
	assert.Contains(t, text, "```json")
}

func TestLiteLLMRespectsContextCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	ex := iexec.NewLiteLLM("http://invalid.example.com", "", 1*time.Second)
	_, _, err := ex.Complete(ctx, "model", "sys", "user")
	assert.Error(t, err)
}