MuyueWorkspace/internal/orchestrator/orchestrator_test.go

package orchestrator

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

	"github.com/muyue/muyue/internal/config"
)

func TestCleanAIResponse(t *testing.T) {
	tests := []struct {
		name     string
		input    string
		expected string
	}{
		{
			"malformed think tags pass through",
			"<think internal reasoning</think Hello world",
			"<think internal reasoning</think Hello world",
		},
		{
			"removes Think tags",
			"<Think>reasoning</Think>response",
			"response",
		},
		{
			"think with attrs, no closing bracket",
			"<think type=re>reasoning</think result",
			"<think type=re>reasoning</think result",
		},
		{
			"removes stream markers",
			"text\n<<\ninternal\n>>\nvisible",
			"text\nvisible",
		},
		{
			"removes triple markers",
			"text\n<<<\ninternal\n>>>\nvisible",
			"text\nvisible",
		},
		{
			"plain text unchanged",
			"Hello world",
			"Hello world",
		},
		{
			"empty input",
			"",
			"",
		},
		{
			"malformed think block no closing bracket",
			"<think some reasoning here</think rest",
			"<think some reasoning here</think rest",
		},
		{
			"malformed simple think no closing bracket",
			"before<think reasoning</think after",
			"before<think reasoning</think after",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := CleanAIResponse(tt.input)
			result = strings.TrimSpace(result)
			expected := strings.TrimSpace(tt.expected)
			if result != expected {
				t.Errorf("CleanAIResponse(%q) = %q, want %q", tt.input, result, expected)
			}
		})
	}
}

func TestCleanAIResponseThinkRegex(t *testing.T) {
	input2 := "<Think>some reasoning</Think>actual response"
	result2 := CleanAIResponse(input2)
	if result2 != "actual response" {
		t.Errorf("Valid Think tags should be removed: %q", result2)
	}

	input3 := "<think\nmultiline\nreasoning</think visible"
	result3 := CleanAIResponse(input3)
	// No closing > on opening tag, so won't match regex
	if result3 != "<think\nmultiline\nreasoning</think visible" {
		t.Errorf("Malformed think should not be removed: %q", result3)
	}

	input4 := "<think type=re>reasoning</think visible"
	result4 := CleanAIResponse(input4)
	// </think followed by space, not >, so won't match
	if result4 != "<think type=re>reasoning</think visible" {
		t.Errorf("Malformed closing should not be removed: %q", result4)
	}

	input_real := "prefix<think reasoning here</think suffix"
	result_real := CleanAIResponse(input_real)
	// The closing </think has no > after it, so won't match
	if result_real != "prefix<think reasoning here</think suffix" {
		t.Errorf("Malformed tags should pass through: %q", result_real)
	}

	input_valid := "<Think>reasoning</Think>result"
	result_valid := CleanAIResponse(input_valid)
	if result_valid != "result" {
		t.Errorf("Valid tags should be removed: %q", result_valid)
	}
}

func TestGetProviderBaseURL(t *testing.T) {
	tests := []struct {
		provider string
		want     string
	}{
		{"minimax", "https://api.minimax.io/v1"},
		{"anthropic", "https://api.anthropic.com/v1"},
		{"openai", "https://api.openai.com/v1"},
		{"zai", "https://api.z.ai/v1"},
		{"unknown", ""},
	}
	for _, tt := range tests {
		got := getProviderBaseURL(tt.provider)
		if got != tt.want {
			t.Errorf("getProviderBaseURL(%q) = %q, want %q", tt.provider, got, tt.want)
		}
	}
}

func TestNewNoProvider(t *testing.T) {
	cfg := config.Default()
	for i := range cfg.AI.Providers {
		cfg.AI.Providers[i].Active = false
	}
	_, err := New(cfg)
	if err == nil {
		t.Error("Should fail with no active provider")
	}
}

func TestNewNoAPIKey(t *testing.T) {
	cfg := config.Default()
	cfg.AI.Providers[0].Active = true
	cfg.AI.Providers[0].APIKey = ""
	_, err := New(cfg)
	if err == nil {
		t.Error("Should fail with no API key")
	}
}

func TestSendStreamChunks(t *testing.T) {
	sseBody := `data: {"choices":[{"delta":{"content":"Hello"}}]}
data: {"choices":[{"delta":{"content":" world"}}]}
data: {"choices":[{"delta":{"content":"!"}}]}
data: [DONE]
`
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer test-key" {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		var reqBody ChatRequest
		if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		if !reqBody.Stream {
			http.Error(w, "stream must be true", http.StatusBadRequest)
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte(sseBody))
	}))
	defer ts.Close()

	cfg := config.Default()
	cfg.AI.Providers[0].Active = true
	cfg.AI.Providers[0].APIKey = "test-key"
	cfg.AI.Providers[0].BaseURL = ts.URL

	orb, err := New(cfg)
	if err != nil {
		t.Fatalf("New: %v", err)
	}

	var chunks []string
	result, err := orb.SendStream("hi", func(chunk string) {
		chunks = append(chunks, chunk)
	})
	if err != nil {
		t.Fatalf("SendStream: %v", err)
	}
	if result != "Hello world!" {
		t.Errorf("SendStream result = %q, want %q", result, "Hello world!")
	}
	if len(chunks) != 3 {
		t.Fatalf("expected 3 chunks, got %d: %v", len(chunks), chunks)
	}
	if strings.Join(chunks, "") != "Hello world!" {
		t.Errorf("chunks joined = %q, want %q", strings.Join(chunks, ""), "Hello world!")
	}
}

func TestSendStreamHistory(t *testing.T) {
	callCount := 0
	sseBody := `data: {"choices":[{"delta":{"content":"reply"}}]}
data: [DONE]
`
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		callCount++
		var reqBody ChatRequest
		if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		if callCount == 1 {
			if len(reqBody.Messages) != 2 {
				t.Errorf("first call: expected 2 messages (system + 1 user), got %d", len(reqBody.Messages))
			}
		} else {
			if len(reqBody.Messages) != 4 {
				t.Errorf("second call: expected 4 messages (system + 3 history), got %d", len(reqBody.Messages))
			}
		}
		w.Header().Set("Content-Type", "text/event-stream")
		w.Write([]byte(sseBody))
	}))
	defer ts.Close()

	cfg := config.Default()
	cfg.AI.Providers[0].Active = true
	cfg.AI.Providers[0].APIKey = "test-key"
	cfg.AI.Providers[0].BaseURL = ts.URL

	orb, err := New(cfg)
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	orb.SetSystemPrompt("you are helpful")

	_, _ = orb.SendStream("first", func(string) {})
	_, _ = orb.SendStream("second", func(string) {})

	orb.histMu.Lock()
	if len(orb.history) != 4 {
		t.Errorf("expected 4 history entries (2 user + 2 assistant), got %d", len(orb.history))
	}
	orb.histMu.Unlock()
}

func TestSendStreamAPIError(t *testing.T) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, `{"error":"rate limited"}`, http.StatusTooManyRequests)
	}))
	defer ts.Close()

	cfg := config.Default()
	cfg.AI.Providers[0].Active = true
	cfg.AI.Providers[0].APIKey = "test-key"
	cfg.AI.Providers[0].BaseURL = ts.URL

	orb, err := New(cfg)
	if err != nil {
		t.Fatalf("New: %v", err)
	}

	_, err = orb.SendStream("hi", func(string) {})
	if err == nil {
		t.Error("expected error for non-200 response")
	}
	if !strings.Contains(err.Error(), "429") {
		t.Errorf("error should mention status code, got: %v", err)
	}
}