dynagate/gateway_test.go

package main

import (
  "bytes"
  "context"
  "encoding/json"
  "fmt"
  "io"
  "net/http"
  "net/http/httptest"
  "os"
  "path/filepath"
  "reflect"
  "strings"
  "testing"
  "time"
)

func TestURLBuilding(t *testing.T) {
  tests := []struct {
    endpoint string
    reqPath  string
    expected string
  }{
    {"https://api.openai.com/v1", "/v1/chat/completions", "https://api.openai.com/v1/chat/completions"},
    {"https://api.openai.com/v1/", "/v1/chat/completions", "https://api.openai.com/v1/chat/completions"},
    {"https://api.openai.com", "/v1/chat/completions", "https://api.openai.com/v1/chat/completions"},
    {"http://localhost:8000/api/v1", "/v1/chat/completions", "http://localhost:8000/api/v1/chat/completions"},
    {"http://localhost:11434", "/v1/chat/completions", "http://localhost:11434/v1/chat/completions"},
    {"https://custom.endpoint/v1/custom", "/v1/chat/completions", "https://custom.endpoint/v1/custom/v1/chat/completions"},
  }

  for _, tc := range tests {
    t.Run(tc.endpoint+"+"+tc.reqPath, func(t *testing.T) {
      got := buildURL(tc.endpoint, tc.reqPath)
      if got != tc.expected {
        t.Errorf("buildURL(%q, %q) = %q; want %q", tc.endpoint, tc.reqPath, got, tc.expected)
      }
    })
  }
}

func TestTrialOrdering(t *testing.T) {
  configs := []ModelConfig{
    {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
    {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
    {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
    {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
  }
  uniqueModels := []string{"gpt-4", "gpt-3.5", "claude"}

  // Test case 1: Requested model matches the first model "gpt-4"
  {
    trials := getTrialConfigs(configs, uniqueModels, "gpt-4")
    expected := []ModelConfig{
      {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
      {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
      {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
      {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
    }
    if !reflect.DeepEqual(trials, expected) {
      t.Errorf("Trial ordering for gpt-4 mismatch.\nGot: %+v\nWant: %+v", trials, expected)
    }
  }

  // Test case 2: Requested model matches the second model "gpt-3.5"
  {
    trials := getTrialConfigs(configs, uniqueModels, "gpt-3.5")
    expected := []ModelConfig{
      {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
      {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
      {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
      {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
    }
    if !reflect.DeepEqual(trials, expected) {
      t.Errorf("Trial ordering for gpt-3.5 mismatch.\nGot: %+v\nWant: %+v", trials, expected)
    }
  }

  // Test case 3: Requested model matches the last model "claude"
  {
    trials := getTrialConfigs(configs, uniqueModels, "claude")
    expected := []ModelConfig{
      {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
      {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
      {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
      {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
    }
    if !reflect.DeepEqual(trials, expected) {
      t.Errorf("Trial ordering for claude mismatch.\nGot: %+v\nWant: %+v", trials, expected)
    }
  }

  // Test case 4: Requested model doesn't match any config
  {
    trials := getTrialConfigs(configs, uniqueModels, "non-existent")
    expected := []ModelConfig{
      {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
      {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
      {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
      {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
    }
    if !reflect.DeepEqual(trials, expected) {
      t.Errorf("Trial ordering for non-existent model mismatch.\nGot: %+v\nWant: %+v", trials, expected)
    }
  }
}

func TestConfigManagerCSVWatcher(t *testing.T) {
  tmpDir, err := os.MkdirTemp("", "dynagate-test")
  if err != nil {
    t.Fatalf("failed to create temp dir: %v", err)
  }
  defer os.RemoveAll(tmpDir)

  csvFile := filepath.Join(tmpDir, "models.csv")

  // Write initial content
  initialContent := `model,key,endpoint
gpt-4,key-1,http://localhost:8001
gpt-3.5,key-2,http://localhost:8002
`
  if err := os.WriteFile(csvFile, []byte(initialContent), 0644); err != nil {
    t.Fatalf("failed to write csv file: %v", err)
  }

  cm := NewConfigManager(csvFile)
  configs := cm.GetConfigs()
  if len(configs) != 2 || configs[0].Model != "gpt-4" || configs[1].Model != "gpt-3.5" {
    t.Fatalf("unexpected initial configs: %+v", configs)
  }

  // Start watching in background
  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()
  go cm.Watch(ctx)

  // Wait briefly, modify file, and check if reload occurs
  time.Sleep(100 * time.Millisecond)

  updatedContent := `model,key,endpoint
gpt-4,key-1,http://localhost:8001
gpt-3.5,key-2,http://localhost:8002
claude,key-3,http://localhost:8003
`
  // Make sure we sleep slightly to ensure file modtime actually changes on modern filesystems
  time.Sleep(1 * time.Second)
  if err := os.WriteFile(csvFile, []byte(updatedContent), 0644); err != nil {
    t.Fatalf("failed to update csv file: %v", err)
  }

  // Poll GetConfigs until it gets updated (timeout after 3 seconds)
  deadline := time.Now().Add(3 * time.Second)
  var updatedConfigs []ModelConfig
  for time.Now().Before(deadline) {
    updatedConfigs = cm.GetConfigs()
    if len(updatedConfigs) == 3 {
      break
    }
    time.Sleep(100 * time.Millisecond)
  }

  if len(updatedConfigs) != 3 || updatedConfigs[2].Model != "claude" {
    t.Errorf("Hot-reload failed. Got configs: %+v", updatedConfigs)
  }
}

func TestGatewayCompletionsFailover(t *testing.T) {
  // We will setup 3 mock upstream servers.
  // Server A: Always returns 500
  // Server B: Always returns 401
  // Server C: Succeeds and returns a chat completion response

  serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusInternalServerError)
    _, _ = w.Write([]byte(`{"error": "Internal Server Error A"}`))
  }))
  defer serverA.Close()

  serverB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusUnauthorized)
    _, _ = w.Write([]byte(`{"error": "Unauthorized B"}`))
  }))
  defer serverB.Close()

  serverC := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    // Read request body to verify model was modified
    var body map[string]any
    _ = json.NewDecoder(r.Body).Decode(&body)
    modelSent := body["model"].(string)

    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(fmt.Sprintf(`{"choices":[{"message":{"content":"Hello from %s"}}]}`, modelSent)))
  }))
  defer serverC.Close()

  // Configure gateway ConfigManager with endpoints pointing to these servers
  // We want to verify that when we call completions:
  // - First it tries model-1 (pointing to serverA) -> fails
  // - Then tries model-1 next key (pointing to serverB) -> fails
  // - Then retries with next model model-2 (pointing to serverC) -> succeeds!
  configs := []ModelConfig{
    {Model: "model-1", Key: "key-fail-1", Endpoint: serverA.URL},
    {Model: "model-1", Key: "key-fail-2", Endpoint: serverB.URL},
    {Model: "model-2", Key: "key-success", Endpoint: serverC.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-1", "model-2"},
  }

  handler := handleChatCompletions(cm, nil)

  // 1. Test non-streaming failover request starting with model-1
  reqBodyObj := map[string]any{
    "model": "model-1",
    "messages": []map[string]string{
      {"role": "user", "content": "Hi"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  resp := w.Result()
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200, got %d", resp.StatusCode)
  }

  var respObj map[string]any
  _ = json.NewDecoder(resp.Body).Decode(&respObj)

  choices, ok := respObj["choices"].([]any)
  if !ok || len(choices) == 0 {
    t.Fatalf("Invalid response structure: %+v", respObj)
  }

  content := choices[0].(map[string]any)["message"].(map[string]any)["content"].(string)
  expectedContent := "Hello from model-2"
  if content != expectedContent {
    t.Errorf("Expected content %q, got %q (failover to model-2 did not happen or model was not updated)", expectedContent, content)
  }
}

func TestGatewayCompletionsStreamingFailover(t *testing.T) {
  // We will setup 2 mock upstream servers.
  // Server A: Always returns 500
  // Server B: Streams chat completions chunks

  serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusInternalServerError)
  }))
  defer serverA.Close()

  serverB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")
    w.WriteHeader(http.StatusOK)

    flusher, _ := w.(http.Flusher)
    _, _ = w.Write([]byte("data: chunk1\n\n"))
    flusher.Flush()
    _, _ = w.Write([]byte("data: chunk2\n\n"))
    flusher.Flush()
  }))
  defer serverB.Close()

  configs := []ModelConfig{
    {Model: "model-1", Key: "key-fail", Endpoint: serverA.URL},
    {Model: "model-2", Key: "key-success", Endpoint: serverB.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-1", "model-2"},
  }

  handler := handleChatCompletions(cm, nil)

  reqBodyObj := map[string]any{
    "model":  "model-1",
    "stream": true,
    "messages": []map[string]string{
      {"role": "user", "content": "Hi"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  resp := w.Result()
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200, got %d", resp.StatusCode)
  }

  respBytes, _ := io.ReadAll(resp.Body)
  respStr := string(respBytes)

  if !strings.Contains(respStr, "data: chunk1") || !strings.Contains(respStr, "data: chunk2") {
    t.Errorf("Expected streaming chunks to be forwarded, got: %q", respStr)
  }
}

func TestModelsEndpoint(t *testing.T) {
  configs := []ModelConfig{
    {Model: "model-1", Key: "key-1", Endpoint: "ep-1"},
    {Model: "model-1", Key: "key-2", Endpoint: "ep-1"},
    {Model: "model-2", Key: "key-3", Endpoint: "ep-2"},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-1", "model-2"},
  }

  handler := handleModels(cm, nil)

  req := httptest.NewRequest("GET", "/v1/models", nil)
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  resp := w.Result()
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200, got %d", resp.StatusCode)
  }

  var respObj map[string]any
  _ = json.NewDecoder(resp.Body).Decode(&respObj)

  data, ok := respObj["data"].([]any)
  if !ok || len(data) != 2 {
    t.Fatalf("Expected 2 models in response data, got: %+v", respObj)
  }

  model1 := data[0].(map[string]any)["id"].(string)
  model2 := data[1].(map[string]any)["id"].(string)

  if model1 != "model-1" || model2 != "model-2" {
    t.Errorf("Unexpected models. Got: %s, %s", model1, model2)
  }
}

func TestGatewayAuthentication(t *testing.T) {
  configs := []ModelConfig{
    {Model: "model-1", Key: "key-1", Endpoint: "ep-1"},
  }
  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-1"},
  }

  // Set up handlers expecting multiple secret tokens
  expectedTokens := []string{"secret-token-1", "secret-token-2"}
  modelsHandler := handleModels(cm, expectedTokens)
  completionsHandler := handleChatCompletions(cm, expectedTokens)

  // Case 1: No Authorization header
  {
    req := httptest.NewRequest("GET", "/v1/models", nil)
    w := httptest.NewRecorder()
    modelsHandler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusUnauthorized {
      t.Errorf("Expected 401 Unauthorized for missing token, got %d", w.Result().StatusCode)
    }
  }

  // Case 2: Incorrect Authorization header
  {
    req := httptest.NewRequest("GET", "/v1/models", nil)
    req.Header.Set("Authorization", "Bearer wrong-token")
    w := httptest.NewRecorder()
    modelsHandler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusUnauthorized {
      t.Errorf("Expected 401 Unauthorized for wrong token, got %d", w.Result().StatusCode)
    }
  }

  // Case 3: Correct Authorization header (first token)
  {
    req := httptest.NewRequest("GET", "/v1/models", nil)
    req.Header.Set("Authorization", "Bearer secret-token-1")
    w := httptest.NewRecorder()
    modelsHandler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusOK {
      t.Errorf("Expected 200 OK for correct token 1, got %d", w.Result().StatusCode)
    }
  }

  // Case 3b: Correct Authorization header (second token)
  {
    req := httptest.NewRequest("GET", "/v1/models", nil)
    req.Header.Set("Authorization", "Bearer secret-token-2")
    w := httptest.NewRecorder()
    modelsHandler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusOK {
      t.Errorf("Expected 200 OK for correct token 2, got %d", w.Result().StatusCode)
    }
  }

  // Case 4: Completions endpoint with correct token
  {
    reqBodyObj := map[string]any{
      "model": "model-1",
    }
    reqBytes, _ := json.Marshal(reqBodyObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Authorization", "Bearer secret-token-2")
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()
    completionsHandler.ServeHTTP(w, req)
    // It should proceed past auth and return 502/503/500 because "ep-1" is not a valid endpoint,
    // but critically, it should NOT return 401.
    status := w.Result().StatusCode
    if status == http.StatusUnauthorized {
      t.Errorf("Expected completions request to pass authentication, but got 401 Unauthorized")
    }
  }
}

func TestCSVUpdater(t *testing.T) {
  tmpDir, err := os.MkdirTemp("", "dynagate-updater-test")
  if err != nil {
    t.Fatalf("failed to create temp dir: %v", err)
  }
  defer os.RemoveAll(tmpDir)

  csvFile := filepath.Join(tmpDir, "models.csv")
  initialContent := "model,key,endpoint\ngpt-4,key-1,http://localhost:8001\n"
  if err := os.WriteFile(csvFile, []byte(initialContent), 0644); err != nil {
    t.Fatalf("failed to write csv file: %v", err)
  }

  // Create a command to append a model row to the CSV file
  command := fmt.Sprintf("echo 'gpt-3.5,key-2,http://localhost:8002' >> %s", csvFile)

  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()

  // Start background updater with a fast interval (100ms)
  go startCSVUpdater(ctx, command, 100*time.Millisecond)

  // The updater runs once immediately at startup, so we expect the file to have the appended line
  // almost immediately. Let's poll to check.
  deadline := time.Now().Add(2 * time.Second)
  var contentBytes []byte
  var containsNewModel bool

  for time.Now().Before(deadline) {
    contentBytes, err = os.ReadFile(csvFile)
    if err == nil && strings.Contains(string(contentBytes), "gpt-3.5") {
      containsNewModel = true
      break
    }
    time.Sleep(50 * time.Millisecond)
  }

  if !containsNewModel {
    t.Fatalf("Expected CSV file to contain the updater-appended model, but it didn't. Content: %s", string(contentBytes))
  }
}

func TestGatewayImageGenerations(t *testing.T) {
  // We will setup a mock upstream server that handles image generation.
  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path != "/v1/images/generations" {
      w.WriteHeader(http.StatusNotFound)
      return
    }

    var body map[string]any
    _ = json.NewDecoder(r.Body).Decode(&body)
    modelSent := body["model"].(string)

    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(fmt.Sprintf(`{"created":1589478378,"data":[{"url":"http://image-url/generated-by-%s"}]}`, modelSent)))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "dall-e-3", Key: "dalle-key", Endpoint: server.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"dall-e-3"},
  }

  handler := handleImageGenerations(cm, nil)

  reqBodyObj := map[string]any{
    "prompt": "a beautiful kitten",
    "model": "dall-e-3",
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/images/generations", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  resp := w.Result()
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200, got %d", resp.StatusCode)
  }

  var respObj map[string]any
  _ = json.NewDecoder(resp.Body).Decode(&respObj)

  data, ok := respObj["data"].([]any)
  if !ok || len(data) == 0 {
    t.Fatalf("Expected data field to contain image urls, got: %+v", respObj)
  }

  urlVal := data[0].(map[string]any)["url"].(string)
  expectedURL := "http://image-url/generated-by-dall-e-3"
  if urlVal != expectedURL {
    t.Errorf("Expected URL %q, got %q", expectedURL, urlVal)
  }
}

func TestGatewayNoAuthHeaderOnEmptyKey(t *testing.T) {
  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if auth := r.Header.Get("Authorization"); auth != "" {
      t.Errorf("Expected no Authorization header, but got %q", auth)
      w.WriteHeader(http.StatusBadRequest)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hello"}}]}n`))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "model-no-key", Key: "", Endpoint: server.URL},
    {Model: "model-no-key-2", Key: "-", Endpoint: server.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-no-key", "model-no-key-2"},
  }

  handler := handleChatCompletions(cm, nil)

  // Case 1: Empty Key
  {
    reqBodyObj := map[string]any{
      "model": "model-no-key",
      "messages": []map[string]string{
        {"role": "user", "content": "Hi"},
      },
    }
    reqBytes, _ := json.Marshal(reqBodyObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    req.Header.Set("Authorization", "Bearer some-client-token")
    w := httptest.NewRecorder()
    handler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusOK {
      t.Errorf("Expected status 200 for empty key, got %d", w.Result().StatusCode)
    }
  }

  // Case 2: Dash Key
  {
    reqBodyObj := map[string]any{
      "model": "model-no-key-2",
      "messages": []map[string]string{
        {"role": "user", "content": "Hi"},
      },
    }
    reqBytes, _ := json.Marshal(reqBodyObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    req.Header.Set("Authorization", "Bearer some-client-token")
    w := httptest.NewRecorder()
    handler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusOK {
      t.Errorf("Expected status 200 for dash key, got %d", w.Result().StatusCode)
    }
  }
}

func TestGatewayBlankKey(t *testing.T) {
  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    auth := r.Header.Get("Authorization")
    if auth != "Bearer" {
      t.Errorf("Expected Authorization header to be exactly 'Bearer', but got %q", auth)
      w.WriteHeader(http.StatusBadRequest)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hello"}}]}n`))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "model-blank-key", Key: "-blank-", Endpoint: server.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-blank-key"},
  }

  handler := handleChatCompletions(cm, nil)

  reqBodyObj := map[string]any{
    "model": "model-blank-key",
    "messages": []map[string]string{
      {"role": "user", "content": "Hi"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  req.Header.Set("Authorization", "Bearer some-client-token")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  if w.Result().StatusCode != http.StatusOK {
    t.Errorf("Expected status 200 for -blank- key, got %d", w.Result().StatusCode)
  }
}

func TestGatewayExtraHeaders(t *testing.T) {
  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if h := r.Header.Get("X-My-Custom-Header"); h != "HeaderValue" {
      t.Errorf("Expected X-My-Custom-Header to be 'HeaderValue', got %q", h)
      w.WriteHeader(http.StatusBadRequest)
      return
    }
    if h := r.Header.Get("X-Another-Header"); h != "123" {
      t.Errorf("Expected X-Another-Header to be '123', got %q", h)
      w.WriteHeader(http.StatusBadRequest)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hello"}}]}n`))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "model-extra-headers", Key: "my-key", Endpoint: server.URL, Extra: `{"X-My-Custom-Header":"HeaderValue","X-Another-Header":123}`},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-extra-headers"},
  }

  handler := handleChatCompletions(cm, nil)

  reqBodyObj := map[string]any{
    "model": "model-extra-headers",
    "messages": []map[string]string{
      {"role": "user", "content": "Hi"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  if w.Result().StatusCode != http.StatusOK {
    t.Errorf("Expected status 200, got %d", w.Result().StatusCode)
  }
}

func TestMultipleTokensKeyFile(t *testing.T) {
  // Create a temporary key file with multiple tokens, some empty lines, and comments/spaces
  tmpDir, err := os.MkdirTemp("", "dynagate-keyfile-test")
  if err != nil {
    t.Fatalf("failed to create temp dir: %v", err)
  }
  defer os.RemoveAll(tmpDir)

  keyFilePath := filepath.Join(tmpDir, "keys.txt")
  keyFileContent := "\n  token-a  \n\ntoken-b\n   \ntoken-c\n"
  if err := os.WriteFile(keyFilePath, []byte(keyFileContent), 0644); err != nil {
    t.Fatalf("failed to write key file: %v", err)
  }

  // Parse using the same logic as main.go
  content, err := os.ReadFile(keyFilePath)
  if err != nil {
    t.Fatalf("failed to read key file: %v", err)
  }
  var authTokens []string
  for _, line := range strings.Split(string(content), "\n") {
    token := strings.TrimSpace(line)
    if token != "" {
      authTokens = append(authTokens, token)
    }
  }

  expected := []string{"token-a", "token-b", "token-c"}
  if !reflect.DeepEqual(authTokens, expected) {
    t.Errorf("Expected parsed tokens to be %v, got %v", expected, authTokens)
  }

  // Verify that the helper checkAuth works with these parsed tokens
  req := httptest.NewRequest("GET", "/v1/models", nil)
  req.Header.Set("Authorization", "Bearer token-b")
  if !checkAuth(authTokens, req) {
    t.Errorf("Expected token-b to authenticate successfully")
  }

  reqWrong := httptest.NewRequest("GET", "/v1/models", nil)
  reqWrong.Header.Set("Authorization", "Bearer token-wrong")
  if checkAuth(authTokens, reqWrong) {
    t.Errorf("Expected token-wrong to fail authentication")
  }

  // Test empty key file content behavior
  emptyKeyFilePath := filepath.Join(tmpDir, "empty_keys.txt")
  emptyKeyFileContent := "\n   \n\n  \n"
  if err := os.WriteFile(emptyKeyFilePath, []byte(emptyKeyFileContent), 0644); err != nil {
    t.Fatalf("failed to write empty key file: %v", err)
  }

  // Parse using empty key file logic as in main.go
  emptyContent, err := os.ReadFile(emptyKeyFilePath)
  if err != nil {
    t.Fatalf("failed to read empty key file: %v", err)
  }
  var emptyAuthTokens []string
  var tokens []string
  for _, line := range strings.Split(string(emptyContent), "\n") {
    token := strings.TrimSpace(line)
    if token != "" {
      tokens = append(tokens, token)
    }
  }
  if len(tokens) > 0 {
    emptyAuthTokens = tokens
  } else {
    emptyAuthTokens = nil
  }

  if emptyAuthTokens != nil {
    t.Errorf("Expected emptyAuthTokens to be nil, got %v", emptyAuthTokens)
  }

  // Verify that checkAuth returns true when expectedTokens is nil
  reqNoAuth := httptest.NewRequest("GET", "/v1/models", nil)
  if !checkAuth(emptyAuthTokens, reqNoAuth) {
    t.Errorf("Expected checkAuth to return true for nil expectedTokens")
  }
}

func TestRequestMessagesAutodetectionAndContentArrays(t *testing.T) {
  var receivedBody map[string]any

  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    var b map[string]any
    _ = json.NewDecoder(r.Body).Decode(&b)
    receivedBody = b
    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"OK"}}]} `))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "test-model", Key: "test-key", Endpoint: server.URL},
  }
  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"test-model"},
  }
  handler := handleChatCompletions(cm, nil)

  t.Run("Pi agent format - array content", func(t *testing.T) {
    reqObj := map[string]any{
      "model": "test-model",
      "messages": []any{
        map[string]any{
          "role": "user",
          "content": []any{
            map[string]any{"type": "text", "text": "Hello from Pi agent"},
          },
        },
      },
    }
    reqBytes, _ := json.Marshal(reqObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()

    handler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusOK {
      t.Fatalf("Expected status 200, got %d", w.Result().StatusCode)
    }

    msgs, ok := receivedBody["messages"].([]any)
    if !ok || len(msgs) != 1 {
      t.Fatalf("Expected 1 message in forwarded body, got %+v", receivedBody)
    }
    userMsg := msgs[0].(map[string]any)
    if userMsg["role"] != "user" || userMsg["content"] != "Hello from Pi agent" {
      t.Errorf("Expected content 'Hello from Pi agent', got %+v", userMsg)
    }
  })

  t.Run("Top-level system prompt", func(t *testing.T) {
    reqObj := map[string]any{
      "model":  "test-model",
      "system": "You are a helpful coding assistant",
      "messages": []any{
        map[string]any{
          "role":    "user",
          "content": "Write a test",
        },
      },
    }
    reqBytes, _ := json.Marshal(reqObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()

    handler.ServeHTTP(w, req)

    msgs, ok := receivedBody["messages"].([]any)
    if !ok || len(msgs) != 2 {
      t.Fatalf("Expected 2 messages in forwarded body, got %+v", receivedBody)
    }
    sysMsg := msgs[0].(map[string]any)
    userMsg := msgs[1].(map[string]any)
    if sysMsg["role"] != "system" || sysMsg["content"] != "You are a helpful coding assistant" {
      t.Errorf("Unexpected system message: %+v", sysMsg)
    }
    if userMsg["role"] != "user" || userMsg["content"] != "Write a test" {
      t.Errorf("Unexpected user message: %+v", userMsg)
    }
    if _, exists := receivedBody["system"]; exists {
      t.Errorf("Expected top-level system field to be removed")
    }
  })

  t.Run("Prompt field autodetection", func(t *testing.T) {
    reqObj := map[string]any{
      "model":  "test-model",
      "prompt": "Explain recursion",
    }
    reqBytes, _ := json.Marshal(reqObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()

    handler.ServeHTTP(w, req)

    msgs, ok := receivedBody["messages"].([]any)
    if !ok || len(msgs) != 1 {
      t.Fatalf("Expected 1 message in forwarded body, got %+v", receivedBody)
    }
    userMsg := msgs[0].(map[string]any)
    if userMsg["role"] != "user" || userMsg["content"] != "Explain recursion" {
      t.Errorf("Unexpected user message: %+v", userMsg)
    }
  })

  t.Run("Multimodal array content preserved", func(t *testing.T) {
    reqObj := map[string]any{
      "model": "test-model",
      "messages": []any{
        map[string]any{
          "role": "user",
          "content": []any{
            map[string]any{"type": "text", "text": "Describe this image"},
            map[string]any{"type": "image_url", "image_url": map[string]any{"url": "http://example.com/img.png"}},
          },
        },
      },
    }
    reqBytes, _ := json.Marshal(reqObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()

    handler.ServeHTTP(w, req)

    msgs, ok := receivedBody["messages"].([]any)
    if !ok || len(msgs) != 1 {
      t.Fatalf("Expected 1 message in forwarded body, got %+v", receivedBody)
    }
    userMsg := msgs[0].(map[string]any)
    cntArray, isArray := userMsg["content"].([]any)
    if !isArray || len(cntArray) != 2 {
      t.Fatalf("Expected multimodal array content to be preserved, got %+v", userMsg["content"])
    }
  })
}

func TestFibonacciBackoff(t *testing.T) {
  expectedSeq := []struct {
    n        int
    expected int64
  }{
    {0, 0},
    {1, 1},
    {2, 1},
    {3, 2},
    {4, 3},
    {5, 5},
    {6, 8},
    {7, 13},
    {8, 21},
    {9, 34},
    {10, 55},
  }

  for _, tc := range expectedSeq {
    got := fibonacci(tc.n)
    if got != tc.expected {
      t.Errorf("fibonacci(%d) = %d; want %d", tc.n, got, tc.expected)
    }
  }

  // Verify failover works with zero-delay in handleChatCompletions
  serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusInternalServerError)
  }))
  defer serverA.Close()

  serverB := 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":[{"message":{"content":"Success from B"}}]} `))
  }))
  defer serverB.Close()

  configs := []ModelConfig{
    {Model: "m1", Key: "k1", Endpoint: serverA.URL},
    {Model: "m1", Key: "k2", Endpoint: serverB.URL},
  }
  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"m1"},
  }
  handler := handleChatCompletions(cm, nil, 0) // 0 delay for fast test

  reqBodyObj := map[string]any{
    "model": "m1",
    "messages": []map[string]string{
      {"role": "user", "content": "test"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)
  if w.Result().StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200 after Fibonacci retry, got %d", w.Result().StatusCode)
  }
}

func TestAnthropicMessagesEndpoint(t *testing.T) {
  var receivedOpenAIBody map[string]any

  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path != "/v1/chat/completions" {
      t.Errorf("Expected request to be routed to /v1/chat/completions, got %s", r.URL.Path)
    }

    var body map[string]any
    _ = json.NewDecoder(r.Body).Decode(&body)
    receivedOpenAIBody = body

    isStream, _ := body["stream"].(bool)
    w.Header().Set("Content-Type", "application/json")
    if isStream {
      w.Header().Set("Content-Type", "text/event-stream")
      w.WriteHeader(http.StatusOK)
      flusher, _ := w.(http.Flusher)
      _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n"))
      flusher.Flush()
      _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Anthropic!\"},\"finish_reason\":\"stop\"}]}\n\n"))
      flusher.Flush()
      _, _ = w.Write([]byte("data: [DONE]\n\n"))
      flusher.Flush()
      return
    }

    w.WriteHeader(http.StatusOK)
    resp := map[string]any{
      "id":    "chatcmpl-test",
      "model": "claude-model",
      "choices": []any{
        map[string]any{
          "index": 0,
          "message": map[string]any{
            "role":    "assistant",
            "content": "Hello from Anthropic proxy",
          },
          "finish_reason": "stop",
        },
      },
      "usage": map[string]any{
        "prompt_tokens":     10,
        "completion_tokens": 15,
        "total_tokens":      25,
      },
    }
    _ = json.NewEncoder(w).Encode(resp)
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "claude-model", Key: "test-key", Endpoint: server.URL},
  }
  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"claude-model"},
  }
  handler := handleMessages(cm, nil, 0)

  // 1. Non-streaming test
  t.Run("Non-streaming Anthropic message", func(t *testing.T) {
    reqObj := map[string]any{
      "model":  "claude-model",
      "system": "You are a helpful bot",
      "messages": []any{
        map[string]any{
          "role":    "user",
          "content": "Hello",
        },
      },
      "max_tokens": 100,
    }
    reqBytes, _ := json.Marshal(reqObj)
    req := httptest.NewRequest("POST", "/v1/messages", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()

    handler.ServeHTTP(w, req)

    if w.Result().StatusCode != http.StatusOK {
      t.Fatalf("Expected 200 OK, got %d", w.Result().StatusCode)
    }

    var anthropicResp map[string]any
    _ = json.NewDecoder(w.Body).Decode(&anthropicResp)

    if anthropicResp["type"] != "message" {
      t.Errorf("Expected response type 'message', got %v", anthropicResp["type"])
    }
    if anthropicResp["role"] != "assistant" {
      t.Errorf("Expected role 'assistant', got %v", anthropicResp["role"])
    }
    content, ok := anthropicResp["content"].([]any)
    if !ok || len(content) == 0 {
      t.Fatalf("Expected non-empty content array in Anthropic response, got %+v", anthropicResp)
    }
    textVal := content[0].(map[string]any)["text"].(string)
    if textVal != "Hello from Anthropic proxy" {
      t.Errorf("Unexpected text in response: %s", textVal)
    }

    // Verify system message was prepended into OpenAI messages format
    openAIMsgs, ok := receivedOpenAIBody["messages"].([]any)
    if !ok || len(openAIMsgs) != 2 {
      t.Fatalf("Expected 2 messages (system + user) in OpenAI body, got %+v", receivedOpenAIBody)
    }
    if openAIMsgs[0].(map[string]any)["role"] != "system" {
      t.Errorf("Expected first message to be system prompt")
    }
  })

  // 2. Streaming test
  t.Run("Streaming Anthropic message", func(t *testing.T) {
    reqObj := map[string]any{
      "model":  "claude-model",
      "stream": true,
      "messages": []any{
        map[string]any{
          "role":    "user",
          "content": "Stream me",
        },
      },
    }
    reqBytes, _ := json.Marshal(reqObj)
    req := httptest.NewRequest("POST", "/v1/messages", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()

    handler.ServeHTTP(w, req)

    if w.Result().StatusCode != http.StatusOK {
      t.Fatalf("Expected 200 OK, got %d", w.Result().StatusCode)
    }

    bodyStr := w.Body.String()
    if !strings.Contains(bodyStr, "event: message_start") || !strings.Contains(bodyStr, "event: content_block_delta") || !strings.Contains(bodyStr, "event: message_stop") {
      t.Errorf("Expected Anthropic SSE stream events in response, got: %s", bodyStr)
    }
  })
}

func TestOpenAIResponsesEndpoint(t *testing.T) {
  var receivedOpenAIBody map[string]any

  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path != "/v1/chat/completions" {
      t.Errorf("Expected request to be routed to /v1/chat/completions, got %s", r.URL.Path)
    }

    var body map[string]any
    _ = json.NewDecoder(r.Body).Decode(&body)
    receivedOpenAIBody = body

    isStream, _ := body["stream"].(bool)
    w.Header().Set("Content-Type", "application/json")
    if isStream {
      w.Header().Set("Content-Type", "text/event-stream")
      w.WriteHeader(http.StatusOK)
      flusher, _ := w.(http.Flusher)
      _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Response \"}}]}\n\n"))
      flusher.Flush()
      _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Chunk\"},\"finish_reason\":\"stop\"}]}\n\n"))
      flusher.Flush()
      _, _ = w.Write([]byte("data: [DONE]\n\n"))
      flusher.Flush()
      return
    }

    w.WriteHeader(http.StatusOK)
    resp := map[string]any{
      "id":    "chatcmpl-resp-test",
      "model": "gpt-5-model",
      "choices": []any{
        map[string]any{
          "index": 0,
          "message": map[string]any{
            "role":    "assistant",
            "content": "Hello from Responses proxy",
          },
          "finish_reason": "stop",
        },
      },
      "usage": map[string]any{
        "prompt_tokens":     12,
        "completion_tokens": 18,
        "total_tokens":      30,
      },
    }
    _ = json.NewEncoder(w).Encode(resp)
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "gpt-5-model", Key: "test-key", Endpoint: server.URL},
  }
  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"gpt-5-model"},
  }
  handler := handleResponses(cm, nil, 0)

  // 1. Non-streaming test
  t.Run("Non-streaming OpenAI response", func(t *testing.T) {
    reqObj := map[string]any{
      "model":        "gpt-5-model",
      "instructions": "Be accurate",
      "input":        "What is 2+2?",
    }
    reqBytes, _ := json.Marshal(reqObj)
    req := httptest.NewRequest("POST", "/v1/responses", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()

    handler.ServeHTTP(w, req)

    if w.Result().StatusCode != http.StatusOK {
      t.Fatalf("Expected 200 OK, got %d", w.Result().StatusCode)
    }

    var respObj map[string]any
    _ = json.NewDecoder(w.Body).Decode(&respObj)

    if respObj["object"] != "response" {
      t.Errorf("Expected object 'response', got %v", respObj["object"])
    }
    if respObj["status"] != "completed" {
      t.Errorf("Expected status 'completed', got %v", respObj["status"])
    }
    output, ok := respObj["output"].([]any)
    if !ok || len(output) == 0 {
      t.Fatalf("Expected output array in Responses API response, got %+v", respObj)
    }
    firstOutput := output[0].(map[string]any)
    contentSlice, ok := firstOutput["content"].([]any)
    if !ok || len(contentSlice) == 0 {
      t.Fatalf("Expected output message content slice, got %+v", firstOutput)
    }
    textVal := contentSlice[0].(map[string]any)["text"].(string)
    if textVal != "Hello from Responses proxy" {
      t.Errorf("Unexpected text in response output: %s", textVal)
    }

    // Verify instructions and input were parsed into OpenAI messages format
    openAIMsgs, ok := receivedOpenAIBody["messages"].([]any)
    if !ok || len(openAIMsgs) != 2 {
      t.Fatalf("Expected 2 messages in OpenAI body, got %+v", receivedOpenAIBody)
    }
    if openAIMsgs[0].(map[string]any)["role"] != "system" {
      t.Errorf("Expected system role for instructions")
    }
  })

  // 2. Streaming test
  t.Run("Streaming OpenAI response", func(t *testing.T) {
    reqObj := map[string]any{
      "model":  "gpt-5-model",
      "stream": true,
      "input":  "Stream test",
    }
    reqBytes, _ := json.Marshal(reqObj)
    req := httptest.NewRequest("POST", "/v1/responses", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()

    handler.ServeHTTP(w, req)

    if w.Result().StatusCode != http.StatusOK {
      t.Fatalf("Expected 200 OK, got %d", w.Result().StatusCode)
    }

    bodyStr := w.Body.String()
    if !strings.Contains(bodyStr, "event: response.created") || !strings.Contains(bodyStr, "event: response.output_text.delta") || !strings.Contains(bodyStr, "event: response.completed") {
      t.Errorf("Expected OpenAI Responses SSE events in response stream, got: %s", bodyStr)
    }
  })
}